diff --git a/.github/dependabot.yml b/.github/dependabot.yml index f79facf..0effc24 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -16,9 +16,7 @@ # graph -> Dependabot should list this manifest. If it does not detect the PEP # 621 setuptools [project.optional-dependencies] table, nothing here runs and # the exact pins freeze silently — which is the exact rot this file exists to -# prevent. On non-detection, fall back to a manual quarterly bump and record -# that in the coding_guidelines.md Decisions row, so the obligation survives -# outside this file. +# prevent. On non-detection, fall back to a manual quarterly bump. version: 2 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 87e3f88..990718d 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -71,6 +71,21 @@ jobs: with: python-version: "3.12" + # A mistyped tag would silently publish an artifact whose version does + # not match the tag it was released as. Compare refs/tags/vX.Y.Z against + # pyproject's project.version BEFORE building, so the failure is cheap + # and nothing mismatched reaches the artifact. Skipped on + # workflow_dispatch (the ref is a branch, not a tag). + - name: Verify tag matches pyproject version + if: startsWith(github.ref, 'refs/tags/v') + run: | + tag="${GITHUB_REF_NAME#v}" + version="$(grep -m1 '^version = ' pyproject.toml | cut -d'\"' -f2)" + if [ "$tag" != "$version" ]; then + echo "::error::tag ${GITHUB_REF_NAME} does not match pyproject version ${version}" + exit 1 + fi + - name: Build sdist + wheel run: | python -m pip install --upgrade pip build diff --git a/CHANGELOG.md b/CHANGELOG.md index 5ddda42..5954811 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,82 @@ All notable changes to `fifty-agent-sdk` are documented here. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added +- `OpenAICompatibleClient` gains `aclose()` and async context-manager support, + so the underlying httpx client can be disposed deterministically. +- `ChatRequest.tool_choice` accepts the specific-tool dict form (forcing one + named tool), alongside the existing string forms. +- `temperature=None` omits the parameter from the request body instead of + sending it, letting a provider's own default apply. + +### Fixed +- Python 3.11 package imports work with the named-tool `tool_choice` form: + `TypedDict` now comes from the directly-declared `typing-extensions` + dependency, as required by Pydantic on Python versions below 3.12. (BR-014) +- All five direct `json.loads` boundaries now contain bare `ValueError` as + `ParserError` or `LLMError`, including CPython's oversized-integer guard, + while preserving the existing malformed-syntax phases, messages, context, + and exception chaining. (BR-013) +- Parser recursion-limit regressions now inject `RecursionError` + deterministically instead of relying on interpreter-specific behavior from + a 100,000-level JSON value, restoring portable Python 3.14 coverage. + (BR-017) +- **Audit payload shape change (consumer-visible for `AuditSink` + implementors):** the `tool_invocation` payload's `args` field no longer + embeds the raw argument dict — it is now per-key metadata (sorted argument + keys with each value's type name and length, `len=None` for unsized + values), closing a leak of secrets/PII into persisted audit payloads. The + `on_tool_start` hook still receives the full args. +- The runner now correlates tool invocations by `call_id` across + `MultiAction` batches — previously the first terminal event inherited the + last call's `call_id`/args and every other call was audited with + `args={}`. +- A fatal `AgentSdkError` escaping the loop (e.g. `MCPError`) is now + surfaced to hooks and audit before re-raising: the error audit event is + emitted, `on_error` fires, and the run records + `terminated_by="sdk_error"` instead of exiting as `"interrupted"` with + `on_run_end(error=None)`. +- Hook-failure logs now carry `hook_name` and `error_type` only, never + `str(exc)` — a raising hook could previously dump conversation content + into a WARNING log line. +- `SqlStateStore` persists `ChatMessage.tool_calls` (nullable JSON column; + JSONB on Postgres), so a persisted native-tool-calling assistant turn no + longer loses its `tool_calls` — and orphans the paired `role="tool"` + replies — on session resume. Additive, via the existing consumer-owned + migration path; the SDK still ships no migrations. +- `RedisStateStore` rejects non-positive `ttl_seconds` at construction — + `ttl_seconds=0` previously made the append's `EXPIRE` delete every session + key immediately. +- Every Redis mutation (`append`, `fork`, `switch_branch`, and + `truncate_after`) now runs through one optimistic transaction that refreshes + every session key to the same sliding TTL. Registry/active/message-key + conflicts retry from a fresh snapshot up to a bounded limit, preventing + metadata keys from outliving their message lists. (BR-015) +- The loop rejects `stream=True` combined with `native_tools_enabled` at + construction instead of misbehaving later. +- A `RecursionError` from pathologically nested JSON is translated into + `ParserError` (`error_phase` `json_decode` / `action_input_decode`) + instead of escaping the parser contract. +- The JSON-mode parser strips `tool_name` and rejects a whitespace-only one, + which previously produced a `ThoughtAction` the registry could never + match. +- An MCP auth callable that raises (e.g. a down token endpoint) is + translated into `MCPError` — only the exception type name is captured, + never its text — instead of escaping the MCPError-only contract and being + downgraded to a model-recoverable `ToolResult`. +- A `CancelledError` arriving as a leaf of a mixed transport + `BaseExceptionGroup` re-raises untouched instead of being translated into + `MCPError`, restoring the cancellation contract. +- Tool schemas sent to the LLM now have `#/$defs` references inlined, so a + nested `BaseModel` parameter no longer reaches the model as a dangling + `$ref`; recursive models are rejected at decoration time for `@tool` and + fall back to an empty schema for untrusted MCP server schemas. Expansion is + additionally capped at 10,000 total resolver visits, preventing acyclic + fan-out from exhausting memory; MCP fallback logs contain only stable + reason/type metadata, never remote schema text. (BR-016) + ## [1.5.0] - 2026-07-30 ### Added @@ -217,7 +293,7 @@ extracted with its full commit history from the monorepo it was first built in. - Import root is now `fifty_agent_sdk` (was `agent_sdk`). - Distributed and published as `fifty-agent-sdk` on PyPI. -## [1.0.0] +## [1.0.0] - 2026-06-19 Initial production release: custom ReACT loop, JSON-mode tool calling, a pluggable LLM client (any OpenAI-compatible endpoint), in-process + MCP tool diff --git a/MAINTAINING.md b/MAINTAINING.md index 7b8fa2e..df5902b 100644 --- a/MAINTAINING.md +++ b/MAINTAINING.md @@ -33,14 +33,15 @@ tells you **that** you must notify a consumer and **what** to tell them, but not **which files to grep** — for that, open the tracker. **Authoring rule (the anti-rot test depends on it).** Every path is backticked. -Every path in this file is a path in **this** repo, starting `src/` or `tests/`, -optionally with a `:line` suffix. `tests/test_maintaining.py` asserts that every -one of them still exists. **Do not add consumer-repo paths** — see above. +Every path in this file is a path in **this** repo, starting `src/` or +`tests/` — or a root-level `*.md` doc — optionally with a `:line` suffix. +`tests/test_maintaining.py` asserts that every one of them still exists. **Do +not add consumer-repo paths** — see above. **A row requires a real, verified consumer**, even though it is described by role rather than named. Not "a seam somebody might implement" — an integration -confirmed by grepping an actual checkout. This repo has three standing decisions -of the form *never on speculation* (see `coding_guidelines.md` Decisions), and a +confirmed by grepping an actual checkout. This repo has standing decisions +of the form *never on speculation*, and a map padded with maybe-consumers is exactly the doc that rots. Seams with no confirmed consumer go in [Not mapped](#not-mapped), so their absence is a recorded decision rather than an oversight. @@ -60,7 +61,7 @@ invented one. |----------|---------|-----------|-----------------| | **`state/sql.py` physical database schema** — **highest-consequence row** | The table, column, unique-constraint and index NAMES emitted by the `sql` extra's ORM: `agent_sessions` (`src/fifty_agent_sdk/state/sql.py:199`), `agent_messages` (`:280`), `agent_branches` (`:372`), `uq_agent_messages_session_branch_sequence` (`src/fifty_agent_sdk/state/sql.py:286`), `ix_agent_messages_session_branch_sequence` (`:289`), and the `ON DELETE CASCADE` FKs to `agent_sessions.session_id`. Exposed as `fifty_agent_sdk.sql_metadata`. The SDK ships ORM models and **deliberately does not own migrations** — `SqlStateStore` has no `create_all` / bootstrap and expects the tables to pre-exist. | **Consumer B** — hand-authors an Alembic chain against a **production** database, transcribed from this ORM column-for-column, and pins itself to an SDK version in that migration's own docstring. Its migrations drop and recreate the `uq_`/`ix_` pair **by name**. | **Any DDL-name change is a MAJOR break, and this repo catches only some of them.** `tests/state/test_sql.py` asserts against `sql_metadata` directly, not only through the ORM: a **table** rename fails `test_metadata_exposes_both_tables` (`tests/state/test_sql.py:640`) and a **column** rename fails `test_metadata_columns_match_schema` (`:648`), which compares the full column-name set of all three tables. **The constraint and index NAMES were the uncovered gap** — `test_metadata_unique_constraint_on_session_branch_sequence` (`:691`) matches on the constraint's *columns* and never on its `name=`. TD-003 closed that with `test_metadata_pins_constraint_and_index_names`, so a rename now fails here instead of silently desynchronising a production database. On any rename/drop/constraint change: bump MAJOR, say so in `CHANGELOG.md` with the old→new names, and notify Consumer B so a follow-on migration lands before the version bump. Additive nullable columns are safe (their table simply lacks them until they migrate). | | **MCP public surface** — pinned by `tests/mcp/test_interface_stability.py` | `MCPClient`, `MCPClientConfig`, `MCPToolDef`, `MCPToolErrorHook` (`src/fifty_agent_sdk/mcp/client.py`), `MCPProvider`, `RefreshSummary` (`src/fifty_agent_sdk/tools/mcp_provider.py`) — plus two couplings the pin does not express: the module path `fifty_agent_sdk.mcp.client`, and the fact that `on_tool_error` fires *inside* `MCPClient.invoke`. | **Consumer A** — **subclasses** `MCPClient` and overrides the public `invoke`, constructing it with the public `on_tool_error` hook; patches `fifty_agent_sdk.mcp.client.StreamableHttpTransport` by **string path** in its tests. **Consumer B** — *calls* `MCPClient` / `MCPClientConfig` / `MCPProvider` but does not subclass. | Update the pin in the SAME change that widens the surface, then sweep both consumers. **Subclassing is a tighter contract than calling:** A overriding `invoke` makes the method's name, its `(self, name, args)` shape, and where `on_tool_error` fires relative to it all load-bearing for them, while B would survive an internal re-order. Moving `mcp/client.py` breaks A's string-path test patch with an `AttributeError` naming no SDK file — keep a re-export at the old module path, or major-bump. | -| **Submodule import paths** — no pin | `api_pattern.md` says everything is imported from the package root and "submodule paths work but are not the contract". In practice both consumers import through them in **src-tier** code: `fifty_agent_sdk.tools.protocol`, `fifty_agent_sdk.llm.types`, `fifty_agent_sdk.state.protocol`, `fifty_agent_sdk.streaming`, `fifty_agent_sdk.errors`. Two adjacent cases: `fifty_agent_sdk.observability.hooks` is reached this way only in a consumer **test**, and `fifty_agent_sdk.mcp.client` is never imported by submodule path — it appears only as a `patch()` **string**, which the MCP row covers. | **Consumer A** — three src-tier sites importing from `tools.protocol`. **Consumer B** — sites importing from `state.protocol`, `tools.protocol`, `llm.types` and `errors`; `streaming` is the widest at five src-tier sites. | The stated contract is root-only, so this coupling is **the consumers' risk, not a promise this repo made** — recorded here so the cost of moving a module is visible rather than surprising. **Do not infer a file's import style from the symbols it uses:** the same names are imported by root elsewhere in the same repos, so a sweep of this row must re-grep for `fifty_agent_sdk.` **with a trailing dot**, never for the symbol names. (TD-003's own first draft got this wrong on four files and CI stayed green — see [Known limitations](#known-limitations).) When relocating a module under `src/fifty_agent_sdk/`, prefer leaving a re-export at the old path for one MINOR; otherwise treat it as MAJOR and name the moved paths in `CHANGELOG.md`. | +| **Submodule import paths** — no pin | The package docstring (`src/fifty_agent_sdk/__init__.py`) declares `fifty_agent_sdk.__all__` the API contract: everything is imported from the package root, and submodule paths work but are not the contract. In practice both consumers import through them in **src-tier** code: `fifty_agent_sdk.tools.protocol`, `fifty_agent_sdk.llm.types`, `fifty_agent_sdk.state.protocol`, `fifty_agent_sdk.streaming`, `fifty_agent_sdk.errors`. Two adjacent cases: `fifty_agent_sdk.observability.hooks` is reached this way only in a consumer **test**, and `fifty_agent_sdk.mcp.client` is never imported by submodule path — it appears only as a `patch()` **string**, which the MCP row covers. | **Consumer A** — three src-tier sites importing from `tools.protocol`. **Consumer B** — sites importing from `state.protocol`, `tools.protocol`, `llm.types` and `errors`; `streaming` is the widest at five src-tier sites. | The stated contract is root-only, so this coupling is **the consumers' risk, not a promise this repo made** — recorded here so the cost of moving a module is visible rather than surprising. **Do not infer a file's import style from the symbols it uses:** the same names are imported by root elsewhere in the same repos, so a sweep of this row must re-grep for `fifty_agent_sdk.` **with a trailing dot**, never for the symbol names. (TD-003's own first draft got this wrong on four files and CI stayed green — see [Known limitations](#known-limitations).) When relocating a module under `src/fifty_agent_sdk/`, prefer leaving a re-export at the old path for one MINOR; otherwise treat it as MAJOR and name the moved paths in `CHANGELOG.md`. | | **Protocol seams with a confirmed external implementor** — no pin | Structural typing means an external class satisfies `Tool` (`src/fifty_agent_sdk/tools/protocol.py`), `AuditSink` (`src/fifty_agent_sdk/audit/protocol.py`) or `LLMClient` (`src/fifty_agent_sdk/llm/protocol.py`) without importing or inheriting anything — and `Hooks` (`src/fifty_agent_sdk/observability/hooks.py`) is constructed by keyword. Widening a signature or adding a required protocol method breaks the implementor with **no** failure in this repo and no `runtime_checkable` failure there (`runtime_checkable` checks method *presence* only). | `Tool` — both consumers (A wraps tools for caching; B wraps one for gateway resilience). `AuditSink` — Consumer B. `LLMClient` — Consumer B, in a src-tier scripted client (not a test double). `Hooks` — both, each constructing it by keyword. | The only guard is the **consumer's own** `mypy` run, which happens after they upgrade. So: adding a protocol method or widening an existing signature is MAJOR; adding a keyword-only parameter with a default is MINOR and safe. New `Hooks` slots must default to `None` (the existing convention) or every keyword construction site breaks. Name any protocol change in `CHANGELOG.md` — it is the only notice these implementors get. | ## Private symbols diff --git a/README.md b/README.md index a4dca5d..8809505 100644 --- a/README.md +++ b/README.md @@ -213,7 +213,7 @@ depends (→): legend: ▶ entry ▢ package name module → depends ``` -## What's new in 1.2.0 +## Highlights - **branching** — first-class conversation branching on `StateStore`: `fork`, `list_branches`, `switch_branch`, branch-scoped `get_messages(..., branch_id=...)`, plus `BranchInfo` and `TRUNK_BRANCH_ID`. a session is now a tree of branches with an active head, and `append` writes to the active branch (the edit-a-message / regenerate model). implemented across memory, SQL, and Redis backends, data-additive and zero-migration: existing sessions read as the trunk branch. breaking for custom `StateStore` implementations: they must add the new methods. - **`StateStore.truncate_after(session_id, sequence, *, branch_id=None)`** — a destructive hard-delete of a branch's tail (messages with sequence > N), for redaction, retention, and rollback. only the target branch's own messages are removed (a `fork`'s inherited prefix is never touched), and it is idempotent: a no-op on an unknown session or branch. diff --git a/pyproject.toml b/pyproject.toml index bc20ca0..6291a26 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,9 @@ classifiers = [ dependencies = [ "pydantic>=2.5.0", + # ChatRequest's named-tool choice uses TypedDict. Importing it from + # typing_extensions is required by Pydantic on Python 3.11. + "typing-extensions>=4.6.1", "httpx>=0.26.0", "structlog>=24.1.0", "openai>=1.30.0,<3.0.0", @@ -128,6 +131,10 @@ exclude = ["tests/", "build/", ".venv/"] [tool.pytest.ini_options] asyncio_mode = "auto" +# Pin the fixture loop scope explicitly (TD-004): "function" is pytest-asyncio's +# current default, and setting it silences the unset-option deprecation warning +# while pinning the behavior against a future default change. +asyncio_default_fixture_loop_scope = "function" testpaths = ["tests"] python_files = ["test_*.py"] addopts = ["-v", "--tb=short", "--strict-markers", "--import-mode=importlib"] diff --git a/src/fifty_agent_sdk/llm/__init__.py b/src/fifty_agent_sdk/llm/__init__.py index b5bd02c..6d664dd 100644 --- a/src/fifty_agent_sdk/llm/__init__.py +++ b/src/fifty_agent_sdk/llm/__init__.py @@ -14,6 +14,8 @@ FinishReason, Role, ToolCall, + ToolChoiceFunction, + ToolChoiceFunctionName, Usage, ) @@ -26,5 +28,7 @@ "OpenAICompatibleClient", "Role", "ToolCall", + "ToolChoiceFunction", + "ToolChoiceFunctionName", "Usage", ] diff --git a/src/fifty_agent_sdk/llm/openai_compat.py b/src/fifty_agent_sdk/llm/openai_compat.py index 410b5fa..82dafa6 100644 --- a/src/fifty_agent_sdk/llm/openai_compat.py +++ b/src/fifty_agent_sdk/llm/openai_compat.py @@ -14,6 +14,7 @@ import json from collections.abc import AsyncIterator +from types import TracebackType from typing import Any import httpx @@ -58,6 +59,13 @@ class OpenAICompatibleClient: Provider variation is absorbed by ``base_url`` — the same client class drives OpenAI itself, GDC, and local OSS servers. + Async context management + The client is an async context manager: ``async with + OpenAICompatibleClient(...) as client:`` closes the underlying + connection pool on exit via :meth:`aclose`. Without the context + manager, call :meth:`aclose` explicitly — an unclosed client leaks + its ``httpx.AsyncClient`` pool until GC. + Args: api_key: API key passed to the upstream provider. Required even for local servers that ignore it; pass any non-empty string. @@ -71,7 +79,11 @@ class OpenAICompatibleClient: errors surface immediately, which is what tests want. http_client: Optional pre-configured ``httpx.AsyncClient``. Useful for tests that need to inject a mock transport. When omitted, the - SDK builds its own client. + SDK builds its own client, OWNS it, and closes it on + :meth:`aclose`. An INJECTED client is NOT closed by + :meth:`aclose` — its lifecycle belongs to the caller. This + owned-vs-injected discipline mirrors + :class:`fifty_agent_sdk.mcp.client.MCPClient`. """ def __init__( @@ -95,6 +107,39 @@ def __init__( kwargs["http_client"] = http_client self._client = AsyncOpenAI(**kwargs) self._default_model = model + # Ownership discipline mirrors MCPClient: an injected http_client is + # the caller's to close; only the client built here is closed by + # aclose(). (`AsyncOpenAI.close()` closes the underlying httpx client + # unconditionally, so it must only be called on the owned path.) + self._owns_client = http_client is None + self._closed = False + + async def aclose(self) -> None: + """Close the underlying ``openai`` client and its connection pool. + + Only an OWNED client is closed: when ``http_client`` was injected at + construction, its lifecycle belongs to the caller and ``aclose()`` + leaves it open. Idempotent: a second ``aclose()`` is a no-op and does + NOT raise. The client MUST NOT be used after ``aclose()`` returns. + """ + if self._closed: + return + if self._owns_client: + await self._client.close() + self._closed = True + + async def __aenter__(self) -> OpenAICompatibleClient: + """Enter the async context manager, returning ``self``.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + """Exit the async context manager, closing via :meth:`aclose`.""" + await self.aclose() async def complete(self, request: ChatRequest) -> ChatResponse: """Run a single non-streaming completion. @@ -278,13 +323,20 @@ def _serialize_message(msg: ChatMessage) -> dict[str, Any]: @classmethod def _build_body(cls, request: ChatRequest, *, model: str, stream: bool) -> dict[str, Any]: - """Build the kwargs passed to ``client.chat.completions.create``.""" + """Build the kwargs passed to ``client.chat.completions.create``. + + ``temperature`` is omitted entirely when :attr:`ChatRequest.temperature` + is ``None`` (the opt-out for providers that reject a non-default + temperature); at the ``0.0`` default it is sent, preserving the + pre-existing wire shape for every caller that does not opt out. + """ body: dict[str, Any] = { "model": model, "messages": [cls._serialize_message(m) for m in request.messages], - "temperature": request.temperature, "stream": stream, } + if request.temperature is not None: + body["temperature"] = request.temperature if request.max_tokens is not None: body["max_tokens"] = request.max_tokens if request.response_format is not None: @@ -383,7 +435,10 @@ def _map_tool_calls(cls, raw: Any, *, model: str) -> list[ToolCall] | None: # n arguments = function.arguments if function.arguments is not None else "" try: args = json.loads(arguments) if arguments else {} - except json.JSONDecodeError as e: + except ValueError as e: + # json.loads also raises bare ValueError for interpreter + # limits such as oversized integers. Provider payloads must + # remain contained as LLMError without parsing exception text. raise LLMError( "provider tool_call arguments is not valid JSON", context={ diff --git a/src/fifty_agent_sdk/llm/types.py b/src/fifty_agent_sdk/llm/types.py index b80b672..fea86a9 100644 --- a/src/fifty_agent_sdk/llm/types.py +++ b/src/fifty_agent_sdk/llm/types.py @@ -16,6 +16,7 @@ from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field +from typing_extensions import TypedDict Role = Literal["system", "user", "assistant", "tool"] """Discriminator for the speaker of a :class:`ChatMessage`.""" @@ -115,6 +116,24 @@ class Usage(BaseModel): total_tokens: int = Field(ge=0) +class ToolChoiceFunctionName(TypedDict): + """The ``function`` member of :class:`ToolChoiceFunction`.""" + + name: str + + +class ToolChoiceFunction(TypedDict): + """The specific-tool form of :attr:`ChatRequest.tool_choice`. + + Mirrors the OpenAI ``ChatCompletionNamedToolChoiceParam`` shape — + ``{"type": "function", "function": {"name": ...}}`` — and forces the + model to call the named function instead of choosing freely. + """ + + type: Literal["function"] + function: ToolChoiceFunctionName + + class ChatRequest(BaseModel): """Provider-agnostic chat-completion request. @@ -122,6 +141,11 @@ class ChatRequest(BaseModel): messages: Ordered list of conversation messages. model: Model identifier. Adapters may override this with a default. temperature: Sampling temperature in ``[0.0, 2.0]``. Default ``0.0``. + Set to ``None`` to OMIT the parameter from the request body + entirely — for providers/models (e.g. OpenAI's reasoning-model + family) that reject a non-default temperature. The default is + ``0.0`` (NOT ``None``) so existing callers' wire behavior is + unchanged: the key is sent unless explicitly set to ``None``. max_tokens: Optional cap on completion tokens. Must be ``>= 1`` if set. response_format: Optional provider-format hint. Common values are ``{"type": "json_object"}`` or ``{"type": "text"}``. Adapters @@ -135,10 +159,11 @@ class ChatRequest(BaseModel): default) means NO tools are declared and the request wire is byte-for-byte the pre-BR-008 shape. tool_choice: Optional steering for native tool-calling when - :attr:`tools` is set. Common values are ``"auto"`` (the default - the adapter emits when this is ``None``), ``"required"``, - ``"none"``, or a specific-tool object - ``{"type": "function", "function": {"name": ...}}``. Ignored when + :attr:`tools` is set. Accepts a mode string — ``"auto"`` (the + default the adapter emits when this is ``None``), + ``"required"``, or ``"none"`` — or a :class:`ToolChoiceFunction` + object ``{"type": "function", "function": {"name": ...}}`` that + forces the model to call the named function. Ignored when :attr:`tools` is ``None``. """ @@ -146,11 +171,11 @@ class ChatRequest(BaseModel): messages: list[ChatMessage] model: str - temperature: float = Field(default=0.0, ge=0.0, le=2.0) + temperature: float | None = Field(default=0.0, ge=0.0, le=2.0) max_tokens: int | None = Field(default=None, ge=1) response_format: dict[str, Any] | None = None tools: list[dict[str, Any]] | None = None - tool_choice: str | None = None + tool_choice: str | ToolChoiceFunction | None = None class ChatResponse(BaseModel): @@ -184,5 +209,7 @@ class ChatResponse(BaseModel): "FinishReason", "Role", "ToolCall", + "ToolChoiceFunction", + "ToolChoiceFunctionName", "Usage", ] diff --git a/src/fifty_agent_sdk/loop.py b/src/fifty_agent_sdk/loop.py index 534d2c0..b80bc85 100644 --- a/src/fifty_agent_sdk/loop.py +++ b/src/fifty_agent_sdk/loop.py @@ -12,10 +12,15 @@ retry orchestration are a higher Runner's job (see BR-007). System prompt snapshot - The system prompt is constructed ONCE at :meth:`AgentLoop.run` start + The system prompt is constructed ONCE at ``AgentLoop.__init__`` from a snapshot of :meth:`fifty_agent_sdk.tools.registry.Registry.list`. - Tools registered AFTER ``run()`` begins are NOT visible to the - model. Dynamic-tool consumers must rebuild the loop. + In the default text/JSON mode, tools registered AFTER construction are + NOT visible to the model. Dynamic-tool consumers must rebuild the loop. + Under native-tools mode + (:attr:`fifty_agent_sdk.safety.SafetyConfig.native_tools_enabled`) the + prompt-side tool block is suppressed and the OpenAI ``tools`` request + param is rebuilt from the registry on EVERY iteration, so late + registrations ARE declared to the provider on subsequent turns. Cancellation :class:`asyncio.CancelledError` propagates untouched. The most @@ -30,6 +35,13 @@ is emitted — the parser requires the complete structured completion to disambiguate ``tool`` versus ``final``. This is a deliberate contract trade-off. + + Streaming is INCOMPATIBLE with native-tools mode: a streamed response + carries no structured ``tool_calls`` (the adapter deliberately drops + streamed tool-call deltas), so a ``stream=True`` + + ``native_tools_enabled=True`` loop would parse empty completions and + terminate on the fallback. ``AgentLoop(...)`` raises ``ValueError`` + at construction when both are enabled. """ from __future__ import annotations @@ -283,7 +295,10 @@ class AgentLoop: stream: If ``True``, use :meth:`LLMClient.stream` and emit :class:`fifty_agent_sdk.streaming.TokenEvent` for the FINAL answer only. Default ``False`` — uses :meth:`LLMClient.complete`, - no token events. + no token events. MUST NOT be combined with + ``SafetyConfig(native_tools_enabled=True)`` — streamed responses + carry no structured ``tool_calls``, so that combination is + rejected at construction (see Raises). output_format: Optional override for the system prompt's ``output_format`` slot. When non-empty, replaces whatever is in ``prompts.output_format``. Useful for callers that want @@ -310,6 +325,16 @@ class AgentLoop: receive ``hooks`` from a Runner — the two are wired independently by the consumer). A raising hook never aborts the loop. When ``None`` (default), hook dispatch is zero-overhead. + + Raises: + ValueError: When ``stream=True`` is combined with + ``SafetyConfig(native_tools_enabled=True)``. A streamed + completion carries no structured ``tool_calls`` (the LLM + adapter deliberately drops streamed tool-call deltas), so the + combination can never dispatch a native tool call — every + streamed turn would parse empty content, burn the parser + retry, and terminate on the fallback message. The + misconfiguration is rejected at construction instead. """ def __init__( @@ -328,6 +353,14 @@ def __init__( tool_message_role: Literal["tool", "user", "assistant"] = "tool", hooks: Hooks | None = None, ) -> None: + if stream and safety.native_tools_enabled: + raise ValueError( + "stream=True is incompatible with " + "SafetyConfig(native_tools_enabled=True): a streamed " + "completion carries no structured tool_calls, so native " + "tool dispatch can never fire. Disable streaming or turn " + "native_tools_enabled off." + ) self._llm = llm self._registry = registry self._parser = parser diff --git a/src/fifty_agent_sdk/mcp/client.py b/src/fifty_agent_sdk/mcp/client.py index 604f8ea..2e29d10 100644 --- a/src/fifty_agent_sdk/mcp/client.py +++ b/src/fifty_agent_sdk/mcp/client.py @@ -53,7 +53,14 @@ Every public method either returns successfully or raises :class:`MCPError`. No ``mcp`` SDK exceptions (``McpError``) and no ``httpx`` exceptions leak — they are unwrapped from anyio - ``ExceptionGroup``s and translated into :class:`MCPError`. The + ``ExceptionGroup``s and translated into :class:`MCPError`. An exception + raised by a callable ``auth`` provider (e.g. a token endpoint being down) + is translated the same way, with only the exception TYPE captured in + ``context`` — never the exception text, which may carry credential + material. Cancellation is never translated: a + :class:`asyncio.CancelledError` propagates untouched, including when it + arrives as a leaf of a mixed ``BaseExceptionGroup`` (cancellation racing + a real transport error inside the anyio task group). The :class:`fifty_agent_sdk.tools.mcp_provider._MCPToolAdapter` deliberately does NOT catch :class:`MCPError` so the :class:`fifty_agent_sdk.tools.registry.Registry`'s ``AgentSdkError`` branch @@ -439,6 +446,10 @@ async def discover(self) -> list[MCPToolDef]: or server-returned JSON-RPC error. Also raised with ``message == "MCP client is closed"`` when invoked after :meth:`aclose`. + asyncio.CancelledError: Propagates untouched — including when it + arrives as a leaf of a mixed anyio ``BaseExceptionGroup`` + (cancellation racing a real transport error). See + :meth:`_mcp_error_from_transport`. """ if self._closed: raise MCPError( @@ -513,6 +524,10 @@ async def invoke(self, name: str, args: dict[str, Any]) -> Any: :meth:`aclose`. A per-call ``isError=True`` result no longer raises — it is returned as a :class:`_MCPCallError` (see ``Returns``). + asyncio.CancelledError: Propagates untouched — including when it + arrives as a leaf of a mixed anyio ``BaseExceptionGroup`` + (cancellation racing a real transport error). See + :meth:`_mcp_error_from_transport`. """ if self._closed: raise MCPError( @@ -605,7 +620,15 @@ async def _resolve_auth(self) -> tuple[Mapping[str, str], frozenset[str]]: sensitive (used to extend redaction). Raises: - MCPError: When a callable provider returns a non-Mapping value. + MCPError: When a callable provider returns a non-Mapping value, or + when the callable itself raises (e.g. a token endpoint being + down surfacing as :class:`httpx.ConnectError`). A raising + callable is translated so the module's "returns or raises + MCPError" contract holds; without it the raw exception would + reach :class:`fifty_agent_sdk.tools.registry.Registry`'s + non-``AgentSdkError`` branch and be DOWNGRADED to a + model-recoverable ``ToolResult(is_error=True)`` instead of + staying the fatal infrastructure failure it is. """ auth = self._auth if auth is None: @@ -613,7 +636,29 @@ async def _resolve_auth(self) -> tuple[Mapping[str, str], frozenset[str]]: if isinstance(auth, Mapping): return auth, frozenset(k.lower() for k in auth) # auth is a callable - resolved = await auth() + try: + resolved = await auth() + # ``except Exception`` (not bare/BaseException) keeps + # ``asyncio.CancelledError`` — a BaseException on 3.11+ — propagating + # untouched: consumer cancellation must never be re-labelled as an + # auth failure. This is the same classification-ordering rule the + # ``on_tool_error`` hook documents. + except Exception as exc: + # Type name only, never str(exc): an auth error message may echo + # endpoint material or credentials, and this module's standing + # rule is that such material never reaches an MCPError context. + _log.warning( + "mcp.auth_callable_failed", + server_url=self._config.base_url, + error_type=type(exc).__name__, + ) + raise MCPError( + f"auth callable raised: {type(exc).__name__}", + context={ + "server_url": self._config.base_url, + "wrapped": type(exc).__name__, + }, + ) from exc if not isinstance(resolved, Mapping): raise MCPError( "auth callable must return a Mapping[str, str]", @@ -693,8 +738,24 @@ def _mcp_error_from_transport( (e.g. an ``initialize`` protocol-version failure raised inside the task group) is routed through the session mapping. Never carries headers. + + Cancellation — re-raised, never translated: + A leaf that is a :class:`BaseException` but NOT an + :class:`Exception` (:class:`asyncio.CancelledError`, + :class:`KeyboardInterrupt`) is re-raised untouched. Such a leaf + is reachable when consumer cancellation races a real transport + error inside the anyio task group and both arrive wrapped in one + ``BaseExceptionGroup``; translating it would surface + ``MCPError("MCP transport error: CancelledError")`` and break the + SDK's cancellation contract. This is the same + classification-ordering rule the ``on_tool_error`` hook documents + (and loop.py's classifier follows): cancellation always wins over + error translation. """ leaves = _iter_leaf_exceptions(exc) + for leaf in leaves: + if not isinstance(leaf, Exception): + raise leaf # Prefer a nested protocol error (a 4xx/handshake failure can surface # as an McpError inside the transport task group). for leaf in leaves: @@ -864,12 +925,13 @@ async def _apply_tool_error_hook(self, err: _MCPCallError, *, tool_name: str) -> except asyncio.CancelledError: raise except Exception as exc: - # DELIBERATE DIVERGENCE from `invoke_hook`, which logs `str(exc)`: - # a hook screening untrusted server content may embed that content - # in its own exception message (`ValueError(f"bad: {content}")`), and - # this module's standing rule is that untrusted material never - # reaches a log line. Log the exception TYPE only — never its text, - # never `content`, never the message. + # Same discipline as `invoke_hook` (which now also logs the + # exception TYPE only): a hook screening untrusted server content + # may embed that content in its own exception message + # (`ValueError(f"bad: {content}")`), and this module's standing + # rule is that untrusted material never reaches a log line. Log + # the exception TYPE only — never its text, never `content`, + # never the message. _log.warning( "mcp.tool_error_hook_failed", tool_name=tool_name, diff --git a/src/fifty_agent_sdk/observability/hooks.py b/src/fifty_agent_sdk/observability/hooks.py index 0f04735..e1c8faf 100644 --- a/src/fifty_agent_sdk/observability/hooks.py +++ b/src/fifty_agent_sdk/observability/hooks.py @@ -39,8 +39,11 @@ A raising hook never breaks a run. :func:`invoke_hook` catches :class:`Exception`, logs a ``WARNING`` (event ``hook.invoke_failed``) under the fixed ``fifty_agent_sdk.observability`` logger, and swallows it. - :class:`asyncio.CancelledError` is the one exception re-raised - untouched so consumer cancellation still propagates. + The log line carries the hook name and the exception TYPE only — never + ``str(exc)`` — because hook arguments are high-value (user message, + ``ChatRequest``, tool args) and a hook may embed them in its own + exception message. :class:`asyncio.CancelledError` is the one exception + re-raised untouched so consumer cancellation still propagates. Hot-path latency Hooks are awaited INLINE. A slow ``on_iteration`` or ``on_llm_call`` @@ -125,7 +128,9 @@ class Hooks: ``duration_ms`` is a monotonic wall measurement of the whole run. ``error`` is typed ``BaseException | None`` — non-``None`` ONLY when an exception terminated the run (a - :class:`~fifty_agent_sdk.errors.StateStoreError`, or a surfaced + :class:`~fifty_agent_sdk.errors.StateStoreError`, a fatal + :class:`~fifty_agent_sdk.errors.AgentSdkError` escaping the loop, + or a surfaced :class:`asyncio.CancelledError` which is a :class:`BaseException`, not an :class:`Exception`, hence the wider type). A loop-internal failure surfaces an @@ -154,9 +159,11 @@ class Hooks: tool's start and terminal events. Runner tier. on_error: ``(session_id, error, context) -> Any``. Fires on a loop-internal failure (a synthesized exception built from the - :class:`~fifty_agent_sdk.streaming.ErrorEvent`) and on a state-store + :class:`~fifty_agent_sdk.streaming.ErrorEvent`), on a state-store durability failure (the caught - :class:`~fifty_agent_sdk.errors.StateStoreError`). ``context`` is a + :class:`~fifty_agent_sdk.errors.StateStoreError`), and on a fatal + :class:`~fifty_agent_sdk.errors.AgentSdkError` escaping the loop + (the exception itself, before it is re-raised). ``context`` is a structured detail dict. Runner tier. """ @@ -201,7 +208,10 @@ async def invoke_hook( untouched (consumer cancellation must propagate); any other :class:`Exception` is logged at ``WARNING`` (event ``hook.invoke_failed``) under the ``fifty_agent_sdk.observability`` logger - and swallowed — a hook failure NEVER aborts a run. + and swallowed — a hook failure NEVER aborts a run. The log line + carries the exception TYPE only, never ``str(exc)``: hooks receive + high-value objects (the raw user message, the full ``ChatRequest``, + tool args) and a hook may embed them in its exception message. Args: hook: The callable to invoke, or ``None`` for a no-op. @@ -221,11 +231,17 @@ async def invoke_hook( except asyncio.CancelledError: raise except Exception as exc: + # Log the exception TYPE only — never str(exc). Hooks receive + # high-value objects (on_run_start gets the raw user_message, + # on_llm_call the full ChatRequest, on_tool_start the tool args), + # and a hook may embed them in its own exception message + # (`ValueError(f"unexpected: {request.messages}")`); logging + # str(exc) would dump that content into a WARNING line. This + # matches the mcp.client tool-error-hook logging discipline. _log.warning( "hook.invoke_failed", hook_name=hook_name, error_type=type(exc).__name__, - error_message=str(exc), ) diff --git a/src/fifty_agent_sdk/parser/json_mode.py b/src/fifty_agent_sdk/parser/json_mode.py index cf65ee0..7270536 100644 --- a/src/fifty_agent_sdk/parser/json_mode.py +++ b/src/fifty_agent_sdk/parser/json_mode.py @@ -114,7 +114,10 @@ class JsonModeParser: * ``error_phase="empty_completion"`` — empty/whitespace-only input. * ``error_phase="json_decode"`` — both the strict and the recovery pass - failed to produce valid JSON. + failed to produce valid JSON, or the input nested deeper than the + interpreter recursion limit (``json.loads`` raises + :class:`RecursionError` there, translated into :class:`ParserError` + so the loop's ``ParserError`` contract holds). * ``error_phase="schema_validation"`` — JSON decoded but the envelope does not match the schema (unknown key, wrong ``action`` value, missing required field for the chosen action). @@ -140,10 +143,31 @@ def parse(self, completion: str) -> ParseResult: # ------------------------------------------------------------------ # def _load_json(self, completion: str) -> Any: - """Strict-then-recover JSON decode. Raises on total failure.""" + """Strict-then-recover JSON decode. Raises on total failure. + + :class:`RecursionError` from a pathologically nested input is caught + at BOTH decode attempts and translated into :class:`ParserError` + (``error_phase="json_decode"``): the :class:`Parser` protocol mandates + ``ParserError`` on malformed input, and the loop catches only + ``ParserError`` — a raw ``RecursionError`` would escape the async + generator with no ``ErrorEvent`` and no terminal ``FinalEvent``. + """ try: return json.loads(completion.strip()) - except json.JSONDecodeError as first_err: + except RecursionError as depth_err: + raise ParserError( + "could not decode JSON envelope: nesting depth exceeds the recursion limit", + context={ + "parser": "JsonModeParser", + "error_phase": "json_decode", + "completion_excerpt": completion[:_MAX_EXCERPT], + "cause": repr(depth_err), + }, + ) from depth_err + except ValueError as first_err: + # CPython also raises bare ValueError for implementation limits + # such as oversized integer literals. Keep it in the established + # decode phase instead of parsing version-specific error text. recovered = _strip_code_fences(completion) if recovered is None: raise ParserError( @@ -157,7 +181,17 @@ def _load_json(self, completion: str) -> Any: ) from first_err try: return json.loads(recovered) - except json.JSONDecodeError as second_err: + except RecursionError as depth_err: + raise ParserError( + "could not decode JSON envelope: nesting depth exceeds the recursion limit", + context={ + "parser": "JsonModeParser", + "error_phase": "json_decode", + "completion_excerpt": completion[:_MAX_EXCERPT], + "cause": repr(depth_err), + }, + ) from depth_err + except ValueError as second_err: raise ParserError( "could not decode JSON envelope after fence recovery", context={ @@ -186,7 +220,13 @@ def _validate(self, raw: Any, completion: str) -> _RawEnvelope: def _to_parse_result(self, env: _RawEnvelope, completion: str) -> ParseResult: """Convert a validated envelope into the public :data:`ParseResult`.""" if env.action == "tool": - if not env.tool_name: + # Strip before the emptiness check (and before emitting the + # ToolCall): a whitespace-only name must take the same + # schema_validation path as an empty one, and a padded name + # must reach the registry in the same stripped form the prose + # parser produces (prose_mode strips the Action: header). + tool_name = env.tool_name.strip() if env.tool_name is not None else None + if not tool_name: raise ParserError( "action='tool' requires non-empty tool_name", context={ @@ -197,7 +237,7 @@ def _to_parse_result(self, env: _RawEnvelope, completion: str) -> ParseResult: }, ) tool_call = ToolCall( - name=env.tool_name, + name=tool_name, args=env.tool_args if env.tool_args is not None else {}, ) return ThoughtAction(thought=env.thought, tool_call=tool_call) diff --git a/src/fifty_agent_sdk/parser/prose_mode.py b/src/fifty_agent_sdk/parser/prose_mode.py index ca7511c..1531c6f 100644 --- a/src/fifty_agent_sdk/parser/prose_mode.py +++ b/src/fifty_agent_sdk/parser/prose_mode.py @@ -64,7 +64,9 @@ class ProseModeParser: * ``error_phase="header_match"`` — neither pattern matched. * ``error_phase="action_input_decode"`` — the ``Action Input:`` body could not be parsed as JSON, even after the shared fence-stripping - recovery pass. + recovery pass, or it nested deeper than the interpreter recursion + limit (:class:`RecursionError` from ``json.loads``, translated into + :class:`ParserError` so the loop's ``ParserError`` contract holds). """ def parse(self, completion: str) -> ParseResult: @@ -118,10 +120,29 @@ def _parse_final(self, match: re.Match[str]) -> FinalAnswer: return FinalAnswer(thought=thought, content=answer) def _decode_action_input(self, body: str, completion: str) -> dict[str, object]: - """Decode the ``Action Input:`` body as JSON, with one fence retry.""" + """Decode the ``Action Input:`` body as JSON, with one fence retry. + + :class:`RecursionError` from a pathologically nested body is caught at + BOTH decode attempts and translated into :class:`ParserError` + (``error_phase="action_input_decode"``) — the same ``ParserError``-only + contract the JSON-mode parser upholds on its own decode path. + """ try: decoded = json.loads(body) - except json.JSONDecodeError as first_err: + except RecursionError as depth_err: + raise ParserError( + "could not decode Action Input JSON: nesting depth exceeds the recursion limit", + context={ + "parser": "ProseModeParser", + "error_phase": "action_input_decode", + "completion_excerpt": completion[:_MAX_EXCERPT], + "cause": repr(depth_err), + }, + ) from depth_err + except ValueError as first_err: + # json.loads may raise bare ValueError for interpreter limits + # (for example oversized integers). The public contract classifies + # every decode failure identically without inspecting error text. recovered = _strip_code_fences(body) if recovered is None: raise ParserError( @@ -135,7 +156,17 @@ def _decode_action_input(self, body: str, completion: str) -> dict[str, object]: ) from first_err try: decoded = json.loads(recovered) - except json.JSONDecodeError as second_err: + except RecursionError as depth_err: + raise ParserError( + "could not decode Action Input JSON: nesting depth exceeds the recursion limit", + context={ + "parser": "ProseModeParser", + "error_phase": "action_input_decode", + "completion_excerpt": completion[:_MAX_EXCERPT], + "cause": repr(depth_err), + }, + ) from depth_err + except ValueError as second_err: raise ParserError( "could not decode Action Input JSON after fence recovery", context={ diff --git a/src/fifty_agent_sdk/runner.py b/src/fifty_agent_sdk/runner.py index 9c84639..bac08a7 100644 --- a/src/fifty_agent_sdk/runner.py +++ b/src/fifty_agent_sdk/runner.py @@ -51,7 +51,9 @@ When an optional :class:`fifty_agent_sdk.audit.protocol.AuditSink` is wired in, the Runner emits an :class:`fifty_agent_sdk.audit.protocol.AuditEvent` at four points of every ``run()``: session start, each tool invocation - (args plus a bounded result summary), the final answer, and any error. + (argument metadata — sorted keys with per-value type names and lengths, + never the values — plus a bounded result summary), the final answer, + and any error. Audit emission is best-effort and isolated from the run: a raising sink is caught by :meth:`_emit_audit`, logged at ``WARNING`` under the @@ -66,8 +68,9 @@ When an optional :class:`fifty_agent_sdk.observability.Hooks` is wired in, the Runner fires five of the seven hooks: ``on_run_start`` once at run start, ``on_tool_start`` / ``on_tool_end`` per tool invocation, - ``on_error`` on a loop or durability failure, and ``on_run_end`` once - from the ``finally`` block on EVERY exit path. The remaining two hooks + ``on_error`` on a loop-internal, durability, or fatal-SDK-error + failure, and ``on_run_end`` once from the ``finally`` block on EVERY + exit path. The remaining two hooks (``on_iteration``, ``on_llm_call``) are Loop-tier — the consumer must wire the SAME :class:`Hooks` instance into :class:`fifty_agent_sdk.loop. AgentLoop` as well. The Runner does NOT forward ``hooks`` into the loop; @@ -98,15 +101,26 @@ the load/append calls. A companion ``phase`` field names which site failed: ``"load"``, ``"persist_system"``, ``"persist_user"``, or ``"persist_assistant"``. - * ``"cancelled"`` — caller cancelled the consumer task or broke out - of the ``async for`` via ``aclose()``. + * ``"cancelled"`` — the consumer task was cancelled while the run was + in flight; :class:`asyncio.CancelledError` propagates untouched. + * ``"interrupted"`` — fallback for every other exit path not + attributable to the categories above: an unexpected non-SDK + exception escaping the loop, or the consumer breaking out of the + ``async for`` / closing the generator via ``aclose()`` (which + surfaces as :class:`GeneratorExit`, not cancellation). + * ``"sdk_error"`` — a fatal :class:`fifty_agent_sdk.errors.AgentSdkError` + subclass escaped the loop as an exception (for example an + :class:`~fifty_agent_sdk.errors.MCPError` re-raised by the tool + registry). The Runner audits the error, fires ``on_error``, passes + the exception to ``on_run_end``, and re-raises it. """ from __future__ import annotations import asyncio import time -from collections.abc import AsyncIterator +from collections import deque +from collections.abc import AsyncIterator, Sized from datetime import UTC, datetime from typing import Any, Final from uuid import uuid4 @@ -114,7 +128,7 @@ import structlog from fifty_agent_sdk.audit import AuditEvent, AuditSink -from fifty_agent_sdk.errors import StateStoreError +from fifty_agent_sdk.errors import AgentSdkError, StateStoreError from fifty_agent_sdk.llm.types import ChatMessage from fifty_agent_sdk.loop import AgentLoop from fifty_agent_sdk.observability import Hooks @@ -155,6 +169,30 @@ def _bounded_repr(value: object) -> str: return text[:_RESULT_SUMMARY_CAP] + _TRUNCATION_MARKER +def _args_metadata(args: dict[str, Any]) -> dict[str, Any]: + """Return a non-content summary of a tool's argument dict. + + Tool args routinely carry secrets and PII, and an audit payload is + persisted and logged verbatim (see the "no secrets in the payload" + contract on :class:`fifty_agent_sdk.audit.protocol.AuditEvent`), so the + ``tool_invocation`` payload MUST NOT embed argument values. The summary + keeps the argument keys (sorted, for determinism) and, per key, the + value's type name and its length when the value is sized (``len`` is + ``None`` otherwise) — enough for shape-level debugging without leaking + content. This mirrors the value-replacement discipline of + :func:`fifty_agent_sdk.mcp.client._redact_headers`, which keeps header + names and drops header values. + """ + summary: dict[str, Any] = {} + for key in sorted(args): + value = args[key] + summary[key] = { + "type": type(value).__name__, + "len": len(value) if isinstance(value, Sized) else None, + } + return summary + + class AgentRunner: """User-facing orchestrator that drives an :class:`AgentLoop` with state. @@ -328,7 +366,9 @@ async def _emit_audit( event_type: One of ``"session_start"``, ``"tool_invocation"``, ``"final_answer"``, ``"error"``. payload: Structured, event-specific detail (lengths/counts and - tool metadata only — never message or prompt content). + tool metadata only — never message or prompt content, and + never tool-argument VALUES: ``tool_invocation`` carries the + non-content summary built by :func:`_args_metadata`). """ if self._audit is None: return @@ -377,26 +417,37 @@ async def _invoke_hook(self, hook_name: str, *args: Any) -> None: @staticmethod def _tool_invocation_payload( event: ObservationEvent | ToolFailedEvent, - pending_action: ActionEvent | None, + args: dict[str, Any], pending_call: ToolStartedEvent | None, ) -> dict[str, Any]: """Build the ``payload`` for a ``tool_invocation`` audit event. - Correlates the terminal tool event with the - :class:`ActionEvent` that carried the ``args`` and the - :class:`ToolStartedEvent` that carried the ``call_id``. The loop is - strictly sequential, so the single pending slots are the correct - pair; they are still treated as optional for robustness. + ``args`` and ``pending_call`` are the per-call state the event loop + correlated under this call's ``call_id``: ``args`` was claimed FIFO + from the pending :class:`ActionEvent` queue when the matching + :class:`ToolStartedEvent` arrived (an :class:`ActionEvent` carries + no ``call_id``; both the single-call branch and the ``MultiAction`` + batch emit actions and starts in call order), and ``pending_call`` + is the :class:`ToolStartedEvent` popped under the same key. Both + are already resolved per-call by the caller — this function only + shapes the payload. ``result_summary`` is a bounded ``repr`` of the tool's output (on success) or the failure string (on a recoverable failure), capped so a large or binary result cannot bloat the audit row. + ``args`` is NEVER embedded verbatim: argument values routinely carry + secrets and PII, and the payload is persisted and logged as-is, so + the payload's ``"args"`` field is the non-content summary built by + :func:`_args_metadata` (sorted keys, per-value type names and + lengths — never values). + Args: event: The :class:`ObservationEvent` or :class:`ToolFailedEvent` that ended the tool call. - pending_action: The most recent :class:`ActionEvent`, if seen. - pending_call: The most recent :class:`ToolStartedEvent`, if seen. + args: The tool's argument dict, correlated per call; ``{}`` + when no :class:`ActionEvent` could be claimed for the call. + pending_call: The correlated :class:`ToolStartedEvent`, if seen. Returns: The structured ``payload`` dict for the audit event. @@ -410,7 +461,7 @@ def _tool_invocation_payload( return { "tool_name": event.tool_name, "call_id": (pending_call.call_id if pending_call is not None else event.call_id), - "args": (pending_action.args if pending_action is not None else {}), + "args": _args_metadata(args), "outcome": outcome, "result_summary": result_summary, } @@ -436,9 +487,18 @@ async def run(self, session_id: str, user_message: str) -> AsyncIterator[AgentEv parsed text on the safety paths. Otherwise skip — the fallback final answer is yielded but not committed. - On consumer cancellation (the consumer breaks out of the - ``async for`` loop): :class:`asyncio.CancelledError` propagates - untouched. The user message persisted in step 3 survives; no + On consumer cancellation (the consumer task is cancelled): + :class:`asyncio.CancelledError` propagates untouched. The user + message persisted in step 3 survives; no assistant message is + persisted. + + On a fatal :class:`fifty_agent_sdk.errors.AgentSdkError` escaping the + loop (for example an :class:`~fifty_agent_sdk.errors.MCPError` + re-raised by the tool registry — :meth:`AgentLoop.run` documents + that non-recoverable SDK errors propagate): the Runner emits the + ``error`` audit event, fires ``on_error`` with the exception, + records it for ``on_run_end``, sets ``terminated_by="sdk_error"``, + and re-raises so the exception still reaches the caller. No assistant message is persisted. On :class:`fifty_agent_sdk.errors.StateStoreError` raised by the state @@ -452,13 +512,19 @@ async def run(self, session_id: str, user_message: str) -> AsyncIterator[AgentEv Yields: :class:`AgentEvent` values forwarded from the inner - :class:`AgentLoop` in monotonic ``sequence`` order. The - terminal event is always a :class:`FinalEvent`. + :class:`AgentLoop` in monotonic ``sequence`` order. On a clean + or loop-internal-failure termination the terminal event is a + :class:`FinalEvent`; a fatal :class:`AgentSdkError` escaping the + loop ends the stream by raising instead. Raises: fifty_agent_sdk.errors.StateStoreError: If any state-store operation fails. The error is logged before being re-raised. + fifty_agent_sdk.errors.AgentSdkError: Any fatal SDK error the + loop lets propagate. It is audited, reported to + ``on_error``, and passed to ``on_run_end`` before being + re-raised. asyncio.CancelledError: Propagated untouched from the loop or from the consumer's cancellation. @@ -474,8 +540,9 @@ async def run(self, session_id: str, user_message: str) -> AsyncIterator[AgentEv ``run_id``, ``terminated_by``, ``assistant_message_persisted``, ``event_count``, ``final_event_type``, ``phase``. ``terminated_by`` is one - of ``"final_answer"``, ``"error"``, - ``"state_store_error"``, or ``"cancelled"``. + of ``"final_answer"``, ``"error"``, ``"sdk_error"``, + ``"state_store_error"``, ``"cancelled"``, or + ``"interrupted"`` (see the module docstring). ``runner.persist_failed`` (ERROR): Emitted at each of the four state-store boundaries — load, system-prompt persist, user persist, assistant persist — when the underlying @@ -536,11 +603,14 @@ async def run(self, session_id: str, user_message: str) -> AsyncIterator[AgentEv await self._invoke_hook("on_run_start", session_id, user_message) # Initial value is "interrupted" — neutral and applies to any - # unexpected exit path (e.g. an exception escaping the loop that - # we did not catch explicitly). The dedicated + # unexpected exit path (e.g. a non-SDK exception escaping the loop + # that we did not catch explicitly, or the consumer closing the + # generator via ``aclose()``). The dedicated # ``except asyncio.CancelledError`` branch upgrades this to # ``"cancelled"`` ONLY when we can attribute exit to an actual - # task/consumer cancellation. + # task/consumer cancellation, and the ``except AgentSdkError`` + # branch upgrades it to ``"sdk_error"`` for a fatal SDK error + # escaping the loop. terminated_by = "interrupted" state_store_error_phase: str | None = None saw_error = False @@ -549,8 +619,9 @@ async def run(self, session_id: str, user_message: str) -> AsyncIterator[AgentEv event_count = 0 # `run_error` carries the exception that terminated the run, for the # `on_run_end` hook. It is set ONLY by an exception that escaped the - # run — a `StateStoreError` from a persist site or a surfaced - # `CancelledError`. Typed `BaseException | None` because + # run — a `StateStoreError` from a persist site, a fatal + # `AgentSdkError` escaping the loop, or a surfaced `CancelledError`. + # Typed `BaseException | None` because # `asyncio.CancelledError` is a `BaseException`, not an `Exception`. # A loop-internal failure surfaces an `ErrorEvent` (not a Python # exception) and is reported via `on_error`; for that path @@ -641,23 +712,25 @@ async def run(self, session_id: str, user_message: str) -> AsyncIterator[AgentEv # coexist in the prompt; that is intentional. loop_messages = list(history) # defensive copy for the loop - # Single-slot correlation for `tool_invocation` audit events. - # The ReACT loop is strictly sequential — one tool in flight at - # a time — so a single pending `ActionEvent` (carries `args`) - # and pending `ToolStartedEvent` (carries `call_id`) is - # sufficient; the paired Observation/ToolFailed clears them. + # Per-call correlation for `tool_invocation` audit events and the + # on_tool_start/on_tool_end hooks, keyed by `call_id`. The + # MultiAction branch of AgentLoop emits N ActionEvents, then N + # ToolStartedEvents, then N terminal events — all in call order — + # so single pending slots would mis-correlate a batch (the first + # terminal event would inherit the LAST call's pairing). + # ActionEvents carry no `call_id`, so their `args` are claimed + # FIFO from `pending_actions` as each ToolStartedEvent arrives + # (both branches emit actions and starts in call order), then + # stored under the started call's `call_id`. The paired terminal + # event pops its entry, leaving no residue. The single-call path + # is the N=1 case of the same flow and behaves exactly as before. # `last_error` holds the most recent ErrorEvent for the error # branch below. - pending_action: ActionEvent | None = None - pending_call: ToolStartedEvent | None = None + pending_actions: deque[ActionEvent] = deque() + pending_calls: dict[str, ToolStartedEvent] = {} + pending_args: dict[str, dict[str, Any]] = {} + tool_started_at: dict[str, float] = {} last_error: ErrorEvent | None = None - # Monotonic stamp set when a `ToolStartedEvent` is seen and - # diffed on the terminal tool event for the `on_tool_end` - # `duration_ms`. A single slot is sufficient — the ReACT loop - # runs one tool at a time — and it lives alongside the existing - # `pending_action`/`pending_call` single-slot correlation, not - # as a second correlation pass. - tool_started_at: float | None = None async for event in self._loop.run(loop_messages, session_id=session_id): event_count += 1 @@ -668,37 +741,39 @@ async def run(self, session_id: str, user_message: str) -> AsyncIterator[AgentEv final_text = event.text raw_final_completion = event.raw_completion elif isinstance(event, ActionEvent): - pending_action = event + pending_actions.append(event) elif isinstance(event, ToolStartedEvent): - pending_call = event - tool_started_at = time.perf_counter() + action = pending_actions.popleft() if pending_actions else None + pending_calls[event.call_id] = event + pending_args[event.call_id] = action.args if action is not None else {} + tool_started_at[event.call_id] = time.perf_counter() yield event # `on_tool_start` fires once the `ToolStartedEvent` is seen; - # `args` come from the correlated `pending_action`. Fired - # AFTER yielding so consumer delivery is never blocked. + # `args` were correlated under this call's `call_id` above. + # Fired AFTER yielding so consumer delivery is never blocked. if isinstance(event, ToolStartedEvent): await self._invoke_hook( "on_tool_start", session_id, event.tool_name, - pending_action.args if pending_action is not None else {}, + pending_args[event.call_id], ) # Emit `tool_invocation` AFTER yielding so consumer event # delivery is never blocked on audit latency. if isinstance(event, ObservationEvent | ToolFailedEvent): + pending_call = pending_calls.pop(event.call_id, None) + args = pending_args.pop(event.call_id, {}) + started_at = tool_started_at.pop(event.call_id, None) await self._emit_audit( session_id, "tool_invocation", - self._tool_invocation_payload(event, pending_action, pending_call), + self._tool_invocation_payload(event, args, pending_call), ) - # `on_tool_end` fires beside the audit emission, BEFORE - # the pending slots are cleared. `result` is the tool's - # output on success or the failure string on a - # recoverable failure. + # `on_tool_end` fires beside the audit emission. `result` + # is the tool's output on success or the failure string + # on a recoverable failure. tool_duration_ms = ( - (time.perf_counter() - tool_started_at) * 1000 - if tool_started_at is not None - else 0.0 + (time.perf_counter() - started_at) * 1000 if started_at is not None else 0.0 ) tool_result = ( event.result.output if isinstance(event, ObservationEvent) else event.error @@ -710,9 +785,6 @@ async def run(self, session_id: str, user_message: str) -> AsyncIterator[AgentEv tool_result, tool_duration_ms, ) - pending_action = None - pending_call = None - tool_started_at = None # ── PHASE 5: PERSIST ASSISTANT (SUCCESS PATH) ───────────── if not saw_error and final_text is not None: @@ -814,12 +886,44 @@ async def run(self, session_id: str, user_message: str) -> AsyncIterator[AgentEv }, ) except asyncio.CancelledError as exc: - # Caller cancelled the consumer task (or broke out via - # ``aclose()``). Attribute exit to cancellation and let the - # exception propagate untouched. + # The consumer task was cancelled while the run was in flight. + # Attribute exit to cancellation and let the exception propagate + # untouched. (``aclose()`` surfaces as GeneratorExit instead and + # leaves `terminated_by` at its "interrupted" fallback.) terminated_by = "cancelled" run_error = exc raise + except AgentSdkError as exc: + if isinstance(exc, StateStoreError): + # Persist-site failures already emitted their `error` audit + # event and fired `on_error` at the site; the `finally` + # block upgrades `terminated_by` to "state_store_error". + raise + # A fatal SDK error escaped the loop (e.g. an MCPError re-raised + # by the tool registry — AgentLoop.run documents that + # non-recoverable SDK errors propagate). Audit it and fire + # `on_error` BEFORE re-raising, mirroring the persist-site + # convention (the `finally` block must not await the sink while + # an exception is in flight); `run_error` hands it to + # `on_run_end`. + terminated_by = "sdk_error" + run_error = exc + await self._emit_audit( + session_id, + "error", + { + "run_id": run_id, + "error_type": type(exc).__name__, + "error_message": str(exc), + }, + ) + await self._invoke_hook( + "on_error", + session_id, + exc, + {"error_type": type(exc).__name__, **dict(exc.context)}, + ) + raise finally: if state_store_error_phase is not None: terminated_by = "state_store_error" @@ -836,9 +940,10 @@ async def run(self, session_id: str, user_message: str) -> AsyncIterator[AgentEv ) # `on_run_end` fires on EVERY exit path. `run_error` is # non-`None` only for an exception that escaped the run (a - # `StateStoreError` or the surfaced `CancelledError`); a - # loop-internal `terminated_by == "error"` keeps it `None` - # (`on_error` already fired for that). Awaiting a hook in + # `StateStoreError`, a fatal `AgentSdkError` from the loop, or + # the surfaced `CancelledError`); a loop-internal + # `terminated_by == "error"` keeps it `None` (`on_error` already + # fired for that). Awaiting a hook in # `finally` is safe: `_invoke_hook`/`invoke_hook` swallow every # `Exception` and re-raise only `CancelledError`, so a raising # `on_run_end` cannot mask an in-flight `StateStoreError`. diff --git a/src/fifty_agent_sdk/state/redis.py b/src/fifty_agent_sdk/state/redis.py index 3d2e230..e188541 100644 --- a/src/fifty_agent_sdk/state/redis.py +++ b/src/fifty_agent_sdk/state/redis.py @@ -35,21 +35,28 @@ sessions read as the trunk and gain the extra keys only when first forked. TTL semantics - When ``ttl_seconds`` is a positive integer, every :meth:`append` re-issues + When ``ttl_seconds`` is a positive integer, every state mutation re-issues ``EXPIRE`` across ALL of the session's keys (trunk, every fork list, the registry, and the active pointer) so the whole session's expiry window slides forward together — a "hot session stays alive" cache, and a fork's - parent line never expires out from under it. When ``ttl_seconds`` is - ``None`` no ``EXPIRE`` is ever issued and the session is durable until - :meth:`delete`. :meth:`get_messages` NEVER sets or refreshes a TTL — + parent line never expires out from under it. ``append``, ``fork``, + ``switch_branch``, and ``truncate_after`` share one optimistic transaction: + registry, active pointer, and message lists are watched, the mutation and + every ``EXPIRE`` execute together, and conflicts retry from a fresh snapshot + up to a bounded limit. When ``ttl_seconds`` is ``None`` no ``EXPIRE`` is + ever issued and the session is durable until :meth:`delete`. A non-positive + ``ttl_seconds`` is rejected with + :class:`ValueError` at construction: ``EXPIRE`` with ``0`` deletes a key + immediately, so accepting it would wipe every session key on the first + append. :meth:`get_messages` NEVER sets or refreshes a TTL — reading a session does not keep it alive. Atomicity - :meth:`append` issues ``RPUSH`` and (when applicable) ``EXPIRE`` inside a - single ``MULTI``/``EXEC`` transaction via a ``transaction=True`` pipeline. - A concurrent reader therefore sees either the pre-append or post-append - state of the list, never a half-written one — satisfying the - :class:`StateStore` ``append``-vs-read atomicity invariant. + Every state mutation and its optional whole-session TTL refresh execute in + one ``MULTI``/``EXEC`` transaction. Optimistic ``WATCH`` covers the branch + registry, active pointer, and all enumerated message lists; sustained + contention fails after three attempts through the normal + :class:`StateStoreError` wrapping contract. Error wrapping contract Every public method wraps :class:`redis.exceptions.RedisError` (the @@ -83,14 +90,16 @@ import json import uuid +from collections.abc import Awaitable, Callable +from dataclasses import dataclass from datetime import UTC, datetime -from typing import Any, Final, cast +from typing import Any, Final, Generic, TypeVar, cast import structlog try: import redis.asyncio as aioredis - from redis.exceptions import RedisError + from redis.exceptions import RedisError, WatchError except ImportError as exc: # pragma: no cover - exercised via importlib in tests raise ImportError( "fifty_agent_sdk.state.redis requires redis-py. Install with: pip install 'fifty-agent-sdk[redis]'" @@ -127,6 +136,31 @@ and SQL have no such constraint — they do not derive structured keys. """ +_MAX_MUTATION_RETRIES: Final[int] = 3 +"""Maximum optimistic-lock attempts for one session mutation (BR-015).""" + +_T = TypeVar("_T") + + +@dataclass(frozen=True) +class _SessionSnapshot: + """Consistent session metadata captured under Redis ``WATCH``.""" + + exists: bool + active: str + registry: dict[str, dict[str, Any]] + branch_map: dict[str, tuple[str | None, int | None]] + keys: tuple[str, ...] + + +@dataclass(frozen=True) +class _MutationPlan(Generic[_T]): + """Commands and result prepared from one watched session snapshot.""" + + apply: Callable[[Any], None] | None + result: _T + created_keys: tuple[str, ...] = () + def _now() -> datetime: """Current timezone-aware UTC time (branch creation stamp).""" @@ -174,16 +208,14 @@ class RedisStateStore: order. See the module docstring's "Key layout" section. TTL model: - With a positive ``ttl_seconds`` every :meth:`append` slides the - session's expiry window forward (``EXPIRE`` re-issued on each - write) — a hot-session cache. With ``ttl_seconds=None`` the list - never expires. :meth:`get_messages` never touches the TTL. + With a positive ``ttl_seconds`` every successful mutation slides every + extant session key's expiry window forward together — a hot-session + cache. With ``ttl_seconds=None`` no expiry command is emitted. + :meth:`get_messages` never touches the TTL. Atomicity: - :meth:`append` runs ``RPUSH`` + ``EXPIRE`` in a single - ``MULTI``/``EXEC`` transaction, so a concurrent reader sees only - pre- or post-append state — the :class:`StateStore` atomicity - invariant. + Mutations and their whole-session TTL refresh run in one watched + ``MULTI``/``EXEC`` transaction with bounded conflict retry. Connection ownership: The store owns the connection pool created from the URL. @@ -246,13 +278,19 @@ def __init__( key_prefix: Namespace prepended to each ``session_id`` to form the Redis key. Defaults to :data:`_DEFAULT_KEY_PREFIX` (``"fifty_agent_sdk:state:"``). - ttl_seconds: Per-session time-to-live, in seconds. When a - positive integer, every :meth:`append` re-issues an - ``EXPIRE`` so the session's expiry window slides forward - on each write. When ``None`` (the default), no ``EXPIRE`` - is ever issued and the session list is durable until - :meth:`delete`. + ttl_seconds: Per-session time-to-live, in seconds. When positive, + every mutation re-issues ``EXPIRE`` across all session keys so + the window slides forward as one unit. When ``None`` (the + default), no ``EXPIRE`` is issued. + + Raises: + ValueError: If ``ttl_seconds`` is not positive. ``EXPIRE`` + with a non-positive value deletes a key immediately, so + accepting ``0`` or less would wipe every session key on + the first append. """ + if ttl_seconds is not None and ttl_seconds <= 0: + raise ValueError(f"ttl_seconds must be a positive integer or None, got {ttl_seconds!r}") # ``decode_responses=True`` makes list members come back as ``str`` # (rather than ``bytes``), ready to hand straight to # ``ChatMessage.model_validate_json``. @@ -453,14 +491,96 @@ async def _session_keys(self, session_id: str) -> list[str]: keys.extend(self._msgs_key(session_id, fid) for fid in fork_ids if fid != TRUNK_BRANCH_ID) return keys - async def _ensure_trunk(self, session_id: str) -> None: - """Idempotently record the trunk in the registry with a creation stamp - (lazy, on first fork) via ``HSETNX``.""" - meta = json.dumps( - {"parent_branch_id": None, "forked_from_sequence": None, "created_at": _now_iso()} + async def _mutation_snapshot(self, pipe: Any, session_id: str) -> _SessionSnapshot: # noqa: ANN401 + """Read and watch every key needed by a session mutation. + + The registry is watched before it is enumerated. Any concurrent fork + changes that enumeration and therefore invalidates ``EXEC``. The active + pointer and all known message lists are watched as well, making routing, + validation, mutation, and the subsequent whole-session TTL refresh one + optimistic transaction. + """ + trunk_key = self._msgs_key(session_id, TRUNK_BRANCH_ID) + branches_key = self._branches_key(session_id) + active_key = self._active_key(session_id) + base_keys = (trunk_key, branches_key, active_key) + await pipe.watch(*base_keys) + + raw_registry = await pipe.hgetall(branches_key) + raw_active = await pipe.get(active_key) + exists = bool(await pipe.exists(*base_keys)) + registry: dict[str, dict[str, Any]] = { + str(branch_id): json.loads(meta) for branch_id, meta in raw_registry.items() + } + if TRUNK_BRANCH_ID not in registry: + registry[TRUNK_BRANCH_ID] = { + "parent_branch_id": None, + "forked_from_sequence": None, + "created_at": None, + } + active = str(raw_active) if raw_active is not None else TRUNK_BRANCH_ID + if active not in registry: + active = TRUNK_BRANCH_ID + + fork_keys = tuple( + self._msgs_key(session_id, branch_id) + for branch_id in registry + if branch_id != TRUNK_BRANCH_ID ) - await cast( - "Any", self._client.hsetnx(self._branches_key(session_id), TRUNK_BRANCH_ID, meta) + if fork_keys: + await pipe.watch(*fork_keys) + keys = tuple(dict.fromkeys((*base_keys, *fork_keys))) + return _SessionSnapshot( + exists=exists, + active=active, + registry=registry, + branch_map=self._branch_map(registry), + keys=keys, + ) + + async def _mutate_session( + self, + session_id: str, + prepare: Callable[[Any, _SessionSnapshot], Awaitable[_MutationPlan[_T]]], + ) -> _T: + """Apply one session mutation and slide every session key's TTL atomically. + + A ``WatchError`` retries from a fresh registry/active/key snapshot. The + retry count is deliberately bounded so sustained contention surfaces + through the existing ``RedisError`` -> ``StateStoreError`` translation. + With ``ttl_seconds=None`` the same mutation transaction runs but queues + no expiry commands. + """ + last_error: WatchError | None = None + for _attempt in range(_MAX_MUTATION_RETRIES): + try: + pipeline = cast("Any", self._client.pipeline(transaction=True)) + async with pipeline as pipe: + snapshot = await self._mutation_snapshot(pipe, session_id) + plan = await prepare(pipe, snapshot) + if plan.apply is None: + await pipe.unwatch() + return plan.result + + pipe.multi() + plan.apply(pipe) + ttl = self._ttl_seconds + if ttl is not None: + ttl_keys = dict.fromkeys((*snapshot.keys, *plan.created_keys)) + for key in ttl_keys: + pipe.expire(key, ttl) + await pipe.execute() + return plan.result + except WatchError as exc: + last_error = exc + assert last_error is not None + raise last_error + + @staticmethod + def _trunk_meta() -> str: + """Serialize the lazy trunk-registry entry queued by every write.""" + return json.dumps( + {"parent_branch_id": None, "forked_from_sequence": None, "created_at": _now_iso()} ) async def get_messages( @@ -545,29 +665,26 @@ async def append(self, session_id: str, message: ChatMessage) -> None: """ payload = message.model_dump_json() try: - active = await self._get_active(session_id) - # Record the trunk in the registry on first write so the session - # durably "exists" even after its trunk list is truncated to empty - # (an empty Redis list auto-deletes its key). - await self._ensure_trunk(session_id) - key = self._msgs_key(session_id, active) - ttl = self._ttl_seconds - ttl_keys: list[str] = [] - if ttl is not None: - ttl_keys = await self._session_keys(session_id) - if key not in ttl_keys: - ttl_keys.append(key) - async with self._client.pipeline(transaction=True) as pipe: - pipe.rpush(key, payload) - if ttl is not None: - for k in ttl_keys: - pipe.expire(k, ttl) - await pipe.execute() + active = TRUNK_BRANCH_ID + + async def prepare(_pipe: Any, snapshot: _SessionSnapshot) -> _MutationPlan[None]: + nonlocal active + active = snapshot.active + key = self._msgs_key(session_id, active) + trunk_meta = self._trunk_meta() + + def apply(pipe: Any) -> None: + pipe.hsetnx(self._branches_key(session_id), TRUNK_BRANCH_ID, trunk_meta) + pipe.rpush(key, payload) + + return _MutationPlan(apply=apply, result=None, created_keys=(key,)) + + await self._mutate_session(session_id, prepare) _log.debug( "redis_state_store.append", session_id=session_id, branch_id=active, - ttl=ttl, + ttl=self._ttl_seconds, ) except RedisError as exc: raise _wrap_state_store_error(exc, session_id=session_id, operation="append") from exc @@ -607,38 +724,44 @@ async def fork(self, session_id: str, from_sequence: int) -> str: Records a new entry in the branch-registry hash whose parent is the active branch. The new branch's own list is created lazily on its - first :meth:`append`. Does NOT change the active head. + first :meth:`append`. Does NOT change the active head. The registry + write and optional whole-session TTL refresh are one transaction. Raises: ValueError: If the session is unknown, or ``from_sequence`` is outside ``0..head`` of the active branch. fifty_agent_sdk.errors.StateStoreError: On backend failure. """ + new_id = uuid.uuid4().hex try: - if not await self._session_exists(session_id): - raise ValueError(f"cannot fork unknown session {session_id!r}") - # Stamp the trunk's creation time so list_branches is stable. - await self._ensure_trunk(session_id) - registry = await self._load_registry(session_id) - branch_map = self._branch_map(registry) - active = await self._get_active(session_id) - if active not in branch_map: - active = TRUNK_BRANCH_ID - head = await self._materialized_len(session_id, active, branch_map) - if not 0 <= from_sequence <= head: - raise ValueError( - f"from_sequence={from_sequence} out of range 0..{head} " - f"for active branch {active!r} of session {session_id!r}" + + async def prepare(_pipe: Any, snapshot: _SessionSnapshot) -> _MutationPlan[str]: + if not snapshot.exists: + raise ValueError(f"cannot fork unknown session {session_id!r}") + head = await self._materialized_len( + session_id, snapshot.active, snapshot.branch_map ) - new_id = uuid.uuid4().hex - meta = json.dumps( - { - "parent_branch_id": active, - "forked_from_sequence": from_sequence, - "created_at": _now_iso(), - } - ) - await cast("Any", self._client.hset(self._branches_key(session_id), new_id, meta)) + if not 0 <= from_sequence <= head: + raise ValueError( + f"from_sequence={from_sequence} out of range 0..{head} " + f"for active branch {snapshot.active!r} of session {session_id!r}" + ) + meta = json.dumps( + { + "parent_branch_id": snapshot.active, + "forked_from_sequence": from_sequence, + "created_at": _now_iso(), + } + ) + trunk_meta = self._trunk_meta() + + def apply(pipe: Any) -> None: + pipe.hsetnx(self._branches_key(session_id), TRUNK_BRANCH_ID, trunk_meta) + pipe.hset(self._branches_key(session_id), new_id, meta) + + return _MutationPlan(apply=apply, result=new_id) + + await self._mutate_session(session_id, prepare) _log.debug("redis_state_store.fork", session_id=session_id, branch_id=new_id) return new_id except RedisError as exc: @@ -690,21 +813,33 @@ async def list_branches(self, session_id: str) -> list[BranchInfo]: async def switch_branch(self, session_id: str, branch_id: str) -> None: """Set the session's active head to ``branch_id``. + The pointer write and optional whole-session TTL refresh are one + optimistic transaction, so the pointer cannot outlive its messages. + Raises: ValueError: If ``branch_id`` does not exist for this session. fifty_agent_sdk.errors.StateStoreError: On backend failure. """ try: - if not await self._session_exists(session_id): - raise ValueError( - f"branch_id={branch_id!r} does not exist for session {session_id!r}" - ) - registry = await self._load_registry(session_id) - if branch_id not in registry: - raise ValueError( - f"branch_id={branch_id!r} does not exist for session {session_id!r}" + + async def prepare(_pipe: Any, snapshot: _SessionSnapshot) -> _MutationPlan[None]: + if not snapshot.exists or branch_id not in snapshot.registry: + raise ValueError( + f"branch_id={branch_id!r} does not exist for session {session_id!r}" + ) + trunk_meta = self._trunk_meta() + + def apply(pipe: Any) -> None: + pipe.hsetnx(self._branches_key(session_id), TRUNK_BRANCH_ID, trunk_meta) + pipe.set(self._active_key(session_id), branch_id) + + return _MutationPlan( + apply=apply, + result=None, + created_keys=(self._active_key(session_id),), ) - await cast("Any", self._client.set(self._active_key(session_id), branch_id)) + + await self._mutate_session(session_id, prepare) _log.debug( "redis_state_store.switch_branch", session_id=session_id, branch_id=branch_id ) @@ -725,28 +860,32 @@ async def truncate_after( is refreshed across all of its keys. """ try: - if not await self._session_exists(session_id): - return - target = branch_id if branch_id is not None else await self._get_active(session_id) - registry = await self._load_registry(session_id) - if target not in registry: - return - anchor = registry[target]["forked_from_sequence"] or 0 - # Own message at index i has materialized sequence anchor + 1 + i; - # keep the first ``N - anchor`` (those with sequence <= N). - keep = sequence - anchor - key = self._msgs_key(session_id, target) - ttl = self._ttl_seconds - ttl_keys = await self._session_keys(session_id) if ttl is not None else [] - async with self._client.pipeline(transaction=True) as pipe: - if keep <= 0: - pipe.ltrim(key, 1, 0) # start > end empties the list - else: - pipe.ltrim(key, 0, keep - 1) - if ttl is not None: - for k in ttl_keys: - pipe.expire(k, ttl) - await pipe.execute() + target = branch_id or TRUNK_BRANCH_ID + + async def prepare(_pipe: Any, snapshot: _SessionSnapshot) -> _MutationPlan[None]: + nonlocal target + if not snapshot.exists: + return _MutationPlan(apply=None, result=None) + target = branch_id if branch_id is not None else snapshot.active + if target not in snapshot.registry: + return _MutationPlan(apply=None, result=None) + anchor = snapshot.registry[target]["forked_from_sequence"] or 0 + # Own message at index i has materialized sequence anchor + 1 + i; + # keep the first ``N - anchor`` (those with sequence <= N). + keep = sequence - anchor + key = self._msgs_key(session_id, target) + trunk_meta = self._trunk_meta() + + def apply(pipe: Any) -> None: + pipe.hsetnx(self._branches_key(session_id), TRUNK_BRANCH_ID, trunk_meta) + if keep <= 0: + pipe.ltrim(key, 1, 0) # start > end empties the list + else: + pipe.ltrim(key, 0, keep - 1) + + return _MutationPlan(apply=apply, result=None) + + await self._mutate_session(session_id, prepare) _log.debug( "redis_state_store.truncate_after", session_id=session_id, diff --git a/src/fifty_agent_sdk/state/sql.py b/src/fifty_agent_sdk/state/sql.py index 2d1a6be..80b168d 100644 --- a/src/fifty_agent_sdk/state/sql.py +++ b/src/fifty_agent_sdk/state/sql.py @@ -85,6 +85,20 @@ :meth:`append` / :meth:`fork`. Review the generated revision before applying. +Additive columns & older databases + ``agent_messages.tool_calls`` (nullable JSON) was added after the + initial schema so native tool-calling assistant turns survive a + reload. Like the BR-004 additions it is **additive and zero-loss**: + ``alembic revision --autogenerate`` against :data:`sql_metadata` + emits a plain ``ADD COLUMN``, and pre-existing rows hold ``NULL``, + which :meth:`SqlStateStore.get_messages` maps back to + ``tool_calls=None``. The SDK still does NOT own migrations: a + database created by an older SDK version simply lacks the column + until the consumer applies the additive migration, and the ORM + ``SELECT`` names every column explicitly — so upgrade the schema + together with the SDK (an unmigrated database fails the read with + a wrapped ``OperationalError``, not with silent data loss). + Error wrapping contract Every public method wraps :class:`sqlalchemy.exc.SQLAlchemyError` (the SQLAlchemy base class) into @@ -159,7 +173,7 @@ ) from exc from fifty_agent_sdk.errors import StateStoreError -from fifty_agent_sdk.llm.types import ChatMessage +from fifty_agent_sdk.llm.types import ChatMessage, ToolCall from fifty_agent_sdk.state.protocol import TRUNK_BRANCH_ID, BranchInfo _log: Final = structlog.get_logger(__name__) @@ -270,9 +284,9 @@ class AgentSession(Base): class AgentMessage(Base): """One persisted :class:`ChatMessage` row. - Round-trips all four :class:`ChatMessage` fields (``role``, - ``content``, ``name``, ``tool_call_id``) plus a per-session - ``sequence`` and the surrogate ``id`` primary key. The + Round-trips all five :class:`ChatMessage` fields (``role``, + ``content``, ``name``, ``tool_call_id``, ``tool_calls``) plus a + per-session ``sequence`` and the surrogate ``id`` primary key. The ``sequence`` is the durable ordering key; ``id`` is only used for row identity. """ @@ -343,6 +357,22 @@ class AgentMessage(Base): tool_call_id: Mapped[str | None] = mapped_column(String(255), nullable=True) """Optional id echoed back on ``role="tool"`` replies.""" + tool_calls: Mapped[list[dict[str, Any]] | None] = mapped_column( + JSON().with_variant(JSONB(), "postgresql"), + nullable=True, + default=None, + ) + """Native tool invocations carried on an ``assistant`` turn (nullable). + + Holds :class:`ChatMessage.tool_calls` as a list of JSON objects in the + same shape the Redis backend persists via ``model_dump_json`` (one + ``{"name", "args", "id"}`` object per :class:`ToolCall`), so all three + backends round-trip identical payloads. ``JSONB`` on Postgres, + JSON-encoded ``TEXT`` on SQLite. ``NULL`` — including every row written + before this column existed — reads back as ``tool_calls=None``; see the + module docstring's "Additive columns" section for the migration story. + """ + created_at: Mapped[Any] = mapped_column( DateTime(timezone=True), nullable=False, @@ -669,6 +699,12 @@ async def get_messages( content=row.content, name=row.name, tool_call_id=row.tool_call_id, + # NULL (incl. pre-column legacy rows) -> None. + tool_calls=( + [ToolCall.model_validate(tc) for tc in row.tool_calls] + if row.tool_calls is not None + else None + ), ) ) messages = self._materialize_positional(branch_map, by_branch, target) @@ -761,6 +797,13 @@ async def append(self, session_id: str, message: ChatMessage) -> None: content=message.content, name=message.name, tool_call_id=message.tool_call_id, + # Same JSON shape the Redis backend persists via + # model_dump_json, so all backends agree on the wire. + tool_calls=( + [tc.model_dump(mode="json") for tc in message.tool_calls] + if message.tool_calls is not None + else None + ), ) ) # commit on session.begin() context exit diff --git a/src/fifty_agent_sdk/tools/_schema_refs.py b/src/fifty_agent_sdk/tools/_schema_refs.py new file mode 100644 index 0000000..d4c0b44 --- /dev/null +++ b/src/fifty_agent_sdk/tools/_schema_refs.py @@ -0,0 +1,137 @@ +"""Shared ``$ref`` inlining for tool-facing JSON Schemas. + +Both tool providers emit a :class:`fifty_agent_sdk.tools.protocol.ToolSchema`, +which carries only ``type`` / ``properties`` / ``required`` / +``additionalProperties`` — the loop copies exactly those four fields into the +system prompt and the native function-calling ``parameters`` envelope, and +provider function-calling APIs reject unknown top-level keys, so a ``$defs`` +block cannot be shipped alongside. Pydantic v2 (and some MCP servers) emit +nested models as ``{"$ref": "#/$defs/Foo"}`` plus a top-level ``$defs`` dict; +keeping only ``properties``/``required`` would hand the LLM a dangling +pointer to a definition that does not exist. :func:`inline_local_refs` +resolves those local references before the schema is emitted. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Final + +_DEFS_PREFIX: Final[str] = "#/$defs/" + +_MAX_EXPANSION_DEPTH: Final[int] = 32 +"""Cap on ``$ref`` expansion depth. + +Guards against pathological reference chains where distinct definitions +reference one another deeply. Genuine cycles are caught earlier and precisely +by the expansion stack in :func:`inline_local_refs`; broad or diamond fan-out +is bounded separately by :data:`_MAX_EXPANDED_NODES`. +""" + +_MAX_EXPANDED_NODES: Final[int] = 10_000 +"""Maximum resolver visits for one schema expansion. + +Depth alone cannot bound an acyclic schema whose references fan out. Each +dictionary, list, and scalar visit consumes one unit so untrusted schemas fail +before exponential expansion can exhaust process memory. +""" + + +def inline_local_refs(node: Any, defs: Mapping[str, Any]) -> Any: + """Return a deep-resolved copy of ``node`` with local ``$defs`` refs inlined. + + Every ``{"$ref": "#/$defs/"}`` mapping whose ```` is present + in ``defs`` is replaced by the (recursively resolved) definition; sibling + keys next to a ``$ref`` are preserved by merging them over the resolved + definition (JSON Schema 2020-12 sibling semantics, which Pydantic v2 + relies on for e.g. a field ``description``). References that cannot be + resolved locally — a missing ````, or a non-``#/$defs/`` URI — are + left untouched: they were dangling before and inlining must not invent + information the source schema did not carry. + + Args: + node: The schema fragment to resolve (dicts and lists are walked + recursively; scalars pass through). + defs: The top-level ``$defs`` mapping the refs resolve against. + + Returns: + A new structure; the inputs are never mutated. + + Raises: + ValueError: On a reference cycle (a def that transitively references + itself — a recursive model has no finite inline expansion), when + expansion exceeds the depth cap, or when the shared node budget is + exhausted. Limit and traversal messages never include schema + content; cycle diagnostics name the involved definition ids. + """ + budget = [_MAX_EXPANDED_NODES] + try: + return _resolve(node, defs, _stack=(), _depth=0, _budget=budget) + except RecursionError as exc: + # Ordinary dict/list nesting does not advance the $ref depth counter. + # Contain the interpreter's traversal limit at this untrusted-data + # boundary without exposing any schema content in the error message. + raise ValueError("schema traversal exceeded the safe nesting limit") from exc + + +def _resolve( + node: Any, + defs: Mapping[str, Any], + *, + _stack: tuple[str, ...], + _depth: int, + _budget: list[int], +) -> Any: + _budget[0] -= 1 + if _budget[0] < 0: + raise ValueError(f"$ref expansion exceeded the {_MAX_EXPANDED_NODES}-node budget") + if isinstance(node, dict): + ref = node.get("$ref") + if isinstance(ref, str) and ref.startswith(_DEFS_PREFIX): + name = ref[len(_DEFS_PREFIX) :] + target = defs.get(name) + if target is None: + return { + k: _resolve(v, defs, _stack=_stack, _depth=_depth, _budget=_budget) + for k, v in node.items() + } + if name in _stack: + cycle = " -> ".join([*_stack, name]) + raise ValueError( + f"recursive $ref cycle detected ({cycle}); a recursive model has no " + "finite inline expansion and cannot be emitted as a $defs-free schema" + ) + if _depth >= _MAX_EXPANSION_DEPTH: + raise ValueError( + f"$ref expansion exceeded {_MAX_EXPANSION_DEPTH} levels; the schema is " + "too deeply chained to inline" + ) + resolved = _resolve( + target, + defs, + _stack=(*_stack, name), + _depth=_depth + 1, + _budget=_budget, + ) + siblings = { + k: _resolve(v, defs, _stack=_stack, _depth=_depth, _budget=_budget) + for k, v in node.items() + if k != "$ref" + } + if not siblings: + return resolved + if isinstance(resolved, dict): + return {**resolved, **siblings} + return resolved + return { + k: _resolve(v, defs, _stack=_stack, _depth=_depth, _budget=_budget) + for k, v in node.items() + } + if isinstance(node, list): + return [ + _resolve(item, defs, _stack=_stack, _depth=_depth, _budget=_budget) for item in node + ] + return node + + +__all__ = ["inline_local_refs"] diff --git a/src/fifty_agent_sdk/tools/inproc_provider.py b/src/fifty_agent_sdk/tools/inproc_provider.py index eaedc8d..79fae3d 100644 --- a/src/fifty_agent_sdk/tools/inproc_provider.py +++ b/src/fifty_agent_sdk/tools/inproc_provider.py @@ -8,7 +8,11 @@ Rich types are supported via Pydantic's JSON Schema emitter: ``int | None``, ``list[str]``, ``dict[str, int]``, ``Literal[...]``, ``Enum``, nested ``BaseModel`` — anything Pydantic can validate it can also -emit a schema for. +emit a schema for. Nested-model ``#/$defs/...`` references are INLINED into +the emitted :class:`fifty_agent_sdk.tools.protocol.ToolSchema` (see +:mod:`fifty_agent_sdk.tools._schema_refs`) so the schema the LLM receives is +self-contained; a recursive model has no finite inline expansion and is +rejected at decoration time with a :class:`ValueError`. """ from __future__ import annotations @@ -22,6 +26,7 @@ from pydantic import BaseModel, ConfigDict, Field, ValidationError, create_model from pydantic_core import PydanticUndefined +from fifty_agent_sdk.tools._schema_refs import inline_local_refs from fifty_agent_sdk.tools.protocol import Tool, ToolResult, ToolSchema from fifty_agent_sdk.tools.registry import Registry @@ -110,6 +115,9 @@ async def get_customer(customer_id: str) -> dict[str, Any]: ``*args``/``**kwargs``, has a positional-only parameter, or has any unannotated parameter. The error message names the offending parameter. + ValueError: At decoration time if a parameter's model is recursive + (a ``$ref`` cycle): recursive models have no finite inline + expansion and cannot be emitted as a ``$defs``-free schema. """ def decorator(fn: ToolFn) -> Tool: @@ -210,11 +218,26 @@ def _schema_from_model(model: type[BaseModel]) -> ToolSchema: Pydantic emits ``{"type": "object", "properties": {...}, "required": [...]}`` for a plain model; these pass through. ``additionalProperties`` is forced to ``False`` to mirror the model's ``extra="forbid"`` config. + + Nested models come out of Pydantic v2 as ``{"$ref": "#/$defs/Foo"}`` plus + a top-level ``$defs`` dict. ``ToolSchema`` has no ``$defs`` slot and the + loop ships only the four function-calling fields (providers reject unknown + top-level keys), so the refs are INLINED via + :func:`fifty_agent_sdk.tools._schema_refs.inline_local_refs` — shipping the + ``$ref`` without its definition would hand the LLM a dangling pointer. + Decoration raises :class:`ValueError` when a recursive model has no finite + expansion, a reference chain exceeds the depth cap, or expansion exhausts + the shared node budget. Cycle failures name the cycle; limit failures carry + bounded content-free messages. """ raw = model.model_json_schema() + defs_raw = raw.get("$defs", {}) + defs: dict[str, Any] = dict(defs_raw) if isinstance(defs_raw, dict) else {} properties_raw = raw.get("properties", {}) required_raw = raw.get("required", []) - properties: dict[str, Any] = properties_raw if isinstance(properties_raw, dict) else {} + properties: dict[str, Any] = ( + inline_local_refs(properties_raw, defs) if isinstance(properties_raw, dict) else {} + ) required: list[str] = list(required_raw) if isinstance(required_raw, list) else [] return ToolSchema( type=str(raw.get("type", "object")), diff --git a/src/fifty_agent_sdk/tools/mcp_provider.py b/src/fifty_agent_sdk/tools/mcp_provider.py index 2371467..ba10f7d 100644 --- a/src/fifty_agent_sdk/tools/mcp_provider.py +++ b/src/fifty_agent_sdk/tools/mcp_provider.py @@ -46,6 +46,7 @@ from fifty_agent_sdk.mcp import MCPClient, MCPToolDef from fifty_agent_sdk.mcp.client import _MCPCallError +from fifty_agent_sdk.tools._schema_refs import inline_local_refs from fifty_agent_sdk.tools.protocol import ToolResult, ToolSchema from fifty_agent_sdk.tools.registry import Registry @@ -88,16 +89,25 @@ def _to_tool_schema(input_schema: dict[str, Any]) -> ToolSchema: logs a WARNING — the resulting tool will reject any args via Pydantic validation downstream, which is the desired defensive behavior. - Unknown top-level JSON-Schema keys (``$defs``, ``examples``, …) are - dropped: :class:`ToolSchema` is ``extra="forbid"`` and exists to - describe the parameter surface the LLM sees, not to round-trip the - full JSON-Schema vocabulary. + ``$defs`` handling: local ``#/$defs/...`` references inside + ``properties`` are INLINED via + :func:`fifty_agent_sdk.tools._schema_refs.inline_local_refs` before the + ``$defs`` block is discarded — :class:`ToolSchema` is ``extra="forbid"`` + and provider function-calling envelopes reject unknown top-level keys, so + the block cannot be shipped, and keeping a ``$ref`` without it would hand + the LLM a dangling pointer. A server schema whose refs cannot be inlined + (a reference cycle, an over-deep chain, or an exhausted node budget) takes the same defensive + fallback as a non-object schema: WARNING plus an empty object schema. + Other unknown top-level JSON-Schema keys (``examples``, …) are still + dropped: :class:`ToolSchema` exists to describe the parameter surface the + LLM sees, not to round-trip the full JSON-Schema vocabulary. """ raw_type = input_schema.get("type", "object") if raw_type != "object": _log.warning( "mcp.input_schema.non_object", - type=raw_type, + reason="non_object_type", + received_type=type(raw_type).__name__, ) return ToolSchema( type="object", @@ -105,9 +115,29 @@ def _to_tool_schema(input_schema: dict[str, Any]) -> ToolSchema: required=[], additionalProperties=False, ) + defs_raw = input_schema.get("$defs", {}) + defs: dict[str, Any] = dict(defs_raw) if isinstance(defs_raw, dict) else {} properties_raw = input_schema.get("properties", {}) required_raw = input_schema.get("required", []) - properties: dict[str, Any] = dict(properties_raw) if isinstance(properties_raw, dict) else {} + try: + properties: dict[str, Any] = ( + inline_local_refs(properties_raw, defs) if isinstance(properties_raw, dict) else {} + ) + except ValueError as exc: + # Untrusted server data takes the defensive fallback rather than + # raising out of adapter construction (which would abort the whole + # discover/attach batch, not just the offending tool's schema). + _log.warning( + "mcp.input_schema.unresolvable_refs", + reason="invalid_local_refs", + error_type=type(exc).__name__, + ) + return ToolSchema( + type="object", + properties={}, + required=[], + additionalProperties=False, + ) required: list[str] = [str(r) for r in required_raw] if isinstance(required_raw, list) else [] return ToolSchema( type="object", diff --git a/src/fifty_agent_sdk/tools/protocol.py b/src/fifty_agent_sdk/tools/protocol.py index 92d5dc5..157b7d1 100644 --- a/src/fifty_agent_sdk/tools/protocol.py +++ b/src/fifty_agent_sdk/tools/protocol.py @@ -38,7 +38,13 @@ class ToolSchema(BaseModel): properties: Mapping of parameter name to its JSON-Schema description. Pydantic's ``model_json_schema()`` populates this when the schema is derived from a function signature (see - :func:`fifty_agent_sdk.tools.inproc_provider.tool`). + :func:`fifty_agent_sdk.tools.inproc_provider.tool`). Local + ``#/$defs/...`` references are INLINED by the providers before the + schema is emitted (see + :mod:`fifty_agent_sdk.tools._schema_refs`) — the dict must be + self-contained because the loop ships it into provider + function-calling envelopes that reject unknown top-level keys, so + a ``$defs`` block cannot ride along. required: Names of parameters without defaults. Pydantic's schema emitter populates this automatically. additionalProperties: Always ``False`` — the schema is the complete diff --git a/tests/llm/test_openai_compat.py b/tests/llm/test_openai_compat.py index e2568f4..e2c67b7 100644 --- a/tests/llm/test_openai_compat.py +++ b/tests/llm/test_openai_compat.py @@ -7,6 +7,7 @@ from __future__ import annotations import json +import sys from typing import Any import httpx @@ -16,7 +17,7 @@ from fifty_agent_sdk.errors import LLMError from fifty_agent_sdk.llm.openai_compat import OpenAICompatibleClient from fifty_agent_sdk.llm.protocol import LLMClient -from fifty_agent_sdk.llm.types import ChatMessage, ChatRequest, ToolCall +from fifty_agent_sdk.llm.types import ChatMessage, ChatRequest, ToolCall, ToolChoiceFunction # --------------------------------------------------------------------------- # Helpers @@ -232,6 +233,28 @@ async def test_complete_omits_optional_fields_when_unset(httpx_mock: HTTPXMock) assert "response_format" not in body +async def test_build_body_sends_default_temperature(httpx_mock: HTTPXMock) -> None: + """The ``0.0`` default IS sent — the key is present unless explicitly ``None``.""" + httpx_mock.add_response(method="POST", url=ENDPOINT, json=_canonical_response()) + client = _make_client() + await client.complete(_basic_request()) + raw = httpx_mock.get_request() + assert raw is not None + body = json.loads(raw.read()) + assert body["temperature"] == 0.0 + + +async def test_build_body_omits_temperature_when_none(httpx_mock: HTTPXMock) -> None: + """``temperature=None`` removes the key entirely (reasoning-model providers).""" + httpx_mock.add_response(method="POST", url=ENDPOINT, json=_canonical_response()) + client = _make_client() + await client.complete(_basic_request(temperature=None)) + raw = httpx_mock.get_request() + assert raw is not None + body = json.loads(raw.read()) + assert "temperature" not in body + + @pytest.mark.parametrize( "upstream,expected", [ @@ -342,6 +365,33 @@ async def test_complete_malformed_arguments_raises_llm_error(httpx_mock: HTTPXMo assert exc.value.__cause__ is not None +async def test_complete_oversized_integer_arguments_raise_llm_error( + httpx_mock: HTTPXMock, +) -> None: + """BR-013 contains bare ValueError from native tool-call argument decoding.""" + digits = "9" * (sys.get_int_max_str_digits() + 1) + arguments = f'{{"n":{digits}}}' + payload = _canonical_response( + finish_reason="tool_calls", + tool_calls=[ + { + "id": "call_big", + "type": "function", + "function": {"name": "calculate", "arguments": arguments}, + } + ], + ) + httpx_mock.add_response(method="POST", url=ENDPOINT, json=payload) + client = _make_client() + with pytest.raises(LLMError) as exc: + await client.complete(_basic_request()) + assert str(exc.value) == "provider tool_call arguments is not valid JSON" + assert exc.value.context["type"] == "MalformedResponse" + assert exc.value.context["tool_call_id"] == "call_big" + assert len(str(exc.value.context["arguments_excerpt"])) <= 200 + assert type(exc.value.__cause__) is ValueError + + async def test_complete_non_object_arguments_raises_llm_error(httpx_mock: HTTPXMock) -> None: """A non-object JSON `arguments` (e.g. a bare array) is rejected.""" payload = _canonical_response( @@ -446,6 +496,23 @@ async def test_build_body_tool_choice_override(httpx_mock: HTTPXMock) -> None: assert body["tool_choice"] == "required" +async def test_build_body_tool_choice_dict_form_on_wire(httpx_mock: HTTPXMock) -> None: + """The specific-tool `tool_choice` object is emitted verbatim on the wire.""" + choice: ToolChoiceFunction = {"type": "function", "function": {"name": "search"}} + httpx_mock.add_response(method="POST", url=ENDPOINT, json=_canonical_response()) + client = _make_client() + await client.complete( + _basic_request( + tools=[{"type": "function", "function": {"name": "search"}}], + tool_choice=choice, + ) + ) + raw = httpx_mock.get_request() + assert raw is not None + body = json.loads(raw.read()) + assert body["tool_choice"] == {"type": "function", "function": {"name": "search"}} + + async def test_assistant_tool_calls_envelope_on_wire(httpx_mock: HTTPXMock) -> None: """An assistant turn carrying tool_calls serializes to the OpenAI envelope shape. @@ -792,6 +859,43 @@ async def test_stream_malformed_chunk_raises_llm_error(httpx_mock: HTTPXMock) -> pass +# --------------------------------------------------------------------------- +# Lifecycle: aclose() and the async context manager +# --------------------------------------------------------------------------- + + +async def test_aclose_closes_owned_client() -> None: + """An owned client (no injected http_client) is closed by aclose().""" + client = _make_client() + await client.aclose() + assert client._client.is_closed() + + +async def test_aclose_is_idempotent() -> None: + """A second ``aclose()`` is a no-op and MUST NOT raise.""" + client = _make_client() + await client.aclose() + await client.aclose() + assert client._client.is_closed() + + +async def test_aclose_does_not_close_injected_http_client() -> None: + """An injected ``http_client`` stays open after aclose() — the caller owns it.""" + injected = httpx.AsyncClient() + client = _make_client(http_client=injected) + await client.aclose() + assert not injected.is_closed + await injected.aclose() + + +async def test_async_context_manager_closes_owned_client_on_exit() -> None: + """``async with`` returns the client itself and aclose()s it on exit.""" + async with _make_client() as client: + assert isinstance(client, OpenAICompatibleClient) + assert not client._client.is_closed() + assert client._client.is_closed() + + # --------------------------------------------------------------------------- # Protocol conformance # --------------------------------------------------------------------------- diff --git a/tests/llm/test_types.py b/tests/llm/test_types.py index 2b57815..f6d7be3 100644 --- a/tests/llm/test_types.py +++ b/tests/llm/test_types.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import Any + import pytest from pydantic import ValidationError @@ -10,6 +12,7 @@ ChatRequest, ChatResponse, ToolCall, + ToolChoiceFunction, Usage, ) @@ -124,6 +127,16 @@ def test_chat_request_accepts_boundary_temperatures() -> None: assert req.temperature == t +def test_chat_request_accepts_none_temperature() -> None: + """``None`` is the explicit opt-out: the adapter omits the key from the body.""" + req = ChatRequest( + messages=[ChatMessage(role="user", content="hi")], + model="gpt-4o", + temperature=None, + ) + assert req.temperature is None + + def test_chat_request_rejects_zero_max_tokens() -> None: with pytest.raises(ValidationError): ChatRequest( @@ -170,6 +183,37 @@ def test_chat_request_accepts_tools_and_tool_choice() -> None: assert req.tool_choice == "auto" +def test_chat_request_accepts_tool_choice_dict_form() -> None: + """BR-014 named-tool typing imports and validates across the Python matrix.""" + choice: ToolChoiceFunction = {"type": "function", "function": {"name": "search"}} + req = ChatRequest( + messages=[ChatMessage(role="user", content="hi")], + model="gpt-4o", + tools=[{"type": "function", "function": {"name": "search"}}], + tool_choice=choice, + ) + assert req.tool_choice == {"type": "function", "function": {"name": "search"}} + + +@pytest.mark.parametrize( + "bad_choice", + [ + {"type": "function"}, # missing required "function" member + {"type": "not-a-function", "function": {"name": "search"}}, + {"function": {"name": "search"}}, # missing required "type" + {"type": "function", "function": {}}, # missing required "name" + 42, + ], +) +def test_chat_request_rejects_malformed_tool_choice(bad_choice: Any) -> None: + with pytest.raises(ValidationError): + ChatRequest( + messages=[ChatMessage(role="user", content="hi")], + model="gpt-4o", + tool_choice=bad_choice, + ) + + def test_chat_request_forbids_extra_fields() -> None: with pytest.raises(ValidationError): ChatRequest( diff --git a/tests/loop/test_loop.py b/tests/loop/test_loop.py index 177d376..3992034 100644 --- a/tests/loop/test_loop.py +++ b/tests/loop/test_loop.py @@ -124,6 +124,45 @@ async def _collect( return [event async for event in iterator] +# --------------------------------------------------------------------------- +# Construction-time validation +# --------------------------------------------------------------------------- + + +def test_stream_with_native_tools_enabled_raises_value_error() -> None: + """stream=True + native_tools_enabled=True is rejected at construction. + + A streamed completion carries no structured ``tool_calls`` (the adapter + drops streamed tool-call deltas), so the combination could never dispatch + a native tool call — it must fail fast at ``__init__`` with a ValueError + instead of silently burning the parser retry and ending on the fallback. + """ + llm = FakeLLMClient([make_response(_final_json("never reached"))]) + with pytest.raises(ValueError, match="stream=True"): + _make_loop(llm=llm, stream=True, safety=SafetyConfig(native_tools_enabled=True)) + + +@pytest.mark.parametrize( + ("stream", "native_tools_enabled"), + [ + (False, False), + (False, True), + (True, False), + ], +) +def test_other_stream_native_combinations_construct( + stream: bool, native_tools_enabled: bool +) -> None: + """Every other stream/native combination remains constructible.""" + llm = FakeLLMClient([make_response(_final_json("ok"))]) + loop = _make_loop( + llm=llm, + stream=stream, + safety=SafetyConfig(native_tools_enabled=native_tools_enabled), + ) + assert isinstance(loop, AgentLoop) + + # --------------------------------------------------------------------------- # Happy single-step # --------------------------------------------------------------------------- diff --git a/tests/mcp/test_client.py b/tests/mcp/test_client.py index 181b768..51d25ea 100644 --- a/tests/mcp/test_client.py +++ b/tests/mcp/test_client.py @@ -16,8 +16,11 @@ from __future__ import annotations +import asyncio import json import re +from collections.abc import AsyncIterator +from contextlib import AbstractAsyncContextManager, asynccontextmanager from typing import Any import httpx @@ -27,7 +30,7 @@ from fifty_agent_sdk.errors import MCPError from fifty_agent_sdk.mcp import MCPClient, MCPClientConfig -from .conftest import MCP_URL, make_compat_client, make_strict_http_client +from .conftest import MCP_URL, StrictTransportMock, make_compat_client, make_strict_http_client def _config(**overrides: Any) -> MCPClientConfig: @@ -94,6 +97,76 @@ def boom(_request: httpx.Request) -> httpx.Response: assert exc.value.context["method"] == "tools/list" +# --------------------------------------------------------------------------- +# Cancellation leaves inside exception groups (never translated to MCPError) +# --------------------------------------------------------------------------- + + +def _transport_raising(exc: BaseException) -> Any: + """A fake Transport whose connect() raises ``exc`` on __aenter__.""" + + class _FailingTransport: + def connect(self) -> AbstractAsyncContextManager[Any]: + @asynccontextmanager + async def _cm() -> AsyncIterator[Any]: + raise exc + yield # pragma: no cover — unreachable; keeps this a generator + + return _cm() + + return _FailingTransport() + + +async def test_cancelled_error_leaf_in_exception_group_propagates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A CancelledError leaf in a mixed BaseExceptionGroup is re-raised, not + translated into MCPError. + + Regression: consumer cancellation racing a real transport error can arrive + as ``BaseExceptionGroup([ConnectError, CancelledError])`` out of the anyio + task group; translating the group wholesale surfaced + ``MCPError("MCP transport error: CancelledError")`` and broke the SDK's + cancellation contract. The race itself is impractical to stage, so the + mixed group is constructed directly and driven through ``discover``. + """ + race = BaseExceptionGroup("race", [httpx.ConnectError("nope"), asyncio.CancelledError()]) + client = MCPClient(_config()) + monkeypatch.setattr(client, "_build_transport", lambda _headers: _transport_raising(race)) + with pytest.raises(asyncio.CancelledError): + await client.discover() + + +async def test_cancelled_error_leaf_in_nested_exception_group_propagates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The same guarantee when the CancelledError is buried one group deeper + (``_iter_leaf_exceptions`` flattens recursively before classification).""" + race = BaseExceptionGroup( + "outer", + [ + BaseExceptionGroup("inner", [asyncio.CancelledError()]), + httpx.ConnectError("nope"), + ], + ) + client = MCPClient(_config()) + monkeypatch.setattr(client, "_build_transport", lambda _headers: _transport_raising(race)) + with pytest.raises(asyncio.CancelledError): + await client.invoke("search", {"q": "x"}) + + +async def test_pure_exception_group_still_translates_to_mcp_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A group whose leaves are all plain Exceptions still maps to MCPError.""" + group = ExceptionGroup("transport", [httpx.ConnectError("nope"), httpx.ReadTimeout("slow")]) + client = MCPClient(_config()) + monkeypatch.setattr(client, "_build_transport", lambda _headers: _transport_raising(group)) + with pytest.raises(MCPError) as exc: + await client.discover() + assert exc.value.context["wrapped"] == "ConnectError" + + # --------------------------------------------------------------------------- # initialize() RuntimeError -> MCPError (uniform-MCPError contract) # --------------------------------------------------------------------------- @@ -265,33 +338,117 @@ async def bad_auth() -> Any: assert mock.observed_requests == [], "auth must fail before any request" +async def test_raising_auth_callable_raises_mcp_error() -> None: + """A RAISING auth callable (e.g. token endpoint down) surfaces as MCPError. + + Regression: a bad return TYPE was already validated into MCPError, but an + auth callable that raised propagated the raw exception — past the + "returns or raises MCPError" contract and into the Registry's + non-``AgentSdkError`` branch, which downgrades it to a model-recoverable + ``ToolResult(is_error=True)`` instead of the fatal failure it is. The + exception TEXT must not leak into the MCPError (an auth error message may + carry credential material); only the type name is captured. + """ + + def _never_called(_request: httpx.Request) -> httpx.Response: + raise AssertionError("transport should not be invoked") + + http_client, mock = make_strict_http_client(_never_called) + + async def down_token_endpoint() -> dict[str, str]: + raise httpx.ConnectError("https://secret-token-endpoint.internal unreachable") + + client = MCPClient(_config(), auth=down_token_endpoint, client=http_client) + with pytest.raises(MCPError) as exc: + await client.invoke("x", {}) + assert exc.value.context["wrapped"] == "ConnectError" + assert exc.value.context["server_url"] == MCP_URL + serialized = json.dumps(exc.value.context, default=str) + exc.value.message + assert "secret-token-endpoint" not in serialized + assert mock.observed_requests == [], "auth must fail before any request" + + +async def test_cancelled_auth_callable_propagates_cancellation() -> None: + """A CancelledError out of the auth callable is NOT translated into + MCPError — consumer cancellation must propagate (CancelledError is a + BaseException on 3.11+, so the ``except Exception`` guard skips it).""" + + def _never_called(_request: httpx.Request) -> httpx.Response: + raise AssertionError("transport should not be invoked") + + http_client, _ = make_strict_http_client(_never_called) + + async def cancelled_auth() -> dict[str, str]: + raise asyncio.CancelledError + + client = MCPClient(_config(), auth=cancelled_auth, client=http_client) + with pytest.raises(asyncio.CancelledError): + await client.invoke("x", {}) + + # --------------------------------------------------------------------------- # user_agent config field set on the owned httpx client # --------------------------------------------------------------------------- -async def test_user_agent_header_set_on_request() -> None: - """``MCPClientConfig.user_agent`` is set on the owned httpx client.""" +async def test_user_agent_header_set_on_owned_client( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """``MCPClientConfig.user_agent`` reaches the wire on the owned-client path. + + No client is injected, so the transport builds its own + ``httpx.AsyncClient`` with the configured ``User-Agent``. The construction + is intercepted — the network transport swapped for the strict mock, the + headers the transport passes left untouched — and the captured request + must carry the configured value. (The previous version of this test + injected a client that ALREADY carried the header, so it passed even with + the ``user_agent`` plumbing deleted.) + """ + captured: dict[str, str] = {} + + def capture(request: httpx.Request) -> httpx.Response: + captured["user-agent"] = request.headers.get("user-agent", "") + # Fail fast after capture — we only care about the header. + raise httpx.ConnectError("captured") + + mock = StrictTransportMock(capture) + real_async_client = httpx.AsyncClient + + def owned_client_factory(**kwargs: Any) -> httpx.AsyncClient: + # Keep the timeout/headers the transport passes; swap only the network + # transport so the request is capturable without real I/O. + kwargs["transport"] = httpx.MockTransport(mock.handle) + return real_async_client(**kwargs) + monkeypatch.setattr(httpx, "AsyncClient", owned_client_factory) + + client = MCPClient(_config(user_agent="my-agent/2.0")) # owned-client path + with pytest.raises(MCPError): + await client.discover() + assert mock.observed_requests, "expected the request to reach the mock transport" + assert captured["user-agent"] == "my-agent/2.0" + + +async def test_user_agent_not_applied_to_injected_client() -> None: + """An injected client's own headers win — ``user_agent`` is NOT applied. + + Documents the injected-client contract: the transport merges only the + resolved auth headers onto an externally-provided client, so the + ``User-Agent`` the server sees is the injected client's own (here the + httpx default), never ``MCPClientConfig.user_agent``. + """ captured: dict[str, str] = {} def capture(request: httpx.Request) -> httpx.Response: captured["user-agent"] = request.headers.get("user-agent", "") raise httpx.ConnectError("captured") - # Inject a client built from our config so the transport sets User-Agent. - # We use the strict mock for the transport but build the client ourselves - # via the owned-client path by NOT injecting one — instead patch the - # transport's httpx client construction is internal, so assert through an - # injected client carrying the UA the transport would set. - http_client = httpx.AsyncClient( - transport=httpx.MockTransport(capture), - headers={"User-Agent": "my-agent/2.0"}, - ) + http_client, _ = make_strict_http_client(capture) client = MCPClient(_config(user_agent="my-agent/2.0"), client=http_client) with pytest.raises(MCPError): await client.discover() - assert captured["user-agent"] == "my-agent/2.0" + assert captured["user-agent"] == http_client.headers["user-agent"] + assert captured["user-agent"] != "my-agent/2.0" # --------------------------------------------------------------------------- diff --git a/tests/observability/test_hooks.py b/tests/observability/test_hooks.py index fda9788..9ddb022 100644 --- a/tests/observability/test_hooks.py +++ b/tests/observability/test_hooks.py @@ -186,7 +186,31 @@ def hook() -> None: assert failures[0]["log_level"] == "warning" assert failures[0]["hook_name"] == "on_tool_start" assert failures[0]["error_type"] == "RuntimeError" - assert failures[0]["error_message"] == "hook boom" + # Only the exception TYPE is logged — never str(exc) (see below). + assert "error_message" not in failures[0] + + +async def test_invoke_hook_log_never_leaks_hook_arguments() -> None: + """A secret embedded in a hook's exception message never reaches the log. + + Hooks receive high-value objects (``on_run_start`` gets the raw user + message, ``on_llm_call`` the full ``ChatRequest``, ``on_tool_start`` the + tool args), so a hook raising ``ValueError(f"unexpected: {args}")`` would + dump that content into a WARNING line if ``str(exc)`` were logged. The + failure log carries the exception TYPE only. + """ + secret = "SECRET-hook-arg-DO-NOT-LEAK" + + def hook(*args: object) -> None: + raise ValueError(f"unexpected: {args}") + + with structlog.testing.capture_logs() as logs: + await invoke_hook(hook, "on_tool_start", "s1", "search", {"token": secret}) + + failures = [e for e in logs if e.get("event") == "hook.invoke_failed"] + assert len(failures) == 1 + assert failures[0]["error_type"] == "ValueError" + assert secret not in str(failures[0]) async def test_invoke_hook_reraises_cancelled_error() -> None: diff --git a/tests/parser/test_json_mode.py b/tests/parser/test_json_mode.py index 50e8611..f196704 100644 --- a/tests/parser/test_json_mode.py +++ b/tests/parser/test_json_mode.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import sys import pytest @@ -13,6 +14,7 @@ Parser, ThoughtAction, ) +from fifty_agent_sdk.parser import json_mode as json_mode_module from fifty_agent_sdk.parser.json_mode import _RawEnvelope from fifty_agent_sdk.prompts import JSON_MODE_OUTPUT_FORMAT @@ -161,6 +163,31 @@ def test_action_tool_empty_tool_name_raises() -> None: assert excinfo.value.context["missing"] == "tool_name" +def test_action_tool_whitespace_only_tool_name_raises() -> None: + """A whitespace-only name is blank-after-strip and takes the same + schema_validation path as an empty one (the prose parser strips the + ``Action:`` header, so the JSON parser must match).""" + completion = json.dumps( + {"thought": "t", "action": "tool", "tool_name": " \t ", "tool_args": {}} + ) + with pytest.raises(ParserError) as excinfo: + _parser().parse(completion) + ctx = excinfo.value.context + assert ctx["error_phase"] == "schema_validation" + assert ctx["missing"] == "tool_name" + + +def test_action_tool_padded_tool_name_is_stripped() -> None: + """Surrounding whitespace is stripped so the registry sees the same clean + name the prose parser would emit.""" + completion = json.dumps( + {"thought": "t", "action": "tool", "tool_name": " search ", "tool_args": {}} + ) + result = _parser().parse(completion) + assert isinstance(result, ThoughtAction) + assert result.tool_call.name == "search" + + def test_action_final_missing_answer_raises() -> None: completion = json.dumps({"thought": "t", "action": "final"}) with pytest.raises(ParserError) as excinfo: @@ -218,6 +245,94 @@ def test_parser_error_chains_cause_via_raise_from() -> None: assert excinfo.value.__cause__ is not None +def test_strict_recursion_error_is_translated_to_parser_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """BR-017 deterministically pins strict-pass ``RecursionError`` translation.""" + sentinel = RecursionError("deterministic depth failure") + calls = 0 + + def raise_recursion(_payload: str) -> object: + nonlocal calls + calls += 1 + raise sentinel + + monkeypatch.setattr(json_mode_module.json, "loads", raise_recursion) + completion = '{"thought":"t","action":"final","answer":"ok"}' + with pytest.raises(ParserError) as excinfo: + _parser().parse(completion) + ctx = excinfo.value.context + assert ctx["parser"] == "JsonModeParser" + assert ctx["error_phase"] == "json_decode" + assert "RecursionError" in str(ctx["cause"]) + assert len(str(ctx["completion_excerpt"])) <= 200 + assert excinfo.value.__cause__ is sentinel + assert calls == 1 + + +def test_recovery_recursion_error_is_translated_to_parser_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """BR-017 deterministically pins recovery-pass ``RecursionError`` translation.""" + first = json.JSONDecodeError("strict failed", "prefix", 0) + sentinel = RecursionError("deterministic recovery depth failure") + errors = iter((first, sentinel)) + calls = 0 + + def raise_scripted(_payload: str) -> object: + nonlocal calls + calls += 1 + raise next(errors) + + monkeypatch.setattr(json_mode_module.json, "loads", raise_scripted) + completion = 'prefix {"thought":"t","action":"final","answer":"ok"} suffix' + with pytest.raises(ParserError) as excinfo: + _parser().parse(completion) + ctx = excinfo.value.context + assert ctx["error_phase"] == "json_decode" + assert "RecursionError" in str(ctx["cause"]) + assert len(str(ctx["completion_excerpt"])) <= 200 + assert excinfo.value.__cause__ is sentinel + assert calls == 2 + + +def test_oversized_integer_strict_decode_is_contained() -> None: + """BR-013 contains bare ValueError from strict ``json.loads`` decoding.""" + digits = "9" * (sys.get_int_max_str_digits() + 1) + with pytest.raises(ParserError) as excinfo: + _parser().parse(digits) + assert str(excinfo.value) == "could not decode JSON envelope" + assert excinfo.value.context["error_phase"] == "json_decode" + assert len(str(excinfo.value.context["completion_excerpt"])) <= 200 + assert type(excinfo.value.__cause__) is ValueError + + +def test_oversized_integer_recovery_decode_is_contained() -> None: + """BR-013 contains bare ValueError from the JSON recovery decode.""" + digits = "9" * (sys.get_int_max_str_digits() + 1) + completion = f'prefix {{"thought":"t","action":"final","answer":{digits}}} suffix' + with pytest.raises(ParserError) as excinfo: + _parser().parse(completion) + assert str(excinfo.value) == "could not decode JSON envelope after fence recovery" + assert excinfo.value.context["error_phase"] == "json_decode" + assert type(excinfo.value.__cause__) is ValueError + + +def test_malformed_json_preserves_decode_message_and_json_cause() -> None: + """BR-013 leaves the common malformed-syntax contract unchanged.""" + completion = "not json" + with pytest.raises(ParserError) as excinfo: + _parser().parse(completion) + assert str(excinfo.value) == "could not decode JSON envelope" + assert excinfo.value.context == { + "parser": "JsonModeParser", + "error_phase": "json_decode", + "completion_excerpt": completion, + "cause": repr(excinfo.value.__cause__), + } + assert isinstance(excinfo.value.__cause__, json.JSONDecodeError) + + # ---------------------------------------------------------------------- # # Protocol / cross-brief contract # # ---------------------------------------------------------------------- # diff --git a/tests/parser/test_prose_mode.py b/tests/parser/test_prose_mode.py index 89997eb..40e1b62 100644 --- a/tests/parser/test_prose_mode.py +++ b/tests/parser/test_prose_mode.py @@ -2,10 +2,14 @@ from __future__ import annotations +import json +import sys + import pytest from fifty_agent_sdk.errors import ParserError from fifty_agent_sdk.parser import FinalAnswer, Parser, ProseModeParser, ThoughtAction +from fifty_agent_sdk.parser import prose_mode as prose_mode_module def _parser() -> ProseModeParser: @@ -190,6 +194,94 @@ def test_huge_whitespace_payload_does_not_hang() -> None: assert excinfo.value.context["error_phase"] == "empty_completion" +def test_strict_action_input_recursion_error_is_translated_to_parser_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """BR-017 deterministically pins strict Action Input recursion containment.""" + sentinel = RecursionError("deterministic depth failure") + calls = 0 + + def raise_recursion(_payload: str) -> object: + nonlocal calls + calls += 1 + raise sentinel + + monkeypatch.setattr(prose_mode_module.json, "loads", raise_recursion) + completion = 'Thought: T\nAction: search\nAction Input: {"q":"x"}' + with pytest.raises(ParserError) as excinfo: + _parser().parse(completion) + ctx = excinfo.value.context + assert ctx["parser"] == "ProseModeParser" + assert ctx["error_phase"] == "action_input_decode" + assert "RecursionError" in str(ctx["cause"]) + assert len(str(ctx["completion_excerpt"])) <= 200 + assert excinfo.value.__cause__ is sentinel + assert calls == 1 + + +def test_recovery_action_input_recursion_error_is_translated_to_parser_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """BR-017 deterministically pins recovery Action Input recursion containment.""" + first = json.JSONDecodeError("strict failed", "fence", 0) + sentinel = RecursionError("deterministic recovery depth failure") + errors = iter((first, sentinel)) + calls = 0 + + def raise_scripted(_payload: str) -> object: + nonlocal calls + calls += 1 + raise next(errors) + + monkeypatch.setattr(prose_mode_module.json, "loads", raise_scripted) + completion = 'Thought: T\nAction: search\nAction Input: ```json\n{"q":"x"}\n```' + with pytest.raises(ParserError) as excinfo: + _parser().parse(completion) + ctx = excinfo.value.context + assert ctx["error_phase"] == "action_input_decode" + assert "RecursionError" in str(ctx["cause"]) + assert len(str(ctx["completion_excerpt"])) <= 200 + assert excinfo.value.__cause__ is sentinel + assert calls == 2 + + +def test_oversized_integer_action_input_strict_decode_is_contained() -> None: + """BR-013 contains strict Action Input bare ValueError.""" + digits = "9" * (sys.get_int_max_str_digits() + 1) + completion = f"Thought: T\nAction: search\nAction Input: {digits}" + with pytest.raises(ParserError) as excinfo: + _parser().parse(completion) + assert str(excinfo.value) == "could not decode Action Input JSON" + assert excinfo.value.context["error_phase"] == "action_input_decode" + assert type(excinfo.value.__cause__) is ValueError + + +def test_oversized_integer_action_input_recovery_decode_is_contained() -> None: + """BR-013 contains recovery Action Input bare ValueError.""" + digits = "9" * (sys.get_int_max_str_digits() + 1) + completion = f'Thought: T\nAction: search\nAction Input: ```json\n{{"n":{digits}}}\n```' + with pytest.raises(ParserError) as excinfo: + _parser().parse(completion) + assert str(excinfo.value) == "could not decode Action Input JSON after fence recovery" + assert excinfo.value.context["error_phase"] == "action_input_decode" + assert type(excinfo.value.__cause__) is ValueError + + +def test_malformed_action_input_preserves_decode_message_and_json_cause() -> None: + """BR-013 leaves malformed Action Input syntax behavior unchanged.""" + completion = "Thought: T\nAction: search\nAction Input: not json" + with pytest.raises(ParserError) as excinfo: + _parser().parse(completion) + assert str(excinfo.value) == "could not decode Action Input JSON" + assert excinfo.value.context == { + "parser": "ProseModeParser", + "error_phase": "action_input_decode", + "completion_excerpt": completion, + "cause": repr(excinfo.value.__cause__), + } + assert isinstance(excinfo.value.__cause__, json.JSONDecodeError) + + # ---------------------------------------------------------------------- # # Protocol # # ---------------------------------------------------------------------- # diff --git a/tests/runner/test_runner_audit.py b/tests/runner/test_runner_audit.py index f3ba643..fc158a4 100644 --- a/tests/runner/test_runner_audit.py +++ b/tests/runner/test_runner_audit.py @@ -18,8 +18,16 @@ import structlog -from fifty_agent_sdk import AuditEvent, AuditSink, Registry, SafetyConfig, ToolResult -from tests.loop.conftest import FakeLLMClient, FakeTool, make_response +from fifty_agent_sdk import ( + ActionEvent, + AuditEvent, + AuditSink, + Registry, + SafetyConfig, + ToolResult, + ToolStartedEvent, +) +from tests.loop.conftest import FakeLLMClient, FakeTool, make_multi_tool_response, make_response from tests.runner.conftest import collect, final_json, make_runner, tool_json # --------------------------------------------------------------------------- @@ -94,7 +102,9 @@ async def test_tool_invocation_payload_carries_correlation_fields() -> None: tool_event = next(e for e in spy.events if e.event_type == "tool_invocation") assert tool_event.payload["tool_name"] == "search" - assert tool_event.payload["args"] == {"q": "weather"} + # Args are reduced to non-content metadata: sorted keys, per-value type + # and length — never the values themselves. + assert tool_event.payload["args"] == {"q": {"type": "str", "len": len("weather")}} assert tool_event.payload["outcome"] == "ok" assert isinstance(tool_event.payload["call_id"], str) assert tool_event.payload["call_id"] != "" @@ -164,9 +174,93 @@ async def test_multi_tool_run_emits_one_event_per_tool() -> None: ] tool_events = [e for e in spy.events if e.event_type == "tool_invocation"] assert tool_events[0].payload["tool_name"] == "alpha" - assert tool_events[0].payload["args"] == {"n": 1} + assert tool_events[0].payload["args"] == {"n": {"type": "int", "len": None}} assert tool_events[1].payload["tool_name"] == "beta" - assert tool_events[1].payload["args"] == {"n": 2} + assert tool_events[1].payload["args"] == {"n": {"type": "int", "len": None}} + + +async def test_multi_action_batch_audits_each_call_with_own_args_and_call_id() -> None: + """A native MultiAction batch emits one tool_invocation per call, each with + its OWN args and call_id. + + Regression gate for per-call correlation: the loop's MultiAction branch + emits N ActionEvents, then N ToolStartedEvents, then N terminal events — + all in call order. The old single-slot correlation gave the FIRST + terminal event the LAST call's call_id/args and left every other call + with ``args={}``. With per-call-id correlation each invocation's payload + carries its own pair. + """ + registry = Registry() + registry.register(FakeTool("alpha", result=ToolResult(output="A-result"))) + registry.register(FakeTool("beta", result=ToolResult(output="B-result"))) + llm = FakeLLMClient( + replies=[ + make_multi_tool_response([("alpha", {"n": 1}), ("beta", {"n": 2})]), + make_response(final_json("done")), + ] + ) + spy = SpyAuditSink() + runner, _store = make_runner( + llm=llm, + registry=registry, + safety=SafetyConfig(native_tools_enabled=True, max_concurrent_tool_calls=2), + audit=spy, + ) + + events = await collect(runner.run("s1", "Hi")) + + # Sanity: this run really went through the MultiAction batch shape — two + # ActionEvents followed by two ToolStartedEvents. + actions = [e for e in events if isinstance(e, ActionEvent)] + started = [e for e in events if isinstance(e, ToolStartedEvent)] + assert len(actions) == len(started) == 2 + + call_id_by_tool = {e.tool_name: e.call_id for e in started} + tool_events = [e for e in spy.events if e.event_type == "tool_invocation"] + assert len(tool_events) == 2 + by_tool = {e.payload["tool_name"]: e.payload for e in tool_events} + # Correlation is per call; args in the payload are the non-content + # metadata summary (see _args_metadata), not the raw values. + assert by_tool["alpha"]["args"] == {"n": {"type": "int", "len": None}} + assert by_tool["alpha"]["call_id"] == call_id_by_tool["alpha"] + assert by_tool["beta"]["args"] == {"n": {"type": "int", "len": None}} + assert by_tool["beta"]["call_id"] == call_id_by_tool["beta"] + + +# --------------------------------------------------------------------------- +# Args redaction — payload never carries argument values +# --------------------------------------------------------------------------- + + +async def test_tool_invocation_args_never_leak_secret_values() -> None: + """A secret-looking arg VALUE never appears in the tool_invocation payload. + + Mirrors the ``"SECRET" not in json.dumps(...)`` redaction-proof pattern + of the MCP auth tests: argument values routinely carry credentials and + PII, so the payload carries only keys, per-value type names and lengths. + """ + import json + + secret = "SECRET-api-key-DO-NOT-LEAK" + registry = Registry() + registry.register(FakeTool("auth_call", result=ToolResult(output="ok"))) + llm = FakeLLMClient( + replies=[ + make_response(tool_json("t", "auth_call", {"token": secret, "retries": 3})), + make_response(final_json("done")), + ] + ) + spy = SpyAuditSink() + runner, _store = make_runner(llm=llm, registry=registry, audit=spy) + + await collect(runner.run("s1", "Hi")) + + tool_event = next(e for e in spy.events if e.event_type == "tool_invocation") + assert secret not in json.dumps(tool_event.payload, default=str) + assert tool_event.payload["args"] == { + "retries": {"type": "int", "len": None}, + "token": {"type": "str", "len": len(secret)}, + } # --------------------------------------------------------------------------- @@ -224,6 +318,49 @@ async def test_llm_error_run_emits_session_start_then_error() -> None: assert "run_id" in error_event.payload +async def test_fatal_sdk_error_escaping_loop_is_audited() -> None: + """An MCPError escaping the loop emits an error audit event, terminates + with ``terminated_by="sdk_error"``, and still re-raises to the caller. + + The registry re-raises non-recoverable :class:`AgentSdkError` subclasses + untouched, so an MCP transport failure propagates out of the loop as an + exception rather than as an ErrorEvent. Before this fix such a run exited + invisibly: no error audit event and ``terminated_by="interrupted"``. + """ + import pytest + + from fifty_agent_sdk.errors import MCPError + + registry = Registry() + registry.register( + FakeTool( + "mcp_tool", + raises=MCPError("transport down", context={"server_url": "https://mcp.example.com"}), + ) + ) + llm = FakeLLMClient(replies=[make_response(tool_json("t", "mcp_tool", {}))]) + spy = SpyAuditSink() + runner, store = make_runner(llm=llm, registry=registry, audit=spy) + + with structlog.testing.capture_logs() as logs, pytest.raises(MCPError): + await collect(runner.run("s1", "Hi")) + + # The escaped SDK error is audited like any other error. + assert [e.event_type for e in spy.events] == ["session_start", "error"] + error_event = spy.events[-1] + assert error_event.payload["error_type"] == "MCPError" + assert "transport down" in error_event.payload["error_message"] + + # The run_completed log attributes the exit to the escaped SDK error. + completed = [e for e in logs if e.get("event") == "runner.run_completed"] + assert len(completed) == 1 + assert completed[0]["terminated_by"] == "sdk_error" + + # No assistant message was committed; the durable user message survives. + history = await store.get_messages("s1") + assert [m.role for m in history] == ["user"] + + async def test_state_store_error_on_assistant_persist_is_audited() -> None: """A persist_assistant durability failure emits an error audit event.""" from fifty_agent_sdk import ( diff --git a/tests/runner/test_runner_hooks.py b/tests/runner/test_runner_hooks.py index ad962ed..682e837 100644 --- a/tests/runner/test_runner_hooks.py +++ b/tests/runner/test_runner_hooks.py @@ -37,8 +37,8 @@ StateStoreError, ToolResult, ) -from fifty_agent_sdk.errors import LLMError -from tests.loop.conftest import FakeLLMClient, FakeTool, make_response +from fifty_agent_sdk.errors import LLMError, MCPError +from tests.loop.conftest import FakeLLMClient, FakeTool, make_multi_tool_response, make_response from tests.runner.conftest import collect, final_json, make_runner, tool_json # --------------------------------------------------------------------------- @@ -218,6 +218,45 @@ async def test_counting_hooks_multi_tool_run() -> None: assert [a[1] for a in rec.args["on_tool_end"]] == ["alpha", "beta"] +async def test_multi_action_batch_tool_hooks_carry_per_call_args_and_results() -> None: + """A native MultiAction batch correlates each tool hook with its own call. + + Regression gate for per-call correlation: the loop's MultiAction branch + emits N ActionEvents, then N ToolStartedEvents, then N terminal events — + all in call order. ``on_tool_start`` for the FIRST call must carry the + FIRST call's args (the old single-slot correlation handed it the last + call's args), and ``on_tool_end`` must pair each tool with its own + result and a non-negative duration. + """ + registry = Registry() + registry.register(FakeTool("alpha", result=ToolResult(output="A-result"))) + registry.register(FakeTool("beta", result=ToolResult(output="B-result"))) + llm = FakeLLMClient( + replies=[ + make_multi_tool_response([("alpha", {"n": 1}), ("beta", {"n": 2})]), + make_response(final_json("done")), + ] + ) + rec = HookRecorder() + runner, _store = make_runner( + llm=llm, + registry=registry, + safety=SafetyConfig(native_tools_enabled=True, max_concurrent_tool_calls=2), + hooks=rec.hooks(), + ) + + await collect(runner.run("s1", "Hi")) + + assert rec.counts["on_tool_start"] == 2 + assert rec.counts["on_tool_end"] == 2 + starts = {a[1]: a[2] for a in rec.args["on_tool_start"]} + assert starts == {"alpha": {"n": 1}, "beta": {"n": 2}} + ends = {a[1]: (a[2], a[3]) for a in rec.args["on_tool_end"]} + assert ends["alpha"][0] == "A-result" + assert ends["beta"][0] == "B-result" + assert all(duration >= 0.0 for _, duration in ends.values()) + + # --------------------------------------------------------------------------- # on_run_end — duration + error # --------------------------------------------------------------------------- @@ -262,6 +301,45 @@ async def test_on_error_fires_on_llm_failure() -> None: assert rec.args["on_run_end"][0][2] is None +# --------------------------------------------------------------------------- +# on_error — fatal AgentSdkError escaping the loop +# --------------------------------------------------------------------------- + + +async def test_on_error_fires_on_fatal_sdk_error_escaping_loop() -> None: + """An MCPError escaping the loop fires ``on_error``, is handed to + ``on_run_end``, and still re-raises to the caller. + + The registry re-raises non-recoverable :class:`AgentSdkError` subclasses + untouched, so the failure arrives as an exception rather than an + ErrorEvent. Before this fix the run exited with no ``on_error`` call and + ``on_run_end(error=None)``. + """ + registry = Registry() + registry.register( + FakeTool( + "mcp_tool", + raises=MCPError("transport down", context={"server_url": "https://mcp.example.com"}), + ) + ) + llm = FakeLLMClient(replies=[make_response(tool_json("t", "mcp_tool", {}))]) + rec = HookRecorder() + runner, _store = make_runner(llm=llm, registry=registry, hooks=rec.hooks()) + + with pytest.raises(MCPError): + await collect(runner.run("s1", "Hi")) + + assert rec.counts["on_error"] == 1 + session_id, error, context = rec.args["on_error"][0] + assert session_id == "s1" + assert isinstance(error, MCPError) + assert context["error_type"] == "MCPError" + # The escaping exception is carried to on_run_end (unlike a loop-internal + # failure, which leaves it None). + assert rec.counts["on_run_end"] == 1 + assert isinstance(rec.args["on_run_end"][0][2], MCPError) + + # --------------------------------------------------------------------------- # on_error — state-store failure # --------------------------------------------------------------------------- diff --git a/tests/state/test_redis.py b/tests/state/test_redis.py index e176ba0..5da4159 100644 --- a/tests/state/test_redis.py +++ b/tests/state/test_redis.py @@ -15,8 +15,8 @@ * Fresh-list-per-call (the defensive-copy invariant). * Idempotent delete; delete is scoped to one session. * The configured ``key_prefix`` is applied (default ``fifty_agent_sdk:state:``). -* TTL is *set* on append when ``ttl_seconds`` is configured, refreshed on - every append, and absent when ``ttl_seconds`` is ``None``. +* Every mutation atomically refreshes all session keys to one sliding TTL; + TTL-disabled mode emits no expiry commands. * Every backend failure (:class:`redis.exceptions.RedisError`) is wrapped into :class:`StateStoreError` with the documented context shape. * :class:`RedisStateStore` satisfies the :class:`StateStore` protocol. @@ -33,12 +33,14 @@ from __future__ import annotations +import json from collections.abc import AsyncIterator +from typing import Any import fakeredis import pytest import pytest_asyncio -from redis.exceptions import RedisError +from redis.exceptions import RedisError, WatchError from fifty_agent_sdk import ChatMessage, RedisStateStore, StateStore, StateStoreError @@ -273,6 +275,261 @@ async def test_no_ttl_when_ttl_seconds_is_none(store: RedisStateStore) -> None: assert ttl == -1 +async def test_switch_branch_applies_ttl_to_active_pointer( + store_with_ttl: RedisStateStore, +) -> None: + """``switch_branch`` must not leave the ``:active`` pointer without a TTL. + + Regression: the pointer was written with a bare ``SET``, which both + stripped any TTL the key already carried and left a first-time key with + none — so the active head could outlive the session it belongs to. + """ + await store_with_ttl.append("s1", ChatMessage(role="user", content="a")) + branch = await store_with_ttl.fork("s1", from_sequence=1) + await store_with_ttl.switch_branch("s1", branch) + ttl = await store_with_ttl._client.ttl("fifty_agent_sdk:state:s1:active") + assert 0 < ttl <= 3600 + # Switching back re-applies the TTL just the same. + await store_with_ttl.switch_branch("s1", "trunk") + ttl = await store_with_ttl._client.ttl("fifty_agent_sdk:state:s1:active") + assert 0 < ttl <= 3600 + + +async def test_switch_branch_without_ttl_leaves_pointer_durable( + store: RedisStateStore, +) -> None: + """With ``ttl_seconds=None`` the ``:active`` pointer gets no expiry.""" + await store.append("s1", ChatMessage(role="user", content="a")) + branch = await store.fork("s1", from_sequence=1) + await store.switch_branch("s1", branch) + ttl = await store._client.ttl("fifty_agent_sdk:state:s1:active") + assert ttl == -1 + + +async def test_fork_applies_ttl_to_branches_registry( + store_with_ttl: RedisStateStore, +) -> None: + """Forking a pre-BR-004 session gives the new ``:branches`` hash the TTL. + + Regression: the registry hash was created via HSETNX/HSET with no + ``EXPIRE``, so forking a legacy single-list session (whose bare list + predates the registry) left the registry durable while the rest of the + session expired. + """ + # Pre-BR-004 layout: a bare message list with no registry hash. + await store_with_ttl._client.rpush( + "fifty_agent_sdk:state:legacy", + ChatMessage(role="user", content="old").model_dump_json(), + ) + await store_with_ttl.fork("legacy", from_sequence=1) + ttl = await store_with_ttl._client.ttl("fifty_agent_sdk:state:legacy:branches") + assert 0 < ttl <= 3600 + + +async def _session_pttls(store: RedisStateStore, session_id: str) -> list[int]: + keys = await store._client.keys(f"fifty_agent_sdk:state:{session_id}*") + return [int(await store._client.pttl(key)) for key in keys] + + +def _assert_synchronized_ttls(ttls: list[int]) -> None: + assert ttls + assert min(ttls) > 3_500_000 + assert max(ttls) - min(ttls) <= 100 + + +async def _seed_branched_session(store: RedisStateStore) -> str: + await store.append("sync", ChatMessage(role="user", content="trunk")) + branch = await store.fork("sync", from_sequence=1) + await store.switch_branch("sync", branch) + await store.append("sync", ChatMessage(role="user", content="fork")) + return branch + + +async def _shorten_session_ttls(store: RedisStateStore) -> None: + keys = await store._client.keys("fifty_agent_sdk:state:sync*") + for key in keys: + await store._client.pexpire(key, 10_000) + + +async def test_append_refreshes_every_session_key_to_one_ttl( + store_with_ttl: RedisStateStore, +) -> None: + """BR-015 append atomically slides all branch/session keys together.""" + await _seed_branched_session(store_with_ttl) + await _shorten_session_ttls(store_with_ttl) + await store_with_ttl.append("sync", ChatMessage(role="user", content="again")) + _assert_synchronized_ttls(await _session_pttls(store_with_ttl, "sync")) + + +async def test_fork_refreshes_legacy_and_registry_keys_to_one_ttl( + store_with_ttl: RedisStateStore, +) -> None: + """BR-015 forking a legacy session cannot create a longer-lived registry.""" + key = "fifty_agent_sdk:state:legacy-sync" + await store_with_ttl._client.rpush( + key, ChatMessage(role="user", content="old").model_dump_json() + ) + await store_with_ttl._client.pexpire(key, 10_000) + await store_with_ttl.fork("legacy-sync", from_sequence=1) + _assert_synchronized_ttls(await _session_pttls(store_with_ttl, "legacy-sync")) + + +async def test_switch_refreshes_every_session_key_to_one_ttl( + store_with_ttl: RedisStateStore, +) -> None: + """BR-015 switching cannot create an active pointer that outlives messages.""" + branch = await _seed_branched_session(store_with_ttl) + await _shorten_session_ttls(store_with_ttl) + await store_with_ttl.switch_branch("sync", branch) + _assert_synchronized_ttls(await _session_pttls(store_with_ttl, "sync")) + + +async def test_truncate_refreshes_every_session_key_to_one_ttl( + store_with_ttl: RedisStateStore, +) -> None: + """BR-015 truncate uses the same whole-session sliding TTL transaction.""" + await _seed_branched_session(store_with_ttl) + await _shorten_session_ttls(store_with_ttl) + await store_with_ttl.truncate_after("sync", 1) + _assert_synchronized_ttls(await _session_pttls(store_with_ttl, "sync")) + + +async def test_ttl_disabled_mode_keeps_all_session_keys_durable(store: RedisStateStore) -> None: + """BR-015 ttl_seconds=None emits no expiry for any mutation-created key.""" + branch = await _seed_branched_session(store) + await store.switch_branch("sync", branch) + await store.truncate_after("sync", 1) + assert set(await _session_pttls(store, "sync")) == {-1} + + +async def test_mutation_retries_from_fresh_snapshot_after_watch_conflict( + store_with_ttl: RedisStateStore, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """BR-015 a WatchError retries without duplicating the state mutation.""" + original = store_with_ttl._mutation_snapshot + calls = 0 + + async def conflict_once(pipe: Any, session_id: str) -> Any: + nonlocal calls + calls += 1 + snapshot = await original(pipe, session_id) + if calls == 1: + await store_with_ttl._client.hset( + store_with_ttl._branches_key(session_id), + "trunk", + json.dumps( + { + "parent_branch_id": None, + "forked_from_sequence": None, + "created_at": None, + } + ), + ) + return snapshot + + monkeypatch.setattr(store_with_ttl, "_mutation_snapshot", conflict_once) + await store_with_ttl.append("retry", ChatMessage(role="user", content="once")) + assert calls == 2 + assert [message.content for message in await store_with_ttl.get_messages("retry")] == ["once"] + + +async def test_append_retries_and_reroutes_when_active_pointer_changes( + store_with_ttl: RedisStateStore, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """BR-015 active-pointer conflicts retry before choosing the append target.""" + await store_with_ttl.append("route", ChatMessage(role="user", content="trunk")) + branch = await store_with_ttl.fork("route", from_sequence=1) + original = store_with_ttl._mutation_snapshot + calls = 0 + + async def switch_during_first_snapshot(pipe: Any, session_id: str) -> Any: + nonlocal calls + calls += 1 + snapshot = await original(pipe, session_id) + if calls == 1: + await store_with_ttl._client.set(store_with_ttl._active_key(session_id), branch) + return snapshot + + monkeypatch.setattr(store_with_ttl, "_mutation_snapshot", switch_during_first_snapshot) + await store_with_ttl.append("route", ChatMessage(role="user", content="routed")) + + assert calls == 2 + trunk = await store_with_ttl.get_messages("route", branch_id="trunk") + fork = await store_with_ttl.get_messages("route", branch_id=branch) + assert [message.content for message in trunk] == ["trunk"] + assert [message.content for message in fork] == ["trunk", "routed"] + + +async def test_append_retries_when_existing_fork_list_changes( + store_with_ttl: RedisStateStore, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """BR-015 fork-list conflicts retry without lost, reordered, or duplicate writes.""" + branch = await _seed_branched_session(store_with_ttl) + await _shorten_session_ttls(store_with_ttl) + original = store_with_ttl._mutation_snapshot + calls = 0 + + async def append_during_first_snapshot(pipe: Any, session_id: str) -> Any: + nonlocal calls + calls += 1 + snapshot = await original(pipe, session_id) + if calls == 1: + await store_with_ttl._client.rpush( + store_with_ttl._msgs_key(session_id, branch), + ChatMessage(role="user", content="concurrent").model_dump_json(), + ) + return snapshot + + monkeypatch.setattr(store_with_ttl, "_mutation_snapshot", append_during_first_snapshot) + await store_with_ttl.append("sync", ChatMessage(role="user", content="requested")) + + assert calls == 2 + trunk = await store_with_ttl.get_messages("sync", branch_id="trunk") + fork = await store_with_ttl.get_messages("sync", branch_id=branch) + assert [message.content for message in trunk] == ["trunk"] + assert [message.content for message in fork] == [ + "trunk", + "fork", + "concurrent", + "requested", + ] + _assert_synchronized_ttls(await _session_pttls(store_with_ttl, "sync")) + + +async def test_mutation_wraps_watch_error_after_bounded_retries( + store_with_ttl: RedisStateStore, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """BR-015 sustained optimistic-lock contention fails after three attempts.""" + calls = 0 + + async def always_conflict(_pipe: Any, _session_id: str) -> Any: + nonlocal calls + calls += 1 + raise WatchError("forced conflict") + + monkeypatch.setattr(store_with_ttl, "_mutation_snapshot", always_conflict) + with pytest.raises(StateStoreError) as excinfo: + await store_with_ttl.append("retry", ChatMessage(role="user", content="never")) + assert calls == 3 + assert excinfo.value.context["wrapped"] == "WatchError" + + +def test_non_positive_ttl_seconds_rejected() -> None: + """``ttl_seconds <= 0`` raises at construction instead of deleting data. + + Regression: the check was ``is not None``, so ``ttl_seconds=0`` issued + ``EXPIRE 0`` inside the append transaction — deleting every session key + on the first append. + """ + for bad in (0, -1, -3600): + with pytest.raises(ValueError, match="positive integer or None"): + RedisStateStore("redis://localhost:6379/0", ttl_seconds=bad) + + # --------------------------------------------------------------------------- # Error wrapping # --------------------------------------------------------------------------- diff --git a/tests/state/test_redis_integration.py b/tests/state/test_redis_integration.py index c08fe6b..3651ae4 100644 --- a/tests/state/test_redis_integration.py +++ b/tests/state/test_redis_integration.py @@ -61,7 +61,7 @@ async def store() -> AsyncIterator[RedisStateStore]: try: yield s finally: - for session_id in ("s1", "s2", "expiring"): + for session_id in ("s1", "s2", "expiring", "branched"): await s.delete(session_id) await s.aclose() @@ -92,6 +92,25 @@ async def test_redis_ttl_is_set(store: RedisStateStore) -> None: assert 0 < ttl <= 3600 +async def test_redis_branch_mutation_synchronizes_all_session_pttls( + store: RedisStateStore, +) -> None: + """BR-015 real Redis applies one transaction's TTL to every session key.""" + await store.append("branched", ChatMessage(role="user", content="a")) + branch = await store.fork("branched", from_sequence=1) + await store.switch_branch("branched", branch) + await store.append("branched", ChatMessage(role="user", content="b")) + + keys = await store._client.keys(f"{_TEST_KEY_PREFIX}branched*") + for key in keys: + await store._client.pexpire(key, 10_000) + await store.switch_branch("branched", "trunk") + + pttls = [int(await store._client.pttl(key)) for key in keys] + assert min(pttls) > 3_500_000 + assert max(pttls) - min(pttls) <= 100 + + async def test_redis_delete_removes_key(store: RedisStateStore) -> None: """``delete`` removes the session key from the real server.""" await store.append("s2", ChatMessage(role="user", content="a")) diff --git a/tests/state/test_sql.py b/tests/state/test_sql.py index a113e08..c808ab8 100644 --- a/tests/state/test_sql.py +++ b/tests/state/test_sql.py @@ -43,7 +43,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from sqlalchemy.pool import StaticPool -from fifty_agent_sdk import ChatMessage, StateStore, StateStoreError +from fifty_agent_sdk import ChatMessage, StateStore, StateStoreError, ToolCall from fifty_agent_sdk.state.sql import ( AgentMessage, AgentSession, @@ -166,7 +166,7 @@ async def test_round_trip_preserves_ordering(store: SqlStateStore) -> None: async def test_round_trip_preserves_all_chat_message_fields( store: SqlStateStore, ) -> None: - """All four optional/required ChatMessage fields survive a round-trip.""" + """All five optional/required ChatMessage fields survive a round-trip.""" msg = ChatMessage( role="tool", content="result-body", @@ -180,16 +180,71 @@ async def test_round_trip_preserves_all_chat_message_fields( assert got[0].content == "result-body" assert got[0].name == "search" assert got[0].tool_call_id == "call-abc" + assert got[0].tool_calls is None + + +async def test_round_trip_preserves_tool_calls(store: SqlStateStore) -> None: + """A native tool-calling assistant turn round-trips its ``tool_calls``. + + Regression: ``AgentMessage`` used to persist only four of the five + :class:`ChatMessage` fields, silently dropping ``tool_calls`` on reload + and orphaning the paired ``role="tool"`` reply on session resume. + """ + assistant_turn = ChatMessage( + role="assistant", + content="", + tool_calls=[ + ToolCall(name="search", args={"q": "x"}, id="call-abc"), + ToolCall(name="lookup", args={"id": 7}), + ], + ) + tool_reply = ChatMessage( + role="tool", + content="result-body", + name="search", + tool_call_id="call-abc", + ) + await store.append("s1", assistant_turn) + await store.append("s1", tool_reply) + + got = await store.get_messages("s1") + assert got == [assistant_turn, tool_reply] + assert got[0].tool_calls is not None + assert got[0].tool_calls[0].id == "call-abc" + # The pairing key survived the reload: the tool reply still matches the + # assistant turn's first native call. + assert got[1].tool_call_id == got[0].tool_calls[0].id + + +async def test_legacy_null_tool_calls_reads_back_as_none( + store: SqlStateStore, engine: AsyncEngine +) -> None: + """A row with ``tool_calls`` NULL (every pre-column row) reads as None.""" + await store.append("s1", ChatMessage(role="user", content="seed")) + async with engine.begin() as conn: + await conn.execute( + insert(AgentMessage).values( + session_id="s1", + sequence=2, + role="assistant", + content="legacy", + tool_calls=None, + ) + ) + got = await store.get_messages("s1") + assert [m.content for m in got] == ["seed", "legacy"] + assert got[1].tool_calls is None async def test_round_trip_handles_optional_fields_as_none( store: SqlStateStore, ) -> None: - """``name`` and ``tool_call_id`` are nullable and round-trip as None.""" + """``name``, ``tool_call_id`` and ``tool_calls`` are nullable and round-trip as None.""" await store.append("s1", ChatMessage(role="user", content="hi")) got = await store.get_messages("s1") assert got[0].name is None assert got[0].tool_call_id is None + assert got[0].tool_calls is None async def test_round_trip_allows_empty_content(store: SqlStateStore) -> None: @@ -667,6 +722,7 @@ def test_metadata_columns_match_schema() -> None: "content", "name", "tool_call_id", + "tool_calls", "created_at", } assert {c.name for c in branches.columns} == { diff --git a/tests/test_maintaining.py b/tests/test_maintaining.py index e0c73e4..72b14a8 100644 --- a/tests/test_maintaining.py +++ b/tests/test_maintaining.py @@ -31,11 +31,12 @@ MAP_PATH: Final[Path] = REPO_ROOT / "MAINTAINING.md" # A backticked span whose whole content is a repo-relative path under ``src/`` -# or ``tests/``, with an optional ``:`` suffix. MAINTAINING.md carries -# LOCAL paths only — consumer identities and their file paths live in the -# private brief tracker, because this repo is public and they are not. +# or ``tests/``, or a root-level ``*.md`` doc, with an optional ``:`` +# suffix. MAINTAINING.md carries LOCAL paths only — consumer identities and +# their file paths live in the private brief tracker, because this repo is +# public and they are not. _LOCAL_PATH_IN_BACKTICKS: Final[re.Pattern[str]] = re.compile( - r"`((?:src|tests)/[A-Za-z0-9_./-]+?)(?::\d+)?`" + r"`((?:src|tests)/[A-Za-z0-9_./-]+?|[A-Za-z0-9_-]+\.md)(?::\d+)?`" ) _PREFIX_HINT: Final[str] = ( @@ -91,8 +92,9 @@ def test_every_repo_path_named_by_the_map_exists() -> None: cited = sorted({match.group(1) for match in _LOCAL_PATH_IN_BACKTICKS.finditer(text)}) assert cited, ( - "MAINTAINING.md cites no local src/ or tests/ paths — the map is either empty " - "or its paths stopped being backticked, and this check has gone vacuous" + "MAINTAINING.md cites no local src/, tests/, or root-level *.md paths — " + "the map is either empty or its paths stopped being backticked, and this " + "check has gone vacuous" ) missing = [rel for rel in cited if not REPO_ROOT.joinpath(rel).exists()] diff --git a/tests/tools/test_inproc_provider.py b/tests/tools/test_inproc_provider.py index 81e437d..90658dc 100644 --- a/tests/tools/test_inproc_provider.py +++ b/tests/tools/test_inproc_provider.py @@ -348,3 +348,81 @@ async def describe(address: _Address) -> dict[str, Any]: "street": "1 Infinite Loop", "city": "Cupertino", } + + +# --------------------------------------------------------------------------- +# @tool — nested BaseModel schemas: $defs inlined, no dangling $ref +# --------------------------------------------------------------------------- + + +def _collect_refs(node: Any) -> list[str]: + """Recursively collect every ``$ref`` value in a schema structure.""" + refs: list[str] = [] + if isinstance(node, dict): + for key, value in node.items(): + if key == "$ref" and isinstance(value, str): + refs.append(value) + else: + refs.extend(_collect_refs(value)) + elif isinstance(node, list): + for item in node: + refs.extend(_collect_refs(item)) + return refs + + +def test_tool_schema_nested_basemodel_refs_are_inlined() -> None: + """A nested BaseModel parameter must not emit a dangling ``$ref``. + + Pydantic v2 emits nested models as ``{"$ref": "#/$defs/Foo"}`` plus a + top-level ``$defs``; ``ToolSchema`` has no ``$defs`` slot and the loop + ships only the four function-calling fields, so the refs are inlined. + """ + + @tool() + async def lookup(address: _Address) -> str: + return address.city + + schema = lookup.schema + assert _collect_refs(schema.properties) == [] + prop = schema.properties["address"] + assert prop["type"] == "object" + assert prop["properties"]["street"]["type"] == "string" + assert prop["properties"]["city"]["type"] == "string" + assert schema.required == ["address"] + + +class _Order(BaseModel): + """Doubly-nested model: its schema refs ``_Address``, which is itself a + def — inlining must recurse through defs, not just expand one level.""" + + shipping: _Address + qty: int + + +def test_tool_schema_doubly_nested_basemodel_refs_are_inlined() -> None: + @tool() + async def place(order: _Order) -> str: + return order.shipping.city + + schema = place.schema + assert _collect_refs(schema.properties) == [] + shipping = schema.properties["order"]["properties"]["shipping"] + assert shipping["properties"]["city"]["type"] == "string" + + +class _Node(BaseModel): + """Recursive model: ``children`` refs ``_Node`` itself, so the def has no + finite inline expansion.""" + + name: str + children: list[_Node] = [] + + +def test_tool_schema_recursive_model_raises_clear_error() -> None: + """A recursive model cannot be inlined; decoration must fail loudly rather + than emit a dangling ``$ref``.""" + with pytest.raises(ValueError, match=r"recursive \$ref cycle"): + + @tool() + async def walk(node: _Node) -> int: + return 0 diff --git a/tests/tools/test_mcp_provider.py b/tests/tools/test_mcp_provider.py index 6016ab7..694555e 100644 --- a/tests/tools/test_mcp_provider.py +++ b/tests/tools/test_mcp_provider.py @@ -91,26 +91,224 @@ def test_schema_translation_identity() -> None: assert schema.additionalProperties is False -def test_schema_translation_drops_unknown_top_level_keys() -> None: +def _collect_refs(node: Any) -> list[str]: + """Recursively collect every ``$ref`` value in a schema structure.""" + refs: list[str] = [] + if isinstance(node, dict): + for key, value in node.items(): + if key == "$ref" and isinstance(value, str): + refs.append(value) + else: + refs.extend(_collect_refs(value)) + elif isinstance(node, list): + for item in node: + refs.extend(_collect_refs(item)) + return refs + + +def test_schema_translation_inlines_defs_and_drops_unknown_top_level_keys() -> None: + """``$defs`` is not just dropped: local ``#/$defs/...`` refs are inlined + first, so no dangling pointer reaches the LLM-facing schema. Other unknown + top-level keys (``examples``, …) are still silently dropped — ToolSchema is + ``extra="forbid"``.""" schema = _to_tool_schema( { "type": "object", - "properties": {}, - "required": [], - "$defs": {"X": {"type": "object"}}, + "properties": {"addr": {"$ref": "#/$defs/Address"}}, + "required": ["addr"], + "$defs": { + "Address": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + } + }, "examples": [{}], } ) - # No KeyError; ToolSchema is extra="forbid" so we silently drop the - # unknown top-level keys at translation time. assert schema.type == "object" + assert _collect_refs(schema.properties) == [] + addr = schema.properties["addr"] + assert addr["type"] == "object" + assert addr["properties"]["city"] == {"type": "string"} + assert addr["required"] == ["city"] + assert schema.required == ["addr"] + + +def test_schema_translation_inlines_nested_defs_recursively() -> None: + """A def that itself refs another def must resolve to full depth.""" + schema = _to_tool_schema( + { + "type": "object", + "properties": {"order": {"$ref": "#/$defs/Order"}}, + "$defs": { + "Order": { + "type": "object", + "properties": {"shipping": {"$ref": "#/$defs/Address"}}, + }, + "Address": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } + ) + assert _collect_refs(schema.properties) == [] + shipping = schema.properties["order"]["properties"]["shipping"] + assert shipping["properties"]["city"] == {"type": "string"} + + +def test_schema_translation_ref_sibling_keys_are_preserved() -> None: + """Sibling keys next to a ``$ref`` (JSON Schema 2020-12 conjunctive + semantics) are merged over the resolved definition, not discarded.""" + schema = _to_tool_schema( + { + "type": "object", + "properties": {"addr": {"$ref": "#/$defs/Address", "description": "where"}}, + "$defs": {"Address": {"type": "object", "properties": {}}}, + } + ) + addr = schema.properties["addr"] + assert addr["description"] == "where" + assert addr["type"] == "object" + + +def test_schema_translation_leaves_unresolvable_refs_untouched() -> None: + """A ref the resolver cannot resolve locally (missing def, external URI) + passes through — inlining must not invent information the server schema + did not carry.""" + schema = _to_tool_schema( + { + "type": "object", + "properties": { + "x": {"$ref": "#/$defs/Missing"}, + "y": {"$ref": "https://schemas.example.com/ext"}, + }, + } + ) + assert schema.properties["x"]["$ref"] == "#/$defs/Missing" + assert schema.properties["y"]["$ref"] == "https://schemas.example.com/ext" + + +def test_schema_translation_defs_cycle_falls_back_to_empty_schema() -> None: + """A cyclic server schema (a def that transitively refs itself) has no + finite inline expansion; untrusted server data takes the same defensive + fallback as a non-object schema rather than aborting discovery.""" + schema = _to_tool_schema( + { + "type": "object", + "properties": {"node": {"$ref": "#/$defs/Node"}}, + "$defs": { + "Node": { + "type": "object", + "properties": {"child": {"$ref": "#/$defs/Node"}}, + } + }, + } + ) + assert schema.type == "object" + assert schema.properties == {} + assert schema.required == [] + + +def _exponential_schema(secret: str = "") -> dict[str, Any]: + """Build a compact acyclic schema whose inline form grows exponentially.""" + defs: dict[str, Any] = {"D14": {"type": "string", "description": secret}} + for index in range(13, -1, -1): + child = {"$ref": f"#/$defs/D{index + 1}"} + defs[f"D{index}"] = {"left": child, "right": child} + return { + "type": "object", + "properties": {"payload": {"$ref": "#/$defs/D0"}}, + "$defs": defs, + } + + +def test_schema_translation_over_budget_falls_back_without_logging_schema() -> None: + """BR-016 safely rejects over-budget remote schemas without content leakage.""" + secret = "SCHEMA_SECRET_MUST_NOT_BE_LOGGED" + with structlog.testing.capture_logs() as logs: + schema = _to_tool_schema(_exponential_schema(secret)) + + assert schema.properties == {} + failures = [ + entry for entry in logs if entry.get("event") == "mcp.input_schema.unresolvable_refs" + ] + assert failures == [ + { + "event": "mcp.input_schema.unresolvable_refs", + "reason": "invalid_local_refs", + "error_type": "ValueError", + "log_level": "warning", + } + ] + assert secret not in repr(logs) + + +def test_schema_translation_deep_ordinary_nesting_falls_back_safely() -> None: + """BR-016 raw traversal recursion takes the safe content-free MCP fallback.""" + secret = "DEEP_SCHEMA_SECRET_MUST_NOT_BE_LOGGED" + nested: Any = secret + for index in range(1_100): + nested = {"level": index, "child": [nested]} + input_schema = {"type": "object", "properties": {"payload": nested}} + + with structlog.testing.capture_logs() as logs: + schema = _to_tool_schema(input_schema) + + assert schema.properties == {} + failures = [ + entry for entry in logs if entry.get("event") == "mcp.input_schema.unresolvable_refs" + ] + assert failures == [ + { + "event": "mcp.input_schema.unresolvable_refs", + "reason": "invalid_local_refs", + "error_type": "ValueError", + "log_level": "warning", + } + ] + assert secret not in repr(logs) + + +async def test_attach_continues_after_over_budget_sibling_schema( + controllable_server: ControllableServer, +) -> None: + """BR-016 one hostile schema cannot abort discovery of sibling tools.""" + controllable_server.set_tool_catalog( + [ + _tool_def("hostile", schema=_exponential_schema()), + _tool_def( + "healthy", + schema={"type": "object", "properties": {"q": {"type": "string"}}}, + ), + ] + ) + registry = Registry() + await MCPProvider(make_controllable_client(controllable_server)).attach(registry) + + assert registry.get("hostile").schema.properties == {} + assert registry.get("healthy").schema.properties == {"q": {"type": "string"}} def test_schema_translation_falls_back_for_non_object_top_level() -> None: - schema = _to_tool_schema({"type": "string"}) + """BR-016 non-object fallback logs type metadata, never remote content.""" + secret = "REMOTE_TYPE_SECRET_MUST_NOT_BE_LOGGED" + with structlog.testing.capture_logs() as logs: + schema = _to_tool_schema({"type": secret}) assert schema.type == "object" assert schema.properties == {} assert schema.required == [] + failures = [entry for entry in logs if entry.get("event") == "mcp.input_schema.non_object"] + assert failures == [ + { + "event": "mcp.input_schema.non_object", + "reason": "non_object_type", + "received_type": "str", + "log_level": "warning", + } + ] + assert secret not in repr(logs) # --------------------------------------------------------------------------- diff --git a/tests/tools/test_schema_refs.py b/tests/tools/test_schema_refs.py new file mode 100644 index 0000000..3a2c94d --- /dev/null +++ b/tests/tools/test_schema_refs.py @@ -0,0 +1,102 @@ +"""Tests for bounded local JSON-Schema reference expansion (BR-016).""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any + +import pytest + +from fifty_agent_sdk.tools._schema_refs import _MAX_EXPANDED_NODES, inline_local_refs + + +def test_inline_local_refs_resolves_shared_refs_without_mutating_inputs() -> None: + """BR-016 preserves ordinary shared-ref expansion and input immutability.""" + node = { + "left": {"$ref": "#/$defs/Leaf"}, + "right": {"$ref": "#/$defs/Leaf", "description": "right leaf"}, + } + defs = {"Leaf": {"type": "object", "properties": {"value": {"type": "string"}}}} + original_node = deepcopy(node) + original_defs = deepcopy(defs) + + resolved = inline_local_refs(node, defs) + + assert resolved["left"]["properties"]["value"] == {"type": "string"} + assert resolved["right"]["description"] == "right leaf" + assert node == original_node + assert defs == original_defs + + +def test_inline_local_refs_leaves_missing_and_external_refs_untouched() -> None: + """BR-016 does not invent definitions for refs outside the local map.""" + node = { + "missing": {"$ref": "#/$defs/Missing"}, + "external": {"$ref": "https://example.invalid/schema"}, + } + assert inline_local_refs(node, {}) == node + + +def test_inline_local_refs_rejects_cycles_before_budget_exhaustion() -> None: + """BR-016 preserves the precise recursive-cycle failure.""" + defs = {"Node": {"properties": {"child": {"$ref": "#/$defs/Node"}}}} + with pytest.raises(ValueError, match=r"recursive \$ref cycle.*Node -> Node"): + inline_local_refs({"$ref": "#/$defs/Node"}, defs) + + +def _reference_chain(length: int) -> tuple[dict[str, str], dict[str, Any]]: + defs: dict[str, Any] = {} + for index in range(length): + defs[f"D{index}"] = ( + {"$ref": f"#/$defs/D{index + 1}"} if index + 1 < length else {"type": "string"} + ) + return {"$ref": "#/$defs/D0"}, defs + + +def test_inline_local_refs_accepts_depth_boundary() -> None: + """BR-016 preserves the documented 32-reference depth boundary.""" + node, defs = _reference_chain(32) + assert inline_local_refs(node, defs) == {"type": "string"} + + +def test_inline_local_refs_rejects_over_depth_boundary() -> None: + """BR-016 preserves rejection beyond the 32-reference depth boundary.""" + node, defs = _reference_chain(33) + with pytest.raises(ValueError, match="exceeded 32 levels"): + inline_local_refs(node, defs) + + +def test_inline_local_refs_accepts_exact_node_budget() -> None: + """BR-016 permits a schema whose resolver visits exactly the node budget.""" + node = [0] * (_MAX_EXPANDED_NODES - 1) + assert inline_local_refs(node, {}) == node + + +def test_inline_local_refs_rejects_one_node_over_budget() -> None: + """BR-016 rejects before visiting beyond the shared work/output budget.""" + node = [0] * _MAX_EXPANDED_NODES + with pytest.raises(ValueError, match="node budget"): + inline_local_refs(node, {}) + + +def test_inline_local_refs_bounds_acyclic_exponential_fanout() -> None: + """BR-016 bounds acyclic diamond fan-out that a depth limit cannot contain.""" + defs: dict[str, Any] = {"D14": {"type": "string"}} + for index in range(13, -1, -1): + child = {"$ref": f"#/$defs/D{index + 1}"} + defs[f"D{index}"] = {"left": child, "right": child} + + with pytest.raises(ValueError, match="node budget"): + inline_local_refs({"$ref": "#/$defs/D0"}, defs) + + +def test_inline_local_refs_contains_ordinary_nesting_recursion_error() -> None: + """BR-016 ordinary dict/list depth cannot escape as raw RecursionError.""" + node: Any = "SCHEMA_SECRET_MUST_NOT_ESCAPE" + for index in range(1_100): + node = {"level": index, "child": [node]} + + with pytest.raises(ValueError, match="schema traversal exceeded the safe nesting limit") as exc: + inline_local_refs(node, {}) + assert isinstance(exc.value.__cause__, RecursionError) + assert "SCHEMA_SECRET_MUST_NOT_ESCAPE" not in str(exc.value)