From 7ddef1ac994889036312f666c760217c631d4fc9 Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Sun, 17 May 2026 19:57:17 -0300 Subject: [PATCH 1/4] Stabilize approval gate execution resume The approval gate must fail closed when its hook infrastructure is missing, pause write-capable calls until an operator resolves them, and resume execution without duplicating tool results in the provider transcript. This keeps the minimal approval path releaseable without carrying unrelated branch work. Constraint: Keep the delivery scoped to the minimal approval-gate path from origin/main Constraint: Approval-gate resolution must execute allowed calls through the existing function dispatch path Rejected: Execute the target function inside approval::resolve | would bypass orchestrator lifecycle and dispatch invariants Rejected: Publish run::resume for active turns | created duplicate step execution and repeated tool results Confidence: high Scope-risk: moderate Directive: Do not change approval resume behavior without rerunning the browser allow/deny E2E against a real provider Tested: git diff --check Tested: cargo test in approval-gate, hook-fanout, turn-orchestrator, provider-router Tested: npm test -- ApprovalRow.test.tsx reducer.test.ts in harness/web Tested: npm run build in harness/web Tested: npm run e2e -- tests/e2e/approval.spec.ts --output=/tmp/harness-approval-e2e-results with claude-haiku-4-5 and Anthropic env credential Not-tested: separate fresh-worktree patch replay after the final E2E-driven resume fixes --- approval-gate/src/lib.rs | 682 +++++++++++++----- approval-gate/tests/integration.rs | 127 ++-- harness/web/src/App.tsx | 1 + .../web/src/components/ApprovalRow.test.tsx | 68 ++ harness/web/src/components/ApprovalRow.tsx | 17 +- harness/web/src/reducer.test.ts | 41 +- harness/web/src/reducer.ts | 12 + harness/web/src/types.ts | 11 + harness/web/tests/e2e/approval.spec.ts | 64 +- hook-fanout/src/handler.rs | 66 +- hook-fanout/src/lib.rs | 28 +- provider-router/src/register.rs | 80 +- turn-orchestrator/src/bootstrap.rs | 5 +- turn-orchestrator/src/manifest.rs | 3 +- turn-orchestrator/src/run_start.rs | 75 ++ turn-orchestrator/src/states/assistant.rs | 41 ++ turn-orchestrator/src/states/functions.rs | 522 +++++++++++++- turn-orchestrator/src/states/provisioning.rs | 5 +- turn-orchestrator/src/states/tearing_down.rs | 24 + turn-orchestrator/src/system_prompt.rs | 29 +- 20 files changed, 1571 insertions(+), 330 deletions(-) create mode 100644 harness/web/src/components/ApprovalRow.test.tsx diff --git a/approval-gate/src/lib.rs b/approval-gate/src/lib.rs index f6225e662..39d21dee1 100644 --- a/approval-gate/src/lib.rs +++ b/approval-gate/src/lib.rs @@ -1,6 +1,6 @@ //! Approval gate. Subscribes to `agent::before_function_call` and blocks calls //! whose `function_call.function_id` appears in the run's `approval_required` list, -//! waiting for the UI to call `approval::resolve` (or for a timeout). +//! then lets `approval::consume` drain resolved decisions back into the turn. pub mod config; pub mod manifest; @@ -16,6 +16,8 @@ use serde_json::{json, Value}; pub const FN_RESOLVE: &str = "approval::resolve"; pub const FN_LIST_PENDING: &str = "approval::list_pending"; +pub const FN_CONSUME: &str = "approval::consume"; +pub const FN_SWEEP_SESSION: &str = "approval::sweep_session"; /// Default `approval_state_scope` (matches [`WorkerConfig::default`]). pub const STATE_SCOPE: &str = "approvals"; @@ -102,6 +104,7 @@ pub fn extract_call(envelope: &Value) -> Option { } pub fn build_pending_record( + session_id: &str, function_call_id: &str, function_id: &str, args: &Value, @@ -109,20 +112,28 @@ pub fn build_pending_record( timeout_ms: u64, ) -> Value { json!({ + "session_id": session_id, "function_call_id": function_call_id, "function_id": function_id, "args": args, "status": "pending", + "created_at": now_ms, "expires_at": now_ms.saturating_add(timeout_ms), }) } pub fn block_reply_for(decision: &Decision) -> Value { match decision { - Decision::Allow => json!({ "block": false }), + Decision::Allow => json!({ + "block": false, + "subscriber": "approval-gate", + "approval_gate": true, + }), Decision::Deny { reason } => json!({ "block": true, "reason": format!("approval-gate: {reason}"), + "subscriber": "approval-gate", + "approval_gate": true, }), } } @@ -130,6 +141,8 @@ pub fn block_reply_for(decision: &Decision) -> Value { pub struct Refs { pub resolve: FunctionRef, pub list_pending: FunctionRef, + pub consume: FunctionRef, + pub sweep_session: FunctionRef, pub subscriber_fn: FunctionRef, pub subscriber_trigger: iii_sdk::Trigger, } @@ -168,8 +181,23 @@ pub async fn handle_resolve(bus: &dyn StateBus, state_scope: &str, payload: Valu if existing.get("status").and_then(Value::as_str) != Some("pending") { return json!({ "ok": false, "error": "already_resolved" }); } - existing["status"] = + let now_ms = now_ms(); + if existing + .get("expires_at") + .and_then(Value::as_u64) + .is_some_and(|expires_at| now_ms >= expires_at) + { + existing["status"] = json!("resolved"); + existing["decision"] = json!("deny"); + existing["reason"] = json!("timed_out"); + existing["resolved_at"] = json!(now_ms); + let _ = bus.set(state_scope, &key, existing).await; + return json!({ "ok": false, "error": "timed_out" }); + } + existing["status"] = json!("resolved"); + existing["decision"] = serde_json::to_value(decision).expect("WireDecision serializes via Serialize"); + existing["resolved_at"] = json!(now_ms); if let Some(reason) = payload.get("reason").cloned() { existing["reason"] = reason; } @@ -180,6 +208,157 @@ pub async fn handle_resolve(bus: &dyn StateBus, state_scope: &str, payload: Valu json!({ "ok": true }) } +pub async fn handle_intercept( + bus: &dyn StateBus, + state_scope: &str, + call: &IncomingCall, + now_ms: u64, + timeout_ms: u64, +) -> Value { + if !call.requires_approval() { + return block_reply_for(&Decision::Allow); + } + + let record = build_pending_record( + &call.session_id, + &call.function_call_id, + &call.function_id, + &call.args, + now_ms, + timeout_ms, + ); + if let Err(err) = bus + .set( + state_scope, + &pending_key(&call.session_id, &call.function_call_id), + record, + ) + .await + { + tracing::error!( + "approval-gate: failed to write pending record for {}/{}: {err}", + call.session_id, + call.function_call_id + ); + return json!({ + "block": true, + "status": "denied", + "reason": "approval-gate: state_write_failed", + "subscriber": "approval-gate", + "approval_gate": true, + "denial": { + "kind": "state_error", + "detail": { + "phase": "pending_write", + "error": err.to_string(), + } + } + }); + } + + json!({ + "block": true, + "status": "pending", + "reason": "approval required", + "function_call_id": call.function_call_id, + "tool_call_id": call.function_call_id, + "function_id": call.function_id, + "subscriber": "approval-gate", + "approval_gate": true, + }) +} + +pub async fn handle_consume(bus: &dyn StateBus, state_scope: &str, payload: Value) -> Value { + let session_id = payload + .get("session_id") + .and_then(Value::as_str) + .unwrap_or(""); + if session_id.is_empty() { + return json!({ "ok": false, "error": "missing_session_id", "entries": [] }); + } + + let prefix = format!("{session_id}/"); + let rows = bus.list_prefix(state_scope, &prefix).await; + let mut entries = Vec::new(); + for mut row in rows { + if row.get("session_id").and_then(Value::as_str) != Some(session_id) { + continue; + } + if row.get("status").and_then(Value::as_str) != Some("resolved") { + continue; + } + let Some(function_call_id) = row.get("function_call_id").and_then(Value::as_str) else { + continue; + }; + let key = pending_key(session_id, function_call_id); + entries.push(json!({ + "function_call_id": function_call_id, + "tool_call_id": function_call_id, + "function_id": row.get("function_id").cloned().unwrap_or(Value::Null), + "args": row.get("args").cloned().unwrap_or_else(|| json!({})), + "decision": row.get("decision").cloned().unwrap_or_else(|| json!("deny")), + "reason": row.get("reason").cloned().unwrap_or(Value::Null), + })); + row["status"] = json!("consumed"); + row["consumed_at"] = json!(now_ms()); + let _ = bus.set(state_scope, &key, row).await; + } + json!({ "ok": true, "entries": entries }) +} + +pub async fn handle_sweep_session(bus: &dyn StateBus, state_scope: &str, payload: Value) -> Value { + let session_id = payload + .get("session_id") + .and_then(Value::as_str) + .unwrap_or(""); + if session_id.is_empty() { + return json!({ "ok": false, "error": "missing_session_id", "swept": 0 }); + } + + let prefix = format!("{session_id}/"); + let rows = bus.list_prefix(state_scope, &prefix).await; + let now_ms = now_ms(); + let mut swept = 0u64; + for mut row in rows { + if row.get("session_id").and_then(Value::as_str) != Some(session_id) { + continue; + } + if row.get("status").and_then(Value::as_str) != Some("pending") { + continue; + } + let Some(function_call_id) = row + .get("function_call_id") + .and_then(Value::as_str) + .map(str::to_string) + else { + continue; + }; + row["status"] = json!("resolved"); + row["decision"] = json!("deny"); + row["reason"] = json!("timed_out"); + row["resolved_at"] = json!(now_ms); + if bus + .set( + state_scope, + &pending_key(session_id, &function_call_id), + row, + ) + .await + .is_ok() + { + swept += 1; + } + } + json!({ "ok": true, "swept": swept }) +} + +fn now_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} + pub async fn handle_list_pending(bus: &dyn StateBus, state_scope: &str, payload: Value) -> Value { let session_id = payload .get("session_id") @@ -192,50 +371,29 @@ pub async fn handle_list_pending(bus: &dyn StateBus, state_scope: &str, payload: let all = bus.list_prefix(state_scope, &prefix).await; let pending: Vec = all .into_iter() + .filter(|v| v.get("session_id").and_then(Value::as_str) == Some(session_id)) .filter(|v| v.get("status").and_then(Value::as_str) == Some("pending")) .collect(); json!({ "pending": pending }) } -const POLL_INTERVAL_MS: u64 = 250; +fn state_list_values(resp: Value) -> Vec { + let entries = match resp { + Value::Array(entries) => entries, + Value::Object(mut obj) => obj + .remove("items") + .and_then(|items| match items { + Value::Array(entries) => Some(entries), + _ => None, + }) + .unwrap_or_default(), + _ => Vec::new(), + }; -pub async fn await_decision( - bus: &dyn StateBus, - state_scope: &str, - session_id: &str, - function_call_id: &str, - expires_at: u64, -) -> Decision { - let key = pending_key(session_id, function_call_id); - loop { - let Some(rec) = bus.get(state_scope, &key).await else { - return Decision::Deny { - reason: "state_unavailable".into(), - }; - }; - match rec.get("status").and_then(Value::as_str) { - Some("allow") => return Decision::Allow, - Some("deny") => { - let reason = rec - .get("reason") - .and_then(Value::as_str) - .unwrap_or("user") - .to_string(); - return Decision::Deny { reason }; - } - _ => {} - } - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(expires_at); - if now >= expires_at { - return Decision::Deny { - reason: "timeout".into(), - }; - } - tokio::time::sleep(std::time::Duration::from_millis(POLL_INTERVAL_MS)).await; - } + entries + .into_iter() + .map(|entry| entry.get("value").cloned().unwrap_or(entry)) + .collect() } /// Production [`StateBus`] backed by a real iii-sdk [`III`] connection. @@ -277,12 +435,7 @@ impl StateBus for IiiStateBus { }) .await .unwrap_or_else(|_| json!({ "items": [] })); - resp.get("items") - .and_then(|v| v.as_array().cloned()) - .unwrap_or_default() - .into_iter() - .map(|entry| entry.get("value").cloned().unwrap_or(entry)) - .collect() + state_list_values(resp) } } @@ -341,18 +494,58 @@ pub fn register(iii: &III, cfg: &WorkerConfig) -> anyhow::Result { let bus_for_resolve = bus.clone(); let scope_resolve = state_scope.clone(); - let resolve = - iii.register_function(( - RegisterFunctionMessage::with_id(FN_RESOLVE.into()) - .with_description("Flip a pending approval entry to allow or deny.".into()), - move |payload: Value| { - let bus = bus_for_resolve.clone(); - let scope_resolve = scope_resolve.clone(); - async move { - Ok::<_, IIIError>(handle_resolve(bus.as_ref(), &scope_resolve, payload).await) + let iii_for_resolve = iii.clone(); + let resolve = iii.register_function(( + RegisterFunctionMessage::with_id(FN_RESOLVE.into()) + .with_description("Flip a pending approval entry to allow or deny.".into()), + move |payload: Value| { + let bus = bus_for_resolve.clone(); + let scope_resolve = scope_resolve.clone(); + let iii = iii_for_resolve.clone(); + async move { + let session_id = payload + .get("session_id") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + let function_call_id = payload + .get("function_call_id") + .or_else(|| payload.get("tool_call_id")) + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + let out = handle_resolve(bus.as_ref(), &scope_resolve, payload).await; + if out.get("ok").and_then(Value::as_bool) == Some(true) + && !session_id.is_empty() + && !function_call_id.is_empty() + { + if let Some(record) = bus + .get(&scope_resolve, &pending_key(&session_id, &function_call_id)) + .await + { + write_event( + &iii, + &session_id, + &approval_resolved_event(&function_call_id, &record), + ) + .await; + } + if let Err(err) = resume_session(&iii, &session_id).await { + write_event( + &iii, + &session_id, + &json!({ + "type": "approval_wake_failed", + "error": err, + }), + ) + .await; + } } - }, - )); + Ok::<_, IIIError>(out) + } + }, + )); let bus_for_list = bus.clone(); let scope_list = state_scope.clone(); @@ -368,6 +561,35 @@ pub fn register(iii: &III, cfg: &WorkerConfig) -> anyhow::Result { }, )); + let bus_for_consume = bus.clone(); + let scope_consume = state_scope.clone(); + let consume = + iii.register_function(( + RegisterFunctionMessage::with_id(FN_CONSUME.into()) + .with_description("Return resolved approval decisions for a session once.".into()), + move |payload: Value| { + let bus = bus_for_consume.clone(); + let scope_consume = scope_consume.clone(); + async move { + Ok::<_, IIIError>(handle_consume(bus.as_ref(), &scope_consume, payload).await) + } + }, + )); + + let bus_for_sweep = bus.clone(); + let scope_sweep = state_scope.clone(); + let sweep_session = iii.register_function(( + RegisterFunctionMessage::with_id(FN_SWEEP_SESSION.into()) + .with_description("Resolve a session's pending approvals as denied.".into()), + move |payload: Value| { + let bus = bus_for_sweep.clone(); + let scope_sweep = scope_sweep.clone(); + async move { + Ok::<_, IIIError>(handle_sweep_session(bus.as_ref(), &scope_sweep, payload).await) + } + }, + )); + let iii_for_sub = iii.clone(); let bus_for_sub = bus.clone(); let subscriber_scope = state_scope.clone(); @@ -380,81 +602,27 @@ pub fn register(iii: &III, cfg: &WorkerConfig) -> anyhow::Result { let sc = subscriber_scope.clone(); async move { let Some(call) = extract_call(&envelope) else { - return Ok::<_, IIIError>(json!({ "block": false })); - }; - if !call.requires_approval() { - let reply = json!({ "block": false }); - write_hook_reply(&iii, &call.reply_stream, &call.event_id, &reply).await; + let reply = block_reply_for(&Decision::Allow); return Ok(reply); - } - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0); - let expires_at = now.saturating_add(timeout_ms); - let record = build_pending_record( - &call.function_call_id, - &call.function_id, - &call.args, - now, - timeout_ms, - ); - if let Err(err) = bus - .set( - &sc, - &pending_key(&call.session_id, &call.function_call_id), - record, + }; + let now = now_ms(); + let reply = handle_intercept(bus.as_ref(), &sc, &call, now, timeout_ms).await; + if reply.get("status").and_then(Value::as_str) == Some("pending") { + write_event( + &iii, + &call.session_id, + &json!({ + "type": "approval_requested", + "function_call_id": call.function_call_id, + "tool_call_id": call.function_call_id, + "function_id": call.function_id, + "tool_name": call.function_id, + "args": call.args, + "expires_at": now.saturating_add(timeout_ms), + }), ) - .await - { - tracing::error!( - "approval-gate: failed to write pending record for {}/{}: {err}", - call.session_id, - call.function_call_id - ); - let reply = json!({ "block": false }); - write_hook_reply(&iii, &call.reply_stream, &call.event_id, &reply).await; - return Ok(reply); + .await; } - write_event( - &iii, - &call.session_id, - &json!({ - "type": "approval_requested", - "function_call_id": call.function_call_id, - "tool_call_id": call.function_call_id, - "function_id": call.function_id, - "tool_name": call.function_id, - "args": call.args, - "expires_at": expires_at, - }), - ) - .await; - let decision = await_decision( - bus.as_ref(), - &sc, - &call.session_id, - &call.function_call_id, - expires_at, - ) - .await; - let (decision_str, reason_for_event) = match &decision { - Decision::Allow => ("allow", None), - Decision::Deny { reason } => ("deny", Some(reason.clone())), - }; - write_event( - &iii, - &call.session_id, - &json!({ - "type": "approval_resolved", - "function_call_id": call.function_call_id, - "tool_call_id": call.function_call_id, - "decision": decision_str, - "reason": reason_for_event, - }), - ) - .await; - let reply = block_reply_for(&decision); write_hook_reply(&iii, &call.reply_stream, &call.event_id, &reply).await; Ok(reply) } @@ -473,11 +641,47 @@ pub fn register(iii: &III, cfg: &WorkerConfig) -> anyhow::Result { Ok(Refs { resolve, list_pending, + consume, + sweep_session, subscriber_fn, subscriber_trigger, }) } +fn approval_resolved_event(function_call_id: &str, record: &Value) -> Value { + json!({ + "type": "approval_resolved", + "function_call_id": function_call_id, + "tool_call_id": function_call_id, + "decision": record.get("decision").cloned().unwrap_or_else(|| json!("deny")), + "reason": record.get("reason").cloned().unwrap_or(Value::Null), + }) +} + +async fn resume_session(iii: &III, session_id: &str) -> Result<(), String> { + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(30); + loop { + let response = iii + .trigger(TriggerRequest { + function_id: "run::resume".into(), + payload: json!({ "session_id": session_id }), + action: None, + timeout_ms: Some(5_000), + }) + .await + .map_err(|err| err.to_string())?; + + if response.get("resumed").and_then(Value::as_bool) == Some(true) { + return Ok(()); + } + + if std::time::Instant::now() >= deadline { + return Err("run::resume did not reopen approval turn before timeout".to_string()); + } + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + } +} + #[cfg(test)] mod tests { use super::*; @@ -547,7 +751,7 @@ mod tests { #[test] fn build_pending_record_sets_status_and_expiry() { let now = 1_000_000; - let rec = build_pending_record("tc-1", "write", &json!({"x": 1}), now, 60_000); + let rec = build_pending_record("s1", "tc-1", "write", &json!({"x": 1}), now, 60_000); assert_eq!(rec["status"], "pending"); assert_eq!(rec["function_call_id"], "tc-1"); assert_eq!(rec["expires_at"], 1_060_000); @@ -645,7 +849,7 @@ mod tests { bus.set( STATE_SCOPE, &pending_key("s1", "tc-1"), - build_pending_record("tc-1", "write", &json!({}), 0, 60_000), + build_pending_record("s1", "tc-1", "write", &json!({}), now_ms(), 60_000), ) .await .unwrap(); @@ -666,7 +870,117 @@ mod tests { .get(STATE_SCOPE, &pending_key("s1", "tc-1")) .await .unwrap(); - assert_eq!(stored["status"], "allow"); + assert_eq!(stored["status"], "resolved"); + assert_eq!(stored["decision"], "allow"); + } + + #[tokio::test] + async fn intercept_required_call_writes_pending_and_returns_marked_pending_block() { + let bus = InMemoryStateBus::new(); + let call = IncomingCall { + session_id: "s1".into(), + function_call_id: "tc-1".into(), + function_id: "shell::exec".into(), + args: json!({"command": "date"}), + approval_required: vec!["shell::exec".into()], + event_id: "evt".into(), + reply_stream: "replies".into(), + }; + + let reply = handle_intercept(&bus, STATE_SCOPE, &call, 1_000, 60_000).await; + + assert_eq!(reply["block"], true); + assert_eq!(reply["status"], "pending"); + assert_eq!(reply["subscriber"], "approval-gate"); + assert_eq!(reply["approval_gate"], true); + let stored = bus + .get(STATE_SCOPE, &pending_key("s1", "tc-1")) + .await + .unwrap(); + assert_eq!(stored["status"], "pending"); + assert_eq!(stored["session_id"], "s1"); + assert_eq!(stored["function_id"], "shell::exec"); + } + + #[tokio::test] + async fn intercept_non_required_call_returns_marked_allow_reply() { + let bus = InMemoryStateBus::new(); + let call = IncomingCall { + session_id: "s1".into(), + function_call_id: "tc-1".into(), + function_id: "shell::fs::ls".into(), + args: json!({}), + approval_required: vec!["shell::exec".into()], + event_id: "evt".into(), + reply_stream: "replies".into(), + }; + + let reply = handle_intercept(&bus, STATE_SCOPE, &call, 1_000, 60_000).await; + + assert_eq!(reply["block"], false); + assert_eq!(reply["subscriber"], "approval-gate"); + assert_eq!(reply["approval_gate"], true); + assert!(bus + .get(STATE_SCOPE, &pending_key("s1", "tc-1")) + .await + .is_none()); + } + + #[tokio::test] + async fn consume_returns_resolved_entries_once_and_marks_consumed() { + let bus = InMemoryStateBus::new(); + let mut rec = build_pending_record( + "s1", + "tc-1", + "shell::exec", + &json!({"command": "date"}), + 0, + 60_000, + ); + rec["status"] = json!("resolved"); + rec["decision"] = json!("allow"); + bus.set(STATE_SCOPE, &pending_key("s1", "tc-1"), rec) + .await + .unwrap(); + + let first = handle_consume(&bus, STATE_SCOPE, json!({ "session_id": "s1" })).await; + assert_eq!(first["ok"], true); + assert_eq!(first["entries"].as_array().unwrap().len(), 1); + assert_eq!(first["entries"][0]["decision"], "allow"); + + let second = handle_consume(&bus, STATE_SCOPE, json!({ "session_id": "s1" })).await; + assert_eq!(second["entries"].as_array().unwrap().len(), 0); + + let stored = bus + .get(STATE_SCOPE, &pending_key("s1", "tc-1")) + .await + .unwrap(); + assert_eq!(stored["status"], "consumed"); + } + + #[tokio::test] + async fn sweep_session_resolves_pending_as_deny_and_prevents_later_allow() { + let bus = InMemoryStateBus::new(); + bus.set( + STATE_SCOPE, + &pending_key("s1", "tc-1"), + build_pending_record("s1", "tc-1", "shell::exec", &json!({}), 0, 60_000), + ) + .await + .unwrap(); + + let sweep = handle_sweep_session(&bus, STATE_SCOPE, json!({ "session_id": "s1" })).await; + assert_eq!(sweep["ok"], true); + assert_eq!(sweep["swept"], 1); + + let allow = handle_resolve( + &bus, + STATE_SCOPE, + json!({ "session_id": "s1", "function_call_id": "tc-1", "decision": "allow" }), + ) + .await; + assert_eq!(allow["ok"], false); + assert_eq!(allow["error"], "already_resolved"); } #[tokio::test] @@ -675,7 +989,7 @@ mod tests { bus.set( STATE_SCOPE, &pending_key("s1", "tc-1"), - build_pending_record("tc-1", "write", &json!({}), 0, 60_000), + build_pending_record("s1", "tc-1", "write", &json!({}), now_ms(), 60_000), ) .await .unwrap(); @@ -697,7 +1011,7 @@ mod tests { #[tokio::test] async fn resolve_rejects_already_resolved_entry() { let bus = InMemoryStateBus::new(); - let mut rec = build_pending_record("tc-1", "write", &json!({}), 0, 60_000); + let mut rec = build_pending_record("s1", "tc-1", "write", &json!({}), 0, 60_000); rec["status"] = json!("allow"); bus.set(STATE_SCOPE, &pending_key("s1", "tc-1"), rec) .await @@ -719,11 +1033,11 @@ mod tests { bus.set( STATE_SCOPE, &pending_key("s1", "tc-1"), - build_pending_record("tc-1", "write", &json!({}), 0, 60_000), + build_pending_record("s1", "tc-1", "write", &json!({}), 0, 60_000), ) .await .unwrap(); - let mut resolved = build_pending_record("tc-2", "write", &json!({}), 0, 60_000); + let mut resolved = build_pending_record("s1", "tc-2", "write", &json!({}), 0, 60_000); resolved["status"] = json!("allow"); bus.set(STATE_SCOPE, &pending_key("s1", "tc-2"), resolved) .await @@ -731,7 +1045,7 @@ mod tests { bus.set( STATE_SCOPE, &pending_key("other", "tc-3"), - build_pending_record("tc-3", "write", &json!({}), 0, 60_000), + build_pending_record("other", "tc-3", "write", &json!({}), 0, 60_000), ) .await .unwrap(); @@ -742,67 +1056,44 @@ mod tests { assert_eq!(items[0]["function_call_id"], "tc-1"); } - use std::sync::Arc; - use std::time::Duration; + #[test] + fn state_list_values_accepts_raw_state_array() { + let out = state_list_values(json!([ + { + "session_id": "s1", + "function_call_id": "tc-1", + "status": "pending" + } + ])); - fn now_ms() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64 + assert_eq!(out.len(), 1); + assert_eq!(out[0]["function_call_id"], "tc-1"); } - #[tokio::test] - async fn await_decision_returns_allow_when_status_flips() { - let bus = Arc::new(InMemoryStateBus::new()); - let key = pending_key("s1", "tc-1"); - bus.set( - STATE_SCOPE, - &key, - build_pending_record("tc-1", "write", &json!({}), now_ms(), 5_000), - ) - .await - .unwrap(); - - let bus2 = bus.clone(); - let writer = tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(50)).await; - let mut rec = bus2.get(STATE_SCOPE, &key).await.unwrap(); - rec["status"] = json!("allow"); - bus2.set(STATE_SCOPE, &key, rec).await.unwrap(); - }); - - let decision = await_decision(&*bus, STATE_SCOPE, "s1", "tc-1", now_ms() + 5_000).await; - writer.await.unwrap(); - assert_eq!(decision, Decision::Allow); - } + #[test] + fn state_list_values_accepts_items_envelope_and_unwraps_values() { + let out = state_list_values(json!({ + "items": [ + { + "key": "s1/tc-1", + "value": { + "session_id": "s1", + "function_call_id": "tc-1", + "status": "pending" + } + } + ] + })); - #[tokio::test] - async fn await_decision_returns_deny_timeout_when_expired() { - let bus = InMemoryStateBus::new(); - let key = pending_key("s1", "tc-1"); - let _ = bus - .set( - STATE_SCOPE, - &key, - build_pending_record("tc-1", "write", &json!({}), 0, 0), - ) - .await; - let decision = await_decision(&bus, STATE_SCOPE, "s1", "tc-1", now_ms() - 10).await; - match decision { - Decision::Deny { reason } => assert_eq!(reason, "timeout"), - other => panic!("expected Deny(timeout), got {other:?}"), - } + assert_eq!(out.len(), 1); + assert_eq!(out[0]["function_call_id"], "tc-1"); } - #[tokio::test] - async fn await_decision_fail_closed_on_missing_record() { - let bus = InMemoryStateBus::new(); - let decision = await_decision(&bus, STATE_SCOPE, "s1", "tc-1", now_ms() + 1_000).await; - match decision { - Decision::Deny { reason } => assert_eq!(reason, "state_unavailable"), - other => panic!("expected Deny(state_unavailable), got {other:?}"), - } + fn now_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64 } #[tokio::test] @@ -812,7 +1103,7 @@ mod tests { .set( STATE_SCOPE, &pending_key("s1", "tc-1"), - build_pending_record("tc-1", "write", &json!({}), 0, 60_000), + build_pending_record("s1", "tc-1", "write", &json!({}), now_ms(), 60_000), ) .await; @@ -833,7 +1124,8 @@ mod tests { .get(STATE_SCOPE, &pending_key("s1", "tc-1")) .await .unwrap(); - assert_eq!(stored["status"], "deny"); + assert_eq!(stored["status"], "resolved"); + assert_eq!(stored["decision"], "deny"); assert_eq!(stored["reason"], "user clicked cancel"); } } diff --git a/approval-gate/tests/integration.rs b/approval-gate/tests/integration.rs index 7997a0447..66639fbcb 100644 --- a/approval-gate/tests/integration.rs +++ b/approval-gate/tests/integration.rs @@ -1,14 +1,13 @@ //! Engine-backed test for approval-gate. Connects to an in-process / //! local iii engine, registers the gate, fires a `before_function_call` -//! envelope on a per-test topic, posts `approval::resolve`, and asserts -//! the subscriber unblocks under 1 s. +//! envelope on a per-test topic, verifies the pending record is visible, +//! posts `approval::resolve`, and asserts `approval::consume` returns the +//! resolved entry once. //! //! Skips cleanly when no engine is reachable so `cargo test` stays green //! in CI without a running engine. -use std::time::Duration; - -use approval_gate::{register, WorkerConfig, FN_RESOLVE, STATE_SCOPE}; +use approval_gate::{register, WorkerConfig, FN_CONSUME, FN_LIST_PENDING, FN_RESOLVE, STATE_SCOPE}; use iii_sdk::{register_worker, InitOptions, TriggerRequest}; use serde_json::json; @@ -16,7 +15,7 @@ const DEFAULT_ENGINE_URL: &str = "ws://127.0.0.1:49134"; const ENGINE_PROBE_TIMEOUT_MS: u64 = 500; #[tokio::test] -async fn round_trip_allow_unblocks_under_one_second() { +async fn pending_allow_resolves_and_consumes_once() { let url = std::env::var("III_URL").unwrap_or_else(|_| DEFAULT_ENGINE_URL.to_string()); let iii = register_worker(&url, InitOptions::default()); @@ -72,43 +71,51 @@ async fn round_trip_allow_unblocks_under_one_second() { } }); - // Drive the subscriber by directly triggering its function id. - let subscriber_call = tokio::spawn({ - let iii = iii.clone(); - async move { - iii.trigger(TriggerRequest { - function_id: "policy::approval_gate".into(), - payload: envelope, - action: None, - timeout_ms: Some(10_000), - }) - .await - } - }); + let reply = iii + .trigger(TriggerRequest { + function_id: "policy::approval_gate".into(), + payload: envelope, + action: None, + timeout_ms: Some(10_000), + }) + .await + .expect("subscriber returned ok"); + assert_eq!(reply["block"], true, "subscriber reply: {reply}"); + assert_eq!(reply["status"], "pending", "subscriber reply: {reply}"); + assert_eq!( + reply["subscriber"], "approval-gate", + "subscriber reply: {reply}" + ); + assert_eq!(reply["approval_gate"], true, "subscriber reply: {reply}"); - // Wait for the gate to write the pending record before we resolve. let key = format!("{session_id}/{function_call_id}"); - let mut tries = 0; - loop { - let v = iii - .trigger(TriggerRequest { - function_id: "state::get".into(), - payload: json!({ "scope": STATE_SCOPE, "key": key }), - action: None, - timeout_ms: Some(1_000), - }) - .await - .unwrap_or(json!(null)); - if v.get("status").and_then(|s| s.as_str()) == Some("pending") { - break; - } - tries += 1; - assert!(tries < 40, "pending entry never appeared (key={key})"); - tokio::time::sleep(Duration::from_millis(50)).await; - } + let stored = iii + .trigger(TriggerRequest { + function_id: "state::get".into(), + payload: json!({ "scope": STATE_SCOPE, "key": key }), + action: None, + timeout_ms: Some(1_000), + }) + .await + .expect("state::get pending record"); + assert_eq!(stored["status"], "pending", "stored record: {stored}"); + + let pending = iii + .trigger(TriggerRequest { + function_id: FN_LIST_PENDING.into(), + payload: json!({ "session_id": session_id }), + action: None, + timeout_ms: Some(5_000), + }) + .await + .expect("list pending trigger"); + let pending_items = pending["pending"].as_array().expect("pending array"); + assert_eq!(pending_items.len(), 1, "pending response: {pending}"); + assert_eq!( + pending_items[0]["function_call_id"], function_call_id, + "pending response: {pending}" + ); - // Post the allow decision and time the unblock. - let started = std::time::Instant::now(); let resolve = iii .trigger(TriggerRequest { function_id: FN_RESOLVE.into(), @@ -124,14 +131,38 @@ async fn round_trip_allow_unblocks_under_one_second() { .expect("resolve trigger"); assert_eq!(resolve["ok"], true, "resolve response: {resolve}"); - let reply = subscriber_call + let consumed = iii + .trigger(TriggerRequest { + function_id: FN_CONSUME.into(), + payload: json!({ "session_id": session_id }), + action: None, + timeout_ms: Some(5_000), + }) + .await + .expect("consume trigger"); + let entries = consumed["entries"].as_array().expect("entries array"); + assert_eq!(entries.len(), 1, "consume response: {consumed}"); + assert_eq!( + entries[0]["decision"], "allow", + "consume response: {consumed}" + ); + assert_eq!( + entries[0]["function_call_id"], function_call_id, + "consume response: {consumed}" + ); + + let consumed_again = iii + .trigger(TriggerRequest { + function_id: FN_CONSUME.into(), + payload: json!({ "session_id": session_id }), + action: None, + timeout_ms: Some(5_000), + }) .await - .expect("subscriber task join") - .expect("subscriber returned ok"); - let elapsed = started.elapsed(); - assert!( - elapsed < Duration::from_millis(1_000), - "allow round-trip took {elapsed:?}, expected <1s", + .expect("second consume trigger"); + assert_eq!( + consumed_again["entries"].as_array().map(Vec::len), + Some(0), + "second consume response: {consumed_again}" ); - assert_eq!(reply["block"], false, "subscriber reply: {reply}"); } diff --git a/harness/web/src/App.tsx b/harness/web/src/App.tsx index 93638c16f..05b863b00 100644 --- a/harness/web/src/App.tsx +++ b/harness/web/src/App.tsx @@ -546,6 +546,7 @@ export default function App() { { + class BridgeError extends Error { + constructor(message: string) { + super(message); + this.name = "BridgeError"; + } + } + return { + BridgeError, + bridge: vi.fn(), + }; +}); + +const mockedBridge = vi.mocked(bridge); + +describe("ApprovalRow", () => { + beforeEach(() => { + mockedBridge.mockReset(); + }); + + afterEach(() => { + cleanup(); + }); + + it("surfaces approval resolve ok:false responses", async () => { + mockedBridge.mockResolvedValueOnce({ + ok: false, + error: "already_resolved", + }); + + render( + , + ); + + fireEvent.click(screen.getByRole("button", { name: "allow" })); + + await waitFor(() => { + expect(screen.getByRole("alert").textContent).toContain("already_resolved"); + }); + }); + + it("renders wake failures even when no approval is pending", () => { + render( + , + ); + + expect(screen.getByRole("alert").textContent).toContain("run::resume timed out"); + }); +}); diff --git a/harness/web/src/components/ApprovalRow.tsx b/harness/web/src/components/ApprovalRow.tsx index 25ef1b995..636d4d9d1 100644 --- a/harness/web/src/components/ApprovalRow.tsx +++ b/harness/web/src/components/ApprovalRow.tsx @@ -1,28 +1,32 @@ import { useState } from "react"; import { bridge, BridgeError } from "../bridge"; -import type { PendingApproval } from "../types"; +import type { PendingApproval, WakeFailure } from "../types"; interface Props { sessionId: string; pending: PendingApproval[]; + wakeFailures?: WakeFailure[]; } -export function ApprovalRow({ sessionId, pending }: Props) { +export function ApprovalRow({ sessionId, pending, wakeFailures = [] }: Props) { const [busyId, setBusyId] = useState(null); const [err, setErr] = useState(null); - if (pending.length === 0) return null; + if (pending.length === 0 && wakeFailures.length === 0) return null; const resolve = async (functionCallId: string, decision: "allow" | "deny") => { setBusyId(functionCallId); setErr(null); try { - await bridge<{ ok: boolean }>("approval::resolve", { + const response = await bridge<{ ok: boolean; error?: string }>("approval::resolve", { session_id: sessionId, function_call_id: functionCallId, tool_call_id: functionCallId, decision, }); + if (response.ok === false) { + setErr(response.error || "approval resolve failed"); + } } catch (e) { setErr(e instanceof BridgeError ? e.message : String(e)); } finally { @@ -69,6 +73,11 @@ export function ApprovalRow({ sessionId, pending }: Props) { {err}

) : null} + {wakeFailures.map((failure) => ( +

+ {failure.error} +

+ ))} ); } diff --git a/harness/web/src/reducer.test.ts b/harness/web/src/reducer.test.ts index ed5405b04..388087f64 100644 --- a/harness/web/src/reducer.test.ts +++ b/harness/web/src/reducer.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect } from "vitest"; +import { describe, it, expect, vi } from "vitest"; import { applyEvent } from "./reducer"; import { INITIAL_STREAM_STATE, type AgentEvent, type AgentMessage } from "./types"; @@ -89,3 +89,42 @@ describe("reducer (entry-id keyed)", () => { expect(s.messageMap.size).toBe(0); }); }); + +describe("reducer (approval wake failures)", () => { + it("collapses repeated same-error wake failures to one alert", () => { + let s = applyEvent(INITIAL_STREAM_STATE, { + type: "approval_wake_failed", + error: "run::resume timed out", + }); + s = applyEvent(s, { + type: "approval_wake_failed", + error: "run::resume timed out", + }); + + expect(s.wakeFailures).toHaveLength(1); + expect(s.wakeFailures[0].error).toBe("run::resume timed out"); + }); + + it("refreshes timestamp when the same wake failure repeats", () => { + vi.useFakeTimers(); + try { + vi.setSystemTime(new Date("2026-05-17T10:00:00Z")); + let s = applyEvent(INITIAL_STREAM_STATE, { + type: "approval_wake_failed", + error: "run::resume timed out", + }); + const firstTs = s.wakeFailures[0].ts; + + vi.setSystemTime(new Date("2026-05-17T10:00:05Z")); + s = applyEvent(s, { + type: "approval_wake_failed", + error: "run::resume timed out", + }); + + expect(s.wakeFailures).toHaveLength(1); + expect(s.wakeFailures[0].ts).toBeGreaterThan(firstTs); + } finally { + vi.useRealTimers(); + } + }); +}); diff --git a/harness/web/src/reducer.ts b/harness/web/src/reducer.ts index 3c1fba765..5aefb8c0b 100644 --- a/harness/web/src/reducer.ts +++ b/harness/web/src/reducer.ts @@ -138,6 +138,18 @@ export function applyEvent(state: StreamState, event: AgentEvent): StreamState { }; } + case "approval_wake_failed": { + const error = event.error || "run::resume failed"; + const nextFailure = { error, ts: Date.now() }; + const existing = state.wakeFailures.findIndex((f) => f.error === error); + if (existing === -1) { + return { ...state, wakeFailures: [...state.wakeFailures, nextFailure] }; + } + const wakeFailures = [...state.wakeFailures]; + wakeFailures[existing] = nextFailure; + return { ...state, wakeFailures }; + } + case "turn_start": case "tool_execution_start": case "tool_execution_update": diff --git a/harness/web/src/types.ts b/harness/web/src/types.ts index cfdccfb0c..7b495a726 100644 --- a/harness/web/src/types.ts +++ b/harness/web/src/types.ts @@ -241,6 +241,10 @@ export type AgentEvent = decision: "allow" | "deny"; reason?: string | null; tool_call_id?: string; + } + | { + type: "approval_wake_failed"; + error?: string | null; }; export interface PendingApproval { @@ -252,12 +256,18 @@ export interface PendingApproval { expires_at?: number; } +export interface WakeFailure { + error: string; + ts: number; +} + export interface StreamState { messageMap: Map; unkeyedMessages: AgentMessage[]; // backwards-compat for events without entry_id messageOrder: EntryId[]; lastEntryId: EntryId | null; pendingApprovals: PendingApproval[]; + wakeFailures: WakeFailure[]; status: "idle" | "running" | "ended"; } @@ -267,5 +277,6 @@ export const INITIAL_STREAM_STATE: StreamState = { messageOrder: [], lastEntryId: null, pendingApprovals: [], + wakeFailures: [], status: "idle", }; diff --git a/harness/web/tests/e2e/approval.spec.ts b/harness/web/tests/e2e/approval.spec.ts index 628ad2938..752328214 100644 --- a/harness/web/tests/e2e/approval.spec.ts +++ b/harness/web/tests/e2e/approval.spec.ts @@ -1,27 +1,67 @@ import { expect, test } from "@playwright/test"; +import { existsSync, rmSync, statSync } from "node:fs"; -const PROMPT = "create /tmp/harness-e2e.md with the body 'hi'"; +const MODEL = "claude-haiku-4-5"; +const HARNESS_CALL_URL = "http://127.0.0.1:3111/harness/call"; +type ModelRow = { id?: string }; + +function isDirectory(path: string): boolean { + return existsSync(path) && statSync(path).isDirectory(); +} test.describe("approval flow", () => { - test.beforeEach(async ({ page }) => { + test.beforeEach(async ({ page, request }) => { + await expect + .poll( + async () => { + const response = await request.post(HARNESS_CALL_URL, { + data: { function_id: "models::list", payload: {} }, + }); + if (!response.ok()) return false; + const body = await response.json(); + return body.models?.some((model: ModelRow) => model.id === MODEL) === true; + }, + { timeout: 30_000 }, + ) + .toBe(true); + await page.goto("/"); + await expect(page.locator(".model-select")).toContainText(MODEL, { timeout: 30_000 }); + await page.locator(".model-select").selectOption(MODEL); }); - test("allow path writes the file and renders function result block", async ({ page }) => { - await page.getByPlaceholder(/say something/i).fill(PROMPT); + test("allow path runs the approved function and renders function result block", async ({ + page, + }) => { + const path = `/tmp/harness-e2e-allow-${test.info().parallelIndex}`; + rmSync(path, { force: true, recursive: true }); + + await page.locator(".composer-input").fill( + `Use shell::fs::mkdir to create the directory ${path}.`, + ); await page.getByRole("button", { name: /send/i }).click(); const approval = page.locator(".approval"); - await expect(approval).toBeVisible({ timeout: 30_000 }); - await page.getByRole("button", { name: "allow" }).click(); - await expect(page.locator(".block-tool-result")).toBeVisible({ timeout: 30_000 }); + await expect(approval).toBeVisible({ timeout: 90_000 }); + await approval.locator(".approval-allow").click(); + await expect(page.locator(".block-tool-result").filter({ hasText: "created" })).toBeVisible({ + timeout: 90_000, + }); + await expect.poll(() => isDirectory(path), { timeout: 90_000 }).toBe(true); }); - test("deny path renders denied tool_result and does not write", async ({ page }) => { - await page.getByPlaceholder(/say something/i).fill(PROMPT); + test("deny path renders denied tool_result and does not run function", async ({ page }) => { + const path = `/tmp/harness-e2e-deny-${test.info().parallelIndex}`; + rmSync(path, { force: true, recursive: true }); + + await page.locator(".composer-input").fill( + `Use shell::fs::mkdir to create the directory ${path}.`, + ); await page.getByRole("button", { name: /send/i }).click(); - await expect(page.locator(".approval")).toBeVisible({ timeout: 30_000 }); - await page.getByRole("button", { name: "deny" }).click(); + const approval = page.locator(".approval"); + await expect(approval).toBeVisible({ timeout: 90_000 }); + await approval.locator(".approval-deny").click(); const result = page.locator(".block-tool-result[data-error='true']"); - await expect(result).toBeVisible({ timeout: 30_000 }); + await expect(result).toBeVisible({ timeout: 90_000 }); + expect(existsSync(path)).toBe(false); }); }); diff --git a/hook-fanout/src/handler.rs b/hook-fanout/src/handler.rs index c1fe238d2..a7385ac95 100644 --- a/hook-fanout/src/handler.rs +++ b/hook-fanout/src/handler.rs @@ -48,6 +48,7 @@ pub async fn execute( let envelope = build_publish_envelope(&topic, &event_id, inner.clone()); let started_at = Instant::now(); let mut publish_failed = false; + let mut publish_error: Option = None; if let Err(e) = iii .trigger(TriggerRequest { function_id: "iii::durable::publish".into(), @@ -59,6 +60,7 @@ pub async fn execute( { tracing::warn!(error = %e, %topic, "hook-fanout::publish_collect: publish trigger failed"); publish_failed = true; + publish_error = Some(e.to_string()); } let deadline = started_at + Duration::from_millis(timeout_ms.max(min_timeout)); @@ -131,11 +133,39 @@ pub async fn execute( MergeRule::PipelineLastWins => merge_pipeline_last_wins(inner.clone(), &replies), }; - Ok(json!({ + Ok(build_response( + &event_id, + replies, + merged, + publish_failed, + publish_error.as_deref(), + )) +} + +pub(crate) fn build_response( + event_id: &str, + replies: Vec, + merged: Value, + publish_failed: bool, + publish_error: Option<&str>, +) -> Value { + let mut out = json!({ "event_id": event_id, "replies": replies, "merged": merged, - })) + "publish": if publish_failed { + match publish_error { + Some(error) => json!({ "ok": false, "error": error }), + None => json!({ "ok": false }), + } + } else { + json!({ "ok": true }) + }, + }); + if publish_failed { + out["publish_failed"] = json!(true); + } + out } pub fn register(iii: &Arc, config: &Arc) { @@ -393,4 +423,36 @@ mod tests { ); assert_eq!(reason, Some("expected_replies")); } + + #[test] + fn build_response_marks_publish_ok_on_success() { + let out = build_response( + "evt-1", + vec![json!({ "block": false })], + json!({ "block": false }), + false, + None, + ); + + assert_eq!(out["event_id"], "evt-1"); + assert_eq!(out["publish"]["ok"], true); + assert!(out["publish"]["error"].is_null()); + assert!(!out.as_object().unwrap().contains_key("publish_failed")); + } + + #[test] + fn build_response_marks_publish_failed_with_error_text() { + let out = build_response( + "evt-2", + vec![], + json!({ "block": false }), + true, + Some("ws closed"), + ); + + assert_eq!(out["publish"]["ok"], false); + assert_eq!(out["publish"]["error"], "ws closed"); + assert_eq!(out["publish_failed"], true); + assert_eq!(out["merged"]["block"], false); + } } diff --git a/hook-fanout/src/lib.rs b/hook-fanout/src/lib.rs index c3cdf67ac..465a1b14f 100644 --- a/hook-fanout/src/lib.rs +++ b/hook-fanout/src/lib.rs @@ -47,10 +47,7 @@ pub struct PublishCollectResponse { pub fn merge_first_block_wins(replies: &[Value]) -> Value { for reply in replies { if reply.get("block").and_then(Value::as_bool).unwrap_or(false) { - return serde_json::json!({ - "block": true, - "reason": reply.get("reason").cloned().unwrap_or(Value::Null), - }); + return reply.clone(); } } serde_json::json!({ "block": false }) @@ -150,6 +147,29 @@ mod tests { assert_eq!(merged["reason"], "first"); } + #[test] + fn first_block_wins_preserves_full_blocking_reply() { + let replies = vec![ + json!({ "block": false }), + json!({ + "block": true, + "status": "pending", + "function_call_id": "call-1", + "function_id": "shell::exec", + "subscriber": "approval-gate", + "approval_gate": true + }), + ]; + + let merged = merge_first_block_wins(&replies); + + assert_eq!(merged["block"], true); + assert_eq!(merged["status"], "pending"); + assert_eq!(merged["function_call_id"], "call-1"); + assert_eq!(merged["subscriber"], "approval-gate"); + assert_eq!(merged["approval_gate"], true); + } + #[test] fn first_block_wins_defaults_to_no_block() { let replies = vec![json!({}), json!({ "block": false })]; diff --git a/provider-router/src/register.rs b/provider-router/src/register.rs index 762a704fc..46c580625 100644 --- a/provider-router/src/register.rs +++ b/provider-router/src/register.rs @@ -385,29 +385,30 @@ fn build_routing_request(payload: &Value) -> Option { fn register_abort(iii: &III) { let iii_for_handler = iii.clone(); iii.register_function(( - RegisterFunctionMessage::with_id("router::abort".to_string()) - .with_description("Set abort signal in iii state.".to_string()), + RegisterFunctionMessage::with_id("router::abort".to_string()).with_description( + "Set abort signal and sweep pending approvals for a session.".to_string(), + ), move |payload: Value| { let iii = iii_for_handler.clone(); async move { let session_id = required_str(&payload, "session_id")?; - // Direct state::set (was flag::set via the deleted state-flag - // worker). Convention: name "abort" maps to key - // session//abort_signal under scope "agent". - if let Err(e) = iii - .trigger(TriggerRequest { - function_id: "state::set".to_string(), - payload: json!({ - "scope": STATE_SCOPE, - "key": format!("session/{session_id}/abort_signal"), - "value": true, - }), - action: None, - timeout_ms: None, - }) - .await - { - tracing::warn!(error = %e, %session_id, "router::abort: state::set failed"); + for effect in abort_side_effects(&session_id) { + if let Err(e) = iii + .trigger(TriggerRequest { + function_id: effect.function_id.to_string(), + payload: effect.payload, + action: None, + timeout_ms: None, + }) + .await + { + tracing::warn!( + error = %e, + %session_id, + function_id = effect.function_id, + "router::abort side effect failed" + ); + } } Ok(json!({ "ok": true })) } @@ -415,6 +416,29 @@ fn register_abort(iii: &III) { )); } +#[derive(Debug, Clone, PartialEq)] +struct AbortSideEffect { + function_id: &'static str, + payload: Value, +} + +fn abort_side_effects(session_id: &str) -> Vec { + vec![ + AbortSideEffect { + function_id: "state::set", + payload: json!({ + "scope": STATE_SCOPE, + "key": format!("session/{session_id}/abort_signal"), + "value": true, + }), + }, + AbortSideEffect { + function_id: "approval::sweep_session", + payload: json!({ "session_id": session_id }), + }, + ] +} + fn register_push_steering(iii: &III) { let iii_for_handler = iii.clone(); iii.register_function(( @@ -513,4 +537,22 @@ mod tests { assert_eq!(a.error_message.as_deref(), Some("boom")); assert!(matches!(a.stop_reason, StopReason::Error)); } + + #[test] + fn abort_side_effects_set_abort_flag_then_sweep_approvals() { + let effects = abort_side_effects("sess-a"); + + assert_eq!(effects.len(), 2); + assert_eq!(effects[0].function_id, "state::set"); + assert_eq!( + effects[0].payload, + json!({ + "scope": STATE_SCOPE, + "key": "session/sess-a/abort_signal", + "value": true, + }) + ); + assert_eq!(effects[1].function_id, "approval::sweep_session"); + assert_eq!(effects[1].payload, json!({ "session_id": "sess-a" })); + } } diff --git a/turn-orchestrator/src/bootstrap.rs b/turn-orchestrator/src/bootstrap.rs index d9ca60187..b09d4b110 100644 --- a/turn-orchestrator/src/bootstrap.rs +++ b/turn-orchestrator/src/bootstrap.rs @@ -191,10 +191,7 @@ mod tests { #[test] fn leading_namespace_handles_missing_prefix() { - assert_eq!( - leading_namespace("shell/index"), - Some("shell".to_string()) - ); + assert_eq!(leading_namespace("shell/index"), Some("shell".to_string())); } #[test] diff --git a/turn-orchestrator/src/manifest.rs b/turn-orchestrator/src/manifest.rs index 3dcd2b3e8..101471a1f 100644 --- a/turn-orchestrator/src/manifest.rs +++ b/turn-orchestrator/src/manifest.rs @@ -62,8 +62,7 @@ mod tests { #[test] fn default_config_matches_struct_default() { let m = build_manifest(); - let from_struct = - serde_json::to_value(TurnOrchestratorConfig::default()).unwrap(); + let from_struct = serde_json::to_value(TurnOrchestratorConfig::default()).unwrap(); assert_eq!(m.default_config, from_struct); } } diff --git a/turn-orchestrator/src/run_start.rs b/turn-orchestrator/src/run_start.rs index ad3697fdf..8518eae6b 100644 --- a/turn-orchestrator/src/run_start.rs +++ b/turn-orchestrator/src/run_start.rs @@ -13,6 +13,7 @@ use crate::state::TurnStateRecord; pub const FUNCTION_ID: &str = "run::start"; pub const SYNC_FUNCTION_ID: &str = "run::start_and_wait"; +pub const RESUME_FUNCTION_ID: &str = "run::resume"; pub const STEP_TOPIC: &str = "turn::step_requested"; pub async fn execute(iii: III, payload: Value) -> Result { @@ -132,6 +133,39 @@ pub async fn execute_sync( } } +pub(crate) fn build_resume_record(existing: &TurnStateRecord) -> TurnStateRecord { + let mut resumed = TurnStateRecord::new(existing.session_id.clone(), existing.max_turns); + resumed.turn_count = existing.turn_count; + resumed +} + +pub(crate) fn build_resume_plan(existing: &TurnStateRecord) -> Option { + existing.is_terminal().then(|| build_resume_record(existing)) +} + +pub async fn execute_resume(iii: III, payload: Value) -> Result { + let session_id = required_str(&payload, "session_id")?; + let existing = persistence::load_record(&iii, &session_id) + .await + .ok_or_else(|| IIIError::Handler(format!("unknown session: {session_id}")))?; + + let Some(resume_record) = build_resume_plan(&existing) else { + return Ok(json!({ + "ok": true, + "session_id": session_id, + "resumed": false, + })); + }; + + persistence::save_record(&iii, &resume_record).await; + publish_step(&iii, &session_id).await; + Ok(json!({ + "ok": true, + "session_id": session_id, + "resumed": true, + })) +} + pub async fn publish_step(iii: &III, session_id: &str) { if let Err(e) = iii .trigger(TriggerRequest { @@ -172,6 +206,16 @@ pub fn register(iii: &III, cfg: &Arc) { async move { execute_sync(iii, cfg, payload).await } }, )); + let iii_resume = iii.clone(); + iii.register_function(( + RegisterFunctionMessage::with_id(RESUME_FUNCTION_ID.to_string()).with_description( + "Resume an idled approval session and publish a turn step.".to_string(), + ), + move |payload: Value| { + let iii = iii_resume.clone(); + async move { execute_resume(iii, payload).await } + }, + )); } fn required_str(payload: &Value, field: &str) -> Result { @@ -237,6 +281,37 @@ mod tests { assert_eq!(request["cwd_hash"], Value::Null); } + #[test] + fn build_resume_record_reopens_terminal_record_without_resetting_budget() { + let mut stopped = TurnStateRecord::new("sess-pending", Some(4)); + stopped.turn_count = 3; + stopped.transition_to(crate::state::TurnState::Stopped); + + let resumed = build_resume_record(&stopped); + + assert_eq!(resumed.session_id, "sess-pending"); + assert_eq!(resumed.state, crate::state::TurnState::Provisioning); + assert_eq!(resumed.turn_count, 3); + assert_eq!(resumed.max_turns, Some(4)); + assert!(!resumed.is_terminal()); + assert!(resumed.last_assistant.is_none()); + assert!(resumed.pending_function_calls.is_empty()); + } + + #[test] + fn build_resume_plan_reopens_only_terminal_records() { + let mut stopped = TurnStateRecord::new("sess-stopped", None); + stopped.transition_to(crate::state::TurnState::Stopped); + assert!(build_resume_plan(&stopped).is_some()); + + let mut active = TurnStateRecord::new("sess-active", None); + active.transition_to(crate::state::TurnState::FunctionExecute); + assert!( + build_resume_plan(&active).is_none(), + "run::resume must not publish another step while a turn is already active" + ); + } + #[test] fn build_run_request_propagates_approval_required() { let request = build_run_request(&json!({ diff --git a/turn-orchestrator/src/states/assistant.rs b/turn-orchestrator/src/states/assistant.rs index 41f020592..a00e7fd27 100644 --- a/turn-orchestrator/src/states/assistant.rs +++ b/turn-orchestrator/src/states/assistant.rs @@ -11,6 +11,29 @@ use crate::persistence; use crate::state::{TurnState, TurnStateRecord}; pub async fn handle_awaiting(iii: &III, record: &mut TurnStateRecord) -> anyhow::Result<()> { + let request = persistence::load_run_request(iii, &record.session_id).await; + if approval_required_enabled(&request) { + match crate::states::functions::consume_resolved_approval_entries(iii, &record.session_id) + .await + { + Ok(prepared) if !prepared.is_empty() => { + let executed: Vec<(FunctionCall, harness_types::FunctionResult, bool)> = Vec::new(); + persistence::save_executed_calls(iii, &record.session_id, &executed).await; + persistence::save_prepared_calls(iii, &record.session_id, &prepared).await; + record.transition_to(TurnState::FunctionExecute); + return Ok(()); + } + Ok(_) => {} + Err(err) => { + tracing::warn!( + %err, + session_id = %record.session_id, + "approval::consume failed; continuing without resolved approvals" + ); + } + } + } + if record .max_turns .map_or(false, |cap| record.turn_count >= cap) @@ -74,6 +97,13 @@ pub async fn handle_awaiting(iii: &III, record: &mut TurnStateRecord) -> anyhow: Ok(()) } +pub(crate) fn approval_required_enabled(request: &serde_json::Value) -> bool { + request + .get("approval_required") + .and_then(serde_json::Value::as_array) + .is_some_and(|items| !items.is_empty()) +} + pub async fn handle_streaming(iii: &III, record: &mut TurnStateRecord) -> anyhow::Result<()> { persistence::maybe_reload_after_compaction(iii, &record.session_id).await; let request = persistence::load_run_request(iii, &record.session_id).await; @@ -236,4 +266,15 @@ mod tests { harness_types::AgentEvent::MessageEnd { .. } )); } + + #[test] + fn approval_consume_is_only_enabled_when_run_has_required_functions() { + assert!(approval_required_enabled(&json!({ + "approval_required": ["shell::exec"] + }))); + assert!(!approval_required_enabled(&json!({ + "approval_required": [] + }))); + assert!(!approval_required_enabled(&json!({}))); + } } diff --git a/turn-orchestrator/src/states/functions.rs b/turn-orchestrator/src/states/functions.rs index 7af573fb3..4dd943d8c 100644 --- a/turn-orchestrator/src/states/functions.rs +++ b/turn-orchestrator/src/states/functions.rs @@ -6,6 +6,7 @@ use harness_types::{ }; use iii_sdk::{TriggerRequest, Value, III}; use serde_json::json; +use std::collections::HashSet; use crate::agent_call::TOOL_NAME as AGENT_CALL_TOOL_NAME; use crate::events; @@ -133,6 +134,100 @@ fn char_boundary_ceil(s: &str, idx: usize) -> usize { i } +pub(crate) fn prefilled_result_for_block( + merged: &Value, + call_id: &str, + function_id: &str, +) -> FunctionResult { + if merged.get("status").and_then(Value::as_str) == Some("pending") { + let body = json!({ + "status": "pending_approval", + "call_id": call_id, + "function_id": function_id, + "message": "Awaiting human approval. The result will be reported in a future turn." + }); + return FunctionResult { + content: vec![ContentBlock::Text(TextContent { + text: serde_json::to_string_pretty(&body).unwrap_or_else(|_| body.to_string()), + })], + details: json!({ "pending_approval": true, "call_id": call_id }), + terminate: true, + truncated: None, + }; + } + + let reason = merged + .get("reason") + .and_then(Value::as_str) + .unwrap_or("blocked") + .to_string(); + FunctionResult { + content: vec![ContentBlock::Text(TextContent { text: reason })], + details: json!({ "blocked": true }), + terminate: false, + truncated: None, + } +} + +pub(crate) fn prefilled_result_is_error(result: &FunctionResult) -> bool { + !result + .details + .get("pending_approval") + .and_then(Value::as_bool) + .unwrap_or(false) +} + +pub(crate) fn fail_closed_block_reply(phase: &str, error: &str) -> Value { + json!({ + "block": true, + "status": "denied", + "denial": { + "kind": "state_error", + "detail": { + "phase": phase, + "error": error, + }, + }, + "reason": format!("hook bus unavailable during {phase}: {error}"), + }) +} + +pub(crate) fn publish_failure_from_response( + response: &Value, + require_approval_gate_reply: bool, +) -> Option { + if let Some(publish) = response.get("publish") { + if publish.get("ok").and_then(Value::as_bool) == Some(false) { + return Some( + publish + .get("error") + .and_then(Value::as_str) + .unwrap_or("publish failed") + .to_string(), + ); + } + } + if response.get("publish_failed").and_then(Value::as_bool) == Some(true) { + return Some("publish failed".to_string()); + } + if require_approval_gate_reply { + let approval_gate_replied = response + .get("replies") + .and_then(Value::as_array) + .map(|replies| replies.iter().any(is_approval_gate_reply)) + .unwrap_or(false); + if !approval_gate_replied { + return Some("publish succeeded but approval-gate did not reply".to_string()); + } + } + None +} + +fn is_approval_gate_reply(reply: &Value) -> bool { + reply.get("approval_gate").and_then(Value::as_bool) == Some(true) + || reply.get("subscriber").and_then(Value::as_str) == Some("approval-gate") +} + /// Map `tool_use {name: "agent_call", input: {function, payload}}` back to /// a normal [`FunctionCall`] carrying the inner function id. Non-`agent_call` /// calls pass through unchanged so legacy/test fixtures keep working. @@ -183,33 +278,25 @@ pub async fn handle_prepare(iii: &III, record: &mut TurnStateRecord) -> anyhow:: let mut prepared: Vec<(FunctionCall, Option)> = Vec::with_capacity(record.pending_function_calls.len()); for fc in record.pending_function_calls.iter().cloned() { - let merged = publish_collect( + let merged = match publish_collect_checked( iii, TOPIC_BEFORE, - build_before_function_call_payload(&fc, &approval_required), + build_before_function_call_payload(&record.session_id, &fc, &approval_required), "first_block_wins", HOOK_TIMEOUT_MS, + true, ) - .await; + .await + { + Ok(merged) => merged, + Err(err) => fail_closed_block_reply("hook_publish", &err), + }; let blocked = merged .get("block") .and_then(Value::as_bool) .unwrap_or(false); - let prefilled = if blocked { - let reason = merged - .get("reason") - .and_then(Value::as_str) - .unwrap_or("blocked") - .to_string(); - Some(FunctionResult { - content: vec![ContentBlock::Text(TextContent { text: reason })], - details: json!({ "blocked": true }), - terminate: false, - truncated: None, - }) - } else { - None - }; + let prefilled = + blocked.then(|| prefilled_result_for_block(&merged, &fc.id, &fc.function_id)); prepared.push((fc, prefilled)); } @@ -237,9 +324,13 @@ pub async fn handle_execute(iii: &III, record: &mut TurnStateRecord) -> anyhow:: ) .await; if let Some(blocked) = prefilled { - persistence::upsert_executed_call(&mut results, (fc.clone(), blocked.clone(), true)); + let is_error = prefilled_result_is_error(&blocked); + persistence::upsert_executed_call( + &mut results, + (fc.clone(), blocked.clone(), is_error), + ); persistence::save_executed_calls(iii, &record.session_id, &results).await; - let evt = build_function_execution_event(&fc, &blocked, true); + let evt = build_function_execution_event(&fc, &blocked, is_error); events::emit(iii, &record.session_id, &evt).await; continue; } @@ -303,7 +394,14 @@ pub async fn handle_finalize(iii: &III, record: &mut TurnStateRecord) -> anyhow: "field_merge", HOOK_TIMEOUT_MS, ) - .await; + .await + .unwrap_or_else(|err| { + tracing::warn!( + error = %err, + "after-hook publish failed; preserving original function result", + ); + json!({}) + }); if let Ok(after) = serde_json::from_value::(merged.clone()) { result = after; } @@ -321,6 +419,7 @@ pub async fn handle_finalize(iii: &III, record: &mut TurnStateRecord) -> anyhow: } let mut messages = persistence::load_messages(iii, &record.session_id).await; + replace_pending_approval_placeholders(&mut messages, &function_results); for r in &function_results { messages.push(AgentMessage::FunctionResult(r.clone())); } @@ -329,11 +428,11 @@ pub async fn handle_finalize(iii: &III, record: &mut TurnStateRecord) -> anyhow: let Some(last_assistant) = record.last_assistant.clone() else { tracing::warn!( session_id = %record.session_id, - "FunctionFinalize reached without last_assistant; tearing down without lifecycle emit" + "FunctionFinalize reached without last_assistant; skipping lifecycle emit" ); record.function_results = function_results; record.pending_function_calls.clear(); - record.transition_to(TurnState::TearingDown); + record.transition_to(next_state_after_finalize(false, all_terminate)); return Ok(()); }; for evt in build_finalize_lifecycle(&last_assistant, &function_results) { @@ -343,12 +442,19 @@ pub async fn handle_finalize(iii: &III, record: &mut TurnStateRecord) -> anyhow: record.function_results = function_results; record.pending_function_calls.clear(); + record.transition_to(next_state_after_finalize(true, all_terminate)); + Ok(()) +} + +pub(crate) fn next_state_after_finalize( + _has_last_assistant: bool, + all_terminate: bool, +) -> TurnState { if all_terminate { - record.transition_to(TurnState::TearingDown); + TurnState::TearingDown } else { - record.transition_to(TurnState::SteeringCheck); + TurnState::SteeringCheck } - Ok(()) } pub(crate) fn executed_staging_for_new_prepare_batch( @@ -357,17 +463,107 @@ pub(crate) fn executed_staging_for_new_prepare_batch( Vec::new() } +pub(crate) fn prepared_calls_from_approval_entries( + entries: &[Value], +) -> Vec<(FunctionCall, Option)> { + entries + .iter() + .filter_map(prepared_call_from_approval_entry) + .collect() +} + +fn prepared_call_from_approval_entry( + entry: &Value, +) -> Option<(FunctionCall, Option)> { + let function_call_id = entry + .get("function_call_id") + .or_else(|| entry.get("tool_call_id")) + .and_then(Value::as_str)? + .to_string(); + let function_id = entry + .get("function_id") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + let args = entry.get("args").cloned().unwrap_or_else(|| json!({})); + let decision = entry + .get("decision") + .and_then(Value::as_str) + .unwrap_or("deny"); + let fc = FunctionCall { + id: function_call_id.clone(), + function_id, + arguments: args, + }; + if decision == "allow" { + return Some((fc, None)); + } + + let reason = entry + .get("reason") + .and_then(Value::as_str) + .unwrap_or("denied"); + let text = if reason == "timed_out" || decision == "timed_out" { + "approval timed out before resolution".to_string() + } else { + format!("approval denied: {reason}") + }; + let result = FunctionResult { + content: vec![ContentBlock::Text(TextContent { text })], + details: json!({ + "approval_denied": true, + "decision": decision, + "reason": reason, + "resolved_via_approval_gate": true, + "call_id": function_call_id, + }), + terminate: false, + truncated: None, + }; + Some((fc, Some(result))) +} + /// Pure helper: inner payload for the `agent::before_function_call` topic. pub(crate) fn build_before_function_call_payload( + session_id: &str, fc: &FunctionCall, approval_required: &[String], ) -> Value { json!({ + "session_id": session_id, "function_call": fc, "approval_required": approval_required, }) } +pub(crate) async fn consume_resolved_approval_entries( + iii: &III, + session_id: &str, +) -> Result)>, String> { + let response = iii + .trigger(TriggerRequest { + function_id: "approval::consume".into(), + payload: json!({ "session_id": session_id }), + action: None, + timeout_ms: Some(5_000), + }) + .await + .map_err(|err| err.to_string())?; + if response.get("ok").and_then(Value::as_bool) == Some(false) { + return Err(response + .get("error") + .and_then(Value::as_str) + .unwrap_or("approval::consume failed") + .to_string()); + } + let entries = response + .get("entries") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + Ok(prepared_calls_from_approval_entries(&entries)) +} + /// Pure helper: build [`AgentEvent::FunctionExecutionEnd`] for one call. pub(crate) fn build_function_execution_event( fc: &FunctionCall, @@ -400,13 +596,49 @@ pub(crate) fn build_finalize_lifecycle( out } +pub(crate) fn replace_pending_approval_placeholders( + messages: &mut Vec, + replacements: &[FunctionResultMessage], +) { + let replacement_ids = replacements + .iter() + .map(|r| r.function_call_id.as_str()) + .collect::>(); + if replacement_ids.is_empty() { + return; + } + messages.retain(|message| match message { + AgentMessage::FunctionResult(result) => { + let is_replaced_call = replacement_ids.contains(result.function_call_id.as_str()); + let is_pending_placeholder = result + .details + .get("pending_approval") + .and_then(Value::as_bool) + .unwrap_or(false); + !(is_replaced_call && is_pending_placeholder) + } + _ => true, + }); +} + async fn publish_collect( iii: &III, topic: &str, inner: Value, merge_rule: &str, timeout_ms: u64, -) -> Value { +) -> Result { + publish_collect_checked(iii, topic, inner, merge_rule, timeout_ms, false).await +} + +async fn publish_collect_checked( + iii: &III, + topic: &str, + inner: Value, + merge_rule: &str, + timeout_ms: u64, + require_approval_gate_reply: bool, +) -> Result { let payload = json!({ "topic": topic, "payload": inner, @@ -420,9 +652,14 @@ async fn publish_collect( timeout_ms: None, }) .await - .ok() - .and_then(|v| v.get("merged").cloned()) - .unwrap_or_else(|| json!({})) + .map_err(|err| err.to_string()) + .and_then(|response| { + if let Some(err) = publish_failure_from_response(&response, require_approval_gate_reply) { + Err(err) + } else { + Ok(response.get("merged").cloned().unwrap_or_else(|| json!({}))) + } + }) } #[cfg(test)] @@ -623,6 +860,82 @@ mod tests { } } + #[test] + fn pending_block_prefill_terminates_without_error_flag() { + let merged = json!({ + "block": true, + "status": "pending", + "reason": "approval required", + }); + + let result = prefilled_result_for_block(&merged, "tc-1", "shell::exec"); + + assert!(result.terminate, "pending approvals must stop the turn"); + assert_eq!(result.details["pending_approval"], true); + assert!(!prefilled_result_is_error(&result)); + assert!(matches!( + result.content.first(), + Some(ContentBlock::Text(text)) + if text.text.contains("\"status\": \"pending_approval\"") + && text.text.contains("\"call_id\": \"tc-1\"") + )); + } + + #[test] + fn hard_block_prefill_remains_error_without_terminating() { + let merged = json!({ + "block": true, + "status": "denied", + "reason": "blocked by policy", + }); + + let result = prefilled_result_for_block(&merged, "tc-2", "shell::exec"); + + assert!(!result.terminate); + assert_eq!(result.details["blocked"], true); + assert!(prefilled_result_is_error(&result)); + } + + #[test] + fn approval_allow_entry_prepares_dispatchable_function_call() { + let prepared = prepared_calls_from_approval_entries(&[json!({ + "function_call_id": "tc-1", + "function_id": "shell::exec", + "args": { "command": "date" }, + "decision": "allow", + })]); + + assert_eq!(prepared.len(), 1); + assert_eq!(prepared[0].0.id, "tc-1"); + assert_eq!(prepared[0].0.function_id, "shell::exec"); + assert_eq!(prepared[0].0.arguments, json!({ "command": "date" })); + assert!( + prepared[0].1.is_none(), + "allow must execute through the normal dispatch path" + ); + } + + #[test] + fn approval_deny_entry_prepares_prefilled_result_without_dispatch() { + let prepared = prepared_calls_from_approval_entries(&[json!({ + "function_call_id": "tc-1", + "function_id": "shell::exec", + "args": { "command": "date" }, + "decision": "deny", + "reason": "timed_out", + })]); + + let result = prepared[0].1.as_ref().expect("deny should prefill"); + assert_eq!(prepared[0].0.id, "tc-1"); + assert!(prefilled_result_is_error(result)); + assert_eq!(result.details["approval_denied"], true); + assert_eq!(result.details["reason"], "timed_out"); + assert!(matches!( + result.content.first(), + Some(ContentBlock::Text(text)) if text.text.contains("approval timed out") + )); + } + #[test] fn before_function_call_payload_carries_approval_required() { let fc = FunctionCall { @@ -631,7 +944,8 @@ mod tests { arguments: json!({"path": "/tmp/x"}), }; let approval_required = vec!["shell::fs::write".to_string()]; - let inner = build_before_function_call_payload(&fc, &approval_required); + let inner = build_before_function_call_payload("sess-a", &fc, &approval_required); + assert_eq!(inner["session_id"], "sess-a"); assert_eq!(inner["function_call"]["id"], "tc-1"); assert_eq!(inner["approval_required"], json!(["shell::fs::write"]),); } @@ -643,7 +957,7 @@ mod tests { function_id: "shell::fs::ls".into(), arguments: json!({}), }; - let inner = build_before_function_call_payload(&fc, &[]); + let inner = build_before_function_call_payload("sess-a", &fc, &[]); assert_eq!(inner["approval_required"], json!([])); } @@ -660,6 +974,75 @@ mod tests { assert!(matches!(evs.last(), Some(AgentEvent::TurnEnd { .. }))); } + #[test] + fn approval_resume_replaces_pending_placeholder_for_same_call_id() { + let mut messages = vec![ + AgentMessage::User(harness_types::UserMessage { + content: vec![ContentBlock::Text(TextContent { text: "run".into() })], + timestamp: 0, + }), + AgentMessage::FunctionResult(FunctionResultMessage { + function_call_id: "tc-1".into(), + function_id: "shell::fs::mkdir".into(), + content: vec![ContentBlock::Text(TextContent { + text: "pending approval".into(), + })], + details: json!({ "pending_approval": true }), + is_error: false, + timestamp: 1, + }), + AgentMessage::FunctionResult(FunctionResultMessage { + function_call_id: "tc-2".into(), + function_id: "shell::fs::ls".into(), + content: vec![ContentBlock::Text(TextContent { text: "ok".into() })], + details: json!({}), + is_error: false, + timestamp: 2, + }), + ]; + let replacement = FunctionResultMessage { + function_call_id: "tc-1".into(), + function_id: "shell::fs::mkdir".into(), + content: vec![ContentBlock::Text(TextContent { + text: "created".into(), + })], + details: json!({ "created": true }), + is_error: false, + timestamp: 3, + }; + + replace_pending_approval_placeholders(&mut messages, &[replacement]); + + assert_eq!(messages.len(), 2); + assert!(!messages.iter().any(|message| matches!( + message, + AgentMessage::FunctionResult(result) + if result.function_call_id == "tc-1" + && result.details.get("pending_approval").and_then(Value::as_bool) == Some(true) + ))); + assert!(messages.iter().any(|message| matches!( + message, + AgentMessage::FunctionResult(result) if result.function_call_id == "tc-2" + ))); + } + + #[test] + fn approval_resume_keeps_non_placeholder_result_with_same_call_id() { + let mut messages = vec![AgentMessage::FunctionResult(FunctionResultMessage { + function_call_id: "tc-1".into(), + function_id: "shell::fs::mkdir".into(), + content: vec![ContentBlock::Text(TextContent { text: "old".into() })], + details: json!({ "created": true }), + is_error: false, + timestamp: 1, + })]; + let replacement = function_result_msg("shell::fs::mkdir", false); + + replace_pending_approval_placeholders(&mut messages, &[replacement]); + + assert_eq!(messages.len(), 1); + } + /// policy-denylist subscribes to this topic by exact name. #[test] fn topic_constants_are_stable() { @@ -676,13 +1059,68 @@ mod tests { function_id: "shell::fs::ls".into(), arguments: json!({"path": "/tmp"}), }; - let inner = build_before_function_call_payload(&fc, &[]); + let inner = build_before_function_call_payload("sess-a", &fc, &[]); + assert_eq!(inner["session_id"], "sess-a"); assert_eq!(inner["function_call"]["id"], "tc-1"); assert_eq!(inner["function_call"]["function_id"], "shell::fs::ls"); assert_eq!(inner["function_call"]["arguments"], json!({"path": "/tmp"})); assert!(inner.get("approval_required").is_some()); } + #[test] + fn publish_failure_from_response_fails_closed_on_publish_error() { + let response = json!({ + "event_id": "evt", + "replies": [], + "merged": { "block": false }, + "publish": { "ok": false, "error": "ws closed" } + }); + + assert_eq!( + publish_failure_from_response(&response, true).as_deref(), + Some("ws closed"), + ); + } + + #[test] + fn publish_failure_from_response_requires_approval_gate_reply_for_before_hook() { + let empty = json!({ + "event_id": "evt", + "replies": [], + "merged": { "block": false }, + "publish": { "ok": true } + }); + assert!(publish_failure_from_response(&empty, true).is_some()); + + let non_gate = json!({ + "event_id": "evt", + "replies": [{ "block": false, "subscriber": "policy-denylist" }], + "merged": { "block": false }, + "publish": { "ok": true } + }); + assert!(publish_failure_from_response(&non_gate, true).is_some()); + + let gate = json!({ + "event_id": "evt", + "replies": [{ "block": false, "subscriber": "approval-gate", "approval_gate": true }], + "merged": { "block": false }, + "publish": { "ok": true } + }); + assert!(publish_failure_from_response(&gate, true).is_none()); + } + + #[test] + fn publish_failure_from_response_allows_zero_replies_for_after_hook() { + let response = json!({ + "event_id": "evt", + "replies": [], + "merged": { "block": false }, + "publish": { "ok": true } + }); + + assert!(publish_failure_from_response(&response, false).is_none()); + } + #[test] fn handle_finalize_does_not_expect_last_assistant() { let src = include_str!("functions.rs"); @@ -788,4 +1226,22 @@ mod tests { assert!(t >= 1024, "threshold should never round down below 1 KB"); assert!(t <= 1_000_000, "threshold should be sane"); } + + #[test] + fn finalize_without_last_assistant_still_continues_after_function_results() { + assert_eq!( + next_state_after_finalize(false, false), + TurnState::SteeringCheck, + "approval resume records have no last_assistant, but allowed function results \ + must still flow through steering into the next assistant turn" + ); + } + + #[test] + fn finalize_without_last_assistant_tears_down_when_all_results_terminate() { + assert_eq!( + next_state_after_finalize(false, true), + TurnState::TearingDown + ); + } } diff --git a/turn-orchestrator/src/states/provisioning.rs b/turn-orchestrator/src/states/provisioning.rs index 70e60db61..960835f76 100644 --- a/turn-orchestrator/src/states/provisioning.rs +++ b/turn-orchestrator/src/states/provisioning.rs @@ -213,6 +213,9 @@ mod tests { "body": "real body content here", "modified_at": "2026-05-13T00:00:00+00:00" }); - assert_eq!(response_to_string(&resp).as_deref(), Some("real body content here")); + assert_eq!( + response_to_string(&resp).as_deref(), + Some("real body content here") + ); } } diff --git a/turn-orchestrator/src/states/tearing_down.rs b/turn-orchestrator/src/states/tearing_down.rs index 0ef8ca2d3..4d4d446f1 100644 --- a/turn-orchestrator/src/states/tearing_down.rs +++ b/turn-orchestrator/src/states/tearing_down.rs @@ -9,6 +9,30 @@ use crate::persistence; use crate::state::{TurnState, TurnStateRecord}; pub async fn handle(iii: &III, record: &mut TurnStateRecord) -> anyhow::Result<()> { + let request = persistence::load_run_request(iii, &record.session_id).await; + if crate::states::assistant::approval_required_enabled(&request) { + match crate::states::functions::consume_resolved_approval_entries(iii, &record.session_id) + .await + { + Ok(prepared) if !prepared.is_empty() => { + let executed = + crate::states::functions::executed_staging_for_new_prepare_batch(&[]); + persistence::save_executed_calls(iii, &record.session_id, &executed).await; + persistence::save_prepared_calls(iii, &record.session_id, &prepared).await; + record.transition_to(TurnState::FunctionExecute); + return Ok(()); + } + Ok(_) => {} + Err(err) => { + tracing::warn!( + %err, + session_id = %record.session_id, + "approval::consume failed during teardown; stopping session" + ); + } + } + } + if let Some(sandbox_id) = persistence::load_sandbox_id(iii, &record.session_id).await { if let Err(e) = iii .trigger(TriggerRequest { diff --git a/turn-orchestrator/src/system_prompt.rs b/turn-orchestrator/src/system_prompt.rs index 3bfc1f917..efd3b34df 100644 --- a/turn-orchestrator/src/system_prompt.rs +++ b/turn-orchestrator/src/system_prompt.rs @@ -124,13 +124,21 @@ mod tests { #[test] fn override_returns_verbatim_when_non_empty() { - let out = build(&[skill("iii://iii", "body")], Some(Path::new("/tmp")), Some("custom")); + let out = build( + &[skill("iii://iii", "body")], + Some(Path::new("/tmp")), + Some("custom"), + ); assert_eq!(out, "custom"); } #[test] fn empty_override_falls_through_to_canonical() { - let out = build(&[skill("iii://iii", "body")], Some(Path::new("/tmp")), Some("")); + let out = build( + &[skill("iii://iii", "body")], + Some(Path::new("/tmp")), + Some(""), + ); assert!(out.contains("You are an iii agent worker")); assert!(out.contains("/tmp")); assert!(out.contains("body")); @@ -179,7 +187,10 @@ mod tests { ); let pos_iii = out.find("AAA").expect("first skill body must be present"); let pos_shell = out.find("BBB").expect("second skill body must be present"); - assert!(pos_iii < pos_shell, "skills must appear in config-list order"); + assert!( + pos_iii < pos_shell, + "skills must appear in config-list order" + ); } #[test] @@ -192,7 +203,11 @@ mod tests { #[test] fn cwd_appears_between_preamble_and_skills() { - let out = build(&[skill("iii://iii", "BODY")], Some(Path::new("/work/proj")), None); + let out = build( + &[skill("iii://iii", "BODY")], + Some(Path::new("/work/proj")), + None, + ); let pos_preamble = out.find("iii agent worker").unwrap(); let pos_cwd = out.find("/work/proj").unwrap(); let pos_body = out.find("BODY").unwrap(); @@ -225,7 +240,11 @@ mod tests { #[test] fn large_override_returns_same_length() { let huge = "a".repeat(1_000_000); - let out = build(&[skill("iii://iii", "body")], Some(Path::new("/tmp")), Some(&huge)); + let out = build( + &[skill("iii://iii", "body")], + Some(Path::new("/tmp")), + Some(&huge), + ); assert_eq!(out.len(), 1_000_000); assert_eq!(out, huge); } From 55891f2ec9623eb3035e188743c6df575c0531c6 Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Sun, 17 May 2026 23:15:14 -0300 Subject: [PATCH 2/4] Keep resumed approval responses on the active session path The iii-state session store was reconstructing entry order by UUID/id after state::list, but active_path(None) treats the last loaded entry as the current leaf. Approval resume can append the final assistant response with an id that sorts before earlier entries, so the backend state contains the response while the web console reload picks the wrong path. Sorting loaded entries by entry timestamp preserves append order and keeps the resumed response visible. Constraint: state::list returns values without storage keys or insertion order Rejected: Add web-console polling as the fix | the missing result was already present in state and the root issue was active-path reconstruction Confidence: high Scope-risk: narrow Directive: Do not change active_path(None) or state-store ordering without verifying approval resume reloads in the web console Tested: git diff --check Tested: cargo test --lib load_entries_lists_session_scope_and_sorts_by_timestamp -- --nocapture Tested: cargo test library suite passed as part of cargo test before integration tests Not-tested: full session cargo test; tests/inbox_integration.rs end_to_end_push_via_iii_sdk failed because session-inbox::push was not registered in the local iii runtime --- session/src/tree/mod.rs | 9 ++++++++ session/src/tree/store_iii_state.rs | 35 ++++++++++++++++++++--------- 2 files changed, 34 insertions(+), 10 deletions(-) diff --git a/session/src/tree/mod.rs b/session/src/tree/mod.rs index c3a0c47b7..4ff38035b 100644 --- a/session/src/tree/mod.rs +++ b/session/src/tree/mod.rs @@ -100,6 +100,15 @@ impl SessionEntry { } } + pub fn timestamp(&self) -> i64 { + match self { + Self::Message { timestamp, .. } + | Self::CustomMessage { timestamp, .. } + | Self::BranchSummary { timestamp, .. } + | Self::Compaction { timestamp, .. } => *timestamp, + } + } + /// Replace the entry's id, returning a new entry. fn with_id(mut self, new_id: String) -> Self { match &mut self { diff --git a/session/src/tree/store_iii_state.rs b/session/src/tree/store_iii_state.rs index 109c38bb6..507bdbfb2 100644 --- a/session/src/tree/store_iii_state.rs +++ b/session/src/tree/store_iii_state.rs @@ -5,8 +5,8 @@ //! - Scope `session_tree:`, key ``, value `SessionEntry` //! - Scope `session_tree_meta`, key ``, value `SessionMeta` //! -//! `state::list` returns values without keys, so each entry's `id` field is -//! used to recover ordering when loading entries (Task E12). +//! `state::list` returns values without keys, so each entry's timestamp is +//! used to recover append ordering when loading entries (Task E12). use std::sync::Arc; @@ -109,7 +109,11 @@ impl SessionStore for IiiStateSessionStore { .map_err(|e| SessionError::Storage(format!("deserialize SessionEntry: {e}"))) }) .collect::, _>>()?; - entries.sort_by(|a, b| a.id().cmp(b.id())); + entries.sort_by(|a, b| { + a.timestamp() + .cmp(&b.timestamp()) + .then_with(|| a.id().cmp(b.id())) + }); Ok(entries) } @@ -364,11 +368,22 @@ mod tests { } #[tokio::test] - async fn load_entries_lists_session_scope_and_sorts() -> anyhow::Result<()> { - let e1 = sample_entry("01"); - let e2 = sample_entry("02"); - // Mock returns entries in REVERSE order; load_entries must sort by id. - let response = serde_json::json!([serde_json::to_value(&e2)?, serde_json::to_value(&e1)?,]); + async fn load_entries_lists_session_scope_and_sorts_by_timestamp() -> anyhow::Result<()> { + let mut later = sample_entry("aaa-later-id"); + let mut earlier = sample_entry("zzz-earlier-id"); + if let SessionEntry::CustomMessage { timestamp, .. } = &mut later { + *timestamp = 2; + } + if let SessionEntry::CustomMessage { timestamp, .. } = &mut earlier { + *timestamp = 1; + } + // Mock returns entries in reverse append order. The ids intentionally + // sort opposite the timestamps, so UUID/id ordering would pick the + // wrong active leaf after reload. + let response = serde_json::json!([ + serde_json::to_value(&later)?, + serde_json::to_value(&earlier)?, + ]); let mock = Arc::new(MockTrigger::new(vec![Ok(response)])); let store = IiiStateSessionStore::new(mock.clone()); @@ -377,8 +392,8 @@ mod tests { .await .map_err(|e| anyhow::anyhow!("{e}"))?; assert_eq!(entries.len(), 2); - assert_eq!(entries[0].id(), "01"); - assert_eq!(entries[1].id(), "02"); + assert_eq!(entries[0].id(), "zzz-earlier-id"); + assert_eq!(entries[1].id(), "aaa-later-id"); let calls = mock.calls.lock().unwrap(); assert_eq!(calls.len(), 1); From 9c954efe4f0007648ad33636592fb9ada6120402 Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Mon, 18 May 2026 08:02:35 -0300 Subject: [PATCH 3/4] refactor(approval-gate): wake run::resume via notifier instead of polling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The resolver's 250 ms `run::resume` retry loop racing the executor's park is replaced with an in-process `AwaitingApproval` notifier. The executor signals on terminal save in `subscriber::execute`; `execute_resume` does fast-path → arm Notified → recheck → bounded wait. The resolver fires one `run::resume` and waits server-side. Median resume latency drops from up-to-250ms quantized polling to the microseconds it takes Notify::notify_waiters to wake the parked task. The 30 s ceiling is unchanged from the old loop's deadline; it's the failure timeout, not the expected wait. --- approval-gate/src/lib.rs | 39 +++++----- turn-orchestrator/src/awaiting.rs | 111 ++++++++++++++++++++++++++++ turn-orchestrator/src/lib.rs | 3 + turn-orchestrator/src/register.rs | 6 +- turn-orchestrator/src/run_start.rs | 76 +++++++++++++++---- turn-orchestrator/src/subscriber.rs | 24 ++++-- 6 files changed, 215 insertions(+), 44 deletions(-) create mode 100644 turn-orchestrator/src/awaiting.rs diff --git a/approval-gate/src/lib.rs b/approval-gate/src/lib.rs index 39d21dee1..865abe9d9 100644 --- a/approval-gate/src/lib.rs +++ b/approval-gate/src/lib.rs @@ -530,7 +530,7 @@ pub fn register(iii: &III, cfg: &WorkerConfig) -> anyhow::Result { ) .await; } - if let Err(err) = resume_session(&iii, &session_id).await { + if let Err(err) = trigger_resume(&iii, &session_id).await { write_event( &iii, &session_id, @@ -658,28 +658,25 @@ fn approval_resolved_event(function_call_id: &str, record: &Value) -> Value { }) } -async fn resume_session(iii: &III, session_id: &str) -> Result<(), String> { - let deadline = std::time::Instant::now() + std::time::Duration::from_secs(30); - loop { - let response = iii - .trigger(TriggerRequest { - function_id: "run::resume".into(), - payload: json!({ "session_id": session_id }), - action: None, - timeout_ms: Some(5_000), - }) - .await - .map_err(|err| err.to_string())?; - - if response.get("resumed").and_then(Value::as_bool) == Some(true) { - return Ok(()); - } +/// Fire `run::resume` once and rely on the handler to block in-process +/// until the executor parks the session. The 35 s ceiling sits just +/// above the server-side 30 s wait so a server timeout surfaces as +/// `resumed=false` rather than as a transport timeout here. +async fn trigger_resume(iii: &III, session_id: &str) -> Result<(), String> { + let response = iii + .trigger(TriggerRequest { + function_id: "run::resume".into(), + payload: json!({ "session_id": session_id }), + action: None, + timeout_ms: Some(35_000), + }) + .await + .map_err(|err| err.to_string())?; - if std::time::Instant::now() >= deadline { - return Err("run::resume did not reopen approval turn before timeout".to_string()); - } - tokio::time::sleep(std::time::Duration::from_millis(250)).await; + if response.get("resumed").and_then(Value::as_bool) == Some(true) { + return Ok(()); } + Err("run::resume did not reopen approval turn".to_string()) } #[cfg(test)] diff --git a/turn-orchestrator/src/awaiting.rs b/turn-orchestrator/src/awaiting.rs new file mode 100644 index 000000000..46af6735b --- /dev/null +++ b/turn-orchestrator/src/awaiting.rs @@ -0,0 +1,111 @@ +//! In-process notifier that wakes `run::resume` the moment the FSM +//! persists a terminal record. Replaces the 250 ms poll loop the +//! approval-gate resolver used to run. +//! +//! The same process owns both the executor (which signals on terminal +//! save) and `execute_resume` (which waits). Cross-process resume +//! requires switching to a state-bus subscriber instead — see the PR +//! description for the variant we deferred. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use tokio::sync::Notify; + +#[derive(Clone, Default)] +pub struct AwaitingApproval { + inner: Arc>>>, +} + +impl AwaitingApproval { + pub fn new() -> Self { + Self::default() + } + + pub fn slot(&self, session_id: &str) -> Arc { + let mut guard = self.inner.lock().expect("AwaitingApproval mutex poisoned"); + guard + .entry(session_id.to_string()) + .or_default() + .clone() + } + + pub fn signal(&self, session_id: &str) { + let slot = self.slot(session_id); + slot.notify_waiters(); + } + + pub fn clear(&self, session_id: &str) { + let mut guard = self.inner.lock().expect("AwaitingApproval mutex poisoned"); + guard.remove(session_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + #[tokio::test] + async fn signal_after_arming_wakes_waiter() { + let awaiting = AwaitingApproval::new(); + let slot = awaiting.slot("s1"); + let notified = slot.notified(); + + let signaller = { + let awaiting = awaiting.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(10)).await; + awaiting.signal("s1"); + }) + }; + + tokio::time::timeout(Duration::from_secs(1), notified) + .await + .expect("waiter should wake when signal fires"); + signaller.await.unwrap(); + } + + #[tokio::test] + async fn signal_without_waiter_is_noop() { + let awaiting = AwaitingApproval::new(); + awaiting.signal("nobody-home"); + // Should not panic, should not block. + } + + #[tokio::test] + async fn clear_drops_slot_so_next_call_gets_fresh_notify() { + let awaiting = AwaitingApproval::new(); + let slot_a = awaiting.slot("s1"); + awaiting.clear("s1"); + let slot_b = awaiting.slot("s1"); + // Different Arc identity after clear. + assert!(!Arc::ptr_eq(&slot_a, &slot_b)); + } + + #[tokio::test] + async fn slot_returns_same_arc_for_same_session() { + let awaiting = AwaitingApproval::new(); + let a = awaiting.slot("s1"); + let b = awaiting.slot("s1"); + assert!(Arc::ptr_eq(&a, &b)); + } + + /// Documents the constraint that drives `execute_resume`'s + /// arm-then-recheck pattern: a `signal` that fires before any + /// waiter has called `notified()` is lost. Without re-checking + /// persistence state between arming and awaiting, a parked + /// session whose signal arrived during the first load would + /// hang until the 30 s timeout. Don't "simplify" the recheck + /// away. + #[tokio::test] + async fn signal_before_notified_is_lost() { + let awaiting = AwaitingApproval::new(); + awaiting.signal("s1"); + let slot = awaiting.slot("s1"); + let notified = slot.notified(); + + let elapsed = tokio::time::timeout(Duration::from_millis(50), notified).await; + assert!(elapsed.is_err(), "signal fired before arming must not wake"); + } +} diff --git a/turn-orchestrator/src/lib.rs b/turn-orchestrator/src/lib.rs index 3436033b1..d65a0e719 100644 --- a/turn-orchestrator/src/lib.rs +++ b/turn-orchestrator/src/lib.rs @@ -1,6 +1,7 @@ //! Durable session state machine. See plan doc for details. pub mod agent_call; +pub mod awaiting; pub mod bootstrap; pub mod config; pub mod events; @@ -14,6 +15,8 @@ pub mod subscriber; pub mod system_prompt; pub mod transitions; +pub use awaiting::AwaitingApproval; + pub use config::TurnOrchestratorConfig; pub use register::register_with_iii; pub use state::{ diff --git a/turn-orchestrator/src/register.rs b/turn-orchestrator/src/register.rs index 13f4f1f88..e280ff274 100644 --- a/turn-orchestrator/src/register.rs +++ b/turn-orchestrator/src/register.rs @@ -7,6 +7,7 @@ use iii_sdk::{RegisterTriggerInput, III}; use serde_json::json; use crate::agent_call; +use crate::awaiting::AwaitingApproval; use crate::config::TurnOrchestratorConfig; use crate::run_start::{self, STEP_TOPIC}; use crate::subscriber::{self, FUNCTION_ID as STEP_FN_ID}; @@ -15,9 +16,10 @@ pub async fn register_with_iii( iii: &Arc, cfg: &Arc, ) -> anyhow::Result<()> { - run_start::register(iii, cfg); + let awaiting = AwaitingApproval::new(); + run_start::register(iii, cfg, awaiting.clone()); agent_call::register(iii); - subscriber::register(iii, cfg); + subscriber::register(iii, cfg, awaiting); iii.register_trigger(RegisterTriggerInput { trigger_type: "durable:subscriber".into(), diff --git a/turn-orchestrator/src/run_start.rs b/turn-orchestrator/src/run_start.rs index 8518eae6b..2400d5c9f 100644 --- a/turn-orchestrator/src/run_start.rs +++ b/turn-orchestrator/src/run_start.rs @@ -6,6 +6,7 @@ use harness_types::{AgentEvent, AgentMessage}; use iii_sdk::{IIIError, RegisterFunctionMessage, TriggerRequest, Value, III}; use serde_json::json; +use crate::awaiting::AwaitingApproval; use crate::config::TurnOrchestratorConfig; use crate::events; use crate::persistence; @@ -143,22 +144,67 @@ pub(crate) fn build_resume_plan(existing: &TurnStateRecord) -> Option Result { +/// Max time `execute_resume` will wait for the executor to park the +/// session. The happy path completes within milliseconds — this is the +/// failure ceiling, not the expected latency. +const RESUME_WAIT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + +pub async fn execute_resume( + iii: III, + awaiting: AwaitingApproval, + payload: Value, +) -> Result { let session_id = required_str(&payload, "session_id")?; - let existing = persistence::load_record(&iii, &session_id) + + if let Some(plan) = try_resume_now(&iii, &session_id).await? { + return finish_resume(&iii, &awaiting, &session_id, plan).await; + } + + // Arm a Notified future *before* the second load so a signal that + // fires between the two checks is captured. Same lost-wake guard + // pattern as tokio::sync::Notify's documented usage. + let slot = awaiting.slot(&session_id); + let notified = slot.notified(); + tokio::pin!(notified); + + if let Some(plan) = try_resume_now(&iii, &session_id).await? { + return finish_resume(&iii, &awaiting, &session_id, plan).await; + } + + match tokio::time::timeout(RESUME_WAIT_TIMEOUT, notified).await { + Ok(()) => { + if let Some(plan) = try_resume_now(&iii, &session_id).await? { + finish_resume(&iii, &awaiting, &session_id, plan).await + } else { + Ok(json!({ + "ok": true, + "session_id": session_id, + "resumed": false, + })) + } + } + Err(_) => Err(IIIError::Handler(format!( + "run::resume timed out waiting for session {session_id} to park" + ))), + } +} + +async fn try_resume_now(iii: &III, session_id: &str) -> Result, IIIError> { + let existing = persistence::load_record(iii, session_id) .await .ok_or_else(|| IIIError::Handler(format!("unknown session: {session_id}")))?; + Ok(build_resume_plan(&existing)) +} - let Some(resume_record) = build_resume_plan(&existing) else { - return Ok(json!({ - "ok": true, - "session_id": session_id, - "resumed": false, - })); - }; - - persistence::save_record(&iii, &resume_record).await; - publish_step(&iii, &session_id).await; +async fn finish_resume( + iii: &III, + awaiting: &AwaitingApproval, + session_id: &str, + plan: TurnStateRecord, +) -> Result { + persistence::save_record(iii, &plan).await; + publish_step(iii, session_id).await; + awaiting.clear(session_id); Ok(json!({ "ok": true, "session_id": session_id, @@ -183,7 +229,7 @@ pub async fn publish_step(iii: &III, session_id: &str) { } } -pub fn register(iii: &III, cfg: &Arc) { +pub fn register(iii: &III, cfg: &Arc, awaiting: AwaitingApproval) { let iii_async = iii.clone(); iii.register_function(( RegisterFunctionMessage::with_id(FUNCTION_ID.to_string()) @@ -207,13 +253,15 @@ pub fn register(iii: &III, cfg: &Arc) { }, )); let iii_resume = iii.clone(); + let awaiting_resume = awaiting.clone(); iii.register_function(( RegisterFunctionMessage::with_id(RESUME_FUNCTION_ID.to_string()).with_description( "Resume an idled approval session and publish a turn step.".to_string(), ), move |payload: Value| { let iii = iii_resume.clone(); - async move { execute_resume(iii, payload).await } + let awaiting = awaiting_resume.clone(); + async move { execute_resume(iii, awaiting, payload).await } }, )); } diff --git a/turn-orchestrator/src/subscriber.rs b/turn-orchestrator/src/subscriber.rs index 792d97ad6..3b1e9cd02 100644 --- a/turn-orchestrator/src/subscriber.rs +++ b/turn-orchestrator/src/subscriber.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use iii_sdk::{IIIError, RegisterFunctionMessage, Value, III}; use serde_json::json; +use crate::awaiting::AwaitingApproval; use crate::config::TurnOrchestratorConfig; use crate::persistence; use crate::run_start::publish_step; @@ -17,6 +18,7 @@ pub const FUNCTION_ID: &str = "turn::step"; pub async fn execute( iii: III, cfg: Arc, + awaiting: AwaitingApproval, payload: Value, ) -> Result { let session_id = extract_session_id(&payload).ok_or_else(|| { @@ -46,7 +48,11 @@ pub async fn execute( })?; persistence::save_record(&iii, &record).await; - if !record.is_terminal() { + if record.is_terminal() { + // Wake any `run::resume` racing the executor — it can now + // observe the terminal record and rebuild a resume plan. + awaiting.signal(&session_id); + } else { publish_step(&iii, &session_id).await; } Ok(json!({ @@ -56,9 +62,10 @@ pub async fn execute( })) } -pub fn register(iii: &III, cfg: &Arc) { +pub fn register(iii: &III, cfg: &Arc, awaiting: AwaitingApproval) { let iii_for_handler = iii.clone(); let cfg_for_handler = Arc::clone(cfg); + let awaiting_for_handler = awaiting; iii.register_function(( RegisterFunctionMessage::with_id(FUNCTION_ID.to_string()).with_description( "Run one durable state machine transition for a session.".to_string(), @@ -66,7 +73,8 @@ pub fn register(iii: &III, cfg: &Arc) { move |payload: Value| { let iii = iii_for_handler.clone(); let cfg = Arc::clone(&cfg_for_handler); - async move { execute(iii, cfg, payload).await } + let awaiting = awaiting_for_handler.clone(); + async move { execute(iii, cfg, awaiting, payload).await } }, )); } @@ -89,14 +97,16 @@ mod tests { use super::*; #[test] - fn subscriber_register_accepts_config_arc() { - // Compile-time pin: register() must take an Arc. - // This guards against silently dropping config plumbing. + fn subscriber_register_accepts_config_arc_and_awaiting() { + // Compile-time pin: register() must take an Arc + // and an AwaitingApproval so the signal-on-terminal-save plumbing + // can't get silently dropped. fn _assert_signature( iii: &iii_sdk::III, cfg: &std::sync::Arc, + awaiting: crate::awaiting::AwaitingApproval, ) { - super::register(iii, cfg); + super::register(iii, cfg, awaiting); } } From a5f8ef635728f2bef42b9bab7f15b0f8891e4adc Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Mon, 18 May 2026 08:07:54 -0300 Subject: [PATCH 4/4] test(harness): update trace_correlation for harness::call rename MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `bridge::trigger` was renamed to `harness::call` (and span name `harness.bridge.trigger` → `harness.call`) in a95f2ad, and `bridge::events` was removed entirely in 9b9b015. The trace_correlation tests still referenced the old names and failed with `function_not_found`. - Update the three remaining tests to call `harness::call` and assert the new span name `harness.call`. - Rename `bridge_trigger_*` test fn names to `harness_call_*` to match. - Delete `bridge_events_sse_emits_traceparent_and_message_id_headers` — the function it exercised no longer exists. --- harness/tests/trace_correlation.rs | 104 ++++++----------------------- 1 file changed, 21 insertions(+), 83 deletions(-) diff --git a/harness/tests/trace_correlation.rs b/harness/tests/trace_correlation.rs index e2087cee7..a399fdd7a 100644 --- a/harness/tests/trace_correlation.rs +++ b/harness/tests/trace_correlation.rs @@ -12,7 +12,7 @@ use serde_json::{json, Value}; const DEFAULT_ENGINE_URL: &str = "ws://127.0.0.1:49134"; /// Per-call timeout for engine triggers in this test suite. 2s is enough for -/// the engine to dispatch and return for `harness::status` / `bridge::trigger` +/// the engine to dispatch and return for `harness::status` / `harness::call` /// / `engine::traces::*` in a local engine; tests that need a longer budget /// should declare their own override. const TRIGGER_TIMEOUT_MS: u64 = 2_000; @@ -24,7 +24,7 @@ const EXPORTER_FLUSH_INTERVAL_MS: u64 = 100; /// Same retry pattern but expanded for the test that runs under parallel-test /// pressure on the shared exporter. 30×200ms = 6s tolerance for eviction and -/// dispatch jitter; matched to the `bridge_trigger_missing_function_id` +/// dispatch jitter; matched to the `harness_call_missing_function_id` /// error path which competes with other tests for exporter capacity. const EXPORTER_FLUSH_RETRIES_UNDER_LOAD: u32 = 30; const EXPORTER_FLUSH_INTERVAL_UNDER_LOAD_MS: u64 = 200; @@ -118,8 +118,8 @@ async fn call_status(iii: &iii_sdk::III, session: &str, message: Option<&str>) - .expect("call harness::status") } -/// Helper: invoke `harness::status` via `bridge::trigger` (the HTTP envelope -/// path that browsers actually use). The outer wrapper around bridge::trigger +/// Helper: invoke `harness::status` via `harness::call` (the HTTP envelope +/// path that browsers actually use). The outer wrapper around `harness::call` /// echoes `traceparent` / `x-iii-message-id` into response headers; the inner /// `harness::status` returns its raw shape directly. async fn call_status_via_bridge(iii: &iii_sdk::III, session: &str, message: Option<&str>) -> Value { @@ -129,18 +129,18 @@ async fn call_status_via_bridge(iii: &iii_sdk::III, session: &str, message: Opti payload["message_id"] = Value::String(m.to_string()); } iii.trigger(TriggerRequest { - function_id: "bridge::trigger".into(), + function_id: "harness::call".into(), payload, action: None, timeout_ms: Some(TRIGGER_TIMEOUT_MS), }) .await - .expect("call bridge::trigger -> harness::status") + .expect("call harness::call -> harness::status") } #[tokio::test] #[serial_test::serial(harness_trace)] -async fn bridge_trigger_echoes_message_id_when_provided() { +async fn harness_call_echoes_message_id_when_provided() { let Some((iii, _url)) = boot_or_skip().await else { return; }; @@ -151,7 +151,7 @@ async fn bridge_trigger_echoes_message_id_when_provided() { #[tokio::test] #[serial_test::serial(harness_trace)] -async fn bridge_trigger_omits_message_id_header_when_absent() { +async fn harness_call_omits_message_id_header_when_absent() { // Option A: when the caller doesn't supply a message_id and baggage // is empty, no x-iii-message-id header is emitted. Plumbing calls // stay out of `Group by message`. @@ -335,20 +335,20 @@ async fn engine_traces_tree_returns_root_and_children() { #[tokio::test] #[serial_test::serial(harness_trace)] -async fn bridge_trigger_missing_function_id_still_emits_traced_error() { +async fn harness_call_missing_function_id_still_emits_traced_error() { let Some((iii, _url)) = boot_or_skip().await else { return; }; let session = format!("err-session-{}", uuid::Uuid::new_v4()); - // Probe: send a well-formed bridge::trigger first and confirm we get a + // Probe: send a well-formed harness::call first and confirm we get a // traceparent back. If not, observability is down and the rest of this // test cannot record/query spans — skip. Use a different session so the // probe span doesn't collide with the error span during disambiguation. let probe_session = format!("probe-{}", uuid::Uuid::new_v4()); let probe = iii .trigger(TriggerRequest { - function_id: "bridge::trigger".into(), + function_id: "harness::call".into(), payload: json!({ "function_id": "harness::status", "payload": {}, @@ -359,18 +359,18 @@ async fn bridge_trigger_missing_function_id_still_emits_traced_error() { timeout_ms: Some(TRIGGER_TIMEOUT_MS), }) .await - .expect("probe bridge::trigger"); + .expect("probe harness::call"); if probe["headers"].get("traceparent").is_none() { skip_or_panic_otel("observability worker not active"); return; } - // bridge::trigger expects { function_id, payload }. We deliberately omit + // harness::call expects { function_id, payload }. We deliberately omit // function_id. The handler returns IIIError::Handler("missing function_id"), // but the span we opened around the handler should still be recorded. let result = iii .trigger(TriggerRequest { - function_id: "bridge::trigger".into(), + function_id: "harness::call".into(), payload: json!({ "session_id": session, "message_id": "M-err" }), action: None, timeout_ms: Some(TRIGGER_TIMEOUT_MS), @@ -387,7 +387,7 @@ async fn bridge_trigger_missing_function_id_still_emits_traced_error() { .trigger(TriggerRequest { function_id: "engine::traces::list".into(), payload: json!({ - "name": "harness.bridge.trigger", + "name": "harness.call", "search_all_spans": true, "limit": 500, }), @@ -409,14 +409,14 @@ async fn bridge_trigger_missing_function_id_still_emits_traced_error() { let Some(spans) = found else { skip_or_panic_otel( - "no harness.bridge.trigger spans landed in memory exporter (harness OTel exporter \ + "no harness.call spans landed in memory exporter (harness OTel exporter \ may not be wired to the engine in this run)", ); return; }; // `name + search_all_spans` may match many traces; disambiguate by walking - // each candidate's tree and finding the one whose harness.bridge.trigger - // child carries our test session_id. + // each candidate's tree and finding the one whose harness.call child + // carries our test session_id. #[allow(clippy::items_after_statements)] fn find_named<'a>(node: &'a Value, name: &str) -> Option<&'a Value> { @@ -454,7 +454,7 @@ async fn bridge_trigger_missing_function_id_still_emits_traced_error() { continue; }; for root in roots { - if let Some(child) = find_named(root, "harness.bridge.trigger") { + if let Some(child) = find_named(root, "harness.call") { if span_has_attr(child, "iii.session.id", &session) { matched_child = Some(child.clone()); break 'outer; @@ -465,7 +465,7 @@ async fn bridge_trigger_missing_function_id_still_emits_traced_error() { let Some(child) = matched_child else { skip_or_panic_otel(&format!( - "matched no harness.bridge.trigger span with session_id={session} (memory exporter \ + "matched no harness.call span with session_id={session} (memory exporter \ likely evicted it under parallel test traffic; the span DID record per the \ retry-loop hit on name)" )); @@ -475,72 +475,10 @@ async fn bridge_trigger_missing_function_id_still_emits_traced_error() { let status = child["status"].as_str().unwrap_or(""); assert!( status.eq_ignore_ascii_case("error"), - "expected Status::error on the failed harness.bridge.trigger span; got {status:?}" + "expected Status::error on the failed harness.call span; got {status:?}" ); } -/// `bridge::events` is the SSE path. The handler returns an envelope -/// (`status_code` + `headers` + body string with seed events), and -/// `with_envelope_span` merges `traceparent` + `x-iii-message-id` into the -/// envelope's `headers`. This test asserts both headers arrive verbatim for -/// the seed-pump phase. When T12 wires live-tail, the wrapper-shrinking -/// follow-up must keep this assertion passing — if it fails after live-tail, -/// the wrapper is no longer recording the right scope. -#[tokio::test] -#[serial_test::serial(harness_trace)] -async fn bridge_events_sse_emits_traceparent_and_message_id_headers() { - let Some((iii, _url)) = boot_or_skip().await else { - return; - }; - let session = format!("sse-session-{}", uuid::Uuid::new_v4()); - - let env = iii - .trigger(TriggerRequest { - function_id: "bridge::events".into(), - payload: json!({ - "query_params": { "session_id": &session, "message_id": "M-sse" }, - }), - action: None, - timeout_ms: Some(TRIGGER_TIMEOUT_MS), - }) - .await - .expect("call bridge::events"); - - // bridge::events returns the HTTP envelope. - assert_eq!( - env["status_code"].as_u64(), - Some(200), - "bridge::events envelope shape; got {env:?}" - ); - let headers = &env["headers"]; - assert_eq!( - headers["content-type"], "text/event-stream", - "SSE content-type preserved alongside the trace headers" - ); - - // x-iii-message-id is echoed verbatim whenever the envelope path is taken. - assert_eq!( - headers["x-iii-message-id"], "M-sse", - "x-iii-message-id should echo the caller-supplied id" - ); - - // traceparent is only present when the iii-observability worker is active. - // The require-OTel gate matches the other SSE-adjacent tests. - let Some(tp) = headers.get("traceparent").and_then(Value::as_str) else { - skip_or_panic_otel("bridge::events traceparent: observability worker not active"); - return; - }; - // Format: 00-<32 hex>-<16 hex>-01 - assert!( - tp.starts_with("00-"), - "traceparent should start with version 00-; got {tp}" - ); - let parts: Vec<&str> = tp.split('-').collect(); - assert_eq!(parts.len(), 4, "traceparent has 4 hyphen-separated parts"); - assert_eq!(parts[1].len(), 32, "trace_id is 32 hex chars; got {tp}"); - assert_eq!(parts[2].len(), 16, "span_id is 16 hex chars; got {tp}"); -} - #[cfg(test)] mod unit_tests { use super::{skip_or_panic_inner, REQUIRE_ENGINE_ENV, REQUIRE_OTEL_ENV};