Add CodeMode dynamic_catalog flag for cache-stable tool disclosure (#243)

* feat(code_mode): add CodeModeDynamicCatalog for cache-stable tool disclosure

Moves CodeMode's per-tool signature catalog out of `run_code.description`
(which lives in the prompt-cache-keyed tool-definitions block) and into
agent instructions as a dynamic `InstructionPart`, then announces newly-
discovered tools via `RunContext.enqueue` rather than by mutating the
cached description. The result: the tool-defs block stays byte-stable
across Tool Search discoveries / per-step toolset swaps, only the dynamic
instructions and append-only system-prompt announcements grow.

This is the Tier-2 reshape from pydantic/pydantic-ai-harness#232. Opt-in
(default `CodeMode` keeps the catalog in `run_code.description`, which is
slightly cheaper on prompt size when the toolset never changes); pair
with `ToolSearch` or a churning toolset for the cache win.

Depends on pydantic/pydantic-ai#4980 for the pending message queue;
pinned to the `background-tools` branch of pydantic-ai-slim until that
lands. Cross-provider cache safety completes once
pydantic/pydantic-ai#5437 (XML-wrapped mid-conversation SystemPromptPart
mapping) ships.

* Re-lock pydantic-ai-slim to background-tools HEAD after main merge

* Adapt to pydantic-ai refactor: PendingMessage.payload→request

* Add e2e test exercising eager tool through run_code (closes coverage gap)

* Fold CodeModeDynamicCatalog into CodeMode as a `dynamic_catalog` flag

Replace the separate `CodeModeDynamicCatalog` capability (which surgically
wrapped `CodeMode`'s assembled toolset) with a `dynamic_catalog: bool` flag on
`CodeMode` itself. When set, `CodeModeToolset` keeps only the static base prose
in `run_code.description` and surfaces the sandboxed-tool catalog as a dynamic
`InstructionPart` via `get_instructions`, and `CodeMode` announces newly
discovered tools by enqueuing a `SystemPromptPart` through `ctx.enqueue`.

This drops the shim package and its dedicated test module, merging the behavior
and its tests into `code_mode`. The search addendum now stays in
`run_code.description` in both modes (it's cache-stable), an improvement over
the prior wrapper which dropped it.

* Enqueue discovery announcement as a bare SystemPromptPart

Re-lock pydantic-ai-slim to the latest background-tools HEAD, which now accepts
`ModelRequestPart`s directly in `enqueue` and renders mid-conversation
`SystemPromptPart`s inline rather than hoisting them (pydantic/pydantic-ai#5509).
Drop the `ModelRequest` wrapper workaround and enqueue the `SystemPromptPart`
directly; on the wire it renders as an XML-wrapped user prompt, so the catalog
cache stays intact across providers.

* deps: re-lock pydantic-ai-slim; give test RunContext a live pending-message queue

`RunContext.enqueue` now raises when the context has no pending-message queue (pydantic-ai
removed the silent-drop path). Pass `pending_messages=[]` in the test helper so the
discovery-announcement unit tests still exercise enqueue, and narrow `ctx.pending_messages`
before reading it. Real runs wire the queue automatically.

* code_mode: drop cast/Any from discovered-name extraction

_extract_discovered_names walked the tool-search return with cast() over Any,
which the no-typecast rule disallows. Validate the untrusted return through two
lenient pydantic TypeAdapters instead: a malformed catalog yields [] and a
malformed entry is skipped, matching the prior behavior with a fully-typed path.
Also rewrite the one em-dash in the get_instructions docstring to --.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* code_mode: import TypedDict from typing_extensions for py<3.12

pydantic rejects `typing.TypedDict` on Python < 3.12 (PydanticUserError at
TypeAdapter construction), which broke the 3.10/3.11 CI jobs. Match _toolset.py
and use `typing_extensions.TypedDict`.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: David SF <david.sanchez@pydantic.dev>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

pydantic_ai_harness/code_mode/README.md

@@ -90,13 +90,22 @@ Tools that match the selector are wrapped inside `run_code`. Non-matching tools
 
 When you mark tools or whole toolsets `defer_loading=True` ([Tool Search](https://ai.pydantic.dev/tools-advanced/#tool-search)), `CodeMode` keeps them out of `run_code` while they're undiscovered -- they pass straight through, so Tool Search drives them as usual (sent on the wire with `defer_loading` on providers with native tool search; otherwise dropped until discovered, with a `search_tools` tool alongside `run_code`). Once the model discovers a tool it comes back with `defer_loading=False`, and from then on `CodeMode` folds it into `run_code` like any other tool, so it's callable from generated code.
 
-That fold-in grows `run_code`'s description, which invalidates the prompt-cache prefix once at the moment of discovery (turns with no discovery stay cache-warm). To instead keep a Tool Search corpus fully native -- never folded into `run_code`, fully cache-stable, but not callable from inside it -- exclude it with a `tools` selector; corpus members carry `with_native` set to the managing native tool:
+That fold-in grows `run_code`'s description, which invalidates the prompt-cache prefix once at the moment of discovery (turns with no discovery stay cache-warm). Two ways to avoid the bust:
+
+- Pass `dynamic_catalog=True` to keep `run_code.description` static across discoveries -- the catalog of sandboxed-tool signatures moves into agent instructions (as a dynamic [`InstructionPart`](https://ai.pydantic.dev/api/messages/#pydantic_ai.messages.InstructionPart)) and newly-discovered tools are announced via [`ctx.enqueue`](https://ai.pydantic.dev/api/tools/#pydantic_ai.tools.RunContext.enqueue) instead of by rebuilding the description:
+
+```python
+CodeMode(dynamic_catalog=True)
+```
+
+  This pays off when paired with Tool Search: the tool-definitions block stays byte-stable so the prefix cache survives discoveries, at the cost of a larger (but cache-friendly) system prompt. With a fixed toolset and no Tool Search, the default keeps the system prompt shorter and is the better choice.
+
+- To instead keep a Tool Search corpus fully native -- never folded into `run_code`, but not callable from inside it -- exclude it with a `tools` selector; corpus members carry `with_native` set to the managing native tool:
 
 ```python
 CodeMode(tools=lambda ctx, td: td.with_native is None)
 ```
 
-A future Pydantic AI change will let `run_code`'s description stay static -- newly discovered tools announced separately -- so the fold-in costs nothing; until then, the selector above is the escape hatch.
 
 ### Metadata-based selection
 

pydantic_ai_harness/code_mode/_capability.py

@@ -2,15 +2,31 @@
 
 from __future__ import annotations
 
-from dataclasses import KW_ONLY, dataclass, field
+from collections.abc import Sequence
+from dataclasses import KW_ONLY, dataclass, field, replace
+from typing import TYPE_CHECKING, Any
 
+from pydantic import TypeAdapter, ValidationError
 from pydantic_ai import AbstractToolset
 from pydantic_ai.capabilities import AbstractCapability, CapabilityOrdering
 from pydantic_ai.capabilities._tool_search import ToolSearch as _ToolSearch
-from pydantic_ai.tools import AgentDepsT, ToolSelector
+from pydantic_ai.messages import ModelResponse, NativeToolSearchReturnPart, SystemPromptPart
+from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition, ToolSelector
+from typing_extensions import TypedDict
 
 from pydantic_ai_harness.code_mode._toolset import CodeModeMount, CodeModeOS, CodeModeToolset
 
+if TYPE_CHECKING:
+    from pydantic_ai.capabilities.abstract import ValidatedToolArgs
+    from pydantic_ai.messages import ToolCallPart
+    from pydantic_ai.models import ModelRequestContext
+
+
+_DISCOVERY_ANNOUNCEMENT_PREFIX = (
+    'New functions are now available inside `run_code`. Their signatures have been '
+    'added to the available-functions catalog in the system prompt'
+)
+
 
 @dataclass
 class CodeMode(AbstractCapability[AgentDepsT]):
@@ -73,16 +89,141 @@ class CodeMode(AbstractCapability[AgentDepsT]):
     mount: CodeModeMount | None = None
     """Host directories to expose to sandboxed `pathlib` code; each mount's `mode` controls whether writes reach the host."""
 
+    dynamic_catalog: bool = False
+    """Keep the `run_code` tool definition cache-stable as the sandboxed toolset grows.
+
+    By default the signatures of all sandboxed tools are rendered into `run_code`'s
+    description, which lives in the prompt-cache-keyed tool-definitions block. When the
+    toolset changes mid-run -- e.g. [`ToolSearch`][pydantic_ai.capabilities.ToolSearch]
+    reveals a new tool that then gets folded into `run_code` -- the description changes and
+    busts the prefix cache from that point on.
+
+    Set `dynamic_catalog=True` to instead:
+
+    - keep only the static base prose (sandbox restrictions, return-value contract) in
+      `run_code.description`, so the tool-definitions block stays byte-stable across
+      discoveries;
+    - move the "available functions" catalog (TypedDict definitions + signatures) into
+      agent instructions as a dynamic
+      [`InstructionPart`][pydantic_ai.messages.InstructionPart], which providers with
+      static/dynamic instruction splitting (Anthropic, Bedrock) place after the cache
+      breakpoint;
+    - announce newly-discovered tools via a short
+      [`SystemPromptPart`][pydantic_ai.messages.SystemPromptPart] enqueued through
+      [`RunContext.enqueue`][pydantic_ai.tools.RunContext.enqueue], so the model knows the
+      new functions are callable without rewriting the cached description.
+
+    This pays off when paired with [`ToolSearch`][pydantic_ai.capabilities.ToolSearch]: the
+    tool-definitions cache survives discoveries at the cost of a larger (but
+    cache-friendly) system prompt. With a fixed toolset and no `ToolSearch`, the default
+    keeps the system prompt shorter and is the better choice.
+    """
+
+    _announced_tools: set[str] = field(default_factory=set[str], init=False, repr=False)
+
     def get_ordering(self) -> CapabilityOrdering:
         """CodeMode wraps around ToolSearch so that search_tools stays native."""
         return CapabilityOrdering(position='outermost', wraps=[_ToolSearch])
 
+    async def for_run(self, ctx: RunContext[AgentDepsT]) -> CodeMode[AgentDepsT]:
+        """Return a fresh instance so concurrent runs don't share `_announced_tools`."""
+        if not self.dynamic_catalog:
+            return self
+        return replace(self)
+
     def get_wrapper_toolset(self, toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT] | None:
         """Wrap the agent's assembled toolset, splitting it into native + sandboxed subsets if needed."""
         return CodeModeToolset(
             wrapped=toolset,
             tool_selector=self.tools,
             max_retries=self.max_retries,
+            dynamic_catalog=self.dynamic_catalog,
             os_access=self.os_access,
             mount=self.mount,
         )
+
+    async def after_tool_execute(
+        self,
+        ctx: RunContext[AgentDepsT],
+        *,
+        call: ToolCallPart,
+        tool_def: ToolDefinition,
+        args: ValidatedToolArgs,
+        result: Any,
+    ) -> Any:
+        """Announce newly-discovered tools from a local `search_tools` return.
+
+        Only active with `dynamic_catalog=True`. The native-search path is handled by
+        [`after_model_request`][pydantic_ai_harness.CodeMode.after_model_request] instead
+        (server-side search emits a `NativeToolSearchReturnPart` rather than a regular tool
+        execute result).
+        """
+        if self.dynamic_catalog and tool_def.tool_kind == 'tool-search':
+            self._announce_newly_discovered(ctx, _extract_discovered_names(result))
+        return result
+
+    async def after_model_request(
+        self,
+        ctx: RunContext[AgentDepsT],
+        *,
+        request_context: ModelRequestContext,
+        response: ModelResponse,
+    ) -> ModelResponse:
+        """Announce newly-discovered tools from a native (server-side) tool-search return.
+
+        Only active with `dynamic_catalog=True`.
+        """
+        if self.dynamic_catalog:
+            for part in response.parts:
+                if isinstance(part, NativeToolSearchReturnPart):
+                    self._announce_newly_discovered(ctx, _extract_discovered_names(part.content))
+        return response
+
+    def _announce_newly_discovered(self, ctx: RunContext[AgentDepsT], names: Sequence[str]) -> None:
+        """Enqueue a system-prompt announcement for any names we haven't already announced."""
+        fresh = [n for n in names if n not in self._announced_tools]
+        if not fresh:
+            return
+        self._announced_tools.update(fresh)
+        listing = ', '.join(f'`{name}`' for name in fresh)
+        # Enqueue a `SystemPromptPart` so the announcement is framed as system-level context.
+        # Mid-conversation `SystemPromptPart`s are rendered inline (not hoisted to the top-level
+        # system prompt) on all providers since pydantic/pydantic-ai#5509, so this is cache-safe.
+        ctx.enqueue(SystemPromptPart(content=f'{_DISCOVERY_ANNOUNCEMENT_PREFIX}: {listing}.'))
+
+
+class _DiscoveredCatalog(TypedDict):
+    """Lenient view of a tool-search return: just the entry list, items left unvalidated."""
+
+    discovered_tools: list[object]
+
+
+class _DiscoveredEntry(TypedDict):
+    """Lenient view of one discovered-tool entry: only the name we announce."""
+
+    name: str
+
+
+_CATALOG_ADAPTER = TypeAdapter(_DiscoveredCatalog)
+_ENTRY_ADAPTER = TypeAdapter(_DiscoveredEntry)
+
+
+def _extract_discovered_names(content: object) -> list[str]:
+    """Read newly-discovered tool names from a tool-search return content.
+
+    Carried on both the local `ToolSearchReturnPart` and the native
+    `NativeToolSearchReturnPart`. Validated leniently: a malformed catalog yields `[]` and a
+    malformed entry is skipped, since the announcement is a courtesy nudge, not load-bearing
+    logic.
+    """
+    try:
+        catalog = _CATALOG_ADAPTER.validate_python(content)
+    except ValidationError:
+        return []
+    names: list[str] = []
+    for entry in catalog['discovered_tools']:
+        try:
+            names.append(_ENTRY_ADAPTER.validate_python(entry)['name'])
+        except ValidationError:
+            continue
+    return names

pydantic_ai_harness/code_mode/_toolset.py

@@ -6,7 +6,7 @@ import asyncio
 import keyword
 import re
 import warnings
-from collections.abc import Callable, Coroutine
+from collections.abc import Callable, Coroutine, Sequence
 from dataclasses import dataclass, field, replace
 from typing import Annotated, Any
 
@@ -14,7 +14,14 @@ from pydantic import Field, TypeAdapter
 from pydantic_ai import AbstractToolset, RunContext, ToolDefinition, WrapperToolset
 from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
 from pydantic_ai.function_signature import FunctionSignature
-from pydantic_ai.messages import ToolCallPart, ToolReturn, ToolReturnContent, ToolReturnPart, is_multi_modal_content
+from pydantic_ai.messages import (
+    InstructionPart,
+    ToolCallPart,
+    ToolReturn,
+    ToolReturnContent,
+    ToolReturnPart,
+    is_multi_modal_content,
+)
 from pydantic_ai.tool_manager import ToolManager
 from pydantic_ai.tools import AgentDepsT, ToolDenied, ToolSelector, matches_tool_selector
 from pydantic_ai.toolsets.abstract import SchemaValidatorProt, ToolsetTool
@@ -241,10 +248,25 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
     mount: CodeModeMount | None = None
     """Host directories to expose to sandboxed `pathlib` code; each mount's `mode` controls whether writes reach the host."""
 
+    dynamic_catalog: bool = False
+    """Move the sandboxed-tool catalog out of `run_code.description` and into instructions.
+
+    When `False` (default), every sandboxed tool's signature is rendered into the
+    `run_code` description, which lives in the prompt-cache-keyed tool-definitions block.
+    When `True`, the description keeps only the static base prose and the catalog is
+    surfaced as a dynamic [`InstructionPart`][pydantic_ai.messages.InstructionPart] via
+    [`get_instructions`][pydantic_ai_harness.code_mode.CodeModeToolset.get_instructions],
+    so Tool Search discoveries don't bust the tool-definitions cache prefix.
+    """
+
     # init=False so `replace()` in `for_run` produces a fresh instance with _repl=None,
     # giving each agent run isolated REPL state. Lazy-initialized on first call_tool.
     _repl: MontyRepl | None = field(default=None, init=False, repr=False)
 
+    # Catalog string stashed during `get_tools` (when `dynamic_catalog`) and read back by
+    # `get_instructions` in the same step. Empty when there's nothing to surface.
+    _last_catalog: str = field(default='', init=False, repr=False)
+
     # Tracks deferred-tool names we've already warned about so we don't spam the
     # logs every step. Reset on `for_run` because each run gets a fresh instance.
     _warned_deferred: set[str] = field(default_factory=set[str], init=False, repr=False)
@@ -262,8 +284,30 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
         new_self = replace(self, wrapped=new_wrapped)
         new_self._repl = self._repl
         new_self._warned_deferred = self._warned_deferred
+        new_self._last_catalog = self._last_catalog
         return new_self
 
+    async def get_instructions(
+        self, ctx: RunContext[AgentDepsT]
+    ) -> str | InstructionPart | Sequence[str | InstructionPart] | None:
+        """Surface the tool catalog as a dynamic instruction when `dynamic_catalog` is set.
+
+        The catalog is stashed by `get_tools` earlier in the same step. `dynamic=True` so
+        providers that split static/dynamic instructions (Anthropic, Bedrock) place a cache
+        breakpoint *before* the catalog -- discoveries change it but leave the static prefix
+        cache intact. When `dynamic_catalog` is off (or there are no sandboxed tools) the
+        stash is empty and we defer entirely to the wrapped toolset.
+        """
+        upstream = await self.wrapped.get_instructions(ctx)
+        if not self._last_catalog:
+            return upstream
+        catalog_part = InstructionPart(content=self._last_catalog, dynamic=True)
+        if upstream is None:
+            return catalog_part
+        if isinstance(upstream, (str, InstructionPart)):
+            return [upstream, catalog_part]
+        return [*upstream, catalog_part]
+
     async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
         """Return the `run_code` tool plus any native (non-sandboxed) tools."""
         wrapped_tools = await self.wrapped.get_tools(ctx)
@@ -293,9 +337,19 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
 
         callable_defs, sanitized_to_original = self._partition_callable_tools(sandboxed_tools)
 
-        description = self._build_description(
-            callable_defs, has_os=self.os_access is not None, has_mount=self.mount is not None
-        )
+        # `dynamic_catalog` keeps the catalog out of `run_code.description` (cache-stable
+        # tool-defs block) and surfaces it via `get_instructions` instead. Stash it for the
+        # `get_instructions` call later this step; empty string means "nothing to surface".
+        # The base prose stays host-aware in both modes -- its OS/mount restriction line is
+        # static (it doesn't change per discovery), so it belongs in the cached description.
+        has_os = self.os_access is not None
+        has_mount = self.mount is not None
+        if self.dynamic_catalog:
+            description = _base_description(has_os=has_os, has_mount=has_mount)
+            self._last_catalog = self._render_catalog(callable_defs)
+        else:
+            description = self._build_description(callable_defs, has_os=has_os, has_mount=has_mount)
+            self._last_catalog = ''
 
         if _RUN_CODE_TOOL_NAME in native_tools:
             raise UserError(
@@ -568,8 +622,22 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
     def _build_description(callable_defs: dict[str, ToolDefinition], *, has_os: bool, has_mount: bool) -> str:
         """Render the `run_code` description: base prose + TypedDicts + function signatures."""
         base = _base_description(has_os=has_os, has_mount=has_mount)
-        if not callable_defs:
+        catalog = CodeModeToolset._render_catalog(callable_defs)
+        if not catalog:
             return base
+        return base + '\n\n' + catalog
+
+    @staticmethod
+    def _render_catalog(callable_defs: dict[str, ToolDefinition]) -> str:
+        """Render the functions-header + TypedDict + function-signature blocks, or `''` if no defs.
+
+        Excludes the `run_code` base prose; the catalog is the discovery-driven portion that's
+        cache-hostile when carried in `run_code.description`. Used by `_build_description`
+        (default static-description path) and by `get_instructions` (the `dynamic_catalog`
+        path, which moves it into instructions instead).
+        """
+        if not callable_defs:
+            return ''
 
         sigs, conflicting = _get_sigs_and_conflicting(callable_defs)
         type_blocks = FunctionSignature.render_type_definitions(sigs, conflicting)
@@ -580,9 +648,7 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
 
         has_sync = any(td.sequential for td in callable_defs.values())
         has_async = any(not td.sequential for td in callable_defs.values())
-        header = _functions_header(has_sync=has_sync, has_async=has_async)
-
-        sections = [base, header]
+        sections = [_functions_header(has_sync=has_sync, has_async=has_async)]
         if type_blocks:
             sections.append('```python\n' + '\n\n'.join(type_blocks) + '\n```')
         sections.append('```python\n' + '\n\n'.join(function_blocks) + '\n```')

tests/code_mode/test_code_mode.py

@@ -60,6 +60,8 @@ def build_run_context(deps: T, run_step: int = 0) -> RunContext[T]:
         prompt=None,
         messages=[],
         run_step=run_step,
+        # A live queue so `ctx.enqueue` works in tests; a real run wires this to the run's queue.
+        pending_messages=[],
     )
 
 
@@ -1858,6 +1860,439 @@ class TestToolSearchIntegration:
         assert ToolSearch in ordering.wraps
 
 
+class TestDynamicCatalog:
+    """`CodeMode(dynamic_catalog=True)`: move the catalog to instructions + announce discoveries.
+
+    Two surfaces:
+
+    1. **Catalog placement** — `CodeModeToolset` strips signatures from `run_code.description`
+       and re-exposes them as a dynamic `InstructionPart` via `get_instructions`.
+    2. **Discovery announcements** — `CodeMode.after_tool_execute` (local search) and
+       `after_model_request` (native search) enqueue a `SystemPromptPart` so the model
+       learns that freshly-discovered tools are callable.
+    """
+
+    # -- catalog placement -------------------------------------------------
+
+    async def test_description_drops_signatures_keeps_base_prose(self) -> None:
+        toolset = CodeModeToolset(wrapped=_build_function_toolset(add), tool_selector='all', dynamic_catalog=True)
+        tools = await toolset.get_tools(build_run_context(None))
+
+        description = tools['run_code'].tool_def.description
+        assert description is not None
+        # The signature is gone from the description...
+        assert 'async def add' not in description
+        # ...but the static base prose remains.
+        assert 'sandboxed environment' in description
+
+    async def test_catalog_surfaces_as_dynamic_instruction_part(self) -> None:
+        toolset = CodeModeToolset(wrapped=_build_function_toolset(add), tool_selector='all', dynamic_catalog=True)
+        ctx = build_run_context(None)
+        await toolset.get_tools(ctx)
+        instructions = await toolset.get_instructions(ctx)
+
+        from pydantic_ai.messages import InstructionPart
+
+        # No upstream instructions → the catalog is the only InstructionPart returned.
+        assert isinstance(instructions, InstructionPart)
+        assert 'async def add' in instructions.content
+        # `dynamic=True` so Anthropic/Bedrock place the cache breakpoint before this block.
+        assert instructions.dynamic is True
+
+    async def test_get_instructions_appends_to_upstream_string(self) -> None:
+        from pydantic_ai.messages import InstructionPart
+
+        class _UpstreamToolset(FunctionToolset[None]):
+            async def get_instructions(self, ctx: RunContext[None]) -> str:  # pyright: ignore[reportIncompatibleMethodOverride]
+                return 'wrapped instructions'
+
+        toolset = CodeModeToolset(
+            wrapped=_UpstreamToolset(tools=[Tool(add)]), tool_selector='all', dynamic_catalog=True
+        )
+        ctx = build_run_context(None)
+        await toolset.get_tools(ctx)
+        instructions = await toolset.get_instructions(ctx)
+
+        assert isinstance(instructions, list)
+        assert instructions[0] == 'wrapped instructions'
+        assert isinstance(instructions[1], InstructionPart)
+        assert 'async def add' in instructions[1].content
+
+    async def test_get_instructions_appends_to_upstream_sequence(self) -> None:
+        from pydantic_ai.messages import InstructionPart
+
+        class _UpstreamToolset(FunctionToolset[None]):
+            async def get_instructions(  # pyright: ignore[reportIncompatibleMethodOverride]
+                self, ctx: RunContext[None]
+            ) -> list[str | InstructionPart]:
+                return ['a', InstructionPart(content='b')]
+
+        toolset = CodeModeToolset(
+            wrapped=_UpstreamToolset(tools=[Tool(add)]), tool_selector='all', dynamic_catalog=True
+        )
+        ctx = build_run_context(None)
+        await toolset.get_tools(ctx)
+        instructions = await toolset.get_instructions(ctx)
+
+        assert isinstance(instructions, list)
+        assert instructions[0] == 'a'
+        assert isinstance(instructions[1], InstructionPart) and instructions[1].content == 'b'
+        # The catalog is appended at the end.
+        assert isinstance(instructions[2], InstructionPart) and 'async def add' in instructions[2].content
+
+    async def test_default_keeps_catalog_in_description_and_no_instructions(self) -> None:
+        """With `dynamic_catalog=False` (default) the catalog stays in the description."""
+        toolset = CodeModeToolset(wrapped=_build_function_toolset(add), tool_selector='all')
+        ctx = build_run_context(None)
+        tools = await toolset.get_tools(ctx)
+
+        description = tools['run_code'].tool_def.description
+        assert description is not None
+        assert 'async def add' in description
+        # Nothing stashed → defer to upstream (None for FunctionToolset).
+        assert await toolset.get_instructions(ctx) is None
+
+    async def test_empty_catalog_emits_no_instruction(self) -> None:
+        """No sandboxed tools → empty catalog → defer to upstream instructions."""
+        toolset = CodeModeToolset(wrapped=_build_function_toolset(), tool_selector='all', dynamic_catalog=True)
+        ctx = build_run_context(None)
+        tools = await toolset.get_tools(ctx)
+
+        description = tools['run_code'].tool_def.description
+        assert description is not None
+        assert 'sandboxed environment' in description
+        assert await toolset.get_instructions(ctx) is None
+
+    async def test_search_addendum_stays_in_description(self) -> None:
+        """The (cache-stable) search addendum stays in `run_code.description` even in dynamic mode."""
+        toolset = CodeModeToolset(
+            wrapped=_StaticToolset([_search_tool_def()]), tool_selector='all', dynamic_catalog=True
+        )
+        ctx = build_run_context(None)
+        tools = await toolset.get_tools(ctx)
+
+        description = tools['run_code'].tool_def.description
+        assert description is not None
+        assert _TOOL_SEARCH_ADDENDUM.strip() in description
+
+    async def test_for_run_step_preserves_catalog_stash(self) -> None:
+        """A per-step rebuild must carry `_last_catalog` so instructions stay populated."""
+
+        class _ChangingToolset(FunctionToolset[None]):
+            async def for_run_step(self, ctx: RunContext[None]) -> AbstractToolset[None]:
+                # Force `CodeModeToolset.for_run_step` down the `new_wrapped is not self.wrapped`
+                # branch by returning a distinct (but equivalent) wrapped instance.
+                return type(self)(tools=list(self.tools.values()))
+
+        toolset = CodeModeToolset(
+            wrapped=_ChangingToolset(tools=[Tool(add)]), tool_selector='all', dynamic_catalog=True
+        )
+        ctx = build_run_context(None)
+        await toolset.get_tools(ctx)
+        stashed = toolset._last_catalog  # pyright: ignore[reportPrivateUsage]
+        assert stashed  # populated
+
+        new_toolset = await toolset.for_run_step(ctx)
+        assert isinstance(new_toolset, CodeModeToolset)
+        assert new_toolset is not toolset
+        assert new_toolset._last_catalog == stashed  # pyright: ignore[reportPrivateUsage]
+
+    # -- capability per-run state -----------------------------------------
+
+    async def test_for_run_returns_fresh_state_when_enabled(self) -> None:
+        cap = CodeMode[None](dynamic_catalog=True)
+        cap._announced_tools.add('foo')  # pyright: ignore[reportPrivateUsage]
+        fresh = await cap.for_run(build_run_context(None))
+        assert fresh is not cap
+        assert fresh._announced_tools == set()  # pyright: ignore[reportPrivateUsage]
+
+    async def test_for_run_returns_self_when_disabled(self) -> None:
+        cap = CodeMode[None]()
+        assert await cap.for_run(build_run_context(None)) is cap
+
+    # -- discovery announcement: local search path ------------------------
+
+    async def test_announce_on_local_search_return(self) -> None:
+        from pydantic_ai.messages import ModelRequest, SystemPromptPart, ToolCallPart
+
+        cap = CodeMode[None](dynamic_catalog=True)
+        ctx = build_run_context(None)
+        await cap.after_tool_execute(
+            ctx,
+            call=ToolCallPart(tool_name='search_tools', args={}, tool_call_id='c1'),
+            tool_def=_search_tool_def_with_kind(),
+            args={},
+            result={'discovered_tools': [{'name': 'weather', 'description': '...'}]},
+        )
+
+        assert ctx.pending_messages is not None
+        assert len(ctx.pending_messages) == 1
+        [request] = ctx.pending_messages[0].messages
+        assert isinstance(request, ModelRequest)
+        [part] = request.parts
+        assert isinstance(part, SystemPromptPart)
+        assert '`weather`' in part.content
+
+    async def test_no_announce_when_disabled(self) -> None:
+        """With `dynamic_catalog=False`, the hooks are inert even on a real search return."""
+        from pydantic_ai.messages import ToolCallPart
+
+        cap = CodeMode[None]()
+        ctx = build_run_context(None)
+        await cap.after_tool_execute(
+            ctx,
+            call=ToolCallPart(tool_name='search_tools', args={}, tool_call_id='c1'),
+            tool_def=_search_tool_def_with_kind(),
+            args={},
+            result={'discovered_tools': [{'name': 'weather'}]},
+        )
+        assert ctx.pending_messages == []
+
+    async def test_announce_skipped_when_no_discoveries(self) -> None:
+        from pydantic_ai.messages import ToolCallPart
+
+        cap = CodeMode[None](dynamic_catalog=True)
+        ctx = build_run_context(None)
+        await cap.after_tool_execute(
+            ctx,
+            call=ToolCallPart(tool_name='search_tools', args={}, tool_call_id='c1'),
+            tool_def=_search_tool_def_with_kind(),
+            args={},
+            result={'discovered_tools': []},
+        )
+        assert ctx.pending_messages == []
+
+    async def test_no_announce_for_non_search_tool(self) -> None:
+        """`tool_kind != 'tool-search'` short-circuits before reading the result."""
+        from pydantic_ai.messages import ToolCallPart
+
+        cap = CodeMode[None](dynamic_catalog=True)
+        ctx = build_run_context(None)
+        await cap.after_tool_execute(
+            ctx,
+            call=ToolCallPart(tool_name='add', args={}, tool_call_id='c1'),
+            tool_def=ToolDefinition(name='add', description='', parameters_json_schema={}),
+            args={},
+            # Even a `discovered_tools`-shaped result doesn't trigger an announcement:
+            # the `tool_kind` guard is the source of truth.
+            result={'discovered_tools': [{'name': 'spurious'}]},
+        )
+        assert ctx.pending_messages == []
+
+    async def test_no_duplicate_announcement_for_same_tool(self) -> None:
+        from pydantic_ai.messages import ToolCallPart
+
+        cap = CodeMode[None](dynamic_catalog=True)
+        ctx = build_run_context(None)
+        result = {'discovered_tools': [{'name': 'weather'}]}
+        for cid in ('c1', 'c2'):
+            await cap.after_tool_execute(
+                ctx,
+                call=ToolCallPart(tool_name='search_tools', args={}, tool_call_id=cid),
+                tool_def=_search_tool_def_with_kind(),
+                args={},
+                result=result,
+            )
+        # Only the first discovery of `weather` announces.
+        assert ctx.pending_messages is not None
+        assert len(ctx.pending_messages) == 1
+
+    # -- discovery announcement: native search path -----------------------
+
+    async def test_announce_on_native_search_return_part(self) -> None:
+        from pydantic_ai.messages import ModelRequest, ModelResponse, NativeToolSearchReturnPart, SystemPromptPart
+        from pydantic_ai.usage import RequestUsage
+
+        cap = CodeMode[None](dynamic_catalog=True)
+        ctx = build_run_context(None)
+        response = ModelResponse(
+            parts=[
+                NativeToolSearchReturnPart(
+                    tool_name='tool_search',
+                    content={'discovered_tools': [{'name': 'weather', 'description': 'Get the weather.'}]},
+                    tool_call_id='c1',
+                )
+            ],
+            usage=RequestUsage(input_tokens=1, output_tokens=1),
+        )
+        await cap.after_model_request(ctx, request_context=None, response=response)  # pyright: ignore[reportArgumentType]
+
+        assert ctx.pending_messages is not None
+        assert len(ctx.pending_messages) == 1
+        [request] = ctx.pending_messages[0].messages
+        assert isinstance(request, ModelRequest)
+        [part] = request.parts
+        assert isinstance(part, SystemPromptPart) and '`weather`' in part.content
+
+    async def test_no_announce_for_unrelated_response_parts(self) -> None:
+        from pydantic_ai.messages import ModelResponse, NativeToolReturnPart, TextPart
+        from pydantic_ai.usage import RequestUsage
+
+        cap = CodeMode[None](dynamic_catalog=True)
+        ctx = build_run_context(None)
+        response = ModelResponse(
+            parts=[
+                TextPart('hi'),
+                NativeToolReturnPart(tool_name='whatever', content='ignored', tool_call_id='c1'),
+            ],
+            usage=RequestUsage(input_tokens=1, output_tokens=1),
+        )
+        await cap.after_model_request(ctx, request_context=None, response=response)  # pyright: ignore[reportArgumentType]
+        assert ctx.pending_messages == []
+
+    # -- `_extract_discovered_names` edge cases ---------------------------
+
+    @pytest.mark.parametrize(
+        ('content', 'expected'),
+        [
+            ('not a dict', []),
+            ({}, []),
+            ({'discovered_tools': 'not a list'}, []),
+            ({'discovered_tools': [{'name': 'a'}, 'not a dict', {'no_name': 1}, {'name': 42}]}, ['a']),
+        ],
+    )
+    def test_extract_discovered_names_handles_malformed(self, content: Any, expected: list[str]) -> None:
+        from pydantic_ai_harness.code_mode._capability import (
+            _extract_discovered_names,  # pyright: ignore[reportPrivateUsage]
+        )
+
+        assert _extract_discovered_names(content) == expected
+
+    # -- end-to-end via `Agent.run` ---------------------------------------
+
+    async def test_agent_run_announces_discovery_and_lists_catalog_in_instructions(self) -> None:
+        """`Agent.run` end-to-end: catalog in instructions, discovery enqueues an announcement.
+
+        Two-step run:
+          1. Model calls `search_tools(['weather'])` (the discovery surface).
+          2. After the local tool-search returns, `CodeMode.after_tool_execute` enqueues a
+             `SystemPromptPart`; the pending-message queue drains it into the next request.
+             On the wire it renders as an (XML-wrapped) `UserPromptPart` — mid-conversation
+             system content is no longer hoisted (pydantic/pydantic-ai#5509) — so the model
+             sees the announcement inline and replies.
+        """
+        from pydantic_ai.capabilities import ToolSearch
+        from pydantic_ai.messages import (
+            ModelMessage,
+            ModelRequest,
+            ModelResponse,
+            SystemPromptPart,
+            TextPart,
+            ToolCallPart,
+            ToolReturnPart,
+            ToolSearchReturnPart,
+            UserPromptPart,
+        )
+        from pydantic_ai.models.function import AgentInfo, FunctionModel
+        from pydantic_ai.usage import RequestUsage
+
+        captured_prompt_texts: list[list[str]] = []
+        captured_descriptions: list[str] = []
+
+        def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+            run_code_def = next(td for td in info.function_tools if td.name == 'run_code')
+            assert run_code_def.description is not None
+            captured_descriptions.append(run_code_def.description)
+
+            last_request = messages[-1]
+            assert isinstance(last_request, ModelRequest)
+            # The announcement may arrive as a `SystemPromptPart` or, after wire-rendering of
+            # mid-conversation system content, an (XML-wrapped) `UserPromptPart` — capture both.
+            captured_prompt_texts.append(
+                [
+                    p.content
+                    for p in last_request.parts
+                    if isinstance(p, (SystemPromptPart, UserPromptPart)) and isinstance(p.content, str)
+                ]
+            )
+
+            if len(captured_descriptions) == 1:
+                return ModelResponse(
+                    parts=[ToolCallPart(tool_name='search_tools', args={'queries': ['weather']}, tool_call_id='c1')],
+                    usage=RequestUsage(input_tokens=1, output_tokens=1),
+                )
+            return ModelResponse(parts=[TextPart('done')], usage=RequestUsage(input_tokens=1, output_tokens=1))
+
+        def weather(city: str) -> str:
+            """Get the weather."""
+            return f'sunny in {city}'  # pragma: no cover — only the signature matters.
+
+        agent: Agent[None, str] = Agent(
+            FunctionModel(model_fn),
+            tools=[Tool(weather, defer_loading=True)],
+            capabilities=[ToolSearch[None](), CodeMode[None](dynamic_catalog=True)],
+        )
+        result = await agent.run('please find a weather tool')
+
+        # `run_code.description` stayed static across both turns — no signature in the tool-defs block.
+        assert all('async def' not in d for d in captured_descriptions)
+        # The discovery announcement landed in turn 2's request (system- or user-framed).
+        assert len(captured_prompt_texts) >= 2
+        assert 'weather' in '\n'.join(captured_prompt_texts[1])
+        # The local `ToolSearchReturnPart` is in history.
+        history = result.all_messages()
+        assert any(
+            isinstance(p, ToolSearchReturnPart) for msg in history if isinstance(msg, ModelRequest) for p in msg.parts
+        )
+        assert any(
+            isinstance(p, ToolReturnPart) and p.tool_name == 'search_tools'
+            for msg in history
+            if isinstance(msg, ModelRequest)
+            for p in msg.parts
+        )
+        assert result.output == 'done'
+
+    async def test_run_code_calls_eager_tool_with_catalog_in_instructions(self) -> None:
+        """An eager tool whose signature lives in instructions is still callable via `run_code`."""
+        from pydantic_ai.messages import (
+            ModelMessage,
+            ModelRequest,
+            ModelResponse,
+            TextPart,
+            ToolCallPart,
+            ToolReturnPart,
+        )
+        from pydantic_ai.models.function import AgentInfo, FunctionModel
+        from pydantic_ai.usage import RequestUsage
+
+        def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+            run_code_def = next(td for td in info.function_tools if td.name == 'run_code')
+            assert run_code_def.description is not None
+            assert 'async def add' not in run_code_def.description
+            if not any(isinstance(msg, ModelResponse) for msg in messages):
+                return ModelResponse(
+                    parts=[
+                        ToolCallPart(
+                            tool_name='run_code',
+                            # CodeMode renders all tools as `async def` by default — use `await`.
+                            args={'code': 'result = await add(a=3, b=4)\nresult'},
+                            tool_call_id='c1',
+                        )
+                    ],
+                    usage=RequestUsage(input_tokens=1, output_tokens=1),
+                )
+            last_request = messages[-1]
+            assert isinstance(last_request, ModelRequest)
+            run_code_return = next(p for p in last_request.parts if isinstance(p, ToolReturnPart))
+            return ModelResponse(
+                parts=[TextPart(f'got {run_code_return.content}')],
+                usage=RequestUsage(input_tokens=1, output_tokens=1),
+            )
+
+        agent: Agent[None, str] = Agent(
+            FunctionModel(model_fn),
+            tools=[Tool(add)],
+            capabilities=[CodeMode[None](dynamic_catalog=True)],
+        )
+        result = await agent.run('add 3 and 4 via run_code')
+        assert result.output == 'got 7'
+
+
+def _search_tool_def_with_kind(name: str = 'search_tools') -> ToolDefinition:
+    """A `search_tools` ToolDefinition carrying `tool_kind='tool-search'` (the announcement trigger)."""
+    return ToolDefinition(name=name, description='', parameters_json_schema={}, tool_kind='tool-search')
+
+
 def _unused_os_callback(fn: OsFunction, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
     """An `os` callback for tests that only assert description/forwarding, never run code."""
     return NOT_HANDLED  # pragma: no cover - never invoked by these tests