fix(core,api): harden cancellation and tool telemetry

.env.example

@@ -52,7 +52,8 @@ LLM_RETRY_BASE_DELAY_S=0.5
 LLM_RETRY_MAX_DELAY_S=10
 
 # External Tools & Sandboxing
-SEARCH_API_KEY=your_search_api_key_here
+# Optional: web_search is omitted from the tool registry when this is unset.
+SEARCH_API_KEY=
 FILE_TOOL_BASE_DIR=/tmp/agentkit_sandbox
 
 # Security & API Authentication

README.md

@@ -24,13 +24,14 @@ Designed from first principles for reliability, auditability, and production dep
   - `sql_readonly`: Read-only SQL executor over a dedicated non-superuser role (`DATABASE_READONLY_URL`, `default_transaction_read_only`, `SELECT`-only grants, fail-closed when unconfigured). A statement lexer accepts a single `SELECT`/`WITH`, rejects comments, DDL/DML and privileged built-ins (`pg_read_file`, `setval`, `query_to_xml`, …), and clamps `LIMIT ALL`/oversized limits in SQL.
   - `read_file`: Path traversal and symlink escape defended sandboxed file reader.
   - `http_fetch`: SSRF-protected HTTP client blocking private subnets, loopbacks, and cloud metadata endpoints.
-  - `web_search`: Search provider adapter with result truncation.
+  - `web_search`: Search provider adapter with result truncation; registered only when `SEARCH_API_KEY` is configured.
 - **Dual-Store Persistence**:
   - **Redis**: Sliding-window session memory and high-throughput sliding window rate limiter.
   - **PostgreSQL**: Durable execution audit trails with step-by-step latency, token counts, and tool outcomes.
 - **Enterprise-Grade Reliability**:
   - Exponential backoff retry with full jitter for transient LLM errors.
   - Global run timeout, per-tool timeout guards, and consecutive duplicate tool call loop detection.
+  - Client disconnects never leave zombie runs: cancellation persists a terminal `canceled` status plus repaired session history inside `asyncio.shield` before re-raising.
   - Structured contextual JSON logging (`run_id`, `step_no`, duration) with zero raw confidential prompt leakage.
 - **Containerized & CI-Ready**: Multi-stage production `Dockerfile`, `docker-compose.yml` with healthchecks, and full GitHub Actions CI matrix (`ruff`, `mypy`, `pytest` >= 85% coverage).
 
@@ -287,6 +288,8 @@ View run traces programmatically via REST API:
 curl -H "X-API-Key: YOUR_API_KEY" http://localhost:8000/runs/{run_id}/trace | jq .
 ```
 
+Token attribution note: when one model turn requests several tool calls, the turn's token usage is recorded on that step's *first* successfully recorded tool-call row and sibling rows report `0` on the token columns — so `sum(steps.*) <= runs.total_*` always holds and zero-token rows mean "this row does not carry the turn's usage", not "no tokens were used" (see [docs/DECISIONS.md](docs/DECISIONS.md) ADR-021).
+
 ---
 
 ## Architectural Decisions (ADR Summary)

agentkit/api/deps.py

@@ -116,14 +116,15 @@ def get_trace_sink(request: Request) -> TraceSink:
     return PostgresTraceSink(factory)
 
 
-def create_default_tool_registry() -> ToolRegistry:
-    """Construct standard ToolRegistry populated with all AgentKit built-in tools."""
+def create_default_tool_registry(search_api_key: str | None = None) -> ToolRegistry:
+    """Construct the built-in registry, including search only when configured."""
     registry = ToolRegistry()
     registry.register(calculator)
     registry.register(read_file)
     registry.register(http_fetch)
     registry.register(sql_readonly)
-    registry.register(web_search)
+    if search_api_key:
+        registry.register(web_search)
     return registry
 
 
@@ -133,7 +134,10 @@ def get_tool_registry(request: Request) -> ToolRegistry:
         reg: ToolRegistry = request.app.state.tool_registry
         return reg
 
-    registry = create_default_tool_registry()
+    settings: Settings = getattr(request.app.state, "settings", None) or get_current_settings(
+        request
+    )
+    registry = create_default_tool_registry(settings.SEARCH_API_KEY)
     request.app.state.tool_registry = registry
     return registry
 

agentkit/core/agent.py

@@ -195,8 +195,18 @@ async def run(
                         run_id=run_id,
                     ) from exc
 
-            async def _finalize_and_return(result: RunResult) -> RunResult:
-                if session_id and self.memory is not None and hasattr(self.memory, "save_messages"):
+            async def _finalize_and_return(
+                result: RunResult, *, save_history: bool = True
+            ) -> RunResult:
+                # save_history=False covers cancellation while the history load is still in
+                # flight: at that point the local list is empty and saving it would overwrite
+                # the stored conversation with nothing instead of repairing it.
+                if (
+                    save_history
+                    and session_id
+                    and self.memory is not None
+                    and hasattr(self.memory, "save_messages")
+                ):
                     _repair_trailing_tool_calls(messages)
                     try:
                         await self.memory.save_messages(session_id, messages)
@@ -257,10 +267,60 @@ async def _finalize_and_return(result: RunResult) -> RunResult:
                         ) from exc
                 return result
 
+            async def _finalize_canceled(*, save_history: bool) -> None:
+                """Persist a terminal canceled run state before the cancellation re-raises.
+
+                create_run() has already persisted a running row, and cancellation can land
+                on any awaited step that follows it — the session-history load or an await
+                inside the reasoning loop. Without this finalization the row would be
+                stranded in 'running' even though the caller received the cancellation.
+                """
+                duration_ms = max(0, int((time.perf_counter() - start_time) * 1000))
+                canceled_result = RunResult(
+                    run_id=run_id,
+                    session_id=session_id,
+                    goal=goal,
+                    status=RunStatus.CANCELED,
+                    failure_reason="Run was canceled by the caller.",
+                    steps_count=steps_count,
+                    total_input_tokens=total_input_tokens,
+                    total_output_tokens=total_output_tokens,
+                    duration_ms=duration_ms,
+                )
+                logger.info("Agent run canceled")
+                try:
+                    # shield keeps the terminal write alive against further cancellation so
+                    # the store never reads back a run that stayed 'running' behind a cancel.
+                    await asyncio.shield(
+                        _finalize_and_return(canceled_result, save_history=save_history)
+                    )
+                except PersistenceError:
+                    logger.error("Could not persist canceled run state for run %s", run_id)
+
+            # Initialized before the history load so a cancellation landing during setup
+            # can still read zeroed counters when composing the canceled RunResult.
+            total_input_tokens = 0
+            total_output_tokens = 0
+            steps_count = 0
+            traced_turn_usage: set[int] = set()
+            # Sliding window of recent call signatures. Remembering only the immediately
+            # preceding signature lets an alternating pattern (A, B, A, B, ...) reset the
+            # counter on every step, so the guard never fires and the run burns max_steps.
+            recent_tool_signatures: deque[tuple[str, str]] = deque(
+                maxlen=self.config.duplicate_tool_call_window
+            )
+
             messages: list[Message] = []
             if session_id and self.memory is not None and hasattr(self.memory, "get_messages"):
                 try:
                     loaded = await self.memory.get_messages(session_id)
+                except asyncio.CancelledError:
+                    # The history load is the only awaited setup the CancelledError handler
+                    # below does not reach, and it runs after the run row already exists.
+                    # Finalize here so the row cannot stay 'running'; history was never read
+                    # or mutated in this window, hence save_history is deliberately False.
+                    await _finalize_canceled(save_history=False)
+                    raise
                 except Exception as exc:
                     # Continuing here would answer with no context and then overwrite the real
                     # history with the shorter one: silent, permanent conversation loss.
@@ -297,16 +357,6 @@ async def _finalize_and_return(result: RunResult) -> RunResult:
                     messages.insert(0, Message.system(self.config.system_prompt))
                 messages.append(Message.user(goal))
 
-            total_input_tokens = 0
-            total_output_tokens = 0
-            steps_count = 0
-            # Sliding window of recent call signatures. Remembering only the immediately
-            # preceding signature lets an alternating pattern (A, B, A, B, ...) reset the
-            # counter on every step, so the guard never fires and the run burns max_steps.
-            recent_tool_signatures: deque[tuple[str, str]] = deque(
-                maxlen=self.config.duplicate_tool_call_window
-            )
-
             try:
                 async with asyncio.timeout(self.config.run_timeout_s):
                     for _ in range(self.config.max_steps):
@@ -401,10 +451,19 @@ async def _finalize_and_return(result: RunResult) -> RunResult:
                                             result=result.model_dump(),
                                             error=result.error,
                                             latency_ms=result.latency_ms,
-                                            input_tokens=response.usage.input_tokens,
-                                            output_tokens=response.usage.output_tokens,
+                                            input_tokens=(
+                                                response.usage.input_tokens
+                                                if steps_count not in traced_turn_usage
+                                                else 0
+                                            ),
+                                            output_tokens=(
+                                                response.usage.output_tokens
+                                                if steps_count not in traced_turn_usage
+                                                else 0
+                                            ),
                                         )
                                         await self.trace.record(trace_record)
+                                        traced_turn_usage.add(steps_count)
                         else:
                             # No tool calls: this is the final answer, a provider refusal, or
                             # an empty completion. Only a non-empty answer is a success; the
@@ -485,6 +544,12 @@ async def _finalize_and_return(result: RunResult) -> RunResult:
                         raise MaxStepsExceeded(self.config.max_steps)
                     return res_max
 
+            except asyncio.CancelledError:
+                # Loop-time cancellation: every step up to here may have left an assistant
+                # tool-call turn without its results, so the history save (with repair) and
+                # the terminal run write must both run before the cancellation propagates.
+                await _finalize_canceled(save_history=True)
+                raise
             except PersistenceError:
                 # Already logged with its run id; re-raise rather than converting it into a
                 # FAILED result that the caller could mistake for an executed-but-failed run.

agentkit/core/types.py

@@ -10,6 +10,7 @@ class RunStatus(StrEnum):
     RUNNING = "running"
     SUCCEEDED = "succeeded"
     FAILED = "failed"
+    CANCELED = "canceled"
     #: The provider refused or filtered the response (safety block, content filter, refusal).
     REFUSED = "refused"
     MAX_STEPS_EXCEEDED = "max_steps_exceeded"

docs/DECISIONS.md

@@ -2,6 +2,14 @@
 
 This log records significant architectural, technical, and design decisions made throughout the lifecycle of AgentKit.
 
+## ADR-021: Run Cancellation and Turn-Level Trace Usage
+
+- **Date**: 2026-11-01
+- **Status**: Accepted
+- **Context**: Cancellation can land on any awaited step that follows creation of the durable `running` row — the session-history load before the loop included — and a model response that requests multiple tools has one usage count shared by all calls.
+- **Decision**: On any such cancellation, finalize the run as `canceled` through a shielded finalization before re-raising; while the history load is still pending, history is left untouched because the empty local snapshot would overwrite the stored conversation, so only the run row is written. Once the loop is running, the same finalization also persists repaired session history. Store turn-level token usage on only the first successfully recorded tool trace for that reasoning step; subsequent call traces of the same step carry zero on the token columns, keeping per-call latency and outcome. Register the optional `web_search` tool only when `SEARCH_API_KEY` is present.
+- **Consequences**: Callers still observe task cancellation, persisted run state is terminal, and `sum(steps.input_tokens/output_tokens) <= runs.total_input_tokens/total_output_tokens` for any run. Equality holds only when every LLM turn either requested tools (so its usage landed on that step's first recorded tool trace) or consumed no tokens; the last, tool-free answer turn contributes to `runs.total_*` only, never to `steps`. Trace consumers must not treat zero-token rows as "the model used no tokens": they mean "this row does not carry the turn's usage". Search availability is consistent between runtime schemas and `GET /tools`.
+
 ---
 
 ## ADR-001: Framework-less Agent Runtime from Scratch

docs/PROGRESS.md

@@ -4,6 +4,7 @@ This document tracks execution progress across all backlog tasks, including pull
 
 | Task ID | Milestone | Title | Branch | PR # | Commits | Status | Notes & Blockers |
 |---------|-----------|-------|--------|------|---------|--------|------------------|
+| ROADMAP-01/06/07 | M14 | Run cancellation, token accounting, optional search registration | `fix/related-run-tool-hardening` | [#122](https://github.com/getdownandcode/AgentKit/pull/122) | 5 | In review | Fixes #93, #98 and #99 together: canceled runs persist terminal state, turn usage is attributed once across tool traces, and web_search is registered only with SEARCH_API_KEY. Review-fix round: cancellation during the awaited history load (after create_run but before the guarded loop) also finalizes the run as canceled with history untouched (tests/unit/test_agent_cancellation.py, ADR-021), and the real-infra /tools test now pins SEARCH_API_KEY per case so it is deterministic with the key set or unset. Quality gates pass (ruff check/format, mypy agentkit tests, pytest 508 hermetic passed, 31 real-infra passed, pytest-cov 93.4% >= 85%). |
 | M0-T01  | M0        | Repository Setup & Context Scaffolding | `main` | N/A | 7 | Done | Initial warm-up repository scaffolding |
 | M0-T02  | M0        | Configuration Management | `feat/config-settings` | #1 | 2 | Done | Pydantic BaseSettings with env validation |
 | M0-T03  | M0        | Domain Errors & Core Base Types | `feat/core-errors-types` | #2 | 2 | Done | AgentKitError hierarchy, to_dict, StrEnum types |

tests/real/test_api_e2e_real.py

@@ -31,8 +31,13 @@
 API_KEY = "ak_real_infra_key"
 
 
-def _build_settings() -> Settings:
-    """Settings pointing the app at the live infrastructure under test."""
+def _build_settings(search_api_key: str | None = None) -> Settings:
+    """Settings pointing the app at the live infrastructure under test.
+
+    The search key is always part of the constructed settings object so the served
+    tool set is pinned by the test case itself and cannot depend on whether the CI
+    or the developer environment happens to export SEARCH_API_KEY.
+    """
     return Settings(
         DATABASE_URL=database_test_url(),
         REDIS_URL=redis_test_url(),
@@ -42,6 +47,7 @@ def _build_settings() -> Settings:
         MAX_STEPS=5,
         RUN_TIMEOUT_S=30,
         SESSION_TTL_S=3600,
+        SEARCH_API_KEY=search_api_key,
     )
 
 
@@ -92,6 +98,33 @@ def llm(real_client: tuple[TestClient, FakeLLMClient]) -> FakeLLMClient:
     return real_client[1]
 
 
+@pytest.fixture(
+    params=[
+        pytest.param("fake_search_key", id="search-enabled"),
+        pytest.param(None, id="search-disabled"),
+    ]
+)
+def tools_client(
+    request: pytest.FixtureRequest,
+    real_engine: AsyncEngine,
+    real_redis: Redis[str],
+) -> Iterator[tuple[TestClient, bool]]:
+    """TestClient built with a deterministic SEARCH_API_KEY for each case.
+
+    A separate application per case is what makes both behaviors observable: the
+    registry is memoized on app.state at first /tools access, so one shared client
+    could serve whichever case ran first. The fixtures are requested for the same
+    guarantee as real_client: truncate/cleanup happen before the app starts.
+    """
+    del real_engine, real_redis
+
+    search_key: str | None = request.param
+    app = create_app(settings=_build_settings(search_api_key=search_key))
+
+    with TestClient(app, raise_server_exceptions=False) as client:
+        yield client, search_key is not None
+
+
 async def test_health_reports_live_dependencies(client: TestClient) -> None:
     """The health endpoint must probe the real engine and Redis pool."""
     response = client.get("/health")
@@ -203,13 +236,23 @@ async def test_run_details_and_missing_run(client: TestClient) -> None:
     assert client.get(f"/runs/{uuid.uuid4()}", headers=headers).status_code == 404
 
 
-async def test_tools_endpoint_exposes_builtin_registry(client: TestClient) -> None:
-    """The built-in tool registry is served from the real dependency graph."""
+async def test_tools_endpoint_matches_configured_search_key(
+    tools_client: tuple[TestClient, bool],
+) -> None:
+    """The served registry depends only on the explicitly pinned SEARCH_API_KEY.
+
+    Runs twice: with a key configured (web_search present) and without one
+    (web_search omitted) — the latter being exactly what CI sees, where
+    SEARCH_API_KEY is never set for the real-infrastructure job.
+    """
+    client, web_search_expected = tools_client
+
     response = client.get("/tools", headers={"X-API-Key": API_KEY})
 
     assert response.status_code == 200
     names = {t["name"] for t in response.json()["tools"]}
-    assert {"calculator", "read_file", "http_fetch", "sql_readonly", "web_search"} <= names
+    assert {"calculator", "read_file", "http_fetch", "sql_readonly"} <= names
+    assert ("web_search" in names) == web_search_expected
 
 
 async def test_unauthorized_request_is_rejected(client: TestClient) -> None: