From 814a9a6c2cb721078abd3995250c263084430a18 Mon Sep 17 00:00:00 2001 From: Sebastien Kawada Date: Sat, 12 Sep 2026 14:25:25 -0700 Subject: [PATCH] feat(agents): add explicit frozen-context delivery sessions --- docs/durable-delivery.md | 115 +++++++ mantis_sdk/__init__.py | 2 + mantis_sdk/agents.py | 20 ++ mantis_sdk/delivery.py | 428 ++++++++++++++++++++++++++ tests/test_delivery.py | 637 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 1202 insertions(+) create mode 100644 docs/durable-delivery.md create mode 100644 mantis_sdk/delivery.py create mode 100644 tests/test_delivery.py diff --git a/docs/durable-delivery.md b/docs/durable-delivery.md new file mode 100644 index 0000000..f5abf72 --- /dev/null +++ b/docs/durable-delivery.md @@ -0,0 +1,115 @@ +# Frozen-context delivery sessions + +`client.agents.delivery_session(...)` creates an explicit Cartographer client for +the version 1 delivery protocol. It requires a model ID supported by the +deployment and uses the configured signed-in session cookie for its WebSocket. +The existing `session(...)` and `run_sync(...)` APIs keep their defaults and wire +format. + +```python +session = client.agents.delivery_session( + space_id, + model_id=deployment_model_id, + user_email=user_email, + space_state_id=space_state_id, + reasoning_effort="high", +) +async with session: + async for event in session.ask( + "Compare the selected samples and explain the evidence.", + active_map_id=map_id, + point_ids=selected_point_ids, + ): + print(event.type, event.item_id) + +result = session.result() +if result.completed: + print(result.text) +``` + +The client prepares a server snapshot before sending the turn. Alternatively, +pass the object returned by `prepare_context()` as `context_snapshot`; only its +version and ID are sent. A snapshot cannot be combined with live source choices. +The request pins the message UUID, runtime, model, reasoning effort, and snapshot. +The first turn must receive a matching `run.accepted` snapshot before any event +is applied or acknowledged. Early frames are buffered within the replay bound. +The supported effort values are `low`, `high`, and `max`. + +```mermaid +sequenceDiagram + participant Caller + participant SDK as DeliverySession + participant Server + Caller->>SDK: ask with source selection + SDK->>Server: Prepare frozen context + Server-->>SDK: Snapshot reference + SDK->>Server: Send pinned request and message ID + Server-->>SDK: run.accepted with run ID and snapshot + Server-->>SDK: runtime_state with acknowledged floor + loop Contiguous event delivery + Server-->>SDK: Sequenced text or native activity + SDK->>SDK: Apply state and deduplicate replay + SDK->>Server: Acknowledge applied sequence + SDK-->>Caller: Normalized event + end + Server-->>SDK: Sequenced terminal event + SDK->>Server: Acknowledge terminal + SDK-->>Caller: Confirmed result +``` + +For scientific workflows, stable text segments and native agent identities keep +replayed progress from duplicating the analysis output. The pinned snapshot +identifies the source selection for the turn, while separate terminal flags keep +a partial answer from being mistaken for a completed analysis. + +## Reconnect and resume + +After an interrupted stream, retain the same session object, reconnect, and call +`resume()`. It resends the exact saved request instead of preparing new sources or +starting another message. Changes to the original snapshot dictionary do not +change the saved request. + +```python +await session.close() +async with session: + async for event in session.resume(): + print(event.type, event.item_id) +result = session.result() +``` + +Event application precedes acknowledgment. If an acknowledgment is lost, the +client retains the applied state and acknowledges replay without duplicating +its result. Repairing a terminal acknowledgment waits for the current +connection's matching cursor handshake. A notification might not reach the caller if transport failure +occurs between state application and yielding the event; inspect `result()` after +an interruption. This protocol does not promise exactly-once execution of caller +callbacks. + +The current result, request, and acknowledgment cursor are held in memory. This +API does not restore them after process restart. A server acknowledgment floor +beyond the retained local state is rejected instead of silently skipping output. +Out-of-order buffering and recent replay-identity tracking are bounded to 2,048 +entries. Conflicting sequences or reused delivery identities raise a protocol +error. + +## Completion and cancellation + +| Observation | Result | +| --- | --- | +| Contiguously applied `chat_complete` | `completed=True` | +| Admission rejection before a run exists, or an applied failure terminal | `failed=True`, with `error` | +| Authoritative cancellation terminal | `cancelled=True` | +| Partial text, completed tool, or completed child agent | No whole-run completion | +| Timeout, disconnect, closed socket, or stopped iterator | Does not itself confirm a terminal or cancel the server turn | +| Rejected resume after a run was admitted | Raises an error; does not claim that the original run failed | + +`await session.cancel()` requests cancellation of the exact message/run. Continue +consuming `ask()` or `resume()` until a terminal arrives. A cancellation rejection +keeps the turn active. Only the server's canonical cancellation terminal can +authorize discarding a missing progress prefix; ordinary completion cannot skip +missing output. + +One consumer can stream a session at a time. Closing an iterator releases that +consumer and preserves the unfinished request for `resume()`. Start a new message +only after the previous turn has a confirmed terminal or was rejected before +admission. Reusing an earlier message UUID for a new turn is rejected. diff --git a/mantis_sdk/__init__.py b/mantis_sdk/__init__.py index 1e81053..6e25fb0 100644 --- a/mantis_sdk/__init__.py +++ b/mantis_sdk/__init__.py @@ -4,6 +4,7 @@ from .agents import AgentEvent, AgentResult, AgentSession from .client import MantisClient from .config import ConfigurationManager +from .delivery import DeliverySession from .enums import AIProvider, DataType, Provider, ReducerModels, SpacePrivacy from .exceptions import ( AgentRunError, @@ -45,6 +46,7 @@ def __getattr__(name: str): "Cell", # agents "AgentSession", + "DeliverySession", "AgentEvent", "AgentResult", # enums diff --git a/mantis_sdk/agents.py b/mantis_sdk/agents.py index 1c0b4e5..44827d0 100644 --- a/mantis_sdk/agents.py +++ b/mantis_sdk/agents.py @@ -28,6 +28,7 @@ from .exceptions import AgentRunError, ConfigurationError, ProviderUnavailableError if TYPE_CHECKING: + from .delivery import DeliverySession from .resources import MantisClientProtocol logger = logging.getLogger("mantis_sdk") @@ -465,6 +466,25 @@ def _assert_available(self, user_email: str, provider: Provider) -> None: ) # --- runs --- + def delivery_session(self, space_id: str | None = None, *, model_id: str, + user_email: str | None = None, chat_id: str | None = None, + timeout: float = 180.0, check_capability: bool = True, + space_state_id: str | None = None, auto_space_state: bool = True, + reasoning_effort: str = "high") -> DeliverySession: + """Create an explicit Cartographer session using frozen context and delivery v1. + + Supply a model ID supported by the deployment. Legacy ``session`` and + ``run_sync`` continue to use their established defaults and protocol. + """ + from .delivery import create_delivery_session + + return create_delivery_session( + self, space_id=space_id, model_id=model_id, user_email=user_email, + chat_id=chat_id, timeout=timeout, check_capability=check_capability, + space_state_id=space_state_id, auto_space_state=auto_space_state, + reasoning_effort=reasoning_effort, + ) + def session(self, space_id: str | None = None, *, provider: Provider | str = DEFAULT_PROVIDER, user_email: str | None = None, diff --git a/mantis_sdk/delivery.py b/mantis_sdk/delivery.py new file mode 100644 index 0000000..8283281 --- /dev/null +++ b/mantis_sdk/delivery.py @@ -0,0 +1,428 @@ +"""Explicit frozen-context turns with ordered, acknowledged event delivery.""" +from __future__ import annotations + +import asyncio +import copy +import hashlib +import json +import math +import time +import uuid +from collections import OrderedDict +from collections.abc import AsyncIterator +from contextlib import aclosing +from typing import Any +from urllib.parse import quote, urlencode + +from .agents import AgentEvent, AgentResult, AgentSession, AgentsResource, _context_uuid, _text +from .enums import Provider +from .exceptions import AgentRunError, ConfigurationError, ProviderUnavailableError + +_EFFORTS = frozenset({"low", "high", "max"}) +_MAX_REPLAY_GAP = 2048 + + +def _uuid(value: str | None, name: str) -> str: + return str(uuid.uuid4()) if value is None else _context_uuid(value, name) + + +def _is_cancellation(data: dict) -> bool: + return data.get("type") == "run_cancelled" or ( + data.get("type") == "chat_fail" and str(data.get("reason") or "").lower() in {"cancelled", "canceled"} + ) + + +def _fingerprint(data: dict) -> str: + return hashlib.sha256(json.dumps(data, sort_keys=True).encode()).hexdigest() + + +class DeliverySession(AgentSession): + """A Cartographer chat with stable request identity and replay-safe delivery. + + ``ask`` prepares frozen context before sending. ``resume`` resends the exact + previous request after an interrupted stream. Keep this object to retain its + in-memory text, event and ACK state across reconnects. + """ + + def __init__(self, resource: AgentsResource, *, user_email: str, provider: Provider, + space_id: str | None, chat_id: str, timeout: float, + model_id: str, space_state_id: str | None = None, reasoning_effort: str = "high"): + self.model_id = model_id + self._resource = resource + self.user_email = user_email + self.provider = provider + self.space_id = space_id + self.space_state_id = space_state_id + self.chat_id = chat_id + self.server_chat_id = chat_id + self._delivery_session_id = str(uuid.uuid4()) + self.timeout = timeout + self.reasoning_effort = reasoning_effort + self._message_id: str | None = None + self._runtime_run_id: str | None = None + self._context_snapshot: dict | None = None + self._ws: Any = None + self._events: list[AgentEvent] = [] + self._segments: dict[str, str] = {} + self._payload: dict | None = None + self._terminal: AgentEvent | None = None + self._streaming = False + self._preparing = False + self._watermark = 0 + self._pending_frames: dict[int, dict] = {} + self._delivery_ids: OrderedDict[str, int] = OrderedDict() + self._seen_sequences: OrderedDict[int, str] = OrderedDict() + self._used_message_ids: set[str] = set() + self._accepted = False + self._before_accept: list[dict] = [] + + @property + def message_id(self) -> str | None: + return self._message_id + + @property + def runtime_run_id(self) -> str | None: + return self._runtime_run_id + + @property + def context_snapshot(self) -> dict | None: + return dict(self._context_snapshot) if self._context_snapshot is not None else None + + @property + def delivery_session_id(self) -> str: + return self._delivery_session_id + + def _ws_url(self) -> str: + cfg = self._resource.http.config + host = cfg.backend_host or cfg.host + ws_base = host.replace("https://", "wss://").replace("http://", "ws://").rstrip("/") + model_id = self._payload["model_id"] if self._payload is not None else self.model_id + query = {"model_id": model_id, "runtime_delivery_v": "1", + "runtime_delivery_session_id": self.delivery_session_id} + if self.space_id: + query["space_id"] = self.space_id + if self.space_state_id: + query["space_state_id"] = self.space_state_id + return f"{ws_base}/ws/chat/{self.chat_id}/{quote(self.user_email, safe='')}/default/?{urlencode(query)}" + + async def _connect(self) -> None: + import inspect + + import websockets + + cookie = self._resource.http.cookie + if not cookie: + raise ConfigurationError("Composer sessions require a signed-in session cookie; internal-service headers do not authenticate this WebSocket") + # Select the supported keyword before connecting, without retrying an + # unrelated TypeError as a second handshake. + header_arg = "additional_headers" if "additional_headers" in inspect.signature(websockets.connect).parameters else "extra_headers" + options: dict[str, Any] = {header_arg: {"Cookie": cookie}, "max_size": None, "open_timeout": self.timeout} + self._ws = await websockets.connect(self._ws_url(), **options) + + async def __aenter__(self) -> DeliverySession: + if self._ws is None: + await self._connect() + return self + + async def __aexit__(self, *exc: Any) -> None: + await self.close() + + async def close(self) -> None: + """Close the transport. Closing a socket does not cancel a server turn.""" + if self._ws is not None: + await self._ws.close() + self._ws = None + + async def ask(self, message: str, *, active_map_id: str | None = None, + bag_ids: list[str] | None = None, cluster_ids: list[str] | None = None, + point_ids: list[str] | None = None, space_ids: list[str] | None = None, + files: list[dict] | None = None, source_snapshot_ids: list[str] | None = None, + open_notebook_path: str | None = None, context_snapshot: dict | None = None, + message_id: str | None = None, reasoning_effort: str | None = None, + generate_suggestions: bool = False) -> AsyncIterator[AgentEvent]: + """Start one turn and yield events until an authoritative terminal. + + Reuse ``resume()`` after an interruption, preserving the original UUID, + model, effort and snapshot. A final text segment alone never ends a turn. + """ + if self._ws is None: + raise ConfigurationError("DeliverySession must be used as an async context manager") + if self._preparing or self._streaming or (self._payload is not None and self._terminal is None): + raise AgentRunError("The previous turn is unfinished; resume it or cancel it before asking again") + if not isinstance(message, str) or not message.strip(): + raise ConfigurationError("message must be nonempty text") + effort = reasoning_effort if reasoning_effort is not None else self.reasoning_effort + if not isinstance(effort, str) or effort not in _EFFORTS: + raise ConfigurationError("reasoning_effort must be low, high, or max") + next_message_id = _uuid(message_id, "message_id") + if next_message_id in self._used_message_ids: + raise ConfigurationError("message_id was already used; resume the existing turn or use a new UUID") + if context_snapshot is not None: + if any(value is not None for value in (active_map_id, bag_ids, cluster_ids, point_ids, space_ids, + files, source_snapshot_ids, open_notebook_path)): + raise ConfigurationError("Pass context_snapshot without live source selections") + if (not isinstance(context_snapshot, dict) or type(context_snapshot.get("version")) is not int + or context_snapshot["version"] != 1 or not context_snapshot.get("id")): + raise ConfigurationError("context_snapshot must contain version=1 and a UUID id") + snapshot_id = _uuid(context_snapshot["id"], "context_snapshot.id") + else: + self._preparing = True + try: + sources = copy.deepcopy({"active_map_id": active_map_id, "bag_ids": bag_ids, + "cluster_ids": cluster_ids, "point_ids": point_ids, + "space_ids": space_ids, "files": files, + "source_snapshot_ids": source_snapshot_ids, + "open_notebook_path": open_notebook_path}) + snapshot = await asyncio.to_thread(self.prepare_context, **sources) + finally: + self._preparing = False + snapshot_id = snapshot["id"] + self._message_id = next_message_id + self._used_message_ids.add(next_message_id) + self._runtime_run_id = None + self._context_snapshot = {"version": 1, "id": snapshot_id} + self._events = [] + self._segments = {} + self._terminal = None + self._watermark = 0 + self._pending_frames = {} + self._delivery_ids.clear() + self._seen_sequences.clear() + self._accepted = False + self._before_accept = [] + self._payload = {"message": message, "message_id": self.message_id, "runtime": self.provider.value, + "model_id": self.model_id, "reasoning_effort": effort, + "context_snapshot": dict(self.context_snapshot), "generate_suggestions": generate_suggestions} + async with aclosing(self._stream()) as events: + async for event in events: + yield event + + async def resume(self) -> AsyncIterator[AgentEvent]: + """Resume the latest turn with its exact saved request and ACK cursor. + + Reconnect by entering this same session object again after ``close()``. + Already applied events remain deduplicated even if an ACK was lost. + """ + if self._payload is None: + raise AgentRunError("There is no prepared turn to resume") + if self._terminal is not None and self.runtime_run_id is None: + raise AgentRunError("This request was rejected before admission; ask again with a new message ID") + async with aclosing(self._stream()) as events: + async for event in events: + yield event + + async def cancel(self) -> None: + """Request cancellation; continue consuming until ``cancelled`` or another terminal.""" + if self._ws is None or not self.message_id or self._terminal is not None: + raise AgentRunError("There is no connected active turn to cancel") + payload = {"type": "cancel_run", "message_id": self.message_id} + if self.runtime_run_id: + payload["runtime_run_id"] = self.runtime_run_id + await self._ws.send(json.dumps(payload)) + + async def _stream(self) -> AsyncIterator[AgentEvent]: + from websockets.exceptions import ConnectionClosed + + if self._streaming: + raise AgentRunError("Only one consumer may stream a session at a time") + if self._ws is None: + raise ConfigurationError("Reconnect by entering this DeliverySession before resuming") + self._streaming = True + try: + await self._ws.send(json.dumps(self._payload)) + repairing_ack = self._terminal is not None + deadline = time.monotonic() + self.timeout + while self._terminal is None or repairing_ack: + try: + raw = await asyncio.wait_for(self._ws.recv(), timeout=max(0, deadline - time.monotonic())) + except asyncio.TimeoutError as exc: + if repairing_ack: + raise AgentRunError("Composer became idle before acknowledgment synchronization; reconnect and resume this session.") from exc + raise AgentRunError("Composer became idle before an authoritative terminal; the turn may still be running. Use resume() with this session.") from exc + try: + data = json.loads(raw) + except (TypeError, ValueError): + continue + if not isinstance(data, dict): + continue + if data.get("type") == "error": + raise AgentRunError(_text(data.get("message") or data.get("code") or "Composer connection failed")) + if data.get("message_id") != self.message_id: + continue + if data.get("chat_id") and data["chat_id"] != self.chat_id: + continue + deadline = time.monotonic() + self.timeout + if repairing_ack and data.get("type") not in {"run.accepted", "run.rejected", "runtime_state"}: + continue + for event in await self._receive(data): + yield event + if (repairing_ack and data.get("type") == "runtime_state" + and data.get("delivery_session_id") == self.delivery_session_id): + # A matching, validated handshake proves this connection + # has claimed the retained cursor before its ACK is resent. + await self._ack() + return + except ConnectionClosed as exc: + raise AgentRunError("Composer disconnected during delivery; reconnect this session and call resume()") from exc + finally: + self._streaming = False + + def _adopt_run(self, data: dict) -> None: + run_id = data.get("runtime_run_id") + if not isinstance(run_id, str) or not run_id: + raise AgentRunError("Composer omitted runtime_run_id") + try: + run_id = str(uuid.UUID(run_id)) + except ValueError as exc: + raise AgentRunError("Composer supplied an invalid runtime_run_id") from exc + if self.runtime_run_id and run_id != self.runtime_run_id: + raise AgentRunError("Composer changed runtime_run_id for the same message") + provider = data.get("provider") or data.get("runtime") + if provider and provider != Provider.Cartographer.value: + raise ProviderUnavailableError(f"Requested Cartographer but the run reported {provider!r}") + self._runtime_run_id = run_id + + def _apply(self, data: dict) -> AgentEvent: + event = AgentEvent.from_wire(copy.deepcopy(data), text_updates=True) + if _is_cancellation(data): + event.type = "cancelled" + self._events.append(event) + if event.type in {"text", "text_update"}: + if event.item_id: + self._segments[event.item_id] = event.text + else: + self._segments[f"segment-{len(self._events)}"] = event.text + if event.is_terminal: + self._terminal = event + return copy.deepcopy(event) + + async def _ack(self) -> None: + await self._ws.send(json.dumps({"type": "runtime_delivery_ack", "runtime_run_id": self.runtime_run_id, + "delivery_session_id": self.delivery_session_id, + "through_sequence": self._watermark})) + + async def _receive(self, data: dict) -> list[AgentEvent]: + wire = data.get("type") + if wire == "run.rejected": + if self.runtime_run_id is not None: + raise AgentRunError("Composer rejected resuming an admitted turn: " + + _text(data.get("message") or data.get("reason"))) + return [self._apply(data)] + if wire == "run.accepted": + snapshot = data.get("context_snapshot") + if (not isinstance(snapshot, dict) or type(snapshot.get("version")) is not int + or snapshot != self.context_snapshot): + raise AgentRunError("Composer accepted a different context snapshot") + self._adopt_run(data) + if self._accepted: + return [] + self._accepted = True + events = [self._apply(data)] + buffered, self._before_accept = self._before_accept, [] + for frame in buffered: + events.extend(await self._receive(frame)) + if self._terminal is not None: + break + return events + if not self._accepted: + if wire == "runtime_state": + if data.get("delivery_session_id") != self.delivery_session_id: + return [] + elif (type(data.get("delivery_sequence")) is not int or data["delivery_sequence"] <= 0 + or not isinstance(data.get("delivery_id"), str) or not data["delivery_id"]): + return [] + if len(self._before_accept) >= _MAX_REPLAY_GAP: + raise AgentRunError("Composer exceeded the pre-acceptance buffer; reconnect and resume") + self._before_accept.append(data) + return [] + if wire == "runtime_state": + if data.get("delivery_session_id") != self.delivery_session_id: + return [] + self._adopt_run(data) + floor = data.get("delivery_ack_floor") + if type(floor) is not int or floor < 0 or floor > self._watermark: + raise AgentRunError("The server ACK floor exceeds this session's applied state; resume with the original session object") + # A handshake never substitutes for unapplied replay frames. + return [] if self._terminal is not None else [self._apply(data)] + sequence = data.get("delivery_sequence") + delivery_id = data.get("delivery_id") + # Historical and unsequenced frames cannot prove progress or completion. + if type(sequence) is not int or sequence <= 0 or not isinstance(delivery_id, str) or not delivery_id: + return [] + self._adopt_run(data) + if sequence <= self._watermark: + if sequence in self._seen_sequences and self._seen_sequences[sequence] != _fingerprint(data): + raise AgentRunError("Composer changed a previously delivered sequence") + await self._ack() + return [] + discarded = data.get("discard_through_sequence") + if _is_cancellation(data) and discarded is not None: + if type(discarded) is not int or discarded != sequence - 1: + raise AgentRunError("Composer supplied an invalid cancellation discard marker") + self._watermark = discarded + self._pending_frames = {key: frame for key, frame in self._pending_frames.items() if key > discarded} + if sequence > self._watermark + _MAX_REPLAY_GAP or len(self._pending_frames) >= _MAX_REPLAY_GAP: + raise AgentRunError("Composer delivery gap exceeded the replay window; reconnect and resume") + previous = self._pending_frames.get(sequence) + if previous is not None and previous != data: + raise AgentRunError("Composer changed a previously delivered sequence") + self._pending_frames[sequence] = data + events = [] + while self._watermark + 1 in self._pending_frames: + frame = self._pending_frames.pop(self._watermark + 1) + identity = frame["delivery_id"] + if identity in self._delivery_ids: + raise AgentRunError("Composer reused a delivery identity for a different sequence") + events.append(self._apply(frame)) + self._delivery_ids[identity] = frame["delivery_sequence"] + self._seen_sequences[frame["delivery_sequence"]] = _fingerprint(frame) + if len(self._delivery_ids) > _MAX_REPLAY_GAP: + self._delivery_ids.popitem(last=False) + self._seen_sequences.popitem(last=False) + self._watermark += 1 + if self._terminal is not None: + break + # State is applied before acknowledgement. Lost ACKs are replay-safe. + await self._ack() + return events + + def result(self) -> AgentResult: + """Return the latest turn, with ``completed=False`` until server confirmation.""" + terminal = self._terminal + failed = bool(terminal and terminal.type in {"fail", "rejected"}) + return AgentResult(text="".join(self._segments.values()), provider=self.provider.value, + events=copy.deepcopy(self._events), failed=failed, error=terminal.text if terminal and failed else None, + completed=bool(terminal and terminal.type == "complete"), + cancelled=bool(terminal and terminal.type == "cancelled"), + message_id=self.message_id, runtime_run_id=self.runtime_run_id) + + +def create_delivery_session(resource: AgentsResource, *, model_id: str, space_id: str | None, + user_email: str | None, chat_id: str | None, timeout: float, + check_capability: bool, space_state_id: str | None, + auto_space_state: bool, reasoning_effort: str) -> DeliverySession: + if not isinstance(model_id, str) or not model_id.strip(): + raise ConfigurationError("model_id must name a model supported by the deployment") + if not isinstance(reasoning_effort, str) or reasoning_effort not in _EFFORTS: + raise ConfigurationError("reasoning_effort must be low, high, or max") + if (isinstance(timeout, bool) or not isinstance(timeout, (int, float)) + or not math.isfinite(timeout) or timeout <= 0): + raise ConfigurationError("timeout must be positive and finite") + user_email = user_email or getattr(resource.http.config, "user_email", None) + if not isinstance(user_email, str) or not user_email.strip(): + raise ConfigurationError("delivery_session requires user_email matching the signed-in session cookie") + canonical_chat_id = _uuid(None if chat_id == "new" else chat_id, "chat_id") + canonical_space_id = _context_uuid(space_id, "space_id") if space_id is not None else None + if space_state_id is not None and canonical_space_id is None: + raise ConfigurationError("space_state_id requires its matching space_id") + canonical_state_id = _context_uuid(space_state_id, "space_state_id") if space_state_id is not None else None + if canonical_space_id and canonical_state_id is None and not auto_space_state: + raise ConfigurationError("Supply space_state_id when disabling auto_space_state for a scoped session") + if check_capability: + resource._assert_available(user_email, Provider.Cartographer) + if canonical_space_id and canonical_state_id is None: + canonical_state_id = _context_uuid(resource._client.space_states.create(canonical_space_id, name="SDK agent"), + "space_state_id") + return DeliverySession(resource, user_email=user_email, provider=Provider.Cartographer, + model_id=model_id, space_id=canonical_space_id, chat_id=canonical_chat_id, + timeout=timeout, space_state_id=canonical_state_id, reasoning_effort=reasoning_effort) diff --git a/tests/test_delivery.py b/tests/test_delivery.py new file mode 100644 index 0000000..011c578 --- /dev/null +++ b/tests/test_delivery.py @@ -0,0 +1,637 @@ +"""Delivery v1 contracts with synthetic snapshots and scripted WebSockets.""" +from __future__ import annotations + +import asyncio +import copy +import json +import uuid +from urllib.parse import parse_qs, urlparse + +import pytest +from websockets.exceptions import ConnectionClosed + +from mantis_sdk import AgentRunError, ConfigurationError, DeliverySession, Provider, ProviderUnavailableError +from mantis_sdk.agents import AgentsResource + +MODEL_ID = "example/analysis-model" +MESSAGE, RUN, SNAPSHOT_ID, CHAT, SPACE, STATE, MAP, CLUSTER, SECOND_RUN = ( + str(uuid.UUID(int=value)) for value in range(1, 10) +) +SNAPSHOT = {"version": 1, "id": SNAPSHOT_ID} + + +def _frame(sequence, **data): + return {"runtime_run_id": RUN, "delivery_sequence": sequence, + "delivery_id": str(uuid.uuid5(uuid.NAMESPACE_OID, str(sequence))), **data} + + +class FakeWS: + def __init__(self, session, script=()): + self.session = session + self.script = list(script) + self.sent = [] + self.closed = False + + async def send(self, data): + self.sent.append(json.loads(data)) + + async def recv(self): + if not self.script: + await asyncio.Event().wait() + data = self.script.pop(0) + if isinstance(data, Exception): + raise data + if isinstance(data, str): + return data + data = {"message_id": self.session.message_id, "chat_id": self.session.chat_id, **data} + if data.get("type") == "run.accepted": + data = {"runtime_run_id": RUN, "context_snapshot": SNAPSHOT, **data} + if data.get("type") == "runtime_state": + data = {"runtime_run_id": RUN, "delivery_session_id": self.session.delivery_session_id, + "delivery_ack_floor": 0, "run_phase": "active", "terminal_sequence": None, **data} + return json.dumps(data) + + async def close(self): + self.closed = True + + +def _session(client, script=(), auto_accept=True, **kwargs): + session = client.agents.delivery_session(SPACE, model_id=MODEL_ID, user_email="u@example.com", chat_id=CHAT, + space_state_id=STATE, check_capability=False, timeout=0.02, **kwargs) + session._ws = FakeWS(session, ([{"type": "run.accepted"}] if auto_accept else []) + list(script)) + return session + + +async def _ask(session, **kwargs): + return [event async for event in session.ask("hello", message_id=MESSAGE, + context_snapshot=SNAPSHOT, **kwargs)] + + +def _acks(session): + return [frame for frame in session._ws.sent if frame.get("type") == "runtime_delivery_ack"] + + +async def test_ask_prepares_snapshot_before_sending_and_omits_live_selection(client, transport): + session = _session(client, [_frame(1, type="chat_complete")]) + transport.queue = [SNAPSHOT] + events = [event async for event in session.ask("hello", active_map_id=MAP, cluster_ids=[CLUSTER], reasoning_effort="max")] + sent = session._ws.sent[0] + assert events[-1].is_terminal + assert sent["runtime"] == "cartographer" and sent["model_id"] == MODEL_ID and sent["reasoning_effort"] == "max" + assert sent["context_snapshot"] == SNAPSHOT + uuid.UUID(sent["message_id"]) + assert not {"bagIds", "clusterIds", "selection", "all_spaces_mode", "composer_mode", "agent_id"} & sent.keys() + assert len(transport.calls) == 1 + + +async def test_snapshot_failure_sends_nothing(client, transport): + session = _session(client) + transport.queue = [{"version": 1}] + with pytest.raises(AgentRunError, match="snapshot"): + async for _ in session.ask("hello"): + pass + assert session._ws.sent == [] and session.message_id is None + + +async def test_frozen_snapshot_cannot_be_mixed_with_mutable_sources(client): + session = _session(client) + with pytest.raises(ConfigurationError, match="without live"): + await _ask(session, active_map_id=MAP) + assert session._ws.sent == [] + + +async def test_text_snapshots_replace_by_segment_and_terminal_never_duplicates_text(client): + session = _session(client, [ + {"type": "run.accepted"}, {"type": "runtime_state"}, + _frame(1, sender="ai", id="segment-1", message="He", partial=True), + _frame(2, sender="ai", id="segment-1", message="Hello ", partial=False), + _frame(3, event="tool_call", id="tool-1", tool_name="commandExecution", arguments={"command": "pwd"}), + _frame(4, sender="ai", id="segment-2", message="world", partial=False), + _frame(5, type="chat_complete", content="Hello world", provider="cartographer"), + ]) + events = await _ask(session) + assert [event.type for event in events] == ["accepted", "init", "text_update", "text", "tool_use", "text", "complete"] + assert session.result().text == "Hello world" and session.result().completed + assert not session.result().failed and session.result().runtime_run_id == RUN + assert _acks(session)[-1] == {"type": "runtime_delivery_ack", "runtime_run_id": RUN, + "delivery_session_id": session.delivery_session_id, "through_sequence": 5} + + +async def test_replay_dedupes_native_progress_and_holds_terminal_behind_gap(client): + progress = _frame(2, event="tool_progress", id="tool", output_delta="working", progress="step 1") + session = _session(client, [ + _frame(1, event="tool_call", id="tool", tool_name="commandExecution"), progress, progress, + _frame(5, type="chat_complete"), + _frame(4, event="tool_output", id="tool", output={"exitCode": 0, "output": "done"}), + _frame(3, event="context_compaction", id="compact", payload={"status": "completed"}), + ]) + events = await _ask(session) + assert [event.type for event in events] == ["accepted", "tool_use", "tool_progress", "context_compaction", "tool_result", "complete"] + assert [ack["through_sequence"] for ack in _acks(session)] == [1, 2, 2, 2, 2, 5] + + +async def test_committed_text_or_unsequenced_complete_cannot_claim_success(client): + session = _session(client, [ + _frame(1, sender="ai", id="segment", message="This is a segment", partial=False), + {"type": "chat_complete", "runtime_run_id": RUN}, + ]) + with pytest.raises(AgentRunError, match="authoritative terminal"): + await _ask(session) + assert session.result().text == "This is a segment" and not session.result().completed + assert not session.result().failed + + +async def test_foreign_run_frames_do_not_change_result_or_cursor(client): + session = _session(client, [ + _frame(1, type="chat_complete", message_id="foreign", provider="claude_code"), + _frame(1, type="chat_complete", chat_id="foreign", provider="claude_code"), + _frame(1, type="chat_complete", provider="cartographer"), + ]) + assert [event.type for event in await _ask(session)] == ["accepted", "complete"] + assert len(_acks(session)) == 1 + + +async def test_same_message_cannot_change_runtime_or_provider(client): + session = _session(client, [ + _frame(1, typing=True), + _frame(2, type="chat_complete", runtime_run_id=SECOND_RUN), + ]) + with pytest.raises(AgentRunError, match="changed runtime_run_id"): + await _ask(session) + session = _session(client, [_frame(1, type="chat_complete", provider="claude_code")]) + with pytest.raises(ProviderUnavailableError, match="Cartographer"): + await _ask(session) + + +@pytest.mark.parametrize("wire, field", [("chat_fail", "failed"), ("run_cancelled", "cancelled")]) +async def test_failure_and_cancellation_are_authoritative_outcomes(client, wire, field): + session = _session(client, [_frame(1, type=wire, content="stopped")]) + await _ask(session) + assert getattr(session.result(), field) and not session.result().completed + assert _acks(session)[-1]["through_sequence"] == 1 + + +async def test_admission_rejection_retains_error_and_sends_no_ack(client): + session = _session(client, [{"type": "run.rejected", "reason": "model_unavailable", "message": "Model unavailable"}], auto_accept=False) + assert [event.type for event in await _ask(session)] == ["rejected"] + assert session.result().failed and session.result().error == "Model unavailable" + assert not _acks(session) + + +async def test_resume_reuses_message_snapshot_effort_and_applied_cursor(client, transport): + first = _frame(1, sender="ai", id="segment", message="partial", partial=True) + session = _session(client, [first]) + transport.queue = [SNAPSHOT] + with pytest.raises(AgentRunError, match="authoritative terminal"): + async for _ in session.ask("hello", message_id=MESSAGE, reasoning_effort="max"): + pass + original = session._ws.sent[0] + delivery_session = session.delivery_session_id + await session.close() + session._ws = FakeWS(session, [ + {"type": "runtime_state", "delivery_ack_floor": 1}, first, + _frame(2, sender="ai", id="segment", message="complete answer", partial=False), + _frame(3, type="chat_complete"), + ]) + events = [event async for event in session.resume()] + assert session._ws.sent[0] == original and session.delivery_session_id == delivery_session + assert original["reasoning_effort"] == "max" and len(transport.calls) == 1 + assert [event.type for event in events] == ["init", "text", "complete"] + assert session.result().text == "complete answer" and session.result().completed + assert _acks(session)[-1]["through_sequence"] == 3 + + +async def test_terminal_handshake_waits_for_unapplied_text_and_terminal_replay(client): + session = _session(client, [ + {"type": "runtime_state", "run_phase": "terminal", "terminal_sequence": 2, + "terminal_payload": {"type": "chat_complete"}}, + _frame(1, sender="ai", id="segment", message="replayed", partial=False), + _frame(2, type="chat_complete"), + ]) + events = await _ask(session) + assert [event.type for event in events] == ["accepted", "init", "text", "complete"] + assert session.result().text == "replayed" + + +async def test_server_cannot_skip_unapplied_prefix(client): + session = _session(client, [{"type": "runtime_state", "delivery_ack_floor": 4}]) + with pytest.raises(AgentRunError, match="ACK floor"): + await _ask(session) + assert not session.result().completed and not _acks(session) + + +async def test_unbounded_gap_fails_without_ack(client): + session = _session(client, [_frame(3000, type="chat_complete")]) + with pytest.raises(AgentRunError, match="delivery gap"): + await _ask(session) + assert not _acks(session) + + +async def test_disconnect_is_incomplete_and_preserves_prepared_request(client): + session = _session(client, [ConnectionClosed(None, None)]) + with pytest.raises(AgentRunError, match="reconnect"): + await _ask(session) + assert session.message_id == MESSAGE and session.context_snapshot == SNAPSHOT + assert not session.result().completed + + +async def test_new_ask_cannot_replace_an_unfinished_turn(client): + session = _session(client) + with pytest.raises(AgentRunError, match="authoritative terminal"): + await _ask(session) + with pytest.raises(AgentRunError, match="previous turn"): + async for _ in session.ask("replacement", context_snapshot=SNAPSHOT): + pass + assert len(session._ws.sent) == 1 + + +async def test_second_turn_resets_result_but_retains_chat_and_transport_identity(client): + session = _session(client, [_frame(1, sender="ai", id="segment", message="one", partial=False), + _frame(2, type="chat_complete")]) + await _ask(session) + session._ws.script.extend([{"type": "run.accepted"}, _frame(1, sender="ai", id="segment", message="two", partial=False), + _frame(2, type="chat_complete")]) + async for _ in session.ask("next", context_snapshot=SNAPSHOT): + pass + assert session.result().text == "two" and session.message_id != MESSAGE + requests = [frame for frame in session._ws.sent if "message" in frame] + assert len(requests) == 2 and requests[0]["message_id"] != requests[1]["message_id"] + + +async def test_cancel_is_a_request_not_terminal_confirmation(client): + session = _session(client, [_frame(1, typing=True)]) + with pytest.raises(AgentRunError): + await _ask(session) + await session.cancel() + assert session._ws.sent[-1] == {"type": "cancel_run", "message_id": MESSAGE, "runtime_run_id": RUN} + assert not session.result().cancelled + session._ws.script.append(_frame(2, type="run_cancelled")) + _ = [event async for event in session.resume()] + assert session.result().cancelled + + +async def test_connect_requires_session_cookie_and_never_sends_initialization(client, monkeypatch): + session = _session(client) + await session.close() + with pytest.raises(ConfigurationError, match="session cookie"): + await session.__aenter__() + client.http.cookie = "sessionid=test" + calls = [] + + async def connect(url, *, additional_headers=None, **kwargs): + calls.append((url, additional_headers)) + return FakeWS(session) + + monkeypatch.setattr("websockets.connect", connect) + async with session: + assert session._ws.sent == [] + assert calls == [(session._ws_url(), {"Cookie": "sessionid=test"})] + + +def test_explicit_factory_preserves_legacy_default_and_separates_model_from_runtime(client): + session = _session(client) + assert isinstance(session, DeliverySession) and session.provider is Provider.Cartographer + assert AgentsResource.DEFAULT_PROVIDER is Provider.OpenCode + assert session.reasoning_effort == "high" + first_url = session._ws_url() + assert first_url == session._ws_url() + query = parse_qs(urlparse(first_url).query) + assert query == {"model_id": [MODEL_ID], "runtime_delivery_v": ["1"], + "runtime_delivery_session_id": [session.delivery_session_id], + "space_id": [SPACE], "space_state_id": [STATE]} + uuid.UUID(session.delivery_session_id) + fresh = client.agents.delivery_session(model_id=MODEL_ID, user_email="a+b@example.com", check_capability=False) + uuid.UUID(fresh.chat_id) + assert "/new/" not in fresh._ws_url() and "a%2Bb%40example.com" in fresh._ws_url() + + +@pytest.mark.parametrize("options,field", [ + ({"model_id": ""}, "model_id"), ({"model_id": None}, "model_id"), + ({"reasoning_effort": "ultra"}, "reasoning_effort"), + ({"timeout": 0}, "timeout"), ({"timeout": -1}, "timeout"), + ({"timeout": float("nan")}, "timeout"), ({"timeout": float("inf")}, "timeout"), + ({"timeout": True}, "timeout"), ({"chat_id": "invalid"}, "chat_id"), + ({"space_id": "invalid"}, "space_id"), ({"space_state_id": STATE}, "space_id"), + ({"space_id": SPACE, "auto_space_state": False}, "space_state_id"), +]) +def test_invalid_delivery_options_fail_before_network(client, transport, options, field): + with pytest.raises(ConfigurationError, match=field): + client.agents.delivery_session(user_email="researcher@example.com", **{"model_id": MODEL_ID, **options}) + assert transport.calls == [] + + +def test_model_and_identity_are_explicit_requirements(client, transport): + with pytest.raises(TypeError, match="model_id"): + client.agents.delivery_session(user_email="researcher@example.com") + with pytest.raises(ConfigurationError, match="user_email"): + client.agents.delivery_session(model_id=MODEL_ID) + assert transport.calls == [] + + +def test_capability_checks_and_automatic_space_state_use_existing_rest_contracts(client, transport): + transport.queue = [{"providers": ["cartographer"]}, [], {"id": STATE, "name": "SDK agent"}] + session = client.agents.delivery_session(SPACE, model_id=MODEL_ID, user_email="researcher@example.com") + assert session.space_state_id == STATE and len(transport.calls) == 3 + assert transport.calls[0]["url"].endswith("/api/agent_execution/providers/") + assert transport.calls[-1]["url"].endswith("/api/space-state/") + transport.calls.clear() + transport.queue = [{"providers": ["opencode"]}] + with pytest.raises(ProviderUnavailableError, match="cartographer"): + client.agents.delivery_session(model_id=MODEL_ID, user_email="researcher@example.com") + assert len(transport.calls) == 1 + + +async def test_handshake_without_replayed_terminal_is_not_completion(client): + session = _session(client, [{"type": "runtime_state", "run_phase": "terminal", "terminal_sequence": 2, + "terminal_payload": {"type": "chat_complete"}}]) + with pytest.raises(AgentRunError, match="authoritative terminal"): + await _ask(session) + assert not session.result().completed and not _acks(session) + + +@pytest.mark.parametrize("reference", [{"version": 1, "id": CHAT}, {"version": True, "id": SNAPSHOT_ID}]) +async def test_accepted_snapshot_must_match_the_prepared_identity(client, reference): + session = _session(client, [{"type": "run.accepted", "context_snapshot": reference}], auto_accept=False) + with pytest.raises(AgentRunError, match="different context"): + await _ask(session) + assert not session.result().completed and not _acks(session) + + +async def test_duplicate_acceptance_and_terminal_content_do_not_duplicate_result(client): + session = _session(client, [{"type": "run.accepted"}, {"type": "run.accepted"}, + _frame(1, sender="ai", id="segment", message="answer", partial=False), + _frame(2, type="chat_complete", content="answer")]) + events = await _ask(session) + assert [event.type for event in events] == ["accepted", "text", "complete"] + assert session.result().text == "answer" + + +@pytest.mark.parametrize("pending", [False, True]) +async def test_replayed_sequence_cannot_change_content(client, pending): + sequence = 2 if pending else 1 + original = _frame(sequence, event="tool_progress", id="tool", output_delta="first") + session = _session(client, [original, {**original, "output_delta": "changed"}]) + with pytest.raises(AgentRunError, match="changed a previously delivered sequence"): + await _ask(session) + assert _acks(session)[-1]["through_sequence"] == (0 if pending else 1) + + +async def test_delivery_identity_cannot_be_reassigned_to_another_sequence(client): + first = _frame(1, event="tool_progress", id="tool", output_delta="first") + session = _session(client, [first, _frame(2, type="chat_complete", delivery_id=first["delivery_id"])]) + with pytest.raises(AgentRunError, match="reused a delivery identity"): + await _ask(session) + assert not session.result().completed and _acks(session)[-1]["through_sequence"] == 1 + + +@pytest.mark.parametrize("wire,reason", [("run_cancelled", None), ("chat_fail", "cancelled"), ("chat_fail", "CANCELED")]) +async def test_authoritative_cancellation_discards_only_the_server_named_prefix(client, wire, reason): + session = _session(client, [ + _frame(1, sender="ai", id="segment", message="applied", partial=True), + _frame(3, sender="ai", id="segment", message="queued", partial=True), + _frame(5000, type=wire, reason=reason, discard_through_sequence=4999), + ]) + events = await _ask(session) + assert [event.type for event in events] == ["accepted", "text_update", "cancelled"] + assert session.result().cancelled and not session.result().failed + assert not session.result().completed and session.result().text == "applied" + assert _acks(session)[-1]["through_sequence"] == 5000 + + +@pytest.mark.parametrize("discarded", [True, -1, 1, 3]) +async def test_invalid_cancellation_discard_marker_never_advances_ack(client, discarded): + session = _session(client, [_frame(3, type="run_cancelled", discard_through_sequence=discarded)]) + with pytest.raises(AgentRunError, match="invalid cancellation discard"): + await _ask(session) + assert not session.result().cancelled and not _acks(session) + + +async def test_non_cancellation_cannot_discard_missing_output(client): + session = _session(client, [_frame(3, type="chat_complete", discard_through_sequence=2)]) + with pytest.raises(AgentRunError, match="authoritative terminal"): + await _ask(session) + assert not session.result().completed and _acks(session)[-1]["through_sequence"] == 0 + + +async def test_cancel_rejection_is_not_a_run_terminal(client): + session = _session(client, [_frame(1, type="run_cancel_pending"), + _frame(2, type="run_cancel_rejected"), _frame(3, type="chat_complete")]) + await _ask(session) + assert session.result().completed and not session.result().cancelled + + +async def test_rejection_after_admission_does_not_report_the_running_turn_as_failed(client): + session = _session(client, [{"type": "run.accepted"}]) + with pytest.raises(AgentRunError, match="authoritative terminal"): + await _ask(session) + session._ws = FakeWS(session, [{"type": "run.rejected", "reason": "context_forbidden"}]) + with pytest.raises(AgentRunError, match="rejected resuming"): + _ = [event async for event in session.resume()] + result = session.result() + assert not result.failed and not result.completed and not result.cancelled and result.runtime_run_id == RUN + + +async def test_rejected_admission_cannot_be_retried_by_resume(client): + session = _session(client, [{"type": "run.rejected", "reason": "unavailable"}], auto_accept=False) + await _ask(session) + sent = copy.deepcopy(session._ws.sent) + with pytest.raises(AgentRunError, match="rejected before admission"): + _ = [event async for event in session.resume()] + assert session._ws.sent == sent + + +@pytest.mark.parametrize("terminal", [False, True]) +async def test_ack_loss_preserves_applied_state_and_resume_repairs_it(client, terminal): + first = _frame(1, type="chat_complete") if terminal else _frame(1, sender="ai", id="segment", message="answer", partial=False) + session = _session(client) + + class LostAck(FakeWS): + async def send(self, data): + if json.loads(data).get("type") == "runtime_delivery_ack": + raise ConnectionClosed(None, None) + await super().send(data) + + session._ws = LostAck(session, [{"type": "run.accepted"}, first]) + with pytest.raises(AgentRunError, match="disconnected"): + await _ask(session) + assert session.result().completed is terminal + previous_count = len(session.result().events) + session._ws = FakeWS(session, [{"type": "runtime_state"}] if terminal else [first, _frame(2, type="chat_complete")]) + _ = [event async for event in session.resume()] + assert session.result().completed + assert len(session.result().events) == previous_count + (0 if terminal else 1) + assert _acks(session)[-1]["through_sequence"] == (1 if terminal else 2) + + +async def test_resume_pins_snapshot_model_effort_and_message_after_caller_changes(client): + supplied = {**SNAPSHOT, "digest": "synthetic-digest"} + session = _session(client) + with pytest.raises(AgentRunError): + _ = [event async for event in session.ask("hello", message_id=MESSAGE, context_snapshot=supplied, + reasoning_effort="max")] + original = copy.deepcopy(session._ws.sent[0]) + supplied["id"] = CHAT + session.context_snapshot["id"] = CHAT + session.model_id = "example/changed-model" + session.reasoning_effort = "low" + session._ws = FakeWS(session, [_frame(1, type="chat_complete")]) + _ = [event async for event in session.resume()] + assert session._ws.sent[0] == original + assert session.context_snapshot == SNAPSHOT + assert parse_qs(urlparse(session._ws_url()).query)["model_id"] == [MODEL_ID] + + +async def test_yielded_and_result_events_cannot_mutate_authoritative_outcome(client): + session = _session(client, [_frame(1, type="chat_complete")]) + events = await _ask(session) + events[-1].type = "fail" + result = session.result() + result.events[-1].type = "cancelled" + assert session.result().completed and not session.result().failed and not session.result().cancelled + + +async def test_closing_iterator_releases_consumer_without_cancelling_server_turn(client): + session = _session(client, [_frame(1, typing=True)]) + stream = session.ask("hello", message_id=MESSAGE, context_snapshot=SNAPSHOT) + await anext(stream) + with pytest.raises(AgentRunError, match="Only one consumer"): + _ = [event async for event in session.resume()] + await stream.aclose() + assert not session._streaming and not any(frame.get("type") == "cancel_run" for frame in session._ws.sent) + session._ws.script.append(_frame(2, type="chat_complete")) + _ = [event async for event in session.resume()] + assert session.result().completed + + +async def test_finished_message_ids_cannot_be_reused_for_a_later_turn(client): + session = _session(client, [_frame(1, type="chat_complete")]) + await _ask(session) + session._ws.script.extend([{"type": "run.accepted"}, _frame(1, type="chat_complete")]) + _ = [event async for event in session.ask("second", context_snapshot=SNAPSHOT)] + with pytest.raises(ConfigurationError, match="already used"): + _ = [event async for event in session.ask("third", message_id=MESSAGE, context_snapshot=SNAPSHOT)] + + +async def test_completed_child_activity_keeps_structure_without_completing_the_analysis(client): + activity = _frame(1, event="subagent_update", id="child-1", + payload={"status": "completed"}, metadata={"agent_name": "Sample Reader"}) + session = _session(client, [activity, activity, _frame(2, event="plan_update", payload={"plan": []})]) + with pytest.raises(AgentRunError, match="authoritative terminal"): + await _ask(session) + result = session.result() + assert [event.type for event in result.events] == ["accepted", "subagent_update", "plan_update"] + assert result.events[1].metadata["agent_name"] == "Sample Reader" + assert result.events[1].payload["status"] == "completed" + assert not result.completed and not result.failed and not result.cancelled + + +async def test_websocket_legacy_header_keyword_is_selected_without_retry(client, monkeypatch): + session = _session(client) + await session.close() + client.http.cookie = "sessionid=test" + calls = [] + + async def connect(url, *, extra_headers=None, **kwargs): + calls.append(extra_headers) + return FakeWS(session) + + monkeypatch.setattr("websockets.connect", connect) + async with session: + assert session._ws.sent == [] + assert calls == [{"Cookie": "sessionid=test"}] + + +async def test_connect_typeerror_does_not_trigger_a_second_handshake(client, monkeypatch): + session = _session(client) + await session.close() + client.http.cookie = "sessionid=test" + calls = [] + + async def connect(url, *, additional_headers=None, **kwargs): + calls.append(url) + raise TypeError("unexpected transport failure") + + monkeypatch.setattr("websockets.connect", connect) + with pytest.raises(TypeError, match="unexpected transport failure"): + await session.__aenter__() + assert len(calls) == 1 + + +async def test_early_sequenced_terminal_cannot_replace_verified_snapshot_acceptance(client): + session = _session(client, [_frame(1, type="chat_complete")], auto_accept=False) + with pytest.raises(AgentRunError, match="authoritative terminal"): + await _ask(session) + assert not session.result().completed and session.result().events == [] + assert session.runtime_run_id is None and not _acks(session) + + +async def test_early_terminal_with_foreign_accepted_snapshot_never_applies_or_acks(client): + session = _session(client, [_frame(1, type="chat_complete"), + {"type": "run.accepted", "context_snapshot": {"version": 1, "id": CHAT}}], + auto_accept=False) + with pytest.raises(AgentRunError, match="different context snapshot"): + await _ask(session) + assert not session.result().completed and session.result().events == [] + assert session.runtime_run_id is None and not _acks(session) + + +async def test_early_frames_apply_after_acceptance_and_only_in_sequence(client): + session = _session(client, [_frame(2, type="chat_complete"), {"type": "run.accepted"}, + _frame(1, sender="ai", id="segment", message="verified", partial=False)], + auto_accept=False) + events = await _ask(session) + assert [event.type for event in events] == ["accepted", "text", "complete"] + assert session.result().text == "verified" and session.result().completed + assert [ack["through_sequence"] for ack in _acks(session)] == [0, 2] + + +async def test_initial_state_and_replay_wait_for_matching_acceptance(client): + session = _session(client, [{"type": "runtime_state"}, _frame(1, type="chat_complete"), + {"type": "run.accepted"}], auto_accept=False) + events = await _ask(session) + assert [event.type for event in events] == ["accepted", "init", "complete"] + assert session.result().completed and _acks(session)[-1]["through_sequence"] == 1 + + +async def test_pre_acceptance_buffer_is_bounded(client, monkeypatch): + monkeypatch.setattr("mantis_sdk.delivery._MAX_REPLAY_GAP", 2) + session = _session(client, [_frame(1, typing=True), _frame(2, typing=True), _frame(3, typing=True)], auto_accept=False) + with pytest.raises(AgentRunError, match="pre-acceptance buffer"): + await _ask(session) + assert session.result().events == [] and not _acks(session) + + +async def test_terminal_ack_repair_waits_for_current_connection_cursor_claim(client): + session = _session(client, [_frame(1, type="chat_complete")]) + await _ask(session) + + class ClaimedSocket(FakeWS): + claimed = False + + async def recv(self): + raw = await super().recv() + frame = json.loads(raw) + if frame.get("type") == "runtime_state" and frame.get("delivery_session_id") == session.delivery_session_id: + self.claimed = True + return raw + + async def send(self, data): + if json.loads(data).get("type") == "runtime_delivery_ack": + assert self.claimed, "ACK sent before this connection received its cursor handshake" + await super().send(data) + + session._ws = ClaimedSocket(session, [_frame(1, type="chat_complete"), + {"type": "run.accepted"}, + {"type": "runtime_state", "delivery_session_id": "foreign"}, + {"type": "runtime_state"}]) + assert [event async for event in session.resume()] == [] + assert _acks(session)[-1]["through_sequence"] == 1 and session.result().completed + + +async def test_terminal_ack_repair_without_current_handshake_keeps_confirmed_outcome(client): + session = _session(client, [_frame(1, type="chat_complete")]) + await _ask(session) + session._ws = FakeWS(session, [{"type": "run.accepted"}]) + with pytest.raises(AgentRunError, match="acknowledgment synchronization"): + _ = [event async for event in session.resume()] + assert session.result().completed and not _acks(session)