diff --git a/dev-notes/architecture/phase-21-extensions.md b/dev-notes/architecture/phase-21-extensions.md index 0ff9222cee..27b3efd19a 100644 --- a/dev-notes/architecture/phase-21-extensions.md +++ b/dev-notes/architecture/phase-21-extensions.md @@ -4,8 +4,9 @@ title: "Phase 21: Extensions" Tau extensions are Python modules that customize a coding session: they add tools and slash commands, observe the agent event stream, and intercept tool -calls, tool results, and user input. The design is a deliberate port of Pi's -extension system (`packages/coding-agent/src/core/extensions/` in +calls, tool results, user input, and per-run system prompts. The design is a +deliberate port of Pi's extension system +(`packages/coding-agent/src/core/extensions/` in `earendil-works/pi`) onto Tau's Python architecture, scoped so the core is small while still supporting real extensions such as a Claude Code-style subagents extension. @@ -21,11 +22,11 @@ are called out inline as **Ruling:** notes. `register_tool`, `register_command`, `on(event)`, `send_user_message`, `append_entry`, and read access to session context. - Support Pi's load-bearing hook semantics: `tool_call` (block/mutate), - `tool_result` (transform), `input` (transform/handle), plus observation of - every portable `AgentEvent`. -- Keep `tau_agent` untouched: the extension machinery lives entirely in - `tau_coding`, using existing seams (`AgentHarness.subscribe`, executor - wrapping, `CommandRegistry`, `CustomEntry`). + `tool_result` (transform), `input` (transform/handle), + `before_agent_start` (replace the run prompt), plus observation of every + portable `AgentEvent`. +- Keep extension policy in `tau_coding`; `tau_agent` exposes only portable + seams such as `AgentHarness.subscribe` and the run-scoped prompt argument. - Isolate failures: a broken extension is a `ResourceDiagnostic`, never a crashed session. @@ -34,8 +35,8 @@ are called out inline as **Ruling:** notes. - npm-style package management (`pi install`), provider registration, custom TUI components/widgets (extension-authored Textual widgets), custom **entry** renderers (`registerEntryRenderer`/`appendEntry`-rendered, - non-LLM-context cards), shortcut and flag registration, system-prompt - replacement, `context`/`before_provider_request` rewriting, and a project + non-LLM-context cards), shortcut and flag registration, + `context`/`before_provider_request` rewriting, and a project trust store. These have reserved names and documented extension points but no implementation yet. @@ -240,6 +241,7 @@ Lifecycle events, dispatched by the runtime: | `session_start` | `SessionStartEvent(reason: "startup" \| "reload" \| "new" \| "resume" \| "branch")` | — | | `session_shutdown` | `SessionShutdownEvent(reason)` | — | | `input` | `InputEvent(text, source="interactive" \| "extension", streaming_behavior="steer" \| "follow_up" \| None)` | `InputHookResult(action="continue" \| "transform" \| "handled", text=None, message=None)` | +| `before_agent_start` | `BeforeAgentStartEvent(system_prompt, system_prompt_inputs)` | `BeforeAgentStartHookResult(system_prompt=None)` | | `tool_call` | `ToolCallHookEvent(tool_name, arguments)` | `ToolCallHookResult(block=False, reason=None, arguments=None)` | | `tool_result` | `ToolResultHookEvent(tool_name, arguments, result)` | `ToolResultHookResult(content=None, ok=None, details=None)` | @@ -559,7 +561,7 @@ the extension runtime) and the **queued-message preview** (`harness.py` `queue_update_event` reports queued content strings). Both are raw-text views by design; only live transcripts (TUI + print mode) render. -## Hook wiring (how interception works without touching tau_agent) +## Hook wiring and the portable tau_agent seam - **Observation** — the runtime subscribes one listener via `AgentHarness.subscribe` and fans events out to extension handlers. @@ -578,6 +580,13 @@ raw-text views by design; only live transcripts (TUI + print mode) render. consumes the input: the prompt generator returns without yielding run events, and the optional `message` is delivered through the UI bridge notification channel. +- **`before_agent_start`** — `CodingSession` chains prompt replacements before + each `prompt` or `continue_` run, then passes the final value through a + keyword-only `AgentHarness` run argument. The harness keeps that local value + for every provider call in the tool loop and clears it in `finally`; its base + config and transcript are unchanged. The hook receives frozen, repr-hidden + prompt inputs with skill content omitted. Its fresh `ExtensionContext` also + exposes the current chained prompt, matching Pi's `ctx.getSystemPrompt()`. - **`send_user_message` / `send_custom_message`** — both funnel through one `_deliver_message` path. When a run is active, they map to `queue_steering_message` / `queue_follow_up_message` (which build a @@ -709,8 +718,9 @@ of the newer API seams (manifest, dialogs, renderers, `on_update`, falling back to `send_user_message` on older builds — which also exercises the idle `turn_requested` path. -Smaller examples: `hello_tool.py` (minimal tool) and `permission_gate.py` -(`tool_call` blocking for dangerous bash commands). +Smaller examples: `hello_tool.py` (minimal tool), `prompt_customizer.py` +(`before_agent_start` replacement), and `permission_gate.py` (`tool_call` +blocking for dangerous bash commands). ## Verification @@ -719,7 +729,9 @@ Smaller examples: `hello_tool.py` (minimal tool) and `permission_gate.py` isolation, sync-only `setup` enforcement, tool registration/override, command registration/duplicate handling, event fan-out, `tool_call` block + argument mutation, `tool_result` transform, `input` - transform/handled, send_user_message queueing and idle turn-request, + transform/handled, `before_agent_start` chaining/failure isolation and + run reset/tool-loop/transcript behavior, send_user_message queueing and idle + turn-request, append_entry persistence and on-path replay, reload including module purge and stale-listener replacement, runtime survival across resume/new. diff --git a/dev-notes/architecture/phase-4-agent-harness.md b/dev-notes/architecture/phase-4-agent-harness.md index b0c7e20161..70f0148345 100644 --- a/dev-notes/architecture/phase-4-agent-harness.md +++ b/dev-notes/architecture/phase-4-agent-harness.md @@ -66,7 +66,10 @@ config = AgentHarnessConfig( ) ``` -The harness receives this config and uses it for every prompt or continuation. +The harness receives this config as the stable default. `prompt()` and +`continue_()` may receive a keyword-only run prompt; when omitted, they use the +configured value. A run prompt is cleared when that run settles and never +changes the config. ## Prompt flow diff --git a/examples/extensions/prompt_customizer.py b/examples/extensions/prompt_customizer.py new file mode 100644 index 0000000000..86038e9d0f --- /dev/null +++ b/examples/extensions/prompt_customizer.py @@ -0,0 +1,32 @@ +"""Tau extension that adds a run-scoped system-prompt instruction. + +Install by copying into `~/.tau/extensions/`, or run: + + tau -e examples/extensions/prompt_customizer.py +""" + +from typing import cast + +from tau_coding.extensions import ( + BeforeAgentStartEvent, + BeforeAgentStartHookResult, + ExtensionAPI, + ExtensionContext, + ExtensionHandler, +) + + +def _customize_prompt( + event: BeforeAgentStartEvent, + context: ExtensionContext, +) -> BeforeAgentStartHookResult: + del context + tools = ", ".join(event.system_prompt_inputs.tools) or "none" + return BeforeAgentStartHookResult( + system_prompt=f"{event.system_prompt}\n\nActive tools for this run: {tools}." + ) + + +def setup(tau: ExtensionAPI) -> None: + """Customize each agent run without changing the saved base prompt.""" + tau.on("before_agent_start", cast(ExtensionHandler, _customize_prompt)) diff --git a/src/tau_agent/harness.py b/src/tau_agent/harness.py index 4ead48aa63..831cb9c75d 100644 --- a/src/tau_agent/harness.py +++ b/src/tau_agent/harness.py @@ -72,6 +72,7 @@ def __init__( self._messages = list(messages) self._listeners: list[EventListener] = [] self._current_signal: SimpleCancellationToken | None = None + self._active_system_prompt: str | None = None self._running = False self._steering_queue: deque[AgentMessage] = deque() self._follow_up_queue: deque[AgentMessage] = deque() @@ -84,6 +85,13 @@ def messages(self) -> tuple[AgentMessage, ...]: def config(self) -> AgentHarnessConfig: return self._config + @property + def system_prompt(self) -> str: + """Return the active run prompt, or the configured base while idle.""" + if self._active_system_prompt is not None: + return self._active_system_prompt + return self._config.system + @property def is_running(self) -> bool: return self._running @@ -144,26 +152,34 @@ def pop_latest_follow_up(self) -> AgentMessage | None: def pop_latest_steering(self) -> AgentMessage | None: return self._steering_queue.pop() if self._steering_queue else None - def prompt_message(self, message: AgentMessage) -> AsyncIterator[AgentEvent]: + def prompt_message( + self, + message: AgentMessage, + *, + system: str | None = None, + ) -> AsyncIterator[AgentEvent]: self._ensure_not_running() self._running = True - return self._run(prompts=(message,)) + return self._run(prompts=(message,), system=system) - def prompt(self, content: str) -> AsyncIterator[AgentEvent]: - return self.prompt_message(UserMessage(content=content)) + def prompt(self, content: str, *, system: str | None = None) -> AsyncIterator[AgentEvent]: + return self.prompt_message(UserMessage(content=content), system=system) - def continue_(self) -> AsyncIterator[AgentEvent]: + def continue_(self, *, system: str | None = None) -> AsyncIterator[AgentEvent]: self._ensure_not_running() self._running = True - return self._run() + return self._run(system=system) async def _run( self, *, prompts: Sequence[AgentMessage] = (), + system: str | None = None, ) -> AsyncIterator[AgentEvent]: signal = SimpleCancellationToken() self._current_signal = signal + effective_system = self._config.system if system is None else system + self._active_system_prompt = effective_system try: # Repair dangling tool calls here, not in prompt()/continue_(), # so the synthetic results flow through events and reach push @@ -174,7 +190,7 @@ async def _run( async for event in run_agent_loop( provider=self._config.provider, model=self._config.model, - system=self._config.system, + system=effective_system, messages=self._messages, prompts=prompts, prelude_messages=repairs, @@ -190,19 +206,22 @@ async def _run( await self._notify(event) yield event finally: - if signal.is_cancelled(): - repaired_from = len(self._messages) - self._append_interrupted_tool_results() - # The consumer is usually gone here; push the repairs to - # subscribers. Listener errors are suppressed; cancellation - # itself is not. - for message in self._messages[repaired_from:]: - with suppress(Exception): - await self._notify(MessageStartEvent(message=message)) - await self._notify(MessageEndEvent(message=message)) - if self._current_signal is signal: - self._current_signal = None - self._running = False + try: + if signal.is_cancelled(): + repaired_from = len(self._messages) + self._append_interrupted_tool_results() + # The consumer is usually gone here; push the repairs to + # subscribers. Listener errors are suppressed; cancellation + # itself is not. + for message in self._messages[repaired_from:]: + with suppress(Exception): + await self._notify(MessageStartEvent(message=message)) + await self._notify(MessageEndEvent(message=message)) + finally: + if self._current_signal is signal: + self._current_signal = None + self._active_system_prompt = None + self._running = False async def _notify(self, event: AgentEvent) -> None: for listener in list(self._listeners): diff --git a/src/tau_coding/data/docs/extensions.md b/src/tau_coding/data/docs/extensions.md index bed4431c98..183f43306e 100644 --- a/src/tau_coding/data/docs/extensions.md +++ b/src/tau_coding/data/docs/extensions.md @@ -37,6 +37,26 @@ Project extensions cannot approve themselves. They execute arbitrary Python and remain disabled without both approval and the explicit code opt-in. Trust is not a process/filesystem/network/tool/model sandbox. +## Per-run system prompts + +Register `before_agent_start` to replace the system prompt for one agent run. +Handlers receive `BeforeAgentStartEvent.system_prompt` and a typed +`system_prompt_inputs` snapshot, then may return +`BeforeAgentStartHookResult(system_prompt=...)`. Handlers run in registration +order and each sees the prior replacement. The final prompt remains active for +tool-loop requests, then the next run starts again from the session's base +prompt. It is never added to the transcript. + +During this hook, both `event.system_prompt` and `context.system_prompt` expose +the current chained value. + +Skill metadata includes `disable_model_invocation`, matching whether a skill is +eligible for model invocation. Prompt inputs can contain project instructions +and paths. Their container types hide values from `repr` and Tau diagnostics, +but extensions should still treat +explicitly accessed fields as sensitive. See +`examples/extensions/prompt_customizer.py` for a complete extension. + ## Development checklist 1. Read this document and the closest installed example under `examples/extensions/` completely before implementing. diff --git a/src/tau_coding/data/examples/extensions/prompt_customizer.py b/src/tau_coding/data/examples/extensions/prompt_customizer.py new file mode 100644 index 0000000000..86038e9d0f --- /dev/null +++ b/src/tau_coding/data/examples/extensions/prompt_customizer.py @@ -0,0 +1,32 @@ +"""Tau extension that adds a run-scoped system-prompt instruction. + +Install by copying into `~/.tau/extensions/`, or run: + + tau -e examples/extensions/prompt_customizer.py +""" + +from typing import cast + +from tau_coding.extensions import ( + BeforeAgentStartEvent, + BeforeAgentStartHookResult, + ExtensionAPI, + ExtensionContext, + ExtensionHandler, +) + + +def _customize_prompt( + event: BeforeAgentStartEvent, + context: ExtensionContext, +) -> BeforeAgentStartHookResult: + del context + tools = ", ".join(event.system_prompt_inputs.tools) or "none" + return BeforeAgentStartHookResult( + system_prompt=f"{event.system_prompt}\n\nActive tools for this run: {tools}." + ) + + +def setup(tau: ExtensionAPI) -> None: + """Customize each agent run without changing the saved base prompt.""" + tau.on("before_agent_start", cast(ExtensionHandler, _customize_prompt)) diff --git a/src/tau_coding/data/release-notes/releases.json b/src/tau_coding/data/release-notes/releases.json index b704469120..f839bbbe5b 100644 --- a/src/tau_coding/data/release-notes/releases.json +++ b/src/tau_coding/data/release-notes/releases.json @@ -22,7 +22,8 @@ "date": "2026-08-17", "sections": { "New": [ - "Keep user-triggered skills out of the model-facing system prompt with disable-model-invocation frontmatter while retaining explicit invocation, picker, autocomplete, reload, and sidebar support." + "Keep user-triggered skills out of the model-facing system prompt with disable-model-invocation frontmatter while retaining explicit invocation, picker, autocomplete, reload, and sidebar support.", + "Let extensions replace the system prompt for one agent run through a chained before_agent_start hook without persisting the replacement." ], "Changed": [ "Refine sidebar skill and prompt headings with bold labels, muted metadata, and no hover or focus background highlights." diff --git a/src/tau_coding/extensions/__init__.py b/src/tau_coding/extensions/__init__.py index 9446fdeb83..1dae8db17f 100644 --- a/src/tau_coding/extensions/__init__.py +++ b/src/tau_coding/extensions/__init__.py @@ -4,6 +4,8 @@ AGENT_EVENT_TYPES, AGENT_EVENT_WILDCARD, LIFECYCLE_EVENT_TYPES, + BeforeAgentStartEvent, + BeforeAgentStartHookResult, ComponentBridge, CustomMessageMarkup, CustomMessageView, @@ -49,11 +51,14 @@ ExtensionRuntime, InputHookOutcome, ) +from tau_coding.system_prompt import SystemPromptInputs, SystemPromptSkill __all__ = [ "AGENT_EVENT_TYPES", "AGENT_EVENT_WILDCARD", "LIFECYCLE_EVENT_TYPES", + "BeforeAgentStartEvent", + "BeforeAgentStartHookResult", "BoundSession", "ComponentBridge", "CustomMessageMarkup", @@ -84,6 +89,8 @@ "SessionShutdownEvent", "SessionStartEvent", "StderrUiBridge", + "SystemPromptInputs", + "SystemPromptSkill", "ToolCallHookEvent", "ToolCallHookResult", "ToolResultHookEvent", diff --git a/src/tau_coding/extensions/api.py b/src/tau_coding/extensions/api.py index 8677248617..881e2cbc95 100644 --- a/src/tau_coding/extensions/api.py +++ b/src/tau_coding/extensions/api.py @@ -11,6 +11,7 @@ from tau_agent.messages import AgentMessage, ToolResultMessage from tau_agent.tools import AgentTool, AgentToolResult from tau_agent.types import JSONValue +from tau_coding.system_prompt import SystemPromptInputs if TYPE_CHECKING: from textual import events @@ -49,6 +50,7 @@ "session_start", "session_shutdown", "input", + "before_agent_start", "tool_call", "tool_result", "project_trust", @@ -386,6 +388,21 @@ class InputHookResult: message: str | None = None +@dataclass(frozen=True, slots=True, repr=False) +class BeforeAgentStartEvent: + """Sensitive inputs for a run-scoped system-prompt replacement hook.""" + + system_prompt: str + system_prompt_inputs: SystemPromptInputs + + +@dataclass(frozen=True, slots=True, repr=False) +class BeforeAgentStartHookResult: + """Optional system-prompt replacement for the current agent run.""" + + system_prompt: str | None = None + + @dataclass(frozen=True, slots=True) class ToolCallHookEvent: """Payload for the `tool_call` hook, before a tool executes. @@ -758,10 +775,13 @@ def __init__( self, runtime: ExtensionRuntime, generation: ExtensionGeneration | None = None, + *, + system_prompt: str | None = None, ) -> None: self._runtime = runtime self._generation = generation if generation is not None else ExtensionGeneration() self._ui = ExtensionUi(runtime, self._generation) + self._system_prompt_override = system_prompt @property def cwd(self) -> Path: @@ -795,10 +815,15 @@ def session_id(self) -> str | None: @property def system_prompt(self) -> str: - """Return the active system prompt.""" + """Return the active prompt, including the current pre-run hook chain.""" self._generation.assert_active() + if self._system_prompt_override is not None: + return self._system_prompt_override return self._runtime.session_view.system_prompt + def _clear_system_prompt_override(self) -> None: + self._system_prompt_override = None + @property def is_running(self) -> bool: """Return whether an agent run is currently active.""" diff --git a/src/tau_coding/extensions/runtime.py b/src/tau_coding/extensions/runtime.py index 631a46aecb..a8975c1c54 100644 --- a/src/tau_coding/extensions/runtime.py +++ b/src/tau_coding/extensions/runtime.py @@ -31,6 +31,8 @@ AGENT_EVENT_TYPES, AGENT_EVENT_WILDCARD, LIFECYCLE_EVENT_TYPES, + BeforeAgentStartEvent, + BeforeAgentStartHookResult, CustomMessageView, ExtensionAPI, ExtensionCommandContext, @@ -63,6 +65,7 @@ ) from tau_coding.project_trust import ExtensionTrustResult, ProjectTrustEvent from tau_coding.resources import ResourceDiagnostic, TauResourcePaths +from tau_coding.system_prompt import SystemPromptInputs # Host callback that delivers a message through the frontend's serialized run # path when the session is idle. Carries the same presentation metadata as a @@ -850,6 +853,60 @@ async def run_input_hooks( current = result.text return InputHookOutcome(handled=False, text=current) + async def run_before_agent_start( + self, + system_prompt: str, + system_prompt_inputs: SystemPromptInputs, + ) -> str: + """Chain run-scoped system-prompt replacements in registration order.""" + current = system_prompt + handlers = tuple(self._handlers_for("before_agent_start")) + for extension, handler in handlers: + context = self._fresh_context(extension, system_prompt=current) + try: + result = await _resolve( + handler( + BeforeAgentStartEvent( + system_prompt=current, + system_prompt_inputs=system_prompt_inputs, + ), + context, + ) + ) + except Exception as exc: # noqa: BLE001 - extensions are an isolation boundary + self._record_runtime_failure( + extension, + "before_agent_start", + exc, + sensitive=True, + ) + continue + finally: + context._clear_system_prompt_override() + if result is None: + continue + if not isinstance(result, BeforeAgentStartHookResult): + self._record_bad_result( + extension, + "before_agent_start", + result, + sensitive=True, + ) + continue + replacement = result.system_prompt + if replacement is None: + continue + if not isinstance(replacement, str): + self._record_bad_result( + extension, + "before_agent_start.system_prompt", + replacement, + sensitive=True, + ) + continue + current = replacement + return current + async def emit_event(self, event: object) -> None: """Dispatch one canonical agent or coding-session event to extensions.""" event_type = getattr(event, "type", None) @@ -903,10 +960,15 @@ def _extension_by_name(self, name: str) -> RegisteredExtension | None: return extension return None - def _fresh_context(self, extension_name: str) -> ExtensionContext: + def _fresh_context( + self, + extension_name: str, + *, + system_prompt: str | None = None, + ) -> ExtensionContext: """Return a fresh context for one handler invocation.""" api = self._api_for(extension_name) - return ExtensionContext(self, api._generation) + return ExtensionContext(self, api._generation, system_prompt=system_prompt) def _api_for(self, extension_name: str) -> ExtensionAPI: extension = self._extension_by_name(extension_name) @@ -914,25 +976,42 @@ def _api_for(self, extension_name: str) -> ExtensionAPI: raise ExtensionError(f"unknown extension: {extension_name}") return extension.api - def _record_runtime_failure(self, extension: str, event: str, exc: Exception) -> None: + def _record_runtime_failure( + self, + extension: str, + event: str, + exc: Exception, + *, + sensitive: bool = False, + ) -> None: + detail = "details omitted for sensitive hook" if sensitive else repr(exc) self._runtime_diagnostics.append( ResourceDiagnostic( kind="extension", name=extension, - message=f"handler for `{event}` raised: {exc!r}", + message=f"handler for `{event}` raised: {detail}", severity="error", ) ) - def _record_bad_result(self, extension: str, event: str, result: object) -> None: + def _record_bad_result( + self, + extension: str, + event: str, + result: object, + *, + sensitive: bool = False, + ) -> None: + detail = ( + "unsupported result; details omitted for sensitive hook" + if sensitive + else f"unsupported result type {type(result).__name__}; ignored" + ) self._runtime_diagnostics.append( ResourceDiagnostic( kind="extension", name=extension, - message=( - f"handler for `{event}` returned unsupported" - f" result type {type(result).__name__}; ignored" - ), + message=f"handler for `{event}` returned {detail}", ) ) diff --git a/src/tau_coding/session.py b/src/tau_coding/session.py index ad0f79d45a..c2c4c1da66 100644 --- a/src/tau_coding/session.py +++ b/src/tau_coding/session.py @@ -3,7 +3,7 @@ from __future__ import annotations import string -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextlib import suppress from dataclasses import dataclass, replace from pathlib import Path @@ -127,7 +127,9 @@ from tau_coding.system_prompt import ( BuildSystemPromptOptions, ProjectContextFile, + SystemPromptInputs, build_system_prompt, + snapshot_system_prompt_inputs, ) from tau_coding.thinking import ( DEFAULT_THINKING_LEVEL, @@ -298,6 +300,7 @@ def __init__( state: SessionState, harness: AgentHarness, last_parent_id: str | None, + system_prompt_inputs: SystemPromptInputs | None = None, skills: tuple[Skill, ...] = (), prompt_templates: tuple[PromptTemplate, ...] = (), context_files: tuple[ProjectContextFile, ...] = (), @@ -316,6 +319,16 @@ def __init__( self._state = state self._harness = harness self._extension_runtime = extension_runtime or ExtensionRuntime() + if system_prompt_inputs is None: + system_prompt_inputs = snapshot_system_prompt_inputs( + BuildSystemPromptOptions( + cwd=config.cwd, + tools=config.tools or (), + custom_prompt=config.system, + extra_guidelines=self._extension_runtime.prompt_guidelines, + ) + ) + self._system_prompt_inputs = system_prompt_inputs self._image_support = image_support or ImageSupportState() self._session_start_pending = False self._last_parent_id = last_parent_id @@ -471,28 +484,40 @@ async def load(cls, config: CodingSessionConfig) -> CodingSession: ) ) tools = extension_runtime.compose_tools(base_tools) + system_prompt_options = BuildSystemPromptOptions( + cwd=config.cwd, + tools=tools, + skills=resources.skills if config.system is None else (), + custom_prompt=( + config.system + if config.system is not None + else ( + config.custom_system_prompt + if config.custom_system_prompt is not None + else resources.custom_system_prompt + ) + ), + append_system_prompt=( + None + if config.system is not None + else ( + config.append_system_prompt + if config.append_system_prompt is not None + else resources.append_system_prompt + ) + ), + context_files=resources.context_files if config.system is None else (), + extra_guidelines=extension_runtime.prompt_guidelines, + ) + system_prompt_inputs = snapshot_system_prompt_inputs(system_prompt_options) + system_prompt_options = replace( + system_prompt_options, + current_date=system_prompt_inputs.current_date, + ) system = ( config.system if config.system is not None - else build_system_prompt( - BuildSystemPromptOptions( - cwd=config.cwd, - tools=tools, - skills=resources.skills, - custom_prompt=( - config.custom_system_prompt - if config.custom_system_prompt is not None - else resources.custom_system_prompt - ), - append_system_prompt=( - config.append_system_prompt - if config.append_system_prompt is not None - else resources.append_system_prompt - ), - context_files=resources.context_files, - extra_guidelines=extension_runtime.prompt_guidelines, - ) - ) + else build_system_prompt(system_prompt_options) ) harness = AgentHarness( AgentHarnessConfig( @@ -511,6 +536,7 @@ async def load(cls, config: CodingSessionConfig) -> CodingSession: state=state, harness=harness, last_parent_id=_last_parent_id_from_state(state), + system_prompt_inputs=system_prompt_inputs, skills=resources.skills, prompt_templates=resources.prompt_templates, context_files=resources.context_files, @@ -779,7 +805,7 @@ async def export( title=_session_export_title(self), source=str(session_path) if session_path is not None else self.session_id, format=export_format, - system_prompt=self.system_prompt, + system_prompt=self._harness.config.system, ) @property @@ -812,7 +838,7 @@ def context_usage(self) -> ContextUsageEstimate: """Return structured context accounting for the active provider context.""" if self._context_usage_cache is None: self._context_usage_cache = estimate_context_usage( - system=self._harness.config.system, + system=self.system_prompt, messages=self._harness.messages, tools=tuple(self._harness.config.tools), ) @@ -821,7 +847,7 @@ def context_usage(self) -> ContextUsageEstimate: @property def system_prompt(self) -> str: """Return the effective system prompt sent to the model.""" - return self._harness.config.system + return self._harness.system_prompt @property def auto_compact_token_threshold(self) -> int | None: @@ -1478,7 +1504,7 @@ async def reload(self) -> CodingReloadSummary: append_system_prompt_path=self._append_system_prompt_path, ) before_extensions = _extension_signatures(self._extension_runtime) - before_tool_names = tuple(tool.name for tool in self._harness.config.tools) + before_tool_prompts = _tool_prompt_signatures(self._harness.config.tools) before_guidelines = self._extension_runtime.prompt_guidelines # Nothing below mutates the live session. Eligible extensions are loaded @@ -1569,29 +1595,46 @@ async def reload(self) -> CodingReloadSummary: append_system_prompt_path=resources.append_system_prompt_path, ) after_guidelines = staged_runtime.prompt_guidelines - system_prompt_rebuilt = self._config.system is None and ( + system_prompt_inputs_changed = ( before_system_prompt_inputs != after_system_prompt_inputs - or before_tool_names != tuple(tool.name for tool in staged_tools) + or before_tool_prompts != _tool_prompt_signatures(staged_tools) or before_guidelines != after_guidelines ) + system_prompt_rebuilt = self._config.system is None and system_prompt_inputs_changed staged_system = self._harness.config.system + staged_system_prompt_inputs = self._system_prompt_inputs if system_prompt_rebuilt: - staged_system = build_system_prompt( + staged_system_prompt_options = BuildSystemPromptOptions( + cwd=self._config.cwd, + tools=staged_tools, + skills=resources.skills, + custom_prompt=( + self._config.custom_system_prompt + if self._config.custom_system_prompt is not None + else resources.custom_system_prompt + ), + append_system_prompt=( + self._config.append_system_prompt + if self._config.append_system_prompt is not None + else resources.append_system_prompt + ), + context_files=resources.context_files, + extra_guidelines=after_guidelines, + ) + staged_system_prompt_inputs = snapshot_system_prompt_inputs( + staged_system_prompt_options + ) + staged_system_prompt_options = replace( + staged_system_prompt_options, + current_date=staged_system_prompt_inputs.current_date, + ) + staged_system = build_system_prompt(staged_system_prompt_options) + elif self._config.system is not None and system_prompt_inputs_changed: + staged_system_prompt_inputs = snapshot_system_prompt_inputs( BuildSystemPromptOptions( cwd=self._config.cwd, tools=staged_tools, - skills=resources.skills, - custom_prompt=( - self._config.custom_system_prompt - if self._config.custom_system_prompt is not None - else resources.custom_system_prompt - ), - append_system_prompt=( - self._config.append_system_prompt - if self._config.append_system_prompt is not None - else resources.append_system_prompt - ), - context_files=resources.context_files, + custom_prompt=self._config.system, extra_guidelines=after_guidelines, ) ) @@ -1632,6 +1675,7 @@ async def reload(self) -> CodingReloadSummary: self._extension_runtime = staged_runtime self._harness.config.tools = staged_tools self._harness.config.system = staged_system + self._system_prompt_inputs = staged_system_prompt_inputs if system_prompt_rebuilt: self._invalidate_context_usage_cache() staged_runtime.attach_harness_listener(self._harness.subscribe) @@ -1857,6 +1901,7 @@ async def _adopt_replacement( self._skills = replacement._skills self._prompt_templates = replacement._prompt_templates self._context_files = replacement._context_files + self._system_prompt_inputs = replacement._system_prompt_inputs self._custom_system_prompt = replacement._custom_system_prompt self._custom_system_prompt_path = replacement._custom_system_prompt_path self._append_system_prompt = replacement._append_system_prompt @@ -2080,7 +2125,14 @@ async def prompt( ) else: prompt_message = UserMessage(content=expanded_content) - events = self._harness.prompt_message(prompt_message) + effective_system_prompt = await self._extension_runtime.run_before_agent_start( + self._harness.config.system, + self._system_prompt_inputs, + ) + events = self._harness.prompt_message( + prompt_message, + system=effective_system_prompt, + ) self._invalidate_context_usage_cache() async for event in events: auto_name_message: str | None = None @@ -2136,7 +2188,7 @@ async def prompt( ) await self._extension_runtime.emit_event(retry_start) yield retry_start - events = self._harness.continue_() + events = self._harness.continue_(system=effective_system_prompt) self._invalidate_context_usage_cache() async for retry_event in events: if isinstance(retry_event, ToolExecutionEndEvent): @@ -2176,6 +2228,7 @@ async def prompt( try: await self._reconcile_run_persistence(events, context=context) finally: + self._invalidate_context_usage_cache() if events is not None: settled_event = await self._dispatch_agent_settled() if settled_event is not None: @@ -2192,7 +2245,11 @@ async def continue_(self) -> AsyncIterator[CodingSessionEvent]: events: AsyncIterator[AgentEvent] | None = None settled_event: AgentSettledEvent | None = None try: - events = self._harness.continue_() + effective_system_prompt = await self._extension_runtime.run_before_agent_start( + self._harness.config.system, + self._system_prompt_inputs, + ) + events = self._harness.continue_(system=effective_system_prompt) self._invalidate_context_usage_cache() async for event in events: if isinstance(event, ToolExecutionEndEvent): @@ -2223,6 +2280,7 @@ async def continue_(self) -> AsyncIterator[CodingSessionEvent]: try: await self._reconcile_run_persistence(events, context=context) finally: + self._invalidate_context_usage_cache() if events is not None: settled_event = await self._dispatch_agent_settled() if settled_event is not None: @@ -3269,6 +3327,10 @@ def _extension_signatures(runtime: ExtensionRuntime) -> tuple[tuple[object, ...] return tuple((name,) for name in runtime.extension_names) +def _tool_prompt_signatures(tools: Sequence[AgentTool]) -> tuple[tuple[object, ...], ...]: + return tuple((tool.name, tool.prompt_snippet, tool.prompt_guidelines) for tool in tools) + + def _system_prompt_resource_signatures( *, skills: tuple[Skill, ...], diff --git a/src/tau_coding/system_prompt.py b/src/tau_coding/system_prompt.py index 178d9a1255..1acb7461c9 100644 --- a/src/tau_coding/system_prompt.py +++ b/src/tau_coding/system_prompt.py @@ -21,6 +21,35 @@ class ProjectContextFile: content: str +@dataclass(frozen=True, slots=True, repr=False) +class SystemPromptSkill: + """Read-only skill-index metadata included in a generated system prompt.""" + + name: str + description: str | None + path: Path + disable_model_invocation: bool = False + + +@dataclass(frozen=True, slots=True, repr=False) +class SystemPromptInputs: + """Immutable extension-facing snapshot of system-prompt inputs. + + Prompt text and project paths may be sensitive, so the generated repr is + deliberately disabled. Extensions can inspect fields explicitly without + diagnostics accidentally rendering the complete prompt inputs. + """ + + custom_prompt: str | None + append_system_prompt: str | None + tools: tuple[str, ...] + guidelines: tuple[str, ...] + context_files: tuple[ProjectContextFile, ...] + skills: tuple[SystemPromptSkill, ...] + cwd: Path + current_date: date + + @dataclass(frozen=True, slots=True) class BuildSystemPromptOptions: """Options used to build Tau's system prompt.""" @@ -142,6 +171,28 @@ def add(value: str) -> None: return guidelines +def snapshot_system_prompt_inputs(options: BuildSystemPromptOptions) -> SystemPromptInputs: + """Freeze the structured inputs extensions may inspect before an agent run.""" + return SystemPromptInputs( + custom_prompt=options.custom_prompt, + append_system_prompt=options.append_system_prompt, + tools=tuple(tool.name for tool in options.tools), + guidelines=tuple(collect_prompt_guidelines(options.tools, options.extra_guidelines)), + context_files=tuple(options.context_files), + skills=tuple( + SystemPromptSkill( + name=skill.name, + description=skill.description, + path=skill.path, + disable_model_invocation=skill.disable_model_invocation, + ) + for skill in options.skills + ), + cwd=options.cwd, + current_date=options.current_date or date.today(), + ) + + def format_guidelines(tools: Sequence[AgentTool], extra_guidelines: Sequence[str] = ()) -> str: """Format prompt guidelines as markdown bullets.""" return "\n".join( diff --git a/tests/test_agent_harness.py b/tests/test_agent_harness.py index 73d7a39c8c..29cc73b334 100644 --- a/tests/test_agent_harness.py +++ b/tests/test_agent_harness.py @@ -189,6 +189,55 @@ async def execute( assert provider.calls[0][3] == [tool] +@pytest.mark.anyio +async def test_run_system_override_survives_tool_loop_and_resets() -> None: + harness: AgentHarness + + async def execute( + tool_call_id: str, + arguments: Mapping[str, JSONValue], + signal=None, # noqa: ANN001 + on_update=None, # noqa: ANN001 + ) -> AgentToolResult: + del tool_call_id, arguments, signal, on_update + harness.config.tools.append(tool) + return AgentToolResult(content="ran") + + tool = AgentTool( + name="work", + label="Work", + description="Do work.", + parameters={"type": "object"}, + execute_fn=execute, + ) + call = ToolCall(id="call-1", name="work", arguments={}) + provider = FakeProvider( + [ + [ + assistant_start(), + tool_call_end(call), + assistant_done(AssistantMessage(content=[call]), "toolUse"), + ], + [assistant_start(), assistant_done(AssistantMessage(content="Done"))], + [assistant_start(), assistant_done(AssistantMessage(content="Again"))], + ] + ) + harness = AgentHarness( + AgentHarnessConfig(provider=provider, model="fake", system="Base prompt", tools=[tool]) + ) + + _ = [event async for event in harness.prompt("start", system="Run-only prompt")] + + assert [call[1] for call in provider.calls[:2]] == ["Run-only prompt", "Run-only prompt"] + assert [len(call[3]) for call in provider.calls[:2]] == [1, 2] + assert harness.system_prompt == "Base prompt" + + _ = [event async for event in harness.prompt("next")] + + assert provider.calls[2][1] == "Base prompt" + assert harness.system_prompt == "Base prompt" + + def test_queue_mutators_return_canonical_snapshots() -> None: harness = AgentHarness( AgentHarnessConfig(provider=FakeProvider([]), model="fake", system="You are Tau.") @@ -254,7 +303,7 @@ async def test_cancelled_run_notifies_listeners_of_interrupted_tool_repair() -> harness.subscribe(seen.append) async def consume() -> None: - async for _event in harness.prompt("go"): + async for _event in harness.prompt("go", system="Run-only prompt"): pass task = asyncio.create_task(consume()) @@ -273,6 +322,7 @@ async def consume() -> None: if isinstance(event, MessageEndEvent) and isinstance(event.message, ToolResultMessage) ] assert [event.message.tool_call_id for event in repair_ends] == ["call-1"] + assert harness.system_prompt == "You are Tau." @pytest.mark.anyio @@ -301,6 +351,33 @@ async def consume() -> None: assert isinstance(harness.messages[-1], ToolResultMessage) +@pytest.mark.anyio +async def test_teardown_base_exception_still_resets_run_state() -> None: + tool_started = asyncio.Event() + release = asyncio.Event() + harness = _blocking_run_harness(tool_started, release) + + def cancel_cleanup(event: object) -> None: + if isinstance(event, MessageEndEvent) and isinstance(event.message, ToolResultMessage): + raise asyncio.CancelledError + + harness.subscribe(cancel_cleanup) + + async def consume() -> None: + async for _event in harness.prompt("go", system="Run-only prompt"): + pass + + task = asyncio.create_task(consume()) + await asyncio.wait_for(tool_started.wait(), timeout=5) + harness.cancel() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert harness.is_running is False + assert harness.system_prompt == "You are Tau." + + def test_harness_repairs_interrupted_tool_calls() -> None: call = ToolCall(id="call-1", name="read", arguments={"path": "README.md"}) harness = AgentHarness( diff --git a/tests/test_example_extensions.py b/tests/test_example_extensions.py index 013d39ad63..8570ccaed4 100644 --- a/tests/test_example_extensions.py +++ b/tests/test_example_extensions.py @@ -7,6 +7,7 @@ for how extension authors can test their own extensions. """ +from datetime import date from pathlib import Path import pytest @@ -14,7 +15,7 @@ from tau_agent.messages import TextContent from tau_agent.tools import AgentTool, AgentToolResult from tau_coding import TauResourcePaths -from tau_coding.extensions import ExtensionRuntime +from tau_coding.extensions import ExtensionRuntime, SystemPromptInputs pytestmark = pytest.mark.anyio @@ -87,6 +88,27 @@ async def test_hello_tool_defaults_to_world(tmp_path: Path) -> None: assert result.text == "Hello, world!" +# -- prompt_customizer.py ------------------------------------------------------- + + +async def test_prompt_customizer_replaces_prompt_for_run(tmp_path: Path) -> None: + runtime = _runtime_with_examples(tmp_path, "prompt_customizer.py") + inputs = SystemPromptInputs( + custom_prompt=None, + append_system_prompt=None, + tools=("read", "bash"), + guidelines=(), + context_files=(), + skills=(), + cwd=tmp_path, + current_date=date(2026, 8, 17), + ) + + result = await runtime.run_before_agent_start("Base prompt", inputs) + + assert result == "Base prompt\n\nActive tools for this run: read, bash." + + # -- permission_gate.py --------------------------------------------------------- diff --git a/tests/test_extensions.py b/tests/test_extensions.py index 0043b2915b..94a7a2ae67 100644 --- a/tests/test_extensions.py +++ b/tests/test_extensions.py @@ -3,11 +3,14 @@ import asyncio import sys import time +from dataclasses import replace +from datetime import date from pathlib import Path from typing import cast import pytest +import tau_coding.system_prompt as system_prompt_module from pi_event_helpers import assistant_done, assistant_start from tau_agent import AssistantMessage, ToolCall, UserMessage from tau_agent.messages import AgentMessage, assistant_content @@ -15,15 +18,26 @@ from tau_agent.tools import AgentTool, AgentToolResult from tau_agent.types import JSONValue from tau_ai import FakeProvider -from tau_coding import CodingSession, CodingSessionConfig, ResourceError, TauResourcePaths +from tau_coding import ( + CodingSession, + CodingSessionConfig, + ProjectContextFile, + ResourceError, + TauResourcePaths, +) from tau_coding.extensions import ( + BeforeAgentStartEvent, + BeforeAgentStartHookResult, CustomMessageView, ExtensionAPI, + ExtensionContext, ExtensionError, ExtensionRuntime, InputEvent, InputHookResult, MessageRenderOptions, + SystemPromptInputs, + SystemPromptSkill, ToolCallHookResult, ToolResultHookResult, discover_extensions, @@ -894,6 +908,150 @@ async def test_input_hook_receives_source_and_streaming_behavior(tmp_path: Path) assert seen[0].streaming_behavior == "steer" +async def test_before_agent_start_hooks_chain_and_isolate_failures(tmp_path: Path) -> None: + runtime = ExtensionRuntime() + runtime.bind(RecordingSession(tmp_path)) + first = cast(ExtensionAPI, _register_inline_extension(runtime, "first")) + broken = cast(ExtensionAPI, _register_inline_extension(runtime, "broken")) + last = cast(ExtensionAPI, _register_inline_extension(runtime, "last")) + seen: list[tuple[str, str]] = [] + captured_contexts: list[ExtensionContext] = [] + + def replace_prompt( + event: BeforeAgentStartEvent, context: ExtensionContext + ) -> BeforeAgentStartHookResult: + captured_contexts.append(context) + seen.append((event.system_prompt, context.system_prompt)) + return BeforeAgentStartHookResult(system_prompt=f"{event.system_prompt} one") + + def replace_from_context( + event: BeforeAgentStartEvent, context: ExtensionContext + ) -> BeforeAgentStartHookResult: + captured_contexts.append(context) + seen.append((event.system_prompt, context.system_prompt)) + return BeforeAgentStartHookResult(system_prompt=f"{context.system_prompt} two") + + SensitiveFailure = type("PrivatePromptLeakFailure", (RuntimeError,), {}) + SensitiveResult = type("PrivatePromptLeakResult", (), {}) + + def fail(event: BeforeAgentStartEvent, context: ExtensionContext) -> None: + captured_contexts.append(context) + raise SensitiveFailure(event.system_prompt) + + def invalid_result(event: BeforeAgentStartEvent, context: ExtensionContext) -> object: + del event + captured_contexts.append(context) + return SensitiveResult() + + def invalid_prompt( + event: BeforeAgentStartEvent, context: ExtensionContext + ) -> BeforeAgentStartHookResult: + del event + captured_contexts.append(context) + return BeforeAgentStartHookResult(system_prompt=SensitiveResult()) # type: ignore[arg-type] + + def finish_prompt( + event: BeforeAgentStartEvent, context: ExtensionContext + ) -> BeforeAgentStartHookResult: + captured_contexts.append(context) + seen.append((event.system_prompt, context.system_prompt)) + return BeforeAgentStartHookResult(system_prompt=f"{event.system_prompt} three") + + first.on("before_agent_start", replace_prompt) + first.on("before_agent_start", replace_from_context) + broken.on("before_agent_start", fail) + broken.on("before_agent_start", invalid_result) + broken.on("before_agent_start", invalid_prompt) + last.on("before_agent_start", finish_prompt) + inputs = SystemPromptInputs( + custom_prompt="sensitive custom base", + append_system_prompt="sensitive append", + tools=("read",), + guidelines=("sensitive effective guideline",), + context_files=(ProjectContextFile(path="/secret/AGENTS.md", content="sensitive context"),), + skills=( + SystemPromptSkill( + name="private-skill", + description="sensitive skill description", + path=Path("/secret/SKILL.md"), + disable_model_invocation=False, + ), + ), + cwd=tmp_path, + current_date=date(2026, 8, 17), + ) + + effective = await runtime.run_before_agent_start("base", inputs) + + assert effective == "base one two three" + assert seen == [ + ("base", "base"), + ("base one", "base one"), + ("base one two", "base one two"), + ] + assert [context.system_prompt for context in captured_contexts] == [ + "You are Tau.", + "You are Tau.", + "You are Tau.", + "You are Tau.", + "You are Tau.", + "You are Tau.", + ] + hook_diagnostics = [ + diagnostic + for diagnostic in runtime.diagnostics + if "before_agent_start" in diagnostic.message + ] + assert len(hook_diagnostics) == 3 + assert all("details omitted" in diagnostic.message for diagnostic in hook_diagnostics) + assert not any("PrivatePromptLeak" in diagnostic.message for diagnostic in hook_diagnostics) + assert not any("base one two" in diagnostic.message for diagnostic in hook_diagnostics) + assert "sensitive" not in repr(inputs) + assert "sensitive" not in repr( + BeforeAgentStartEvent(system_prompt="sensitive prompt", system_prompt_inputs=inputs) + ) + assert "sensitive" not in repr( + BeforeAgentStartHookResult(system_prompt="sensitive replacement") + ) + + +async def test_before_agent_start_uses_a_registration_snapshot(tmp_path: Path) -> None: + runtime = ExtensionRuntime() + api = cast(ExtensionAPI, _register_inline_extension(runtime, "dynamic")) + seen: list[str] = [] + registered = False + + def late(event: BeforeAgentStartEvent, context: object) -> None: + del event, context + seen.append("late") + + def register_late(event: BeforeAgentStartEvent, context: object) -> None: + nonlocal registered + del event, context + seen.append("initial") + if not registered: + registered = True + api.on("before_agent_start", late) + + api.on("before_agent_start", register_late) + inputs = SystemPromptInputs( + custom_prompt=None, + append_system_prompt=None, + tools=(), + guidelines=(), + context_files=(), + skills=(), + cwd=tmp_path, + current_date=date(2026, 8, 17), + ) + + await runtime.run_before_agent_start("base", inputs) + assert seen == ["initial"] + + await runtime.run_before_agent_start("base", inputs) + assert seen == ["initial", "initial", "late"] + + async def test_agent_event_fan_out_and_wildcard(tmp_path: Path) -> None: runtime = ExtensionRuntime() api = _register_inline_extension(runtime, "observer") @@ -1493,22 +1651,228 @@ async def test_extension_guideline_reaches_system_prompt(tmp_path: Path) -> None assert "Never commit directly to main" in session.system_prompt -async def test_reload_picks_up_guideline_changes(tmp_path: Path) -> None: - provider = FakeProvider([]) - session = await CodingSession.load(_session_config(tmp_path, provider)) +async def test_per_run_system_prompt_replacement_reaches_provider_without_persisting( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + body = ( + "from tau_coding.extensions import BeforeAgentStartHookResult\n\n" + "INPUTS = []\n" + "ACTIVE_PROMPTS = []\n\n\n" + "def _hook(event, context):\n" + " INPUTS.append(event.system_prompt_inputs)\n" + " return BeforeAgentStartHookResult(\n" + " system_prompt=f'{event.system_prompt}\\nRun marker: {len(INPUTS)}'\n" + " )\n\n\n" + "def setup(tau):\n" + " tau.add_prompt_guideline('Static extension guideline')\n" + " tau.on('before_agent_start', _hook)\n" + " tau.on(\n" + " 'agent_start',\n" + " lambda event, context: ACTIVE_PROMPTS.append(context.system_prompt),\n" + " )\n" + ) + paths = _paths(tmp_path) + skill_path = paths.root / "skills" / "testing" / "SKILL.md" + skill_path.parent.mkdir(parents=True, exist_ok=True) + skill_path.write_text( + "---\ndescription: Test code\ndisable-model-invocation: true\n---\n# Testing", + encoding="utf-8", + ) + call = ToolCall(id="call-1", name="read", arguments={"path": "missing.txt"}) + provider = FakeProvider( + [ + [ + assistant_start(model="fake"), + assistant_done( + message=AssistantMessage(content=assistant_content("", [call])), + finish_reason="toolUse", + ), + ], + [ + assistant_start(model="fake"), + assistant_done(message=AssistantMessage(content="first done")), + ], + [ + assistant_start(model="fake"), + assistant_done(message=AssistantMessage(content="second done")), + ], + [ + assistant_start(model="fake"), + assistant_done(message=AssistantMessage(content="continued")), + ], + ] + ) + config = replace( + _session_config(tmp_path, provider, extension_body=body), + custom_system_prompt="Sensitive custom base", + append_system_prompt="Sensitive append", + context_files=( + ProjectContextFile(path="virtual/AGENTS.md", content="Sensitive project context"), + ), + ) + + class LoadDate(date): + @classmethod + def today(cls) -> date: + return date(2026, 8, 17) + + class RunDate(date): + @classmethod + def today(cls) -> date: + return date(2026, 8, 18) + + monkeypatch.setattr(system_prompt_module, "date", LoadDate) + session = await CodingSession.load(config) + base_prompt = session.system_prompt + module = _loaded_extension_module("integration") + monkeypatch.setattr(system_prompt_module, "date", RunDate) + + first_run = session.prompt("first") + first_event = await anext(first_run) + assert first_event.type == "agent_start" + active_export = await session.export(tmp_path / "active-run.html") + assert "Run marker:" not in active_export.read_text(encoding="utf-8") + _ = [event async for event in first_run] + + assert provider.calls[0][1] == provider.calls[1][1] + assert provider.calls[0][1].endswith("Run marker: 1") + assert [provider.calls[0][1]] == module.ACTIVE_PROMPTS # type: ignore[attr-defined] + assert session.system_prompt == base_prompt + assert len(module.INPUTS) == 1 # type: ignore[attr-defined] + inputs = module.INPUTS[0] # type: ignore[attr-defined] + assert isinstance(inputs, SystemPromptInputs) + assert inputs.custom_prompt == "Sensitive custom base" + assert inputs.append_system_prompt == "Sensitive append" + assert inputs.tools[:4] == ("read", "write", "edit", "bash") + assert "Static extension guideline" in inputs.guidelines + assert inputs.context_files == ( + ProjectContextFile( + path="virtual/AGENTS.md", + content="Sensitive project context", + ), + ) + assert [ + ( + skill.name, + skill.description, + skill.path, + skill.disable_model_invocation, + ) + for skill in inputs.skills + ] == [("testing", "Test code", skill_path, True)] + assert inputs.cwd == paths.cwd + assert inputs.current_date == date(2026, 8, 17) + assert "Current date: 2026-08-17" in base_prompt + assert "testing" not in base_prompt + + _ = [event async for event in session.prompt("second")] + + assert provider.calls[2][1].endswith("Run marker: 2") + assert "Run marker: 1" not in provider.calls[2][1] + assert provider.calls[2][1].count("Run marker:") == 1 + assert session.system_prompt == base_prompt + + _ = [event async for event in session.continue_()] + + assert provider.calls[3][1].endswith("Run marker: 3") + assert [ + provider.calls[0][1], + provider.calls[2][1], + provider.calls[3][1], + ] == module.ACTIVE_PROMPTS # type: ignore[attr-defined] + assert not any("Run marker:" in message.text for message in session.messages) + assert "Run marker:" not in (tmp_path / "session.jsonl").read_text(encoding="utf-8") + + +@pytest.mark.parametrize( + ("static_system", "expected_rebuilt"), + [(None, True), ("Static system prompt", False)], +) +async def test_reload_picks_up_guideline_changes( + tmp_path: Path, + static_system: str | None, + expected_rebuilt: bool, +) -> None: + provider = FakeProvider( + [[assistant_start(model="fake"), assistant_done(message=AssistantMessage(content="done"))]] + ) + config = replace(_session_config(tmp_path, provider), system=static_system) + session = await CodingSession.load(config) assert "Prefer uv over pip" not in session.system_prompt paths = _paths(tmp_path) _write_extension( _user_extensions_dir(paths), "late_guideline", - "def setup(tau):\n tau.add_prompt_guideline('Prefer uv over pip')\n", + "INPUTS = []\n\n\n" + "def _capture(event, context):\n" + " INPUTS.append(event.system_prompt_inputs)\n\n\n" + "def setup(tau):\n" + " tau.add_prompt_guideline('Prefer uv over pip')\n" + " tau.on('before_agent_start', _capture)\n", ) summary = await session.reload() + _ = [event async for event in session.prompt("check reload")] + module = _loaded_extension_module("late_guideline") + + assert summary.system_prompt_rebuilt is expected_rebuilt + if static_system is None: + assert "Prefer uv over pip" in session.system_prompt + else: + assert session.system_prompt == static_system + inputs = module.INPUTS[0] # type: ignore[attr-defined] + assert inputs.custom_prompt == static_system + assert "Prefer uv over pip" in inputs.guidelines + + +async def test_reload_rebuilds_prompt_when_same_tool_metadata_changes(tmp_path: Path) -> None: + def extension_body(snippet: str, guideline: str) -> str: + return ( + "from tau_agent.tools import AgentTool, AgentToolResult\n\n" + "INPUTS = []\n\n\n" + "async def _run(tool_call_id, arguments, signal=None, on_update=None):\n" + " return AgentToolResult(content='done')\n\n\n" + "def _capture(event, context):\n" + " INPUTS.append(event.system_prompt_inputs)\n\n\n" + "def setup(tau):\n" + " tau.register_tool(AgentTool(\n" + " name='changing', label='Changing', description='Changes',\n" + " parameters={}, execute_fn=_run,\n" + f" prompt_snippet={snippet!r},\n" + f" prompt_guidelines=({guideline!r},),\n" + " ))\n" + " tau.on('before_agent_start', _capture)\n" + ) + + paths = _paths(tmp_path) + extension_path = _write_extension( + _user_extensions_dir(paths), + "changing_tool", + extension_body("First tool summary", "First tool guideline"), + ) + provider = FakeProvider( + [[assistant_start(model="fake"), assistant_done(message=AssistantMessage(content="done"))]] + ) + session = await CodingSession.load(_session_config(tmp_path, provider)) + assert "First tool summary" in session.system_prompt + assert "First tool guideline" in session.system_prompt + + extension_path.write_text( + extension_body("Second tool summary", "Second tool guideline"), + encoding="utf-8", + ) + summary = await session.reload() + _ = [event async for event in session.prompt("check metadata")] + module = _loaded_extension_module("changing_tool") + inputs = module.INPUTS[0] # type: ignore[attr-defined] assert summary.system_prompt_rebuilt is True - assert "Prefer uv over pip" in session.system_prompt + assert "Second tool summary" in session.system_prompt + assert "First tool summary" not in session.system_prompt + assert "Second tool guideline" in inputs.guidelines + assert "First tool guideline" not in inputs.guidelines async def test_session_start_deferred_until_host_emits(tmp_path: Path) -> None: diff --git a/tests/test_package_metadata.py b/tests/test_package_metadata.py index 1cb382548d..094a8f08f6 100644 --- a/tests/test_package_metadata.py +++ b/tests/test_package_metadata.py @@ -11,6 +11,7 @@ "tau_coding/data/docs/README.md", "tau_coding/data/docs/extensions.md", "tau_coding/data/examples/extensions/hello_tool.py", + "tau_coding/data/examples/extensions/prompt_customizer.py", } diff --git a/website/content/guides/extensions.md b/website/content/guides/extensions.md index d23728c3b5..d295e2aff5 100644 --- a/website/content/guides/extensions.md +++ b/website/content/guides/extensions.md @@ -406,6 +406,7 @@ Lifecycle and intercepting hooks: | `session_start` | `SessionStartEvent(reason)` | — | | `session_shutdown` | `SessionShutdownEvent(reason)` | — | | `input` | `InputEvent(text)` | `InputHookResult(action, text, message)` | +| `before_agent_start` | `BeforeAgentStartEvent(system_prompt, system_prompt_inputs)` | `BeforeAgentStartHookResult(system_prompt)` | | `tool_call` | `ToolCallHookEvent(tool_name, arguments)` | `ToolCallHookResult(block, reason, arguments)` | | `tool_result` | `ToolResultHookEvent(tool_name, arguments, result)` | `ToolResultHookResult(content, details)` | @@ -415,6 +416,12 @@ Lifecycle and intercepting hooks: - `input` runs on the raw prompt text before skill/template expansion. `action="transform"` rewrites it (transforms chain), `action="handled"` consumes it without an agent run and shows `message` as a notification. +- `before_agent_start` runs immediately before an agent run. Replacements chain + in registration order and apply to the initial provider request and every + tool-loop continuation in that run. The next run starts from the session's base + prompt, and replacements are not transcript entries. Inside the hook, + `event.system_prompt` and `context.system_prompt` both expose the current + chained value. - `tool_call` runs before a tool executes. `block=True` prevents execution and reports `reason` to the model; returning `arguments` rewrites the call. A crashing `tool_call` handler blocks the tool (fail-safe). @@ -424,6 +431,34 @@ Lifecycle and intercepting hooks: All other handler failures are contained: they are recorded as diagnostics (visible in `/session`) and never crash the session. +#### Per-run system prompts + +Return a `BeforeAgentStartHookResult` to replace the prompt for the current run: + +```python +from tau_coding.extensions import ( + BeforeAgentStartEvent, + BeforeAgentStartHookResult, +) + + +def setup(tau): + @tau.on("before_agent_start") + def customize(event: BeforeAgentStartEvent, context): + del context + return BeforeAgentStartHookResult( + system_prompt=f"{event.system_prompt}\n\nKeep the final answer concise." + ) +``` + +`event.system_prompt_inputs` is a frozen snapshot containing `custom_prompt`, +`append_system_prompt`, active tool names and effective `guidelines`, project +`context_files`, skill metadata (including `disable_model_invocation`), `cwd`, +and `current_date`. Prompt text, +instructions, and paths can be sensitive; Tau keeps these event and input +values out of diagnostic representations and autocomplete metadata, and +extensions should avoid logging them. + ### Messages and persistence `send_user_message` delivers a user message into the conversation. During a @@ -534,6 +569,7 @@ runtime only sees the former. See [`examples/extensions/`](https://github.com/huggingface/tau/tree/main/examples/extensions): - **`hello_tool.py`** — minimal custom tool. +- **`prompt_customizer.py`** — replaces the system prompt for each agent run. - **`permission_gate.py`** — blocks dangerous bash commands with the `tool_call` hook. @@ -557,7 +593,10 @@ tau -e ./tau-subagents Compared to Pi's extension system, Tau does not yet include a complete package manager (the installer has no registry, dependency installation, remove, or -package-update commands), custom providers, custom entry renderers (non-context -cards), declarative keyboard-shortcut registration, CLI flag registration, -system-prompt replacement, or context rewriting. The architecture document -(`dev-notes/architecture/phase-21-extensions.md`) tracks the extension design. +package-update commands), custom providers, extension-authored TUI widgets +(custom *message* rendering via `register_message_renderer` is supported; the +host-provided `context.ui` dialogs are supported), custom entry renderers +(non-context cards), declarative keyboard-shortcut registration, CLI flag +registration, context rewriting, or a project trust store. The architecture +document (`dev-notes/architecture/phase-21-extensions.md`) tracks the +extension design.