---
phase: 16-plugin-bridge-and-action-executor
plan: 02
type: execute
wave: 1
depends_on: []
files_modified:
  - src/ztlctl/mcp/server.py
  - src/ztlctl/commands/serve.py
  - tests/mcp/test_shutdown.py
autonomous: true
requirements: [DEBT-04]

must_haves:
  truths:
    - "When ztlctl serve exits (normal or signal), the vault is closed and EventBus drains pending WAL events"
    - "vault.close(wait_for_events=True) runs even if server.run() raises SystemExit"
    - "create_server returns both the FastMCP server and the Vault so callers can manage vault lifecycle"
  artifacts:
    - path: "src/ztlctl/mcp/server.py"
      provides: "ServerContext dataclass returned from create_server"
      contains: "ServerContext"
    - path: "src/ztlctl/commands/serve.py"
      provides: "try/finally block around server.run() calling vault.close()"
      contains: "vault.close"
    - path: "tests/mcp/test_shutdown.py"
      provides: "Tests for graceful shutdown on normal exit and SystemExit"
      contains: "test_vault_closed"
  key_links:
    - from: "src/ztlctl/commands/serve.py"
      to: "src/ztlctl/mcp/server.py"
      via: "ServerContext.server and ServerContext.vault"
      pattern: "ctx\\.vault\\.close"
    - from: "src/ztlctl/commands/serve.py"
      to: "src/ztlctl/infrastructure/vault.py"
      via: "vault.close(wait_for_events=True) drains EventBus"
      pattern: "wait_for_events=True"
---

<objective>
Implement MCP server graceful shutdown so the vault and EventBus are properly closed when `ztlctl serve` exits.

Purpose: Currently `create_server()` creates a Vault but never returns it, and `serve.py` never calls `vault.close()`. When the MCP client disconnects, the EventBus ThreadPoolExecutor and DB engine remain open, and pending WAL events are not drained. This plan fixes that by returning a ServerContext from create_server and wrapping server.run() in try/finally.

Output: ServerContext dataclass in server.py, vault cleanup in serve.py, shutdown tests.
</objective>

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

<context>
@.planning/PROJECT.md
@.planning/ROADMAP.md
@.planning/STATE.md
@.planning/phases/16-plugin-bridge-and-action-executor/16-RESEARCH.md

<interfaces>
<!-- Key types and contracts. Extracted from codebase. -->

From src/ztlctl/mcp/server.py:
```python
def create_server(
    *,
    vault_root: Path | None = None,
    host: str = "127.0.0.1",
    port: int = 8000,
) -> Any:
    """Create and configure the MCP server. Returns the FastMCP instance."""
    # Creates vault internally but does not return it
    settings = ZtlSettings.from_cli(vault_root=vault_root)
    vault = Vault(settings)
    vault.init_event_bus(sync=settings.sync)
    server = _FastMCP("ztlctl", host=host, port=port)
    generate_tools(server, vault)
    register_resources(server, vault)
    register_prompts(server, vault)
    return server
```

From src/ztlctl/commands/serve.py:
```python
server = create_server(vault_root=app.settings.vault_root, host=host, port=port)
server.run(transport=transport)  # vault never closed
```

From src/ztlctl/infrastructure/vault.py:
```python
class Vault:
    def close(self, *, wait_for_events: bool = False) -> None:
        """Close the vault, optionally waiting for event bus to drain."""
        ...
```
</interfaces>
</context>

<tasks>

<task type="auto">
  <name>Task 1: Add ServerContext and update create_server return type (DEBT-04)</name>
  <files>src/ztlctl/mcp/server.py, src/ztlctl/commands/serve.py</files>
  <read_first>src/ztlctl/mcp/server.py, src/ztlctl/commands/serve.py</read_first>
  <action>
In `src/ztlctl/mcp/server.py`:
1. Add a `ServerContext` dataclass at module level (after imports, before `create_server`):
```python
from dataclasses import dataclass

@dataclass
class ServerContext:
    """MCP server and its backing vault for lifecycle management."""
    server: Any      # FastMCP instance
    vault: Any       # Vault instance (Any to avoid import when mcp not installed)
```
2. Update `__all__` to include `ServerContext`: `__all__ = ["create_server", "mcp_available", "ServerContext"]`
3. Change `create_server()` return type from `-> Any` to `-> ServerContext`.
4. At the end of `create_server()`, instead of `return server`, do `return ServerContext(server=server, vault=vault)`.

In `src/ztlctl/commands/serve.py`:
1. Change the serve command to use ServerContext and add try/finally:
```python
ctx = create_server(vault_root=app.settings.vault_root, host=host, port=port)
try:
    ctx.server.run(transport=transport)
finally:
    ctx.vault.close(wait_for_events=True)
```
2. Update the import: `from ztlctl.mcp.server import create_server, mcp_available` (no change needed since ServerContext is not imported directly -- it comes via `ctx`).

Anti-patterns to avoid:
- Do NOT store vault on the FastMCP server object (e.g. `server._vault = vault`). Use the clean ServerContext pattern.
- Do NOT add signal handlers. FastMCP handles SIGTERM/SIGINT internally before returning from run(). The finally block runs after run() returns regardless.
- Do NOT call vault.close() inside asyncio callbacks. vault.close() is synchronous (SQLAlchemy).
  </action>
  <verify>
    <automated>uv run pytest tests/mcp/ -x -q && uv run mypy src/ztlctl/mcp/server.py src/ztlctl/commands/serve.py</automated>
  </verify>
  <acceptance_criteria>
    - grep -c "ServerContext" src/ztlctl/mcp/server.py returns at least 2
    - grep -c "vault.close" src/ztlctl/commands/serve.py returns at least 1
    - grep -c "wait_for_events=True" src/ztlctl/commands/serve.py returns at least 1
    - grep -c "finally" src/ztlctl/commands/serve.py returns at least 1
    - uv run mypy src/ztlctl/mcp/server.py src/ztlctl/commands/serve.py passes
  </acceptance_criteria>
  <done>create_server returns ServerContext with both server and vault. serve.py wraps server.run() in try/finally that calls vault.close(wait_for_events=True).</done>
</task>

<task type="auto">
  <name>Task 2: Add shutdown tests for MCP vault cleanup (DEBT-04)</name>
  <files>tests/mcp/test_shutdown.py</files>
  <read_first>src/ztlctl/mcp/server.py, src/ztlctl/commands/serve.py, tests/mcp/</read_first>
  <action>
Create new file `tests/mcp/test_shutdown.py` with tests for the graceful shutdown behavior:

1. `test_vault_closed_after_server_run_returns`: Mock `create_server` to return a ServerContext with a mock server (whose `run()` returns normally) and a mock vault. Invoke the serve Click command (or call the function directly). Assert `mock_vault.close` called once with `wait_for_events=True`.

2. `test_vault_closed_on_system_exit`: Same as above but mock server's `run()` to raise `SystemExit(0)`. Assert `mock_vault.close` still called (Python's try/finally handles SystemExit). Use `pytest.raises(SystemExit)` around the invocation.

3. `test_vault_closed_on_runtime_error`: Mock server's `run()` to raise `RuntimeError("connection lost")`. Assert `mock_vault.close` still called. Use `pytest.raises(RuntimeError)`.

4. `test_server_context_has_server_and_vault`: Import `ServerContext` from `ztlctl.mcp.server`, verify it is a dataclass with `server` and `vault` fields.

Implementation notes:
- Use `unittest.mock.patch` to mock `create_server` in `ztlctl.commands.serve` (the import target).
- For Click command testing, use `click.testing.CliRunner` or invoke the function directly.
- Since the `serve` command has `@click.pass_obj`, the simplest approach is to directly test the vault close behavior by extracting the logic, OR by mocking at the `create_server` level and using CliRunner with an AppContext mock.
- Alternative: test the ServerContext + try/finally pattern directly without going through Click, since the Click wiring is trivial.
  </action>
  <verify>
    <automated>uv run pytest tests/mcp/test_shutdown.py -x -q</automated>
  </verify>
  <acceptance_criteria>
    - test file tests/mcp/test_shutdown.py exists
    - grep -c "test_vault_closed" tests/mcp/test_shutdown.py returns at least 2
    - grep -c "wait_for_events" tests/mcp/test_shutdown.py returns at least 1
    - grep -c "SystemExit" tests/mcp/test_shutdown.py returns at least 1
    - uv run pytest tests/mcp/test_shutdown.py -x -q passes
  </acceptance_criteria>
  <done>Shutdown tests verify vault.close(wait_for_events=True) runs on normal exit, SystemExit, and RuntimeError. ServerContext structure verified.</done>
</task>

</tasks>

<verification>
After both tasks:
1. `uv run pytest tests/mcp/ -x -q` -- all pass
2. `uv run ruff check src/ztlctl/mcp/server.py src/ztlctl/commands/serve.py` -- clean
3. `uv run mypy src/ztlctl/mcp/server.py src/ztlctl/commands/serve.py` -- clean
</verification>

<success_criteria>
- create_server returns ServerContext (not bare FastMCP instance)
- serve command wraps server.run() in try/finally with vault.close()
- Tests prove vault cleanup happens on normal exit, SystemExit, and exceptions
- DEBT-04 complete
</success_criteria>

<output>
After completion, create `.planning/phases/16-plugin-bridge-and-action-executor/16-02-SUMMARY.md`
</output>
