---
phase: 03-mcp-surface-generation
plan: 02
type: execute
wave: 2
depends_on: [03-01]
files_modified:
  - src/ztlctl/mcp/generator.py
  - tests/mcp/test_generator.py
  - tests/mcp/test_parity.py
autonomous: true
requirements: [AGNT-02, PLUG-04]

must_haves:
  truths:
    - "list_items, search, vault_review, and decision_support MCP tools accept a token_budget parameter"
    - "When token_budget is set, response data with large lists is truncated to fit the budget"
    - "When token_budget is None (default), response is returned untruncated"
    - "Every ActionDefinition in the registry has a corresponding MCP tool (parity test proves 59/59)"
    - "Previously missing tools (archive, supersede, upgrade actions) are all present in the generated MCP surface"
  artifacts:
    - path: "tests/mcp/test_parity.py"
      provides: "Parity verification: every ActionDefinition has a matching MCP tool"
      contains: "def test_all_actions_have_mcp_tools"
    - path: "src/ztlctl/mcp/generator.py"
      provides: "Token budget support via _apply_token_budget and BUDGET_AWARE_ACTIONS"
      contains: "BUDGET_AWARE_ACTIONS"
    - path: "tests/mcp/test_generator.py"
      provides: "Token budget tests"
      contains: "def test_token_budget"
  key_links:
    - from: "src/ztlctl/mcp/generator.py"
      to: "ActionDefinition params"
      via: "token_budget injected into __annotations__ and __kwdefaults__ for budget-aware tools"
      pattern: "BUDGET_AWARE_ACTIONS"
    - from: "tests/mcp/test_parity.py"
      to: "src/ztlctl/actions/registry.py"
      via: "get_action_registry().list_actions() compared against DummyServer registered tools"
      pattern: "get_action_registry.*list_actions"
---

<objective>
Add token-budget-aware truncation for high-volume MCP tools and create a parity test suite proving every ActionDefinition is exposed as an MCP tool.

Purpose: AGNT-02 (token-budget responses for agents) and PLUG-04 (complete CLI/MCP parity verified by construction).

Output: Token budget support in generator.py, parity test in test_parity.py, budget tests in test_generator.py.
</objective>

<execution_context>
@/Users/shparki/.claude/get-shit-done/workflows/execute-plan.md
@/Users/shparki/.claude/get-shit-done/templates/summary.md
</execution_context>

<context>
@.planning/PROJECT.md
@.planning/ROADMAP.md
@.planning/STATE.md
@.planning/phases/03-mcp-surface-generation/03-CONTEXT.md
@.planning/phases/03-mcp-surface-generation/03-RESEARCH.md
@.planning/phases/03-mcp-surface-generation/03-01-SUMMARY.md

<interfaces>
<!-- From Plan 01 outputs -->

From src/ztlctl/mcp/generator.py (created in Plan 01):
```python
_vault_ref: Any = None

def set_vault(vault: Any) -> None: ...
def generate_tools(server: Any, vault: Any) -> None: ...
def _make_tool_fn(action: ActionDefinition, vault: Any) -> Callable[..., dict[str, Any]]: ...
def _build_annotations(params: tuple[ActionParam, ...]) -> dict[str, Any]: ...
def _build_defaults(params: tuple[ActionParam, ...]) -> dict[str, Any]: ...
def _render_action_doc(action: ActionDefinition) -> str: ...
def _register_plugin_tools(server: Any, vault: Any) -> None: ...
```

From src/ztlctl/mcp/response.py (created in Plan 01):
```python
COMMON_ERROR_RECOVERY: dict[str, str] = { ... }

class McpError(BaseModel): ...
class McpResponse(BaseModel):
    ok: bool
    op: str
    data: dict[str, Any] = ...
    warnings: list[str] = ...
    error: McpError | None = None

    @classmethod
    def from_result(cls, result: ServiceResult) -> "McpResponse": ...
```

From src/ztlctl/actions/registry.py:
```python
def get_action_registry() -> ActionRegistry: ...
# ActionRegistry.list_actions() returns all 59 ActionDefinitions
```
</interfaces>
</context>

<tasks>

<task type="auto" tdd="true">
  <name>Task 1: Add token-budget truncation for high-volume MCP tools</name>
  <files>
    src/ztlctl/mcp/generator.py,
    tests/mcp/test_generator.py
  </files>
  <read_first>
    src/ztlctl/mcp/generator.py,
    tests/mcp/test_generator.py,
    src/ztlctl/actions/_register_core.py
  </read_first>
  <behavior>
    - _apply_token_budget(data, None) returns data unchanged
    - _apply_token_budget({"items": [1,2,3]}, budget=10) returns truncated data with "truncated": True when items exceed budget
    - _apply_token_budget({"items": [1]}, budget=10000) returns data unchanged (within budget)
    - _apply_token_budget({"count": 5}, budget=10) returns data unchanged (no list field to truncate)
    - Budget-aware tool "list_items" has "token_budget" in __annotations__ with type int | None
    - Budget-aware tool "list_items" has "token_budget" in __kwdefaults__ with default None
    - Non-budget-aware tool "create_note" does NOT have "token_budget" in __annotations__
    - Calling a budget-aware tool with token_budget=50 produces truncated output
    - Calling a budget-aware tool with token_budget=None produces full output
  </behavior>
  <action>
    **Add to `src/ztlctl/mcp/generator.py`:**

    1. Define `BUDGET_AWARE_ACTIONS: frozenset[str]` containing exactly: `"list_items"`, `"search"`, `"vault_review"`, `"decision_support"`.

    2. Create `_apply_token_budget(data: dict[str, Any], budget: int | None) -> dict[str, Any]`:
       - If `budget is None`, return `data` unchanged.
       - Serialize `data` to JSON via `json.dumps(data)`.
       - Estimate tokens as `len(serialized) // 4`.
       - If estimated tokens <= budget, return `data` unchanged.
       - Find the first list-valued field in `data.items()` that has elements.
       - Iteratively remove the last element from that list (creating new dicts, since we shouldn't mutate) until the estimated token count fits within budget or the list is empty.
       - Add `"truncated": True` and `"token_budget": budget` to the returned dict.
       - `import json` at module top level (stdlib, fine for module-level).

    3. Modify `_make_tool_fn()` or create `_make_budget_aware_tool_fn()`:
       - In `generate_tools()`, check `if action.name in BUDGET_AWARE_ACTIONS`.
       - If so, use a variant that:
         a. Adds `"token_budget": int | None` to the annotations dict.
         b. Adds `"token_budget": None` to the kwdefaults dict.
         c. In the inner `tool_fn`, pops `token_budget` from kwargs before calling handler, then wraps result: `response = McpResponse.from_result(result).model_dump(exclude_none=True)`, then `response["data"] = _apply_token_budget(response.get("data", {}), token_budget)`, then returns response.
       - The cleanest approach: add budget support directly in `_make_tool_fn` conditionally. Check `is_budget_aware = action.name in BUDGET_AWARE_ACTIONS`. If true, inject the extra param into annotations/defaults and handle in the wrapper.

    **Add tests to `tests/mcp/test_generator.py`:**

    - `test_apply_token_budget_none`: `_apply_token_budget({"items": list(range(100))}, None)` returns data unchanged.
    - `test_apply_token_budget_within`: Small data + large budget returns unchanged.
    - `test_apply_token_budget_truncates`: Large data + small budget returns truncated data with `"truncated": True`.
    - `test_apply_token_budget_no_list`: Data with no list field returns unchanged regardless of budget.
    - `test_budget_aware_tool_has_token_budget_param`: Register tools, check "list_items" has `token_budget` in `__annotations__`.
    - `test_non_budget_tool_no_token_budget_param`: Check "create_note" does NOT have `token_budget`.
    - `test_budget_aware_actions_set`: Verify `BUDGET_AWARE_ACTIONS == {"list_items", "search", "vault_review", "decision_support"}`.
  </action>
  <verify>
    <automated>cd /Users/shparki/Documents/Workspace/thatdev/ztlctl && uv run pytest tests/mcp/test_generator.py -x -q -k "budget"</automated>
  </verify>
  <acceptance_criteria>
    - src/ztlctl/mcp/generator.py contains `BUDGET_AWARE_ACTIONS = frozenset(` with "list_items", "search", "vault_review", "decision_support"
    - src/ztlctl/mcp/generator.py contains `def _apply_token_budget(data: dict`
    - src/ztlctl/mcp/generator.py contains `"token_budget"` string (injected into annotations)
    - tests/mcp/test_generator.py contains `def test_apply_token_budget_none` and `def test_apply_token_budget_truncates` and `def test_budget_aware_tool_has_token_budget_param`
    - `uv run pytest tests/mcp/test_generator.py -x -q -k "budget"` exits 0
  </acceptance_criteria>
  <done>Four high-volume MCP tools accept token_budget parameter. When set, large list responses are truncated. Non-budget tools are unaffected. Tests prove truncation behavior.</done>
</task>

<task type="auto" tdd="true">
  <name>Task 2: Create parity test suite verifying all ActionDefinitions have MCP tools</name>
  <files>
    tests/mcp/test_parity.py
  </files>
  <read_first>
    src/ztlctl/mcp/generator.py,
    src/ztlctl/actions/registry.py,
    src/ztlctl/actions/_register_core.py,
    tests/mcp/test_generator.py
  </read_first>
  <behavior>
    - Every action name from get_action_registry().list_actions() has a matching key in DummyServer.tools after generate_tools()
    - Tool count equals registry action count (59)
    - archive, supersede, upgrade_apply, upgrade_check_pending, check_integrity, init_vault, init_workflow, update_workflow are all present (previously missing from hand-written MCP surface)
    - Every registered tool function has a non-empty __doc__
    - Every registered tool function has __annotations__ dict with at least "return" key
    - Category coverage: all 13 categories from the registry produce at least one MCP tool
  </behavior>
  <action>
    **Create `tests/mcp/test_parity.py`:**

    Use the same DummyServer pattern from test_generator.py. Import `generate_tools` from `ztlctl.mcp.generator` and `get_action_registry` from `ztlctl.actions`.

    1. Fixture `parity_server()`:
       - Create DummyServer instance.
       - Create a mock vault (MagicMock with `plugin_manager=None` so plugin shim is skipped).
       - Call `generate_tools(server, vault)`.
       - Return (server, registry).

    2. `test_all_actions_have_mcp_tools(parity_server)`:
       - `registry_names = {a.name for a in get_action_registry().list_actions()}`.
       - `tool_names = set(parity_server[0].tools.keys())`.
       - Assert `registry_names <= tool_names` (every registry action is a registered tool).
       - Assert `len(registry_names) >= 59`.

    3. `test_tool_count_matches_registry(parity_server)`:
       - Assert `len(parity_server[0].tools) == len(get_action_registry().list_actions())`.

    4. `test_previously_missing_tools_present(parity_server)`:
       - Define `PREVIOUSLY_MISSING = {"archive", "supersede", "upgrade_apply", "upgrade_check_pending", "upgrade_stamp_current", "check_integrity", "init_vault", "init_workflow", "update_workflow"}`.
       - Assert each name is in `parity_server[0].tools`.

    5. `test_every_tool_has_doc(parity_server)`:
       - For each tool fn in `parity_server[0].tools.values()`, assert `fn.__doc__` is not None and len > 0.

    6. `test_every_tool_has_annotations(parity_server)`:
       - For each tool fn, assert `"return"` in `fn.__annotations__`.

    7. `test_category_coverage(parity_server)`:
       - `categories = {a.category for a in get_action_registry().list_actions()}`.
       - Assert `len(categories) == 13`.
       - For each category, assert at least one action from that category has a tool registered.

    **Run full validation after this task:**
    ```bash
    uv run ruff check . && uv run ruff format --check . && uv run pytest -x -q && uv run mypy src/
    ```
  </action>
  <verify>
    <automated>cd /Users/shparki/Documents/Workspace/thatdev/ztlctl && uv run pytest tests/mcp/test_parity.py -x -q && uv run pytest -x -q && uv run mypy src/</automated>
  </verify>
  <acceptance_criteria>
    - tests/mcp/test_parity.py exists and contains `def test_all_actions_have_mcp_tools` and `def test_tool_count_matches_registry` and `def test_previously_missing_tools_present`
    - tests/mcp/test_parity.py contains `def test_every_tool_has_doc` and `def test_every_tool_has_annotations` and `def test_category_coverage`
    - tests/mcp/test_parity.py references `PREVIOUSLY_MISSING` set containing "archive", "supersede", "upgrade_apply"
    - `uv run pytest tests/mcp/test_parity.py -x -q` exits 0
    - `uv run pytest -x -q` exits 0 (full suite still green)
    - `uv run mypy src/` exits 0
  </acceptance_criteria>
  <done>Parity test suite proves all 59 ActionDefinitions have MCP tools. Previously missing tools confirmed present. All tools have docstrings and annotations. Full suite green.</done>
</task>

</tasks>

<verification>
After both tasks:
- `uv run pytest tests/mcp/ -x -q` — all MCP tests pass
- `uv run pytest -x -q` — full suite passes
- `uv run mypy src/` — zero type errors
- `uv run ruff check .` — zero lint errors
- Token budget: `list_items`, `search`, `vault_review`, `decision_support` all accept `token_budget` param
- Parity: 59/59 ActionDefinitions have MCP tools registered
</verification>

<success_criteria>
1. BUDGET_AWARE_ACTIONS contains exactly list_items, search, vault_review, decision_support
2. _apply_token_budget truncates large responses when budget is set, passes through when None
3. Parity test asserts 59/59 actions have MCP tools
4. Previously missing tools (archive, supersede, upgrade, check, init, workflow) confirmed present
5. Every MCP tool has __doc__, __annotations__, and correct __kwdefaults__
6. Full validation green: pytest + mypy + ruff
</success_criteria>

<output>
After completion, create `.planning/phases/03-mcp-surface-generation/03-02-SUMMARY.md`
</output>
