feat(planning): add Planning capability with cache-safe plan reminders (#266)

Planning gives the model a single write_plan(items) tool that owns the plan via whole-plan replacement (no fragile indices). The current plan is surfaced back as an ephemeral reminder appended to the request tail in wrap_model_request, behind a CachePoint: - the reminder reaches the model but is never written to the durable message history, so no stale copies accumulate - the CachePoint sits before the reminder, so the cached prefix (tools + system + real conversation) stays byte-identical across turns and only the small reminder falls outside the cache Static usage guidance goes in the system prompt (cache-stable); the mutable plan is never injected there. Co-authored-by: Claude <noreply@anthropic.com>
2026-07-21 02:45:34 +00:00 · 2026-06-09 11:05:19 -05:00
parent 442e9ff493
commit 3d3247110c
6 changed files with 571 additions and 0 deletions
@@ -0,0 +1,102 @@
+# Planning
+
+Give an agent a structured, self-updating task plan -- without ever invalidating the prompt cache.
+
+## The problem
+
+Long agentic runs drift: the model loses track of what it set out to do and what's left. The usual fix -- keep a running plan and re-inject it into the system prompt each turn -- invalidates the prompt cache. The system prompt sits at the front of the request, so every plan edit changes the cached prefix and forces the whole conversation to be re-processed at full token price.
+
+## The solution
+
+`Planning` gives the model one tool, `write_plan`, that owns the plan (whole-plan replacement -- pass the full list every call, no indices). The current plan is surfaced back to the model as an ephemeral reminder appended to the tail of each request, behind a cache breakpoint:
+
+- The reminder is added in `wrap_model_request`, which runs *after* the durable history is persisted, so it reaches the model but is never written to `message_history`. No reminders accumulate across turns.
+- A `CachePoint` is placed immediately *before* the reminder, so the cached prefix (tools + system + real conversation) stays byte-identical turn over turn. Only the reminder falls outside the cache.
+
+So the plan stays current in the model's view while the cached prefix is never invalidated; the only added cost is re-reading the reminder each turn.
+
+```python
+from pydantic_ai import Agent
+from pydantic_ai_harness.experimental.planning import Planning
+
+agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[Planning()])
+
+result = agent.run_sync('Refactor the auth module and add tests.')
+print(result.output)
+```
+
+## The tool
+
+| Tool | Purpose |
+|---|---|
+| `write_plan(items)` | Create or replace the full plan. The model passes the entire ordered list every time, including unchanged, completed, and cancelled steps. |
+
+Each item is a `content` string plus a `status` (`pending`, `in_progress`, `completed`, `cancelled`). The convention -- stated in the guidance and noted in the tool's reply -- is to keep exactly one step `in_progress`.
+
+There is no `get_plan` tool: the current plan is already in the model's context via the tail reminder every turn.
+
+## Why whole-plan replacement
+
+Addressing steps by mutable integer index (insert/remove/reorder) is error-prone for both the code (index bookkeeping) and the model (indices it just saw can go stale within a turn). Restating the whole plan each call removes that: there are no indices to track, and a later call can't corrupt partial state. For short plans the token cost is negligible.
+
+## Caching guarantee
+
+The plan is never injected into the system prompt or instructions. Static usage guidance goes there (cache-stable); only the mutable plan rides the ephemeral tail reminder. Across turns:
+
+- the durable history grows append-only and is replayed byte-identically, so the whole prefix is a cache hit;
+- the reminder and its `CachePoint` live only in the per-request copy, so they can't invalidate anything and aren't persisted.
+
+`CachePoint` is supported on Anthropic and Amazon Bedrock; on providers without prompt caching it's simply ignored (nothing to bust).
+
+## Configuration
+
+```python
+Planning(
+    guidance=None,      # static system-prompt guidance; None = default, '' = omit
+    cache_ttl='5m',     # TTL for the cache breakpoint before the reminder ('5m' | '1h')
+)
+```
+
+## Observing the plan
+
+Plan state is per-run (a fresh, isolated plan each run via `for_run`), so it
+doesn't live on the `Planning()` instance you construct. To see the final
+plan, read the most recent `write_plan` tool return from the run's messages --
+its content is the rendered plan:
+
+```python
+from pydantic_ai.messages import ToolReturnPart
+
+result = agent.run_sync('...')
+plans = [
+    part.content
+    for message in result.all_messages()
+    for part in message.parts
+    if isinstance(part, ToolReturnPart) and part.tool_name == 'write_plan'
+]
+latest_plan = plans[-1] if plans else None
+```
+
+## Agent spec (YAML/JSON)
+
+`Planning` works with Pydantic AI's [agent spec](https://ai.pydantic.dev/agent-spec/):
+
+```yaml
+# agent.yaml
+model: anthropic:claude-sonnet-4-6
+capabilities:
+  - Planning: {}
+```
+
+```python
+from pydantic_ai import Agent
+from pydantic_ai_harness.experimental.planning import Planning
+
+agent = Agent.from_file('agent.yaml', custom_capability_types=[Planning])
+```
+
+## Further reading
+
+- [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/)
+- [Hooks](https://ai.pydantic.dev/hooks/) -- `wrap_model_request` is the ephemeral injection point used here
+- [Anthropic prompt caching](https://docs.claude.com/en/docs/build-with-claude/prompt-caching)
@@ -0,0 +1,9 @@
+"""Planning capability: model-owned, cache-friendly task planning for agents."""
+
+from pydantic_ai_harness.experimental._warn import warn_experimental
+from pydantic_ai_harness.experimental.planning._capability import Planning
+from pydantic_ai_harness.experimental.planning._toolset import PlanItem, PlanningToolset, TaskStatus
+
+warn_experimental('planning')
+
+__all__ = ['PlanItem', 'Planning', 'PlanningToolset', 'TaskStatus']
@@ -0,0 +1,106 @@
+"""Planning capability: model-owned task plans surfaced without busting the cache."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field, replace
+from typing import TYPE_CHECKING, Literal
+
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.messages import CachePoint, ModelRequest, ModelResponse, UserPromptPart
+from pydantic_ai.tools import AgentDepsT, RunContext
+from pydantic_ai.toolsets import AgentToolset
+
+from pydantic_ai_harness.experimental.planning._toolset import PlanningToolset, PlanState, render_plan
+
+if TYPE_CHECKING:
+    from pydantic_ai._instructions import AgentInstructions
+    from pydantic_ai.capabilities.abstract import WrapModelRequestHandler
+    from pydantic_ai.models import ModelRequestContext
+
+
+_DEFAULT_GUIDANCE = (
+    'You have a planning tool, `write_plan`. For multi-step work, call it first to lay out the '
+    'steps, then call it again to update statuses as you start and finish each step. Pass the '
+    'full plan every time and keep exactly one step `in_progress`.'
+)
+
+
+@dataclass
+class Planning(AbstractCapability[AgentDepsT]):
+    """Structured task planning that never invalidates the prompt cache.
+
+    The plan is owned by the model through a single `write_plan` tool. The
+    current plan is surfaced back as an *ephemeral* reminder appended to the tail
+    of each request (after the latest message), with a cache breakpoint placed
+    in front of it. Because the reminder always sits after the breakpoint and is
+    never written to the durable message history, the cached prefix stays
+    byte-identical across turns -- only the small reminder is re-read each turn.
+
+    Static usage guidance goes into the system prompt via `get_instructions`,
+    which is cache-stable; the mutable plan is *never* injected there.
+
+    ```python
+    from pydantic_ai import Agent
+    from pydantic_ai_harness.experimental.planning import Planning
+
+    agent = Agent('anthropic:claude-sonnet-4-6', capabilities=[Planning()])
+    ```
+    """
+
+    guidance: str | None = None
+    """Static planning guidance for the system prompt. Cache-stable (identical every
+    request). Leave as `None` for the default, or set `''` to omit guidance entirely."""
+
+    cache_ttl: Literal['5m', '1h'] = '5m'
+    """TTL for the cache breakpoint placed before the plan reminder."""
+
+    _state: PlanState = field(default_factory=PlanState, init=False, repr=False, compare=False)
+
+    async def for_run(self, ctx: RunContext[AgentDepsT]) -> Planning[AgentDepsT]:
+        """Return a fresh per-run instance with isolated plan state (config preserved)."""
+        return replace(self)
+
+    def get_instructions(self) -> AgentInstructions[AgentDepsT] | None:
+        """Static, cache-stable guidance on using the planning tool."""
+        guidance = _DEFAULT_GUIDANCE if self.guidance is None else self.guidance
+        return guidance or None
+
+    def get_toolset(self) -> AgentToolset[AgentDepsT] | None:
+        """Toolset providing `write_plan` over this run's plan state."""
+        return PlanningToolset[AgentDepsT](self._state)
+
+    async def wrap_model_request(
+        self,
+        ctx: RunContext[AgentDepsT],
+        *,
+        request_context: ModelRequestContext,
+        handler: WrapModelRequestHandler,
+    ) -> ModelResponse:
+        """Append the current plan as an ephemeral tail reminder behind a cache breakpoint.
+
+        This runs *after* core has persisted the durable history, and the
+        per-request message list it mutates is never written back. So the
+        reminder and its `CachePoint` reach the model but never enter
+        `ctx.state.message_history` -- the cached prefix stays byte-identical
+        across turns and no stale reminders accumulate. The `CachePoint` sits
+        before the reminder text, so the reminder falls outside the cached
+        region and cannot invalidate it.
+        """
+        items = self._state.items
+        if not items:
+            return await handler(request_context)
+        messages = request_context.messages
+        last = messages[-1]
+        if isinstance(last, ModelRequest):
+            reminder = UserPromptPart(content=[CachePoint(ttl=self.cache_ttl), _reminder_text(render_plan(items))])
+            messages[-1] = replace(last, parts=[*last.parts, reminder])
+        return await handler(request_context)
+
+    @classmethod
+    def get_serialization_name(cls) -> str | None:
+        """Serialization name for agent-spec support."""
+        return 'Planning'
+
+
+def _reminder_text(plan: str) -> str:
+    return f'<plan-reminder>\nYour current plan (keep it updated with `write_plan`):\n\n{plan}\n</plan-reminder>'
@@ -0,0 +1,76 @@
+"""Planning toolset: a single `write_plan` tool over a shared, per-run plan."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+from pydantic import BaseModel, Field
+from pydantic_ai.tools import AgentDepsT
+from pydantic_ai.toolsets import FunctionToolset
+
+
+class TaskStatus(str, Enum):
+    """Lifecycle status of a single plan step."""
+
+    pending = 'pending'
+    in_progress = 'in_progress'
+    completed = 'completed'
+    cancelled = 'cancelled'
+
+
+_STATUS_ICONS = {
+    TaskStatus.pending: '[ ]',
+    TaskStatus.in_progress: '[~]',
+    TaskStatus.completed: '[x]',
+    TaskStatus.cancelled: '[-]',
+}
+
+
+class PlanItem(BaseModel):
+    """A single step in the plan."""
+
+    content: str = Field(description='Imperative description of the step, e.g. "Add the database migration".')
+    status: TaskStatus = Field(default=TaskStatus.pending, description='Current status of this step.')
+
+
+@dataclass
+class PlanState:
+    """Mutable per-run plan storage shared between the toolset and the capability hook."""
+
+    items: list[PlanItem] = field(default_factory=list[PlanItem])
+
+
+def render_plan(items: list[PlanItem]) -> str:
+    """Render the plan as a checklist with a one-line progress summary."""
+    if not items:
+        return 'No plan yet.'
+    lines = [f'{i + 1}. {_STATUS_ICONS[item.status]} {item.content}' for i, item in enumerate(items)]
+    completed = sum(1 for item in items if item.status is TaskStatus.completed)
+    lines.append(f'({completed}/{len(items)} completed)')
+    return '\n'.join(lines)
+
+
+class PlanningToolset(FunctionToolset[AgentDepsT]):
+    """Exposes a single `write_plan` tool that overwrites the shared `PlanState`."""
+
+    def __init__(self, state: PlanState) -> None:
+        super().__init__()
+        self._state = state
+        self.add_function(self.write_plan, name='write_plan')
+
+    async def write_plan(self, items: list[PlanItem]) -> str:
+        """Create or replace the full task plan.
+
+        Pass the entire ordered plan every time -- including steps that are
+        unchanged, completed, or cancelled. Keep exactly one step `in_progress`.
+        Call this when you start and when you finish a step so your progress
+        stays visible.
+
+        Args:
+            items: The complete ordered list of plan steps.
+        """
+        self._state.items = list(items)
+        in_progress = sum(1 for item in items if item.status is TaskStatus.in_progress)
+        note = '' if in_progress <= 1 else '\n\nNote: keep only one step in_progress at a time.'
+        return f'Plan updated: {len(items)} step(s).\n\n{render_plan(self._state.items)}{note}'
@@ -0,0 +1,278 @@
+"""Tests for the Planning capability."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+from pydantic_ai import Agent
+from pydantic_ai.messages import (
+    CachePoint,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    TextPart,
+    ToolCallPart,
+    UserPromptPart,
+)
+from pydantic_ai.models import ModelRequestContext, ModelRequestParameters
+from pydantic_ai.models.function import AgentInfo, FunctionModel
+from pydantic_ai.models.test import TestModel
+
+from pydantic_ai_harness.experimental.planning import PlanItem, Planning, PlanningToolset, TaskStatus
+from pydantic_ai_harness.experimental.planning._toolset import PlanState, render_plan
+
+pytestmark = pytest.mark.anyio
+
+
+@pytest.fixture
+def anyio_backend() -> str:
+    """Run async tests on the asyncio backend (matching upstream pydantic-ai)."""
+    return 'asyncio'
+
+
+def _make_request_context(messages: list[ModelMessage]) -> ModelRequestContext:
+    return ModelRequestContext(
+        model=TestModel(),
+        messages=messages,
+        model_settings=None,
+        model_request_parameters=ModelRequestParameters(),
+    )
+
+
+def _all_text(messages: list[ModelMessage]) -> str:
+    """Flatten user-prompt and assistant text fragments across all messages."""
+    out: list[str] = []
+    for msg in messages:
+        for part in msg.parts:
+            if isinstance(part, UserPromptPart):
+                if isinstance(part.content, str):
+                    out.append(part.content)
+                else:
+                    out.extend(c for c in part.content if isinstance(c, str))
+            elif isinstance(part, TextPart):
+                out.append(part.content)
+    return '\n'.join(out)
+
+
+class TestTaskStatus:
+    def test_values(self) -> None:
+        assert TaskStatus.pending == 'pending'
+        assert TaskStatus.in_progress == 'in_progress'
+        assert TaskStatus.completed == 'completed'
+        assert TaskStatus.cancelled == 'cancelled'
+
+
+class TestPlanItem:
+    def test_default_status(self) -> None:
+        item = PlanItem(content='Do something')
+        assert item.status is TaskStatus.pending
+
+
+class TestRenderPlan:
+    def test_empty(self) -> None:
+        assert render_plan([]) == 'No plan yet.'
+
+    def test_checkboxes_and_progress(self) -> None:
+        result = render_plan(
+            [
+                PlanItem(content='First', status=TaskStatus.completed),
+                PlanItem(content='Second', status=TaskStatus.in_progress),
+                PlanItem(content='Third'),
+                PlanItem(content='Fourth', status=TaskStatus.cancelled),
+            ]
+        )
+        assert result == ('1. [x] First\n2. [~] Second\n3. [ ] Third\n4. [-] Fourth\n(1/4 completed)')
+
+
+class TestPlanningToolset:
+    async def test_write_plan_replaces_state(self) -> None:
+        state = PlanState(items=[PlanItem(content='old', status=TaskStatus.completed)])
+        toolset = PlanningToolset(state)
+        result = await toolset.write_plan([PlanItem(content='A'), PlanItem(content='B', status=TaskStatus.in_progress)])
+        assert [i.content for i in state.items] == ['A', 'B']
+        assert 'Plan updated: 2 step(s).' in result
+        assert '2. [~] B' in result
+
+    async def test_write_plan_warns_on_multiple_in_progress(self) -> None:
+        toolset = PlanningToolset(PlanState())
+        result = await toolset.write_plan(
+            [
+                PlanItem(content='A', status=TaskStatus.in_progress),
+                PlanItem(content='B', status=TaskStatus.in_progress),
+            ]
+        )
+        # Exact tail so a reworded/wrapped note can't slip through.
+        assert result.endswith('\n\nNote: keep only one step in_progress at a time.')
+
+    async def test_write_plan_single_in_progress_no_warning(self) -> None:
+        toolset = PlanningToolset(PlanState())
+        result = await toolset.write_plan([PlanItem(content='A', status=TaskStatus.in_progress)])
+        # No note is appended -- the reply ends exactly with the rendered plan.
+        assert result == 'Plan updated: 1 step(s).\n\n1. [~] A\n(0/1 completed)'
+
+
+class TestPlanningCapability:
+    def test_serialization_name(self) -> None:
+        assert Planning.get_serialization_name() == 'Planning'
+
+    def test_default_instructions(self) -> None:
+        assert Planning[None]().get_instructions() == Planning[None]().get_instructions()
+        instructions = Planning[None]().get_instructions()
+        assert isinstance(instructions, str)
+        assert 'write_plan' in instructions
+
+    def test_custom_guidance(self) -> None:
+        assert Planning[None](guidance='Custom guidance.').get_instructions() == 'Custom guidance.'
+
+    def test_empty_guidance_omitted(self) -> None:
+        assert Planning[None](guidance='').get_instructions() is None
+
+    def test_get_toolset_type(self) -> None:
+        assert isinstance(Planning[None]().get_toolset(), PlanningToolset)
+
+    async def test_for_run_isolates_state_and_preserves_config(self) -> None:
+        cap = Planning[None](guidance='G', cache_ttl='1h')
+        cap._state.items.append(PlanItem(content='leftover'))
+
+        run = await cap.for_run(MagicMock())
+
+        assert run is not cap
+        assert run._state.items == []
+        assert run.guidance == 'G'
+        assert run.cache_ttl == '1h'
+        assert len(cap._state.items) == 1  # original untouched
+
+    async def test_two_runs_are_independent(self) -> None:
+        cap = Planning[None]()
+        run1 = await cap.for_run(MagicMock())
+        run2 = await cap.for_run(MagicMock())
+        run1._state.items.append(PlanItem(content='only run1'))
+        assert run2._state.items == []
+
+
+class TestEphemeralReminder:
+    async def _run_hook(
+        self, cap: Planning[None], messages: list[ModelMessage]
+    ) -> tuple[list[ModelMessage], ModelResponse]:
+        """Invoke `wrap_model_request` with a recording handler.
+
+        Returns the messages the handler was actually given (i.e. what would be
+        sent to the model) and the handler's response.
+        """
+        captured: dict[str, list[ModelMessage]] = {}
+
+        async def handler(rc: ModelRequestContext) -> ModelResponse:
+            captured['messages'] = list(rc.messages)
+            return ModelResponse(parts=[TextPart('ok')])
+
+        ctx = _make_request_context(messages)
+        response = await cap.wrap_model_request(MagicMock(), request_context=ctx, handler=handler)
+        return captured['messages'], response
+
+    async def test_no_reminder_when_plan_empty(self) -> None:
+        cap = Planning[None]()
+        original = ModelRequest(parts=[UserPromptPart('hello')])
+        seen, response = await self._run_hook(cap, [original])
+        assert seen[-1] is original
+        assert len(original.parts) == 1
+        assert isinstance(response.parts[0], TextPart)
+        assert response.parts[0].content == 'ok'  # handler was still called
+
+    async def test_reminder_appended_behind_cachepoint(self) -> None:
+        cap = Planning[None]()
+        cap._state.items = [PlanItem(content='Do X', status=TaskStatus.in_progress)]
+        original = ModelRequest(parts=[UserPromptPart('hello')])
+
+        seen, _ = await self._run_hook(cap, [original])
+
+        # Append-only: the original object is never mutated in place.
+        assert len(original.parts) == 1
+        last = seen[-1]
+        assert isinstance(last, ModelRequest)
+        assert last is not original
+        assert len(last.parts) == 2
+        reminder = last.parts[-1]
+        assert isinstance(reminder, UserPromptPart)
+        content = reminder.content
+        assert isinstance(content, list)
+        # The cache breakpoint precedes the reminder text, so the reminder
+        # falls outside the cached prefix.
+        assert isinstance(content[0], CachePoint)
+        assert content[0].ttl == '5m'
+        assert isinstance(content[1], str)
+        assert '<plan-reminder>' in content[1]
+        assert 'Do X' in content[1]
+
+    async def test_cache_ttl_is_forwarded(self) -> None:
+        cap = Planning[None](cache_ttl='1h')
+        cap._state.items = [PlanItem(content='Do X')]
+        seen, _ = await self._run_hook(cap, [ModelRequest(parts=[UserPromptPart('hi')])])
+        reminder = seen[-1].parts[-1]
+        assert isinstance(reminder, UserPromptPart)
+        assert isinstance(reminder.content, list)
+        cache_point = reminder.content[0]
+        assert isinstance(cache_point, CachePoint)
+        assert cache_point.ttl == '1h'
+
+    async def test_no_injection_when_last_is_not_model_request(self) -> None:
+        # Defensive: core guarantees a trailing ModelRequest, but if the last
+        # message isn't one we leave it untouched and still call the handler.
+        cap = Planning[None]()
+        cap._state.items = [PlanItem(content='Do X')]
+        prior = ModelResponse(parts=[TextPart('prior')])
+        seen, response = await self._run_hook(cap, [prior])
+        assert seen[-1] is prior
+        assert len(prior.parts) == 1
+        assert isinstance(response.parts[0], TextPart)
+        assert response.parts[0].content == 'ok'
+
+
+class TestEndToEnd:
+    async def test_write_plan_runs_and_plan_is_visible(self) -> None:
+        agent = Agent(TestModel(), capabilities=[Planning()])
+        # TestModel calls every tool once, including write_plan.
+        result = await agent.run('plan and do the work')
+        assert result.output is not None
+
+    async def test_reminder_reaches_model_then_is_ephemeral(self) -> None:
+        captured: dict[str, list[ModelMessage]] = {}
+        calls = 0
+
+        def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+            nonlocal calls
+            calls += 1
+            if calls == 1:
+                return ModelResponse(
+                    parts=[
+                        ToolCallPart(
+                            'write_plan',
+                            {'items': [{'content': 'Step A', 'status': 'in_progress'}]},
+                            tool_call_id='c1',
+                        )
+                    ]
+                )
+            captured['messages'] = messages
+            return ModelResponse(parts=[TextPart('done')])
+
+        agent: Agent[None, str] = Agent(FunctionModel(model_fn), capabilities=[Planning()])
+        result = await agent.run('go')
+        assert result.output == 'done'
+
+        # The plan reminder reached the model on the second request...
+        sent = _all_text(captured['messages'])
+        assert '<plan-reminder>' in sent
+        assert 'Step A' in sent
+        # ...with a CachePoint marking the boundary before it.
+        has_cache_point = any(
+            isinstance(part, UserPromptPart)
+            and not isinstance(part.content, str)
+            and any(isinstance(c, CachePoint) for c in part.content)
+            for msg in captured['messages']
+            if isinstance(msg, ModelRequest)
+            for part in msg.parts
+        )
+        assert has_cache_point
+        # ...but neither the reminder nor its CachePoint is ever written to the
+        # durable message history (the CachePoint only ever rides the reminder).
+        assert '<plan-reminder>' not in _all_text(result.all_messages())