From c739a918015a0c90130a9d1168d557333275732f Mon Sep 17 00:00:00 2001 From: Amar K P Date: Sun, 1 Nov 2026 18:40:11 +0530 Subject: [PATCH 1/5] fix(core): harden cancellation and tool telemetry --- .env.example | 3 ++- README.md | 5 ++++- agentkit/api/deps.py | 12 ++++++++---- agentkit/core/agent.py | 33 +++++++++++++++++++++++++++++++-- agentkit/core/types.py | 1 + docs/DECISIONS.md | 8 ++++++++ docs/PROGRESS.md | 1 + 7 files changed, 55 insertions(+), 8 deletions(-) diff --git a/.env.example b/.env.example index 2edb690..06c9126 100644 --- a/.env.example +++ b/.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 diff --git a/README.md b/README.md index 304b9fd..4a93a3b 100644 --- a/README.md +++ b/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) diff --git a/agentkit/api/deps.py b/agentkit/api/deps.py index 2254a12..3669bed 100644 --- a/agentkit/api/deps.py +++ b/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 diff --git a/agentkit/core/agent.py b/agentkit/core/agent.py index 185ac5c..b6d5b54 100644 --- a/agentkit/core/agent.py +++ b/agentkit/core/agent.py @@ -300,6 +300,7 @@ async def _finalize_and_return(result: RunResult) -> 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. @@ -401,10 +402,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 +495,25 @@ async def _finalize_and_return(result: RunResult) -> RunResult: raise MaxStepsExceeded(self.config.max_steps) return res_max + except asyncio.CancelledError: + duration_ms = max(0, int((time.perf_counter() - start_time) * 1000)) + logger.info("Agent run canceled") + 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, + ) + try: + await asyncio.shield(_finalize_and_return(canceled_result)) + except PersistenceError: + logger.error("Could not persist canceled run state for run %s", run_id) + 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. diff --git a/agentkit/core/types.py b/agentkit/core/types.py index 7966940..d771a0c 100644 --- a/agentkit/core/types.py +++ b/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" diff --git a/docs/DECISIONS.md b/docs/DECISIONS.md index 2c8b408..b455311 100644 --- a/docs/DECISIONS.md +++ b/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 interrupt a run after its durable `running` row is created, and a model response that requests multiple tools has one usage count shared by all calls. +- **Decision**: On cancellation, finalize the run as `canceled` and persist repaired session history through a shielded finalization before re-raising cancellation. 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 diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 91a52d1..57ef97d 100644 --- a/docs/PROGRESS.md +++ b/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` | Pending | 0 | In progress | 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. Quality gates pass (ruff check/format, mypy agentkit); tests not added or run per task instruction. Staged for review: awaiting task approval to commit, push, and open the single review PR (`gh` authentication verified). | | 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 | From 144bc97c76b1fb0530afa94bbce10cc9a253cb8a Mon Sep 17 00:00:00 2001 From: Amar K P Date: Sun, 1 Nov 2026 18:41:13 +0530 Subject: [PATCH 2/5] docs(progress): record run hardening PR #122 --- docs/PROGRESS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 57ef97d..6ae2b92 100644 --- a/docs/PROGRESS.md +++ b/docs/PROGRESS.md @@ -4,7 +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` | Pending | 0 | In progress | 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. Quality gates pass (ruff check/format, mypy agentkit); tests not added or run per task instruction. Staged for review: awaiting task approval to commit, push, and open the single review PR (`gh` authentication verified). | +| 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) | 1 | 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. Quality gates pass (ruff check/format, mypy agentkit); tests not added or run per task instruction. | | 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 | From e2765efd7ebedcf8b823965328b6e9f60fc9b7fd Mon Sep 17 00:00:00 2001 From: Amar K P Date: Sun, 1 Nov 2026 19:11:22 +0530 Subject: [PATCH 3/5] fix(core): finalize canceled runs when the history load is interrupted Cancellation delivered while get_messages() was still awaiting escaped finalization entirely: the run row already existed in 'running', but the only CancelledError handler sat inside the reasoning-loop try block, so no terminal write (and no history repair happened at all) before re-raising. Route both cancellation windows through one shielded _finalize_canceled() helper. During setup the history save is skipped: the local snapshot is still empty and writing it would destroy the stored conversation rather than repair it. Adds the core-loop regression tests pinned to the post-create_run window and records the broadened rule in ADR-021. --- agentkit/core/agent.py | 96 ++++++++++----- docs/DECISIONS.md | 4 +- tests/unit/test_agent_cancellation.py | 171 ++++++++++++++++++++++++++ 3 files changed, 239 insertions(+), 32 deletions(-) create mode 100644 tests/unit/test_agent_cancellation.py diff --git a/agentkit/core/agent.py b/agentkit/core/agent.py index b6d5b54..5f27f18 100644 --- a/agentkit/core/agent.py +++ b/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,17 +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 - 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 - ) - try: async with asyncio.timeout(self.config.run_timeout_s): for _ in range(self.config.max_steps): @@ -496,23 +545,10 @@ async def _finalize_and_return(result: RunResult) -> RunResult: return res_max except asyncio.CancelledError: - duration_ms = max(0, int((time.perf_counter() - start_time) * 1000)) - logger.info("Agent run canceled") - 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, - ) - try: - await asyncio.shield(_finalize_and_return(canceled_result)) - except PersistenceError: - logger.error("Could not persist canceled run state for run %s", run_id) + # 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 diff --git a/docs/DECISIONS.md b/docs/DECISIONS.md index b455311..b48f7a5 100644 --- a/docs/DECISIONS.md +++ b/docs/DECISIONS.md @@ -6,8 +6,8 @@ This log records significant architectural, technical, and design decisions made - **Date**: 2026-11-01 - **Status**: Accepted -- **Context**: Cancellation can interrupt a run after its durable `running` row is created, and a model response that requests multiple tools has one usage count shared by all calls. -- **Decision**: On cancellation, finalize the run as `canceled` and persist repaired session history through a shielded finalization before re-raising cancellation. 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. +- **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`. --- diff --git a/tests/unit/test_agent_cancellation.py b/tests/unit/test_agent_cancellation.py new file mode 100644 index 0000000..9e661e9 --- /dev/null +++ b/tests/unit/test_agent_cancellation.py @@ -0,0 +1,171 @@ +"""Regression tests for cancellation finalization (PR #122 review finding). + +A cancellation can land on any awaited step that follows create_run(): the +session-history load before the reasoning loop, or an await inside the loop. In +both windows the run row already exists, so the run must reach a terminal canceled +status before the exception propagates; once the loop is running, the session +history must also be repaired so no orphaned assistant tool call stays stored. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from agentkit.core.agent import _ABORTED_TOOL_RESULT_CONTENT, Agent, AgentConfig +from agentkit.core.types import Role, RunStatus +from agentkit.llm.base import Message +from agentkit.llm.fake import FakeLLMClient +from agentkit.memory.base import InMemoryMemoryStore +from agentkit.tools.registry import ToolRegistry + + +class _RecordingStore(InMemoryMemoryStore): + """In-memory store recording the writes the agent performs during a run.""" + + def __init__(self) -> None: + super().__init__() + self.update_statuses: list[RunStatus] = [] + self.save_calls: list[list[Message]] = [] + + async def update_run( + self, + run_id: str, + status: RunStatus | str, + final_answer: str | None = None, + failure_reason: str | None = None, + total_input_tokens: int = 0, + total_output_tokens: int = 0, + ) -> None: + """Record the requested status, then delegate to the in-memory store.""" + self.update_statuses.append(RunStatus(status)) + await super().update_run( + run_id, + status, + final_answer=final_answer, + failure_reason=failure_reason, + total_input_tokens=total_input_tokens, + total_output_tokens=total_output_tokens, + ) + + async def save_messages(self, session_id: str, messages: list[Message]) -> None: + """Record each saved snapshot, then delegate to the in-memory store.""" + self.save_calls.append(list(messages)) + await super().save_messages(session_id, messages) + + +class _SuspendingHistoryStore(_RecordingStore): + """Store whose get_messages() stays pending until the test opens the gate. + + A pending first load reproduces the review finding: the await starts after + create_run() persisted the run row but before the reasoning loop — where the + inner CancelledError handler historically did not even exist. Opening the gate + after cancellation lets the test read the untouched history via the public API. + """ + + def __init__(self) -> None: + super().__init__() + self.run_created = asyncio.Event() + self._history_gate = asyncio.Event() + + def release_history(self) -> None: + """Open the history gate so the test can read the store back after cancellation.""" + self._history_gate.set() + + async def create_run(self, run_id: str, goal: str, session_id: str | None = None) -> None: + """Signal that the run row exists, mirroring the persisted 'running' state.""" + await super().create_run(run_id, goal, session_id=session_id) + self.run_created.set() + + async def get_messages(self, session_id: str) -> list[Message]: + """Stay pending until the gate opens, then load through the in-memory store.""" + await self._history_gate.wait() + return await super().get_messages(session_id) + + +@pytest.mark.asyncio +async def test_cancellation_during_history_load_finalizes_run_canceled() -> None: + """Cancel while get_messages() awaits: terminal write happens before re-raise. + + The gate opens only after create_run() persisted the row, so the cancellation is + delivered in exactly the window the review flagged: after the run exists and + before the handler guarded by the reasoning loop. + """ + store = _SuspendingHistoryStore() + session_id = "sess-cancel-setup" + history = [Message.user("Earlier conversation.")] + await store.save_messages(session_id, history) + # The pre-run seed writes through the same spy: the run being tested has not + # started yet, so its own save activity starts from an empty recording. + store.save_calls.clear() + + agent = Agent( + llm=FakeLLMClient(), + registry=ToolRegistry(), + memory=store, + config=AgentConfig(max_steps=3), + ) + + task = asyncio.create_task(agent.run(goal="Do the thing.", session_id=session_id)) + await asyncio.wait_for(store.run_created.wait(), timeout=2.0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # Exactly one terminal write, and it is the canceled one: nothing between + # create_run and the guard could have marked the run failed or finished. + assert store.update_statuses == [RunStatus.CANCELED] + # History was still loading, so the local list is empty: saving it would + # overwrite the stored session with nothing instead of repairing it. + assert store.save_calls == [] + store.release_history() + assert await store.get_messages(session_id) == history + + +@pytest.mark.asyncio +async def test_cancellation_during_tool_execution_repairs_history() -> None: + """Cancel while a tool executes: canceled status plus repaired history persist.""" + started = asyncio.Event() + gate = asyncio.Event() + registry = ToolRegistry() + + async def blocking_tool() -> str: + """Hold the reasoning step open until the run is cancelled.""" + started.set() + await gate.wait() + return "unreachable: the run is cancelled before the tool completes" + + registry.register(blocking_tool) + + store = _RecordingStore() + llm = FakeLLMClient() + llm.queue_tool_call( + tool_name="blocking_tool", + arguments={}, + call_id="call-blocked", + text="Calling the tool.", + ) + + agent = Agent( + llm=llm, + registry=registry, + memory=store, + config=AgentConfig(max_steps=3), + ) + + task = asyncio.create_task(agent.run(goal="Use the tool.", session_id="sess-cancel-loop")) + await asyncio.wait_for(started.wait(), timeout=2.0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert store.update_statuses == [RunStatus.CANCELED] + assert len(store.save_calls) == 1 + snapshot = store.save_calls[0] + assert [m.role for m in snapshot] == [Role.SYSTEM, Role.USER, Role.ASSISTANT, Role.TOOL] + # The categorical repair rule: every requested tool call keeps a matching tool + # result, so the next run in this session gets a provider-valid history. + assert snapshot[-2].tool_calls is not None and snapshot[-2].tool_calls[0].id == "call-blocked" + assert snapshot[-1].tool_call_id == "call-blocked" + assert snapshot[-1].content == _ABORTED_TOOL_RESULT_CONTENT From 277905281e22e2360955a9ae0d37bd48f6754888 Mon Sep 17 00:00:00 2001 From: Amar K P Date: Sun, 1 Nov 2026 19:11:28 +0530 Subject: [PATCH 4/5] test(real): pin /tools registry to an explicit SEARCH_API_KEY per case test_tools_endpoint_exposes_builtin_registry asserted web_search unconditionally, but the real-infra CI job never exports SEARCH_API_KEY, so the conditional registration change failed that suite at 144bc97. The new parametrized tools_client fixture builds one application per case with SEARCH_API_KEY passed explicitly: web_search is expected present with a key configured and absent without one, independent of the ambient environment, so a developer with the variable set cannot mask the key-missing path (or vice versa). --- tests/real/test_api_e2e_real.py | 53 +++++++++++++++++++++++++++++---- 1 file changed, 48 insertions(+), 5 deletions(-) diff --git a/tests/real/test_api_e2e_real.py b/tests/real/test_api_e2e_real.py index 7a66e85..c2cd35a 100644 --- a/tests/real/test_api_e2e_real.py +++ b/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: From c7e04686a9756271d4145780dcb0dbeea442d17d Mon Sep 17 00:00:00 2001 From: Amar K P Date: Sun, 1 Nov 2026 19:11:30 +0530 Subject: [PATCH 5/5] docs(progress): record PR #122 review-fix round --- docs/PROGRESS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 6ae2b92..f8f18ad 100644 --- a/docs/PROGRESS.md +++ b/docs/PROGRESS.md @@ -4,7 +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) | 1 | 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. Quality gates pass (ruff check/format, mypy agentkit); tests not added or run per task instruction. | +| 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 |