diff --git a/approval-gate/src/lib.rs b/approval-gate/src/lib.rs index f6225e662..865abe9d9 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) = trigger_resume(&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,44 @@ 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), + }) +} + +/// 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 response.get("resumed").and_then(Value::as_bool) == Some(true) { + return Ok(()); + } + Err("run::resume did not reopen approval turn".to_string()) +} + #[cfg(test)] mod tests { use super::*; @@ -547,7 +748,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 +846,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 +867,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 +986,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 +1008,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 +1030,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 +1042,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 +1053,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(); - }); + #[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" + } + } + ] + })); - let decision = await_decision(&*bus, STATE_SCOPE, "s1", "tc-1", now_ms() + 5_000).await; - writer.await.unwrap(); - assert_eq!(decision, Decision::Allow); + assert_eq!(out.len(), 1); + assert_eq!(out[0]["function_call_id"], "tc-1"); } - #[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:?}"), - } - } - - #[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 +1100,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 +1121,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/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}; 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/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); 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/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/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/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/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 ad3697fdf..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; @@ -13,6 +14,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 +134,84 @@ 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)) +} + +/// 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")?; + + 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)) +} + +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, + "resumed": true, + })) +} + pub async fn publish_step(iii: &III, session_id: &str) { if let Err(e) = iii .trigger(TriggerRequest { @@ -149,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()) @@ -172,6 +252,18 @@ pub fn register(iii: &III, cfg: &Arc) { async move { execute_sync(iii, cfg, payload).await } }, )); + 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(); + let awaiting = awaiting_resume.clone(); + async move { execute_resume(iii, awaiting, payload).await } + }, + )); } fn required_str(payload: &Value, field: &str) -> Result { @@ -237,6 +329,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/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); } } 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); }