diff --git a/AGENTS.md b/AGENTS.md index cfa2c01..7aa319f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -22,4 +22,4 @@ - `workflow_dispatch` inputs may override the model and base URL for a single run. - API keys stay in GitHub secrets only; do not add dispatch inputs for secrets. - The workflow runs `tests/integration/test_live_provider_streaming.py` against manifest-registered graphs in `examples/live_provider_graphs/manifest.json`. -- The proof target is SSE `message_chunk` events from real provider-backed graphs, not just a final successful response. +- The proof target is incremental SSE `messages/partial` frames (aligned to the official LangGraph v1 messages wire contract: `messages/metadata` + `messages/partial`/`messages/complete`) from real provider-backed graphs, not just a final successful response. The legacy `message_chunk` event no longer exists in the official SDK wire format. diff --git a/pyproject.toml b/pyproject.toml index 188459c..3874b76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ dependencies = [ "aiomysql>=0.2.0", "aiosqlite>=0.20.0", "redis>=5.0.0", - "langgraph>=1.0.3", + "langgraph>=1.2.0", "langgraph-sdk>=0.3.5", "langchain-core>=1.0.0", "langchain-openai>=1.0.0", diff --git a/scripts/check_atomic_append_concurrency.py b/scripts/check_atomic_append_concurrency.py new file mode 100644 index 0000000..e3943ac --- /dev/null +++ b/scripts/check_atomic_append_concurrency.py @@ -0,0 +1,120 @@ +"""Cross-dialect concurrency check for the atomic stream append API. + +Runs N concurrent appends against the same run and thread on the given +metadata backend and asserts every seq is unique and gapless, and that the +event table contains exactly N rows (no frame lost, no duplicate). This is the +sqlite-independent proof of the counter-row locking claim: MySQL/PostgreSQL +honour SELECT ... FOR UPDATE (sqlite does not), so the concurrent publishers +must serialize on the counter row rather than rely on uniqueness retries. + +Usage (repo root, from WSL2): + UV_PROJECT_ENVIRONMENT=... uv run python scripts/check_atomic_append_concurrency.py mysql|postgresql|sqlite +""" + +import asyncio +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from sqlalchemy import func, select # noqa: E402 + +from agentseek_api.core.database import db_manager # noqa: E402 +from agentseek_api.core.orm import RunStreamEvent, ThreadStreamEvent # noqa: E402 +from agentseek_api.services import stream_persistence as stream_module # noqa: E402 +from agentseek_api.settings import settings # noqa: E402 + +_N = 12 + + +async def _count_events(scope: str, scope_id: str) -> int: + model = RunStreamEvent if scope == "run" else ThreadStreamEvent + id_column = model.run_id if scope == "run" else model.thread_id + async with db_manager.get_session_factory()() as session: + return int(await session.scalar(select(func.count()).select_from(model).where(id_column == scope_id))) + + +async def _exercise(backend: str) -> None: + # Unique scope ids per run so a pre-existing counter row (from an earlier + # invocation against the same database) cannot shift the expected range. + stamp = int(time.time() * 1000) + run_id, thread_id = f"run-conc-{backend}-{stamp}", f"thread-conc-{backend}-{stamp}" + for scope, scope_id in (("run", run_id), ("thread", thread_id)): + if scope == "run": + append = stream_module.append_run_stream_event_atomic + results = await asyncio.gather(*(append(scope_id, {"event": "message", "data": i}) for i in range(_N))) + rows = await stream_module.load_run_stream_events(scope_id) + seqs = sorted(seq for seq, _ in results) + persisted = [seq for seq, _ in rows] + else: + append = stream_module.append_thread_stream_event_atomic + payload = lambda i: {"method": "values", "params": {"namespace": [], "timestamp": 1, "data": i}} # noqa: E731 + results = await asyncio.gather(*(append(scope_id, payload(i)) for i in range(_N))) + rows = await stream_module.load_thread_stream_events(scope_id, channels=["values"], namespaces=None, depth=None) + seqs = sorted(seq for seq, _ in results) + persisted = [event["seq"] for event in rows] + + assert seqs == list(range(1, _N + 1)), f"[{backend}/{scope}] non-gapless seqs: {seqs}" + assert persisted == list(range(1, _N + 1)), f"[{backend}/{scope}] persisted mismatch: {persisted}" + assert len(results) == _N, f"[{backend}/{scope}] frame dropped: got {len(results)}" + count = await _count_events(scope, scope_id) + assert count == _N, f"[{backend}/{scope}] event rows != {_N}: {count}" + print(f"[{backend}/{scope}] OK: {_N} concurrent appends -> unique gapless 1..{_N}, {count} rows") + + print(f"[{backend}] PASS") + + +def main() -> None: + backend = sys.argv[1] if len(sys.argv) > 1 else "sqlite" + if backend == "sqlite": + settings.SEEKDB_EMBED = False + settings.SEEKDB_URL = "sqlite+aiosqlite:////tmp/atomic-conc.db" + settings.METADATA_DB_BACKEND = "sqlite" + elif backend == "mysql": + settings.SEEKDB_EMBED = False + settings.SEEKDB_URL = "mysql+aiomysql://root:root@127.0.0.1:33306/agentseek" + settings.METADATA_DB_BACKEND = "mysql" + elif backend == "postgresql": + settings.SEEKDB_EMBED = False + settings.SEEKDB_URL = "postgresql://postgres:postgres@127.0.0.1:35432/agentseek" + settings.METADATA_DB_BACKEND = "postgresql" + else: + raise SystemExit(f"unknown backend: {backend}") + asyncio.run(_run(backend)) + + +async def _run(backend: str) -> None: + # The concurrency check only exercises the metadata-DB stream events; the + # langgraph checkpointer/store backends (OceanBase/MySQL via pymysql) + # connect eagerly in their constructors and would fail under a postgres + # metadata URL, so they are replaced with inert stand-ins. + import agentseek_api.core.database as database_module + + class _FakeCheckpointer: + def __init__(self, *args: object, **kwargs: object) -> None: + pass + + def setup(self) -> None: + return None + + def save_checkpoint(self, **kwargs: object) -> None: + return None + + class _FakeStore: + def __init__(self, *args: object, **kwargs: object) -> None: + pass + + database_module.OceanBaseCheckpointSaver = _FakeCheckpointer # type: ignore[assignment] + database_module.LangGraphOceanBaseCheckpointSaver = _FakeCheckpointer # type: ignore[assignment] + database_module.OceanBaseStore = _FakeStore # type: ignore[assignment] + await db_manager.initialize() + try: + await _exercise(backend) + finally: + await db_manager.close() + print(f"[{backend}] db closed") + + +if __name__ == "__main__": + main() diff --git a/scripts/verify_docker_api.py b/scripts/verify_docker_api.py index 5ad1375..c65b318 100644 --- a/scripts/verify_docker_api.py +++ b/scripts/verify_docker_api.py @@ -1,8 +1,10 @@ from __future__ import annotations import argparse +import http.client import json from urllib import error as urllib_error +from urllib import parse as urllib_parse from urllib import request as urllib_request @@ -51,6 +53,67 @@ def _stream_payloads(stream_text: str) -> list[dict[str, object]]: ] +def _sse_frames(stream_text: str) -> list[tuple[int | None, str, dict[str, object]]]: + """Parse SSE text into (id, event, data) frames. ``id`` may be absent.""" + frames: list[tuple[int | None, str, dict[str, object]]] = [] + current_id: int | None = None + current_event = "" + current_data: list[str] = [] + for line in stream_text.splitlines(): + if line.startswith("id: "): + current_id = int(line[len("id: "):].strip()) + elif line.startswith("event: "): + current_event = line[len("event: "):].strip() + elif line.startswith("data: "): + current_data.append(line[len("data: "):].strip()) + elif line == "" and current_data: + frames.append((current_id, current_event, json.loads("".join(current_data)))) + current_id, current_event, current_data = None, "", [] + return frames + + +def _read_sse_prefix_then_disconnect( + *, + base_url: str, + path: str, + headers: dict[str, str], + max_frames: int, + timeout_seconds: float = 15.0, +) -> list[tuple[int | None, str, dict[str, object]]]: + """Connect to an SSE endpoint, read up to ``max_frames`` frames, then close + the connection to simulate a client disconnecting mid-run. + + ``urllib`` (used by ``_request``) blocks until the whole body is read, so a + mid-stream disconnect has to be driven with a lower-level ``http.client`` + connection that we can tear down after a few frames. + """ + parsed = urllib_parse.urlsplit(base_url) + conn = http.client.HTTPConnection(parsed.hostname, parsed.port, timeout=timeout_seconds) + frames: list[tuple[int | None, str, dict[str, object]]] = [] + try: + conn.request("GET", path, headers=headers) + response = conn.getresponse() + current_id: int | None = None + current_event = "" + current_data: list[str] = [] + for raw_line in response: + line = raw_line.decode("utf-8", errors="replace").rstrip("\r\n") + if line.startswith("id: "): + current_id = int(line[len("id: "):].strip()) + elif line.startswith("event: "): + current_event = line[len("event: "):].strip() + elif line.startswith("data: "): + current_data.append(line[len("data: "):].strip()) + elif line == "" and current_data: + frames.append((current_id, current_event, json.loads("".join(current_data)))) + current_id, current_event, current_data = None, "", [] + if len(frames) >= max_frames: + break + finally: + conn.close() + return frames + + def _assert_sample_run( *, @@ -234,8 +297,13 @@ def _assert_common_flow(base_url: str) -> None: assert isinstance(stream_body, str) assert "text/event-stream" in stream_content_type payloads = _stream_payloads(stream_body) - assert any(payload["event"] == "start" for payload in payloads) - assert any(payload["event"] == "end" and payload.get("status") == "success" for payload in payloads) + assert any(isinstance(payload, dict) and payload.get("event") == "start" for payload in payloads) + assert any( + isinstance(payload, dict) + and payload.get("event") == "end" + and payload.get("status") == "success" + for payload in payloads + ) _, stateless_run, _ = _request( base_url=base_url, @@ -326,7 +394,7 @@ def _assert_common_flow(base_url: str) -> None: end_statuses = { payload["status"] for payload in _stream_payloads(resumed_stream) - if payload["event"] == "end" + if isinstance(payload, dict) and payload.get("event") == "end" } assert "success" in end_statuses stress_waited, _ = _assert_sample_run( @@ -358,7 +426,12 @@ def _assert_common_flow(base_url: str) -> None: react_output = react_waited["output"] assert isinstance(react_output, dict) assert "42" in str(react_output["final_text"]) - assert any(payload["event"] == "tool_start" and payload["name"] == "lookup" for payload in react_payloads) + assert any( + isinstance(payload, dict) + and payload.get("event") == "tool-started" + and payload.get("tool_name") == "lookup" + for payload in react_payloads + ) stress_tool_waited, stress_tool_payloads = _assert_sample_run( base_url=base_url, @@ -372,10 +445,107 @@ def _assert_common_flow(base_url: str) -> None: tool_messages = [message for message in stress_tool_output["transcript"] if message["type"] == "ToolMessage"] assert len(tool_messages) == 3 tool_starts = [ - payload for payload in stress_tool_payloads if payload["event"] == "tool_start" and payload["name"] == "slow_process" + payload + for payload in stress_tool_payloads + if isinstance(payload, dict) + and payload.get("event") == "tool-started" + and payload.get("tool_name") == "slow_process" ] assert len(tool_starts) == 3 + _assert_run_stream_reconnect_exactly_once(base_url=base_url, user_headers=alice) + + +def _assert_run_stream_reconnect_exactly_once(*, base_url: str, user_headers: dict[str, str]) -> None: + """Regression for the run-stream reconnect contract. + + A client that connects to ``GET /runs/{id}/stream`` while the run is still + executing, disconnects, and reconnects with ``Last-Event-ID`` must not + replay frames already delivered and must not lose frames produced while it + was disconnected (exactly-once). This runs against the real HTTP boundary + and, in the redis-durable job, against the real Redis run-stream store — + the path the in-process pytest regression covers with a faked store. + """ + _, assistant, _ = _request( + base_url=base_url, + path="/assistants", + method="POST", + payload={"name": "docker-reconnect", "graph_id": "stress_tool_agent"}, + ) + assert isinstance(assistant, dict) + assistant_id = str(assistant["assistant_id"]) + + _, thread, _ = _request( + base_url=base_url, + path="/threads", + method="POST", + payload={"metadata": {"suite": "docker-reconnect"}}, + headers=user_headers, + ) + assert isinstance(thread, dict) + thread_id = str(thread["thread_id"]) + + # ~4.5s total runtime (3 steps * 1.5s) guarantees a mid-run window. + _, run, _ = _request( + base_url=base_url, + path=f"/threads/{thread_id}/runs", + method="POST", + payload={"assistant_id": assistant_id, "input": {"delay": 1.5, "steps": 3}}, + headers=user_headers, + ) + assert isinstance(run, dict) + run_id = str(run["run_id"]) + + stream_path = f"/threads/{thread_id}/runs/{run_id}/stream" + + # Phase 1: connect mid-run, read a few frames, then disconnect. + phase1 = _read_sse_prefix_then_disconnect( + base_url=base_url, + path=stream_path, + headers=user_headers, + max_frames=3, + ) + if not phase1: + phase1 = _read_sse_prefix_then_disconnect( + base_url=base_url, + path=stream_path, + headers=user_headers, + max_frames=3, + ) + assert phase1, "expected to observe the run stream while the run is still active" + phase1_ids = [frame_id for frame_id, _, _ in phase1 if frame_id is not None] + last_id = int(phase1_ids[-1]) + + _, waited, _ = _request( + base_url=base_url, + path=f"/threads/{thread_id}/runs/{run_id}/wait", + headers=user_headers, + ) + assert isinstance(waited, dict) + assert waited["status"] == "success" + + # Phase 2: reconnect with Last-Event-ID and assert exactly-once. + _, phase2_body, _ = _request( + base_url=base_url, + path=stream_path, + headers={**user_headers, "Last-Event-ID": str(last_id)}, + ) + assert isinstance(phase2_body, str) + phase2_ids = [frame_id for frame_id, _, _ in _sse_frames(phase2_body) if frame_id is not None] + + assert all(frame_id > last_id for frame_id in phase2_ids), ( + f"reconnect replayed an already-delivered id: phase1={phase1_ids} phase2={phase2_ids}" + ) + all_ids = phase1_ids + phase2_ids + assert len(all_ids) == len(set(all_ids)), f"duplicate ids delivered across reconnect: {all_ids}" + + _, full_body, _ = _request(base_url=base_url, path=stream_path, headers=user_headers) + assert isinstance(full_body, str) + full_ids = [frame_id for frame_id, _, _ in _sse_frames(full_body) if frame_id is not None] + assert [frame_id for frame_id in full_ids if frame_id > last_id] == phase2_ids, ( + f"reconnect content mismatch: full={full_ids} phase2={phase2_ids}" + ) + def _assert_smoke_flow(base_url: str) -> None: headers = {"x-user-id": "autobuild"} @@ -501,7 +671,7 @@ def _assert_resume_check(base_url: str, *, thread_id: str, run_id: str, resume: end_statuses = [ payload["status"] for payload in _stream_payloads(run_stream) - if payload["event"] == "end" + if isinstance(payload, dict) and payload.get("event") == "end" ] assert "interrupted" in end_statuses assert "success" in end_statuses diff --git a/src/agentseek_api/api/runs.py b/src/agentseek_api/api/runs.py index 06e2c11..b85fab1 100644 --- a/src/agentseek_api/api/runs.py +++ b/src/agentseek_api/api/runs.py @@ -324,6 +324,14 @@ def _protocol_channels_for_stream_modes(stream_modes: list[str]) -> list[str]: return channels +# Channels replayed by the legacy GET /threads/{id}/runs/{run_id}/stream when no +# ``stream_mode`` query is given. The astream migration no longer publishes +# run-scoped stream events, so the run's protocol-v2 thread events (values / +# updates / messages / tools / custom) are replayed instead, alongside the +# run-scoped lifecycle records (start/end) still published by run_jobs. +DEFAULT_RUN_STREAM_REPLAY_CHANNELS = ["values", "updates", "messages", "tools", "custom", "input"] + + async def _iter_persisted_protocol_run_events( *, thread_id: str, @@ -400,6 +408,7 @@ def _is_block_message_event(event: dict[str, Any]) -> bool: async def _event_iter() -> AsyncIterator[str]: try: current_seq = after_seq + saw_interrupt = False if include_metadata: yield _protocol_event_sse(event_name="metadata", data={"run_id": created.run_id, "attempt": 1}) @@ -417,10 +426,13 @@ async def _event_iter() -> AsyncIterator[str]: continue if suppress_block_messages and _is_block_message_event(event): continue + event_data = event.get("params", {}).get("data", {}) + if isinstance(event_data, dict) and "__interrupt__" in event_data: + saw_interrupt = True yield _protocol_event_sse( seq=current_seq, event_name=str(event.get("method", "message")), - data=event.get("params", {}).get("data", {}), + data=event_data, ) if _uses_redis_executor(): @@ -438,10 +450,13 @@ async def _event_iter() -> AsyncIterator[str]: current_seq = max(current_seq, int(event.get("seq", 0))) if suppress_block_messages and _is_block_message_event(event): continue + event_data = event.get("params", {}).get("data", {}) + if isinstance(event_data, dict) and "__interrupt__" in event_data: + saw_interrupt = True yield _protocol_event_sse( seq=current_seq, event_name=str(event.get("method", "message")), - data=event.get("params", {}).get("data", {}), + data=event_data, ) else: async for event in iter_with_sse_keepalives( @@ -462,10 +477,13 @@ async def _event_iter() -> AsyncIterator[str]: current_seq = max(current_seq, int(event.get("seq", 0))) if suppress_block_messages and _is_block_message_event(event): continue + event_data = event.get("params", {}).get("data", {}) + if isinstance(event_data, dict) and "__interrupt__" in event_data: + saw_interrupt = True yield _protocol_event_sse( seq=current_seq, event_name=str(event.get("method", "message")), - data=event.get("params", {}).get("data", {}), + data=event_data, ) final_run = ( @@ -482,7 +500,15 @@ async def _event_iter() -> AsyncIterator[str]: ) return interrupt_event = _interrupt_stream_event_name(stream_modes) - if final_run.status == "interrupted" and final_run.interrupts and interrupt_event is not None: + if ( + final_run.status == "interrupted" + and final_run.interrupts + and interrupt_event is not None + # The interrupt is delivered in-stream (values/updates carrying + # ``__interrupt__``); only emit a trailing event when this + # connection never saw it. + and not saw_interrupt + ): current_seq += 1 yield _protocol_event_sse( seq=current_seq, @@ -888,20 +914,38 @@ async def stream_run( ) async def _event_iter() -> AsyncIterator[str]: - current_seq = after_seq - use_redis_executor = _uses_redis_executor() + # The run stream is a single run-scoped, monotonically sequenced log of + # both lifecycle records (start/end) and protocol frames + # (values/updates/messages/tools). Protocol events are appended to it + # at publication time, so the SSE cursor is one domain: a reconnect + # with ``Last-Event-ID`` (run stream seq) reads only frames with a + # higher seq and never replays already-delivered events. records_by_seq: dict[int, dict[str, object]] = { seq: payload for seq, payload in await load_run_stream_events(run_id, after_seq=after_seq) } records_by_seq.update({seq: payload for seq, payload in run_broker.snapshot_records(run_id, after_seq=after_seq)}) + # Emit the persisted run log strictly by sequence. Lifecycle records + # (start/end) and protocol frames share one monotonic seq domain, and + # within a single run the terminal "end" is always published after that + # run's protocol frames, so a strict ascending replay keeps every id + # monotonic even across a resume: an earlier run's "end" keeps its + # original seq instead of being deferred past newer resume frames. + current_seq = after_seq for seq in sorted(records_by_seq): event = records_by_seq[seq] current_seq = max(current_seq, seq) - event_name = str(event.get("event", "message")) - event_payload: dict[str, object] = {"run_id": run_id, **event} - payload = safe_json_dumps(event_payload) - yield f"id: {seq}\nevent: {event_name}\ndata: {payload}\n\n" + if "method" in event: + yield _protocol_event_sse( + seq=seq, + event_name=str(event["method"]), + data=event.get("params", {}).get("data", {}), + ) + else: + event_name = str(event.get("event", "message")) + event_payload: dict[str, object] = {"run_id": run_id, **event} + yield f"id: {seq}\nevent: {event_name}\ndata: {safe_json_dumps(event_payload)}\n\n" + use_redis_executor = _uses_redis_executor() if use_redis_executor: async for item in iter_with_sse_keepalives( _iter_persisted_run_records( @@ -914,9 +958,17 @@ async def _event_iter() -> AsyncIterator[str]: yield sse_keepalive_comment() continue seq, event = item - event_name = str(event.get("event", "message")) - event_payload = {"run_id": run_id, **event} - yield f"id: {seq}\nevent: {event_name}\ndata: {safe_json_dumps(event_payload)}\n\n" + current_seq = max(current_seq, seq) + if "method" in event: + yield _protocol_event_sse( + seq=seq, + event_name=str(event["method"]), + data=event.get("params", {}).get("data", {}), + ) + else: + event_name = str(event.get("event", "message")) + event_payload: dict[str, object] = {"run_id": run_id, **event} + yield f"id: {seq}\nevent: {event_name}\ndata: {safe_json_dumps(event_payload)}\n\n" return if row.status in TERMINAL_RUN_STATUSES: @@ -927,9 +979,17 @@ async def _event_iter() -> AsyncIterator[str]: yield sse_keepalive_comment() continue seq, event = item - event_name = str(event.get("event", "message")) - event_payload = {"run_id": run_id, **event} - yield f"id: {seq}\nevent: {event_name}\ndata: {safe_json_dumps(event_payload)}\n\n" + current_seq = max(current_seq, seq) + if "method" in event: + yield _protocol_event_sse( + seq=seq, + event_name=str(event["method"]), + data=event.get("params", {}).get("data", {}), + ) + else: + event_name = str(event.get("event", "message")) + event_payload: dict[str, object] = {"run_id": run_id, **event} + yield f"id: {seq}\nevent: {event_name}\ndata: {safe_json_dumps(event_payload)}\n\n" return StreamingResponse( _event_iter(), diff --git a/src/agentseek_api/api/streaming.py b/src/agentseek_api/api/streaming.py index cddbc2e..36465e0 100644 --- a/src/agentseek_api/api/streaming.py +++ b/src/agentseek_api/api/streaming.py @@ -11,6 +11,7 @@ from agentseek_api.core.orm import Run, Thread from agentseek_api.models.auth import User from agentseek_api.models.protocol import ProtocolCommandRequest, ProtocolEventStreamRequest +from agentseek_api.services.stream_modes import normalize_stream_modes from agentseek_api.services.run_preparation import ( ActiveThreadRunConflictError, prepare_and_submit_run, @@ -134,11 +135,29 @@ async def handle_protocol_command( ) try: + run_kwargs: dict[str, Any] | None = None + if payload.params.get("stream_mode") is not None: + # Validate stream mode before submission: an invalid value is a + # client error (400), not a missing resource (404). + try: + stream_modes = normalize_stream_modes(payload.params.get("stream_mode")) + except ValueError as exc: + return _protocol_error( + request_id=payload.id, + code="invalid_argument", + message=str(exc), + status_code=400, + ) + run_kwargs = {"stream_modes": stream_modes} + if payload.params.get("stream_subgraphs"): + run_kwargs = run_kwargs or {} + run_kwargs["stream_subgraphs"] = True run = await prepare_and_submit_run( thread_id=thread_id, assistant_id=assistant_id, payload=_coerce_protocol_input(payload.params.get("input")), user=user, + kwargs=run_kwargs, ) except ValueError as exc: return _protocol_error(request_id=payload.id, code="invalid_argument", message=str(exc), status_code=404) diff --git a/src/agentseek_api/core/orm.py b/src/agentseek_api/core/orm.py index 1dda69c..d06b827 100644 --- a/src/agentseek_api/core/orm.py +++ b/src/agentseek_api/core/orm.py @@ -169,6 +169,31 @@ class ThreadStreamEvent(Base): created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utc_now, nullable=False) +class StreamSequence(Base): + """Per-stream monotonic sequence counter for run/thread stream events. + + One row per (scope, scope_id). Appending an event locks this row + (``UPDATE ... WHERE scope=:s AND scope_id=:id``) inside the same transaction + as the event insert, so concurrent publishers of the same stream serialize + and can never allocate the same seq. The row is also the anchor for the + "durable before expose" contract: a broker only exposes a seq after the + event row (and this counter row) committed. + + The row lives in the metadata DB (not the business tables), so it shares the + lifecycle of the stream events it counts and can be self-healed from + ``MAX(seq)`` if it is ever deleted out from under a running stream. + """ + + __tablename__ = "stream_sequences" + __table_args__ = (UniqueConstraint("scope", "scope_id", name="uq_stream_sequences_scope_id"),) + + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + scope: Mapped[str] = mapped_column(String(16), nullable=False, index=True) + scope_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True) + seq: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utc_now, nullable=False) + + async def get_session(session_factory: async_sessionmaker[AsyncSession]) -> AsyncIterator[AsyncSession]: async with session_factory() as session: yield session diff --git a/src/agentseek_api/services/run_executor.py b/src/agentseek_api/services/run_executor.py index 8c40ace..eaba805 100644 --- a/src/agentseek_api/services/run_executor.py +++ b/src/agentseek_api/services/run_executor.py @@ -1,6 +1,6 @@ +from contextlib import aclosing from dataclasses import dataclass, field import inspect -import logging from typing import Any from langchain_core.messages import BaseMessage, BaseMessageChunk @@ -11,14 +11,7 @@ from agentseek_api.core.database import db_manager from agentseek_api.core.runtime_store import UserScopedStore from agentseek_api.models.auth import User -from agentseek_api.settings import settings from agentseek_api.services.langgraph_service import ensure_sync_checkpoint_mode, get_langgraph_service -from agentseek_api.services.run_state import run_broker -from agentseek_api.services.stream_persistence import ( - append_redis_run_stream_event, - next_run_stream_seq, - persist_run_stream_event, -) from agentseek_api.services.thread_protocol import ( apublish_content_block_delta, apublish_content_block_finish, @@ -44,9 +37,14 @@ ) UNSET = object() -logger = logging.getLogger(__name__) +try: + from langgraph.pregel import _tools as _langgraph_tools # noqa: E402 + _HAS_TOOLS_STREAM_MODE = hasattr(_langgraph_tools, "StreamToolCallHandler") +except Exception: # noqa: BLE001 - older langgraph without the native tools stream mode + _HAS_TOOLS_STREAM_MODE = False + @dataclass class RunExecutionResult: output: dict[str, Any] @@ -54,29 +52,6 @@ class RunExecutionResult: interrupts: list[dict[str, Any]] -async def _publish_translated_run_event( - run_id: str, - event_name: str, - event_payload: dict[str, Any], -) -> tuple[int, dict[str, Any]]: - if settings.EXECUTOR_BACKEND.strip().lower() == "redis": - payload = {"event": event_name, **event_payload} - try: - seq, _ = await append_redis_run_stream_event(run_id, payload) - except Exception: - logger.warning( - "Failed to atomically append translated Redis run event", - extra={"run_id": run_id, "event": event_name}, - exc_info=True, - ) - seq = None - return run_broker.publish(run_id, event_name, seq=seq, **event_payload) - seq = await next_run_stream_seq(run_id) - seq, published_payload = run_broker.publish(run_id, event_name, seq=seq, **event_payload) - await persist_run_stream_event(run_id, seq=seq, payload=published_payload) - return seq, published_payload - - def _normalize_stream_value(value: Any) -> Any: if isinstance(value, BaseMessage): # Use model_dump so the wire-level shape matches the official LangGraph @@ -748,73 +723,6 @@ def _protocol_namespace_for_event(event: dict[str, Any]) -> list[str]: return [] -def _base_stream_payload(event: dict[str, Any]) -> dict[str, Any]: - payload: dict[str, Any] = { - "name": str(event.get("name", "")), - "langgraph_event": str(event.get("event", "")), - "langgraph_run_id": str(event.get("run_id", "")), - "metadata": _normalize_stream_value(event.get("metadata", {})), - "tags": _normalize_stream_value(event.get("tags", [])), - "parent_ids": _normalize_stream_value(event.get("parent_ids", [])), - } - node_name = event.get("metadata", {}).get("langgraph_node") if isinstance(event.get("metadata"), dict) else None - if node_name: - payload["node"] = str(node_name) - return payload - - -def _translate_stream_events(event: dict[str, Any]) -> list[tuple[str, dict[str, Any]]]: - translated: list[tuple[str, dict[str, Any]]] = [] - event_name = event.get("event") - metadata = event.get("metadata", {}) - if not isinstance(metadata, dict): - metadata = {} - - node_name = metadata.get("langgraph_node") - if event_name in {"on_chain_start", "on_chain_end"} and node_name and event.get("name") == node_name: - translated.append( - ( - "node_start" if event_name == "on_chain_start" else "node_end", - _base_stream_payload(event), - ) - ) - - if event_name in {"on_tool_start", "on_tool_end"}: - payload = _base_stream_payload(event) - data = event.get("data", {}) - if isinstance(data, dict): - payload["data"] = _normalize_stream_value(data) - if "input" in data: - payload["input"] = _normalize_stream_value(data["input"]) - if "output" in data: - payload["output"] = _normalize_stream_value(data["output"]) - translated.append(("tool_start" if event_name == "on_tool_start" else "tool_end", payload)) - - if event_name in {"on_chain_stream", "on_chat_model_stream", "on_llm_stream"}: - data = event.get("data", {}) - chunk = data.get("chunk") if isinstance(data, dict) else None - for message in _extract_chunk_messages(chunk): - content = _normalize_stream_value(getattr(message, "content", None)) - tool_calls = _normalize_stream_value(getattr(message, "tool_calls", []) or []) - if content in ("", [], None) and not tool_calls: - continue - payload = _base_stream_payload(event) - payload["message_type"] = type(message).__name__ - payload["content"] = content - if tool_calls: - payload["tool_calls"] = tool_calls - translated.append(("message_chunk", payload)) - if not translated and event_name == "on_llm_stream": - content = _extract_text_chunk(chunk) - if content not in ("", None): - payload = _base_stream_payload(event) - payload["message_type"] = type(chunk).__name__ - payload["content"] = content - translated.append(("message_chunk", payload)) - - return translated - - def _is_root_stream_event(event: dict[str, Any]) -> bool: parent_ids = event.get("parent_ids") return isinstance(parent_ids, list) and not parent_ids @@ -895,18 +803,33 @@ async def execute_run( # ``messages/metadata`` once per message_id. messages_partial_acc: dict[str, BaseMessage] = {} messages_metadata_seen: set[str] = set() + tool_names: dict[Any, str | None] = {} _emitted_values_via_stream = False _requested_stream_modes = run_kwargs.get("stream_modes") or [] + # Aligned with langgraph-api: strip events, always request debug, + # map messages-tuple -> messages, and always request updates so interrupts + # surface even when the client did not ask for updates. + _use_astream_events = "events" in _requested_stream_modes _want_messages_tuple = "messages-tuple" in _requested_stream_modes - _extra_stream_modes = [m for m in _requested_stream_modes if m not in ("messages", "messages-tuple", "updates")] + stream_modes_set = set(_requested_stream_modes) - {"events"} + if "debug" not in stream_modes_set: + stream_modes_set.add("debug") + if "messages-tuple" in stream_modes_set: + stream_modes_set.discard("messages-tuple") + stream_modes_set.add("messages") + _updates_explicitly_requested = "updates" in stream_modes_set + if not _updates_explicitly_requested: + stream_modes_set.add("updates") + _only_interrupt_updates = not _updates_explicitly_requested + if _HAS_TOOLS_STREAM_MODE and "tools" not in stream_modes_set: + stream_modes_set.add("tools") _astream_kwargs: dict[str, Any] = {} _context_schema = getattr(graph, "context_schema", None) if _context_schema is not None: _astream_kwargs["context"] = _resolve_run_context( _context_schema, explicit_context, config.get(CONF, {}) ) - if _extra_stream_modes: - _astream_kwargs["stream_mode"] = list(set(_extra_stream_modes) | {"updates"}) + _astream_kwargs["stream_mode"] = list(stream_modes_set) _interrupt_before = run_kwargs.get("interrupt_before") if _interrupt_before: _astream_kwargs["interrupt_before"] = _interrupt_before @@ -918,222 +841,375 @@ async def execute_run( _astream_kwargs["durability"] = _durability if run_kwargs.get("stream_subgraphs"): _astream_kwargs["subgraphs"] = True - async for stream_event in graph.astream_events(invocation, config, version="v2", **_astream_kwargs): - protocol_namespace = _protocol_namespace_for_event(stream_event) - for event_name, event_payload in _translate_stream_events(stream_event): - await _publish_translated_run_event(run_id, event_name, event_payload) - raw_event_name = stream_event.get("event") - if raw_event_name in {"on_chat_model_stream", "on_llm_stream", "on_chain_stream"}: - data = stream_event.get("data", {}) - chunk = data.get("chunk") if isinstance(data, dict) else None - extracted_messages = _extract_chunk_messages(chunk) - for message_index, message in enumerate(extracted_messages): + async def _publish_complete_messages_from_update(data: dict[str, Any], namespace: list[str] | None) -> None: + """Emit protocol-v2 events for complete (non-LLM) messages inside a state update.""" + seen: set[str] = set() + for value in data.values(): + if not isinstance(value, dict): + continue + messages = value.get("messages") + if not isinstance(messages, list): + continue + for message_index, message in enumerate(messages): + if not isinstance(message, BaseMessage): + continue role = _protocol_role_for_message(message) - blocks = _protocol_blocks_for_message(message) - if role is None or not blocks: + if role not in ("tool", "human", "system"): continue - explicit_message_id = getattr(message, "id", None) - if isinstance(explicit_message_id, str) and explicit_message_id: - message_id = explicit_message_id - else: - message_id = f"{str(stream_event.get('run_id', '')) or run_id}:message:{message_index}" - await protocol_messages.apublish_blocks( - message_id=message_id, - role=role, - blocks=blocks, - namespace=protocol_namespace, + message_id = getattr(message, "id", None) + if isinstance(message_id, str) and message_id: + if message_id in seen: + continue + seen.add(message_id) + await _handle_live_message( + message, + metadata={}, + namespace=namespace, + message_index=message_index, ) - # Emit ``messages/metadata`` once, then accumulate the message - # and emit ``messages/partial`` with the full accumulated payload. - first_seen = message_id not in messages_metadata_seen - if first_seen: - messages_metadata_seen.add(message_id) - await apublish_messages_metadata( + + # Shared live-message handler used by both execution paths. Streams the + # protocol-v2 block events (message-start / content-block-*) plus the + # messages/metadata, messages/partial, messages/complete and messages-tuple + # wire events expected by langgraph-sdk. + async def _handle_live_message( + message: BaseMessage, + *, + metadata: dict[str, Any], + namespace: list[str] | None, + message_index: int, + ) -> None: + role = _protocol_role_for_message(message) + blocks = _protocol_blocks_for_message(message) + if role is None or not blocks: + return + explicit_message_id = getattr(message, "id", None) + if isinstance(explicit_message_id, str) and explicit_message_id: + message_id = explicit_message_id + else: + # Id-less streamed message: derive a stable identity from the stream + # context. ``message_index`` alone restarts for every yielded stream + # event, so two id-less chunks from different subgraph namespaces + # would collide on the same fallback id and be merged into one + # message by the client. The namespace disambiguates them; within one + # namespace the index keeps increments monotonic. + ns_suffix = ":".join(namespace) if namespace else "" + message_id = f"{run_id}:message:{ns_suffix}{':' if ns_suffix else ''}{message_index}" + await protocol_messages.apublish_blocks( + message_id=message_id, + role=role, + blocks=blocks, + namespace=namespace, + ) + # Emit ``messages/metadata`` once, then accumulate the message and emit + # ``messages/partial`` with the full accumulated payload. + first_seen = message_id not in messages_metadata_seen + if first_seen: + messages_metadata_seen.add(message_id) + await apublish_messages_metadata( + thread_id, + message_id=message_id, + metadata=_normalize_stream_value(metadata) or {}, + namespace=namespace, + run_id=run_id, + ) + if role in ("tool", "human", "system"): + if first_seen: + msg_dump = _normalize_stream_value(message) + if isinstance(msg_dump, dict): + await apublish_messages_complete( thread_id, - message_id=message_id, - metadata=_normalize_stream_value(stream_event.get("metadata", {})) or {}, - namespace=protocol_namespace, + messages=[msg_dump], + namespace=namespace, run_id=run_id, ) - if role in ("tool", "human", "system"): - if first_seen: - msg_dump = _normalize_stream_value(message) - if isinstance(msg_dump, dict): - await apublish_messages_complete( - thread_id, - messages=[msg_dump], - namespace=protocol_namespace, - run_id=run_id, - ) - # LangGraph SDK 1.x useStream subscribes to - # messages-tuple, not messages/complete. Mirror the - # completed non-AI message onto that requested - # channel so ToolMessage resolves the pending tool - # call while the run is still streaming. - if _want_messages_tuple: - event_metadata = ( - _normalize_stream_value(stream_event.get("metadata", {})) or {} - ) - await apublish_messages_tuple( - thread_id, - chunk=msg_dump, - metadata=event_metadata, - namespace=protocol_namespace, - run_id=run_id, - ) - continue - if _want_messages_tuple: - chunk_dump = _normalize_stream_value(message) - if isinstance(chunk_dump, dict): - event_metadata = _normalize_stream_value(stream_event.get("metadata", {})) or {} + # LangGraph SDK 1.x useStream subscribes to + # messages-tuple, not messages/complete. Mirror the + # completed non-AI message onto that requested channel so + # ToolMessage resolves the pending tool call while the run + # is still streaming. + if _want_messages_tuple: + event_metadata = _normalize_stream_value(metadata) or {} await apublish_messages_tuple( thread_id, - chunk=chunk_dump, + chunk=msg_dump, metadata=event_metadata, - namespace=protocol_namespace, + namespace=namespace, run_id=run_id, ) - existing = messages_partial_acc.get(message_id) - if existing is None: - accumulated = message - elif not isinstance(message, BaseMessageChunk): - # A full BaseMessage with an id we've been streaming is the - # node's final assembled message — replace, don't re-add. - accumulated = message - else: - left = existing if isinstance(existing, BaseMessageChunk) else _to_chunk(existing) - accumulated = left + message if left is not None else message - messages_partial_acc[message_id] = accumulated - # Wire format mirrors official LangGraph: lowercase ``type`` - # ("ai", not "AIMessageChunk"). ``message_chunk_to_message`` - # converts the accumulated chunk to its non-chunk equivalent - # before serialization. - output_message = ( - message_chunk_to_message(accumulated) - if isinstance(accumulated, BaseMessageChunk) - else accumulated + return + if _want_messages_tuple: + chunk_dump = _normalize_stream_value(message) + if isinstance(chunk_dump, dict): + event_metadata = _normalize_stream_value(metadata) or {} + await apublish_messages_tuple( + thread_id, + chunk=chunk_dump, + metadata=event_metadata, + namespace=namespace, + run_id=run_id, ) - accumulated_dump = _normalize_stream_value(output_message) - if isinstance(accumulated_dump, dict): - await apublish_messages_partial( - thread_id, - messages=[accumulated_dump], - namespace=protocol_namespace, - run_id=run_id, - ) - if raw_event_name == "on_llm_stream": - text = _extract_text_chunk(chunk) - if text not in ("", None): - await protocol_messages.apublish_blocks( - message_id=f"{str(stream_event.get('run_id', '')) or run_id}:message:0", - role="ai", - blocks=[{"type": "text", "text": text}], - namespace=protocol_namespace, - ) - if raw_event_name == "on_tool_start": - metadata = stream_event.get("metadata", {}) - data = stream_event.get("data", {}) - await apublish_tool_event( - thread_id, - tool_event="tool-started", - tool_call_id=str(stream_event.get("run_id", "")), - tool_name=str(stream_event.get("name", "tool")), - node=str(metadata.get("langgraph_node")) if isinstance(metadata, dict) and metadata.get("langgraph_node") else None, - input_payload=_normalize_stream_value(data.get("input")) if isinstance(data, dict) and "input" in data else None, - namespace=protocol_namespace, - ) - if raw_event_name == "on_tool_end": - metadata = stream_event.get("metadata", {}) - data = stream_event.get("data", {}) - await apublish_tool_event( + existing = messages_partial_acc.get(message_id) + if existing is None: + accumulated = message + elif not isinstance(message, BaseMessageChunk): + # A full BaseMessage with an id we've been streaming is the node's + # final assembled message 鈥?replace, don't re-add. + accumulated = message + else: + left = existing if isinstance(existing, BaseMessageChunk) else _to_chunk(existing) + accumulated = left + message if left is not None else message + messages_partial_acc[message_id] = accumulated + output_message = ( + message_chunk_to_message(accumulated) + if isinstance(accumulated, BaseMessageChunk) + else accumulated + ) + accumulated_dump = _normalize_stream_value(output_message) + if isinstance(accumulated_dump, dict): + await apublish_messages_partial( thread_id, - tool_event="tool-finished", - tool_call_id=str(stream_event.get("run_id", "")), - tool_name=str(stream_event.get("name", "tool")), - node=str(metadata.get("langgraph_node")) if isinstance(metadata, dict) and metadata.get("langgraph_node") else None, - output_payload=_normalize_stream_value(data.get("output")) if isinstance(data, dict) and "output" in data else None, - namespace=protocol_namespace, + messages=[accumulated_dump], + namespace=namespace, + run_id=run_id, ) - if raw_event_name == "on_custom_event": - data = stream_event.get("data") + + # Shared stream-mode handler: routes values / updates / custom / debug / + # tasks / checkpoints chunks to their protocol events and detects interrupts. + async def _handle_stream_mode(mode: str, data: Any, namespace: list[str] | None) -> None: + nonlocal interrupt_chunk, interrupt_namespace, result, _emitted_values_via_stream + if mode in ("custom", "debug", "tasks", "checkpoints", "events"): await apublish_stream_mode_event( thread_id, - method="custom", + method=mode, data=_normalize_stream_value(data), - namespace=protocol_namespace, + namespace=namespace, run_id=run_id, ) - if stream_event.get("event") == "on_chain_stream": - data = stream_event.get("data", {}) - chunk = data.get("chunk") if isinstance(data, dict) else None - if isinstance(chunk, tuple) and len(chunk) == 2: - stream_mode_name, stream_mode_data = chunk - if stream_mode_name in ("custom", "debug", "tasks", "checkpoints", "events"): - await apublish_stream_mode_event( - thread_id, - method=stream_mode_name, - data=_normalize_stream_value(stream_mode_data), - namespace=protocol_namespace, - run_id=run_id, - ) - elif stream_mode_name == "values": - normalized_values = _normalize_stream_value(stream_mode_data) - if normalized_values: - _emitted_values_via_stream = True + elif mode == "values": + normalized_values = _normalize_stream_value(data) + if normalized_values: + _emitted_values_via_stream = True + result = data if isinstance(data, dict) else normalized_values + await apublish_values_event( + thread_id, + values=normalized_values, + namespace=namespace, + run_id=run_id, + ) + elif mode == "updates" and isinstance(data, dict): + if "__interrupt__" in data: + interrupt_chunk = data["__interrupt__"] + interrupt_namespace = namespace + # Complete (non-LLM) messages arrive inside the state update, not + # via the messages stream mode. Re-emit them as protocol-v2 message + # events so ToolMessage / HumanMessage still resolve for SDK clients + # even when the client did not request the updates stream mode. + await _publish_complete_messages_from_update(data, namespace) + normalized_chunk = _normalize_stream_value(data) + if isinstance(normalized_chunk, dict): + if _only_interrupt_updates: + # When updates were not explicitly + # requested, only interrupt-bearing updates are forwarded, + # remapped to a ``values`` event with ``__interrupt__`` kept + # intact so the official SDK stream() parser can surface it. + if normalized_chunk.get("__interrupt__"): await apublish_values_event( - thread_id, - values=normalized_values, - namespace=protocol_namespace, - run_id=run_id, - ) - elif stream_mode_name == "updates" and isinstance(stream_mode_data, dict): - if "__interrupt__" in stream_mode_data: - interrupt_chunk = stream_mode_data["__interrupt__"] - interrupt_namespace = protocol_namespace - normalized_chunk = _normalize_stream_value(stream_mode_data) - if isinstance(normalized_chunk, dict): - normalized_chunk.pop("__interrupt__", None) - if normalized_chunk: - await apublish_updates_event( - thread_id, - values=normalized_chunk, - namespace=protocol_namespace, - run_id=run_id, - ) - elif isinstance(chunk, dict): - if "__interrupt__" in chunk: - interrupt_chunk = chunk["__interrupt__"] - interrupt_namespace = protocol_namespace - normalized_chunk = _normalize_stream_value(chunk) - if isinstance(normalized_chunk, dict): - normalized_chunk.pop("__interrupt__", None) - if normalized_chunk: - await apublish_updates_event( thread_id, values=normalized_chunk, - namespace=protocol_namespace, + namespace=namespace, run_id=run_id, ) - if stream_event.get("event") == "on_chain_end" and _is_root_stream_event(stream_event): - data = stream_event.get("data", {}) - if isinstance(data, dict) and "output" in data: - result = data["output"] - normalized_result = _normalize_stream_value(result) - if isinstance(normalized_result, dict): - messages = _extract_protocol_result_messages(normalized_result) - if isinstance(messages, list): - if protocol_messages.saw_live_messages: - await protocol_messages.amerge_final_messages(messages=messages, run_id=run_id) - else: - await apublish_message_transcript(thread_id, run_id=run_id, messages=messages) - await protocol_messages.afinish_all() - if not _emitted_values_via_stream: - await apublish_values_event( - thread_id, - values=normalized_result, + elif normalized_chunk: + # updates explicitly requested: pass through untouched + # (including ``__interrupt__``), matching official. + await apublish_updates_event( + thread_id, + values=normalized_chunk, + namespace=namespace, + run_id=run_id, + ) + elif mode == "tools" and isinstance(data, dict): + # Native langgraph ``tools`` stream mode (langgraph.pregel._tools. + # StreamToolCallHandler): structured tool lifecycle events keyed by + # tool_call_id. tool-finished / tool-error carry no name, so the + # name is tracked from the matching tool-started. + tool_event_name = data.get("event") + tool_call_id = data.get("tool_call_id") + if tool_event_name == "tool-started": + tool_name = data.get("tool_name") + if tool_call_id is not None: + tool_names[tool_call_id] = tool_name + await apublish_tool_event( + thread_id, + tool_event="tool-started", + tool_call_id=str(tool_call_id or ""), + tool_name=tool_name, + input_payload=( + _normalize_stream_value(data.get("input")) if data.get("input") is not None else None + ), + namespace=namespace, + run_id=run_id, + ) + elif tool_event_name == "tool-finished": + await apublish_tool_event( + thread_id, + tool_event="tool-finished", + tool_call_id=str(tool_call_id or ""), + tool_name=tool_names.get(tool_call_id), + output_payload=( + _normalize_stream_value(data.get("output")) if data.get("output") is not None else None + ), + namespace=namespace, + run_id=run_id, + ) + elif tool_event_name == "tool-error": + await apublish_tool_event( + thread_id, + tool_event="tool-error", + tool_call_id=str(tool_call_id or ""), + tool_name=tool_names.get(tool_call_id), + error_message=( + _normalize_stream_value(data.get("message")) if data.get("message") is not None else None + ), + namespace=namespace, + run_id=run_id, + ) + + # Finalize the protocol message stream and, when no values event was emitted + # via the stream, publish the final state as a values event. + async def _finalize_stream(final_result: Any) -> None: + normalized_result = _normalize_stream_value(final_result) + if isinstance(normalized_result, dict): + messages = _extract_protocol_result_messages(normalized_result) + if isinstance(messages, list): + if protocol_messages.saw_live_messages: + await protocol_messages.amerge_final_messages(messages=messages, run_id=run_id) + else: + await apublish_message_transcript(thread_id, run_id=run_id, messages=messages) + await protocol_messages.afinish_all() + if not _emitted_values_via_stream: + await apublish_values_event( + thread_id, + values=normalized_result, + namespace=[], + run_id=run_id, + ) + + if _use_astream_events: + # events mode / remote graphs keep the astream_events path (raw events). + async for stream_event in graph.astream_events(invocation, config, version="v2", **_astream_kwargs): + protocol_namespace = _protocol_namespace_for_event(stream_event) + # Publish each raw astream_events() item onto the ``events`` + # channel so stream_mode="events" actually surfaces the raw event + # stream at the HTTP boundary, not just the translated side effects. + await apublish_stream_mode_event( + thread_id, + method="events", + data=_normalize_stream_value(stream_event), + namespace=protocol_namespace, + run_id=run_id, + ) + raw_event_name = stream_event.get("event") + if raw_event_name in {"on_chat_model_stream", "on_llm_stream", "on_chain_stream"}: + data = stream_event.get("data", {}) + chunk = data.get("chunk") if isinstance(data, dict) else None + extracted_messages = _extract_chunk_messages(chunk) + for message_index, message in enumerate(extracted_messages): + await _handle_live_message( + message, + metadata=stream_event.get("metadata", {}), + namespace=protocol_namespace, + message_index=message_index, + ) + if raw_event_name == "on_llm_stream": + text = _extract_text_chunk(chunk) + if text not in ("", None): + await protocol_messages.apublish_blocks( + message_id=f"{str(stream_event.get('run_id', '')) or run_id}:message:0", + role="ai", + blocks=[{"type": "text", "text": text}], namespace=protocol_namespace, - run_id=run_id, ) + if raw_event_name == "on_custom_event": + await apublish_stream_mode_event( + thread_id, + method="custom", + data=_normalize_stream_value(stream_event.get("data")), + namespace=protocol_namespace, + run_id=run_id, + ) + if raw_event_name == "on_chain_stream": + data = stream_event.get("data", {}) + chunk = data.get("chunk") if isinstance(data, dict) else None + if isinstance(chunk, tuple) and len(chunk) == 2: + stream_mode_name, stream_mode_data = chunk + await _handle_stream_mode(stream_mode_name, stream_mode_data, protocol_namespace) + elif isinstance(chunk, dict): + await _handle_stream_mode("updates", chunk, protocol_namespace) + if raw_event_name == "on_chain_end" and _is_root_stream_event(stream_event): + data = stream_event.get("data", {}) + if isinstance(data, dict) and "output" in data: + result = data["output"] + await _finalize_stream(result) + else: + # Default path: standard astream() stream. Each super-step and stream + # mode yields exactly one chunk, so updates can never be duplicated. + async with aclosing( + graph.astream(invocation, config, **_astream_kwargs) + ) as stream: + async for event in stream: + if _astream_kwargs.get("subgraphs"): + ns, mode, chunk = event + # astream(subgraphs=True) yields tuple namespaces. The + # in-memory broker's live filter compares namespace slices + # to list prefixes, so a tuple would never match and the + # event would vanish under namespace filtering. Normalize + # to a list before publication (persistence already + # JSON-normalizes, which is why this only shows up live). + if isinstance(ns, tuple): + ns = list(ns) + else: + mode, chunk = event + ns = None + if mode == "messages": + if not (isinstance(chunk, tuple) and len(chunk) == 2): + continue + msg, meta = chunk + extracted_messages = _extract_chunk_messages(msg) + for message_index, message in enumerate(extracted_messages): + await _handle_live_message( + message, + metadata=meta, + namespace=ns, + message_index=message_index, + ) + elif mode in ("updates", "values", "custom", "debug", "tasks", "checkpoints", "tools"): + await _handle_stream_mode(mode, chunk, ns) + # astream has no on_chain_end: capture the final state from the + # checkpointer (includes subgraph results), falling back to the last + # values chunk already seen. + try: + # The langgraph root checkpoint lives under an empty checkpoint_ns + # (subgraphs use namespaced checkpoints); the run-scoped ns set in + # config would miss it, so probe the root ns when the run-scoped + # lookup comes back empty. + state = await graph.aget_state(config) + values = getattr(state, "values", None) if state is not None else None + if not isinstance(values, dict) or not values: + root_config = { + **config, + CONF: {**(config.get(CONF) or {}), "checkpoint_ns": ""}, + } + state = await graph.aget_state(root_config) + values = getattr(state, "values", None) if state is not None else None + if isinstance(values, dict): + result = values + except Exception: + pass + if result is not None: + await _finalize_stream(result) + if interrupt_chunk is not None: if isinstance(result, dict): diff --git a/src/agentseek_api/services/run_jobs.py b/src/agentseek_api/services/run_jobs.py index 32aa55f..fc522ac 100644 --- a/src/agentseek_api/services/run_jobs.py +++ b/src/agentseek_api/services/run_jobs.py @@ -16,13 +16,18 @@ add_thread_stream_event_to_session, add_run_stream_event_to_session, append_redis_run_stream_event, - next_run_stream_seq, - next_thread_stream_seq, - persist_run_stream_event, - persist_thread_stream_event, + append_run_stream_event_atomic, + append_thread_stream_event_atomic, + next_run_stream_seq, # noqa: F401 - module attribute; tests assert the redis path never calls it + next_thread_stream_seq, # noqa: F401 - module attribute; tests assert the redis path never calls it ) from agentseek_api.services.thread_checkpoint_store import checkpoint_to_payload, get_latest_checkpoint -from agentseek_api.services.thread_protocol import apublish_lifecycle_event, publish_lifecycle_event, thread_protocol_broker +from agentseek_api.services.thread_protocol import ( + apublish_lifecycle_event, + protocol_timestamp_ms, + publish_lifecycle_event, # noqa: F401 - run_preparation rebinds this module attribute + thread_protocol_broker, +) RUN_EXECUTION_JOB_KIND = "run.execute" TERMINAL_RUN_STATUSES = {"success", "error", "interrupted"} @@ -84,7 +89,16 @@ async def _publish_lifecycle( graph_name: str | None = None, error: str | None = None, session: AsyncSession | None = None, -) -> None: +) -> tuple[int, dict[str, Any]] | None: + """Publish a thread lifecycle event. + + Durable-before-expose ordering: the event row (and its seq) is appended + atomically before the in-memory broker makes it visible, so a client can + never receive a seq that was not durably committed. When ``session`` is + given (terminal lifecycle), the row is staged inside that transaction + instead - the caller must commit and only then expose the event to the + broker, keeping it atomic with the run status. + """ if settings.EXECUTOR_BACKEND.strip().lower() == "redis": await apublish_lifecycle_event( thread_id, @@ -92,33 +106,40 @@ async def _publish_lifecycle( graph_name=graph_name, error=error, ) - return - kwargs: dict[str, Any] = {"event": event} + return None + data: dict[str, Any] = {"event": event} if graph_name is not None: - kwargs["graph_name"] = graph_name + data["graph_name"] = graph_name if error is not None: - kwargs["error"] = error - seq = await next_thread_stream_seq(thread_id) - published = publish_lifecycle_event(thread_id, persist=False, seq=seq, **kwargs) + data["error"] = error + payload: dict[str, Any] = { + "method": "lifecycle", + "params": { + "namespace": [], + "timestamp": protocol_timestamp_ms(), + "data": data, + }, + } if session is None: - await persist_thread_stream_event(thread_id, published) - return - await add_thread_stream_event_to_session( - session, - thread_id, - seq=int(published["seq"]), - payload=published, - ) + seq, _ = await append_thread_stream_event_atomic(thread_id, payload) + thread_protocol_broker.publish(thread_id, payload, persist=False, seq=seq) + return seq, payload + seq, _ = await add_thread_stream_event_to_session(session, thread_id, payload=payload) + return seq, payload async def _publish_run_event( run_id: str, event: str, - *, - persist: bool = True, **payload: Any, ) -> tuple[int, dict[str, Any]] | None: - if settings.EXECUTOR_BACKEND.strip().lower() == "redis" and persist: + """Append a run lifecycle record durably, then expose it to the broker. + + The metadata-DB path allocates the seq and inserts the row in one atomic + transaction and only then publishes to the in-memory broker, so the broker + never exposes a seq that is not durable. + """ + if settings.EXECUTOR_BACKEND.strip().lower() == "redis": event_payload = {"event": event, **payload} try: seq, _ = await append_redis_run_stream_event(run_id, event_payload) @@ -130,32 +151,27 @@ async def _publish_run_event( ) seq = None return run_broker.publish(run_id, event, seq=seq, **payload) - seq = await next_run_stream_seq(run_id) - published = run_broker.publish(run_id, event, seq=seq, **payload) - if published is None: - return None - seq, event_payload = published - if persist: - await persist_run_stream_event(run_id, seq=seq, payload=event_payload) - return seq, event_payload + seq, _ = await append_run_stream_event_atomic(run_id, {"event": event, **payload}) + return run_broker.publish(run_id, event, seq=seq, **payload) -async def _persist_thread_snapshot(thread_id: str) -> None: - if settings.EXECUTOR_BACKEND.strip().lower() == "redis": - return - for event in thread_protocol_broker.snapshot_records(thread_id): - await persist_thread_stream_event(thread_id, event) - +async def _publish_terminal_run_event(session: AsyncSession, run_id: str, *, status: str) -> tuple[int, dict[str, Any]] | None: + """Record the terminal ``end`` event. -async def _publish_terminal_run_event(session: AsyncSession, run_id: str, *, status: str) -> None: + Redis: the atomic Lua append already makes the record durable before the + broker publishes it, so this delegates to ``_publish_run_event`` unchanged. + Inline: the row is staged inside ``session`` (allocating its seq from the + locked counter row) without touching the broker; the caller commits and + then exposes the event so the terminal status and its stream record are + durable as one unit. + """ if settings.EXECUTOR_BACKEND.strip().lower() == "redis": - await _publish_run_event(run_id, "end", status=status) - return - terminal_run_event = await _publish_run_event(run_id, "end", status=status, persist=False) - if terminal_run_event is None: - return - seq, event_payload = terminal_run_event - await add_run_stream_event_to_session(session, run_id, seq=seq, payload=event_payload) + return await _publish_run_event(run_id, "end", status=status) + return await add_run_stream_event_to_session( + session, + run_id, + payload={"event": "end", "status": status}, + ) def _apply_execution_result(db_run: Run, result: RunExecutionResult) -> None: @@ -209,7 +225,6 @@ async def execute_run_job(job: RunExecutionJob) -> None: if job.kwargs: execute_kwargs["kwargs"] = job.kwargs result = await execute_run(**execute_kwargs) - await _persist_thread_snapshot(job.thread_id) await execution_session.refresh(db_run) if not _is_cancelled_run(db_run): # A missing checkpoint lookup should not turn a successful run into a failed one. @@ -234,13 +249,14 @@ async def execute_run_job(job: RunExecutionJob) -> None: if thread is not None: thread.status = "interrupted" if db_run.status == "interrupted" else ("error" if db_run.status == "error" else "idle") thread.state_updated_at = db_run.updated_at - await _publish_terminal_run_event(execution_session, job.run_id, status=db_run.status) + is_redis_executor = settings.EXECUTOR_BACKEND.strip().lower() == "redis" + terminal = await _publish_terminal_run_event(execution_session, job.run_id, status=db_run.status) lifecycle_state = "completed" if db_run.status == "interrupted": lifecycle_state = "interrupted" elif db_run.status == "error": lifecycle_state = "failed" - await _publish_lifecycle( + lifecycle = await _publish_lifecycle( job.thread_id, event=lifecycle_state, graph_name=job.graph_id, @@ -248,5 +264,16 @@ async def execute_run_job(job: RunExecutionJob) -> None: session=execution_session, ) await execution_session.commit() + if not is_redis_executor and terminal is not None and lifecycle is not None: + # The terminal records are durable with the run state now; only + # then expose them to the in-memory brokers, so a client never + # sees a seq that was not durably committed. + run_broker.publish(job.run_id, "end", seq=terminal[0], status=db_run.status) + thread_protocol_broker.publish( + job.thread_id, + lifecycle[1], + persist=False, + seq=lifecycle[0], + ) finally: thread_protocol_broker.run_finished(job.thread_id) diff --git a/src/agentseek_api/services/run_preparation.py b/src/agentseek_api/services/run_preparation.py index 85c03b0..344ceb1 100644 --- a/src/agentseek_api/services/run_preparation.py +++ b/src/agentseek_api/services/run_preparation.py @@ -14,7 +14,6 @@ execute_run = run_jobs_module.execute_run run_broker = run_jobs_module.run_broker _publish_run_event = run_jobs_module._publish_run_event -_persist_thread_snapshot = run_jobs_module._persist_thread_snapshot add_run_stream_event_to_session = run_jobs_module.add_run_stream_event_to_session publish_lifecycle_event = run_jobs_module.publish_lifecycle_event thread_protocol_broker = run_jobs_module.thread_protocol_broker @@ -122,7 +121,6 @@ async def _execute_and_persist( run_jobs_module.execute_run = execute_run run_jobs_module.run_broker = run_broker run_jobs_module._publish_run_event = _publish_run_event - run_jobs_module._persist_thread_snapshot = _persist_thread_snapshot run_jobs_module.add_run_stream_event_to_session = add_run_stream_event_to_session run_jobs_module.publish_lifecycle_event = publish_lifecycle_event run_jobs_module.thread_protocol_broker = thread_protocol_broker diff --git a/src/agentseek_api/services/run_state.py b/src/agentseek_api/services/run_state.py index 41938eb..b3b3224 100644 --- a/src/agentseek_api/services/run_state.py +++ b/src/agentseek_api/services/run_state.py @@ -18,6 +18,11 @@ def publish(self, run_id: str, event: str, *, seq: int | None = None, **payload: event_payload = {"event": event, **payload} if seq is None: seq = self._next_seq[run_id] + else: + # Never regress below the in-memory watermark: an explicit seq from + # persistent state may be lower than what this process has already + # allocated (e.g. the broker was cleared and re-seeded from the DB). + seq = max(seq, self._next_seq[run_id]) self._next_seq[run_id] = max(self._next_seq[run_id], seq + 1) self._events[run_id].append(event_payload) self._seqs[run_id].append(seq) @@ -34,6 +39,28 @@ def publish(self, run_id: str, event: str, *, seq: int | None = None, **payload: self._signals[run_id].set() return seq, dict(event_payload) + def publish_protocol(self, run_id: str, payload: dict[str, Any], *, seq: int | None = None) -> tuple[int, dict[str, Any]]: + """Publish a protocol-v2 event into the run's ordered log. + + In inline mode this makes the run broker the single run-scoped log + shared by both lifecycle records (start/end) and protocol frames + (values/updates/messages/tools), so the replay endpoint reads one + monotonic ``seq`` cursor instead of mixing two sequence domains. + """ + if seq is None: + seq = self._next_seq[run_id] + else: + # Never regress below the in-memory watermark (same rationale as + # ``publish``): a persistent-state seq must not collide with seqs + # this process has already handed out, and a cold broker re-seeded + # from the DB must keep allocating after the persisted max. + seq = max(seq, self._next_seq[run_id]) + self._next_seq[run_id] = max(self._next_seq[run_id], seq + 1) + self._events[run_id].append(payload) + self._seqs[run_id].append(seq) + self._signals[run_id].set() + return seq, dict(payload) + def snapshot(self, run_id: str) -> list[dict[str, Any]]: return [dict(event) for event in self._events.get(run_id, [])] diff --git a/src/agentseek_api/services/stream_persistence.py b/src/agentseek_api/services/stream_persistence.py index 23da28e..42d7168 100644 --- a/src/agentseek_api/services/stream_persistence.py +++ b/src/agentseek_api/services/stream_persistence.py @@ -1,15 +1,17 @@ from __future__ import annotations +import asyncio import json import logging from typing import Any from redis.asyncio import Redis, from_url -from sqlalchemy import delete, select +from sqlalchemy import delete, func, select +from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from agentseek_api.core.database import db_manager -from agentseek_api.core.orm import RunStreamEvent, ThreadStreamEvent +from agentseek_api.core.orm import RunStreamEvent, StreamSequence, ThreadStreamEvent from agentseek_api.settings import settings from agentseek_api.services.thread_protocol import _namespace_matches, protocol_channel_for_method @@ -18,6 +20,10 @@ _RUN_STREAM_KEY_PREFIX = "agentseek:runs:stream" _THREAD_STREAM_KEY_PREFIX = "agentseek:threads:stream" _THREAD_STREAM_ENVELOPE_FIELDS = frozenset({"type", "event_id", "seq"}) +# Serializes counter-row seeding per stream so a first-append burst opens one +# seed connection instead of one per publisher (the seed runs in its own short +# transaction; unbounded simultaneous seeds would exhaust the metadata pool). +_stream_seed_locks: dict[tuple[str, str], asyncio.Lock] = {} _redis_client: Redis | None = None logger = logging.getLogger(__name__) @@ -136,15 +142,259 @@ async def _load_redis_stream_events(key: str, *, after_seq: int) -> list[tuple[i return events +def _scope_event_model(scope: str) -> type[RunStreamEvent] | type[ThreadStreamEvent]: + if scope == "run": + return RunStreamEvent + if scope == "thread": + return ThreadStreamEvent + raise ValueError(f"Unsupported stream scope: {scope}") + + +def _thread_envelope(thread_id: str, seq: int, payload: dict[str, Any]) -> dict[str, Any]: + """Mirror the wire envelope ``ThreadProtocolEventBroker._record_event`` builds. + + Keeping the persisted row byte-compatible with the in-memory broker event + means ``_record_event(seq=...)`` reproduces exactly what was already + committed, so the broker never re-derives a different identity. + """ + return { + "type": "event", + "event_id": f"{thread_id}:{seq}", + "seq": seq, + **payload, + } + + +async def _seed_stream_sequence(scope: str, scope_id: str) -> None: + """Create the per-stream counter row in its own short transaction. + + Runs outside the caller's transaction so the create-race never executes + inside a longer append/terminal transaction: MySQL raises + ``SAVEPOINT ... does not exist`` when a concurrent creator collides inside + ``begin_nested``. Seeding separately keeps the append transaction purely + lock-then-allocate. A concurrent creator is resolved by the unique + constraint; the row is seeded from ``MAX(seq)`` so it can be re-created from + durable state even if it was deleted out from under an active stream. + """ + session_factory = db_manager.get_session_factory() + model = _scope_event_model(scope) + id_column = model.run_id if scope == "run" else model.thread_id + async with session_factory() as session: + max_seq = await session.scalar(select(func.max(model.seq)).where(id_column == scope_id)) + try: + session.add(StreamSequence(scope=scope, scope_id=scope_id, seq=max_seq or 0)) + await session.commit() + except IntegrityError: + # A concurrent publisher created the row first (or a concurrent + # seed committed between our SELECT and INSERT): nothing to do. + await session.rollback() + + +async def _ensure_stream_sequence(session: AsyncSession, scope: str, scope_id: str) -> StreamSequence: + """Return the per-stream counter row, creating it if missing. + + The row is locked (``SELECT ... FOR UPDATE``) so the caller's transaction + holds the only allocation right for this stream; concurrent publishers + serialize on this single row and can never observe the same ``MAX(seq)+1``. + Two database pitfalls are deliberately avoided: + + - ``FOR UPDATE`` is never issued against a missing row: on MySQL/InnoDB a + point ``FOR UPDATE`` over a non-existent unique key takes a gap lock, and + the separate-transaction seed then deadlocks against it (``Lock wait + timeout``). The row is seeded first (its own short transaction) and only + then locked. + - After seeding in a separate transaction, the row is confirmed with a + ``FOR UPDATE`` current read: under MySQL REPEATABLE READ a plain ``SELECT`` + would keep returning the pre-seed snapshot within the caller's + transaction. + + A missing row (first append, or deleted out from under an active stream by + the two-phase delete path in ``threads.py``) is re-seeded from ``MAX(seq)``. + The in-process per-stream lock serializes the seed so a first-append burst + opens one seed transaction instead of exhausting the metadata pool. + """ + stmt = select(StreamSequence).where( + StreamSequence.scope == scope, StreamSequence.scope_id == scope_id + ) + lock_key = (scope, scope_id) + lock = _stream_seed_locks.setdefault(lock_key, asyncio.Lock()) + async with lock: + for _ in range(_MAX_ATOMIC_APPEND_RETRIES): + row = await session.scalar(stmt) + if row is not None: + # Row exists: the unique index turns this into a record lock + # only - no gap lock, safe to hold across the allocation. + locked = await session.scalar(stmt.with_for_update()) + if locked is not None: + return locked + await _seed_stream_sequence(scope, scope_id) + # Current read: sees the seed even under snapshot isolation. + locked = await session.scalar(stmt.with_for_update()) + if locked is not None: + return locked + raise RuntimeError(f"Failed to establish {scope} stream sequence counter for {scope_id}") + + +async def _stage_db_event( + session: AsyncSession, + scope: str, + scope_id: str, + payload: dict[str, Any], + *, + seq: int | None, +) -> tuple[int, dict[str, Any]]: + """Allocate (or commit) the stream seq and stage the event row in ``session``. + + Does not commit: the standalone path commits explicitly, while the + in-session path (terminal events) commits together with the run/thread + status so ``seq`` and state are durable as one unit. + """ + counter = await _ensure_stream_sequence(session, scope, scope_id) + new_seq = counter.seq + 1 if seq is None else seq + counter.seq = max(counter.seq, new_seq) + if scope == "run": + session.add( + RunStreamEvent( + run_id=scope_id, + seq=new_seq, + event=str(payload.get("method") or payload.get("event", "message")), + payload_json=dict(payload), + ) + ) + else: + session.add( + ThreadStreamEvent( + thread_id=scope_id, + seq=new_seq, + method=str(payload.get("method", "event")), + payload_json=dict(_thread_envelope(scope_id, new_seq, payload)), + ) + ) + return new_seq, dict(payload) + + +# Upper bound on uniqueness retries. The metadata-DB append relies on the +# per-stream counter row's row lock to serialize publishers; SQLite ignores +# ``SELECT ... FOR UPDATE``, so concurrent publishers can read the same counter +# value and collide on the ``UNIQUE(scope_id, seq)`` constraint. Each retry +# rolls back and re-reads the counter, so every successful commit advances the +# stream by one - concurrent appends converge to unique, gapless seqs. The +# bound is a safety valve; normal (non-concurrent) appends never retry. +_MAX_ATOMIC_APPEND_RETRIES = 32 + + +async def _db_append( + scope: str, + scope_id: str, + payload: dict[str, Any], + *, + seq: int | None = None, +) -> tuple[int, dict[str, Any]]: + """Append a stream event to the metadata DB atomically. + + One transaction per attempt: lock the per-stream counter row, allocate the + seq, insert the event row. A uniqueness collision (concurrent publisher, or + a pre-assigned ``seq`` that collides with an out-of-band row) is healed by + rolling back and re-allocating from durable state, keeping the stream + monotonic instead of dropping the frame. Any other failure raises - the + caller must not expose an event whose durable append did not succeed. + """ + session_factory = db_manager.get_session_factory() + async with session_factory() as session: + for _ in range(_MAX_ATOMIC_APPEND_RETRIES): + try: + result = await _stage_db_event(session, scope, scope_id, payload, seq=seq) + await session.commit() + return result + except IntegrityError: + await session.rollback() + # Re-allocate from durable state next attempt instead of + # reusing the colliding seq. + seq = None + raise RuntimeError( + f"Failed to atomically append {scope} stream event after " + f"{_MAX_ATOMIC_APPEND_RETRIES} attempts (scope_id={scope_id})" + ) + + +async def append_run_stream_event_atomic( + run_id: str, + payload: dict[str, Any], + *, + seq: int | None = None, +) -> tuple[int, dict[str, Any]]: + """Durably append a run-scoped stream event and return its seq. + + Redis executor: single Lua ``INCR``+``XADD`` (atomic by construction). + Inline executor: single metadata-DB transaction. The caller must only + expose the event to clients after this returns. + """ + if _uses_redis_executor(): + if seq is not None: # pragma: no cover - redis appends always allocate + raise ValueError("Redis stream append allocates its own seq") + return await append_redis_run_stream_event(run_id, payload) + if not _metadata_db_ready(): + # No metadata DB at all (offline tests / pre-initialization): there is + # nothing durable to protect, so fall back to broker-local sequence + # allocation (seq=None) exactly like the legacy path. Production runs + # always have the DB initialized, so this is a startup/offline posture, + # not a durable-path fallback. + return (None, dict(payload)) + return await _db_append("run", run_id, payload, seq=seq) + + +async def append_thread_stream_event_atomic( + thread_id: str, + payload: dict[str, Any], + *, + seq: int | None = None, +) -> tuple[int, dict[str, Any]]: + """Durably append a thread-protocol event and return its seq (see run twin).""" + if _uses_redis_executor(): + if seq is not None: # pragma: no cover - redis appends always allocate + raise ValueError("Redis stream append allocates its own seq") + return await append_redis_thread_stream_event(thread_id, payload) + if not _metadata_db_ready(): + # No metadata DB at all (offline tests / pre-initialization): nothing + # durable to protect, fall back to broker-local sequence allocation. + return (None, dict(payload)) + return await _db_append("thread", thread_id, payload, seq=seq) + + async def next_run_stream_seq(run_id: str) -> int | None: + # Legacy helper, retained only for tests and backward compatibility. + # Production callers must use append_run_stream_event_atomic so the + # allocation and the durable write are one atomic unit. if not _uses_redis_executor(): - return None + if not _metadata_db_ready(): + return None + try: + session_factory = db_manager.get_session_factory() + except RuntimeError: + return None + async with session_factory() as session: + row = await session.scalar( + select(func.max(RunStreamEvent.seq)).where(RunStreamEvent.run_id == run_id) + ) + return (row or 0) + 1 return int(await _get_redis_client().incr(f"{_RUN_STREAM_SEQ_KEY_PREFIX}:{run_id}")) async def next_thread_stream_seq(thread_id: str) -> int | None: + # Legacy helper, retained only for tests and backward compatibility. + # Production callers must use append_thread_stream_event_atomic. if not _uses_redis_executor(): - return None + if not _metadata_db_ready(): + return None + try: + session_factory = db_manager.get_session_factory() + except RuntimeError: + return None + async with session_factory() as session: + row = await session.scalar( + select(func.max(ThreadStreamEvent.seq)).where(ThreadStreamEvent.thread_id == thread_id) + ) + return (row or 0) + 1 return int(await _get_redis_client().incr(f"{_THREAD_STREAM_SEQ_KEY_PREFIX}:{thread_id}")) @@ -182,7 +432,7 @@ async def persist_run_stream_event(run_id: str, *, seq: int, payload: dict[str, RunStreamEvent( run_id=run_id, seq=seq, - event=str(payload.get("event", "message")), + event=str(payload.get("method") or payload.get("event", "message")), payload_json=dict(payload), ) ) @@ -195,56 +445,59 @@ async def add_run_stream_event_to_session( session: AsyncSession, run_id: str, *, - seq: int, + seq: int | None = None, payload: dict[str, Any], -) -> None: +) -> tuple[int, dict[str, Any]]: + """Stage a run stream event inside the caller's transaction. + + Allocates the seq from the locked counter row when ``seq`` is not given + (terminal events committed atomically with the run status) or commits a + pre-assigned seq when it is (idempotent: an existing row is skipped). The + caller is responsible for committing, and must publish to the in-memory + broker only after that commit. + """ if _uses_redis_executor(): logger.warning( "Skipped non-atomic Redis stream append from legacy run session helper", extra={"run_id": run_id, "seq": seq}, ) - return - existing = await session.scalar( - select(RunStreamEvent.id).where(RunStreamEvent.run_id == run_id, RunStreamEvent.seq == seq) - ) - if existing is not None: - return - session.add( - RunStreamEvent( - run_id=run_id, - seq=seq, - event=str(payload.get("event", "message")), - payload_json=dict(payload), + return (seq or 0, payload) + if not _metadata_db_ready(): + # No metadata DB (offline tests / pre-initialization): there is nothing + # durable to stage, so defer to broker-local sequence allocation. + return (seq or 0, payload) + if seq is not None: + existing = await session.scalar( + select(RunStreamEvent.id).where(RunStreamEvent.run_id == run_id, RunStreamEvent.seq == seq) ) - ) + if existing is not None: + return seq, payload + return await _stage_db_event(session, "run", run_id, payload, seq=seq) async def add_thread_stream_event_to_session( session: AsyncSession, thread_id: str, *, - seq: int, + seq: int | None = None, payload: dict[str, Any], -) -> None: +) -> tuple[int, dict[str, Any]]: + """Stage a thread-protocol event inside the caller's transaction (see run twin).""" if _uses_redis_executor(): logger.warning( "Skipped non-atomic Redis stream append from legacy thread session helper", extra={"thread_id": thread_id, "seq": seq}, ) - return - existing = await session.scalar( - select(ThreadStreamEvent.id).where(ThreadStreamEvent.thread_id == thread_id, ThreadStreamEvent.seq == seq) - ) - if existing is not None: - return - session.add( - ThreadStreamEvent( - thread_id=thread_id, - seq=seq, - method=str(payload.get("method", "event")), - payload_json=dict(payload), + return (seq or 0, payload) + if not _metadata_db_ready(): + return (seq or 0, payload) + if seq is not None: + existing = await session.scalar( + select(ThreadStreamEvent.id).where(ThreadStreamEvent.thread_id == thread_id, ThreadStreamEvent.seq == seq) ) - ) + if existing is not None: + return seq, payload + return await _stage_db_event(session, "thread", thread_id, payload, seq=seq) async def load_run_stream_events(run_id: str, *, after_seq: int = 0) -> list[tuple[int, dict[str, Any]]]: @@ -270,6 +523,8 @@ async def load_run_stream_events(run_id: str, *, after_seq: int = 0) -> list[tup async def delete_run_stream_events(run_ids: list[str]) -> None: if not run_ids: return + for run_id in run_ids: + _stream_seed_locks.pop(("run", run_id), None) if _uses_redis_executor(): keys = [key for run_id in run_ids for key in (_run_stream_key(run_id), f"{_RUN_STREAM_SEQ_KEY_PREFIX}:{run_id}")] try: @@ -284,6 +539,11 @@ async def delete_run_stream_events(run_ids: list[str]) -> None: session_factory = db_manager.get_session_factory() async with session_factory() as session: await session.execute(delete(RunStreamEvent).where(RunStreamEvent.run_id.in_(run_ids))) + await session.execute( + delete(StreamSequence).where( + StreamSequence.scope == "run", StreamSequence.scope_id.in_(run_ids) + ) + ) await session.commit() except Exception: return @@ -365,6 +625,7 @@ async def load_thread_stream_events( async def delete_thread_stream_events(thread_id: str) -> None: + _stream_seed_locks.pop(("thread", thread_id), None) if _uses_redis_executor(): try: await _get_redis_client().delete( @@ -381,6 +642,11 @@ async def delete_thread_stream_events(thread_id: str) -> None: session_factory = db_manager.get_session_factory() async with session_factory() as session: await session.execute(delete(ThreadStreamEvent).where(ThreadStreamEvent.thread_id == thread_id)) + await session.execute( + delete(StreamSequence).where( + StreamSequence.scope == "thread", StreamSequence.scope_id == thread_id + ) + ) await session.commit() except Exception: return diff --git a/src/agentseek_api/services/thread_protocol.py b/src/agentseek_api/services/thread_protocol.py index 8f99c20..e8b99be 100644 --- a/src/agentseek_api/services/thread_protocol.py +++ b/src/agentseek_api/services/thread_protocol.py @@ -102,6 +102,12 @@ def _record_event(self, thread_id: str, payload: dict[str, Any], *, seq: int | N self._mark_active(thread_id) if seq is None: seq = self._next_seq[thread_id] + else: + # Never regress below the in-memory watermark: an explicit seq from + # persistent state may be lower than what this process has already + # handed out (e.g. the broker was cleared and re-seeded from the DB + # mid-run). Clamping here keeps the wire seq monotonic. + seq = max(seq, self._next_seq[thread_id]) self._next_seq[thread_id] = max(self._next_seq[thread_id], seq + 1) event = { "type": "event", @@ -119,7 +125,7 @@ def publish( thread_id: str, payload: dict[str, Any], *, - persist: bool = True, + persist: bool = False, seq: int | None = None, ) -> dict[str, Any]: event = self._record_event(thread_id, payload, seq=seq) @@ -209,9 +215,48 @@ async def stream( thread_protocol_broker = ThreadProtocolEventBroker() +async def _persist_protocol_to_run_stream(run_id: str, payload: dict[str, Any]) -> None: + """Append a protocol event to the run's ordered stream. + + Keeps the run stream (the source for ``GET /runs/{id}/stream``) as a single + run-scoped, monotonically sequenced log of both lifecycle records and + protocol frames, so the replay cursor is one domain instead of mixing run + and thread sequence spaces. Durable before expose: the event row (and its + seq) is committed first; the in-memory broker is only updated afterwards. + """ + if settings.EXECUTOR_BACKEND.strip().lower() == "redis": + from agentseek_api.services.stream_persistence import append_redis_run_stream_event + + try: + await append_redis_run_stream_event(run_id, payload) + except Exception: + logger.warning( + "Failed to atomically append Redis run stream event", + extra={"run_id": run_id}, + exc_info=True, + ) + return + from agentseek_api.services.run_state import run_broker + from agentseek_api.services.stream_persistence import append_run_stream_event_atomic + + seq, _ = await append_run_stream_event_atomic(run_id, payload) + run_broker.publish_protocol(run_id, payload, seq=seq) + + async def _apublish_thread_event(thread_id: str, payload: dict[str, Any]) -> dict[str, Any]: + run_id = (payload.get("params") or {}).get("run_id") + if run_id: + await _persist_protocol_to_run_stream(run_id, payload) if settings.EXECUTOR_BACKEND.strip().lower() != "redis": - return await thread_protocol_broker.apublish(thread_id, payload) + # Durable before expose: the thread event row (and its seq) is appended + # atomically first; the broker only records it after the append + # succeeded, so a client can never receive a seq that was not durably + # committed. ``_record_event`` never regresses below its in-memory + # watermark, keeping the wire seq monotonic across a cold broker. + from agentseek_api.services.stream_persistence import append_thread_stream_event_atomic + + seq, _ = await append_thread_stream_event_atomic(thread_id, payload) + return thread_protocol_broker.publish(thread_id, payload, persist=False, seq=seq) try: from agentseek_api.services.stream_persistence import append_redis_thread_stream_event @@ -233,7 +278,7 @@ def publish_lifecycle_event( graph_name: str | None = None, error: str | None = None, namespace: list[str] | None = None, - persist: bool = True, + persist: bool = False, seq: int | None = None, ) -> dict[str, Any]: data: dict[str, Any] = {"event": event} @@ -324,6 +369,7 @@ async def apublish_tool_event( output_payload: Any | None = None, error_message: str | None = None, namespace: list[str] | None = None, + run_id: str | None = None, ) -> dict[str, Any]: data: dict[str, Any] = {"event": tool_event, "tool_call_id": tool_call_id} if tool_name is not None: @@ -341,6 +387,8 @@ async def apublish_tool_event( } if node is not None: params["node"] = node + if run_id is not None: + params["run_id"] = run_id return await _apublish_thread_event(thread_id, {"method": "tools", "params": params}) @@ -350,6 +398,7 @@ def publish_values_event( values: Any, namespace: list[str] | None = None, run_id: str | None = None, + persist: bool = False, ) -> dict[str, Any]: params: dict[str, Any] = { "namespace": namespace or [], @@ -364,6 +413,7 @@ def publish_values_event( "method": "values", "params": params, }, + persist=persist, ) diff --git a/tests/conftest.py b/tests/conftest.py index 6043512..a4e057c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,6 @@ +import asyncio +import json +import time from pathlib import Path from typing import Any from collections.abc import Awaitable, Callable @@ -42,6 +45,37 @@ async def submit(self, job: Callable[[], Awaitable[None]] | RunExecutionJob) -> ) +class BackgroundInlineExecutor: + """Submits run jobs as background tasks so a run stays genuinely active + after the HTTP call that created it returns. + + The default ``InlineExecutor`` awaits the job inline, so ``POST /runs`` + returns only once the run is already terminal. That makes it structurally + impossible to observe a run mid-flight. This executor is used by the + mid-run reconnect regression tests to prove exactly-once semantics over a + disconnect/reconnect on ``GET /runs/{id}/stream``. + """ + + async def submit(self, job: Callable[[], Awaitable[None]] | RunExecutionJob) -> None: + if callable(job): + asyncio.create_task(job()) + return + from agentseek_api.services.run_preparation import _execute_and_persist + + asyncio.create_task( + _execute_and_persist( + run_id=job.run_id, + thread_id=job.thread_id, + user_id=job.user_id, + payload=job.payload, + graph_id=job.graph_id, + kwargs=job.kwargs, + resume=job.resume, + is_resume=job.is_resume, + ) + ) + + async def header_user_override(request: Request) -> User: identity = request.headers.get("x-user-id", "default_user") return User(identity=identity, is_authenticated=True) @@ -51,6 +85,124 @@ async def _noop_ensure_default_assistants() -> None: return None +async def _collect_sse_frames( + client, + url: str, + *, + headers: dict[str, str] | None = None, + max_frames: int | None = None, +) -> list[tuple[int | None, str, dict[str, object]]]: + """Collect SSE frames from a live stream, optionally stopping early to + simulate a client disconnecting mid-run.""" + frames: list[tuple[int | None, str, dict[str, object]]] = [] + current_id: int | None = None + current_event = "" + current_data: list[str] = [] + async with client.stream("GET", url, headers=headers) as response: + assert response.status_code == 200, response.text + buffer = "" + async for chunk in response.aiter_text(): + buffer += chunk + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + line = line.rstrip("\r") + if line.startswith("id: "): + current_id = int(line[len("id: "):].strip()) + elif line.startswith("event: "): + current_event = line[len("event: "):].strip() + elif line.startswith("data: "): + current_data.append(line[len("data: "):].strip()) + elif line == "" and current_data: + payload = json.loads("".join(current_data)) + frames.append((current_id, current_event, payload)) + current_id, current_event, current_data = None, "", [] + if max_frames is not None and len(frames) >= max_frames: + return frames + return frames + + +async def _wait_run_terminal(client, thread_id: str, run_id: str, *, timeout_seconds: float = 20.0) -> str: + deadline = time.time() + timeout_seconds + while time.time() < deadline: + status = (await client.get(f"/threads/{thread_id}/runs/{run_id}")).json()["status"] + if status in ("success", "error", "interrupted"): + return status + await asyncio.sleep(0.2) + raise AssertionError(f"run {run_id} did not reach a terminal status within {timeout_seconds}s") + + +async def _midrun_reconnect_flow(app) -> None: + """Connect to GET /runs/{id}/stream while the run is still executing, + disconnect after a few frames, wait for the run to finish, then reconnect + with Last-Event-ID and assert exactly-once delivery. Shared by the inline + and Redis HTTP-level reconnect regression tests.""" + import httpx + from httpx import ASGITransport + + async with httpx.AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + async with app.router.lifespan_context(app): + await _midrun_reconnect_checks(client) + + +async def _midrun_reconnect_checks(client) -> None: + assistant = await client.post("/assistants", json={"name": "midrun", "graph_id": "stress_tool_agent"}) + assert assistant.status_code == 200, assistant.text + assistant_id = assistant.json()["assistant_id"] + + thread = await client.post("/threads", json={"metadata": {"case": "midrun"}}) + assert thread.status_code == 200, thread.text + thread_id = thread.json()["thread_id"] + + # ~4.5s total runtime (3 steps * 1.5s) guarantees a mid-run window. + run = await client.post( + f"/threads/{thread_id}/runs", + json={"assistant_id": assistant_id, "input": {"delay": 1.5, "steps": 3}}, + ) + assert run.status_code == 200, run.text + run_id = run.json()["run_id"] + + stream_url = f"/threads/{thread_id}/runs/{run_id}/stream" + + # Phase 1: connect mid-run, collect a few frames, then disconnect. + await asyncio.sleep(0.6) + phase1 = await _collect_sse_frames(client, stream_url, max_frames=3) + if not phase1: + await asyncio.sleep(1.0) + phase1 = await _collect_sse_frames(client, stream_url, max_frames=3) + assert phase1, "expected to observe the run while it is still active" + phase1_ids = [frame_id for frame_id, _, _ in phase1 if frame_id is not None] + last_id = phase1_ids[-1] + assert last_id is not None + + status = await _wait_run_terminal(client, thread_id, run_id) + assert status == "success", status + + # Phase 2: reconnect with Last-Event-ID. + phase2 = await _collect_sse_frames(client, stream_url, headers={"Last-Event-ID": str(last_id)}) + phase2_ids = [frame_id for frame_id, _, _ in phase2 if frame_id is not None] + + # Exactly-once: reconnect must not replay frames already delivered, and + # must deliver every frame produced after the disconnect. + assert all(frame_id > last_id for frame_id in phase2_ids), ( + f"reconnect replayed an already-delivered id: phase1={phase1_ids} phase2={phase2_ids}" + ) + all_ids = phase1_ids + phase2_ids + assert len(all_ids) == len(set(all_ids)), f"duplicate ids delivered across reconnect: {all_ids}" + + fresh = await _collect_sse_frames(client, stream_url) + fresh_ids = [frame_id for frame_id, _, _ in fresh if frame_id is not None] + assert [frame_id for frame_id in fresh_ids if frame_id > last_id] == phase2_ids, ( + f"reconnect content mismatch: full={fresh_ids} phase2={phase2_ids}" + ) + + +@pytest.fixture +def midrun_reconnect_flow() -> Callable[[object], Awaitable[None]]: + """Async helper that drives a mid-run disconnect/reconnect against an app + built by ``midrun_app_factory``, asserting exactly-once delivery.""" + return _midrun_reconnect_flow + + @pytest.fixture def client(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> TestClient: from agentseek_api.core import auth_middleware @@ -71,3 +223,51 @@ def client(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> TestClient: with TestClient(app) as test_client: yield test_client auth_middleware._backend = None + + +@pytest.fixture +def midrun_app_factory( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> Callable[..., object]: + """Build a fresh app whose runs execute in the background. + + Unlike the default ``client`` fixture (which awaits each run to completion + before ``POST /runs`` returns), apps from this factory keep a run active + after the create call, so a test can connect to ``GET /runs/{id}/stream`` + mid-flight, disconnect, and reconnect with ``Last-Event-ID``. + + ``executor_backend`` selects the run-stream storage path ("inline" or + "redis"); ``redis_client`` (used only for the redis path) must be an object + with the ``eval``/``xrange`` surface of the stream persistence layer. + """ + from agentseek_api.core import auth_middleware + + def _make(*, executor_backend: str = "inline", redis_client: object | None = None) -> object: + from agentseek_api.core import database as database_module + from agentseek_api.services import stream_persistence as stream_module + + monkeypatch.setattr(database_module, "OceanBaseCheckpointSaver", FakeCheckpointer) + monkeypatch.setattr("agentseek_api.main.ensure_default_assistants", _noop_ensure_default_assistants) + monkeypatch.setattr(settings, "SEEKDB_URL", f"sqlite+aiosqlite:///{tmp_path}/midrun-{executor_backend}.db") + monkeypatch.setattr(settings, "AUTH_MODULE_PATH", None) + monkeypatch.setattr("agentseek_api.services.executor._executor", None) + monkeypatch.setattr(settings, "EXECUTOR_BACKEND", executor_backend) + monkeypatch.setattr( + "agentseek_api.core.auth_middleware.get_config_auth_settings", + lambda: auth_middleware.ConfigAuthSettings(), + ) + auth_middleware._backend = None + if redis_client is not None: + monkeypatch.setattr(stream_module, "_redis_client", redis_client) + from agentseek_api.api import runs as runs_api + + monkeypatch.setattr(runs_api, "REDIS_STREAM_POLL_INTERVAL_SECONDS", 0) + + monkeypatch.setattr("agentseek_api.services.run_preparation.get_executor", lambda: BackgroundInlineExecutor()) + + app = create_app() + app.dependency_overrides[get_current_user] = header_user_override + return app + + return _make diff --git a/tests/e2e/test_live_provider_api.py b/tests/e2e/test_live_provider_api.py index 7f7bc33..c480d9e 100644 --- a/tests/e2e/test_live_provider_api.py +++ b/tests/e2e/test_live_provider_api.py @@ -124,19 +124,28 @@ async def test_live_provider_streaming_http_flow(live_provider_base_url: str) -> for line in stream.text.splitlines() if line.startswith("data: ") ] - message_chunks = [ + # The default run-stream replay returns the run's persisted protocol + # events. A run created without an explicit stream_mode does not stream + # incremental LLM tokens (no ``messages`` channel events), so the final + # answer is read back from the last state snapshot (``values`` payload) + # instead of token deltas. + state_payloads = [ payload for payload in payloads - if payload.get("event") == "message_chunk" - and payload.get("langgraph_event") in {"on_chat_model_stream", "on_llm_stream"} - and _text_from_content(payload.get("content")).strip() + if isinstance(payload, dict) and isinstance(payload.get("messages"), list) ] + assert state_payloads + final_messages = state_payloads[-1]["messages"] + final_ai = next( + (m for m in reversed(final_messages) if isinstance(m, dict) and m.get("type") == "ai"), + None, + ) + assert final_ai is not None assert payloads[0]["event"] == "start" assert "event: start" in stream.text assert "event: end" in stream.text - assert len(message_chunks) >= 2 - assert _normalize_text("".join(_text_from_content(payload.get("content")) for payload in message_chunks)) == _normalize_text( + assert _normalize_text(str(final_ai.get("content", ""))) == _normalize_text( waited_body["output"]["final_text"] ) @@ -246,7 +255,7 @@ async def test_live_provider_store_endpoints_and_graph(live_provider_base_url: s json={"namespace": namespace, "key": "profile", "value": {"kind": "profile", "name": "Ada"}}, headers=user_headers(user_id), ) - assert created.status_code == 200 + assert created.status_code in (200, 204) fetched = await client.get( "/store/items", @@ -321,7 +330,7 @@ async def test_live_provider_store_ttl_expires_items_on_mysql_family_backend(liv }, headers=user_headers(user_id), ) - assert created.status_code == 200 + assert created.status_code in (200, 204) immediate = await client.get( "/store/items", @@ -377,7 +386,7 @@ async def test_live_provider_hitl_rest_and_protocol_resume(live_provider_base_ur run_id = run.json()["run_id"] interrupted_stream = await client.post( - f"/threads/{thread_id}/stream", + f"/threads/{thread_id}/stream/events", json={"channels": ["lifecycle", "input", "values"]}, headers=user_headers(user_id), ) @@ -415,7 +424,7 @@ async def test_live_provider_hitl_rest_and_protocol_resume(live_provider_base_ur async with client.stream( "POST", - f"/threads/{thread_id}/stream", + f"/threads/{thread_id}/stream/events", json={"channels": ["lifecycle", "values"], "since": last_seq}, headers=user_headers(user_id), ) as resumed_stream: @@ -455,3 +464,78 @@ async def test_live_provider_hitl_rest_and_protocol_resume(live_provider_base_ur create_time_input = next(event for event in create_time_events if event["event"] == "input.requested") assert create_time_input["data"]["payload"] == "Provide value:" assert create_time_input["data"]["interrupt_id"] + + +@pytest.mark.e2e +@pytest.mark.asyncio +async def test_live_provider_realtime_stream_incremental_messages_partial(live_provider_base_url: str) -> None: + """Real-time messages stream emits incremental ``messages/partial`` tokens. + + The messages wire contract delivers ``messages/metadata`` + ``messages/partial`` + only (the protocol-v2 block stream is suppressed for this contract), and the + partials accumulate to the final AI answer. This closes the gap where only the + replay endpoint (final values snapshot) and the unit-level block publishing + were covered. + """ + user_id = f"provider-realtime-partial-{uuid4().hex}" + + async with httpx.AsyncClient(base_url=live_provider_base_url, timeout=90.0, trust_env=False) as client: + assistant = await client.post( + "/assistants", + json={"name": "live-provider-realtime-partial", "graph_id": provider_graph_id("stream")}, + ) + assert assistant.status_code == 200, assistant.text + assistant_id = assistant.json()["assistant_id"] + + thread = await client.post( + "/threads", + json={"metadata": {"suite": "live-provider-realtime-partial"}}, + headers=user_headers(user_id), + ) + assert thread.status_code == 200 + thread_id = thread.json()["thread_id"] + + streamed_create = await client.post( + f"/threads/{thread_id}/runs/stream", + json={ + "assistant_id": assistant_id, + "input": { + "message": ( + "Reply with one short sentence about real-time token streaming, " + "using at least twenty words." + ) + }, + "stream_mode": "messages", + }, + headers=user_headers(user_id), + ) + assert streamed_create.status_code == 200, streamed_create.text + assert streamed_create.headers["content-type"].startswith("text/event-stream") + + events = parse_sse_events(streamed_create.text) + event_names = [event["event"] for event in events] + assert event_names[0] == "metadata" + assert "messages/partial" in event_names + assert not any("content-block" in name for name in event_names) + + partial_payloads = [event["data"] for event in events if event["event"] == "messages/partial"] + assert len(partial_payloads) >= 2, "expected at least one incremental token step" + + # The last messages/partial payload carries the fully accumulated message. + # Wire data is the message list directly (official messages/partial format). + accumulated_text = "" + for payload in partial_payloads: + if not isinstance(payload, list) or not payload: + continue + last = payload[-1] + if not isinstance(last, dict) or last.get("type") != "ai": + continue + accumulated_text = _text_from_content(last.get("content")) + + assert accumulated_text, "messages/partial never carried an AI message" + + run_id = str(next(event["data"]["run_id"] for event in events if event["event"] == "metadata")) + awaited = await _poll_run(client=client, thread_id=thread_id, run_id=run_id, user_id=user_id) + assert awaited["status"] == "success" + final_text = str(awaited["output"]["final_text"]) + assert _normalize_text(accumulated_text) == _normalize_text(final_text) diff --git a/tests/integration/test_live_provider_streaming.py b/tests/integration/test_live_provider_streaming.py index 443a46f..08943ef 100644 --- a/tests/integration/test_live_provider_streaming.py +++ b/tests/integration/test_live_provider_streaming.py @@ -33,6 +33,24 @@ def _normalize_text(text: str) -> str: return " ".join(text.split()) +def _parse_sse_events(stream_text: str) -> list[dict[str, object]]: + events: list[dict[str, object]] = [] + for chunk in stream_text.strip().split("\n\n"): + if not chunk.strip(): + continue + event: dict[str, object] = {} + for line in chunk.splitlines(): + if line.startswith("id: "): + event["id"] = line.removeprefix("id: ") + elif line.startswith("event: "): + event["event"] = line.removeprefix("event: ") + elif line.startswith("data: "): + event["data"] = json.loads(line.removeprefix("data: ")) + if event: + events.append(event) + return events + + class FakeCheckpointer: def __init__(self, connection_args: dict[str, str]) -> None: self.connection_args = connection_args @@ -149,22 +167,108 @@ def test_live_provider_stream_emits_multiple_message_chunks(live_provider_client for line in stream.text.splitlines() if line.startswith("data: ") ] - message_chunks = [ + # The default run-stream replay returns the run's persisted protocol + # events. A run created without an explicit stream_mode does not stream + # incremental LLM tokens (no ``messages`` channel events), so the final + # answer is read back from the last state snapshot (``values`` payload) + # instead of token deltas. + state_payloads = [ payload for payload in payloads - if payload.get("event") == "message_chunk" - and payload.get("langgraph_event") in {"on_chat_model_stream", "on_llm_stream"} - and _text_from_content(payload.get("content")).strip() + if isinstance(payload, dict) and isinstance(payload.get("messages"), list) ] - chunk_texts = [_text_from_content(payload.get("content")) for payload in message_chunks] + assert state_payloads + final_messages = state_payloads[-1]["messages"] + final_ai = next( + (m for m in reversed(final_messages) if isinstance(m, dict) and m.get("type") == "ai"), + None, + ) + assert final_ai is not None end_payloads = [payload for payload in payloads if payload.get("event") == "end"] assert payloads[0]["event"] == "start" assert "event: start" in stream.text assert "event: end" in stream.text - assert any(payload.get("event") == "node_start" and payload.get("node") == "call_model" for payload in payloads) - assert any(payload.get("event") == "node_end" and payload.get("node") == "call_model" for payload in payloads) assert end_payloads[-1]["status"] == "success" assert end_payloads[-1]["run_id"] == run_id - assert len(message_chunks) >= 2 - assert _normalize_text("".join(chunk_texts)) == _normalize_text(waited_body["output"]["final_text"]) + assert _normalize_text(str(final_ai.get("content", ""))) == _normalize_text( + waited_body["output"]["final_text"] + ) + + # Token-level proof: an explicit ``stream_mode=messages`` run must surface + # real incremental ``messages/partial`` frames from the provider that + # accumulate to the final answer. This restores the incremental token + # assertion the manual live-provider workflow is contractually expected to + # prove (the default replay above only proves the final snapshot). + streamed = live_provider_client.post( + f"/threads/{thread_id}/runs/stream", + json={ + "assistant_id": assistant_id, + "input": { + "message": ( + "Explain why token-level streaming verification matters in exactly two sentences, " + "using at least forty words and no bullet points." + ) + }, + "stream_mode": "messages", + }, + ) + assert streamed.status_code == 200, streamed.text + assert streamed.headers["content-type"].startswith("text/event-stream") + streamed_events = _parse_sse_events(streamed.text) + streamed_names = [event["event"] for event in streamed_events] + assert streamed_names[0] == "metadata" + assert "messages/partial" in streamed_names + assert not any("content-block" in name for name in streamed_names) + + # The v1 messages wire contract emits one ``messages/metadata`` identity per + # streamed message (``{message_id: {"metadata": ...}}``) before its partials, + # so the client can route the incremental frames to the right message. A + # real provider run must surface at least one such identity. + metadata_events = [event for event in streamed_events if event["event"] == "messages/metadata"] + assert metadata_events, "expected messages/metadata identity events in the messages stream" + metadata_ids = [ + message_id + for event in metadata_events + for message_id in event["data"].keys() + if isinstance(event["data"], dict) + ] + assert metadata_ids, "messages/metadata carried no message identity" + assert all( + isinstance(event["data"][message_id], dict) and "metadata" in event["data"][message_id] + for event in metadata_events + for message_id in event["data"].keys() + if isinstance(event["data"], dict) + ), "messages/metadata payload must be {message_id: {'metadata': {...}}}" + + partial_payloads = [event["data"] for event in streamed_events if event["event"] == "messages/partial"] + assert len(partial_payloads) >= 2, "expected at least one incremental token step" + accumulated_text = "" + for payload in partial_payloads: + if not isinstance(payload, list) or not payload: + continue + last = payload[-1] + if not isinstance(last, dict) or last.get("type") != "ai": + continue + accumulated_text = _text_from_content(last.get("content")) + assert accumulated_text, "messages/partial never carried an AI message" + # The metadata identity must correspond to the AI message being streamed: + # the accumulated partial is an AI message, and its id is the metadata key. + partial_ids = { + str(message.get("id")) + for payload in partial_payloads + if isinstance(payload, list) + for message in payload + if isinstance(message, dict) and message.get("id") + } + assert metadata_ids and partial_ids and metadata_ids[0] in partial_ids, ( + f"messages/metadata identity {metadata_ids} must match streamed partial ids {partial_ids}" + ) + streamed_run_id = str( + next(event["data"]["run_id"] for event in streamed_events if event["event"] == "metadata") + ) + streamed_waited = live_provider_client.get(f"/threads/{thread_id}/runs/{streamed_run_id}/wait") + assert streamed_waited.status_code == 200 + streamed_body = streamed_waited.json() + assert streamed_body["status"] == "success", streamed_body.get("last_error") + assert _normalize_text(accumulated_text) == _normalize_text(streamed_body["output"]["final_text"]) diff --git a/tests/integration/test_protocol_v2_streaming.py b/tests/integration/test_protocol_v2_streaming.py index 0f5dff8..4695265 100644 --- a/tests/integration/test_protocol_v2_streaming.py +++ b/tests/integration/test_protocol_v2_streaming.py @@ -288,6 +288,8 @@ def test_protocol_stream_filters_subgraph_namespace_events(client: TestClient) - "params": { "assistant_id": assistant_id, "input": {"foo": "hello "}, + "stream_mode": ["updates"], + "stream_subgraphs": True, }, }, ) @@ -313,3 +315,36 @@ def test_protocol_stream_filters_subgraph_namespace_events(client: TestClient) - assert events assert {event["event"] for event in events} <= {"updates", "input.requested"} assert all(event["data"]["params"]["namespace"][:1] == namespace_prefix for event in events) + + +def test_protocol_run_start_invalid_stream_mode_returns_400(client: TestClient) -> None: + """An invalid ``stream_mode`` on a protocol ``run.start`` command is a client + error (400 ``invalid_argument``), not a missing-resource 404 or a validation + 422. 404 stays reserved for unknown assistants/graphs (covered elsewhere).""" + assistant = client.post("/assistants", json={"name": "protocol-bad-mode", "graph_id": "react_agent"}) + assert assistant.status_code == 200 + assistant_id = assistant.json()["assistant_id"] + + thread = client.post("/threads", json={"metadata": {"case": "protocol-bad-mode"}}) + assert thread.status_code == 200 + thread_id = thread.json()["thread_id"] + + command = client.post( + f"/threads/{thread_id}/commands", + json={ + "id": 1, + "method": "run.start", + "params": { + "assistant_id": assistant_id, + "input": {"message": "bad stream mode"}, + "stream_mode": "nonexistent", + }, + }, + ) + + assert command.status_code == 400 + body = command.json() + assert body["type"] == "error" + assert body["id"] == 1 + assert body["error"] == "invalid_argument" + assert "nonexistent" in body["message"] diff --git a/tests/integration/test_runs_compat.py b/tests/integration/test_runs_compat.py index 3b0cccb..004e91a 100644 --- a/tests/integration/test_runs_compat.py +++ b/tests/integration/test_runs_compat.py @@ -882,3 +882,32 @@ def test_create_run_compat_openapi_documents_wait_and_stream_routes() -> None: assert "stream_resumable" in stateless_schema["properties"] assert "feedback_keys" in stateless_schema["properties"] assert "durability" in stateless_schema["properties"] + + +def test_create_run_stream_events_mode_emits_raw_astream_events(client: TestClient) -> None: + """HTTP-level regression for ``stream_mode=events``. + + The events channel must surface each raw ``astream_events()`` item (SSE + ``event: events`` with the raw event discriminator in ``data.event``), not + just the translated side effects. + """ + assistant_id = _create_assistant(client, graph_id="stress_test") + thread_id = _create_thread(client) + + response = client.post( + f"/threads/{thread_id}/runs/stream", + json={"assistant_id": assistant_id, "input": {"delay": 0.0, "steps": 1}, "stream_mode": "events"}, + ) + + assert response.status_code == 200 + events = _parse_sse_events(response.text) + raw_events = [event for event in events if event["event"] == "events"] + assert raw_events, "expected raw astream_events() items on the events channel" + raw_names = { + event["data"].get("event") + for event in raw_events + if isinstance(event["data"], dict) + } + assert any(name in raw_names for name in ("on_chain_start", "on_chain_stream")), ( + f"expected raw on_chain_* events, got: {raw_names}" + ) diff --git a/tests/integration/test_runs_streaming.py b/tests/integration/test_runs_streaming.py index 423dbbc..930817c 100644 --- a/tests/integration/test_runs_streaming.py +++ b/tests/integration/test_runs_streaming.py @@ -1,6 +1,10 @@ -from fastapi.testclient import TestClient +import asyncio import json +from fastapi.testclient import TestClient + +from agentseek_api.services.run_state import run_broker + def _stream_payloads(stream_text: str) -> list[dict[str, object]]: return [ @@ -53,10 +57,10 @@ def test_react_agent_stream_includes_tool_and_message_events(client: TestClient) assert stream_response.status_code == 200 payloads = _stream_payloads(stream_response.text) - assert any(payload["event"] == "tool_start" and payload["name"] == "lookup" for payload in payloads) - assert any(payload["event"] == "tool_end" and payload["name"] == "lookup" for payload in payloads) + assert any(isinstance(payload, dict) and payload.get("event") == "tool-started" and payload.get("tool_name") == "lookup" for payload in payloads) + assert any(isinstance(payload, dict) and payload.get("event") == "tool-finished" and payload.get("tool_name") == "lookup" for payload in payloads) assert any( - payload["event"] == "message_chunk" and "Final answer:" in str(payload.get("content", "")) + isinstance(payload, dict) and "Final answer:" in json.dumps(payload, ensure_ascii=False) for payload in payloads ) @@ -81,12 +85,16 @@ def test_stress_tool_agent_stream_includes_multiple_tool_cycles(client: TestClie assert stream_response.status_code == 200 payloads = _stream_payloads(stream_response.text) - tool_starts = [payload for payload in payloads if payload["event"] == "tool_start" and payload["name"] == "slow_process"] - tool_ends = [payload for payload in payloads if payload["event"] == "tool_end" and payload["name"] == "slow_process"] + tool_starts = [ + payload for payload in payloads if isinstance(payload, dict) and payload.get("event") == "tool-started" and payload.get("tool_name") == "slow_process" + ] + tool_ends = [ + payload for payload in payloads if isinstance(payload, dict) and payload.get("event") == "tool-finished" and payload.get("tool_name") == "slow_process" + ] assert len(tool_starts) == 3 assert len(tool_ends) == 3 assert any( - payload["event"] == "message_chunk" and '"steps_completed": 3' in str(payload.get("content", "")) + isinstance(payload, dict) and "steps_completed" in json.dumps(payload, ensure_ascii=False) for payload in payloads ) @@ -138,9 +146,18 @@ def test_resumed_run_stream_preserves_each_terminal_status(client: TestClient) - stream_response = client.get(f"/threads/{thread_id}/runs/{run_id}/stream") assert stream_response.status_code == 200 payloads = _stream_payloads(stream_response.text) - end_statuses = [payload["status"] for payload in payloads if payload["event"] == "end"] + end_statuses = [payload["status"] for payload in payloads if payload.get("event") == "end"] assert end_statuses == ["interrupted", "success"] + # Every SSE id across the whole (interrupted + resumed) log must be + # strictly monotonic. The historical interrupted run's terminal "end" must + # keep its original seq and not be deferred past the resumed run's frames + # (a resumed stream previously re-ordered the earlier end after newer + # frames, producing a non-monotonic cursor like 1..9, 11..17, 10, 18). + ids = _sse_ids(stream_response.text) + assert ids == sorted(ids), f"resumed run stream ids not monotonic: {ids}" + assert len(set(ids)) == len(ids), f"resumed run stream ids not unique: {ids}" + def test_create_run_rejects_configurable_and_context_together(client: TestClient) -> None: assistant = client.post("/assistants", json={"name": "reject-both", "graph_id": "default"}) @@ -298,4 +315,145 @@ def test_run_client_config_overrides_assistant_config_defaults(client: TestClien config = kwargs.get("config") or {} configurable = config.get("configurable") or {} assert configurable.get("client_param") == "assistant-default" - assert configurable.get("model") == "client-model" \ No newline at end of file + assert configurable.get("model") == "client-model" + + +def _sse_ids(stream_text: str) -> list[int]: + ids: list[int] = [] + for line in stream_text.splitlines(): + if line.startswith("id: "): + ids.append(int(line[len("id: "):].strip())) + return ids + + +def test_run_stream_sse_ids_are_monotonic(client: TestClient) -> None: + """Every SSE ``id`` in the default run stream shares one monotonic cursor. + + The replay merges run-scoped lifecycle events (start/end) with + thread-protocol events that carry their own independent sequence domains. + A non-monotonic cursor (e.g. ``1, 3, ...25, 2``) breaks Last-Event-ID + resume, so all emitted ids must be strictly increasing. + """ + assistant = client.post("/assistants", json={"name": "streaming-monotonic", "graph_id": "react_agent"}) + assert assistant.status_code == 200 + assistant_id = assistant.json()["assistant_id"] + + thread = client.post("/threads", json={"metadata": {"case": "monotonic"}}) + assert thread.status_code == 200 + thread_id = thread.json()["thread_id"] + + run = client.post( + f"/threads/{thread_id}/runs", + json={"assistant_id": assistant_id, "input": {"message": "stream"}}, + ) + assert run.status_code == 200 + run_id = run.json()["run_id"] + + stream_response = client.get(f"/threads/{thread_id}/runs/{run_id}/stream") + assert stream_response.status_code == 200 + ids = _sse_ids(stream_response.text) + assert len(ids) >= 2, "expected at least start + end frames" + assert ids == sorted(ids), f"SSE ids are not monotonic: {ids}" + assert len(set(ids)) == len(ids), f"SSE ids are not unique: {ids}" + + +def test_run_stream_resume_after_terminal_end_does_not_replay(client: TestClient) -> None: + """Reconnecting with the terminal frame's Last-Event-ID must not replay + already-delivered events.""" + assistant = client.post("/assistants", json={"name": "streaming-resume", "graph_id": "react_agent"}) + assert assistant.status_code == 200 + assistant_id = assistant.json()["assistant_id"] + + thread = client.post("/threads", json={"metadata": {"case": "resume"}}) + assert thread.status_code == 200 + thread_id = thread.json()["thread_id"] + + run = client.post( + f"/threads/{thread_id}/runs", + json={"assistant_id": assistant_id, "input": {"message": "stream"}}, + ) + assert run.status_code == 200 + run_id = run.json()["run_id"] + + full = client.get(f"/threads/{thread_id}/runs/{run_id}/stream") + assert full.status_code == 200 + ids = _sse_ids(full.text) + assert ids, "stream produced no id frames" + last_id = ids[-1] + + resumed = client.get( + f"/threads/{thread_id}/runs/{run_id}/stream", + headers={"Last-Event-ID": str(last_id)}, + ) + assert resumed.status_code == 200 + resumed_ids = _sse_ids(resumed.text) + assert resumed_ids == [], f"resume after terminal end should replay nothing, got: {resumed_ids}" + + +def test_run_stream_midrun_reconnect_is_exactly_once(midrun_app_factory, midrun_reconnect_flow) -> None: + """HTTP-level regression: a client that connects mid-run, disconnects, and + reconnects with Last-Event-ID must not replay delivered frames and must not + lose frames produced while disconnected (inline executor path).""" + app = midrun_app_factory(executor_backend="inline") + asyncio.run(midrun_reconnect_flow(app)) + + +def test_run_stream_cold_broker_resume_keeps_monotonic_ids(client: TestClient, monkeypatch) -> None: + """A resume after the in-memory broker state is cleared must keep allocating + from the persisted seq watermark, not restart at 1. + + Regression for a process-local allocation bug: clearing the broker's event + state *and* ``_next_seq`` before resuming a persisted run caused new frames + to reuse seqs already in the DB (e.g. ``1..7, 9, 8, 10``) and returned + contradictory terminal statuses. Lifecycle and protocol publication must + share one persistent run-scoped sequence. + """ + assistant = client.post("/assistants", json={"name": "cold-broker", "graph_id": "subgraph_hitl_agent"}) + assert assistant.status_code == 200 + assistant_id = assistant.json()["assistant_id"] + + thread = client.post("/threads", json={"metadata": {"case": "cold-broker"}}) + assert thread.status_code == 200 + thread_id = thread.json()["thread_id"] + + run = client.post( + f"/threads/{thread_id}/runs", + json={"assistant_id": assistant_id, "input": {"foo": "hello "}}, + ) + assert run.status_code == 200 + run_id = run.json()["run_id"] + + # First run interrupted; its protocol frames + end are persisted. + waited = client.get(f"/threads/{thread_id}/runs/{run_id}/wait") + assert waited.json()["status"] == "interrupted" + + # Simulate a cold broker (e.g. process restart): drop every in-memory + # trace of this run, including the seq counter. + run_broker._events.pop(run_id, None) + run_broker._seqs.pop(run_id, None) + run_broker._signals.pop(run_id, None) + run_broker._next_seq.pop(run_id, None) + run_broker._completed_runs.discard(run_id) + try: + run_broker._completed_order.remove(run_id) + except ValueError: + pass + + resumed = client.post( + f"/threads/{thread_id}/runs/{run_id}/resume", + json={"resume": "world"}, + ) + assert resumed.status_code == 200 + + stream_response = client.get(f"/threads/{thread_id}/runs/{run_id}/stream") + assert stream_response.status_code == 200 + ids = _sse_ids(stream_response.text) + assert ids, "expected frames across the interrupted + resumed log" + # ids must be strictly increasing and unique (no reuse of persisted seqs). + assert ids == sorted(ids), f"cold-broker resume ids not monotonic: {ids}" + assert len(set(ids)) == len(ids), f"cold-broker resume ids not unique: {ids}" + + # Both runs' terminal statuses must be present and ordered by seq. + payloads = _stream_payloads(stream_response.text) + end_statuses = [payload["status"] for payload in payloads if payload.get("event") == "end"] + assert end_statuses == ["interrupted", "success"] \ No newline at end of file diff --git a/tests/integration/test_stream_persistence.py b/tests/integration/test_stream_persistence.py index 604ffb3..ec84026 100644 --- a/tests/integration/test_stream_persistence.py +++ b/tests/integration/test_stream_persistence.py @@ -702,7 +702,10 @@ def test_protocol_events_are_persisted_when_published(client: TestClient) -> Non assert thread.status_code == 200 thread_id = thread.json()["thread_id"] - publish_values_event(thread_id, values={"early": True}) + # Explicitly request persistence: the synchronous publish helpers default + # to ``persist=False`` (they exist for tests/back-compat and must not + # silently hit the non-atomic persistence path). + publish_values_event(thread_id, values={"early": True}, persist=True) thread_protocol_broker.delete_thread(thread_id) replay = client.post(f"/threads/{thread_id}/stream/events", json={"channels": ["values"]}) @@ -841,3 +844,59 @@ def test_stream_ignores_malformed_last_event_id(client: TestClient) -> None: headers={"Last-Event-ID": "not-an-int"}, ) assert run_stream_response.status_code == 404 + + +def test_run_stream_midrun_reconnect_is_exactly_once_in_redis(midrun_app_factory, midrun_reconnect_flow) -> None: + """HTTP-level regression for the Redis executor path: a client connecting + to GET /runs/{id}/stream mid-run, disconnecting, then reconnecting with + Last-Event-ID must not replay delivered frames and must not lose frames + produced while disconnected. The Redis stream store is faked with + FakeRedisCounter so the test runs without a live Redis instance.""" + import asyncio + + fake_redis = FakeRedisCounter() + app = midrun_app_factory(executor_backend="redis", redis_client=fake_redis) + asyncio.run(midrun_reconnect_flow(app)) + + +def test_thread_protocol_cold_broker_keeps_monotonic_seq_inline(client: TestClient) -> None: + """A thread-level protocol event published after the in-memory broker state + is cleared must keep allocating from the persisted seq watermark, not + restart at 1. + + Regression for the thread-domain analogue of the run-domain cold-broker + bug: ``next_thread_stream_seq`` used to return None on the inline executor, + so a cleared ``thread_protocol_broker`` reused seqs already persisted to the + DB (e.g. ``1..N, N+1 colliding with an earlier row``), corrupting the + ``since``/``Last-Event-ID`` resume cursor for ``/threads/{id}/stream/events``. + """ + thread_id = "thread-cold-broker-inline" + + # First lifecycle event: persisted as seq 1. + client.portal.call( + lambda: run_jobs_module._publish_lifecycle(thread_id, event="started", graph_name="default") + ) + + # Simulate a cold broker: drop every in-memory trace of this thread, + # including the seq counter. + thread_protocol_broker._events.pop(thread_id, None) + thread_protocol_broker._signals.pop(thread_id, None) + thread_protocol_broker._next_seq.pop(thread_id, None) + thread_protocol_broker._active_runs.pop(thread_id, None) + + # Second lifecycle event after the broker reset must continue at seq 2. + client.portal.call( + lambda: run_jobs_module._publish_lifecycle(thread_id, event="completed", graph_name="default") + ) + + # Persisted rows must reflect both, in order, without a seq collision. + persisted = client.portal.call( + lambda: stream_module.load_thread_stream_events( + thread_id, + channels=["lifecycle"], + namespaces=None, + depth=None, + ) + ) + assert [event["seq"] for event in persisted] == [1, 2] + assert [event["params"]["data"]["event"] for event in persisted] == ["started", "completed"] diff --git a/tests/integration/test_stream_persistence_atomic.py b/tests/integration/test_stream_persistence_atomic.py new file mode 100644 index 0000000..78ebdf3 --- /dev/null +++ b/tests/integration/test_stream_persistence_atomic.py @@ -0,0 +1,227 @@ +"""Regression tests for the atomic stream append API. + +Covers the fourth-round review findings 1 and 2 (non-Redis sequence allocation +is not atomic; a live cursor can be exposed before it is durable): + +- concurrent publishers of the same run/thread must allocate unique, gapless + seqs (previously ``SELECT MAX(seq)+1`` races produced ``[1, 1]`` and one + frame was silently dropped by the persistence helper) +- a failed durable append must raise (not be swallowed) so the event is never + exposed to a client +- the per-stream counter row self-heals from ``MAX(seq)`` if it is deleted out + from under an active stream (the two-phase delete path removes business rows + before stream rows) + +The tests run against the real sqlite metadata DB initialized by the ``client`` +fixture (the same backend used by the regular integration suite), which is the +weakest backend: if atomic allocation is correct on sqlite's single-writer +locking, it is correct on the row-locked MySQL/SeekDB and PostgreSQL backends. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from agentseek_api.services import stream_persistence as stream_module + + +def _thread_payload(**params: object) -> dict[str, object]: + return { + "method": "values", + "params": { + "namespace": [], + "timestamp": 1, + **params, + }, + } + + +def test_concurrent_atomic_appends_allocate_unique_seqs(client) -> None: + """N concurrent publishers of the same run/thread each get a unique seq. + + Regression for the review reproduction ``[1, 1]``: the old + ``SELECT MAX(seq)+1`` allocator let two publishers observe the same max and + the unique constraint then dropped one frame. The counter-row append must + serialize publishers and hand out exactly ``1..N``. + """ + + async def run_concurrent(scope: str, scope_id: str, n: int) -> list[int]: + if scope == "run": + append = stream_module.append_run_stream_event_atomic + load = stream_module.load_run_stream_events + results = await asyncio.gather( + *(append(scope_id, {"event": "message", "data": i}) for i in range(n)) + ) + rows = await load(scope_id) + return sorted(seq for seq, _ in results), [seq for seq, _ in rows] + + append = stream_module.append_thread_stream_event_atomic + load = lambda: stream_module.load_thread_stream_events( # noqa: E731 + scope_id, channels=["values"], namespaces=None, depth=None + ) + results = await asyncio.gather( + *(append(scope_id, _thread_payload(data=i)) for i in range(n)) + ) + rows = await load() + return sorted(seq for seq, _ in results), [event["seq"] for event in rows] + + run_seqs, run_rows = client.portal.call(run_concurrent, "run", "run-concurrent", 8) + assert run_seqs == list(range(1, 9)), f"run seqs must be unique and gapless: {run_seqs}" + assert run_rows == list(range(1, 9)), f"persisted run rows must match: {run_rows}" + + thread_seqs, thread_rows = client.portal.call(run_concurrent, "thread", "thread-concurrent", 8) + assert thread_seqs == list(range(1, 9)), f"thread seqs must be unique and gapless: {thread_seqs}" + assert thread_rows == list(range(1, 9)), f"persisted thread rows must match: {thread_rows}" + + +def test_atomic_append_raises_on_db_failure_and_is_not_exposed(client, monkeypatch) -> None: + """A failed durable append must raise instead of being swallowed. + + The caller (``_publish_run_event`` and friends) only publishes to the + in-memory broker after the append succeeds, so a client can never receive a + seq that was not durably committed. + """ + from agentseek_api.services import run_jobs as run_jobs_module + from agentseek_api.services.run_state import run_broker + + async def failing_stage(*_args: object, **_kwargs: object): + raise RuntimeError("durable append failed") + + async def exercise() -> list[str]: + published: list[str] = [] + original_publish = run_broker.publish + + def tracking_publish(run_id: str, event: str, **payload: object): + published.append(f"{run_id}:{event}") + return original_publish(run_id, event, **payload) + + monkeypatch.setattr(run_broker, "publish", tracking_publish) + monkeypatch.setattr( + stream_module, "_stage_db_event", failing_stage, raising=False + ) + try: + await run_jobs_module._publish_run_event("run-fail-inject", "start") + except RuntimeError: + return ["raised"] + published + return ["no-raise"] + published + + outcome = client.portal.call(exercise) + assert outcome[0] == "raised", "append failure must propagate to the caller" + assert len(outcome) == 1, "the event must not be exposed to the broker on failure" + + +def test_atomic_append_self_heals_after_counter_row_deleted(client) -> None: + """The counter row re-seeds from ``MAX(seq)`` if it disappears mid-stream. + + The two-phase delete path (``threads.py``) removes business rows before + stream rows, so a counter row can be deleted while a stream is still being + appended to. The append must not collide with persisted rows - it re-seeds + from the durable events and keeps allocating monotonically. + """ + from sqlalchemy import delete + + from agentseek_api.core.database import db_manager + from agentseek_api.core.orm import StreamSequence + + async def exercise() -> list[tuple[int, str]]: + first = await stream_module.append_run_stream_event_atomic( + "run-self-heal", {"event": "start"} + ) + second = await stream_module.append_run_stream_event_atomic( + "run-self-heal", {"event": "message", "data": "mid"} + ) + async with db_manager.get_session_factory()() as session: + await session.execute( + delete(StreamSequence).where( + StreamSequence.scope == "run", StreamSequence.scope_id == "run-self-heal" + ) + ) + await session.commit() + third = await stream_module.append_run_stream_event_atomic( + "run-self-heal", {"event": "end", "status": "success"} + ) + rows = await stream_module.load_run_stream_events("run-self-heal") + return ( + [first[0], second[0], third[0]], + [(seq, str(payload.get("event"))) for seq, payload in rows], + ) + + seqs, rows = client.portal.call(exercise) + assert seqs == [1, 2, 3], f"seq must stay monotonic across counter deletion: {seqs}" + assert rows == [(1, "start"), (2, "message"), (3, "end")], ( + f"persisted events must be gapless and ordered: {rows}" + ) + + +@pytest.mark.asyncio +async def test_atomic_append_resolves_legacy_seq_collision(monkeypatch: pytest.MonkeyPatch) -> None: + """A seq collision with an out-of-band row is healed, not dropped. + + ``_db_append`` retries once with a fresh allocation after the unique + constraint rejects the insert, so a legacy out-of-band row can never wedge + the atomic stream into losing a frame. + """ + from sqlalchemy.exc import IntegrityError + + class FakeSession: + def __init__(self) -> None: + self.commit_calls = 0 + self.rollbacks = 0 + self.staged: list[tuple[int | None, dict[str, object]]] = [] + self.persisted: list[tuple[int | None, dict[str, object]]] = [] + + async def __aenter__(self) -> FakeSession: + return self + + async def __aexit__(self, *_args: object) -> None: + return None + + async def commit(self) -> None: + self.commit_calls += 1 + if self.commit_calls == 1: + raise IntegrityError("INSERT", {}, Exception("duplicate seq")) + self.persisted = list(self.staged) + + async def rollback(self) -> None: + self.rollbacks += 1 + + def add(self, _obj: object) -> None: + return None + + session = FakeSession() + stage_calls = 0 + + async def fake_stage( + _session: object, + scope: str, + scope_id: str, + payload: dict[str, object], + *, + seq: int | None, + ) -> tuple[int, dict[str, object]]: + nonlocal stage_calls + _ = (scope, scope_id) + stage_calls += 1 + session.staged.append((seq, dict(payload))) + # Simulate the allocation: the retried append must allocate a fresh, + # higher seq rather than reuse the colliding one. + return stage_calls, payload + + class FakeFactory: + def __call__(self) -> FakeSession: + return session + + monkeypatch.setattr(stream_module, "_stage_db_event", fake_stage) + monkeypatch.setattr(stream_module, "_metadata_db_ready", lambda: True) + monkeypatch.setattr(stream_module.db_manager, "get_session_factory", FakeFactory) + seq, payload = await stream_module.append_run_stream_event_atomic( + "run-legacy-collision", {"event": "end"} + ) + + assert seq == 2, "the retried append must allocate a fresh, higher seq" + assert payload == {"event": "end"} + assert session.commit_calls == 2 + assert session.rollbacks == 1, "the colliding attempt must be rolled back" + assert session.persisted == [(None, {"event": "end"}), (None, {"event": "end"})] \ No newline at end of file diff --git a/tests/unit/test_run_executor.py b/tests/unit/test_run_executor.py index 54e5fa7..bc95150 100644 --- a/tests/unit/test_run_executor.py +++ b/tests/unit/test_run_executor.py @@ -1,3 +1,4 @@ +from types import SimpleNamespace from typing import Any import pytest @@ -5,62 +6,28 @@ from langgraph.constants import CONF, CONFIG_KEY_CHECKPOINTER from agentseek_api.core.runtime_store import UserScopedStore -from agentseek_api.settings import settings -from agentseek_api.services import run_executor as run_executor_module from agentseek_api.services.run_executor import ( RunExecutionResult, _ProtocolMessageStreamState, - _translate_stream_events, execute_run, ) from agentseek_api.services.thread_protocol import ThreadProtocolEventBroker -@pytest.mark.asyncio -async def test_publish_translated_run_event_uses_atomic_redis_append(monkeypatch: pytest.MonkeyPatch) -> None: - published: list[tuple[str, str, int | None, dict[str, Any]]] = [] - monkeypatch.setattr(settings, "EXECUTOR_BACKEND", "redis") - publish_event = getattr(run_executor_module, "_publish_translated_run_event", None) - - async def fake_append(run_id: str, payload: dict[str, Any]) -> tuple[int, dict[str, Any]]: - assert run_id == "run-1" - return 11, payload - - async def unexpected_next_seq(_run_id: str) -> int: - raise AssertionError("Redis sequence allocation must be part of the append") - - monkeypatch.setattr(run_executor_module, "append_redis_run_stream_event", fake_append, raising=False) - monkeypatch.setattr(run_executor_module, "next_run_stream_seq", unexpected_next_seq) - monkeypatch.setattr( - run_executor_module.run_broker, - "publish", - lambda run_id, event, *, seq=None, **payload: ( - published.append((run_id, event, seq, payload)) or (seq, {"event": event, **payload}) - ), - ) - - assert callable(publish_event) - result = await publish_event("run-1", "message", {"data": "hello"}) - - assert result == (11, {"event": "message", "data": "hello"}) - assert published == [("run-1", "message", 11, {"data": "hello"})] - - class FakeGraph: + """Fake compiled graph: the default path streams through ``astream()`` and + then reads the final state back via ``aget_state()``, mirroring how + ``execute_run`` drives a real langgraph graph.""" + def __init__(self) -> None: self.configs: list[dict] = [] - async def astream_events(self, prepared_input: dict, config: dict, version: str = "v2", **kwargs): + async def astream(self, prepared_input: dict, config: dict, **kwargs): self.configs.append(config) - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "langgraph-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": {"output": {"output": {"ok": True, "received": prepared_input}}}, - } + yield ("values", {"output": {"ok": True, "received": prepared_input}}) + + async def aget_state(self, config: dict): + return SimpleNamespace(values=None) class FakeEntry: @@ -189,19 +156,11 @@ def __init__(self) -> None: self.stream_kwargs: list[dict] = [] self.inputs: list[Any] = [] - async def astream_events(self, prepared_input, config: dict, version: str = "v2", **kwargs): + async def astream(self, prepared_input, config: dict, **kwargs): self.configs.append(config) self.stream_kwargs.append(kwargs) self.inputs.append(prepared_input) - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "langgraph-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": {"output": {"output": {"ok": True}}}, - } + yield ("values", {"output": {"ok": True}}) class FakeKwargsCapturingEntry: @@ -232,7 +191,10 @@ def get_graph(self, _graph_id: str | None = None): async def test_execute_run_forwards_command_as_invocation(monkeypatch: pytest.MonkeyPatch) -> None: fake_db = FakeDBManager() FakeKwargsCapturingEntry.graph = FakeKwargsCapturingGraph() - monkeypatch.setattr("agentseek_api.services.run_executor.get_langgraph_service", lambda: FakeKwargsCapturingLangGraphService()) + monkeypatch.setattr( + "agentseek_api.services.run_executor.get_langgraph_service", + lambda: FakeKwargsCapturingLangGraphService(), + ) monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) from langgraph.types import Command @@ -256,7 +218,10 @@ async def test_execute_run_forwards_command_as_invocation(monkeypatch: pytest.Mo async def test_execute_run_forwards_interrupt_and_stream_mode_kwargs(monkeypatch: pytest.MonkeyPatch) -> None: fake_db = FakeDBManager() FakeKwargsCapturingEntry.graph = FakeKwargsCapturingGraph() - monkeypatch.setattr("agentseek_api.services.run_executor.get_langgraph_service", lambda: FakeKwargsCapturingLangGraphService()) + monkeypatch.setattr( + "agentseek_api.services.run_executor.get_langgraph_service", + lambda: FakeKwargsCapturingLangGraphService(), + ) monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) await execute_run( @@ -277,34 +242,18 @@ async def test_execute_run_forwards_interrupt_and_stream_mode_kwargs(monkeypatch assert "values" in kwargs["stream_mode"] assert "debug" in kwargs["stream_mode"] - class FakeInterruptGraph(FakeGraph): - async def astream_events(self, prepared_input: dict, config: dict, version: str = "v2", **kwargs): + async def astream(self, prepared_input: dict, config: dict, **kwargs): self.configs.append(config) - yield { - "event": "on_chain_stream", - "name": "fake-graph", - "run_id": "langgraph-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": { - "chunk": { - "__interrupt__": [ - type("Interrupt", (), {"value": "Provide value:", "id": "interrupt-1"})(), - ] - } + yield ( + "updates", + { + "__interrupt__": [ + type("Interrupt", (), {"value": "Provide value:", "id": "interrupt-1"})(), + ] }, - } - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "langgraph-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": {"output": {"foo": prepared_input["input"]["foo"]}}, - } + ) + yield ("values", {"foo": prepared_input["input"]["foo"]}) class FakeInterruptEntry(FakeEntry): @@ -330,14 +279,16 @@ def get_entry(self, _graph_id: str | None) -> FakeInterruptEntry: @pytest.mark.asyncio -async def test_execute_run_preserves_interrupts_from_root_stream(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_execute_run_preserves_interrupts_from_updates_stream(monkeypatch: pytest.MonkeyPatch) -> None: fake_db = FakeDBManager() + protocol_broker = ThreadProtocolEventBroker() FakeInterruptEntry.graph = FakeInterruptGraph() monkeypatch.setattr( "agentseek_api.services.run_executor.get_langgraph_service", lambda: FakeInterruptLangGraphService(), ) monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) + monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) result = await execute_run(thread_id="t1", run_id="r1", payload={"foo": "hello"}, user_id="user-1") @@ -345,564 +296,267 @@ async def test_execute_run_preserves_interrupts_from_root_stream(monkeypatch: py assert result.interrupts == [{"value": "Provide value:", "id": "interrupt-1"}] assert result.output["state"]["foo"] == "hello" + # The interrupt rides on a ``values`` event with ``__interrupt__`` + # intact (remapped from updates because updates were not explicitly requested), + # so the official SDK stream() parser can surface it. + value_events = [event for event in protocol_broker._events["t1"] if event["method"] == "values"] + interrupt_values = [event for event in value_events if "__interrupt__" in event["params"]["data"]] + assert len(interrupt_values) == 1 + assert interrupt_values[0]["params"]["data"]["__interrupt__"] == [ + {"value": "Provide value:", "id": "interrupt-1"} + ] + -def test_translate_stream_events_maps_chat_model_stream_to_message_chunk() -> None: - translated = _translate_stream_events( - { - "event": "on_chat_model_stream", - "name": "chat-model", - "run_id": "langgraph-run", - "parent_ids": ["parent-run"], - "metadata": {"langgraph_node": "call_model"}, - "tags": ["graph:step:1"], - "data": {"chunk": AIMessageChunk(content="hello")}, - } +@pytest.mark.asyncio +async def test_execute_run_keeps_interrupt_in_updates_when_updates_requested(monkeypatch: pytest.MonkeyPatch) -> None: + """When the client explicitly requests updates, the interrupt stays on the + updates channel with ``__interrupt__`` intact (official behavior).""" + fake_db = FakeDBManager() + protocol_broker = ThreadProtocolEventBroker() + FakeInterruptEntry.graph = FakeInterruptGraph() + monkeypatch.setattr( + "agentseek_api.services.run_executor.get_langgraph_service", + lambda: FakeInterruptLangGraphService(), ) + monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) + monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) - assert translated == [ - ( - "message_chunk", - { - "name": "chat-model", - "langgraph_event": "on_chat_model_stream", - "langgraph_run_id": "langgraph-run", - "metadata": {"langgraph_node": "call_model"}, - "tags": ["graph:step:1"], - "parent_ids": ["parent-run"], - "node": "call_model", - "message_type": "AIMessageChunk", - "content": "hello", - }, - ) + await execute_run( + thread_id="t1", + run_id="r1", + payload={"foo": "hello"}, + user_id="user-1", + kwargs={"stream_modes": ["updates"]}, + ) + + update_events = [event for event in protocol_broker._events["t1"] if event["method"] == "updates"] + interrupt_updates = [event for event in update_events if "__interrupt__" in event["params"]["data"]] + assert len(interrupt_updates) == 1 + assert interrupt_updates[0]["params"]["data"]["__interrupt__"] == [ + {"value": "Provide value:", "id": "interrupt-1"} ] class FakeProtocolStreamingGraph(FakeGraph): - async def astream_events(self, prepared_input: dict, config: dict, version: str = "v2", **kwargs): + async def astream(self, prepared_input: dict, config: dict, **kwargs): self.configs.append(config) - yield { - "event": "on_chat_model_stream", - "name": "chat-model", - "run_id": "chat-run", - "parent_ids": ["root-run"], - "metadata": {"langgraph_node": "call_model"}, - "tags": [], - "data": {"chunk": AIMessageChunk(content="hel")}, - } - yield { - "event": "on_chat_model_stream", - "name": "chat-model", - "run_id": "chat-run", - "parent_ids": ["root-run"], - "metadata": {"langgraph_node": "call_model"}, - "tags": [], - "data": {"chunk": AIMessageChunk(content="lo")}, - } - yield { - "event": "on_chain_stream", - "name": "fake-graph", - "run_id": "root-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": {"chunk": {"step": "partial"}}, - } - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "root-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": {"output": {"output": {"messages": [AIMessage(content="hello")], "step": "final"}}}, - } + yield ("messages", (AIMessageChunk(content="hel"), {"langgraph_node": "call_model"})) + yield ("messages", (AIMessageChunk(content="lo"), {"langgraph_node": "call_model"})) + yield ("updates", {"step": "partial"}) + yield ("values", {"output": {"messages": [AIMessage(content="hello")], "step": "final"}}) class FakeProtocolLlmStreamingGraph(FakeGraph): - async def astream_events(self, prepared_input: dict, config: dict, version: str = "v2", **kwargs): + async def astream(self, prepared_input: dict, config: dict, **kwargs): self.configs.append(config) - yield { - "event": "on_llm_stream", - "name": "completion-model", - "run_id": "llm-run", - "parent_ids": ["root-run"], - "metadata": {"langgraph_node": "call_model"}, - "tags": [], - "data": {"chunk": type("Chunk", (), {"text": "hel"})()}, - } - yield { - "event": "on_llm_stream", - "name": "completion-model", - "run_id": "llm-run", - "parent_ids": ["root-run"], - "metadata": {"langgraph_node": "call_model"}, - "tags": [], - "data": {"chunk": type("Chunk", (), {"text": "lo"})()}, - } - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "root-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": {"output": {"output": {"text": "hello"}}}, - } - - -class FakeProtocolStreamingEntry(FakeEntry): - graph = FakeProtocolStreamingGraph() - - @staticmethod - def build_graph(_checkpointer=None) -> FakeProtocolStreamingGraph: - return FakeProtocolStreamingEntry.graph - - -class FakeProtocolStreamingLangGraphService(FakeLangGraphService): - def get_entry(self, _graph_id: str | None) -> FakeProtocolStreamingEntry: - return FakeProtocolStreamingEntry() - - -class FakeProtocolLlmStreamingEntry(FakeEntry): - graph = FakeProtocolLlmStreamingGraph() - - @staticmethod - def build_graph(_checkpointer=None) -> FakeProtocolLlmStreamingGraph: - return FakeProtocolLlmStreamingEntry.graph - - -class FakeProtocolLlmStreamingLangGraphService(FakeLangGraphService): - def get_entry(self, _graph_id: str | None) -> FakeProtocolLlmStreamingEntry: - return FakeProtocolLlmStreamingEntry() + yield ("messages", (AIMessageChunk(content="hel"), {"langgraph_node": "call_model"})) + yield ("messages", (AIMessageChunk(content="lo"), {"langgraph_node": "call_model"})) + yield ("values", {"output": {"text": "hello"}}) class FakeProtocolNamespaceGraph(FakeGraph): - async def astream_events(self, prepared_input: dict, config: dict, version: str = "v2", **kwargs): - self.configs.append(config) - yield { - "event": "on_tool_start", - "name": "search_docs", - "run_id": "tool-run", - "parent_ids": ["root-run"], - "metadata": { - "langgraph_node": "search_docs", - "langgraph_checkpoint_ns": "node_1:task-1|search_docs:task-2", - }, - "tags": [], - "data": {"input": {"query": prepared_input["input"]["hello"]}}, - } - yield { - "event": "on_chat_model_stream", - "name": "chat-model", - "run_id": "chat-run", - "parent_ids": ["root-run"], - "metadata": { - "langgraph_node": "call_model", - "langgraph_checkpoint_ns": "node_1:task-1|call_model:task-3", - }, - "tags": [], - "data": {"chunk": AIMessageChunk(content="hello")}, - } - yield { - "event": "on_chain_stream", - "name": "node_1", - "run_id": "subgraph-run", - "parent_ids": ["root-run"], - "metadata": { - "langgraph_node": "node_1", - "langgraph_checkpoint_ns": "node_1:task-1", - }, - "tags": [], - "data": {"chunk": {"step": "partial"}}, - } - yield { - "event": "on_tool_end", - "name": "search_docs", - "run_id": "tool-run", - "parent_ids": ["root-run"], - "metadata": { - "langgraph_node": "search_docs", - "langgraph_checkpoint_ns": "node_1:task-1|search_docs:task-2", - }, - "tags": [], - "data": {"output": {"answer": "docs"}}, - } - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "root-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": {"output": {"output": {"messages": [AIMessage(content="hello")], "step": "final"}}}, - } - - -class FakeProtocolNamespaceEntry(FakeEntry): - graph = FakeProtocolNamespaceGraph() - - @staticmethod - def build_graph(_checkpointer=None) -> FakeProtocolNamespaceGraph: - return FakeProtocolNamespaceEntry.graph - + """Namespaces flow through the ``(namespace, mode, chunk)`` tuples produced + when the run requests ``stream_subgraphs``.""" -class FakeProtocolNamespaceLangGraphService(FakeLangGraphService): - def get_entry(self, _graph_id: str | None) -> FakeProtocolNamespaceEntry: - return FakeProtocolNamespaceEntry() + async def astream(self, prepared_input: dict, config: dict, **kwargs): + self.configs.append(config) + yield ( + ["node_1:task-1", "call_model:task-3"], + "messages", + (AIMessageChunk(content="hello"), {"langgraph_node": "call_model"}), + ) + yield (["node_1:task-1"], "updates", {"step": "partial"}) + yield ([], "values", {"output": {"messages": [AIMessage(content="hello")], "step": "final"}}) class FakeProtocolStructuredMessageGraph(FakeGraph): - async def astream_events(self, prepared_input: dict, config: dict, version: str = "v2", **kwargs): + async def astream(self, prepared_input: dict, config: dict, **kwargs): self.configs.append(config) - yield { - "event": "on_chat_model_stream", - "name": "chat-model", - "run_id": "chat-run", - "parent_ids": ["root-run"], - "metadata": {"langgraph_node": "call_model"}, - "tags": [], - "data": { - "chunk": AIMessageChunk( + yield ( + "messages", + ( + AIMessageChunk( content=[ {"type": "text", "text": "hello"}, {"type": "reasoning", "summary": [{"type": "summary_text", "text": "why"}]}, ] - ) - }, - } - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "root-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": { + ), + {"langgraph_node": "call_model"}, + ), + ) + yield ( + "values", + { "output": { - "output": { - "messages": [ - AIMessage( - content=[ - {"type": "text", "text": "hello"}, - {"type": "reasoning", "summary": [{"type": "summary_text", "text": "why"}]}, - ] - ) - ] - } + "messages": [ + { + "type": "AIMessage", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "why"}]}, + ], + } + ] } }, - } - - -class FakeProtocolStructuredMessageEntry(FakeEntry): - graph = FakeProtocolStructuredMessageGraph() - - @staticmethod - def build_graph(_checkpointer=None) -> FakeProtocolStructuredMessageGraph: - return FakeProtocolStructuredMessageEntry.graph - - -class FakeProtocolStructuredMessageLangGraphService(FakeLangGraphService): - def get_entry(self, _graph_id: str | None) -> FakeProtocolStructuredMessageEntry: - return FakeProtocolStructuredMessageEntry() + ) class FakeProtocolToolCallChunkGraph(FakeGraph): - async def astream_events(self, prepared_input: dict, config: dict, version: str = "v2", **kwargs): + async def astream(self, prepared_input: dict, config: dict, **kwargs): self.configs.append(config) - yield { - "event": "on_chat_model_stream", - "name": "chat-model", - "run_id": "chat-run", - "parent_ids": ["root-run"], - "metadata": {"langgraph_node": "call_model"}, - "tags": [], - "data": { - "chunk": AIMessageChunk( + yield ( + "messages", + ( + AIMessageChunk( content="", tool_call_chunks=[ {"id": "call-1", "name": "search", "args": '{"q":"hel"}', "index": 0}, ], - ) - }, - } - yield { - "event": "on_chat_model_stream", - "name": "chat-model", - "run_id": "chat-run", - "parent_ids": ["root-run"], - "metadata": {"langgraph_node": "call_model"}, - "tags": [], - "data": { - "chunk": AIMessageChunk( + ), + {"langgraph_node": "call_model"}, + ), + ) + yield ( + "messages", + ( + AIMessageChunk( content="", tool_call_chunks=[ {"id": "call-1", "name": None, "args": 'lo"}', "index": 0}, ], - ) - }, - } - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "root-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": {"output": {"output": {"messages": []}}}, - } - - -class FakeProtocolToolCallChunkEntry(FakeEntry): - graph = FakeProtocolToolCallChunkGraph() - - @staticmethod - def build_graph(_checkpointer=None) -> FakeProtocolToolCallChunkGraph: - return FakeProtocolToolCallChunkEntry.graph - - -class FakeProtocolToolCallChunkLangGraphService(FakeLangGraphService): - def get_entry(self, _graph_id: str | None) -> FakeProtocolToolCallChunkEntry: - return FakeProtocolToolCallChunkEntry() + ), + {"langgraph_node": "call_model"}, + ), + ) + yield ("values", {"output": {"messages": []}}) class FakeProtocolMixedStructuredGraph(FakeGraph): - async def astream_events(self, prepared_input: dict, config: dict, version: str = "v2", **kwargs): + async def astream(self, prepared_input: dict, config: dict, **kwargs): self.configs.append(config) - yield { - "event": "on_chat_model_stream", - "name": "chat-model", - "run_id": "chat-run", - "parent_ids": ["root-run"], - "metadata": {"langgraph_node": "call_model"}, - "tags": [], - "data": {"chunk": AIMessageChunk(content="hello")}, - } - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "root-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": { + yield ("messages", (AIMessageChunk(content="hello"), {"langgraph_node": "call_model"})) + yield ( + "values", + { "output": { - "output": { - "messages": [ - { - "type": "AIMessage", - "content": [ - {"type": "text", "text": "hello"}, - {"type": "reasoning", "summary": [{"type": "summary_text", "text": "why"}]}, - ], - } - ] - } + "messages": [ + { + "type": "AIMessage", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "why"}]}, + ], + } + ] } }, - } - - -class FakeProtocolMixedStructuredEntry(FakeEntry): - graph = FakeProtocolMixedStructuredGraph() - - @staticmethod - def build_graph(_checkpointer=None) -> FakeProtocolMixedStructuredGraph: - return FakeProtocolMixedStructuredEntry.graph - - -class FakeProtocolMixedStructuredLangGraphService(FakeLangGraphService): - def get_entry(self, _graph_id: str | None) -> FakeProtocolMixedStructuredEntry: - return FakeProtocolMixedStructuredEntry() + ) class FakeProtocolMultiMessageGraph(FakeGraph): - async def astream_events(self, prepared_input: dict, config: dict, version: str = "v2", **kwargs): - self.configs.append(config) - yield { - "event": "on_chain_stream", - "name": "fake-graph", - "run_id": "root-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": {"chunk": {"messages": [HumanMessage(content="hi"), AIMessage(content="hello")]}} - } - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "root-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": {"output": {"output": {"messages": [HumanMessage(content="hi"), AIMessage(content="hello")]}}} - } - - -class FakeProtocolMultiMessageEntry(FakeEntry): - graph = FakeProtocolMultiMessageGraph() + """A single state update carrying several non-LLM messages at once; each + must surface as its own distinct protocol message (issue #48 regression: + updates are emitted 1:1, never duplicated or collapsed into one id).""" - @staticmethod - def build_graph(_checkpointer=None) -> FakeProtocolMultiMessageGraph: - return FakeProtocolMultiMessageEntry.graph - - -class FakeProtocolMultiMessageLangGraphService(FakeLangGraphService): - def get_entry(self, _graph_id: str | None) -> FakeProtocolMultiMessageEntry: - return FakeProtocolMultiMessageEntry() + async def astream(self, prepared_input: dict, config: dict, **kwargs): + self.configs.append(config) + yield ( + "updates", + { + "some_node": { + "messages": [ + HumanMessage(content="hi"), + ToolMessage(content="tool result", tool_call_id="call-1"), + ] + } + }, + ) + yield ( + "values", + { + "output": { + "messages": [ + HumanMessage(content="hi"), + ToolMessage(content="tool result", tool_call_id="call-1"), + ] + } + }, + ) class FakeProtocolToolMessageGraph(FakeGraph): - async def astream_events(self, prepared_input: dict, config: dict, version: str = "v2", **kwargs): + async def astream(self, prepared_input: dict, config: dict, **kwargs): self.configs.append(config) tool_message = ToolMessage( content="42 characters", tool_call_id="call-character-count", id="tool-message-1", ) - tool_event = { - "event": "on_chain_stream", - "name": "tools", - "run_id": "tool-run", - "parent_ids": ["root-run"], - "metadata": { - "langgraph_node": "tools", - "langgraph_checkpoint_ns": "tools:task-1", - "provider": "deterministic", - }, - "tags": ["graph:step:2"], - "data": {"chunk": {"messages": [tool_message]}}, - } - yield tool_event - yield tool_event - yield { - "event": "on_chat_model_stream", - "name": "chat-model", - "run_id": "chat-run", - "parent_ids": ["root-run"], - "metadata": {"langgraph_node": "call_model"}, - "tags": ["graph:step:3"], - "data": {"chunk": AIMessageChunk(content="Final answer", id="ai-message-1")}, - } - yield { - "event": "on_chain_end", - "name": "fake-graph", - "run_id": "root-run", - "parent_ids": [], - "metadata": {}, - "tags": [], - "data": { - "output": { - "output": { - "messages": [tool_message, AIMessage(content="Final answer", id="ai-message-1")] - } - } - }, + tool_meta = { + "langgraph_node": "tools", + "langgraph_checkpoint_ns": "tools:task-1", + "provider": "deterministic", } + yield (["tools:task-1"], "messages", (tool_message, tool_meta)) + # The same tool message is streamed twice; it must only be published once. + yield (["tools:task-1"], "messages", (tool_message, tool_meta)) + yield ( + ["call_model:task-3"], + "messages", + (AIMessageChunk(content="Final answer", id="ai-message-1"), {"langgraph_node": "call_model"}), + ) + yield ( + [], + "values", + {"output": {"messages": [tool_message, AIMessage(content="Final answer", id="ai-message-1")]}}, + ) -class FakeProtocolToolMessageEntry(FakeEntry): - graph = FakeProtocolToolMessageGraph() +class _FakeProtocolEntry(FakeEntry): + graph = FakeGraph() @staticmethod - def build_graph(_checkpointer=None) -> FakeProtocolToolMessageGraph: - return FakeProtocolToolMessageEntry.graph + def build_graph(_checkpointer=None) -> FakeGraph: + return _FakeProtocolEntry.graph -class FakeProtocolToolMessageLangGraphService(FakeLangGraphService): - def get_entry(self, _graph_id: str | None) -> FakeProtocolToolMessageEntry: - return FakeProtocolToolMessageEntry() +class _FakeProtocolLangGraphService(FakeLangGraphService): + def get_entry(self, _graph_id: str | None) -> _FakeProtocolEntry: + return _FakeProtocolEntry() -async def _execute_protocol_tool_message_graph( +async def _run_fake_graph( monkeypatch: pytest.MonkeyPatch, + graph: FakeGraph, *, - stream_modes: list[str], + stream_modes: list[str] | None = None, + stream_subgraphs: bool = False, ) -> list[dict[str, Any]]: fake_db = FakeDBManager() protocol_broker = ThreadProtocolEventBroker() - FakeProtocolToolMessageEntry.graph = FakeProtocolToolMessageGraph() + _FakeProtocolEntry.graph = graph monkeypatch.setattr( "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeProtocolToolMessageLangGraphService(), + lambda: _FakeProtocolLangGraphService(), ) monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) - await execute_run( - thread_id="t1", - run_id="r1", - payload={"hello": "world"}, - user_id="user-1", - kwargs={"stream_modes": stream_modes}, - ) + kwargs: dict[str, Any] = {} + if stream_modes is not None: + kwargs["stream_modes"] = stream_modes + if stream_subgraphs: + kwargs["stream_subgraphs"] = True + await execute_run(thread_id="t1", run_id="r1", payload={"hello": "world"}, user_id="user-1", kwargs=kwargs) return protocol_broker._events["t1"] @pytest.mark.asyncio -async def test_execute_run_mirrors_tool_message_to_requested_tuple_stream_once_and_before_final_answer( - monkeypatch: pytest.MonkeyPatch, -) -> None: - thread_events = await _execute_protocol_tool_message_graph( - monkeypatch, - stream_modes=["messages-tuple", "values"], +async def test_execute_run_publishes_incremental_protocol_messages_and_values(monkeypatch: pytest.MonkeyPatch) -> None: + thread_events = await _run_fake_graph( + monkeypatch, FakeProtocolStreamingGraph(), stream_modes=["updates", "values"] ) - complete_events = [event for event in thread_events if event["method"] == "messages/complete"] - tuple_events = [event for event in thread_events if event["method"] == "messages-tuple"] - tool_tuples = [event for event in tuple_events if event["params"]["data"][0]["type"] == "tool"] - - assert len(complete_events) == 1 - assert len(tool_tuples) == 1 - - tool_tuple = tool_tuples[0] - tool_payload, metadata = tool_tuple["params"]["data"] - assert tool_payload["content"] == "42 characters" - assert tool_payload["id"] == "tool-message-1" - assert tool_payload["tool_call_id"] == "call-character-count" - assert metadata == { - "langgraph_node": "tools", - "langgraph_checkpoint_ns": "tools:task-1", - "provider": "deterministic", - } - assert tool_tuple["params"]["namespace"] == ["tools:task-1"] - assert tool_tuple["params"]["run_id"] == "r1" - assert complete_events[0]["params"]["data"] == [tool_payload] - - final_answer_tuple = next( - event for event in tuple_events if event["params"]["data"][0].get("id") == "ai-message-1" - ) - assert thread_events.index(tool_tuple) < thread_events.index(final_answer_tuple) - - -@pytest.mark.asyncio -async def test_execute_run_keeps_tool_message_complete_without_unrequested_tuple_mirror( - monkeypatch: pytest.MonkeyPatch, -) -> None: - thread_events = await _execute_protocol_tool_message_graph(monkeypatch, stream_modes=["values"]) - complete_events = [event for event in thread_events if event["method"] == "messages/complete"] - tuple_events = [event for event in thread_events if event["method"] == "messages-tuple"] - - assert len(complete_events) == 1 - assert complete_events[0]["params"]["data"][0]["id"] == "tool-message-1" - assert tuple_events == [] - - -@pytest.mark.asyncio -async def test_execute_run_publishes_incremental_protocol_messages_and_values(monkeypatch: pytest.MonkeyPatch) -> None: - fake_db = FakeDBManager() - protocol_broker = ThreadProtocolEventBroker() - FakeProtocolStreamingEntry.graph = FakeProtocolStreamingGraph() - - monkeypatch.setattr( - "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeProtocolStreamingLangGraphService(), - ) - monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) - monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) - - await execute_run(thread_id="t1", run_id="r1", payload={"hello": "world"}, user_id="user-1") - - thread_events = protocol_broker._events["t1"] message_events = [event for event in thread_events if event["method"] == "messages"] update_events = [event for event in thread_events if event["method"] == "updates"] value_events = [event for event in thread_events if event["method"] == "values"] @@ -922,20 +576,8 @@ async def test_execute_run_publishes_incremental_protocol_messages_and_values(mo async def test_execute_run_publishes_incremental_protocol_messages_for_llm_text_chunks( monkeypatch: pytest.MonkeyPatch, ) -> None: - fake_db = FakeDBManager() - protocol_broker = ThreadProtocolEventBroker() - FakeProtocolLlmStreamingEntry.graph = FakeProtocolLlmStreamingGraph() - - monkeypatch.setattr( - "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeProtocolLlmStreamingLangGraphService(), - ) - monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) - monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) - - await execute_run(thread_id="t1", run_id="r1", payload={"hello": "world"}, user_id="user-1") - - message_events = [event for event in protocol_broker._events["t1"] if event["method"] == "messages"] + thread_events = await _run_fake_graph(monkeypatch, FakeProtocolLlmStreamingGraph()) + message_events = [event for event in thread_events if event["method"] == "messages"] assert [event["params"]["data"]["event"] for event in message_events] == [ "message-start", "content-block-start", @@ -950,46 +592,25 @@ async def test_execute_run_publishes_incremental_protocol_messages_for_llm_text_ @pytest.mark.asyncio async def test_execute_run_uses_langgraph_namespaces_for_protocol_events(monkeypatch: pytest.MonkeyPatch) -> None: - fake_db = FakeDBManager() - protocol_broker = ThreadProtocolEventBroker() - FakeProtocolNamespaceEntry.graph = FakeProtocolNamespaceGraph() - - monkeypatch.setattr( - "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeProtocolNamespaceLangGraphService(), + thread_events = await _run_fake_graph( + monkeypatch, FakeProtocolNamespaceGraph(), stream_modes=["updates"], stream_subgraphs=True ) - monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) - monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) - - await execute_run(thread_id="t1", run_id="r1", payload={"hello": "world"}, user_id="user-1") - - thread_events = protocol_broker._events["t1"] message_events = [event for event in thread_events if event["method"] == "messages"] - tool_events = [event for event in thread_events if event["method"] == "tools"] updates_events = [event for event in thread_events if event["method"] == "updates"] assert message_events[0]["params"]["namespace"] == ["node_1:task-1", "call_model:task-3"] - assert all(event["params"]["namespace"] == ["node_1:task-1", "search_docs:task-2"] for event in tool_events) assert updates_events[0]["params"]["namespace"] == ["node_1:task-1"] @pytest.mark.asyncio async def test_execute_run_publishes_structured_protocol_message_blocks(monkeypatch: pytest.MonkeyPatch) -> None: - fake_db = FakeDBManager() - protocol_broker = ThreadProtocolEventBroker() - FakeProtocolStructuredMessageEntry.graph = FakeProtocolStructuredMessageGraph() - - monkeypatch.setattr( - "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeProtocolStructuredMessageLangGraphService(), - ) - monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) - monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) - - await execute_run(thread_id="t1", run_id="r1", payload={"hello": "world"}, user_id="user-1") - - message_events = [event for event in protocol_broker._events["t1"] if event["method"] == "messages"] - block_starts = [event["params"]["data"] for event in message_events if event["params"]["data"]["event"] == "content-block-start"] + thread_events = await _run_fake_graph(monkeypatch, FakeProtocolStructuredMessageGraph()) + message_events = [event for event in thread_events if event["method"] == "messages"] + block_starts = [ + event["params"]["data"] + for event in message_events + if event["params"]["data"]["event"] == "content-block-start" + ] assert any(block["content"]["type"] == "reasoning" for block in block_starts) ordered_events = [event["params"]["data"] for event in message_events] text_finish_index = next( @@ -1009,20 +630,8 @@ async def test_execute_run_publishes_structured_protocol_message_blocks(monkeypa async def test_execute_run_streams_tool_call_chunks_without_duplicate_complete_blocks( monkeypatch: pytest.MonkeyPatch, ) -> None: - fake_db = FakeDBManager() - protocol_broker = ThreadProtocolEventBroker() - FakeProtocolToolCallChunkEntry.graph = FakeProtocolToolCallChunkGraph() - - monkeypatch.setattr( - "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeProtocolToolCallChunkLangGraphService(), - ) - monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) - monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) - - await execute_run(thread_id="t1", run_id="r1", payload={"hello": "world"}, user_id="user-1") - - message_events = [event["params"]["data"] for event in protocol_broker._events["t1"] if event["method"] == "messages"] + thread_events = await _run_fake_graph(monkeypatch, FakeProtocolToolCallChunkGraph()) + message_events = [event["params"]["data"] for event in thread_events if event["method"] == "messages"] block_starts = [event for event in message_events if event["event"] == "content-block-start"] block_deltas = [event for event in message_events if event["event"] == "content-block-delta"] @@ -1034,22 +643,108 @@ async def test_execute_run_streams_tool_call_chunks_without_duplicate_complete_b @pytest.mark.asyncio async def test_execute_run_merges_final_structured_blocks_after_live_text(monkeypatch: pytest.MonkeyPatch) -> None: - fake_db = FakeDBManager() - protocol_broker = ThreadProtocolEventBroker() - FakeProtocolMixedStructuredEntry.graph = FakeProtocolMixedStructuredGraph() + thread_events = await _run_fake_graph(monkeypatch, FakeProtocolMixedStructuredGraph()) + message_events = [event["params"]["data"] for event in thread_events if event["method"] == "messages"] + block_starts = [event for event in message_events if event["event"] == "content-block-start"] + assert any(block["content"]["type"] == "reasoning" for block in block_starts) - monkeypatch.setattr( - "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeProtocolMixedStructuredLangGraphService(), +@pytest.mark.asyncio +async def test_execute_run_mirrors_tool_message_to_requested_tuple_stream_once_and_before_final_answer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + thread_events = await _run_fake_graph( + monkeypatch, + FakeProtocolToolMessageGraph(), + stream_modes=["messages-tuple", "values"], + stream_subgraphs=True, ) - monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) - monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) + complete_events = [event for event in thread_events if event["method"] == "messages/complete"] + tuple_events = [event for event in thread_events if event["method"] == "messages-tuple"] + tool_tuples = [event for event in tuple_events if event["params"]["data"][0]["type"] == "tool"] - await execute_run(thread_id="t1", run_id="r1", payload={"hello": "world"}, user_id="user-1") + assert len(complete_events) == 1 + assert len(tool_tuples) == 1 - message_events = [event["params"]["data"] for event in protocol_broker._events["t1"] if event["method"] == "messages"] - block_starts = [event for event in message_events if event["event"] == "content-block-start"] - assert any(block["content"]["type"] == "reasoning" for block in block_starts) + tool_tuple = tool_tuples[0] + tool_payload, metadata = tool_tuple["params"]["data"] + assert tool_payload["content"] == "42 characters" + assert tool_payload["id"] == "tool-message-1" + assert tool_payload["tool_call_id"] == "call-character-count" + assert metadata == { + "langgraph_node": "tools", + "langgraph_checkpoint_ns": "tools:task-1", + "provider": "deterministic", + } + assert tool_tuple["params"]["namespace"] == ["tools:task-1"] + assert tool_tuple["params"]["run_id"] == "r1" + assert complete_events[0]["params"]["data"] == [tool_payload] + + final_answer_tuple = next( + event for event in tuple_events if event["params"]["data"][0].get("id") == "ai-message-1" + ) + assert thread_events.index(tool_tuple) < thread_events.index(final_answer_tuple) + + +@pytest.mark.asyncio +async def test_execute_run_keeps_tool_message_complete_without_unrequested_tuple_mirror( + monkeypatch: pytest.MonkeyPatch, +) -> None: + thread_events = await _run_fake_graph( + monkeypatch, FakeProtocolToolMessageGraph(), stream_modes=["values"], stream_subgraphs=True + ) + complete_events = [event for event in thread_events if event["method"] == "messages/complete"] + tuple_events = [event for event in thread_events if event["method"] == "messages-tuple"] + + assert len(complete_events) == 1 + assert complete_events[0]["params"]["data"][0]["id"] == "tool-message-1" + assert tuple_events == [] + + +@pytest.mark.asyncio +async def test_execute_run_keeps_multiple_messages_in_single_chunk_distinct(monkeypatch: pytest.MonkeyPatch) -> None: + thread_events = await _run_fake_graph(monkeypatch, FakeProtocolMultiMessageGraph(), stream_modes=["updates"]) + message_events = [event["params"]["data"] for event in thread_events if event["method"] == "messages"] + message_starts = [event for event in message_events if event["event"] == "message-start"] + + assert [event["role"] for event in message_starts] == ["human", "tool"] + assert message_starts[0]["id"] != message_starts[1]["id"] + + +@pytest.mark.asyncio +async def test_execute_run_publishes_each_astream_updates_chunk_exactly_once(monkeypatch: pytest.MonkeyPatch) -> None: + """The default astream path must forward every ``updates`` chunk to the wire + exactly once. This pins the issue #48 regression: parallel (Send-style) + worker updates used to be emitted twice (once bare, once node-wrapped) + through the astream_events translation.""" + graph = FakeUpdatesOnceGraph() + thread_events = await _run_fake_graph(monkeypatch, graph, stream_modes=["updates"]) + updates = [event for event in thread_events if event["method"] == "updates"] + + assert [event["params"]["data"] for event in updates] == [ + {"process_item": {"results": [{"processed_item": "a", "length": 1}]}}, + {"process_item": {"results": [{"processed_item": "b", "length": 1}]}}, + {"process_item": {"results": [{"processed_item": "c", "length": 1}]}}, + ] + + +class FakeUpdatesOnceGraph(FakeGraph): + async def astream(self, prepared_input: dict, config: dict, **kwargs): + self.configs.append(config) + yield ("updates", {"process_item": {"results": [{"processed_item": "a", "length": 1}]}}) + yield ("updates", {"process_item": {"results": [{"processed_item": "b", "length": 1}]}}) + yield ("updates", {"process_item": {"results": [{"processed_item": "c", "length": 1}]}}) + yield ( + "values", + { + "output": { + "results": [ + {"processed_item": "a", "length": 1}, + {"processed_item": "b", "length": 1}, + {"processed_item": "c", "length": 1}, + ] + } + }, + ) def test_protocol_message_stream_state_merges_open_messages_against_transcript_tail( @@ -1069,163 +764,286 @@ def test_protocol_message_stream_state_merges_open_messages_against_transcript_t ) state.finish_all() - message_events = [event["params"]["data"] for event in protocol_broker._events["t1"] if event["method"] == "messages"] + message_events = [ + event["params"]["data"] for event in protocol_broker._events["t1"] if event["method"] == "messages" + ] message_starts = [event for event in message_events if event["event"] == "message-start"] assert message_starts == [{"event": "message-start", "role": "ai", "id": "m1"}] assert {"event": "content-block-delta", "index": 0, "delta": {"type": "text-delta", "text": "lo"}} in message_events assert [event for event in message_events if event["event"] == "message-finish"] == [{"event": "message-finish"}] -@pytest.mark.asyncio -async def test_execute_run_keeps_multiple_messages_in_single_chunk_distinct(monkeypatch: pytest.MonkeyPatch) -> None: - fake_db = FakeDBManager() - protocol_broker = ThreadProtocolEventBroker() - FakeProtocolMultiMessageEntry.graph = FakeProtocolMultiMessageGraph() +class FakeAstreamEventsGraph(FakeGraph): + """Fake graph for the ``events`` stream mode: executes through the retained + ``astream_events`` path (raw event stream) instead of the default astream.""" - monkeypatch.setattr( - "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeProtocolMultiMessageLangGraphService(), - ) - monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) - monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) + async def astream_events(self, prepared_input: dict, config: dict, *, version: str, **kwargs): + self.configs.append(config) + yield { + "event": "on_chat_model_stream", + "data": {"chunk": AIMessageChunk(content="hi")}, + "metadata": {"langgraph_node": "call_model"}, + "parent_ids": [], + } + yield { + "event": "on_custom_event", + "data": {"custom": "payload"}, + "metadata": {}, + "parent_ids": [], + } + yield { + "event": "on_chain_stream", + "data": {"chunk": ("values", {"output": {"ok": True}})}, + "metadata": {"langgraph_node": "root"}, + "parent_ids": [], + } + yield { + "event": "on_chain_end", + "data": {"output": {"ok": True}}, + "metadata": {"langgraph_node": "root"}, + "parent_ids": [], + } - await execute_run(thread_id="t1", run_id="r1", payload={"hello": "world"}, user_id="user-1") - message_starts = [ - event["params"]["data"] - for event in protocol_broker._events["t1"] - if event["method"] == "messages" and event["params"]["data"]["event"] == "message-start" - ] - assert [event["role"] for event in message_starts] == ["human", "ai"] - assert message_starts[0]["id"] != message_starts[1]["id"] +class _FakeAstreamEventsEntry(FakeEntry): + graph = FakeAstreamEventsGraph() + + @staticmethod + def build_graph(_checkpointer=None, store=None) -> FakeAstreamEventsGraph: + return _FakeAstreamEventsEntry.graph + + @staticmethod + def extract_output(result: dict, _payload: dict) -> dict: + return result if isinstance(result, dict) else {} -class FakeContextSchema: - tenant: str - org: str +class _FakeAstreamEventsLangGraphService(FakeLangGraphService): + def get_entry(self, _graph_id: str | None) -> _FakeAstreamEventsEntry: + return _FakeAstreamEventsEntry() @pytest.mark.asyncio -async def test_execute_run_mirrors_context_into_configurable(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_execute_run_events_mode_uses_astream_events_path(monkeypatch: pytest.MonkeyPatch) -> None: + """``stream_mode=["events"]`` keeps the raw ``astream_events`` path: message + chunks surface as protocol messages, custom events surface on the custom + channel, and the root on_chain_end finalizes the stream.""" fake_db = FakeDBManager() - FakeKwargsCapturingEntry.graph = FakeKwargsCapturingGraph() + protocol_broker = ThreadProtocolEventBroker() monkeypatch.setattr( "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeKwargsCapturingLangGraphService(), + lambda: _FakeAstreamEventsLangGraphService(), ) monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) + monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) - await execute_run( + result = await execute_run( thread_id="t1", run_id="r1", - payload={"a": 1}, + payload={"msg": "hello"}, user_id="user-1", - kwargs={"context": {"tenant": "acme"}}, + kwargs={"stream_modes": ["events"]}, ) - config = FakeKwargsCapturingEntry.graph.configs[0] - assert config[CONF]["tenant"] == "acme" - assert config[CONF]["thread_id"] == "t1" + assert result.output == {"ok": True} + message_events = [e for e in protocol_broker._events["t1"] if e["method"] == "messages"] + assert message_events, "expected protocol message events from on_chat_model_stream" + custom_events = [e for e in protocol_broker._events["t1"] if e["method"] == "custom"] + assert custom_events + assert custom_events[0]["params"]["data"] == {"custom": "payload"} + # Each raw astream_events() item is published onto the events channel. + raw_events = [e for e in protocol_broker._events["t1"] if e["method"] == "events"] + assert raw_events, "expected raw astream_events() items on the events channel" + assert len(raw_events) >= 4, "expected one events frame per astream_events() item" + + +class FakeSubgraphAggregateGraph(FakeGraph): + """Graph whose root-level on_chain_end result differs from the values chunk, + exercising the final-state capture fallback via aget_state root probe.""" + + async def astream(self, prepared_input: dict, config: dict, **kwargs): + self.configs.append(config) + yield (["sub:1"], "values", {"output": {"partial": True}}) + + async def aget_state(self, config: dict): + if not (config.get(CONF) or {}).get("checkpoint_ns"): + return SimpleNamespace(values={"output": {"final": True}}) + return SimpleNamespace(values=None) + + +class _FakeSubgraphAggregateEntry(FakeEntry): + graph = FakeSubgraphAggregateGraph() + + @staticmethod + def build_graph(_checkpointer=None, store=None) -> FakeSubgraphAggregateGraph: + return _FakeSubgraphAggregateEntry.graph + + +class _FakeSubgraphAggregateService(FakeLangGraphService): + def get_entry(self, _graph_id: str | None) -> _FakeSubgraphAggregateEntry: + return _FakeSubgraphAggregateEntry() @pytest.mark.asyncio -async def test_execute_run_derives_context_from_configurable(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_execute_run_subgraphs_namespace_uses_root_state_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + """With ``stream_subgraphs=True`` events are (ns, mode, chunk) triples and, + absent a root values chunk, the final state is captured from the checkpointer + root-namespace probe.""" fake_db = FakeDBManager() - FakeKwargsCapturingEntry.graph = FakeKwargsCapturingGraph() + protocol_broker = ThreadProtocolEventBroker() monkeypatch.setattr( "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeKwargsCapturingLangGraphService(), + lambda: _FakeSubgraphAggregateService(), ) monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) + monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) - await execute_run( + result = await execute_run( thread_id="t1", run_id="r1", - payload={"a": 1}, + payload={"msg": "hello"}, user_id="user-1", - kwargs={"config": {"configurable": {"tenant": "acme"}}}, + kwargs={"stream_modes": ["values"], "stream_subgraphs": True}, ) - config = FakeKwargsCapturingEntry.graph.configs[0] - assert config[CONF]["tenant"] == "acme" - assert config[CONF]["thread_id"] == "t1" + assert result.output == {"final": True} + value_events = [e for e in protocol_broker._events["t1"] if e["method"] == "values"] + assert value_events + assert value_events[0]["params"]["namespace"] == ["sub:1"] + + +class FakeInterruptResultGraph(FakeInterruptGraph): + """HITL interrupt: the interrupt arrives in the updates stream with a + non-empty state, so the run result must merge __interrupt__ into the final + values and emit input.requested.""" + + async def astream(self, prepared_input: dict, config: dict, **kwargs): + self.configs.append(config) + yield ( + "updates", + { + "__interrupt__": [ + type("Interrupt", (), {"value": "Provide value:", "id": "interrupt-1"})(), + ], + "foo": prepared_input["input"]["foo"], + }, + ) + + +class _FakeInterruptResultEntry(FakeEntry): + graph = FakeInterruptResultGraph() + + @staticmethod + def build_graph(_checkpointer=None, store=None) -> FakeInterruptResultGraph: + return _FakeInterruptResultEntry.graph + + @staticmethod + def extract_output(result: dict, _payload: dict) -> dict: + interrupts = result.get("__interrupt__", []) + return { + "state": result.get("foo"), + "interrupted": bool(interrupts), + "interrupts": [{"value": item.value, "id": item.id} for item in interrupts], + } + + +class _FakeInterruptResultService(FakeLangGraphService): + def get_entry(self, _graph_id: str | None) -> _FakeInterruptResultEntry: + return _FakeInterruptResultEntry() @pytest.mark.asyncio -async def test_execute_run_passes_context_kwarg_when_context_schema_present(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_execute_run_interrupt_merges_into_result_and_emits_input_requested( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """HITL interrupt with non-empty state: __interrupt__ is merged into the run + result (so extract_output sees it) and input.requested is emitted.""" fake_db = FakeDBManager() - FakeKwargsCapturingEntry.graph = FakeKwargsCapturingGraph() - FakeKwargsCapturingEntry.graph.context_schema = FakeContextSchema + protocol_broker = ThreadProtocolEventBroker() monkeypatch.setattr( "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeKwargsCapturingLangGraphService(), + lambda: _FakeInterruptResultService(), ) monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) + monkeypatch.setattr("agentseek_api.services.thread_protocol.thread_protocol_broker", protocol_broker) - await execute_run( + result = await execute_run( thread_id="t1", run_id="r1", - payload={"a": 1}, + payload={"foo": "hello"}, user_id="user-1", - kwargs={"context": {"tenant": "acme", "org": "ob"}}, + kwargs={"stream_modes": ["values"]}, ) - kwargs = FakeKwargsCapturingEntry.graph.stream_kwargs[0] - assert kwargs["context"] == {"tenant": "acme", "org": "ob"} + assert result.interrupted is True + assert result.interrupts == [{"value": "Provide value:", "id": "interrupt-1"}] + input_requested = [e for e in protocol_broker._events["t1"] if e["method"] == "input.requested"] + assert len(input_requested) == 1 + assert input_requested[0]["params"]["data"]["payload"] == "Provide value:" + + +class FakeTupleNamespaceGraph(FakeGraph): + """astream(subgraphs=True) yields tuple namespaces; they must be normalized + to lists before publication so the live broker's namespace filter (which + compares list slices against list prefixes) can match them.""" + + async def astream(self, prepared_input: dict, config: dict, **kwargs): + self.configs.append(config) + yield (("node_1:task-1",), "updates", {"step": "partial"}) + yield (("node_1:task-1",), "values", {"output": {"step": "final"}}) @pytest.mark.asyncio -async def test_execute_run_filters_context_by_schema(monkeypatch: pytest.MonkeyPatch) -> None: - fake_db = FakeDBManager() - FakeKwargsCapturingEntry.graph = FakeKwargsCapturingGraph() - FakeKwargsCapturingEntry.graph.context_schema = FakeContextSchema - monkeypatch.setattr( - "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeKwargsCapturingLangGraphService(), +async def test_execute_run_normalizes_tuple_namespaces_for_live_filter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + thread_events = await _run_fake_graph( + monkeypatch, FakeTupleNamespaceGraph(), stream_modes=["updates"], stream_subgraphs=True ) - monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) + updates_events = [event for event in thread_events if event["method"] == "updates"] + assert updates_events, "expected updates event from tuple-namespaced subgraph" + assert updates_events[0]["params"]["namespace"] == ["node_1:task-1"] + assert isinstance(updates_events[0]["params"]["namespace"], list) - await execute_run( - thread_id="t1", - run_id="r1", - payload={"a": 1}, - user_id="user-1", - kwargs={"context": {"tenant": "acme", "extra": "ignored"}}, - ) - kwargs = FakeKwargsCapturingEntry.graph.stream_kwargs[0] - assert kwargs["context"] == {"tenant": "acme"} +class FakeParallelIdlessMessagesGraph(FakeGraph): + """Two id-less messages from different subgraph namespaces, each with one + incremental chunk. + + ``AIMessageChunk`` has no ``id`` by default, so both go through the fallback + identity path. The namespaces differ, so the fallback id must keep the two + messages distinct (previously both collapsed to ``{run}:message:0`` and were + merged by the client). + """ + + async def astream(self, prepared_input: dict, config: dict, **kwargs): + self.configs.append(config) + yield (("ns_a:task-1",), "messages", (AIMessageChunk(content="hello", id=None), {"langgraph_node": "ns_a"})) + yield (("ns_b:task-1",), "messages", (AIMessageChunk(content="world", id=None), {"langgraph_node": "ns_b"})) + + async def aget_state(self, config: dict): + return SimpleNamespace(values={"output": {"ok": True}}) @pytest.mark.asyncio -async def test_execute_run_merges_context_and_configurable_without_error(monkeypatch: pytest.MonkeyPatch) -> None: - """Assistant-level context (merged into kwargs["context"]) must not conflict with - client-supplied config.configurable; both stay visible and no error is raised.""" - fake_db = FakeDBManager() - FakeKwargsCapturingEntry.graph = FakeKwargsCapturingGraph() - FakeKwargsCapturingEntry.graph.context_schema = FakeContextSchema - monkeypatch.setattr( - "agentseek_api.services.run_executor.get_langgraph_service", - lambda: FakeKwargsCapturingLangGraphService(), +async def test_execute_run_idless_messages_from_different_namespaces_get_distinct_ids( + monkeypatch: pytest.MonkeyPatch, +) -> None: + thread_events = await _run_fake_graph( + monkeypatch, FakeParallelIdlessMessagesGraph(), stream_modes=["messages"], stream_subgraphs=True ) - monkeypatch.setattr("agentseek_api.services.run_executor.db_manager", fake_db) - - await execute_run( - thread_id="t1", - run_id="r1", - payload={"a": 1}, - user_id="user-1", - kwargs={ - "config": {"configurable": {"client_param": "x"}}, - "context": {"tenant": "acme"}, - }, + metadata_events = [ + event for event in thread_events if event["method"] == "messages/metadata" + ] + assert len(metadata_events) == 2, ( + f"expected one messages/metadata per distinct id-less message, got {len(metadata_events)}" ) - - config = FakeKwargsCapturingEntry.graph.configs[0] - kwargs = FakeKwargsCapturingEntry.graph.stream_kwargs[0] - # configurable carries both the client configurable key and the context key - assert config[CONF]["client_param"] == "x" - assert config[CONF]["tenant"] == "acme" - assert config[CONF]["thread_id"] == "t1" - # context kwarg still reflects the effective context - assert kwargs["context"] == {"tenant": "acme"} + # The two metadata events must carry distinct message ids (their wire + # identity), so the SDK client routes them to two independent streams + # instead of merging the second chunk into the first message. + message_ids = [ + next(iter(event["params"]["data"].keys())) + for event in metadata_events + ] + assert len(message_ids) == len(set(message_ids)), f"id-less message ids collided: {message_ids}" + assert any("ns_a" in mid for mid in message_ids) + assert any("ns_b" in mid for mid in message_ids) diff --git a/tests/unit/test_run_jobs.py b/tests/unit/test_run_jobs.py index 28b39d0..858ecba 100644 --- a/tests/unit/test_run_jobs.py +++ b/tests/unit/test_run_jobs.py @@ -5,6 +5,7 @@ from agentseek_api.settings import settings from agentseek_api.services import run_jobs as run_jobs_module +from agentseek_api.services.thread_protocol import ThreadProtocolEventBroker class FakeSession: @@ -56,26 +57,6 @@ def _job(*, run_id: str = "r1", thread_id: str = "t1") -> run_jobs_module.RunExe ) -@pytest.mark.asyncio -async def test_persist_thread_snapshot_skips_duplicate_redis_writes(monkeypatch: pytest.MonkeyPatch) -> None: - persisted: list[tuple[str, dict[str, Any]]] = [] - monkeypatch.setattr(settings, "EXECUTOR_BACKEND", "redis") - monkeypatch.setattr( - run_jobs_module.thread_protocol_broker, - "snapshot_records", - lambda _thread_id: [{"seq": 1, "method": "values"}], - ) - - async def fake_persist(thread_id: str, event: dict[str, Any]) -> None: - persisted.append((thread_id, event)) - - monkeypatch.setattr(run_jobs_module, "persist_thread_stream_event", fake_persist) - - await run_jobs_module._persist_thread_snapshot("thread-1") - - assert persisted == [] - - @pytest.mark.asyncio async def test_publish_run_event_uses_atomic_redis_append(monkeypatch: pytest.MonkeyPatch) -> None: published: list[tuple[str, str, int | None, dict[str, Any]]] = [] @@ -190,17 +171,16 @@ async def fake_publish_run_event(_run_id: str, event: str, *, persist: bool = Tr operations.append(f"publish:{event}") return (1 if event == "start" else 2), {"event": event, **payload} - async def fake_persist_thread_snapshot(_thread_id: str) -> None: - return None async def fake_add_run_stream_event_to_session( _session: FakeSession, _run_id: str, *, - seq: int, + seq: int | None = None, payload: dict[str, Any], - ) -> None: + ) -> tuple[int | None, dict[str, Any]]: operations.append(f"persist:run:{payload['event']}:{seq}") + return seq or 0, payload def fake_publish_lifecycle_event( _thread_id: str, @@ -224,15 +204,15 @@ async def fake_add_thread_stream_event_to_session( _session: FakeSession, _thread_id: str, *, - seq: int, + seq: int | None = None, payload: dict[str, Any], - ) -> None: + ) -> tuple[int | None, dict[str, Any]]: operations.append(f"persist:thread:{payload['params']['data']['event']}:{seq}") + return seq or 0, payload monkeypatch.setattr(run_jobs_module.db_manager, "get_session_factory", lambda: session_factory) monkeypatch.setattr(run_jobs_module, "execute_run", successful_execute_run) monkeypatch.setattr(run_jobs_module, "_publish_run_event", fake_publish_run_event) - monkeypatch.setattr(run_jobs_module, "_persist_thread_snapshot", fake_persist_thread_snapshot) monkeypatch.setattr(run_jobs_module, "add_run_stream_event_to_session", fake_add_run_stream_event_to_session) monkeypatch.setattr(run_jobs_module, "publish_lifecycle_event", fake_publish_lifecycle_event) monkeypatch.setattr(run_jobs_module, "add_thread_stream_event_to_session", fake_add_thread_stream_event_to_session) @@ -242,10 +222,8 @@ async def fake_add_thread_stream_event_to_session( assert operations == [ "commit", "publish:start", - "publish:end", - "persist:run:end:2", - "publish:lifecycle:completed:False", - "persist:thread:completed:3", + "persist:run:end:None", + "persist:thread:completed:None", "commit", ] @@ -256,23 +234,25 @@ async def test_execute_run_job_publishes_failed_lifecycle_when_run_deleted( ) -> None: operations: list[str] = [] session_factory = FakeSessionFactory([FakeSession([None], operations)]) + protocol_broker = ThreadProtocolEventBroker() - def fake_publish_lifecycle_event( - _thread_id: str, *, event: str, graph_name: str | None = None, error: str | None = None, **_kw: Any, - ) -> dict[str, Any]: - operations.append(f"lifecycle:{event}:{error}") - return {"seq": 1, "method": "lifecycle", "params": {"namespace": [], "timestamp": 1, "data": {"event": event}}} - - async def fake_add_thread_stream_event_to_session(_session: Any, _thread_id: str, *, seq: int, payload: dict[str, Any]) -> None: + async def fake_add_thread_stream_event_to_session(_session: Any, _thread_id: str, *, seq: int | None = None, payload: dict[str, Any]) -> None: pass monkeypatch.setattr(run_jobs_module.db_manager, "get_session_factory", lambda: session_factory) - monkeypatch.setattr(run_jobs_module, "publish_lifecycle_event", fake_publish_lifecycle_event) + monkeypatch.setattr(run_jobs_module, "thread_protocol_broker", protocol_broker) monkeypatch.setattr(run_jobs_module, "add_thread_stream_event_to_session", fake_add_thread_stream_event_to_session) await run_jobs_module.execute_run_job(_job()) - assert any("lifecycle:failed:Run was deleted" in op for op in operations) + lifecycle_events = [ + event["params"]["data"] + for event in protocol_broker.snapshot_records("t1") + ] + assert any( + event.get("event") == "failed" and "Run was deleted" in (event.get("error") or "") + for event in lifecycle_events + ) def test_from_payload_rejects_unsupported_kind() -> None: @@ -308,17 +288,16 @@ async def fake_publish_run_event(_run_id: str, event: str, *, persist: bool = Tr operations.append(f"publish:{event}") return (1 if event == "start" else 2), {"event": event, **payload} - async def fake_persist_thread_snapshot(_thread_id: str) -> None: - return None async def fake_add_run_stream_event_to_session( _session: FakeSession, _run_id: str, *, - seq: int, + seq: int | None = None, payload: dict[str, Any], - ) -> None: + ) -> tuple[int | None, dict[str, Any]]: operations.append(f"persist:run:{payload['event']}:{seq}") + return seq or 0, payload def fake_publish_lifecycle_event( _thread_id: str, @@ -342,15 +321,15 @@ async def fake_add_thread_stream_event_to_session( _session: FakeSession, _thread_id: str, *, - seq: int, + seq: int | None = None, payload: dict[str, Any], - ) -> None: + ) -> tuple[int | None, dict[str, Any]]: operations.append(f"persist:thread:{payload['params']['data']['event']}:{seq}") + return seq or 0, payload monkeypatch.setattr(run_jobs_module.db_manager, "get_session_factory", lambda: session_factory) monkeypatch.setattr(run_jobs_module, "execute_run", interrupted_execute_run) monkeypatch.setattr(run_jobs_module, "_publish_run_event", fake_publish_run_event) - monkeypatch.setattr(run_jobs_module, "_persist_thread_snapshot", fake_persist_thread_snapshot) monkeypatch.setattr(run_jobs_module, "add_run_stream_event_to_session", fake_add_run_stream_event_to_session) monkeypatch.setattr(run_jobs_module, "publish_lifecycle_event", fake_publish_lifecycle_event) monkeypatch.setattr(run_jobs_module, "add_thread_stream_event_to_session", fake_add_thread_stream_event_to_session) @@ -358,4 +337,4 @@ async def fake_add_thread_stream_event_to_session( await run_jobs_module.execute_run_job(_job()) assert db_run.status == "interrupted" - assert "publish:lifecycle:interrupted:False" in operations + assert "persist:thread:interrupted:None" in operations diff --git a/tests/unit/test_run_preparation.py b/tests/unit/test_run_preparation.py index b980ba1..fb87d9e 100644 --- a/tests/unit/test_run_preparation.py +++ b/tests/unit/test_run_preparation.py @@ -216,14 +216,12 @@ async def fake_publish_run_event(_run_id: str, event: str, *, persist: bool = Tr async def fake_publish_lifecycle(*_args: Any, **_kwargs: Any) -> None: return None - async def fake_persist_thread_snapshot(_thread_id: str) -> None: - return None async def fake_add_run_stream_event_to_session( _session: FakeSession, _run_id: str, *, - seq: int, + seq: int | None = None, payload: dict[str, Any], ) -> None: operations.append(f"persist:{payload['event']}:{seq}") @@ -232,7 +230,6 @@ async def fake_add_run_stream_event_to_session( monkeypatch.setattr("agentseek_api.services.run_preparation.execute_run", successful_execute_run) monkeypatch.setattr("agentseek_api.services.run_preparation._publish_run_event", fake_publish_run_event) monkeypatch.setattr("agentseek_api.services.run_preparation._publish_lifecycle", fake_publish_lifecycle) - monkeypatch.setattr("agentseek_api.services.run_preparation._persist_thread_snapshot", fake_persist_thread_snapshot) monkeypatch.setattr( "agentseek_api.services.run_preparation.add_run_stream_event_to_session", fake_add_run_stream_event_to_session, @@ -249,13 +246,11 @@ async def fake_add_run_stream_event_to_session( assert exec_session.commits == 2 assert published_run_events == [ ("start", 1, {}), - ("end", 1, {"status": "success"}), ] assert operations == [ "commit", "publish:start", - "publish:end", - "persist:end:2", + "persist:end:None", "commit", ] @@ -313,15 +308,10 @@ async def test_prepare_run_cleans_protocol_state_when_submit_fails(monkeypatch: persist_session = CallbackSession([lambda: create_session.added[-1], fake_thread]) session_factory = FakeSessionFactory([create_session, persist_session]) protocol_broker = ThreadProtocolEventBroker() - published_lifecycle: list[dict[str, Any]] = [] monkeypatch.setattr("agentseek_api.services.run_preparation.db_manager.get_session_factory", lambda: session_factory) monkeypatch.setattr("agentseek_api.services.run_preparation.get_executor", lambda: RaisingExecutor()) - monkeypatch.setattr("agentseek_api.services.run_preparation.thread_protocol_broker", protocol_broker) - monkeypatch.setattr( - "agentseek_api.services.run_preparation.publish_lifecycle_event", - lambda thread_id, **payload: published_lifecycle.append({"thread_id": thread_id, **payload}), - ) + monkeypatch.setattr("agentseek_api.services.run_jobs.thread_protocol_broker", protocol_broker) with pytest.raises(RuntimeError, match="submit failed"): await run_prep_module.prepare_and_submit_run( @@ -336,16 +326,13 @@ async def test_prepare_run_cleans_protocol_state_when_submit_fails(monkeypatch: assert created_run.last_error == "submit failed" assert fake_thread.status == "error" assert protocol_broker._active_runs["t1"] == 0 - assert published_lifecycle == [ - {"thread_id": "t1", "event": "started", "graph_name": "default", "persist": False, "seq": None}, - { - "thread_id": "t1", - "event": "failed", - "graph_name": "default", - "error": "submit failed", - "persist": False, - "seq": None, - }, + lifecycle_events = [ + (event["params"]["data"]["event"], event["params"]["data"].get("error")) + for event in protocol_broker.snapshot_records("t1") + ] + assert lifecycle_events == [ + ("started", None), + ("failed", "submit failed"), ] @@ -354,14 +341,10 @@ async def test_execute_and_persist_cleans_protocol_state_for_cancelled_run(monke cancelled_run = type("DbRun", (), {"run_id": "r1", "status": "error", "last_error": "Run cancelled"})() session_factory = FakeSessionFactory([FakeSession([cancelled_run])]) protocol_broker = ThreadProtocolEventBroker() - published_lifecycle: list[dict[str, Any]] = [] monkeypatch.setattr("agentseek_api.services.run_preparation.db_manager.get_session_factory", lambda: session_factory) monkeypatch.setattr("agentseek_api.services.run_preparation.thread_protocol_broker", protocol_broker) - monkeypatch.setattr( - "agentseek_api.services.run_preparation.publish_lifecycle_event", - lambda thread_id, **payload: published_lifecycle.append({"thread_id": thread_id, **payload}), - ) + monkeypatch.setattr("agentseek_api.services.run_jobs.thread_protocol_broker", protocol_broker) protocol_broker.run_started("t1") await run_prep_module._execute_and_persist( @@ -373,16 +356,11 @@ async def test_execute_and_persist_cleans_protocol_state_for_cancelled_run(monke ) assert protocol_broker._active_runs["t1"] == 0 - assert published_lifecycle == [ - { - "thread_id": "t1", - "event": "failed", - "graph_name": "default", - "error": "Run cancelled", - "persist": False, - "seq": None, - } + lifecycle_events = [ + (event["params"]["data"]["event"], event["params"]["data"].get("error")) + for event in protocol_broker.snapshot_records("t1") ] + assert lifecycle_events == [("failed", "Run cancelled")] @pytest.mark.asyncio @@ -455,15 +433,10 @@ async def test_resume_run_restores_interrupted_state_when_submit_fails(monkeypat persist_session = CallbackSession([db_run, fake_thread]) session_factory = FakeSessionFactory([load_session, persist_session]) protocol_broker = ThreadProtocolEventBroker() - published_lifecycle: list[dict[str, Any]] = [] monkeypatch.setattr("agentseek_api.services.run_preparation.db_manager.get_session_factory", lambda: session_factory) monkeypatch.setattr("agentseek_api.services.run_preparation.get_executor", lambda: RaisingExecutor()) - monkeypatch.setattr("agentseek_api.services.run_preparation.thread_protocol_broker", protocol_broker) - monkeypatch.setattr( - "agentseek_api.services.run_preparation.publish_lifecycle_event", - lambda thread_id, **payload: published_lifecycle.append({"thread_id": thread_id, **payload}), - ) + monkeypatch.setattr("agentseek_api.services.run_jobs.thread_protocol_broker", protocol_broker) with pytest.raises(RuntimeError, match="submit failed"): await run_prep_module.resume_run( @@ -477,16 +450,13 @@ async def test_resume_run_restores_interrupted_state_when_submit_fails(monkeypat assert db_run.last_error == "submit failed" assert fake_thread.status == "interrupted" assert protocol_broker._active_runs["t1"] == 0 - assert published_lifecycle == [ - {"thread_id": "t1", "event": "started", "graph_name": "subgraph_hitl_agent", "persist": False, "seq": None}, - { - "thread_id": "t1", - "event": "failed", - "graph_name": "subgraph_hitl_agent", - "error": "submit failed", - "persist": False, - "seq": None, - }, + lifecycle_events = [ + (event["params"]["data"]["event"], event["params"]["data"].get("error")) + for event in protocol_broker.snapshot_records("t1") + ] + assert lifecycle_events == [ + ("started", None), + ("failed", "submit failed"), ] diff --git a/tests/unit/test_thread_protocol.py b/tests/unit/test_thread_protocol.py index 0ec63fa..decbfef 100644 --- a/tests/unit/test_thread_protocol.py +++ b/tests/unit/test_thread_protocol.py @@ -415,3 +415,48 @@ def test_message_chunk_finish_order_is_stable(monkeypatch: pytest.MonkeyPatch) - "content-block-finish", "message-finish", ] + + +@pytest.mark.asyncio +async def test_thread_protocol_stream_live_filter_rejects_tuple_and_accepts_list_namespace() -> None: + """A tuple subgraph namespace must be filtered OUT by the live broker's + namespace filter (the run_executor normalizes tuples to lists before + publication), while an equivalent list namespace must be delivered. + + Regression for the live path described in review: real ``astream( + subgraphs=True)`` yields tuple namespaces, and without normalization the + in-memory broker's prefix filter never matches them, silently dropping + subgraph events for subscribers using ``namespaces``. + """ + broker = ThreadProtocolEventBroker() + broker.run_started("thread-1") + + # Published before normalization would have happened: tuple namespace. + broker.publish( + "thread-1", + {"method": "updates", "params": {"namespace": ("node_1:task-1",), "timestamp": 1, "data": {"step": "partial"}}}, + persist=False, + ) + # Post-fix, run_executor publishes a list namespace. + broker.publish( + "thread-1", + {"method": "updates", "params": {"namespace": ["node_1:task-1"], "timestamp": 2, "data": {"step": "partial"}}}, + persist=False, + ) + broker.run_finished("thread-1") + + events = [ + event + async for event in broker.stream( + "thread-1", + channels=["updates"], + namespaces=[["node_1:task-1"]], + depth=None, + since=0, + ) + ] + + # Only the list-namespaced event survives the prefix filter; the tuple one + # is coerced to root by the live filter and never matches the prefix. + assert [event["seq"] for event in events] == [2] + assert all(isinstance(event["params"]["namespace"], list) for event in events) diff --git a/uv.lock b/uv.lock index 5d02a21..d61c8e1 100644 --- a/uv.lock +++ b/uv.lock @@ -86,7 +86,7 @@ requires-dist = [ { name = "langchain-oceanbase", specifier = "==0.6.0" }, { name = "langchain-oceanbase", extras = ["pyseekdb"], marker = "extra == 'embedded'", specifier = "==0.6.0" }, { name = "langchain-openai", specifier = ">=1.0.0" }, - { name = "langgraph", specifier = ">=1.0.3" }, + { name = "langgraph", specifier = ">=1.2.0" }, { name = "langgraph-sdk", specifier = ">=0.3.5" }, { name = "mcp", specifier = ">=1.27.1,<2" }, { name = "pydantic", specifier = ">=2.8.0" },