From e96ff8658c8409edac5f65400312520135d482ff Mon Sep 17 00:00:00 2001 From: Alex Date: Sat, 22 Aug 2026 09:44:09 +0100 Subject: [PATCH] fix: more stability for durable tasks, retry strategy, refactor dead code --- .gitignore | 1 + application/agents/tool_executor.py | 60 +++++++---- .../agents/tools/artifact_generator.py | 39 ++++--- .../versions/0030_superseded_messages.py | 55 ++++++++++ application/api/answer/routes/base.py | 78 +++++++++++--- .../answer/services/continuation_service.py | 21 ++++ .../answer/services/conversation_service.py | 14 +++ application/api/user/reconciliation.py | 78 +++++++++----- application/api/user/tasks.py | 73 ++++++++++--- application/llm/handlers/base.py | 40 ++++++- application/storage/db/models.py | 22 ++++ .../storage/db/repositories/conversations.py | 89 +++++++++++++-- .../db/repositories/pending_tool_state.py | 31 ++++++ .../src/conversation/ConversationBubble.tsx | 9 +- .../src/conversation/ConversationMessages.tsx | 12 +++ .../conversation/sandboxLinks.live.test.ts | 49 ++++++++- scripts/e2e/mock_llm.py | 19 +++- .../test_tool_executor_hallucinated_calls.py | 75 ++++++++++++- .../tools/test_artifact_generator_unit.py | 39 ++++--- tests/api/user/test_reconciliation.py | 101 ++++++++++++++++++ tests/api/user/test_tasks.py | 10 +- tests/e2e/specs/tier-a/tool-arg-merge.spec.ts | 32 +++--- tests/llm/handlers/test_llm_handlers.py | 74 +++++++++++++ .../repositories/test_pending_tool_state.py | 53 +++++++++ .../repositories/test_truncate_supersede.py | 83 ++++++++++++++ tests/storage/db/test_migration_0029.py | 17 ++- tests/test_celery.py | 48 +++++++-- 27 files changed, 1060 insertions(+), 162 deletions(-) create mode 100644 application/alembic/versions/0030_superseded_messages.py diff --git a/.gitignore b/.gitignore index b6a4a4e1..47995b39 100644 --- a/.gitignore +++ b/.gitignore @@ -204,3 +204,4 @@ tests/e2e/test-results/ tests/e2e/.e2e-last-run.json application/core/models/foundry_test_internal.yaml bench/ +/benchmark/ diff --git a/application/agents/tool_executor.py b/application/agents/tool_executor.py index 477cc8cf..0b3243e3 100644 --- a/application/agents/tool_executor.py +++ b/application/agents/tool_executor.py @@ -1,4 +1,3 @@ -import json import logging import re import uuid @@ -398,7 +397,7 @@ class ToolExecutor: self.attachments: List[Dict] = [] self.client_tools: Optional[List[Dict]] = None self._name_to_tool: Dict[str, Tuple[str, str]] = {} - # Per-signature failure counts for invented tool names this turn. After + # Per-NAME failure counts for invented tool names this turn. After # ``UNRESOLVABLE_CALL_LIMIT`` the model is handed a directive message # instead of the generic error; this is a prompt-level nudge, not a hard # bound — the loop bound remains ``MAX_TOOL_ITERATIONS``. @@ -895,27 +894,38 @@ class ToolExecutor: a tool call. Fall back to tool names when the mapping has not been built (headless paths, tests). Never the internal tool ids: quoting those tells the model nothing about what to call instead. + + Narrowed to ``tools_dict`` so a call made with a restricted toolset is + never told to call something out of scope, and capped: this string is + returned as a tool RESULT, so it joins the message history and is + re-sent on every later round. Uncapped, a large MCP fleet turns one + failed call into kilobytes of prose duplicating the tool schema the + provider already received. """ - names = sorted(self._tool_to_name.values()) + scope = set(tools_dict or {}) + names = sorted( + name + for (tool_id, _action), name in self._tool_to_name.items() + if not scope or tool_id in scope + ) if not names: names = sorted( {(tool or {}).get("name") or tool_id for tool_id, tool in (tools_dict or {}).items()} ) names = [str(name) for name in names if name != exclude] + if len(names) > self.MAX_ADVERTISED_TOOL_NAMES: + hidden = len(names) - self.MAX_ADVERTISED_TOOL_NAMES + names = names[: self.MAX_ADVERTISED_TOOL_NAMES] + [f"and {hidden} more"] return ", ".join(names) if names else "(none available)" - @staticmethod - def _call_signature(llm_name: str, call_args: Any) -> str: - """Stable key for "the model just made this exact call again".""" - try: - rendered = json.dumps(call_args, sort_keys=True, default=str) - except (TypeError, ValueError): - rendered = str(call_args) - return f"{llm_name}::{rendered}" - - # After this many identical unresolvable calls, refuse rather than re-run. + # After this many unresolvable calls to the same name, refuse rather than + # re-run. UNRESOLVABLE_CALL_LIMIT = 2 + # Cap on names rendered into a failed-call result; see + # :meth:`_available_tool_names`. + MAX_ADVERTISED_TOOL_NAMES = 30 + def execute(self, tools_dict: Dict, call, llm_class_name: str): """Execute a tool call. Yields status events, returns (result, call_id).""" parser = ToolActionParser(llm_class_name, name_mapping=self._name_to_tool) @@ -930,18 +940,26 @@ class ToolExecutor: # third attempt on, hand back a directive message instead of the generic # error. Scoped to an unknown *name*: a registered tool whose arguments # merely failed to decode is recoverable, and refusing it would strand a - # working tool for the rest of the turn. Keyed on the raw payload, so two - # different malformed bodies stay two different signatures. + # working tool for the rest of the turn. + # + # Keyed on the NAME ALONE. Keying on the payload too made the guard a + # no-op in the case it exists for: told a call failed, a model varies + # its arguments on the next attempt, minting a fresh counter every + # round and never reaching the limit. Arguments cannot rescue an + # unknown name anyway — ``ToolActionParser`` resolves from ``call.name`` + # only — and the "two different malformed bodies" case this used to + # protect belongs to REGISTERED tools, which the + # ``llm_name not in self._name_to_tool`` gate already excludes. if unresolvable and llm_name not in self._name_to_tool: - signature = self._call_signature(llm_name, getattr(call, "arguments", None)) - failures = self._unresolvable_calls.get(signature, 0) + failures = self._unresolvable_calls.get(llm_name, 0) # Count before refusing, so the message and the log escalate rather # than freezing at the limit for the rest of the turn. - self._unresolvable_calls[signature] = failures + 1 + self._unresolvable_calls[llm_name] = failures + 1 if failures >= self.UNRESOLVABLE_CALL_LIMIT: repeated = ( - f"'{llm_name}' has already failed {failures} times with these arguments and " - f"will keep failing. Stop calling it and either use a different tool " + f"'{llm_name}' has already failed {failures} times and will keep " + f"failing — it is not a tool that exists. Stop calling it and " + f"either use a different tool " f"({self._available_tool_names(tools_dict, exclude=llm_name)}) or answer without one." ) logger.warning( @@ -987,7 +1005,6 @@ class ToolExecutor: if name_is_known else "the tool name could not be resolved and its arguments were not a JSON object" ) - error_message = f"Error: Failed to parse LLM tool call. Tool name: {llm_name}" logger.error( "tool_call_parse_failed", extra={ @@ -1029,7 +1046,6 @@ class ToolExecutor: return tool_call_data["result"], call_id if tool_id not in tools_dict: - error_message = f"Error: Tool ID '{tool_id}' extracted from LLM call not found in available tools_dict. Available IDs: {list(tools_dict.keys())}" logger.error( "tool_id_not_found", extra={ diff --git a/application/agents/tools/artifact_generator.py b/application/agents/tools/artifact_generator.py index 4aee4298..5ec60dcb 100644 --- a/application/agents/tools/artifact_generator.py +++ b/application/agents/tools/artifact_generator.py @@ -15,7 +15,7 @@ import copy import json import logging import uuid -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple from application.agents.tools.artifact_ref import resolve_artifact_id from application.agents.tools.base import Tool @@ -597,11 +597,10 @@ class ArtifactGeneratorTool(Tool): def _create(self, **kwargs: Any) -> Dict[str, Any]: """Validate, render, and persist a new artifact at version 1.""" kind = kwargs.get("kind") - spec = self._coerce_spec(kwargs.get("spec")) title = kwargs.get("title") if kind not in _KIND_INFO: return {"status": "error", "error": f"unsupported kind: {kind!r}; expected one of {sorted(_KIND_INFO)}."} - valid = self._validate(kind, spec) + spec, valid = self._validate(kind, kwargs.get("spec")) if valid is not None: return valid @@ -674,8 +673,7 @@ class ArtifactGeneratorTool(Tool): self, artifact_id: str, kind: str, spec: Any, action: str, title: Optional[str] = None ) -> Dict[str, Any]: """Validate the new spec, re-render, and append the next version of an existing artifact.""" - spec = self._coerce_spec(spec) - valid = self._validate(kind, spec) + spec, valid = self._validate(kind, spec) if valid is not None: return valid rendered = self._render(kind, spec) @@ -736,16 +734,30 @@ class ArtifactGeneratorTool(Tool): return spec return parsed if isinstance(parsed, dict) else spec - def _validate(self, kind: str, spec: Any) -> Optional[Dict[str, Any]]: - """Return an error payload when ``spec`` is invalid for ``kind``, else None.""" + def _validate(self, kind: str, spec: Any) -> Tuple[Any, Optional[Dict[str, Any]]]: + """Validate ``spec`` for ``kind``, returning ``(spec, error_payload)``. + + The coerced spec comes back with the verdict so a caller cannot + validate one value and then render another: returning only the error + left the caller holding the original, and a stringified spec that + passed validation here would be written to spec.json as a JSON + *string*, which the renderer cannot read. + + Args: + kind: Artifact kind, a key of ``_SCHEMAS``. + spec: The spec as received, possibly JSON-encoded. + + Returns: + ``(coerced_spec, None)`` when valid, else ``(coerced_spec, error)``. + """ spec = self._coerce_spec(spec) if not isinstance(spec, dict): - return {"status": "error", "error": "spec must be a JSON object."} + return spec, {"status": "error", "error": "spec must be a JSON object."} try: jsonschema.validate(spec, _SCHEMAS[kind]) except jsonschema.ValidationError as exc: - return {"status": "error", "error": f"invalid {kind} spec: {exc.message}"} - return None + return spec, {"status": "error", "error": f"invalid {kind} spec: {exc.message}"} + return spec, None def _load_current(self, raw_id: Any) -> Dict[str, Any]: """Resolve a short ref/uuid to its parent-scoped artifact and current-version spec for edit/rewrite.""" @@ -804,10 +816,9 @@ class ArtifactGeneratorTool(Tool): def _render(self, kind: str, spec: Any) -> Dict[str, Any]: """Run the fixed renderer in the sandbox and return the produced file bytes.""" - # ``_validate`` coerces a stringified spec internally but returns only an - # error payload, so validating and rendering the same value is the - # caller's job. Coerce here too: a caller that skipped it would otherwise - # pass validation and then write a JSON *string* to spec.json. + # Callers reach here with the spec ``_validate`` handed back, which is + # already coerced. Kept as a cheap guard for direct callers: it is an + # ``isinstance`` check, not a re-parse, once the value is a dict. spec = self._coerce_spec(spec) session_id = self._resolve_session_id() if session_id is None: diff --git a/application/alembic/versions/0030_superseded_messages.py b/application/alembic/versions/0030_superseded_messages.py new file mode 100644 index 00000000..9a5f8940 --- /dev/null +++ b/application/alembic/versions/0030_superseded_messages.py @@ -0,0 +1,55 @@ +"""0030 superseded_messages — distinguish a replaced turn from a lost answer. + +``truncate_after`` deletes the tail of a conversation on retry/edit, but a +stream can still own one of those rows: the SSE keepalive pump drains the +generator after a client disconnect, so the old stream runs to completion and +then finalizes into a row that no longer exists. That produced an ERROR-level +``answer_persist_failed`` — indistinguishable from a genuinely orphaned answer, +at the rate users hit retry mid-answer. + +The supersede and the late finalize usually happen in DIFFERENT worker +processes, so the signal has to be durable. This table is that signal: one row +per superseded message, written in the same transaction as the delete. There is +deliberately NO foreign key to ``conversation_messages`` — the row it describes +is being deleted microseconds later, and a CASCADE would take the tombstone +with it. + +Rows are swept by the existing retention beat; ``superseded_at`` is indexed for +that. Idempotent both ways. + +Revision ID: 0030_superseded_messages +Revises: 0029_agent_guardrails +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0030_superseded_messages" +down_revision: Union[str, None] = "0029_agent_guardrails" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + """ + CREATE TABLE IF NOT EXISTS superseded_messages ( + message_id uuid PRIMARY KEY, + conversation_id uuid NOT NULL, + superseded_at timestamptz NOT NULL DEFAULT clock_timestamp() + ); + """ + ) + op.execute( + """ + CREATE INDEX IF NOT EXISTS ix_superseded_messages_superseded_at + ON superseded_messages (superseded_at); + """ + ) + + +def downgrade() -> None: + op.execute("DROP INDEX IF EXISTS ix_superseded_messages_superseded_at;") + op.execute("DROP TABLE IF EXISTS superseded_messages;") diff --git a/application/api/answer/routes/base.py b/application/api/answer/routes/base.py index 60acef4c..38791923 100644 --- a/application/api/answer/routes/base.py +++ b/application/api/answer/routes/base.py @@ -1047,22 +1047,50 @@ class BaseAnswerResource: # `ok`, because activity logging never observes the DB. # Emit a distinct, countable signal instead. if finalize_outcome is MessageUpdateOutcome.NOT_FOUND: - logger.error( - "answer_persist_failed: message row %s no longer " - "exists; %d chars of answer were produced and " - "could not be saved (conversation=%s)", - reserved_message_id, - len(response_full or ""), - conversation_id, - extra={ - "alert": "answer_persist_failed", - "message_id": reserved_message_id, - "conversation_id": ( - str(conversation_id) if conversation_id else None - ), - "answer_length": len(response_full or ""), - }, + # The row can be gone for two very different reasons. + # A retry/edit deliberately replaced this turn — the + # user asked for the answer to be discarded, and the + # cancel flag could not reach us because the ticker + # that sets it lives in the superseding request's own + # process. That is routine and must not page anyone. + # Anything else is a genuinely orphaned answer. + superseded = self.conversation_service.was_superseded( + reserved_message_id ) + if superseded: + logger.info( + "stream superseded: message row %s was replaced " + "by a newer turn after %d chars; discarding " + "(conversation=%s)", + reserved_message_id, + len(response_full or ""), + conversation_id, + extra={ + "alert": "stream_superseded", + "message_id": reserved_message_id, + "conversation_id": ( + str(conversation_id) if conversation_id else None + ), + "answer_length": len(response_full or ""), + }, + ) + else: + logger.error( + "answer_persist_failed: message row %s no longer " + "exists; %d chars of answer were produced and " + "could not be saved (conversation=%s)", + reserved_message_id, + len(response_full or ""), + conversation_id, + extra={ + "alert": "answer_persist_failed", + "message_id": reserved_message_id, + "conversation_id": ( + str(conversation_id) if conversation_id else None + ), + "answer_length": len(response_full or ""), + }, + ) else: conversation_id = self.conversation_service.save_conversation( conversation_id, @@ -1379,6 +1407,26 @@ class BaseAnswerResource: return except Exception as e: logger.error(f"Error in stream: {str(e)}", exc_info=True) + # This process took the resume claim, so it owns releasing it. The + # only other way back is ``revert_stale_resuming``'s 600 s grace, + # which leaves the user locked out of their own conversation for + # ten minutes after a resume that errored: ``load_state`` sees + # only ``pending`` rows, so the paused turn is still resumable but + # invisible, and every retry inside the window gets another 409. + # Not on the ``StreamSuperseded`` path above — there the turn was + # deliberately replaced and the row is legitimately gone. + if _continuation and conversation_id: + try: + ContinuationService().release_claim( + str(conversation_id), + decoded_token.get("sub", "local"), + ) + except Exception as release_err: + logger.error( + f"Failed to release resume claim after a failed " + f"resume: {release_err}", + exc_info=True, + ) if reserved_message_id is not None: try: self.conversation_service.finalize_message( diff --git a/application/api/answer/services/continuation_service.py b/application/api/answer/services/continuation_service.py index 20ce8d93..0aa6d645 100644 --- a/application/api/answer/services/continuation_service.py +++ b/application/api/answer/services/continuation_service.py @@ -174,6 +174,27 @@ class ContinuationService: ) return deleted + def release_claim(self, conversation_id: str, user: str) -> bool: + """Return a claimed row to ``pending`` after a resume failed. + + Returns: + True if a claim was released. + """ + with db_session() as conn: + conv = ConversationsRepository(conn).get_by_legacy_id(conversation_id) + if conv is not None: + pg_conv_id = conv["id"] + elif looks_like_uuid(conversation_id): + pg_conv_id = conversation_id + else: + return False + released = PendingToolStateRepository(conn).release_claim(pg_conv_id, user) + if released: + logger.info( + f"Released resume claim for conversation {conversation_id}" + ) + return released + def mark_resuming(self, conversation_id: str, user: str) -> bool: """Flip the pending row to ``resuming`` so a crashed resume can be retried.""" with db_session() as conn: diff --git a/application/api/answer/services/conversation_service.py b/application/api/answer/services/conversation_service.py index 56920e6f..6a0075b9 100644 --- a/application/api/answer/services/conversation_service.py +++ b/application/api/answer/services/conversation_service.py @@ -340,6 +340,20 @@ class ConversationService: message_id, status, ) + def was_superseded(self, message_id: str) -> bool: + """True when a retry/edit truncation deleted this message's row. + + Args: + message_id: The reserved message id the stream was writing to. + + Returns: + True if the row was deliberately superseded rather than lost. + """ + if not message_id: + return False + with db_readonly() as conn: + return ConversationsRepository(conn).was_superseded(message_id) + def heartbeat_message(self, message_id: str) -> bool: """Bump ``message_metadata.last_heartbeat_at`` so the reconciler's staleness sweep counts the row as alive. No-ops on terminal rows. diff --git a/application/api/user/reconciliation.py b/application/api/user/reconciliation.py index b6c02757..4415fdd5 100644 --- a/application/api/user/reconciliation.py +++ b/application/api/user/reconciliation.py @@ -4,10 +4,11 @@ from __future__ import annotations import logging import uuid +from contextlib import contextmanager from datetime import datetime, timezone -from typing import Any, Dict, Optional, TYPE_CHECKING +from typing import Any, Dict, Iterator, Optional, TYPE_CHECKING -from sqlalchemy import Connection +from sqlalchemy import Connection, Engine from application.api.user.idempotency import MAX_TASK_ATTEMPTS from application.core.settings import settings @@ -60,29 +61,54 @@ def run_reconciliation() -> Dict[str, Any]: # ``(user_id, event_type, payload, scope)``. events: list[tuple] = [] - # Each sweep below commits in its own transaction, but the queued - # user-facing events are fanned out once, at the end. Without the - # ``finally``, a transient DB error in a later sweep discards events whose - # writes already landed — and the loss is permanent, because every sweep - # filters out the rows it just made terminal (``find_and_lock_stalled_ - # ingests`` looks only at ``status='active'``, ``find_and_lock_stuck_ - # messages`` only at pending/streaming), so no later tick re-finds them to - # re-emit. The user is left with an upload toast spinning on "Training..." - # forever. Seen live on 2026-08-21 when an Azure DNS blip took Postgres - # unreachable mid-tick. + # Each sweep below commits in its own transaction, and the queued + # user-facing events are fanned out once, at the end. The ``finally`` is + # what makes a LATER sweep's failure non-destructive: without it, a + # transient DB error discards events whose writes already landed — and the + # loss is permanent, because every sweep filters out the rows it just made + # terminal (``find_and_lock_stalled_ingests`` looks only at + # ``status='active'``, ``find_and_lock_stuck_messages`` only at + # pending/streaming), so no later tick re-finds them to re-emit. The user + # is left with an upload toast spinning on "Training..." forever. Seen + # live on 2026-08-21 when an Azure DNS blip took Postgres unreachable + # mid-tick. + # + # Publish-after-commit is preserved by ``_sweep``: each sweep stages its + # own events and they only reach ``events`` once that sweep's transaction + # has committed. A sweep that raises PARTWAY THROUGH takes its staged + # events down with its rolled-back writes, which is the point — otherwise + # the ``finally`` would fan out terminal events for rows that are still + # active, and the next tick would re-find those rows and emit a duplicate. try: return _run_sweeps(engine, summary, events) finally: _publish_events(events) -def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[str, Any]: +@contextmanager +def _sweep(engine: Engine, events: list[tuple]) -> Iterator[tuple[Connection, list[tuple]]]: + """Open one sweep transaction, staging its events until it commits. + + Yields ``(conn, staged)``. Append user-facing events to ``staged``, never + to ``events`` directly: ``staged`` is merged into ``events`` only after the + ``engine.begin()`` block exits cleanly, so a sweep that raises mid-loop + publishes nothing for the writes Postgres just rolled back. + """ + staged: list[tuple] = [] + with engine.begin() as conn: + yield conn, staged + events.extend(staged) + + +def _run_sweeps(engine: Engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[str, Any]: """Run every reconciler sweep, appending user-facing events to ``events``. - Split out of :func:`run_reconciliation` so the caller can guarantee the - queued events are published even when a sweep raises partway through. + Split out of :func:`run_reconciliation` so the caller can guarantee that + events from sweeps that already COMMITTED are published even when a later + sweep raises. Each sweep stages into its own list via :func:`_sweep`; see + the comment in :func:`run_reconciliation` for why staging matters. """ - with engine.begin() as conn: + with _sweep(engine, events) as (conn, staged): repo = ReconciliationRepository(conn) pt_repo = PendingToolStateRepository(conn) for msg in repo.find_and_lock_stuck_messages(): @@ -115,7 +141,7 @@ def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[st # been told the approval was cleared, so landing an answer # here would contradict what the user was shown. repo.mark_message_approval_cleared(msg["id"]) - events.append( + staged.append( ( str(user_id), "tool.approval.cleared", @@ -137,7 +163,7 @@ def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[st }, ) - with engine.begin() as conn: + with _sweep(engine, events) as (conn, staged): repo = ReconciliationRepository(conn) for row in repo.find_and_lock_proposed_tool_calls(): repo.mark_tool_call_failed( @@ -161,7 +187,7 @@ def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[st }, ) - with engine.begin() as conn: + with _sweep(engine, events) as (conn, staged): repo = ReconciliationRepository(conn) for row in repo.find_and_lock_executed_tool_calls(): repo.mark_tool_call_failed( @@ -192,7 +218,7 @@ def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[st # worker kill, no rollback of the partial embed. The 'stalled' flag # ends the re-alert loop and drives the "indexing failed" badge the # sources list derives from this row. - with engine.begin() as conn: + with _sweep(engine, events) as (conn, staged): repo = ReconciliationRepository(conn) for row in repo.find_and_lock_stalled_ingests(): summary["ingests_stalled"] += 1 @@ -214,7 +240,7 @@ def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[st user_id = row.get("user_id") source_id = row.get("source_id") if user_id and source_id: - events.append( + staged.append( ( str(user_id), "source.ingest.failed", @@ -235,7 +261,7 @@ def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[st # via ``_lookup_completed`` returning None for the whole 24 h TTL. # Promote to ``failed`` so a retry can re-claim and either resume # or fail loudly. - with engine.begin() as conn: + with _sweep(engine, events) as (conn, staged): repo = ReconciliationRepository(conn) for row in repo.find_stuck_idempotency_pending( max_attempts=MAX_TASK_ATTEMPTS, @@ -271,7 +297,7 @@ def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[st stuck_age = max( 15, int(_settings.SCHEDULE_RUN_TIMEOUT // 60) + 5, ) - with engine.begin() as conn: + with _sweep(engine, events) as (conn, staged): runs_repo = ScheduleRunsRepository(conn) schedules_repo = SchedulesRepository(conn) for run in runs_repo.list_stuck_running(age_minutes=stuck_age): @@ -292,7 +318,7 @@ def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[st schedules_repo, str(run["schedule_id"]), ) summary["schedule_runs_failed"] += 1 - events.extend( + staged.extend( _schedule_terminal_events( run, error_type="timeout", @@ -315,7 +341,7 @@ def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[st # Q7: scheduler runs orphaned in 'pending' — dispatcher committed but # apply_async failed (broker outage / crash mid-dispatch). - with engine.begin() as conn: + with _sweep(engine, events) as (conn, staged): runs_repo = ScheduleRunsRepository(conn) schedules_repo = SchedulesRepository(conn) for run in runs_repo.list_stuck_pending(age_minutes=stuck_age): @@ -336,7 +362,7 @@ def _run_sweeps(engine, summary: Dict[str, Any], events: list[tuple]) -> Dict[st schedules_repo, str(run["schedule_id"]), ) summary["schedule_runs_failed"] += 1 - events.extend( + staged.extend( _schedule_terminal_events( run, error_type="internal", diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index 8f2e8b4f..b8974226 100644 --- a/application/api/user/tasks.py +++ b/application/api/user/tasks.py @@ -28,24 +28,44 @@ logger = logging.getLogger(__name__) # durability story is grep-able next to the body. Combined with # ``autoretry_for=(Exception,)`` and a bounded ``max_retries`` so a poison # message can't loop forever. +# # ``retry_backoff`` is the factor, NOT a boolean toggle: celery's # ``add_autoretry_behaviour`` OVERWRITES ``retry_kwargs["countdown"]`` whenever # it is truthy (celery/app/autoretry.py). With the old ``retry_backoff=True`` -# the factor was ``int(max(1.0, True)) == 1``, so the declared 60 s countdown -# never reached ``task.retry`` and the three waits were jittered 0-1 s, 0-2 s -# and 0-4 s — a whole retry envelope of at most 7 seconds. A 3.5-minute network -# blip (2026-08-21) therefore exhausted every attempt and published a terminal -# ``source.ingest.failed`` / ``attachment.failed`` to the user. -# Factor 30 gives 30 s / 60 s / 120 s (full-jittered, capped by celery's -# ``retry_backoff_max`` default of 600 s), which actually spans an outage of -# that length. ``countdown`` is deliberately absent: leaving it in only -# documents a value the library discards. +# the factor was ``int(max(1.0, True)) == 1``, so the three waits were jittered +# 0-1 s, 0-2 s and 0-4 s — a whole retry envelope of at most 7 seconds. A +# 3.5-minute network blip (2026-08-21) therefore exhausted every attempt and +# published a terminal ``source.ingest.failed`` / ``attachment.failed``. +# +# celery applies FULL JITTER (``retry_jitter`` defaults to True), so each wait +# is drawn uniformly from ``[0, factor * 2**retries]`` — the ceilings are NOT +# the waits, and the envelope is a distribution, not a guarantee. Factor 60 +# gives ceilings 60/120/240 s: a median envelope of ~210 s, which covers a +# 60 s blip ~98% of the time and a 120 s blip ~85%. Factor 30 would have +# covered the 210 s incident 0% of the time, because its 210 s nominal +# maximum was reachable only by drawing all three waits at their ceiling. +# Raising ``max_retries`` would buy more headroom but is deliberately not +# done here: attempt 6 would newly reach the ``MAX_TASK_ATTEMPTS=5`` +# poison-loop guard in ``idempotency.py``, which is a separate behaviour +# change from widening the waits. +# +# Jitter is deliberately left ON. These nine task types share their +# dependencies (Postgres, object storage, the embedding provider), so an +# outage fails them all at once; ``retry_jitter=False`` would wake every +# in-flight task at the same instant and stampede the service that just +# recovered. +# +# ``max_retries`` is passed at the top level rather than inside +# ``retry_kwargs``: celery captures that dict BY REFERENCE and writes +# ``retry_kwargs["countdown"]`` into it on every retry, so one shared dict +# spread across all the decorators below would let concurrent retries race +# and would permanently mutate this module constant. DURABLE_TASK = dict( bind=True, acks_late=True, autoretry_for=(Exception,), - retry_kwargs={"max_retries": 3}, - retry_backoff=30, + max_retries=3, + retry_backoff=60, ) @@ -706,11 +726,21 @@ def reconciliation_task(self): try: return run_reconciliation() - except Exception as exc: # noqa: BLE001 - housekeeping must never crash the beat loop - logging.getLogger(__name__).warning( - "reconciliation_task failed; the next beat retries: %s", exc, - ) - return {"messages_failed": 0, "tool_calls_failed": 0, "error": True} + except Exception: # noqa: BLE001 - housekeeping must never crash the beat loop + # ``exception`` not ``warning``: without the traceback a programming + # error in a sweep logs one context-free line every 30 s forever and + # there is no way to locate it. Celery still records the tick as + # SUCCESS because the exception is swallowed here, so this log is the + # only channel that carries the stack. + logger.exception("reconciliation_task failed; the next beat retries") + return { + "messages_failed": 0, + "tool_calls_failed": 0, + "ingests_stalled": 0, + "attachments_stalled": 0, + "schedule_runs_failed": 0, + "error": True, + } @celery.task(bind=True, acks_late=False) @@ -732,11 +762,20 @@ def cleanup_message_events(self): MessageEventsRepository, ) + from application.storage.db.repositories.conversations import ( + ConversationsRepository, + ) + ttl_days = settings.MESSAGE_EVENTS_RETENTION_DAYS engine = get_engine() with engine.begin() as conn: deleted = MessageEventsRepository(conn).cleanup_older_than(ttl_days) - return {"deleted": deleted, "ttl_days": ttl_days} + # Supersede tombstones ride the same beat: both are per-stream + # bookkeeping with the same retention story. + tombstones = ConversationsRepository(conn).cleanup_superseded_older_than( + ttl_days + ) + return {"deleted": deleted, "superseded_deleted": tombstones, "ttl_days": ttl_days} @celery.task(bind=True, acks_late=False) diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index b682c0bf..4f5df5ea 100644 --- a/application/llm/handlers/base.py +++ b/application/llm/handlers/base.py @@ -97,6 +97,31 @@ class LLMResponse: return bool(self.tool_calls) and self.finish_reason == "tool_calls" +def _is_restated_payload(existing: str, incoming: str) -> bool: + """True when ``incoming`` restates an already-complete ``existing`` payload. + + Genuine argument deltas are fragments: a strict prefix of a single + top-level JSON object can never parse, because the outer brace balances + only at the final character. So "``existing`` parses on its own" is a + reliable signal that it is finished and ``incoming`` is a restatement + rather than a continuation. + + The character check runs first and is what keeps this O(1) per frame — + parsing the accumulator on every delta would be quadratic in the number + of frames, which is 100 ms+ of CPU for a large artifact spec streamed a + few characters at a time. + """ + if not existing or not incoming: + return False + if existing.rstrip()[-1:] not in ("}", "]") or incoming.lstrip()[:1] not in ("{", "["): + return False + try: + json.loads(existing) + except ValueError: + return False + return True + + class LLMHandler(ABC): """Abstract base class for LLM handlers.""" @@ -1422,7 +1447,20 @@ class LLMHandler(ABC): elif isinstance(existing.arguments, str) and isinstance( call.arguments, str ): - existing.arguments += call.arguments + if _is_restated_payload( + existing.arguments, call.arguments + ): + # Some OpenAI-compatible gateways restate + # the WHOLE argument payload on the finish + # frame for an index instead of streaming + # deltas. Appending produced '{...}{...}', + # which never parses, so the tool call + # failed on every one of the turn's + # MAX_TOOL_ITERATIONS rounds while the turn + # still reported status='complete'. + existing.arguments = call.arguments + else: + existing.arguments += call.arguments else: # Complete (non-delta) payloads: latest wins. existing.arguments = call.arguments diff --git a/application/storage/db/models.py b/application/storage/db/models.py index fd7b8f80..aa80b19b 100644 --- a/application/storage/db/models.py +++ b/application/storage/db/models.py @@ -655,6 +655,28 @@ message_events_table = Table( ) + +# One row per message deleted by ``truncate_after``. A stream that outlives its +# own row — the SSE pump drains a disconnected generator to completion — reads +# this to tell "the user replaced this turn" (expected) from "an answer was +# genuinely orphaned" (an incident). Deliberately no FK to +# ``conversation_messages``: the row it describes is deleted in the same +# transaction, and a CASCADE would take the tombstone with it. +superseded_messages_table = Table( + "superseded_messages", + metadata, + Column("message_id", UUID(as_uuid=True), primary_key=True, nullable=False), + Column("conversation_id", UUID(as_uuid=True), nullable=False), + Column( + "superseded_at", + DateTime(timezone=True), + nullable=False, + server_default=func.clock_timestamp(), + index=True, + ), +) + + shared_conversations_table = Table( "shared_conversations", metadata, diff --git a/application/storage/db/repositories/conversations.py b/application/storage/db/repositories/conversations.py index b3ba3c62..6926f274 100644 --- a/application/storage/db/repositories/conversations.py +++ b/application/storage/db/repositories/conversations.py @@ -1167,16 +1167,23 @@ class ConversationsRepository: turns — so the rows go, but any non-terminal one is first stamped terminal, which leaves an audit trail of what was superseded. - That stamp does NOT survive: the UPDATE and the DELETE below run the - same ``position > :pos`` predicate in one transaction, so every stamped - row is removed microseconds later and the owning stream's late - ``finalize_message`` correctly reports ``NOT_FOUND``, not - ``ALREADY_FAILED``. Do not rely on the stamp for anything but the - ``RETURNING id`` log line. A stream superseded less than - ``STREAM_HEARTBEAT_INTERVAL`` before it finishes is never cancelled - (the ticker in ``api/answer/routes/base.py`` is the only thing that - sets the cancel flag), so it runs to completion and raises - ``answer_persist_failed`` with an answer nothing can save. + That stamp does NOT survive: the UPDATE's predicate is a subset of the + DELETE's (it also filters ``status NOT IN ('complete','failed')``), and + both run in one transaction, so every stamped row is removed + microseconds later and the owning stream's late ``finalize_message`` + correctly reports ``NOT_FOUND``, not ``ALREADY_FAILED``. Do not rely on + the stamp for anything but the ``RETURNING id`` log line. + + What DOES survive is the ``superseded_messages`` tombstone written + below. A stream superseded less than ``STREAM_HEARTBEAT_INTERVAL`` + before it finishes is never cancelled — the ticker in + ``api/answer/routes/base.py`` is the only thing that sets the cancel + flag, and it cannot be set from here because it is a ``threading.Event`` + in another worker process — so the stream runs to completion and + finalizes into a row that is gone. Without the tombstone that logged an + ERROR ``answer_persist_failed``, identical to a genuinely orphaned + answer, every time a user hit retry on an answer shorter than the + heartbeat interval. Args: conversation_id: Conversation whose tail is being trimmed. @@ -1218,6 +1225,22 @@ class ConversationsRepository: conversation_id, ", ".join(str(r[0]) for r in superseded), ) + # Tombstone before the delete, in the same transaction: the late + # finalize usually runs in a different worker process, so the signal + # has to be durable. No FK, or the CASCADE below would remove it. + self._conn.execute( + text( + """ + INSERT INTO superseded_messages (message_id, conversation_id) + SELECT id, conversation_id + FROM conversation_messages + WHERE conversation_id = CAST(:conv_id AS uuid) + AND position > :pos + ON CONFLICT (message_id) DO NOTHING + """ + ), + {"conv_id": conversation_id, "pos": keep_up_to}, + ) result = self._conn.execute( text( "DELETE FROM conversation_messages " @@ -1228,6 +1251,52 @@ class ConversationsRepository: ) return result.rowcount + def was_superseded(self, message_id: str) -> bool: + """True when ``message_id`` was deleted by a retry/edit truncation. + + Lets a stream that outlived its own row tell an expected supersede from + a genuinely orphaned answer. See :meth:`truncate_after`. + + Args: + message_id: The reserved message id the stream was writing to. + + Returns: + True if a tombstone exists for that message. + """ + if not looks_like_uuid(message_id): + return False + row = self._conn.execute( + text( + "SELECT 1 FROM superseded_messages " + "WHERE message_id = CAST(:mid AS uuid)" + ), + {"mid": message_id}, + ).fetchone() + return row is not None + + def cleanup_superseded_older_than(self, days: int) -> int: + """Delete ``superseded_messages`` tombstones older than ``days``. + + The tombstone only has to outlive the stream that was superseded, which + is bounded by the request timeout — days of retention is already + generous, and the table has no other reader. + + Args: + days: Retention window in days. + + Returns: + Number of rows deleted. + """ + result = self._conn.execute( + text( + "DELETE FROM superseded_messages " + "WHERE superseded_at < clock_timestamp() " + "- make_interval(days => :days)" + ), + {"days": days}, + ) + return result.rowcount or 0 + def set_feedback( self, conversation_id: str, position: int, feedback: dict | None, ) -> bool: diff --git a/application/storage/db/repositories/pending_tool_state.py b/application/storage/db/repositories/pending_tool_state.py index 5450e568..09c02e63 100644 --- a/application/storage/db/repositories/pending_tool_state.py +++ b/application/storage/db/repositories/pending_tool_state.py @@ -176,6 +176,37 @@ class PendingToolStateRepository: ) return result.rowcount > 0 + def release_claim(self, conversation_id: str, user_id: str) -> bool: + """Flip a ``resuming`` row back to ``pending`` so a retry can claim it. + + The inverse of :meth:`mark_resuming`, for the process that TOOK the + claim to call when its resume fails. Without it the only way back is + :meth:`revert_stale_resuming`, whose 600 s grace leaves the user + locked out of their own conversation for ten minutes after a resume + that errored — the paused turn stays resumable but invisible, because + :meth:`load_state` only sees ``pending`` rows. + + Deliberately narrow: it matches on ``status = 'resuming'`` only, so it + can never resurrect a row another request has since deleted, and it + does not touch ``expires_at`` (the TTL is unrelated to the claim). + + Returns: + True when a claim was released. + """ + result = self._conn.execute( + text( + """ + UPDATE pending_tool_state + SET status = 'pending', resumed_at = NULL + WHERE conversation_id = CAST(:conv_id AS uuid) + AND user_id = :user_id + AND status = 'resuming' + """ + ), + {"conv_id": conversation_id, "user_id": user_id}, + ) + return result.rowcount > 0 + def revert_stale_resuming( self, grace_seconds: int = 600, diff --git a/frontend/src/conversation/ConversationBubble.tsx b/frontend/src/conversation/ConversationBubble.tsx index f1afef5f..9328b39f 100644 --- a/frontend/src/conversation/ConversationBubble.tsx +++ b/frontend/src/conversation/ConversationBubble.tsx @@ -65,6 +65,12 @@ const ConversationBubble = forwardRef< index?: number, ) => void; filesAttached?: { id: string; fileName: string }[]; + /** + * Every artifact in the conversation, for resolving inline links. Refs + * are conversation-scoped, so a link may point at an earlier turn's file. + * The chip RAIL still uses this turn's artifacts only. + */ + conversationArtifacts?: ReturnType; onOpenArtifact?: (artifact: { id: string; toolName: string }) => void; onToolAction?: ( callId: string, @@ -92,6 +98,7 @@ const ConversationBubble = forwardRef< isStreaming, handleUpdatedQuestionSubmission, filesAttached, + conversationArtifacts, onOpenArtifact, onToolAction, agentId, @@ -431,7 +438,7 @@ const ConversationBubble = forwardRef< // A research run already narrates itself above; the status line // would be a second live indicator away from the point of action. suppressStatusLine={Boolean(research)} - artifacts={completedArtifacts} + artifacts={conversationArtifacts ?? completedArtifacts} onOpenArtifact={onOpenArtifact} renderApproval={(toolCall: ToolCallsType) => (
diff --git a/frontend/src/conversation/ConversationMessages.tsx b/frontend/src/conversation/ConversationMessages.tsx index 30a6cba5..40174b5b 100644 --- a/frontend/src/conversation/ConversationMessages.tsx +++ b/frontend/src/conversation/ConversationMessages.tsx @@ -3,6 +3,7 @@ import { ReactNode, useEffect, useLayoutEffect, + useMemo, useRef, useState, } from 'react'; @@ -19,6 +20,7 @@ import { MessageScrollerViewport, } from '../components/ui/message-scroller'; import Hero from '../Hero'; +import { deriveArtifactChips } from './artifactChips'; import ConversationBubble from './ConversationBubble'; import { FEEDBACK, Query, Status } from './conversationModels'; import StreamingStatusLine from './StreamingStatusLine'; @@ -73,6 +75,15 @@ export default function ConversationMessages({ }: ConversationMessagesProps) { const { t } = useTranslation(); + // Artifact refs (A1, A2, ...) are CONVERSATION-scoped on the backend — a + // ref minted on turn 1 still resolves, and edit_artifact still accepts it, + // on turn 5. Resolving inline links against only the current turn's + // artifacts made "the CSV from earlier" render as dead plain text. + const conversationArtifacts = useMemo( + () => queries.flatMap((query) => deriveArtifactChips(query.tool_calls)), + [queries], + ); + const hasMessages = queries.length > 0; const lastQuery = queries[queries.length - 1]; const lastTurnContentLength = @@ -212,6 +223,7 @@ export default function ConversationMessages({ segments={query.segments} workflowRunId={query.workflow_run_id} research={query.research} + conversationArtifacts={conversationArtifacts} onOpenArtifact={onOpenArtifact} onToolAction={onToolAction} feedback={query.feedback} diff --git a/frontend/src/conversation/sandboxLinks.live.test.ts b/frontend/src/conversation/sandboxLinks.live.test.ts index aee3e0b3..11c3769e 100644 --- a/frontend/src/conversation/sandboxLinks.live.test.ts +++ b/frontend/src/conversation/sandboxLinks.live.test.ts @@ -3,9 +3,14 @@ * * Payloads captured verbatim from `GET /api/get_single_conversation` against a * live run: turn 1's `run_code` wrote a CSV and a PNG (refs A1/A2), turn 2 wrote - * a JSON summary (ref A3). Refs are numbered per *conversation*, so turn 2's - * only artifact is A3 — indexing a turn's artifact list by the ref number picks - * the wrong file, or none. + * a JSON summary (ref A3). Refs are numbered per *conversation*, so indexing a + * turn's artifact list by the ref number picks the wrong file, or none. + * + * `resolveSandboxLink` matches on the ref itself, so what it can resolve is + * decided entirely by the list it is handed. Handed one turn's artifacts it + * degrades an earlier turn's ref to plain text — safe, but wrong, since the + * ref is still live conversation-wide. `ConversationMessages` therefore passes + * the whole conversation's artifacts; both cases are pinned below. */ import { describe, expect, it } from 'vitest'; @@ -46,13 +51,49 @@ describe('conversation-scoped refs against real turn payloads', () => { // The bug: A1 is turn 1's CSV. Positional resolution made it index 0 of // turn 2's list — the JSON summary — and opened the wrong file silently. - it('does not hand turn 1s ref to turn 2s file', () => { + // Never resolving beats resolving to the wrong file, so a list that lacks + // the ref must still degrade to plain text. + it('never substitutes a different file for a ref it does not hold', () => { expect(resolveSandboxLink('artifact:A1', turn2)).toEqual({ kind: 'plain' }); expect(resolveSandboxLink('sandbox:/artifact/A2', turn2)).toEqual({ kind: 'plain', }); }); + // ...but the ref IS still live: `edit_artifact` accepts A1 on turn 2, and + // the model is told so. Given the conversation's artifacts — what + // `ConversationMessages` actually passes — the cross-turn link resolves. + it('resolves an earlier turns ref against the whole conversation', () => { + const conversation = [...turn1, ...turn2]; + expect(resolveSandboxLink('artifact:A1', conversation)).toEqual({ + kind: 'artifact', + artifact: turn1[0], + }); + expect(resolveSandboxLink('sandbox:/artifact/A2', conversation)).toEqual({ + kind: 'artifact', + artifact: turn1[1], + }); + expect(resolveSandboxLink('artifact:A3', conversation)).toEqual({ + kind: 'artifact', + artifact: turn2[0], + }); + }); + + // Last-write-wins must survive widening: a filename re-written on a later + // turn resolves to the newest artifact, not the first match. + it('keeps last-write-wins across the conversation', () => { + const rewritten: SandboxArtifact = { + id: 'a1111111-2222-3333-4444-555555555555', + label: 'monthly_sales.csv', + ref: 'A4', + toolName: 'code_executor', + }; + const conversation = [...turn1, ...turn2, rewritten]; + expect( + resolveSandboxLink('sandbox:/mnt/data/monthly_sales.csv', conversation), + ).toEqual({ kind: 'artifact', artifact: rewritten }); + }); + it('still resolves both of turn 1s files on turn 1', () => { expect(resolveSandboxLink('sandbox:/artifact/A1', turn1)).toEqual({ kind: 'artifact', diff --git a/scripts/e2e/mock_llm.py b/scripts/e2e/mock_llm.py index 10cd19cf..38030f1c 100644 --- a/scripts/e2e/mock_llm.py +++ b/scripts/e2e/mock_llm.py @@ -80,14 +80,20 @@ EMIT_DIRECTIVE = re.compile(r"\[\[MOCK_LLM_EMIT:([A-Za-z0-9_=\-]+)\]\]") # ``repeat`` — TWO frames for the same ``index``, each carrying the COMPLETE # arguments. Some OpenAI-compatible gateways restate a short # argument payload on the finish frame rather than sending a -# delta; DocsGPT's per-index merge appends, so the two combine -# into invalid JSON (``{}`` + ``{}`` -> ``{}{}``). +# delta. The merge recognises the restatement and takes the +# latest, rather than appending into invalid JSON +# (``{}`` + ``{}`` -> ``{}{}``). # ``delta`` — arguments split into genuine partial deltas, which is what the # merge's ``+=`` exists to reassemble. The control case. +# ``truncated`` — a single frame carrying a PREFIX of the arguments, i.e. a +# provider that stopped mid-payload. Unlike ``repeat`` this is +# genuinely unrecoverable: nothing downstream can invent the +# missing bytes, so the turn runs to the iteration cap. # An optional 4th field is a base64url JSON object to send as the arguments; # it defaults to ``{}`` (a zero-parameter action such as ``note_view``). TOOLCALL_DIRECTIVE = re.compile( - r"\[\[MOCK_LLM_TOOLCALL:([A-Za-z0-9_\-]+):(once|repeat|delta)(?::([A-Za-z0-9_=\-]+))?\]\]" + r"\[\[MOCK_LLM_TOOLCALL:([A-Za-z0-9_\-]+):(once|repeat|delta|truncated)" + r"(?::([A-Za-z0-9_=\-]+))?\]\]" ) app = Flask(__name__) @@ -286,6 +292,13 @@ def _toolcall_arg_frames(arguments: str, mode: str) -> list[str]: if mode == "repeat": # The incident shape: the complete payload arrives twice for one index. return [arguments, arguments] + if mode == "truncated": + # A provider that stopped mid-payload. Strip the closing brace so the + # accumulator can never parse, however it is merged. + stripped = arguments.rstrip() + if len(stripped) < 2: + return ['{"'] + return [stripped[:-1]] if mode == "delta": if len(arguments) < 2: return [arguments] diff --git a/tests/agents/test_tool_executor_hallucinated_calls.py b/tests/agents/test_tool_executor_hallucinated_calls.py index f472ae24..a7741b2c 100644 --- a/tests/agents/test_tool_executor_hallucinated_calls.py +++ b/tests/agents/test_tool_executor_hallucinated_calls.py @@ -119,6 +119,62 @@ class TestErrorNamesWhatTheModelCanCall: _events, (result, _call_id) = _drain(executor, _call("make_a_pdf"), tools=tools_dict) assert "create_artifact" in result + def test_only_advertises_tools_in_scope_for_this_call(self): + """A narrowed ``tools_dict`` must not be told about out-of-scope tools. + + The error string is a tool RESULT handed straight back to the model, so + naming a tool it cannot call this round just buys another failed round. + """ + executor = ToolExecutor() + wide = { + f"t{n}": { + "name": f"server_{n}", + "actions": [ + { + "name": f"action_{n}", + "description": "D", + "active": True, + "parameters": {"properties": {}}, + } + ], + } + for n in range(4) + } + executor.prepare_tools_for_llm(wide) + narrowed = {"t1": wide["t1"]} + _events, (result, _call_id) = _drain( + executor, _call("make_a_pdf"), tools=narrowed + ) + assert "action_1" in result, result + for out_of_scope in ("action_0", "action_2", "action_3"): + assert out_of_scope not in result, (out_of_scope, result) + + def test_the_advertised_list_is_capped(self): + """This string joins the message history and is re-sent every round. + + Uncapped, a large MCP fleet turns a single failed call into kilobytes + of prose duplicating the tool schema the provider already has. + """ + executor = ToolExecutor() + many = { + f"t{n}": { + "name": f"server_{n}", + "actions": [ + { + "name": f"action_{n:03d}", + "description": "D", + "active": True, + "parameters": {"properties": {}}, + } + ], + } + for n in range(150) + } + executor.prepare_tools_for_llm(many) + _events, (result, _call_id) = _drain(executor, _call("make_a_pdf"), tools=many) + assert "and 120 more" in result, result + assert len(result) < 1000, len(result) + @pytest.mark.unit class TestThrottleScope: @@ -167,14 +223,27 @@ class TestThrottleScope: assert "has already failed" in result, result assert "(none available)" in result, result - def test_distinct_payloads_for_an_invented_name_stay_distinct(self): + def test_varying_payloads_for_an_invented_name_still_trip_the_guard(self): + """Varying the arguments must not reset the count for an unknown name. + + Told a call failed, a model's natural next move is to adjust its + arguments. Keying the counter on name+payload minted a fresh signature + every round, so the guard never fired and the turn ran to + ``MAX_TOOL_ITERATIONS``. Arguments cannot rescue an unknown name: + ``ToolActionParser`` resolves from ``call.name`` alone. + """ executor = ToolExecutor() + results = [] for index, body in enumerate(['{"a": 1}', '{"b": 2}', '{"c": 3}']): _events, (result, _call_id) = _drain( executor, _call("note_view", arguments=body, call_id=f"c{index}") ) - assert "has already failed" not in result, (index, result) - assert len(executor._unresolvable_calls) == 3 + results.append(result) + assert "has already failed" not in results[0], results[0] + assert "has already failed" not in results[1], results[1] + assert "has already failed 2 times" in results[2], results[2] + # One counter for the name, not one per payload. + assert list(executor._unresolvable_calls) == ["note_view"] def test_the_failure_count_keeps_escalating(self): """A count frozen at the limit makes the message and the ops log useless.""" diff --git a/tests/agents/tools/test_artifact_generator_unit.py b/tests/agents/tools/test_artifact_generator_unit.py index 8ec9da89..5c317ff7 100644 --- a/tests/agents/tools/test_artifact_generator_unit.py +++ b/tests/agents/tools/test_artifact_generator_unit.py @@ -39,17 +39,17 @@ def _tool(): def test_validate_accepts_minimal_presentation(): - assert _tool()._validate("presentation", {"slides": [{"title": "x"}]}) is None + assert _tool()._validate("presentation", {"slides": [{"title": "x"}]})[1] is None def test_validate_rejects_missing_required_key(): - err = _tool()._validate("presentation", {"title": "no slides"}) + _spec, err = _tool()._validate("presentation", {"title": "no slides"}) assert err["status"] == "error" assert "invalid presentation spec" in err["error"] def test_validate_rejects_unknown_key(): - err = _tool()._validate("document", {"sections": [], "bogus": 1}) + _spec, err = _tool()._validate("document", {"sections": [], "bogus": 1}) assert err["status"] == "error" @@ -97,12 +97,12 @@ def test_reversion_without_title_falls_back_to_generic(monkeypatch): def test_validate_rejects_wrong_type(): - err = _tool()._validate("spreadsheet", {"sheets": "not-a-list"}) + _spec, err = _tool()._validate("spreadsheet", {"sheets": "not-a-list"}) assert err["status"] == "error" def test_validate_rejects_non_object_spec(): - err = _tool()._validate("pdf", "just a string") + _spec, err = _tool()._validate("pdf", "just a string") assert err["status"] == "error" assert "spec must be a JSON object" in err["error"] @@ -300,16 +300,16 @@ def test_validate_accepts_html_spec(): {"type": "code", "text": "print(1)"}, ], } - assert _tool()._validate("html", spec) is None + assert _tool()._validate("html", spec)[1] is None def test_validate_rejects_html_unknown_block_key(): - err = _tool()._validate("html", {"blocks": [{"type": "paragraph", "text": "x", "bogus": 1}]}) + _spec, err = _tool()._validate("html", {"blocks": [{"type": "paragraph", "text": "x", "bogus": 1}]}) assert err["status"] == "error" def test_validate_rejects_html_bad_heading_level(): - err = _tool()._validate("html", {"blocks": [{"type": "heading", "text": "x", "level": 9}]}) + _spec, err = _tool()._validate("html", {"blocks": [{"type": "heading", "text": "x", "level": 9}]}) assert err["status"] == "error" @@ -480,19 +480,28 @@ def test_validate_accepts_a_json_encoded_spec_string(): """ tool = _tool() spec = json.dumps({"blocks": [{"type": "heading", "text": "Hi"}]}) - assert tool._validate("pdf", spec) is None + assert tool._validate("pdf", spec)[1] is None def test_validate_coerces_the_string_spec_for_the_caller(): + """The coerced spec must come back, not just the verdict. + + Returning only the error left the caller holding the original string: it + passed validation and was then written to spec.json as a JSON *string*, + which the renderer cannot read. + """ tool = _tool() - spec = json.dumps({"sections": [{"heading": "H", "paragraphs": ["p"]}]}) - assert tool._coerce_spec(spec) == {"sections": [{"heading": "H", "paragraphs": ["p"]}]} + raw = json.dumps({"sections": [{"heading": "H", "paragraphs": ["p"]}]}) + spec, err = tool._validate("document", raw) + assert err is None + assert spec == {"sections": [{"heading": "H", "paragraphs": ["p"]}]} + assert tool._coerce_spec(raw) == spec def test_validate_still_rejects_a_json_string_that_is_not_an_object(): tool = _tool() - assert tool._validate("pdf", json.dumps(["not", "an", "object"]))["status"] == "error" - assert tool._validate("pdf", "not json at all")["status"] == "error" + assert tool._validate("pdf", json.dumps(["not", "an", "object"]))[1]["status"] == "error" + assert tool._validate("pdf", "not json at all")[1]["status"] == "error" def test_create_accepts_a_stringified_spec(monkeypatch): @@ -525,12 +534,12 @@ def test_pdf_schema_accepts_heading_level(): {"type": "paragraph", "text": "Body."}, ], } - assert _tool()._validate("pdf", spec) is None + assert _tool()._validate("pdf", spec)[1] is None def test_pdf_schema_still_rejects_an_out_of_range_level(): spec = {"blocks": [{"type": "heading", "text": "x", "level": 9}]} - assert _tool()._validate("pdf", spec)["status"] == "error" + assert _tool()._validate("pdf", spec)[1]["status"] == "error" def test_pdf_renderer_honours_heading_level(): diff --git a/tests/api/user/test_reconciliation.py b/tests/api/user/test_reconciliation.py index 4830d9b5..db97bda5 100644 --- a/tests/api/user/test_reconciliation.py +++ b/tests/api/user/test_reconciliation.py @@ -1023,3 +1023,104 @@ class TestEventsSurviveAFailedSweep: # ...and so did its event. Without the ``finally`` the user's upload # toast would spin on "Training..." forever. assert [p[1] for p in published] == ["source.ingest.failed"] + + +# --------------------------------------------------------------------------- +# Publish-after-commit — a sweep that raises must not fan out its own events +# --------------------------------------------------------------------------- + + +class TestPublishAfterCommit: + """``run_reconciliation`` fans events out in a ``finally``; staging keeps + that from publishing events whose writes Postgres rolled back. + + These use the real ``pg_engine`` rather than the ``_route_engine_to`` fake, + because the whole point is that ``engine.begin()`` is a genuine transaction + that genuinely rolls back. + """ + + @pytest.mark.unit + def test_rolled_back_sweep_publishes_nothing(self, pg_engine, monkeypatch): + """A sweep raising mid-loop must not emit events for reverted rows. + + Without staging, rows 1 and 2 are marked stalled, queue a terminal + ``source.ingest.failed`` each, and then row 3 raises. The marks roll + back — the sources are still ``active`` and their workers may still be + embedding — but the queued events would still be fanned out, flipping + the user's upload toast to a terminal failure and leaving the next tick + to re-find the same rows and emit a duplicate. + """ + from application.api.user import reconciliation as rec + + # ``get_engine`` caches a process-wide global, so pin it to this + # test's ephemeral DB rather than whichever one was created first. + monkeypatch.setattr(rec, "get_engine", lambda: pg_engine) + + source_ids = [ + f"6ba7b81{n}-9dad-11d1-80b4-00c04fd430c8" for n in range(3) + ] + with pg_engine.begin() as conn: + for n, sid in enumerate(source_ids): + _seed_source(conn, source_id=sid, user_id=f"u-{n}", name=f"doc{n}.pdf") + _seed_ingest_progress(conn, source_id=sid, embedded=1, total=10) + + # Blow up on the third row, after the first two have been marked and + # have queued their events. + real_mark = rec.ReconciliationRepository.mark_ingest_stalled + calls = {"n": 0} + + def _boom(self, *args, **kwargs): + calls["n"] += 1 + if calls["n"] == 3: + raise RuntimeError("connection is bad: Network is unreachable") + return real_mark(self, *args, **kwargs) + + publish_patch, captured = _capture_published(None) + with publish_patch, patch.object( + rec.ReconciliationRepository, "mark_ingest_stalled", _boom + ): + with pytest.raises(RuntimeError, match="Network is unreachable"): + rec.run_reconciliation() + + ingest_events = [e for e in captured if e[1] == "source.ingest.failed"] + assert ingest_events == [], ( + "events were published for a sweep whose writes rolled back: " + f"{ingest_events}" + ) + + # ...and the rows really did roll back, so a later tick can re-find them. + with pg_engine.connect() as conn: + statuses = [_ingest_status(conn, sid) for sid in source_ids] + assert statuses == ["active", "active", "active"] + + @pytest.mark.unit + def test_committed_sweep_survives_a_later_failure(self, pg_engine, monkeypatch): + """The property the ``finally`` buys: an earlier sweep's events still + publish when a LATER sweep raises, because that sweep already committed. + """ + from application.api.user import reconciliation as rec + + # ``get_engine`` caches a process-wide global, so pin it to this + # test's ephemeral DB rather than whichever one was created first. + monkeypatch.setattr(rec, "get_engine", lambda: pg_engine) + + sid = "6ba7b820-9dad-11d1-80b4-00c04fd430c8" + with pg_engine.begin() as conn: + _seed_source(conn, source_id=sid, user_id="u-late", name="late.pdf") + _seed_ingest_progress(conn, source_id=sid, embedded=1, total=10) + + # Fail the idempotency sweep, which runs strictly after the ingest + # sweep has committed. + def _boom(self, *args, **kwargs): + raise RuntimeError("connection is bad: Network is unreachable") + + publish_patch, captured = _capture_published(None) + with publish_patch, patch.object( + rec.ReconciliationRepository, "find_stuck_idempotency_pending", _boom + ): + with pytest.raises(RuntimeError, match="Network is unreachable"): + rec.run_reconciliation() + + assert [e[1] for e in captured] == ["source.ingest.failed"] + with pg_engine.connect() as conn: + assert _ingest_status(conn, sid) == "stalled" diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index 832d8183..59982227 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -414,8 +414,12 @@ class TestDurableTaskRetryPolicy: task = getattr(tasks_module, task_name) assert task.acks_late is True assert Exception in task.autoretry_for - assert task.retry_backoff is True - assert task.retry_kwargs == {"max_retries": 3, "countdown": 60} + assert task.retry_backoff == 60 + assert task.max_retries == 3 + # ``retry_kwargs`` is deliberately unset: celery mutates that dict in + # place on every retry, so sharing one across the decorators would + # race. See the DURABLE_TASK comment in application/api/user/tasks.py. + assert not getattr(task, "retry_kwargs", None) @pytest.mark.unit @pytest.mark.parametrize( @@ -653,6 +657,8 @@ class TestCleanupMessageEventsTask: assert result == { "deleted": 1, + # Supersede tombstones ride the same retention beat. + "superseded_deleted": 0, "ttl_days": settings.MESSAGE_EVENTS_RETENTION_DAYS, } # Only the fresh row survives. diff --git a/tests/e2e/specs/tier-a/tool-arg-merge.spec.ts b/tests/e2e/specs/tier-a/tool-arg-merge.spec.ts index d0ecf7fc..68b00ae8 100644 --- a/tests/e2e/specs/tier-a/tool-arg-merge.spec.ts +++ b/tests/e2e/specs/tier-a/tool-arg-merge.spec.ts @@ -48,7 +48,7 @@ const ACTION = 'memory_view'; /** Directive understood by scripts/e2e/mock_llm.py. */ function toolCallDirective( action: string, - mode: 'once' | 'repeat' | 'delta', + mode: 'once' | 'repeat' | 'delta' | 'truncated', args?: Record, ): string { const base = `[[MOCK_LLM_TOOLCALL:${action}:${mode}`; @@ -88,7 +88,7 @@ test.describe('streaming tool-call argument merge', () => { await resetDb(); }); - test('a provider that restates complete arguments makes the tool unrunnable, and the turn still reports success', async ({ + test('a provider that restates complete arguments still runs the tool', async ({ browser, }) => { const { sub, token } = await newUserContext(browser); @@ -105,20 +105,19 @@ test.describe('streaming tool-call argument merge', () => { expect(res.ok()).toBeTruthy(); const body = await res.text(); - // ---- the failure is silent by every signal a user or a dashboard sees expect(body).toContain('"type": "end"'); expect(await messageStatuses(sub)).toEqual(['complete']); - // ---- ...and total by the only signal that records it + // A restated payload is recognised as a restatement, not a delta, so + // the arguments still parse and the call resolves to the real tool. + // Before the fix these were appended into '{}{}', which never parsed: + // every attempt landed as tool_name='unknown' and the turn burned all + // of MAX_TOOL_ITERATIONS while still reporting status='complete'. const attempts = await attemptsFor(sub); expect(attempts.length).toBeGreaterThan(0); - expect(attempts.every((a) => a.status === 'failed')).toBe(true); - // `tool_name` is 'unknown' because parsing dies before the name is - // resolved — the production fingerprint. - expect(attempts.every((a) => a.tool_name === 'unknown')).toBe(true); + expect(attempts.some((a) => a.tool_name === 'unknown')).toBe(false); + expect(attempts[0].tool_name).toBe('memory'); expect(attempts.every((a) => a.action_name === ACTION)).toBe(true); - // The tool itself never executed: not one attempt reached the registry. - expect(attempts.some((a) => a.tool_name === 'memory')).toBe(false); } finally { await api.dispose(); } @@ -176,7 +175,7 @@ test.describe('streaming tool-call argument merge', () => { } }); - test('a registered tool with unparseable arguments is re-asked every round, up to the iteration cap', async ({ + test('a registered tool with genuinely unparseable arguments is re-asked every round, up to the iteration cap', async ({ browser, }) => { const { sub, token } = await newUserContext(browser); @@ -185,7 +184,7 @@ test.describe('streaming tool-call argument merge', () => { try { const res = await api.post('/stream', { data: { - question: `check my memory ${toolCallDirective(ACTION, 'repeat')}`, + question: `check my memory ${toolCallDirective(ACTION, 'truncated')}`, save_conversation: true, }, timeout: 180_000, @@ -203,11 +202,10 @@ test.describe('streaming tool-call argument merge', () => { expect(body).toContain('its arguments were not a valid JSON object'); expect(body).not.toContain('has already failed'); - // That reasoning holds when the model can fix its own output. It cannot - // here — the arguments are mangled downstream of the model by the merge, - // so the retry is unwinnable and the only thing that ends the turn is - // MAX_TOOL_ITERATIONS (application/llm/handlers/base.py:17). Every one of - // these 25 rounds is a billed provider call that could never succeed. + // A genuinely truncated payload stays unwinnable — the model cannot fix + // what it never finished emitting — so the turn still runs to + // MAX_TOOL_ITERATIONS (application/llm/handlers/base.py:17). This is the + // cap doing its job, not the merge bug: `repeat` above now resolves. const attempts = await attemptsFor(sub); expect(attempts.length).toBe(25); } finally { diff --git a/tests/llm/handlers/test_llm_handlers.py b/tests/llm/handlers/test_llm_handlers.py index 1b068eba..51d9d76f 100644 --- a/tests/llm/handlers/test_llm_handlers.py +++ b/tests/llm/handlers/test_llm_handlers.py @@ -1293,6 +1293,80 @@ class TestHandleStreaming: executed_call = tool_call_args[0][1] assert executed_call.arguments == '{"q":"test"}' + @pytest.mark.parametrize( + "first,second,expected", + [ + # A gateway that restates the whole payload on the finish frame + # instead of streaming deltas. Appending gave '{}{}' / the doubled + # object, which never parses. + ("{}", "{}", "{}"), + ('{"path": "/"}', '{"path": "/"}', '{"path": "/"}'), + # ...but genuine deltas must still concatenate. + ('{"q":', '"test"}', '{"q":"test"}'), + ('{"a": 1', ', "b": 2}', '{"a": 1, "b": 2}'), + ], + ) + def test_restated_tool_arguments_replace_rather_than_append( + self, first, second, expected + ): + """A complete payload restated for the same index must not be appended. + + A strict prefix of one top-level JSON object never parses, so + "the accumulator already parses" cleanly separates a restatement from + a continuation. + """ + import json as _json + + handler = ConcreteHandler() + agent = Mock() + agent.llm = Mock() + agent.model_id = "test" + agent.tools = [] + agent._check_context_limit = Mock(return_value=False) + agent.context_limit_reached = False + agent.llm.__class__.__name__ = "MockLLM" + agent.tool_executor.check_pause = Mock(return_value=None) + + chunk1 = LLMResponse( + content="", + tool_calls=[ToolCall(id="c1", name="search", arguments=first, index=0)], + finish_reason="", + raw_response={}, + ) + chunk2 = LLMResponse( + content="", + tool_calls=[ToolCall(id="", name="", arguments=second, index=0)], + finish_reason="tool_calls", + raw_response={}, + ) + + handler.parse_response = lambda c: c + + def fake_iterate(response): + yield from response + + handler._iterate_stream = fake_iterate + + def fake_execute(tools_dict, call): + yield {"type": "tool_call", "data": {"status": "pending"}} + return ("result", call.id) + + agent._execute_tool_action = Mock(side_effect=fake_execute) + agent.llm.gen_stream = Mock( + return_value=[ + LLMResponse( + content="done", tool_calls=[], finish_reason="stop", raw_response={} + ) + ] + ) + + list(handler.handle_streaming(agent, [chunk1, chunk2], {"1": {"name": "t"}}, [])) + + executed_call = agent._execute_tool_action.call_args[0][1] + assert executed_call.arguments == expected + # The whole point: the merged payload parses. + _json.loads(executed_call.arguments) + def test_context_limit_adds_system_message(self): handler = ConcreteHandler() agent = Mock() diff --git a/tests/storage/db/repositories/test_pending_tool_state.py b/tests/storage/db/repositories/test_pending_tool_state.py index 05f65fd7..3c16c563 100644 --- a/tests/storage/db/repositories/test_pending_tool_state.py +++ b/tests/storage/db/repositories/test_pending_tool_state.py @@ -222,3 +222,56 @@ class TestRevertStaleResuming: loaded = repo.load_state(conv["id"], "user-1") assert loaded is not None assert loaded["status"] == "pending" + + +class TestReleaseClaim: + def test_returns_a_claimed_row_to_pending(self, pg_conn): + """A failed resume must hand the claim back immediately. + + Without this the only way back is ``revert_stale_resuming``'s 600 s + grace, and ``load_state`` ignores ``resuming`` rows — so the paused + turn stays resumable but invisible and the user is locked out of their + own conversation for ten minutes. + """ + conv = _conv(pg_conn) + repo = _repo(pg_conn) + repo.save_state(conv["id"], "user-1", **_sample_state()) + assert repo.claim_state(conv["id"], "user-1") is not None + assert repo.load_state(conv["id"], "user-1") is None + + assert repo.release_claim(conv["id"], "user-1") is True + + reclaimed = repo.load_state(conv["id"], "user-1") + assert reclaimed is not None + assert reclaimed["status"] == "pending" + assert reclaimed["resumed_at"] is None + # ...and the next resume attempt can take it straight away. + assert repo.claim_state(conv["id"], "user-1") is not None + + def test_noop_on_a_pending_row(self, pg_conn): + """Only a claim can be released; a live pending row is left alone.""" + conv = _conv(pg_conn) + repo = _repo(pg_conn) + repo.save_state(conv["id"], "user-1", **_sample_state()) + assert repo.release_claim(conv["id"], "user-1") is False + assert repo.load_state(conv["id"], "user-1")["status"] == "pending" + + def test_noop_when_the_row_is_gone(self, pg_conn): + """Must not resurrect a row another request deleted.""" + conv = _conv(pg_conn) + repo = _repo(pg_conn) + repo.save_state(conv["id"], "user-1", **_sample_state()) + repo.claim_state(conv["id"], "user-1") + repo.delete_state(conv["id"], "user-1") + + assert repo.release_claim(conv["id"], "user-1") is False + assert repo.load_state_any(conv["id"], "user-1") is None + + def test_scoped_to_the_owning_user(self, pg_conn): + conv = _conv(pg_conn) + repo = _repo(pg_conn) + repo.save_state(conv["id"], "user-1", **_sample_state()) + repo.claim_state(conv["id"], "user-1") + + assert repo.release_claim(conv["id"], "someone-else") is False + assert repo.load_state_any(conv["id"], "user-1")["status"] == "resuming" diff --git a/tests/storage/db/repositories/test_truncate_supersede.py b/tests/storage/db/repositories/test_truncate_supersede.py index 8eb800f2..7c5e7c54 100644 --- a/tests/storage/db/repositories/test_truncate_supersede.py +++ b/tests/storage/db/repositories/test_truncate_supersede.py @@ -118,3 +118,86 @@ class TestTruncateAfterSupersede: def test_invalid_conversation_id_is_rejected(self, pg_conn): repo = ConversationsRepository(pg_conn) assert repo.truncate_after("not-a-uuid", keep_up_to=0) == 0 + + +class TestSupersedeTombstone: + """``NOT_FOUND`` alone cannot tell a replaced turn from a lost answer. + + The cancel flag is a ``threading.Event`` in the superseding request's own + process, so a stream in another worker never sees it and runs to + completion. Its late finalize then hits a row that is gone — routine on the + retry/edit path, and previously logged as an ERROR indistinguishable from a + genuinely orphaned answer. + """ + + def test_truncate_tombstones_every_row_it_deletes(self, pg_conn): + repo = ConversationsRepository(pg_conn) + conv = repo.create("u", "tombstone") + live = _seed(pg_conn, conv["id"], 0, "streaming") + done = _seed(pg_conn, conv["id"], 1, "complete") + + repo.truncate_after(conv["id"], keep_up_to=-1) + + # Both, not just the in-flight one: the UPDATE's status filter is + # narrower than the DELETE's, so the stamp alone would miss `done`. + assert repo.was_superseded(live) is True + assert repo.was_superseded(done) is True + + def test_tombstone_survives_the_delete_that_created_it(self, pg_conn): + """No FK to conversation_messages, or the CASCADE would take it.""" + repo = ConversationsRepository(pg_conn) + conv = repo.create("u", "cascade") + live = _seed(pg_conn, conv["id"], 0, "streaming") + + repo.truncate_after(conv["id"], keep_up_to=-1) + + assert _row(pg_conn, live) is None + assert repo.was_superseded(live) is True + + def test_an_untouched_message_has_no_tombstone(self, pg_conn): + """A genuinely orphaned answer must still reach the ERROR path.""" + repo = ConversationsRepository(pg_conn) + conv = repo.create("u", "kept") + kept = _seed(pg_conn, conv["id"], 0, "streaming") + dropped = _seed(pg_conn, conv["id"], 1, "streaming") + + repo.truncate_after(conv["id"], keep_up_to=0) + + assert repo.was_superseded(kept) is False + assert repo.was_superseded(dropped) is True + + def test_was_superseded_shrugs_off_a_non_uuid(self, pg_conn): + """Shape-gate: a legacy ObjectId must not poison the transaction.""" + repo = ConversationsRepository(pg_conn) + assert repo.was_superseded("507f1f77bcf86cd799439011") is False + assert repo.was_superseded("") is False + + def test_truncate_is_repeatable(self, pg_conn): + """A second truncate over the same range must not raise on the PK.""" + repo = ConversationsRepository(pg_conn) + conv = repo.create("u", "twice") + first = _seed(pg_conn, conv["id"], 0, "streaming") + repo.truncate_after(conv["id"], keep_up_to=-1) + _seed(pg_conn, conv["id"], 0, "streaming") + repo.truncate_after(conv["id"], keep_up_to=-1) + assert repo.was_superseded(first) is True + + def test_cleanup_sweeps_old_tombstones_only(self, pg_conn): + repo = ConversationsRepository(pg_conn) + conv = repo.create("u", "sweep") + recent = _seed(pg_conn, conv["id"], 0, "streaming") + repo.truncate_after(conv["id"], keep_up_to=-1) + stale = _seed(pg_conn, conv["id"], 0, "streaming") + repo.truncate_after(conv["id"], keep_up_to=-1) + pg_conn.execute( + text( + "UPDATE superseded_messages " + "SET superseded_at = clock_timestamp() - make_interval(days => 30) " + "WHERE message_id = CAST(:mid AS uuid)" + ), + {"mid": stale}, + ) + + assert repo.cleanup_superseded_older_than(14) == 1 + assert repo.was_superseded(stale) is False + assert repo.was_superseded(recent) is True diff --git a/tests/storage/db/test_migration_0029.py b/tests/storage/db/test_migration_0029.py index 78b9eab9..9987c00c 100644 --- a/tests/storage/db/test_migration_0029.py +++ b/tests/storage/db/test_migration_0029.py @@ -29,6 +29,15 @@ def _run_alembic(url: str, *args: str) -> None: ) +def _alembic_history(url: str) -> str: + return subprocess.check_output( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "history"], + timeout=60, + env={**os.environ, "POSTGRES_URI": url}, + text=True, + ) + + def _alembic_heads(url: str) -> list[str]: out = subprocess.check_output( [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "heads"], @@ -61,10 +70,16 @@ def _table_exists(conn, table: str) -> bool: class TestMigration0029RoundTrip: def test_single_head(self, pg_engine): + """One head, and 0029 still on the path to it. + + Deliberately does NOT assert 0029 *is* the head — every later + migration would break that — only that the chain stays linear and + 0029 has not been orphaned off it. + """ url = pg_engine.url.render_as_string(hide_password=False) heads = _alembic_heads(url) assert len(heads) == 1, f"expected one alembic head, got {heads}" - assert _0029 in heads[0] + assert _0029 in _alembic_history(url) def test_upgrade_creates_column_and_table(self, pg_engine): with pg_engine.connect() as conn: diff --git a/tests/test_celery.py b/tests/test_celery.py index 345c7ca3..f6d3477b 100644 --- a/tests/test_celery.py +++ b/tests/test_celery.py @@ -62,14 +62,17 @@ def test_durable_task_retry_envelope_spans_a_multi_minute_outage(): from application.api.user.tasks import DURABLE_TASK - assert DURABLE_TASK["retry_backoff"] == 30 - # Leaving a ``countdown`` in would only document a value celery discards. - assert "countdown" not in DURABLE_TASK["retry_kwargs"] + assert DURABLE_TASK["retry_backoff"] == 60 + # ``retry_kwargs`` must stay absent: celery captures it by reference and + # writes ``countdown`` into it on every retry, so one dict shared across + # every decorator would race and would mutate the module constant. + assert "retry_kwargs" not in DURABLE_TASK + assert DURABLE_TASK["max_retries"] == 3 seen: list[dict] = [] class _Task: - max_retries = DURABLE_TASK["retry_kwargs"]["max_retries"] + max_retries = DURABLE_TASK["max_retries"] def __init__(self) -> None: self.request = MagicMock() @@ -85,20 +88,43 @@ def test_durable_task_retry_envelope_spans_a_multi_minute_outage(): add_autoretry_behaviour( task, autoretry_for=DURABLE_TASK["autoretry_for"], - retry_kwargs=dict(DURABLE_TASK["retry_kwargs"]), + retry_kwargs={"max_retries": DURABLE_TASK["max_retries"]}, retry_backoff=DURABLE_TASK["retry_backoff"], ) ceilings = [] - for attempt in range(DURABLE_TASK["retry_kwargs"]["max_retries"]): + for attempt in range(DURABLE_TASK["max_retries"]): task.request.retries = attempt with pytest.raises(RuntimeError, match="retry-sentinel"): task.run() - ceilings.append(30 * 2**attempt) + ceilings.append(DURABLE_TASK["retry_backoff"] * 2**attempt) # full jitter: the wait is drawn from [0, factor * 2**retries]. assert 0 <= seen[-1]["countdown"] <= ceilings[-1] - # 30 / 60 / 120 — an envelope that can actually outlast a multi-minute - # blip, versus the 1 / 2 / 4 the old settings produced. - assert ceilings == [30, 60, 120] - assert sum(ceilings) >= 200 + # 60 / 120 / 240 ceilings, versus the 1 / 2 / 4 the old settings produced. + assert ceilings == [60, 120, 240] + + # Full jitter means the ceilings are not the waits, so pin the + # DISTRIBUTION rather than the nominal sum: asserting only + # ``0 <= countdown <= ceiling`` passes even when every draw is 0. + import statistics + + from celery.utils.time import get_exponential_backoff_interval + + envelopes = [ + sum( + get_exponential_backoff_interval( + factor=DURABLE_TASK["retry_backoff"], + retries=attempt, + maximum=600, + full_jitter=True, + ) + for attempt in range(DURABLE_TASK["max_retries"]) + ) + for _ in range(2000) + ] + # Median envelope is half the 420 s nominal; a wide margin keeps this + # from flaking while still failing if the factor or jitter regresses. + assert 150 <= statistics.median(envelopes) <= 270 + # A 60 s outage must be survivable in the large majority of retry runs. + assert sum(e >= 60 for e in envelopes) / len(envelopes) > 0.9