From bb01dcc6fde42b2a09f55736a3bebf88d8a883d6 Mon Sep 17 00:00:00 2001 From: Shadaj Laddad Date: Wed, 8 Jul 2026 20:40:59 +0000 Subject: [PATCH] fix(rap): treat servers that don't support `/tool_call_status` as having given up on the call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, a server responding 404 (or any non-2xx) to a `/tool_call_status` query was treated as *unknown* liveness and its pending tool calls / subscriptions were kept pending forever. Since such a server cannot be tracking the call in a way that survives restarts, the default is now to treat the call/subscription as failed and prune it. ## New liveness semantics (`rap-client`) * `check_tool_call_status` now maps: * 2xx + valid body → `Alive` / `Gone` as reported (unchanged) * **4xx (endpoint unsupported, e.g. 404 from an older server) → `Gone`** (was `Unknown`) * **2xx with unparseable body → `Gone`** (was `Unknown`) * 5xx (transient server error) → `Unknown` (kept pending) * transport error (unreachable) → `Unknown` (kept pending) * `ToolCallLiveness` docs updated to reflect the new contract. The daemon's reconcile logic is unchanged (prunes on `Gone`, keeps on `Alive`/`Unknown`) — the new mapping means unsupported-endpoint servers now get their orphaned work pruned on boot, while genuinely unreachable servers still get the benefit of the doubt. ## Tests * Flipped `keeps_pending_tool_call_when_endpoint_unsupported` → `prunes_pending_tool_call_when_endpoint_unsupported` (404 now injects the failure result). * New `prunes_dead_subscription_when_endpoint_unsupported` (404 → final subscription-failure event). * New `keeps_pending_tool_call_on_transient_server_error` (500 → kept pending). ## Docs * `tool-call-status.md`: rewrote **Runtime Error Semantics** into the two categories (endpoint unsupported → prune; server unavailable → unknown, never prune), with the rationale that a hung conversation is worse than a spurious failure; updated the recovery and dispatch sections (multi-server broadcast querying must not treat a non-owning server's `alive: false`/404 as authoritative when the owning server is unreachable). * Updated `lifecycle.md`, `building-a-rap-tool.md` (endpoint is now effectively required for async/subscription tools), `building-a-runtime.md`, `subscription-events.md`, and `infinity-code/background-agents.md` to match. Co-authored-by: Infinity 🤖 PR: #62 --- .../src/batch_processor.rs | 3 + .../src/tools/rap_http.rs | 43 + crates/infinity-daemon/src/mcp_proxy.rs | 61 +- crates/infinity-daemon/src/rap_tools.rs | 21 + crates/infinity-daemon/src/session/mod.rs | 38 +- .../infinity-daemon/src/session/reconcile.rs | 1015 +++++++++++++++++ crates/rap-client/src/http.rs | 22 + crates/rap-client/src/notifier.rs | 76 ++ crates/rap-github-event-poller/src/lib.rs | 55 + crates/rap-github-event-poller/src/main.rs | 12 +- crates/rap-protocol/src/lib.rs | 28 + crates/sandbox-core/src/server.rs | 109 +- docs/docs/infinity-code/background-agents.md | 11 + docs/docs/rap/about/agent-runtime.md | 2 + docs/docs/rap/about/subscription-events.md | 2 + docs/docs/rap/spec/basic/lifecycle.md | 10 +- docs/docs/rap/spec/basic/migration.md | 2 +- docs/docs/rap/spec/basic/tool-call-status.md | 127 +++ docs/docs/rap/spec/basic/tool-cancellation.md | 2 + docs/docs/rap/spec/overview.md | 6 +- .../rap/spec/server/subscription-events.md | 6 + .../docs/rap/using-rap/building-a-rap-tool.md | 18 + docs/docs/rap/using-rap/building-a-runtime.md | 1 + 23 files changed, 1657 insertions(+), 13 deletions(-) create mode 100644 crates/infinity-daemon/src/session/reconcile.rs create mode 100644 docs/docs/rap/spec/basic/tool-call-status.md diff --git a/crates/infinity-agent-core/src/batch_processor.rs b/crates/infinity-agent-core/src/batch_processor.rs index a971a00b..d3e4cc70 100644 --- a/crates/infinity-agent-core/src/batch_processor.rs +++ b/crates/infinity-agent-core/src/batch_processor.rs @@ -537,6 +537,9 @@ mod tests { async fn post(&self, _: &str, _: &str) -> Result { Ok(200) } + async fn post_read(&self, _: &str, _: &str) -> Result<(u16, Vec), E> { + Ok((200, vec![])) + } async fn get(&self, _: &str) -> Result<(u16, Vec), E> { Ok((200, vec![])) } diff --git a/crates/infinity-agent-lambda/src/tools/rap_http.rs b/crates/infinity-agent-lambda/src/tools/rap_http.rs index a96c9ca7..091185db 100644 --- a/crates/infinity-agent-lambda/src/tools/rap_http.rs +++ b/crates/infinity-agent-lambda/src/tools/rap_http.rs @@ -126,6 +126,49 @@ impl HttpClient for RapHttpClient { Ok(response.status().as_u16()) } + async fn post_read(&self, url: &str, body: &str) -> Result<(u16, Vec), HttpError> { + let parsed = url::Url::parse(url).map_err(|e| HttpError(e.to_string()))?; + let host = parsed + .host_str() + .ok_or(HttpError("missing host".into()))? + .to_owned(); + + let signed_headers = self + .sign_request( + "POST", + url, + std::iter::once(("host", host.as_str())) + .chain(std::iter::once(("content-type", "application/json"))), + SignableBody::Bytes(body.as_bytes()), + ) + .await?; + + let mut request = self + .http_client + .post(url) + .header("host", &host) + .header("content-type", "application/json"); + + for (name, value) in &signed_headers { + request = request.header(name.as_str(), value.as_str()); + } + + let response = request + .body(body.to_owned()) + .send() + .await + .map_err(|e| HttpError(e.to_string()))?; + + let status = response.status().as_u16(); + let body_bytes = response + .bytes() + .await + .map_err(|e| HttpError(e.to_string()))? + .to_vec(); + + Ok((status, body_bytes)) + } + async fn get(&self, url: &str) -> Result<(u16, Vec), HttpError> { let parsed = url::Url::parse(url).map_err(|e| HttpError(e.to_string()))?; let host = parsed diff --git a/crates/infinity-daemon/src/mcp_proxy.rs b/crates/infinity-daemon/src/mcp_proxy.rs index ffa8d06b..cfd0e73f 100644 --- a/crates/infinity-daemon/src/mcp_proxy.rs +++ b/crates/infinity-daemon/src/mcp_proxy.rs @@ -9,7 +9,10 @@ use hyper::body::{Bytes, Incoming}; use hyper::server::conn::http1; use hyper::{Request, Response, StatusCode}; use hyper_util::rt::TokioIo; -use rap_protocol::{DisplaySegment, RapCallback, RapInvocation, RapToolResult}; +use rap_protocol::{ + DisplaySegment, RapCallback, RapInvocation, RapToolCallStatusRequest, + RapToolCallStatusResponse, RapToolResult, +}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::convert::Infallible; @@ -294,6 +297,26 @@ struct ProxyState { client_factory: McpClientFactory, client: Mutex>>, port: u16, + /// Tool call IDs currently being processed, for `/tool_call_status` + /// queries. Uses a std mutex so it can be updated from a drop guard. + in_flight: std::sync::Mutex>, +} + +/// RAII guard that removes a tool call ID from the proxy's in-flight set when +/// the invocation task finishes (on any exit path). +struct InFlightGuard { + id: String, + state: Arc, +} + +impl Drop for InFlightGuard { + fn drop(&mut self) { + self.state + .in_flight + .lock() + .expect("bug: in_flight mutex poisoned") + .remove(&self.id); + } } impl ProxyState { @@ -457,19 +480,52 @@ async fn handle(req: Request, state: Arc) -> Response b.to_bytes(), Err(_) => return text_response(StatusCode::BAD_REQUEST, "bad body"), }; + + // Tool call status query: report whether an invocation is still in flight. + // A freshly restarted proxy has an empty in-flight set, so calls from + // before the restart correctly report `alive: false`. + if path.ends_with("/tool_call_status") { + let request: RapToolCallStatusRequest = match serde_json::from_slice(&body) { + Ok(r) => r, + Err(e) => return text_response(StatusCode::BAD_REQUEST, &format!("bad json: {e}")), + }; + let alive = state + .in_flight + .lock() + .expect("bug: in_flight mutex poisoned") + .contains(&request.tool_call_id); + let response = serde_json::to_string(&RapToolCallStatusResponse { alive }) + .expect("bug: serialize RapToolCallStatusResponse"); + return json_response(StatusCode::OK, &response); + } + + // Parse invocation let inv: RapInvocation = match serde_json::from_slice(&body) { Ok(i) => i, Err(e) => return text_response(StatusCode::BAD_REQUEST, &format!("bad json: {e}")), }; + // Track the invocation as in-flight for the duration of the async task, + // so `/tool_call_status` reports it alive until the callback is sent. + state + .in_flight + .lock() + .expect("bug: in_flight mutex poisoned") + .insert(inv.id.clone()); + let in_flight_guard = InFlightGuard { + id: inv.id.clone(), + state: state.clone(), + }; + // Return immediately, process async let state = state.clone(); tokio::spawn(rap_protocol::log_panic("mcp_proxy_invoke", async move { + let _in_flight_guard = in_flight_guard; let res = if inv.operation.ends_with("_list_tools") { state.list_tools().await } else if inv.operation.ends_with("_invoke_tool") { @@ -580,6 +636,7 @@ pub async fn start_proxy_server(name: String, factory: McpClientFactory) -> Resu client_factory: factory, client: Mutex::new(None), port, + in_flight: std::sync::Mutex::new(std::collections::HashSet::new()), }); tokio::spawn(rap_protocol::log_panic( diff --git a/crates/infinity-daemon/src/rap_tools.rs b/crates/infinity-daemon/src/rap_tools.rs index c460aeba..94d25dbe 100644 --- a/crates/infinity-daemon/src/rap_tools.rs +++ b/crates/infinity-daemon/src/rap_tools.rs @@ -1,5 +1,7 @@ //! RAP tool support for the CLI: loads tools from RAP servers using rap-client. +use std::collections::HashMap; + use infinity_agent_core::tools::Tool; use infinity_agent_core::tools::rap_tool::RapTool; use infinity_agent_core::traits::InputSender; @@ -15,6 +17,10 @@ pub struct LoadedRapTools { pub tools: Vec>>, /// Servers that declared needsMigration: true, as (config_id, url) pairs. pub migration_servers: Vec<(String, String)>, + /// Maps each loaded RAP tool name → the base URL of the server that + /// provides it. Used to route protocol messages (e.g. `/tool_call_status` + /// queries) to the server that originally received an invocation. + pub tool_servers: HashMap, } pub async fn load_rap_tools( @@ -29,8 +35,21 @@ pub async fn load_rap_tools( let mut tools: Vec>> = Vec::new(); let mut migration_servers = Vec::new(); + let mut tool_servers = HashMap::new(); for ts in loaded { let endpoint = ts.manifest.endpoint.clone(); + // Resolve the configured base URL for this toolset's server. Falls + // back to the endpoint with any trailing `/invoke` path stripped. + let base_url = servers + .iter() + .find(|(u, _)| endpoint.starts_with(u.as_str())) + .map(|(u, _)| u.clone()) + .unwrap_or_else(|| { + endpoint + .trim_end_matches('/') + .trim_end_matches("/invoke") + .to_owned() + }); if ts.manifest.needs_migration { // Find the (url, id) entry for this toolset if let Some((url, Some(id))) = servers @@ -42,6 +61,7 @@ pub async fn load_rap_tools( } for def in ts.manifest.tools { tracing::info!("Loaded RAP tool: {} from {}", def.name, endpoint); + tool_servers.insert(def.name.clone(), base_url.clone()); tools.push(Box::new(RapTool { name: def.name, description: def.description, @@ -55,5 +75,6 @@ pub async fn load_rap_tools( Ok(LoadedRapTools { tools, migration_servers, + tool_servers, }) } diff --git a/crates/infinity-daemon/src/session/mod.rs b/crates/infinity-daemon/src/session/mod.rs index 70fd217b..bb37a4d3 100644 --- a/crates/infinity-daemon/src/session/mod.rs +++ b/crates/infinity-daemon/src/session/mod.rs @@ -25,6 +25,7 @@ use crate::sleep_tools::{SleepTool, SleepUntilTool}; pub mod agent_loop; pub mod display; +pub mod reconcile; pub mod thread_worker; pub use agent_loop::agent_loop; @@ -324,7 +325,9 @@ impl SessionManager { let spawned_servers = booted.spawned_servers; let urls = booted.urls; - let rap_tools: Vec>> = if !urls.is_empty() { + let rap_tools: Vec>>; + let rap_tool_servers: HashMap; + if !urls.is_empty() { let servers_with_ids: Vec<(String, Option)> = urls .iter() .map(|u| { @@ -338,16 +341,19 @@ impl SessionManager { emit(info(String::new())).await; - loaded.tools + rap_tools = loaded.tools; + rap_tool_servers = loaded.tool_servers; } Err(e) => { emit(info(format!("Warning: failed to load RAP tools: {e}"))).await; - Vec::new() + rap_tools = Vec::new(); + rap_tool_servers = HashMap::new(); } } } else { - Vec::new() - }; + rap_tools = Vec::new(); + rap_tool_servers = HashMap::new(); + } let extra_system_prompt = Some(format!( "The user's current working directory is: {cwd:?}\n\n\ @@ -361,6 +367,28 @@ impl SessionManager { let (shutdown_tx, shutdown_rx) = oneshot::channel(); + // Reconcile pending RAP tool calls / active subscriptions against + // their servers in the background: if the server gave up on any of + // them while the agent was down, inject failure messages so the + // affected threads don't hang forever. See `session::reconcile`. + if !rap_tool_servers.is_empty() { + let conversation_store = self.conversation_store.clone(); + let reconcile_state_store = self.state_store.clone(); + let reconcile_sender = sender.clone(); + let reconcile_session_id = session_id.clone(); + tokio::task::spawn_local(rap_protocol::log_panic("rap_reconcile", async move { + reconcile::reconcile_rap_state( + &conversation_store, + &reconcile_state_store, + &reconcile_session_id, + &rap_tool_servers, + &rap_tools::SimpleHttpClient::new(), + &reconcile_sender, + ) + .await; + })); + } + let (idle_tx, agent_handle, subscriber_map) = self.start_agent_loop( session_id.clone(), agent_rx, diff --git a/crates/infinity-daemon/src/session/reconcile.rs b/crates/infinity-daemon/src/session/reconcile.rs new file mode 100644 index 00000000..a2eef66f --- /dev/null +++ b/crates/infinity-daemon/src/session/reconcile.rs @@ -0,0 +1,1015 @@ +//! Boot-time reconciliation of pending RAP tool calls and active subscriptions. +//! +//! If the daemon shuts down (or crashes) while a RAP tool call is in flight +//! or a subscription is active, the tool server may give up on it in the +//! meantime (e.g. because the server itself restarted and lost its state). +//! When the agent boots back up, the affected thread would otherwise hang +//! forever waiting for a callback that will never arrive. +//! +//! On session boot, [`reconcile_rap_state`] queries the RAP server that +//! originally received each pending tool call / active subscription via the +//! `/tool_call_status` protocol message and prunes the ones the server gave +//! up on: +//! +//! * A **pending tool call** is answered with a synthetic failed tool result +//! so the model can observe the failure and retry. +//! * An **active subscription** is terminated with a synthetic final +//! subscription-failure event, which also removes it from the thread's +//! active-subscription tracking when processed. +//! +//! Servers that respond but do not support `/tool_call_status` (4xx / +//! invalid body) cannot confirm the call is alive, so their pending work is +//! treated as failed and pruned. Servers that are *unreachable* (or return a +//! transient 5xx error) yield an *unknown* liveness, which is treated +//! conservatively: nothing is pruned, matching the previous behavior of +//! waiting indefinitely. + +use std::collections::HashMap; + +use infinity_agent_core::message::{ + InfinityMessage, InputMessage, InputMessageContent, SyntheticKind, TaggedSyntheticKind, +}; +use infinity_agent_core::traits::{ConversationStore, InputSender, StateStore}; +use rap_client::http::HttpClient; +use rap_client::notifier::{ToolCallLiveness, check_tool_call_status}; +use rig::message::{ToolResult, ToolResultContent, UserContent}; + +use crate::memory_store::{InMemoryConversationStore, InMemoryStateStore}; + +/// What kind of pending RAP work a tool call represents. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PendingKind { + /// A dispatched tool call still waiting for its result. + ToolCall, + /// An active subscription established by a completed tool call. + Subscription, +} + +/// A pending RAP tool call or active subscription found in a session's threads. +#[derive(Debug, Clone)] +pub struct PendingRapCall { + pub thread_id: String, + pub tool_call_id: String, + pub tool_name: String, + /// Base URL of the RAP server that originally received the invocation. + pub server_url: String, + pub kind: PendingKind, +} + +/// Scan all open threads of a session for pending RAP tool calls and active +/// subscriptions. Only tool calls whose tool name maps to a known RAP server +/// (via `tool_servers`) are returned — built-in tools (sleep, spawn_thread, +/// …) never involve a RAP server and must not be pruned. +pub async fn collect_pending_rap_calls( + conversation_store: &InMemoryConversationStore, + state_store: &InMemoryStateStore, + session_id: &str, + tool_servers: &HashMap, +) -> Vec { + let mut thread_ids = vec![session_id.to_owned()]; + thread_ids.extend( + conversation_store + .get_open_subthreads(session_id) + .into_iter() + .map(|t| t.thread_id), + ); + + let mut pending = Vec::new(); + for thread_id in thread_ids { + let history = match conversation_store + .load_history_up_to(&thread_id, None, None) + .await + { + Ok(h) => h, + Err(e) => { + tracing::warn!("reconcile: failed to load history for {thread_id}: {e}"); + continue; + } + }; + + // A thread is waiting on a tool call iff its history ends with an + // unanswered ToolCall (mirrors the thread worker's idle check). + if let Some(InfinityMessage::ToolCall { call, .. }) = history.last() { + let name = call.function.name.clone(); + if let Some(server_url) = tool_servers.get(&name) { + pending.push(PendingRapCall { + thread_id: thread_id.clone(), + tool_call_id: call.id.clone(), + tool_name: name, + server_url: server_url.clone(), + kind: PendingKind::ToolCall, + }); + } + } + + // Active subscriptions are tracked by the tool_call_id that + // established them; recover the tool name from history to find the + // owning server. + let subscriptions = state_store + .get_active_subscriptions(&thread_id) + .await + .unwrap_or_default(); + for tool_call_id in subscriptions { + let tool_name = history.iter().find_map(|m| { + if let InfinityMessage::ToolCall { call, .. } = m + && call.id == tool_call_id + { + Some(call.function.name.clone()) + } else { + None + } + }); + let Some(tool_name) = tool_name else { + tracing::warn!( + "reconcile: subscription {tool_call_id} in thread {thread_id} has no \ + originating tool call in history; skipping" + ); + continue; + }; + let Some(server_url) = tool_servers.get(&tool_name) else { + tracing::warn!( + "reconcile: subscription {tool_call_id} (tool `{tool_name}`) in thread \ + {thread_id} has no known RAP server; skipping" + ); + continue; + }; + pending.push(PendingRapCall { + thread_id: thread_id.clone(), + tool_call_id, + tool_name, + server_url: server_url.clone(), + kind: PendingKind::Subscription, + }); + } + } + pending +} + +/// Build the synthetic tool result injected when a pending tool call was +/// abandoned by its RAP server. +fn tool_call_failure_message(pending: &PendingRapCall) -> InputMessage { + InputMessage { + content: InputMessageContent::User(UserContent::ToolResult(ToolResult { + id: pending.tool_call_id.clone(), + call_id: None, + content: rig::OneOrMany::one(ToolResultContent::Text(rig::agent::Text { + text: format!( + "Error: the `{}` tool call failed — the tool server is no longer \ + processing it (its state was likely lost in a restart) and no result \ + will be delivered. Retry the call if the operation is still needed.", + pending.tool_name + ), + })), + })), + group_id: pending.thread_id.clone(), + metadata: None, + synthetic: None, + display_as: None, + subscription: false, + } +} + +/// Build the synthetic final subscription event injected when an active +/// subscription was abandoned by its RAP server. Marked `final` so that +/// processing it also removes the subscription from active tracking. +fn subscription_failure_message(pending: &PendingRapCall) -> InputMessage { + InputMessage { + content: InputMessageContent::User(UserContent::ToolResult(ToolResult { + id: pending.tool_call_id.clone(), + call_id: None, + content: rig::OneOrMany::one(ToolResultContent::Text(rig::agent::Text { + text: format!( + "This subscription (created by `{}`) failed: the tool server is no \ + longer tracking it (its state was likely lost in a restart) and no \ + further events will be delivered. Re-subscribe if you still need \ + these events.", + pending.tool_name + ), + })), + })), + group_id: pending.thread_id.clone(), + metadata: None, + synthetic: Some(SyntheticKind::Tagged( + TaggedSyntheticKind::SubscriptionEvent { + tool_call_id: pending.tool_call_id.clone(), + associative: false, + r#final: true, + }, + )), + display_as: None, + subscription: false, + } +} + +/// Query the owning RAP server for each pending tool call / active +/// subscription in the session and inject failure messages for the ones the +/// server has given up on (including servers that respond without supporting +/// the status endpoint). Servers reporting *alive* — or servers that cannot +/// be reached at all (*unknown*) — leave the pending work untouched. +pub async fn reconcile_rap_state( + conversation_store: &InMemoryConversationStore, + state_store: &InMemoryStateStore, + session_id: &str, + tool_servers: &HashMap, + client: &H, + sender: &S, +) { + let pending = + collect_pending_rap_calls(conversation_store, state_store, session_id, tool_servers).await; + + for p in pending { + let liveness = + check_tool_call_status(client, &p.server_url, &p.thread_id, &p.tool_call_id).await; + match liveness { + ToolCallLiveness::Alive => { + tracing::debug!( + "reconcile: {:?} {} in thread {} still alive on {}", + p.kind, + p.tool_call_id, + p.thread_id, + p.server_url + ); + } + ToolCallLiveness::Unknown => { + tracing::debug!( + "reconcile: {:?} {} in thread {} has unknown status on {}; keeping", + p.kind, + p.tool_call_id, + p.thread_id, + p.server_url + ); + } + ToolCallLiveness::Gone => { + tracing::info!( + "reconcile: pruning {:?} {} in thread {} — server {} gave up on it", + p.kind, + p.tool_call_id, + p.thread_id, + p.server_url + ); + let msg = match p.kind { + PendingKind::ToolCall => tool_call_failure_message(&p), + PendingKind::Subscription => subscription_failure_message(&p), + }; + let dedup_id = uuid::Uuid::new_v4().to_string(); + if let Err(e) = sender + .send_to_input_queue(msg, &p.thread_id, &dedup_id) + .await + { + tracing::error!( + "reconcile: failed to inject failure message for {} in thread {}: {e}", + p.tool_call_id, + p.thread_id + ); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use super::*; + use crate::memory_store::InMemoryMessageSender; + use async_trait::async_trait; + use infinity_agent_core::traits::ConversationStore; + use tokio::sync::mpsc; + + #[derive(Debug)] + struct MockHttpError(String); + impl std::fmt::Display for MockHttpError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } + } + impl std::error::Error for MockHttpError {} + + /// Mock HTTP client that records `post_read` requests and returns canned + /// responses keyed by URL. URLs without a canned response yield an error + /// (simulating an unreachable server). + #[derive(Clone, Default)] + struct MockHttp { + responses: Arc>>, + requests: Arc>>, + } + + impl MockHttp { + fn respond(&self, url: &str, status: u16, body: &str) { + self.responses + .lock() + .expect("bug: mutex poisoned") + .insert(url.to_owned(), (status, body.to_owned())); + } + + fn requests(&self) -> Vec<(String, String)> { + self.requests.lock().expect("bug: mutex poisoned").clone() + } + } + + #[async_trait] + impl HttpClient for MockHttp { + type Error = MockHttpError; + + async fn post(&self, _url: &str, _body: &str) -> Result { + Ok(200) + } + + async fn post_read(&self, url: &str, body: &str) -> Result<(u16, Vec), MockHttpError> { + self.requests + .lock() + .expect("bug: mutex poisoned") + .push((url.to_owned(), body.to_owned())); + match self.responses.lock().expect("bug: mutex poisoned").get(url) { + Some((status, body)) => Ok((*status, body.clone().into_bytes())), + None => Err(MockHttpError("connection refused".to_owned())), + } + } + + async fn get(&self, _url: &str) -> Result<(u16, Vec), MockHttpError> { + Ok((404, vec![])) + } + } + + fn test_model_ref() -> infinity_protocol::ModelRef { + infinity_protocol::ModelRef { + provider_id: "mock".to_owned(), + model_id: "mock".to_owned(), + } + } + + fn tmp_stores() -> ( + InMemoryConversationStore, + InMemoryStateStore, + tempfile::TempDir, + ) { + let dir = tempfile::tempdir().expect("create temp dir"); + let conv = + InMemoryConversationStore::new_with_dir(dir.path().join("threads"), test_model_ref()); + let state = InMemoryStateStore::new(dir.path().join("state")); + (conv, state, dir) + } + + fn tool_call_msg(id: &str, name: &str) -> InfinityMessage { + InfinityMessage::ToolCall { + call: rig::message::ToolCall { + id: id.to_owned(), + call_id: None, + signature: None, + additional_params: None, + function: rig::message::ToolFunction { + name: name.to_owned(), + arguments: serde_json::json!({}), + }, + }, + display_as: None, + } + } + + fn tool_result_msg(id: &str, text: &str) -> InfinityMessage { + InfinityMessage::ToolResult { + result: ToolResult { + id: id.to_owned(), + call_id: None, + content: rig::OneOrMany::one(ToolResultContent::Text(rig::agent::Text { + text: text.to_owned(), + })), + }, + display_segments: None, + } + } + + async fn append(conv: &InMemoryConversationStore, thread_id: &str, msg: InfinityMessage) { + conv.append_messages(thread_id, vec![(msg, uuid::Uuid::new_v4().to_string())]) + .await + .expect("append message"); + } + + fn capture_sender() -> ( + InMemoryMessageSender, + mpsc::UnboundedReceiver<(InputMessage, String)>, + ) { + let (tx, rx) = mpsc::unbounded_channel(); + (InMemoryMessageSender::new(tx), rx) + } + + fn servers(entries: &[(&str, &str)]) -> HashMap { + entries + .iter() + .map(|(tool, url)| ((*tool).to_owned(), (*url).to_owned())) + .collect() + } + + const STATUS_URL: &str = "http://server-a/tool_call_status"; + + #[tokio::test] + async fn prunes_dead_pending_tool_call() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + append(&conv, "t1", tool_call_msg("tc-1", "my_rap_tool")).await; + + let http = MockHttp::default(); + http.respond(STATUS_URL, 200, r#"{"alive": false}"#); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("my_rap_tool", "http://server-a")]), + &http, + &sender, + ) + .await; + + // The server was asked about the right tool call. + let requests = http.requests(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].0, STATUS_URL); + let req: rap_protocol::RapToolCallStatusRequest = + serde_json::from_str(&requests[0].1).expect("valid status request body"); + assert_eq!(req.thread_id, "t1"); + assert_eq!(req.tool_call_id, "tc-1"); + + // A synthetic failed tool result was injected. + let (msg, _) = rx.try_recv().expect("failure message injected"); + assert_eq!(msg.group_id, "t1"); + assert!(msg.synthetic.is_none()); + let InputMessageContent::User(UserContent::ToolResult(result)) = &msg.content else { + panic!("expected tool result content, got {:?}", msg.content); + }; + assert_eq!(result.id, "tc-1"); + let ToolResultContent::Text(text) = result.content.first() else { + panic!("expected text content"); + }; + assert!(text.text.contains("my_rap_tool"), "text: {}", text.text); + assert!(text.text.contains("no longer"), "text: {}", text.text); + assert!(rx.try_recv().is_err(), "only one message expected"); + } + + #[tokio::test] + async fn keeps_alive_pending_tool_call() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + append(&conv, "t1", tool_call_msg("tc-1", "my_rap_tool")).await; + + let http = MockHttp::default(); + http.respond(STATUS_URL, 200, r#"{"alive": true}"#); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("my_rap_tool", "http://server-a")]), + &http, + &sender, + ) + .await; + + assert_eq!(http.requests().len(), 1); + assert!(rx.try_recv().is_err(), "alive tool call must not be pruned"); + } + + #[tokio::test] + async fn keeps_pending_tool_call_when_server_unreachable() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + append(&conv, "t1", tool_call_msg("tc-1", "my_rap_tool")).await; + + // No canned response → post_read errors (unreachable server). + let http = MockHttp::default(); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("my_rap_tool", "http://server-a")]), + &http, + &sender, + ) + .await; + + assert!( + rx.try_recv().is_err(), + "unknown liveness must not prune the tool call" + ); + } + + #[tokio::test] + async fn prunes_pending_tool_call_when_endpoint_unsupported() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + append(&conv, "t1", tool_call_msg("tc-1", "my_rap_tool")).await; + + // Old server without the endpoint → 404. The server is reachable but + // cannot be tracking the call, so the call is treated as failed. + let http = MockHttp::default(); + http.respond(STATUS_URL, 404, "not found"); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("my_rap_tool", "http://server-a")]), + &http, + &sender, + ) + .await; + + let (msg, _) = rx.try_recv().expect( + "a 404 from the status endpoint must prune the tool call (server can't be tracking it)", + ); + assert_eq!(msg.group_id, "t1"); + let InputMessageContent::User(UserContent::ToolResult(result)) = &msg.content else { + panic!("expected tool result content"); + }; + assert_eq!(result.id, "tc-1"); + } + + #[tokio::test] + async fn prunes_dead_subscription_when_endpoint_unsupported() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + append(&conv, "t1", tool_call_msg("tc-sub", "subscribe_events")).await; + append(&conv, "t1", tool_result_msg("tc-sub", "subscribed")).await; + state + .add_active_subscription("t1", "tc-sub") + .await + .expect("track subscription"); + + // The server doesn't implement the endpoint → the subscription is + // treated as lost. + let http = MockHttp::default(); + http.respond(STATUS_URL, 404, "not found"); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("subscribe_events", "http://server-a")]), + &http, + &sender, + ) + .await; + + let (msg, _) = rx + .try_recv() + .expect("unsupported endpoint must prune the subscription"); + assert!( + matches!( + &msg.synthetic, + Some(SyntheticKind::Tagged( + TaggedSyntheticKind::SubscriptionEvent { r#final: true, .. } + )) + ), + "expected a final synthetic subscription event, got {:?}", + msg.synthetic + ); + } + + #[tokio::test] + async fn keeps_pending_tool_call_on_transient_server_error() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + append(&conv, "t1", tool_call_msg("tc-1", "my_rap_tool")).await; + + // A 5xx is a transient server error, not proof the call was lost. + let http = MockHttp::default(); + http.respond(STATUS_URL, 500, "internal error"); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("my_rap_tool", "http://server-a")]), + &http, + &sender, + ) + .await; + + assert!( + rx.try_recv().is_err(), + "a 5xx from the status endpoint must not prune the tool call" + ); + } + + #[tokio::test] + async fn skips_pending_builtin_tool_call() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + // Pending tool call for a tool that is not provided by any RAP server + // (e.g. a built-in like sleep or spawn_thread). + append(&conv, "t1", tool_call_msg("tc-1", "sleep")).await; + + let http = MockHttp::default(); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("my_rap_tool", "http://server-a")]), + &http, + &sender, + ) + .await; + + assert!(http.requests().is_empty(), "no server should be queried"); + assert!(rx.try_recv().is_err(), "built-in calls must not be pruned"); + } + + #[tokio::test] + async fn prunes_dead_subscription_with_final_event() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + // The subscription's originating tool call completed (has a result), + // so it is not a pending tool call — only an active subscription. + append(&conv, "t1", tool_call_msg("tc-sub", "subscribe_events")).await; + append(&conv, "t1", tool_result_msg("tc-sub", "subscribed")).await; + state + .add_active_subscription("t1", "tc-sub") + .await + .expect("track subscription"); + + let http = MockHttp::default(); + http.respond(STATUS_URL, 200, r#"{"alive": false}"#); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("subscribe_events", "http://server-a")]), + &http, + &sender, + ) + .await; + + let (msg, _) = rx.try_recv().expect("subscription failure event injected"); + assert_eq!(msg.group_id, "t1"); + let Some(SyntheticKind::Tagged(TaggedSyntheticKind::SubscriptionEvent { + tool_call_id, + associative, + r#final, + })) = &msg.synthetic + else { + panic!( + "expected synthetic subscription event, got {:?}", + msg.synthetic + ); + }; + assert_eq!(tool_call_id, "tc-sub"); + assert!(!associative); + assert!( + r#final, + "event must be final so the subscription is removed from tracking" + ); + let InputMessageContent::User(UserContent::ToolResult(result)) = &msg.content else { + panic!("expected tool result content"); + }; + assert_eq!(result.id, "tc-sub"); + let ToolResultContent::Text(text) = result.content.first() else { + panic!("expected text content"); + }; + assert!( + text.text.contains("subscription") && text.text.contains("failed"), + "text: {}", + text.text + ); + assert!(rx.try_recv().is_err(), "only one message expected"); + } + + #[tokio::test] + async fn keeps_alive_subscription() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + append(&conv, "t1", tool_call_msg("tc-sub", "subscribe_events")).await; + append(&conv, "t1", tool_result_msg("tc-sub", "subscribed")).await; + state + .add_active_subscription("t1", "tc-sub") + .await + .expect("track subscription"); + + let http = MockHttp::default(); + http.respond(STATUS_URL, 200, r#"{"alive": true}"#); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("subscribe_events", "http://server-a")]), + &http, + &sender, + ) + .await; + + assert_eq!(http.requests().len(), 1); + assert!(rx.try_recv().is_err(), "alive subscription must be kept"); + } + + #[tokio::test] + async fn collects_from_open_child_threads() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("root").await.expect("ensure root"); + let child = conv + .spawn_thread("root", "tc-spawn", false, None) + .await + .expect("spawn child"); + append(&conv, &child, tool_call_msg("tc-child", "my_rap_tool")).await; + + let http = MockHttp::default(); + http.respond(STATUS_URL, 200, r#"{"alive": false}"#); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "root", + &servers(&[("my_rap_tool", "http://server-a")]), + &http, + &sender, + ) + .await; + + let requests = http.requests(); + assert_eq!(requests.len(), 1); + let req: rap_protocol::RapToolCallStatusRequest = + serde_json::from_str(&requests[0].1).expect("valid status request body"); + assert_eq!(req.thread_id, child, "child thread id must be used"); + assert_eq!(req.tool_call_id, "tc-child"); + + let (msg, _) = rx.try_recv().expect("failure message injected"); + assert_eq!( + msg.group_id, child, + "failure must be routed to the child thread" + ); + } + + /// Multiple pieces of pending work across servers are each routed to the + /// server that owns the tool, and only the dead ones are pruned. + #[tokio::test] + async fn routes_to_owning_server_and_prunes_selectively() { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("root").await.expect("ensure root"); + let child = conv + .spawn_thread("root", "tc-spawn", false, None) + .await + .expect("spawn child"); + + // Root: active subscription owned by server A (alive). + append(&conv, "root", tool_call_msg("tc-sub", "subscribe_events")).await; + append(&conv, "root", tool_result_msg("tc-sub", "subscribed")).await; + state + .add_active_subscription("root", "tc-sub") + .await + .expect("track subscription"); + + // Child: pending tool call owned by server B (dead). + append(&conv, &child, tool_call_msg("tc-b", "tool_b")).await; + + let http = MockHttp::default(); + http.respond(STATUS_URL, 200, r#"{"alive": true}"#); + http.respond( + "http://server-b/tool_call_status", + 200, + r#"{"alive": false}"#, + ); + let (sender, mut rx) = capture_sender(); + + reconcile_rap_state( + &conv, + &state, + "root", + &servers(&[ + ("subscribe_events", "http://server-a"), + ("tool_b", "http://server-b"), + ]), + &http, + &sender, + ) + .await; + + let requests = http.requests(); + assert_eq!(requests.len(), 2); + assert!( + requests + .iter() + .any(|(url, body)| { url == STATUS_URL && body.contains("tc-sub") }) + ); + assert!(requests.iter().any(|(url, body)| { + url == "http://server-b/tool_call_status" && body.contains("tc-b") + })); + + // Only the dead tool call on server B is pruned. + let (msg, _) = rx.try_recv().expect("one failure message"); + assert_eq!(msg.group_id, child); + let InputMessageContent::User(UserContent::ToolResult(result)) = &msg.content else { + panic!("expected tool result content"); + }; + assert_eq!(result.id, "tc-b"); + assert!(rx.try_recv().is_err(), "alive subscription must be kept"); + } + + // ── End-to-end: injected failures flow through the agent loop ── + + async fn test_catalog( + model: rig_mock::MockCompletionModel, + ) -> Arc { + use infinity_agent_core::model_provider::{ModelEntry, SingleModelProvider}; + let entry = ModelEntry { + model_id: "mock".to_owned(), + display_name: "mock".to_owned(), + context_window: 0, + max_output_tokens: None, + }; + Arc::new( + crate::models::ModelCatalog::new(vec![( + "mock".to_owned(), + Arc::new(SingleModelProvider::new(entry, model)) as _, + )]) + .await + .expect("build test catalog"), + ) + } + + /// Spawn a real agent loop wired to an `InMemoryMessageSender`, returning + /// the sender used for reconciliation injection. + async fn spawn_agent_loop_for_reconcile( + session_id: &str, + conv: InMemoryConversationStore, + state: InMemoryStateStore, + model: rig_mock::MockCompletionModel, + ) -> InMemoryMessageSender { + use crate::session::{AgentMessage, SubscriberMap}; + + let (agent_tx, agent_rx) = mpsc::unbounded_channel(); + let (idle_tx, _idle_rx) = mpsc::unbounded_channel(); + let (input_tx, mut input_adapter_rx) = mpsc::unbounded_channel::<(InputMessage, String)>(); + let agent_tx_clone = agent_tx.clone(); + tokio::task::spawn_local(async move { + while let Some((msg, id)) = input_adapter_rx.recv().await { + if agent_tx_clone + .send(AgentMessage::Input(Box::new(msg), id)) + .is_err() + { + break; + } + } + }); + let sender = InMemoryMessageSender::new(input_tx); + let subscriber_map: SubscriberMap = Arc::new(Mutex::new(HashMap::new())); + let active_threads = Arc::new(Mutex::new(std::collections::HashSet::new())); + + tokio::task::spawn_local(crate::session::agent_loop( + session_id.to_owned(), + agent_rx, + test_catalog(model).await, + conv, + state, + sender.clone(), + String::new(), + Arc::new(vec![]), + Arc::new(None), + None, + subscriber_map, + active_threads, + idle_tx, + tokio_util::sync::CancellationToken::new(), + )); + sender + } + + /// A pruned pending tool call wakes the thread worker: the model receives + /// a completion containing the injected failure result and can continue. + #[tokio::test(flavor = "current_thread")] + async fn pruned_tool_call_resumes_thread_worker() { + use rig_mock::mock_model; + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + append(&conv, "t1", tool_call_msg("tc-1", "my_rap_tool")).await; + + let (model, mut ctrl) = mock_model(); + let sender = + spawn_agent_loop_for_reconcile("t1", conv.clone(), state.clone(), model).await; + + let http = MockHttp::default(); + http.respond(STATUS_URL, 200, r#"{"alive": false}"#); + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("my_rap_tool", "http://server-a")]), + &http, + &sender, + ) + .await; + + // The injected failure result triggers a completion whose + // history contains it. + let req = + tokio::time::timeout(std::time::Duration::from_secs(5), ctrl.next_request()) + .await + .expect("model should be woken by the injected failure result"); + let has_failure = req.chat_history.iter().any(|m| { + if let rig::message::Message::User { content } = m + && let UserContent::ToolResult(r) = content.first() + && let ToolResultContent::Text(t) = r.content.first() + { + r.id == "tc-1" && t.text.contains("no longer") + } else { + false + } + }); + assert!(has_failure, "failure result should be in the completion"); + ctrl.send_text("recovered"); + ctrl.finish(); + }) + .await; + } + + /// A pruned subscription's final failure event removes the subscription + /// from active tracking once processed by the thread worker. + #[tokio::test(flavor = "current_thread")] + async fn pruned_subscription_is_removed_from_tracking() { + use rig_mock::mock_model; + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (conv, state, _dir) = tmp_stores(); + conv.ensure_root_thread("t1").await.expect("ensure root"); + append(&conv, "t1", tool_call_msg("tc-sub", "subscribe_events")).await; + append(&conv, "t1", tool_result_msg("tc-sub", "subscribed")).await; + state + .add_active_subscription("t1", "tc-sub") + .await + .expect("track subscription"); + + let (model, mut ctrl) = mock_model(); + let sender = + spawn_agent_loop_for_reconcile("t1", conv.clone(), state.clone(), model).await; + + let http = MockHttp::default(); + http.respond(STATUS_URL, 200, r#"{"alive": false}"#); + reconcile_rap_state( + &conv, + &state, + "t1", + &servers(&[("subscribe_events", "http://server-a")]), + &http, + &sender, + ) + .await; + + // The injected final subscription event triggers a completion + // containing the failure text. + let req = + tokio::time::timeout(std::time::Duration::from_secs(5), ctrl.next_request()) + .await + .expect("model should be woken by the injected subscription failure"); + let has_failure = req.chat_history.iter().any(|m| { + if let rig::message::Message::User { content } = m + && let UserContent::ToolResult(r) = content.first() + && let ToolResultContent::Text(t) = r.content.first() + { + t.text.contains("subscription") && t.text.contains("failed") + } else { + false + } + }); + assert!( + has_failure, + "subscription failure should be in the completion" + ); + ctrl.send_text("acknowledged"); + ctrl.finish(); + + // The final event removes the subscription from tracking. + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + let subs = state + .get_active_subscriptions("t1") + .await + .expect("get subscriptions"); + if subs.is_empty() { + break; + } + assert!( + std::time::Instant::now() < deadline, + "subscription should be removed from tracking, still have: {subs:?}" + ); + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await; + } +} diff --git a/crates/rap-client/src/http.rs b/crates/rap-client/src/http.rs index 90cd0ff2..849108e1 100644 --- a/crates/rap-client/src/http.rs +++ b/crates/rap-client/src/http.rs @@ -11,6 +11,10 @@ pub trait HttpClient: Send + Sync + Clone { /// POST a JSON body to the given URL. Returns the HTTP status code. async fn post(&self, url: &str, body: &str) -> Result; + /// POST a JSON body to the given URL. Returns the HTTP status code and + /// response body. Used for request/response style protocol messages + /// (e.g. `/tool_call_status`) where the runtime needs to read the reply. + async fn post_read(&self, url: &str, body: &str) -> Result<(u16, Vec), Self::Error>; /// GET the given URL. Returns the HTTP status code and response body. async fn get(&self, url: &str) -> Result<(u16, Vec), Self::Error>; } @@ -76,6 +80,24 @@ impl HttpClient for SimpleHttpClient { Ok(resp.status().as_u16()) } + async fn post_read(&self, url: &str, body: &str) -> Result<(u16, Vec), SimpleHttpError> { + let resp = self + .client + .post(url) + .header("content-type", "application/json") + .header("accept", "application/json") + .body(body.to_owned()) + .send() + .await + .map_err(|e| SimpleHttpError(e.to_string()))?; + let status = resp.status().as_u16(); + let bytes = resp + .bytes() + .await + .map_err(|e| SimpleHttpError(e.to_string()))?; + Ok((status, bytes.to_vec())) + } + async fn get(&self, url: &str) -> Result<(u16, Vec), SimpleHttpError> { let resp = self .client diff --git a/crates/rap-client/src/notifier.rs b/crates/rap-client/src/notifier.rs index 0a6aa8e8..f59a37cd 100644 --- a/crates/rap-client/src/notifier.rs +++ b/crates/rap-client/src/notifier.rs @@ -1,8 +1,84 @@ /// Best-effort lifecycle notifications to RAP tool servers. use std::collections::HashMap; +use rap_protocol::{RapToolCallStatusRequest, RapToolCallStatusResponse}; + use crate::http::HttpClient; +/// Liveness of a tool call (or its subscription) as reported by a tool server +/// via the `/tool_call_status` endpoint. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ToolCallLiveness { + /// The server is still processing the tool call, or still maintains an + /// active subscription established by it. + Alive, + /// The server no longer tracks the tool call — either it answered + /// `alive: false`, or it responded without supporting the endpoint + /// (4xx / invalid body), in which case it cannot be tracking the call. + /// No result or further events should be expected. + Gone, + /// Liveness could not be determined: the server was unreachable or + /// returned a transient server error (5xx). Callers SHOULD treat this + /// conservatively (i.e. as if the call may still be alive). + Unknown, +} + +/// Query a tool server's `/tool_call_status` endpoint to check whether a tool +/// call (or the subscription it established) is still alive. +/// +/// `server_url` is the tool server's base URL — the same base used to derive +/// the `/.well-known/rap-toolset` discovery endpoint. +/// +/// A server that responds but does not support the endpoint — a 4xx status +/// (e.g. a 404 from a server that predates it) or an unparseable body — +/// yields [`ToolCallLiveness::Gone`]: such a server cannot confirm the call +/// is alive, so the call is treated as failed. Transport errors and 5xx +/// responses yield [`ToolCallLiveness::Unknown`], since the server may just +/// be temporarily unavailable. +pub async fn check_tool_call_status( + client: &H, + server_url: &str, + thread_id: &str, + tool_call_id: &str, +) -> ToolCallLiveness { + let endpoint = format!("{}/tool_call_status", server_url.trim_end_matches('/')); + let payload = serde_json::to_string(&RapToolCallStatusRequest { + thread_id: thread_id.to_owned(), + tool_call_id: tool_call_id.to_owned(), + }) + .expect("bug: failed to serialize tool_call_status request"); + match client.post_read(&endpoint, &payload).await { + Ok((status, body)) if (200..300).contains(&status) => { + match serde_json::from_slice::(&body) { + Ok(resp) if resp.alive => ToolCallLiveness::Alive, + Ok(_) => ToolCallLiveness::Gone, + Err(e) => { + tracing::warn!( + "invalid tool_call_status response from {endpoint}: {e}; \ + treating the tool call as gone" + ); + ToolCallLiveness::Gone + } + } + } + Ok((status, _)) if (400..500).contains(&status) => { + tracing::info!( + "tool_call_status at {endpoint} returned status {status} \ + (endpoint unsupported); treating the tool call as gone" + ); + ToolCallLiveness::Gone + } + Ok((status, _)) => { + tracing::warn!("tool_call_status at {endpoint} returned status {status}"); + ToolCallLiveness::Unknown + } + Err(e) => { + tracing::warn!("failed to query tool_call_status at {endpoint}: {e}"); + ToolCallLiveness::Unknown + } + } +} + /// Sends best-effort notifications to all configured RAP tool servers. #[derive(Clone)] pub struct RapNotifier { diff --git a/crates/rap-github-event-poller/src/lib.rs b/crates/rap-github-event-poller/src/lib.rs index 7706d56c..1ad6ac4f 100644 --- a/crates/rap-github-event-poller/src/lib.rs +++ b/crates/rap-github-event-poller/src/lib.rs @@ -176,6 +176,16 @@ impl Poller { repos.retain(|_, state| !state.subscriptions.is_empty()); } + /// Whether a subscription with the given tool_call_id is still active. + /// Used to answer `/tool_call_status` queries so a restarted runtime can + /// detect subscriptions this server no longer tracks. + pub async fn is_active(&self, tool_call_id: &str) -> bool { + let repos = self.repos.read().await; + repos + .values() + .any(|state| state.subscriptions.contains_key(tool_call_id)) + } + /// Run the polling loop. This never returns. pub async fn run(&self) -> ! { loop { @@ -522,3 +532,48 @@ fn describe_filters(f: &Filters) -> String { format!("Filters: {}", parts.join(", ")) } } + +#[cfg(test)] +mod tests { + use super::*; + use rap_protocol::PlainCallbackClient; + + fn test_invocation(id: &str) -> RapInvocation { + RapInvocation { + operation: "subscribe_github_events".to_owned(), + arguments: serde_json::json!({"owner": "octocat", "repo": "hello-world"}), + id: id.to_owned(), + call_id: None, + callback_url: "http://127.0.0.1:1/callback".to_owned(), + group_id: "thread-1".to_owned(), + user_id: None, + thread_ancestors: None, + } + } + + #[tokio::test] + async fn is_active_reflects_subscription_lifecycle() { + let poller = Poller::new(PlainCallbackClient::new(), None); + + assert!( + !poller.is_active("tc-1").await, + "unknown tool_call_id should not be active" + ); + + poller.subscribe(&test_invocation("tc-1")).await; + assert!( + poller.is_active("tc-1").await, + "subscription should be active after subscribe" + ); + assert!( + !poller.is_active("tc-other").await, + "other tool_call_ids should stay inactive" + ); + + poller.cancel("tc-1").await; + assert!( + !poller.is_active("tc-1").await, + "subscription should be inactive after cancel" + ); + } +} diff --git a/crates/rap-github-event-poller/src/main.rs b/crates/rap-github-event-poller/src/main.rs index 037d1c9a..ccfd1afc 100644 --- a/crates/rap-github-event-poller/src/main.rs +++ b/crates/rap-github-event-poller/src/main.rs @@ -9,7 +9,8 @@ use tracing_subscriber::EnvFilter; use rap_github_event_poller::Poller; use rap_protocol::{ - PlainCallbackClient, RapInvocation, ToolDef, ToolsetManifest, send_tool_result, + PlainCallbackClient, RapInvocation, RapToolCallStatusRequest, RapToolCallStatusResponse, + ToolDef, ToolsetManifest, send_tool_result, }; #[derive(Parser)] @@ -126,6 +127,14 @@ async fn cancel_handler( StatusCode::OK } +async fn tool_call_status_handler( + State(state): State>, + Json(req): Json, +) -> Json { + let alive = state.is_active(&req.tool_call_id).await; + Json(RapToolCallStatusResponse { alive }) +} + fn main() -> Result<(), Box> { let log_file = std::fs::File::create("./rap-github-event-poller.log").expect("failed to create log file"); @@ -162,6 +171,7 @@ fn main() -> Result<(), Box> { .route("/.well-known/rap-toolset", get(toolset_handler)) .route("/invoke", post(invoke_handler)) .route("/cancel_tool_call", post(cancel_handler)) + .route("/tool_call_status", post(tool_call_status_handler)) .with_state(poller); let embedded = std::env::var("RAP_EMBEDDED").is_ok(); diff --git a/crates/rap-protocol/src/lib.rs b/crates/rap-protocol/src/lib.rs index 72d35ccb..22dc5f6b 100644 --- a/crates/rap-protocol/src/lib.rs +++ b/crates/rap-protocol/src/lib.rs @@ -126,6 +126,34 @@ pub enum RapCallback { ViewUpdate(RapViewUpdate), } +/// Request body for the `/tool_call_status` endpoint (runtime → tool server). +/// +/// Asks the tool server whether a previously dispatched tool call — or the +/// subscription it established — is still alive (i.e. the server still +/// intends to deliver a tool result or further subscription events for it). +/// Runtimes use this after a restart to detect tool calls and subscriptions +/// that the tool server has given up on. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RapToolCallStatusRequest { + /// The conversation thread identifier (`group_id`) of the thread + /// containing the tool call. Matches the `group_id` sent in the original + /// tool invocation. + pub thread_id: String, + /// The unique identifier of the tool call to query. Matches the `id` + /// sent in the original tool invocation. + pub tool_call_id: String, +} + +/// Response body for the `/tool_call_status` endpoint (tool server → runtime). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RapToolCallStatusResponse { + /// `true` if the server is still processing the tool call or still + /// maintains an active subscription for it; `false` if the server has no + /// record of the tool call (e.g. it was lost in a restart, completed + /// long ago, or was cancelled). + pub alive: bool, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolsetManifest { pub name: String, diff --git a/crates/sandbox-core/src/server.rs b/crates/sandbox-core/src/server.rs index 9ada78a0..d8d81ba3 100644 --- a/crates/sandbox-core/src/server.rs +++ b/crates/sandbox-core/src/server.rs @@ -15,8 +15,9 @@ use tokio::task::JoinHandle; use tracing; use rap_protocol::{ - CallbackClient, DiffContent, DisplaySegment, RapCallback, RapInvocation, RapToolResult, - ToolDef, ToolsetManifest, send_subscription_event, send_tool_result, send_user_choice, + CallbackClient, DiffContent, DisplaySegment, RapCallback, RapInvocation, + RapToolCallStatusRequest, RapToolCallStatusResponse, RapToolResult, ToolDef, ToolsetManifest, + send_subscription_event, send_tool_result, send_user_choice, }; type DisplayResult = Result<(String, Option>), SandboxError>; @@ -48,6 +49,43 @@ type PendingTasks = Arc>>>; /// cancellation; the handler receives it and sends SIGTERM to the process. type InFlightMap = Arc>>>; +/// Tool call IDs currently being processed by this server, for answering +/// `/tool_call_status` queries. Unlike [`InFlightMap`] (which only tracks +/// cancellable `execute_command` invocations), this covers every accepted +/// invocation for the full duration of its background task. +/// +/// Uses a `std::sync::Mutex` so entries can be removed from a `Drop` impl. +type ActiveInvocations = Arc>>; + +/// RAII guard that registers a tool call ID as active and removes it on drop, +/// so every exit path of an invocation task (success, error, cancellation, +/// panic) clears the entry. +struct ActiveInvocationGuard { + id: String, + set: ActiveInvocations, +} + +impl ActiveInvocationGuard { + fn register(id: String, set: &ActiveInvocations) -> Self { + set.lock() + .expect("bug: active invocations mutex poisoned") + .insert(id.clone()); + Self { + id, + set: set.clone(), + } + } +} + +impl Drop for ActiveInvocationGuard { + fn drop(&mut self) { + self.set + .lock() + .expect("bug: active invocations mutex poisoned") + .remove(&self.id); + } +} + /// Send SIGTERM to a process group by PID. /// /// The spawned command is expected to have set its PGID to its own PID @@ -69,6 +107,8 @@ struct AppState { callback_client: C, pending_tasks: PendingTasks, in_flight: InFlightMap, + /// All invocation IDs currently being processed (see [`ActiveInvocations`]). + active_invocations: ActiveInvocations, /// Pending user choice responses, keyed by tool call ID. /// The sender delivers the user's selected index. pending_choices: Arc>>>, @@ -143,6 +183,7 @@ where callback_client, pending_tasks: tracker.pending_tasks.clone(), in_flight, + active_invocations: Arc::new(std::sync::Mutex::new(HashSet::new())), pending_choices: Arc::new(Mutex::new(HashMap::new())), server_base_url: std::sync::OnceLock::new(), needs_migration, @@ -158,6 +199,10 @@ where "/cancel_tool_call", post(cancel_tool_call_handler::), ) + .route( + "/tool_call_status", + post(tool_call_status_handler::), + ) .route( "/user_choice_response", post(user_choice_response_handler::), @@ -203,6 +248,12 @@ async fn invoke_handler< State(state): State>>, Json(invocation): Json, ) -> StatusCode { + // Register the invocation as active synchronously (before spawning the + // task) so `/tool_call_status` reports it alive for the entire time the + // server is working on it. The guard removes it when the task finishes. + let active_guard = + ActiveInvocationGuard::register(invocation.id.clone(), &state.active_invocations); + // For execute_command, register the cancellation channel synchronously // (before spawning the task) so that cancel_tool_call always finds an // entry — even if the cancel arrives before the command starts. @@ -216,6 +267,7 @@ async fn invoke_handler< let state_clone = state.clone(); let handle = tokio::spawn(async move { + let _active_guard = active_guard; handle_execute_command_streaming(&state_clone, &invocation, cancel_rx).await; }); state.pending_tasks.lock().await.push(handle); @@ -224,6 +276,7 @@ async fn invoke_handler< let state_clone = state.clone(); let handle = tokio::spawn(async move { + let _active_guard = active_guard; let result_text = match invocation.operation.as_str() { "clone_repo" => handle_clone_repo(&state_clone, &invocation) .await @@ -365,6 +418,36 @@ async fn cancel_tool_call_handler< StatusCode::OK } +/// Query endpoint for the `/tool_call_status` RAP protocol message. +/// +/// The runtime POSTs `{"thread_id":"…","tool_call_id":"…"}` (typically after +/// a restart) to ask whether this server is still processing the tool call. +/// Responds `{"alive": true}` while the invocation task is running (including +/// while it is blocked on a user choice), and `{"alive": false}` for unknown +/// or completed tool calls — signalling the runtime that no result will ever +/// be delivered and the call can be pruned. +async fn tool_call_status_handler< + B: SandboxBackend + 'static, + M: MetadataStore + 'static, + C: CallbackClient + 'static, +>( + State(state): State>>, + Json(request): Json, +) -> Json { + let alive = state + .active_invocations + .lock() + .expect("bug: active invocations mutex poisoned") + .contains(&request.tool_call_id); + tracing::info!( + tool_call_id = %request.tool_call_id, + thread_id = %request.thread_id, + alive, + "answered tool_call_status query" + ); + Json(RapToolCallStatusResponse { alive }) +} + /// Request payload for the `/user_choice_response` endpoint. #[derive(Debug, Deserialize)] struct UserChoiceResponse { @@ -2372,6 +2455,28 @@ fn build_manifest(endpoint: &str, needs_migration: bool) -> ToolsetManifest { mod tests { use super::*; + #[tokio::test] + async fn active_invocation_guard_tracks_liveness() { + let set: ActiveInvocations = Arc::new(std::sync::Mutex::new(HashSet::new())); + + let guard = ActiveInvocationGuard::register("tc-1".to_owned(), &set); + assert!( + set.lock() + .expect("bug: active invocations mutex poisoned") + .contains("tc-1"), + "id should be active while the guard is alive" + ); + + // Dropping the guard (any exit path of the invocation task) removes it. + drop(guard); + assert!( + !set.lock() + .expect("bug: active invocations mutex poisoned") + .contains("tc-1"), + "id should be removed once the guard is dropped" + ); + } + #[test] fn detects_cd_to_exact_repo_path() { let uri = "/Users/foo/my-repo"; diff --git a/docs/docs/infinity-code/background-agents.md b/docs/docs/infinity-code/background-agents.md index 9831fe34..9b4cc3e2 100644 --- a/docs/docs/infinity-code/background-agents.md +++ b/docs/docs/infinity-code/background-agents.md @@ -45,3 +45,14 @@ It will auto-start again the next time you run `infinity`. ## Persistence Conversation history is persisted to disk under `~/.infinity/`. You can shut down the CLI entirely, boot it back up later, and continue right where you left off with all your existing context intact. + +### Recovering pending tool calls + +If the daemon shuts down (or crashes) while a tool call is in flight or a subscription is active, the RAP server handling it may give up in the meantime — for example, an embedded server is restarted along with the session and loses its in-memory state. Without intervention, the conversation would wait forever for a result that will never arrive. + +To prevent this, whenever the daemon boots an agent session it reconciles the session's threads against the RAP servers: for every pending tool call and every active subscription, it asks the originating server whether the call is still alive using the RAP [tool call status check](/docs/rap/spec/basic/tool-call-status). Calls the server has given up on are pruned: + +- A **pending tool call** gets a synthetic error result injected into the conversation, so the agent sees the failure and can retry the call. +- A **lost subscription** gets a synthetic final subscription event and is removed from the thread's active subscriptions, so the agent can re-subscribe. + +Servers that answer `alive: false` — or that respond without supporting the status endpoint at all — are treated as having given up, and their calls are pruned. Only servers that can't be reached (or return a transient server error) are left alone: those calls simply stay pending, and the results are delivered normally if they eventually arrive. diff --git a/docs/docs/rap/about/agent-runtime.md b/docs/docs/rap/about/agent-runtime.md index 596ef6cf..0f226bae 100644 --- a/docs/docs/rap/about/agent-runtime.md +++ b/docs/docs/rap/about/agent-runtime.md @@ -19,6 +19,8 @@ The protocol doesn't prescribe how the runtime is built. The reference implement **Result routing.** Tool results, subscription events, and user messages all arrive through the same input channel. The runtime doesn't distinguish between them at the transport level — it loads state, appends the new message, and runs the LLM again. The `group_id` field routes messages to the correct conversation thread. +**Recovery.** Because tool calls are asynchronous, a conversation can be waiting on a callback when the runtime — or the tool server — restarts. If the tool server lost the call in the meantime, the conversation would hang forever. On boot, the runtime can reconcile its pending tool calls and active subscriptions by querying each tool server's [`/tool_call_status`](/docs/rap/spec/basic/tool-call-status) endpoint, pruning calls the server has given up on by injecting synthetic failure results the LLM can react to. + **Tool definitions.** The runtime maintains the set of available tools and their JSON Schema definitions. These are passed to the LLM on each completion request so it knows what tools it can call. How definitions are stored and loaded is implementation-specific. ## Capabilities provided to tools diff --git a/docs/docs/rap/about/subscription-events.md b/docs/docs/rap/about/subscription-events.md index ccd40100..9c48d637 100644 --- a/docs/docs/rap/about/subscription-events.md +++ b/docs/docs/rap/about/subscription-events.md @@ -55,6 +55,8 @@ Cancellation is tool-specific — the tool exposes a separate operation (e.g. `c The runtime does not automatically cancel subscriptions when a thread closes. Agents should cancel subscriptions explicitly before shutting down. If a subscription isn't cancelled and the subscribing thread is closed, events will still arrive at the callback URL but the runtime may not have a valid thread to process them in. +The reverse failure mode also exists: the tool server can lose a subscription (say, it restarted without durable storage) while the runtime still tracks it — leaving the agent waiting for events that will never come. Runtimes can detect this after a restart by asking the tool server whether the subscription is still alive via the [tool call status check](/docs/rap/spec/basic/tool-call-status), and surface lost subscriptions to the agent as a final event so it can re-subscribe. + :::warning This is an active area of development and subject to change. Future versions of RAP will include a standard protocol for cancelling subscriptions to enable auto-cleanup. diff --git a/docs/docs/rap/spec/basic/lifecycle.md b/docs/docs/rap/spec/basic/lifecycle.md index b6c10526..fb8a50b7 100644 --- a/docs/docs/rap/spec/basic/lifecycle.md +++ b/docs/docs/rap/spec/basic/lifecycle.md @@ -21,6 +21,8 @@ When a tool provider starts, it MUST expose two HTTP endpoints: The tool provider SHOULD also expose a **thread closure endpoint** at `/close_thread` to receive best-effort cleanup notifications from the runtime. See [Thread Closure](/docs/rap/spec/basic/thread-closure) for details. +Tool providers that track in-flight invocations or maintain [subscriptions](/docs/rap/spec/server/subscription-events) SHOULD additionally expose a **status endpoint** at `/tool_call_status`, which lets runtimes ask whether a previously dispatched tool call is still alive. See [Tool Call Status Check](/docs/rap/spec/basic/tool-call-status) for details. + The tool provider MUST be ready to serve both required endpoints before accepting traffic. The discovery endpoint is how runtimes learn what operations the tool supports — if it is unavailable or returns an invalid toolset, no runtime will be able to invoke the tool. ```mermaid @@ -109,7 +111,7 @@ Tool providers MUST handle stale invocations gracefully — either by maintainin The protocol does not define a shutdown handshake for tool providers. When a tool provider shuts down: -- Active invocations that have been acknowledged but not yet completed MAY be lost. Tool providers SHOULD persist in-flight work to allow recovery on restart. +- Active invocations that have been acknowledged but not yet completed MAY be lost. Tool providers SHOULD persist in-flight work to allow recovery on restart. Tool providers that expose a [`/tool_call_status`](/docs/rap/spec/basic/tool-call-status) endpoint allow runtimes to detect invocations that were lost this way and prune them instead of waiting forever. - Active subscriptions SHOULD continue to function if the tool provider restarts. Tools that store subscription state durably can resume event delivery after restart. - The discovery endpoint becomes unavailable. Runtimes that have already cached the toolset definition will continue to send invocations to the invocation endpoint, which will fail with connection errors. The runtime SHOULD record these failures as tool results with error descriptions. @@ -119,6 +121,12 @@ When a runtime closes a conversation thread, it sends a best-effort notification This notification is strictly best-effort — the runtime MUST NOT retry on failure, and tool servers MAY ignore it entirely. Tool servers that do handle the notification MUST always respond with HTTP 200. See [Thread Closure](/docs/rap/spec/basic/thread-closure) for the full specification. +### Tool Call Status Checks + +Because both participants can restart independently, a runtime may find itself waiting on a tool call that the tool provider has given up on — for example, an invocation that was in flight when the tool provider restarted and lost its in-memory state. To recover from this, a runtime that boots with pending tool calls or active subscriptions in its persisted state SHOULD query each originating tool server's `/tool_call_status` endpoint. If the server answers that a call is no longer alive, the runtime SHOULD prune it — injecting a synthetic failed tool result (or a synthetic final subscription event) so the LLM can reason about the failure and retry. + +Unlike the best-effort notifications, the status check is a request/response query: the runtime interprets the response body. An `"alive": false` answer — or a response showing the server does not support the endpoint (4xx / invalid body) — means the call is lost and SHOULD be pruned; only transport errors and transient 5xx responses leave the call pending. Tool providers that perform asynchronous work or maintain subscriptions SHOULD therefore implement `/tool_call_status`, or their in-flight work will be treated as failed whenever a runtime restarts. See [Tool Call Status Check](/docs/rap/spec/basic/tool-call-status) for the full specification. + ## Concurrency ### Runtime Concurrency diff --git a/docs/docs/rap/spec/basic/migration.md b/docs/docs/rap/spec/basic/migration.md index 94e90b54..45cacbe3 100644 --- a/docs/docs/rap/spec/basic/migration.md +++ b/docs/docs/rap/spec/basic/migration.md @@ -1,5 +1,5 @@ --- -sidebar_position: 8 +sidebar_position: 9 title: Migration --- diff --git a/docs/docs/rap/spec/basic/tool-call-status.md b/docs/docs/rap/spec/basic/tool-call-status.md new file mode 100644 index 00000000..6da62746 --- /dev/null +++ b/docs/docs/rap/spec/basic/tool-call-status.md @@ -0,0 +1,127 @@ +--- +sidebar_position: 8 +title: Tool Call Status Check +--- + +# Tool Call Status Check + +A tool call status check is a query sent from the runtime to a tool server to ask whether a previously dispatched tool call is still alive — that is, whether the tool server still intends to deliver a [tool result](/docs/rap/spec/basic/tool-result) for it, or still maintains an active [subscription](/docs/rap/spec/server/subscription-events) established by it. The runtime POSTs to the tool server's `/tool_call_status` endpoint with the identifiers of the original invocation, and the tool server answers with a JSON body indicating whether it still has a record of the call. + +Unlike the [thread closure](/docs/rap/spec/basic/thread-closure) and [tool cancellation](/docs/rap/spec/basic/tool-cancellation) notifications — which are fire-and-forget — the status check is a **request/response** message: the runtime reads and interprets the response body. + +The primary use case is recovery after a runtime restart. RAP's fire-and-forget invocation model means a conversation can be waiting indefinitely for a callback — a pending tool result or the next subscription event. If the runtime crashes or shuts down while a call is in flight, and the tool server meanwhile gives up on the call (for example, because the tool server itself restarted and lost its in-memory state), the conversation would hang forever waiting for a callback that will never arrive. The status check lets a rebooting runtime detect these orphaned calls and prune them, so the agent can observe the failure and recover instead of hanging. + +A single message covers both pending tool calls and subscriptions: a subscription is identified by the `tool_call_id` of the tool call that established it. + +## Request + +The runtime MUST send the status check as an HTTP POST with `Content-Type: application/json` to the tool server's `/tool_call_status` endpoint. + +```http +POST https://tool.example.com/tool_call_status +Content-Type: application/json + +{ + "thread_id": "thread_xyz", + "tool_call_id": "call_abc123" +} +``` + +The `/tool_call_status` path is relative to the tool server's base URL — the same base URL used to derive the `/.well-known/rap-toolset` [discovery endpoint](/docs/rap/spec/basic/toolsets#discovery-endpoint), `/close_thread`, and `/cancel_tool_call`. For example, if the tool server's base URL is `https://tool.example.com`, the runtime POSTs to `https://tool.example.com/tool_call_status`. + +### Fields + +| Field | Type | Required | Description | +|---|---|---|---| +| `thread_id` | `string` | Yes | The conversation thread identifier (`group_id`) of the thread containing the tool call. This is the same value that was sent as `group_id` in the original [tool invocation](/docs/rap/spec/basic/tool-invocation). | +| `tool_call_id` | `string` | Yes | The unique identifier of the tool call to query. This is the same value that was sent as `id` in the original [tool invocation](/docs/rap/spec/basic/tool-invocation). For subscriptions, this is the `id` of the tool call that established the subscription. | + +## Response + +A tool server that implements the endpoint MUST respond with HTTP 200 and a JSON body: + +```http +HTTP/1.1 200 OK +Content-Type: application/json + +{ + "alive": true +} +``` + +### Fields + +| Field | Type | Required | Description | +|---|---|---|---| +| `alive` | `boolean` | Yes | Whether the tool server still tracks the tool call. See [Query Semantics](#query-semantics). | + +## Query Semantics + +`"alive": true` means the tool server still tracks the tool call, in either of two senses: + +- **Pending tool call** — the server is still processing the invocation and a [tool result](/docs/rap/spec/basic/tool-result) will eventually be delivered to the callback URL. +- **Active subscription** — the server still maintains a subscription established by that tool call, and further [subscription events](/docs/rap/spec/server/subscription-events) may be delivered. + +`"alive": false` means the tool server has no record of the tool call — for example, its state was lost in a restart, the call completed long ago, or it was cancelled — and the runtime SHOULD NOT expect any further callbacks for it. + +Tool servers MUST answer `"alive": false` for unknown `tool_call_id` values. Detecting calls the server has given up on (or never heard of) is the entire purpose of this message; answering anything else for an unknown identifier would defeat it. + +The status check is read-only: tool servers MUST NOT cancel, complete, or otherwise mutate the state of a tool call in response to a status query. + +:::note +There is an inherent race between a status check and callback delivery: a tool result or subscription event MAY be in flight while the server reports on the call's status. Runtimes MUST tolerate receiving a callback for a call they have already pruned (typically by ignoring it), and tool servers MUST NOT assume a status query implies the runtime has or has not received prior callbacks. +::: + +## Runtime Error Semantics + +Responses that do not carry a valid `"alive": true` answer fall into two categories with opposite handling. + +**Endpoint unsupported — treat the call as dead.** A tool server that responds but does not support the status check cannot be tracking the tool call in a way that survives restarts, so the runtime SHOULD treat the call as lost and prune it, exactly as if the server had answered `"alive": false`: + +- A 4xx response status — including a 404 from a tool server that predates this endpoint +- A 2xx response whose body cannot be parsed as the JSON object above + +This is the default because a hung conversation is worse than a spurious failure: a pruned call surfaces an explicit error the agent can retry, whereas an unpruned dead call hangs the thread forever. Tool servers that perform asynchronous work or maintain subscriptions SHOULD therefore implement `/tool_call_status` — servers that do not will have their pending calls and subscriptions treated as failed whenever the runtime restarts. + +**Server unavailable — treat the status as unknown.** The runtime MUST NOT prune based on a failure to get an answer at all, since the server may only be temporarily unavailable and may still deliver callbacks: + +- A transport error (connection failure, timeout) +- A 5xx response status (transient server error) + +In the unknown case the call simply stays pending, exactly as it would have before this message existed. + +## Recovery After Runtime Restart + +On boot, a runtime SHOULD reconcile its persisted conversation state against the tool servers: for every pending tool call and every active subscription it is still waiting on, it SHOULD query the tool server that originally received the invocation. When the server answers `"alive": false` — or responds without supporting the endpoint — the runtime SHOULD prune the orphaned call: + +- **Pending tool call** — inject a synthetic failed [tool result](/docs/rap/spec/basic/tool-result) into the conversation (e.g. `"Error: the tool server is no longer processing this call"`), so the LLM can reason about the failure and retry if appropriate. +- **Active subscription** — inject a synthetic final [subscription event](/docs/rap/spec/server/subscription-events) reporting the loss and remove the subscription from active tracking, so the agent can re-subscribe if appropriate. + +```mermaid +sequenceDiagram + participant R as Agent Runtime + participant T as Tool Server + + R->>T: POST /invoke {"id": "call_abc123", ...} + T-->>R: HTTP 200 OK + Note over R: Runtime crashes / shuts down + Note over T: Tool server restarts,
loses in-flight state + Note over R: Runtime boots, finds pending
tool call "call_abc123" + R->>T: POST /tool_call_status {"tool_call_id": "call_abc123", "thread_id": "thread_xyz"} + T-->>R: HTTP 200 {"alive": false} + Note over R: Injects synthetic failed tool result
so the agent can retry +``` + +Pruning is a SHOULD, not a MUST — runtimes MAY apply additional heuristics (e.g. grace periods, retry-before-prune) as long as they never prune on unknown status (server unreachable or 5xx). + +## Dispatch Behavior + +Unlike [tool cancellation](/docs/rap/spec/basic/tool-cancellation) — which is broadcast to every configured tool server — the status check SHOULD be sent only to the tool server that originally received the invocation, since only that server can answer authoritatively. Runtimes that cannot determine the originating server MAY query multiple servers instead; in that case they SHOULD treat the call as alive if any server reports `"alive": true`, and as unknown (rather than dead) if any queried server is unavailable — a `"alive": false` from a server that never received the invocation is not authoritative. + +Runtimes SHOULD issue status checks for independent tool calls concurrently and SHOULD NOT block conversation processing on their completion. + +## Security Considerations + +Tool servers MUST validate that `/tool_call_status` requests are authentic — for example, by requiring the same authentication mechanism used for [tool invocations](/docs/rap/spec/basic/tool-invocation) (AWS SigV4, bearer tokens, mutual TLS, etc.). An unauthenticated status endpoint would allow an attacker to probe for the existence of tool calls and enumerate active subscriptions. + +Tool servers MUST treat the `tool_call_id` and `thread_id` as untrusted input and MUST validate them before using them to look up state. The response MUST NOT expose any information about the tool call beyond the `alive` boolean — no arguments, partial results, or internal state. Tool servers SHOULD rate-limit the `/tool_call_status` endpoint to prevent abuse. diff --git a/docs/docs/rap/spec/basic/tool-cancellation.md b/docs/docs/rap/spec/basic/tool-cancellation.md index e8bfbc52..2c9d80ac 100644 --- a/docs/docs/rap/spec/basic/tool-cancellation.md +++ b/docs/docs/rap/spec/basic/tool-cancellation.md @@ -54,6 +54,8 @@ Tool cancellation notifications are **best-effort** by design. The protocol expl Because this notification is advisory, tool servers SHOULD NOT rely on it as the sole mechanism for managing in-flight operations. Tool servers SHOULD implement independent timeout and cleanup strategies to handle cases where the notification is never received. +Conversely, runtimes SHOULD NOT rely on cancellation notifications to keep their view of in-flight work consistent with the tool server's. A runtime can ask a tool server whether a dispatched call is still being processed at all — for example, after a restart — using the [tool call status check](/docs/rap/spec/basic/tool-call-status), which detects calls the tool server has lost or given up on entirely. + ## Cancellation Behavior When a tool server receives a `/cancel_tool_call` notification, it MAY attempt to abort the identified operation. The specific behavior depends on the tool: diff --git a/docs/docs/rap/spec/overview.md b/docs/docs/rap/spec/overview.md index fb792fc8..5aa9a184 100644 --- a/docs/docs/rap/spec/overview.md +++ b/docs/docs/rap/spec/overview.md @@ -56,6 +56,7 @@ The protocol defines two roles that communicate through the HTTP message contrac - **User ID**: An optional end-user identity that tools MAY use for authorization decisions, personalization, or audit logging. - **Thread closure notifications**: A best-effort signal sent to tool servers when a conversation thread is closed, allowing them to clean up thread-specific resources. See [Thread Closure](/docs/rap/spec/basic/thread-closure). - **Tool cancellation notifications**: A best-effort signal sent to tool servers when a tool call is interrupted, allowing them to abort in-flight operations. See [Tool Cancellation](/docs/rap/spec/basic/tool-cancellation). +- **Tool call status checks**: A query sent to tool servers to ask whether a pending tool call or active subscription is still alive, allowing runtimes to detect and prune calls that were lost (e.g. across restarts). See [Tool Call Status Check](/docs/rap/spec/basic/tool-call-status). **Tools** are independent HTTP services that receive invocations, process them on their own schedule, and return results through the callback mechanism. Tools provide the following capabilities to runtimes: @@ -67,11 +68,12 @@ The protocol defines two roles that communicate through the HTTP message contrac ### Message Types -The protocol defines four message types that cover the full range of communication between runtimes and tools. Tool invocations flow from runtime to tool, while the remaining three message types flow from tool to runtime through the callback URL. +The protocol defines seven message types that cover the full range of communication between runtimes and tools. Tool invocations and tool call status checks flow from runtime to tool, while the remaining message types flow from tool to runtime through the callback URL. | Message | Direction | Description | |---|---|---| | [Tool Invocation](/docs/rap/spec/basic/tool-invocation) | Runtime → Tool | Invoke a tool operation. Contains the operation name, arguments, callback URL, and routing identifiers. | +| [Tool Call Status Check](/docs/rap/spec/basic/tool-call-status) | Runtime → Tool | Ask whether a pending tool call or active subscription is still alive. The tool answers with an `alive` boolean, letting the runtime prune calls the tool server has given up on. | | [Tool Result](/docs/rap/spec/basic/tool-result) | Tool → Runtime | Return the result of a completed operation. Contains the result text and the identifiers needed to match it to the original invocation. | | [Subscription Event](/docs/rap/spec/server/subscription-events) | Tool → Runtime | Deliver an event from an active subscription. References the original subscription tool call so the runtime can associate the event with the correct context. | | [OAuth](/docs/rap/spec/server/oauth) | Tool → Runtime | Initiate a user authorization flow. Contains an authorization URL that the runtime surfaces to the user. The tool retries the original operation after authorization completes. | @@ -121,6 +123,8 @@ Explore the detailed specification for each protocol component: - [Toolsets](/docs/rap/spec/basic/toolsets) — Declaring and discovering tool definitions - [Thread Closure](/docs/rap/spec/basic/thread-closure) — Best-effort thread cleanup notifications - [Tool Cancellation](/docs/rap/spec/basic/tool-cancellation) — Best-effort tool call cancellation notifications + - [Tool Call Status Check](/docs/rap/spec/basic/tool-call-status) — Querying whether a tool call or subscription is still alive + - [Migration](/docs/rap/spec/basic/migration) — Migrating tool server state between servers - **Server Features** - [Subscription Events](/docs/rap/spec/server/subscription-events) — Event-driven subscriptions diff --git a/docs/docs/rap/spec/server/subscription-events.md b/docs/docs/rap/spec/server/subscription-events.md index d492e3c9..1e8684bb 100644 --- a/docs/docs/rap/spec/server/subscription-events.md +++ b/docs/docs/rap/spec/server/subscription-events.md @@ -128,6 +128,12 @@ Tools that support subscriptions MUST store the `callback_url`, `group_id`, and To enable cancellation, tools SHOULD include a subscription identifier in the initial `tool_result` (e.g., `"Subscribed to pull_request events. Subscription ID: sub_abc"`). The tool result MUST include `"subscription": true` so that the runtime can [track the subscription](/docs/rap/spec/basic/tool-result). +## Liveness + +Subscriptions are long-lived, so the runtime and the tool can lose sync — most commonly when the tool server restarts without durable storage and silently drops its subscriptions. The runtime would then track a subscription that will never produce another event. + +Runtimes MAY verify that a subscription is still alive — typically after a runtime restart — by sending a [tool call status check](/docs/rap/spec/basic/tool-call-status) to the tool server with the subscription's originating `tool_call_id`. Tools that maintain subscriptions SHOULD implement the `/tool_call_status` endpoint and report `"alive": true` for as long as the subscription is active — runtimes treat a server that responds without supporting the endpoint as unable to vouch for the subscription, and prune it just as if the server had answered `"alive": false`. When a subscription is reported (or presumed) lost, the runtime SHOULD remove it from its active tracking and surface the loss to the agent (e.g. as a synthetic final event) so it can re-subscribe if needed. + ## Cancellation ### Subscription tracking diff --git a/docs/docs/rap/using-rap/building-a-rap-tool.md b/docs/docs/rap/using-rap/building-a-rap-tool.md index d3b8f952..ebd9e205 100644 --- a/docs/docs/rap/using-rap/building-a-rap-tool.md +++ b/docs/docs/rap/using-rap/building-a-rap-tool.md @@ -204,6 +204,24 @@ Include `subscription: true` in the [tool result](/docs/rap/spec/basic/tool-resu The runtime spawns a child thread for each subscription event, giving each event a clean context window. The subscription remains active until explicitly cancelled. +## Answering status checks + +Runtimes that restart while one of your tool calls is pending — or while one of your subscriptions is active — may ask whether the call is still alive by POSTing to `/tool_call_status` (see the [Tool Call Status Check spec](/docs/rap/spec/basic/tool-call-status)). Answer `alive: true` if you're still working on the call or still hold the subscription, and `alive: false` if you have no record of it. That lets the runtime prune calls you've given up on (e.g. lost in a restart) instead of waiting forever. + +```javascript +// POST /tool_call_status +async function handleToolCallStatus(body, res) { + const { tool_call_id } = body; + const stillProcessing = inFlight.has(tool_call_id); + const stillSubscribed = Boolean(await db.get(tool_call_id)); + + res.writeHead(200, { 'Content-Type': 'application/json' }); + res.end(JSON.stringify({ alive: stillProcessing || stillSubscribed })); +} +``` + +Implement this endpoint if your tool does any asynchronous work or maintains subscriptions: runtimes treat a server that responds without supporting the status check (a 404, or any non-JSON answer) as unable to vouch for the call, and will prune your pending calls and subscriptions as failed whenever the runtime restarts. Only unreachable servers (connection errors, 5xx) are given the benefit of the doubt and left pending. + ## Schema evolution Runtimes cache your toolset definition for the duration of an agent session. If you deploy a breaking schema change while agents hold cached definitions, they'll send invocations with stale arguments. Your tool should handle this gracefully — either maintain backward compatibility or return a clear error via the normal tool result path. See [Loading Toolsets](/docs/rap/spec/basic/toolsets#loading-toolsets) for details on caching behavior. diff --git a/docs/docs/rap/using-rap/building-a-runtime.md b/docs/docs/rap/using-rap/building-a-runtime.md index ae911a6c..e8ad3327 100644 --- a/docs/docs/rap/using-rap/building-a-runtime.md +++ b/docs/docs/rap/using-rap/building-a-runtime.md @@ -319,6 +319,7 @@ This is a minimal runtime to demonstrate the protocol. A production runtime woul - **Subscription event handling** — requires generating [synthetic tool calls](/docs/rap/about/subscription-events#synthetic-tool-calls) to present events to the LLM in a way it can reason about. See [Subscription Events](/docs/rap/about/subscription-events) for the full design. - **Concurrency control** — serialize messages within a thread (e.g. with a queue or database lock). See [Agent Runtime](/docs/rap/about/agent-runtime#interruption-model). +- **Recovery after restarts** — durable state means a restarted runtime can find itself waiting on tool calls or subscriptions that the tool server has meanwhile given up on. On boot, query each originating server's `/tool_call_status` endpoint and prune dead calls by injecting synthetic failure results. Servers that respond without supporting the endpoint are treated as having given up too; only unreachable servers leave the call pending. See the [Tool Call Status Check spec](/docs/rap/spec/basic/tool-call-status). - **Hibernation** — for a serverless deployment, replace the Express server with a Lambda triggered by SQS, and use scheduled messages for sleep. See [Agent Hibernation](/docs/rap/about/architecture#hibernation). - **Authentication** — sign requests to tool servers with SigV4 or bearer tokens, and authenticate callback requests to prevent unauthorized message injection. - **Streaming** — stream LLM responses to the user instead of waiting for the full completion.