From e1a9fd03a794236f92398e5985acec8ce5d80e57 Mon Sep 17 00:00:00 2001 From: eandualem Date: Fri, 2 Oct 2026 11:20:58 +0300 Subject: [PATCH 1/5] refactor: restore module boundaries and remove obsolete execution paths Separate turn snapshots, task scheduling, provider transports and database repositories so each module owns one responsibility. Enforce application layering and remove execution helpers superseded by native Pydantic AI runs. Preserve validated state and accumulated usage, roll back partially started services, and propagate database failures consistently. Add regressions for the confirmed defects. --- AGENTS.md | 20 +- README.md | 7 + .../app/assistant/_prompt_builder.py | 3 - .../app/assistant/_serialization.py | 13 +- .../app/assistant/_session_store.py | 2 +- .../app/assistant/interface.py | 10 +- src/assistant_runtime/app/event_log/_store.py | 15 +- src/assistant_runtime/app/ingress/factory.py | 2 +- .../app/ingress/interface.py | 15 +- src/assistant_runtime/app/routes/artifacts.py | 27 + src/assistant_runtime/app/socketio_server.py | 15 +- .../app/streaming/_runner.py | 494 +------ .../app/streaming/_snapshot.py | 237 ++++ .../app/streaming/_telemetry.py | 223 +++ src/assistant_runtime/app/streaming/_turn.py | 1 + .../app/streaming/interface.py | 8 +- src/assistant_runtime/app/tasks/_runner.py | 215 +++ src/assistant_runtime/app/tasks/_store.py | 37 +- src/assistant_runtime/app/tasks/interface.py | 205 +-- src/assistant_runtime/app/voice/_codex.py | 283 +--- src/assistant_runtime/app/voice/_codex_rpc.py | 202 +++ .../app/voice/_codex_schema.py | 94 ++ src/assistant_runtime/app/voice/_state.py | 5 +- src/assistant_runtime/app/voice/_transport.py | 25 +- src/assistant_runtime/app/voice/interface.py | 12 +- src/assistant_runtime/base/__init__.py | 7 +- src/assistant_runtime/base/_logging.py | 6 - src/assistant_runtime/base/instrument.py | 57 - src/assistant_runtime/base/lifecycle.py | 7 +- src/assistant_runtime/main.py | 7 +- .../services/actions/_store.py | 27 +- .../services/database/repositories.py | 1219 ----------------- .../database/repositories/__init__.py | 34 + .../database/repositories/_actions.py | 108 ++ .../database/repositories/_artifacts.py | 241 ++++ .../services/database/repositories/_events.py | 63 + .../database/repositories/_host_state.py | 104 ++ .../services/database/repositories/_inbox.py | 80 ++ .../services/database/repositories/_oauth.py | 78 ++ .../database/repositories/_ordering.py | 15 + .../database/repositories/_sessions.py | 345 +++++ .../database/repositories/_settings.py | 54 + .../services/database/repositories/_tasks.py | 165 +++ .../services/database/repositories/_traces.py | 65 + .../services/llm/__init__.py | 3 +- .../services/llm/interface.py | 133 +- .../services/oauth/_credentials.py | 120 ++ .../services/oauth/interface.py | 192 +-- .../services/oauth/models.py | 65 + .../tools/builtin/_subagent_executor.py | 31 +- tests/compatibility/test_codex_fast.py | 8 +- tests/compatibility/test_usage.py | 23 + tests/integration/test_validation.py | 5 +- tests/unit/app/assistant/test_interface.py | 10 +- .../unit/app/assistant/test_prompt_builder.py | 65 +- .../unit/app/assistant/test_prompt_record.py | 12 +- tests/unit/app/event_log/test_interface.py | 19 + tests/unit/app/ingress/test_factory.py | 6 +- tests/unit/app/ingress/test_interface.py | 16 +- tests/unit/app/routes/test_artifacts.py | 18 + .../app/streaming/test_client_error_detail.py | 16 +- tests/unit/app/streaming/test_interface.py | 26 +- tests/unit/app/tasks/test_interface.py | 26 + tests/unit/app/test_socketio_server.py | 13 + tests/unit/app/voice/test_codex.py | 6 +- tests/unit/base/test_instrument.py | 51 - tests/unit/base/test_lifecycle.py | 27 + tests/unit/base/test_logging.py | 13 +- tests/unit/services/actions/test_interface.py | 26 + tests/unit/services/llm/test_codex_only.py | 13 +- tests/unit/services/llm/test_interface.py | 250 +--- .../tools/builtin/test_subagent_executor.py | 16 +- tests/unit/test_imports.py | 61 +- tests/unit/test_main_lifespan.py | 2 +- 74 files changed, 3075 insertions(+), 3049 deletions(-) create mode 100644 src/assistant_runtime/app/streaming/_snapshot.py create mode 100644 src/assistant_runtime/app/streaming/_telemetry.py create mode 100644 src/assistant_runtime/app/tasks/_runner.py create mode 100644 src/assistant_runtime/app/voice/_codex_rpc.py create mode 100644 src/assistant_runtime/app/voice/_codex_schema.py delete mode 100644 src/assistant_runtime/base/instrument.py delete mode 100644 src/assistant_runtime/services/database/repositories.py create mode 100644 src/assistant_runtime/services/database/repositories/__init__.py create mode 100644 src/assistant_runtime/services/database/repositories/_actions.py create mode 100644 src/assistant_runtime/services/database/repositories/_artifacts.py create mode 100644 src/assistant_runtime/services/database/repositories/_events.py create mode 100644 src/assistant_runtime/services/database/repositories/_host_state.py create mode 100644 src/assistant_runtime/services/database/repositories/_inbox.py create mode 100644 src/assistant_runtime/services/database/repositories/_oauth.py create mode 100644 src/assistant_runtime/services/database/repositories/_ordering.py create mode 100644 src/assistant_runtime/services/database/repositories/_sessions.py create mode 100644 src/assistant_runtime/services/database/repositories/_settings.py create mode 100644 src/assistant_runtime/services/database/repositories/_tasks.py create mode 100644 src/assistant_runtime/services/database/repositories/_traces.py create mode 100644 src/assistant_runtime/services/oauth/_credentials.py create mode 100644 src/assistant_runtime/services/oauth/models.py delete mode 100644 tests/unit/base/test_instrument.py diff --git a/AGENTS.md b/AGENTS.md index cabdbd2a..b929cc16 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -105,7 +105,9 @@ execution, history, serialization, or the upstream dependency; see `app/access`, `app/assistant`, `app/streaming`, `app/voice`, `app/ingress`, `app/tasks`, `app/heartbeat`, `app/event_log`; the leaf modules `base`, `artifacts`, `host_context`, `principal`, `config` - and the `app/routes` package are exempt). + and the `app/routes` package are exempt). Ingress has no configurable + options and omits `config.py`; MCP uses its server configuration file and + omits `config.py` and `deps.py`. `config.py` (a frozen pydantic model nested into `AppSettings`), `deps.py` (FastAPI `Depends` accessors reading `app.state`), `factory.py` (`register_(app_state, @@ -118,7 +120,8 @@ execution, history, serialization, or the upstream dependency; see database, oauth, llm, history, media, decisions, actions, host_state, mcp, artifacts, tools, assistant, streaming, voice, ingress, tasks, heartbeat, event_log. `LifecycleManager` starts in that order, stops in reverse, and rolls back - on a failed start. `RuntimeSettings` is created after `start_all()` and + on a failed or cancelled start, including the partially started component. + `RuntimeSettings` is created after `start_all()` and attached through each service's `set_runtime_settings()`. - **Layering, bottom up.** `base` (lifecycle, protocols, resilience, exceptions, the domain event hub), `artifacts` (assistant profiles: the artifact schema, the @@ -137,8 +140,11 @@ execution, history, serialization, or the upstream dependency; see are the top. `tests/unit/test_imports.py` asserts that nothing below the top layer imports `app` and that `_`-prefixed files stay inside their module (what another module needs is re-exported from the module's - `__init__.py` or a public file such as `services/tools/request_context.py`); - a new cross-package import must keep it green. `main.py` builds nothing at + `__init__.py` or a public file such as `services/tools/request_context.py`). + It also enforces the application-module dependency graph, including that + transport edges are never imported by application services. Ingress receives + its event sink from composition; Socket.IO delivery stays at the edge. + A new cross-package import must keep it green. `main.py` builds nothing at import: uvicorn runs the `create_asgi_app` factory, and loading `.env` is the CLI's job. - **Configuration has three tiers**, resolved once per request by @@ -311,6 +317,8 @@ execution, history, serialization, or the upstream dependency; see `TASKS__MAX_CONCURRENT` at once; a turn that asks the host for an action fails its task. `task_finished` is published on `app.state.events` and is not forwarded to Socket.IO; nothing is steered into the parent session. + Its private `_runner.py` owns queue reservations, ordering, capacity and + cancellation; the service owns admission, records and request construction. With Postgres a restart marks unfinished tasks `interrupted` and replays nothing; without it tasks live in process memory. A *persistent agent* is the continuing form: the host starts it for a (profile, subject), one @@ -456,7 +464,9 @@ names, ids or private hostnames). - **Shared turn pipeline.** Every request kind (new message, host-tool continuation, promoted steering) is described by `TurnPlanner` (`app/streaming/_turn.py`) and executed by `TurnRunner` - (`app/streaming/_runner.py`); `_agent_run.py` maps public Pydantic AI events + (`app/streaming/_runner.py`), which delegates native message capture and + persistence to `_snapshot.py` and diagnostics to `_telemetry.py`; + `_agent_run.py` maps public Pydantic AI events and applies steering/screen policy through native capability hooks; `_host_tool.py` keeps the pending host-tool state. The non-streaming `POST /api/chat` collects the same stream through diff --git a/README.md b/README.md index 44a90f27..49df07c0 100644 --- a/README.md +++ b/README.md @@ -108,6 +108,13 @@ src/assistant_runtime/cli/ chat, serve, doctor, docs and migrations docs/ and alembic/ documentation and database migrations ``` +Within each service, `interface.py` owns the public operations and lifecycle; +private modules hold the implementation. Database queries are grouped by record +domain behind `services/database/repositories`. The turn runner delegates native +message snapshots and diagnostics to separate modules, and background tasks use +one scheduler for ordering, capacity and cancellation. Transport adapters own +Socket.IO and voice protocol details. The import checks enforce these boundaries. + ## License MIT — see [LICENSE](https://github.com/eandualem/assistant-runtime/blob/main/LICENSE). diff --git a/src/assistant_runtime/app/assistant/_prompt_builder.py b/src/assistant_runtime/app/assistant/_prompt_builder.py index ff3d4d5e..343ab4f9 100644 --- a/src/assistant_runtime/app/assistant/_prompt_builder.py +++ b/src/assistant_runtime/app/assistant/_prompt_builder.py @@ -22,7 +22,6 @@ NavigationTarget, ) from assistant_runtime.services.history.models import WorkingMemory -from assistant_runtime.services.tools.models import ToolSet # --- Fragment builders --- @@ -236,7 +235,6 @@ def host_context_prompt(host_context: dict[str, Any] | None) -> str: def build_system_prompt( *, - available_tools: ToolSet, session_context: dict[str, Any], working_memory: bool = False, artifacts: dict[str, str], @@ -249,7 +247,6 @@ def build_system_prompt( memory (dynamic, last) when ``working_memory`` is set. Args: - available_tools: Tools available for this request. session_context: Session context dict (may contain working memory). working_memory: Whether the session's working memory is rendered. artifacts: Artifact name→text map, normally ``ArtifactService.active_texts()``. diff --git a/src/assistant_runtime/app/assistant/_serialization.py b/src/assistant_runtime/app/assistant/_serialization.py index 342730fd..34b97e3d 100644 --- a/src/assistant_runtime/app/assistant/_serialization.py +++ b/src/assistant_runtime/app/assistant/_serialization.py @@ -17,6 +17,7 @@ from pydantic_ai.messages import ( BinaryContent, DocumentUrl, + FileUrl, ImageUrl, ModelMessage, ModelMessagesTypeAdapter, @@ -364,16 +365,8 @@ def _omit_tool_output_media(tool_name: str, content: Any) -> Any: return "[Inspected current screen]" if isinstance(content, BinaryContent): return "[Binary content omitted]" - try: - import pydantic_ai.messages as pydantic_messages - except ImportError: # pragma: no cover - pydantic_messages = None # type: ignore[assignment] - file_url_cls = getattr(pydantic_messages, "FileUrl", None) - if file_url_cls is not None and isinstance(content, file_url_cls): - result: dict[str, Any] = {"url": content.url} - if hasattr(content, "kind"): - result["kind"] = content.kind - return result + if isinstance(content, FileUrl): + return {"url": content.url, "kind": content.kind} if isinstance(content, dict): return {key: _omit_tool_output_media(tool_name, value) for key, value in content.items()} if isinstance(content, (list, tuple)): diff --git a/src/assistant_runtime/app/assistant/_session_store.py b/src/assistant_runtime/app/assistant/_session_store.py index ad9073a1..7dd579e6 100644 --- a/src/assistant_runtime/app/assistant/_session_store.py +++ b/src/assistant_runtime/app/assistant/_session_store.py @@ -417,7 +417,7 @@ async def save_prompt_snapshot(self, snapshot_hash: str, content: str) -> None: """Keep the stable text a prompt record names (stored once per hash).""" self._snapshots[snapshot_hash] = content self._snapshots.move_to_end(snapshot_hash) - while len(self._snapshots) > _SNAPSHOT_CACHE_SIZE: + while self._db is not None and len(self._snapshots) > _SNAPSHOT_CACHE_SIZE: self._snapshots.popitem(last=False) if self._db is not None: await self._db.save_prompt_snapshot(snapshot_hash, content) diff --git a/src/assistant_runtime/app/assistant/interface.py b/src/assistant_runtime/app/assistant/interface.py index ba9983e4..cfda25cb 100644 --- a/src/assistant_runtime/app/assistant/interface.py +++ b/src/assistant_runtime/app/assistant/interface.py @@ -6,7 +6,6 @@ from __future__ import annotations -import asyncio import inspect import time from typing import TYPE_CHECKING, Any @@ -177,7 +176,6 @@ async def preview_prompt(self, profile: str | None, subject: str | None) -> dict ).for_subject(subject) texts, extras, versions = await artifacts.prompt_inputs() prompt = build_system_prompt( - available_tools=self._tools.get_available_tools(None), session_context={}, artifacts=texts, profile=artifacts.profile, @@ -200,7 +198,6 @@ async def profile_prompt(self, profile: str) -> tuple[str, dict[str, int | None] artifacts = self._artifacts.for_profile(profile) texts, extras, versions = await artifacts.prompt_inputs() prompt = build_system_prompt( - available_tools=self._tools.get_available_tools(None), session_context={}, artifacts=texts, profile=artifacts.profile, @@ -268,9 +265,7 @@ async def prepare_agent_context( if inspect.isawaitable(deps): deps = await deps - inputs_task = asyncio.create_task(artifacts_service.prompt_inputs()) - - # 2. Config resolution can run while prompt inputs load. + # 2. Resolve synchronous configuration before loading prompt inputs. effective = resolve_effective_config( self._config, self._runtime_settings, request_config ) @@ -282,10 +277,9 @@ async def prepare_agent_context( ) # 3. Artifacts + system prompt - artifacts, artifact_extras, artifact_versions = await inputs_task + artifacts, artifact_extras, artifact_versions = await artifacts_service.prompt_inputs() prompt_result = build_system_prompt( - available_tools=available_tools, session_context=session_context, working_memory=effective.enable_working_memory, artifacts=artifacts, diff --git a/src/assistant_runtime/app/event_log/_store.py b/src/assistant_runtime/app/event_log/_store.py index 34bcf938..beb8dc4e 100644 --- a/src/assistant_runtime/app/event_log/_store.py +++ b/src/assistant_runtime/app/event_log/_store.py @@ -8,6 +8,7 @@ import asyncio import contextlib from collections.abc import AsyncIterator +from copy import deepcopy from dataclasses import fields, replace from datetime import UTC, datetime from typing import Any, Protocol @@ -52,19 +53,19 @@ async def transaction(self) -> AsyncIterator[InMemoryEventStore]: async def create_if_new(self, **values: Any) -> tuple[EventRecord, bool]: key = (values["source"], values["event_id"]) if key in self._keys: - return self._events[self._keys[key]], False - record = EventRecord(id=self._next, created_at=datetime.now(UTC), **values) + return deepcopy(self._events[self._keys[key]]), False + record = EventRecord(id=self._next, created_at=datetime.now(UTC), **deepcopy(values)) self._events[record.id] = record self._keys[key] = record.id self._next += 1 - return record, True + return deepcopy(record), True async def get(self, event_id: int, *, lock: bool = False) -> EventRecord | None: - return self._events.get(event_id) + return deepcopy(self._events.get(event_id)) async def update(self, event_id: int, **values: Any) -> EventRecord: - record = self._events[event_id] = replace(self._events[event_id], **values) - return record + record = self._events[event_id] = replace(self._events[event_id], **deepcopy(values)) + return deepcopy(record) async def list(self, *, after: int, limit: int, **filters: Any) -> list[EventRecord]: records = [ @@ -73,7 +74,7 @@ async def list(self, *, after: int, limit: int, **filters: Any) -> list[EventRec if r.id > after and all(value is None or getattr(r, name) == value for name, value in filters.items()) ] - return sorted(records, key=lambda r: r.id)[:limit] + return deepcopy(sorted(records, key=lambda r: r.id)[:limit]) class DatabaseEventStore: diff --git a/src/assistant_runtime/app/ingress/factory.py b/src/assistant_runtime/app/ingress/factory.py index efacbc7f..0260e9ab 100644 --- a/src/assistant_runtime/app/ingress/factory.py +++ b/src/assistant_runtime/app/ingress/factory.py @@ -17,7 +17,7 @@ async def register_ingress(app_state: Any, lifecycle: LifecycleManager) -> None: assistant_service=app_state.assistant_service, streaming_service=app_state.streaming_service, database_service=getattr(app_state, "database_service", None), - socket_server=getattr(app_state, "sio", None), + event_sink=getattr(app_state, "ingress_event_sink", None), ) app_state.ingress_service = service app_state.streaming_service.attach_ingress(service) diff --git a/src/assistant_runtime/app/ingress/interface.py b/src/assistant_runtime/app/ingress/interface.py index f20c1bac..9afc78df 100644 --- a/src/assistant_runtime/app/ingress/interface.py +++ b/src/assistant_runtime/app/ingress/interface.py @@ -14,13 +14,13 @@ import asyncio import contextlib import uuid +from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any from loguru import logger from assistant_runtime.app.assistant.exceptions import SessionError from assistant_runtime.app.assistant.models import AssistantRequest -from assistant_runtime.app.socketio_server import socket_event_name if TYPE_CHECKING: from assistant_runtime.app.assistant.interface import AssistantService @@ -42,12 +42,12 @@ def __init__( assistant_service: AssistantService, streaming_service: StreamingService, database_service: DatabaseService | None = None, - socket_server: Any | None = None, + event_sink: Callable[[str, dict[str, Any]], Awaitable[None]] | None = None, ) -> None: self._assistant = assistant_service self._streaming = streaming_service self._db = database_service - self._sio = socket_server + self._event_sink = event_sink self._queued: list[dict[str, Any]] = [] # the inbox when there is no database self._background: set[asyncio.Task[None]] = set() self._started = False @@ -186,15 +186,12 @@ def _run_in_background(self, request: AssistantRequest) -> None: task.add_done_callback(self._background.discard) async def _run(self, request: AssistantRequest) -> None: - """Run the promoted turn; events go to the session's Socket.IO room, if any.""" - room = f"session:{request.session_id}" + """Run the promoted turn and forward its events through the configured sink.""" try: async for event in self._streaming.stream_message(request): - if self._sio is not None: + if self._event_sink is not None: with contextlib.suppress(Exception): - await self._sio.emit( - socket_event_name(event), event, room=room, namespace="/assistant" - ) + await self._event_sink(request.session_id, event) except asyncio.CancelledError: raise except Exception as e: diff --git a/src/assistant_runtime/app/routes/artifacts.py b/src/assistant_runtime/app/routes/artifacts.py index 1d295daf..63dcbba7 100644 --- a/src/assistant_runtime/app/routes/artifacts.py +++ b/src/assistant_runtime/app/routes/artifacts.py @@ -29,6 +29,7 @@ ) from assistant_runtime.services.artifacts.interface import ArtifactService from assistant_runtime.services.artifacts.models import Actor, MutationResult +from assistant_runtime.services.database.exceptions import DatabaseError router = APIRouter(prefix="/artifacts", tags=["artifacts"]) proposals_router = APIRouter(prefix="/artifact-proposals", tags=["artifacts"]) @@ -272,6 +273,8 @@ async def get_artifact(name: str, artifacts: ScopedArtifactDep, principal: Princ """The active version of an artifact, or its default text when none is active.""" try: row = await artifacts.get_active(name) + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -302,6 +305,8 @@ async def get_artifact_history( """Version history for an artifact, newest first.""" try: return [row.to_dict() for row in await artifacts.history(name, limit=limit)] + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -316,6 +321,8 @@ async def get_artifact_version( """One version as a proposal record: its content, the active text and a unified diff.""" try: return await artifacts.version_record(name, version) + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -333,6 +340,8 @@ async def propose_artifact( return await _propose( artifacts, name, body.content, who, body.expected_version, body.rationale ) + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -348,6 +357,8 @@ async def update_artifact( who = _who(admin, body) try: return await _update(artifacts, name, body.content, who, body.expected_version) + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -367,6 +378,8 @@ async def approve_artifact( who = _who(admin, body) try: return await _activate(artifacts, name, version, who, rollback=False) + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -386,6 +399,8 @@ async def reject_artifact( who = _who(admin, body) try: return await _reject(artifacts, name, version, who, body.reason if body else None) + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -400,6 +415,8 @@ async def rollback_artifact( """Reactivate an earlier version of an artifact.""" try: return await _activate(artifacts, name, version, admin.id, rollback=True) + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -442,6 +459,8 @@ async def artifact_action( ) except HTTPException: raise + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -454,6 +473,8 @@ async def delete_artifact(name: str, artifacts: ScopedArtifactDep, admin: AdminD """Delete every stored version; the profile's default text applies again.""" try: count = await artifacts.delete(name, actor=Actor("host", admin.id)) + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -487,6 +508,8 @@ async def list_proposals( records: list[dict] = [] for name in names: records.extend(await service.for_profile(name).proposals(status, limit + 1, before_id)) + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -508,6 +531,8 @@ async def list_subjects( try: artifacts = get_artifact_service(request).for_profile(profile) return {"profile": artifacts.profile.name, "subjects": await artifacts.subjects()} + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except Exception as e: @@ -529,6 +554,8 @@ async def preview_prompt( """ try: return await request.app.state.assistant_service.preview_prompt(profile, subject) + except DatabaseError: + raise except ArtifactError as e: raise _http_error(e) from e except ValueError as e: # a required artifact without text diff --git a/src/assistant_runtime/app/socketio_server.py b/src/assistant_runtime/app/socketio_server.py index 6b61ea3e..2249097a 100644 --- a/src/assistant_runtime/app/socketio_server.py +++ b/src/assistant_runtime/app/socketio_server.py @@ -9,7 +9,7 @@ import asyncio import contextlib import inspect -from collections.abc import Callable +from collections.abc import Awaitable, Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Any from weakref import WeakValueDictionary @@ -90,6 +90,19 @@ async def forward(event: dict[str, Any]) -> None: return forward +def forward_ingress_events( + sio: socketio.AsyncServer, +) -> Callable[[str, dict[str, Any]], Awaitable[None]]: + """Adapt promoted ingress turns to the session's Socket.IO room.""" + + async def forward(session_id: str, event: dict[str, Any]) -> None: + await sio.emit( + socket_event_name(event), event, room=f"session:{session_id}", namespace="/assistant" + ) + + return forward + + def create_sio(access: AccessConfig | None = None) -> socketio.AsyncServer: """Create and configure the Socket.IO async server. diff --git a/src/assistant_runtime/app/streaming/_runner.py b/src/assistant_runtime/app/streaming/_runner.py index b0150552..d6a893eb 100644 --- a/src/assistant_runtime/app/streaming/_runner.py +++ b/src/assistant_runtime/app/streaming/_runner.py @@ -25,7 +25,6 @@ import time import uuid from collections.abc import AsyncIterator, Awaitable, Callable -from dataclasses import dataclass, field, replace from datetime import UTC, datetime from typing import TYPE_CHECKING, Any @@ -34,49 +33,34 @@ from pydantic_ai.exceptions import RunCancelled, UsageLimitExceeded from pydantic_ai.messages import ( ModelMessage, - ModelRequest, - ModelResponse, - TextPart, - ThinkingPart, - ToolCallPart, - ToolReturnPart, - UserPromptPart, ) from pydantic_ai.usage import RunUsage from assistant_runtime.app.access.exceptions import AccessDeniedError -from assistant_runtime.app.assistant import ( - ActionStatus, - build_assistant_message_content, - dump_model_messages, -) from assistant_runtime.app.assistant.exceptions import SessionError from assistant_runtime.app.assistant.models import HoldDecision from assistant_runtime.app.assistant.usage import ( - cache_counts, - merge_usage, - response_service_tiers, usage_dict, with_auxiliary, ) +from assistant_runtime.app.streaming import _telemetry from assistant_runtime.app.streaming._agent_run import TurnPolicy, iterate_run from assistant_runtime.app.streaming._control import TurnControl from assistant_runtime.app.streaming._coordinator import EventCoordinator from assistant_runtime.app.streaming._event_builder import ( - make_debug_agent_config_event, make_debug_completed_event, - make_debug_error_event, make_debug_final_response_event, - make_debug_history_event, make_debug_request_event, - make_debug_system_prompt_event, - make_debug_tool_selection_event, - make_debug_usage_event, - make_error_event, ) from assistant_runtime.app.streaming._host_tool import pending_call_payloads, queued_payload +from assistant_runtime.app.streaming._snapshot import ( + RunSnapshot, + _IncompleteRun, + capture_result, + mark_unanswered_calls, + persist_snapshot, +) from assistant_runtime.app.streaming.exceptions import InvalidDecisionError -from assistant_runtime.services.llm.exceptions import LLMCallError, classify_llm_error from assistant_runtime.services.tools.request_context import ( assistant_request_context, get_current_telegram_chat_binding, @@ -90,92 +74,9 @@ from assistant_runtime.app.streaming._turn import TurnPlan from assistant_runtime.app.streaming.config import StreamingConfig from assistant_runtime.services.database.interface import DatabaseService - from assistant_runtime.services.history.interface import HistoryProcessor, HistoryService + from assistant_runtime.services.history.interface import HistoryService from assistant_runtime.services.tools.interface import ToolService -_LLM_ERROR_TYPES = { - "RATE_LIMIT": "rate_limit", - "SERVER_ERROR": "provider_error", - "CONNECTION_ERROR": "connection_error", - "TIMEOUT": "timeout", - "AUTH_ERROR": "provider_auth", - "CLIENT_ERROR": "provider_client_error", -} - - -def format_error_message(message: str, trace_id: str | None) -> str: - """Append the trace id so client-visible errors can be correlated with traces.""" - if trace_id is None or "[trace_id:" in message: - return message - return f"{message} [trace_id: {trace_id}]" - - -def _describe_error(exc: Exception, *, detail: bool = True) -> tuple[str, str, bool]: - """``(message, error_type, retry_allowed)`` for an exception raised by the run. - - ``detail`` (``STREAMING__CLIENT_ERROR_DETAIL``) decides whether the - exception text itself reaches the client; the log always has it. - """ - if isinstance(exc, InvalidDecisionError): - return str(exc), "invalid_decision", False - if isinstance(exc, LLMCallError): - llm_error: LLMCallError | None = exc - else: - classified = classify_llm_error(exc) - llm_error = None if classified.error_category == "UNKNOWN" else classified - if llm_error is None: - message = f"Request failed: {exc.__class__.__name__}: {exc}" if detail else "Request failed" - return message, "internal", False - error_type = _LLM_ERROR_TYPES.get(llm_error.error_category, "provider_error") - message = str(llm_error) if detail else f"LLM call failed ({llm_error.error_category})" - return message, error_type, llm_error.retry_allowed - - -def _keep_row_auxiliary( - usage: dict[str, Any] | None, row: dict[str, Any] | None -) -> dict[str, Any] | None: - """Keep row-owned memory totals and auxiliary sections missing from a snapshot.""" - existing = ((row or {}).get("usage") or {}).get("auxiliary") or {} - for name, part in existing.items(): - auxiliary = (usage or {}).get("auxiliary") or {} - if name == "working_memory" and name in auxiliary: - # Background extraction owns this cumulative total; the turn's - # earlier copy must neither overwrite it nor be added a second time. - usage = {**usage, "auxiliary": {**auxiliary, name: part}} - elif name not in auxiliary: - usage = with_auxiliary(usage, name, part) - return usage - - -def _persistence_failure(what: str, exc: BaseException, *, detail: bool) -> str: - """The client-visible text for a snapshot that could not be saved.""" - return f"{what} could not be saved: {exc}" if detail else f"{what} could not be saved" - - -@dataclass -class _RunState: - """What one agent run left behind, carried across steering follow-ups.""" - - all_messages: list[ModelMessage] = field(default_factory=list) - assistant_messages: list[ModelMessage] = field(default_factory=list) - assistant_segments: list[dict[str, Any]] | None = None - run_usage: Any = None - usage: dict[str, int] | None = None - service_tiers: list[dict[str, Any]] | None = None - output: Any = None - final_output: str | None = None - pending_tool_call: dict[str, Any] | None = None - persisted: bool = False - cancelled: bool = False - # How calls left without a result by this turn are recorded (see ActionStatus). - interrupted_status: ActionStatus = "superseded" - history: HistoryProcessor | None = None - # The usage written on the assistant row: the run's own plus auxiliary work. - stored_usage: dict[str, Any] | None = None - # What the latest run of the turn was given (app/assistant/prompt_record). - prompt_record: dict[str, Any] | None = None - prompt_snapshot: str | None = None - class TurnRunner: """Run ``TurnPlan``s. One instance per ``StreamingService``.""" @@ -212,7 +113,7 @@ async def run( request = plan.request session_id = plan.session_id session_context = plan.session_context - state = _RunState( + state = RunSnapshot( assistant_messages=list(plan.prior_assistant_messages), usage=plan.prior_usage ) resolved_model = "unknown" @@ -228,7 +129,7 @@ async def run( # before dependencies/model setup can be cancelled, using the same row. if plan.accepted_tool_result is not None: state.assistant_messages.append(plan.accepted_tool_result) - await self._persist(plan, state) + await persist_snapshot(self._sessions, plan, state) await self._clear_pending(plan) if plan.next_pending is not None: # More host calls from the same response wait; hand the next one @@ -292,7 +193,7 @@ async def run( else self._history.processor(session_context, collect_debug=emit_debug) ) if emit_debug: - for event in self._setup_debug_events(ctx, session_id): + for event in _telemetry._setup_debug_events(ctx, session_id): yield coordinator.track_debug(event) logger.info( "[STREAM] Turn setup ready", @@ -325,12 +226,12 @@ async def run( ): yield event if emit_debug: - history_event = self._history_debug_event(state) + history_event = _telemetry._history_debug_event(state) if history_event is not None: yield coordinator.track_debug(history_event) - yield coordinator.track_debug(self._usage_debug_event(state.run_usage)) + yield coordinator.track_debug(_telemetry._usage_debug_event(state.run_usage)) control.check_cancelled() - await self._persist(plan, state) + await persist_snapshot(self._sessions, plan, state) await self._clear_pending(plan) async for event in self._handle_output(plan, ctx, coordinator, state, control=control): yield event @@ -358,10 +259,10 @@ async def run( logger.exception( "Cancelled turn snapshot could not be saved", session_id=session_id ) - events = self._outcome_events( + events = _telemetry._outcome_events( coordinator, session_id=session_id, - message=_persistence_failure( + message=_telemetry._persistence_failure( "Cancelled turn", exc, detail=self._config.client_error_detail ), error_type="persistence_error", @@ -373,7 +274,7 @@ async def run( ) else: if isinstance(cancellation, TimeoutError): - events = self._outcome_events( + events = _telemetry._outcome_events( coordinator, session_id=session_id, message=f"Request timed out after {self._config.stream_timeout_seconds}s", @@ -385,7 +286,7 @@ async def run( emit_debug=emit_debug, ) else: - events = self._outcome_events( + events = _telemetry._outcome_events( coordinator, session_id=session_id, message="Request cancelled", @@ -409,10 +310,10 @@ async def run( await self._persist_cancelled(plan, state, interrupted=False) except Exception as persist_exc: logger.exception("Usage-limited turn could not be saved", session_id=session_id) - events = self._outcome_events( + events = _telemetry._outcome_events( coordinator, session_id=session_id, - message=_persistence_failure( + message=_telemetry._persistence_failure( "Usage-limited turn", persist_exc, detail=self._config.client_error_detail ), error_type="persistence_error", @@ -427,7 +328,7 @@ async def run( logger.info( "[STREAM] Turn stopped by usage limit", session_id=session_id, error=str(exc) ) - events = self._outcome_events( + events = _telemetry._outcome_events( coordinator, session_id=session_id, message=str(exc), @@ -460,7 +361,7 @@ async def run( ) retry_allowed = not own_request else: - message, error_type, retry_allowed = _describe_error(exc, detail=detail) + message, error_type, retry_allowed = _telemetry._describe_error(exc, detail=detail) # The log always carries the exception; only the client text is redacted. logger.warning( "[STREAM] Turn failed", @@ -474,9 +375,9 @@ async def run( await self._resolve_unanswered_calls(plan, state) except Exception as persist_exc: logger.exception("Failed turn snapshot could not be saved", session_id=session_id) - message = _persistence_failure("Failed turn", persist_exc, detail=detail) + message = _telemetry._persistence_failure("Failed turn", persist_exc, detail=detail) error_type, retry_allowed = "persistence_error", False - for event in self._outcome_events( + for event in _telemetry._outcome_events( coordinator, session_id=session_id, message=message, @@ -519,7 +420,8 @@ async def run( if trace_cm is not None: with contextlib.suppress(Exception): trace_cm.__exit__(None, None, None) - await self.save_trace( + await _telemetry.save_trace( + self._db, session_id, coordinator.debug_events, trace_id=trace_id, @@ -540,7 +442,7 @@ async def run( ) async def _update_working_memory( - self, plan: TurnPlan, state: _RunState, *, release_pin: bool = False + self, plan: TurnPlan, state: RunSnapshot, *, release_pin: bool = False ) -> None: """Extract the working-memory delta and account its usage on the assistant row.""" try: @@ -570,50 +472,25 @@ async def _clear_pending(self, plan: TurnPlan) -> None: await self._sessions.clear_pending_action(plan.session_id) async def _persist_cancelled( - self, plan: TurnPlan, state: _RunState, *, interrupted: bool = True + self, plan: TurnPlan, state: RunSnapshot, *, interrupted: bool = True ) -> None: """Retain native work and explicitly mark unresolved calls interrupted.""" state.cancelled = interrupted state.interrupted_status = "cancelled" - self._mark_unanswered_calls(state) - await self._persist(plan, state) + mark_unanswered_calls(state) + await persist_snapshot(self._sessions, plan, state) await self._clear_pending(plan) - async def _resolve_unanswered_calls(self, plan: TurnPlan, state: _RunState) -> None: + async def _resolve_unanswered_calls(self, plan: TurnPlan, state: RunSnapshot) -> None: """Save completed work and mark unanswered calls interrupted after a failure.""" if not state.assistant_messages: return state.cancelled = True - self._mark_unanswered_calls(state) + mark_unanswered_calls(state) state.interrupted_status = "cancelled" - await self._persist(plan, state) + await persist_snapshot(self._sessions, plan, state) await self._clear_pending(plan) - @staticmethod - def _mark_unanswered_calls(state: _RunState) -> bool: - """Append interrupted returns for calls without a result; True when any were added.""" - returned = { - part.tool_call_id - for message in state.assistant_messages - for part in message.parts - if isinstance(part, ToolReturnPart) - } - interrupted = [ - ToolReturnPart( - tool_name=part.tool_name, - tool_call_id=part.tool_call_id, - content="[Tool execution was interrupted; its external outcome is unknown.]", - outcome="interrupted", - ) - for message in state.assistant_messages - for part in message.parts - if isinstance(part, ToolCallPart) and part.tool_call_id not in returned - ] - if not interrupted: - return False - state.assistant_messages.append(ModelRequest(parts=interrupted)) - return True - # --- pieces ----------------------------------------------------------------- async def _run_agent( @@ -621,7 +498,7 @@ async def _run_agent( plan: TurnPlan, ctx: AgentSetupContext, coordinator: EventCoordinator, - state: _RunState, + state: RunSnapshot, *, control: TurnControl, user_prompt: str | None, @@ -691,15 +568,15 @@ async def _run_agent( yield event finally: telegram_chat_id = get_current_telegram_chat_binding() - self._capture_result(plan, state, run.result) + capture_result(plan, state, run.result) except (RunCancelled, asyncio.CancelledError, TimeoutError) as exc: snapshot = RunCancelled.from_cancellation(exc) if snapshot is not None: - self._capture_result(plan, state, snapshot) + capture_result(plan, state, snapshot) raise except Exception as exc: snapshot = RunCancelled.from_cancellation(exc) - self._capture_result(plan, state, snapshot or _IncompleteRun(captured, usage, run_id)) + capture_result(plan, state, snapshot or _IncompleteRun(captured, usage, run_id)) raise finally: if telegram_chat_id: @@ -712,105 +589,12 @@ async def _run_agent( # or prevent the response snapshot from being persisted. logger.exception("Failed to save Telegram binding", session_id=session_id) - @staticmethod - def _capture_result(plan: TurnPlan, state: _RunState, result: Any) -> None: - """Persist native history, including snapshots attached during teardown.""" - state.all_messages = list(result.all_messages()) - # new_messages(), not index slicing: pydantic-ai may merge consecutive - # ModelRequests while cleaning the history, shrinking the list. - new_messages = result.new_messages() - if plan.kind == "message" and not state.assistant_messages: - new_messages = _without_user_prompt(new_messages) - if plan.accepted_tool_result is not None and any( - isinstance(part, ToolReturnPart) and part.tool_call_id == plan.request.tool_call_id - for message in new_messages - for part in message.parts - ): - # The result was saved before setup; native resumption now includes it. - state.assistant_messages = [ - message - for message in state.assistant_messages - if message is not plan.accepted_tool_result - ] - state.assistant_messages.extend(new_messages) - state.run_usage = result.usage - state.usage = merge_usage(plan.prior_usage, usage_dict(result)) - # Counts are cumulative within a turn; tiers arrive per response. Restored host - # history may lack provider_details, so retain the prior usage prefix explicitly. - if state.service_tiers is None: - state.service_tiers = list((plan.prior_usage or {}).get("service_tiers") or []) - state.service_tiers.extend(response_service_tiers(new_messages)) - if state.usage is not None and state.service_tiers: - state.usage["service_tiers"] = list(state.service_tiers) - if not isinstance(result, RunCancelled): - state.output = result.output - - async def _persist(self, plan: TurnPlan, state: _RunState) -> None: - """Create or extend the assistant row with everything run so far.""" - if plan.request.output_mode == "host_tools": - state.assistant_messages = [ - replace( - m, parts=[p for p in m.parts if not isinstance(p, (TextPart, ThinkingPart))] - ) - if isinstance(m, ModelResponse) - else m - for m in state.assistant_messages - ] - # Steering whose model request did not complete stays queued and is sent - # again; the snapshot must not keep a copy of it as well. - state.assistant_messages = _without_steering( - state.assistant_messages, set(plan.session_context.get("pending_steering_ids") or ()) - ) - content, segments, timestamp = build_assistant_message_content( - state.assistant_messages, interrupted_status=state.interrupted_status - ) - state.assistant_segments = segments - # The history policy's summarisation calls are this turn's auxiliary - # model work; state.usage stays the run's own so this is idempotent. - summarisation = ( - usage_dict(state.history.usage) - if state.history is not None and state.history.usage.has_values() - else None - ) - state.stored_usage = with_auxiliary(state.usage, "summarization", summarisation) - # Auxiliary sections already on the row (a background extraction that - # finished after this turn took its snapshot) are kept, not erased. - if state.prompt_record is not None and state.prompt_snapshot is not None: - await self._sessions.save_prompt_snapshot( - state.prompt_record["snapshot_hash"], state.prompt_snapshot - ) - model_messages = dump_model_messages(state.assistant_messages) - if plan.assistant_parent_id is not None and not state.persisted: - await self._sessions.register_assistant_message( - plan.session_id, - message_id=plan.assistant_message_id, - parent_id=plan.assistant_parent_id, - content=content, - segments=segments, - usage=state.stored_usage, - created_at=timestamp, - prompt=state.prompt_record, - model_messages=model_messages, - ) - else: - record = await self._sessions.update_message( - plan.session_id, - plan.assistant_message_id, - content=content, - segments=segments, - usage=lambda current: _keep_row_auxiliary(state.stored_usage, {"usage": current}), - prompt=state.prompt_record, - model_messages=model_messages, - ) - state.stored_usage = record.get("usage") - state.persisted = True - async def _handle_output( self, plan: TurnPlan, ctx: AgentSetupContext, coordinator: EventCoordinator, - state: _RunState, + state: RunSnapshot, *, control: TurnControl, ) -> AsyncIterator[dict[str, Any]]: @@ -839,11 +623,13 @@ async def _handle_output( ): yield event control.check_cancelled() - await self._persist(plan, state) + await persist_snapshot(self._sessions, plan, state) if await self._take_output(plan, state, ctx): return - async def _take_output(self, plan: TurnPlan, state: _RunState, ctx: AgentSetupContext) -> bool: + async def _take_output( + self, plan: TurnPlan, state: RunSnapshot, ctx: AgentSetupContext + ) -> bool: """Record the run output on ``state``; True when it is a deferred host-tool call.""" session_context = plan.session_context if isinstance(state.output, DeferredToolRequests): @@ -882,200 +668,6 @@ async def _take_output(self, plan: TurnPlan, state: _RunState, ctx: AgentSetupCo state.final_output = str(state.output) return False - @staticmethod - def _setup_debug_events(ctx: AgentSetupContext, session_id: str) -> list[dict[str, Any]]: - tools = ctx.available_tools - return [ - make_debug_tool_selection_event( - page=tools.page, - backend_count=len(tools.backend_tools), - filtered_out=tools.filtered_out_count, - tool_names=tools.tool_names, - tools=[ - { - "name": t.name, - "description": t.description, - "parameters_schema": t.parameters_schema, - } - for t in tools.backend_tools - ], - ), - make_debug_system_prompt_event( - total_length=len(ctx.prompt_result.content), - fragment_count=len(ctx.prompt_result.fragments), - fragments=ctx.prompt_result.fragments, - content=ctx.prompt_result.content, - ), - make_debug_agent_config_event( - model=ctx.resolved_model, - output_type="str", - thinking_budget=ctx.effective_config.thinking_budget, - temperature=ctx.effective_config.temperature, - session_id=session_id, - ), - ] - - @staticmethod - def _history_debug_event(state: _RunState) -> dict[str, Any] | None: - """Describe the model input the history policy last produced, if it ran.""" - result = state.history.result if state.history is not None else None - if result is None: - return None - return make_debug_history_event( - message_count=result.message_count, - estimated_tokens=result.estimated_tokens, - was_compacted=result.was_compacted, - compacted_from=result.compacted_from, - messages=result.message_summaries, - ) - - @staticmethod - def _usage_debug_event(run_usage: Any) -> dict[str, Any]: - snapshot = usage_dict(run_usage) or { - "input_tokens": 0, - "output_tokens": 0, - "total_tokens": 0, - } - cache_read, cache_write = cache_counts(run_usage) - return make_debug_usage_event( - input_tokens=snapshot["input_tokens"], - output_tokens=snapshot["output_tokens"], - cache_read=cache_read, - cache_write=cache_write, - requests=getattr(run_usage, "requests", 0) or 0, - total=snapshot["total_tokens"], - ) - - @staticmethod - def _outcome_events( - coordinator: EventCoordinator, - *, - session_id: str, - message: str, - error_type: str, - retry_allowed: bool, - trace_id: str, - model: str, - phase: str, - emit_debug: bool, - message_id: str | None = None, - content: str | None = None, - usage: dict[str, Any] | None = None, - ) -> list[dict[str, Any]]: - """The one terminal lifecycle of a turn that did not finish normally. - - ``final_response(error=True)``, the debug error (always traced, sent - only with ``emit_debug``) and the terminal ``error``. A saved snapshot - (cancellation, usage limit) names its ``message_id`` and ``usage``; - ``content`` stays empty, only native middleware decides what text a - client sees. Like the other terminal events, the error is not counted - against the event limit: the client must always receive it. - """ - final = coordinator.try_final_response( - content, - model, - session_id=session_id, - message_id=message_id, - trace_id=trace_id, - error=True, - error_type=error_type, - usage=usage, - ) - debug = coordinator.track_debug( - make_debug_error_event( - message, - error_type=error_type, - retry_allowed=retry_allowed, - trace_id=trace_id, - model=model, - phase=phase, - ) - ) - events: list[dict[str, Any]] = [] - if final: - events.append(final) - if emit_debug: - events.append(debug) - error = make_error_event( - format_error_message(message, trace_id), - error_type=error_type, - trace_id=trace_id, - terminal=True, - retry_allowed=retry_allowed, - ) - events.append(error) - return events - - async def save_trace( - self, - session_id: str, - trace_events: list[dict[str, Any]], - *, - trace_id: str | None = None, - user_message: str | None = None, - duration_ms: float | None = None, - ) -> None: - """Persist collected debug events as a trace row (never the screenshot). Best-effort.""" - if self._db is None or not self._db.healthy or not trace_events: - return - try: - from assistant_runtime.services.database.repositories import TraceRepository - - async with self._db.session_context() as db_session: - await TraceRepository(db_session).create( - trace_id=trace_id or str(uuid.uuid4()), - session_id=session_id, - events=trace_events, - user_message=user_message, - duration_ms=duration_ms, - ) - except Exception as e: - logger.warning("Failed to persist trace", session_id=session_id, error=str(e)) - - -class _IncompleteRun: - """Captured native messages and usage when a run raises without a result.""" - - output = None - - def __init__(self, captured: list[ModelMessage], usage: Any, run_id: str) -> None: - self._messages = list(captured) - self._run_id = run_id - self.usage = usage - - def all_messages(self) -> list[ModelMessage]: - return list(self._messages) - - def new_messages(self) -> list[ModelMessage]: - return [m for m in self._messages if getattr(m, "run_id", None) == self._run_id] - - -def _without_user_prompt(messages: list[ModelMessage]) -> list[ModelMessage]: - """A message turn's messages after its user prompt, which the user row keeps. - - Pydantic AI sends the prompt in the run's first request; steering it - delivers follows in requests of its own. - """ - if not messages or not isinstance(messages[0], ModelRequest): - return messages - first = messages[0] - rest = [part for part in first.parts if not isinstance(part, UserPromptPart)] - return [replace(first, parts=rest), *messages[1:]] if rest else messages[1:] - - -def _without_steering(messages: list[ModelMessage], pending_ids: set[str]) -> list[ModelMessage]: - """``messages`` without the steering requests (``build_steering_request``) still pending.""" - if not pending_ids: - return messages - return [ - message - for message in messages - if not ( - isinstance(message, ModelRequest) - and pending_ids.intersection((message.metadata or {}).get("steering_ids") or ()) - ) - ] - def _host_tool_names(ctx: AgentSetupContext) -> set[str]: """The host tools of this turn: configured ones plus request-declared actions.""" diff --git a/src/assistant_runtime/app/streaming/_snapshot.py b/src/assistant_runtime/app/streaming/_snapshot.py new file mode 100644 index 00000000..40d63d06 --- /dev/null +++ b/src/assistant_runtime/app/streaming/_snapshot.py @@ -0,0 +1,237 @@ +"""Native run snapshots, assistant-row persistence and interrupted tool results.""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from typing import TYPE_CHECKING, Any + +from pydantic_ai.exceptions import RunCancelled +from pydantic_ai.messages import ( + ModelMessage, + ModelRequest, + ModelResponse, + TextPart, + ThinkingPart, + ToolCallPart, + ToolReturnPart, + UserPromptPart, +) + +from assistant_runtime.app.assistant import ( + ActionStatus, + build_assistant_message_content, + dump_model_messages, +) +from assistant_runtime.app.assistant.usage import ( + merge_usage, + response_service_tiers, + usage_dict, + with_auxiliary, +) + +if TYPE_CHECKING: + from assistant_runtime.app.assistant import SessionStore + from assistant_runtime.app.streaming._turn import TurnPlan + from assistant_runtime.services.history.interface import HistoryProcessor + + +def _keep_row_auxiliary( + usage: dict[str, Any] | None, row: dict[str, Any] | None +) -> dict[str, Any] | None: + """Keep row-owned memory totals and auxiliary sections missing from a snapshot.""" + existing = ((row or {}).get("usage") or {}).get("auxiliary") or {} + for name, part in existing.items(): + auxiliary = (usage or {}).get("auxiliary") or {} + if name == "working_memory" and name in auxiliary: + # Background extraction owns this cumulative total; the turn's + # earlier copy must neither overwrite it nor be added a second time. + usage = {**usage, "auxiliary": {**auxiliary, name: part}} + elif name not in auxiliary: + usage = with_auxiliary(usage, name, part) + return usage + + +@dataclass +class RunSnapshot: + """What one agent run left behind, carried across steering follow-ups.""" + + all_messages: list[ModelMessage] = field(default_factory=list) + assistant_messages: list[ModelMessage] = field(default_factory=list) + assistant_segments: list[dict[str, Any]] | None = None + run_usage: Any = None + usage: dict[str, int] | None = None + service_tiers: list[dict[str, Any]] | None = None + output: Any = None + final_output: str | None = None + pending_tool_call: dict[str, Any] | None = None + persisted: bool = False + cancelled: bool = False + # How calls left without a result by this turn are recorded (see ActionStatus). + interrupted_status: ActionStatus = "superseded" + history: HistoryProcessor | None = None + # The usage written on the assistant row: the run's own plus auxiliary work. + stored_usage: dict[str, Any] | None = None + # What the latest run of the turn was given (app/assistant/prompt_record). + prompt_record: dict[str, Any] | None = None + prompt_snapshot: str | None = None + + +class _IncompleteRun: + """Captured native messages and usage when a run raises without a result.""" + + output = None + + def __init__(self, captured: list[ModelMessage], usage: Any, run_id: str) -> None: + self._messages = list(captured) + self._run_id = run_id + self.usage = usage + + def all_messages(self) -> list[ModelMessage]: + return list(self._messages) + + def new_messages(self) -> list[ModelMessage]: + return [m for m in self._messages if getattr(m, "run_id", None) == self._run_id] + + +def _without_user_prompt(messages: list[ModelMessage]) -> list[ModelMessage]: + """A message turn's messages after its user prompt, which the user row keeps. + + Pydantic AI sends the prompt in the run's first request; steering it + delivers follows in requests of its own. + """ + if not messages or not isinstance(messages[0], ModelRequest): + return messages + first = messages[0] + rest = [part for part in first.parts if not isinstance(part, UserPromptPart)] + return [replace(first, parts=rest), *messages[1:]] if rest else messages[1:] + + +def _without_steering(messages: list[ModelMessage], pending_ids: set[str]) -> list[ModelMessage]: + """``messages`` without the steering requests (``build_steering_request``) still pending.""" + if not pending_ids: + return messages + return [ + message + for message in messages + if not ( + isinstance(message, ModelRequest) + and pending_ids.intersection((message.metadata or {}).get("steering_ids") or ()) + ) + ] + + +def capture_result(plan: TurnPlan, state: RunSnapshot, result: Any) -> None: + """Persist native history, including snapshots attached during teardown.""" + state.all_messages = list(result.all_messages()) + # new_messages(), not index slicing: pydantic-ai may merge consecutive + # ModelRequests while cleaning the history, shrinking the list. + new_messages = result.new_messages() + if plan.kind == "message" and not state.assistant_messages: + new_messages = _without_user_prompt(new_messages) + if plan.accepted_tool_result is not None and any( + isinstance(part, ToolReturnPart) and part.tool_call_id == plan.request.tool_call_id + for message in new_messages + for part in message.parts + ): + # The result was saved before setup; native resumption now includes it. + state.assistant_messages = [ + message + for message in state.assistant_messages + if message is not plan.accepted_tool_result + ] + state.assistant_messages.extend(new_messages) + state.run_usage = result.usage + state.usage = merge_usage(plan.prior_usage, usage_dict(result)) + # Counts are cumulative within a turn; tiers arrive per response. Restored host + # history may lack provider_details, so retain the prior usage prefix explicitly. + if state.service_tiers is None: + state.service_tiers = list((plan.prior_usage or {}).get("service_tiers") or []) + state.service_tiers.extend(response_service_tiers(new_messages)) + if state.usage is not None and state.service_tiers: + state.usage["service_tiers"] = list(state.service_tiers) + if not isinstance(result, RunCancelled): + state.output = result.output + + +async def persist_snapshot(sessions: SessionStore, plan: TurnPlan, state: RunSnapshot) -> None: + """Create or extend the assistant row with everything run so far.""" + if plan.request.output_mode == "host_tools": + state.assistant_messages = [ + replace(m, parts=[p for p in m.parts if not isinstance(p, (TextPart, ThinkingPart))]) + if isinstance(m, ModelResponse) + else m + for m in state.assistant_messages + ] + # Steering whose model request did not complete stays queued and is sent + # again; the snapshot must not keep a copy of it as well. + state.assistant_messages = _without_steering( + state.assistant_messages, set(plan.session_context.get("pending_steering_ids") or ()) + ) + content, segments, timestamp = build_assistant_message_content( + state.assistant_messages, interrupted_status=state.interrupted_status + ) + state.assistant_segments = segments + # The history policy's summarisation calls are this turn's auxiliary + # model work; state.usage stays the run's own so this is idempotent. + summarisation = ( + usage_dict(state.history.usage) + if state.history is not None and state.history.usage.has_values() + else None + ) + state.stored_usage = with_auxiliary(state.usage, "summarization", summarisation) + # Auxiliary sections already on the row (a background extraction that + # finished after this turn took its snapshot) are kept, not erased. + if state.prompt_record is not None and state.prompt_snapshot is not None: + await sessions.save_prompt_snapshot( + state.prompt_record["snapshot_hash"], state.prompt_snapshot + ) + model_messages = dump_model_messages(state.assistant_messages) + if plan.assistant_parent_id is not None and not state.persisted: + await sessions.register_assistant_message( + plan.session_id, + message_id=plan.assistant_message_id, + parent_id=plan.assistant_parent_id, + content=content, + segments=segments, + usage=state.stored_usage, + created_at=timestamp, + prompt=state.prompt_record, + model_messages=model_messages, + ) + else: + record = await sessions.update_message( + plan.session_id, + plan.assistant_message_id, + content=content, + segments=segments, + usage=lambda current: _keep_row_auxiliary(state.stored_usage, {"usage": current}), + prompt=state.prompt_record, + model_messages=model_messages, + ) + state.stored_usage = record.get("usage") + state.persisted = True + + +def mark_unanswered_calls(state: RunSnapshot) -> bool: + """Append interrupted returns for calls without a result; True when any were added.""" + returned = { + part.tool_call_id + for message in state.assistant_messages + for part in message.parts + if isinstance(part, ToolReturnPart) + } + interrupted = [ + ToolReturnPart( + tool_name=part.tool_name, + tool_call_id=part.tool_call_id, + content="[Tool execution was interrupted; its external outcome is unknown.]", + outcome="interrupted", + ) + for message in state.assistant_messages + for part in message.parts + if isinstance(part, ToolCallPart) and part.tool_call_id not in returned + ] + if not interrupted: + return False + state.assistant_messages.append(ModelRequest(parts=interrupted)) + return True diff --git a/src/assistant_runtime/app/streaming/_telemetry.py b/src/assistant_runtime/app/streaming/_telemetry.py new file mode 100644 index 00000000..dd1c5ecb --- /dev/null +++ b/src/assistant_runtime/app/streaming/_telemetry.py @@ -0,0 +1,223 @@ +"""Streaming diagnostics, failure envelopes and best-effort trace persistence.""" + +from __future__ import annotations + +import uuid +from typing import TYPE_CHECKING, Any + +from loguru import logger + +from assistant_runtime.app.assistant.usage import ( + cache_counts, + usage_dict, +) +from assistant_runtime.app.streaming._coordinator import EventCoordinator +from assistant_runtime.app.streaming._event_builder import ( + make_debug_agent_config_event, + make_debug_error_event, + make_debug_history_event, + make_debug_system_prompt_event, + make_debug_tool_selection_event, + make_debug_usage_event, + make_error_event, +) +from assistant_runtime.app.streaming.exceptions import InvalidDecisionError +from assistant_runtime.services.llm.exceptions import LLMCallError, classify_llm_error + +if TYPE_CHECKING: + from assistant_runtime.app.assistant.models import AgentSetupContext + from assistant_runtime.app.streaming._snapshot import RunSnapshot + from assistant_runtime.services.database.interface import DatabaseService + +_LLM_ERROR_TYPES = { + "RATE_LIMIT": "rate_limit", + "SERVER_ERROR": "provider_error", + "CONNECTION_ERROR": "connection_error", + "TIMEOUT": "timeout", + "AUTH_ERROR": "provider_auth", + "CLIENT_ERROR": "provider_client_error", +} + + +def format_error_message(message: str, trace_id: str | None) -> str: + """Append the trace id so client-visible errors can be correlated with traces.""" + if trace_id is None or "[trace_id:" in message: + return message + return f"{message} [trace_id: {trace_id}]" + + +def _describe_error(exc: Exception, *, detail: bool = True) -> tuple[str, str, bool]: + """``(message, error_type, retry_allowed)`` for an exception raised by the run. + + ``detail`` (``STREAMING__CLIENT_ERROR_DETAIL``) decides whether the + exception text itself reaches the client; the log always has it. + """ + if isinstance(exc, InvalidDecisionError): + return str(exc), "invalid_decision", False + if isinstance(exc, LLMCallError): + llm_error: LLMCallError | None = exc + else: + classified = classify_llm_error(exc) + llm_error = None if classified.error_category == "UNKNOWN" else classified + if llm_error is None: + message = f"Request failed: {exc.__class__.__name__}: {exc}" if detail else "Request failed" + return message, "internal", False + error_type = _LLM_ERROR_TYPES.get(llm_error.error_category, "provider_error") + message = str(llm_error) if detail else f"LLM call failed ({llm_error.error_category})" + return message, error_type, llm_error.retry_allowed + + +def _persistence_failure(what: str, exc: BaseException, *, detail: bool) -> str: + """The client-visible text for a snapshot that could not be saved.""" + return f"{what} could not be saved: {exc}" if detail else f"{what} could not be saved" + + +def _setup_debug_events(ctx: AgentSetupContext, session_id: str) -> list[dict[str, Any]]: + tools = ctx.available_tools + return [ + make_debug_tool_selection_event( + page=tools.page, + backend_count=len(tools.backend_tools), + filtered_out=tools.filtered_out_count, + tool_names=tools.tool_names, + tools=[ + { + "name": t.name, + "description": t.description, + "parameters_schema": t.parameters_schema, + } + for t in tools.backend_tools + ], + ), + make_debug_system_prompt_event( + total_length=len(ctx.prompt_result.content), + fragment_count=len(ctx.prompt_result.fragments), + fragments=ctx.prompt_result.fragments, + content=ctx.prompt_result.content, + ), + make_debug_agent_config_event( + model=ctx.resolved_model, + output_type="str", + thinking_budget=ctx.effective_config.thinking_budget, + temperature=ctx.effective_config.temperature, + session_id=session_id, + ), + ] + + +def _history_debug_event(state: RunSnapshot) -> dict[str, Any] | None: + """Describe the model input the history policy last produced, if it ran.""" + result = state.history.result if state.history is not None else None + if result is None: + return None + return make_debug_history_event( + message_count=result.message_count, + estimated_tokens=result.estimated_tokens, + was_compacted=result.was_compacted, + compacted_from=result.compacted_from, + messages=result.message_summaries, + ) + + +def _usage_debug_event(run_usage: Any) -> dict[str, Any]: + snapshot = usage_dict(run_usage) or { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + } + cache_read, cache_write = cache_counts(run_usage) + return make_debug_usage_event( + input_tokens=snapshot["input_tokens"], + output_tokens=snapshot["output_tokens"], + cache_read=cache_read, + cache_write=cache_write, + requests=getattr(run_usage, "requests", 0) or 0, + total=snapshot["total_tokens"], + ) + + +def _outcome_events( + coordinator: EventCoordinator, + *, + session_id: str, + message: str, + error_type: str, + retry_allowed: bool, + trace_id: str, + model: str, + phase: str, + emit_debug: bool, + message_id: str | None = None, + content: str | None = None, + usage: dict[str, Any] | None = None, +) -> list[dict[str, Any]]: + """The one terminal lifecycle of a turn that did not finish normally. + + ``final_response(error=True)``, the debug error (always traced, sent + only with ``emit_debug``) and the terminal ``error``. A saved snapshot + (cancellation, usage limit) names its ``message_id`` and ``usage``; + ``content`` stays empty, only native middleware decides what text a + client sees. Like the other terminal events, the error is not counted + against the event limit: the client must always receive it. + """ + final = coordinator.try_final_response( + content, + model, + session_id=session_id, + message_id=message_id, + trace_id=trace_id, + error=True, + error_type=error_type, + usage=usage, + ) + debug = coordinator.track_debug( + make_debug_error_event( + message, + error_type=error_type, + retry_allowed=retry_allowed, + trace_id=trace_id, + model=model, + phase=phase, + ) + ) + events: list[dict[str, Any]] = [] + if final: + events.append(final) + if emit_debug: + events.append(debug) + error = make_error_event( + format_error_message(message, trace_id), + error_type=error_type, + trace_id=trace_id, + terminal=True, + retry_allowed=retry_allowed, + ) + events.append(error) + return events + + +async def save_trace( + database: DatabaseService | None, + session_id: str, + trace_events: list[dict[str, Any]], + *, + trace_id: str | None = None, + user_message: str | None = None, + duration_ms: float | None = None, +) -> None: + """Persist collected debug events as a trace row (never the screenshot). Best-effort.""" + if database is None or not database.healthy or not trace_events: + return + try: + from assistant_runtime.services.database.repositories import TraceRepository + + async with database.session_context() as db_session: + await TraceRepository(db_session).create( + trace_id=trace_id or str(uuid.uuid4()), + session_id=session_id, + events=trace_events, + user_message=user_message, + duration_ms=duration_ms, + ) + except Exception as e: + logger.warning("Failed to persist trace", session_id=session_id, error=str(e)) diff --git a/src/assistant_runtime/app/streaming/_turn.py b/src/assistant_runtime/app/streaming/_turn.py index 93247fd1..dc995db1 100644 --- a/src/assistant_runtime/app/streaming/_turn.py +++ b/src/assistant_runtime/app/streaming/_turn.py @@ -419,6 +419,7 @@ async def _plan_steering(self, request: AssistantRequest, principal: Principal) prior_assistant_messages=( path_records_to_model_history([active_leaf]) if extends_assistant else [] ), + prior_usage=active_leaf.get("usage") if extends_assistant else None, # The current context comes just before the steering. user_prompt=_current_host_context(request, session_context), history=self._sessions.get_history(session_id), diff --git a/src/assistant_runtime/app/streaming/interface.py b/src/assistant_runtime/app/streaming/interface.py index 329345b7..ec725e25 100644 --- a/src/assistant_runtime/app/streaming/interface.py +++ b/src/assistant_runtime/app/streaming/interface.py @@ -23,6 +23,7 @@ from assistant_runtime.app.access.interface import AccessService from assistant_runtime.app.assistant.exceptions import AgentRunError, SessionError from assistant_runtime.app.assistant.models import AssistantRequest, AssistantResult +from assistant_runtime.app.streaming import _telemetry from assistant_runtime.app.streaming._control import TurnControl from assistant_runtime.app.streaming._event_builder import ( make_agent_status_event, @@ -32,7 +33,7 @@ make_voice_event, ) from assistant_runtime.app.streaming._host_tool import clear_stale_pending_call -from assistant_runtime.app.streaming._runner import TurnRunner, format_error_message +from assistant_runtime.app.streaming._runner import TurnRunner from assistant_runtime.app.streaming._turn import TurnPlanner, apply_binding from assistant_runtime.app.streaming.config import StreamingConfig from assistant_runtime.app.streaming.exceptions import StreamingError, StreamSetupError @@ -630,7 +631,8 @@ async def _setup_failure_envelope( message = "Setup failed" # A denied request leaves nothing under the session it could not reach. if not isinstance(exc, AccessDeniedError): - await self._runner.save_trace( + await _telemetry.save_trace( + self._db, request.session_id, [ make_debug_error_event( @@ -655,7 +657,7 @@ async def _setup_failure_envelope( error_type=error_type, ) yield make_error_event( - format_error_message(message, trace_id), + _telemetry.format_error_message(message, trace_id), error_type=error_type, trace_id=trace_id, terminal=True, diff --git a/src/assistant_runtime/app/tasks/_runner.py b/src/assistant_runtime/app/tasks/_runner.py new file mode 100644 index 00000000..01362420 --- /dev/null +++ b/src/assistant_runtime/app/tasks/_runner.py @@ -0,0 +1,215 @@ +"""Scheduling and terminal bookkeeping for task and persistent-agent turns.""" + +from __future__ import annotations + +import asyncio +import contextlib +from collections.abc import Awaitable, Callable, Coroutine, Hashable, Iterator +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any, NamedTuple + +from loguru import logger + +from assistant_runtime.app.assistant.models import AssistantRequest +from assistant_runtime.app.tasks.config import TasksConfig +from assistant_runtime.app.tasks.exceptions import TaskError, TaskLimitError +from assistant_runtime.principal import Principal + +if TYPE_CHECKING: + from assistant_runtime.app.streaming.interface import StreamingService + +INTERRUPTED = "The runtime stopped while it was queued or running" +_UNFINISHED = ("queued", "running") + + +class _Job(NamedTuple): + """A queued or running background turn: a task, or a message to an agent.""" + + worker: asyncio.Task[None] + update: Callable[..., Awaitable[Any]] + publish: Callable[[Any], Awaitable[None]] + agent_id: str | None = None + + +@dataclass +class _Order: + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + users: int = 0 + + +class BackgroundTurnRunner: + """Own queue capacity, per-key ordering, shared slots and worker lifetime.""" + + def __init__(self, config: TasksConfig, streaming: StreamingService) -> None: + self._config = config + self._streaming = streaming + self._reserved = 0 + self._slots = asyncio.Semaphore(config.max_concurrent) + self._order_locks: dict[Hashable, _Order] = {} + self._running: dict[str, _Job] = {} + self._cancel_requested: set[str] = set() + + @property + def count(self) -> int: + return len(self._running) + + @contextlib.contextmanager + def reserve(self) -> Iterator[None]: + """Reserve queue capacity across asynchronous record creation.""" + count = self.count + self._reserved + if count >= self._config.max_waiting: + raise TaskLimitError( + f"{count} background turns are queued or running; wait for some to finish" + ) + self._reserved += 1 + try: + yield + finally: + self._reserved -= 1 + + def submit( + self, + job_id: str, + work: Coroutine[Any, Any, None], + update: Callable[..., Awaitable[Any]], + publish: Callable[[Any], Awaitable[None]], + *, + name: str, + agent_id: str | None = None, + ) -> None: + self._running[job_id] = _Job( + asyncio.create_task(work, name=name), update, publish, agent_id + ) + + def mark_cancelled(self, job_id: str) -> None: + self._cancel_requested.add(job_id) + + async def cancel_agent(self, agent_id: str) -> None: + for job_id in [i for i, job in self._running.items() if job.agent_id == agent_id]: + await self.cancel(job_id) + + async def stop(self) -> None: + running = list(self._running.items()) + for _, job in running: + job.worker.cancel() + for job_id, job in running: + with contextlib.suppress(asyncio.CancelledError, Exception): + await job.worker + await self._finish_unstarted(job_id, "interrupted") + + async def execute( + self, + job_id: str, + order_key: Hashable | None, + start: Callable[[], Awaitable[AssistantRequest | None]], + update: Callable[..., Awaitable[Any]], + publish: Callable[[Any], Awaitable[None]], + principal: Principal, + ) -> None: + """Run one background turn in its order and a free slot; record and publish its end.""" + final: Any = None + try: + async with self._in_order(order_key), self._slots: + request = await start() + if request is None: + return + async with asyncio.timeout(self._config.timeout_seconds): + result = await self._streaming.run_message(request, principal=principal) + if result.pending_tool_call: + # A background turn has no host to perform an action and send its result back. + raise TaskError( + "The turn asked the host to perform " + f"'{result.pending_tool_call.get('tool_name')}', which a background " + "turn cannot answer" + ) + final = await update( + job_id, + only_from=_UNFINISHED, + status="done", + result=(result.content or "")[: self._config.result_max_chars], + usage=result.usage, + finished_at=_now(), + ) + except asyncio.CancelledError: + requested = job_id in self._cancel_requested + final = await update( + job_id, + only_from=_UNFINISHED, + status="cancelled" if requested else "interrupted", + error=None if requested else INTERRUPTED, + finished_at=_now(), + ) + except TimeoutError: + final = await update( + job_id, + only_from=_UNFINISHED, + status="failed", + error=f"Timed out after {self._config.timeout_seconds} s", + finished_at=_now(), + ) + except Exception as e: + logger.warning("Background turn failed", job_id=job_id, error=str(e)) + final = await update( + job_id, + only_from=_UNFINISHED, + status="failed", + error=str(e) or type(e).__name__, + finished_at=_now(), + ) + finally: + self._running.pop(job_id, None) + self._cancel_requested.discard(job_id) + if final is not None: + await publish(final) + + async def cancel(self, job_id: str) -> None: + """Cancel a queued or running background turn and wait until its end is recorded.""" + job = self._running.get(job_id) + if job is None: + return + self._cancel_requested.add(job_id) + job.worker.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await job.worker + await self._finish_unstarted(job_id, "cancelled") + + async def _finish_unstarted(self, job_id: str, status: str) -> None: + """Record the end of a background turn whose worker was cancelled before it ever ran. + + asyncio does not enter a coroutine cancelled before its first step, so + its own handler and ``finally`` never run; this does their work. + """ + job = self._running.pop(job_id, None) + if job is None: + return # the worker ran and finished its record itself + self._cancel_requested.discard(job_id) + final = await job.update( + job_id, + only_from=_UNFINISHED, + status=status, + error=None if status == "cancelled" else INTERRUPTED, + finished_at=_now(), + ) + if final is not None: + await job.publish(final) + + @contextlib.asynccontextmanager + async def _in_order(self, key: Hashable | None): + """One background turn at a time per key, in arrival order; no key, no order.""" + if key is None: + yield + return + entry = self._order_locks.setdefault(key, _Order()) + entry.users += 1 + try: + async with entry.lock: + yield + finally: + entry.users -= 1 + if entry.users == 0: + self._order_locks.pop(key, None) + + +def _now() -> datetime: + return datetime.now(UTC) diff --git a/src/assistant_runtime/app/tasks/_store.py b/src/assistant_runtime/app/tasks/_store.py index 78173cc3..ad1f28c5 100644 --- a/src/assistant_runtime/app/tasks/_store.py +++ b/src/assistant_runtime/app/tasks/_store.py @@ -5,6 +5,7 @@ from __future__ import annotations +from copy import deepcopy from dataclasses import fields, replace from datetime import UTC, datetime from typing import Any, Protocol @@ -86,9 +87,9 @@ def __init__(self) -> None: self._messages: dict[str, AgentMessageRecord] = {} async def create(self, record: TaskRecord) -> TaskRecord: - record = replace(record, created_at=record.created_at or datetime.now(UTC)) + record = replace(deepcopy(record), created_at=record.created_at or datetime.now(UTC)) self._tasks[record.id] = record - return record + return deepcopy(record) async def update( self, task_id: str, *, only_from: tuple[str, ...] | None = None, **values: Any @@ -96,11 +97,11 @@ async def update( record = self._tasks.get(task_id) if record is None or (only_from is not None and record.status not in only_from): return None - record = self._tasks[task_id] = replace(record, **values) - return record + record = self._tasks[task_id] = replace(record, **deepcopy(values)) + return deepcopy(record) async def get(self, task_id: str) -> TaskRecord | None: - return self._tasks.get(task_id) + return deepcopy(self._tasks.get(task_id)) async def list( self, @@ -117,25 +118,25 @@ async def list( and (parent_session_id is None or r.parent_session_id == parent_session_id) and (status is None or r.status == status) ] - return list(reversed(records))[:limit] + return deepcopy(list(reversed(records))[:limit]) async def mark_unfinished(self, status: str, error: str) -> list[TaskRecord]: return [] # nothing survives a restart in memory async def create_agent(self, record: AgentRecord) -> AgentRecord: - record = replace(record, created_at=record.created_at or datetime.now(UTC)) + record = replace(deepcopy(record), created_at=record.created_at or datetime.now(UTC)) self._agents[record.id] = record - return record + return deepcopy(record) async def update_agent(self, agent_id: str, **values: Any) -> AgentRecord | None: record = self._agents.get(agent_id) if record is None: return None - record = self._agents[agent_id] = replace(record, **values) - return record + record = self._agents[agent_id] = replace(record, **deepcopy(values)) + return deepcopy(record) async def get_agent(self, agent_id: str) -> AgentRecord | None: - return self._agents.get(agent_id) + return deepcopy(self._agents.get(agent_id)) async def list_agents(self, *, created_by: str | None, status: str | None) -> list[AgentRecord]: records = [ @@ -144,12 +145,12 @@ async def list_agents(self, *, created_by: str | None, status: str | None) -> li if (created_by is None or r.created_by == created_by) and (status is None or r.status == status) ] - return list(reversed(records)) + return deepcopy(list(reversed(records))) async def create_message(self, record: AgentMessageRecord) -> AgentMessageRecord: - record = replace(record, created_at=record.created_at or datetime.now(UTC)) + record = replace(deepcopy(record), created_at=record.created_at or datetime.now(UTC)) self._messages[record.id] = record - return record + return deepcopy(record) async def update_message( self, message_id: str, *, only_from: tuple[str, ...] | None = None, **values: Any @@ -157,11 +158,11 @@ async def update_message( record = self._messages.get(message_id) if record is None or (only_from is not None and record.status not in only_from): return None - record = self._messages[message_id] = replace(record, **values) - return record + record = self._messages[message_id] = replace(record, **deepcopy(values)) + return deepcopy(record) async def get_message(self, message_id: str) -> AgentMessageRecord | None: - return self._messages.get(message_id) + return deepcopy(self._messages.get(message_id)) async def list_messages( self, @@ -180,7 +181,7 @@ async def list_messages( and (parent_session_id is None or r.parent_session_id == parent_session_id) and (status is None or r.status == status) ] - return list(reversed(records))[:limit] + return deepcopy(list(reversed(records))[:limit]) async def mark_unfinished_messages(self, status: str, error: str) -> list[AgentMessageRecord]: return [] # nothing survives a restart in memory diff --git a/src/assistant_runtime/app/tasks/interface.py b/src/assistant_runtime/app/tasks/interface.py index 433d920f..e1ef635b 100644 --- a/src/assistant_runtime/app/tasks/interface.py +++ b/src/assistant_runtime/app/tasks/interface.py @@ -24,11 +24,9 @@ from __future__ import annotations import asyncio -import contextlib import uuid -from collections.abc import Awaitable, Callable, Hashable from datetime import UTC, datetime -from typing import TYPE_CHECKING, Any, NamedTuple +from typing import TYPE_CHECKING, Any from loguru import logger from pydantic import ValidationError @@ -36,6 +34,7 @@ from assistant_runtime.app.assistant.config import AssistantConfig, TunableOverrides from assistant_runtime.app.assistant.models import AssistantRequest from assistant_runtime.app.settings import outside_request_allowance +from assistant_runtime.app.tasks._runner import INTERRUPTED, BackgroundTurnRunner from assistant_runtime.app.tasks._store import ( AgentTakenError, DatabaseTaskStore, @@ -48,7 +47,6 @@ AgentConflictError, AgentNotFoundError, TaskError, - TaskLimitError, TaskNotFoundError, TasksDisabledError, ) @@ -61,19 +59,6 @@ from assistant_runtime.services.database.interface import DatabaseService from assistant_runtime.services.tools.interface import ToolService -_INTERRUPTED = "The runtime stopped while it was queued or running" -# An end is recorded once: later terminal updates find it finished and change nothing. -_UNFINISHED = ("queued", "running") - - -class _Job(NamedTuple): - """A queued or running background turn: a task, or a message to an agent.""" - - worker: asyncio.Task[None] - update: Callable[..., Awaitable[Any]] - publish: Callable[[Any], Awaitable[None]] - agent_id: str | None = None - class TaskService: """Start, list, read and cancel background tasks. Implements LifecycleAware.""" @@ -93,15 +78,11 @@ def __init__( self._assistant_config = assistant_config or AssistantConfig() self._tool_service = tool_service self._default_profile = default_profile - self._reserved = 0 self._streaming = streaming_service self._database = database_service self._events = events self._store: TaskStore | None = None - self._slots = asyncio.Semaphore(config.max_concurrent) - self._order_locks: dict[Hashable, list[Any]] = {} - self._running: dict[str, _Job] = {} - self._cancel_requested: set[str] = set() + self._runner = BackgroundTurnRunner(config, streaming_service) # Agent starts, moves, stops and message admission see one another's changes. self._agent_admission = asyncio.Lock() @@ -111,9 +92,9 @@ async def start(self) -> None: database = self._database if database is not None and getattr(database, "healthy", False): self._store = DatabaseTaskStore(database) - for record in await self._store.mark_unfinished("interrupted", _INTERRUPTED): + for record in await self._store.mark_unfinished("interrupted", INTERRUPTED): await self._publish(record) - for message in await self._store.mark_unfinished_messages("interrupted", _INTERRUPTED): + for message in await self._store.mark_unfinished_messages("interrupted", INTERRUPTED): await self._publish_message(message) for agent in await self._store.list_agents(created_by=None, status="active"): self._streaming.pin_session(agent.session_id) @@ -125,13 +106,7 @@ async def start(self) -> None: logger.info("Task service started", enabled=self._config.enabled) async def stop(self) -> None: - running = list(self._running.items()) - for _, job in running: - job.worker.cancel() - for job_id, job in running: - with contextlib.suppress(asyncio.CancelledError, Exception): - await job.worker - await self._finish_unstarted(job_id, "interrupted") + await self._runner.stop() self._store = None logger.info("Task service stopped") @@ -139,7 +114,7 @@ async def health_check(self) -> dict[str, Any]: return { "healthy": self._store is not None, "enabled": self._config.enabled, - "running": len(self._running), + "running": self._runner.count, "durable": self._store.durable if self._store is not None else None, } @@ -162,16 +137,9 @@ async def start_task( if not task: raise TaskError("A task needs a description of the work") self._check_names(profile, subject) - # Reserve a place before any await, so concurrent starts cannot all pass. - if len(self._running) + self._reserved >= self._config.max_waiting: - raise TaskLimitError( - f"{len(self._running) + self._reserved} tasks are queued or running; " - "wait for some to finish" - ) - self._reserved += 1 principal = principal or LOCAL_PRINCIPAL task_id = str(uuid.uuid4()) - try: + with self._runner.reserve(): record = await store.create( TaskRecord( id=task_id, @@ -186,13 +154,13 @@ async def start_task( parent_session_id=parent_session_id, ) ) - self._running[task_id] = _Job( - asyncio.create_task(self._run(record, principal), name=f"task-{task_id}"), + self._runner.submit( + task_id, + self._run(record, principal), store.update, self._publish, + name=f"task-{task_id}", ) - finally: - self._reserved -= 1 logger.info("Task queued", task_id=task_id, profile=profile, subject=subject) return record @@ -222,9 +190,9 @@ async def list( async def cancel(self, task_id: str, principal: Principal | None = None) -> TaskRecord: """Stop a queued or running task; a finished one is returned as it is.""" record = await self.get(task_id, principal) - if task_id not in self._running or record.finished: + if record.finished: return record - await self._cancel_job(task_id) + await self._runner.cancel(task_id) return await self.get(task_id, principal) # --- persistent agents --- @@ -324,8 +292,7 @@ async def stop_agent(self, agent_id: str, principal: Principal | None = None) -> if agent.status != "active": return agent stopped = await store.update_agent(agent_id, status="stopped", stopped_at=_now()) - for job_id in [i for i, job in self._running.items() if job.agent_id == agent_id]: - await self._cancel_job(job_id) + await self._runner.cancel_agent(agent_id) self._streaming.unpin_session(agent.session_id) logger.info("Agent stopped", agent_id=agent_id) return stopped or agent @@ -370,14 +337,8 @@ async def message_agent( agent = await self._find_agent(agent_id, profile, subject, principal) if agent.status != "active": raise AgentConflictError(f"Agent '{agent.id}' is stopped") - if len(self._running) + self._reserved >= self._config.max_waiting: - raise TaskLimitError( - f"{len(self._running) + self._reserved} background turns are queued or " - "running; wait for some to finish" - ) - self._reserved += 1 message_id = str(uuid.uuid4()) - try: + with self._runner.reserve(): record = await store.create_message( AgentMessageRecord( id=message_id, @@ -388,16 +349,14 @@ async def message_agent( parent_session_id=parent_session_id, ) ) - self._running[message_id] = _Job( - asyncio.create_task( - self._run_message(record, principal), name=f"agent-message-{message_id}" - ), + self._runner.submit( + message_id, + self._run_message(record, principal), store.update_message, self._publish_message, - agent.id, + name=f"agent-message-{message_id}", + agent_id=agent.id, ) - finally: - self._reserved -= 1 logger.info("Agent message queued", agent_id=agent.id, message_id=message_id) return record @@ -448,8 +407,8 @@ async def start() -> AssistantRequest | None: ) subject_key = (record.profile, record.subject) if record.subject is not None else None - await self._execute( - record.id, self._in_order(subject_key), start, store.update, self._publish, principal + await self._runner.execute( + record.id, subject_key, start, store.update, self._publish, principal ) async def _run_message(self, record: AgentMessageRecord, principal: Principal) -> None: @@ -461,7 +420,7 @@ async def start() -> AssistantRequest | None: agent = await store.get_agent(record.agent_id) if agent is None or agent.status != "active": # Stopped elsewhere while this waited: end it as the stop would have. - self._cancel_requested.add(record.id) + self._runner.mark_cancelled(record.id) raise asyncio.CancelledError # The agent's session now, so a move applies to every message not yet started. started = await store.update_message( @@ -483,127 +442,15 @@ async def start() -> AssistantRequest | None: config=TunableOverrides(**agent.config), ) - await self._execute( + await self._runner.execute( record.id, - self._in_order(("agent", record.agent_id)), + ("agent", record.agent_id), start, store.update_message, self._publish_message, principal, ) - async def _execute( - self, - job_id: str, - order: contextlib.AbstractAsyncContextManager[None], - start: Callable[[], Awaitable[AssistantRequest | None]], - update: Callable[..., Awaitable[Any]], - publish: Callable[[Any], Awaitable[None]], - principal: Principal, - ) -> None: - """Run one background turn in its order and a free slot; record and publish its end.""" - final: Any = None - try: - async with order, self._slots: - request = await start() - if request is None: - return - async with asyncio.timeout(self._config.timeout_seconds): - result = await self._streaming.run_message(request, principal=principal) - if result.pending_tool_call: - # A background turn has no host to perform an action and send its result back. - raise TaskError( - "The turn asked the host to perform " - f"'{result.pending_tool_call.get('tool_name')}', which a background " - "turn cannot answer" - ) - final = await update( - job_id, - only_from=_UNFINISHED, - status="done", - result=(result.content or "")[: self._config.result_max_chars], - usage=result.usage, - finished_at=_now(), - ) - except asyncio.CancelledError: - requested = job_id in self._cancel_requested - final = await update( - job_id, - only_from=_UNFINISHED, - status="cancelled" if requested else "interrupted", - error=None if requested else _INTERRUPTED, - finished_at=_now(), - ) - except TimeoutError: - final = await update( - job_id, - only_from=_UNFINISHED, - status="failed", - error=f"Timed out after {self._config.timeout_seconds} s", - finished_at=_now(), - ) - except Exception as e: - logger.warning("Background turn failed", job_id=job_id, error=str(e)) - final = await update( - job_id, - only_from=_UNFINISHED, - status="failed", - error=str(e) or type(e).__name__, - finished_at=_now(), - ) - finally: - self._running.pop(job_id, None) - self._cancel_requested.discard(job_id) - if final is not None: - await publish(final) - - async def _cancel_job(self, job_id: str) -> None: - """Cancel a queued or running background turn and wait until its end is recorded.""" - job = self._running.get(job_id) - if job is None: - return - self._cancel_requested.add(job_id) - job.worker.cancel() - with contextlib.suppress(asyncio.CancelledError, Exception): - await job.worker - await self._finish_unstarted(job_id, "cancelled") - - async def _finish_unstarted(self, job_id: str, status: str) -> None: - """Record the end of a background turn whose worker was cancelled before it ever ran. - - asyncio does not enter a coroutine cancelled before its first step, so - its own handler and ``finally`` never run; this does their work. - """ - job = self._running.pop(job_id, None) - if job is None: - return # the worker ran and finished its record itself - self._cancel_requested.discard(job_id) - final = await job.update( - job_id, - only_from=_UNFINISHED, - status=status, - error=None if status == "cancelled" else _INTERRUPTED, - finished_at=_now(), - ) - if final is not None: - await job.publish(final) - - @contextlib.asynccontextmanager - async def _in_order(self, key: Hashable | None): - """One background turn at a time per key, in arrival order; no key, no order.""" - if key is None: - yield - return - entry = self._order_locks.setdefault(key, [asyncio.Lock(), 0]) - entry[1] += 1 - try: - async with entry[0]: - yield - finally: - entry[1] -= 1 - if entry[1] == 0: - self._order_locks.pop(key, None) - async def _publish(self, record: TaskRecord) -> None: if self._events is None: return diff --git a/src/assistant_runtime/app/voice/_codex.py b/src/assistant_runtime/app/voice/_codex.py index c7cdffe4..5f557a4f 100644 --- a/src/assistant_runtime/app/voice/_codex.py +++ b/src/assistant_runtime/app/voice/_codex.py @@ -20,111 +20,20 @@ import asyncio import contextlib -import itertools import json -import os import shutil import tempfile import time -from dataclasses import dataclass, field from pathlib import Path from typing import Any from loguru import logger +from assistant_runtime.app.voice._codex_rpc import _AppServer +from assistant_runtime.app.voice._codex_schema import REALTIME_VERSION, missing_from_schema from assistant_runtime.app.voice.config import CODEX_MODEL, CODEX_VOICES from assistant_runtime.app.voice.exceptions import VoiceError -REALTIME_VERSION = "v3" -# What this transport calls and listens for; a CLI without them is incompatible. -_REQUIRED_METHODS = { - "ClientRequest.json": { - "account/read", - "thread/start", - "thread/realtime/start", - "thread/realtime/stop", - "thread/realtime/appendSpeech", - "thread/realtime/appendText", - "turn/interrupt", - }, - "ServerNotification.json": { - "thread/realtime/sdp", - "thread/realtime/started", - "thread/realtime/closed", - "thread/realtime/error", - "thread/realtime/transcript/delta", - "thread/realtime/transcript/done", - "thread/realtime/itemAdded", - "turn/started", - }, -} -_REQUIRED_THREAD_PARAMS = {"approvalPolicy", "cwd", "developerInstructions", "ephemeral", "sandbox"} -_REQUIRED_START_PARAMS = { - "clientManagedHandoffs", - "includeStartupContext", - "initialItems", - "outputModality", - "prompt", - "transport", - "version", - "voice", -} - - -def _methods(schema: dict) -> set[str]: - methods = set() - for variant in schema.get("oneOf", []): - method = (variant.get("properties") or {}).get("method") or {} - # One unusual variant must not hide the others. - names = method.get("enum") or ([method["const"]] if "const" in method else []) - methods.update(name for name in names if isinstance(name, str)) - return methods - - -def missing_from_schema(directory: Path) -> list[str]: - """Names this transport needs that the generated app-server schema lacks.""" - missing = [] - for file, needed in _REQUIRED_METHODS.items(): - try: - present = _methods(json.loads((directory / file).read_text())) - except (OSError, ValueError, KeyError, IndexError, TypeError): - present = set() - missing += sorted(needed - present) - try: - start = json.loads((directory / "v2" / "ThreadRealtimeStartParams.json").read_text()) - except (OSError, ValueError): - start = {} - try: - params = set(start.get("properties", {})) - except (AttributeError, TypeError): - params = set() - try: - versions = list(start["definitions"]["RealtimeConversationVersion"]["enum"]) - except (KeyError, TypeError): - versions = [] - missing += [f"thread/realtime/start.{name}" for name in sorted(_REQUIRED_START_PARAMS - params)] - if REALTIME_VERSION not in versions: - missing.append(f"thread/realtime/start.version={REALTIME_VERSION}") - try: - started = json.loads( - (directory / "v2" / "ThreadRealtimeStartedNotification.json").read_text() - ) - started_fields = set(started.get("properties", {})) - except (OSError, ValueError, TypeError, AttributeError): - started_fields = set() - if "version" not in started_fields: # the session's version is checked on start - missing.append("thread/realtime/started.version") - for file, method, needed in ( - ("ThreadStartParams.json", "thread/start", _REQUIRED_THREAD_PARAMS), - ("ThreadRealtimeAppendTextParams.json", "thread/realtime/appendText", {"role"}), - ): - try: - fields = set(json.loads((directory / "v2" / file).read_text()).get("properties", {})) - except (OSError, ValueError, AttributeError, TypeError): - fields = set() - missing += [f"{method}.{name}" for name in sorted(needed - fields)] - return missing - def usage_report(limits: dict, ceiling_percent: int) -> dict: """Normalise account diagnostics without deciding whether voice may run. @@ -209,194 +118,6 @@ def _error_code(message: str) -> str: return "usage_limit_reached" if "usage limit" in message.lower() else "codex_realtime_error" -@dataclass -class _Link: - """One spawned app-server, with the requests and calls that depend on it.""" - - proc: asyncio.subprocess.Process - pending: dict[int, asyncio.Future] = field(default_factory=dict) - threads: dict[str, asyncio.Queue] = field(default_factory=dict) - reader: asyncio.Task | None = None - - -class _AppServer: - """Newline-delimited JSON-RPC over ``codex app-server --stdio``.""" - - def __init__(self, command: str, timeout: float) -> None: - self.command = command - self.timeout = timeout - self._link: _Link | None = None - self._ids = itertools.count(1) - self._lock = asyncio.Lock() - - @property - def running(self) -> bool: - return self._link is not None and self._link.proc.returncode is None - - async def ensure_started(self) -> None: - async with self._lock: - if self.running: - return - env = { - k: v for k, v in os.environ.items() if k not in ("OPENAI_API_KEY", "CODEX_API_KEY") - } - try: - proc = await asyncio.create_subprocess_exec( - self.command, - "app-server", - "--stdio", - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.DEVNULL, - env=env, - limit=2**24, - ) - except OSError as exc: - raise VoiceError( - "Codex voice needs the Codex CLI; set VOICE__CODEX_COMMAND", - 503, - allocation_status="rejected", - ) from exc - link = self._link = _Link(proc) - link.reader = asyncio.create_task(self._read(link)) - try: - await self.request( - "initialize", - { - "clientInfo": {"name": "assistant-runtime", "version": "voice"}, - # thread/realtime/* is experimental; without this the - # app-server drops its notifications silently. - "capabilities": {"experimentalApi": True}, - }, - ) - await self._write(link, {"jsonrpc": "2.0", "method": "initialized"}) - await self.require_chatgpt() - except BaseException as exc: - # A half-initialised server, or one on another login, must not - # serve a later call: the next call starts and checks a new one. - with contextlib.suppress(ProcessLookupError): - proc.kill() - self._link = None - if not isinstance(exc, asyncio.CancelledError): - await asyncio.wait([link.reader], timeout=5) - raise - - async def require_chatgpt(self) -> None: - """Refuse unless the CLI is signed in with ChatGPT (never an API key).""" - account = await self.request("account/read", {}) - if (account.get("account") or {}).get("type") != "chatgpt": - raise VoiceError( - "Codex voice requires the Codex CLI to be signed in with ChatGPT; " - "run `codex login`", - 503, - allocation_status="rejected", - ) - - @staticmethod - async def _write(link: _Link, message: dict) -> None: - assert link.proc.stdin is not None - link.proc.stdin.write((json.dumps(message) + "\n").encode()) - await link.proc.stdin.drain() - - async def _read(self, link: _Link) -> None: - assert link.proc.stdout is not None - try: - while line := await link.proc.stdout.readline(): - try: - message = json.loads(line) - except ValueError: - continue - if not isinstance(message, dict): - continue - ident = message.get("id") - if "method" in message and ident is not None: - # A server request (approval, tool call): this client serves none. - with contextlib.suppress(Exception): - await self._write( - link, - { - "jsonrpc": "2.0", - "id": ident, - "error": {"code": -32601, "message": "unsupported by this client"}, - }, - ) - continue - if ident in link.pending and ("result" in message or "error" in message): - future = link.pending.pop(ident) - if not future.done(): - future.set_result(message) - continue - thread = (message.get("params") or {}).get("threadId") - if thread in link.threads: - link.threads[thread].put_nowait(message) - finally: - # Also when reading fails (an oversized line): nothing reads this - # server any more, so end it; its requests and calls end with it, - # and the next call starts a new one. - if self._link is link: - self._link = None - if link.proc.returncode is None: - with contextlib.suppress(ProcessLookupError): - link.proc.kill() - for future in link.pending.values(): - if not future.done(): - future.set_exception(VoiceError("Codex app-server exited", 502)) - link.pending.clear() - for queue in link.threads.values(): - queue.put_nowait( - {"method": "thread/realtime/closed", "params": {"reason": "app_server_exit"}} - ) - - async def request(self, method: str, params: dict) -> dict: - link = self._link - if link is None: - raise VoiceError("Codex app-server is not running", 502) - ident = next(self._ids) - future = asyncio.get_running_loop().create_future() - link.pending[ident] = future - try: - await self._write( - link, {"jsonrpc": "2.0", "id": ident, "method": method, "params": params} - ) - message = await asyncio.wait_for(future, self.timeout) - except (OSError, TimeoutError) as exc: - raise VoiceError(f"Codex {method} failed", 502) from exc - finally: - link.pending.pop(ident, None) - if "error" in message: - # Error text can carry account or prompt detail; keep it out of responses. - logger.warning("Codex app-server {} failed: {}", method, str(message["error"])[:300]) - raise VoiceError(f"Codex {method} failed", 502) - result = message.get("result") - return result if isinstance(result, dict) else {} - - def subscribe(self, thread_id: str) -> asyncio.Queue: - if self._link is None: - raise VoiceError("Codex app-server is not running", 502) - return self._link.threads.setdefault(thread_id, asyncio.Queue()) - - def unsubscribe(self, thread_id: str) -> None: - if self._link is not None: - self._link.threads.pop(thread_id, None) - - async def stop(self) -> None: - link, self._link = self._link, None - if link is None: - return - if link.proc.returncode is None: - with contextlib.suppress(Exception): - assert link.proc.stdin is not None - link.proc.stdin.close() - try: - await asyncio.wait_for(link.proc.wait(), 5) - except TimeoutError: - link.proc.kill() - await link.proc.wait() - if link.reader is not None: - link.reader.cancel() - await asyncio.gather(link.reader, return_exceptions=True) - - class _CodexConnection: """The sideband the voice service reads, as Live-style JSON events.""" diff --git a/src/assistant_runtime/app/voice/_codex_rpc.py b/src/assistant_runtime/app/voice/_codex_rpc.py new file mode 100644 index 00000000..f078608b --- /dev/null +++ b/src/assistant_runtime/app/voice/_codex_rpc.py @@ -0,0 +1,202 @@ +"""Codex app-server subprocess ownership and JSON-RPC request correlation.""" + +from __future__ import annotations + +import asyncio +import contextlib +import itertools +import json +import os +from dataclasses import dataclass, field + +from loguru import logger + +from assistant_runtime.app.voice.exceptions import VoiceError + + +@dataclass +class _Link: + """One spawned app-server, with the requests and calls that depend on it.""" + + proc: asyncio.subprocess.Process + pending: dict[int, asyncio.Future] = field(default_factory=dict) + threads: dict[str, asyncio.Queue] = field(default_factory=dict) + reader: asyncio.Task | None = None + + +class _AppServer: + """Newline-delimited JSON-RPC over ``codex app-server --stdio``.""" + + def __init__(self, command: str, timeout: float) -> None: + self.command = command + self.timeout = timeout + self._link: _Link | None = None + self._ids = itertools.count(1) + self._lock = asyncio.Lock() + + @property + def running(self) -> bool: + return self._link is not None and self._link.proc.returncode is None + + async def ensure_started(self) -> None: + async with self._lock: + if self.running: + return + env = { + k: v for k, v in os.environ.items() if k not in ("OPENAI_API_KEY", "CODEX_API_KEY") + } + try: + proc = await asyncio.create_subprocess_exec( + self.command, + "app-server", + "--stdio", + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, + env=env, + limit=2**24, + ) + except OSError as exc: + raise VoiceError( + "Codex voice needs the Codex CLI; set VOICE__CODEX_COMMAND", + 503, + allocation_status="rejected", + ) from exc + link = self._link = _Link(proc) + link.reader = asyncio.create_task(self._read(link)) + try: + await self.request( + "initialize", + { + "clientInfo": {"name": "assistant-runtime", "version": "voice"}, + # thread/realtime/* is experimental; without this the + # app-server drops its notifications silently. + "capabilities": {"experimentalApi": True}, + }, + ) + await self._write(link, {"jsonrpc": "2.0", "method": "initialized"}) + await self.require_chatgpt() + except BaseException as exc: + # A half-initialised server, or one on another login, must not + # serve a later call: the next call starts and checks a new one. + with contextlib.suppress(ProcessLookupError): + proc.kill() + self._link = None + if not isinstance(exc, asyncio.CancelledError): + await asyncio.wait([link.reader], timeout=5) + raise + + async def require_chatgpt(self) -> None: + """Refuse unless the CLI is signed in with ChatGPT (never an API key).""" + account = await self.request("account/read", {}) + if (account.get("account") or {}).get("type") != "chatgpt": + raise VoiceError( + "Codex voice requires the Codex CLI to be signed in with ChatGPT; " + "run `codex login`", + 503, + allocation_status="rejected", + ) + + @staticmethod + async def _write(link: _Link, message: dict) -> None: + assert link.proc.stdin is not None + link.proc.stdin.write((json.dumps(message) + "\n").encode()) + await link.proc.stdin.drain() + + async def _read(self, link: _Link) -> None: + assert link.proc.stdout is not None + try: + while line := await link.proc.stdout.readline(): + try: + message = json.loads(line) + except ValueError: + continue + if not isinstance(message, dict): + continue + ident = message.get("id") + if "method" in message and ident is not None: + # A server request (approval, tool call): this client serves none. + with contextlib.suppress(Exception): + await self._write( + link, + { + "jsonrpc": "2.0", + "id": ident, + "error": {"code": -32601, "message": "unsupported by this client"}, + }, + ) + continue + if ident in link.pending and ("result" in message or "error" in message): + future = link.pending.pop(ident) + if not future.done(): + future.set_result(message) + continue + thread = (message.get("params") or {}).get("threadId") + if thread in link.threads: + link.threads[thread].put_nowait(message) + finally: + # Also when reading fails (an oversized line): nothing reads this + # server any more, so end it; its requests and calls end with it, + # and the next call starts a new one. + if self._link is link: + self._link = None + if link.proc.returncode is None: + with contextlib.suppress(ProcessLookupError): + link.proc.kill() + for future in link.pending.values(): + if not future.done(): + future.set_exception(VoiceError("Codex app-server exited", 502)) + link.pending.clear() + for queue in link.threads.values(): + queue.put_nowait( + {"method": "thread/realtime/closed", "params": {"reason": "app_server_exit"}} + ) + + async def request(self, method: str, params: dict) -> dict: + link = self._link + if link is None: + raise VoiceError("Codex app-server is not running", 502) + ident = next(self._ids) + future = asyncio.get_running_loop().create_future() + link.pending[ident] = future + try: + await self._write( + link, {"jsonrpc": "2.0", "id": ident, "method": method, "params": params} + ) + message = await asyncio.wait_for(future, self.timeout) + except (OSError, TimeoutError) as exc: + raise VoiceError(f"Codex {method} failed", 502) from exc + finally: + link.pending.pop(ident, None) + if "error" in message: + # Error text can carry account or prompt detail; keep it out of responses. + logger.warning("Codex app-server {} failed: {}", method, str(message["error"])[:300]) + raise VoiceError(f"Codex {method} failed", 502) + result = message.get("result") + return result if isinstance(result, dict) else {} + + def subscribe(self, thread_id: str) -> asyncio.Queue: + if self._link is None: + raise VoiceError("Codex app-server is not running", 502) + return self._link.threads.setdefault(thread_id, asyncio.Queue()) + + def unsubscribe(self, thread_id: str) -> None: + if self._link is not None: + self._link.threads.pop(thread_id, None) + + async def stop(self) -> None: + link, self._link = self._link, None + if link is None: + return + if link.proc.returncode is None: + with contextlib.suppress(Exception): + assert link.proc.stdin is not None + link.proc.stdin.close() + try: + await asyncio.wait_for(link.proc.wait(), 5) + except TimeoutError: + link.proc.kill() + await link.proc.wait() + if link.reader is not None: + link.reader.cancel() + await asyncio.gather(link.reader, return_exceptions=True) diff --git a/src/assistant_runtime/app/voice/_codex_schema.py b/src/assistant_runtime/app/voice/_codex_schema.py new file mode 100644 index 00000000..29ba4d53 --- /dev/null +++ b/src/assistant_runtime/app/voice/_codex_schema.py @@ -0,0 +1,94 @@ +"""Offline compatibility checks for the Codex realtime protocol.""" + +import json +from pathlib import Path + +REALTIME_VERSION = "v3" +# What this transport calls and listens for; a CLI without them is incompatible. +_REQUIRED_METHODS = { + "ClientRequest.json": { + "account/read", + "thread/start", + "thread/realtime/start", + "thread/realtime/stop", + "thread/realtime/appendSpeech", + "thread/realtime/appendText", + "turn/interrupt", + }, + "ServerNotification.json": { + "thread/realtime/sdp", + "thread/realtime/started", + "thread/realtime/closed", + "thread/realtime/error", + "thread/realtime/transcript/delta", + "thread/realtime/transcript/done", + "thread/realtime/itemAdded", + "turn/started", + }, +} +_REQUIRED_THREAD_PARAMS = {"approvalPolicy", "cwd", "developerInstructions", "ephemeral", "sandbox"} +_REQUIRED_START_PARAMS = { + "clientManagedHandoffs", + "includeStartupContext", + "initialItems", + "outputModality", + "prompt", + "transport", + "version", + "voice", +} + + +def _methods(schema: dict) -> set[str]: + methods = set() + for variant in schema.get("oneOf", []): + method = (variant.get("properties") or {}).get("method") or {} + # One unusual variant must not hide the others. + names = method.get("enum") or ([method["const"]] if "const" in method else []) + methods.update(name for name in names if isinstance(name, str)) + return methods + + +def missing_from_schema(directory: Path) -> list[str]: + """Names this transport needs that the generated app-server schema lacks.""" + missing = [] + for file, needed in _REQUIRED_METHODS.items(): + try: + present = _methods(json.loads((directory / file).read_text())) + except (OSError, ValueError, KeyError, IndexError, TypeError): + present = set() + missing += sorted(needed - present) + try: + start = json.loads((directory / "v2" / "ThreadRealtimeStartParams.json").read_text()) + except (OSError, ValueError): + start = {} + try: + params = set(start.get("properties", {})) + except (AttributeError, TypeError): + params = set() + try: + versions = list(start["definitions"]["RealtimeConversationVersion"]["enum"]) + except (KeyError, TypeError): + versions = [] + missing += [f"thread/realtime/start.{name}" for name in sorted(_REQUIRED_START_PARAMS - params)] + if REALTIME_VERSION not in versions: + missing.append(f"thread/realtime/start.version={REALTIME_VERSION}") + try: + started = json.loads( + (directory / "v2" / "ThreadRealtimeStartedNotification.json").read_text() + ) + started_fields = set(started.get("properties", {})) + except (OSError, ValueError, TypeError, AttributeError): + started_fields = set() + if "version" not in started_fields: # the session's version is checked on start + missing.append("thread/realtime/started.version") + for file, method, needed in ( + ("ThreadStartParams.json", "thread/start", _REQUIRED_THREAD_PARAMS), + ("ThreadRealtimeAppendTextParams.json", "thread/realtime/appendText", {"role"}), + ): + try: + fields = set(json.loads((directory / "v2" / file).read_text()).get("properties", {})) + except (OSError, ValueError, AttributeError, TypeError): + fields = set() + missing += [f"{method}.{name}" for name in sorted(needed - fields)] + return missing diff --git a/src/assistant_runtime/app/voice/_state.py b/src/assistant_runtime/app/voice/_state.py index 9e7dd235..f3c4c0f1 100644 --- a/src/assistant_runtime/app/voice/_state.py +++ b/src/assistant_runtime/app/voice/_state.py @@ -4,8 +4,9 @@ import time from collections import deque from dataclasses import dataclass, field -from typing import Any, Literal +from typing import Literal +from assistant_runtime.app.voice._transport import VoiceConnection from assistant_runtime.app.voice.models import VoiceOffer from assistant_runtime.principal import Principal @@ -37,7 +38,7 @@ class VoiceCall: pending: dict | None = None cursor: int = 0 events: deque = field(default_factory=deque) - connection: Any = None + connection: VoiceConnection | None = None task: asyncio.Task | None = None work: asyncio.Task | None = None work_delegation: str | None = None # the delegation the running work belongs to diff --git a/src/assistant_runtime/app/voice/_transport.py b/src/assistant_runtime/app/voice/_transport.py index 2b295348..fff62ba9 100644 --- a/src/assistant_runtime/app/voice/_transport.py +++ b/src/assistant_runtime/app/voice/_transport.py @@ -13,7 +13,8 @@ import json import os import re -from typing import Any, Protocol +from collections.abc import AsyncIterator +from typing import Protocol from urllib.parse import quote import httpx @@ -26,6 +27,22 @@ _REJECTED_CREATE_STATUSES = {400, 401, 402, 403, 404, 405, 409, 413, 415, 422, 429} +class VoiceConnection(Protocol): + """The sideband operations shared by the provider adapters.""" + + def __aiter__(self) -> AsyncIterator[str | bytes]: ... + + async def send(self, frame: str) -> None: ... + + async def close(self) -> None: ... + + +class FactConnection(VoiceConnection, Protocol): + """The extra operation supplied when a transport accepts host facts.""" + + async def append_fact(self, text: str, *, speak: bool) -> None: ... + + class VoiceTransport(Protocol): """What the voice service needs from a provider.""" @@ -52,7 +69,7 @@ async def usage(self) -> dict: async def create(self, api_key: str | None, session: dict, sdp: str) -> tuple[str, str]: ... - async def attach(self, api_key: str | None, provider_id: str) -> Any: ... + async def attach(self, api_key: str | None, provider_id: str) -> VoiceConnection: ... def abandon(self, provider_id: str, api_key: str | None) -> None: """Stop a session that was created but never attached; best effort, not awaited.""" @@ -207,7 +224,7 @@ async def create(self, api_key: str, session: dict, sdp: str) -> tuple[str, str] allocation_status="unknown", ) from exc - async def attach(self, api_key: str, session_id: str) -> Any: + async def attach(self, api_key: str, session_id: str) -> VoiceConnection: # Import only when voice is enabled, with a declared optional extra. from websockets.asyncio.client import connect @@ -246,6 +263,6 @@ async def stop(self) -> None: await self.http.aclose() -async def send(connection: Any, event: dict) -> None: +async def send(connection: VoiceConnection, event: dict) -> None: async with asyncio.timeout(10): await connection.send(json.dumps(event)) diff --git a/src/assistant_runtime/app/voice/interface.py b/src/assistant_runtime/app/voice/interface.py index 1c1daafb..8206c3f3 100644 --- a/src/assistant_runtime/app/voice/interface.py +++ b/src/assistant_runtime/app/voice/interface.py @@ -14,6 +14,7 @@ import time import uuid from collections.abc import AsyncIterator +from typing import cast from loguru import logger @@ -22,7 +23,12 @@ from assistant_runtime.app.streaming.interface import StreamingService from assistant_runtime.app.voice._persistence import VoicePersistence from assistant_runtime.app.voice._state import VoiceCall -from assistant_runtime.app.voice._transport import VoiceTransport, build_transport, send +from assistant_runtime.app.voice._transport import ( + FactConnection, + VoiceTransport, + build_transport, + send, +) from assistant_runtime.app.voice.config import VoiceConfig from assistant_runtime.app.voice.exceptions import VoiceError from assistant_runtime.app.voice.models import VoiceContext, VoiceOffer, VoiceToolResult @@ -276,7 +282,9 @@ async def update_context( claimed = call.last_fact_at = time.monotonic() try: # The host's own words, as context for the voice (not backend context). - await call.connection.append_fact(update.fact, speak=update.speak) + await cast(FactConnection, call.connection).append_fact( + update.fact, speak=update.speak + ) except BaseException: if call.last_fact_at == claimed: # a newer fact may hold the slot now call.last_fact_at = previous diff --git a/src/assistant_runtime/base/__init__.py b/src/assistant_runtime/base/__init__.py index fd9b2b46..876793fc 100644 --- a/src/assistant_runtime/base/__init__.py +++ b/src/assistant_runtime/base/__init__.py @@ -1,12 +1,11 @@ -"""Base module — foundation layer providing lifecycle, instrumentation, exceptions, and resilience.""" +"""Base module — foundation layer providing lifecycle, exceptions, and resilience.""" -from assistant_runtime.base._logging import truncate, truncate_payload +from assistant_runtime.base._logging import truncate from assistant_runtime.base.exceptions import ( AssistantRuntimeError, ConfigurationError, ExternalServiceError, ) -from assistant_runtime.base.instrument import instrument from assistant_runtime.base.lifecycle import LifecycleManager from assistant_runtime.base.protocols import LifecycleAware from assistant_runtime.base.resilience import retry_with_backoff @@ -17,8 +16,6 @@ "LifecycleAware", "LifecycleManager", "AssistantRuntimeError", - "instrument", "retry_with_backoff", "truncate", - "truncate_payload", ] diff --git a/src/assistant_runtime/base/_logging.py b/src/assistant_runtime/base/_logging.py index 71c0e41e..9a3c8239 100644 --- a/src/assistant_runtime/base/_logging.py +++ b/src/assistant_runtime/base/_logging.py @@ -14,9 +14,3 @@ def truncate(text: str, max_length: int = FULL_LOG_THRESHOLD) -> str: if len(text) <= max_length: return text return text[:max_length] + f"... [{len(text) - max_length} chars truncated]" - - -def truncate_payload(payload: object, max_length: int = FULL_LOG_THRESHOLD) -> str: - """Truncate an arbitrary payload's string representation for logging.""" - text = str(payload) - return truncate(text, max_length) diff --git a/src/assistant_runtime/base/instrument.py b/src/assistant_runtime/base/instrument.py deleted file mode 100644 index 1bce54ca..00000000 --- a/src/assistant_runtime/base/instrument.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Instrumentation — @instrument decorator for async function timing and tracking.""" - -import time -from collections.abc import Callable -from functools import wraps -from typing import Any - -from loguru import logger - - -def instrument( - operation: str | None = None, - module: str | None = None, -) -> Callable: - """Decorator that instruments async functions with timing and success/failure tracking. - - Emits structured log events with operation name, module, duration_ms, and success status. - - Args: - operation: Override the operation name (defaults to function name). - module: Override the module name (defaults to first segment of __module__). - """ - - def decorator(func: Callable) -> Callable: - op_name = operation or func.__name__ - mod_name = module or func.__module__.split(".")[-2] if "." in func.__module__ else "root" - - @wraps(func) - async def wrapper(*args: Any, **kwargs: Any) -> Any: - start = time.perf_counter() - try: - result = await func(*args, **kwargs) - duration_ms = (time.perf_counter() - start) * 1000 - logger.info( - "operation.success", - module=mod_name, - operation=op_name, - duration_ms=round(duration_ms, 2), - success=True, - ) - return result - except Exception as e: - duration_ms = (time.perf_counter() - start) * 1000 - logger.warning( - "operation.failure", - module=mod_name, - operation=op_name, - duration_ms=round(duration_ms, 2), - success=False, - error_type=type(e).__name__, - error=str(e), - ) - raise - - return wrapper - - return decorator diff --git a/src/assistant_runtime/base/lifecycle.py b/src/assistant_runtime/base/lifecycle.py index 85717c59..dbe159d0 100644 --- a/src/assistant_runtime/base/lifecycle.py +++ b/src/assistant_runtime/base/lifecycle.py @@ -1,5 +1,7 @@ """Lifecycle manager — ordered startup, shutdown, and health aggregation.""" +import asyncio + from loguru import logger from assistant_runtime.base.protocols import LifecycleAware @@ -22,11 +24,12 @@ async def register(self, name: str, component: LifecycleAware) -> None: async def start_all(self) -> None: """Start all components in registration order. Rolls back on failure.""" for name, component in self._components.items(): + # Own resources as soon as startup begins, including partial acquisition. + self._started.append(name) try: await component.start() - self._started.append(name) logger.info("Component started", component=name) - except Exception as e: + except (Exception, asyncio.CancelledError) as e: logger.error("Component failed to start", component=name, error=str(e)) await self.stop_all() raise diff --git a/src/assistant_runtime/main.py b/src/assistant_runtime/main.py index 59b6505f..33e85613 100644 --- a/src/assistant_runtime/main.py +++ b/src/assistant_runtime/main.py @@ -296,7 +296,11 @@ def create_asgi_app( *, assistant: AssistantDefinition | None = None, settings: AppSettings | None = None ) -> socketio.ASGIApp: """Create the full ASGI application with Socket.IO wrapper.""" - from assistant_runtime.app.socketio_server import create_sio, forward_session_events + from assistant_runtime.app.socketio_server import ( + create_sio, + forward_ingress_events, + forward_session_events, + ) if settings is None: # Both transports read the origin rule, so resolve settings once. @@ -305,6 +309,7 @@ def create_asgi_app( sio = create_sio(settings.access) sio.fastapi_app = fastapi_app fastapi_app.state.sio = sio + fastapi_app.state.ingress_event_sink = forward_ingress_events(sio) fastapi_app.state.events.subscribe(forward_session_events(sio)) return socketio.ASGIApp(sio, fastapi_app) diff --git a/src/assistant_runtime/services/actions/_store.py b/src/assistant_runtime/services/actions/_store.py index de70035f..ec6b948f 100644 --- a/src/assistant_runtime/services/actions/_store.py +++ b/src/assistant_runtime/services/actions/_store.py @@ -10,6 +10,7 @@ import asyncio import contextlib from collections.abc import AsyncIterator +from copy import deepcopy from dataclasses import fields, replace from datetime import UTC, datetime from typing import Any, Protocol @@ -78,39 +79,41 @@ async def transaction(self) -> AsyncIterator[InMemoryActionStore]: yield self async def create_action(self, **values: Any) -> ActionRecord: - record = ActionRecord(id=self._next_action, created_at=datetime.now(UTC), **values) + record = ActionRecord( + id=self._next_action, created_at=datetime.now(UTC), **deepcopy(values) + ) self._actions[record.id] = record self._next_action += 1 - return record + return deepcopy(record) async def get_action(self, action_id: int, *, lock: bool = False) -> ActionRecord | None: - return self._actions.get(action_id) + return deepcopy(self._actions.get(action_id)) async def update_action(self, action_id: int, **values: Any) -> ActionRecord: - record = self._actions[action_id] = replace(self._actions[action_id], **values) - return record + record = self._actions[action_id] = replace(self._actions[action_id], **deepcopy(values)) + return deepcopy(record) async def list_actions(self, *, limit: int, **filters: Any) -> list[ActionRecord]: records = [r for r in self._actions.values() if _matches(r, filters)] - return sorted(records, key=lambda r: r.id, reverse=True)[:limit] + return deepcopy(sorted(records, key=lambda r: r.id, reverse=True)[:limit]) async def add_confirmation(self, **values: Any) -> ConfirmationRecord | None: if values["id"] in self._confirmations: return None - record = ConfirmationRecord(seq=self._next_seq, **values) + record = ConfirmationRecord(seq=self._next_seq, **deepcopy(values)) self._confirmations[record.id] = record self._next_seq += 1 - return record + return deepcopy(record) async def get_confirmation( self, confirmation_id: str, *, lock: bool = False ) -> ConfirmationRecord | None: - return self._confirmations.get(confirmation_id) + return deepcopy(self._confirmations.get(confirmation_id)) async def update_confirmation(self, confirmation_id: str, **values: Any) -> ConfirmationRecord: - record = replace(self._confirmations[confirmation_id], **values) + record = replace(self._confirmations[confirmation_id], **deepcopy(values)) self._confirmations[confirmation_id] = record - return record + return deepcopy(record) async def list_confirmations( self, @@ -129,7 +132,7 @@ async def list_confirmations( and (signed is None or (r.key_epoch is not None) == signed) and (reconciled is None or (r.reconciled is not None) == reconciled) ] - return sorted(records, key=lambda r: r.seq)[:limit] + return deepcopy(sorted(records, key=lambda r: r.seq)[:limit]) class DatabaseActionStore: diff --git a/src/assistant_runtime/services/database/repositories.py b/src/assistant_runtime/services/database/repositories.py deleted file mode 100644 index 54e7583b..00000000 --- a/src/assistant_runtime/services/database/repositories.py +++ /dev/null @@ -1,1219 +0,0 @@ -"""Database repositories — query layer over ORM models.""" - -from __future__ import annotations - -import json -from datetime import datetime -from typing import Any - -from sqlalchemy import case, delete, func, insert, select, update -from sqlalchemy.dialects.postgresql import insert as pg_insert -from sqlalchemy.ext.asyncio import AsyncSession - -from assistant_runtime.services.database.models import ( - ActionConfirmationORM, - ActionORM, - AgentMessageORM, - AgentORM, - ArtifactORM, - EventORM, - HostStateORM, - InboxItemORM, - MessageORM, - OAuthTokenORM, - PromptSnapshotORM, - SessionORM, - SteeringORM, - TaskORM, - TraceORM, - UserSettingsORM, -) - - -async def _in_commit_order(session: AsyncSession, feed: str) -> None: - """Hold the feed's lock until commit, so its ids become visible in increasing order. - - Identity values are handed out when a row is inserted, not when it - commits; without this a reader continuing ``after`` a later id could - skip an earlier one that commits afterwards. - """ - key = f"assistant_runtime:{feed}" - await session.execute(select(func.pg_advisory_xact_lock(func.hashtextextended(key, 0)))) - - -# Severity ordering for unsurfaced inbox queries (highest priority first). -_SEVERITY_ORDER = {"urgent": 0, "action_needed": 1, "info": 2} - - -class SessionRepository: - """CRUD operations for sessions. Uses flush() — caller owns commit.""" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def create( - self, - session_id: str, - title: str | None = None, - expires_at: datetime | None = None, - owner_id: str | None = None, - profile: str | None = None, - subject: str | None = None, - ) -> SessionORM: - """Create a new session row.""" - values: dict[str, object] = { - "id": session_id, - "title": title, - "turn_number": 0, - "owner_id": owner_id, - "profile": profile, - "subject": subject, - } - if expires_at is not None: - values["expires_at"] = expires_at - - result = await self._session.execute( - insert(SessionORM).values(**values).returning(SessionORM) - ) - await self._session.flush() - return result.scalar_one() - - async def get(self, session_id: str) -> SessionORM | None: - """Get a non-expired session by ID, or None.""" - result = await self._session.execute( - select(SessionORM).where( - SessionORM.id == session_id, - SessionORM.expires_at > func.now(), - ) - ) - return result.scalar_one_or_none() - - async def list_all( - self, limit: int = 50, offset: int = 0, *, owner_id: str | None = None - ) -> list[SessionORM]: - """List non-expired sessions, most recently updated first. - - With ``owner_id``, only that principal's sessions are listed. - """ - stmt = select(SessionORM).where(SessionORM.expires_at > func.now()) - if owner_id is not None: - stmt = stmt.where(SessionORM.owner_id == owner_id) - result = await self._session.execute( - stmt.order_by(SessionORM.updated_at.desc()).limit(limit).offset(offset) - ) - return list(result.scalars().all()) - - async def get_latest_by_telegram_chat_id(self, chat_id: str) -> SessionORM | None: - """Return the newest Telegram-bound session for a chat, or None.""" - result = await self._session.execute( - select(SessionORM) - .where( - SessionORM.telegram_chat_id == chat_id, - SessionORM.expires_at > func.now(), - ) - .order_by(SessionORM.telegram_bound_at.desc(), SessionORM.updated_at.desc()) - .limit(1) - ) - return result.scalar_one_or_none() - - async def update(self, session_id: str, **fields) -> None: - """Update specific fields on a session. No-op if session doesn't exist.""" - if not fields: - return - stmt = ( - update(SessionORM) - .where( - SessionORM.id == session_id, - SessionORM.expires_at > func.now(), - ) - .values(**fields, updated_at=func.now()) - ) - await self._session.execute(stmt) - await self._session.flush() - - async def delete(self, session_id: str) -> bool: - """Delete a session. Returns True if a row was deleted.""" - result = await self._session.execute(delete(SessionORM).where(SessionORM.id == session_id)) - await self._session.flush() - return (result.rowcount or 0) > 0 - - async def exists(self, session_id: str) -> bool: - """Check if a non-expired session exists.""" - result = await self._session.execute( - select(SessionORM.id).where( - SessionORM.id == session_id, - SessionORM.expires_at > func.now(), - ) - ) - return result.scalar_one_or_none() is not None - - async def cleanup_expired(self) -> int: - """Delete all expired sessions. Returns count of deleted rows.""" - result = await self._session.execute( - delete(SessionORM).where(SessionORM.expires_at <= func.now()) - ) - await self._session.flush() - return result.rowcount or 0 - - async def upsert( - self, - session_id: str, - *, - title: str | None = None, - turn_number: int = 0, - working_memory: dict | None = None, - telegram_chat_id: str | None = None, - telegram_bound_at: datetime | None = None, - expires_at: datetime | None = None, - owner_id: str | None = None, - pending_action: dict | None = None, - profile: str | None = None, - subject: str | None = None, - ) -> None: - """Atomic INSERT ... ON CONFLICT DO UPDATE. - - Eliminates the race condition in check-then-insert patterns. The - owner is written on insert only: reassignment goes through - ``update()``, so a stale state save can never restore an owner an - administrator cleared or changed. The profile and subject binding - is written on insert only too: it never changes. - """ - values: dict[str, object] = { - "id": session_id, - "title": title, - "turn_number": turn_number, - "working_memory": working_memory, - "telegram_chat_id": telegram_chat_id, - "telegram_bound_at": telegram_bound_at, - "owner_id": owner_id, - "pending_action": pending_action, - "profile": profile, - "subject": subject, - } - if expires_at is not None: - values["expires_at"] = expires_at - - stmt = pg_insert(SessionORM).values(**values) - - update_fields: dict[str, object] = { - "title": stmt.excluded.title, - "turn_number": stmt.excluded.turn_number, - "working_memory": stmt.excluded.working_memory, - "telegram_chat_id": stmt.excluded.telegram_chat_id, - "telegram_bound_at": stmt.excluded.telegram_bound_at, - "pending_action": stmt.excluded.pending_action, - "updated_at": func.now(), - } - if expires_at is not None: - update_fields["expires_at"] = stmt.excluded.expires_at - - stmt = stmt.on_conflict_do_update(index_elements=["id"], set_=update_fields) - await self._session.execute(stmt) - await self._session.flush() - - -class MessageRepository: - """CRUD operations for conversation messages.""" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def create( - self, - *, - message_id: str, - session_id: str, - parent_id: str | None, - role: str, - message_type: str, - content: str, - segments: list[dict[str, Any]] | None = None, - usage: dict[str, Any] | None = None, - prompt: dict[str, Any] | None = None, - model_messages: list[dict[str, Any]] | None = None, - ) -> MessageORM: - result = await self._session.execute( - insert(MessageORM) - .values( - id=message_id, - session_id=session_id, - parent_id=parent_id, - role=role, - message_type=message_type, - content=content, - segments=segments, - usage=usage, - prompt=prompt, - model_messages=model_messages, - ) - .returning(MessageORM) - ) - await self._session.flush() - return result.scalar_one() - - async def get(self, message_id: str) -> MessageORM | None: - result = await self._session.execute(select(MessageORM).where(MessageORM.id == message_id)) - return result.scalar_one_or_none() - - async def list_by_session(self, session_id: str) -> list[MessageORM]: - result = await self._session.execute( - select(MessageORM) - .where(MessageORM.session_id == session_id) - .order_by(MessageORM.created_at.asc()) - ) - return list(result.scalars().all()) - - async def count_by_sessions(self, session_ids: list[str]) -> dict[str, int]: - if not session_ids: - return {} - result = await self._session.execute( - select(MessageORM.session_id, func.count(MessageORM.id)) - .where(MessageORM.session_id.in_(session_ids)) - .group_by(MessageORM.session_id) - ) - return {session_id: int(count) for session_id, count in result.all()} - - async def update( - self, - message_id: str, - *, - content: str | None = None, - segments: list[dict[str, Any]] | None = None, - usage: dict[str, Any] | None = None, - prompt: dict[str, Any] | None = None, - model_messages: list[dict[str, Any]] | None = None, - ) -> None: - fields: dict[str, Any] = {} - if content is not None: - fields["content"] = content - if segments is not None: - fields["segments"] = segments - if usage is not None: - fields["usage"] = usage - if prompt is not None: - fields["prompt"] = prompt - if model_messages is not None: - fields["model_messages"] = model_messages - if not fields: - return - await self._session.execute( - update(MessageORM).where(MessageORM.id == message_id).values(**fields) - ) - await self._session.flush() - - -class PromptSnapshotRepository: - """Stable prompt texts by hash. Uses flush() — caller owns commit.""" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def put(self, snapshot_hash: str, content: str) -> None: - """Store the text unless its hash is already there.""" - await self._session.execute( - pg_insert(PromptSnapshotORM) - .values(hash=snapshot_hash, content=content) - .on_conflict_do_nothing(index_elements=["hash"]) - ) - await self._session.flush() - - async def get(self, snapshot_hash: str) -> str | None: - result = await self._session.execute( - select(PromptSnapshotORM.content).where(PromptSnapshotORM.hash == snapshot_hash) - ) - return result.scalar_one_or_none() - - -class TaskRepository: - """Background task records. Uses flush() — caller owns commit.""" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def create(self, **values: Any) -> TaskORM: - result = await self._session.execute(insert(TaskORM).values(**values).returning(TaskORM)) - await self._session.flush() - return result.scalar_one() - - async def update( - self, task_id: str, *, only_from: tuple[str, ...] | None = None, **fields: Any - ) -> TaskORM | None: - """Change the task; with ``only_from``, only while its status is one of them.""" - statement = update(TaskORM).where(TaskORM.id == task_id) - if only_from is not None: - statement = statement.where(TaskORM.status.in_(only_from)) - result = await self._session.execute(statement.values(**fields).returning(TaskORM)) - await self._session.flush() - return result.scalar_one_or_none() - - async def get(self, task_id: str) -> TaskORM | None: - result = await self._session.execute(select(TaskORM).where(TaskORM.id == task_id)) - return result.scalar_one_or_none() - - async def list( - self, - *, - created_by: str | None = None, - parent_session_id: str | None = None, - status: str | None = None, - limit: int = 50, - ) -> list[TaskORM]: - """Newest first, filtered by whatever is given.""" - query = select(TaskORM).order_by(TaskORM.created_at.desc(), TaskORM.id.desc()) - if created_by is not None: - query = query.where(TaskORM.created_by == created_by) - if parent_session_id is not None: - query = query.where(TaskORM.parent_session_id == parent_session_id) - if status is not None: - query = query.where(TaskORM.status == status) - result = await self._session.execute(query.limit(limit)) - return list(result.scalars().all()) - - async def mark_unfinished(self, status: str, error: str) -> list[TaskORM]: - """Finish every queued or running task with ``status`` (after a restart).""" - result = await self._session.execute( - update(TaskORM) - .where(TaskORM.status.in_(("queued", "running"))) - .values(status=status, error=error, finished_at=func.now()) - .returning(TaskORM) - ) - await self._session.flush() - return list(result.scalars().all()) - - -class AgentRepository: - """Persistent agents and their messages. Uses flush() — caller owns commit.""" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def create(self, **values: Any) -> AgentORM: - result = await self._session.execute(insert(AgentORM).values(**values).returning(AgentORM)) - await self._session.flush() - return result.scalar_one() - - async def update(self, agent_id: str, **fields: Any) -> AgentORM | None: - result = await self._session.execute( - update(AgentORM).where(AgentORM.id == agent_id).values(**fields).returning(AgentORM) - ) - await self._session.flush() - return result.scalar_one_or_none() - - async def get(self, agent_id: str) -> AgentORM | None: - result = await self._session.execute(select(AgentORM).where(AgentORM.id == agent_id)) - return result.scalar_one_or_none() - - async def list( - self, *, created_by: str | None = None, status: str | None = None - ) -> list[AgentORM]: - """Newest first, filtered by whatever is given.""" - query = select(AgentORM).order_by(AgentORM.created_at.desc(), AgentORM.id.desc()) - if created_by is not None: - query = query.where(AgentORM.created_by == created_by) - if status is not None: - query = query.where(AgentORM.status == status) - result = await self._session.execute(query) - return list(result.scalars().all()) - - async def create_message(self, **values: Any) -> AgentMessageORM: - result = await self._session.execute( - insert(AgentMessageORM).values(**values).returning(AgentMessageORM) - ) - await self._session.flush() - return result.scalar_one() - - async def update_message( - self, message_id: str, *, only_from: tuple[str, ...] | None = None, **fields: Any - ) -> AgentMessageORM | None: - """Change the message; with ``only_from``, only while its status is one of them.""" - statement = update(AgentMessageORM).where(AgentMessageORM.id == message_id) - if only_from is not None: - statement = statement.where(AgentMessageORM.status.in_(only_from)) - result = await self._session.execute(statement.values(**fields).returning(AgentMessageORM)) - await self._session.flush() - return result.scalar_one_or_none() - - async def get_message(self, message_id: str) -> AgentMessageORM | None: - result = await self._session.execute( - select(AgentMessageORM).where(AgentMessageORM.id == message_id) - ) - return result.scalar_one_or_none() - - async def list_messages( - self, - *, - agent_id: str | None = None, - created_by: str | None = None, - parent_session_id: str | None = None, - status: str | None = None, - limit: int = 50, - ) -> list[AgentMessageORM]: - """Newest first, filtered by whatever is given.""" - query = select(AgentMessageORM).order_by( - AgentMessageORM.created_at.desc(), AgentMessageORM.id.desc() - ) - if agent_id is not None: - query = query.where(AgentMessageORM.agent_id == agent_id) - if created_by is not None: - query = query.where(AgentMessageORM.created_by == created_by) - if parent_session_id is not None: - query = query.where(AgentMessageORM.parent_session_id == parent_session_id) - if status is not None: - query = query.where(AgentMessageORM.status == status) - result = await self._session.execute(query.limit(limit)) - return list(result.scalars().all()) - - async def mark_unfinished_messages(self, status: str, error: str) -> list[AgentMessageORM]: - """Finish every queued or running message with ``status`` (after a restart).""" - result = await self._session.execute( - update(AgentMessageORM) - .where(AgentMessageORM.status.in_(("queued", "running"))) - .values(status=status, error=error, finished_at=func.now()) - .returning(AgentMessageORM) - ) - await self._session.flush() - return list(result.scalars().all()) - - -class EventRepository: - """Event records, in arrival order. Uses flush() — caller owns commit.""" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def create_if_new(self, **values: Any) -> tuple[EventORM, bool]: - """Insert the event, or return the stored one with the same (source, event_id).""" - await _in_commit_order(self._session, "events") - result = await self._session.execute( - pg_insert(EventORM) - .values(**values) - .on_conflict_do_nothing(constraint="uq_events_source_event_id") - .returning(EventORM) - ) - row = result.scalar_one_or_none() - await self._session.flush() - if row is not None: - return row, True - existing = await self._session.execute( - select(EventORM).where( - EventORM.source == values["source"], EventORM.event_id == values["event_id"] - ) - ) - return existing.scalar_one(), False - - async def get(self, event_id: int, *, lock: bool = False) -> EventORM | None: - query = select(EventORM).where(EventORM.id == event_id) - result = await self._session.execute(query.with_for_update() if lock else query) - return result.scalar_one_or_none() - - async def update(self, event_id: int, **fields: Any) -> EventORM | None: - result = await self._session.execute( - update(EventORM).where(EventORM.id == event_id).values(**fields).returning(EventORM) - ) - await self._session.flush() - return result.scalar_one_or_none() - - async def list(self, *, after: int, limit: int, **filters: Any) -> list[EventORM]: - """Events after ``after`` in id order; ``history=False`` leaves out imported history.""" - query = select(EventORM).where(EventORM.id > after) - for name, value in filters.items(): - if value is not None: - query = query.where(getattr(EventORM, name) == value) - result = await self._session.execute(query.order_by(EventORM.id).limit(limit)) - return list(result.scalars().all()) - - -class ActionRepository: - """Action records and their confirmations. Uses flush() — caller owns commit.""" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def create(self, **values: Any) -> ActionORM: - result = await self._session.execute( - insert(ActionORM).values(**values).returning(ActionORM) - ) - await self._session.flush() - return result.scalar_one() - - async def get(self, action_id: int, *, lock: bool = False) -> ActionORM | None: - query = select(ActionORM).where(ActionORM.id == action_id) - result = await self._session.execute(query.with_for_update() if lock else query) - return result.scalar_one_or_none() - - async def update(self, action_id: int, **fields: Any) -> ActionORM: - result = await self._session.execute( - update(ActionORM).where(ActionORM.id == action_id).values(**fields).returning(ActionORM) - ) - await self._session.flush() - return result.scalar_one() - - async def list(self, *, limit: int, **filters: Any) -> list[ActionORM]: - """Newest first, filtered by whatever is given.""" - query = select(ActionORM) - for name, value in filters.items(): - if value is not None: - query = query.where(getattr(ActionORM, name) == value) - result = await self._session.execute(query.order_by(ActionORM.id.desc()).limit(limit)) - return list(result.scalars().all()) - - async def add_confirmation(self, **values: Any) -> ActionConfirmationORM | None: - """Insert a confirmation; None when its id is already taken.""" - await _in_commit_order(self._session, "action_confirmations") - result = await self._session.execute( - pg_insert(ActionConfirmationORM) - .values(**values) - .on_conflict_do_nothing(index_elements=["id"]) - .returning(ActionConfirmationORM) - ) - await self._session.flush() - return result.scalar_one_or_none() - - async def get_confirmation( - self, confirmation_id: str, *, lock: bool = False - ) -> ActionConfirmationORM | None: - query = select(ActionConfirmationORM).where(ActionConfirmationORM.id == confirmation_id) - result = await self._session.execute(query.with_for_update() if lock else query) - return result.scalar_one_or_none() - - async def update_confirmation( - self, confirmation_id: str, **fields: Any - ) -> ActionConfirmationORM: - result = await self._session.execute( - update(ActionConfirmationORM) - .where(ActionConfirmationORM.id == confirmation_id) - .values(**fields) - .returning(ActionConfirmationORM) - ) - await self._session.flush() - return result.scalar_one() - - async def list_confirmations( - self, - *, - after: int, - limit: int, - signed: bool | None = None, - reconciled: bool | None = None, - **filters: Any, - ) -> list[ActionConfirmationORM]: - """Confirmations after ``after`` in insertion order.""" - model = ActionConfirmationORM - query = select(model).where(model.seq > after) - for name, value in filters.items(): - if value is not None: - query = query.where(getattr(model, name) == value) - if signed is not None: - query = query.where( - model.key_epoch.is_not(None) if signed else model.key_epoch.is_(None) - ) - if reconciled is not None: - query = query.where( - model.reconciled.is_not(None) if reconciled else model.reconciled.is_(None) - ) - result = await self._session.execute(query.order_by(model.seq).limit(limit)) - return list(result.scalars().all()) - - -class HostStateRepository: - """Versioned host values. Every write is one conditional statement, so none is lost. - - A deleted key is kept as a tombstone with its version, so a version is - never reused and a stale write cannot match a value created later. - """ - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def get(self, namespace: str, key: str) -> HostStateORM | None: - result = await self._session.execute( - select(HostStateORM).where( - HostStateORM.namespace == namespace, - HostStateORM.key == key, - HostStateORM.deleted.is_(False), - ) - ) - return result.scalar_one_or_none() - - async def list(self, namespace: str, after: str | None, limit: int) -> list[HostStateORM]: - query = select(HostStateORM).where( - HostStateORM.namespace == namespace, HostStateORM.deleted.is_(False) - ) - if after is not None: - query = query.where(HostStateORM.key > after) - result = await self._session.execute(query.order_by(HostStateORM.key).limit(limit)) - return list(result.scalars().all()) - - async def put( - self, namespace: str, key: str, value: Any, by: str, expected_version: int | None - ) -> HostStateORM | None: - """Write the value; None when ``expected_version`` does not match what is stored. - - ``expected_version`` 0 means no value may be stored (a tombstone is - revived with its next version); ``None`` writes whatever is there. - """ - model = HostStateORM - if expected_version is None or expected_version == 0: - statement = pg_insert(model).values( - namespace=namespace, key=key, value=value, version=1, updated_by=by - ) - statement = statement.on_conflict_do_update( - index_elements=["namespace", "key"], - set_={ - "value": statement.excluded.value, - "version": model.version + 1, - "deleted": False, - "updated_by": statement.excluded.updated_by, - "updated_at": func.now(), - }, - where=model.deleted.is_(True) if expected_version == 0 else None, - ) - else: - statement = ( - update(model) - .where( - model.namespace == namespace, - model.key == key, - model.version == expected_version, - model.deleted.is_(False), - ) - .values( - value=value, version=model.version + 1, updated_by=by, updated_at=func.now() - ) - ) - result = await self._session.execute(statement.returning(model)) - await self._session.flush() - return result.scalar_one_or_none() - - async def delete(self, namespace: str, key: str, by: str, expected_version: int | None) -> bool: - """Leave a tombstone at the next version; False when nothing matched.""" - model = HostStateORM - statement = update(model).where( - model.namespace == namespace, model.key == key, model.deleted.is_(False) - ) - if expected_version is not None: - statement = statement.where(model.version == expected_version) - result = await self._session.execute( - statement.values( - value=None, - deleted=True, - version=model.version + 1, - updated_by=by, - updated_at=func.now(), - ).returning(model.key) - ) - await self._session.flush() - return result.scalar_one_or_none() is not None - - -class SteeringRepository: - """CRUD operations for steering records.""" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def create( - self, - *, - steering_id: str, - session_id: str, - content: str, - status: str, - profile: str | None = None, - delivered_at: datetime | None = None, - attachments: list[dict[str, Any]] | None = None, - ) -> SteeringORM: - result = await self._session.execute( - insert(SteeringORM) - .values( - id=steering_id, - session_id=session_id, - profile=profile, - content=content, - status=status, - delivered_at=delivered_at, - attachments=attachments, - ) - .returning(SteeringORM) - ) - await self._session.flush() - return result.scalar_one() - - async def list_by_session(self, session_id: str) -> list[SteeringORM]: - result = await self._session.execute( - select(SteeringORM) - .where(SteeringORM.session_id == session_id) - .order_by(SteeringORM.created_at.asc(), SteeringORM.id.asc()) - ) - return list(result.scalars().all()) - - async def mark_status( - self, - steering_ids: list[str], - *, - status: str, - delivered_at: datetime | None, - ) -> None: - if not steering_ids: - return - - await self._session.execute( - update(SteeringORM) - .where(SteeringORM.id.in_(steering_ids)) - .values(status=status, delivered_at=delivered_at) - ) - await self._session.flush() - - -class TraceRepository: - """CRUD operations for debug traces. Uses flush() — caller owns commit.""" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def create( - self, - trace_id: str, - session_id: str, - events: list[dict[str, Any]], - *, - user_message: str | None = None, - is_continuation: bool = False, - duration_ms: float | None = None, - ) -> TraceORM: - """Create a new trace row.""" - result = await self._session.execute( - insert(TraceORM) - .values( - id=trace_id, - session_id=session_id, - events=events, - user_message=user_message, - is_continuation=is_continuation, - duration_ms=duration_ms, - ) - .returning(TraceORM) - ) - await self._session.flush() - return result.scalar_one() - - async def delete_older_than(self, cutoff: datetime) -> int: - """Delete traces created before ``cutoff``. Returns the number of rows removed.""" - result = await self._session.execute(delete(TraceORM).where(TraceORM.created_at < cutoff)) - await self._session.flush() - return result.rowcount or 0 - - async def list_by_session( - self, session_id: str, limit: int = 50, offset: int = 0 - ) -> list[TraceORM]: - """List traces for a session, ordered by creation time.""" - result = await self._session.execute( - select(TraceORM) - .where(TraceORM.session_id == session_id) - .order_by(TraceORM.created_at.asc()) - .limit(limit) - .offset(offset) - ) - return list(result.scalars().all()) - - -class SettingsRepository: - """CRUD operations for user settings. Uses flush() — caller owns commit.""" - - SETTINGS_ID: str = "default" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def get(self) -> UserSettingsORM | None: - """Get the settings row, or None.""" - result = await self._session.execute( - select(UserSettingsORM).where(UserSettingsORM.id == self.SETTINGS_ID) - ) - return result.scalar_one_or_none() - - async def save(self, state: dict[str, Any]) -> UserSettingsORM: - """Upsert settings — create or update the single settings row.""" - if not state: - row = await self.get() - if row is not None: - return row - result = await self._session.execute( - insert(UserSettingsORM).values(id=self.SETTINGS_ID).returning(UserSettingsORM) - ) - await self._session.flush() - return result.scalar_one() - - stmt = pg_insert(UserSettingsORM).values(id=self.SETTINGS_ID, **state) - result = await self._session.execute( - stmt.on_conflict_do_update( - index_elements=["id"], - set_={ - **{key: getattr(stmt.excluded, key) for key in state}, - "updated_at": func.now(), - }, - ).returning(UserSettingsORM) - ) - await self._session.flush() - return result.scalar_one() - - -class InboxRepository: - """CRUD operations for inbox items. Uses flush() — caller owns commit.""" - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def create( - self, - from_agent: str, - message: str, - severity: str = "info", - context: dict[str, Any] | None = None, - ) -> InboxItemORM: - """Create a new inbox item.""" - result = await self._session.execute( - insert(InboxItemORM) - .values( - from_agent=from_agent, - message=message, - severity=severity, - context=context, - ) - .returning(InboxItemORM) - ) - await self._session.flush() - return result.scalar_one() - - async def list_unsurfaced(self, limit: int = 20) -> list[InboxItemORM]: - """List unsurfaced items ordered by severity (urgent first) then created_at desc.""" - severity_sort = case( - _SEVERITY_ORDER, - value=InboxItemORM.severity, - else_=99, - ) - result = await self._session.execute( - select(InboxItemORM) - .where(InboxItemORM.surfaced.is_(False)) - .order_by(severity_sort, InboxItemORM.created_at.desc()) - .limit(limit) - ) - return list(result.scalars().all()) - - async def mark_surfaced(self, item_id: str) -> InboxItemORM | None: - """Mark an item as surfaced. Returns the item or None if not found.""" - result = await self._session.execute( - update(InboxItemORM) - .where(InboxItemORM.id == item_id) - .values(surfaced=True) - .returning(InboxItemORM) - ) - row = result.scalar_one_or_none() - if row is not None: - await self._session.flush() - return row - - async def list_all(self, limit: int = 50, offset: int = 0) -> list[InboxItemORM]: - """List all items ordered by created_at desc.""" - result = await self._session.execute( - select(InboxItemORM) - .order_by(InboxItemORM.created_at.desc()) - .limit(limit) - .offset(offset) - ) - return list(result.scalars().all()) - - -class ArtifactRepository: - """CRUD operations for versioned prompt artifacts. Uses flush() — caller owns commit. - - Every method takes the profile ``scope`` first and a ``subject`` (empty - for a profile-scoped artifact): rows of different assistants or subjects - never mix, even when artifact names coincide. - """ - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - @staticmethod - def _key(scope: str, subject: str, name: str) -> Any: - return ( - (ArtifactORM.assistant == scope) - & (ArtifactORM.subject == subject) - & (ArtifactORM.name == name) - ) - - async def lock(self, scope: str, name: str, *, subject: str = "") -> None: - """Serialize mutations until commit, including an artifact's first version.""" - key = json.dumps([scope, name] if not subject else [scope, subject, name]) - await self._session.execute( - select(func.pg_advisory_xact_lock(func.hashtextextended(key, 0))) - ) - - async def get_active(self, scope: str, name: str, *, subject: str = "") -> ArtifactORM | None: - """Get the active version of an artifact by name.""" - result = await self._session.execute( - select(ArtifactORM).where( - self._key(scope, subject, name), ArtifactORM.is_active.is_(True) - ) - ) - return result.scalar_one_or_none() - - async def get_all_active(self, scope: str, *, subject: str = "") -> list[ArtifactORM]: - """Get all active artifacts of the scope and subject, ordered by name.""" - result = await self._session.execute( - select(ArtifactORM) - .where( - ArtifactORM.assistant == scope, - ArtifactORM.subject == subject, - ArtifactORM.is_active.is_(True), - ) - .order_by(ArtifactORM.name) - ) - return list(result.scalars().all()) - - async def get_subjects(self, scope: str) -> list[str]: - """The subjects with stored versions in the scope, sorted.""" - result = await self._session.execute( - select(ArtifactORM.subject) - .where(ArtifactORM.assistant == scope, ArtifactORM.subject != "") - .distinct() - .order_by(ArtifactORM.subject) - ) - return list(result.scalars().all()) - - async def delete_by_name(self, scope: str, name: str, *, subject: str = "") -> int: - """Delete all versions of an artifact by name. Returns count deleted.""" - result = await self._session.execute( - delete(ArtifactORM).where(self._key(scope, subject, name)) - ) - await self._session.flush() - return result.rowcount or 0 - - async def get_history( - self, scope: str, name: str, limit: int = 20, *, subject: str = "" - ) -> list[ArtifactORM]: - """Get version history for an artifact, newest first.""" - result = await self._session.execute( - select(ArtifactORM) - .where(self._key(scope, subject, name)) - .order_by(ArtifactORM.version.desc()) - .limit(limit) - ) - return list(result.scalars().all()) - - async def get_version( - self, scope: str, name: str, version: int, *, subject: str = "" - ) -> ArtifactORM | None: - """One stored version, or None.""" - result = await self._session.execute( - select(ArtifactORM).where( - self._key(scope, subject, name), ArtifactORM.version == version - ) - ) - return result.scalar_one_or_none() - - async def get_by_status( - self, scope: str, status: str, limit: int | None = None, before_id: int | None = None - ) -> list[ArtifactORM]: - """Versions of the scope (every subject) in ``status``, newest first; ``before_id`` pages back.""" - query = ( - select(ArtifactORM) - .where(ArtifactORM.assistant == scope, ArtifactORM.status == status) - .order_by(ArtifactORM.id.desc()) - ) - if before_id is not None: - query = query.where(ArtifactORM.id < before_id) - if limit is not None: - query = query.limit(limit) - result = await self._session.execute(query) - return list(result.scalars().all()) - - async def propose( - self, - scope: str, - name: str, - content: str, - proposed_by: str, - *, - subject: str = "", - actor_kind: str = "host", - rationale: str | None = None, - ) -> ArtifactORM: - """Create a pending version inside the caller's locked artifact transaction.""" - max_result = await self._session.execute( - select(func.max(ArtifactORM.version)).where(self._key(scope, subject, name)) - ) - result = await self._session.execute( - insert(ArtifactORM) - .values( - assistant=scope, - subject=subject, - name=name, - content=content, - version=(max_result.scalar_one_or_none() or 0) + 1, - is_active=False, - status="pending", - proposed_by=proposed_by, - actor_kind=actor_kind, - rationale=rationale, - ) - .returning(ArtifactORM) - ) - await self._session.flush() - return result.scalar_one() - - async def approve( - self, - scope: str, - name: str, - version: int, - *, - subject: str = "", - decided_by: str | None = None, - ) -> ArtifactORM | None: - """Activate a version; the previously active one becomes superseded. - - ``decided_by`` is recorded when the target was a pending proposal. - Returns None if the target version doesn't exist. - """ - target = await self.get_version(scope, name, version, subject=subject) - if target is None: - return None - target_id, was_pending = target.id, target.status == "pending" - - # Deactivate first: at most one active row per artifact (a unique index). - await self._session.execute( - update(ArtifactORM) - .where( - self._key(scope, subject, name), - ArtifactORM.is_active.is_(True), - ArtifactORM.id != target_id, - ) - .values(is_active=False, status="superseded") - ) - values: dict[str, Any] = {"is_active": True, "status": "active"} - if decided_by is not None and was_pending: - values.update(decided_by=decided_by, decided_at=func.now()) - result = await self._session.execute( - update(ArtifactORM) - .where(ArtifactORM.id == target_id) - .values(**values) - .returning(ArtifactORM) - ) - await self._session.flush() - return result.scalar_one() - - async def reject( - self, - scope: str, - name: str, - version: int, - *, - subject: str = "", - decided_by: str, - reason: str | None, - ) -> ArtifactORM | None: - """Mark a pending version rejected; None when it is not pending (or absent).""" - result = await self._session.execute( - update(ArtifactORM) - .where( - self._key(scope, subject, name), - ArtifactORM.version == version, - ArtifactORM.status == "pending", - ) - .values( - status="rejected", - decided_by=decided_by, - decided_at=func.now(), - decision_reason=reason, - ) - .returning(ArtifactORM) - ) - await self._session.flush() - return result.scalar_one_or_none() - - async def prune_superseded(self, scope: str, name: str, keep: int, *, subject: str = "") -> int: - """Delete superseded versions beyond the newest ``keep``. Returns count deleted.""" - kept = ( - select(ArtifactORM.id) - .where(self._key(scope, subject, name), ArtifactORM.status == "superseded") - .order_by(ArtifactORM.version.desc()) - .limit(keep) - ) - result = await self._session.execute( - delete(ArtifactORM).where( - self._key(scope, subject, name), - ArtifactORM.status == "superseded", - ArtifactORM.id.not_in(kept.scalar_subquery()), - ) - ) - await self._session.flush() - return result.rowcount or 0 - - -class OAuthTokenRepository: - """CRUD operations for OAuth tokens. Uses flush() — caller owns commit. - - ``kind`` is "api_key" for the provider-key store and "login" for a - ChatGPT/Codex login; every operation touches only the row of that kind. - """ - - def __init__(self, session: AsyncSession) -> None: - self._session = session - - async def get(self, provider: str, kind: str) -> OAuthTokenORM | None: - """Get the provider's token of one kind.""" - result = await self._session.execute( - select(OAuthTokenORM).where( - OAuthTokenORM.provider == provider, OAuthTokenORM.kind == kind - ) - ) - return result.scalar_one_or_none() - - async def upsert( - self, - provider: str, - kind: str, - *, - encrypted_api_key: str | None = None, - encrypted_refresh_token: str | None = None, - encrypted_id_token: str | None = None, - expires_at: float | None = None, - email: str | None = None, - ) -> None: - """Insert or update the provider's token of one kind.""" - values: dict[str, object] = { - "provider": provider, - "kind": kind, - "encrypted_api_key": encrypted_api_key, - "encrypted_refresh_token": encrypted_refresh_token, - "encrypted_id_token": encrypted_id_token, - "expires_at": expires_at, - "email": email, - } - - stmt = pg_insert(OAuthTokenORM).values(**values) - stmt = stmt.on_conflict_do_update( - index_elements=["provider", "kind"], - set_={ - "encrypted_api_key": stmt.excluded.encrypted_api_key, - "encrypted_refresh_token": stmt.excluded.encrypted_refresh_token, - "encrypted_id_token": stmt.excluded.encrypted_id_token, - "expires_at": stmt.excluded.expires_at, - "email": stmt.excluded.email, - "updated_at": func.now(), - }, - ) - await self._session.execute(stmt) - await self._session.flush() - - async def delete(self, provider: str, kind: str) -> bool: - """Delete the provider's token of one kind. Returns True if deleted.""" - result = await self._session.execute( - delete(OAuthTokenORM).where( - OAuthTokenORM.provider == provider, OAuthTokenORM.kind == kind - ) - ) - await self._session.flush() - return (result.rowcount or 0) > 0 diff --git a/src/assistant_runtime/services/database/repositories/__init__.py b/src/assistant_runtime/services/database/repositories/__init__.py new file mode 100644 index 00000000..a8154bc0 --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/__init__.py @@ -0,0 +1,34 @@ +"""Database query API, with implementations grouped by record domain.""" + +from ._actions import ActionRepository +from ._artifacts import ArtifactRepository +from ._events import EventRepository +from ._host_state import HostStateRepository +from ._inbox import InboxRepository +from ._oauth import OAuthTokenRepository +from ._sessions import ( + MessageRepository, + PromptSnapshotRepository, + SessionRepository, + SteeringRepository, +) +from ._settings import SettingsRepository +from ._tasks import AgentRepository, TaskRepository +from ._traces import TraceRepository + +__all__ = [ + "SessionRepository", + "MessageRepository", + "PromptSnapshotRepository", + "SteeringRepository", + "TaskRepository", + "AgentRepository", + "EventRepository", + "ActionRepository", + "HostStateRepository", + "TraceRepository", + "SettingsRepository", + "InboxRepository", + "ArtifactRepository", + "OAuthTokenRepository", +] diff --git a/src/assistant_runtime/services/database/repositories/_actions.py b/src/assistant_runtime/services/database/repositories/_actions.py new file mode 100644 index 00000000..8f9007d4 --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_actions.py @@ -0,0 +1,108 @@ +"""Actions queries; the caller owns the transaction.""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import insert, select, update +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession + +from assistant_runtime.services.database.models import ( + ActionConfirmationORM, + ActionORM, +) + +from ._ordering import _in_commit_order + + +class ActionRepository: + """Action records and their confirmations. Uses flush() — caller owns commit.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def create(self, **values: Any) -> ActionORM: + result = await self._session.execute( + insert(ActionORM).values(**values).returning(ActionORM) + ) + await self._session.flush() + return result.scalar_one() + + async def get(self, action_id: int, *, lock: bool = False) -> ActionORM | None: + query = select(ActionORM).where(ActionORM.id == action_id) + result = await self._session.execute(query.with_for_update() if lock else query) + return result.scalar_one_or_none() + + async def update(self, action_id: int, **fields: Any) -> ActionORM: + result = await self._session.execute( + update(ActionORM).where(ActionORM.id == action_id).values(**fields).returning(ActionORM) + ) + await self._session.flush() + return result.scalar_one() + + async def list(self, *, limit: int, **filters: Any) -> list[ActionORM]: + """Newest first, filtered by whatever is given.""" + query = select(ActionORM) + for name, value in filters.items(): + if value is not None: + query = query.where(getattr(ActionORM, name) == value) + result = await self._session.execute(query.order_by(ActionORM.id.desc()).limit(limit)) + return list(result.scalars().all()) + + async def add_confirmation(self, **values: Any) -> ActionConfirmationORM | None: + """Insert a confirmation; None when its id is already taken.""" + await _in_commit_order(self._session, "action_confirmations") + result = await self._session.execute( + pg_insert(ActionConfirmationORM) + .values(**values) + .on_conflict_do_nothing(index_elements=["id"]) + .returning(ActionConfirmationORM) + ) + await self._session.flush() + return result.scalar_one_or_none() + + async def get_confirmation( + self, confirmation_id: str, *, lock: bool = False + ) -> ActionConfirmationORM | None: + query = select(ActionConfirmationORM).where(ActionConfirmationORM.id == confirmation_id) + result = await self._session.execute(query.with_for_update() if lock else query) + return result.scalar_one_or_none() + + async def update_confirmation( + self, confirmation_id: str, **fields: Any + ) -> ActionConfirmationORM: + result = await self._session.execute( + update(ActionConfirmationORM) + .where(ActionConfirmationORM.id == confirmation_id) + .values(**fields) + .returning(ActionConfirmationORM) + ) + await self._session.flush() + return result.scalar_one() + + async def list_confirmations( + self, + *, + after: int, + limit: int, + signed: bool | None = None, + reconciled: bool | None = None, + **filters: Any, + ) -> list[ActionConfirmationORM]: + """Confirmations after ``after`` in insertion order.""" + model = ActionConfirmationORM + query = select(model).where(model.seq > after) + for name, value in filters.items(): + if value is not None: + query = query.where(getattr(model, name) == value) + if signed is not None: + query = query.where( + model.key_epoch.is_not(None) if signed else model.key_epoch.is_(None) + ) + if reconciled is not None: + query = query.where( + model.reconciled.is_not(None) if reconciled else model.reconciled.is_(None) + ) + result = await self._session.execute(query.order_by(model.seq).limit(limit)) + return list(result.scalars().all()) diff --git a/src/assistant_runtime/services/database/repositories/_artifacts.py b/src/assistant_runtime/services/database/repositories/_artifacts.py new file mode 100644 index 00000000..6b2128e7 --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_artifacts.py @@ -0,0 +1,241 @@ +"""Artifacts queries; the caller owns the transaction.""" + +from __future__ import annotations + +import json +from typing import Any + +from sqlalchemy import delete, func, insert, select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from assistant_runtime.services.database.models import ( + ArtifactORM, +) + + +class ArtifactRepository: + """CRUD operations for versioned prompt artifacts. Uses flush() — caller owns commit. + + Every method takes the profile ``scope`` first and a ``subject`` (empty + for a profile-scoped artifact): rows of different assistants or subjects + never mix, even when artifact names coincide. + """ + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + @staticmethod + def _key(scope: str, subject: str, name: str) -> Any: + return ( + (ArtifactORM.assistant == scope) + & (ArtifactORM.subject == subject) + & (ArtifactORM.name == name) + ) + + async def lock(self, scope: str, name: str, *, subject: str = "") -> None: + """Serialize mutations until commit, including an artifact's first version.""" + key = json.dumps([scope, name] if not subject else [scope, subject, name]) + await self._session.execute( + select(func.pg_advisory_xact_lock(func.hashtextextended(key, 0))) + ) + + async def get_active(self, scope: str, name: str, *, subject: str = "") -> ArtifactORM | None: + """Get the active version of an artifact by name.""" + result = await self._session.execute( + select(ArtifactORM).where( + self._key(scope, subject, name), ArtifactORM.is_active.is_(True) + ) + ) + return result.scalar_one_or_none() + + async def get_all_active(self, scope: str, *, subject: str = "") -> list[ArtifactORM]: + """Get all active artifacts of the scope and subject, ordered by name.""" + result = await self._session.execute( + select(ArtifactORM) + .where( + ArtifactORM.assistant == scope, + ArtifactORM.subject == subject, + ArtifactORM.is_active.is_(True), + ) + .order_by(ArtifactORM.name) + ) + return list(result.scalars().all()) + + async def get_subjects(self, scope: str) -> list[str]: + """The subjects with stored versions in the scope, sorted.""" + result = await self._session.execute( + select(ArtifactORM.subject) + .where(ArtifactORM.assistant == scope, ArtifactORM.subject != "") + .distinct() + .order_by(ArtifactORM.subject) + ) + return list(result.scalars().all()) + + async def delete_by_name(self, scope: str, name: str, *, subject: str = "") -> int: + """Delete all versions of an artifact by name. Returns count deleted.""" + result = await self._session.execute( + delete(ArtifactORM).where(self._key(scope, subject, name)) + ) + await self._session.flush() + return result.rowcount or 0 + + async def get_history( + self, scope: str, name: str, limit: int = 20, *, subject: str = "" + ) -> list[ArtifactORM]: + """Get version history for an artifact, newest first.""" + result = await self._session.execute( + select(ArtifactORM) + .where(self._key(scope, subject, name)) + .order_by(ArtifactORM.version.desc()) + .limit(limit) + ) + return list(result.scalars().all()) + + async def get_version( + self, scope: str, name: str, version: int, *, subject: str = "" + ) -> ArtifactORM | None: + """One stored version, or None.""" + result = await self._session.execute( + select(ArtifactORM).where( + self._key(scope, subject, name), ArtifactORM.version == version + ) + ) + return result.scalar_one_or_none() + + async def get_by_status( + self, scope: str, status: str, limit: int | None = None, before_id: int | None = None + ) -> list[ArtifactORM]: + """Versions of the scope (every subject) in ``status``, newest first; ``before_id`` pages back.""" + query = ( + select(ArtifactORM) + .where(ArtifactORM.assistant == scope, ArtifactORM.status == status) + .order_by(ArtifactORM.id.desc()) + ) + if before_id is not None: + query = query.where(ArtifactORM.id < before_id) + if limit is not None: + query = query.limit(limit) + result = await self._session.execute(query) + return list(result.scalars().all()) + + async def propose( + self, + scope: str, + name: str, + content: str, + proposed_by: str, + *, + subject: str = "", + actor_kind: str = "host", + rationale: str | None = None, + ) -> ArtifactORM: + """Create a pending version inside the caller's locked artifact transaction.""" + max_result = await self._session.execute( + select(func.max(ArtifactORM.version)).where(self._key(scope, subject, name)) + ) + result = await self._session.execute( + insert(ArtifactORM) + .values( + assistant=scope, + subject=subject, + name=name, + content=content, + version=(max_result.scalar_one_or_none() or 0) + 1, + is_active=False, + status="pending", + proposed_by=proposed_by, + actor_kind=actor_kind, + rationale=rationale, + ) + .returning(ArtifactORM) + ) + await self._session.flush() + return result.scalar_one() + + async def approve( + self, + scope: str, + name: str, + version: int, + *, + subject: str = "", + decided_by: str | None = None, + ) -> ArtifactORM | None: + """Activate a version; the previously active one becomes superseded. + + ``decided_by`` is recorded when the target was a pending proposal. + Returns None if the target version doesn't exist. + """ + target = await self.get_version(scope, name, version, subject=subject) + if target is None: + return None + target_id, was_pending = target.id, target.status == "pending" + + # Deactivate first: at most one active row per artifact (a unique index). + await self._session.execute( + update(ArtifactORM) + .where( + self._key(scope, subject, name), + ArtifactORM.is_active.is_(True), + ArtifactORM.id != target_id, + ) + .values(is_active=False, status="superseded") + ) + values: dict[str, Any] = {"is_active": True, "status": "active"} + if decided_by is not None and was_pending: + values.update(decided_by=decided_by, decided_at=func.now()) + result = await self._session.execute( + update(ArtifactORM) + .where(ArtifactORM.id == target_id) + .values(**values) + .returning(ArtifactORM) + ) + await self._session.flush() + return result.scalar_one() + + async def reject( + self, + scope: str, + name: str, + version: int, + *, + subject: str = "", + decided_by: str, + reason: str | None, + ) -> ArtifactORM | None: + """Mark a pending version rejected; None when it is not pending (or absent).""" + result = await self._session.execute( + update(ArtifactORM) + .where( + self._key(scope, subject, name), + ArtifactORM.version == version, + ArtifactORM.status == "pending", + ) + .values( + status="rejected", + decided_by=decided_by, + decided_at=func.now(), + decision_reason=reason, + ) + .returning(ArtifactORM) + ) + await self._session.flush() + return result.scalar_one_or_none() + + async def prune_superseded(self, scope: str, name: str, keep: int, *, subject: str = "") -> int: + """Delete superseded versions beyond the newest ``keep``. Returns count deleted.""" + kept = ( + select(ArtifactORM.id) + .where(self._key(scope, subject, name), ArtifactORM.status == "superseded") + .order_by(ArtifactORM.version.desc()) + .limit(keep) + ) + result = await self._session.execute( + delete(ArtifactORM).where( + self._key(scope, subject, name), + ArtifactORM.status == "superseded", + ArtifactORM.id.not_in(kept.scalar_subquery()), + ) + ) + await self._session.flush() + return result.rowcount or 0 diff --git a/src/assistant_runtime/services/database/repositories/_events.py b/src/assistant_runtime/services/database/repositories/_events.py new file mode 100644 index 00000000..b101f880 --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_events.py @@ -0,0 +1,63 @@ +"""Events queries; the caller owns the transaction.""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import select, update +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession + +from assistant_runtime.services.database.models import ( + EventORM, +) + +from ._ordering import _in_commit_order + + +class EventRepository: + """Event records, in arrival order. Uses flush() — caller owns commit.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def create_if_new(self, **values: Any) -> tuple[EventORM, bool]: + """Insert the event, or return the stored one with the same (source, event_id).""" + await _in_commit_order(self._session, "events") + result = await self._session.execute( + pg_insert(EventORM) + .values(**values) + .on_conflict_do_nothing(constraint="uq_events_source_event_id") + .returning(EventORM) + ) + row = result.scalar_one_or_none() + await self._session.flush() + if row is not None: + return row, True + existing = await self._session.execute( + select(EventORM).where( + EventORM.source == values["source"], EventORM.event_id == values["event_id"] + ) + ) + return existing.scalar_one(), False + + async def get(self, event_id: int, *, lock: bool = False) -> EventORM | None: + query = select(EventORM).where(EventORM.id == event_id) + result = await self._session.execute(query.with_for_update() if lock else query) + return result.scalar_one_or_none() + + async def update(self, event_id: int, **fields: Any) -> EventORM | None: + result = await self._session.execute( + update(EventORM).where(EventORM.id == event_id).values(**fields).returning(EventORM) + ) + await self._session.flush() + return result.scalar_one_or_none() + + async def list(self, *, after: int, limit: int, **filters: Any) -> list[EventORM]: + """Events after ``after`` in id order; ``history=False`` leaves out imported history.""" + query = select(EventORM).where(EventORM.id > after) + for name, value in filters.items(): + if value is not None: + query = query.where(getattr(EventORM, name) == value) + result = await self._session.execute(query.order_by(EventORM.id).limit(limit)) + return list(result.scalars().all()) diff --git a/src/assistant_runtime/services/database/repositories/_host_state.py b/src/assistant_runtime/services/database/repositories/_host_state.py new file mode 100644 index 00000000..425ca041 --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_host_state.py @@ -0,0 +1,104 @@ +"""Host state queries; the caller owns the transaction.""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import func, select, update +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession + +from assistant_runtime.services.database.models import ( + HostStateORM, +) + + +class HostStateRepository: + """Versioned host values. Every write is one conditional statement, so none is lost. + + A deleted key is kept as a tombstone with its version, so a version is + never reused and a stale write cannot match a value created later. + """ + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def get(self, namespace: str, key: str) -> HostStateORM | None: + result = await self._session.execute( + select(HostStateORM).where( + HostStateORM.namespace == namespace, + HostStateORM.key == key, + HostStateORM.deleted.is_(False), + ) + ) + return result.scalar_one_or_none() + + async def list(self, namespace: str, after: str | None, limit: int) -> list[HostStateORM]: + query = select(HostStateORM).where( + HostStateORM.namespace == namespace, HostStateORM.deleted.is_(False) + ) + if after is not None: + query = query.where(HostStateORM.key > after) + result = await self._session.execute(query.order_by(HostStateORM.key).limit(limit)) + return list(result.scalars().all()) + + async def put( + self, namespace: str, key: str, value: Any, by: str, expected_version: int | None + ) -> HostStateORM | None: + """Write the value; None when ``expected_version`` does not match what is stored. + + ``expected_version`` 0 means no value may be stored (a tombstone is + revived with its next version); ``None`` writes whatever is there. + """ + model = HostStateORM + if expected_version is None or expected_version == 0: + statement = pg_insert(model).values( + namespace=namespace, key=key, value=value, version=1, updated_by=by + ) + statement = statement.on_conflict_do_update( + index_elements=["namespace", "key"], + set_={ + "value": statement.excluded.value, + "version": model.version + 1, + "deleted": False, + "updated_by": statement.excluded.updated_by, + "updated_at": func.now(), + }, + where=model.deleted.is_(True) if expected_version == 0 else None, + ) + else: + statement = ( + update(model) + .where( + model.namespace == namespace, + model.key == key, + model.version == expected_version, + model.deleted.is_(False), + ) + .values( + value=value, version=model.version + 1, updated_by=by, updated_at=func.now() + ) + ) + result = await self._session.execute(statement.returning(model)) + await self._session.flush() + return result.scalar_one_or_none() + + async def delete(self, namespace: str, key: str, by: str, expected_version: int | None) -> bool: + """Leave a tombstone at the next version; False when nothing matched.""" + model = HostStateORM + statement = update(model).where( + model.namespace == namespace, model.key == key, model.deleted.is_(False) + ) + if expected_version is not None: + statement = statement.where(model.version == expected_version) + result = await self._session.execute( + statement.values( + value=None, + deleted=True, + version=model.version + 1, + updated_by=by, + updated_at=func.now(), + ).returning(model.key) + ) + await self._session.flush() + return result.scalar_one_or_none() is not None diff --git a/src/assistant_runtime/services/database/repositories/_inbox.py b/src/assistant_runtime/services/database/repositories/_inbox.py new file mode 100644 index 00000000..e8c30d09 --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_inbox.py @@ -0,0 +1,80 @@ +"""Inbox queries; the caller owns the transaction.""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import case, insert, select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from assistant_runtime.services.database.models import ( + InboxItemORM, +) + +_SEVERITY_ORDER = {"urgent": 0, "action_needed": 1, "info": 2} + + +class InboxRepository: + """CRUD operations for inbox items. Uses flush() — caller owns commit.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def create( + self, + from_agent: str, + message: str, + severity: str = "info", + context: dict[str, Any] | None = None, + ) -> InboxItemORM: + """Create a new inbox item.""" + result = await self._session.execute( + insert(InboxItemORM) + .values( + from_agent=from_agent, + message=message, + severity=severity, + context=context, + ) + .returning(InboxItemORM) + ) + await self._session.flush() + return result.scalar_one() + + async def list_unsurfaced(self, limit: int = 20) -> list[InboxItemORM]: + """List unsurfaced items ordered by severity (urgent first) then created_at desc.""" + severity_sort = case( + _SEVERITY_ORDER, + value=InboxItemORM.severity, + else_=99, + ) + result = await self._session.execute( + select(InboxItemORM) + .where(InboxItemORM.surfaced.is_(False)) + .order_by(severity_sort, InboxItemORM.created_at.desc()) + .limit(limit) + ) + return list(result.scalars().all()) + + async def mark_surfaced(self, item_id: str) -> InboxItemORM | None: + """Mark an item as surfaced. Returns the item or None if not found.""" + result = await self._session.execute( + update(InboxItemORM) + .where(InboxItemORM.id == item_id) + .values(surfaced=True) + .returning(InboxItemORM) + ) + row = result.scalar_one_or_none() + if row is not None: + await self._session.flush() + return row + + async def list_all(self, limit: int = 50, offset: int = 0) -> list[InboxItemORM]: + """List all items ordered by created_at desc.""" + result = await self._session.execute( + select(InboxItemORM) + .order_by(InboxItemORM.created_at.desc()) + .limit(limit) + .offset(offset) + ) + return list(result.scalars().all()) diff --git a/src/assistant_runtime/services/database/repositories/_oauth.py b/src/assistant_runtime/services/database/repositories/_oauth.py new file mode 100644 index 00000000..c48754cb --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_oauth.py @@ -0,0 +1,78 @@ +"""Oauth queries; the caller owns the transaction.""" + +from __future__ import annotations + +from sqlalchemy import delete, func, select +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession + +from assistant_runtime.services.database.models import ( + OAuthTokenORM, +) + + +class OAuthTokenRepository: + """CRUD operations for OAuth tokens. Uses flush() — caller owns commit. + + ``kind`` is "api_key" for the provider-key store and "login" for a + ChatGPT/Codex login; every operation touches only the row of that kind. + """ + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def get(self, provider: str, kind: str) -> OAuthTokenORM | None: + """Get the provider's token of one kind.""" + result = await self._session.execute( + select(OAuthTokenORM).where( + OAuthTokenORM.provider == provider, OAuthTokenORM.kind == kind + ) + ) + return result.scalar_one_or_none() + + async def upsert( + self, + provider: str, + kind: str, + *, + encrypted_api_key: str | None = None, + encrypted_refresh_token: str | None = None, + encrypted_id_token: str | None = None, + expires_at: float | None = None, + email: str | None = None, + ) -> None: + """Insert or update the provider's token of one kind.""" + values: dict[str, object] = { + "provider": provider, + "kind": kind, + "encrypted_api_key": encrypted_api_key, + "encrypted_refresh_token": encrypted_refresh_token, + "encrypted_id_token": encrypted_id_token, + "expires_at": expires_at, + "email": email, + } + + stmt = pg_insert(OAuthTokenORM).values(**values) + stmt = stmt.on_conflict_do_update( + index_elements=["provider", "kind"], + set_={ + "encrypted_api_key": stmt.excluded.encrypted_api_key, + "encrypted_refresh_token": stmt.excluded.encrypted_refresh_token, + "encrypted_id_token": stmt.excluded.encrypted_id_token, + "expires_at": stmt.excluded.expires_at, + "email": stmt.excluded.email, + "updated_at": func.now(), + }, + ) + await self._session.execute(stmt) + await self._session.flush() + + async def delete(self, provider: str, kind: str) -> bool: + """Delete the provider's token of one kind. Returns True if deleted.""" + result = await self._session.execute( + delete(OAuthTokenORM).where( + OAuthTokenORM.provider == provider, OAuthTokenORM.kind == kind + ) + ) + await self._session.flush() + return (result.rowcount or 0) > 0 diff --git a/src/assistant_runtime/services/database/repositories/_ordering.py b/src/assistant_runtime/services/database/repositories/_ordering.py new file mode 100644 index 00000000..d98408e5 --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_ordering.py @@ -0,0 +1,15 @@ +"""Commit ordering for cursor-based record feeds.""" + +from sqlalchemy import func, select +from sqlalchemy.ext.asyncio import AsyncSession + + +async def _in_commit_order(session: AsyncSession, feed: str) -> None: + """Hold the feed's lock until commit, so its ids become visible in increasing order. + + Identity values are handed out when a row is inserted, not when it + commits; without this a reader continuing ``after`` a later id could + skip an earlier one that commits afterwards. + """ + key = f"assistant_runtime:{feed}" + await session.execute(select(func.pg_advisory_xact_lock(func.hashtextextended(key, 0)))) diff --git a/src/assistant_runtime/services/database/repositories/_sessions.py b/src/assistant_runtime/services/database/repositories/_sessions.py new file mode 100644 index 00000000..95089b82 --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_sessions.py @@ -0,0 +1,345 @@ +"""Sessions queries; the caller owns the transaction.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from sqlalchemy import delete, func, insert, select, update +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession + +from assistant_runtime.services.database.models import ( + MessageORM, + PromptSnapshotORM, + SessionORM, + SteeringORM, +) + + +class SessionRepository: + """CRUD operations for sessions. Uses flush() — caller owns commit.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def create( + self, + session_id: str, + title: str | None = None, + expires_at: datetime | None = None, + owner_id: str | None = None, + profile: str | None = None, + subject: str | None = None, + ) -> SessionORM: + """Create a new session row.""" + values: dict[str, object] = { + "id": session_id, + "title": title, + "turn_number": 0, + "owner_id": owner_id, + "profile": profile, + "subject": subject, + } + if expires_at is not None: + values["expires_at"] = expires_at + + result = await self._session.execute( + insert(SessionORM).values(**values).returning(SessionORM) + ) + await self._session.flush() + return result.scalar_one() + + async def get(self, session_id: str) -> SessionORM | None: + """Get a non-expired session by ID, or None.""" + result = await self._session.execute( + select(SessionORM).where( + SessionORM.id == session_id, + SessionORM.expires_at > func.now(), + ) + ) + return result.scalar_one_or_none() + + async def list_all( + self, limit: int = 50, offset: int = 0, *, owner_id: str | None = None + ) -> list[SessionORM]: + """List non-expired sessions, most recently updated first. + + With ``owner_id``, only that principal's sessions are listed. + """ + stmt = select(SessionORM).where(SessionORM.expires_at > func.now()) + if owner_id is not None: + stmt = stmt.where(SessionORM.owner_id == owner_id) + result = await self._session.execute( + stmt.order_by(SessionORM.updated_at.desc()).limit(limit).offset(offset) + ) + return list(result.scalars().all()) + + async def get_latest_by_telegram_chat_id(self, chat_id: str) -> SessionORM | None: + """Return the newest Telegram-bound session for a chat, or None.""" + result = await self._session.execute( + select(SessionORM) + .where( + SessionORM.telegram_chat_id == chat_id, + SessionORM.expires_at > func.now(), + ) + .order_by(SessionORM.telegram_bound_at.desc(), SessionORM.updated_at.desc()) + .limit(1) + ) + return result.scalar_one_or_none() + + async def update(self, session_id: str, **fields) -> None: + """Update specific fields on a session. No-op if session doesn't exist.""" + if not fields: + return + stmt = ( + update(SessionORM) + .where( + SessionORM.id == session_id, + SessionORM.expires_at > func.now(), + ) + .values(**fields, updated_at=func.now()) + ) + await self._session.execute(stmt) + await self._session.flush() + + async def delete(self, session_id: str) -> bool: + """Delete a session. Returns True if a row was deleted.""" + result = await self._session.execute(delete(SessionORM).where(SessionORM.id == session_id)) + await self._session.flush() + return (result.rowcount or 0) > 0 + + async def cleanup_expired(self) -> int: + """Delete all expired sessions. Returns count of deleted rows.""" + result = await self._session.execute( + delete(SessionORM).where(SessionORM.expires_at <= func.now()) + ) + await self._session.flush() + return result.rowcount or 0 + + async def upsert( + self, + session_id: str, + *, + title: str | None = None, + turn_number: int = 0, + working_memory: dict | None = None, + telegram_chat_id: str | None = None, + telegram_bound_at: datetime | None = None, + expires_at: datetime | None = None, + owner_id: str | None = None, + pending_action: dict | None = None, + profile: str | None = None, + subject: str | None = None, + ) -> None: + """Atomic INSERT ... ON CONFLICT DO UPDATE. + + Eliminates the race condition in check-then-insert patterns. The + owner is written on insert only: reassignment goes through + ``update()``, so a stale state save can never restore an owner an + administrator cleared or changed. The profile and subject binding + is written on insert only too: it never changes. + """ + values: dict[str, object] = { + "id": session_id, + "title": title, + "turn_number": turn_number, + "working_memory": working_memory, + "telegram_chat_id": telegram_chat_id, + "telegram_bound_at": telegram_bound_at, + "owner_id": owner_id, + "pending_action": pending_action, + "profile": profile, + "subject": subject, + } + if expires_at is not None: + values["expires_at"] = expires_at + + stmt = pg_insert(SessionORM).values(**values) + + update_fields: dict[str, object] = { + "title": stmt.excluded.title, + "turn_number": stmt.excluded.turn_number, + "working_memory": stmt.excluded.working_memory, + "telegram_chat_id": stmt.excluded.telegram_chat_id, + "telegram_bound_at": stmt.excluded.telegram_bound_at, + "pending_action": stmt.excluded.pending_action, + "updated_at": func.now(), + } + if expires_at is not None: + update_fields["expires_at"] = stmt.excluded.expires_at + + stmt = stmt.on_conflict_do_update(index_elements=["id"], set_=update_fields) + await self._session.execute(stmt) + await self._session.flush() + + +class MessageRepository: + """CRUD operations for conversation messages.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def create( + self, + *, + message_id: str, + session_id: str, + parent_id: str | None, + role: str, + message_type: str, + content: str, + segments: list[dict[str, Any]] | None = None, + usage: dict[str, Any] | None = None, + prompt: dict[str, Any] | None = None, + model_messages: list[dict[str, Any]] | None = None, + ) -> MessageORM: + result = await self._session.execute( + insert(MessageORM) + .values( + id=message_id, + session_id=session_id, + parent_id=parent_id, + role=role, + message_type=message_type, + content=content, + segments=segments, + usage=usage, + prompt=prompt, + model_messages=model_messages, + ) + .returning(MessageORM) + ) + await self._session.flush() + return result.scalar_one() + + async def get(self, message_id: str) -> MessageORM | None: + result = await self._session.execute(select(MessageORM).where(MessageORM.id == message_id)) + return result.scalar_one_or_none() + + async def list_by_session(self, session_id: str) -> list[MessageORM]: + result = await self._session.execute( + select(MessageORM) + .where(MessageORM.session_id == session_id) + .order_by(MessageORM.created_at.asc()) + ) + return list(result.scalars().all()) + + async def count_by_sessions(self, session_ids: list[str]) -> dict[str, int]: + if not session_ids: + return {} + result = await self._session.execute( + select(MessageORM.session_id, func.count(MessageORM.id)) + .where(MessageORM.session_id.in_(session_ids)) + .group_by(MessageORM.session_id) + ) + return {session_id: int(count) for session_id, count in result.all()} + + async def update( + self, + message_id: str, + *, + content: str | None = None, + segments: list[dict[str, Any]] | None = None, + usage: dict[str, Any] | None = None, + prompt: dict[str, Any] | None = None, + model_messages: list[dict[str, Any]] | None = None, + ) -> None: + fields: dict[str, Any] = {} + if content is not None: + fields["content"] = content + if segments is not None: + fields["segments"] = segments + if usage is not None: + fields["usage"] = usage + if prompt is not None: + fields["prompt"] = prompt + if model_messages is not None: + fields["model_messages"] = model_messages + if not fields: + return + await self._session.execute( + update(MessageORM).where(MessageORM.id == message_id).values(**fields) + ) + await self._session.flush() + + +class PromptSnapshotRepository: + """Stable prompt texts by hash. Uses flush() — caller owns commit.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def put(self, snapshot_hash: str, content: str) -> None: + """Store the text unless its hash is already there.""" + await self._session.execute( + pg_insert(PromptSnapshotORM) + .values(hash=snapshot_hash, content=content) + .on_conflict_do_nothing(index_elements=["hash"]) + ) + await self._session.flush() + + async def get(self, snapshot_hash: str) -> str | None: + result = await self._session.execute( + select(PromptSnapshotORM.content).where(PromptSnapshotORM.hash == snapshot_hash) + ) + return result.scalar_one_or_none() + + +class SteeringRepository: + """CRUD operations for steering records.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def create( + self, + *, + steering_id: str, + session_id: str, + content: str, + status: str, + profile: str | None = None, + delivered_at: datetime | None = None, + attachments: list[dict[str, Any]] | None = None, + ) -> SteeringORM: + result = await self._session.execute( + insert(SteeringORM) + .values( + id=steering_id, + session_id=session_id, + profile=profile, + content=content, + status=status, + delivered_at=delivered_at, + attachments=attachments, + ) + .returning(SteeringORM) + ) + await self._session.flush() + return result.scalar_one() + + async def list_by_session(self, session_id: str) -> list[SteeringORM]: + result = await self._session.execute( + select(SteeringORM) + .where(SteeringORM.session_id == session_id) + .order_by(SteeringORM.created_at.asc(), SteeringORM.id.asc()) + ) + return list(result.scalars().all()) + + async def mark_status( + self, + steering_ids: list[str], + *, + status: str, + delivered_at: datetime | None, + ) -> None: + if not steering_ids: + return + + await self._session.execute( + update(SteeringORM) + .where(SteeringORM.id.in_(steering_ids)) + .values(status=status, delivered_at=delivered_at) + ) + await self._session.flush() diff --git a/src/assistant_runtime/services/database/repositories/_settings.py b/src/assistant_runtime/services/database/repositories/_settings.py new file mode 100644 index 00000000..fa49e3ab --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_settings.py @@ -0,0 +1,54 @@ +"""Settings queries; the caller owns the transaction.""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import func, insert, select +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession + +from assistant_runtime.services.database.models import ( + UserSettingsORM, +) + + +class SettingsRepository: + """CRUD operations for user settings. Uses flush() — caller owns commit.""" + + SETTINGS_ID: str = "default" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def get(self) -> UserSettingsORM | None: + """Get the settings row, or None.""" + result = await self._session.execute( + select(UserSettingsORM).where(UserSettingsORM.id == self.SETTINGS_ID) + ) + return result.scalar_one_or_none() + + async def save(self, state: dict[str, Any]) -> UserSettingsORM: + """Upsert settings — create or update the single settings row.""" + if not state: + row = await self.get() + if row is not None: + return row + result = await self._session.execute( + insert(UserSettingsORM).values(id=self.SETTINGS_ID).returning(UserSettingsORM) + ) + await self._session.flush() + return result.scalar_one() + + stmt = pg_insert(UserSettingsORM).values(id=self.SETTINGS_ID, **state) + result = await self._session.execute( + stmt.on_conflict_do_update( + index_elements=["id"], + set_={ + **{key: getattr(stmt.excluded, key) for key in state}, + "updated_at": func.now(), + }, + ).returning(UserSettingsORM) + ) + await self._session.flush() + return result.scalar_one() diff --git a/src/assistant_runtime/services/database/repositories/_tasks.py b/src/assistant_runtime/services/database/repositories/_tasks.py new file mode 100644 index 00000000..dae87529 --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_tasks.py @@ -0,0 +1,165 @@ +"""Tasks queries; the caller owns the transaction.""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import func, insert, select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from assistant_runtime.services.database.models import ( + AgentMessageORM, + AgentORM, + TaskORM, +) + + +class TaskRepository: + """Background task records. Uses flush() — caller owns commit.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def create(self, **values: Any) -> TaskORM: + result = await self._session.execute(insert(TaskORM).values(**values).returning(TaskORM)) + await self._session.flush() + return result.scalar_one() + + async def update( + self, task_id: str, *, only_from: tuple[str, ...] | None = None, **fields: Any + ) -> TaskORM | None: + """Change the task; with ``only_from``, only while its status is one of them.""" + statement = update(TaskORM).where(TaskORM.id == task_id) + if only_from is not None: + statement = statement.where(TaskORM.status.in_(only_from)) + result = await self._session.execute(statement.values(**fields).returning(TaskORM)) + await self._session.flush() + return result.scalar_one_or_none() + + async def get(self, task_id: str) -> TaskORM | None: + result = await self._session.execute(select(TaskORM).where(TaskORM.id == task_id)) + return result.scalar_one_or_none() + + async def list( + self, + *, + created_by: str | None = None, + parent_session_id: str | None = None, + status: str | None = None, + limit: int = 50, + ) -> list[TaskORM]: + """Newest first, filtered by whatever is given.""" + query = select(TaskORM).order_by(TaskORM.created_at.desc(), TaskORM.id.desc()) + if created_by is not None: + query = query.where(TaskORM.created_by == created_by) + if parent_session_id is not None: + query = query.where(TaskORM.parent_session_id == parent_session_id) + if status is not None: + query = query.where(TaskORM.status == status) + result = await self._session.execute(query.limit(limit)) + return list(result.scalars().all()) + + async def mark_unfinished(self, status: str, error: str) -> list[TaskORM]: + """Finish every queued or running task with ``status`` (after a restart).""" + result = await self._session.execute( + update(TaskORM) + .where(TaskORM.status.in_(("queued", "running"))) + .values(status=status, error=error, finished_at=func.now()) + .returning(TaskORM) + ) + await self._session.flush() + return list(result.scalars().all()) + + +class AgentRepository: + """Persistent agents and their messages. Uses flush() — caller owns commit.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def create(self, **values: Any) -> AgentORM: + result = await self._session.execute(insert(AgentORM).values(**values).returning(AgentORM)) + await self._session.flush() + return result.scalar_one() + + async def update(self, agent_id: str, **fields: Any) -> AgentORM | None: + result = await self._session.execute( + update(AgentORM).where(AgentORM.id == agent_id).values(**fields).returning(AgentORM) + ) + await self._session.flush() + return result.scalar_one_or_none() + + async def get(self, agent_id: str) -> AgentORM | None: + result = await self._session.execute(select(AgentORM).where(AgentORM.id == agent_id)) + return result.scalar_one_or_none() + + async def list( + self, *, created_by: str | None = None, status: str | None = None + ) -> list[AgentORM]: + """Newest first, filtered by whatever is given.""" + query = select(AgentORM).order_by(AgentORM.created_at.desc(), AgentORM.id.desc()) + if created_by is not None: + query = query.where(AgentORM.created_by == created_by) + if status is not None: + query = query.where(AgentORM.status == status) + result = await self._session.execute(query) + return list(result.scalars().all()) + + async def create_message(self, **values: Any) -> AgentMessageORM: + result = await self._session.execute( + insert(AgentMessageORM).values(**values).returning(AgentMessageORM) + ) + await self._session.flush() + return result.scalar_one() + + async def update_message( + self, message_id: str, *, only_from: tuple[str, ...] | None = None, **fields: Any + ) -> AgentMessageORM | None: + """Change the message; with ``only_from``, only while its status is one of them.""" + statement = update(AgentMessageORM).where(AgentMessageORM.id == message_id) + if only_from is not None: + statement = statement.where(AgentMessageORM.status.in_(only_from)) + result = await self._session.execute(statement.values(**fields).returning(AgentMessageORM)) + await self._session.flush() + return result.scalar_one_or_none() + + async def get_message(self, message_id: str) -> AgentMessageORM | None: + result = await self._session.execute( + select(AgentMessageORM).where(AgentMessageORM.id == message_id) + ) + return result.scalar_one_or_none() + + async def list_messages( + self, + *, + agent_id: str | None = None, + created_by: str | None = None, + parent_session_id: str | None = None, + status: str | None = None, + limit: int = 50, + ) -> list[AgentMessageORM]: + """Newest first, filtered by whatever is given.""" + query = select(AgentMessageORM).order_by( + AgentMessageORM.created_at.desc(), AgentMessageORM.id.desc() + ) + if agent_id is not None: + query = query.where(AgentMessageORM.agent_id == agent_id) + if created_by is not None: + query = query.where(AgentMessageORM.created_by == created_by) + if parent_session_id is not None: + query = query.where(AgentMessageORM.parent_session_id == parent_session_id) + if status is not None: + query = query.where(AgentMessageORM.status == status) + result = await self._session.execute(query.limit(limit)) + return list(result.scalars().all()) + + async def mark_unfinished_messages(self, status: str, error: str) -> list[AgentMessageORM]: + """Finish every queued or running message with ``status`` (after a restart).""" + result = await self._session.execute( + update(AgentMessageORM) + .where(AgentMessageORM.status.in_(("queued", "running"))) + .values(status=status, error=error, finished_at=func.now()) + .returning(AgentMessageORM) + ) + await self._session.flush() + return list(result.scalars().all()) diff --git a/src/assistant_runtime/services/database/repositories/_traces.py b/src/assistant_runtime/services/database/repositories/_traces.py new file mode 100644 index 00000000..2c69760a --- /dev/null +++ b/src/assistant_runtime/services/database/repositories/_traces.py @@ -0,0 +1,65 @@ +"""Traces queries; the caller owns the transaction.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from sqlalchemy import delete, insert, select +from sqlalchemy.ext.asyncio import AsyncSession + +from assistant_runtime.services.database.models import ( + TraceORM, +) + + +class TraceRepository: + """CRUD operations for debug traces. Uses flush() — caller owns commit.""" + + def __init__(self, session: AsyncSession) -> None: + self._session = session + + async def create( + self, + trace_id: str, + session_id: str, + events: list[dict[str, Any]], + *, + user_message: str | None = None, + is_continuation: bool = False, + duration_ms: float | None = None, + ) -> TraceORM: + """Create a new trace row.""" + result = await self._session.execute( + insert(TraceORM) + .values( + id=trace_id, + session_id=session_id, + events=events, + user_message=user_message, + is_continuation=is_continuation, + duration_ms=duration_ms, + ) + .returning(TraceORM) + ) + await self._session.flush() + return result.scalar_one() + + async def delete_older_than(self, cutoff: datetime) -> int: + """Delete traces created before ``cutoff``. Returns the number of rows removed.""" + result = await self._session.execute(delete(TraceORM).where(TraceORM.created_at < cutoff)) + await self._session.flush() + return result.rowcount or 0 + + async def list_by_session( + self, session_id: str, limit: int = 50, offset: int = 0 + ) -> list[TraceORM]: + """List traces for a session, ordered by creation time.""" + result = await self._session.execute( + select(TraceORM) + .where(TraceORM.session_id == session_id) + .order_by(TraceORM.created_at.asc()) + .limit(limit) + .offset(offset) + ) + return list(result.scalars().all()) diff --git a/src/assistant_runtime/services/llm/__init__.py b/src/assistant_runtime/services/llm/__init__.py index 3159b043..6cf94b8c 100644 --- a/src/assistant_runtime/services/llm/__init__.py +++ b/src/assistant_runtime/services/llm/__init__.py @@ -3,12 +3,11 @@ from assistant_runtime.services.llm._settings import validate_model_id from assistant_runtime.services.llm.config import LLMConfig, ProviderConfig from assistant_runtime.services.llm.exceptions import LLMCallError, LLMError, ProviderConfigError -from assistant_runtime.services.llm.interface import LLMResult, LlmService +from assistant_runtime.services.llm.interface import LlmService __all__ = [ "validate_model_id", "LLMConfig", - "LLMResult", "LlmService", "ProviderConfig", "LLMError", diff --git a/src/assistant_runtime/services/llm/interface.py b/src/assistant_runtime/services/llm/interface.py index 94cf59ad..ded58ddd 100644 --- a/src/assistant_runtime/services/llm/interface.py +++ b/src/assistant_runtime/services/llm/interface.py @@ -1,7 +1,6 @@ """LlmService — provider-aware Pydantic AI agent factory. -Public facade for the LLM module. Creates configured agents, executes standalone -LLM calls, and manages provider lifecycle (API key loading and export). +Public facade for the LLM module. Creates configured agents and manages provider lifecycle (API key loading and export). """ from __future__ import annotations @@ -15,14 +14,12 @@ import httpx from loguru import logger from openai import AsyncOpenAI -from pydantic import BaseModel from pydantic_ai import Agent from pydantic_ai.capabilities import AgentCapability from pydantic_ai.models import Model from pydantic_ai.providers.openai import OpenAIProvider from pydantic_ai.tools import Tool, ToolFuncEither -from assistant_runtime.base.resilience import retry_with_backoff from assistant_runtime.model_catalog import ( PROVIDER_DEFAULT_MODELS, PROVIDER_DEFAULT_SUMMARIZATION_MODELS, @@ -34,49 +31,8 @@ from assistant_runtime.services.llm.exceptions import ( ProviderConfigError, ProviderKeyStoreUnavailableError, - classify_llm_error, ) -from assistant_runtime.services.tracing import create_span - -def _collect_retryable_llm_exceptions() -> tuple[type[Exception], ...]: - """Collect retryable exception types from installed LLM provider SDKs. - - Lazily imports from anthropic, openai, and google-genai to avoid hard - dependency on packages that may not be installed. - """ - types: list[type[Exception]] = [ConnectionError, TimeoutError] - - for module_path in ("anthropic", "openai"): - try: - mod = __import__(module_path) - for name in ( - "RateLimitError", - "InternalServerError", - "APIConnectionError", - "APITimeoutError", - ): - cls = getattr(mod, name, None) - if cls is not None: - types.append(cls) - except ImportError: - pass - - # google-genai uses a different module path - try: - from google.genai import errors as google_errors - - for name in ("ClientError", "ServerError"): - cls = getattr(google_errors, name, None) - if cls is not None: - types.append(cls) - except ImportError: - pass - - return tuple(types) - - -_LLM_RETRYABLE_EXCEPTIONS = _collect_retryable_llm_exceptions() # The Codex CLI's backend; the same client id and device-auth flow (services/oauth). _CODEX_BACKEND_BASE_URL = "https://chatgpt.com/backend-api/codex/" @@ -107,13 +63,6 @@ def _cerebras_model(model_name: str) -> Model: ) -class LLMResult(BaseModel): - """Result of a standalone LLM call.""" - - content: str - model: str - - class LlmService: """Provider-aware Pydantic AI agent factory. Implements LifecycleAware.""" @@ -654,83 +603,3 @@ def build_agent( ) return agent - - async def execute_llm_call( - self, - *, - system_prompt: str, - user_prompt: str, - model: str | None = None, - thinking_budget: int | None = None, - temperature: float | None = None, - max_tokens: int | None = None, - ) -> LLMResult: - """Execute a standalone one-shot LLM call with retry on transient failures. - - Creates a one-shot Pydantic AI Agent (no tools, no history) and runs it. - Retries up to 3 times on rate limits, server errors, and connection failures. - - Args: - system_prompt: System instructions for the LLM. - user_prompt: User message to send. - model: Model identifier override. - thinking_budget: Optional thinking token budget. - temperature: Temperature override. - max_tokens: Max tokens override. - - Returns: - LLMResult with content and model identifier. - - Raises: - LLMCallError: If the LLM call fails after all retries. - ProviderConfigError: If the model ID is invalid. - """ - resolved_model = self.resolve_model(model) - - logger.info( - "Executing standalone LLM call", - model=resolved_model, - thinking_budget=thinking_budget, - ) - - @retry_with_backoff( - max_attempts=3, - min_wait=1.0, - max_wait=30.0, - retry_on=_LLM_RETRYABLE_EXCEPTIONS, - name="execute_llm_call", - ) - async def _run_with_retry() -> LLMResult: - agent_model = self._resolve_agent_model(resolved_model) - agent = Agent( - model=agent_model, - instructions=system_prompt, - ) - - model_settings = build_model_settings( - model_id=resolved_model, - thinking_budget=thinking_budget, - temperature=temperature, - max_tokens=max_tokens, - ) - model_settings = self._apply_model_transport_defaults(resolved_model, model_settings) - - result = await agent.run(user_prompt, model_settings=model_settings) - - logger.info("Standalone LLM call completed", model=resolved_model) - - return LLMResult(content=result.output, model=resolved_model) - - with create_span( - "standalone-llm-call", - input_data=user_prompt, - metadata={"model": resolved_model}, - ) as span: - try: - llm_result = await _run_with_retry() - span.update_output(llm_result.content) - return llm_result - except ProviderConfigError: - raise - except Exception as e: - raise classify_llm_error(e) from e diff --git a/src/assistant_runtime/services/oauth/_credentials.py b/src/assistant_runtime/services/oauth/_credentials.py new file mode 100644 index 00000000..c99efe7d --- /dev/null +++ b/src/assistant_runtime/services/oauth/_credentials.py @@ -0,0 +1,120 @@ +"""Read local Codex credentials and decode their metadata without authenticating it.""" + +import base64 +import binascii +import json +from pathlib import Path +from typing import Any + +from assistant_runtime.services.oauth.exceptions import OAuthCodexSyncError +from assistant_runtime.services.oauth.models import CodexCliAuth + + +def read_codex_cli_auth(auth_file: str) -> CodexCliAuth: + """Read ChatGPT/Codex OAuth state from the local Codex CLI auth file.""" + auth_path = Path(auth_file).expanduser() + if not auth_path.exists(): + raise OAuthCodexSyncError(f"Codex auth file not found: {auth_path}") + + try: + data = json.loads(auth_path.read_text()) + except OSError as exc: + raise OAuthCodexSyncError(f"Failed to read Codex auth file: {auth_path}") from exc + except (json.JSONDecodeError, UnicodeDecodeError) as exc: + raise OAuthCodexSyncError(f"Invalid JSON in Codex auth file: {auth_path}") from exc + + if not isinstance(data, dict): + raise OAuthCodexSyncError("Codex auth file must contain a JSON object") + auth_mode = data.get("auth_mode") + if auth_mode == "apikey": + raise OAuthCodexSyncError( + "Codex CLI is using API key mode, not ChatGPT OAuth. Run `codex login` first." + ) + + tokens = data.get("tokens") + if not isinstance(tokens, dict): + raise OAuthCodexSyncError(f"Codex auth file is missing OAuth tokens: {auth_path}") + + access_token = tokens.get("access_token") + refresh_token = tokens.get("refresh_token") + id_token = tokens.get("id_token") + account_id = tokens.get("account_id") + + if not isinstance(access_token, str) or not access_token: + raise OAuthCodexSyncError(f"Codex auth file is missing access_token: {auth_path}") + if not isinstance(refresh_token, str) or not refresh_token: + raise OAuthCodexSyncError(f"Codex auth file is missing refresh_token: {auth_path}") + if id_token is not None and not isinstance(id_token, str): + id_token = None + if account_id is not None and not isinstance(account_id, str): + account_id = None + + return CodexCliAuth( + access_token=access_token, + refresh_token=refresh_token, + id_token=id_token, + account_id=account_id, + auth_mode=auth_mode if isinstance(auth_mode, str) else None, + email=extract_email(id_token), + ) + + +def _decode_jwt_claims(token: str | None) -> dict[str, Any] | None: + """Decode JWT claims without verification for local token metadata access.""" + if not token or token.count(".") < 2: + return None + + payload = token.split(".")[1] + padding = "=" * (-len(payload) % 4) + try: + raw = base64.urlsafe_b64decode(payload + padding) + claims = json.loads(raw) + except (ValueError, TypeError, json.JSONDecodeError, binascii.Error): + return None + + return claims if isinstance(claims, dict) else None + + +def extract_email(id_token: str | None) -> str | None: + """Best-effort decode of email claim from the OIDC id_token payload.""" + claims = _decode_jwt_claims(id_token) + if claims is None: + return None + + email = claims.get("email") + if isinstance(email, str) and email: + return email + + profile = claims.get("https://api.openai.com/profile") + if isinstance(profile, dict): + email = profile.get("email") + if isinstance(email, str) and email: + return email + + return None + + +def extract_account_id(id_token: str | None, access_token: str | None) -> str | None: + """Extract the ChatGPT workspace/account id from token claims.""" + for token in (id_token, access_token): + claims = _decode_jwt_claims(token) + if claims is None: + continue + auth_claim = claims.get("https://api.openai.com/auth") + if isinstance(auth_claim, dict): + account_id = auth_claim.get("chatgpt_account_id") + if isinstance(account_id, str) and account_id: + return account_id + return None + + +def extract_token_expiry(token: str | None) -> float | None: + """Extract expiry from a JWT token, if present.""" + claims = _decode_jwt_claims(token) + if claims is None: + return None + + exp = claims.get("exp") + if isinstance(exp, int | float): + return float(exp) + return None diff --git a/src/assistant_runtime/services/oauth/interface.py b/src/assistant_runtime/services/oauth/interface.py index 997f119d..2994313f 100644 --- a/src/assistant_runtime/services/oauth/interface.py +++ b/src/assistant_runtime/services/oauth/interface.py @@ -9,20 +9,20 @@ from __future__ import annotations import asyncio -import base64 -import binascii import contextlib -import json import time -from enum import StrEnum -from pathlib import Path from typing import Any import httpx from cryptography.fernet import Fernet from loguru import logger -from pydantic import BaseModel +from assistant_runtime.services.oauth._credentials import ( + extract_account_id, + extract_email, + extract_token_expiry, + read_codex_cli_auth, +) from assistant_runtime.services.oauth.config import OPENAI_OAUTH_CLIENT_ID, OAuthConfig from assistant_runtime.services.oauth.exceptions import ( OAuthCodexSyncError, @@ -31,6 +31,12 @@ OAuthRefreshError, OAuthTokenExpiredError, ) +from assistant_runtime.services.oauth.models import AuthSource as AuthSource +from assistant_runtime.services.oauth.models import AuthStatus as AuthStatus +from assistant_runtime.services.oauth.models import CodexCliAuth as CodexCliAuth +from assistant_runtime.services.oauth.models import CodexSession as CodexSession +from assistant_runtime.services.oauth.models import DeviceCodeResponse as DeviceCodeResponse +from assistant_runtime.services.oauth.models import DeviceCodeStatus as DeviceCodeStatus CODEX_CLI_LOGIN_EXPIRED = ( "The Codex CLI login has expired. Run `codex login`; the runtime picks up " @@ -38,66 +44,6 @@ ) -class DeviceCodeStatus(StrEnum): - """State of the device code authorization flow.""" - - IDLE = "idle" - POLLING = "polling" - AUTHORIZED = "authorized" - EXPIRED = "expired" - ERROR = "error" - - -class AuthSource(StrEnum): - """How the current OpenAI auth session was obtained.""" - - DEVICE_CODE = "device_code" - CODEX_CLI = "codex_cli" - DATABASE = "database" - - -class DeviceCodeResponse(BaseModel): - """Response from initiating a device code flow.""" - - user_code: str - verification_uri: str - expires_in: int - - -class AuthStatus(BaseModel): - """Current OAuth connection status.""" - - connected: bool = False - status: DeviceCodeStatus = DeviceCodeStatus.IDLE - source: AuthSource | None = None - email: str | None = None - api_key_preview: str | None = None - expires_at: float | None = None - error: str | None = None - persisted: bool = False - - -class CodexSession(BaseModel): - """Current ChatGPT/Codex auth context used for backend requests.""" - - access_token: str - account_id: str - expires_at: float | None = None - email: str | None = None - source: AuthSource | None = None - - -class CodexCliAuth(BaseModel): - """Relevant OAuth state imported from the local Codex CLI.""" - - access_token: str - refresh_token: str - id_token: str | None = None - account_id: str | None = None - auth_mode: str | None = None - email: str | None = None - - class OAuthService: """Manages OpenAI OAuth tokens — Device Code, Codex CLI sync, refresh.""" @@ -591,54 +537,6 @@ def _reload_codex_cli_auth(self) -> None: except (OAuthCodexSyncError, OSError) as exc: logger.debug("Codex CLI auth file unavailable", error=str(exc)) - def _read_codex_cli_auth(self) -> CodexCliAuth: - """Read ChatGPT/Codex OAuth state from the local Codex CLI auth file.""" - auth_path = Path(self._config.codex_auth_file).expanduser() - if not auth_path.exists(): - raise OAuthCodexSyncError(f"Codex auth file not found: {auth_path}") - - try: - data = json.loads(auth_path.read_text()) - except OSError as exc: - raise OAuthCodexSyncError(f"Failed to read Codex auth file: {auth_path}") from exc - except (json.JSONDecodeError, UnicodeDecodeError) as exc: - raise OAuthCodexSyncError(f"Invalid JSON in Codex auth file: {auth_path}") from exc - - if not isinstance(data, dict): - raise OAuthCodexSyncError("Codex auth file must contain a JSON object") - auth_mode = data.get("auth_mode") - if auth_mode == "apikey": - raise OAuthCodexSyncError( - "Codex CLI is using API key mode, not ChatGPT OAuth. Run `codex login` first." - ) - - tokens = data.get("tokens") - if not isinstance(tokens, dict): - raise OAuthCodexSyncError(f"Codex auth file is missing OAuth tokens: {auth_path}") - - access_token = tokens.get("access_token") - refresh_token = tokens.get("refresh_token") - id_token = tokens.get("id_token") - account_id = tokens.get("account_id") - - if not isinstance(access_token, str) or not access_token: - raise OAuthCodexSyncError(f"Codex auth file is missing access_token: {auth_path}") - if not isinstance(refresh_token, str) or not refresh_token: - raise OAuthCodexSyncError(f"Codex auth file is missing refresh_token: {auth_path}") - if id_token is not None and not isinstance(id_token, str): - id_token = None - if account_id is not None and not isinstance(account_id, str): - account_id = None - - return CodexCliAuth( - access_token=access_token, - refresh_token=refresh_token, - id_token=id_token, - account_id=account_id, - auth_mode=auth_mode if isinstance(auth_mode, str) else None, - email=self._extract_email_from_id_token(id_token), - ) - def _parse_interval_seconds(self, value: object) -> int: """Normalize device-auth polling intervals, which may arrive as strings.""" if isinstance(value, int): @@ -650,67 +548,15 @@ def _parse_interval_seconds(self, value: object) -> int: def _sync_token_metadata(self) -> None: """Refresh derived auth metadata from the current token set.""" - self._email = self._extract_email_from_id_token(self._id_token) or self._email - self._account_id = self._extract_account_id(self._id_token, self._access_token) - expires_at = self._extract_token_expiry(self._access_token) + self._email = extract_email(self._id_token) or self._email + self._account_id = ( + extract_account_id(self._id_token, self._access_token) or self._account_id + ) + expires_at = extract_token_expiry(self._access_token) if expires_at is not None: self._expires_at = expires_at elif self._access_token and self._expires_at <= 0: self._expires_at = time.time() + max(self._config.refresh_buffer_seconds + 300, 3600) - def _decode_jwt_claims(self, token: str | None) -> dict[str, Any] | None: - """Decode JWT claims without verification for local token metadata access.""" - if not token or token.count(".") < 2: - return None - - payload = token.split(".")[1] - padding = "=" * (-len(payload) % 4) - try: - raw = base64.urlsafe_b64decode(payload + padding) - claims = json.loads(raw) - except (ValueError, TypeError, json.JSONDecodeError, binascii.Error): - return None - - return claims if isinstance(claims, dict) else None - - def _extract_email_from_id_token(self, id_token: str | None) -> str | None: - """Best-effort decode of email claim from the OIDC id_token payload.""" - claims = self._decode_jwt_claims(id_token) - if claims is None: - return None - - email = claims.get("email") - if isinstance(email, str) and email: - return email - - profile = claims.get("https://api.openai.com/profile") - if isinstance(profile, dict): - email = profile.get("email") - if isinstance(email, str) and email: - return email - - return None - - def _extract_account_id(self, id_token: str | None, access_token: str | None) -> str | None: - """Extract the ChatGPT workspace/account id from token claims.""" - for token in (id_token, access_token): - claims = self._decode_jwt_claims(token) - if claims is None: - continue - auth_claim = claims.get("https://api.openai.com/auth") - if isinstance(auth_claim, dict): - account_id = auth_claim.get("chatgpt_account_id") - if isinstance(account_id, str) and account_id: - return account_id - return self._account_id - - def _extract_token_expiry(self, token: str | None) -> float | None: - """Extract expiry from a JWT token, if present.""" - claims = self._decode_jwt_claims(token) - if claims is None: - return None - - exp = claims.get("exp") - if isinstance(exp, int | float): - return float(exp) - return None + def _read_codex_cli_auth(self) -> CodexCliAuth: + return read_codex_cli_auth(self._config.codex_auth_file) diff --git a/src/assistant_runtime/services/oauth/models.py b/src/assistant_runtime/services/oauth/models.py new file mode 100644 index 00000000..242d1847 --- /dev/null +++ b/src/assistant_runtime/services/oauth/models.py @@ -0,0 +1,65 @@ +"""OAuth flow state and subscription credential models.""" + +from enum import StrEnum + +from pydantic import BaseModel + + +class DeviceCodeStatus(StrEnum): + """State of the device code authorization flow.""" + + IDLE = "idle" + POLLING = "polling" + AUTHORIZED = "authorized" + EXPIRED = "expired" + ERROR = "error" + + +class AuthSource(StrEnum): + """How the current OpenAI auth session was obtained.""" + + DEVICE_CODE = "device_code" + CODEX_CLI = "codex_cli" + DATABASE = "database" + + +class DeviceCodeResponse(BaseModel): + """Response from initiating a device code flow.""" + + user_code: str + verification_uri: str + expires_in: int + + +class AuthStatus(BaseModel): + """Current OAuth connection status.""" + + connected: bool = False + status: DeviceCodeStatus = DeviceCodeStatus.IDLE + source: AuthSource | None = None + email: str | None = None + api_key_preview: str | None = None + expires_at: float | None = None + error: str | None = None + persisted: bool = False + + +class CodexSession(BaseModel): + """Current ChatGPT/Codex auth context used for backend requests.""" + + access_token: str + account_id: str + expires_at: float | None = None + email: str | None = None + source: AuthSource | None = None + + +class CodexCliAuth(BaseModel): + """Relevant OAuth state imported from the local Codex CLI.""" + + access_token: str + refresh_token: str + id_token: str | None = None + account_id: str | None = None + auth_mode: str | None = None + email: str | None = None diff --git a/src/assistant_runtime/services/tools/builtin/_subagent_executor.py b/src/assistant_runtime/services/tools/builtin/_subagent_executor.py index 7c434119..e835b3bf 100644 --- a/src/assistant_runtime/services/tools/builtin/_subagent_executor.py +++ b/src/assistant_runtime/services/tools/builtin/_subagent_executor.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any from loguru import logger -from pydantic_ai.messages import ToolCallPart, ToolReturnPart +from pydantic_ai.messages import ModelMessage, ModelResponse, ToolCallPart from pydantic_ai.usage import RunUsage, UsageLimits if TYPE_CHECKING: @@ -123,7 +123,7 @@ def _subagent_limits(host: UsageLimits | None, max_iterations: int) -> UsageLimi ) -def _extract_metadata(messages: list, duration: float) -> dict[str, Any]: +def _extract_metadata(messages: list[ModelMessage], duration: float) -> dict[str, Any]: """Extract lightweight metadata from the subagent's message history.""" tool_call_ids: set[str] = set() tools_used: set[str] = set() @@ -132,24 +132,15 @@ def _extract_metadata(messages: list, duration: float) -> dict[str, Any]: output_tokens = 0 for msg in messages: - # Count model responses as iterations - if hasattr(msg, "parts"): - for part in msg.parts: - if isinstance(part, ToolCallPart): - tool_call_ids.add(part.tool_call_id) - tools_used.add(part.tool_name) - elif isinstance(part, ToolReturnPart): - # Already counted via ToolCallPart - pass - - # Count model requests as iterations - if hasattr(msg, "kind") and msg.kind == "response": - iterations += 1 - - # Extract usage if available - if hasattr(msg, "usage") and msg.usage: - input_tokens += getattr(msg.usage, "input_tokens", 0) or 0 - output_tokens += getattr(msg.usage, "output_tokens", 0) or 0 + if not isinstance(msg, ModelResponse): + continue + iterations += 1 + input_tokens += msg.usage.input_tokens + output_tokens += msg.usage.output_tokens + for part in msg.parts: + if isinstance(part, ToolCallPart): + tool_call_ids.add(part.tool_call_id) + tools_used.add(part.tool_name) return { "iterations": iterations, diff --git a/tests/compatibility/test_codex_fast.py b/tests/compatibility/test_codex_fast.py index ea583943..c115f5b6 100644 --- a/tests/compatibility/test_codex_fast.py +++ b/tests/compatibility/test_codex_fast.py @@ -229,16 +229,16 @@ async def test_native_stream_cancellation_closes_transport_and_keeps_actual_unkn async def test_tiers_survive_cumulative_followup_without_double_counting(codex): from pydantic_ai.usage import RunUsage - from assistant_runtime.app.streaming._runner import TurnRunner, _RunState + from assistant_runtime.app.streaming._snapshot import RunSnapshot, capture_result agent = codex.service("fast").build_agent(system_prompt="Decision.", thinking_budget=4000) cumulative = RunUsage() first = await agent.run("First.", usage=cumulative) plan = SimpleNamespace(kind="continuation", accepted_tool_result=None, prior_usage=None) - state = _RunState() - TurnRunner._capture_result(plan, state, first) + state = RunSnapshot() + capture_result(plan, state, first) second = await agent.run("Follow up.", message_history=first.all_messages(), usage=cumulative) - TurnRunner._capture_result(plan, state, second) + capture_result(plan, state, second) assert state.usage["requests"] == 2 assert len(state.usage["service_tiers"]) == 2 assert state.usage["input_tokens"] == 6 diff --git a/tests/compatibility/test_usage.py b/tests/compatibility/test_usage.py index d23e71d3..12812501 100644 --- a/tests/compatibility/test_usage.py +++ b/tests/compatibility/test_usage.py @@ -169,3 +169,26 @@ async def extract(current, recent, turn, *, usage=None): def test_compaction_result_import_keeps_the_summariser_contract(): assert CompactionResult(summary="x").summary == "x" + + +async def test_promoted_steering_accumulates_usage_on_the_existing_row( + isolated_services, script, monkeypatch +): + script.steps = [["First reply"], ["Steered reply"]] + monkeypatch.setattr(LlmService, "_resolve_agent_model", lambda self, name: script.model()) + async with create_runtime(settings=isolated_services) as runtime: + first = await runtime.run_message( + AssistantRequest(id="m1", session_id="s", content="Start") + ) + second = await runtime.run_message( + AssistantRequest( + id="steer", session_id="s", content="Continue", message_type="steering" + ) + ) + path = await runtime._sessions.get_message_path("s") + assert second.message_id == first.message_id + assert len(path) == 2 + assert second.usage["requests"] == 2 + assert second.usage["input_tokens"] > first.usage["input_tokens"] + assert second.usage["output_tokens"] > first.usage["output_tokens"] + assert path[-1]["usage"] == second.usage diff --git a/tests/integration/test_validation.py b/tests/integration/test_validation.py index 3a6fb8db..530055a7 100644 --- a/tests/integration/test_validation.py +++ b/tests/integration/test_validation.py @@ -265,11 +265,10 @@ class TestPublicApiExports: """Verify all modules export their public API correctly.""" def test_llm_exports(self): - from assistant_runtime.services.llm import LLMConfig, LLMResult, LlmService + from assistant_runtime.services.llm import LLMConfig, LlmService assert LlmService is not None assert LLMConfig is not None - assert LLMResult is not None def test_history_exports(self): from assistant_runtime.services.history import ( @@ -319,10 +318,8 @@ def test_base_exports(self): AssistantRuntimeError, LifecycleAware, LifecycleManager, - instrument, ) assert LifecycleManager is not None assert LifecycleAware is not None assert AssistantRuntimeError is not None - assert instrument is not None diff --git a/tests/unit/app/assistant/test_interface.py b/tests/unit/app/assistant/test_interface.py index c07919ba..30807f95 100644 --- a/tests/unit/app/assistant/test_interface.py +++ b/tests/unit/app/assistant/test_interface.py @@ -94,7 +94,6 @@ def tool_service() -> MagicMock: service = MagicMock() service.get_available_tools = MagicMock(return_value=_tool_set()) service.build_toolset = MagicMock(return_value=[]) - service.get_mcp_summary = AsyncMock(return_value=None) return service @@ -212,3 +211,12 @@ async def test_a_profile_prompt_is_its_artifacts_and_includes_with_versions( text, versions = await service.profile_prompt("voice") assert text == "Speak calmly\n\nBe brief" # no time, host context or memory assert versions == {"persona": 1, "owner.preferences": None} + + +async def test_invalid_setup_does_not_start_prompt_loading(service, artifact_service, llm_service): + artifact_service.for_subject.return_value = artifact_service + artifact_service.prompt_inputs = AsyncMock() + llm_service.resolve_model.side_effect = ValueError("Unknown model") + with pytest.raises(ValueError, match="Unknown model"): + await service.prepare_agent_context(_request(message_id="bad"), {}) + artifact_service.prompt_inputs.assert_not_called() diff --git a/tests/unit/app/assistant/test_prompt_builder.py b/tests/unit/app/assistant/test_prompt_builder.py index 759dffec..31e848db 100644 --- a/tests/unit/app/assistant/test_prompt_builder.py +++ b/tests/unit/app/assistant/test_prompt_builder.py @@ -21,7 +21,6 @@ technical_operator_profile, ) from assistant_runtime.services.history.models import WorkingMemory -from assistant_runtime.services.tools.models import ToolCategory, ToolDefinition, ToolSet REQUIRED_ARTIFACTS = { "soul": "The assistant exists to increase the operator's leverage in a live AI workbench.", @@ -36,7 +35,6 @@ class TestProfiles: def test_default_profile_adds_nothing(self): result = _build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=neutral_profile().defaults, ) @@ -53,7 +51,6 @@ def test_profile_order_is_the_prompt_order(self): ), ) result = _build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={"notes": "n", "tone": "Be kind", "policies": "Refund within 30 days"}, profile=profile, @@ -66,7 +63,6 @@ def test_artifact_extras_follow_the_profile_artifacts(self): artifacts=(ArtifactDefinition(name="instructions", required=True, default="i"),) ) result = _build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={"instructions": "Help"}, profile=profile, @@ -82,7 +78,6 @@ def test_optional_artifact_without_text_is_omitted(self): ) ) result = _build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={"instructions": "Help", "notes": " "}, profile=profile, @@ -94,9 +89,7 @@ def test_missing_required_artifact_of_a_custom_profile_raises(self): artifacts=(ArtifactDefinition(name="policies", required=True, default="p"),) ) with pytest.raises(ValueError, match="Missing required artifact: policies"): - _build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={}, profile=profile - ) + _build_system_prompt(session_context={}, artifacts={}, profile=profile) class TestWorkingMemoryFragment: @@ -342,7 +335,6 @@ def test_non_serializable_values_stringified(self): class TestBuildSystemPrompt: def test_contains_soul(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=REQUIRED_ARTIFACTS, ) @@ -350,7 +342,6 @@ def test_contains_soul(self): def test_contains_persona(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=REQUIRED_ARTIFACTS, ) @@ -358,7 +349,6 @@ def test_contains_persona(self): def test_contains_communication_protocol(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=REQUIRED_ARTIFACTS, ) @@ -366,36 +356,13 @@ def test_contains_communication_protocol(self): def test_no_current_time(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=REQUIRED_ARTIFACTS, ) assert "Current time" not in result.content - def test_tools_not_in_prompt_text(self): - """Tools are registered natively with the agent, not duplicated in the system prompt.""" - ts = ToolSet( - backend_tools=[ - ToolDefinition( - name="test_tool", - description="A test", - parameters_schema={}, - category=ToolCategory.BACKEND, - ) - ] - ) - result = build_system_prompt( - available_tools=ts, - session_context={}, - artifacts=REQUIRED_ARTIFACTS, - ) - assert "test_tool" not in result.content - fragment_names = [f["name"] for f in result.fragments] - assert "tools" not in fragment_names - def test_includes_working_memory_when_enabled(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={"working_memory": WorkingMemory(active_goal="Deploy v2")}, working_memory=True, artifacts=REQUIRED_ARTIFACTS, @@ -404,7 +371,6 @@ def test_includes_working_memory_when_enabled(self): def test_stored_working_memory_is_not_rendered_when_disabled(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={"working_memory": WorkingMemory(active_goal="Deploy v2")}, artifacts=REQUIRED_ARTIFACTS, ) @@ -413,7 +379,6 @@ def test_stored_working_memory_is_not_rendered_when_disabled(self): def test_fragments_separated(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=REQUIRED_ARTIFACTS, ) @@ -422,7 +387,6 @@ def test_fragments_separated(self): def test_returns_prompt_result(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=REQUIRED_ARTIFACTS, ) @@ -433,18 +397,7 @@ def test_returns_prompt_result(self): assert len(result.fragments) > 0 def test_fragments_metadata(self): - ts = ToolSet( - backend_tools=[ - ToolDefinition( - name="test_tool", - description="A test", - parameters_schema={}, - category=ToolCategory.BACKEND, - ) - ] - ) result = build_system_prompt( - available_tools=ts, session_context={"working_memory": WorkingMemory(active_goal="Deploy v2")}, working_memory=True, artifacts=REQUIRED_ARTIFACTS, @@ -465,7 +418,6 @@ def test_fragments_metadata(self): def test_working_memory_in_fragments(self): """Working memory content appears as a named fragment in the system prompt.""" result = build_system_prompt( - available_tools=ToolSet(), session_context={"working_memory": WorkingMemory(active_goal="Ship v3")}, working_memory=True, artifacts=REQUIRED_ARTIFACTS, @@ -477,7 +429,6 @@ def test_working_memory_in_fragments(self): def test_minimal_fragments(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=REQUIRED_ARTIFACTS, ) @@ -495,7 +446,6 @@ def test_minimal_fragments(self): def test_ecosystem_included_from_artifact(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=REQUIRED_ARTIFACTS, ) @@ -505,7 +455,6 @@ def test_ecosystem_included_from_artifact(self): def test_required_artifacts_follow_catalog_order(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=REQUIRED_ARTIFACTS, ) @@ -520,7 +469,6 @@ class TestArtifactIntegration: def test_all_artifacts_appear_in_prompt(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={ "soul": "Deeper identity guidance", @@ -545,7 +493,6 @@ def test_all_artifacts_appear_in_prompt(self): def test_scratchpad_appears_when_present(self): artifacts = {**REQUIRED_ARTIFACTS, "scratchpad": "Remember: the user prefers dark mode"} result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=artifacts, ) @@ -556,7 +503,6 @@ def test_scratchpad_appears_when_present(self): def test_empty_scratchpad_excluded(self): artifacts = {**REQUIRED_ARTIFACTS, "scratchpad": ""} result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts=artifacts, ) @@ -565,7 +511,6 @@ def test_empty_scratchpad_excluded(self): def test_communication_protocol_content_appears(self): result = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={ **REQUIRED_ARTIFACTS, @@ -577,7 +522,6 @@ def test_communication_protocol_content_appears(self): def test_missing_persona_raises_error(self): with pytest.raises(ValueError, match="Missing required artifact: persona"): build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={ "soul": "soul", @@ -589,7 +533,6 @@ def test_missing_persona_raises_error(self): def test_missing_communication_protocol_raises_error(self): with pytest.raises(ValueError, match="Missing required artifact: communication_protocol"): build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={ "soul": "soul", @@ -601,7 +544,6 @@ def test_missing_communication_protocol_raises_error(self): def test_missing_soul_raises_error(self): with pytest.raises(ValueError, match="Missing required artifact: soul"): build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={ "persona": "persona", @@ -613,7 +555,6 @@ def test_missing_soul_raises_error(self): def test_missing_ecosystem_raises_error(self): with pytest.raises(ValueError, match="Missing required artifact: ecosystem"): build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={ "soul": "soul", @@ -625,7 +566,6 @@ def test_missing_ecosystem_raises_error(self): def test_empty_persona_raises(self): with pytest.raises(ValueError, match="Missing required artifact: persona"): build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={ "soul": "soul", @@ -638,7 +578,6 @@ def test_empty_persona_raises(self): def test_empty_communication_protocol_raises(self): with pytest.raises(ValueError, match="Missing required artifact: communication_protocol"): build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={ "soul": "soul", @@ -651,7 +590,6 @@ def test_empty_communication_protocol_raises(self): def test_empty_soul_raises(self): with pytest.raises(ValueError, match="Missing required artifact: soul"): build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={ "soul": "", @@ -664,7 +602,6 @@ def test_empty_soul_raises(self): def test_empty_ecosystem_raises(self): with pytest.raises(ValueError, match="Missing required artifact: ecosystem"): build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={ "soul": "soul", diff --git a/tests/unit/app/assistant/test_prompt_record.py b/tests/unit/app/assistant/test_prompt_record.py index 8bd03da6..08ae25c4 100644 --- a/tests/unit/app/assistant/test_prompt_record.py +++ b/tests/unit/app/assistant/test_prompt_record.py @@ -4,7 +4,6 @@ from assistant_runtime.app.assistant.models import PromptResult from assistant_runtime.app.assistant.prompt_record import prompt_record, prompt_text from assistant_runtime.artifacts import ArtifactDefinition, AssistantProfile -from assistant_runtime.services.tools.models import ToolSet PROFILE = AssistantProfile( artifacts=(ArtifactDefinition(name="instructions", required=True, default="i"),) @@ -13,7 +12,6 @@ def _prompt(**kwargs) -> PromptResult: return build_system_prompt( - available_tools=ToolSet(), artifacts={"instructions": "Help the owner"}, profile=PROFILE, artifact_extras=[("owner.preferences", "Short answers")], @@ -67,7 +65,6 @@ def test_an_artifact_named_like_a_dynamic_fragment_keeps_its_place(): ) ) prompt = build_system_prompt( - available_tools=ToolSet(), session_context={}, artifacts={"working_memory": "REMEMBER", "instructions": "HELP"}, profile=profile, @@ -76,3 +73,12 @@ def test_an_artifact_named_like_a_dynamic_fragment_keeps_its_place(): record, snapshot = prompt_record(appended, profile="p", subject=None, artifact_versions={}) assert snapshot == "REMEMBER\n\nHELP" assert prompt_text(record, snapshot) == appended.content + + +async def test_memory_only_snapshots_remain_recoverable_after_the_cache_limit(): + from assistant_runtime.app.assistant import SessionStore + + sessions = SessionStore() + for number in range(257): + await sessions.save_prompt_snapshot(str(number), f"Instructions {number}") + assert await sessions.prompt_snapshot("0") == "Instructions 0" diff --git a/tests/unit/app/event_log/test_interface.py b/tests/unit/app/event_log/test_interface.py index 627f582b..d117b648 100644 --- a/tests/unit/app/event_log/test_interface.py +++ b/tests/unit/app/event_log/test_interface.py @@ -159,3 +159,22 @@ async def test_notices_are_outbound_and_have_no_target(self): await service.record(**_event(direction="outbound", target_session_id="chat")) with pytest.raises(EventNotFoundError): await service.update_notice(99, "heard") + + +async def test_records_do_not_share_payload_or_delivery_state_with_callers(): + service = await _service() + payload = {"details": {"count": 1}} + event, _ = await service.record(**_event(payload=payload, target_session_id="chat")) + payload["details"]["count"] = 2 + event.payload.clear() + event.delivery.clear() + duplicate, created = await service.record(**_event()) + assert not created + assert duplicate.payload == {"details": {"count": 1}} + assert duplicate.delivery == {"session_id": "chat", "how": "queued"} + duplicate.payload.clear() + listed = await service.list() + listed[0].delivery.clear() + reread = await service.get(event.id) + assert reread.payload == {"details": {"count": 1}} + assert reread.delivery == {"session_id": "chat", "how": "queued"} diff --git a/tests/unit/app/ingress/test_factory.py b/tests/unit/app/ingress/test_factory.py index 52fa4edc..df3ba6c4 100644 --- a/tests/unit/app/ingress/test_factory.py +++ b/tests/unit/app/ingress/test_factory.py @@ -12,7 +12,10 @@ async def test_register_ingress_attaches_the_service_to_streaming(): streaming = SimpleNamespace(attach_ingress=MagicMock()) app_state = SimpleNamespace( - assistant_service=MagicMock(), streaming_service=streaming, database_service=None, sio=None + assistant_service=MagicMock(), + streaming_service=streaming, + database_service=None, + ingress_event_sink=AsyncMock(), ) lifecycle = MagicMock() lifecycle.register = AsyncMock() @@ -20,5 +23,6 @@ async def test_register_ingress_attaches_the_service_to_streaming(): await register_ingress(app_state, lifecycle) assert isinstance(app_state.ingress_service, IngressService) + assert app_state.ingress_service._event_sink is app_state.ingress_event_sink streaming.attach_ingress.assert_called_once_with(app_state.ingress_service) lifecycle.register.assert_awaited_once_with("ingress_service", app_state.ingress_service) diff --git a/tests/unit/app/ingress/test_interface.py b/tests/unit/app/ingress/test_interface.py index 37358557..f9771dde 100644 --- a/tests/unit/app/ingress/test_interface.py +++ b/tests/unit/app/ingress/test_interface.py @@ -32,7 +32,9 @@ async def _seeded_store() -> SessionStore: return store -def _service(store: SessionStore | None, *, streaming=None, db=None, sio=None) -> IngressService: +def _service( + store: SessionStore | None, *, streaming=None, db=None, event_sink=None +) -> IngressService: assistant = MagicMock() assistant.get_session_store.return_value = store streaming = streaming or MagicMock() @@ -40,7 +42,7 @@ def _service(store: SessionStore | None, *, streaming=None, db=None, sio=None) - assistant_service=assistant, streaming_service=streaming, database_service=db, - socket_server=sio, + event_sink=event_sink, ) @@ -67,9 +69,8 @@ async def _stream(_request): streaming = MagicMock() streaming.accept_steering = AsyncMock(return_value="promoted") streaming.stream_message = _stream - sio = MagicMock() - sio.emit = AsyncMock() - service = _service(store, streaming=streaming, sio=sio) + sink = AsyncMock() + service = _service(store, streaming=streaming, event_sink=sink) await service.start() result = await service.deliver( @@ -81,8 +82,9 @@ async def _stream(_request): request = streaming.accept_steering.await_args.args[0] assert request.is_steering assert request.content == "[via:tmux from:planner] done" - assert sio.emit.await_count == 2 - assert sio.emit.await_args.kwargs["room"] == "session:sess-1" + assert [call.args for call in sink.await_args_list] == [ + ("sess-1", event) for event in events + ] async def test_queued_into_a_live_turn_runs_nothing(self): store = await _seeded_store() diff --git a/tests/unit/app/routes/test_artifacts.py b/tests/unit/app/routes/test_artifacts.py index 1d3b31ff..00a870f5 100644 --- a/tests/unit/app/routes/test_artifacts.py +++ b/tests/unit/app/routes/test_artifacts.py @@ -502,3 +502,21 @@ async def test_a_text_over_max_chars_is_422_and_the_profile_reports_the_bound(se assert refused.status_code == 422 assert "at most 5 characters" in refused.json()["detail"] assert described["artifacts"][0]["max_chars"] == 5 + + +async def test_database_failure_reaches_the_shared_503_handler(artifacts, monkeypatch): + from unittest.mock import AsyncMock + + from assistant_runtime.config import AppSettings + from assistant_runtime.main import create_app + from assistant_runtime.services.database.exceptions import DatabaseError + + monkeypatch.setattr( + ArtifactService, "update", AsyncMock(side_effect=DatabaseError("Database not reachable")) + ) + app = create_app(settings=AppSettings(_env_file=None)) + app.state.artifact_service = artifacts + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + response = await client.patch("/api/artifacts/scratchpad", json={"content": "New text"}) + assert response.status_code == 503 + assert response.json()["type"] == "DatabaseError" diff --git a/tests/unit/app/streaming/test_client_error_detail.py b/tests/unit/app/streaming/test_client_error_detail.py index 21caaa91..05ba414a 100644 --- a/tests/unit/app/streaming/test_client_error_detail.py +++ b/tests/unit/app/streaming/test_client_error_detail.py @@ -7,7 +7,7 @@ import pytest -from assistant_runtime.app.streaming._runner import _describe_error +from assistant_runtime.app.streaming._telemetry import _describe_error from assistant_runtime.app.streaming.config import StreamingConfig from assistant_runtime.app.streaming.interface import StreamingService from assistant_runtime.services.llm.exceptions import LLMCallError @@ -40,7 +40,7 @@ def test_detail_off_keeps_the_provider_category_only(self): class TestPersistenceFailureText: def test_detail_decides_whether_the_exception_is_shown(self): - from assistant_runtime.app.streaming._runner import _persistence_failure + from assistant_runtime.app.streaming._telemetry import _persistence_failure exc = RuntimeError("connection to db-host refused") assert _persistence_failure("Cancelled turn", exc, detail=True) == ( @@ -54,18 +54,10 @@ def test_detail_decides_whether_the_exception_is_shown(self): class TestTraceRetention: @pytest.mark.asyncio async def test_save_trace_skips_an_unreachable_database(self): - from assistant_runtime.app.streaming._runner import TurnRunner + from assistant_runtime.app.streaming._telemetry import save_trace unreachable = MagicMock(healthy=False) - runner = TurnRunner( - config=StreamingConfig(), - sessions=MagicMock(), - tools=MagicMock(), - history=MagicMock(), - assistant_service=MagicMock(), - database_service=unreachable, - ) - await runner.save_trace("s1", [{"type": "debug"}]) + await save_trace(unreachable, "s1", [{"type": "debug"}]) unreachable.session_context.assert_not_called() @pytest.mark.asyncio diff --git a/tests/unit/app/streaming/test_interface.py b/tests/unit/app/streaming/test_interface.py index f739cefe..fbcc7e64 100644 --- a/tests/unit/app/streaming/test_interface.py +++ b/tests/unit/app/streaming/test_interface.py @@ -555,7 +555,9 @@ async def __aexit__(self, *_args: Any) -> None: ) await service.start() - with patch.object(TurnRunner, "save_trace", new=AsyncMock()) as save_trace: + with patch( + "assistant_runtime.app.streaming._telemetry.save_trace", new=AsyncMock() + ) as save_trace: events = [ event async for event in service.stream_message(_request(message_id="user-1")) ] @@ -575,7 +577,7 @@ async def __aexit__(self, *_args: Any) -> None: trace_args = save_trace.await_args assert trace_args is not None assert trace_args.kwargs["trace_id"] == final["trace_id"] - assert any(event["type"] == "debug_error" for event in trace_args.args[1]) + assert any(event["type"] == "debug_error" for event in trace_args.args[2]) @pytest.mark.asyncio async def test_provider_client_error_persists_debug_error_with_trace_id(self) -> None: @@ -603,7 +605,9 @@ async def __aexit__(self, *_args: Any) -> None: ) await service.start() - with patch.object(TurnRunner, "save_trace", new=AsyncMock()) as save_trace: + with patch( + "assistant_runtime.app.streaming._telemetry.save_trace", new=AsyncMock() + ) as save_trace: events = [ event async for event in service.stream_message(_request(message_id="user-1")) ] @@ -626,7 +630,7 @@ async def __aexit__(self, *_args: Any) -> None: event["type"] == "debug_error" and event["trace_id"] == final["trace_id"] and event["error_type"] == "provider_client_error" - for event in trace_args.args[1] + for event in trace_args.args[2] ) @@ -1392,7 +1396,7 @@ class TestAuxiliaryUsageMerge: @pytest.mark.asyncio async def test_extraction_merges_into_the_rows_current_usage(self) -> None: """A continuation that saved newer usage on the row is not overwritten.""" - from assistant_runtime.app.streaming._runner import TurnRunner, _RunState + from assistant_runtime.app.streaming._snapshot import RunSnapshot service = _make_service() await service.start() @@ -1404,7 +1408,7 @@ async def test_extraction_merges_into_the_rows_current_usage(self) -> None: "sess-1", result.message_id, usage={"input_tokens": 50, "output_tokens": 5} ) runner = service._runner - state = _RunState(stored_usage={"input_tokens": 5, "output_tokens": 1}) + state = RunSnapshot(stored_usage={"input_tokens": 5, "output_tokens": 1}) plan = MagicMock( session_id="sess-1", assistant_message_id=result.message_id, @@ -1444,7 +1448,7 @@ async def extract(session_id, session_context, turn_number, *, leaf_id=None): class TestAuxiliaryKeptAtPersistence: def test_row_sections_the_snapshot_lacks_are_kept(self) -> None: - from assistant_runtime.app.streaming._runner import _keep_row_auxiliary + from assistant_runtime.app.streaming._snapshot import _keep_row_auxiliary row = {"usage": {"input_tokens": 9, "auxiliary": {"working_memory": {"input_tokens": 3}}}} snapshot = {"input_tokens": 12, "auxiliary": {"summarization": {"input_tokens": 1}}} @@ -1461,7 +1465,7 @@ def test_row_sections_the_snapshot_lacks_are_kept(self) -> None: class TestLatestBackgroundUsage: async def test_continuation_keeps_newer_memory_total_without_double_counting(self): - from assistant_runtime.app.streaming._runner import _RunState + from assistant_runtime.app.streaming._snapshot import RunSnapshot, persist_snapshot from assistant_runtime.app.streaming._turn import TurnPlan service = _make_service() @@ -1488,13 +1492,13 @@ async def test_continuation_keeps_newer_memory_total_without_double_counting(sel assistant_parent_id=None, prior_usage=old, ) - continuation = _RunState(usage=plan.prior_usage) + continuation = RunSnapshot(usage=plan.prior_usage) service._assistant_service.update_working_memory = AsyncMock( return_value={"input_tokens": 7} ) - await service._runner._update_working_memory(plan, _RunState(stored_usage=old)) + await service._runner._update_working_memory(plan, RunSnapshot(stored_usage=old)) for _ in range(2): - await service._runner._persist(plan, continuation) + await persist_snapshot(service._sessions, plan, continuation) usage = sessions.get_message("sess-1", "answer")["usage"] assert usage["input_tokens"] == 10 assert usage["auxiliary"]["working_memory"]["input_tokens"] == 10 diff --git a/tests/unit/app/tasks/test_interface.py b/tests/unit/app/tasks/test_interface.py index 1902d1c0..80401ed4 100644 --- a/tests/unit/app/tasks/test_interface.py +++ b/tests/unit/app/tasks/test_interface.py @@ -631,3 +631,29 @@ async def test_the_tools_message_an_agent_for_the_calling_session(self): "done: Status?", "main-1", ) + + +async def test_agent_config_and_completed_usage_cannot_be_changed_through_returned_records(): + from assistant_runtime.app.assistant.config import TunableOverrides + + service, streaming, _ = await _service() + try: + agent = await service.start_agent(config=TunableOverrides(temperature=0.3)) + agent.config["temperature"] = 1.5 + reread = await service.get_agent(agent.id) + assert reread.config == {"temperature": 0.3} + reread.config.clear() + configured = await service.configure_agent(agent.id, TunableOverrides(temperature=0.4)) + configured.config.clear() + listed = await service.list_agents() + listed[0].config.clear() + assert (await service.get_agent(agent.id)).config == {"temperature": 0.4} + task = await service.start_task("One task") + streaming.gate("One task").set() + await _settle() + result = await service.get(task.id) + assert result.status == "done" + result.usage.clear() + assert (await service.get(task.id)).usage == {"requests": 1} + finally: + await service.stop() diff --git a/tests/unit/app/test_socketio_server.py b/tests/unit/app/test_socketio_server.py index c542e42b..b3ba9f77 100644 --- a/tests/unit/app/test_socketio_server.py +++ b/tests/unit/app/test_socketio_server.py @@ -793,3 +793,16 @@ async def test_events_without_a_session_or_unknown_types_stay_in_process(self) - await forward({"type": "artifact_decision", "session_id": None}) await forward({"type": "task_finished", "session_id": "s-1"}) sio.emit.assert_not_awaited() + + +async def test_composed_ingress_sink_uses_the_session_room(): + from assistant_runtime.main import create_asgi_app + + app = create_asgi_app() + sio = app.engineio_server + sio.emit = AsyncMock() + event = {"type": "text_delta", "content": "A background reply"} + await sio.fastapi_app.state.ingress_event_sink("s1", event) + sio.emit.assert_awaited_once_with( + "assistant:text_delta", event, room="session:s1", namespace="/assistant" + ) diff --git a/tests/unit/app/voice/test_codex.py b/tests/unit/app/voice/test_codex.py index 6d2f81ab..0e93d077 100644 --- a/tests/unit/app/voice/test_codex.py +++ b/tests/unit/app/voice/test_codex.py @@ -10,10 +10,10 @@ from assistant_runtime.app.voice._codex import ( CodexTransport, - _AppServer, - missing_from_schema, usage_report, ) +from assistant_runtime.app.voice._codex_rpc import _AppServer +from assistant_runtime.app.voice._codex_schema import missing_from_schema from assistant_runtime.app.voice.config import VoiceConfig from assistant_runtime.app.voice.exceptions import VoiceError from assistant_runtime.app.voice.interface import VoiceService @@ -818,7 +818,7 @@ async def counted(): async def test_a_servers_cleanup_ends_its_own_calls_even_after_it_was_replaced(): from types import SimpleNamespace - from assistant_runtime.app.voice._codex import _Link + from assistant_runtime.app.voice._codex_rpc import _Link server = _AppServer("codex", 1) old_output = asyncio.StreamReader() diff --git a/tests/unit/base/test_instrument.py b/tests/unit/base/test_instrument.py deleted file mode 100644 index 9f267f3b..00000000 --- a/tests/unit/base/test_instrument.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Tests for the @instrument decorator.""" - -import pytest - -from assistant_runtime.base.instrument import instrument - - -async def test_instrument_returns_result(): - @instrument(operation="test_op", module="test") - async def my_func(): - return 42 - - assert await my_func() == 42 - - -async def test_instrument_reraises_exceptions(): - @instrument(operation="failing_op", module="test") - async def failing_func(): - raise ValueError("boom") - - with pytest.raises(ValueError, match="boom"): - await failing_func() - - -async def test_instrument_logs_success(capfd): - @instrument(operation="logged_op", module="test_mod") - async def my_func(): - return "ok" - - await my_func() - # The decorator logs via loguru — we verify it doesn't error. - # Detailed log content testing would require loguru sink capture. - - -async def test_instrument_defaults_to_function_name(): - @instrument(module="test") - async def specific_name(): - return True - - # Verify the wrapper preserves the function name - assert specific_name.__name__ == "specific_name" - assert await specific_name() is True - - -async def test_instrument_preserves_args(): - @instrument(operation="echo", module="test") - async def echo(a, b, key=None): - return (a, b, key) - - result = await echo(1, 2, key="val") - assert result == (1, 2, "val") diff --git a/tests/unit/base/test_lifecycle.py b/tests/unit/base/test_lifecycle.py index 59103b60..74380cdc 100644 --- a/tests/unit/base/test_lifecycle.py +++ b/tests/unit/base/test_lifecycle.py @@ -150,3 +150,30 @@ async def test_health_with_no_components(): lm = LifecycleManager() result = await lm.health() assert result == {"healthy": True, "status": "ok", "components": {}} + + +@pytest.mark.parametrize("cancelled", [False, True]) +async def test_partial_startup_is_rolled_back_in_reverse_order(cancelled): + import asyncio + + manager = LifecycleManager() + stopped = [] + failure = asyncio.CancelledError() if cancelled else RuntimeError("startup failed") + + class Component(FakeComponent): + async def start(self): + self.started = True # resources can exist before startup finishes + if self.name == "partial": + raise failure + + async def stop(self): + stopped.append(self.name) + + for name in ("first", "partial", "unstarted"): + await manager.register(name, Component(name)) + with pytest.raises(type(failure)) as raised: + await manager.start_all() + assert raised.value is failure + assert stopped == ["partial", "first"] + await manager.stop_all() + assert stopped == ["partial", "first"] diff --git a/tests/unit/base/test_logging.py b/tests/unit/base/test_logging.py index 67986a10..0ee1aaf5 100644 --- a/tests/unit/base/test_logging.py +++ b/tests/unit/base/test_logging.py @@ -1,6 +1,6 @@ """Tests for structured logging helpers — tag constants and payload truncation.""" -from assistant_runtime.base._logging import FULL_LOG_THRESHOLD, truncate, truncate_payload +from assistant_runtime.base._logging import FULL_LOG_THRESHOLD, truncate class TestTruncate: @@ -18,14 +18,3 @@ def test_long_text_truncated(self): def test_exact_length_unchanged(self): text = "a" * FULL_LOG_THRESHOLD assert truncate(text) == text - - -class TestTruncatePayload: - def test_converts_non_string_to_string(self): - result = truncate_payload(12345) - assert result == "12345" - - def test_truncates_long_payload(self): - payload = {"key": "v" * 1000} - result = truncate_payload(payload) - assert "chars truncated" in result diff --git a/tests/unit/services/actions/test_interface.py b/tests/unit/services/actions/test_interface.py index 0f599b1b..13997eb2 100644 --- a/tests/unit/services/actions/test_interface.py +++ b/tests/unit/services/actions/test_interface.py @@ -148,3 +148,29 @@ async def test_confirmations_read_in_insertion_order_for_reconciliation(self): assert [r.id for r in await service.confirmations(signed=True, reconciled=False)] == [b.id] _, listed = await service.get(second.id) assert [r.id for r in listed] == [a.id, c.id] + + +async def test_records_do_not_share_mutable_state_with_callers(): + service = await _service() + arguments = {"target": {"name": "original"}} + action = await service.create(kind="message", by="local", arguments=arguments) + arguments["target"]["name"] = "changed" + action.arguments["target"]["name"] = "changed again" + action.history.clear() + stored, _ = await service.get(action.id) + assert stored.arguments == {"target": {"name": "original"}} + assert len(stored.history) == 1 + stored.arguments.clear() + listed = await service.list() + listed[0].history.clear() + confirmed = await service.confirm(action.id, **_confirmation()) + result = {"receipt": {"id": 7}} + settled = await service.update_confirmation(confirmed.id, {"status": "sent", "result": result}) + result["receipt"]["id"] = 8 + settled.result.clear() + confirmations = await service.confirmations() + confirmations[0].result.clear() + reread, confirmations = await service.get(action.id) + assert reread.arguments == {"target": {"name": "original"}} + assert len(reread.history) == 1 + assert confirmations[0].result == {"receipt": {"id": 7}} diff --git a/tests/unit/services/llm/test_codex_only.py b/tests/unit/services/llm/test_codex_only.py index 4c9f7484..4e2dde38 100644 --- a/tests/unit/services/llm/test_codex_only.py +++ b/tests/unit/services/llm/test_codex_only.py @@ -12,9 +12,8 @@ @pytest.mark.parametrize("tier", [None, "fast"]) -@pytest.mark.parametrize("path", ["build_agent", "execute"]) @pytest.mark.parametrize("scenario", ["missing", "disconnected", "expired", "excluded"]) -async def test_rejects_before_agent_or_api_transport(path, scenario, tier, monkeypatch): +async def test_rejects_before_agent_or_api_transport(scenario, tier, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "test-api-key-must-not-be-used") service = LlmService( LLMConfig(codex_only=True, codex_models=["gpt-5.6-sol"], codex_service_tier=tier) @@ -31,14 +30,8 @@ async def test_rejects_before_agent_or_api_transport(path, scenario, tier, monke patch("assistant_runtime.services.llm.interface.Agent") as agent, patch("assistant_runtime.services.llm.interface.AsyncOpenAI") as client, ): - if path == "execute": - with pytest.raises(ProviderConfigError): - await service.execute_llm_call( - model=model, system_prompt="test", user_prompt="test" - ) - else: - with pytest.raises(ProviderConfigError): - service.build_agent(model=model, system_prompt="test") + with pytest.raises(ProviderConfigError): + service.build_agent(model=model, system_prompt="test") agent.assert_not_called() client.assert_not_called() diff --git a/tests/unit/services/llm/test_interface.py b/tests/unit/services/llm/test_interface.py index 0b5aa0dc..92a40476 100644 --- a/tests/unit/services/llm/test_interface.py +++ b/tests/unit/services/llm/test_interface.py @@ -1,4 +1,4 @@ -"""Tests for LlmService lifecycle, build_agent, and execute_llm_call.""" +"""Tests for LlmService lifecycle, agent construction, and provider keys.""" from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -9,8 +9,8 @@ from assistant_runtime.model_catalog import PROVIDER_ENV_VARS from assistant_runtime.services.llm._codex_model import OpenAICodexResponsesModel from assistant_runtime.services.llm.config import LLMConfig -from assistant_runtime.services.llm.exceptions import LLMCallError, ProviderConfigError -from assistant_runtime.services.llm.interface import LLMResult, LlmService +from assistant_runtime.services.llm.exceptions import ProviderConfigError +from assistant_runtime.services.llm.interface import LlmService class TestLlmServiceLifecycle: @@ -284,250 +284,6 @@ def test_build_agent_rejects_invalid_model_format(self, service): ) -class TestExecuteLlmCall: - """LlmService.execute_llm_call tests — all mocked, no live API.""" - - @pytest.fixture - def service(self): - return LlmService(config=LLMConfig()) - - async def test_execute_returns_llm_result(self, service): - mock_result = MagicMock() - mock_result.output = "Hello world" - - with patch("assistant_runtime.services.llm.interface.Agent") as mock_agent_cls: - mock_agent_instance = MagicMock() - mock_agent_instance.run = AsyncMock(return_value=mock_result) - mock_agent_cls.return_value = mock_agent_instance - - result = await service.execute_llm_call( - system_prompt="You are helpful.", - user_prompt="Say hello", - ) - - assert isinstance(result, LLMResult) - assert result.content == "Hello world" - assert result.model == "anthropic:claude-opus-5" - - async def test_execute_with_model_override(self, service): - mock_result = MagicMock() - mock_result.output = "Response" - - with patch("assistant_runtime.services.llm.interface.Agent") as mock_agent_cls: - mock_agent_instance = MagicMock() - mock_agent_instance.run = AsyncMock(return_value=mock_result) - mock_agent_cls.return_value = mock_agent_instance - - result = await service.execute_llm_call( - system_prompt="Test", - user_prompt="Test", - model="openai:gpt-4o", - ) - - assert result.model == "openai:gpt-4o" - - async def test_execute_with_invalid_model_raises(self, service): - with pytest.raises(ProviderConfigError, match="Invalid model ID"): - await service.execute_llm_call( - system_prompt="Test", - user_prompt="Test", - model="bad-model", - ) - - async def test_execute_wraps_exceptions_in_llm_call_error(self, service): - with patch("assistant_runtime.services.llm.interface.Agent") as mock_agent_cls: - mock_agent_instance = MagicMock() - mock_agent_instance.run = AsyncMock(side_effect=RuntimeError("API timeout")) - mock_agent_cls.return_value = mock_agent_instance - - with pytest.raises(LLMCallError, match="LLM call failed"): - await service.execute_llm_call( - system_prompt="Test", - user_prompt="Test", - ) - - async def test_execute_passes_model_settings(self, service): - mock_result = MagicMock() - mock_result.output = "OK" - - with patch("assistant_runtime.services.llm.interface.Agent") as mock_agent_cls: - mock_agent_instance = MagicMock() - mock_agent_instance.run = AsyncMock(return_value=mock_result) - mock_agent_cls.return_value = mock_agent_instance - - await service.execute_llm_call( - system_prompt="Test", - user_prompt="Test", - thinking_budget=5000, - temperature=0.5, - max_tokens=2000, - ) - - # Verify Agent.run was called with model_settings - mock_agent_instance.run.assert_called_once() - call_kwargs = mock_agent_instance.run.call_args - assert "model_settings" in call_kwargs.kwargs - - async def test_execute_uses_codex_model_for_supported_openai_models(self, service): - mock_result = MagicMock() - mock_result.output = "OK" - service.set_oauth_service( - SimpleNamespace( - get_codex_session=lambda: SimpleNamespace( - access_token="access-token", - account_id="acct_123", - ) - ) - ) - - with patch("assistant_runtime.services.llm.interface.Agent") as mock_agent_cls: - mock_agent_instance = MagicMock() - mock_agent_instance.run = AsyncMock(return_value=mock_result) - mock_agent_cls.return_value = mock_agent_instance - - await service.execute_llm_call( - system_prompt="Test", - user_prompt="Test", - model="openai:gpt-5.4", - ) - - model_arg = mock_agent_cls.call_args.kwargs["model"] - assert isinstance(model_arg, OpenAICodexResponsesModel) - - async def test_execute_passes_openai_reasoning_settings_to_codex_run(self, service): - mock_result = MagicMock() - mock_result.output = "OK" - service.set_oauth_service( - SimpleNamespace( - get_codex_session=lambda: SimpleNamespace( - access_token="access-token", - account_id="acct_123", - ) - ) - ) - - with patch("assistant_runtime.services.llm.interface.Agent") as mock_agent_cls: - mock_agent_instance = MagicMock() - mock_agent_instance.run = AsyncMock(return_value=mock_result) - mock_agent_cls.return_value = mock_agent_instance - - await service.execute_llm_call( - system_prompt="Test", - user_prompt="Explain briefly.", - model="openai:gpt-5.4", - thinking_budget=5_000, - ) - - call_kwargs = mock_agent_instance.run.call_args.kwargs - assert call_kwargs["model_settings"]["openai_store"] is False - assert call_kwargs["model_settings"]["openai_send_reasoning_ids"] is False - assert "openai_previous_response_id" not in call_kwargs["model_settings"] - assert call_kwargs["model_settings"]["openai_reasoning_effort"] == "medium" - assert call_kwargs["model_settings"]["openai_reasoning_summary"] == "detailed" - - -class TestExecuteLlmCallRetry: - """Tests for retry behavior in execute_llm_call.""" - - @pytest.fixture - def service(self): - return LlmService(config=LLMConfig()) - - async def test_retries_on_connection_error_then_succeeds(self, service): - """execute_llm_call retries on ConnectionError and succeeds.""" - call_count = 0 - mock_result = MagicMock() - mock_result.output = "recovered" - - with patch("assistant_runtime.services.llm.interface.Agent") as mock_agent_cls: - mock_agent_instance = MagicMock() - - async def _run_side_effect(*args, **kwargs): - nonlocal call_count - call_count += 1 - if call_count < 2: - raise ConnectionError("transient") - return mock_result - - mock_agent_instance.run = _run_side_effect - mock_agent_cls.return_value = mock_agent_instance - - # Patch wait to avoid real sleeps - with patch( - "assistant_runtime.base.resilience.wait_random_exponential", - return_value=MagicMock(return_value=0), - ): - result = await service.execute_llm_call( - system_prompt="Test", - user_prompt="Test", - ) - - assert result.content == "recovered" - assert call_count == 2 - - async def test_gives_up_after_max_retries(self, service): - """execute_llm_call raises LLMCallError after exhausting retries.""" - with patch("assistant_runtime.services.llm.interface.Agent") as mock_agent_cls: - mock_agent_instance = MagicMock() - mock_agent_instance.run = AsyncMock(side_effect=ConnectionError("permanent")) - mock_agent_cls.return_value = mock_agent_instance - - with pytest.raises(LLMCallError) as exc_info: - await service.execute_llm_call( - system_prompt="Test", - user_prompt="Test", - ) - - assert exc_info.value.error_category == "CONNECTION_ERROR" - assert exc_info.value.is_retryable is True - - async def test_does_not_retry_on_non_retryable_error(self, service): - """Non-retryable errors (like RuntimeError) are not retried.""" - call_count = 0 - - with patch("assistant_runtime.services.llm.interface.Agent") as mock_agent_cls: - mock_agent_instance = MagicMock() - - async def _run_side_effect(*args, **kwargs): - nonlocal call_count - call_count += 1 - raise RuntimeError("not retryable") - - mock_agent_instance.run = _run_side_effect - mock_agent_cls.return_value = mock_agent_instance - - with pytest.raises(LLMCallError) as exc_info: - await service.execute_llm_call( - system_prompt="Test", - user_prompt="Test", - ) - - assert exc_info.value.error_category == "UNKNOWN" - assert exc_info.value.is_retryable is False - - # Should be called only once — no retries - assert call_count == 1 - - async def test_error_classified_via_classify_llm_error(self, service): - """Errors are classified with error_category via classify_llm_error.""" - - class RateLimitError(Exception): - pass - - with patch("assistant_runtime.services.llm.interface.Agent") as mock_agent_cls: - mock_agent_instance = MagicMock() - mock_agent_instance.run = AsyncMock(side_effect=RateLimitError("slow down")) - mock_agent_cls.return_value = mock_agent_instance - - with pytest.raises(LLMCallError) as exc_info: - await service.execute_llm_call( - system_prompt="Test", - user_prompt="Test", - ) - - assert exc_info.value.error_category == "RATE_LIMIT" - - class TestDbProviderKeys: """Tests for database-sourced provider API keys — load, reload, remove, status.""" diff --git a/tests/unit/services/tools/builtin/test_subagent_executor.py b/tests/unit/services/tools/builtin/test_subagent_executor.py index 40c96e43..74d347ed 100644 --- a/tests/unit/services/tools/builtin/test_subagent_executor.py +++ b/tests/unit/services/tools/builtin/test_subagent_executor.py @@ -5,7 +5,8 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from pydantic_ai.usage import RunUsage, UsageLimits +from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, ToolCallPart, ToolReturnPart +from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits from assistant_runtime.services.tools.builtin._subagent_executor import execute_subagent from assistant_runtime.services.tools.builtin.subagent import SubagentDefinition @@ -32,7 +33,14 @@ def mock_agent_result() -> MagicMock: """A mock agent run result with sensible defaults.""" result = MagicMock() result.output = "Research findings here" - result.all_messages.return_value = [] + result.all_messages.return_value = [ + ModelResponse( + parts=[ToolCallPart("lookup", {}, "call-1")], + usage=RequestUsage(input_tokens=11, output_tokens=13), + ), + ModelRequest(parts=[ToolReturnPart("lookup", "Found", "call-1")]), + ModelResponse(parts=[TextPart("Research findings here")]), + ] return result @@ -97,6 +105,10 @@ async def test_success_includes_metadata(self, subagent_definition, mock_llm_ser "duration_seconds", } assert set(metadata.keys()) == expected_keys + assert metadata["iterations"] == 2 + assert metadata["tool_calls_count"] == 1 + assert metadata["tools_used"] == ["lookup"] + assert (metadata["input_tokens"], metadata["output_tokens"]) == (11, 13) async def test_exception_returns_error_dict( self, subagent_definition, mock_llm_service, mock_agent diff --git a/tests/unit/test_imports.py b/tests/unit/test_imports.py index bddc1183..4154144b 100644 --- a/tests/unit/test_imports.py +++ b/tests/unit/test_imports.py @@ -89,7 +89,7 @@ def _all_offenders(forbidden: tuple[str, ...], allowed_top_level: tuple[str, ... def _module_of(parts: tuple[str, ...]) -> str: """The module a file or import path belongs to: ``app/``, ``services/``, or the top dir.""" if parts[0] in ("app", "services") and len(parts) > 1: - return f"{parts[0]}/{parts[1]}" + return f"{parts[0]}/{parts[1].removesuffix('.py')}" return parts[0] @@ -113,6 +113,50 @@ def _private_cross_module_imports() -> list[str]: return result +# These are module dependencies, including type-only imports. Settings compose +# the assistant config schema; all transport adapters stay above application services. +_APP_DEPENDENCIES = { + "access": set(), + "assistant": {"settings"}, + "settings": {"assistant"}, + "streaming": {"access", "assistant"}, + "voice": {"access", "assistant", "streaming"}, + "ingress": {"assistant", "streaming"}, + "tasks": {"assistant", "settings", "streaming"}, + "heartbeat": {"ingress"}, + "event_log": {"ingress"}, + "routes": { + "access", + "assistant", + "settings", + "streaming", + "voice", + "ingress", + "tasks", + "event_log", + }, + "socketio_server": {"access", "assistant", "streaming"}, +} + + +def _app_layer_violations() -> list[str]: + violations = [] + prefix = f"{PACKAGE}.app." + for path in sorted((SRC / "app").rglob("*.py")): + owner = path.relative_to(SRC / "app").parts[0].removesuffix(".py") + for imported in _imported_modules(path): + if not imported.startswith(prefix): + continue + target = imported.removeprefix(prefix).split(".")[0] + if target != owner and target not in _APP_DEPENDENCIES.get(owner, set()): + violations.append(f"{path.relative_to(SRC)}: {imported}") + return sorted(violations) + + +def test_application_dependencies_follow_the_module_graph() -> None: + assert _app_layer_violations() == [] + + def test_private_modules_stay_inside_their_module() -> None: """A ``_``-prefixed file is imported only by files of the same module (AGENTS.md).""" assert _private_cross_module_imports() == [] @@ -130,7 +174,9 @@ def test_base_does_not_import_app_or_services() -> None: assert _offenders("base", (f"{PACKAGE}.app", f"{PACKAGE}.services")) == [] -@pytest.mark.parametrize("leaf", ["artifacts.py", "model_catalog.py"]) +@pytest.mark.parametrize( + "leaf", ["artifacts.py", "model_catalog.py", "host_context.py", "principal.py"] +) def test_leaf_modules_import_nothing_from_the_package(leaf: str) -> None: project_imports = {m for m in _imported_modules(SRC / leaf) if m.startswith(PACKAGE)} assert project_imports == set() @@ -166,3 +212,14 @@ def test_plain_import_resolves(self, check) -> None: def test_init_relative_import_resolves(self, check) -> None: found = check("services/tools/__init__.py", "from ..llm import z\n") assert f"{PACKAGE}.services.llm" in found + + def test_transport_dependencies_are_rejected_inside_services(self, check) -> None: + check( + "app/ingress/interface.py", + f"from {PACKAGE}.app.socketio_server import socket_event_name\n", + ) + assert _app_layer_violations() + + def test_lower_application_dependencies_are_allowed(self, check) -> None: + check("app/ingress/interface.py", f"from {PACKAGE}.app.streaming import StreamingService\n") + assert _app_layer_violations() == [] diff --git a/tests/unit/test_main_lifespan.py b/tests/unit/test_main_lifespan.py index 927a97aa..68a6bb18 100644 --- a/tests/unit/test_main_lifespan.py +++ b/tests/unit/test_main_lifespan.py @@ -152,5 +152,5 @@ async def open_runtime(): assert started == ( ["first", "second"] if stage == "component-start" else ["first", "second", "third"] ) - assert stopped == list(reversed(started)) + assert stopped == ["third", "second", "first"] shutdown.assert_called_once_with() From ada6f7a185d82c13b69f30cb614ea89e652abccf Mon Sep 17 00:00:00 2001 From: eandualem Date: Fri, 2 Oct 2026 11:38:55 +0300 Subject: [PATCH 2/5] fix: address cleanup review findings Keep memory-only prompt snapshots with their session, drain lifecycle shutdown despite repeated cancellation, and let shared database error handling serve every artifact route. Restrict the settings import exception to the assistant configuration schema. --- .../app/assistant/_session_store.py | 24 +++--- .../app/assistant/interface.py | 2 +- src/assistant_runtime/app/routes/artifacts.py | 78 +------------------ .../app/streaming/_snapshot.py | 2 +- src/assistant_runtime/base/lifecycle.py | 14 +++- tests/compatibility/test_prompt_record.py | 2 +- .../unit/app/assistant/test_prompt_record.py | 7 +- tests/unit/app/routes/test_artifacts.py | 14 +++- tests/unit/base/test_lifecycle.py | 36 +++++++++ tests/unit/test_imports.py | 12 ++- 10 files changed, 92 insertions(+), 99 deletions(-) diff --git a/src/assistant_runtime/app/assistant/_session_store.py b/src/assistant_runtime/app/assistant/_session_store.py index 7dd579e6..faf185d4 100644 --- a/src/assistant_runtime/app/assistant/_session_store.py +++ b/src/assistant_runtime/app/assistant/_session_store.py @@ -45,7 +45,7 @@ _MAX_MEMORY_SESSIONS = 200 -# Stable prompt texts kept in process; without a database they are the only copy. +# Reloadable database prompt cache; memory-only texts live with their session. _SNAPSHOT_CACHE_SIZE = 256 @@ -413,22 +413,24 @@ async def update_message( ] return record - async def save_prompt_snapshot(self, snapshot_hash: str, content: str) -> None: - """Keep the stable text a prompt record names (stored once per hash).""" + async def save_prompt_snapshot(self, session_id: str, snapshot_hash: str, content: str) -> None: + """Keep stable prompt text for as long as its session can reference it.""" + if self._db is None: + self.get_context(session_id).setdefault("prompt_snapshots", {})[snapshot_hash] = content + return self._snapshots[snapshot_hash] = content self._snapshots.move_to_end(snapshot_hash) - while self._db is not None and len(self._snapshots) > _SNAPSHOT_CACHE_SIZE: + while len(self._snapshots) > _SNAPSHOT_CACHE_SIZE: self._snapshots.popitem(last=False) - if self._db is not None: - await self._db.save_prompt_snapshot(snapshot_hash, content) + await self._db.save_prompt_snapshot(snapshot_hash, content) - async def prompt_snapshot(self, snapshot_hash: str) -> str | None: - """The stable text under ``snapshot_hash``, from the cache or the database.""" + async def prompt_snapshot(self, session_id: str, snapshot_hash: str) -> str | None: + """The stable text from its session, or the reloadable database cache.""" + if self._db is None: + return self._sessions.get(session_id, {}).get("prompt_snapshots", {}).get(snapshot_hash) if snapshot_hash in self._snapshots: return self._snapshots[snapshot_hash] - if self._db is not None: - return await self._db.prompt_snapshot(snapshot_hash) - return None + return await self._db.prompt_snapshot(snapshot_hash) def get_message(self, session_id: str, message_id: str) -> MessageRecord | None: """The cached record of one message, or None when the session or message is not cached.""" diff --git a/src/assistant_runtime/app/assistant/interface.py b/src/assistant_runtime/app/assistant/interface.py index cfda25cb..cc3226c6 100644 --- a/src/assistant_runtime/app/assistant/interface.py +++ b/src/assistant_runtime/app/assistant/interface.py @@ -162,7 +162,7 @@ async def message_prompt(self, session_id: str, message_id: str) -> dict[str, An prompt = (record or {}).get("prompt") if not prompt: return None - snapshot = await self._sessions.prompt_snapshot(prompt["snapshot_hash"]) + snapshot = await self._sessions.prompt_snapshot(session_id, prompt["snapshot_hash"]) return { **prompt, "message_id": message_id, diff --git a/src/assistant_runtime/app/routes/artifacts.py b/src/assistant_runtime/app/routes/artifacts.py index 63dcbba7..1b9a973e 100644 --- a/src/assistant_runtime/app/routes/artifacts.py +++ b/src/assistant_runtime/app/routes/artifacts.py @@ -12,7 +12,6 @@ from typing import Annotated, Any from fastapi import APIRouter, Depends, HTTPException, Query, Request -from loguru import logger from pydantic import BaseModel, Field from assistant_runtime.app.access.deps import AdminDep, PrincipalDep @@ -29,7 +28,6 @@ ) from assistant_runtime.services.artifacts.interface import ArtifactService from assistant_runtime.services.artifacts.models import Actor, MutationResult -from assistant_runtime.services.database.exceptions import DatabaseError router = APIRouter(prefix="/artifacts", tags=["artifacts"]) proposals_router = APIRouter(prefix="/artifact-proposals", tags=["artifacts"]) @@ -211,21 +209,13 @@ async def _reject( @router.get("") async def list_artifacts(artifacts: ScopedArtifactDep, principal: PrincipalDep) -> list[dict]: """Active versions of the profile's artifacts, in prompt order.""" - try: - return [row.to_dict() for row in await artifacts.list_active()] - except Exception as e: - logger.error("Failed to list artifacts", error=str(e)) - raise HTTPException(status_code=503, detail="Artifact store unavailable") from e + return [row.to_dict() for row in await artifacts.list_active()] @router.get("/profile") async def get_profile(artifacts: ScopedArtifactDep, principal: PrincipalDep) -> dict: """The profile: every artifact, its role, policy and whether a version is active.""" - try: - active = {row.name: row.version for row in await artifacts.list_active()} - except Exception as e: - logger.error("Failed to read artifacts", error=str(e)) - raise HTTPException(status_code=503, detail="Artifact store unavailable") from e + active = {row.name: row.version for row in await artifacts.list_active()} profile = artifacts.profile return { "name": profile.name, @@ -273,13 +263,8 @@ async def get_artifact(name: str, artifacts: ScopedArtifactDep, principal: Princ """The active version of an artifact, or its default text when none is active.""" try: row = await artifacts.get_active(name) - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to get artifact", name=name, error=str(e)) - raise HTTPException(status_code=503, detail="Artifact store unavailable") from e if row is None: definition = artifacts.profile.get(name) return { @@ -305,13 +290,8 @@ async def get_artifact_history( """Version history for an artifact, newest first.""" try: return [row.to_dict() for row in await artifacts.history(name, limit=limit)] - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to get artifact history", name=name, error=str(e)) - raise HTTPException(status_code=503, detail="Artifact store unavailable") from e @router.get("/{name}/versions/{version}") @@ -321,13 +301,8 @@ async def get_artifact_version( """One version as a proposal record: its content, the active text and a unified diff.""" try: return await artifacts.version_record(name, version) - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to get artifact version", name=name, version=version, error=str(e)) - raise HTTPException(status_code=503, detail="Artifact store unavailable") from e @router.post("/{name}/propose", status_code=201) @@ -340,13 +315,8 @@ async def propose_artifact( return await _propose( artifacts, name, body.content, who, body.expected_version, body.rationale ) - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to propose artifact", name=name, error=str(e)) - raise HTTPException(status_code=500, detail="Failed to propose artifact") from e @router.patch("/{name}") @@ -357,13 +327,8 @@ async def update_artifact( who = _who(admin, body) try: return await _update(artifacts, name, body.content, who, body.expected_version) - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to update artifact", name=name, error=str(e)) - raise HTTPException(status_code=500, detail="Failed to update artifact") from e @router.post("/{name}/approve/{version}") @@ -378,13 +343,8 @@ async def approve_artifact( who = _who(admin, body) try: return await _activate(artifacts, name, version, who, rollback=False) - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to approve artifact", name=name, version=version, error=str(e)) - raise HTTPException(status_code=500, detail="Failed to approve artifact") from e @router.post("/{name}/reject/{version}") @@ -399,13 +359,8 @@ async def reject_artifact( who = _who(admin, body) try: return await _reject(artifacts, name, version, who, body.reason if body else None) - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to reject artifact", name=name, version=version, error=str(e)) - raise HTTPException(status_code=500, detail="Failed to reject artifact") from e @router.post("/{name}/rollback/{version}") @@ -415,13 +370,8 @@ async def rollback_artifact( """Reactivate an earlier version of an artifact.""" try: return await _activate(artifacts, name, version, admin.id, rollback=True) - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to rollback artifact", name=name, version=version, error=str(e)) - raise HTTPException(status_code=500, detail="Failed to rollback artifact") from e @router.post("/{name}/actions") @@ -457,15 +407,8 @@ async def artifact_action( return await _update( artifacts, name, body.content, _who(admin, body), body.expected_version ) - except HTTPException: - raise - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Artifact action failed", name=name, action=body.action, error=str(e)) - raise HTTPException(status_code=500, detail=f"Artifact action failed: {body.action}") from e @router.delete("/{name}") @@ -473,13 +416,8 @@ async def delete_artifact(name: str, artifacts: ScopedArtifactDep, admin: AdminD """Delete every stored version; the profile's default text applies again.""" try: count = await artifacts.delete(name, actor=Actor("host", admin.id)) - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to delete artifact", name=name, error=str(e)) - raise HTTPException(status_code=500, detail="Failed to delete artifact") from e if count == 0: raise HTTPException(status_code=404, detail=f"No stored versions for artifact: {name}") return {"success": True, "name": name, "deleted_versions": count, "durable": artifacts.durable} @@ -508,13 +446,8 @@ async def list_proposals( records: list[dict] = [] for name in names: records.extend(await service.for_profile(name).proposals(status, limit + 1, before_id)) - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to list proposals", error=str(e)) - raise HTTPException(status_code=503, detail="Artifact store unavailable") from e records.sort(key=lambda record: record["id"] or 0, reverse=True) page = records[:limit] more = len(records) > limit @@ -531,13 +464,8 @@ async def list_subjects( try: artifacts = get_artifact_service(request).for_profile(profile) return {"profile": artifacts.profile.name, "subjects": await artifacts.subjects()} - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e - except Exception as e: - logger.error("Failed to list subjects", error=str(e)) - raise HTTPException(status_code=503, detail="Artifact store unavailable") from e @prompt_router.get("") @@ -554,8 +482,6 @@ async def preview_prompt( """ try: return await request.app.state.assistant_service.preview_prompt(profile, subject) - except DatabaseError: - raise except ArtifactError as e: raise _http_error(e) from e except ValueError as e: # a required artifact without text diff --git a/src/assistant_runtime/app/streaming/_snapshot.py b/src/assistant_runtime/app/streaming/_snapshot.py index 40d63d06..9cad65e6 100644 --- a/src/assistant_runtime/app/streaming/_snapshot.py +++ b/src/assistant_runtime/app/streaming/_snapshot.py @@ -183,7 +183,7 @@ async def persist_snapshot(sessions: SessionStore, plan: TurnPlan, state: RunSna # finished after this turn took its snapshot) are kept, not erased. if state.prompt_record is not None and state.prompt_snapshot is not None: await sessions.save_prompt_snapshot( - state.prompt_record["snapshot_hash"], state.prompt_snapshot + plan.session_id, state.prompt_record["snapshot_hash"], state.prompt_snapshot ) model_messages = dump_model_messages(state.assistant_messages) if plan.assistant_parent_id is not None and not state.persisted: diff --git a/src/assistant_runtime/base/lifecycle.py b/src/assistant_runtime/base/lifecycle.py index dbe159d0..bb878c8b 100644 --- a/src/assistant_runtime/base/lifecycle.py +++ b/src/assistant_runtime/base/lifecycle.py @@ -30,12 +30,24 @@ async def start_all(self) -> None: await component.start() logger.info("Component started", component=name) except (Exception, asyncio.CancelledError) as e: - logger.error("Component failed to start", component=name, error=str(e)) + logger.error("Component failed to start", component=name, error=repr(e)) await self.stop_all() raise async def stop_all(self) -> None: """Stop all started components in reverse order.""" + cleanup = asyncio.create_task(self._stop_all()) + cancelled = None + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError as exc: + cancelled = exc + cleanup.result() + if cancelled is not None: + raise cancelled + + async def _stop_all(self) -> None: for name in reversed(self._started): try: await self._components[name].stop() diff --git a/tests/compatibility/test_prompt_record.py b/tests/compatibility/test_prompt_record.py index 9c9105ca..29d9acbd 100644 --- a/tests/compatibility/test_prompt_record.py +++ b/tests/compatibility/test_prompt_record.py @@ -29,7 +29,7 @@ async def capture(messages, info): assistant = (await runtime.sessions.get_message_path("compat"))[-1] record = assistant["prompt"] - snapshot = await runtime.sessions.prompt_snapshot(record["snapshot_hash"]) + snapshot = await runtime.sessions.prompt_snapshot("compat", record["snapshot_hash"]) assert received[-1] == "Answer briefly." assert prompt_text(record, snapshot) == received[-1] assert record["profile"] == "neutral" diff --git a/tests/unit/app/assistant/test_prompt_record.py b/tests/unit/app/assistant/test_prompt_record.py index 08ae25c4..d70b8f3c 100644 --- a/tests/unit/app/assistant/test_prompt_record.py +++ b/tests/unit/app/assistant/test_prompt_record.py @@ -80,5 +80,8 @@ async def test_memory_only_snapshots_remain_recoverable_after_the_cache_limit(): sessions = SessionStore() for number in range(257): - await sessions.save_prompt_snapshot(str(number), f"Instructions {number}") - assert await sessions.prompt_snapshot("0") == "Instructions 0" + await sessions.save_prompt_snapshot("session", str(number), f"Instructions {number}") + assert await sessions.prompt_snapshot("session", "0") == "Instructions 0" + + await sessions.delete_session("session") + assert await sessions.prompt_snapshot("session", "0") is None diff --git a/tests/unit/app/routes/test_artifacts.py b/tests/unit/app/routes/test_artifacts.py index 00a870f5..212d4c1c 100644 --- a/tests/unit/app/routes/test_artifacts.py +++ b/tests/unit/app/routes/test_artifacts.py @@ -504,7 +504,10 @@ async def test_a_text_over_max_chars_is_422_and_the_profile_reports_the_bound(se assert described["artifacts"][0]["max_chars"] == 5 -async def test_database_failure_reaches_the_shared_503_handler(artifacts, monkeypatch): +@pytest.mark.parametrize( + "path", ["/api/artifacts", "/api/artifacts/profile", "/api/artifacts/scratchpad"] +) +async def test_database_failure_reaches_the_shared_503_handler(artifacts, monkeypatch, path): from unittest.mock import AsyncMock from assistant_runtime.config import AppSettings @@ -512,11 +515,16 @@ async def test_database_failure_reaches_the_shared_503_handler(artifacts, monkey from assistant_runtime.services.database.exceptions import DatabaseError monkeypatch.setattr( - ArtifactService, "update", AsyncMock(side_effect=DatabaseError("Database not reachable")) + ArtifactService, + "update" if path.endswith("scratchpad") else "list_active", + AsyncMock(side_effect=DatabaseError("Database not reachable")), ) app = create_app(settings=AppSettings(_env_file=None)) app.state.artifact_service = artifacts async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: - response = await client.patch("/api/artifacts/scratchpad", json={"content": "New text"}) + if path.endswith("scratchpad"): + response = await client.patch(path, json={"content": "New text"}) + else: + response = await client.get(path) assert response.status_code == 503 assert response.json()["type"] == "DatabaseError" diff --git a/tests/unit/base/test_lifecycle.py b/tests/unit/base/test_lifecycle.py index 74380cdc..d8adbd59 100644 --- a/tests/unit/base/test_lifecycle.py +++ b/tests/unit/base/test_lifecycle.py @@ -177,3 +177,39 @@ async def stop(self): assert stopped == ["partial", "first"] await manager.stop_all() assert stopped == ["partial", "first"] + + +async def test_repeated_cancellation_does_not_interrupt_startup_rollback(): + import asyncio + + manager = LifecycleManager() + starting, stopping, release = asyncio.Event(), asyncio.Event(), asyncio.Event() + stopped = [] + + class Component(FakeComponent): + async def start(self): + if self.name == "partial": + starting.set() + await asyncio.Event().wait() + + async def stop(self): + if self.name == "partial": + stopping.set() + await release.wait() + stopped.append(self.name) + + for name in ("first", "partial"): + await manager.register(name, Component(name)) + task = asyncio.create_task(manager.start_all()) + await starting.wait() + task.cancel() + await stopping.wait() + task.cancel() + await asyncio.sleep(0) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert stopped == ["partial", "first"] + await manager.stop_all() + assert stopped == ["partial", "first"] diff --git a/tests/unit/test_imports.py b/tests/unit/test_imports.py index 4154144b..864f9b4a 100644 --- a/tests/unit/test_imports.py +++ b/tests/unit/test_imports.py @@ -118,7 +118,7 @@ def _private_cross_module_imports() -> list[str]: _APP_DEPENDENCIES = { "access": set(), "assistant": {"settings"}, - "settings": {"assistant"}, + "settings": {"assistant.config"}, "streaming": {"access", "assistant"}, "voice": {"access", "assistant", "streaming"}, "ingress": {"assistant", "streaming"}, @@ -147,8 +147,9 @@ def _app_layer_violations() -> list[str]: for imported in _imported_modules(path): if not imported.startswith(prefix): continue - target = imported.removeprefix(prefix).split(".")[0] - if target != owner and target not in _APP_DEPENDENCIES.get(owner, set()): + target = imported.removeprefix(prefix) + allowed = {owner, *_APP_DEPENDENCIES.get(owner, set())} + if not any(target == edge or target.startswith(edge + ".") for edge in allowed): violations.append(f"{path.relative_to(SRC)}: {imported}") return sorted(violations) @@ -223,3 +224,8 @@ def test_transport_dependencies_are_rejected_inside_services(self, check) -> Non def test_lower_application_dependencies_are_allowed(self, check) -> None: check("app/ingress/interface.py", f"from {PACKAGE}.app.streaming import StreamingService\n") assert _app_layer_violations() == [] + + @pytest.mark.parametrize(("module", "allowed"), [("config", True), ("interface", False)]) + def test_settings_only_depends_on_the_assistant_schema(self, check, module, allowed) -> None: + check("app/settings.py", f"import {PACKAGE}.app.assistant.{module}\n") + assert bool(_app_layer_violations()) is not allowed From 71786aa62d0ceb3f60c968273e3cf9127fb0c573 Mon Sep 17 00:00:00 2001 From: eandualem Date: Fri, 2 Oct 2026 11:41:48 +0300 Subject: [PATCH 3/5] fix: reap failed Codex RPC subprocesses before detaching An oversized response could leave the reader task's exception unobserved and detach the process before its pipes were closed. Drain and reap the failed subprocess before waking dependent calls, and assert complete cleanup in the existing failure regression. --- src/assistant_runtime/app/voice/_codex_rpc.py | 9 +++++++-- tests/unit/app/voice/test_codex.py | 8 +++++++- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/assistant_runtime/app/voice/_codex_rpc.py b/src/assistant_runtime/app/voice/_codex_rpc.py index f078608b..0f4974d4 100644 --- a/src/assistant_runtime/app/voice/_codex_rpc.py +++ b/src/assistant_runtime/app/voice/_codex_rpc.py @@ -134,15 +134,20 @@ async def _read(self, link: _Link) -> None: thread = (message.get("params") or {}).get("threadId") if thread in link.threads: link.threads[thread].put_nowait(message) + except Exception as exc: + logger.warning("Codex app-server reader failed: {}", type(exc).__name__) finally: # Also when reading fails (an oversized line): nothing reads this # server any more, so end it; its requests and calls end with it, # and the next call starts a new one. - if self._link is link: - self._link = None if link.proc.returncode is None: with contextlib.suppress(ProcessLookupError): link.proc.kill() + # Drain and reap before releasing waiters or losing the only link to + # this process; stop() can no longer find it after detachment. + await link.proc.communicate() + if self._link is link: + self._link = None for future in link.pending.values(): if not future.done(): future.set_exception(VoiceError("Codex app-server exited", 502)) diff --git a/tests/unit/app/voice/test_codex.py b/tests/unit/app/voice/test_codex.py index 0e93d077..7fe04ff8 100644 --- a/tests/unit/app/voice/test_codex.py +++ b/tests/unit/app/voice/test_codex.py @@ -658,12 +658,17 @@ async def test_a_failed_read_ends_the_server_and_releases_waiters(tmp_path): server = _AppServer(_fake_codex(tmp_path, OVERSIZED), 5) try: await server.ensure_started() + link = server._link calls = server.subscribe("thread-1") with pytest.raises(VoiceError, match="exited"): await server.request("boom", {}) closed = await asyncio.wait_for(calls.get(), 5) assert closed["params"]["reason"] == "app_server_exit" assert not server.running # the next call starts a new server + assert link.proc.returncode is not None + assert link.proc.stdin.is_closing() + assert link.reader.done() + assert link.reader.exception() is None finally: await server.stop() @@ -817,12 +822,13 @@ async def counted(): async def test_a_servers_cleanup_ends_its_own_calls_even_after_it_was_replaced(): from types import SimpleNamespace + from unittest.mock import AsyncMock from assistant_runtime.app.voice._codex_rpc import _Link server = _AppServer("codex", 1) old_output = asyncio.StreamReader() - old = _Link(SimpleNamespace(stdout=old_output, returncode=0)) + old = _Link(SimpleNamespace(stdout=old_output, returncode=0, communicate=AsyncMock())) new = _Link(SimpleNamespace(returncode=None)) server._link = new # a newer server took over before the old one's cleanup loop = asyncio.get_running_loop() From 31f5d018af0a6ef33bb904c25553fde9e158ca84 Mon Sep 17 00:00:00 2001 From: eandualem Date: Fri, 2 Oct 2026 12:07:11 +0300 Subject: [PATCH 4/5] fix: finish RPC reader notifications during shutdown Keep waiter failure and closure delivery unconditional if draining is interrupted, and let normal server shutdown await reader finalization after process exit. Cover shutdown with an outstanding RPC and subscriber. --- src/assistant_runtime/app/voice/_codex_rpc.py | 30 +++++++++++-------- tests/unit/app/voice/test_codex.py | 20 +++++++++++++ 2 files changed, 38 insertions(+), 12 deletions(-) diff --git a/src/assistant_runtime/app/voice/_codex_rpc.py b/src/assistant_runtime/app/voice/_codex_rpc.py index 0f4974d4..c439e6de 100644 --- a/src/assistant_runtime/app/voice/_codex_rpc.py +++ b/src/assistant_runtime/app/voice/_codex_rpc.py @@ -145,17 +145,22 @@ async def _read(self, link: _Link) -> None: link.proc.kill() # Drain and reap before releasing waiters or losing the only link to # this process; stop() can no longer find it after detachment. - await link.proc.communicate() - if self._link is link: - self._link = None - for future in link.pending.values(): - if not future.done(): - future.set_exception(VoiceError("Codex app-server exited", 502)) - link.pending.clear() - for queue in link.threads.values(): - queue.put_nowait( - {"method": "thread/realtime/closed", "params": {"reason": "app_server_exit"}} - ) + try: + await link.proc.communicate() + finally: + if self._link is link: + self._link = None + for future in link.pending.values(): + if not future.done(): + future.set_exception(VoiceError("Codex app-server exited", 502)) + link.pending.clear() + for queue in link.threads.values(): + queue.put_nowait( + { + "method": "thread/realtime/closed", + "params": {"reason": "app_server_exit"}, + } + ) async def request(self, method: str, params: dict) -> dict: link = self._link @@ -203,5 +208,6 @@ async def stop(self) -> None: link.proc.kill() await link.proc.wait() if link.reader is not None: - link.reader.cancel() + # Process exit delivers EOF; let the reader finish draining and + # notifying its waiters instead of interrupting its finalization. await asyncio.gather(link.reader, return_exceptions=True) diff --git a/tests/unit/app/voice/test_codex.py b/tests/unit/app/voice/test_codex.py index 7fe04ff8..90ecfe3f 100644 --- a/tests/unit/app/voice/test_codex.py +++ b/tests/unit/app/voice/test_codex.py @@ -673,6 +673,26 @@ async def test_a_failed_read_ends_the_server_and_releases_waiters(tmp_path): await server.stop() +async def test_stopping_the_server_releases_an_outstanding_rpc(tmp_path): + server = _AppServer(_fake_codex(tmp_path, OVERSIZED), 5) + await server.ensure_started() + link = server._link + calls = server.subscribe("thread-1") + pending = asyncio.create_task(server.request("wait", {})) + try: + await until(lambda: bool(link.pending)) + await server.stop() + with pytest.raises(VoiceError, match="exited"): + await asyncio.wait_for(pending, 1) + assert (await asyncio.wait_for(calls.get(), 1))["params"]["reason"] == "app_server_exit" + assert link.reader.done() + assert link.reader.exception() is None + finally: + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + await server.stop() + + async def test_a_cli_that_cannot_run_is_checked_again(tmp_path): transport = CodexTransport( str(tmp_path / "missing-codex"), From ebe3b7d5b72dc037fe45146ebcc7ce0a198513b5 Mon Sep 17 00:00:00 2001 From: eandualem Date: Fri, 2 Oct 2026 12:22:51 +0300 Subject: [PATCH 5/5] fix: handle Codex shutdown races and malformed schema variants Tolerate a process exiting between the shutdown timeout and kill, then reap it. Skip malformed schema variants without hiding supported methods in valid entries. Cover both CodeRabbit findings with focused regressions. --- src/assistant_runtime/app/voice/_codex_rpc.py | 3 +- .../app/voice/_codex_schema.py | 7 ++++- tests/unit/app/voice/test_codex.py | 28 ++++++++++++++++++- 3 files changed, 35 insertions(+), 3 deletions(-) diff --git a/src/assistant_runtime/app/voice/_codex_rpc.py b/src/assistant_runtime/app/voice/_codex_rpc.py index c439e6de..894f14dd 100644 --- a/src/assistant_runtime/app/voice/_codex_rpc.py +++ b/src/assistant_runtime/app/voice/_codex_rpc.py @@ -205,7 +205,8 @@ async def stop(self) -> None: try: await asyncio.wait_for(link.proc.wait(), 5) except TimeoutError: - link.proc.kill() + with contextlib.suppress(ProcessLookupError): + link.proc.kill() await link.proc.wait() if link.reader is not None: # Process exit delivers EOF; let the reader finish draining and diff --git a/src/assistant_runtime/app/voice/_codex_schema.py b/src/assistant_runtime/app/voice/_codex_schema.py index 29ba4d53..637e8ccd 100644 --- a/src/assistant_runtime/app/voice/_codex_schema.py +++ b/src/assistant_runtime/app/voice/_codex_schema.py @@ -42,7 +42,12 @@ def _methods(schema: dict) -> set[str]: methods = set() for variant in schema.get("oneOf", []): - method = (variant.get("properties") or {}).get("method") or {} + if not isinstance(variant, dict): + continue + properties = variant.get("properties") + method = properties.get("method") if isinstance(properties, dict) else None + if not isinstance(method, dict): + continue # One unusual variant must not hide the others. names = method.get("enum") or ([method["const"]] if "const" in method else []) methods.update(name for name in names if isinstance(name, str)) diff --git a/tests/unit/app/voice/test_codex.py b/tests/unit/app/voice/test_codex.py index 90ecfe3f..b05818ec 100644 --- a/tests/unit/app/voice/test_codex.py +++ b/tests/unit/app/voice/test_codex.py @@ -804,7 +804,14 @@ def test_credit_balance_diagnostics_do_not_gate_usage(): def test_one_unusual_schema_variant_does_not_hide_the_others(tmp_path): _schema(tmp_path) requests = json.loads((tmp_path / "ClientRequest.json").read_text()) - requests["oneOf"].append({"properties": {"method": {"const": "future/method"}}}) + requests["oneOf"].extend( + [ + None, + {"properties": True}, + {"properties": {"method": True}}, + {"properties": {"method": {"const": "future/method"}}}, + ] + ) requests["oneOf"].append({"properties": {"params": {}}}) (tmp_path / "ClientRequest.json").write_text(json.dumps(requests)) assert missing_from_schema(tmp_path) == [] @@ -1203,3 +1210,22 @@ async def test_an_abandoned_session_is_stopped_and_released(codex): assert codex._server.sent("thread/realtime/stop") == [{"threadId": "thread-9"}] assert codex._server.sent("thread/unsubscribe") == [{"threadId": "thread-9"}] assert "thread-9" not in codex._server.queues + + +async def test_a_process_exiting_between_timeout_and_kill_is_reaped(): + from types import SimpleNamespace + from unittest.mock import AsyncMock, Mock + + from assistant_runtime.app.voice._codex_rpc import _Link + + proc = SimpleNamespace( + returncode=None, + stdin=SimpleNamespace(close=Mock()), + wait=AsyncMock(side_effect=[TimeoutError, 0]), + kill=Mock(side_effect=ProcessLookupError), + ) + server = _AppServer("codex", 1) + server._link = _Link(proc) + await server.stop() + assert proc.wait.await_count == 2 + proc.kill.assert_called_once()