diff --git a/src/cli_agent_orchestrator/api/main.py b/src/cli_agent_orchestrator/api/main.py index 411423129..20c8c1bd7 100644 --- a/src/cli_agent_orchestrator/api/main.py +++ b/src/cli_agent_orchestrator/api/main.py @@ -41,9 +41,12 @@ from cli_agent_orchestrator.backends.registry import get_backend from cli_agent_orchestrator.clients.database import ( create_inbox_message, + delete_old_handoff_results, + get_handoff_result, get_inbox_messages, get_terminal_metadata, init_db, + upsert_handoff_result, ) from cli_agent_orchestrator.constants import ( ALLOWED_HOSTS, @@ -313,6 +316,28 @@ def validate_env_var_shape(self) -> "RunStepRequest": ) return self + job_id: Optional[str] = Field( + default=None, + description=( + "Caller-generated opaque identifier for this job (issue #447). " + "When present, the server persists state and result under this key " + "so the caller can retrieve the result via GET /handoff-results/{job_id} " + "if the transport times out before the response arrives. " + "Must be a 32-character lowercase hex string (uuid4().hex)." + ), + ) + + @field_validator("job_id") + @classmethod + def _validate_job_id(cls, v: Optional[str]) -> Optional[str]: + if v is None: + return v + if not re.fullmatch(r"[0-9a-f]{32}", v): + raise ValueError( + "job_id must be a 32-character lowercase hex string (e.g. uuid4().hex)" + ) + return v + class RunStepResponse(BaseModel): """Response wrapping an ``AgentStepResult`` from ``run_agent_step``.""" @@ -1805,7 +1830,20 @@ async def run_step( The plugin registry is threaded so teardown's ``post_kill_terminal`` hooks fire (parity with the DELETE endpoint). + + Durability (issue #447): when the caller supplies a ``job_id``, the handler + records state="running" at request start. On success, ``run_agent_step`` + itself persists state="completed" BETWEEN extraction and teardown (not + here, and not after it returns) — the terminal being torn down is the only + other place the result lives, so persistence must land before that + happens, not merely before the HTTP response. On failure, this handler + persists state="error". A caller that misses the response due to an + MCP-transport timeout can retrieve the result via + ``GET /handoff-results/{job_id}`` at any point thereafter. Requests without + a ``job_id`` behave exactly as before (no persistence overhead). """ + job_id = body.job_id # None when the caller omits it (backward-compatible) + # BR-31: for a script-tier run-step call, record the created terminal into the # shared ScriptRunRecord's step_states AT creation time, so U4's orphan sweep # can tear it down if the subprocess dies mid-call. No-op for YAML/handoff @@ -1857,6 +1895,16 @@ def _settle_step(terminal_id: Optional[str], error: Optional[str]) -> None: except KeyError as e: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) + # Mark the job as in-progress before the long-running substrate starts. + # This is best-effort; a failure here must not block execution. Placed + # AFTER the generation fence above so a fenced-out (stale-generation) + # call never leaves a job_id stuck at "running" with no terminal state. + if job_id: + try: + upsert_handoff_result(job_id, "running") + except Exception: + logger.warning("run_step: failed to record job_id=%s as running", job_id) + try: result = await run_agent_step( provider=body.provider, @@ -1872,11 +1920,19 @@ def _settle_step(terminal_id: Optional[str], error: Optional[str]) -> None: registry=get_plugin_registry(request), env_vars=body.env_vars, on_terminal_created=on_terminal_created, + job_id=job_id, ) # Success -> transition the script step RUNNING->COMPLETED (no-op for # non-script callers). Before building the response so a settle failure # is logged, not raised. _settle_step(result.terminal_id, None) + + # NOTE (issue #447 / PR #453 review): the "completed" persist happens + # INSIDE run_agent_step, between extraction and teardown — not here. + # Persisting only after this call returns would run after the terminal + # (the only other copy of the result) has already been torn down, + # contradicting the "persist result, then tear down" requirement. + return RunStepResponse( terminal_id=result.terminal_id, last_message=result.last_message, @@ -1891,6 +1947,16 @@ def _settle_step(terminal_id: Optional[str], error: Optional[str]) -> None: # rather than regex-scraping the message (the future engine reads it too). # Transition the script step RUNNING->FAILED (no-op for non-script callers). _settle_step(e.terminal_id, str(e)) + if job_id: + try: + upsert_handoff_result( + job_id, + "error", + terminal_id=e.terminal_id, + error_message=str(e), + ) + except Exception: + logger.warning("run_step: failed to persist error result for job_id=%s", job_id) code = status.HTTP_502_BAD_GATEWAY if e.kind == "error" else status.HTTP_504_GATEWAY_TIMEOUT raise HTTPException( status_code=code, @@ -1898,15 +1964,30 @@ def _settle_step(terminal_id: Optional[str], error: Optional[str]) -> None: ) except TimeoutError as e: _settle_step(None, str(e)) + if job_id: + try: + upsert_handoff_result(job_id, "error", error_message=str(e)) + except Exception: + logger.warning("run_step: failed to persist TimeoutError for job_id=%s", job_id) raise HTTPException( status_code=status.HTTP_504_GATEWAY_TIMEOUT, detail={"message": str(e), "kind": "timeout", "terminal_id": None}, ) except ValueError as e: # Unknown terminal / bad input surfaced by the terminal layer. + if job_id: + try: + upsert_handoff_result(job_id, "error", error_message=str(e)) + except Exception: + logger.warning("run_step: failed to persist ValueError for job_id=%s", job_id) raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) except Exception as e: _settle_step(None, str(e)) + if job_id: + try: + upsert_handoff_result(job_id, "error", error_message=str(e)) + except Exception: + logger.warning("run_step: failed to persist Exception for job_id=%s", job_id) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to run step: {str(e)}", @@ -2426,6 +2507,42 @@ async def export_graph_endpoint( return {"written_files": written_files, "sink": body.sink, "dest": body.dest} +@app.get( + "/handoff-results/{job_id}", + summary="Retrieve a durable handoff step result (issue #447)", + description=( + "Returns the persisted state and result for a handoff job identified by " + "``job_id`` (generated by the MCP client and passed to ``POST /terminals/run-step`` " + "in the ``job_id`` field). Use this to recover a result that was not delivered " + "over the original request because the MCP transport timed out. " + "Possible ``state`` values: ``running`` (step still in progress), " + "``completed`` (result in ``last_message``), ``error`` (failure in ``error_message``)." + ), +) +async def get_handoff_result_endpoint( + job_id: str, + _scopes: List[str] = Depends(require_any_scope(SCOPE_READ, SCOPE_WRITE, SCOPE_ADMIN)), +) -> Dict: + """Retrieve a durable handoff step result by job_id (issue #447). + + Returns the persisted record when found; 404 when the job_id is unknown + (either the caller never sent a job_id, or the record has been purged by + the retention sweep). Scope-gated (PR #453 review finding 4): ``last_message`` + can carry worker output (prompts/secrets), so this follows the same + ``require_any_scope`` posture as other content-serving GETs (``/events``, + ``/memory/export``) rather than staying open. A no-op when auth is + disabled (the default) — ``require_any_scope`` only enforces when an IdP + is configured. + """ + record = get_handoff_result(job_id) + if record is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Handoff result not found for job_id '{job_id}'", + ) + return record + + @app.delete("/terminals/{terminal_id}") async def delete_terminal( request: Request, diff --git a/src/cli_agent_orchestrator/clients/database.py b/src/cli_agent_orchestrator/clients/database.py index ed9a72107..02dc113ea 100644 --- a/src/cli_agent_orchestrator/clients/database.py +++ b/src/cli_agent_orchestrator/clients/database.py @@ -146,6 +146,31 @@ class FlowModel(Base): enabled = Column(Boolean, default=True) +class HandoffResultModel(Base): + """Durable record of a handoff step result (issue #447). + + Written by the run-step handler BEFORE returning the HTTP response, so the + result survives an MCP-transport timeout. The caller supplies a ``job_id`` + (generated client-side so retries use the same key); the server upserts on + that key. + + ``state``: + - ``"running"`` — step in progress (written at request start) + - ``"completed"`` — step finished successfully; ``last_message`` populated + - ``"error"`` — step failed; ``error_message`` populated + """ + + __tablename__ = "handoff_results" + + job_id = Column(String, primary_key=True) + state = Column(String, nullable=False) # "running" | "completed" | "error" + terminal_id = Column(String, nullable=True) + last_message = Column(Text, nullable=True) + error_message = Column(Text, nullable=True) + created_at = Column(DateTime(timezone=True), default=_utcnow) + updated_at = Column(DateTime(timezone=True), default=_utcnow, onupdate=_utcnow) + + def _ensure_db_dir() -> None: """Create the DB dir owner-only (0o700). @@ -181,6 +206,7 @@ def init_db() -> None: _migrate_workflow_index() _migrate_workflow_run() _migrate_workflow_run_step() + _migrate_add_handoff_results() def _restrict_db_file_permissions() -> None: @@ -206,6 +232,36 @@ def _restrict_db_file_permissions() -> None: logger.warning(f"Could not restrict DB file permissions on {path}: {e}") +def _migrate_add_handoff_results() -> None: + """Create the handoff_results table on existing databases (issue #447). + + ``Base.metadata.create_all`` already handles fresh databases; this + idempotent migration handles existing ones where the table does not + exist yet. SQLite supports ``CREATE TABLE IF NOT EXISTS``, so we + delegate to raw SQL rather than a full schema rebuild. + """ + import sqlite3 + + from cli_agent_orchestrator.constants import DATABASE_FILE + + try: + with sqlite3.connect(str(DATABASE_FILE)) as conn: + conn.execute(""" + CREATE TABLE IF NOT EXISTS handoff_results ( + job_id TEXT PRIMARY KEY, + state TEXT NOT NULL, + terminal_id TEXT, + last_message TEXT, + error_message TEXT, + created_at DATETIME, + updated_at DATETIME + ) + """) + conn.commit() + except Exception as e: + logger.warning(f"Migration check for handoff_results failed: {e}") + + def _migrate_project_aliases_schema() -> None: """Rebuild project_aliases if it predates the alias-only primary key. @@ -827,6 +883,86 @@ def list_aliases_for_project(project_id: str) -> List[Dict[str, Any]]: return [] +# --------------------------------------------------------------------------- +# Handoff result durability helpers (issue #447) +# --------------------------------------------------------------------------- + + +def upsert_handoff_result( + job_id: str, + state: str, + *, + terminal_id: Optional[str] = None, + last_message: Optional[str] = None, + error_message: Optional[str] = None, +) -> None: + """Create or update the durable record for a handoff step (issue #447). + + Called at two points in the run-step handler: + 1. Request start — ``state="running"`` (marks the job as in-progress so a + concurrent duplicate knows to wait rather than start a second worker). + 2. Before sending the HTTP response — ``state="completed"`` or ``"error"`` + (makes the result retrievable even if the transport closes before the + response arrives). + + Idempotent: a second call for the same ``job_id`` updates the existing row. + """ + now = _utcnow() + with SessionLocal() as db: + row = db.query(HandoffResultModel).filter(HandoffResultModel.job_id == job_id).first() + if row is None: + row = HandoffResultModel( + job_id=job_id, + state=state, + terminal_id=terminal_id, + last_message=last_message, + error_message=error_message, + created_at=now, + updated_at=now, + ) + db.add(row) + else: + row.state = state + if terminal_id is not None: + row.terminal_id = terminal_id + if last_message is not None: + row.last_message = last_message + if error_message is not None: + row.error_message = error_message + row.updated_at = now + db.commit() + + +def get_handoff_result(job_id: str) -> Optional[dict]: + """Return the handoff result record for ``job_id``, or None if not found.""" + with SessionLocal() as db: + row = db.query(HandoffResultModel).filter(HandoffResultModel.job_id == job_id).first() + if row is None: + return None + return { + "job_id": row.job_id, + "state": row.state, + "terminal_id": row.terminal_id, + "last_message": row.last_message, + "error_message": row.error_message, + "created_at": row.created_at, + "updated_at": row.updated_at, + } + + +def delete_old_handoff_results(cutoff: datetime) -> int: + """Delete handoff result rows older than ``cutoff`` (retention sweep). + + Returns the number of rows deleted. + """ + with SessionLocal() as db: + deleted = ( + db.query(HandoffResultModel).filter(HandoffResultModel.created_at < cutoff).delete() + ) + db.commit() + return deleted + + def update_message_status(message_id: int, status: MessageStatus) -> bool: """Update message status to MessageStatus.DELIVERED or MessageStatus.FAILED.""" with SessionLocal() as db: diff --git a/src/cli_agent_orchestrator/mcp_server/models.py b/src/cli_agent_orchestrator/mcp_server/models.py index 4ca70315e..5a6c67148 100644 --- a/src/cli_agent_orchestrator/mcp_server/models.py +++ b/src/cli_agent_orchestrator/mcp_server/models.py @@ -6,9 +6,31 @@ class HandoffResult(BaseModel): - """Result of a handoff operation.""" + """Result of a handoff operation. + + When the MCP transport times out before the step completes (issue #447), + ``pending=True`` and ``job_id`` identify the in-flight job. The caller + can retrieve the final result later via ``GET /handoff-results/{job_id}``. + """ success: bool = Field(description="Whether the handoff was successful") message: str = Field(description="A message describing the result of the handoff") output: Optional[str] = Field(None, description="The output from the target agent") terminal_id: Optional[str] = Field(None, description="The terminal ID used for the handoff") + # Async-retrieval fields (issue #447). Present only when the transport timed + # out before the result arrived — absent (None) on the normal synchronous path. + job_id: Optional[str] = Field( + None, + description=( + "Opaque identifier for this handoff job. Present when pending=True. " + "Use with GET /handoff-results/{job_id} to retrieve the result later." + ), + ) + pending: Optional[bool] = Field( + None, + description=( + "True when the transport timed out before the result was delivered. " + "The job is still running (or completed) server-side; poll " + "GET /handoff-results/{job_id} until state='completed'." + ), + ) diff --git a/src/cli_agent_orchestrator/mcp_server/server.py b/src/cli_agent_orchestrator/mcp_server/server.py index dc5e679a8..3651a4517 100644 --- a/src/cli_agent_orchestrator/mcp_server/server.py +++ b/src/cli_agent_orchestrator/mcp_server/server.py @@ -4,6 +4,7 @@ import os import re import time +import uuid from typing import Any, Dict, NamedTuple, Optional, Tuple, Union import requests @@ -48,6 +49,11 @@ def _mcp_timeout() -> float: TERMINAL_CLEANUP_NUDGE_THRESHOLD = 10 MAX_USER_PROMPT_ANSWER_LENGTH = 4000 +# Extra seconds added on top of the caller's ``timeout`` for the HTTP client +# deadline. Covers the server-side ready-wait (up to DEFAULT_READY_TIMEOUT=120s) +# plus a small scheduling margin (issue #447: also used in the Timeout message). +_CLIENT_TIMEOUT_HEADROOM = 180 + def _get_cleanup_nudge() -> str: """Return a cleanup nudge string if the session has too many terminals, else empty string.""" @@ -709,6 +715,15 @@ async def _handoff_impl( # when provider == codex; otherwise returns the message unchanged). shaped_message = _shape_handoff_message(provider, message) + # Generate a durable job_id BEFORE posting. The server records this in + # handoff_results so the result is retrievable even when the MCP + # transport closes before the response arrives (issue #447). Client-side + # generation means the job_id is available for the Timeout catch block. + # Note: the DB record is idempotent on this key (upsert), but the server + # does NOT deduplicate execution — a retry with the same job_id would + # launch a second worker and race on the final upsert. + job_id = uuid.uuid4().hex + # Single combined call: create -> ready-wait -> input -> complete-wait -> # extract -> teardown, all server-side via run_agent_step. session_name # places the worker in the supervisor's session; caller_id/allowed_tools @@ -719,6 +734,7 @@ async def _handoff_impl( "prompt": shaped_message, "teardown": True, "timeout": float(timeout), + "job_id": job_id, } if ctx.session_name: payload["session_name"] = ctx.session_name @@ -731,7 +747,7 @@ async def _handoff_impl( # Allow the full step time plus the server-side ready-wait (up to 120s) # plus headroom; the server enforces the per-step timeout internally. - client_timeout = float(timeout) + 180.0 + client_timeout = float(timeout) + _CLIENT_TIMEOUT_HEADROOM try: response = requests.post( f"{API_BASE_URL}/terminals/run-step", @@ -739,9 +755,22 @@ async def _handoff_impl( timeout=client_timeout, ) except requests.Timeout: + # The MCP transport timed out, but the step may still be running + # (or may have already completed) server-side. The server persists + # the result in handoff_results under job_id before sending the + # HTTP response, so the caller can retrieve it via the + # get_handoff_result MCP tool (issue #447 / PR #453 review finding + # 3 — a raw "GET /handoff-results/{job_id}" instruction gives the + # supervisor LLM no callable path: no base URL, no auth). return HandoffResult( success=False, - message=f"Handoff timed out after {timeout} seconds", + pending=True, + job_id=job_id, + message=( + f"Handoff transport timed out after {timeout + _CLIENT_TIMEOUT_HEADROOM} seconds. " + f"The job may still be running server-side. " + f"Retrieve the result with the get_handoff_result tool, job_id={job_id}" + ), output=None, terminal_id=None, ) @@ -1301,6 +1330,50 @@ def delete_terminal( return {"success": False, "message": f"Failed to delete terminal: {str(e)}"} +@mcp.tool() +def get_handoff_result( + job_id: str = Field( + description=( + "The job_id returned by handoff when pending=True (transport timed " + "out but the job may still be running or already finished server-side)." + ) + ), +) -> Dict[str, Any]: + """Retrieve a durably persisted handoff result by job_id (issue #447). + + Call this when a prior ``handoff`` call returned ``pending=True`` — the + transport timed out before the result arrived, but the work continues + server-side under ``job_id``. Poll this tool until ``state`` is no longer + ``"running"``. + + Args: + job_id: The job_id from the pending handoff result. + + Returns: + Dict with ``success``, ``state`` ("running"|"completed"|"error"), + ``terminal_id``, ``last_message`` (populated when completed), and + ``error_message`` (populated when errored). ``success=False`` with a + ``message`` when the job_id is unknown or the request failed. + """ + try: + response = requests.get(f"{API_BASE_URL}/handoff-results/{job_id}", timeout=_mcp_timeout()) + response.raise_for_status() + data = response.json() + return { + "success": True, + "state": data.get("state"), + "terminal_id": data.get("terminal_id"), + "last_message": data.get("last_message"), + "error_message": data.get("error_message"), + } + except requests.HTTPError as e: + if e.response is not None and e.response.status_code == 404: + return {"success": False, "message": f"No handoff result found for job_id {job_id}"} + return {"success": False, "message": f"Failed to retrieve handoff result: {str(e)}"} + except Exception as e: + return {"success": False, "message": f"Failed to retrieve handoff result: {str(e)}"} + + # ============================================================================= # Memory Tools # ============================================================================= diff --git a/src/cli_agent_orchestrator/services/agent_step.py b/src/cli_agent_orchestrator/services/agent_step.py index 48ef85323..7ae440340 100644 --- a/src/cli_agent_orchestrator/services/agent_step.py +++ b/src/cli_agent_orchestrator/services/agent_step.py @@ -26,6 +26,7 @@ import time from typing import Callable, Optional +from cli_agent_orchestrator.clients.database import upsert_handoff_result from cli_agent_orchestrator.models.terminal import AgentStepResult, TerminalStatus from cli_agent_orchestrator.plugins import PluginRegistry from cli_agent_orchestrator.services import terminal_service @@ -213,6 +214,7 @@ async def run_agent_step( env_vars: Optional[dict[str, str]] = None, on_terminal_created: Optional[Callable[[str], None]] = None, cancel_event: Optional[asyncio.Event] = None, + job_id: Optional[str] = None, ) -> AgentStepResult: """Run one agent step and return its result (success only). @@ -222,7 +224,14 @@ async def run_agent_step( 3. Send ``prompt`` (sync, bracketed-paste — the existing input path). 4. Wait until COMPLETED (in-process status poll). 5. Extract the last agent message (provider-specific extraction). - 6. Tear the terminal down unless ``teardown=False`` or it was reused. + 6. Persist the result durably (issue #447, if ``job_id`` given). + 7. Tear the terminal down unless ``teardown=False`` or it was reused. + + Step 6 runs BEFORE step 7 specifically so a crash during teardown cannot + lose an already-extracted result (PR #453 review: the prior placement in + the HTTP handler persisted AFTER this function had already torn the + terminal down, contradicting issue #447's "persist result, then tear + down" requirement). Args: provider: Provider type string (e.g. "kiro_cli", "claude_code"). @@ -288,6 +297,15 @@ async def run_agent_step( provider never emits a completion signal is exactly the run that could not otherwise be killed. Default None = no cancellation seam (the handoff caller passes nothing) — behavior unchanged. + job_id: Optional durable-result key (issue #447). When given, the + extracted result is persisted via ``upsert_handoff_result`` BETWEEN + extraction and teardown — before the terminal that carries the only + other copy of the result is destroyed. The persistence write is + best-effort (a DB failure is logged, never raised) so it cannot turn + a successful step into a reported failure. Default None = behavior + unchanged (no persistence). Failure-path persistence (state="error") + remains the HTTP handler's responsibility, since only the handler + can distinguish which exception type occurred. Returns: ``AgentStepResult`` with status COMPLETED — ONLY on success. @@ -428,6 +446,24 @@ async def run_agent_step( status=TerminalStatus.COMPLETED, ) + # Persist BEFORE teardown (issue #447): the terminal about to be destroyed + # is the only other place this result lives, so a crash during teardown + # must not be able to lose it. Off the loop (sqlite I/O); best-effort — a + # write failure must not turn a successful step into a reported failure. + if job_id: + try: + await asyncio.to_thread( + upsert_handoff_result, + job_id, + "completed", + terminal_id=terminal_id, + last_message=last_message, + ) + except Exception: # noqa: BLE001 — persistence is best-effort; step already succeeded + logger.warning( + "run_agent_step: failed to persist completed result for job_id=%s", job_id + ) + if teardown and created_here: await _best_effort_teardown(terminal_id, registry) diff --git a/src/cli_agent_orchestrator/services/cleanup_service.py b/src/cli_agent_orchestrator/services/cleanup_service.py index 97ded72ff..1382cba6e 100644 --- a/src/cli_agent_orchestrator/services/cleanup_service.py +++ b/src/cli_agent_orchestrator/services/cleanup_service.py @@ -4,7 +4,12 @@ from datetime import datetime, timedelta, timezone from pathlib import Path -from cli_agent_orchestrator.clients.database import InboxModel, SessionLocal, TerminalModel +from cli_agent_orchestrator.clients.database import ( + InboxModel, + SessionLocal, + TerminalModel, + delete_old_handoff_results, +) from cli_agent_orchestrator.constants import ( LOG_DIR, MEMORY_BASE_DIR, @@ -67,6 +72,14 @@ def cleanup_old_data(): server_logs_deleted += 1 logger.info(f"Deleted {server_logs_deleted} old server log files") + # Clean up old handoff result records (issue #447). + # Same cutoff as terminals/messages: RETENTION_DAYS. + try: + deleted_handoff = delete_old_handoff_results(cutoff_date) + logger.info(f"Deleted {deleted_handoff} old handoff result records") + except Exception as e: + logger.warning(f"Failed to clean up old handoff results: {e}") + logger.info("Cleanup completed successfully") except Exception as e: diff --git a/test/api/test_handoff_durability.py b/test/api/test_handoff_durability.py new file mode 100644 index 000000000..69160cf31 --- /dev/null +++ b/test/api/test_handoff_durability.py @@ -0,0 +1,227 @@ +"""Tests for durable handoff result persistence (issue #447). + +Verifies the run-step handler writes to handoff_results before responding, +and the GET /handoff-results/{job_id} retrieval endpoint works correctly. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from cli_agent_orchestrator.constants import TERMINALS_RUN_STEP_ROUTE +from cli_agent_orchestrator.models.terminal import AgentStepResult, TerminalStatus +from cli_agent_orchestrator.services.agent_step import StepExecutionError + +_RUN_STEP = "cli_agent_orchestrator.api.main.run_agent_step" +_UPSERT = "cli_agent_orchestrator.api.main.upsert_handoff_result" +_GET = "cli_agent_orchestrator.api.main.get_handoff_result" + + +def _body(**overrides): + base = {"provider": "kiro_cli", "agent": "developer", "prompt": "do it"} + base.update(overrides) + return base + + +class TestRunStepDurability: + def test_success_upserts_running_and_passes_job_id_through(self, client): + """Happy path: handler writes 'running' at request start and forwards + job_id into run_agent_step, which persists 'completed' itself BEFORE + teardown (issue #447 / PR #453 review finding 2) — not here after + run_agent_step has already returned (and torn the terminal down).""" + result = AgentStepResult( + terminal_id="abc12345", + last_message="all done", + status=TerminalStatus.COMPLETED, + ) + calls = [] + with ( + patch(_RUN_STEP, new=AsyncMock(return_value=result)) as m_run_step, + patch(_UPSERT, side_effect=lambda *a, **kw: calls.append((a, kw))), + ): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="cafe1234" * 4)) + + assert resp.status_code == 200 + # Only "running" is written here; "completed" is run_agent_step's job. + assert len(calls) == 1 + assert calls[0][0][1] == "running" + assert m_run_step.await_args.kwargs["job_id"] == "cafe1234" * 4 + + def test_no_job_id_skips_upsert(self, client): + """Without job_id, nothing is persisted (backward-compat).""" + result = AgentStepResult( + terminal_id="abc12345", + last_message="all done", + status=TerminalStatus.COMPLETED, + ) + with ( + patch(_RUN_STEP, new=AsyncMock(return_value=result)), + patch(_UPSERT) as m_upsert, + ): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body()) + + assert resp.status_code == 200 + m_upsert.assert_not_called() + + def test_step_execution_error_persists_error_state(self, client): + """A StepExecutionError is persisted as state=error before the 504.""" + calls = [] + with ( + patch( + _RUN_STEP, + new=AsyncMock( + side_effect=StepExecutionError( + "timed out", kind="timeout", terminal_id="abc12345" + ) + ), + ), + patch(_UPSERT, side_effect=lambda *a, **kw: calls.append((a, kw))), + ): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="aabbccdd" * 4)) + + assert resp.status_code == 504 + # running + error + assert len(calls) == 2 + assert calls[1][0][1] == "error" + assert "timed out" in calls[1][1]["error_message"] + + def test_upsert_failure_does_not_break_successful_step(self, client): + """A DB write failure is logged but must not turn a successful step into + a failure — the work is done and the response must still be 200.""" + result = AgentStepResult( + terminal_id="abc12345", + last_message="ok", + status=TerminalStatus.COMPLETED, + ) + with ( + patch(_RUN_STEP, new=AsyncMock(return_value=result)), + patch(_UPSERT, side_effect=Exception("db boom")), + ): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="deadbeef" * 4)) + + assert resp.status_code == 200 + assert resp.json()["last_message"] == "ok" + + +class TestRunStepDurabilityErrorBranches: + """S-001: ValueError and generic Exception branches must also persist error state.""" + + def test_value_error_persists_error_state(self, client): + """ValueError (e.g. unknown terminal) must transition job to 'error', not leave it + stuck at 'running'.""" + calls = [] + with ( + patch(_RUN_STEP, new=AsyncMock(side_effect=ValueError("Terminal 'x' not found"))), + patch(_UPSERT, side_effect=lambda *a, **kw: calls.append((a, kw))), + ): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="11223344" * 4)) + + assert resp.status_code == 404 + # running then error — not stuck at running. + assert len(calls) == 2 + assert calls[0][0][1] == "running" + assert calls[1][0][1] == "error" + assert "Terminal 'x'" in calls[1][1]["error_message"] + + def test_generic_exception_persists_error_state(self, client): + """An unanticipated Exception must also transition job to 'error'.""" + calls = [] + with ( + patch(_RUN_STEP, new=AsyncMock(side_effect=RuntimeError("unexpected boom"))), + patch(_UPSERT, side_effect=lambda *a, **kw: calls.append((a, kw))), + ): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="55667788" * 4)) + + assert resp.status_code == 500 + assert len(calls) == 2 + assert calls[1][0][1] == "error" + assert "unexpected boom" in calls[1][1]["error_message"] + + +class TestJobIdValidation: + """S-002: job_id field must reject non-hex and wrong-length values.""" + + def test_valid_32_char_hex_accepted(self, client): + result = AgentStepResult( + terminal_id="abc12345", last_message="ok", status=TerminalStatus.COMPLETED + ) + with ( + patch(_RUN_STEP, new=AsyncMock(return_value=result)), + patch(_UPSERT), + ): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="a" * 32)) + assert resp.status_code == 200 + + def test_empty_string_rejected(self, client): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="")) + assert resp.status_code == 422 + + def test_31_char_hex_rejected(self, client): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="a" * 31)) + assert resp.status_code == 422 + + def test_33_char_hex_rejected(self, client): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="a" * 33)) + assert resp.status_code == 422 + + def test_uppercase_hex_rejected(self, client): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="A" * 32)) + assert resp.status_code == 422 + + def test_non_hex_char_rejected(self, client): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body(job_id="g" * 32)) + assert resp.status_code == 422 + + def test_none_omitted_is_fine(self, client): + """Omitting job_id entirely (None) should still be accepted.""" + result = AgentStepResult( + terminal_id="abc12345", last_message="ok", status=TerminalStatus.COMPLETED + ) + with ( + patch(_RUN_STEP, new=AsyncMock(return_value=result)), + patch(_UPSERT), + ): + resp = client.post(TERMINALS_RUN_STEP_ROUTE, json=_body()) + assert resp.status_code == 200 + + +class TestGetHandoffResult: + def test_returns_record_when_found(self, client): + record = { + "job_id": "cafe1234" * 4, + "state": "completed", + "terminal_id": "abc12345", + "last_message": "done", + "error_message": None, + "created_at": "2025-01-01T00:00:00+00:00", + "updated_at": "2025-01-01T00:01:00+00:00", + } + with patch(_GET, return_value=record): + resp = client.get("/handoff-results/cafe1234cafe1234cafe1234cafe1234") + + assert resp.status_code == 200 + data = resp.json() + assert data["state"] == "completed" + assert data["last_message"] == "done" + + def test_returns_404_when_not_found(self, client): + with patch(_GET, return_value=None): + resp = client.get("/handoff-results/unknown-job-id") + + assert resp.status_code == 404 + + def test_returns_running_state(self, client): + record = { + "job_id": "aabbccdd" * 4, + "state": "running", + "terminal_id": None, + "last_message": None, + "error_message": None, + "created_at": "2025-01-01T00:00:00+00:00", + "updated_at": "2025-01-01T00:00:30+00:00", + } + with patch(_GET, return_value=record): + resp = client.get("/handoff-results/aabbccddaabbccddaabbccddaabbccdd") + + assert resp.status_code == 200 + assert resp.json()["state"] == "running" diff --git a/test/clients/test_handoff_result_db.py b/test/clients/test_handoff_result_db.py new file mode 100644 index 000000000..2bcc16db1 --- /dev/null +++ b/test/clients/test_handoff_result_db.py @@ -0,0 +1,111 @@ +"""Unit tests for the HandoffResultModel CRUD helpers (issue #447). + +Uses an in-memory SQLite database so no file system state is required. +""" + +from datetime import datetime, timedelta, timezone + +import pytest +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + +from cli_agent_orchestrator.clients.database import ( + Base, + HandoffResultModel, + delete_old_handoff_results, + get_handoff_result, + upsert_handoff_result, +) + + +@pytest.fixture(autouse=True) +def _use_test_db(monkeypatch): + """Redirect all DB calls to a fresh in-memory SQLite DB.""" + engine = create_engine("sqlite:///:memory:") + Base.metadata.create_all(bind=engine) + TestSession = sessionmaker(bind=engine, autocommit=False, autoflush=False) + monkeypatch.setattr("cli_agent_orchestrator.clients.database.SessionLocal", TestSession) + yield + + +def _utcnow(): + return datetime.now(timezone.utc) + + +class TestUpsertHandoffResult: + def test_creates_new_record_on_first_call(self): + upsert_handoff_result("job-1", "running") + record = get_handoff_result("job-1") + assert record is not None + assert record["state"] == "running" + assert record["last_message"] is None + + def test_updates_existing_record_on_second_call(self): + upsert_handoff_result("job-2", "running") + upsert_handoff_result("job-2", "completed", last_message="ok", terminal_id="abc12345") + record = get_handoff_result("job-2") + assert record["state"] == "completed" + assert record["last_message"] == "ok" + assert record["terminal_id"] == "abc12345" + + def test_error_state_stored_with_message(self): + upsert_handoff_result("job-3", "running") + upsert_handoff_result("job-3", "error", error_message="worker crashed") + record = get_handoff_result("job-3") + assert record["state"] == "error" + assert record["error_message"] == "worker crashed" + + def test_partial_update_does_not_overwrite_nones(self): + upsert_handoff_result("job-4", "completed", last_message="original", terminal_id="t1") + # Calling with no last_message must not clear the existing value. + upsert_handoff_result("job-4", "completed") + record = get_handoff_result("job-4") + assert record["last_message"] == "original" + assert record["terminal_id"] == "t1" + + +class TestGetHandoffResult: + def test_returns_none_for_unknown_job(self): + assert get_handoff_result("no-such-job") is None + + def test_returns_dict_with_expected_keys(self): + upsert_handoff_result("job-5", "running") + record = get_handoff_result("job-5") + assert set(record.keys()) == { + "job_id", + "state", + "terminal_id", + "last_message", + "error_message", + "created_at", + "updated_at", + } + + +class TestDeleteOldHandoffResults: + def test_deletes_records_older_than_cutoff(self): + upsert_handoff_result("old-1", "completed") + upsert_handoff_result("old-2", "error") + upsert_handoff_result("new-1", "running") + # Backdate old-1 and old-2 by reaching into the DB directly. + from cli_agent_orchestrator.clients.database import HandoffResultModel, SessionLocal + + past = _utcnow() - timedelta(days=20) + with SessionLocal() as db: + for jid in ("old-1", "old-2"): + row = db.query(HandoffResultModel).filter(HandoffResultModel.job_id == jid).first() + row.created_at = past + db.commit() + + cutoff = _utcnow() - timedelta(days=10) + deleted = delete_old_handoff_results(cutoff) + assert deleted == 2 + assert get_handoff_result("old-1") is None + assert get_handoff_result("old-2") is None + assert get_handoff_result("new-1") is not None + + def test_returns_zero_when_nothing_to_delete(self): + upsert_handoff_result("recent", "completed") + cutoff = _utcnow() - timedelta(days=30) + deleted = delete_old_handoff_results(cutoff) + assert deleted == 0 diff --git a/test/mcp_server/test_handoff_durability.py b/test/mcp_server/test_handoff_durability.py new file mode 100644 index 000000000..e3e31004a --- /dev/null +++ b/test/mcp_server/test_handoff_durability.py @@ -0,0 +1,167 @@ +"""Tests for durable handoff result retrieval by the MCP client (issue #447). + +Verifies: +- _handoff_impl generates a job_id and passes it in the run-step payload. +- On requests.Timeout, the HandoffResult carries pending=True and job_id + so the caller can retrieve the result later. +- The normal synchronous path still works and job_id/pending are absent. +""" + +import asyncio +import uuid +from unittest.mock import MagicMock, patch + +import requests + +from cli_agent_orchestrator.mcp_server.server import ( + HandoffContext, + _handoff_impl, + get_handoff_result, +) + + +def _ctx(provider="kiro_cli", session_name=None, caller_id=None, allowed_tools=None): + return HandoffContext( + provider=provider, + session_name=session_name, + caller_id=caller_id, + allowed_tools=allowed_tools, + ) + + +def _ok_response(terminal_id="dev-t1", last_message="done"): + resp = MagicMock() + resp.status_code = 200 + resp.json.return_value = { + "terminal_id": terminal_id, + "last_message": last_message, + "status": "completed", + } + return resp + + +class TestHandoffJobId: + @patch("cli_agent_orchestrator.mcp_server.server._get_cleanup_nudge", return_value="") + @patch("cli_agent_orchestrator.mcp_server.server._resolve_handoff_provider") + def test_job_id_included_in_payload(self, mock_provider, _nudge): + """Every handoff POST must include a job_id.""" + mock_provider.return_value = _ctx() + with patch("cli_agent_orchestrator.mcp_server.server.requests") as mock_requests: + mock_requests.post.return_value = _ok_response() + mock_requests.Timeout = Exception + asyncio.run(_handoff_impl("developer", "do task")) + + payload = mock_requests.post.call_args[1]["json"] + assert "job_id" in payload + # Must be a 32-char hex string (uuid4().hex format). + jid = payload["job_id"] + assert isinstance(jid, str) and len(jid) == 32 + int(jid, 16) # raises if not valid hex + + @patch("cli_agent_orchestrator.mcp_server.server._resolve_handoff_provider") + def test_transport_timeout_returns_pending_result_with_job_id(self, mock_provider): + """On requests.Timeout the HandoffResult must carry pending=True and + a non-None job_id so the caller can poll the retrieval endpoint.""" + mock_provider.return_value = _ctx() + + class FakeTimeout(Exception): + pass + + with patch("cli_agent_orchestrator.mcp_server.server.requests") as mock_requests: + mock_requests.post.side_effect = FakeTimeout("timed out") + mock_requests.Timeout = FakeTimeout + result = asyncio.run(_handoff_impl("developer", "do task", timeout=600)) + + assert result.success is False + assert result.pending is True + assert result.job_id is not None + assert len(result.job_id) == 32 + # Message must explain how to retrieve, via the get_handoff_result MCP + # tool (PR #453 review finding 3) — a raw HTTP path gives the + # supervisor LLM no callable path (no base URL, no auth). + assert "get_handoff_result" in result.message + assert result.job_id in result.message + + @patch("cli_agent_orchestrator.mcp_server.server._get_cleanup_nudge", return_value="") + @patch("cli_agent_orchestrator.mcp_server.server._resolve_handoff_provider") + def test_success_path_has_no_pending(self, mock_provider, _nudge): + """Normal synchronous success must not set pending=True on the result.""" + mock_provider.return_value = _ctx() + with patch("cli_agent_orchestrator.mcp_server.server.requests") as mock_requests: + mock_requests.post.return_value = _ok_response() + mock_requests.Timeout = Exception + result = asyncio.run(_handoff_impl("developer", "do task")) + + assert result.success is True + # pending is None (not set) on the normal path — not True. + assert result.pending is not True + + @patch("cli_agent_orchestrator.mcp_server.server._resolve_handoff_provider") + def test_each_call_generates_unique_job_id(self, mock_provider): + """Separate calls must generate distinct job_ids (no collision).""" + mock_provider.return_value = _ctx() + ids_seen = set() + + class FakeTimeout(Exception): + pass + + for _ in range(5): + with patch("cli_agent_orchestrator.mcp_server.server.requests") as mock_requests: + mock_requests.post.side_effect = FakeTimeout("timed out") + mock_requests.Timeout = FakeTimeout + result = asyncio.run(_handoff_impl("developer", "do task")) + ids_seen.add(result.job_id) + + assert len(ids_seen) == 5 + + +class TestGetHandoffResultTool: + """The MCP tool half of the polling contract (PR #453 review finding 3): + a pending handoff's job_id must be retrievable through a callable tool, + not just a bare HTTP path the supervisor LLM cannot reach.""" + + def test_completed_result_returned(self): + with patch("cli_agent_orchestrator.mcp_server.server.requests.get") as mock_get: + mock_get.return_value.raise_for_status.return_value = None + mock_get.return_value.json.return_value = { + "state": "completed", + "terminal_id": "dev-t1", + "last_message": "done", + "error_message": None, + } + result = get_handoff_result("cafe1234" * 4) + assert result["success"] is True + assert result["state"] == "completed" + assert result["last_message"] == "done" + + def test_running_result_returned(self): + with patch("cli_agent_orchestrator.mcp_server.server.requests.get") as mock_get: + mock_get.return_value.raise_for_status.return_value = None + mock_get.return_value.json.return_value = { + "state": "running", + "terminal_id": None, + "last_message": None, + "error_message": None, + } + result = get_handoff_result("cafe1234" * 4) + assert result["success"] is True + assert result["state"] == "running" + + def test_unknown_job_id_returns_not_found(self): + with patch("cli_agent_orchestrator.mcp_server.server.requests.get") as mock_get: + http_err = requests.HTTPError() + http_err.response = MagicMock() + http_err.response.status_code = 404 + mock_get.return_value.raise_for_status.side_effect = http_err + result = get_handoff_result("deadbeef" * 4) + assert result["success"] is False + assert "No handoff result found" in result["message"] + + def test_generic_failure_returns_false(self): + with patch( + "cli_agent_orchestrator.mcp_server.server.requests.get", + side_effect=Exception("connection refused"), + ): + result = get_handoff_result("deadbeef" * 4) + assert result["success"] is False + assert "Failed" in result["message"] diff --git a/test/services/test_agent_step.py b/test/services/test_agent_step.py index aeb6336f7..a8ee548d1 100644 --- a/test/services/test_agent_step.py +++ b/test/services/test_agent_step.py @@ -336,6 +336,83 @@ def test_create_failure_propagates(self): asyncio.run(run_agent_step("kiro_cli", "dev", "x")) +class TestDurablePersistenceBeforeTeardown: + """Issue #447 / PR #453 review finding 2: the durable result must be + written BEFORE teardown, not after run_agent_step returns — a crash + during teardown must not be able to lose an already-extracted result.""" + + def test_job_id_persists_completed_before_delete(self): + """upsert_handoff_result must be called, and its call must be + observed to happen strictly before delete_terminal — proving the + ordering, not just that both happened.""" + create, send, delete, get_output, exit_cli, get_wd, wait, status = _patch_terminal_layer() + order = [] + upsert = patch( + f"{_MODULE}.upsert_handoff_result", + side_effect=lambda *a, **kw: order.append(("upsert", a, kw)), + ) + with ( + create, + send, + delete as m_delete, + get_output, + exit_cli, + get_wd, + wait, + status, + upsert, + ): + m_delete.side_effect = lambda *a, **kw: order.append(("delete", a, kw)) + result = asyncio.run(run_agent_step("kiro_cli", "dev", "x", job_id="cafe1234" * 4)) + + assert result.status == TerminalStatus.COMPLETED + assert [step for step, _, _ in order] == ["upsert", "delete"] + upsert_call = order[0] + assert upsert_call[1][0] == "cafe1234" * 4 + assert upsert_call[1][1] == "completed" + assert upsert_call[2]["terminal_id"] == "abc12345" + assert upsert_call[2]["last_message"] == "the answer" + + def test_no_job_id_skips_persistence(self): + """Without job_id (e.g. the run engine's in-process caller), no + persistence call is made — behavior unchanged from before #447.""" + create, send, delete, get_output, exit_cli, get_wd, wait, status = _patch_terminal_layer() + with ( + create, + send, + delete, + get_output, + exit_cli, + get_wd, + wait, + status, + patch(f"{_MODULE}.upsert_handoff_result") as m_upsert, + ): + asyncio.run(run_agent_step("kiro_cli", "dev", "x")) + m_upsert.assert_not_called() + + def test_persistence_failure_does_not_fail_successful_step(self): + """A DB write failure during the pre-teardown persist must be logged, + not raised — the step already succeeded and must still return.""" + create, send, delete, get_output, exit_cli, get_wd, wait, status = _patch_terminal_layer() + with ( + create, + send, + delete as m_delete, + get_output, + exit_cli, + get_wd, + wait, + status, + patch(f"{_MODULE}.upsert_handoff_result", side_effect=Exception("db boom")), + ): + result = asyncio.run(run_agent_step("kiro_cli", "dev", "x", job_id="a" * 32)) + assert result.status == TerminalStatus.COMPLETED + assert result.last_message == "the answer" + # Teardown must still proceed after a persistence failure. + m_delete.assert_called_once() + + class TestTeardownIsBestEffort: def test_teardown_failure_does_not_fail_successful_step(self): """A delete failure after a successful step is logged, not raised — the