From 28d62f4199fcff6a315e5634878d55667a1fdeb5 Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 9 Jun 2026 18:20:14 +0100 Subject: [PATCH 1/2] feat: better tool call followup on compat mode --- application/api/answer/routes/base.py | 196 +++- application/api/v1/idempotency.py | 260 +++++ application/api/v1/routes.py | 140 ++- tests/api/answer/test_base_routes.py | 30 +- .../v1/test_reasoning_heartbeat_liveness.py | 599 ++++++++++ .../api/v1/test_v1_continuation_ownership.py | 235 ++++ tests/api/v1/test_v1_idempotency.py | 439 +++++++ .../api/v1/test_v1_tool_pause_finalization.py | 1020 +++++++++++++++++ 8 files changed, 2864 insertions(+), 55 deletions(-) create mode 100644 application/api/v1/idempotency.py create mode 100644 tests/api/v1/test_reasoning_heartbeat_liveness.py create mode 100644 tests/api/v1/test_v1_continuation_ownership.py create mode 100644 tests/api/v1/test_v1_idempotency.py create mode 100644 tests/api/v1/test_v1_tool_pause_finalization.py diff --git a/application/api/answer/routes/base.py b/application/api/answer/routes/base.py index 4b4a5d03..5afa1ead 100644 --- a/application/api/answer/routes/base.py +++ b/application/api/answer/routes/base.py @@ -192,6 +192,7 @@ class BaseAnswerResource: model_id: Optional[str] = None, model_user_id: Optional[str] = None, _continuation: Optional[Dict] = None, + finalize_tool_pause_as_complete: bool = False, ) -> Generator[str, None, None]: """ Generator function that streams the complete conversation response. @@ -213,6 +214,22 @@ class BaseAnswerResource: shared_token: Token for shared agent model_id: Model ID used for the request retrieved_docs: Pre-fetched documents for sources (optional) + finalize_tool_pause_as_complete: Stateless-tool-round mode for + the OpenAI-compatible ``/v1/chat/completions`` endpoint. + OpenAI clients resume a tool call by re-POSTing the full + message history (no slot for our ``reserved_message_id``), + so the server cannot rely on a *native* resume to finalize + a paused assistant turn. When ``True`` and the agent pauses + for a client-executed tool, the reserved row is finalized + as ``status="complete"`` (recording the emitted + ``tool_calls``) and the stream ends, instead of writing a + ``pending_tool_state`` record and early-returning a + non-terminal row. This guarantees a ``/v1`` tool round + never strands a ``pending``/``streaming`` row for the + reconciler to fail. Defaults to ``False``, which preserves + the native ``/stream`` + ``/api/answer`` pause/resume UX + byte-for-byte (still writes ``pending_tool_state``, leaves + the row non-terminal, and resumes natively). Yields: Server-sent event strings @@ -234,6 +251,9 @@ class BaseAnswerResource: # mid-stream still leaves the question queryable. Continuations # reuse the original placeholder. reserved_message_id: Optional[str] = None + # Intentional: a continuation round reserves no new WAL row, so on the + # stateless ``/v1`` path the intermediate tool rounds aren't persisted + # (only the first turn + the final answer turn are). Accepted as-is. wal_eligible = should_persist and not _continuation if wal_eligible: try: @@ -265,16 +285,32 @@ class BaseAnswerResource: if primary_llm is not None: primary_llm._request_id = request_id - # Flipped to ``streaming`` on first chunk; reconciler uses this - # to tell "never started" from "in flight". + # Flipped to ``streaming`` on the first ``answer``/``sources`` chunk; + # the reconciler reads ``status`` to tell "never started" from "in + # flight". This is a *status* signal only — it is intentionally + # decoupled from the heartbeat below, which is an "agent is alive / + # producing output" signal (a reasoning model can stream ``thought`` + # chunks for minutes before its first answer token, never marking + # ``streaming``, yet must still count as live). streaming_marked = False # Heartbeat goes into ``metadata.last_heartbeat_at`` (not # ``updated_at``, which reconciler-side writes share) and uses # ``time.monotonic`` so a blocked event loop can't fake fresh. + # ``heartbeat_message`` only touches non-terminal rows, so stamping a + # still-``pending`` row is safe and does NOT change its status. STREAM_HEARTBEAT_INTERVAL = 60 last_heartbeat_at = time.monotonic() def _mark_streaming_once() -> None: + """Flip the reserved row ``pending → streaming`` exactly once. + + Status-only: called on the first ``answer``/``sources`` chunk so + the reconciler can distinguish "never started" from "in flight". + It also re-stamps the heartbeat here for good measure, but the + heartbeat liveness no longer depends on this transition (see + ``_heartbeat_streaming``), so a thought-only reasoning phase that + never reaches this point still stays live. + """ nonlocal streaming_marked, last_heartbeat_at if streaming_marked or not reserved_message_id: return @@ -287,8 +323,8 @@ class BaseAnswerResource: "update_message_status streaming failed for %s", reserved_message_id, ) - # Seed last_heartbeat_at so watchdog doesn't fall back to `timestamp` - # (creation time) before the first STREAM_HEARTBEAT_INTERVAL tick. + # Re-stamp last_heartbeat_at on the transition too; harmless given + # the seed at generation start and the per-interval pump below. try: self.conversation_service.heartbeat_message( reserved_message_id, @@ -302,8 +338,26 @@ class BaseAnswerResource: last_heartbeat_at = time.monotonic() def _heartbeat_streaming() -> None: + """Pump the liveness heartbeat once per ``STREAM_HEARTBEAT_INTERVAL``. + + Deliberately gated on ``reserved_message_id`` only — NOT on + ``streaming_marked``. The loop calls this for *every* chunk + (including ``thought``/``metadata``), so a reasoning model that + streams only ``thought`` chunks while it "thinks" keeps a still- + ``pending`` row's ``last_heartbeat_at`` fresh and stays out of the + reconciler's staleness sweep. ``heartbeat_message`` only updates + non-terminal rows, so this never resurrects or restatuses a + terminal row. + + Residual: a model that emits NO chunks at all (not even + ``thought``) for longer than the reconciler threshold still goes + stale, because this pump only ticks when a chunk flows. Covering a + fully-silent stream would require a background-thread heartbeat or + a higher staleness threshold; both are out of scope here. The + realistic reasoning case (``thought`` chunks streaming) is covered. + """ nonlocal last_heartbeat_at - if not reserved_message_id or not streaming_marked: + if not reserved_message_id: return now_mono = time.monotonic() if now_mono - last_heartbeat_at < STREAM_HEARTBEAT_INTERVAL: @@ -423,9 +477,29 @@ class BaseAnswerResource: else: gen_iter = agent.gen(query=question) + # Seed a liveness heartbeat the moment generation starts, before + # the first chunk. The row is still ``pending`` here; this stamps a + # fresh ``last_heartbeat_at`` so a model that takes a while to emit + # its first token (or only streams ``thought`` chunks) is protected + # from the reconciler's staleness sweep from t=0 — not only from the + # first interval tick after the first answer chunk. + if reserved_message_id: + try: + self.conversation_service.heartbeat_message( + reserved_message_id, + ) + except Exception: + logger.exception( + "generation-start heartbeat seed failed for %s", + reserved_message_id, + ) + last_heartbeat_at = time.monotonic() + for line in gen_iter: - # Cheap closure check that only hits the DB when the - # heartbeat interval has elapsed. + # Cheap closure check that only hits the DB when the heartbeat + # interval has elapsed. Runs for *every* chunk (incl. ``thought`` + # / ``metadata``), so a still-``pending`` reasoning stream stays + # live without waiting for the ``streaming`` status flip. _heartbeat_streaming() if "metadata" in line: query_metadata.update(line["metadata"]) @@ -490,6 +564,35 @@ class BaseAnswerResource: # ---- Paused: save continuation state and end stream early ---- if paused: continuation = getattr(agent, "_pending_continuation", None) + + # ---- Stateless-tool-round mode (OpenAI-compatible /v1) ---- + # OpenAI clients resume by re-POSTing the whole message + # history with ``{role:"tool"}`` results — there is no slot + # for our ``reserved_message_id``, so a *native* resume can't + # finalize this paused turn. Finalize the reserved row as + # ``complete`` here (recording the emitted tool_calls) and end + # the stream, so the reconciler never sees a non-terminal row. + # The client still gets ``finish_reason:"tool_calls"`` + the + # calls from the ``tool_calls_pending`` event yielded above. + if finalize_tool_pause_as_complete: + yield from self._finalize_stateless_tool_pause( + continuation=continuation, + reserved_message_id=reserved_message_id, + conversation_id=conversation_id, + question=question, + response_full=response_full, + thought=thought, + source_log_docs=source_log_docs, + tool_calls=tool_calls, + query_metadata=query_metadata, + model_id=model_id, + should_persist=should_persist, + emit=_emit, + ) + if journal_writer is not None: + journal_writer.close() + return + if continuation: # First-turn pause needs a conversation row to attach to. if not conversation_id and should_persist: @@ -1026,6 +1129,85 @@ class BaseAnswerResource: journal_writer.close() return + def _finalize_stateless_tool_pause( + self, + *, + continuation: Optional[Dict[str, Any]], + reserved_message_id: Optional[str], + conversation_id: Optional[str], + question: str, + response_full: str, + thought: str, + source_log_docs: List[Dict[str, Any]], + tool_calls: List[Dict[str, Any]], + query_metadata: Dict[str, Any], + model_id: Optional[str], + should_persist: bool, + emit: Any, + ) -> Generator[str, None, None]: + """Finalize a client-tool pause as ``complete`` for the ``/v1`` path. + + Used only when ``complete_stream`` runs with + ``finalize_tool_pause_as_complete=True`` (the OpenAI-compatible + ``/v1/chat/completions`` endpoint). Records the emitted/pending + ``tool_calls`` on the reserved row and flips it to ``complete`` so the + reconciler never sweeps it, then yields the terminal ``id``/``end`` + events. No ``pending_tool_state`` is written: an OpenAI client resumes + statelessly (re-POSTing the full history) rather than via a native + resume, so there is no server-side continuation record to load. + + Args: + continuation: The agent's ``_pending_continuation`` (may be None). + reserved_message_id: WAL placeholder row id, if one was reserved. + conversation_id: The conversation id to surface to the client. + question: The user's question for this turn. + response_full: Any assistant text produced before the pause. + thought: Reasoning tokens produced before the pause. + source_log_docs: Retrieval sources gathered before the pause. + tool_calls: Tool-call events emitted during this turn. + query_metadata: Accumulated stream metadata. + model_id: Model id used for the request. + should_persist: Whether persistence is enabled for this request. + emit: The stream's ``_emit`` callable for SSE framing/journaling. + + Yields: + The terminal ``id`` and ``end`` SSE event strings. + """ + # Prefer the structured pending tool calls (carry call_id / name / + # arguments) so the persisted row is a coherent record of what the + # client was asked to execute; fall back to whatever ``tool_calls`` + # events were emitted. + pending_tool_calls = ( + continuation.get("pending_tool_calls") if continuation else None + ) + tool_calls_to_persist = pending_tool_calls or tool_calls or [] + + if should_persist and reserved_message_id is not None: + try: + self.conversation_service.finalize_message( + reserved_message_id, + response_full, + thought=thought, + sources=source_log_docs, + tool_calls=tool_calls_to_persist, + model_id=model_id or self.default_model_id, + metadata=query_metadata if query_metadata else None, + status="complete", + ) + except Exception as e: + logger.error( + f"Failed to finalize stateless tool pause as complete " + f"for message_id={reserved_message_id}: {e}", + exc_info=True, + ) + # When there is no reserved row (stateless OpenAI round with no + # conversation_id — the translator sets persist=false), there is + # nothing durable to finalize and nothing stranded: just end cleanly + # without writing an empty-prompt orphan conversation. + + yield emit({"type": "id", "id": str(conversation_id)}) + yield emit({"type": "end"}) + def process_response_stream(self, stream) -> Dict[str, Any]: """Process the stream response for non-streaming endpoint. diff --git a/application/api/v1/idempotency.py b/application/api/v1/idempotency.py new file mode 100644 index 00000000..7e1b4d7b --- /dev/null +++ b/application/api/v1/idempotency.py @@ -0,0 +1,260 @@ +"""Layer-1 idempotency for the OpenAI-compatible ``/v1/chat/completions`` route. + +The ``/v1`` tool round-trip is fully stateless (the pause finalizes the prior +turn as ``complete`` and resumes via ``build_continuation_from_messages`` with +no ``pending_tool_state``). Dropping the native ``resume_from_tool_actions`` +path also dropped its ``mark_resuming`` guard, so a duplicated/retried POST +could re-run the agent → a duplicate answer row + double token billing. + +This module restores protection the OpenAI-compatible way: a client-supplied +``Idempotency-Key`` header makes retries return the *stored first response* +instead of re-running the agent. It is opt-in (no header → today's behavior, +byte-for-byte) and scoped to **non-streaming** requests only (the b2b client +and the actual regression); streaming replay is intentionally unsupported. + +Storage reuses the existing ``task_dedup`` table via +:class:`~application.storage.db.repositories.idempotency.IdempotencyRepository` +— no new table or migration. The contract maps onto its claim/finalize +semantics: + +- **No record** → ``claim_task`` inserts a ``pending`` row (we run + finalize). +- **``completed`` within 24h TTL** → return the cached body + status code. +- **Fresh ``pending``** (in-flight) → HTTP 409 idempotency conflict. +- **Stale ``pending``** (older than :data:`STALE_PENDING_SECONDS` — the + original request likely died) → release and re-claim. +- **``failed`` / past-TTL** → ``claim_task`` re-claims automatically. + +Only successful (2xx) responses are cached; a 4xx/5xx releases the claim so a +genuine retry can still succeed (matches OpenAI). +""" + +from __future__ import annotations + +import logging +import uuid +from typing import Any, Dict, Optional, Tuple + +from flask import jsonify, make_response, request, Response +from sqlalchemy import text as sql_text + +from application.storage.db.repositories.idempotency import IdempotencyRepository +from application.storage.db.session import db_readonly, db_session + +logger = logging.getLogger(__name__) + +# Distinct ``task_name`` so v1 chat dedup rows never collide with ingest / +# webhook rows that share the ``task_dedup`` table. +TASK_NAME = "v1_chat_completion" + +_IDEMPOTENCY_KEY_MAX_LEN = 256 + +# A ``pending`` claim older than this is treated as a dead in-flight request +# (the process crashed before finalize), so a genuine retry may re-claim it +# rather than waiting out the full 24h TTL or getting a permanent 409. Kept +# short enough to unblock retries quickly, long enough that a normal +# non-streaming completion (which finalizes on the same request) never trips +# it while still running. +STALE_PENDING_SECONDS = 300 + + +def read_idempotency_key() -> Tuple[Optional[str], Optional[Response]]: + """Read and validate the ``Idempotency-Key`` request header. + + Returns: + ``(key, error_response)``. An absent/empty header yields + ``(None, None)`` (idempotency is opt-in). An oversized header yields + ``(None, <400 response>)`` so the caller can short-circuit. + """ + key = request.headers.get("Idempotency-Key") + if not key: + return None, None + if len(key) > _IDEMPOTENCY_KEY_MAX_LEN: + return None, make_response( + jsonify( + { + "error": { + "message": ( + f"Idempotency-Key exceeds maximum length of " + f"{_IDEMPOTENCY_KEY_MAX_LEN} characters" + ), + "type": "invalid_request", + } + } + ), + 400, + ) + return key, None + + +def scoped_key(idempotency_key: Optional[str], agent_id: Optional[str]) -> Optional[str]: + """Compose ``{agent_id}:{idempotency_key}`` so tenants never collide. + + Two agents replaying the same key value resolve to distinct stored rows. + Falls back to ``api_key`` scoping at the call site when no agent id is + available; returns ``None`` when either component is missing (idempotency + is then skipped, preserving today's behavior). + """ + if not idempotency_key or not agent_id: + return None + return f"{agent_id}:{idempotency_key}" + + +def _release_stale_pending(key: str) -> None: + """Delete a stale ``pending`` claim so the caller can re-claim it. + + Scoped to ``status = 'pending'`` and the staleness window so we never + clobber a live in-flight claim or a ``completed`` cache row. + """ + try: + with db_session() as conn: + conn.execute( + sql_text( + "DELETE FROM task_dedup " + "WHERE idempotency_key = :k " + "AND status = 'pending' " + "AND created_at <= now() - make_interval(secs => :secs)" + ), + {"k": key, "secs": STALE_PENDING_SECONDS}, + ) + except Exception: + logger.exception("Failed to release stale v1 idempotency claim for key=%s", key) + + +def claim_or_replay(key: str) -> Tuple[bool, Optional[Response]]: + """Claim ``key`` for this request, or return the prior outcome. + + Claim-before-process: atomically insert a ``pending`` row. The three + outcomes map onto the existing ``task_dedup`` contract: + + - **claimed** → ``(True, None)``: this caller runs the agent and must call + :func:`finalize` (success) or :func:`release` (error) afterwards. + - **``completed`` within TTL** → ``(False, )``: replay the + stored body + status code without re-running. + - **fresh ``pending``** → ``(False, <409 response>)``: a same-key request is + already in progress. + + A ``pending`` row older than :data:`STALE_PENDING_SECONDS` is released and + re-claimed (the original request likely died). ``failed`` / past-TTL rows + are re-claimed by ``claim_task`` itself. + + Args: + key: The tenant-scoped idempotency key. + + Returns: + ``(claimed, response)``. When ``claimed`` is True the caller owns the + run; otherwise ``response`` is the replay/409 to return immediately. + """ + predetermined_id = str(uuid.uuid4()) + with db_session() as conn: + claimed = IdempotencyRepository(conn).claim_task( + key=key, task_name=TASK_NAME, task_id=predetermined_id, + ) + if claimed is not None: + return True, None + + # Lost the claim — resolve why against the within-TTL row. + with db_readonly() as conn: + existing = IdempotencyRepository(conn).get_task(key) + + if existing is not None and existing.get("status") == "completed": + return False, _replay_response(existing.get("result_json")) + + if existing is not None and existing.get("status") == "pending": + # In-flight? Re-claim only if the prior claim is stale (dead request). + _release_stale_pending(key) + with db_session() as conn: + reclaimed = IdempotencyRepository(conn).claim_task( + key=key, task_name=TASK_NAME, task_id=predetermined_id, + ) + if reclaimed is not None: + return True, None + return False, _conflict_response() + + # Row vanished between claim and read (TTL cleanup / release race) — one + # more claim attempt; treat a persistent loss as a conflict. + with db_session() as conn: + reclaimed = IdempotencyRepository(conn).claim_task( + key=key, task_name=TASK_NAME, task_id=predetermined_id, + ) + if reclaimed is not None: + return True, None + with db_readonly() as conn: + existing = IdempotencyRepository(conn).get_task(key) + if existing is not None and existing.get("status") == "completed": + return False, _replay_response(existing.get("result_json")) + return False, _conflict_response() + + +def finalize(key: str, response: Response) -> None: + """Cache a successful (2xx) response under ``key``; release otherwise. + + Stores ``{"status_code", "body"}`` in ``task_dedup.result_json`` so a + retry replays byte-for-byte. Non-2xx responses are not cached — the claim + is released so a genuine retry can still succeed (matches OpenAI). + + Args: + key: The tenant-scoped idempotency key claimed by :func:`claim_or_replay`. + response: The Flask response produced by running the request. + """ + status_code = response.status_code + if not (200 <= status_code < 300): + release(key) + return + try: + body = response.get_json(silent=True) + except Exception: + body = None + result_json = {"status_code": status_code, "body": body} + try: + with db_session() as conn: + IdempotencyRepository(conn).finalize_task( + key=key, result_json=result_json, status="completed", + ) + except Exception: + logger.exception("Failed to finalize v1 idempotency record for key=%s", key) + + +def release(key: str) -> None: + """Drop this request's ``pending`` claim so a retry can re-claim it. + + Used on the error path (non-2xx or an exception before finalize) so a + failed first attempt never blocks a legitimate retry for the full TTL. + """ + try: + with db_session() as conn: + conn.execute( + sql_text( + "DELETE FROM task_dedup " + "WHERE idempotency_key = :k AND status = 'pending'" + ), + {"k": key}, + ) + except Exception: + logger.exception("Failed to release v1 idempotency claim for key=%s", key) + + +def _replay_response(result_json: Optional[Dict[str, Any]]) -> Response: + """Rebuild a Flask response from a cached ``result_json`` row.""" + status_code = 200 + body: Any = None + if isinstance(result_json, dict): + status_code = int(result_json.get("status_code", 200)) + body = result_json.get("body") + return make_response(jsonify(body), status_code) + + +def _conflict_response() -> Response: + """OpenAI-shaped 409 for a same-key request already in progress.""" + return make_response( + jsonify( + { + "error": { + "message": ( + "A request with this Idempotency-Key is already in progress" + ), + "type": "idempotency_conflict", + } + } + ), + 409, + ) diff --git a/application/api/v1/routes.py b/application/api/v1/routes.py index e0084187..e40afddf 100644 --- a/application/api/v1/routes.py +++ b/application/api/v1/routes.py @@ -17,6 +17,7 @@ from flask import Blueprint, jsonify, make_response, request, Response from application.api.answer.routes.base import BaseAnswerResource from application.api.answer.services.persistence_policy import resolve_persistence from application.api.answer.services.stream_processor import StreamProcessor +from application.api.v1 import idempotency as v1_idempotency from application.api.v1.translator import ( translate_request, translate_response, @@ -81,6 +82,28 @@ def chat_completions(): agent_doc = _lookup_agent(api_key) model_name = _get_model_name(agent_doc, api_key) + # ---- Layer-1 idempotency (opt-in, non-streaming only) ---- + # An ``Idempotency-Key`` header makes a retried non-streaming request + # return the stored first response instead of re-running the agent + # (restoring the guard lost when the v1 tool round dropped the native + # ``resume_from_tool_actions`` / ``mark_resuming`` path → would otherwise + # duplicate the answer row and double-bill tokens). Streaming replay is + # intentionally NOT supported (see the ``is_stream`` branch below), so we + # only resolve a key for non-streaming requests. No header → byte-for-byte + # today's behavior. + idem_key: Optional[str] = None + if not is_stream: + raw_key, key_error = v1_idempotency.read_idempotency_key() + if key_error is not None: + return key_error + # Scope per tenant: ``{agent_id}:{key}`` so two agents using the same + # key value never collide. Fall back to api_key scoping when the agent + # has no resolvable id (idempotency still keyed, just per api_key). + agent_scope = None + if agent_doc is not None: + agent_scope = str(agent_doc.get("id") or agent_doc.get("_id") or "") or None + idem_key = v1_idempotency.scoped_key(raw_key, agent_scope or api_key) + try: internal_data = translate_request(data, api_key) except Exception as e: @@ -105,34 +128,40 @@ def chat_completions(): processor = StreamProcessor(internal_data, decoded_token) if internal_data.get("tool_actions"): - # Continuation mode + # Continuation mode — coherent Option B: the v1 tool round-trip is + # fully stateless. The pause finalized the prior turn's row as + # ``complete`` and wrote NO ``pending_tool_state`` (see + # ``complete_stream(finalize_tool_pause_as_complete=True)``), so we + # ALWAYS rebuild the agent + pending calls from the re-POSTed + # message history — even when the client threads back the + # ``conversation_id`` it got from the first response. + # + # We deliberately do NOT call ``resume_from_tool_actions`` here: + # its ``load_state`` would find no pending state and raise (→ HTTP + # 400), since OpenAI clients resume statelessly rather than via a + # native resume. ``resume_from_tool_actions`` stays in place for + # the native ``/stream`` + ``/api/answer`` routes, which are + # unchanged. conversation_id = internal_data.get("conversation_id") + ( + agent, + messages, + tools_dict, + pending_tool_calls, + tool_actions, + reasoning_content, + ) = processor.build_continuation_from_messages( + internal_data.get("messages", []), + internal_data["tool_actions"], + ) + # When a conversation_id is carried, target it for persistence so + # the final answer appends as a NEW terminal turn in that + # conversation (``save_conversation`` keys off ``conversation_id``) + # rather than creating an orphan sibling. ``build_continuation_from_messages`` + # leaves the processor's ``conversation_id`` (set from the request + # in ``__init__``) intact; pin it explicitly for clarity. if conversation_id: - ( - agent, - messages, - tools_dict, - pending_tool_calls, - tool_actions, - reasoning_content, - ) = processor.resume_from_tool_actions( - internal_data["tool_actions"], conversation_id - ) - else: - # Stateless continuation: OpenAI-compatible clients (opencode, - # etc.) resend the full messages array but no conversation_id, - # so rebuild the agent + pending calls from the request itself. - ( - agent, - messages, - tools_dict, - pending_tool_calls, - tool_actions, - reasoning_content, - ) = processor.build_continuation_from_messages( - internal_data.get("messages", []), - internal_data["tool_actions"], - ) + processor.conversation_id = conversation_id continuation = { "messages": messages, "tools_dict": tools_dict, @@ -174,6 +203,10 @@ def chat_completions(): ) if is_stream: + # Idempotency replay is NOT supported for streaming: there is no + # safe way to re-emit a recorded SSE stream (and the regression / + # b2b client is non-streaming), so a streaming request never + # claims a key. This is a known, accepted limitation. return Response( _stream_response( helper, @@ -192,20 +225,41 @@ def chat_completions(): "X-Accel-Buffering": "no", }, ) - else: - return _non_stream_response( - helper, - question, - agent, - processor, - model_name, - continuation, - should_persist, - visibility, - strip_reasoning_leak, - ) + + # ---- Non-streaming: claim-before-process, then finalize/release ---- + # Claim happens here (after auth + agent resolution + continuation + # build, immediately before running the agent) so a duplicate retry + # short-circuits to the cached body / 409 instead of re-running. + if idem_key: + claimed, replay = v1_idempotency.claim_or_replay(idem_key) + if not claimed: + # ``completed`` cache hit, or a 409 for an in-flight same-key + # request — either way return without re-running the agent. + return replay + + # An exception from the agent run propagates to the ``except`` handlers + # below, which release the claim so a genuine retry can re-claim. + response = _non_stream_response( + helper, + question, + agent, + processor, + model_name, + continuation, + should_persist, + visibility, + strip_reasoning_leak, + ) + + # Cache only successful (2xx) responses; ``finalize`` releases the + # claim on a non-2xx so a real retry can still succeed (matches OpenAI). + if idem_key: + v1_idempotency.finalize(idem_key, response) + return response except ValueError as e: + if idem_key: + v1_idempotency.release(idem_key) logger.error( f"/v1/chat/completions error: {e} - {traceback.format_exc()}", extra={"error": str(e)}, @@ -215,6 +269,8 @@ def chat_completions(): 400, ) except Exception as e: + if idem_key: + v1_idempotency.release(idem_key) logger.error( f"/v1/chat/completions error: {e} - {traceback.format_exc()}", extra={"error": str(e)}, @@ -251,6 +307,10 @@ def _stream_response( should_persist=should_persist, visibility=visibility, _continuation=continuation, + # OpenAI clients resume tool calls statelessly (no slot for our + # reserved_message_id), so a tool pause must finalize the row as + # ``complete`` here rather than stranding it for a native resume. + finalize_tool_pause_as_complete=True, ) for line in internal_stream: @@ -314,6 +374,10 @@ def _non_stream_response( should_persist=should_persist, visibility=visibility, _continuation=continuation, + # OpenAI clients resume tool calls statelessly (no slot for our + # reserved_message_id), so a tool pause must finalize the row as + # ``complete`` here rather than stranding it for a native resume. + finalize_tool_pause_as_complete=True, ) result = helper.process_response_stream(stream) diff --git a/tests/api/answer/test_base_routes.py b/tests/api/answer/test_base_routes.py index 8da90cc3..7b6fdaaa 100644 --- a/tests/api/answer/test_base_routes.py +++ b/tests/api/answer/test_base_routes.py @@ -1181,12 +1181,18 @@ class TestCompleteStreamWalAcceptance: @pytest.mark.unit class TestStreamingHeartbeatSeed: - """Regression guard: when the row flips to ``streaming`` we must seed - ``last_heartbeat_at`` so the watchdog doesn't fall back to - ``timestamp`` (creation time) on slow LLM cold-starts (>idle_secs). + """Regression guard: the reserved row must carry a fresh + ``last_heartbeat_at`` from generation start so the reconciler doesn't fall + back to ``timestamp`` (creation time) on slow LLM cold-starts or while a + reasoning model streams only ``thought`` chunks. The heartbeat is seeded + once before the first chunk and re-stamped when the row flips to + ``streaming``; the ``pending → streaming`` status transition itself still + fires exactly once on the first ``answer``/``sources`` chunk. """ - def test_heartbeat_seeded_on_first_chunk(self, mock_mongo_db, flask_app): + def test_heartbeat_seeded_at_generation_start_and_on_first_chunk( + self, mock_mongo_db, flask_app, + ): from application.api.answer.routes.base import BaseAnswerResource with flask_app.app_context(): @@ -1221,7 +1227,7 @@ class TestStreamingHeartbeatSeed: ) # update_message_status flips the row to ``streaming`` exactly - # once (idempotent via streaming_marked). + # once (idempotent via streaming_marked) on the first answer chunk. status_calls = [ c for c in resource.conversation_service .update_message_status.call_args_list @@ -1230,14 +1236,18 @@ class TestStreamingHeartbeatSeed: assert len(status_calls) == 1 assert status_calls[0].args[0] == "msg1" - # heartbeat seed runs exactly once at the same flip — multiple - # chunks don't re-stamp inside this window (the throttled - # _heartbeat_streaming path is gated by STREAM_HEARTBEAT_INTERVAL). + # Heartbeat is stamped twice for this 2-answer-chunk stream: once at + # the generation-start seed (before the first chunk, while still + # ``pending``) and once when the first answer chunk flips the row to + # ``streaming``. The second answer chunk does NOT re-stamp — the + # per-chunk ``_heartbeat_streaming`` pump is throttled by + # STREAM_HEARTBEAT_INTERVAL and the two chunks fall inside one + # interval under real ``time.monotonic``. hb_calls = ( resource.conversation_service.heartbeat_message.call_args_list ) - assert len(hb_calls) == 1 - assert hb_calls[0].args[0] == "msg1" + assert len(hb_calls) == 2 + assert all(c.args[0] == "msg1" for c in hb_calls) def test_heartbeat_seed_skipped_without_reserved_message_id( self, mock_mongo_db, flask_app, diff --git a/tests/api/v1/test_reasoning_heartbeat_liveness.py b/tests/api/v1/test_reasoning_heartbeat_liveness.py new file mode 100644 index 00000000..2fc843ab --- /dev/null +++ b/tests/api/v1/test_reasoning_heartbeat_liveness.py @@ -0,0 +1,599 @@ +"""Streaming-heartbeat liveness for reasoning models that "think" before +answering (``complete_stream`` in ``application/api/answer/routes/base.py``). + +Background +---------- +The reconciler (``ReconciliationRepository.find_and_lock_stuck_messages``) +fails a ``pending``/``streaming`` row whose effective freshness — +``GREATEST(timestamp, message_metadata.last_heartbeat_at)`` — is older than +5 minutes. ``complete_stream`` keeps that freshness up via a streaming +heartbeat. + +The bug these tests pin: the heartbeat pump used to early-return until the row +had been flipped to ``streaming``, and the row is only flipped on the first +``answer``/``sources`` chunk. A reasoning model (e.g. ``reasoning_effort: +high``) that streams only ``thought`` chunks for minutes before its first +answer token therefore left the row ``pending`` with a *frozen* heartbeat, so +the reconciler would falsely fail a live request. + +The fix makes the heartbeat a true "agent is producing output / is alive" +signal: it pumps whenever a ``reserved_message_id`` exists (regardless of the +``streaming`` status flip), and is seeded once at generation start. The +``pending → streaming`` status transition itself is unchanged (still driven by +the first ``answer``/``sources`` chunk). + +Residual (documented, intentionally not covered here): a model that emits *no* +chunks at all — not even ``thought`` — for >5 min would still go stale, since +the pump only ticks when a chunk flows. Covering a fully-silent stream would +need a background-thread heartbeat or a higher threshold; both are out of +scope for this surgical fix. + +These tests drive the real ``complete_stream`` against an ephemeral Postgres +database (``pg_engine``) and reuse the wiring/fakes style of +``tests/api/v1/test_v1_tool_pause_finalization.py``. +""" + +from __future__ import annotations + +import uuid +from contextlib import contextmanager +from typing import Any, Dict, List, Optional + +import pytest +from sqlalchemy import text + +from application.api.answer.routes.base import BaseAnswerResource +from application.api.answer.services.conversation_service import ConversationService +from application.storage.db.repositories.conversations import ConversationsRepository +from application.storage.db.repositories.reconciliation import ( + ReconciliationRepository, +) + + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class _FakeLLM: + """Minimal LLM stand-in so ``complete_stream`` can stamp ``_request_id``.""" + + def __init__(self) -> None: + self._request_id: Optional[str] = None + self.model_id = "gpt-5" + + +class _ThoughtOnlyAgent: + """Agent that streams only ``thought`` chunks and never answers. + + Models a reasoning model mid-"thinking" phase: the stream is alive and + producing output, but no ``answer``/``sources`` chunk has arrived yet, so + the row is still ``pending`` (``streaming`` not marked). + """ + + def __init__(self, n_thoughts: int = 4) -> None: + self.llm = _FakeLLM() + self._n_thoughts = n_thoughts + + def gen(self, query: str = ""): + for i in range(self._n_thoughts): + yield {"thought": f"reasoning step {i}..."} + + +class _ThoughtThenAnswerAgent: + """Agent that thinks (``thought`` chunks) then emits a final ``answer``. + + Used for the regression case: the row must still flip to ``streaming`` on + the first ``answer`` chunk and finalize ``complete``. + """ + + ANSWER_TEXT = "The answer is 42." + + def __init__(self, n_thoughts: int = 3) -> None: + self.llm = _FakeLLM() + self._n_thoughts = n_thoughts + + def gen(self, query: str = ""): + for i in range(self._n_thoughts): + yield {"thought": f"reasoning step {i}..."} + yield {"answer": self.ANSWER_TEXT} + + +class _NoopJournalWriter: + """No-op journal writer so the test exercises DB row state, not the + message-events journal / Redis broadcast (orthogonal to heartbeats).""" + + def __init__(self, *args, **kwargs) -> None: + pass + + def record(self, *args, **kwargs) -> None: + pass + + def flush(self) -> None: + pass + + def close(self) -> None: + pass + + +# --------------------------------------------------------------------------- +# Test harness +# --------------------------------------------------------------------------- + + +def _seed_user(conn, user_id: str) -> None: + conn.execute( + text( + "INSERT INTO users (user_id) VALUES (:u) " + "ON CONFLICT (user_id) DO NOTHING" + ), + {"u": user_id}, + ) + + +def _row(conn, message_id: str) -> Dict[str, Any]: + """Fetch ``(status, message_metadata)`` for a reserved row.""" + r = conn.execute( + text( + "SELECT status, message_metadata FROM conversation_messages " + "WHERE id = CAST(:m AS uuid)" + ), + {"m": message_id}, + ).fetchone() + return {"status": r[0], "metadata": r[1]} if r is not None else {} + + +@contextmanager +def _wire_db(engine, monkeypatch): + """Point conversation_service / continuation_service / base at ``engine``. + + Each helper opens its own short-lived connection (matching production), + so we hand out fresh connections from the same ephemeral engine and swap + the journal writer for a no-op. Mirrors the helper in + ``test_v1_tool_pause_finalization.py``. + """ + from application.api.answer.services import conversation_service as conv_mod + from application.api.answer.services import continuation_service as cont_mod + from application.api.answer.routes import base as base_mod + + @contextmanager + def _session(): + conn = engine.connect() + txn = conn.begin() + try: + yield conn + txn.commit() + except Exception: + txn.rollback() + raise + finally: + conn.close() + + @contextmanager + def _readonly(): + conn = engine.connect() + try: + yield conn + finally: + conn.close() + + monkeypatch.setattr(conv_mod, "db_session", _session) + monkeypatch.setattr(conv_mod, "db_readonly", _readonly) + monkeypatch.setattr(cont_mod, "db_session", _session) + monkeypatch.setattr(cont_mod, "db_readonly", _readonly) + monkeypatch.setattr(base_mod, "db_session", _session) + monkeypatch.setattr(base_mod, "db_readonly", _readonly) + monkeypatch.setattr(base_mod, "BatchedJournalWriter", _NoopJournalWriter) + monkeypatch.setattr(base_mod, "record_event", lambda *a, **kw: None) + monkeypatch.setattr(base_mod, "publish_user_event", lambda *a, **kw: None) + yield + + +class _FakeMonotonic: + """Deterministic ``time.monotonic`` that jumps forward on every call. + + ``complete_stream`` reads ``time.monotonic()`` once per loop iteration + (via ``_heartbeat_streaming``) plus at seed/mark points. Advancing by more + than ``STREAM_HEARTBEAT_INTERVAL`` (60s) on each read guarantees the + interval gate fires every iteration, so a thought-only stream attempts a + heartbeat on every ``thought`` chunk. + """ + + def __init__(self, step: float = 120.0) -> None: + self._t = 0.0 + self._step = step + + def __call__(self) -> float: + self._t += self._step + return self._t + + +def _make_base() -> BaseAnswerResource: + base = BaseAnswerResource.__new__(BaseAnswerResource) + base.default_model_id = "gpt-5" + base.conversation_service = ConversationService() + return base + + +def _drain(gen) -> List[str]: + return list(gen) + + +# --------------------------------------------------------------------------- +# 1. Pump-during-reasoning: a thought-only stream heartbeats while pending +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestHeartbeatPumpsDuringReasoning: + """A reasoning model that only streams ``thought`` chunks (no answer yet) + must still keep its ``last_heartbeat_at`` fresh while the row is ``pending`` + — otherwise the reconciler falsely fails a live request. + + Pre-fix the pump early-returns while ``streaming`` is unmarked, so a + thought-only phase produces NO heartbeat and these assertions fail. + """ + + def test_thought_only_stream_heartbeats_while_pending( + self, pg_engine, monkeypatch + ): + from application.api.answer.routes import base as base_mod + + user_id = f"user-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + + # Count heartbeat invocations directly on the service so the test is + # explicit that the pump fired during the thought-only phase. + heartbeat_calls: List[str] = [] + real_heartbeat = ConversationService.heartbeat_message + + def _counting_heartbeat(self, message_id: str) -> bool: + heartbeat_calls.append(message_id) + return real_heartbeat(self, message_id) + + with _wire_db(pg_engine, monkeypatch): + monkeypatch.setattr( + ConversationService, "heartbeat_message", _counting_heartbeat + ) + # Make every loop iteration cross the heartbeat interval. + monkeypatch.setattr(base_mod.time, "monotonic", _FakeMonotonic()) + + base = _make_base() + agent = _ThoughtOnlyAgent(n_thoughts=4) + frames = _drain( + base.complete_stream( + question="hard reasoning question", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": user_id}, + should_persist=True, + model_id="gpt-5", + ) + ) + + # Sanity: the stream only ever emitted thoughts (plus framing/end), + # never an answer — so ``streaming`` was never marked during the loop. + joined = "\n".join(frames) + assert '"type": "thought"' in joined + assert '"type": "answer"' not in joined + + # Resolve the reserved row. + with pg_engine.connect() as conn: + msg_id = conn.execute( + text( + "SELECT cm.id FROM conversation_messages cm " + "JOIN conversations c ON c.id = cm.conversation_id " + "WHERE c.user_id = :u ORDER BY cm.timestamp DESC LIMIT 1" + ), + {"u": user_id}, + ).scalar() + assert msg_id is not None + row = _row(conn, str(msg_id)) + + # Core assertion: the heartbeat fired for this row during a thought-only + # phase. Pre-fix the pump is gated behind ``streaming_marked`` and never + # runs, so ``heartbeat_calls`` is empty. + assert str(msg_id) in heartbeat_calls, ( + "heartbeat_message was never called for the pending reasoning row; " + "the heartbeat pump is still gated behind the streaming flip" + ) + # And it is observable on the row: a non-null last_heartbeat_at. + assert row["metadata"] is not None + assert row["metadata"].get("last_heartbeat_at") is not None + + def test_heartbeat_advances_last_heartbeat_at_while_row_stays_pending( + self, pg_engine, monkeypatch + ): + """End-of-stream the answer arrives and the row finalizes, so to observe + the *pending* heartbeat we snapshot the row mid-stream: after the first + ``thought`` chunk the row must be ``pending`` with a fresh + ``last_heartbeat_at`` already stamped (seed-at-start), and a subsequent + ``thought`` must bump it further — all before any ``answer``. + """ + from application.api.answer.routes import base as base_mod + + user_id = f"user-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + + observed: List[Dict[str, Any]] = [] + + class _ObservingAgent: + """Yields thoughts; after each, the harness inspects the DB row.""" + + def __init__(self) -> None: + self.llm = _FakeLLM() + + def gen(self, query: str = ""): + # message_id is surfaced before this generator runs, so by the + # time we yield, the reserved row already exists. We can't read + # it from inside gen() (no msg id handle), so just emit; the + # snapshots happen in the consuming loop below via a wrapper. + for i in range(3): + yield {"thought": f"step {i}"} + + with _wire_db(pg_engine, monkeypatch): + monkeypatch.setattr(base_mod.time, "monotonic", _FakeMonotonic()) + base = _make_base() + agent = _ObservingAgent() + + gen = base.complete_stream( + question="reason about it", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": user_id}, + should_persist=True, + model_id="gpt-5", + ) + + # Pull the first frame (the ``message_id`` event) to learn the row. + first = next(gen) + assert "message_id" in first + msg_id = None + with pg_engine.connect() as conn: + msg_id = conn.execute( + text( + "SELECT cm.id FROM conversation_messages cm " + "JOIN conversations c ON c.id = cm.conversation_id " + "WHERE c.user_id = :u ORDER BY cm.timestamp DESC LIMIT 1" + ), + {"u": user_id}, + ).scalar() + assert msg_id is not None + + # Drain the rest, snapshotting the row's heartbeat after each frame. + for _frame in gen: + with pg_engine.connect() as conn: + observed.append(_row(conn, str(msg_id))) + + # While the stream was thought-only the row stayed ``pending`` and the + # heartbeat advanced (was non-null on each snapshot). + pending_snaps = [s for s in observed if s.get("status") == "pending"] + assert pending_snaps, "expected at least one pending snapshot mid-stream" + for snap in pending_snaps: + assert snap["metadata"] is not None + assert snap["metadata"].get("last_heartbeat_at") is not None + + +# --------------------------------------------------------------------------- +# 2. Seed-at-start: the reserved row has a heartbeat from generation start +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestHeartbeatSeededAtGenerationStart: + """Right after ``complete_stream`` reserves the row and begins consuming the + generator (before any ``answer``), the row must already carry a non-null + ``last_heartbeat_at`` — so even a model that takes a while to emit its first + chunk is covered from t=0, not only from the first interval tick. + """ + + def test_reserved_row_has_heartbeat_before_first_answer( + self, pg_engine, monkeypatch + ): + # This case keeps real ``time.monotonic`` so it asserts the *seed* + # (which runs before the loop, with no time advance), not an + # interval-driven pump — so no ``base.time`` monkeypatch is needed. + user_id = f"user-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + + captured: Dict[str, Any] = {} + + class _SlowFirstChunkAgent: + """Yields a single thought then stops — emulates a model that has + produced no answer yet by the time we inspect the row.""" + + def __init__(self) -> None: + self.llm = _FakeLLM() + + def gen(self, query: str = ""): + # Inspect the reserved row at the very first generator step, + # before any answer/sources chunk has been processed. + with pg_engine.connect() as conn: + mid = conn.execute( + text( + "SELECT cm.id FROM conversation_messages cm " + "JOIN conversations c ON c.id = cm.conversation_id " + "WHERE c.user_id = :u " + "ORDER BY cm.timestamp DESC LIMIT 1" + ), + {"u": user_id}, + ).scalar() + captured["row"] = _row(conn, str(mid)) if mid else {} + yield {"thought": "starting to reason"} + + with _wire_db(pg_engine, monkeypatch): + # Keep real monotonic time so this asserts the *seed*, not an + # interval-driven pump (no time advance happens before gen runs). + base = _make_base() + agent = _SlowFirstChunkAgent() + _drain( + base.complete_stream( + question="seed check", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": user_id}, + should_persist=True, + model_id="gpt-5", + ) + ) + + row = captured.get("row") or {} + # The seed runs before the loop, so the row is still ``pending`` and + # already has a heartbeat stamped. + assert row.get("status") == "pending" + assert row.get("metadata") is not None + assert row["metadata"].get("last_heartbeat_at") is not None + + +# --------------------------------------------------------------------------- +# 3. End-to-end reconciler guarantee: a fresh-heartbeat pending row is safe +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestReconcilerProtectsFreshlyHeartbeatedPendingRow: + """The whole point of the fix: a ``pending`` row older than 5 min by + ``timestamp`` but with a fresh ``last_heartbeat_at`` (the reasoning pump) + must NOT be swept by ``find_and_lock_stuck_messages``; a stale-heartbeat + ``pending`` row still IS swept. + + This mirrors the ``streaming`` cases in + ``tests/storage/db/repositories/test_reconciliation.py`` but for the + ``pending`` (still-reasoning) status the fix is about. + """ + + def _seed_pending_with_heartbeat( + self, conn, *, hb_minutes_ago: float, ts_minutes_ago: int = 20, + ) -> str: + conv = ConversationsRepository(conn).create("u-recon", "reasoning recon") + row = conn.execute( + text( + """ + INSERT INTO conversation_messages ( + conversation_id, position, prompt, response, status, + user_id, timestamp, message_metadata + ) + VALUES ( + CAST(:cid AS uuid), 0, 'p', '', 'pending', 'u-recon', + clock_timestamp() - make_interval(mins => :ts), + jsonb_build_object( + 'last_heartbeat_at', + to_jsonb( + clock_timestamp() + - make_interval(secs => :hb_secs) + ) + ) + ) + RETURNING id + """ + ), + { + "cid": conv["id"], + "ts": ts_minutes_ago, + "hb_secs": int(hb_minutes_ago * 60), + }, + ).fetchone() + return str(row[0]) + + def test_fresh_heartbeat_pending_row_is_not_swept(self, pg_conn): + # timestamp 20 min old (well past the 5-min age gate) but the reasoning + # pump heartbeat is only 30s old → the row is live, must be skipped. + msg_id = self._seed_pending_with_heartbeat( + pg_conn, hb_minutes_ago=0.5, ts_minutes_ago=20 + ) + rows = ReconciliationRepository(pg_conn).find_and_lock_stuck_messages() + assert all(str(r["id"]) != msg_id for r in rows), ( + "a still-reasoning pending row with a fresh heartbeat was " + "incorrectly selected for reconciliation" + ) + + def test_stale_heartbeat_pending_row_is_swept(self, pg_conn): + # Both timestamp and heartbeat are older than 5 min → genuinely stuck. + msg_id = self._seed_pending_with_heartbeat( + pg_conn, hb_minutes_ago=10, ts_minutes_ago=20 + ) + rows = ReconciliationRepository(pg_conn).find_and_lock_stuck_messages() + assert any(str(r["id"]) == msg_id for r in rows), ( + "a pending row stale by both timestamp and heartbeat should be " + "selected for reconciliation" + ) + + +# --------------------------------------------------------------------------- +# 4. Regression: a normal answer turn still marks streaming + heartbeats +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestNormalAnswerTurnUnchanged: + """A normal turn that streams an ``answer`` must still flip the row to + ``streaming`` on the first answer chunk (status semantics unchanged) and + finalize ``complete``. The heartbeat liveness change must not regress this. + """ + + def test_answer_turn_marks_streaming_and_finalizes_complete( + self, pg_engine, monkeypatch + ): + from application.api.answer.routes import base as base_mod + + user_id = f"user-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + + # Observe the status transition: snapshot after the first answer frame. + statuses_seen: List[str] = [] + real_update = ConversationService.update_message_status + + def _recording_update(self, message_id: str, status: str) -> bool: + statuses_seen.append(status) + return real_update(self, message_id, status) + + with _wire_db(pg_engine, monkeypatch): + monkeypatch.setattr( + ConversationService, "update_message_status", _recording_update + ) + monkeypatch.setattr(base_mod.time, "monotonic", _FakeMonotonic()) + base = _make_base() + agent = _ThoughtThenAnswerAgent(n_thoughts=2) + frames = _drain( + base.complete_stream( + question="normal question", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": user_id}, + should_persist=True, + model_id="gpt-5", + ) + ) + + joined = "\n".join(frames) + assert '"type": "answer"' in joined + assert '"type": "end"' in joined + + # The first answer chunk marked the row ``streaming`` exactly once. + assert statuses_seen == ["streaming"] + + with pg_engine.connect() as conn: + row = conn.execute( + text( + "SELECT status, response, message_metadata " + "FROM conversation_messages cm " + "JOIN conversations c ON c.id = cm.conversation_id " + "WHERE c.user_id = :u ORDER BY cm.timestamp DESC LIMIT 1" + ), + {"u": user_id}, + ).fetchone() + + # Final state: terminal ``complete`` with the answer, and a heartbeat + # was stamped along the way (liveness intact). + assert row[0] == "complete" + assert row[1] == _ThoughtThenAnswerAgent.ANSWER_TEXT + assert row[2] is not None + assert row[2].get("last_heartbeat_at") is not None diff --git a/tests/api/v1/test_v1_continuation_ownership.py b/tests/api/v1/test_v1_continuation_ownership.py new file mode 100644 index 00000000..757e5396 --- /dev/null +++ b/tests/api/v1/test_v1_continuation_ownership.py @@ -0,0 +1,235 @@ +"""Foreign-``conversation_id`` ownership gate on the ``/v1`` continuation path. + +A ``/v1`` tool continuation persists its final answer into a client-supplied +``conversation_id``. Ownership IS enforced deep in the stack: + + route → ``build_continuation_from_messages`` → ``build_agent("")`` + → ``initialize`` → ``_load_conversation_history`` + → ``get_conversation(conversation_id, owner)`` ← owner-scoped read + → ``None`` for a foreign conversation + → ``ValueError("Conversation not found or unauthorized")`` + → route's ``except ValueError`` → HTTP 400 + +…but the existing E2E suite mocks ``build_agent``, so this security gate is +never actually exercised. This test seeds a conversation owned by user A and an +agent owned by user B, then POSTs a continuation with B's api_key targeting A's +conversation. The ownership read runs against the **real** DB (not mocked): we +only short-circuit ``StreamProcessor.initialize`` to the +``_load_conversation_history`` step so the test doesn't have to stand up the +full agent/model/source/retriever machinery to reach the gate. + +Asserted: the route returns 400 AND no assistant turn is appended to A's +conversation (the foreign write is rejected before persistence). +""" + +from __future__ import annotations + +import json +import uuid +from typing import Any, Dict, List + +import pytest +from flask import Flask +from sqlalchemy import text + +from application.api.answer.services import stream_processor as sp_mod +from application.api.v1.routes import v1_bp +from application.storage.db.repositories.agents import AgentsRepository +from application.storage.db.repositories.conversations import ConversationsRepository + +# Reuse the route-level DB wiring + seed helpers from the tool-pause suite. +from tests.api.v1.test_v1_tool_pause_finalization import ( + _seed_user, + _wire_v1_route_db, +) + + +CLIENT_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} + + +def _seed_conversation(conn, user_id: str, name: str) -> str: + """Create a conversation owned by ``user_id`` with one finalized turn.""" + repo = ConversationsRepository(conn) + conv = repo.create(user_id, name) + conv_id = str(conv["id"]) + repo.reserve_message( + conv_id, + prompt="hello", + placeholder_response="hi there", + status="complete", + ) + return conv_id + + +def _seed_agent(conn, user_id: str, key: str) -> None: + AgentsRepository(conn).create(user_id, "Weather Agent", "published", key=key) + + +def _build_app() -> Flask: + app = Flask(__name__) + app.register_blueprint(v1_bp) + return app + + +def _row_count(conn, conv_id: str) -> int: + return conn.execute( + text( + "SELECT count(*) FROM conversation_messages " + "WHERE conversation_id = CAST(:c AS uuid)" + ), + {"c": conv_id}, + ).scalar() + + +def _only_run_history_load(self) -> None: + """Stand-in for ``StreamProcessor.initialize`` that runs ONLY the + ownership-checked history load. + + The real ownership gate (``_load_conversation_history`` → + ``get_conversation(conversation_id, owner)``) is NOT mocked — it runs + against the real DB. We skip the surrounding agent/model/source/retriever + configuration only so the test can reach the gate without provisioning a + full agent runtime; the ``ValueError`` raised by the gate for a foreign + conversation propagates exactly as in production. + """ + self._load_conversation_history() + + +@pytest.mark.integration +class TestForeignConversationOwnership: + """A continuation targeting another user's conversation is rejected at the + ownership gate (400) and never appends a turn to that conversation.""" + + PENDING_MESSAGES: List[Dict[str, Any]] = [ + {"role": "user", "content": "weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"city": "SF"}), + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call_abc", "content": "72F sunny"}, + ] + + def test_continuation_into_foreign_conversation_is_rejected( + self, pg_engine, monkeypatch + ): + user_a = f"userA-{uuid.uuid4().hex[:8]}" # owns the conversation + user_b = f"userB-{uuid.uuid4().hex[:8]}" # owns the agent (the caller) + api_key = f"key-{uuid.uuid4().hex[:8]}" + + with pg_engine.begin() as conn: + _seed_user(conn, user_a) + _seed_user(conn, user_b) + # Agent owned by B → the v1 route resolves decoded_token={"sub": B}. + _seed_agent(conn, user_b, api_key) + # Conversation owned by A, with one existing complete turn. + conv_id = _seed_conversation(conn, user_a, "A's private chat") + + with pg_engine.connect() as conn: + rows_before = _row_count(conn, conv_id) + assert rows_before == 1 # the seeded turn + + app = _build_app() + + # Drive the REAL ownership read; only short-circuit the surrounding + # initialize steps so we reach the gate without a full agent runtime. + with _wire_v1_route_db(pg_engine, monkeypatch): + monkeypatch.setattr( + sp_mod.StreamProcessor, "initialize", _only_run_history_load + ) + with app.test_client() as c: + resp = c.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + # B (the agent owner) tries to write into A's conversation. + "conversation_id": conv_id, + "messages": self.PENDING_MESSAGES, + "tools": [CLIENT_TOOL], + "docsgpt": {"save_conversation": True}, + }, + ) + + # The ownership gate raised ValueError → route returns 400. + assert resp.status_code == 400, resp.get_data(as_text=True) + body = resp.get_json() + assert body["error"]["type"] == "invalid_request" + + # No assistant turn was appended to A's conversation — the foreign + # write was rejected before any persistence. + with pg_engine.connect() as conn: + rows_after = _row_count(conn, conv_id) + assert rows_after == rows_before == 1 + + def test_owner_can_continue_into_own_conversation(self, pg_engine, monkeypatch): + """Control: the SAME continuation into the caller's OWN conversation + passes the ownership gate (no 400 from the gate). Proves the 400 above + is the ownership check firing, not an unrelated error in the path. + """ + owner = f"owner-{uuid.uuid4().hex[:8]}" + api_key = f"key-{uuid.uuid4().hex[:8]}" + + with pg_engine.begin() as conn: + _seed_user(conn, owner) + _seed_agent(conn, owner, api_key) + # Conversation owned by the agent owner (the caller). + conv_id = _seed_conversation(conn, owner, "owner's chat") + + app = _build_app() + + # Capture whether the gate passed: a flag the stubbed initialize sets + # only if ``_load_conversation_history`` returned without raising. + gate = {"passed": False} + + def _history_then_flag(self): + self._load_conversation_history() + gate["passed"] = True + # Stop here: returning lets build_continuation_from_messages proceed + # to create_agent, which needs a full runtime we deliberately skip. + # Raise a sentinel the test recognises, distinct from the ownership + # ValueError, so we can assert the gate itself did not reject us. + raise RuntimeError("__gate_passed_sentinel__") + + with _wire_v1_route_db(pg_engine, monkeypatch): + monkeypatch.setattr( + sp_mod.StreamProcessor, "initialize", _history_then_flag + ) + with app.test_client() as c: + resp = c.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "conversation_id": conv_id, + "messages": ( + TestForeignConversationOwnership.PENDING_MESSAGES + ), + "tools": [CLIENT_TOOL], + "docsgpt": {"save_conversation": True}, + }, + ) + + # The ownership read accepted the owner (gate passed); the request then + # fails on the deliberately-skipped runtime (500), NOT on a 400 from the + # gate. This isolates the ownership decision from the rest of the path. + assert gate["passed"] is True + assert resp.status_code == 500, resp.get_data(as_text=True) diff --git a/tests/api/v1/test_v1_idempotency.py b/tests/api/v1/test_v1_idempotency.py new file mode 100644 index 00000000..7e5b7d6a --- /dev/null +++ b/tests/api/v1/test_v1_idempotency.py @@ -0,0 +1,439 @@ +"""Layer-1 idempotency for ``/v1/chat/completions`` (non-streaming, OpenAI-style). + +Dropping the native ``resume_from_tool_actions`` path removed its +``mark_resuming`` guard, so a duplicated/retried non-streaming POST could +re-run the agent → a duplicate answer row + double token billing. The fix +honors a client-supplied ``Idempotency-Key`` header: a retry returns the +*stored first response* instead of re-running. + +These tests pin both layers against an ephemeral Postgres (``pg_engine``): + +- the ``application.api.v1.idempotency`` helper contract directly + (completed → cached, fresh-pending → 409, stale-pending → re-claim, + non-2xx → not cached); and +- the real ``/v1/chat/completions`` route end-to-end: the same key twice + returns the cached body and does NOT re-run the agent or append a second + conversation row; no key / different keys run independently; a 5xx first + response is not cached so a retry re-runs. +""" + +from __future__ import annotations + +import uuid +from contextlib import contextmanager +from typing import Any, Dict, List +from unittest.mock import patch + +import pytest +from flask import Flask +from sqlalchemy import text + +from application.api.v1 import idempotency as v1_idem +from application.api.v1.idempotency import ( + STALE_PENDING_SECONDS, + TASK_NAME, + claim_or_replay, + finalize, + release, + scoped_key, +) +from application.api.v1.routes import v1_bp + +# Reuse the route-level DB wiring + fake answering agent from the tool-pause +# suite so a real two-POST round-trip runs against the ephemeral Postgres. +from tests.api.v1.test_v1_tool_pause_finalization import ( + _seed_agent, + _seed_user, + _wire_v1_route_db, +) + + +# --------------------------------------------------------------------------- +# Helper-contract tests (drive the real task_dedup row directly) +# --------------------------------------------------------------------------- + + +def _make_response(app: Flask, body: Dict[str, Any], status: int): + """Build a real Flask Response inside an app context (for ``finalize``).""" + from flask import jsonify, make_response + + with app.test_request_context(): + return make_response(jsonify(body), status) + + +def _wire_idem_db(engine, monkeypatch) -> None: + """Point the v1 idempotency helper's ``db_session``/``db_readonly`` at + the ephemeral engine (each call opens its own short-lived connection).""" + + @contextmanager + def _session(): + conn = engine.connect() + txn = conn.begin() + try: + yield conn + txn.commit() + except Exception: + txn.rollback() + raise + finally: + conn.close() + + @contextmanager + def _readonly(): + conn = engine.connect() + try: + yield conn + finally: + conn.close() + + monkeypatch.setattr(v1_idem, "db_session", _session) + monkeypatch.setattr(v1_idem, "db_readonly", _readonly) + + +def _backdate_pending(engine, key: str, secs_ago: int) -> None: + with engine.begin() as conn: + conn.execute( + text( + "UPDATE task_dedup " + "SET created_at = clock_timestamp() - make_interval(secs => :s) " + "WHERE idempotency_key = :k" + ), + {"s": secs_ago, "k": key}, + ) + + +@pytest.mark.integration +class TestIdempotencyHelperContract: + """The three-outcome claim-before-process contract, against a real row.""" + + def test_scoped_key_namespaces_by_agent(self): + assert scoped_key("abc", "agent-1") == "agent-1:abc" + # Two agents with the same key value never collide. + assert scoped_key("abc", "agent-1") != scoped_key("abc", "agent-2") + # Opt-in: missing either component disables idempotency. + assert scoped_key(None, "agent-1") is None + assert scoped_key("abc", None) is None + + def test_first_claim_then_completed_replays_cached(self, pg_engine, monkeypatch): + # ``claim_or_replay`` builds Flask responses (``jsonify``) for the + # replay/409 paths, so it runs inside an app context — exactly as it + # does in the real route. + app = Flask(__name__) + _wire_idem_db(pg_engine, monkeypatch) + key = f"agent:{uuid.uuid4().hex}" + + with app.app_context(): + claimed, replay = claim_or_replay(key) + assert claimed is True + assert replay is None + + # Finalize a 200 with a body, then a retry replays it byte-for-byte + # without re-claiming. + body = {"id": "chatcmpl-1", "choices": [{"x": 1}]} + finalize(key, _make_response(app, body, 200)) + + claimed2, replay2 = claim_or_replay(key) + assert claimed2 is False + assert replay2 is not None + assert replay2.status_code == 200 + assert replay2.get_json() == body + + def test_fresh_pending_returns_409(self, pg_engine, monkeypatch): + app = Flask(__name__) + _wire_idem_db(pg_engine, monkeypatch) + key = f"agent:{uuid.uuid4().hex}" + + with app.app_context(): + claimed, _ = claim_or_replay(key) + assert claimed is True # first request in flight, never finalized + + # A concurrent same-key request sees a fresh pending row → 409. + claimed2, replay2 = claim_or_replay(key) + assert claimed2 is False + assert replay2 is not None + assert replay2.status_code == 409 + assert replay2.get_json()["error"]["type"] == "idempotency_conflict" + + def test_stale_pending_is_reclaimed(self, pg_engine, monkeypatch): + app = Flask(__name__) + _wire_idem_db(pg_engine, monkeypatch) + key = f"agent:{uuid.uuid4().hex}" + + with app.app_context(): + claimed, _ = claim_or_replay(key) + assert claimed is True + # The original request "died" before finalize; age the claim past + # the safety window so a retry may re-claim instead of 409-ing + # forever. + _backdate_pending(pg_engine, key, secs_ago=STALE_PENDING_SECONDS + 5) + + claimed2, replay2 = claim_or_replay(key) + assert claimed2 is True + assert replay2 is None + + def test_non_2xx_response_is_not_cached(self, pg_engine, monkeypatch): + app = Flask(__name__) + _wire_idem_db(pg_engine, monkeypatch) + key = f"agent:{uuid.uuid4().hex}" + + with app.app_context(): + claimed, _ = claim_or_replay(key) + assert claimed is True + # A 5xx must release the claim (not cache the error). + finalize(key, _make_response(app, {"error": {"message": "boom"}}, 500)) + + with pg_engine.connect() as conn: + row = conn.execute( + text("SELECT * FROM task_dedup WHERE idempotency_key = :k"), + {"k": key}, + ).fetchone() + assert row is None # released → a genuine retry can re-claim + + def test_release_drops_pending_claim(self, pg_engine, monkeypatch): + app = Flask(__name__) + _wire_idem_db(pg_engine, monkeypatch) + key = f"agent:{uuid.uuid4().hex}" + with app.app_context(): + claim_or_replay(key) + release(key) + with pg_engine.connect() as conn: + row = conn.execute( + text("SELECT * FROM task_dedup WHERE idempotency_key = :k"), + {"k": key}, + ).fetchone() + assert row is None + + +# --------------------------------------------------------------------------- +# Route-level tests — same key twice does NOT re-run the agent +# --------------------------------------------------------------------------- + + +class _CountingAnswerAgent: + """Agent that answers once per ``gen`` and records how often it ran. + + A class-level counter survives across the fresh instances handed back by + the patched ``build_agent``, so the route-level test can assert the agent + ran exactly once across two same-key POSTs. + """ + + ANSWER_TEXT = "The answer is 42." + gen_calls = 0 + + def __init__(self) -> None: + from tests.api.v1.test_v1_tool_pause_finalization import ( + _FakeClientToolExecutor, + _FakeLLM, + ) + + self.llm = _FakeLLM() + self.tool_executor = _FakeClientToolExecutor() + self.conversation_id = None + self.initial_user_id = None + self.tools: List[Dict[str, Any]] = [] + + def gen(self, query: str = ""): + # Write the base-class counter explicitly so subclasses (e.g. a + # failing variant) share one run-count rather than shadowing it. + _CountingAnswerAgent.gen_calls += 1 + yield {"answer": self.ANSWER_TEXT} + + +def _build_app() -> Flask: + app = Flask(__name__) + app.register_blueprint(v1_bp) + return app + + +def _post_chat(client, body, api_key, idem_key=None): + headers = {"Authorization": f"Bearer {api_key}"} + if idem_key is not None: + headers["Idempotency-Key"] = idem_key + return client.post("/v1/chat/completions", headers=headers, json=body) + + +def _conv_count(conn, user_id: str) -> int: + return conn.execute( + text("SELECT count(*) FROM conversations WHERE user_id = :u"), + {"u": user_id}, + ).scalar() + + +def _row_count(conn, user_id: str) -> int: + return conn.execute( + text( + "SELECT count(*) FROM conversation_messages cm " + "JOIN conversations c ON c.id = cm.conversation_id " + "WHERE c.user_id = :u" + ), + {"u": user_id}, + ).scalar() + + +@pytest.mark.integration +class TestV1IdempotencyRoute: + """Drive the real ``/v1/chat/completions`` route with/without a key.""" + + QUESTION = {"messages": [{"role": "user", "content": "what is the answer?"}], + "docsgpt": {"save_conversation": True}} + + def _fake_build_agent(self, question): # noqa: ARG002 + return _CountingAnswerAgent() + + def test_same_key_twice_replays_without_rerunning_agent( + self, pg_engine, monkeypatch + ): + user_id = f"user-{uuid.uuid4().hex[:8]}" + api_key = f"key-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + _seed_agent(conn, user_id, api_key) + + app = _build_app() + idem_key = f"idem-{uuid.uuid4().hex}" + _CountingAnswerAgent.gen_calls = 0 + + with _wire_v1_route_db(pg_engine, monkeypatch), patch( + "application.api.answer.services.stream_processor.StreamProcessor" + ".build_agent", + self._fake_build_agent, + ): + _wire_idem_db(pg_engine, monkeypatch) + with app.test_client() as c: + resp1 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key) + resp2 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key) + + assert resp1.status_code == 200, resp1.get_data(as_text=True) + assert resp2.status_code == 200, resp2.get_data(as_text=True) + body1, body2 = resp1.get_json(), resp2.get_json() + # The retry returns the cached first response byte-for-byte. + assert body2 == body1 + assert body2["choices"][0]["message"]["content"] == ( + _CountingAnswerAgent.ANSWER_TEXT + ) + + # Real dedup: the agent ran exactly once across the two POSTs... + assert _CountingAnswerAgent.gen_calls == 1 + # ...and no duplicate conversation / answer row was appended. + with pg_engine.connect() as conn: + assert _conv_count(conn, user_id) == 1 + # One conversation, one terminal answer row (not two). + assert _row_count(conn, user_id) == 1 + + def test_no_key_runs_each_time(self, pg_engine, monkeypatch): + """Opt-in: without a key, behavior is today's — each POST runs.""" + user_id = f"user-{uuid.uuid4().hex[:8]}" + api_key = f"key-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + _seed_agent(conn, user_id, api_key) + + app = _build_app() + _CountingAnswerAgent.gen_calls = 0 + + with _wire_v1_route_db(pg_engine, monkeypatch), patch( + "application.api.answer.services.stream_processor.StreamProcessor" + ".build_agent", + self._fake_build_agent, + ): + _wire_idem_db(pg_engine, monkeypatch) + with app.test_client() as c: + resp1 = _post_chat(c, self.QUESTION, api_key) + resp2 = _post_chat(c, self.QUESTION, api_key) + + assert resp1.status_code == 200 + assert resp2.status_code == 200 + # Two independent runs → two conversations. + assert _CountingAnswerAgent.gen_calls == 2 + with pg_engine.connect() as conn: + assert _conv_count(conn, user_id) == 2 + + def test_different_keys_run_independently(self, pg_engine, monkeypatch): + user_id = f"user-{uuid.uuid4().hex[:8]}" + api_key = f"key-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + _seed_agent(conn, user_id, api_key) + + app = _build_app() + _CountingAnswerAgent.gen_calls = 0 + + with _wire_v1_route_db(pg_engine, monkeypatch), patch( + "application.api.answer.services.stream_processor.StreamProcessor" + ".build_agent", + self._fake_build_agent, + ): + _wire_idem_db(pg_engine, monkeypatch) + with app.test_client() as c: + _post_chat(c, self.QUESTION, api_key, idem_key="k-A") + _post_chat(c, self.QUESTION, api_key, idem_key="k-B") + + assert _CountingAnswerAgent.gen_calls == 2 + with pg_engine.connect() as conn: + assert _conv_count(conn, user_id) == 2 + + def test_5xx_first_response_is_not_cached_and_retry_reruns( + self, pg_engine, monkeypatch + ): + """A failed first response must NOT be cached: a retry re-runs.""" + user_id = f"user-{uuid.uuid4().hex[:8]}" + api_key = f"key-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + _seed_agent(conn, user_id, api_key) + + app = _build_app() + idem_key = f"idem-{uuid.uuid4().hex}" + + class _BoomAgent(_CountingAnswerAgent): + def gen(self, query: str = ""): + _CountingAnswerAgent.gen_calls += 1 + raise RuntimeError("upstream exploded") + yield # pragma: no cover - keeps gen a generator + + boom_state = {"fail": True} + + def _build(self, question): # noqa: ARG001 + return _BoomAgent() if boom_state["fail"] else _CountingAnswerAgent() + + _CountingAnswerAgent.gen_calls = 0 + + with _wire_v1_route_db(pg_engine, monkeypatch), patch( + "application.api.answer.services.stream_processor.StreamProcessor" + ".build_agent", + _build, + ): + _wire_idem_db(pg_engine, monkeypatch) + with app.test_client() as c: + resp1 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key) + # First attempt blew up (500) and released the claim. The key + # is agent-id scoped at the route, so confirm release via a + # table-wide count for this task_name (no live claim left). + assert resp1.status_code == 500, resp1.get_data(as_text=True) + with pg_engine.connect() as conn: + rows = conn.execute( + text( + "SELECT count(*) FROM task_dedup " + "WHERE task_name = :n" + ), + {"n": TASK_NAME}, + ).scalar() + assert rows == 0 # error was not cached; claim released + + # Retry: the agent now answers; idempotency did NOT cache the error. + boom_state["fail"] = False + resp2 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key) + + assert resp2.status_code == 200, resp2.get_data(as_text=True) + # Agent ran twice total: the failing first + the successful retry. + assert _CountingAnswerAgent.gen_calls == 2 + # The successful retry IS now cached. + with pg_engine.connect() as conn: + completed = conn.execute( + text( + "SELECT count(*) FROM task_dedup " + "WHERE task_name = :n AND status = 'completed'" + ), + {"n": TASK_NAME}, + ).scalar() + assert completed == 1 diff --git a/tests/api/v1/test_v1_tool_pause_finalization.py b/tests/api/v1/test_v1_tool_pause_finalization.py new file mode 100644 index 00000000..f1cd4502 --- /dev/null +++ b/tests/api/v1/test_v1_tool_pause_finalization.py @@ -0,0 +1,1020 @@ +"""Tool-pause finalization + stateless resume for the OpenAI-compatible ``/v1`` +path. + +The OpenAI protocol has no slot for DocsGPT's ``reserved_message_id``: clients +resume a tool call by re-POSTing the whole message history with +``{role:"tool"}`` results (optionally threading the ``conversation_id`` they got +back), not via a native resume. The ``/v1`` tool round-trip is therefore fully +stateless on both ends: + +- **Pause side.** When an agent pauses for a client-executed tool, + ``complete_stream`` (run with ``finalize_tool_pause_as_complete=True``) + finalizes the reserved ``conversation_messages`` row as ``complete`` + (recording the tool_calls) instead of writing a ``pending_tool_state`` record + and leaving the row non-terminal. The reconciler never sees an orphaned row. +- **Resume side.** The ``/v1`` route rebuilds the agent + pending calls from the + re-POSTed history via ``StreamProcessor.build_continuation_from_messages`` + (which has *no* ``pending_tool_state`` dependency) — **regardless of whether a + ``conversation_id`` is present**. It never calls ``resume_from_tool_actions`` + (whose ``load_state`` would 400, since the pause wrote no state). When a + ``conversation_id`` is carried, the final answer persists as a NEW terminal + turn appended to that conversation. + +Net result per tool turn: a ``complete`` tool-call turn + a ``complete`` answer +turn — both terminal, no orphan, no 400, OpenAI-faithful. + +These tests drive the real ``/v1/chat/completions`` route and real +``complete_stream`` against an ephemeral Postgres database (``pg_engine``): + +- a real two-POST round-trip (pause then answer, threading the + ``conversation_id``) returns 200 on *both* POSTs and persists the answer into + the SAME conversation, with nothing left ``pending``/``streaming`` — the + regression catcher (the pre-fix route 400s on POST #2 via + ``resume_from_tool_actions``); +- a v1 tool round WITH a conversation context finalizes the reserved row as + ``complete`` and leaves no ``pending``/``streaming`` row behind; +- a stateless v1 tool round (no conversation_id, ``should_persist=False``) + leaves nothing non-terminal and writes no orphan conversation; +- the native ``/stream`` pause (flag defaulted False) is byte-for-byte + unchanged: it still writes ``pending_tool_state`` and leaves the row + non-terminal awaiting a native resume. +""" + +from __future__ import annotations + +import json +import uuid +from contextlib import contextmanager +from typing import Any, Dict, List, Optional +from unittest.mock import MagicMock, patch + +import pytest +from flask import Flask + +from application.api.answer.routes.base import BaseAnswerResource +from application.api.answer.services.conversation_service import ConversationService +from application.api.v1.routes import v1_bp + + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class _FakeLLM: + """Minimal LLM stand-in so ``complete_stream`` can stamp ``_request_id``.""" + + def __init__(self) -> None: + self._request_id: Optional[str] = None + self.model_id = "gpt-4" + + +class _FakeToolExecutor: + """Tool executor stub the native pause path reads ``client_tools`` off of.""" + + def __init__(self) -> None: + self.client_tools: Optional[List[Dict[str, Any]]] = None + self.message_id: Optional[str] = None + self.conversation_id: Optional[str] = None + + +class _PausingAgent: + """Agent whose ``gen``/``gen_continuation`` pauses for a client tool. + + Mirrors what the real handler does on a client-side / approval pause: + yield a ``tool_calls_pending`` event and stash ``_pending_continuation`` + on the agent. ``complete_stream`` keys its pause handling off both. + """ + + def __init__( + self, + pending_tool_calls: List[Dict[str, Any]], + *, + with_tool_executor: bool = False, + ) -> None: + self.llm = _FakeLLM() + self._pending_tool_calls = pending_tool_calls + self._pending_continuation: Optional[Dict[str, Any]] = None + # The native pause path reads ``agent.tool_executor.client_tools`` and + # ``agent.tools`` to persist continuation state; the v1 finalize path + # touches neither. Only attach a tool_executor when the test needs the + # native branch (``getattr(agent, "tool_executor", None)`` stays None + # otherwise). + if with_tool_executor: + self.tool_executor = _FakeToolExecutor() + + def _emit_pause(self): + self._pending_continuation = { + "messages": [{"role": "system", "content": "sys"}], + "pending_tool_calls": self._pending_tool_calls, + "tools_dict": {"0": {"name": "get_weather", "client_side": True}}, + "reasoning_content": "", + } + yield { + "type": "tool_calls_pending", + "data": {"pending_tool_calls": self._pending_tool_calls}, + } + + def gen(self, query: str = ""): + yield from self._emit_pause() + + def gen_continuation(self, **kwargs): + yield from self._emit_pause() + + +class _NoopJournalWriter: + """No-op journal writer so the test exercises DB row state, not the + message-events journal / Redis broadcast (orthogonal to finalization).""" + + def __init__(self, *args, **kwargs) -> None: + pass + + def record(self, *args, **kwargs) -> None: + pass + + def flush(self) -> None: + pass + + def close(self) -> None: + pass + + +# --------------------------------------------------------------------------- +# Test harness +# --------------------------------------------------------------------------- + + +PENDING_TOOL_CALLS = [ + { + "call_id": "call_abc", + "name": "get_weather_0", + "tool_name": "get_weather", + "action_name": "get_weather", + "arguments": {"city": "SF"}, + "pause_type": "requires_client_execution", + } +] + + +def _seed_user(conn, user_id: str) -> None: + from sqlalchemy import text + + conn.execute( + text( + "INSERT INTO users (user_id) VALUES (:u) " + "ON CONFLICT (user_id) DO NOTHING" + ), + {"u": user_id}, + ) + + +def _row_statuses(conn, conversation_id: str) -> List[str]: + from sqlalchemy import text + + rows = conn.execute( + text( + "SELECT status FROM conversation_messages " + "WHERE conversation_id = CAST(:c AS uuid) " + "ORDER BY position ASC" + ), + {"c": conversation_id}, + ).fetchall() + return [r[0] for r in rows] + + +def _row_tool_calls(conn, message_id: str): + from sqlalchemy import text + + row = conn.execute( + text( + "SELECT tool_calls FROM conversation_messages " + "WHERE id = CAST(:m AS uuid)" + ), + {"m": message_id}, + ).fetchone() + return row[0] if row is not None else None + + +def _pending_tool_state_count(conn, conversation_id: str) -> int: + from sqlalchemy import text + + return conn.execute( + text( + "SELECT count(*) FROM pending_tool_state " + "WHERE conversation_id = CAST(:c AS uuid)" + ), + {"c": conversation_id}, + ).scalar() + + +@contextmanager +def _wire_db(engine, monkeypatch): + """Point conversation_service / continuation_service / base at ``engine``. + + Each helper opens its own short-lived connection (matching production), + so we hand out fresh connections from the same ephemeral engine and + swap the journal writer for a no-op. + """ + from application.api.answer.services import conversation_service as conv_mod + from application.api.answer.services import continuation_service as cont_mod + from application.api.answer.routes import base as base_mod + + @contextmanager + def _session(): + conn = engine.connect() + txn = conn.begin() + try: + yield conn + txn.commit() + except Exception: + txn.rollback() + raise + finally: + conn.close() + + @contextmanager + def _readonly(): + conn = engine.connect() + try: + yield conn + finally: + conn.close() + + monkeypatch.setattr(conv_mod, "db_session", _session) + monkeypatch.setattr(conv_mod, "db_readonly", _readonly) + monkeypatch.setattr(cont_mod, "db_session", _session) + monkeypatch.setattr(cont_mod, "db_readonly", _readonly) + monkeypatch.setattr(base_mod, "db_session", _session) + monkeypatch.setattr(base_mod, "db_readonly", _readonly) + monkeypatch.setattr(base_mod, "BatchedJournalWriter", _NoopJournalWriter) + monkeypatch.setattr(base_mod, "record_event", lambda *a, **kw: None) + # The native approval pause publishes ``tool.approval.required`` out of + # band (Redis); no-op it so the test stays DB-only. + monkeypatch.setattr(base_mod, "publish_user_event", lambda *a, **kw: None) + yield + + +def _drain(gen) -> List[str]: + """Consume an SSE generator into a list of frames.""" + return list(gen) + + +# --------------------------------------------------------------------------- +# Guarantee 1 — v1 tool pause never leaves a non-terminal row +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestV1ToolPauseFinalizesRow: + """A first-turn ``/v1`` tool emission reserves a WAL row; the pause must + finalize it as ``complete`` (with the tool_calls) rather than strand it. + """ + + def test_v1_tool_pause_with_conversation_finalizes_row_complete( + self, pg_engine, monkeypatch + ): + user_id = f"user-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + + with _wire_db(pg_engine, monkeypatch): + base = BaseAnswerResource.__new__(BaseAnswerResource) + base.default_model_id = "gpt-4" + base.conversation_service = ConversationService() + + agent = _PausingAgent(PENDING_TOOL_CALLS) + frames = _drain( + base.complete_stream( + question="weather in SF?", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": user_id}, + should_persist=True, + model_id="gpt-4", + finalize_tool_pause_as_complete=True, + ) + ) + + # The client still receives the pause signal + a conversation id. + joined = "\n".join(frames) + assert "tool_calls_pending" in joined + assert '"type": "id"' in joined + assert '"type": "end"' in joined + + # Resolve the conversation that was created for the reserved row. + with pg_engine.connect() as conn: + from sqlalchemy import text + + conv_id = conn.execute( + text( + "SELECT id FROM conversations WHERE user_id = :u " + "ORDER BY created_at DESC LIMIT 1" + ), + {"u": user_id}, + ).scalar() + assert conv_id is not None + statuses = _row_statuses(conn, str(conv_id)) + msg_id = conn.execute( + text( + "SELECT id FROM conversation_messages " + "WHERE conversation_id = CAST(:c AS uuid) LIMIT 1" + ), + {"c": str(conv_id)}, + ).scalar() + tool_calls = _row_tool_calls(conn, str(msg_id)) + pts = _pending_tool_state_count(conn, str(conv_id)) + + # Guarantee 1: exactly one row, terminal ``complete`` — never + # ``pending``/``streaming``. + assert statuses == ["complete"] + assert "pending" not in statuses + assert "streaming" not in statuses + # The emitted/pending tool_calls are recorded on the row. + assert tool_calls + assert tool_calls[0]["call_id"] == "call_abc" + # No native continuation record is written on the v1 path. + assert pts == 0 + + def test_v1_stateless_tool_pause_leaves_no_orphan( + self, pg_engine, monkeypatch + ): + """Pure OpenAI/opencode style: a stateless continuation carries no + conversation_id and ``should_persist=False`` (translator opt-out), so + no WAL row is reserved. The pause must end cleanly with no + non-terminal row and no empty-prompt orphan conversation. + """ + user_id = f"user-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + conv_count_before = conn.execute( + __import__("sqlalchemy").text( + "SELECT count(*) FROM conversations WHERE user_id = :u" + ), + {"u": user_id}, + ).scalar() + + with _wire_db(pg_engine, monkeypatch): + base = BaseAnswerResource.__new__(BaseAnswerResource) + base.default_model_id = "gpt-4" + base.conversation_service = ConversationService() + + agent = _PausingAgent(PENDING_TOOL_CALLS) + frames = _drain( + base.complete_stream( + question="", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": user_id}, + should_persist=False, + model_id="gpt-4", + _continuation={ + "messages": [{"role": "system", "content": "sys"}], + "tools_dict": {"0": {"name": "get_weather"}}, + "pending_tool_calls": PENDING_TOOL_CALLS, + "tool_actions": [{"call_id": "call_abc", "result": "72F"}], + "reserved_message_id": None, + "request_id": None, + "reasoning_content": "", + }, + finalize_tool_pause_as_complete=True, + ) + ) + + joined = "\n".join(frames) + assert "tool_calls_pending" in joined + assert '"type": "end"' in joined + + with pg_engine.connect() as conn: + from sqlalchemy import text + + conv_count_after = conn.execute( + text("SELECT count(*) FROM conversations WHERE user_id = :u"), + {"u": user_id}, + ).scalar() + non_terminal = conn.execute( + text( + "SELECT count(*) FROM conversation_messages cm " + "JOIN conversations c ON c.id = cm.conversation_id " + "WHERE c.user_id = :u " + "AND cm.status IN ('pending', 'streaming')" + ), + {"u": user_id}, + ).scalar() + + # No orphan conversation created, nothing left non-terminal. + assert conv_count_after == conv_count_before + assert non_terminal == 0 + + def test_v1_multi_round_continuation_pause_leaves_nothing_non_terminal( + self, pg_engine, monkeypatch + ): + """Multi-round loop, coherent Option B: a v1 continuation rebuilds + STATELESSLY (the route passes ``reserved_message_id=None`` because the + first turn's row is already ``complete`` and cannot be reused). When + the resume pauses AGAIN for another client tool, no new WAL row is + reserved (``_continuation`` is truthy → ``wal_eligible`` is False), so + there is nothing to strand: the round ends cleanly, the already- + ``complete`` first-turn row is untouched, and nothing is left + ``pending``/``streaming`` for the reconciler. + """ + user_id = f"user-{uuid.uuid4().hex[:8]}" + # Seed a conversation whose first (tool-emitting) turn has already + # been finalized ``complete`` — the coherent-Option-B state after the + # first POST: no lingering ``pending`` row mid-loop. + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + with _wire_db(pg_engine, monkeypatch): + svc = ConversationService() + reservation = svc.save_user_question( + conversation_id=None, + question="multi-step task", + decoded_token={"sub": user_id}, + model_id="gpt-4", + ) + conv_id = reservation["conversation_id"] + first_turn_id = reservation["message_id"] + svc.finalize_message( + first_turn_id, + "", + tool_calls=[{"call_id": "call_abc", "name": "step1"}], + model_id="gpt-4", + status="complete", + ) + + with pg_engine.connect() as conn: + assert _row_statuses(conn, conv_id) == ["complete"] + + second_round_calls = [ + { + "call_id": "call_def", + "name": "lookup_0", + "tool_name": "lookup", + "action_name": "lookup", + "arguments": {"q": "next"}, + "pause_type": "requires_client_execution", + } + ] + + with _wire_db(pg_engine, monkeypatch): + base = BaseAnswerResource.__new__(BaseAnswerResource) + base.default_model_id = "gpt-4" + base.conversation_service = ConversationService() + + agent = _PausingAgent(second_round_calls) + frames = _drain( + base.complete_stream( + question="", + agent=agent, + conversation_id=conv_id, + user_api_key=None, + decoded_token={"sub": user_id}, + should_persist=True, + model_id="gpt-4", + _continuation={ + "messages": [{"role": "system", "content": "sys"}], + "tools_dict": {"0": {"name": "lookup"}}, + "pending_tool_calls": second_round_calls, + "tool_actions": [ + {"call_id": "call_abc", "result": "step1 done"} + ], + # Coherent Option B: the v1 route rebuilds statelessly, + # so it threads no reserved_message_id/request_id. + "reserved_message_id": None, + "request_id": None, + "reasoning_content": "", + }, + finalize_tool_pause_as_complete=True, + ) + ) + + joined = "\n".join(frames) + assert "tool_calls_pending" in joined + assert '"type": "end"' in joined + + with pg_engine.connect() as conn: + statuses = _row_statuses(conn, conv_id) + first_turn_calls = _row_tool_calls(conn, first_turn_id) + pts = _pending_tool_state_count(conn, conv_id) + + # The first-turn row stays exactly as it was (terminal ``complete`` + # with its own tool_calls); no sibling row was reserved or stranded, + # nothing is ``pending``/``streaming``, and no native state was written. + assert statuses == ["complete"] + assert "pending" not in statuses + assert "streaming" not in statuses + assert first_turn_calls + assert first_turn_calls[0]["call_id"] == "call_abc" + assert pts == 0 + + +# --------------------------------------------------------------------------- +# Guarantee 3 — native pause path is unchanged (flag defaults False) +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestNativePauseUnchanged: + """The native ``/stream`` + ``/api/answer`` flow must keep writing + ``pending_tool_state`` and leaving the reserved row non-terminal so a + native resume can finalize it. This is the default (flag omitted). + """ + + def test_native_pause_writes_state_and_leaves_row_non_terminal( + self, pg_engine, monkeypatch + ): + user_id = f"user-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + + approval_calls = [ + { + "call_id": "call_xyz", + "name": "send_msg_0", + "tool_name": "telegram", + "action_name": "send_msg", + "arguments": {"text": "hi"}, + "pause_type": "awaiting_approval", + } + ] + + with _wire_db(pg_engine, monkeypatch): + base = BaseAnswerResource.__new__(BaseAnswerResource) + base.default_model_id = "gpt-4" + base.conversation_service = ConversationService() + + agent = _PausingAgent(approval_calls, with_tool_executor=True) + agent.tools = [] + # Native default: finalize_tool_pause_as_complete omitted (False). + frames = _drain( + base.complete_stream( + question="message my team", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": user_id}, + should_persist=True, + model_id="gpt-4", + ) + ) + + joined = "\n".join(frames) + assert "tool_calls_pending" in joined + assert '"type": "end"' in joined + + with pg_engine.connect() as conn: + from sqlalchemy import text + + conv_id = conn.execute( + text( + "SELECT id FROM conversations WHERE user_id = :u " + "ORDER BY created_at DESC LIMIT 1" + ), + {"u": user_id}, + ).scalar() + assert conv_id is not None + statuses = _row_statuses(conn, str(conv_id)) + pts = _pending_tool_state_count(conn, str(conv_id)) + + # Native UX preserved: row stays non-terminal (awaiting native + # resume) and a continuation record was written. + assert statuses, "expected a reserved row" + assert statuses[0] in ("pending", "streaming") + assert "complete" not in statuses + assert pts == 1 + + +# --------------------------------------------------------------------------- +# Guarantee 2 — coherent Option B routing: a v1 continuation carrying a +# conversation_id rebuilds STATELESSLY via build_continuation_from_messages +# (never resume_from_tool_actions), and threads no reserved_message_id. +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestV1ContinuationRoutesStatelessly: + """Coherent Option B: a v1 tool-result continuation is fully stateless on + the resume side. Because the pause finalized the prior row ``complete`` and + wrote NO ``pending_tool_state``, the route must rebuild from the re-POSTed + history via ``build_continuation_from_messages`` — **even when a + ``conversation_id`` is carried** — and must NOT call + ``resume_from_tool_actions`` (whose ``load_state`` would 400). The + ``_continuation`` dict therefore carries no live ``reserved_message_id`` / + ``request_id`` (id-reuse is inert under this design). + """ + + def _build_app(self) -> Flask: + app = Flask(__name__) + app.register_blueprint(v1_bp) + return app + + def test_continuation_with_conversation_id_uses_stateless_rebuild(self): + app = self._build_app() + + def _fake_translate(data, api_key): + return { + "question": "", + "tool_actions": [{"call_id": "c1", "result": "r"}], + "conversation_id": "conv-1", + "persist": True, + "messages": data["messages"], + } + + fake_processor = MagicMock() + fake_processor.decoded_token = {"sub": "owner"} + fake_processor.conversation_id = "conv-1" + fake_processor.agent_config = {"user_api_key": "k"} + fake_processor.agent_id = None + fake_processor.model_id = "m" + fake_processor.model_user_id = None + # A real processor leaves these at None until a *native* resume hoists + # them; the route must not invent values from them. + fake_processor.reserved_message_id = None + fake_processor.request_id = None + + build_calls: Dict[str, Any] = {} + + def _fake_build_continuation(messages, tool_actions): + build_calls["messages"] = messages + build_calls["tool_actions"] = tool_actions + return (MagicMock(), [], {}, [{"call_id": "c1"}], tool_actions, "") + + fake_processor.build_continuation_from_messages.side_effect = ( + _fake_build_continuation + ) + + captured: Dict[str, Any] = {} + + def _capture_complete_stream(**kw): + captured.update(kw) + return iter(['data: {"type": "end"}']) + + fake_helper = MagicMock() + fake_helper.check_usage.return_value = None + fake_helper.complete_stream.side_effect = _capture_complete_stream + fake_helper.process_response_stream.return_value = { + "error": None, + "conversation_id": "conv-1", + "answer": "done", + "sources": [], + "tool_calls": [], + "thought": "", + } + + @contextmanager + def _yield_conn(): + yield MagicMock() + + with patch( + "application.api.v1.routes.translate_request", + side_effect=_fake_translate, + ), patch( + "application.api.v1.routes.StreamProcessor", + return_value=fake_processor, + ), patch( + "application.api.v1.routes._V1AnswerHelper", + return_value=fake_helper, + ), patch( + "application.api.v1.routes.db_readonly", + _yield_conn, + ), patch( + "application.api.v1.routes.translate_response", + return_value={"id": "x", "choices": []}, + ): + with app.test_client() as c: + resp = c.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer x"}, + json={ + "conversation_id": "conv-1", + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [{"id": "c1", "function": {}}], + }, + {"role": "tool", "tool_call_id": "c1", "content": "r"}, + ], + }, + ) + + assert resp.status_code == 200 + # Coherent Option B: the stateless rebuild ran... + fake_processor.build_continuation_from_messages.assert_called_once() + assert build_calls["tool_actions"] == [{"call_id": "c1", "result": "r"}] + # ...and the stateful native resume was NOT used. + fake_processor.resume_from_tool_actions.assert_not_called() + + continuation = captured.get("_continuation") + assert continuation is not None + # id-reuse is gone from the v1 continuation dict (inert under Option B). + assert "reserved_message_id" not in continuation + assert "request_id" not in continuation + # The v1 path still runs in stateless-finalize mode. + assert captured.get("finalize_tool_pause_as_complete") is True + # The incoming conversation_id is the persistence target. + assert captured.get("conversation_id") == "conv-1" + + +# --------------------------------------------------------------------------- +# Regression catcher — real two-POST /v1/chat/completions round-trip +# --------------------------------------------------------------------------- + + +class _FakeLLMGen: + """LLM stand-in for the answer path's title-gen ``create_llm``. + + The append branch of ``save_conversation`` (conversation_id present) never + calls ``gen``; this just satisfies the ``LLMCreator.create_llm`` call that + ``complete_stream`` makes before persisting. + """ + + model_id = "gpt-4" + _request_id: Optional[str] = None + _token_usage_source: Optional[str] = None + + def gen(self, *args, **kwargs) -> str: + return "Title" + + +class _FakeClientToolExecutor: + """Tool executor stub exposing ``get_tools`` for the real + ``build_continuation_from_messages`` (which reads the agent's tools).""" + + def __init__(self) -> None: + self.client_tools: Optional[List[Dict[str, Any]]] = None + self.message_id: Optional[str] = None + self.conversation_id: Optional[str] = None + + def get_tools(self) -> Dict[str, Any]: + return {"0": {"name": "get_weather", "client_side": True}} + + +class _PauseThenAnswerAgent: + """Agent that pauses for a client tool on ``gen`` (POST #1) and emits a + final answer on ``gen_continuation`` (POST #2). + + A fresh instance is handed back by the patched ``build_agent`` for each + request, so the first POST (normal mode) drives ``gen`` and the second + (continuation mode) drives ``gen_continuation``. + """ + + ANSWER_TEXT = "It is 72F and sunny in SF." + + def __init__(self, pending_tool_calls: List[Dict[str, Any]]) -> None: + self.llm = _FakeLLM() + self.tool_executor = _FakeClientToolExecutor() + self._pending_tool_calls = pending_tool_calls + self._pending_continuation: Optional[Dict[str, Any]] = None + self.conversation_id: Optional[str] = None + self.initial_user_id: Optional[str] = None + self.tools: List[Dict[str, Any]] = [] + + def gen(self, query: str = ""): + # First turn: pause for the client tool. + self._pending_continuation = { + "messages": [{"role": "system", "content": "sys"}], + "pending_tool_calls": self._pending_tool_calls, + "tools_dict": {"0": {"name": "get_weather", "client_side": True}}, + "reasoning_content": "", + } + yield { + "type": "tool_calls_pending", + "data": {"pending_tool_calls": self._pending_tool_calls}, + } + + def gen_continuation(self, **kwargs): + # Resume turn: emit the final answer and finish normally. + yield {"answer": self.ANSWER_TEXT} + + +def _seed_agent(conn, user_id: str, key: str) -> None: + from application.storage.db.repositories.agents import AgentsRepository + + AgentsRepository(conn).create(user_id, "Weather Agent", "published", key=key) + + +@contextmanager +def _wire_v1_route_db(engine, monkeypatch): + """Full route-level DB wiring for the ``/v1/chat/completions`` blueprint. + + Extends ``_wire_db`` (conversation/continuation/base services) with the v1 + routes module's own ``db_readonly`` (used by ``_lookup_agent``) and a fake + title-gen ``LLMCreator`` on the base module, so a real two-POST round-trip + runs entirely against the ephemeral Postgres with no live LLM/provider. + """ + from application.api.v1 import routes as v1_routes_mod + from application.api.answer.routes import base as base_mod + + @contextmanager + def _readonly(): + conn = engine.connect() + try: + yield conn + finally: + conn.close() + + with _wire_db(engine, monkeypatch): + monkeypatch.setattr(v1_routes_mod, "db_readonly", _readonly) + monkeypatch.setattr( + base_mod.LLMCreator, + "create_llm", + staticmethod(lambda *a, **kw: _FakeLLMGen()), + ) + yield + + +def _post_chat(client, body: Dict[str, Any], api_key: str): + return client.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json=body, + ) + + +@pytest.mark.integration +class TestV1ToolRoundTripEndToEnd: + """Drive the real ``/v1/chat/completions`` route twice through the real + ``routes.py`` routing and real ``StreamProcessor.build_continuation_from_messages``. + + Only the agent *creation* is mocked (``StreamProcessor.build_agent`` returns + a fake that pauses then answers) — the route logic, the continuation + rebuild, and ``resume_from_tool_actions`` are NOT mocked. This is the + regression catcher: the pre-fix route sends a conversation_id-carrying + continuation to ``resume_from_tool_actions`` → ``load_state`` returns None → + ``ValueError`` → HTTP 400 on POST #2. + """ + + CLIENT_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + + PENDING = [ + { + "call_id": "call_abc", + "name": "get_weather", + "tool_name": "get_weather", + "action_name": "get_weather", + "arguments": {"city": "SF"}, + "pause_type": "requires_client_execution", + } + ] + + def _build_app(self) -> Flask: + app = Flask(__name__) + app.register_blueprint(v1_bp) + return app + + def _count_non_terminal(self, conn, user_id: str) -> int: + from sqlalchemy import text + + return conn.execute( + text( + "SELECT count(*) FROM conversation_messages cm " + "JOIN conversations c ON c.id = cm.conversation_id " + "WHERE c.user_id = :u " + "AND cm.status IN ('pending', 'streaming')" + ), + {"u": user_id}, + ).scalar() + + def test_pause_then_answer_round_trip_persists_into_same_conversation( + self, pg_engine, monkeypatch + ): + user_id = f"user-{uuid.uuid4().hex[:8]}" + api_key = f"key-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + _seed_agent(conn, user_id, api_key) + + app = self._build_app() + + # ``build_agent`` is the only mock — a fresh pausing/answering agent + # per call. The route's ``build_continuation_from_messages`` calls this + # internally on POST #2, so the rebuild itself still runs for real. + def _fake_build_agent(self, question): # noqa: ARG001 + return _PauseThenAnswerAgent( + TestV1ToolRoundTripEndToEnd.PENDING + ) + + # ---- POST #1: user question -> agent pauses for the client tool ---- + with _wire_v1_route_db(pg_engine, monkeypatch), patch( + "application.api.answer.services.stream_processor.StreamProcessor" + ".build_agent", + _fake_build_agent, + ): + with app.test_client() as c: + resp1 = _post_chat( + c, + { + "messages": [{"role": "user", "content": "weather in SF?"}], + "tools": [self.CLIENT_TOOL], + "docsgpt": {"save_conversation": True}, + }, + api_key, + ) + + assert resp1.status_code == 200, resp1.get_data(as_text=True) + body1 = resp1.get_json() + choice1 = body1["choices"][0] + # OpenAI surfaces the pending client tool call. + assert choice1["finish_reason"] == "tool_calls" + tool_calls1 = choice1["message"]["tool_calls"] + assert tool_calls1[0]["function"]["name"] == "get_weather" + assert tool_calls1[0]["id"] == "call_abc" + # A conversation id is returned for the client to thread back. + conv_id = body1.get("docsgpt", {}).get("conversation_id") + assert conv_id + + # Reserved row finalized ``complete`` (with tool_calls); no + # ``pending_tool_state`` and no non-terminal row. + with pg_engine.connect() as conn: + statuses = _row_statuses(conn, conv_id) + assert statuses == ["complete"] + assert _pending_tool_state_count(conn, conv_id) == 0 + assert self._count_non_terminal(conn, user_id) == 0 + + # ---- POST #2: tool result + conversation_id -> agent answers ---- + with _wire_v1_route_db(pg_engine, monkeypatch), patch( + "application.api.answer.services.stream_processor.StreamProcessor" + ".build_agent", + _fake_build_agent, + ): + with app.test_client() as c: + resp2 = _post_chat( + c, + { + "messages": [ + {"role": "user", "content": "weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"city": "SF"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_abc", + "content": "72F sunny", + }, + ], + "conversation_id": conv_id, + "docsgpt": {"save_conversation": True}, + }, + api_key, + ) + + # The regression: pre-fix this is a 400 (resume_from_tool_actions -> + # load_state None -> ValueError). Post-fix it is a 200 answer. + assert resp2.status_code == 200, resp2.get_data(as_text=True) + body2 = resp2.get_json() + choice2 = body2["choices"][0] + assert choice2["finish_reason"] == "stop" + assert choice2["message"]["content"] == _PauseThenAnswerAgent.ANSWER_TEXT + # The answer persisted into the SAME conversation. + assert body2.get("docsgpt", {}).get("conversation_id") == conv_id + + with pg_engine.connect() as conn: + from sqlalchemy import text + + # The answer is a NEW terminal turn appended to the same + # conversation: the original tool-call turn + the answer turn. + statuses = _row_statuses(conn, conv_id) + assert statuses == ["complete", "complete"] + # Nothing left non-terminal anywhere for this user / conversation. + assert self._count_non_terminal(conn, user_id) == 0 + # The appended turn carries the assistant answer. + answer_rows = conn.execute( + text( + "SELECT response FROM conversation_messages " + "WHERE conversation_id = CAST(:c AS uuid) " + "ORDER BY position ASC" + ), + {"c": conv_id}, + ).fetchall() + assert answer_rows[-1][0] == _PauseThenAnswerAgent.ANSWER_TEXT + # Exactly one conversation was used (no orphan sibling created). + conv_count = conn.execute( + text("SELECT count(*) FROM conversations WHERE user_id = :u"), + {"u": user_id}, + ).scalar() + assert conv_count == 1 + From 5d8187126d9b3472fcf55be0a2c1b1b756939a5e Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 9 Jun 2026 19:08:15 +0100 Subject: [PATCH 2/2] fix: tests --- tests/api/v1/test_reasoning_heartbeat_liveness.py | 2 +- tests/api/v1/test_v1_continuation_ownership.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/api/v1/test_reasoning_heartbeat_liveness.py b/tests/api/v1/test_reasoning_heartbeat_liveness.py index 2fc843ab..a1b198f7 100644 --- a/tests/api/v1/test_reasoning_heartbeat_liveness.py +++ b/tests/api/v1/test_reasoning_heartbeat_liveness.py @@ -365,7 +365,7 @@ class TestHeartbeatPumpsDuringReasoning: assert msg_id is not None # Drain the rest, snapshotting the row's heartbeat after each frame. - for _frame in gen: + for _ in gen: with pg_engine.connect() as conn: observed.append(_row(conn, str(msg_id))) diff --git a/tests/api/v1/test_v1_continuation_ownership.py b/tests/api/v1/test_v1_continuation_ownership.py index 757e5396..45b590e7 100644 --- a/tests/api/v1/test_v1_continuation_ownership.py +++ b/tests/api/v1/test_v1_continuation_ownership.py @@ -179,7 +179,7 @@ class TestForeignConversationOwnership: # write was rejected before any persistence. with pg_engine.connect() as conn: rows_after = _row_count(conn, conv_id) - assert rows_after == rows_before == 1 + assert rows_after == rows_before def test_owner_can_continue_into_own_conversation(self, pg_engine, monkeypatch): """Control: the SAME continuation into the caller's OWN conversation