diff --git a/docker/base-image/agent_server/services/pull_worker.py b/docker/base-image/agent_server/services/pull_worker.py index 337b9d6c28..8f69d28373 100644 --- a/docker/base-image/agent_server/services/pull_worker.py +++ b/docker/base-image/agent_server/services/pull_worker.py @@ -389,6 +389,7 @@ async def _run_and_report( ``persist_session``), matching the SESSION_TAB ``--resume`` convention — the minimal reversible reading of the schema's "presence ⇒ persist/resume".""" from ..services.runtime_adapter import get_runtime + from .process_registry import get_process_registry execution_id = claim.get("execution_id") claim_token = claim.get("claim_token") @@ -398,9 +399,25 @@ async def _run_and_report( session_id = payload.get("session_id") overrides = payload.get("task_overrides") or {} turn_timeout = _resolve_turn_timeout(overrides) - + # #3114: a turn queued by the backend's `execute_task` carries the push + # payload's `persist_session` and `images`. Older rows carry neither, so + # persistence falls back to "a session id means resume and persist". + persist_session = payload.get("persist_session") + if persist_session is None: + persist_session = bool(session_id) + + # #3114: a claimed turn is ACCEPTED-not-spawned until the runtime starts it. + # Registered as pending, `/api/executions/{id}/terminate` reaches it in that + # window (#2433): the cancel is recorded, the turn is skipped (or killed at + # spawn) and reported `cancelled` below. + registry = get_process_registry() + registry.register_pending( + execution_id, timeout_seconds=turn_timeout, metadata={"type": "pull"}, + ) agent_state.record_task_start() try: + if execution_id and registry.was_terminated(execution_id): + raise HTTPException(status_code=409, detail="Execution cancelled before it started") response_text, raw_messages, metadata, ran_session_id = await get_runtime().execute_headless( prompt=message, model=overrides.get("model"), @@ -410,7 +427,8 @@ async def _run_and_report( max_turns=overrides.get("max_turns"), execution_id=execution_id, resume_session_id=session_id, - persist_session=bool(session_id), + persist_session=bool(persist_session), + images=payload.get("images") or None, ) body = _success_result_body(claim_token, response_text, raw_messages, metadata, ran_session_id) finish_success: Optional[bool] = True @@ -424,6 +442,25 @@ async def _run_and_report( except Exception as exc: # noqa: BLE001 — any failure must still report a terminal finish_success = False body = _failed_result_body(claim_token, str(exc) or type(exc).__name__) + finally: + # #2433: promoted at spawn; dropped here if it never spawned. + registry.discard_pending(execution_id) + + # #679 / #3114: a terminated turn reports `cancelled` (result_callback + # parity). Auth and rate-limit terminals keep their class so the breaker + # and SUB-003 still see them. + if ( + execution_id + and registry.was_terminated(execution_id) + and body.get("error_code") not in ("auth", "billing") + ): + body = { + **body, + "status": "cancelled", + "error_code": None, + "error": "Execution terminated by user", + } + finish_success = None agent_state.record_task_finish(success=finish_success) diff --git a/docs/memory/architecture/execution.md b/docs/memory/architecture/execution.md index 6bf5eb10f7..23f1f2cffb 100644 --- a/docs/memory/architecture/execution.md +++ b/docs/memory/architecture/execution.md @@ -154,8 +154,9 @@ The last blocking dispatch callers, moved onto execution terminals so they work - **Fan-out.** `FanOutService` rebuilds the sync aggregate from the batch's `schedule_executions` rows (`fan_out_id` + `fan_out_task_id`), not from `execute_task` return values. Each row is created **at slot grant** inside the `max_concurrency` semaphore — never up front: a pre-created `RUNNING` row waiting behind the semaphore matches the #106 no-session sweep, the watchdog orphan reconcile and the stale sweep (all admission-anchored, #2433/#2435), and a `QUEUED` one is claimable by `claim_next_queued` while the service also dispatches it. The sync caller waits for the (shielded) dispatch, then `sync_waiter.wait_for_fan_out_batch`; with no caller deadline the budget is `ceil(N / min(max_concurrency, max_parallel_tasks)) × execution_timeout + 120s`. `async_mode` returns a receipt; the not-yet-dispatched tail is in-process and does not survive a restart. - **The join** (`fan_out_service.join_fan_out_on_terminal`) and the #2523 loop advance hang off `event_dispatch_service.spawn_task_terminal_event` via `_terminal_side_effects`, each under its own guard — outside the subscriber-gated `emit_task_terminal_event`. One PK read per fleet terminal; a `(fan_out_id, status)`-indexed COUNT only for fan-out rows. -- **Sync edge adapter** (`task_execution_service.dispatch_and_await_terminal` + `result_from_execution_row`): `execute_task`, and on a `QUEUED` return wait for that row's terminal via `sync_waiter.wait_for_sync_terminal` and rebuild the `TaskExecutionResult` from the row. Used by `routers/a2a._run_a2a_task` and `operator_resume_service`. A timed-out wait returns `FAILED/TIMEOUT` without writing a terminal. The pull sink does not signal that registry, so the wake is the 5s DB poll (≤5s added tail latency on a pilot agent). +- **Sync edge adapter** (`task_execution_service.dispatch_and_await_terminal` + `result_from_execution_row`): `execute_task`, and on a `QUEUED` return wait for that row's terminal via `sync_waiter.wait_for_sync_terminal` and rebuild the `TaskExecutionResult` from the row. Used by `routers/a2a._run_a2a_task`, `operator_resume_service` and, since #3114, every sync interactive caller (below). A timed-out wait returns `FAILED/TIMEOUT` ("Execution timed out after …") without writing a terminal. The pull sink signals the in-process waiter when its CAS write wins; a caller on another uvicorn worker wakes on the 5s DB poll. - `PULL_REACHABLE_TRIGGERS` now equals `_AUTONOMOUS_TRIGGERS` but stays an enumerated allow-list (pinned by `test_2048_pull_pilot_reach.py`). +- **Interactive turns on a pilot (#3114).** `pull_owns_dispatch` also returns True for `pull_pilot.PULL_REACHABLE_NON_AUTONOMOUS` = `INTERACTIVE_TRIGGERS - {"chat"} | {"validation"}`; the UI `/chat` path (`dispatch_admission_service` → `run_chat_turn`, trigger `chat`) still pushes. Sync callers — `session_turn_service.run_resumable_turn` (Session tab + Workspace), `public_chat_service`, `adapters/message_router`, `shared_sessions` rooms, `routers/paid`, `mcp_auth_service`, `validation_service`, internal sync run-now — call `dispatch_and_await_terminal` with a `conversation_key` (`session:`/`public:`/`channel:`/`room:`/`paid:`). For interactive and validation triggers the adapter first waits for a claim, bounded by `claim_budget` (default one agent execution timeout; rooms pass 300s); on expiry it stores the row FAILED (same queued-only CAS as `cancel_queued_execution`) and returns FAILED/`CAPACITY`, a caller that went away leaves it CANCELLED, and a lost CAS (just claimed) falls through to the terminal wait. The pull sink spawns `channel_completion_report` on its CAS-won branch like the push terminals; the agent's pull worker registers a claimed turn as pending so a terminate before spawn reports `cancelled`. Room wakes of one agent serialise on a per-(room, agent) `ResumeLock` (`room_wake_lock:`), re-entrant within one mention chain. Images ride `backlog_metadata`. `result_from_execution_row` rebuilds a push-shaped result (`[code]` error prefix → `error_code`, compact events in `raw_response`, dispatch sentinels never a session id). The Session `ResumeLock` and the Workspace in-flight markers stay and span the queued wait: their TTLs add `pull_pilot.pull_queue_allowance` (one agent timeout on a pilot, 0 elsewhere). The live-stream proxies hold while the row is queued (`sync_waiter.wait_while_queued`). Pinned by `test_3114_pull_route_interactive.py` and `test_3114_pull_route_workspace.py`. - **Pull claim order (#2842, #2843).** `claim_next_task` passes `pull_pilot.INTERACTIVE_TRIGGERS` into `claim_next_queued`: interactive rows first, oldest first within each group, no anti-starvation rule. The pull claim also skips a row whose `conversation_key` has a `running` row for the agent; the partial unique index `idx_executions_one_running_turn` (agent, conversation_key WHERE running) turns a concurrent second claim into an `IntegrityError` that the claim retries. The guard is the row's status, so terminal, lease expiry and requeue all release it. `backlog_service.enqueue` stamps the key (`chat_session_id`, else `resume_session_id`) on pilot agents only. The push drain calls `claim_next_queued(agent_name)` with neither, so its order is plain oldest-first. Pinned by `test_2842_2843_pull_claim_order.py`. ### Agent Self-Reminders (#1296) @@ -177,7 +178,7 @@ Durable one-shot deferred self-trigger — the time-deferred sibling of loops ( | `routers/sessions.py` | `agent_sessions` / `agent_session_messages` | platform user | | `client_portal/service.py` (Workspace) | `enterprise_portal_sessions` / `_messages` | verified client email | -`cached_claude_session_id` is the load-bearing field on both tables. The engine: runtime gate (drop the cached id for a runtime without `--resume`, i.e. Codex) → resume lock → `execute_task(persist_session=True, resume_session_id=…)` → **one** cold retry when Claude reports the JSONL missing. Callers pass `on_resume_failure` (clear the cache, count the failure, inside the lock) and may pass a distinct `cold_message` for that retry. `ResumeLockBusy` subclasses `HTTPException(429)` so the Session router surfaces it unchanged while the Workspace catches the precise type and re-raises `ClientPortalError(429)`. +`cached_claude_session_id` is the load-bearing field on both tables. The engine: runtime gate (drop the cached id for a runtime without `--resume`, i.e. Codex) → resume lock → `dispatch_and_await_terminal(persist_session=True, resume_session_id=…, conversation_key="session:")` (plain `execute_task` off a pull pilot; on a pilot the lock is held across the queued wait, #3114) → **one** cold retry when Claude reports the JSONL missing. Callers pass `on_resume_failure` (clear the cache, count the failure, inside the lock) and may pass a distinct `cold_message` for that retry. `ResumeLockBusy` subclasses `HTTPException(429)` so the Session router surfaces it unchanged while the Workspace catches the precise type and re-raises `ClientPortalError(429)`. **Workspace absorbed the Session surface (ent#358).** Agent Detail no longer renders a Session surface: `SessionPanel.vue` and the Session-mode toggle are gone, the Chat tab is stateless-only (plus a "Continue in Workspace" link), and `?tab=session` **redirects** to `/workspace?agent=` (query-preserving, `router.replace`, guarded first in both `onMounted` and `onActivated` since AgentDetail is KeepAlive-cached). The removal was gated on continuity parity, not on streaming — the Session surface never streamed (synchronous POST + reattach poller, #1376/#759), so Workspace streaming (ent#286) was **not** a prerequisite. What *was* a prerequisite: Workspace chat ran stateless with a history prefix, so it moved onto this engine first. `agent_sessions` rows, endpoints and store stay readable; only the entry point went away. Since ent#456 the "Continue in Workspace" link and the NavBar entry open a NEW tab; the redirect stays same-tab. diff --git a/docs/memory/feature-flows/persistent-task-backlog.md b/docs/memory/feature-flows/persistent-task-backlog.md index 70990d4a96..c4dd4f5c6d 100644 --- a/docs/memory/feature-flows/persistent-task-backlog.md +++ b/docs/memory/feature-flows/persistent-task-backlog.md @@ -205,10 +205,25 @@ ON schedule_executions(agent_name, conversation_key) WHERE status = 'running' AND conversation_key IS NOT NULL; ``` -`backlog_service.enqueue` stamps `conversation_key` (`chat_session_id`, else -`resume_session_id`) on pilot agents only. The ordering expression cannot use +`backlog_service.enqueue` stamps `conversation_key` on pilot agents only: the +explicit key `execute_task` was given (`session:`, `public:`, +`channel:`, `room:`, `paid:`, #3114), else `chat_session_id`, else +`resume_session_id`. The ordering expression cannot use `idx_executions_queued`; per-agent queues are small. +**Interactive turns on a pilot (#3114).** `pull_pilot.pull_owns_dispatch` +covers every interactive trigger except `chat` (the UI `/chat` path, which +still pushes), plus `validation`. Their sync callers go through +`task_execution_service.dispatch_and_await_terminal`: phase 1 polls the row +while it is `queued`, for at most one agent execution timeout, and on expiry +cancels it (`cancel_queued_execution`) and returns FAILED/`CAPACITY`; a cancel +that loses to a claim falls through to phase 2, the terminal wait (agent +timeout + 120s from the claim). The pull sink signals the in-process waiter +when its CAS write wins. The live-stream proxies (`routers/chat.py`, +`routers/public.py`, `client_portal/router.py`) hold the SSE connection with +`: queued` comments while the row is queued (`sync_waiter.wait_while_queued`), +then retry an agent 404 briefly while the claimed row is `running`. + ### backlog_metadata JSON Shape Captured at enqueue time, replayed on drain. No credential values, only @@ -237,10 +252,20 @@ identity and request parameters: "triggered_by": "manual", "collaboration_activity_id": null, "is_self_task": false, - "self_task_activity_id": null + "self_task_activity_id": null, + "persist_session": false, + "schedule_context": null, + "attempt": null, + "images": null } ``` +`persist_session`, `schedule_context`, `attempt` and `images` (#3114) are read +by the pull claim only: the envelope hands `persist_session` and `images` +(base64 vision input) to the worker, and the claim-time prompt uses +`schedule_context` and `attempt`. Canary G-04 skips `images` when scanning for +credentials. + ## Lifecycle ``` diff --git a/docs/planning/PULL_MIGRATION_STATUS.md b/docs/planning/PULL_MIGRATION_STATUS.md index dbf73d1fab..b580666718 100644 --- a/docs/planning/PULL_MIGRATION_STATUS.md +++ b/docs/planning/PULL_MIGRATION_STATUS.md @@ -38,7 +38,8 @@ Dispatch topology, not policy. `pull_pilot.PULL_REACHABLE_TRIGGERS` is the sourc |---|---|---| | **On `dev` today** | `agent`, `event`, `schedule`, `webhook`, `reminder`, `loop` | 6 of 9 | | **Adds with Phase 4** | `fan_out`, `a2a`, `operator_response` | → 9 of 9 | -| **Pending** | interactive chat / Session-tab turns | claim order (#2842) and conversation guard (#2843) are in; producers not yet routed (decided 2026-09-16, #1989) | +| **Interactive (#3114)** | `manual`, `mcp`, `session`, `public`, `voice`, `voip`, `room`, `user`, `paid`, `slack`, `telegram`, `whatsapp`, plus `validation` (`pull_pilot.PULL_REACHABLE_NON_AUTONOMOUS`) | routed on pilots; sync callers wait through `dispatch_and_await_terminal` | +| **Pending** | `chat` — the UI `/chat` path (`dispatch_admission_service` → `chat_execution_service.run_chat_turn`) | still pushes on pilots; routed in a later change | | **Unreached, unclassified** | `retry` | in neither trigger set — #2845. 4.7% of `eu2` traffic; looks like an oversight, not a decision | `schedule` / `webhook` / `reminder` landed with #2391; `loop` with #2523. Before #2391 the pilot flag was @@ -62,10 +63,12 @@ The spec names the gates (`TARGET_ARCHITECTURE.md`, §Re-Delivery and Side-Effec and logged. Pilots must run a base image with #2392, or every effect call whose model omits the id is refused. *(Trace fidelity #548/#333 closed Aug/Jun; `prior_trace` injection #1401 closed 2026-07-08; #1402 closed 2026-07-26.)* -3. **A soak on an agent that actually emits.** The current pilot (`cornelius-oracle` on eu2) emits no - messages, calls or shares — measured 2026-09-02: 219 `idempotency_keys` rows, all `agent:*`, zero - `effect:*`. It has therefore never entered the code path item 2 protects, so a clean window on it is not - evidence about that gate. Mechanics for a second, disposable emitting pilot are in the ops repo +3. **A soak on an agent that actually emits.** eu2 runs three pilot arms: `cornelius-oracle` and + `oracle-6-science` emit no messages, calls or shares; `oracle-7-ai-semi` (flipped 2026-09-15) is the + emitting arm, with a 2-hourly schedule that writes a byte-stable file to its public folder and calls + `share_file`, so a re-delivered execution must emit once. What the gate still needs is the pulled-run rate on + that arm: effect rows per pulled execution, with no duplicate share on a re-delivery. All three arms share + one subscription seat, so a quota blackout darkens every arm at once. Mechanics are in the ops repo (`trinity-ops-agent:docs/pull-soak-eu2.md`). System of record for the soak is [#1766](https://github.com/abilityai/trinity/issues/1766)'s comment thread — read it before measuring. 4. **Phase 5: flip default-ON and delete the legacy machinery** — the 9-path cleanup pyramid, the slot ZSET, @@ -84,9 +87,30 @@ The spec names the gates (`TARGET_ARCHITECTURE.md`, §Re-Delivery and Side-Effec row whose `conversation_key` already has a `running` row is skipped, and the unique index `idx_executions_one_running_turn` stops two concurrent claimers. **No worker is reserved** for interactive arrivals: one of N held idle is 33% of a 3-worker agent. Multi-container affinity (`replica_count > 1`, - #927) is out of scope. The claim guard is the one mechanism for pulled turns: when session turns are routed - onto the queue, `session_turn_service`'s Redis `ResumeLock` is removed from that path, not kept beside it. Next: route the interactive producers onto the queue. Until then they still run on - the old path — a **migration state, not a design boundary**. + #927) is out of scope. + + **Routed on pilots (#3114)** for every interactive trigger except `chat` (the UI `/chat` path, still + pushed). Sync callers (Session tab, Workspace, public links, channels, rooms, paid, MCP key auth, + validation, internal run-now) go through `task_execution_service.dispatch_and_await_terminal`: it waits up + to one agent execution timeout for a worker to claim the row (rooms: `ROOM_CLAIM_BUDGET_SECONDS`, 300s, + via `claim_budget`), stores it FAILED/`CAPACITY` if none does (a caller that went away leaves it + CANCELLED), then waits for the terminal. Each producer passes a `conversation_key` (`session:`, `public:`, `channel:`, + `room:`, `paid:`), which is what the claim guard serialises on. Images ride the queue in + `backlog_metadata`, which the #1449 retention sweep NULLs once the row is terminal. Live-stream proxies hold the SSE connection + while the row is queued. + + **The Session-tab `ResumeLock` and the Workspace in-flight markers are retained** and span the queued + wait: their TTLs add `pull_pilot.pull_queue_allowance` (one agent timeout on a pilot). The lock also + keeps the cached resume id fresh across turns, which the claim guard does not do. Rooms take a + per-(room, agent) wake lock (`ResumeLock` under `room_wake_lock::`) from the cursor read + through the posted reply, so a second wake reads what the first left; a mention chain back to the same + agent re-enters it. The room working marker adds the queue allowance (capped at the room claim budget). + The claim guard is the only serialisation for public links and channels. + + **Report-back and stop.** The pull sink (`apply_task_result`) spawns the channel / portal / room + completion report on its CAS-won branch, as the push terminals do. The agent's pull worker registers a + claimed turn as pending in the process registry, so `/api/executions/{id}/terminate` before spawn skips + the turn and reports it `cancelled`. ### The soak duration requirement is mis-cited — correct it when you touch it diff --git a/docs/planning/PULL_MIGRATION_TESTING.md b/docs/planning/PULL_MIGRATION_TESTING.md index 1e606f2359..3f069a9055 100644 --- a/docs/planning/PULL_MIGRATION_TESTING.md +++ b/docs/planning/PULL_MIGRATION_TESTING.md @@ -353,10 +353,10 @@ B6 runtime-verify on the rebuilt image · the ≥2-week soak (#856 / #1766, meas the producer never force-queued (a free slot still meant a push, so rows only queued on overflow) and the backend's own `drain_next` raced the agent's worker for whatever did queue. Two independent capacity counters meant up to 2x `max_parallel_tasks`, invisible to S-02. The flag is now a true either/or for -autonomous triggers; interactive turns keep the synchronous path for now — a **migration state**, not a -boundary. Open Question 7 was **decided 2026-09-16 (#1989)** in favour of the edge adapter: interactive -chat joins the queue once #2842 (queue priority) and #2843 (one turn per conversation) land, after which -the synchronous push dispatch path is deleted. +autonomous triggers and, since #3114, every interactive trigger except `chat` (the UI `/chat` path, which +still pushes). Open Question 7 (#1989) put interactive chat on the queue behind #2842 (queue priority) +and #2843 (one turn per conversation); the synchronous push dispatch path is deleted once `chat` is +routed too. ## 9. Soak measurement set (#1766) @@ -464,10 +464,9 @@ WHERE agent_name = '' AND started_at::timestamptz > now() - interval '24 hours'; ``` -Expected after the #1766 gate + #2391 + #2523 + #2524: **every autonomous -trigger ~100% `pulled`; only the interactive ones (`manual`, `mcp`, `chat`, -`public`, `session`, `voice`, `room`) stay `pushed`, by the Open Question 7 scope -cut (#1989).** Confirm the split per trigger with: +Expected after the #1766 gate + #2391 + #2523 + #2524 + #3114: **every +trigger ~100% `pulled` except `chat` (the UI `/chat` path), which stays +`pushed`.** Confirm the split per trigger with: ```sql SELECT triggered_by, diff --git a/docs/security-reports/cso-diff-2026-09-30-3114-pull-route-interactive.json b/docs/security-reports/cso-diff-2026-09-30-3114-pull-route-interactive.json new file mode 100644 index 0000000000..14d6bd72e2 --- /dev/null +++ b/docs/security-reports/cso-diff-2026-09-30-3114-pull-route-interactive.json @@ -0,0 +1,60 @@ +{ + "version": "1.1", + "date": "2026-09-30", + "mode": "daily", + "scope": "diff", + "confidence_gate": 8, + "branch": "feature/3114-pull-route-interactive", + "issue": [ + "abilityai/trinity#3114" + ], + "phases_run": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 9, + 10, + 12, + 13, + 14 + ], + "attack_surface": { + "new_endpoints": 0, + "changed_endpoints": [ + "stream proxies hold while queued (chat, public, portal)" + ], + "changed_dispatch": "interactive producers enqueue on pull pilots" + }, + "findings": [], + "fixed_during_audit": [], + "appendix": [ + { + "title": "Anonymous public-link turns share max_backlog_depth with scheduled work on pilots", + "class": "DoS (excluded)", + "file": "src/backend/services/backlog_service.py" + }, + { + "title": "Held stream polls the row once per second", + "class": "DoS (excluded)", + "file": "src/backend/services/sync_waiter.py" + } + ], + "supply_chain_summary": "no dependency changes", + "filter_stats": { + "candidates": 2, + "excluded_dos": 2, + "reported": 0 + }, + "totals": { + "critical": 0, + "high": 0, + "medium": 0, + "low": 0 + }, + "trend": "n/a (diff scope)" +} \ No newline at end of file diff --git a/docs/security-reports/cso-diff-2026-09-30-3114-pull-route-interactive.md b/docs/security-reports/cso-diff-2026-09-30-3114-pull-route-interactive.md new file mode 100644 index 0000000000..0b64258f09 --- /dev/null +++ b/docs/security-reports/cso-diff-2026-09-30-3114-pull-route-interactive.md @@ -0,0 +1,21 @@ +# /cso --diff — #3114 interactive producers on the pull queue (2026-09-30) + +**Scope:** `feature/3114-pull-route-interactive` vs `dev`. Daily gate: 8/10. +**Result:** 0 findings (0 critical, 0 high, 0 medium, 0 low). + +## Attack surface added +- **No new endpoint or auth path.** On pull-pilot agents the existing interactive entry points (Session tab, Workspace portal, public link, Slack/Telegram/WhatsApp, rooms, paid, MCP inline auth, validation, run-now) enqueue a row and wait for its terminal instead of pushing. Each entry point keeps its own auth, which runs before dispatch. +- **Claim envelope** now carries `persist_session`, `images`, and the push-equivalent prompt context (`source_user_email`, `source_agent_name`, `source_mcp_key_name`, schedule context, attempt). It is served only to the agent's own scoped key or the internal secret (`_pull_authorized`), so an agent reads only its own queue. +- **Stream proxies** (`routers/chat.py`, `routers/public.py`, `client_portal/router.py`) hold the SSE connection while the row is queued. Each validates its caller and binds the execution to the agent before the hold starts. + +## Phases +- **Secrets:** no key patterns, emails or private IPs in the diff. +- **Enterprise docs guard:** 0 hits across the changed docs. +- **Dependencies, CI, Docker:** unchanged. `docker/base-image/agent_server/services/pull_worker.py` changes; vendored policy mirrors (Invariant #5) untouched. +- **Injection:** no raw SQL added. `conversation_key` is built server-side at every interactive call site (`session:`, `public:`, `channel:`, `room:`, `paid:`); no request model exposes it. +- **Plaintext persistence (G-04 class):** images ride `backlog_metadata` as base64. G-04 skips only each image's `data` field; every other field is scanned. The #1449 sweep clears the column once the row is success/cancelled/skipped; failed rows keep it until the 90-day prune, as for the message text. +- **Paid:** a caller that disconnects while its turn is queued cancels the row, and a turn not claimed within one agent timeout is cancelled before any settlement. + +## Appendix (below the gate) +- **Shared backlog cap (DoS, excluded by rule).** On a pilot, anonymous public-link turns now count toward the agent's `max_backlog_depth` (default 50), which scheduled work also uses. A flood that fills it makes the next scheduled fire fail with `persistent_full` until the queue drains; strict interactive precedence also delays autonomous rows. Before this change public turns on a pilot were pushed against the Redis counter. Follow-up: a separate cap for public-link rows. +- **Stream hold polling.** Each held stream reads the row once per second for up to one agent timeout. diff --git a/src/backend/adapters/message_router.py b/src/backend/adapters/message_router.py index 331d2934fb..920b3edd50 100644 --- a/src/backend/adapters/message_router.py +++ b/src/backend/adapters/message_router.py @@ -663,6 +663,7 @@ async def _handle_message_inner(self, adapter: ChannelAdapter, message: Normaliz result = await self._run_agent_task( adapter, message, agent_name, bot_token, channel, is_group, container, upload_dir, context_prompt, verified_email, image_data, + session_id=session_id, ) if result is None: return @@ -795,6 +796,7 @@ async def _run_agent_task( context_prompt: str, verified_email: Optional[str], image_data: list, + session_id: Optional[str] = None, ): """Step 9: assemble public-channel execution params (source identity, MEM-001 memory injection, restricted tool set) and run the task. @@ -853,12 +855,16 @@ async def _run_agent_task( public_allowed_tools = public_allowed_tools + [_VOICE_REPLY_TOOL] try: - task_execution_service = get_task_execution_service() + # #3114: on a pull pilot the turn is queued under this channel + # conversation's key and awaited here. + from services.task_execution_service import dispatch_and_await_terminal - result = await task_execution_service.execute_task( + result = await dispatch_and_await_terminal( + service=get_task_execution_service(), agent_name=agent_name, message=context_prompt, triggered_by=channel, + conversation_key=f"channel:{session_id}" if session_id else None, source_user_email=source_email, timeout_seconds=None, # Uses agent's configured timeout (TIMEOUT-001) allowed_tools=public_allowed_tools, diff --git a/src/backend/canary/invariants/g04_no_creds_in_backlog_metadata.py b/src/backend/canary/invariants/g04_no_creds_in_backlog_metadata.py index 2804e89d37..8486cdad6c 100644 --- a/src/backend/canary/invariants/g04_no_creds_in_backlog_metadata.py +++ b/src/backend/canary/invariants/g04_no_creds_in_backlog_metadata.py @@ -36,6 +36,7 @@ defect even though it rides the same collected field. """ +import json import re from typing import List @@ -73,6 +74,33 @@ ] +def _scannable(raw: str) -> str: + """``raw`` without the encoded bytes of its ``images`` (#3114). + + A pull pilot's queued turn carries base64 vision input there. Base64 is + random text over ``[A-Za-z0-9+/]``, and ``+``/``/`` give ``\b`` a boundary, + so a multi-megabyte image matches a prefix such as ``AIza`` by chance often + enough to fire a critical violation. Encoded image bytes hold no plaintext + credential to find. Only each image's ``data`` is dropped; every other + field of an image, and any non-dict entry, is still scanned. + """ + if '"images"' not in raw: + return raw + try: + meta = json.loads(raw) + except ValueError: + return raw + if not isinstance(meta, dict) or not meta.get("images"): + return raw + images = meta["images"] + if isinstance(images, list): + images = [ + {k: v for k, v in img.items() if k != "data"} if isinstance(img, dict) else img + for img in images + ] + return json.dumps({**meta, "images": images}) + + def check(snapshot: Snapshot) -> List[ViolationReport]: """One violation per queued row whose backlog_metadata matches a secret pattern.""" violations: List[ViolationReport] = [] @@ -88,6 +116,7 @@ def check(snapshot: Snapshot) -> List[ViolationReport]: # column) has nothing to scan. if not isinstance(raw, str): continue + raw = _scannable(raw) for pattern_name, pattern in _SECRET_PATTERNS: if pattern.search(raw): diff --git a/src/backend/client_portal/models.py b/src/backend/client_portal/models.py index f6df33c63a..29954952fc 100644 --- a/src/backend/client_portal/models.py +++ b/src/backend/client_portal/models.py @@ -106,6 +106,10 @@ class PortalAgentCard(BaseModel): # affordance, so this is the field that decides it. UX, not containment — # the routes re-check with the same predicate. can_manage_canvases: bool = False + # #3114 — this agent's Workspace turns wait on the durable queue for a + # worker (pull pilot). The pending-turn card reads "Waiting for a slot" + # only when this is true. + pulls_turns: bool = False # #138 briefing — ships with the roster at sign-in so the new-chat screen # renders with zero extra fetches. Best-effort live data (a stopped/slow # agent yields None/[]). `playbooks` is the hint-card set (ent#380): the diff --git a/src/backend/client_portal/router.py b/src/backend/client_portal/router.py index d6dd792a21..d0b37e8f55 100644 --- a/src/backend/client_portal/router.py +++ b/src/backend/client_portal/router.py @@ -44,6 +44,8 @@ from services.agent_auth import agent_httpx_client from services.docker_service import get_agent_container from services.platform_audit_service import AuditEventType, platform_audit_service +from services import sync_waiter +from services.pull_pilot import pull_queue_allowance from database import db from . import agent_page, chat_previews, role_card, seat_decisions, service @@ -2083,8 +2085,15 @@ async def proxy_stream(): # cleanly (never `error`) when it simply is not streamable. The client # reads the persisted reply on `stream_end` either way. agent_url = f"http://agent-{agent_name}:8000/api/executions/{execution_id}/stream" - deadline = asyncio.get_event_loop().time() + _STREAM_ATTACH_TIMEOUT_S try: + # #3114: on a pull pilot the turn can wait on the durable queue + # until a worker claims it. Hold while the row is queued; the attach + # window below starts from the claim. + async for tick in sync_waiter.wait_while_queued( + execution_id, pull_queue_allowance(agent_name) + ): + yield tick + deadline = asyncio.get_event_loop().time() + _STREAM_ATTACH_TIMEOUT_S async with agent_httpx_client(agent_name, timeout=None) as client: while True: async with client.stream("GET", agent_url) as response: diff --git a/src/backend/client_portal/service.py b/src/backend/client_portal/service.py index e278730c0c..5d666683e5 100644 --- a/src/backend/client_portal/service.py +++ b/src/backend/client_portal/service.py @@ -745,6 +745,7 @@ def _row_to_card(r: dict, tts_ready: bool, default_voice_id: str | None = None, later silently un-fix #2695 with the whole suite green. """ from services import tts_service + from services.pull_pilot import pull_owns_dispatch name = r["agent_name"] updated = r.get("avatar_updated_at") # Only agents with a generated (non-default) avatar get an image URL; @@ -796,6 +797,8 @@ def _row_to_card(r: dict, tts_ready: bool, default_voice_id: str | None = None, # caller knows its own principal kind and this builder is shared with # the single-agent lookup. can_manage_canvases=can_manage_canvases, + # #3114: an env-var read, so cheap per card. + pulls_turns=pull_owns_dispatch(name, "public"), # ent#403: `None` — no control at all — for every non-platform principal. # The roster payload is the ONLY capability channel an external client # has (#2128): a UI gate written against `GET /api/settings/feature-flags` @@ -3144,7 +3147,13 @@ def _on_resume_failure() -> None: agent_name, message, email, session_id, resolved_model, open_canvas_id=open_canvas_id) if execution_id: - mark_turn_inflight(session_id, execution_id, turn_timeout + 60) + # #3114: plus the claim waits a pull pilot's turn and its cold + # retry may spend. + from services.pull_pilot import pull_queue_allowance + mark_turn_inflight( + session_id, execution_id, + turn_timeout + 60 + 2 * pull_queue_allowance(agent_name), + ) owns_marker = True turn = await _run_sync_turn_and_clear_marker( @@ -3926,7 +3935,10 @@ async def start_portal_turn(agent_name: str, message: str, email: str, # the client's budget and the actual turn disagree about one turn's life. from services.session_turn_service import resolve_turn_timeout turn_timeout = resolve_turn_timeout(agent_name) - wait_budget = portal_max_turn_seconds(turn_timeout) + # #3114: on a pull pilot each of the two attempts may first wait up to one + # agent timeout for a worker to claim it. + from services.pull_pilot import pull_queue_allowance + wait_budget = portal_max_turn_seconds(turn_timeout) + 2 * pull_queue_allowance(agent_name) # #2320: drop any verdict left by the PREVIOUS turn on this thread before # the new marker lands. Without this, a client polling turn N+1 is handed @@ -4438,8 +4450,12 @@ def get_history(agent_name: str, email: str, session_id: str | None = None, # (#2133) — a `lost` verdict never retries, so under-waiting # only costs a premature "check shortly" message, but it is # still the dishonest one. + from services.pull_pilot import pull_queue_allowance from services.session_turn_service import resolve_turn_timeout - wait_budget = portal_max_turn_seconds(resolve_turn_timeout(agent_name)) + wait_budget = ( + portal_max_turn_seconds(resolve_turn_timeout(agent_name)) + + 2 * pull_queue_allowance(agent_name) # #3114 + ) else: wait_budget = ttl except Exception as e: # noqa: BLE001 — budget None → the client falls back diff --git a/src/backend/client_portal/work/service.py b/src/backend/client_portal/work/service.py index 7888f6d52d..2e2ecc3e7b 100644 --- a/src/backend/client_portal/work/service.py +++ b/src/backend/client_portal/work/service.py @@ -93,10 +93,28 @@ def work_kind(row: dict) -> WorkKind: return "other" +def awaits_claim(row: dict) -> bool: + """A pull pilot's ``running`` row that no worker holds yet (#3114). + + A Workspace turn's row is pre-created ``running`` and enqueued a moment + later (after the conversation's resume lock). On a pilot every such turn + reaches the agent only through a worker claim, and a claim writes a lease, + so an un-leased running row is still waiting for a slot. + """ + from services.pull_pilot import pull_owns_dispatch + return ( + (row.get("status") or "").strip().lower() == "running" + and not row.get("lease_expires_at") + and pull_owns_dispatch(row.get("agent_name") or "", row.get("triggered_by")) + ) + + def work_outcome(row: dict, *, stale: bool = False) -> WorkOutcome: """The honest word: `timeout` is a failure whose error says so; `lost` is a running row nothing is watching any more.""" status = (row.get("status") or "").strip().lower() + if awaits_claim(row): + return "lost" if stale else "queued" if status in ("running", "pending_retry"): return "lost" if stale else "running" if status == "queued": diff --git a/src/backend/database.py b/src/backend/database.py index f405c9e6d7..f0e29f962b 100644 --- a/src/backend/database.py +++ b/src/backend/database.py @@ -1588,14 +1588,19 @@ def park_expired_lease(self, execution_id: str, error: str, now_iso: str = None) def count_active_leased_by_agent(self, agent_names): return self._schedule_ops.count_active_leased_by_agent(agent_names) + def execution_awaits_claim(self, execution_id: str) -> bool: + return self._schedule_ops.execution_awaits_claim(execution_id) + def count_active_leased(self, agent_name: str) -> int: return self._schedule_ops.count_active_leased(agent_name) def get_queued_count(self, agent_name: str) -> int: return self._schedule_ops.get_queued_count(agent_name) - def cancel_queued_execution(self, execution_id: str, reason: str = "cancelled") -> bool: - return self._schedule_ops.cancel_queued_execution(execution_id, reason) + def cancel_queued_execution( + self, execution_id: str, reason: str = "cancelled", status: str = "cancelled" + ) -> bool: + return self._schedule_ops.cancel_queued_execution(execution_id, reason, status) def cancel_queued_for_agent(self, agent_name: str, reason: str = "agent_deleted") -> int: return self._schedule_ops.cancel_queued_for_agent(agent_name, reason) diff --git a/src/backend/db/schedules/queue.py b/src/backend/db/schedules/queue.py index 5b5177235c..a6a5d12d64 100644 --- a/src/backend/db/schedules/queue.py +++ b/src/backend/db/schedules/queue.py @@ -1,5 +1,6 @@ """Persistent backlog (BACKLOG-001) + #1081 dark pull/lease/CAS seams.""" +import logging import secrets from datetime import datetime, timedelta, timezone from typing import AbstractSet, Optional, List, Dict @@ -15,6 +16,8 @@ from models import TaskExecutionStatus from utils.helpers import utc_now_iso, to_utc_iso +logger = logging.getLogger(__name__) + class ScheduleQueueMixin: """Queued/backlog accessors + pull-lease reaper seams.""" @@ -219,6 +222,12 @@ def claim_next_queued( return dict(row) if row else None except IntegrityError: if attempt == 2: + # #3114: three straight lost races on one conversation read + # as an empty queue to the worker; log it so a stall shows. + logger.info( + "[Backlog] claim for %s lost the one-running-turn race " + "3 times; returning no row this poll", agent_name, + ) return None if worker_id is not None: stmt = stmt.values(claim_token=secrets.token_urlsafe(32)) @@ -462,6 +471,29 @@ def count_active_leased_by_agent( for row in conn.execute(stmt).mappings() } + def execution_awaits_claim(self, execution_id: str) -> bool: + """True while no worker holds the row yet (#3114): ``queued``, or a + ``running`` row with neither a lease nor a dispatch sentinel, which is a + pre-created portal row in the moment before it is enqueued. A pushed row + carries ``claude_session_id='dispatched'`` before the agent is called and + a claimed row carries a lease, so neither matches. Read-only. + """ + c = schedule_executions.c + stmt = select(c.status, c.lease_expires_at, c.claude_session_id).where( + c.id == execution_id + ) + with get_engine().connect() as conn: + row = conn.execute(stmt).mappings().first() + if row is None: + return False + if row["status"] == TaskExecutionStatus.QUEUED: + return True + return ( + row["status"] == TaskExecutionStatus.RUNNING + and row["lease_expires_at"] is None + and row["claude_session_id"] is None + ) + def count_active_leased(self, agent_name: str) -> int: """Scalar single-agent variant of ``count_active_leased_by_agent``. @@ -491,11 +523,19 @@ def get_queued_count(self, agent_name: str) -> int: row = conn.execute(stmt).mappings().first() return int(row["c"]) if row else 0 - def cancel_queued_execution(self, execution_id: str, reason: str = "cancelled") -> bool: + def cancel_queued_execution( + self, + execution_id: str, + reason: str = "cancelled", + status: str = TaskExecutionStatus.CANCELLED, + ) -> bool: """Cancel a single queued execution. No container interaction. + ``status`` is the terminal written (#3114): CANCELLED by default, FAILED + for an interactive turn no worker claimed in time. Same CAS either way. + Returns: - True if the row was still queued and is now cancelled, False otherwise. + True if the row was still queued and is now terminal, False otherwise. """ now = utc_now_iso() with get_engine().begin() as conn: @@ -508,7 +548,7 @@ def cancel_queued_execution(self, execution_id: str, reason: str = "cancelled") ) ) .values( - status=TaskExecutionStatus.CANCELLED, + status=status, completed_at=now, error=reason, ) diff --git a/src/backend/db/schedules/stats.py b/src/backend/db/schedules/stats.py index 9a354d8e05..0d60baf0dd 100644 --- a/src/backend/db/schedules/stats.py +++ b/src/backend/db/schedules/stats.py @@ -413,7 +413,7 @@ def get_fleet_executions( source_mcp_key_id, source_mcp_key_name, model_used, fan_out_id, business_status, validation_execution_id, turn_integrity, - queued_at, + queued_at, lease_expires_at, -- ent#525: read by the Workspace Work projection; the -- fleet dashboard's model drops them (see FleetExecutionSummary). source_channel, source_channel_chat_id, loop_id diff --git a/src/backend/routers/chat.py b/src/backend/routers/chat.py index 64b4035516..c28601a4ca 100644 --- a/src/backend/routers/chat.py +++ b/src/backend/routers/chat.py @@ -5,6 +5,7 @@ """ from fastapi import APIRouter, Depends, HTTPException, Header from fastapi.responses import StreamingResponse, JSONResponse +import asyncio import httpx import json import logging @@ -21,6 +22,8 @@ ) from services.agent_auth import agent_httpx_client from services.docker_service import get_agent_container +from services import sync_waiter +from services.pull_pilot import pull_queue_allowance from services.model_catalog import InvalidModelError, validate_dispatch_model from services.capacity_manager import ( CapacityFull, @@ -1013,26 +1016,39 @@ async def stream_execution_log( async def proxy_stream(): """Proxy SSE stream from agent container with connect timeout and keepalive.""" agent_url = f"http://agent-{name}:8000/api/executions/{execution_id}/stream" + # #3114: on a pull pilot the turn can wait on the durable queue, and the + # agent knows nothing of it until a worker claims it. Hold the stream + # while the row is queued, then retry a 404 briefly while the row is + # running (claimed, not yet registered on the agent by the worker). + allowance = pull_queue_allowance(name) + attach_retries = sync_waiter.STREAM_ATTACH_RETRIES if allowance else 0 try: + async for tick in sync_waiter.wait_while_queued(execution_id, allowance): + yield tick # Connect timeout prevents hanging if agent is unresponsive, # but read timeout is None since SSE streams are long-lived timeout = httpx.Timeout(connect=10.0, read=None, write=None, pool=None) async with agent_httpx_client(name, timeout=timeout) as client: - async with client.stream("GET", agent_url) as response: - if response.status_code == 404: - # Execution not found on agent (race condition: task not started yet) - yield f"data: {json.dumps({'type': 'error', 'message': 'Execution not yet available on agent', 'retryable': True})}\n\n" - yield f"data: {json.dumps({'type': 'stream_end'})}\n\n" - return - - if response.status_code != 200: - yield f"data: {json.dumps({'type': 'error', 'message': f'Agent returned {response.status_code}'})}\n\n" - yield f"data: {json.dumps({'type': 'stream_end'})}\n\n" - return - - # Stream through data from agent, adding proxy-level keepalive - async for chunk in response.aiter_text(): - yield chunk + while True: + async with client.stream("GET", agent_url) as response: + if response.status_code == 404: + if attach_retries > 0 and sync_waiter.execution_is_running(execution_id): + attach_retries -= 1 + else: + # Execution not found on agent (race condition: task not started yet) + yield f"data: {json.dumps({'type': 'error', 'message': 'Execution not yet available on agent', 'retryable': True})}\n\n" + yield f"data: {json.dumps({'type': 'stream_end'})}\n\n" + return + elif response.status_code != 200: + yield f"data: {json.dumps({'type': 'error', 'message': f'Agent returned {response.status_code}'})}\n\n" + yield f"data: {json.dumps({'type': 'stream_end'})}\n\n" + return + else: + # Stream through data from agent, adding proxy-level keepalive + async for chunk in response.aiter_text(): + yield chunk + return + await asyncio.sleep(sync_waiter.STREAM_ATTACH_RETRY_INTERVAL) except httpx.ConnectError: yield f"data: {json.dumps({'type': 'error', 'message': 'Failed to connect to agent', 'retryable': True})}\n\n" yield f"data: {json.dumps({'type': 'stream_end'})}\n\n" diff --git a/src/backend/routers/internal.py b/src/backend/routers/internal.py index 2119cb22cd..2467746a47 100644 --- a/src/backend/routers/internal.py +++ b/src/backend/routers/internal.py @@ -570,9 +570,13 @@ async def execute_task_internal( idempotency_service.complete(idem, request.execution_id, accepted) return accepted - # Synchronous mode (default, backward compatible) + # Synchronous mode (default, backward compatible). #3114: on a pull pilot + # the turn is queued and awaited here. + from services.task_execution_service import dispatch_and_await_terminal + try: - result = await task_service.execute_task( + result = await dispatch_and_await_terminal( + service=task_service, agent_name=request.agent_name, message=request.message, triggered_by=request.triggered_by, diff --git a/src/backend/routers/paid.py b/src/backend/routers/paid.py index 6d60db34d8..7363588767 100644 --- a/src/backend/routers/paid.py +++ b/src/backend/routers/paid.py @@ -313,12 +313,19 @@ async def paid_chat( ) # Step 3: Execute task - task_service = get_task_execution_service() + from services.task_execution_service import dispatch_and_await_terminal + try: - exec_result = await task_service.execute_task( + # #3114: on a pull pilot the turn is queued and awaited here; a + # resumed session is its conversation key. + exec_result = await dispatch_and_await_terminal( + service=get_task_execution_service(), agent_name=agent_name, message=request_body.message, triggered_by="paid", + conversation_key=( + f"paid:{request_body.session_id}" if request_body.session_id else None + ), system_prompt=build_public_channel_caller_prompt(agent_name), # #1205 resume_session_id=request_body.session_id, # #894: per-agent public-channel model override (None → platform default). diff --git a/src/backend/routers/public.py b/src/backend/routers/public.py index 02b84c62a1..c2c5246e78 100644 --- a/src/backend/routers/public.py +++ b/src/backend/routers/public.py @@ -32,6 +32,8 @@ from services.chat_execution_service import terminate_execution as _terminate_execution from services.chat_signals import ChatDispatchError from services.docker_service import get_agent_container +from services import sync_waiter +from services.pull_pilot import pull_queue_allowance from services.email_service import email_service from services.settings_service import PUBLIC_URL_REACHED_KEY, settings_service from services.task_execution_service import get_task_execution_service @@ -822,16 +824,30 @@ async def public_stream_execution( async def proxy_stream(): """Proxy SSE stream from agent container.""" agent_url = f"http://agent-{agent_name}:8000/api/executions/{execution_id}/stream" + # #3114: hold while a pull pilot's turn is queued, then retry a 404 + # briefly while the claimed row is running (see routers/chat.py). + allowance = pull_queue_allowance(agent_name) + attach_retries = sync_waiter.STREAM_ATTACH_RETRIES if allowance else 0 try: + async for tick in sync_waiter.wait_while_queued(execution_id, allowance): + yield tick async with agent_httpx_client(agent_name, timeout=None) as client: - async with client.stream("GET", agent_url) as response: - if response.status_code != 200: - yield f"data: {json.dumps({'type': 'error', 'message': f'Agent returned {response.status_code}'})}\n\n" - yield f"data: {json.dumps({'type': 'stream_end'})}\n\n" - return - - async for chunk in response.aiter_text(): - yield chunk + while True: + async with client.stream("GET", agent_url) as response: + if response.status_code == 200: + async for chunk in response.aiter_text(): + yield chunk + return + if not ( + response.status_code == 404 + and attach_retries > 0 + and sync_waiter.execution_is_running(execution_id) + ): + yield f"data: {json.dumps({'type': 'error', 'message': f'Agent returned {response.status_code}'})}\n\n" + yield f"data: {json.dumps({'type': 'stream_end'})}\n\n" + return + attach_retries -= 1 + await asyncio.sleep(sync_waiter.STREAM_ATTACH_RETRY_INTERVAL) except httpx.ConnectError: yield f"data: {json.dumps({'type': 'error', 'message': 'Failed to connect to agent'})}\n\n" yield f"data: {json.dumps({'type': 'stream_end'})}\n\n" diff --git a/src/backend/routers/sessions.py b/src/backend/routers/sessions.py index 2dce94abfe..906e5be5c1 100644 --- a/src/backend/routers/sessions.py +++ b/src/backend/routers/sessions.py @@ -427,8 +427,10 @@ def _on_resume_failure() -> None: if result.status != "success": # Bubble execute_task's classified error up. The user message is # already persisted; we don't insert an empty assistant row. + # #3114: an agent at capacity (incl. a pull pilot's turn that no + # worker claimed in time) is a retryable 429. raise HTTPException( - status_code=502, + status_code=429 if result.error_code == "capacity" else 502, detail={ "error": result.error or "Agent execution failed", "execution_id": result.execution_id, diff --git a/src/backend/services/backlog_service.py b/src/backend/services/backlog_service.py index cc995a4971..21a8b18959 100644 --- a/src/backend/services/backlog_service.py +++ b/src/backend/services/backlog_service.py @@ -70,6 +70,11 @@ async def enqueue( collaboration_activity_id: Optional[str], is_self_task: bool = False, self_task_activity_id: Optional[str] = None, + conversation_key: Optional[str] = None, + persist_session: bool = False, + schedule_context: Optional[dict] = None, + attempt: Optional[int] = None, + images: Optional[list] = None, ) -> bool: """Persist an async task request as a QUEUED backlog item. @@ -116,6 +121,15 @@ async def enqueue( "collaboration_activity_id": collaboration_activity_id, "is_self_task": is_self_task, "self_task_activity_id": self_task_activity_id, + # #3114: pull-pilot rows only. The claim hands persist_session and + # images to the worker and composes the prompt from + # schedule_context and attempt. `images` is base64 vision input in + # the push payload's shape; the #1449 scrub NULLs the whole blob + # once the row is authoritative-terminal. + "persist_session": persist_session, + "schedule_context": schedule_context, + "attempt": attempt, + "images": images or None, } queued_at = utc_now_iso() # #2843: a pilot's worker claims these, so tag the conversation a turn @@ -123,9 +137,17 @@ async def enqueue( # Non-pilots are drained by the backend and keep today's behaviour. from services.pull_pilot import is_pull_pilot_agent - conversation_key = None + # #3114: an explicit key from `execute_task` names the conversation + # directly ("session:", "public:", ...) and wins. if is_pull_pilot_agent(agent_name): - conversation_key = request.chat_session_id or request.resume_session_id or None + conversation_key = ( + conversation_key + or request.chat_session_id + or request.resume_session_id + or None + ) + else: + conversation_key = None ok = db.update_execution_to_queued( execution_id, json.dumps(metadata), queued_at, conversation_key ) diff --git a/src/backend/services/capacity_manager.py b/src/backend/services/capacity_manager.py index 75b143b8ea..5edfa3a055 100644 --- a/src/backend/services/capacity_manager.py +++ b/src/backend/services/capacity_manager.py @@ -116,6 +116,13 @@ class PersistentTaskPayload: collaboration_activity_id: Optional[str] is_self_task: bool = False self_task_activity_id: Optional[str] = None + # #3114: pull-pilot rows from `execute_task`. The conversation key feeds + # the claim guard; the rest reach the worker and the claim-time prompt. + conversation_key: Optional[str] = None + persist_session: bool = False + schedule_context: Optional[dict] = None + attempt: Optional[int] = None + images: Optional[list] = None class CapacityFull(Exception): @@ -348,15 +355,16 @@ async def acquire( raise CircuitOpen(agent_name, breaker.retry_after_seconds()) return AcquireResult(state="admitted", execution_id=execution_id) - # ---- #1766: a pull pilot's autonomous work is queue-ONLY ----------- + # ---- #1766: a pull pilot's queued work is queue-ONLY --------------- # The pilot flag used to be purely additive: the agent started pulling, # but the backend kept admitting-and-pushing whenever a slot was free, so # the two paths ran in parallel over one queue with two independent # capacity counters. Skipping admission here makes the durable queue the # single entry point for this agent, so its worker pool IS its capacity - # (#1081 Phase 5, pilot-scoped). `pull_owns_dispatch` excludes - # interactive triggers (temporary, until they are routed — #1989) and fails safe to - # push, so a non-pilot's path is byte-for-byte unchanged. + # (#1081 Phase 5, pilot-scoped). `pull_owns_dispatch` covers the + # autonomous triggers and, since #3114, every interactive trigger but + # `chat`; it fails safe to push, so a non-pilot's path is byte-for-byte + # unchanged. from services.pull_pilot import pull_owns_dispatch pull_exclusive = ( @@ -421,6 +429,11 @@ async def acquire( collaboration_activity_id=overflow_payload.collaboration_activity_id, is_self_task=overflow_payload.is_self_task, self_task_activity_id=overflow_payload.self_task_activity_id, + conversation_key=overflow_payload.conversation_key, + persist_session=overflow_payload.persist_session, + schedule_context=overflow_payload.schedule_context, + attempt=overflow_payload.attempt, + images=overflow_payload.images, ) if not enqueued: raise CapacityFull(agent_name, max_concurrent, "persistent_full") diff --git a/src/backend/services/chat_execution_service.py b/src/backend/services/chat_execution_service.py index 134fcd77a2..9e73c76592 100644 --- a/src/backend/services/chat_execution_service.py +++ b/src/backend/services/chat_execution_service.py @@ -1719,10 +1719,9 @@ def _map_task_failure(name, result, *, idem): """ if result.status in ("failed", "cancelled"): idempotency_service.fail(idem) - # #2889: the immediate path's result carries the producer-side code; - # the backlog-reconstruct path builds its result from the row and has - # none — the header is absent there, except on the at-capacity branch - # below, where `capacity` (#2919) fills the absent code. + # #2889: the producer-side code. The backlog-reconstruct path reads it + # from the row's `[code]` prefix (#3114); a row with no prefix has none, + # and the at-capacity branch below fills `capacity` (#2919). code_headers = _error_code_headers(getattr(result, "error_code", None)) if "at capacity" in (result.error or ""): # #2919: the capacity rejection carries no code, so `capacity` @@ -1906,8 +1905,11 @@ async def _dispatch_sync_backlog(*, name, execution_id, sync_effective_timeout, } idempotency_service.complete(idem, execution_id, receipt) raise ChatDispatchError(504, receipt["message"]) - from services.task_execution_service import TaskExecutionResult + from services.task_execution_service import TaskExecutionResult, _split_error_code + # #3114: a pull-sink terminal stores `[code] text`; keep the code so + # `_map_task_failure` answers the producer's status, not a blanket 503. + error_code, error = _split_error_code(row.error) result = TaskExecutionResult( execution_id=execution_id, status=row.status, @@ -1916,7 +1918,8 @@ async def _dispatch_sync_backlog(*, name, execution_id, sync_effective_timeout, context_used=row.context_used, context_max=row.context_max, session_id=row.claude_session_id, - error=row.error, + error=error, + error_code=error_code, raw_response={ "response": row.response or "", "cost": row.cost, diff --git a/src/backend/services/mcp_auth_service.py b/src/backend/services/mcp_auth_service.py index 85b32b10a1..869fc47d55 100644 --- a/src/backend/services/mcp_auth_service.py +++ b/src/backend/services/mcp_auth_service.py @@ -365,9 +365,10 @@ async def dispatch_chat(email: str, agent: str, message: str) -> "object": assert_email_may_reach_agent(email, agent) email = normalize_email(email) - from services.task_execution_service import get_task_execution_service + from services.task_execution_service import dispatch_and_await_terminal - result = await get_task_execution_service().execute_task( + # #3114: on a pull pilot the turn is queued and awaited here. + result = await dispatch_and_await_terminal( agent_name=agent, message=message, triggered_by=CHAT_TRIGGERED_BY, diff --git a/src/backend/services/public_chat_service.py b/src/backend/services/public_chat_service.py index 919fac0005..74fc89ac9e 100644 --- a/src/backend/services/public_chat_service.py +++ b/src/backend/services/public_chat_service.py @@ -31,7 +31,6 @@ format_user_memory_block, summarize_user_memory_background, ) -from services.task_execution_service import get_task_execution_service from services.upload_service import ( WEB_MAX_FILES, WEB_MAX_FILE_SIZE, @@ -95,11 +94,14 @@ async def _execute_public_chat_background( response in the public chat session. """ try: - task_execution_service = get_task_execution_service() - result = await task_execution_service.execute_task( + # #3114: on a pull pilot the turn is queued and awaited here. + from services.task_execution_service import dispatch_and_await_terminal + + result = await dispatch_and_await_terminal( agent_name=agent_name, message=context_prompt, triggered_by="public", + conversation_key=f"public:{chat_session_id}", source_user_email=source_email, timeout_seconds=900, execution_id=execution_id, @@ -314,7 +316,6 @@ async def run_public_chat(link: dict, chat_request, client_ip: str): # Public executions now get full tracking: execution records, activity stream, # slot management, credential sanitization, and Dashboard timeline visibility. source_email = verified_email or f"anonymous ({client_ip})" - task_execution_service = get_task_execution_service() # Async mode (THINK-001): return execution_id immediately for SSE streaming if chat_request.async_mode: @@ -351,11 +352,15 @@ async def run_public_chat(link: dict, chat_request, client_ip: str): "async_mode": True, } - # Sync mode: wait for result - result = await task_execution_service.execute_task( + # Sync mode: wait for result. #3114: on a pull pilot the turn is queued + # and awaited here. + from services.task_execution_service import dispatch_and_await_terminal + + result = await dispatch_and_await_terminal( agent_name=agent_name, message=context_prompt, triggered_by="public", + conversation_key=f"public:{chat_session.id}", source_user_email=source_email, timeout_seconds=900, # #894: per-agent public-channel model override (None → platform default). diff --git a/src/backend/services/pull_coordination_service.py b/src/backend/services/pull_coordination_service.py index 79b34cc685..bc58227ef3 100644 --- a/src/backend/services/pull_coordination_service.py +++ b/src/backend/services/pull_coordination_service.py @@ -31,6 +31,7 @@ from config import MAX_REDELIVERY from database import db from models import TaskExecutionStatus +from services import channel_completion_report from services import event_dispatch_service from services import subscription_auto_switch from services.activity_service import activity_service @@ -78,6 +79,12 @@ def _compose_pull_system_prompt( *, execution_id: Optional[str], model: Optional[str] = None, + source_user_email: Optional[str] = None, + source_agent_name: Optional[str] = None, + source_mcp_key_name: Optional[str] = None, + timeout_seconds: Optional[int] = None, + schedule_context: Optional[dict] = None, + attempt: Optional[int] = None, ) -> Optional[str]: """Compose platform prompt + execution context + caller override for a pull-claimed turn (#1629). Fail-open: on ANY composition error the turn runs @@ -87,13 +94,26 @@ def _compose_pull_system_prompt( ``model`` (ent#243) selects the prompt tier. It is passed explicitly rather than left to default because this context previously omitted the field entirely — not ``None``-valued, absent — so every pull-claimed turn would - have resolved VERBOSE forever with nothing to indicate why.""" + have resolved VERBOSE forever with nothing to indicate why. + + #3114: the provenance, timeout, schedule and attempt fields match what + ``task_execution_service._compose_effective_system_prompt`` builds on + push, so an interactive turn routed to the queue sees the same context.""" runtime = _resolve_agent_runtime(agent_name) try: + schedule_context = schedule_context or {} exec_ctx = ExecutionContext( agent_name=agent_name, mode=ExecutionContext.derive_mode(triggered_by), triggered_by=triggered_by, + source_user_email=source_user_email, + source_agent_name=source_agent_name, + source_mcp_key_name=source_mcp_key_name, + timeout_seconds=timeout_seconds, + attempt=attempt, + schedule_name=schedule_context.get("name"), + schedule_cron=schedule_context.get("cron"), + schedule_next_run=schedule_context.get("next_run"), execution_id=execution_id, model=model, ) @@ -332,6 +352,12 @@ def _build_claim_response(row: Dict[str, Any]) -> Dict[str, Any]: } if meta.get("file_ids") is not None: payload["file_ids"] = meta.get("file_ids") + # #3114: the push payload's `persist_session` and `images`, carried for + # rows `execute_task` queued. Absent on every other producer's row. + if meta.get("persist_session") is not None: + payload["persist_session"] = bool(meta.get("persist_session")) + if meta.get("images"): + payload["images"] = meta.get("images") # #1629: compose the platform system prompt (parity with the push path) and # hand it to the worker via task_overrides.system_prompt — the field the pull @@ -361,6 +387,12 @@ def _build_claim_response(row: Dict[str, Any]) -> Dict[str, Any]: overrides.get("system_prompt"), execution_id=row["id"], model=overrides.get("model") or row.get("model_used"), + source_user_email=row.get("source_user_email"), + source_agent_name=row.get("source_agent_name"), + source_mcp_key_name=row.get("source_mcp_key_name"), + timeout_seconds=overrides.get("timeout_seconds"), + schedule_context=meta.get("schedule_context") if isinstance(meta.get("schedule_context"), dict) else None, + attempt=meta.get("attempt"), ) payload["task_overrides"] = overrides @@ -520,6 +552,11 @@ def apply_task_result( sanitized_content = sanitize_response(content) if content is not None else None + # #3114: the Session tab reads compact events off the result; persist them + # so a sync caller rebuilding its result from this row still gets them. + compact_events = metadata.get("compact_events") or [] + compact_json = json.dumps(compact_events) if compact_events else None + if row_status == TaskExecutionStatus.SUCCESS: won = db.update_execution_status( execution_id=execution_id, @@ -532,6 +569,7 @@ def apply_task_result( execution_log=log_json, tool_calls=log_json, claude_session_id=session_id, + compact_metadata=compact_json, ), claim_token=claim_token, ) @@ -553,6 +591,7 @@ def apply_task_result( cost=cost, context_used=context_used, context_max=context_max, + compact_metadata=compact_json, ), claim_token=claim_token, ) @@ -580,6 +619,15 @@ def apply_task_result( summary_or_error=summary, cost=cost, ) + # #3114: a delegated child that ran on a pull pilot reports back to the + # channel / portal / room it came from, exactly as the push terminals + # do. `report_completion` skips INLINE_CHANNEL_TRIGGERS turns itself. + channel_completion_report.spawn_completion_report( + execution_id=execution_id, + agent_name=execution.agent_name, + status=str(getattr(row_status, "value", row_status)), + summary_or_error=summary, + ) # #1804: the pull sink is a CAS-won terminal writer, so it owns closing # the paired dispatch activity — the issue names this as one of the next # two victims of the old per-site model. Sync function, async caller: @@ -630,6 +678,11 @@ def apply_task_result( error_message=err_text or f"[{error_code}] pull terminal", failure_kind=switch_kind, ) + # #3114: wake a `dispatch_and_await_terminal` caller in this process + # now; its DB poll covers a caller held by another worker. + from services.sync_waiter import signal_sync_waiter + + signal_sync_waiter(execution_id, None, None) return ResultApplyOutcome("applied", row_status) # CAS lost — reclassify against the freshly-read row. diff --git a/src/backend/services/pull_pilot.py b/src/backend/services/pull_pilot.py index 9cfe53f0b9..9e9ba9a5be 100644 --- a/src/backend/services/pull_pilot.py +++ b/src/backend/services/pull_pilot.py @@ -48,6 +48,25 @@ def is_pull_pilot_agent(agent_name: str) -> bool: return agent_name in _pilot_allowlist() +def pull_queue_allowance(agent_name: str) -> int: + """Seconds a turn may wait on the durable queue before a worker claims it + (#3114): the agent's execution timeout on a pull pilot, 0 elsewhere. + + Sync callers wait up to this long for a claim, and the Redis locks that + span a turn (Session ResumeLock, portal in-flight markers) add it to their + TTL so a turn still waiting for a worker keeps its lock. Fail-safe → 0. + """ + if not is_pull_pilot_agent(agent_name): + return 0 + try: + # Lazy: this module stays stdlib-only at import time. + from database import db + + return max(int(db.get_execution_timeout(agent_name)), 0) + except Exception: # noqa: BLE001 — a config read must not break a turn + return 0 + + # Which autonomous triggers can STRUCTURALLY reach the durable queue # (#2048 named the set; #2391 widened it). # @@ -98,9 +117,8 @@ def is_pull_pilot_agent(agent_name: str) -> bool: # the cron fire (``_execute_retry`` → ``_call_backend_execute_task``), so it is # exactly as pullable as the run it retries. Without it a pilot's scheduled run # was pulled and its retry pushed — the second attempt, which exists because the -# first failed, ran on the path with no lease and no reaper recovery. A retry of -# a MANUAL run still changes system (push → pull) until interactive triggers are -# pullable (#1989); that direction lands on the durable path, not off it. +# first failed, ran on the path with no lease and no reaper recovery. A manual +# run is itself pulled on a pilot (#3114), so its retry stays on the same path. # # **The set is currently equal to ``_AUTONOMOUS_TRIGGERS``, and it stays an # explicit allow-list anyway.** That is the point of it: a trigger added to the @@ -126,6 +144,16 @@ def is_pull_pilot_agent(agent_name: str) -> bool: ) +# Non-autonomous triggers a pull pilot takes ONLY from the durable queue (#3114): +# every interactive trigger except ``chat``, plus ``validation``. Their sync +# callers go through ``task_execution_service.dispatch_and_await_terminal``, +# which waits for a worker to claim the row (bounded by one agent timeout) and +# then for its terminal. ``chat`` is the UI ``/chat`` path +# (``dispatch_admission_service`` → ``chat_execution_service.run_chat_turn``), +# which still pushes; it is routed in a later change. +PULL_REACHABLE_NON_AUTONOMOUS = (INTERACTIVE_TRIGGERS - {"chat"}) | {"validation"} + + def pull_owns_dispatch(agent_name: str, triggered_by: Optional[str]) -> bool: """True when this ``(agent, trigger)`` pair must reach the agent ONLY by the agent claiming it from the durable queue — the backend neither pushes it nor @@ -141,15 +169,12 @@ def pull_owns_dispatch(agent_name: str, triggered_by: Optional[str]) -> bool: S-02, which counts ``ZCARD`` only. Making the pilot flag a true either/or restores one capacity owner per agent. - **Interactive turns are excluded — temporarily.** Only the autonomous - trigger set queues; a human chat / Session-tab turn keeps today's synchronous - push path and today's Redis session lock. ``TARGET_ARCHITECTURE.md`` Open - Question 7 is decided the other way (#1989): the queue carries ALL traffic, - interactive included, and the push path is deleted afterwards. The exclusion - is a migration state, not a design boundary. Its two prerequisites are in - the claim: interactive turns go first (#2842, ``INTERACTIVE_TRIGGERS``) and - one turn per conversation runs at a time (#2843, ``conversation_key``). What - remains is routing the interactive producers onto the queue. + **Interactive turns are pulled too (#3114),** except the UI ``/chat`` path + (trigger ``chat``). ``PULL_REACHABLE_NON_AUTONOMOUS`` lists them. The claim + orders them ahead of batch work (#2842, ``INTERACTIVE_TRIGGERS``) and runs + one turn per conversation at a time (#2843, ``conversation_key``). Session + tab and portal turns keep their Redis locks as well; the lock TTLs carry + ``pull_queue_allowance`` so a turn still waiting for a claim keeps its lock. Fail-safe: any error resolving the trigger set returns ``False``, i.e. the unchanged push behaviour. The dangerous direction would be silently claiming @@ -169,6 +194,8 @@ def pull_owns_dispatch(agent_name: str, triggered_by: Optional[str]) -> bool: # instead of ``"reject"``, so a False here is the exact condition under # which scheduled capacity semantics stay byte-for-byte as they were. # See ``PULL_REACHABLE_TRIGGERS`` for why each trigger is in or out. + if triggered_by in PULL_REACHABLE_NON_AUTONOMOUS: + return True return triggered_by in (_AUTONOMOUS_TRIGGERS & PULL_REACHABLE_TRIGGERS) except Exception: # noqa: BLE001 — unresolvable trigger set ⇒ push, as today logger.warning( diff --git a/src/backend/services/session_turn_service.py b/src/backend/services/session_turn_service.py index a6c04a61af..2f2c30063b 100644 --- a/src/backend/services/session_turn_service.py +++ b/src/backend/services/session_turn_service.py @@ -128,10 +128,18 @@ def resolve_lock_ttl(agent_name: str) -> int: buffer, capped at ``LOCK_TTL_FALLBACK``. Falls back to the cap on any lookup failure — safer to over-TTL than under-TTL since both keys are auto-expiring strings, not state we care to keep precise. + + #3114: on a pull pilot a turn can wait up to one agent timeout for a worker + to claim it, and the cold retry after a resume-not-found waits again, so + twice that allowance is added to both the TTL and its cap. The lock is held + for every wait. """ try: + from services.pull_pilot import pull_queue_allowance + timeout = db.get_execution_timeout(agent_name) - return min(timeout + 30, LOCK_TTL_FALLBACK) + allowance = 2 * pull_queue_allowance(agent_name) + return min(timeout + 30 + allowance, LOCK_TTL_FALLBACK + allowance) except Exception as e: logger.warning( "[SessionTurn] get_execution_timeout failed for %s (%s) — using fallback %ds", @@ -288,13 +296,19 @@ def __init__( claude_session_id: Optional[str], session_id: str, ttl_seconds: int = LOCK_TTL_FALLBACK, + *, + key: Optional[str] = None, + wait_seconds: Optional[float] = None, ): - self._key = ( + # #3114: `key` / `wait_seconds` let a room wake reuse this lock under + # its own key with a wait as long as one turn. + self._key = key or ( session_lock_key(agent_name, claude_session_id) if claude_session_id else f"session_lock:cold:{session_id}" ) self._ttl = ttl_seconds + self._wait = wait_seconds self._token = secrets.token_urlsafe(16) self._redis = None self._held = False @@ -308,7 +322,9 @@ async def __aenter__(self) -> "ResumeLock": ) return self - deadline = asyncio.get_event_loop().time() + LOCK_WAIT_TOTAL_SECONDS + deadline = asyncio.get_event_loop().time() + ( + LOCK_WAIT_TOTAL_SECONDS if self._wait is None else self._wait + ) while True: try: acquired = await self._redis.set( @@ -459,19 +475,23 @@ async def run_resumable_turn( ttl = lock_ttl if lock_ttl is not None else resolve_lock_ttl(agent_name) - from services.task_execution_service import get_task_execution_service - service = get_task_execution_service() + # #3114: both attempts go through the sync adapter, so on a pull pilot the + # turn is queued under this conversation's key and awaited; on push it is + # plain `execute_task`. + from services.task_execution_service import dispatch_and_await_terminal + conversation_key = f"session:{session_key}" fallback_fired = False fallback_reason: Optional[str] = None async with ResumeLock(agent_name, resumed_with, session_key, ttl_seconds=ttl): - result = await service.execute_task( + result = await dispatch_and_await_terminal( agent_name=agent_name, message=message, triggered_by=triggered_by, resume_session_id=resumed_with, persist_session=True, + conversation_key=conversation_key, **execute_kwargs, ) @@ -507,12 +527,13 @@ async def run_resumable_turn( # has already stopped watching. `execute_task` creates the row when # none is supplied. retry_kwargs = {k: v for k, v in execute_kwargs.items() if k != "execution_id"} - result = await service.execute_task( + result = await dispatch_and_await_terminal( agent_name=agent_name, message=cold_message if cold_message is not None else message, triggered_by=triggered_by, resume_session_id=None, persist_session=True, + conversation_key=conversation_key, **retry_kwargs, ) resumed_with = None diff --git a/src/backend/services/sync_waiter.py b/src/backend/services/sync_waiter.py index c8ce2af359..0fcb52f75b 100644 --- a/src/backend/services/sync_waiter.py +++ b/src/backend/services/sync_waiter.py @@ -23,7 +23,7 @@ import asyncio import logging -from typing import Any, Dict, Optional +from typing import Any, AsyncIterator, Dict, Optional from models import TaskExecutionStatus @@ -46,6 +46,17 @@ # Module-level constant so tests can monkeypatch a tighter interval. SYNC_WAITER_POLL_INTERVAL = 5.0 # seconds +# #3114: live-stream proxies on a pull pilot. A queued turn is unknown to the +# agent until a worker claims it. The proxy holds the SSE connection open and +# sends one comment line per tick. Module-level so tests +# can shorten them. +STREAM_QUEUED_POLL_INTERVAL = 1.0 +# After the claim the row is `running` a moment before the worker registers the +# turn with the agent server; a 404 in that window is retried this often. +STREAM_ATTACH_RETRIES = 5 +STREAM_ATTACH_RETRY_INTERVAL = 0.5 +STREAM_QUEUED_TICK = ": queued\n\n" + # Anything that's NOT queued/running counts as terminal for sync waiter purposes. TERMINAL_TASK_STATUSES = frozenset( { @@ -210,3 +221,47 @@ async def _poll_db(): poll_task.cancel() if not fut.done(): fut.cancel() + + +async def wait_while_queued(execution_id: str, budget: float) -> AsyncIterator[str]: + """Yield ``STREAM_QUEUED_TICK`` once per ``STREAM_QUEUED_POLL_INTERVAL`` + while no worker holds the execution row (``db.execution_awaits_claim``: + queued, or pre-created and about to be enqueued), for at most ``budget`` + seconds (#3114). + + Returns once a worker holds the row, it is terminal or gone, or the budget + is spent. ``budget <= 0`` returns at once without reading the row, so a + non-pilot's stream is untouched. A failed row read ends the hold: the + proxy then behaves as it did before. + """ + if budget <= 0: + return + from database import db + + loop = asyncio.get_running_loop() + deadline = loop.time() + budget + while True: + try: + waiting = db.execution_awaits_claim(execution_id) + except Exception as exc: # noqa: BLE001 — a DB blip must not kill the stream + logger.warning("[SyncWaiter] queued-hold read for %s failed: %s", execution_id, exc) + return + if not waiting: + return + remaining = deadline - loop.time() + if remaining <= 0: + return + yield STREAM_QUEUED_TICK + await asyncio.sleep(min(STREAM_QUEUED_POLL_INTERVAL, remaining)) + + +def execution_is_running(execution_id: str) -> bool: + """True when the row is ``running`` (#3114): an agent 404 then means the + claiming worker has not registered the turn yet. False on any read error.""" + from database import db + + try: + row = db.get_execution(execution_id) + except Exception: # noqa: BLE001 + return False + return row is not None and row.status == TaskExecutionStatus.RUNNING diff --git a/src/backend/services/task_execution_service.py b/src/backend/services/task_execution_service.py index 57a2846844..30c3c77090 100644 --- a/src/backend/services/task_execution_service.py +++ b/src/backend/services/task_execution_service.py @@ -24,7 +24,7 @@ from dataclasses import dataclass, field from datetime import datetime from enum import Enum -from typing import Any, Optional +from typing import Any, Optional, Tuple import httpx @@ -54,7 +54,11 @@ from services.platform_audit_service import AuditEventType, platform_audit_service # #2048: stdlib-only leaf by construction, so this cannot cycle back through the # capacity stack at import time (its own reference to this module is lazy). -from services.pull_pilot import note_unreachable_pull_trigger, pull_owns_dispatch +from services.pull_pilot import ( + INTERACTIVE_TRIGGERS, + note_unreachable_pull_trigger, + pull_owns_dispatch, +) from services.settings_service import settings_service from utils.credential_sanitizer import sanitize_dict, sanitize_execution_log, sanitize_response, sanitize_text from services.runtime_secret_scrub import get_staged_values, scrub_obj, scrub_text @@ -604,6 +608,11 @@ def build_pull_queue_payload( source_user_email: Optional[str], source_agent_name: Optional[str], slot_already_held: bool, + conversation_key: Optional[str] = None, + persist_session: bool = False, + schedule_context: Optional[dict] = None, + attempt: Optional[int] = None, + images: Optional[list] = None, ) -> Optional[PersistentTaskPayload]: """The #2391 producer gate: the overflow payload that lets THIS producer put a row on the durable queue, or ``None`` to keep today's ``"reject"`` policy. @@ -691,15 +700,50 @@ def build_pull_queue_payload( x_source_agent=source_agent_name, triggered_by=triggered_by, collaboration_activity_id=None, + # #3114: what the push payload sends beside the request. The worker + # needs persist_session and images to run the same turn; the claim + # composes the prompt from schedule_context and attempt; the claim + # guard serialises turns sharing conversation_key. + conversation_key=conversation_key, + persist_session=persist_session, + schedule_context=schedule_context, + attempt=attempt, + images=images or None, ) +# #3114: triggers whose caller holds a person (or a validation verdict) on the +# answer. On a pull pilot such a turn first waits for a worker to CLAIM it, +# bounded by one agent timeout, and is cancelled if none does in time. +_CLAIM_WAITING_TRIGGERS = INTERACTIVE_TRIGGERS | {"validation"} +# Row-read cadence while a sync turn waits for its claim. Module-level so tests +# can shorten it. +QUEUE_CLAIM_POLL_INTERVAL = 1.0 + + +async def _wait_until_claimed(execution_id: str, budget: float) -> bool: + """Poll the row while it is ``queued``. True once it left ``queued`` + (claimed, terminal, or gone); False when it is still queued at ``budget``.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + budget + while True: + row = db.get_execution(execution_id) + if row is None or row.status != TaskExecutionStatus.QUEUED: + return True + remaining = deadline - loop.time() + if remaining <= 0: + return False + await asyncio.sleep(min(QUEUE_CLAIM_POLL_INTERVAL, remaining)) + + async def dispatch_and_await_terminal( *, agent_name: str, message: str, triggered_by: str, wait_timeout: Optional[float] = None, + claim_budget: Optional[float] = None, + service: Optional["TaskExecutionService"] = None, **execute_kwargs, ) -> TaskExecutionResult: """``execute_task`` for a caller that genuinely needs the answer in-line. @@ -717,17 +761,28 @@ async def dispatch_and_await_terminal( block correctly while the work happens somewhere else, which is a different thing and one the adapter already does. - ⚠️ Under pull, nothing signals the waiter directly (the pull sink writes the - terminal through the CAS, it does not know about this registry), so the wake - comes from ``sync_waiter``'s DB-poll fallback and latency is bounded by - ``SYNC_WAITER_POLL_INTERVAL``. For a turn measured in seconds-to-minutes that - is noise, and it is deliberately not worth a second signalling path. + The pull sink signals the waiter when its CAS write wins (#3114), so a + caller in the same process wakes at once; a caller on another uvicorn + worker wakes on ``sync_waiter``'s DB poll within + ``SYNC_WAITER_POLL_INTERVAL``. + + **Interactive and validation triggers wait in two phases (#3114).** Phase 1 + holds while the row is ``queued``, for at most ``claim_budget`` seconds + (default one agent execution timeout; a room passes a shorter one). If no + worker claimed it by then the row is stored FAILED and the caller gets + FAILED/``CAPACITY``, which is what push answers for an agent with no free + slot. A cancel that loses the race means a worker just claimed the row, so + the wait carries on. Phase 2 is the terminal wait, timed from the claim. + Other triggers keep the single terminal wait. A wait that times out returns a FAILED result with ``TIMEOUT`` rather than raising — the execution keeps running and its real terminal still lands on the row, exactly as for a fan-out deadline (#2524). + + ``service`` is for a caller that holds its own ``TaskExecutionService`` + (``ValidationService`` takes one by injection); default is the singleton. """ - result = await get_task_execution_service().execute_task( + result = await (service or get_task_execution_service()).execute_task( agent_name=agent_name, message=message, triggered_by=triggered_by, @@ -738,16 +793,59 @@ async def dispatch_and_await_terminal( from services.sync_waiter import wait_for_sync_terminal + try: + agent_timeout = float(db.get_execution_timeout(agent_name)) + except Exception: # noqa: BLE001 — a config read must not break dispatch + agent_timeout = 7200.0 if wait_timeout is None: - try: - wait_timeout = float(db.get_execution_timeout(agent_name)) + 120.0 - except Exception: # noqa: BLE001 — a config read must not break dispatch - wait_timeout = 7320.0 + wait_timeout = agent_timeout + 120.0 + if claim_budget is None: + claim_budget = agent_timeout logger.info( "[TaskExecService] %s dispatch for %s queued as %s; awaiting its terminal", triggered_by, agent_name, result.execution_id, ) + if triggered_by in _CLAIM_WAITING_TRIGGERS: + try: + claimed = await _wait_until_claimed(result.execution_id, claim_budget) + except asyncio.CancelledError: + # The caller went away: a turn still queued must not run later for + # nobody (a paid turn would run unsettled). + if db.cancel_queued_execution(result.execution_id, reason="caller went away"): + _close_unclaimed_turn(agent_name, result.execution_id, + TaskExecutionStatus.CANCELLED, "caller went away") + raise + if not claimed: + error = ( + f"Agent at capacity (queued turn not claimed in {int(claim_budget)}s)" + ) + # #3114: stored FAILED, so it reads as a capacity failure; only a + # caller going away (above) leaves a CANCELLED row. + if db.cancel_queued_execution( + result.execution_id, reason=error, status=TaskExecutionStatus.FAILED + ): + logger.info( + "[TaskExecService] %s: queued %s turn %s not claimed in %ss; failed", + agent_name, triggered_by, result.execution_id, int(claim_budget), + ) + _close_unclaimed_turn(agent_name, result.execution_id, + TaskExecutionStatus.FAILED, error) + return TaskExecutionResult( + execution_id=result.execution_id, + status=TaskExecutionStatus.FAILED, + response="", + error=error, + error_code=TaskExecutionErrorCode.CAPACITY, + ) + # Lost the cancel: a worker claimed it just now. Phase 2. + # A short turn can finish inside one claim poll, before the waiter + # registers, and the sink's wake-up is then lost. Read the row first. + from services.sync_waiter import TERMINAL_TASK_STATUSES + + done = result_from_execution_row(result.execution_id) + if done is not None and done.status in TERMINAL_TASK_STATUSES: + return done try: await wait_for_sync_terminal(result.execution_id, wait_timeout) except asyncio.TimeoutError: @@ -755,33 +853,90 @@ async def dispatch_and_await_terminal( execution_id=result.execution_id, status=TaskExecutionStatus.FAILED, response="", - error=f"Timed out after {int(wait_timeout)}s waiting for the queued execution", + error=( + f"Execution timed out after {int(wait_timeout)}s waiting for " + f"the queued execution" + ), error_code=TaskExecutionErrorCode.TIMEOUT, ) return result_from_execution_row(result.execution_id) or result +def _close_unclaimed_turn(agent_name: str, execution_id: str, status, error: str) -> None: + """#3114: a claim-wait CAS winner (unclaimed FAILED, caller-away CANCELLED) + closes its dispatch activity and emits the terminal event, as every other + terminal writer does. Fire-and-forget; a running loop is guaranteed here.""" + activity_service.spawn_close_execution_activity(execution_id, status, error=error) + event_dispatch_service.spawn_task_terminal_event( + agent_name, execution_id, terminal_status=status, summary_or_error=error, + ) + + +# Values `mark_execution_dispatched` writes into `claude_session_id` before a +# real session id exists. Never a resumable session. +_DISPATCH_SESSION_SENTINELS = frozenset({"dispatched", "dispatched_async"}) +_ERROR_CODE_VALUES = {code.value: code for code in TaskExecutionErrorCode} + + +def _split_error_code( + error: Optional[str], +) -> Tuple[Optional[TaskExecutionErrorCode], Optional[str]]: + """``(code, text)`` from an error the pull sink stored as ``[code] text`` + (``apply_task_result``). The prefix is removed only when it names a + ``TaskExecutionErrorCode``; channel adapters post ``text`` to people.""" + if not error or not error.startswith("["): + return None, error + head, sep, rest = error[1:].partition("]") + code = _ERROR_CODE_VALUES.get(head.strip().lower()) if sep else None + return (code, rest.strip()) if code else (None, error) + + def result_from_execution_row(execution_id: str) -> Optional[TaskExecutionResult]: """Rebuild a ``TaskExecutionResult`` from a terminal execution row. The row is the authority once a turn has run somewhere other than inside the caller's await — under pull that is every turn. Returns None when the row is gone, so the caller can fall back to whatever it already had. + + #3114: shaped like the push result its interactive callers already read. + ``response`` is set on success and cancel (#679 returns a cancelled turn's + partial reply), ``error_code`` comes back out of the ``[code]`` prefix the + pull sink writes and the prefix is removed from ``error``, + compact events ride ``raw_response["metadata"]``, and the dispatch + sentinels never pass for a session id. """ execution = db.get_execution(execution_id) if execution is None: return None status = getattr(execution, "status", None) status = status.value if hasattr(status, "value") else str(status) + error = getattr(execution, "error", None) + session_id = getattr(execution, "claude_session_id", None) + if session_id in _DISPATCH_SESSION_SENTINELS: + session_id = None + raw_response: dict = {} + compact = getattr(execution, "compact_metadata", None) + if compact: + try: + raw_response = {"metadata": {"compact_events": json.loads(compact)}} + except (TypeError, ValueError): + raw_response = {} + keeps_response = status in ( + TaskExecutionStatus.SUCCESS.value, TaskExecutionStatus.CANCELLED.value + ) + error_code, error = _split_error_code(error) return TaskExecutionResult( execution_id=execution_id, status=status, - response=getattr(execution, "response", None) or "", + response=(getattr(execution, "response", None) or "") if keeps_response else "", cost=getattr(execution, "cost", None), context_used=getattr(execution, "context_used", None), context_max=getattr(execution, "context_max", None), - session_id=getattr(execution, "claude_session_id", None), - error=getattr(execution, "error", None), + session_id=session_id, + execution_log=getattr(execution, "execution_log", None), + raw_response=raw_response, + error=error, + error_code=error_code, ) @@ -1274,6 +1429,10 @@ async def execute_task( source_channel_client: Optional[str] = None, # ent#555 — which canvas the user had open when they sent this turn. open_canvas_id: Optional[str] = None, + # #3114: the conversation this turn continues, e.g. "session:". + # Only a pull pilot's queue reads it: the claim never runs two queued + # turns with the same key at once. Ignored on the push path. + conversation_key: Optional[str] = None, # #2973: inter-agent chain depth (#2806) the caller already admitted; # stamped on the row created here. Ignored when `execution_id` is given. chain_depth: Optional[int] = None, @@ -1399,6 +1558,11 @@ async def execute_task( source_user_email=source_user_email, source_agent_name=source_agent_name, slot_already_held=slot_already_held, + conversation_key=conversation_key, + persist_session=persist_session, + schedule_context=schedule_context, + attempt=attempt, + images=images, ) # Wrap entire execution flow to ensure execution status is updated on any failure. diff --git a/src/backend/services/validation_service.py b/src/backend/services/validation_service.py index df2980656d..2a6d947b38 100644 --- a/src/backend/services/validation_service.py +++ b/src/backend/services/validation_service.py @@ -275,7 +275,12 @@ async def validate_execution( try: # 4. Run validation via TaskExecutionService - result = await self._task_service.execute_task( + # #3114: on a pull pilot the validation turn is queued and + # awaited here. + from services.task_execution_service import dispatch_and_await_terminal + + result = await dispatch_and_await_terminal( + service=self._task_service, agent_name=agent_name, message=validation_prompt, triggered_by="validation", diff --git a/src/backend/shared_sessions/service.py b/src/backend/shared_sessions/service.py index d703808a98..16ad89eae9 100644 --- a/src/backend/shared_sessions/service.py +++ b/src/backend/shared_sessions/service.py @@ -12,6 +12,7 @@ from __future__ import annotations import asyncio +import contextvars import logging import re import uuid @@ -113,6 +114,12 @@ MAX_CONTENT_CHARS = 8000 ROOM_TURN_TIMEOUT_SECONDS = 300 +# #3114: how long a room turn waits for a pull pilot's worker to claim it. A +# room post is held while it waits, so this is far below the agent timeout a +# Session turn gets; an unclaimed turn fails as at-capacity and the room says so. +ROOM_CLAIM_BUDGET_SECONDS = 300 +# Slack on the per-(room, agent) wake lock beyond the longest turn it covers. +ROOM_WAKE_LOCK_BUFFER_SECONDS = 60 _MENTION_RE = re.compile(r"@([A-Za-z0-9][A-Za-z0-9_-]{0,99})") @@ -1023,8 +1030,16 @@ def _working_key(room_id: str, agent_name: str) -> str: return f"room_working:{room_id}:{agent_name}" -def _working_ttl() -> int: - return ROOM_TURN_TIMEOUT_SECONDS + 30 +def _room_queue_allowance(agent_name: str) -> int: + """Seconds a room turn of ``agent_name`` can sit queued (#3114): the + pilot's queue allowance, bounded by the room claim budget; 0 off-pilot.""" + from services.pull_pilot import pull_queue_allowance + + return min(pull_queue_allowance(agent_name), ROOM_CLAIM_BUDGET_SECONDS) + + +def _working_ttl(agent_name: str) -> int: + return ROOM_TURN_TIMEOUT_SECONDS + 30 + _room_queue_allowance(agent_name) def _mark_agent_working(room_id: str, agent_name: str) -> None: @@ -1032,7 +1047,7 @@ def _mark_agent_working(room_id: str, agent_name: str) -> None: from redis_breaker_util import get_breaker_redis client = get_breaker_redis() if client is not None: - client.set(_working_key(room_id, agent_name), "1", ex=_working_ttl()) + client.set(_working_key(room_id, agent_name), "1", ex=_working_ttl(agent_name)) except Exception as e: # noqa: BLE001 — a missing indicator must not fail a turn logger.warning("room %s: working-marker SET failed for %s: %s", room_id, agent_name, e) @@ -1098,7 +1113,51 @@ def _apply_wake_cap(room_id: str, sender_kind: str, sender_identity: str, return allowed +# #3114: wake-lock keys this task chain already holds. A reply posted under the +# lock can mention its way back to the same agent (A -> B -> A); that nested +# wake runs inside the outer one's chain and must not wait on its own lock. +_held_wake_locks: contextvars.ContextVar[frozenset] = contextvars.ContextVar( + "room_wake_locks_held", default=frozenset() +) + + async def _wake_agent(current_user, room_id: str, agent_name: str, chain_depth: int) -> None: + """Serialise wakes of one pull-pilot agent in one room, then run the turn (#3114). + + Two posts that both mention an agent wake it twice. Unserialised, both + read the same cursor, delta and cached session id and both answer the same + messages. The lock spans the whole turn, so the second wake reads what the + first left; it waits up to one turn's length, then gives up visibly. + """ + from services.pull_pilot import pull_owns_dispatch + from services.session_turn_service import ResumeLock, ResumeLockBusy + + # Pull pilots only. Push agents keep their wake behaviour unchanged. + if not pull_owns_dispatch(agent_name, "room"): + await _wake_agent_locked(current_user, room_id, agent_name, chain_depth) + return + key = f"room_wake_lock:{room_id}:{agent_name}" + held = _held_wake_locks.get() + if key in held: + await _wake_agent_locked(current_user, room_id, agent_name, chain_depth) + return + ttl = (ROOM_TURN_TIMEOUT_SECONDS + _room_queue_allowance(agent_name) + + ROOM_WAKE_LOCK_BUFFER_SECONDS) + try: + async with ResumeLock(agent_name, None, "", ttl, key=key, wait_seconds=ttl): + token = _held_wake_locks.set(held | {key}) + try: + await _wake_agent_locked(current_user, room_id, agent_name, chain_depth) + finally: + _held_wake_locks.reset(token) + except ResumeLockBusy: + logger.warning("room %s: %s still busy with an earlier turn after %ss", + room_id, agent_name, ttl) + _post_system(room_id, f"{agent_name} is still busy with an earlier turn.") + + +async def _wake_agent_locked(current_user, room_id: str, agent_name: str, + chain_depth: int) -> None: """One agent turn: delta -> execute_task -> auto-post the reply. Every wake is an ORDINARY execution through the standard path, so slots, the @@ -1130,7 +1189,7 @@ async def _wake_agent(current_user, room_id: str, agent_name: str, chain_depth: _broadcast("room_participant_state", {"room_id": room_id, "identity": agent_name, "state": "working"}) - from services.task_execution_service import get_task_execution_service + from services.task_execution_service import dispatch_and_await_terminal # trinity-enterprise#363. Derived HERE, per wake, from membership — not # threaded down from `post_message`, which does hold the list: `_wake_agent` @@ -1168,10 +1227,13 @@ async def _wake_agent(current_user, room_id: str, agent_name: str, chain_depth: person_email=client_email, internal_audience=not user_facing)) try: - result = await get_task_execution_service().execute_task( + # #3114: on a pull pilot the turn is queued under this room's key and + # awaited here. + result = await dispatch_and_await_terminal( agent_name=agent_name, message=turn_prefix + manifest_prefix + _build_turn_prompt(room, agent_name, delta, cold, user_facing), triggered_by="room", + conversation_key=f"room:{room_id}", system_prompt=room_prompt, images=images or None, source_user_email=client_email, @@ -1190,6 +1252,8 @@ async def _wake_agent(current_user, room_id: str, agent_name: str, chain_depth: timeout_seconds=ROOM_TURN_TIMEOUT_SECONDS, resume_session_id=cached, persist_session=True, + # #3114: hold the post for a claim at most this long on a pilot. + claim_budget=ROOM_CLAIM_BUDGET_SECONDS, ) except asyncio.CancelledError: # ent#220 item 2: cancellation is a BaseException, so it used to slip diff --git a/src/frontend/src/components/portal/PortalConversation.vue b/src/frontend/src/components/portal/PortalConversation.vue index f8df81f4c6..50195e8480 100644 --- a/src/frontend/src/components/portal/PortalConversation.vue +++ b/src/frontend/src/components/portal/PortalConversation.vue @@ -841,7 +841,7 @@ import { agentDisplayName } from '@/utils/agentName' import PortalAgentBubble from './PortalAgentBubble.vue' import PortalWorkCard from './PortalWorkCard.vue' import { usePortalWorkStore } from '@/stores/portalWork' -import { askAboutItPrefill, childrenForChat, itemById, previewTitle } from './portalWork' +import { askAboutItPrefill, childrenForChat, clockRestartsAt, itemById, pendingTurnOutcome, previewTitle } from './portalWork' import { activityFromStreamEvent, resolveActivityText } from '@/utils/workActivity' import PortalAvatar from './PortalAvatar.vue' import PortalStarButton from './PortalStarButton.vue' @@ -1196,7 +1196,7 @@ const liveCardItem = computed(() => { id: activeExecutionId.value || 'pending', agent_name: props.agent.name, status: 'running', - outcome: 'running', + outcome: pendingTurnOutcome(liveStreamActivity.value, props.agent.pulls_turns), kind: 'turn', // The feed's own title shape, so its row lands without re-wrapping it (#2964). title: previewTitle(pendingUserText.value || lastUserText()), @@ -2141,6 +2141,9 @@ async function cancelTurn() { // stream-json frames never carry, so the card only ever said nothing.) const streaming = ref(false) const liveStreamActivity = ref(null) +watch(() => liveCardItem.value.outcome, (now, before) => { + if (clockRestartsAt(before, now)) elapsed.value = 0 +}) // Bumped after each completed turn; `PortalDeliverables` watches it. const deliverableTick = ref(0) diff --git a/src/frontend/src/components/portal/portalWork.js b/src/frontend/src/components/portal/portalWork.js index 9cbf822f2e..55423a63e4 100644 --- a/src/frontend/src/components/portal/portalWork.js +++ b/src/frontend/src/components/portal/portalWork.js @@ -93,6 +93,24 @@ export function soleStoppableItem(items, stoppingIds = []) { return candidates.length === 1 ? candidates[0] : null } +/** + * #3114: the status of a just-sent turn before the Work feed has read its row. + * On a pull agent (`pullsTurns`), until the agent streams anything the turn may + * still be waiting for a worker, so it reads "Waiting for a slot". Every other + * agent runs the turn on send. + */ +export function pendingTurnOutcome(streamActivity, pullsTurns) { + return pullsTurns && !streamActivity ? 'queued' : 'running' +} + +/** + * #3114: the card's clock restarts when a worker picks the turn up, so + * "Working" counts working time only. + */ +export function clockRestartsAt(before, now) { + return before === 'queued' && now === 'running' +} + /** * The status word a person reads. Honest about WHY it ended (the * `loopStatusLabel` rule, applied to executions): a timeout, a cancel and a diff --git a/src/frontend/tests/unit/portalWork.spec.js b/src/frontend/tests/unit/portalWork.spec.js index b07a87a1d6..413d60c1c4 100644 --- a/src/frontend/tests/unit/portalWork.spec.js +++ b/src/frontend/tests/unit/portalWork.spec.js @@ -54,7 +54,7 @@ vi.mock('@/api', () => ({ default: api })) vi.mock('../../src/api', () => ({ default: api })) import { - EARLIER_PREVIEW, askAboutItPrefill, childrenForChat, earlierSlice, earlierSummary, formatElapsed, + EARLIER_PREVIEW, askAboutItPrefill, childrenForChat, clockRestartsAt, pendingTurnOutcome, earlierSlice, earlierSummary, formatElapsed, holderLine, isHonestTerminal, isLive, itemById, kindLabel, liveElapsedSeconds, liveItems, stageRows, stepsLine, workSignalFromItems, workStatusLabel, workTone, workView, } from '@/components/portal/portalWork' @@ -513,3 +513,23 @@ describe('ent#525 — the tab and the card are wired (source guards)', () => { expect((ws.match(/portalWorkStore\.handleWebSocketEvent\(data\)/g) || []).length).toBe(2) }) }) + +// #3114: a just-sent turn may be waiting for a pull worker. Before the agent +// streams anything the card says "Waiting for a slot", and the clock restarts +// when the turn is picked up, so "Working" counts working time only. +describe('pending turn on a pull agent (#3114)', () => { + it('reads waiting until the agent streams', () => { + expect(workStatusLabel({ outcome: pendingTurnOutcome(null, true) })).toBe('Waiting for a slot') + expect(workStatusLabel({ outcome: pendingTurnOutcome({ text: 'Reading a file' }, true) })).toBe('Working') + }) + it('reads working at once on a push agent', () => { + expect(workStatusLabel({ outcome: pendingTurnOutcome(null, false) })).toBe('Working') + expect(workStatusLabel({ outcome: pendingTurnOutcome(null, undefined) })).toBe('Working') + }) + it('restarts the clock only on queued -> running', () => { + expect(clockRestartsAt('queued', 'running')).toBe(true) + expect(clockRestartsAt('running', 'running')).toBe(false) + expect(clockRestartsAt(undefined, 'running')).toBe(false) + expect(clockRestartsAt('running', 'success')).toBe(false) + }) +}) diff --git a/tests/unit/test_1474_read_boundary_z.py b/tests/unit/test_1474_read_boundary_z.py index 20416f84a9..ccc70a4e10 100644 --- a/tests/unit/test_1474_read_boundary_z.py +++ b/tests/unit/test_1474_read_boundary_z.py @@ -97,6 +97,7 @@ def _make_exec_schema(conn: sqlite3.Connection) -> None: validation_execution_id TEXT, turn_integrity TEXT, queued_at TEXT, + lease_expires_at TEXT, source_channel TEXT, source_channel_chat_id TEXT, loop_id TEXT diff --git a/tests/unit/test_1766_pull_pilot_exclusive.py b/tests/unit/test_1766_pull_pilot_exclusive.py index 50eabb2df1..c3b1843f66 100644 --- a/tests/unit/test_1766_pull_pilot_exclusive.py +++ b/tests/unit/test_1766_pull_pilot_exclusive.py @@ -24,11 +24,10 @@ 2. **Consumer** — the backend never drains a pilot's queue, so the agent's worker pool is the sole claimant. One guard covers every drain path (release callback, 60s orphan sweep, ``drain_on_release``). - 3. **Carve-out** — interactive triggers are excluded and still take the - synchronous push path — temporarily, until #2842/#2843 land (Open - Question 7 is decided for the queue, #1989) — so human chat is not parked - behind N batch tasks and per-session ``--resume`` serialization is - untouched. + 3. **Carve-out** — the UI ``/chat`` trigger (``chat``) still takes the + synchronous push path. Every other interactive trigger is pulled on a + pilot (#3114); the claim orders them ahead of batch work (#2842) and runs + one turn per conversation at a time (#2843). Plus the inertness property the whole dark-ship rests on: with an empty allowlist (the default) every path is byte-for-byte unchanged. @@ -220,9 +219,17 @@ def test_pilot_does_not_own_a_trigger_the_reach_set_omits(self, pilot, monkeypat assert pull_owns_dispatch("alice", "schedule") is False assert pull_owns_dispatch("alice", "agent") is True - @pytest.mark.parametrize("trigger", ["manual", "user", "chat", "voip", "voice", None]) - def test_pilot_does_not_own_interactive_triggers(self, pilot, trigger): - """Human turns keep the synchronous path until #2842/#2843 land (#1989).""" + @pytest.mark.parametrize("trigger", ["manual", "user", "voip", "voice"]) + def test_pilot_owns_interactive_triggers(self, pilot, trigger): + """#3114: interactive turns reach a pilot through the durable queue.""" + from services.agent_service.pull_mode import pull_owns_dispatch + + assert pull_owns_dispatch("alice", trigger) is True + + @pytest.mark.parametrize("trigger", ["chat", None]) + def test_pilot_does_not_own_chat_or_unknown(self, pilot, trigger): + """The UI ``/chat`` path still pushes (#3114); an unclassified trigger + falls back to push.""" from services.agent_service.pull_mode import pull_owns_dispatch assert pull_owns_dispatch("alice", trigger) is False @@ -275,10 +282,12 @@ def test_pilot_autonomous_work_bypasses_admission( slot_service.acquire_slot.assert_not_awaited() backlog_service.enqueue.assert_awaited_once() - def test_pilot_interactive_work_still_pushes( + def test_pilot_interactive_work_is_queued( self, capacity, slot_service, backlog_service, pilot ): - """The carve-out, on the same agent: a human turn is admitted.""" + """#3114: a human turn on a pilot is queued like autonomous work, with + no ZADD, even with a slot free.""" + slot_service.acquire_slot = AsyncMock(return_value=True) result = asyncio.run( capacity.acquire( agent_name="alice", @@ -288,6 +297,23 @@ def test_pilot_interactive_work_still_pushes( overflow_payload=_payload("manual"), ) ) + assert result.state == "queued_persistent" + slot_service.acquire_slot.assert_not_awaited() + backlog_service.enqueue.assert_awaited_once() + + def test_pilot_chat_trigger_still_pushes( + self, capacity, slot_service, backlog_service, pilot + ): + """The remaining carve-out: trigger ``chat`` is admitted (#3114).""" + result = asyncio.run( + capacity.acquire( + agent_name="alice", + execution_id="exec-2b", + max_concurrent=3, + overflow_policy="queue_persistent", + overflow_payload=_payload("chat"), + ) + ) assert result.state == "admitted" slot_service.acquire_slot.assert_awaited_once() backlog_service.enqueue.assert_not_awaited() diff --git a/tests/unit/test_2048_pull_pilot_reach.py b/tests/unit/test_2048_pull_pilot_reach.py index 9e0cfaa1b8..c8f5351e32 100644 --- a/tests/unit/test_2048_pull_pilot_reach.py +++ b/tests/unit/test_2048_pull_pilot_reach.py @@ -210,10 +210,12 @@ def test_narrowing_is_a_runtime_no_op_on_the_only_producer_that_consults_it(pilo only stops the predicate from claiming reach it does not have. """ from services.task_execution_service import _AUTONOMOUS_TRIGGERS - from services.pull_pilot import pull_owns_dispatch + from services.pull_pilot import PULL_REACHABLE_NON_AUTONOMOUS, pull_owns_dispatch for trigger in ("self_task", "agent", "mcp", "manual", "event"): - before = trigger in _AUTONOMOUS_TRIGGERS # the pre-#2048 predicate body + # The pre-#2048 predicate body, widened by #3114 to the interactive + # triggers (`mcp`, `manual` here) a pilot now pulls. + before = trigger in _AUTONOMOUS_TRIGGERS or trigger in PULL_REACHABLE_NON_AUTONOMOUS assert pull_owns_dispatch("pilot-a", trigger) is before, trigger @@ -288,8 +290,9 @@ def test_a_reachable_trigger_is_not_reported(pilot, trigger): @pytest.mark.parametrize("trigger", ["manual", "mcp", "chat", "self_task", "voip"]) def test_an_interactive_trigger_is_not_reported(pilot, trigger): - """Interactive turns are excluded from pull until #2842/#2843 land (#1989), - so they are a known migration state, not a gap, and must not be logged as one.""" + """Interactive triggers are outside the autonomous set, so the stranded + diagnostic never names them: a pilot pulls them (#3114), except ``chat``, + which pushes by design.""" assert pilot.note_unreachable_pull_trigger("pilot-a", trigger) is False diff --git a/tests/unit/test_2317_pull_envelope_parity.py b/tests/unit/test_2317_pull_envelope_parity.py index b398d2e895..6688d0c838 100644 --- a/tests/unit/test_2317_pull_envelope_parity.py +++ b/tests/unit/test_2317_pull_envelope_parity.py @@ -129,6 +129,16 @@ def _import_backend_modules(): "self_task_activity_id", ) +# Keys only a pull pilot's row carries a value for (#3114): `execute_task` +# queues them and the claim envelope / claim-time prompt consume them. The +# push drain never sees a pilot's row (#1766), so it has no counterpart. +_PULL_ONLY_KEYS = ( + "persist_session", + "schedule_context", + "attempt", + "images", +) + # --------------------------------------------------------------------------- # The shared fixture — built by the REAL producer, not by hand @@ -303,15 +313,18 @@ def _applied_by_push(meta: Dict[str, Any], monkeypatch) -> Dict[str, Any]: # =========================================================================== class TestProducerKeyClassification: def test_producer_keys_are_exhaustively_classified(self, enqueued_metadata): - """Every key `enqueue` writes is either runtime-facing or backend-only. + """Every key `enqueue` writes is runtime-facing, backend-only, or + pull-only (#3114). A new key in `backlog_service.enqueue` reds this test until it is added - to one of the two tuples — which is the point at which someone has to + to one of the tuples — which is the point at which someone has to decide whether the pull envelope must carry it. Without this, a setting can be added to the producer and reach the push path only, silently (that is exactly how #2317 shipped). """ - assert set(enqueued_metadata) == set(_RUNTIME_KEYS) | set(_BACKEND_ONLY_KEYS) + assert set(enqueued_metadata) == ( + set(_RUNTIME_KEYS) | set(_BACKEND_ONLY_KEYS) | set(_PULL_ONLY_KEYS) + ) def test_envelope_override_keys_are_real_producer_keys(self, enqueued_metadata): """The envelope's source list names keys the producer actually writes — diff --git a/tests/unit/test_2391_scheduled_pull_reach.py b/tests/unit/test_2391_scheduled_pull_reach.py index 279b3a39bc..3ac9680acc 100644 --- a/tests/unit/test_2391_scheduled_pull_reach.py +++ b/tests/unit/test_2391_scheduled_pull_reach.py @@ -107,6 +107,7 @@ def _run( acquire_result=None, acquire_raises=None, dispatch_async=False, + capacity=None, **execute_kwargs, ): """Drive the real `execute_task` and return `(result, mocks)`. @@ -124,14 +125,17 @@ def _run( mock_db.get_execution.return_value = MagicMock(status="cancelled") mock_db.update_execution_status.return_value = True - mock_capacity = MagicMock() - if acquire_raises is not None: - mock_capacity.acquire = AsyncMock(side_effect=acquire_raises) - else: - mock_capacity.acquire = AsyncMock( - return_value=acquire_result or MagicMock(state="admitted") - ) - mock_capacity.release = AsyncMock() + # `capacity`: a real CapacityManager, for the end-to-end class (#3114). + mock_capacity = capacity + if mock_capacity is None: + mock_capacity = MagicMock() + if acquire_raises is not None: + mock_capacity.acquire = AsyncMock(side_effect=acquire_raises) + else: + mock_capacity.acquire = AsyncMock( + return_value=acquire_result or MagicMock(state="admitted") + ) + mock_capacity.release = AsyncMock() mock_circuit = MagicMock() mock_circuit.allow_request.return_value = True @@ -316,11 +320,18 @@ def test_nothing_autonomous_is_stranded_any_more(self, pilot): assert kw["overflow_policy"] == "queue_persistent", trigger assert kw["overflow_payload"] is not None, trigger - @pytest.mark.parametrize("trigger", ["manual", "mcp", "chat", "public", "voice"]) - def test_interactive_triggers_are_untouched(self, pilot, trigger): - """Until #2842/#2843 land (#1989), a human turn keeps the synchronous - push path and today's Redis session lock.""" + @pytest.mark.parametrize("trigger", ["manual", "mcp", "public", "voice"]) + def test_interactive_triggers_are_queued(self, pilot, trigger): + """#3114: an interactive turn on a pilot reaches the durable queue too; + its sync caller waits on the row via ``dispatch_and_await_terminal``.""" _, m = _run(triggered_by=trigger) + kw = _acquire_kwargs(m["capacity"]) + assert kw["overflow_policy"] == "queue_persistent", trigger + assert kw["overflow_payload"] is not None, trigger + + def test_chat_trigger_is_untouched(self, pilot): + """The UI ``/chat`` trigger keeps the push path (#3114).""" + _, m = _run(triggered_by="chat") assert _acquire_kwargs(m["capacity"])["overflow_policy"] == "reject" @@ -417,6 +428,7 @@ class _FakeQueueDb: def __init__(self): self.queued: dict[str, str] = {} + self.keys: dict[str, str | None] = {} def get_queued_count(self, agent_name): return 0 @@ -426,6 +438,7 @@ def get_max_backlog_depth(self, agent_name): def update_execution_to_queued(self, execution_id, metadata, queued_at, conversation_key=None): self.queued[execution_id] = metadata + self.keys[execution_id] = conversation_key return True @@ -552,6 +565,41 @@ def test_the_persisted_metadata_records_the_scheduled_trigger(self, pilot, monke assert json.loads(fake_db.queued["exec-e2e"])["triggered_by"] == "schedule" +class TestInteractiveTurnReachesTheQueue: + """#3114: the real `execute_task` -> real `CapacityManager.acquire` -> real + `BacklogService.enqueue` for an interactive trigger. The producer's + conversation key must land on the row, or the claim guard stops + serialising rooms, public links and channels.""" + + @pytest.mark.parametrize("trigger, key", [ + ("room", "room:r1"), ("public", "public:s1"), ("slack", "channel:c1"), + ]) + def test_turn_is_enqueued_under_its_conversation_key(self, pilot, monkeypatch, trigger, key): + from services import capacity_manager as cm_module + from services.backlog_service import BacklogService + + monkeypatch.setattr(cm_module.redis, "from_url", lambda *_a, **_kw: MagicMock()) + slots = AsyncMock() + slots.slots_prefix = "agent:slots:" + slots.acquire_slot = AsyncMock(return_value=True) + slots.register_on_release = lambda cb: None + capacity = cm_module.CapacityManager( + redis_url="redis://test", slot_service=slots, backlog_service=BacklogService() + ) + fake_db = _FakeQueueDb() + with ( + patch("database.db", fake_db), + patch("services.settings_service.clamp_to_ceiling", lambda v: v), + ): + result, m = _run(triggered_by=trigger, capacity=capacity, conversation_key=key) + + assert result.status == "queued" + assert fake_db.keys["exec-2391"] == key + assert json.loads(fake_db.queued["exec-2391"])["triggered_by"] == trigger + slots.acquire_slot.assert_not_awaited() + m["post"].assert_not_awaited() + + # --------------------------------------------------------------------------- # 7. Scheduler side — `queued` is not a terminal # --------------------------------------------------------------------------- diff --git a/tests/unit/test_2467_turn_integrity.py b/tests/unit/test_2467_turn_integrity.py index b21f7c0b23..b969add11a 100644 --- a/tests/unit/test_2467_turn_integrity.py +++ b/tests/unit/test_2467_turn_integrity.py @@ -505,6 +505,7 @@ def ops(self, tmp_path, monkeypatch): validation_execution_id TEXT, turn_integrity TEXT, queued_at TEXT, + lease_expires_at TEXT, source_channel TEXT, source_channel_chat_id TEXT, loop_id TEXT diff --git a/tests/unit/test_2661_task_timeout_receipt.py b/tests/unit/test_2661_task_timeout_receipt.py index dfcd54256d..c3908d5ec2 100644 --- a/tests/unit/test_2661_task_timeout_receipt.py +++ b/tests/unit/test_2661_task_timeout_receipt.py @@ -203,6 +203,17 @@ def test_long_poll_timeout_that_missed_a_failure_releases_the_claim(): isvc.complete.assert_not_called() +def test_backlog_row_keeps_the_pull_sinks_error_code(): + """#3114: a pull-sink terminal stores `[code] text`. The backlog path reads + the code back, so a pilot's `/task` answers what the push path answers.""" + out, _ = _run_backlog( + wait=None, row=_row("failed", error="[model_unsupported] model x refused"), + idem=MagicMock(), + ) + assert isinstance(out, ChatDispatchError) and out.status_code == 400 + assert out.detail == "model x refused" + + def test_a_vanished_row_releases_the_claim(): """Nothing is running under this key that anyone can find — hold nothing.""" idem = MagicMock() diff --git a/tests/unit/test_3114_pull_route_callers.py b/tests/unit/test_3114_pull_route_callers.py new file mode 100644 index 0000000000..0b7738d459 --- /dev/null +++ b/tests/unit/test_3114_pull_route_callers.py @@ -0,0 +1,594 @@ +"""Callers of the sync edge adapter on a pull pilot (#3114). + +Every sync ``execute_task`` caller that #3114 moved onto +``dispatch_and_await_terminal`` is driven here through its real function with +collaborators patched. Each test asserts what the caller hands the adapter +(``triggered_by``, ``conversation_key``, ``service``) and how it maps the +``TaskExecutionResult`` it gets back. + +The callers import the adapter lazily (``from services.task_execution_service +import dispatch_and_await_terminal`` inside the function), so it is patched on +``services.task_execution_service``. + +Also covered: the portal in-flight marker TTL and wait budgets grow by the +queue allowance on a pilot, ``CapacityManager.acquire`` forwards the #3114 +payload fields to the backlog, and the public + portal live-stream proxies +hold while the row is queued. +""" +from __future__ import annotations + +import asyncio +import json +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +_BACKEND_STR = str(Path(__file__).resolve().parent.parent.parent / "src" / "backend") +if _BACKEND_STR not in sys.path: + sys.path.insert(0, _BACKEND_STR) + +pytestmark = pytest.mark.unit + +AGENT = "alpha" +CAPACITY_ERR = "Agent at capacity (queued turn not claimed in 900s)" +TIMEOUT_ERR = "Execution timed out after 1020s waiting for the queued turn" + + +def _result(status="success", response="answer", error=None, eid="e1"): + from services.execution_envelope import TaskExecutionResult + + return TaskExecutionResult(execution_id=eid, status=status, response=response, + error=error) + + +@pytest.fixture +def adapter(monkeypatch): + """Patch the adapter where the lazy imports look it up.""" + import services.task_execution_service as tes + + mock = AsyncMock(return_value=_result()) + monkeypatch.setattr(tes, "dispatch_and_await_terminal", mock) + return mock + + +def _timeouts(monkeypatch, seconds: int): + """Every reader of the agent's execution timeout sees ``seconds``.""" + import database + from services import session_turn_service as sts + + for db in {id(database.db): database.db, id(sts.db): sts.db}.values(): + monkeypatch.setattr(db, "get_execution_timeout", lambda _n: seconds) + + +# --------------------------------------------------------------------------- +# 1. public_chat_service +# --------------------------------------------------------------------------- + + +def _public(monkeypatch): + from services import public_chat_service as pcs + + db = MagicMock() + db.get_access_policy.return_value = {} + db.count_recent_messages_by_ip.return_value = 0 + db.count_recent_messages_by_token.return_value = 0 + db.get_or_create_public_chat_session.return_value = SimpleNamespace(id="cs-7") + db.build_public_chat_context.return_value = "ctx" + db.get_public_chat_session.return_value = SimpleNamespace(message_count=2) + db.get_public_channel_model.return_value = None + monkeypatch.setattr(pcs, "db", db) + monkeypatch.setattr(pcs, "get_agent_container", lambda n: SimpleNamespace(status="running")) + monkeypatch.setattr(pcs, "build_public_channel_caller_prompt", lambda *a, **k: None) + return pcs, db + + +def _public_request(): + return SimpleNamespace(session_token=None, session_id="anon-1", files=None, + message="hi", async_mode=False) + + +@pytest.mark.asyncio +async def test_public_sync_turn_is_keyed_by_its_chat_session(monkeypatch, adapter): + pcs, db = _public(monkeypatch) + + out = await pcs.run_public_chat({"id": "link-1", "agent_name": AGENT}, + _public_request(), "1.2.3.4") + + kw = adapter.await_args.kwargs + assert kw["triggered_by"] == "public" + assert kw["conversation_key"] == "public:cs-7" + assert out.response == "answer" + db.add_public_chat_message.assert_called_with( + session_id="cs-7", role="assistant", content="answer", cost=None, + sender_email=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error,status", [(CAPACITY_ERR, 429), (TIMEOUT_ERR, 504), + ("boom", 502)]) +async def test_public_sync_maps_the_adapter_failure(monkeypatch, adapter, error, status): + pcs, _ = _public(monkeypatch) + adapter.return_value = _result(status="failed", response="", error=error) + + with pytest.raises(pcs.PublicChatError) as exc: + await pcs.run_public_chat({"id": "link-1", "agent_name": AGENT}, + _public_request(), "1.2.3.4") + assert exc.value.status_code == status + + +@pytest.mark.asyncio +async def test_public_background_turn_is_keyed_by_its_chat_session(monkeypatch, adapter): + pcs, db = _public(monkeypatch) + + await pcs._execute_public_chat_background( + agent_name=AGENT, context_prompt="ctx", source_email="anonymous (x)", + execution_id="e1", chat_session_id="cs-9", session_identifier="anon", + identifier_type="anonymous") + + kw = adapter.await_args.kwargs + assert (kw["triggered_by"], kw["conversation_key"], kw["execution_id"]) == ( + "public", "public:cs-9", "e1") + assert db.add_public_chat_message.call_args.kwargs["content"] == "answer" + + +@pytest.mark.asyncio +async def test_public_background_posts_nothing_on_capacity(monkeypatch, adapter): + pcs, db = _public(monkeypatch) + adapter.return_value = _result(status="failed", response="", error=CAPACITY_ERR) + + await pcs._execute_public_chat_background( + agent_name=AGENT, context_prompt="ctx", source_email="anonymous (x)", + execution_id="e1", chat_session_id="cs-9", session_identifier="anon", + identifier_type="anonymous") + db.add_public_chat_message.assert_not_called() + + +# --------------------------------------------------------------------------- +# 2. shared_sessions room wake +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_room_wake_is_keyed_by_the_room(monkeypatch, adapter): + from shared_sessions import service + + room = "room-1" + monkeypatch.setattr(service.db, "get_participant", + lambda *a, **k: {"last_read_seq": 0, "cached_session_id": "s1"}) + monkeypatch.setattr(service.db, "get_room", lambda *a, **k: { + "id": room, "name": "Room", "status": "open", "topic": None}) + monkeypatch.setattr(service.db, "get_messages", lambda *a, **k: [ + {"seq": 1, "sender_kind": "user", "sender_identity": "c@example.com", + "content": f"@{AGENT} hi", "kind": "message"}]) + monkeypatch.setattr(service.db, "list_participants", lambda _r: [ + {"kind": "agent", "identity": AGENT, "left_at": None}]) + monkeypatch.setattr(service.db, "advance_read_cursor", lambda *a, **k: None) + for name in ("_mark_agent_working", "_clear_agent_working", "_broadcast", "_post_system"): + monkeypatch.setattr(service, name, lambda *a, **k: None) + monkeypatch.setattr(service, "post_message", AsyncMock(return_value={})) + + await service._wake_agent(SimpleNamespace(email="c@example.com"), room, AGENT, 0) + + kw = adapter.await_args.kwargs + assert kw["triggered_by"] == "room" + assert kw["conversation_key"] == f"room:{room}" + assert kw["resume_session_id"] == "s1" and kw["persist_session"] is True + + +# --------------------------------------------------------------------------- +# 3. routers/paid.py +# --------------------------------------------------------------------------- + + +async def _drive_paid(monkeypatch, adapter, session_id, result): + import routers.paid as paid + + config = SimpleNamespace(enabled=True, nvm_environment="testnet", credits_per_request=1) + db = MagicMock() + db.get_nevermined_config_with_key.return_value = {"config": config, "nvm_api_key": "k"} + payment = MagicMock( + verify_payment=AsyncMock(return_value=SimpleNamespace( + success=True, payer="0xp", agent_request_id="r1", error=None)), + settle_payment_once=AsyncMock(return_value=SimpleNamespace( + success=True, tx_hash="0xt", remaining_balance=1, error=None)), + ) + idem = MagicMock() + idem.begin.return_value = MagicMock(replay=False) + task_service = object() + monkeypatch.setattr(paid, "NEVERMINED_AVAILABLE", True) + monkeypatch.setattr(paid, "db", db) + monkeypatch.setattr(paid, "idempotency_service", idem) + monkeypatch.setattr(paid, "get_nevermined_payment_service", lambda: payment) + monkeypatch.setattr(paid, "get_task_execution_service", lambda: task_service) + monkeypatch.setattr(paid, "build_public_channel_caller_prompt", lambda *a, **k: None) + adapter.return_value = result + + resp = await paid.paid_chat( + AGENT, SimpleNamespace(message="hi", session_id=session_id), + SimpleNamespace(headers={"payment-signature": "tok"}, base_url="http://localhost/")) + body = json.loads(bytes(resp.body).decode()) if hasattr(resp, "body") else resp + return adapter.await_args.kwargs, body, task_service, payment + + +@pytest.mark.asyncio +async def test_paid_resumed_session_is_its_conversation_key(monkeypatch, adapter): + kw, _, task_service, _ = await _drive_paid(monkeypatch, adapter, "sess-3", _result()) + assert kw["triggered_by"] == "paid" + assert kw["conversation_key"] == "paid:sess-3" + assert kw["resume_session_id"] == "sess-3" + assert kw["service"] is task_service + + +@pytest.mark.asyncio +async def test_paid_without_session_has_no_key(monkeypatch, adapter): + kw, _, _, _ = await _drive_paid(monkeypatch, adapter, None, _result()) + assert kw["conversation_key"] is None + + +@pytest.mark.asyncio +async def test_paid_cancelled_returns_partial_and_does_not_settle(monkeypatch, adapter): + _, body, _, payment = await _drive_paid( + monkeypatch, adapter, "sess-3", + _result(status="cancelled", response="partial work", eid="e9")) + assert body["status"] == "cancelled" + assert body["response"] == "partial work" + assert body["execution_id"] == "e9" + assert body["payment"]["settled"] is False + payment.settle_payment_once.assert_not_called() + + +# --------------------------------------------------------------------------- +# 4. mcp_auth_service.dispatch_chat +# --------------------------------------------------------------------------- + + +def _mcp(monkeypatch): + from services import mcp_auth_service as mas + + monkeypatch.setattr(mas, "assert_email_may_reach_agent", lambda *a: {}) + monkeypatch.setattr(mas, "db", MagicMock(get_public_channel_model=lambda a: None)) + return mas + + +@pytest.mark.asyncio +async def test_mcp_inline_chat_has_no_conversation_key(monkeypatch, adapter): + mas = _mcp(monkeypatch) + + out = await mas.dispatch_chat("User@Example.com", AGENT, "hi") + + kw = adapter.await_args.kwargs + assert kw["triggered_by"] == mas.CHAT_TRIGGERED_BY == "mcp" + assert "conversation_key" not in kw + assert kw["source_user_email"] == "user@example.com" + assert out.response == "answer" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error,status", [(CAPACITY_ERR, 429), (TIMEOUT_ERR, 504), + ("boom", 502)]) +async def test_mcp_inline_chat_maps_the_adapter_failure(monkeypatch, adapter, error, status): + from fastapi import HTTPException + + mas = _mcp(monkeypatch) + adapter.return_value = _result(status="failed", response="", error=error) + with pytest.raises(HTTPException) as exc: + await mas.dispatch_chat("u@example.com", AGENT, "hi") + assert exc.value.status_code == status + + +# --------------------------------------------------------------------------- +# 5. validation_service +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_validation_goes_through_the_adapter_with_its_own_service(monkeypatch, adapter): + from services import validation_service as vs + + db = MagicMock() + db.create_validation_execution.return_value = SimpleNamespace(id="v1") + monkeypatch.setattr(vs, "db", db) + monkeypatch.setattr(vs, "get_referee", lambda: None) + svc = MagicMock() + adapter.return_value = _result(response='{"status": "pass", "summary": "ok", "items": []}') + + await vs.ValidationService(svc).validate_execution( + execution_id="e1", agent_name=AGENT, schedule_id="s1", + original_message="do it", execution_response="done") + + kw = adapter.await_args.kwargs + assert kw["service"] is svc + assert kw["triggered_by"] == "validation" + assert kw["execution_id"] == "v1" + assert "conversation_key" not in kw + svc.execute_task.assert_not_called() + + +# --------------------------------------------------------------------------- +# 6. routers/internal.py sync run-now +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_internal_sync_run_goes_through_the_adapter(monkeypatch, adapter): + import routers.internal as internal + import services.cleanup_service as cleanup_service + from models import InternalTaskExecutionRequest + + svc = MagicMock() + monkeypatch.setattr(cleanup_service, "is_startup_recovery_complete", lambda: True) + monkeypatch.setattr(internal, "get_task_execution_service", lambda: svc) + monkeypatch.setattr(internal.idempotency_service, "begin", + lambda *a, **k: SimpleNamespace(replay=False)) + for name in ("attach_execution", "complete", "fail"): + monkeypatch.setattr(internal.idempotency_service, name, lambda *a, **k: None) + monkeypatch.setattr(internal.platform_audit_service, "log", AsyncMock()) + + out = await internal.execute_task_internal( + InternalTaskExecutionRequest(agent_name=AGENT, message="run", triggered_by="manual", + execution_id="e1", schedule_id="sch-1", attempt=2), + idempotency_key=None) + + kw = adapter.await_args.kwargs + assert kw["service"] is svc + assert kw["triggered_by"] == "manual" + assert kw["execution_id"] == "e1" and kw["attempt"] == 2 + assert "conversation_key" not in kw + assert (out["status"], out["response"]) == ("success", "answer") + svc.execute_task.assert_not_called() + + +# --------------------------------------------------------------------------- +# 7. client_portal/service.py — marker TTL and wait budgets +# --------------------------------------------------------------------------- + +TURN = 600 + + +class _Stop(Exception): + pass + + +def _allowance_expected(pilot: bool) -> int: + return 2 * TURN if pilot else 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pilot", [True, False]) +async def test_portal_sync_marker_ttl_adds_the_allowance(monkeypatch, pilot): + from client_portal import service + + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT if pilot else "") + _timeouts(monkeypatch, TURN) + db = MagicMock() + db.get_cached_claude_session_id.return_value = None + db.get_portal_thread_window.return_value = SimpleNamespace(rows=[]) + monkeypatch.setattr(service, "db", db) + monkeypatch.setattr(service, "agent_on_roster", lambda *a: True) + monkeypatch.setattr(service, "_availability_allows_turn", lambda a: True) + monkeypatch.setattr(service, "resolve_turn_model", lambda *a: None) + monkeypatch.setattr(service, "_resolve_session_id", lambda *a, **k: "s1") + monkeypatch.setattr(service, "_refuse_turn_during_voice_call", lambda *a, **k: None) + monkeypatch.setattr(service, "_title_plan", lambda *a: None) + monkeypatch.setattr(service, "_persist_user_turn", lambda *a, **k: None) + monkeypatch.setattr(service, "collect_inbox_context", AsyncMock(return_value=("", []))) + monkeypatch.setattr(service, "_build_portal_system_prompt", lambda *a: None) + monkeypatch.setattr(service, "_precreate_sync_execution", lambda *a, **k: "e1") + marks = [] + monkeypatch.setattr(service, "mark_turn_inflight", lambda *a, **k: marks.append(a)) + monkeypatch.setattr(service, "_run_sync_turn_and_clear_marker", + AsyncMock(side_effect=_Stop())) + + with pytest.raises(_Stop): + await service.portal_chat(AGENT, "hi", "c@example.com", availability="running", + turn_timeout_seconds=TURN) + assert marks == [("s1", "e1", TURN + 60 + _allowance_expected(pilot))] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pilot", [True, False]) +async def test_portal_streaming_budget_adds_the_allowance(monkeypatch, pilot): + import database + from client_portal import service + + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT if pilot else "") + _timeouts(monkeypatch, TURN) + monkeypatch.setattr(service, "agent_on_roster", lambda *a: True) + monkeypatch.setattr(service, "_agent_availability", AsyncMock(return_value="running")) + monkeypatch.setattr(service, "_availability_allows_turn", lambda a: True) + monkeypatch.setattr(service, "resolve_turn_model", lambda *a: None) + monkeypatch.setattr(service, "_resolve_session_id", lambda *a, **k: "s1") + monkeypatch.setattr(service, "_refuse_turn_during_voice_call", lambda *a, **k: None) + monkeypatch.setattr(database.db, "get_agent_subscription_id", lambda a: None) + monkeypatch.setattr(database.db, "create_task_execution", + lambda **k: SimpleNamespace(id="e1")) + marks = [] + monkeypatch.setattr(service, "mark_turn_inflight", lambda *a, **k: marks.append(k)) + for name in ("clear_turn_outcome", "clear_turn_inflight", "record_turn_outcome"): + monkeypatch.setattr(service, name, lambda *a, **k: None) + monkeypatch.setattr(service, "portal_chat", AsyncMock(return_value={})) + + out = await service.start_portal_turn(AGENT, "hi", "c@example.com") + await asyncio.gather(*list(service._INFLIGHT_TURNS)) + + expected = service.portal_max_turn_seconds(TURN) + _allowance_expected(pilot) + assert out["wait_budget_seconds"] == expected + assert marks == [{"ttl_seconds": expected}] + + +@pytest.mark.parametrize("pilot", [True, False]) +def test_portal_history_budget_adds_the_allowance(monkeypatch, pilot): + import redis_breaker_util + from client_portal import service + + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT if pilot else "") + _timeouts(monkeypatch, TURN) + db = MagicMock() + db.get_latest_portal_session_id.return_value = "s1" + db.get_portal_thread_window.return_value = SimpleNamespace(rows=[], truncated=False) + monkeypatch.setattr(service, "db", db) + monkeypatch.setattr(service, "agent_on_roster", lambda *a: True) + monkeypatch.setattr(service, "_attach_own_ratings", lambda *a, **k: None) + monkeypatch.setattr(service, "get_turn_inflight", lambda s: "e1") + monkeypatch.setattr(service, "get_turn_outcome", lambda s: None) + # TTL -1: no expiry on the marker → the full per-agent budget. + monkeypatch.setattr(redis_breaker_util, "get_breaker_redis", + lambda: SimpleNamespace(ttl=lambda k: -1)) + + out = service.get_history(AGENT, "c@example.com") + assert out["in_flight_wait_budget_seconds"] == ( + service.portal_max_turn_seconds(TURN) + _allowance_expected(pilot)) + + +# --------------------------------------------------------------------------- +# 8. CapacityManager.acquire → backlog_service.enqueue +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_acquire_forwards_the_turn_fields_to_enqueue(monkeypatch): + from services import capacity_manager as cm + + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT) + monkeypatch.setattr(cm.redis, "from_url", + lambda *a, **k: SimpleNamespace(set=lambda *a, **k: True)) + slots = MagicMock(acquire_slot=AsyncMock(return_value=True), + register_on_release=lambda cb: None) + backlog = MagicMock(enqueue=AsyncMock(return_value=True)) + manager = cm.CapacityManager(redis_url="redis://test", slot_service=slots, + backlog_service=backlog) + images = [{"media_type": "image/png", "data": "AAAA"}] + + out = await manager.acquire( + agent_name=AGENT, execution_id="e1", max_concurrent=2, + overflow_policy="queue_persistent", + overflow_payload=cm.PersistentTaskPayload( + request=MagicMock(), effective_timeout=900, user_id=1, user_email="u@x", + subscription_id=None, x_source_agent=None, triggered_by="public", + collaboration_activity_id=None, conversation_key="public:cs-1", + persist_session=True, schedule_context={"name": "nightly"}, attempt=3, + images=images)) + + assert out.state == "queued_persistent" + slots.acquire_slot.assert_not_called() + kw = backlog.enqueue.await_args.kwargs + assert kw["conversation_key"] == "public:cs-1" + assert kw["persist_session"] is True + assert kw["schedule_context"] == {"name": "nightly"} + assert kw["attempt"] == 3 + assert kw["images"] == images + + +# --------------------------------------------------------------------------- +# 9. public + portal live-stream proxies hold while queued +# (fakes copied from test_3114_pull_route_interactive.py) +# --------------------------------------------------------------------------- + + +class _FakeResponse: + def __init__(self, status, chunks=()): + self.status_code = status + self._chunks = chunks + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def aiter_text(self): + for c in self._chunks: + yield c + + +class _FakeAgentClient: + def __init__(self, responses): + self._responses = list(responses) + self.calls = 0 + + def __call__(self, *a, **k): + return self + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + def stream(self, method, url): + self.calls += 1 + return self._responses.pop(0) + + +def _stream_env(monkeypatch, router_mod, pilot: bool): + """A mutable execution row every reader sees, plus a fake agent client.""" + import database + from services import sync_waiter + + row = SimpleNamespace(status="queued", agent_name=AGENT) + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT if pilot else "") + _timeouts(monkeypatch, TURN) + monkeypatch.setattr(database.db, "get_execution", lambda eid: row) + # "running" here stands for a claimed row (lease set). + monkeypatch.setattr(database.db, "execution_awaits_claim", lambda eid: row.status == "queued") + monkeypatch.setattr(sync_waiter, "STREAM_QUEUED_POLL_INTERVAL", 0.01) + monkeypatch.setattr(sync_waiter, "STREAM_ATTACH_RETRY_INTERVAL", 0.01) + monkeypatch.setattr(router_mod, "get_agent_container", + lambda n: SimpleNamespace(status="running")) + client = _FakeAgentClient([_FakeResponse(404), _FakeResponse(200, ["data: x\n\n"])]) + monkeypatch.setattr(router_mod, "agent_httpx_client", client) + return row, client + + +async def _public_stream(monkeypatch, pilot): + from routers import public + + row, client = _stream_env(monkeypatch, public, pilot) + monkeypatch.setattr(public.db, "get_execution", lambda eid: row) + monkeypatch.setattr(public, "check_public_link_rate_limit", lambda ip: None) + monkeypatch.setattr(public, "_validate_public_link", lambda t: {"agent_name": AGENT}) + resp = await public.public_stream_execution( + "tok", "e1", SimpleNamespace(client=SimpleNamespace(host="1.2.3.4"), headers={})) + return row, client, resp.body_iterator + + +async def _portal_stream(monkeypatch, pilot): + from client_portal import router as portal_router + from client_portal import service + + row, client = _stream_env(monkeypatch, portal_router, pilot) + monkeypatch.setattr(service, "agent_on_roster", lambda *a: True) + monkeypatch.setattr(service, "execution_belongs_to_caller", lambda *a: True) + monkeypatch.setattr(service, "get_turn_inflight_matches", lambda eid: True) + monkeypatch.setattr(portal_router, "_STREAM_ATTACH_POLL_S", 0.01) + resp = await portal_router.portal_stream_execution( + AGENT, "e1", SimpleNamespace(email="c@example.com", is_platform=False)) + return row, client, resp.body_iterator + + +@pytest.mark.asyncio +@pytest.mark.parametrize("open_stream", [_public_stream, _portal_stream], + ids=["public", "portal"]) +async def test_stream_holds_while_queued_then_attaches(monkeypatch, open_stream): + row, client, body = await open_stream(monkeypatch, pilot=True) + + assert await body.__anext__() == ": queued\n\n" + assert await body.__anext__() == ": queued\n\n" # still queued: still holding + assert client.calls == 0 + row.status = "running" # claimed + rest = [c async for c in body] + assert rest[-1] == "data: x\n\n" # 404 before registration retried + assert client.calls == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("open_stream", [_public_stream, _portal_stream], + ids=["public", "portal"]) +async def test_stream_off_pilot_does_not_hold(monkeypatch, open_stream): + _, client, body = await open_stream(monkeypatch, pilot=False) + chunks = [c async for c in body] + assert not any(c.startswith(":") for c in chunks) + assert client.calls >= 1 diff --git a/tests/unit/test_3114_pull_route_interactive.py b/tests/unit/test_3114_pull_route_interactive.py new file mode 100644 index 0000000000..87ac747cd5 --- /dev/null +++ b/tests/unit/test_3114_pull_route_interactive.py @@ -0,0 +1,893 @@ +"""Interactive producers on the durable queue (#3114). + +What is pinned here: + +* ``pull_owns_dispatch`` covers every interactive trigger but ``chat``, plus + ``validation``, on a pull pilot. + +* ``execute_task``'s queue payload carries what the push payload sends beside + the request (conversation key, persist_session, images, schedule context, + attempt), and ``BacklogService.enqueue`` stores it; an explicit conversation + key wins on a pilot and is dropped elsewhere. +* The claim envelope hands ``persist_session``/``images`` to the worker, and the + claim-time prompt gets the push path's provenance/timeout/schedule/attempt. +* The worker runs the turn with them. +* The pull sink persists compact events and wakes an in-process sync waiter. +* ``result_from_execution_row`` rebuilds a push-shaped result. +* ``dispatch_and_await_terminal`` waits at most one agent timeout for a claim on + interactive triggers, stores FAILED/CAPACITY, and carries on when the + cancel loses to a claim. +* ``run_resumable_turn`` goes through the adapter with a per-conversation key, + and the Session lock TTL grows by the queue allowance on a pilot. +* Live-stream proxies hold while a pilot's turn is queued. +* Concurrency: a burst of queued turns of one conversation is claimed one at a + time (claim guard), and two Session-tab turns of one session never overlap + (the ResumeLock is held across the queued wait). + +Runs against the real db layer through db_harness (SQLite; Postgres too when +TEST_POSTGRES_URL is set). +""" +from __future__ import annotations + +import asyncio +import inspect +import json +import re +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +_BACKEND_STR = str(Path(__file__).resolve().parent.parent.parent / "src" / "backend") +if _BACKEND_STR not in sys.path: + sys.path.insert(0, _BACKEND_STR) + +from db_harness import db_backend, run as _hrun, scalar as _scalar # noqa: E402,F401 + +pytestmark = pytest.mark.unit + +_STUBBED_MODULE_NAMES = [ + "db.connection", + "db.schedules", + "db.agent_settings.resources", + "database", +] + +AGENT = "alpha" + + +@pytest.fixture(autouse=True) +def _restore_sys_modules(): + saved = {name: sys.modules.get(name) for name in _STUBBED_MODULE_NAMES} + try: + yield + finally: + for name, value in saved.items(): + if value is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = value + + +@pytest.fixture +def tmp_db(db_backend): + def _evict(): + for mod in ("db.connection", "db.schedules", "db.agent_settings.resources", "database"): + sys.modules.pop(mod, None) + + _evict() + try: + yield db_backend + finally: + _evict() + + +@pytest.fixture +def seed_agent(tmp_db): + def _seed(name: str = AGENT, timeout: int = 900): + _hrun( + "INSERT INTO agent_ownership (agent_name, owner_id, " + " execution_timeout_seconds, created_at) " + "VALUES (:n, 1, :t, '2026-01-01T00:00:00Z')", + n=name, t=timeout, + ) + + return _seed + + +def _row(eid: str, *, status: str = "queued", trigger: str = "session", meta=None, + error=None, response=None, session=None, compact=None, email=None): + _hrun( + "INSERT INTO schedule_executions " + "(id, schedule_id, agent_name, status, started_at, queued_at, message, " + " triggered_by, backlog_metadata, error, response, claude_session_id, " + " compact_metadata, source_user_email) " + "VALUES (:id, '__manual__', :a, :s, '2026-09-30T09:00:00Z', " + " CASE WHEN :s = 'queued' THEN '2026-09-30T09:00:00Z' END, 'hello', :t, " + " :meta, :err, :resp, :sess, :cm, :em)", + id=eid, a=AGENT, s=status, t=trigger, + meta=json.dumps(meta) if meta is not None else None, + err=error, resp=response, sess=session, cm=compact, em=email, + ) + return eid + + +def _db(): + from database import db + + return db + + +_IMG = [{"media_type": "image/png", "data": "iVBORw0KGgo+AIzaQQQQ"}] + + +# --------------------------------------------------------------------------- +# Predicate +# --------------------------------------------------------------------------- + + +def test_pilot_pulls_interactive_triggers_except_chat(monkeypatch): + from services import pull_pilot as pp + + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT) + assert pp.PULL_REACHABLE_NON_AUTONOMOUS == (pp.INTERACTIVE_TRIGGERS - {"chat"}) | {"validation"} + for trigger in pp.PULL_REACHABLE_NON_AUTONOMOUS: + assert pp.pull_owns_dispatch(AGENT, trigger) is True, trigger + assert pp.pull_owns_dispatch("not-a-pilot", trigger) is False, trigger + assert pp.pull_owns_dispatch(AGENT, "chat") is False + assert pp.pull_owns_dispatch(AGENT, None) is False + + +# --------------------------------------------------------------------------- +# Producer: execute_task → payload → enqueue +# --------------------------------------------------------------------------- + + +def test_queue_payload_carries_the_turn_fields(): + from services import task_execution_service as tes + + with patch.object(tes, "pull_owns_dispatch", return_value=True): + payload = tes.build_pull_queue_payload( + agent_name=AGENT, triggered_by="session", execution_id="e1", + message="hi", model="opus", allowed_tools=None, system_prompt=None, + timeout_seconds=600, resume_session_id="uuid-1", subscription_id=None, + source_user_id=7, source_user_email="u@example.com", + source_agent_name=None, slot_already_held=False, + conversation_key="session:k1", persist_session=True, + schedule_context={"name": "nightly"}, attempt=2, images=_IMG, + ) + assert payload.conversation_key == "session:k1" + assert payload.persist_session is True + assert payload.images == _IMG + assert payload.schedule_context == {"name": "nightly"} + assert payload.attempt == 2 + assert payload.request.resume_session_id == "uuid-1" + + +def test_capacity_manager_forwards_every_payload_field_to_enqueue(): + """Structural: a field added to PersistentTaskPayload must reach enqueue.""" + from dataclasses import fields + + from services import capacity_manager as cm + + src = inspect.getsource(cm.CapacityManager.acquire) + for f in fields(cm.PersistentTaskPayload): + if f.name in ("request", "effective_timeout"): + continue + assert re.search(rf"{f.name}=overflow_payload\.{f.name}\b", src), f.name + + +def _enqueue(eid="e1", **extra): + from models import ParallelTaskRequest + from services.backlog_service import BacklogService + + req = ParallelTaskRequest(message="hi", async_mode=True, chat_session_id="chat-row-1") + return asyncio.run(BacklogService().enqueue( + agent_name=AGENT, execution_id=eid, request=req, effective_timeout=600, + user_id=None, user_email=None, subscription_id=None, x_source_agent=None, + triggered_by="session", collaboration_activity_id=None, **extra, + )) + + +def test_enqueue_stores_turn_fields_and_explicit_key_wins_on_pilot(seed_agent, monkeypatch): + seed_agent() + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT) + _row("e1", status="running") + + assert _enqueue(conversation_key="session:k1", persist_session=True, + schedule_context={"name": "n"}, attempt=3, images=_IMG) is True + + assert _scalar("SELECT conversation_key FROM schedule_executions WHERE id='e1'") == "session:k1" + meta = json.loads(_scalar("SELECT backlog_metadata FROM schedule_executions WHERE id='e1'")) + assert meta["persist_session"] is True + assert meta["images"] == _IMG + assert meta["schedule_context"] == {"name": "n"} + assert meta["attempt"] == 3 + + +def test_enqueue_falls_back_to_chat_session_and_ignores_key_off_pilot(seed_agent, monkeypatch): + seed_agent() + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT) + _row("e1", status="running") + assert _enqueue() is True + assert _scalar("SELECT conversation_key FROM schedule_executions WHERE id='e1'") == "chat-row-1" + + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", "") + _hrun("UPDATE schedule_executions SET status='running', conversation_key=NULL WHERE id='e1'") + assert _enqueue(conversation_key="session:k1") is True + assert _scalar("SELECT conversation_key FROM schedule_executions WHERE id='e1'") is None + + +# --------------------------------------------------------------------------- +# Claim envelope + prompt +# --------------------------------------------------------------------------- + + +def test_claim_envelope_carries_persist_images_and_push_prompt_context(seed_agent): + seed_agent(timeout=900) + _row("e1", meta={ + "message": "hello", "timeout_seconds": 600, "persist_session": True, + "images": _IMG, "schedule_context": {"name": "nightly", "cron": "0 3 * * *"}, + "attempt": 2, "resume_session_id": "uuid-1", + }, email="u@example.com") + from services import pull_coordination_service as pcs + + seen = {} + + def _compose(**kw): + seen["ctx"] = kw["execution_context"] + return "PROMPT" + + with patch.object(pcs, "_resolve_agent_runtime", return_value="claude-code"), \ + patch.object(pcs, "compose_system_prompt", side_effect=_compose), \ + patch.object(pcs, "record_worker_poll"): + claim = pcs.claim_next_task(AGENT, "w1") + + payload = claim["envelope"]["payload"] + assert payload["persist_session"] is True + assert payload["images"] == _IMG + assert payload["session_id"] == "uuid-1" + ctx = seen["ctx"] + assert ctx.source_user_email == "u@example.com" + assert ctx.timeout_seconds == 600 + assert ctx.schedule_name == "nightly" + assert ctx.schedule_cron == "0 3 * * *" + assert ctx.attempt == 2 + + +def test_claim_envelope_omits_the_fields_for_older_rows(seed_agent): + seed_agent() + _row("e1", meta={"message": "hello"}) + from services import pull_coordination_service as pcs + + with patch.object(pcs, "record_worker_poll"): + payload = pcs.claim_next_task(AGENT, "w1")["envelope"]["payload"] + assert "persist_session" not in payload + assert "images" not in payload + + +# --------------------------------------------------------------------------- +# Worker +# --------------------------------------------------------------------------- + + +def _run_worker(monkeypatch, payload): + from agent_server.services import pull_worker as pw + + md = SimpleNamespace(model_dump=lambda: {"cost_usd": 0.0}) + runtime = SimpleNamespace(execute_headless=AsyncMock(return_value=("ok", [], md, "s"))) + stub = types.ModuleType("agent_server.services.runtime_adapter") + stub.get_runtime = lambda: runtime + monkeypatch.setitem(sys.modules, "agent_server.services.runtime_adapter", stub) + monkeypatch.setattr(pw, "_deliver_result", AsyncMock(return_value=True)) + monkeypatch.setattr(pw, "_persist_pull_result", lambda *a, **k: None) + monkeypatch.setattr(pw, "_delete_pull_result", lambda *a, **k: None) + claim = {"execution_id": "e1", "claim_token": "t", "envelope": {"payload": payload}} + asyncio.run(pw._run_and_report(claim, MagicMock(), "http://b", "k", AGENT)) + return runtime.execute_headless.call_args.kwargs + + +def test_worker_takes_persist_session_and_images_from_the_payload(monkeypatch): + kw = _run_worker(monkeypatch, {"message": "m", "session_id": None, + "persist_session": True, "images": _IMG}) + assert kw["persist_session"] is True # cold Session turn still writes the JSONL + assert kw["images"] == _IMG + + +def test_worker_falls_back_to_session_presence(monkeypatch): + kw = _run_worker(monkeypatch, {"message": "m", "session_id": "uuid-1"}) + assert kw["persist_session"] is True + assert kw["images"] is None + kw = _run_worker(monkeypatch, {"message": "m", "session_id": "uuid-1", + "persist_session": False}) + assert kw["persist_session"] is False + + +# --------------------------------------------------------------------------- +# Pull sink + row → result +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_sink_persists_compact_events_and_wakes_the_waiter(seed_agent, monkeypatch): + seed_agent() + _row("e1") + import importlib + + # sys.modules entries: what apply_task_result's lazy import resolves. + pcs = importlib.import_module("services.pull_coordination_service") + sync_waiter = importlib.import_module("services.sync_waiter") + tes = importlib.import_module("services.task_execution_service") + + monkeypatch.setattr(sync_waiter, "SYNC_WAITER_POLL_INTERVAL", 60.0) + row = _db().claim_next_queued(AGENT, worker_id="w1", lease_seconds=900) + waiter = asyncio.create_task(sync_waiter.wait_for_sync_terminal("e1", 5)) + await asyncio.sleep(0) + + events = [{"trigger": "auto", "pre_tokens": 150000}] + outcome = pcs.apply_task_result( + "e1", row["claim_token"], status="success", content="done", + session_id="uuid-9", metadata={"compact_events": events}, + ) + assert outcome.kind == "applied" + assert await asyncio.wait_for(waiter, 1) == {"result": None, "chat_session_id": None} + + result = tes.result_from_execution_row("e1") + assert result.response == "done" + assert result.session_id == "uuid-9" + assert result.raw_response == {"metadata": {"compact_events": events}} + + +def test_row_result_maps_like_push(seed_agent): + seed_agent() + from services import task_execution_service as tes + from services.execution_envelope import TaskExecutionErrorCode + + _row("f1", status="failed", error="[timeout] took too long", response="partial", + session="dispatched") + r = tes.result_from_execution_row("f1") + assert r.error_code == TaskExecutionErrorCode.TIMEOUT + assert r.error == "took too long" # channel adapters post this text to people + assert r.response == "" + assert r.session_id is None + assert r.raw_response == {} + + _row("f2", status="failed", error="[not_a_code] x", session="dispatched_async") + r = tes.result_from_execution_row("f2") + assert r.error_code is None and r.session_id is None + assert r.error == "[not_a_code] x" + + # #679: a cancelled turn's partial reply is returned. + _row("c1", status="cancelled", response="half an answer") + assert tes.result_from_execution_row("c1").response == "half an answer" + + _row("s1", status="success", response="answer", session="uuid-1") + r = tes.result_from_execution_row("s1") + assert (r.response, r.session_id, r.error_code) == ("answer", "uuid-1", None) + + +# --------------------------------------------------------------------------- +# dispatch_and_await_terminal — claim phase +# --------------------------------------------------------------------------- + + +def _queued_service(eid): + from services.execution_envelope import TaskExecutionResult + + svc = MagicMock() + svc.execute_task = AsyncMock(return_value=TaskExecutionResult( + execution_id=eid, status="queued", response="")) + return svc + + +def _terminal_hooks(monkeypatch, tes): + """Record the activity close and terminal event a CAS winner owes.""" + calls = [] + monkeypatch.setattr(tes.activity_service, "spawn_close_execution_activity", + lambda eid, st, **k: calls.append(("close", eid, getattr(st, "value", st)))) + monkeypatch.setattr(tes.event_dispatch_service, "spawn_task_terminal_event", + lambda a, eid, **k: calls.append( + ("event", eid, getattr(k["terminal_status"], "value", k["terminal_status"])))) + return calls + + +@pytest.mark.asyncio +async def test_unclaimed_interactive_turn_fails_as_capacity(seed_agent, monkeypatch): + seed_agent(timeout=1) + _row("e1") + from services import task_execution_service as tes + from services.execution_envelope import TaskExecutionErrorCode + + monkeypatch.setattr(tes, "QUEUE_CLAIM_POLL_INTERVAL", 0.05) + waited = AsyncMock() + monkeypatch.setattr("services.sync_waiter.wait_for_sync_terminal", waited) + hooks = _terminal_hooks(monkeypatch, tes) + + out = await tes.dispatch_and_await_terminal( + agent_name=AGENT, message="m", triggered_by="session", + service=_queued_service("e1"), conversation_key="session:k", + ) + assert hooks == [("close", "e1", "failed"), ("event", "e1", "failed")] + assert out.status == "failed" + assert out.error_code == TaskExecutionErrorCode.CAPACITY + assert out.error == "Agent at capacity (queued turn not claimed in 1s)" + row = _db().get_execution("e1") + assert (row.status, row.error) == ("failed", out.error) + waited.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel_that_loses_to_a_claim_waits_for_the_terminal(seed_agent, monkeypatch): + seed_agent(timeout=1) + _row("e1") + from services import task_execution_service as tes + + db = _db() + real_cancel = db.cancel_queued_execution + + def _claim_then_cancel(eid, reason="cancelled", status="cancelled"): + db.claim_next_queued(AGENT, worker_id="w1", lease_seconds=900) + return real_cancel(eid, reason, status) + + monkeypatch.setattr(tes, "QUEUE_CLAIM_POLL_INTERVAL", 0.05) + monkeypatch.setattr(tes.db, "cancel_queued_execution", _claim_then_cancel) + waited = AsyncMock() + monkeypatch.setattr("services.sync_waiter.wait_for_sync_terminal", waited) + + out = await tes.dispatch_and_await_terminal( + agent_name=AGENT, message="m", triggered_by="mcp", service=_queued_service("e1"), + ) + waited.assert_awaited_once_with("e1", 1 + 120.0) + assert out.status == "running" # rebuilt from the claimed row, never CAPACITY + assert out.error_code is None + + +@pytest.mark.asyncio +async def test_turn_finished_before_the_waiter_registers_returns_at_once(seed_agent, monkeypatch): + """A short turn can be claimed and finished inside one claim poll; the + sink's wake-up then fires before anyone listens.""" + seed_agent(timeout=60) + _row("e1", status="success", response="done") + from services import task_execution_service as tes + + waited = AsyncMock(side_effect=AssertionError("must not wait on a finished row")) + monkeypatch.setattr("services.sync_waiter.wait_for_sync_terminal", waited) + + out = await tes.dispatch_and_await_terminal( + agent_name=AGENT, message="m", triggered_by="session", service=_queued_service("e1"), + ) + assert (out.status, out.response) == ("success", "done") + + +@pytest.mark.asyncio +async def test_caller_going_away_cancels_the_queued_turn(seed_agent, monkeypatch): + seed_agent(timeout=60) + _row("e1") + from services import task_execution_service as tes + + monkeypatch.setattr(tes, "QUEUE_CLAIM_POLL_INTERVAL", 0.05) + hooks = _terminal_hooks(monkeypatch, tes) + task = asyncio.ensure_future(tes.dispatch_and_await_terminal( + agent_name=AGENT, message="m", triggered_by="paid", service=_queued_service("e1"), + )) + await asyncio.sleep(0.2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert _db().get_execution("e1").status == "cancelled" + assert hooks == [("close", "e1", "cancelled"), ("event", "e1", "cancelled")] + + +@pytest.mark.asyncio +async def test_autonomous_trigger_skips_the_claim_phase(seed_agent, monkeypatch): + seed_agent(timeout=1) + _row("e1", trigger="a2a") + from services import task_execution_service as tes + + waited = AsyncMock() + monkeypatch.setattr("services.sync_waiter.wait_for_sync_terminal", waited) + await tes.dispatch_and_await_terminal( + agent_name=AGENT, message="m", triggered_by="a2a", service=_queued_service("e1"), + ) + waited.assert_awaited_once() + assert _db().get_execution("e1").status == "queued" + + +@pytest.mark.asyncio +async def test_terminal_wait_timeout_text_matches_callers(seed_agent, monkeypatch): + """public_chat_service and mcp_auth_service map on the substring "timed out".""" + seed_agent(timeout=1) + _row("e1", status="running") + from services import task_execution_service as tes + + monkeypatch.setattr("services.sync_waiter.wait_for_sync_terminal", + AsyncMock(side_effect=asyncio.TimeoutError())) + out = await tes.dispatch_and_await_terminal( + agent_name=AGENT, message="m", triggered_by="public", service=_queued_service("e1"), + ) + assert "timed out" in out.error + + +# --------------------------------------------------------------------------- +# Session tab: run_resumable_turn + lock TTL +# --------------------------------------------------------------------------- + + +class _NoLock: + def __init__(self, *a, **k): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + +@pytest.mark.asyncio +async def test_resumable_turn_on_pilot_returns_real_uuid_and_retries_cold(seed_agent, monkeypatch): + seed_agent() + import importlib + + from services.execution_envelope import TaskExecutionResult + + # sys.modules entries: what run_resumable_turn's lazy import resolves. + sts = importlib.import_module("services.session_turn_service") + tes = importlib.import_module("services.task_execution_service") + _row("q1", status="failed", error="[agent_error] No conversation found with session ID uuid-old") + _row("q2", status="success", response="hi again", session="uuid-new") + calls = [] + + async def _execute(**kw): + calls.append(kw) + eid = "q1" if len(calls) == 1 else "q2" + return TaskExecutionResult(execution_id=eid, status="queued", response="") + + monkeypatch.setattr(tes, "get_task_execution_service", + lambda: SimpleNamespace(execute_task=_execute)) + monkeypatch.setattr(sts, "ResumeLock", _NoLock) + monkeypatch.setattr(sts, "supports_session_resume", lambda a: True) + monkeypatch.setattr("services.sync_waiter.wait_for_sync_terminal", AsyncMock()) + + turn = await sts.run_resumable_turn( + agent_name=AGENT, session_key="k1", message="hi", cached_uuid="uuid-old", + triggered_by="session", + ) + assert turn.fallback_fired is True + assert turn.real_uuid == "uuid-new" + assert turn.result.response == "hi again" + assert [c["conversation_key"] for c in calls] == ["session:k1", "session:k1"] + assert [c["resume_session_id"] for c in calls] == ["uuid-old", None] + assert all(c["persist_session"] is True for c in calls) + + +def test_lock_ttl_adds_the_queue_allowance_on_a_pilot(seed_agent, monkeypatch): + seed_agent(timeout=600) + seed_agent(name="beta", timeout=7200) + from services import session_turn_service as sts + + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", "") + assert sts.resolve_lock_ttl(AGENT) == 630 + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", f"{AGENT},beta") + # Two claim waits: the turn and its cold retry. + assert sts.resolve_lock_ttl(AGENT) == 630 + 2 * 600 + assert sts.resolve_lock_ttl("beta") == sts.LOCK_TTL_FALLBACK + 2 * 7200 + + +# --------------------------------------------------------------------------- +# Channel router key +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_channel_turn_is_keyed_by_its_chat_session(monkeypatch): + import adapters.message_router as mr + + execute = AsyncMock(return_value=SimpleNamespace( + status="success", response="ok", execution_id="e1", error=None, cost=0.0)) + monkeypatch.setattr(mr, "get_task_execution_service", lambda: MagicMock(execute_task=execute)) + monkeypatch.setattr(mr, "build_public_channel_caller_prompt", lambda *a, **kw: None) + monkeypatch.setattr(mr, "build_voice_capability_prompt", lambda *a, **kw: None) + monkeypatch.setattr(mr, "_get_channel_allowed_tools", lambda: ["WebSearch"]) + monkeypatch.setattr(mr.db, "get_public_channel_model", lambda agent: None) + + adapter = MagicMock() + adapter.get_source_identifier.return_value = "telegram:1:2" + message = SimpleNamespace(channel_id="2", thread_id="7", files=[], metadata={}) + router = mr.ChannelMessageRouter.__new__(mr.ChannelMessageRouter) + await router._run_agent_task( + adapter, message, AGENT, "tok", "telegram", False, None, None, "hi", + None, [], session_id="pcs-42", + ) + assert execute.await_args.kwargs["conversation_key"] == "channel:pcs-42" + + +# --------------------------------------------------------------------------- +# G-04 does not scan base64 image data +# --------------------------------------------------------------------------- + + +def _g04(meta: dict): + from canary.invariants import g04_no_creds_in_backlog_metadata as g04 + + agent = SimpleNamespace(name=AGENT, queued_exec_ids=["e1"], + queued_meta={"e1": {"backlog_metadata": json.dumps(meta)}}) + return g04.check(SimpleNamespace(agents=[agent], snapshot_time="t")) + + +def test_g04_ignores_images_but_still_scans_the_rest(): + assert _g04({"message": "hi", "images": _IMG}) == [] + hits = _g04({"system_prompt": "key AIzaSyExample", "images": _IMG}) + assert [v.observed_state["matched_pattern"] for v in hits] == ["google_api_key"] + # Only each image's encoded bytes are skipped. + for images, pattern in ( + ([{"media_type": "image/png", "data": "x", "note": "ghp_Abcdefgh"}], "github_pat"), + (["xoxb-Abc123"], "slack_bot_token"), + ): + hits = _g04({"message": "hi", "images": images}) + assert [v.observed_state["matched_pattern"] for v in hits] == [pattern] + + +# --------------------------------------------------------------------------- +# Live-stream proxy +# --------------------------------------------------------------------------- + + +class _FakeResponse: + def __init__(self, status, chunks=()): + self.status_code = status + self._chunks = chunks + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def aiter_text(self): + for c in self._chunks: + yield c + + +class _FakeAgentClient: + def __init__(self, responses): + self._responses = list(responses) + self.calls = 0 + + def __call__(self, *a, **k): + return self + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + def stream(self, method, url): + self.calls += 1 + return self._responses.pop(0) + + +@pytest.mark.asyncio +async def test_stream_holds_while_queued_then_attaches_after_the_claim(seed_agent, monkeypatch): + seed_agent() + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT) + _row("e1") + client = _FakeAgentClient([_FakeResponse(404), _FakeResponse(200, ["data: x\n\n"])]) + from routers import chat + from services import sync_waiter + + monkeypatch.setattr(sync_waiter, "STREAM_QUEUED_POLL_INTERVAL", 0.01) + monkeypatch.setattr(sync_waiter, "STREAM_ATTACH_RETRY_INTERVAL", 0.01) + monkeypatch.setattr(chat, "get_agent_container", lambda n: SimpleNamespace(status="running")) + monkeypatch.setattr(chat, "agent_httpx_client", client) + resp = await chat.stream_execution_log( + execution_id="e1", name=AGENT, current_user=SimpleNamespace(id=1)) + body = resp.body_iterator + + assert await body.__anext__() == ": queued\n\n" + assert await body.__anext__() == ": queued\n\n" # still queued: still holding + assert client.calls == 0 + _hrun("UPDATE schedule_executions SET status='running', " + "lease_expires_at='2099-01-01T00:00:00Z' WHERE id='e1'") # claimed + rest = [c async for c in body] + assert rest[-1] == "data: x\n\n" # 404 before registration retried + assert client.calls == 2 + + +@pytest.mark.asyncio +async def test_stream_off_pilot_is_unchanged(seed_agent, monkeypatch): + seed_agent() + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", "") + _row("e1") + client = _FakeAgentClient([_FakeResponse(404)]) + from routers import chat + + monkeypatch.setattr(chat, "get_agent_container", lambda n: SimpleNamespace(status="running")) + monkeypatch.setattr(chat, "agent_httpx_client", client) + resp = await chat.stream_execution_log( + execution_id="e1", name=AGENT, current_user=SimpleNamespace(id=1)) + chunks = [c async for c in resp.body_iterator] + assert not any(c.startswith(":") for c in chunks) + assert '"retryable": true' in chunks[0] and "stream_end" in chunks[-1] + assert client.calls == 1 + + +def test_every_stream_proxy_holds_while_queued(): + """Structural: the three live-stream proxies share the hold.""" + from client_portal import router as portal_router + from routers import chat, public + + for fn in (chat.stream_execution_log, public.public_stream_execution, + portal_router.portal_stream_execution): + src = inspect.getsource(fn) + assert "wait_while_queued(" in src and "pull_queue_allowance(" in src, fn.__name__ + + +# --------------------------------------------------------------------------- +# Concurrency +# --------------------------------------------------------------------------- + + +def _running_with_key(key): + return _scalar( + "SELECT COUNT(*) FROM schedule_executions WHERE conversation_key=:k AND status='running'", + k=key, + ) + + +def test_burst_of_one_conversation_is_claimed_one_at_a_time(seed_agent, monkeypatch): + """AC: on a pilot, queued turns of one conversation never run concurrently, + however many workers poll.""" + seed_agent() + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT) + burst = [f"b{i}" for i in range(3)] + for eid in burst + ["other"]: + _row(eid, status="running") + key = "session:other" if eid == "other" else "session:k" + assert _enqueue(eid, conversation_key=key) is True + from services import pull_coordination_service as pcs + + def _claim(worker): + with patch.object(pcs, "record_worker_poll"): + c = pcs.claim_next_task(AGENT, worker) + return c["execution_id"] if c else None + + done = [] + first = _claim("w1") + assert first in burst + assert _claim("w2") == "other" # another conversation is not held back + assert _claim("w3") is None # the rest of the burst waits + assert _running_with_key("session:k") == 1 + current = first + while True: + _hrun("UPDATE schedule_executions SET status='success' WHERE id=:id", id=current) + done.append(current) + nxt = _claim("w1") + if nxt is None: + break + assert _running_with_key("session:k") == 1 + assert _claim("w2") is None + current = nxt + assert sorted(done) == burst + + +class _FakeRedis: + """SET NX EX + the token-checked release: what ResumeLock uses.""" + + def __init__(self): + self.store = {} + + async def set(self, key, value, nx=False, ex=None): + if nx and key in self.store: + return None + self.store[key] = value + return True + + async def eval(self, _lua, _n, key, token): + if self.store.get(key) == token: + del self.store[key] + return 1 + return 0 + + +@pytest.mark.asyncio +async def test_session_turns_of_one_session_never_overlap_on_a_pilot(seed_agent, monkeypatch): + """AC: the ResumeLock is held across the queued wait, so a second turn of + the same Session cannot even enqueue until the first reached its terminal.""" + seed_agent(timeout=30) + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT) + import importlib + + from services.execution_envelope import TaskExecutionResult + + sts = importlib.import_module("services.session_turn_service") + tes = importlib.import_module("services.task_execution_service") + sync_waiter = importlib.import_module("services.sync_waiter") + db = _db() + events = [] + first_enqueued = asyncio.Event() + + async def _execute(**kw): + eid = f"t{sum(1 for e in events if e[0] == 'enqueue') + 1}" + _row(eid, status="queued", trigger="session") + _hrun("UPDATE schedule_executions SET conversation_key=:k WHERE id=:id", + k=kw["conversation_key"], id=eid) + events.append(("enqueue", eid)) + first_enqueued.set() + return TaskExecutionResult(execution_id=eid, status="queued", response="") + + async def _worker(): + finished = 0 + while finished < 2: + row = db.claim_next_queued(AGENT, worker_id="w1", lease_seconds=900) + if row is None: + await asyncio.sleep(0.01) + continue + events.append(("claim", row["id"])) + await asyncio.sleep(0.05) + _hrun("UPDATE schedule_executions SET status='success', response='ok', " + "claude_session_id='uuid-1' WHERE id=:id", id=row["id"]) + events.append(("done", row["id"])) + finished += 1 + + monkeypatch.setattr(tes, "get_task_execution_service", + lambda: SimpleNamespace(execute_task=_execute)) + monkeypatch.setattr(tes, "QUEUE_CLAIM_POLL_INTERVAL", 0.01) + monkeypatch.setattr(sync_waiter, "SYNC_WAITER_POLL_INTERVAL", 0.01) + fake = _FakeRedis() + monkeypatch.setattr(sts, "get_async_redis", lambda: fake) + monkeypatch.setattr(sts, "LOCK_POLL_INTERVAL_SECONDS", 0.01) + + def _turn(): + return sts.run_resumable_turn( + agent_name=AGENT, session_key="k1", message="hi", cached_uuid=None, + triggered_by="session", + ) + + worker = asyncio.create_task(_worker()) + turn1 = asyncio.create_task(_turn()) + await first_enqueued.wait() + turn2 = asyncio.create_task(_turn()) + r1, r2 = await asyncio.wait_for(asyncio.gather(turn1, turn2), 5) + await asyncio.wait_for(worker, 5) + + assert [e for e in events if e[1] == "t1"] == [("enqueue", "t1"), ("claim", "t1"), ("done", "t1")] + assert events.index(("enqueue", "t2")) > events.index(("done", "t1")) + assert r1.result.response == r2.result.response == "ok" + assert fake.store == {} # both locks released + + +def test_awaits_claim_covers_the_pre_enqueue_moment(seed_agent): + """A pre-created portal row is `running` for a moment before it is + enqueued; the stream must keep holding through it (#3114 live check).""" + seed_agent() + db = _db() + _row("pre", status="running") # pre-created + _row("q", status="queued") + _row("pushed", status="running", session="dispatched") # push sentinel + _row("claimed", status="running") + _hrun("UPDATE schedule_executions SET lease_expires_at='2099-01-01T00:00:00Z' " + "WHERE id='claimed'") + _row("done", status="success") + assert [db.execution_awaits_claim(e) for e in ("pre", "q", "pushed", "claimed", "done", "gone")] == [ + True, True, False, False, False, False] + + +@pytest.mark.asyncio +async def test_stream_holds_from_pre_created_row_until_the_claim(seed_agent, monkeypatch): + seed_agent() + _row("e2", status="running") # not yet enqueued + from services import sync_waiter + + monkeypatch.setattr(sync_waiter, "STREAM_QUEUED_POLL_INTERVAL", 0.01) + hold = sync_waiter.wait_while_queued("e2", budget=5) + assert await hold.__anext__() == ": queued\n\n" + _hrun("UPDATE schedule_executions SET status='queued' WHERE id='e2'") + assert await hold.__anext__() == ": queued\n\n" + _hrun("UPDATE schedule_executions SET status='running', " + "lease_expires_at='2099-01-01T00:00:00Z' WHERE id='e2'") + assert [t async for t in hold] == [] diff --git a/tests/unit/test_3114_pull_route_workspace.py b/tests/unit/test_3114_pull_route_workspace.py new file mode 100644 index 0000000000..02b6030126 --- /dev/null +++ b/tests/unit/test_3114_pull_route_workspace.py @@ -0,0 +1,455 @@ +"""Workspace surfaces on a pull pilot (#3114). + +What is pinned here: + +* The pull sink reports a delegated child back to the channel / portal / room + it came from, as the push terminals do; an inline-trigger turn is not + reported. +* Rooms: the working marker lives as long as a queued room turn can, and a room + turn waits at most ``ROOM_CLAIM_BUDGET_SECONDS`` for a claim + (``dispatch_and_await_terminal(claim_budget=...)``). +* An interactive turn no worker claimed in time is stored FAILED; a caller that + went away leaves it CANCELLED. +* The agent's pull worker registers a claimed turn as pending, so a terminate + before spawn skips the turn and reports it cancelled. +* Two wakes of one agent in one room run one after the other; the second reads + the cursor the first advanced. A mention chain back to the same agent does + not wait on its own lock. + +Runs against the real db layer through db_harness (SQLite; Postgres too when +TEST_POSTGRES_URL is set). +""" +from __future__ import annotations + +import asyncio +import json +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +_BACKEND_STR = str(Path(__file__).resolve().parent.parent.parent / "src" / "backend") +if _BACKEND_STR not in sys.path: + sys.path.insert(0, _BACKEND_STR) + +from db_harness import db_backend, run as _hrun, scalar as _scalar # noqa: E402,F401 + +pytestmark = pytest.mark.unit + +_STUBBED_MODULE_NAMES = [ + "db.connection", + "db.schedules", + "db.agent_settings.resources", + "database", +] + +AGENT = "alpha" + + +@pytest.fixture(autouse=True) +def _restore_sys_modules(): + saved = {name: sys.modules.get(name) for name in _STUBBED_MODULE_NAMES} + try: + yield + finally: + for name, value in saved.items(): + if value is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = value + + +@pytest.fixture +def tmp_db(db_backend): + def _evict(): + for mod in _STUBBED_MODULE_NAMES: + sys.modules.pop(mod, None) + + _evict() + try: + yield db_backend + finally: + _evict() + + +@pytest.fixture +def seed_agent(tmp_db): + def _seed(name: str = AGENT, timeout: int = 900): + _hrun( + "INSERT INTO agent_ownership (agent_name, owner_id, " + " execution_timeout_seconds, created_at) " + "VALUES (:n, 1, :t, '2026-01-01T00:00:00Z')", + n=name, t=timeout, + ) + + return _seed + + +def _row(eid: str, *, status: str = "queued", trigger: str = "mcp", + channel=None, chat_id=None): + _hrun( + "INSERT INTO schedule_executions " + "(id, schedule_id, agent_name, status, started_at, queued_at, message, " + " triggered_by, source_channel, source_channel_chat_id) " + "VALUES (:id, '__manual__', :a, :s, '2026-09-30T09:00:00Z', " + " CASE WHEN :s = 'queued' THEN '2026-09-30T09:00:00Z' END, 'hello', :t, " + " :ch, :cid)", + id=eid, a=AGENT, s=status, t=trigger, ch=channel, cid=chat_id, + ) + return eid + + +def _db(): + from database import db + + return db + + +# --------------------------------------------------------------------------- +# 1. Pull sink → completion report +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("trigger,reported", [("mcp", True), ("public", False)]) +async def test_pull_sink_reports_a_delegated_child_back(seed_agent, monkeypatch, + trigger, reported): + seed_agent() + _row("e1", trigger=trigger, channel="portal", chat_id="thread-1") + import contextlib + import importlib + + pcs = importlib.import_module("services.pull_coordination_service") + ccr = importlib.import_module("services.channel_completion_report") + idem = importlib.import_module("services.idempotency_service") + + delivered = [] + + def _resolver(**kw): + async def _deliver(): + delivered.append(kw) + return True + return _deliver + + @contextlib.asynccontextmanager + async def _guard(*a, **k): + yield SimpleNamespace(replay=False, snapshot=None) + + monkeypatch.setattr(ccr, "_inflight", set()) # other tests' loops are closed + monkeypatch.setitem(ccr._CHANNEL_RESOLVERS, "portal", _resolver) + monkeypatch.setattr(idem, "effect_guard", _guard) + + row = _db().claim_next_queued(AGENT, worker_id="w1", lease_seconds=900) + outcome = pcs.apply_task_result("e1", row["claim_token"], status="success", + content="child done") + assert outcome.kind == "applied" + await asyncio.gather(*list(ccr._inflight)) + + if reported: + assert len(delivered) == 1 + assert delivered[0]["chat_id"] == "thread-1" + assert delivered[0]["status"] == "success" + assert delivered[0]["summary_or_error"] == "child done" + else: + assert delivered == [] + + +# --------------------------------------------------------------------------- +# 2. Rooms: working-marker TTL + claim budget +# --------------------------------------------------------------------------- + + +def test_room_working_marker_ttl_covers_the_queued_wait(seed_agent, monkeypatch): + from shared_sessions import service + + seed_agent(timeout=900) + base = service.ROOM_TURN_TIMEOUT_SECONDS + 30 + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", "") + assert service._working_ttl(AGENT) == base + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT) + assert service._working_ttl(AGENT) == base + service.ROOM_CLAIM_BUDGET_SECONDS + _hrun("UPDATE agent_ownership SET execution_timeout_seconds = 120 WHERE agent_name = :n", + n=AGENT) + assert service._working_ttl(AGENT) == base + 120 + + +@pytest.mark.asyncio +async def test_claim_budget_bounds_the_claim_wait(seed_agent, monkeypatch): + seed_agent(timeout=3600) + _row("e1", trigger="room") + from services import task_execution_service as tes + from services.execution_envelope import TaskExecutionErrorCode, TaskExecutionResult + + monkeypatch.setattr(tes, "QUEUE_CLAIM_POLL_INTERVAL", 0.05) + svc = MagicMock() + svc.execute_task = AsyncMock(return_value=TaskExecutionResult( + execution_id="e1", status="queued", response="")) + + out = await asyncio.wait_for(tes.dispatch_and_await_terminal( + agent_name=AGENT, message="m", triggered_by="room", service=svc, + claim_budget=1, + ), 5) + assert out.error_code == TaskExecutionErrorCode.CAPACITY + assert out.error == "Agent at capacity (queued turn not claimed in 1s)" + assert "claim_budget" not in svc.execute_task.await_args.kwargs + + +# --------------------------------------------------------------------------- +# 3. Unclaimed → FAILED at the db layer (the CAS is unchanged) +# --------------------------------------------------------------------------- + + +def test_cancel_queued_execution_writes_the_requested_terminal(seed_agent): + seed_agent() + db = _db() + _row("f1") + assert db.cancel_queued_execution("f1", reason="not claimed", status="failed") is True + row = db.get_execution("f1") + assert (row.status, row.error) == ("failed", "not claimed") + assert row.completed_at + + _row("c1") + assert db.cancel_queued_execution("c1", reason="gone") is True + assert db.get_execution("c1").status == "cancelled" + + _row("r1", status="running") + assert db.cancel_queued_execution("r1", status="failed") is False + assert db.get_execution("r1").status == "running" + + +# --------------------------------------------------------------------------- +# 4. Pull worker: terminate between claim and spawn +# --------------------------------------------------------------------------- + + +def _drive_worker(monkeypatch, *, before_run=None, run=None): + from agent_server.services import pull_worker as pw + from agent_server.services import process_registry as pr + + registry = pr.ProcessRegistry() + monkeypatch.setattr(pr, "get_process_registry", lambda: registry) + seen = {} + + async def _execute(**kw): + seen["pending_at_spawn"] = kw["execution_id"] in registry.list_pending_ids() + if run: + return run(registry) + md = SimpleNamespace(model_dump=lambda: {"cost_usd": 0.0}) + return ("ok", [], md, "s") + + runtime = SimpleNamespace(execute_headless=AsyncMock(side_effect=_execute)) + stub = types.ModuleType("agent_server.services.runtime_adapter") + stub.get_runtime = lambda: runtime + monkeypatch.setitem(sys.modules, "agent_server.services.runtime_adapter", stub) + if before_run: + real_start = pw.agent_state.record_task_start + + def _start(): + before_run(registry) + real_start() + + monkeypatch.setattr(pw.agent_state, "record_task_start", _start) + deliver = AsyncMock(return_value=True) + monkeypatch.setattr(pw, "_deliver_result", deliver) + monkeypatch.setattr(pw, "_persist_pull_result", lambda *a, **k: None) + monkeypatch.setattr(pw, "_delete_pull_result", lambda *a, **k: None) + claim = {"execution_id": "e1", "claim_token": "t", "envelope": {"payload": {"message": "m"}}} + asyncio.run(pw._run_and_report(claim, MagicMock(), "http://b", "k", AGENT)) + return runtime.execute_headless, deliver.await_args.args[2], registry, seen + + +def test_claimed_turn_is_pending_until_spawn(monkeypatch): + run, body, registry, seen = _drive_worker(monkeypatch) + assert seen["pending_at_spawn"] is True + assert body["status"] == "success" + assert registry.list_pending_ids() == [] + + +def test_terminate_before_spawn_skips_the_turn_and_reports_cancelled(monkeypatch): + outcome = {} + + def _terminate(registry): + outcome["terminate"] = registry.terminate("e1") + + run, body, registry, _ = _drive_worker(monkeypatch, before_run=_terminate) + assert outcome["terminate"]["reason"] == "cancelled_before_start" + run.assert_not_called() + assert (body["status"], body["error_code"], body["claim_token"]) == ("cancelled", None, "t") + assert registry.list_pending_ids() == [] + + +def test_terminated_turn_keeps_an_auth_failure_class(monkeypatch): + def _run(registry): + registry.terminate("e1") + raise HTTPException(status_code=503, detail="auth failed") + + _, body, _, _ = _drive_worker(monkeypatch, run=_run) + assert (body["status"], body["error_code"]) == ("failed", "auth") + + +# --------------------------------------------------------------------------- +# 5. Two wakes of one agent in one room +# --------------------------------------------------------------------------- + + +class _FakeRedis: + """SET NX EX + the token-checked release: what ResumeLock uses.""" + + def __init__(self): + self.store = {} + self.ttls = {} + + async def set(self, key, value, nx=False, ex=None): + if nx and key in self.store: + return None + self.store[key] = value + self.ttls[key] = ex + return True + + async def eval(self, _lua, _n, key, token): + if self.store.get(key) == token: + del self.store[key] + return 1 + return 0 + + +def _fake_room(monkeypatch, *, on_post=None): + """A room with one unread message for AGENT; the cursor is real state.""" + from services import session_turn_service as sts + from shared_sessions import service + + state = {"cursor": 0, "dispatches": [], "since": []} + messages = [{"seq": 1, "sender_kind": "user", "sender_identity": "c@example.com", + "content": f"@{AGENT} hi", "kind": "message"}] + + def _get_messages(_room, since_seq=0): + state["since"].append(since_seq) + return [m for m in messages if m["seq"] > since_seq] + + def _advance(_room, _agent, seq, _sid): + state["cursor"] = seq + + async def _dispatch(**kw): + state["dispatches"].append(kw) + await asyncio.sleep(0.05) + return SimpleNamespace(status="success", response="hello", execution_id="x", + session_id="s1", error=None) + + async def _post(*a, **k): + if on_post: + await on_post() + return {} + + monkeypatch.setattr(service.db, "get_participant", + lambda *a, **k: {"last_read_seq": state["cursor"], + "cached_session_id": "s1"}) + monkeypatch.setattr(service.db, "get_room", lambda *a, **k: { + "id": "r1", "name": "Room", "status": "open", "topic": None}) + monkeypatch.setattr(service.db, "get_messages", _get_messages) + monkeypatch.setattr(service.db, "list_participants", lambda _r: [ + {"kind": "agent", "identity": AGENT, "left_at": None}]) + monkeypatch.setattr(service.db, "advance_read_cursor", _advance) + for name in ("_mark_agent_working", "_clear_agent_working", "_broadcast", "_post_system"): + monkeypatch.setattr(service, name, lambda *a, **k: None) + monkeypatch.setattr(service, "_room_queue_allowance", lambda _a: 0) + monkeypatch.setattr(service, "_room_inbox_context", AsyncMock(return_value=("", []))) + monkeypatch.setattr(service, "post_message", _post) + import services.task_execution_service as tes + + monkeypatch.setattr(tes, "dispatch_and_await_terminal", _dispatch) + fake = _FakeRedis() + monkeypatch.setattr(sts, "get_async_redis", lambda: fake) + monkeypatch.setattr(sts, "LOCK_POLL_INTERVAL_SECONDS", 0.01) + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", AGENT) + return service, state, fake + + +@pytest.mark.asyncio +async def test_second_wake_reads_what_the_first_left(monkeypatch): + service, state, fake = _fake_room(monkeypatch) + user = SimpleNamespace(email="c@example.com") + + await asyncio.wait_for(asyncio.gather( + service._wake_agent(user, "r1", AGENT, 0), + service._wake_agent(user, "r1", AGENT, 0), + ), 5) + + assert len(state["dispatches"]) == 1 # the delta is answered once + assert state["since"] == [0, 1] # second wake saw the advanced cursor + assert state["dispatches"][0]["claim_budget"] == service.ROOM_CLAIM_BUDGET_SECONDS + assert fake.store == {} + assert fake.ttls[f"room_wake_lock:r1:{AGENT}"] == ( + service.ROOM_TURN_TIMEOUT_SECONDS + service.ROOM_WAKE_LOCK_BUFFER_SECONDS) + + +@pytest.mark.asyncio +async def test_mention_chain_back_to_the_same_agent_does_not_wait_on_itself(monkeypatch): + calls = {"n": 0} + user = SimpleNamespace(email="c@example.com") + + async def _reply_mentions_back(): + calls["n"] += 1 + if calls["n"] == 1: # A's reply wakes B, whose reply wakes A again + await service._wake_agent(user, "r1", AGENT, 1) + + service, state, fake = _fake_room(monkeypatch, on_post=_reply_mentions_back) + await asyncio.wait_for(service._wake_agent(user, "r1", AGENT, 0), 5) + assert len(state["dispatches"]) == 2 + assert fake.store == {} + + +@pytest.mark.asyncio +async def test_push_agent_wake_takes_no_lock(monkeypatch): + service, state, fake = _fake_room(monkeypatch) + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", "") + await asyncio.wait_for(service._wake_agent(SimpleNamespace(email="c@example.com"), + "r1", AGENT, 0), 5) + assert len(state["dispatches"]) == 1 + assert fake.ttls == {} + + +# --------------------------------------------------------------------------- +# Live-check fixes: Session capacity status, stop message +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("code, status", [("capacity", 429), ("agent_error", 502), (None, 502)]) +def test_session_turn_maps_capacity_to_429(code, status): + """An unclaimed pilot turn returns FAILED/CAPACITY; the Session tab answers + the retryable 429 for it.""" + import asyncio as _asyncio + from unittest.mock import AsyncMock as _AsyncMock, MagicMock as _MagicMock, patch as _patch + from fastapi import HTTPException + from models import SessionMessageRequest + from routers import sessions as sessions_mod + + user = _MagicMock(id=1, email="u@example.com", username="u", role="admin", + agent_name=None, mcp_key_id=None, mcp_key_name=None) + db = _MagicMock() + db.get_session.return_value = _MagicMock(id="sess1", user_id=1, agent_name="agent1", + subscription_id=None, cached_claude_session_id=None) + turn = _MagicMock(result=_MagicMock(status="failed", error="Agent at capacity (x)", + error_code=code, execution_id="e1"), + fallback_fired=False, fallback_reason=None) + with _patch.object(sessions_mod, "is_session_tab_enabled", return_value=True), \ + _patch.object(sessions_mod, "db", db), \ + _patch.object(sessions_mod, "run_resumable_turn", _AsyncMock(return_value=turn)), \ + _patch.object(sessions_mod, "_resolve_lock_ttl", return_value=60): + with pytest.raises(HTTPException) as exc: + _asyncio.run(sessions_mod.send_session_message( + name="agent1", session_id="sess1", + body=SessionMessageRequest(message="hi"), current_user=user)) + assert exc.value.status_code == status + + +def test_terminated_turn_reports_the_terminate_message(monkeypatch): + def _run(registry): + registry.terminate("e1") + raise RuntimeError("Claude Code reported no error detail") + + _, body, _, _ = _drive_worker(monkeypatch, run=_run) + assert (body["status"], body["error"]) == ("cancelled", "Execution terminated by user") diff --git a/tests/unit/test_ent169_shared_sessions.py b/tests/unit/test_ent169_shared_sessions.py index c683899dee..c8e2ff6b18 100644 --- a/tests/unit/test_ent169_shared_sessions.py +++ b/tests/unit/test_ent169_shared_sessions.py @@ -677,7 +677,7 @@ def test_the_reply_is_posted_before_the_cursor_advances(rooms_db, allow_all, fak import ast, inspect from shared_sessions import service - src = inspect.getsource(service._wake_agent) + src = inspect.getsource(service._wake_agent_locked) tree = ast.parse(src.strip()) order = [] for node in ast.walk(tree): diff --git a/tests/unit/test_ent362_workspace_participant.py b/tests/unit/test_ent362_workspace_participant.py index 75a61c8a38..a05f809680 100644 --- a/tests/unit/test_ent362_workspace_participant.py +++ b/tests/unit/test_ent362_workspace_participant.py @@ -434,7 +434,7 @@ def test_the_marker_is_set_and_cleared_around_the_turn(): import inspect from shared_sessions import service - src = inspect.getsource(service._wake_agent) + src = inspect.getsource(service._wake_agent_locked) assert "_mark_agent_working(" in src assert "_clear_agent_working(" in src # The clear must be in the finally, or a raised turn leaves it stuck. diff --git a/tests/unit/test_ent525_portal_work.py b/tests/unit/test_ent525_portal_work.py index 05b7ea56eb..4374f7fcf4 100644 --- a/tests/unit/test_ent525_portal_work.py +++ b/tests/unit/test_ent525_portal_work.py @@ -631,3 +631,17 @@ def test_the_fleet_dashboard_payload_gains_no_channel_destination_ids(): "started_at": "2026-09-06T10:00:00Z", "message": "m", "triggered_by": "public", "source_channel_chat_id": "sess-1", "loop_id": "l1", "source_channel": "portal"} assert not hasattr(FleetExecutionSummary(**row), "source_channel_chat_id") + + +# #3114: a pilot's Workspace turn is pre-created `running` and enqueued a moment +# later. Until a worker claims it (a claim writes a lease) it is waiting for a +# slot, so the card must not say "Working". +@pytest.mark.parametrize("row, outcome", [ + (dict(agent_name="pilot-a", status="running", triggered_by="public", lease_expires_at=None), "queued"), + (dict(agent_name="pilot-a", status="running", triggered_by="public", + lease_expires_at="2026-10-01T17:00:00Z"), "running"), + (dict(agent_name="push-a", status="running", triggered_by="public", lease_expires_at=None), "running"), +]) +def test_unclaimed_pilot_row_reads_queued(svc, monkeypatch, row, outcome): + monkeypatch.setenv("PULL_MODE_PILOT_AGENTS", "pilot-a") + assert svc.work_outcome(_row(**row)) == outcome