diff --git a/crates/infinity-agent-cli/src/daemon_client.rs b/crates/infinity-agent-cli/src/daemon_client.rs index bd076230..a5bc797a 100644 --- a/crates/infinity-agent-cli/src/daemon_client.rs +++ b/crates/infinity-agent-cli/src/daemon_client.rs @@ -552,7 +552,8 @@ where mpsc::unbounded_channel::>(); let (soft_detach_tx, soft_detach_rx) = mpsc::unbounded_channel::<()>(); let (detach_result_tx, detach_result_rx) = mpsc::unbounded_channel::(); - let (choice_answered_tx, choice_answered_rx) = mpsc::unbounded_channel::<(String, usize)>(); + let (choice_answered_tx, choice_answered_rx) = + mpsc::unbounded_channel::<(infinity_protocol::ChoiceId, usize)>(); if let Some(info) = startup_info { let _ = display_tx.send((None, DisplayEvent::Info(info))); diff --git a/crates/infinity-agent-cli/src/display.rs b/crates/infinity-agent-cli/src/display.rs index c82cfe20..35230fae 100644 --- a/crates/infinity-agent-cli/src/display.rs +++ b/crates/infinity-agent-cli/src/display.rs @@ -36,14 +36,14 @@ pub enum DisplayEvent { auth_url: String, }, UserChoiceRequired { - id: String, + id: infinity_protocol::ChoiceId, prompt: String, choices: Vec, default: usize, response_url: String, }, UserChoiceComplete { - choice_id: String, + choice_id: infinity_protocol::ChoiceId, }, ThinkingStart, ThinkingEnd, diff --git a/crates/infinity-agent-cli/src/terminal.rs b/crates/infinity-agent-cli/src/terminal.rs index cb635adb..15585b4a 100644 --- a/crates/infinity-agent-cli/src/terminal.rs +++ b/crates/infinity-agent-cli/src/terminal.rs @@ -104,7 +104,7 @@ enum SoftDetachAction { /// A queued user choice request waiting to be shown in the TUI. struct PendingChoice { - id: String, + id: infinity_protocol::ChoiceId, prompt: String, choices: Vec, default: usize, @@ -134,7 +134,7 @@ pub async fn run( >, soft_detach_tx: mpsc::UnboundedSender<()>, mut detach_result_rx: mpsc::UnboundedReceiver, - choice_answered_tx: mpsc::UnboundedSender<(String, usize)>, + choice_answered_tx: mpsc::UnboundedSender<(infinity_protocol::ChoiceId, usize)>, ) -> Result where T: TermOut, diff --git a/crates/infinity-agent-cli/tests/common/mod.rs b/crates/infinity-agent-cli/tests/common/mod.rs index 8d41ddd7..bdb5cd59 100644 --- a/crates/infinity-agent-cli/tests/common/mod.rs +++ b/crates/infinity-agent-cli/tests/common/mod.rs @@ -442,7 +442,7 @@ pub struct TuiHarness { pub load_session_rx: mpsc::UnboundedReceiver<(Option, bool)>, pub model_switch_rx: mpsc::UnboundedReceiver, pub soft_detach_rx: mpsc::UnboundedReceiver<()>, - pub choice_answered_rx: mpsc::UnboundedReceiver<(String, usize)>, + pub choice_answered_rx: mpsc::UnboundedReceiver<(infinity_protocol::ChoiceId, usize)>, pub handle: tokio::task::JoinHandle>, } diff --git a/crates/infinity-agent-cli/tests/tui_flow_snapshots.rs b/crates/infinity-agent-cli/tests/tui_flow_snapshots.rs index 56a3788f..f7a696a2 100644 --- a/crates/infinity-agent-cli/tests/tui_flow_snapshots.rs +++ b/crates/infinity-agent-cli/tests/tui_flow_snapshots.rs @@ -315,14 +315,14 @@ async fn queued_choices_and_external_complete() { let mut h = TuiHarness::spawn(80, 16).await; h.display(Evt::UserChoiceRequired { - id: "choice-1".to_owned(), + id: "choice-1".into(), prompt: "First question?".to_owned(), choices: vec!["Yes".to_owned(), "No".to_owned()], default: 0, response_url: String::new(), }); h.display(Evt::UserChoiceRequired { - id: "choice-2".to_owned(), + id: "choice-2".into(), prompt: "Second question?".to_owned(), choices: vec!["Red".to_owned(), "Green".to_owned(), "Blue".to_owned()], default: 1, @@ -341,7 +341,7 @@ async fn queued_choices_and_external_complete() { insta::assert_snapshot!("second_choice_shown", h.screen()); h.display(Evt::UserChoiceComplete { - choice_id: "choice-2".to_owned(), + choice_id: "choice-2".into(), }); h.settle().await; insta::assert_snapshot!("choices_all_dismissed", h.screen()); diff --git a/crates/infinity-agent-cli/tests/tui_viewport_snapshots.rs b/crates/infinity-agent-cli/tests/tui_viewport_snapshots.rs index bd0e9561..24142a5f 100644 --- a/crates/infinity-agent-cli/tests/tui_viewport_snapshots.rs +++ b/crates/infinity-agent-cli/tests/tui_viewport_snapshots.rs @@ -116,7 +116,7 @@ async fn choice_picker_during_stream() { chunk: "Working on it...".to_owned(), }); h.display(Evt::UserChoiceRequired { - id: "choice-1".to_owned(), + id: "choice-1".into(), prompt: "Allow the tool to run?".to_owned(), choices: vec!["Allow".to_owned(), "Deny".to_owned()], default: 0, diff --git a/crates/infinity-agent-core/examples/agent_scale.rs b/crates/infinity-agent-core/examples/agent_scale.rs index 1902df1e..b3642561 100644 --- a/crates/infinity-agent-core/examples/agent_scale.rs +++ b/crates/infinity-agent-core/examples/agent_scale.rs @@ -118,15 +118,15 @@ impl Tool for RunCommand { async fn execute( &self, _args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), Box> { let result = InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id, - call_id, + id: id.into_inner(), + call_id: call_id.map(|c| c.into_inner()), content: vec![ToolResultContent::Text(infinity_provider_protocol::message::Text { text: "test result: ok. 148 passed; 0 failed; 3 ignored; finished in 21.38s" diff --git a/crates/infinity-agent-core/src/event_processor.rs b/crates/infinity-agent-core/src/event_processor.rs index e16f5027..ebddef58 100644 --- a/crates/infinity-agent-core/src/event_processor.rs +++ b/crates/infinity-agent-core/src/event_processor.rs @@ -9,7 +9,7 @@ use infinity_provider_protocol::{ CompletionRequest, StreamChunk, ToolCallDeltaContent, ToolDefinition, message::{AssistantContent, Message, ToolResult, ToolResultContent, UserContent}, }; -use rap_protocol::ThreadId; +use rap_protocol::{ChoiceId, ProviderCallId, ThreadId, ToolCallId}; use serde::Serialize; use tracing; @@ -69,7 +69,7 @@ pub enum PrepareResult { OAuthRequired { auth_url: String }, /// A user choice prompt must be surfaced to the user. UserChoiceRequired { - id: String, + id: ChoiceId, prompt: String, choices: Vec, default: usize, @@ -89,8 +89,8 @@ pub enum CompletionAction { ExecuteToolCall { tool_name: String, tool_args: serde_json::Value, - tool_call_id: String, - call_id: Option, + tool_call_id: ToolCallId, + call_id: Option, display_as: Option, }, } @@ -155,7 +155,7 @@ pub struct HistoryManager { /// Tool call IDs that were interrupted by a new user message during /// `handle_content`. Callers can drain this via `take_interrupted_tool_calls` /// to send best-effort cancellation notifications to RAP tool servers. - interrupted_tool_calls: RefCell>, + interrupted_tool_calls: RefCell>, /// Tracks the absolute store index that the current in-memory compaction /// summary covers up to. Used to compute the correct relative split /// position when a second compaction is applied on top of an existing one. @@ -312,7 +312,7 @@ impl HistoryManager { tracing::info!("Tool call {} interrupted by incoming message", tool_call.id); self.interrupted_tool_calls .borrow_mut() - .push(tool_call.id.clone()); + .push(ToolCallId::from(tool_call.id.clone())); let synthetic_result = InfinityMessage::ToolResult { result: ToolResult { id: tool_call.id.clone(), @@ -596,7 +596,7 @@ impl HistoryManager { /// Drain and return tool call IDs that were interrupted by new user messages. /// Callers use this to send best-effort cancellation notifications to RAP /// tool servers so they can abort in-flight operations. - pub fn take_interrupted_tool_calls(&self) -> Vec { + pub fn take_interrupted_tool_calls(&self) -> Vec { std::mem::take(&mut *self.interrupted_tool_calls.borrow_mut()) } @@ -642,7 +642,7 @@ impl HistoryManager { /// `tool_call_id` is the ID of the tool call whose result had /// `subscription: true`. Ownership is implicit — a subscription is /// stored in the thread that created it. - pub async fn track_subscription(&self, tool_call_id: &str) -> Result<(), BoxError> { + pub async fn track_subscription(&self, tool_call_id: &ToolCallId) -> Result<(), BoxError> { self.state_store .add_active_subscription(&self.thread_id, tool_call_id) .await @@ -650,7 +650,10 @@ impl HistoryManager { } /// Remove a subscription from the current thread's active tracking. - pub async fn remove_subscription(&self, tool_call_id: &str) -> Result<(), BoxError> { + pub async fn remove_subscription( + &self, + tool_call_id: &ToolCallId, + ) -> Result<(), BoxError> { self.state_store .remove_active_subscription(&self.thread_id, tool_call_id) .await @@ -705,7 +708,7 @@ where .as_ref() .is_some_and(SyntheticKind::is_compaction) { - let spawn_call_id = uuid::Uuid::new_v4().to_string(); + let spawn_call_id = ToolCallId::from(uuid::Uuid::new_v4().to_string()); // Compute a safe compaction point: exclude trailing unanswered tool calls // from the compaction range so they aren't lost when apply_compaction runs. @@ -732,7 +735,7 @@ where // inherited history. let spawn_tool_call = InfinityMessage::ToolCall { call: infinity_provider_protocol::message::ToolCall { - id: spawn_call_id.clone(), + id: spawn_call_id.as_str().to_owned(), call_id: None, function: infinity_provider_protocol::message::ToolFunction { name: "__harness_begin_compaction__".to_owned(), @@ -755,7 +758,7 @@ where // Send child its instructions via message sender let child_msg = InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: spawn_call_id.clone(), + id: spawn_call_id.as_str().to_owned(), call_id: None, content: vec![ToolResultContent::Text( infinity_provider_protocol::message::Text { @@ -776,7 +779,7 @@ where subscription: false, }; message_sender - .send_to_input_queue(child_msg, &spawn_call_id) + .send_to_input_queue(child_msg, spawn_call_id.as_str()) .await .map_err(|e| Box::new(e) as BoxError)?; @@ -821,7 +824,7 @@ where // Handle synthetic tool results (subscription events / thread reports) // Capture metadata for SubscriptionEvent variant before synthetic_kind is consumed. - let subscription_event_meta: Option<(String, Option)> = + let subscription_event_meta: Option<(ToolCallId, Option)> = input_msg.synthetic.as_ref().and_then(|s| { if s.is_thread_report() || s.is_associative() || s.is_parent_message() { let child_id = if let SyntheticKind::Tagged(TaggedSyntheticKind::ThreadReport { @@ -852,7 +855,7 @@ where let original_call = current_history.history.borrow().iter().find_map(|msg| { if let InfinityMessage::ToolCall { call, .. } = msg - && call.id == original_tool_call_id + && call.id == original_tool_call_id.as_str() { Some(call.clone()) } else { @@ -925,11 +928,11 @@ where input_msg.group_id ); - let event_call_id = uuid::Uuid::new_v4().to_string(); - let spawn_call_id = uuid::Uuid::new_v4().to_string(); + let event_call_id = ToolCallId::from(uuid::Uuid::new_v4().to_string()); + let spawn_call_id = ToolCallId::from(uuid::Uuid::new_v4().to_string()); let event_content = if let UserContent::ToolResult(mut tool_result) = user_content { - tool_result.id = event_call_id.clone(); + tool_result.id = event_call_id.as_str().to_owned(); tool_result.call_id = None; tool_result } else { @@ -944,7 +947,7 @@ where // Write event + spawn tool calls directly to child's store let event_tool_call = InfinityMessage::ToolCall { call: infinity_provider_protocol::message::ToolCall { - id: event_call_id.clone(), + id: event_call_id.as_str().to_owned(), call_id: None, function: infinity_provider_protocol::message::ToolFunction { name: "receive_event__injected".to_owned(), @@ -959,7 +962,7 @@ where }; let spawn_tool_call = InfinityMessage::ToolCall { call: infinity_provider_protocol::message::ToolCall { - id: spawn_call_id.clone(), + id: spawn_call_id.as_str().to_owned(), call_id: None, function: infinity_provider_protocol::message::ToolFunction { name: "spawn_thread".to_owned(), @@ -972,7 +975,7 @@ where }; let spawn_tool_result = InfinityMessage::ToolResult { result: ToolResult { - id: spawn_call_id.clone(), + id: spawn_call_id.as_str().to_owned(), call_id: None, content: vec![ToolResultContent::Text( infinity_provider_protocol::message::Text { @@ -1005,7 +1008,7 @@ where subscription: false, }; message_sender - .send_to_input_queue(child_msg, &event_call_id) + .send_to_input_queue(child_msg, event_call_id.as_str()) .await .map_err(|e| Box::new(e) as BoxError)?; @@ -1069,7 +1072,9 @@ where tool_call_id, current_history.thread_id ); - current_history.track_subscription(tool_call_id).await?; + current_history + .track_subscription(ToolCallId::from_ref(tool_call_id)) + .await?; } Ok(PrepareResult::Ready) @@ -1094,7 +1099,7 @@ where if let Message::Assistant { content, .. } = h && let Some(AssistantContent::ToolCall(c)) = content.first() { - c.id == synth.tool_call_id() + c.id == synth.tool_call_id().as_str() } else { false } @@ -1487,8 +1492,8 @@ where let res = tool.execute_synchronous( &call.function.arguments, - &call.id, - call.call_id.as_deref(), + ToolCallId::from_ref(&call.id), + call.call_id.as_deref().map(ProviderCallId::from_ref), tool_context, ).await.expect("bug: synchronous tool execution failed"); @@ -1513,8 +1518,8 @@ where yield CompletionEvent::Action(CompletionAction::ExecuteToolCall { tool_name: call.function.name.clone(), tool_args: call.function.arguments.clone(), - tool_call_id: call.id.clone(), - call_id: call.call_id.clone(), + tool_call_id: ToolCallId::from(call.id.clone()), + call_id: call.call_id.clone().map(ProviderCallId::from), display_as: tool_display_as, }); } @@ -2017,7 +2022,7 @@ mod tests { let input = InputMessage { content: InputMessageContent::OAuth(OAuthRequired { content_type: "oauth_required".to_owned(), - id: "oauth-1".to_owned(), + id: "oauth-1".into(), call_id: None, auth_url: "https://example.com/auth".to_owned(), }), @@ -2149,7 +2154,7 @@ mod tests { "tc-sub", "thread report data", Some(SyntheticKind::Tagged(TaggedSyntheticKind::ThreadReport { - tool_call_id: "tc-sub".to_owned(), + tool_call_id: "tc-sub".into(), child_thread_id: "thread-1".into(), })), ); @@ -2187,7 +2192,7 @@ mod tests { "tc-sub", "thread report data", Some(SyntheticKind::Tagged(TaggedSyntheticKind::ThreadReport { - tool_call_id: "tc-sub".to_owned(), + tool_call_id: "tc-sub".into(), child_thread_id: "thread-1".into(), })), ); @@ -2236,7 +2241,7 @@ mod tests { "event payload", Some(SyntheticKind::Tagged( TaggedSyntheticKind::SubscriptionEvent { - tool_call_id: "tc-sub".to_owned(), + tool_call_id: "tc-sub".into(), associative: false, r#final: false, }, @@ -2273,7 +2278,7 @@ mod tests { "event payload", Some(SyntheticKind::Tagged( TaggedSyntheticKind::SubscriptionEvent { - tool_call_id: "tc-sub".to_owned(), + tool_call_id: "tc-sub".into(), associative: false, r#final: false, }, @@ -2300,7 +2305,7 @@ mod tests { "some data", Some(SyntheticKind::Tagged( TaggedSyntheticKind::SubscriptionEvent { - tool_call_id: "nonexistent-tc".to_owned(), + tool_call_id: "nonexistent-tc".into(), associative: false, r#final: false, }, @@ -2368,7 +2373,7 @@ mod tests { "build output chunk\n[exit code: 0]", Some(SyntheticKind::Tagged( TaggedSyntheticKind::SubscriptionEvent { - tool_call_id: "tc-cmd".to_owned(), + tool_call_id: "tc-cmd".into(), associative: true, r#final: false, }, @@ -2411,7 +2416,7 @@ mod tests { "build output chunk\n[exit code: 0]", Some(SyntheticKind::Tagged( TaggedSyntheticKind::SubscriptionEvent { - tool_call_id: "tc-cmd".to_owned(), + tool_call_id: "tc-cmd".into(), associative: true, r#final: false, }, @@ -2753,8 +2758,8 @@ mod tests { async fn execute( &self, _: serde_json::Value, - _: String, - _: Option, + _: ToolCallId, + _: Option, _: &ToolContext, ) -> Result<(), Box> { Ok(()) @@ -2765,14 +2770,14 @@ mod tests { async fn execute_synchronous( &self, args: &serde_json::Value, - id: &str, - call_id: Option<&str>, + id: &ToolCallId, + call_id: Option<&ProviderCallId>, _ctx: &ToolContext, ) -> Option { let text = args["text"].as_str().unwrap_or("?"); Some(ToolResult { - id: id.to_owned(), - call_id: call_id.map(String::from), + id: id.as_str().to_owned(), + call_id: call_id.map(|c| c.as_str().to_owned()), content: vec![ToolResultContent::Text( infinity_provider_protocol::message::Text { text: format!("echo: {}", text), @@ -3130,8 +3135,8 @@ mod tests { async fn execute( &self, _: serde_json::Value, - _: String, - _: Option, + _: ToolCallId, + _: Option, _: &ToolContext, ) -> Result<(), Box> { Ok(()) @@ -3233,8 +3238,8 @@ mod tests { async fn execute( &self, _: serde_json::Value, - _: String, - _: Option, + _: ToolCallId, + _: Option, _: &ToolContext, ) -> Result<(), Box> { Ok(()) @@ -3778,8 +3783,8 @@ mod tests { async fn execute( &self, _: serde_json::Value, - _: String, - _: Option, + _: ToolCallId, + _: Option, _: &ToolContext, ) -> Result<(), Box> { Ok(()) @@ -3988,8 +3993,8 @@ mod tests { async fn execute( &self, _: serde_json::Value, - _: String, - _: Option, + _: ToolCallId, + _: Option, _: &ToolContext, ) -> Result<(), Box> { Err("stub".into()) @@ -4092,7 +4097,7 @@ mod tests { let input = InputMessage { content: InputMessageContent::UserChoice(crate::message::UserChoiceRequired { content_type: "user_choice_required".to_owned(), - id: "choice-1".to_owned(), + id: "choice-1".into(), call_id: None, prompt: "pick one".to_owned(), choices: vec!["a".to_owned(), "b".to_owned()], @@ -4120,7 +4125,7 @@ mod tests { else { panic!("expected UserChoiceRequired, got {result:?}"); }; - assert_eq!(id, "choice-1"); + assert_eq!(id.as_str(), "choice-1"); assert_eq!(prompt, "pick one"); assert_eq!(choices, vec!["a".to_owned(), "b".to_owned()]); assert_eq!(default, 0); diff --git a/crates/infinity-agent-core/src/message.rs b/crates/infinity-agent-core/src/message.rs index 2e3725d5..13b5f9e9 100644 --- a/crates/infinity-agent-core/src/message.rs +++ b/crates/infinity-agent-core/src/message.rs @@ -1,7 +1,7 @@ use infinity_provider_protocol::message::{ AssistantContent, Message, ToolCall, ToolResult, UserContent, }; -use rap_protocol::ThreadId; +use rap_protocol::{ChoiceId, ProviderCallId, ThreadId, ToolCallId}; use serde::{Deserialize, Serialize}; #[derive(Debug, Deserialize, Serialize, Clone)] @@ -16,8 +16,8 @@ pub enum InputMessageContent { pub struct OAuthRequired { #[serde(rename = "type")] pub content_type: String, - pub id: String, - pub call_id: Option, + pub id: ToolCallId, + pub call_id: Option, pub auth_url: String, } @@ -25,8 +25,8 @@ pub struct OAuthRequired { pub struct UserChoiceRequired { #[serde(rename = "type")] pub content_type: String, - pub id: String, - pub call_id: Option, + pub id: ChoiceId, + pub call_id: Option, pub prompt: String, pub choices: Vec, pub default: usize, @@ -41,7 +41,7 @@ pub struct UserChoiceRequired { pub enum SyntheticKind { Tagged(TaggedSyntheticKind), /// Backward compat: a bare string is treated as a subscription event - SubscriptionEvent(String), + SubscriptionEvent(ToolCallId), } #[derive(Debug, Deserialize, Serialize, Clone)] @@ -49,7 +49,7 @@ pub enum SyntheticKind { pub enum TaggedSyntheticKind { #[serde(rename = "subscription_event")] SubscriptionEvent { - tool_call_id: String, + tool_call_id: ToolCallId, #[serde(default)] associative: bool, /// When true, this is the final event — the runtime removes the @@ -59,11 +59,11 @@ pub enum TaggedSyntheticKind { }, #[serde(rename = "thread_report")] ThreadReport { - tool_call_id: String, + tool_call_id: ToolCallId, child_thread_id: ThreadId, }, #[serde(rename = "parent_message")] - ParentMessage { tool_call_id: String }, + ParentMessage { tool_call_id: ToolCallId }, #[serde(rename = "compaction")] Compaction, #[serde(rename = "compaction_complete")] @@ -71,7 +71,7 @@ pub enum TaggedSyntheticKind { } impl SyntheticKind { - pub fn tool_call_id(&self) -> &str { + pub fn tool_call_id(&self) -> &ToolCallId { match self { SyntheticKind::Tagged(TaggedSyntheticKind::SubscriptionEvent { tool_call_id, .. @@ -82,8 +82,10 @@ impl SyntheticKind { SyntheticKind::Tagged(TaggedSyntheticKind::ParentMessage { tool_call_id }) => { tool_call_id } - SyntheticKind::Tagged(TaggedSyntheticKind::Compaction) => "", - SyntheticKind::Tagged(TaggedSyntheticKind::CompactionComplete) => "", + SyntheticKind::Tagged(TaggedSyntheticKind::Compaction) => ToolCallId::from_ref(""), + SyntheticKind::Tagged(TaggedSyntheticKind::CompactionComplete) => { + ToolCallId::from_ref("") + } SyntheticKind::SubscriptionEvent(id) => id, } } @@ -210,7 +212,7 @@ pub enum InfinityMessage { SubscriptionEvent { result: Box, /// The tool_call_id of the original subscription tool call. - tool_call_id: String, + tool_call_id: ToolCallId, /// Set when this is a thread report (used to build the display name). #[serde(default, skip_serializing_if = "Option::is_none")] child_thread_id: Option, diff --git a/crates/infinity-agent-core/src/stores.rs b/crates/infinity-agent-core/src/stores.rs index b104a108..3b2b7cf6 100644 --- a/crates/infinity-agent-core/src/stores.rs +++ b/crates/infinity-agent-core/src/stores.rs @@ -11,7 +11,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::{Arc, Mutex}; use async_trait::async_trait; -use rap_protocol::ThreadId; +use rap_protocol::{ChoiceId, ThreadId, ToolCallId}; use serde::{Deserialize, Serialize}; use crate::message::InfinityMessage; @@ -44,7 +44,7 @@ pub struct ThreadInfo { pub root_thread_id: ThreadId, /// Number of parent messages the child inherits (history cutoff). pub spawn_message_order: Option, - pub spawn_tool_call_id: Option, + pub spawn_tool_call_id: Option, pub closed: bool, pub is_subscription_event: bool, #[serde(default)] @@ -159,7 +159,7 @@ impl InMemoryConversationStore { &self, new_thread_id: &ThreadId, parent_thread_id: &ThreadId, - spawn_tool_call_id: &str, + spawn_tool_call_id: &ToolCallId, is_for_subscription_event: bool, spawn_order_override: Option, ) { @@ -251,7 +251,7 @@ impl ConversationStore for InMemoryConversationStore { async fn spawn_thread( &self, parent_thread_id: &ThreadId, - spawn_tool_call_id: &str, + spawn_tool_call_id: &ToolCallId, is_for_subscription_event: bool, spawn_order_override: Option, ) -> Result { @@ -293,7 +293,7 @@ impl ConversationStore for InMemoryConversationStore { async fn get_thread_parent_info( &self, thread_id: &ThreadId, - ) -> Result, Self::Error> { + ) -> Result, Self::Error> { let threads = self.threads.lock().expect("bug: mutex poisoned"); Ok(threads.get(thread_id).and_then(|t| { match (&t.parent_thread_id, &t.spawn_tool_call_id) { @@ -395,7 +395,7 @@ pub struct ThreadState { #[serde(default)] pub metadata: Option, #[serde(default)] - pub subscriptions: HashSet, + pub subscriptions: HashSet, #[serde(default)] pub pending_user_choices: Vec, } @@ -406,7 +406,7 @@ pub struct ThreadState { pub struct InMemoryStateStore { processed_ids: Arc>>>, metadata: Arc>>, - subscriptions: Arc>>>, + subscriptions: Arc>>>, pending_user_choices: Arc>>>, } @@ -508,7 +508,7 @@ impl StateStore for InMemoryStateStore { async fn get_active_subscriptions( &self, thread_id: &ThreadId, - ) -> Result, Self::Error> { + ) -> Result, Self::Error> { let store = self.subscriptions.lock().expect("bug: mutex poisoned"); Ok(store .get(thread_id) @@ -519,7 +519,7 @@ impl StateStore for InMemoryStateStore { async fn add_active_subscription( &self, thread_id: &ThreadId, - tool_call_id: &str, + tool_call_id: &ToolCallId, ) -> Result<(), Self::Error> { let mut store = self.subscriptions.lock().expect("bug: mutex poisoned"); store @@ -532,7 +532,7 @@ impl StateStore for InMemoryStateStore { async fn remove_active_subscription( &self, thread_id: &ThreadId, - tool_call_id: &str, + tool_call_id: &ToolCallId, ) -> Result<(), Self::Error> { let mut store = self.subscriptions.lock().expect("bug: mutex poisoned"); if let Some(s) = store.get_mut(thread_id) { @@ -562,7 +562,7 @@ impl StateStore for InMemoryStateStore { async fn remove_pending_user_choice( &self, thread_id: &ThreadId, - choice_id: &str, + choice_id: &ChoiceId, ) -> Result<(), Self::Error> { if let Some(choices) = self .pending_user_choices @@ -570,7 +570,7 @@ impl StateStore for InMemoryStateStore { .expect("bug: mutex poisoned") .get_mut(thread_id) { - choices.retain(|choice| choice.id != choice_id); + choices.retain(|choice| choice.id != *choice_id); } Ok(()) } diff --git a/crates/infinity-agent-core/src/system/events.rs b/crates/infinity-agent-core/src/system/events.rs index 48b8e96a..bcd92b3d 100644 --- a/crates/infinity-agent-core/src/system/events.rs +++ b/crates/infinity-agent-core/src/system/events.rs @@ -8,7 +8,7 @@ use crate::message::InfinityMessage; /// `response_url`. #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct UserChoice { - pub id: String, + pub id: rap_protocol::ChoiceId, pub prompt: String, pub choices: Vec, pub default: usize, @@ -63,7 +63,7 @@ pub enum AgentEvent { UserChoiceRequired { choice: UserChoice }, /// A pending user choice became moot and has already been removed from /// persistent state. - UserChoiceDismissed { choice_id: String }, + UserChoiceDismissed { choice_id: rap_protocol::ChoiceId }, /// Human-readable progress/diagnostic information (retries, warnings). Info { text: String }, } diff --git a/crates/infinity-agent-core/src/system/test_support.rs b/crates/infinity-agent-core/src/system/test_support.rs index 33af91c1..89b432eb 100644 --- a/crates/infinity-agent-core/src/system/test_support.rs +++ b/crates/infinity-agent-core/src/system/test_support.rs @@ -158,8 +158,8 @@ impl Tool for AsyncTool { async fn execute( &self, _: serde_json::Value, - _: String, - _: Option, + _: rap_protocol::ToolCallId, + _: Option, _: &ToolContext, ) -> Result<(), Box> { Ok(()) @@ -183,8 +183,8 @@ impl Tool for FailingTool { async fn execute( &self, _: serde_json::Value, - _: String, - _: Option, + _: rap_protocol::ToolCallId, + _: Option, _: &ToolContext, ) -> Result<(), Box> { Err("boom".into()) @@ -354,15 +354,15 @@ impl Tool for SubscribeTool { async fn execute( &self, _: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, ctx: &ToolContext, ) -> Result<(), Box> { let msg = InputMessage { content: InputMessageContent::User(UserContent::ToolResult( infinity_provider_protocol::message::ToolResult { - id: id.clone(), - call_id, + id: id.clone().into_inner(), + call_id: call_id.map(|c| c.into_inner()), content: vec![ infinity_provider_protocol::message::ToolResultContent::Text( infinity_provider_protocol::message::Text { @@ -378,7 +378,9 @@ impl Tool for SubscribeTool { display_as: None, subscription: true, }; - ctx.message_sender.send_to_input_queue(msg, &id).await?; + ctx.message_sender + .send_to_input_queue(msg, id.as_str()) + .await?; Ok(()) } } @@ -420,8 +422,8 @@ impl Tool for NamedTool { async fn execute( &self, _: serde_json::Value, - _: String, - _: Option, + _: rap_protocol::ToolCallId, + _: Option, _: &ToolContext, ) -> Result<(), Box> { Ok(()) diff --git a/crates/infinity-agent-core/src/system/tests.rs b/crates/infinity-agent-core/src/system/tests.rs index 5b231775..06d475be 100644 --- a/crates/infinity-agent-core/src/system/tests.rs +++ b/crates/infinity-agent-core/src/system/tests.rs @@ -40,7 +40,7 @@ async fn pending_choices_are_thread_local_and_replayed() { InputMessage { content: InputMessageContent::UserChoice(UserChoiceRequired { content_type: "user_choice_required".to_owned(), - id: choice_id.to_owned(), + id: choice_id.into(), call_id: None, prompt: "Pick one".to_owned(), choices: vec!["A".to_owned(), "B".to_owned()], @@ -74,7 +74,8 @@ async fn pending_choices_are_thread_local_and_replayed() { .get_pending_user_choices(ThreadId::from_ref("t1")) .await .expect("load first thread choices")[0] - .id, + .id + .as_str(), "choice-1" ); assert_eq!( @@ -82,7 +83,8 @@ async fn pending_choices_are_thread_local_and_replayed() { .get_pending_user_choices(ThreadId::from_ref("t2")) .await .expect("load second thread choices")[0] - .id, + .id + .as_str(), "choice-2" ); @@ -101,7 +103,7 @@ async fn pending_choices_are_thread_local_and_replayed() { panic!("expected replay"); }; assert_eq!(snapshot.pending_choices.len(), 1); - assert_eq!(snapshot.pending_choices[0].id, "choice-1"); + assert_eq!(snapshot.pending_choices[0].id.as_str(), "choice-1"); }) .await; } @@ -342,8 +344,8 @@ async fn driver_idles_after_close_thread_tool_call() { async fn execute( &self, _: serde_json::Value, - _: String, - _: Option, + _: rap_protocol::ToolCallId, + _: Option, _: &ToolContext, ) -> Result<(), Box> { Ok(()) diff --git a/crates/infinity-agent-core/src/system/thread.rs b/crates/infinity-agent-core/src/system/thread.rs index ac2a75be..05d802e0 100644 --- a/crates/infinity-agent-core/src/system/thread.rs +++ b/crates/infinity-agent-core/src/system/thread.rs @@ -367,19 +367,22 @@ where .pending_choices .borrow() .iter() - .any(|choice| choice.id == result.id) + .any(|choice| choice.id.as_str() == result.id) { self.inner .state_store - .remove_pending_user_choice(&thread_id, &result.id) + .remove_pending_user_choice( + &thread_id, + rap_protocol::ChoiceId::from_ref(&result.id), + ) .await?; self.pending_choices .borrow_mut() - .retain(|choice| choice.id != result.id); + .retain(|choice| choice.id.as_str() != result.id); observer.on_event( &thread_id, &AgentEvent::UserChoiceDismissed { - choice_id: result.id.clone(), + choice_id: rap_protocol::ChoiceId::from(result.id.clone()), }, ); } diff --git a/crates/infinity-agent-core/src/tools/cancel_subscription.rs b/crates/infinity-agent-core/src/tools/cancel_subscription.rs index 6a595963..ee0c0470 100644 --- a/crates/infinity-agent-core/src/tools/cancel_subscription.rs +++ b/crates/infinity-agent-core/src/tools/cancel_subscription.rs @@ -49,8 +49,8 @@ impl async fn execute( &self, _args: serde_json::Value, - _id: String, - _call_id: Option, + _id: rap_protocol::ToolCallId, + _call_id: Option, _context: &ToolContext, ) -> Result<(), Box> { // Synchronous-only tool; execute is a no-op. @@ -64,11 +64,15 @@ impl async fn execute_synchronous( &self, args: &serde_json::Value, - id: &str, - call_id: Option<&str>, + id: &rap_protocol::ToolCallId, + call_id: Option<&rap_protocol::ProviderCallId>, context: &ToolContext, ) -> Option { - let Some(tool_call_id) = args.get("tool_call_id").and_then(|v| v.as_str()) else { + let Some(tool_call_id) = args + .get("tool_call_id") + .and_then(|v| v.as_str()) + .map(rap_protocol::ToolCallId::from_ref) + else { return Some(error_result(id, call_id, "Error: tool_call_id is required")); }; @@ -124,8 +128,8 @@ impl ); Some(ToolResult { - id: id.to_owned(), - call_id: call_id.map(String::from), + id: id.as_str().to_owned(), + call_id: call_id.map(|c| c.as_str().to_owned()), content: vec![ToolResultContent::Text(Text { text: format!("Subscription '{}' cancelled successfully.", tool_call_id), })], @@ -133,10 +137,14 @@ impl } } -fn error_result(id: &str, call_id: Option<&str>, text: &str) -> ToolResult { +fn error_result( + id: &rap_protocol::ToolCallId, + call_id: Option<&rap_protocol::ProviderCallId>, + text: &str, +) -> ToolResult { ToolResult { - id: id.to_owned(), - call_id: call_id.map(String::from), + id: id.as_str().to_owned(), + call_id: call_id.map(|c| c.as_str().to_owned()), content: vec![ToolResultContent::Text(Text { text: text.to_owned(), })], diff --git a/crates/infinity-agent-core/src/tools/mod.rs b/crates/infinity-agent-core/src/tools/mod.rs index e67beed1..39c006da 100644 --- a/crates/infinity-agent-core/src/tools/mod.rs +++ b/crates/infinity-agent-core/src/tools/mod.rs @@ -15,8 +15,8 @@ pub(crate) type ToolError = Box; /// Enqueue an error as the result of a tool call so the agent can recover. pub(crate) async fn send_tool_error( context: &ToolContext, - id: &str, - call_id: Option, + id: &rap_protocol::ToolCallId, + call_id: Option, error: impl Into, ) -> Result<(), ToolError> where @@ -24,8 +24,8 @@ where { let message = InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: id.to_owned(), - call_id, + id: id.as_str().to_owned(), + call_id: call_id.map(|c| c.into_inner()), content: vec![ToolResultContent::Text(Text { text: format!("Error: {}", error.into()), })], @@ -39,7 +39,7 @@ where context .message_sender - .send_to_input_queue(message, id) + .send_to_input_queue(message, id.as_str()) .await .map_err(|error| Box::new(error) as ToolError) } @@ -63,8 +63,8 @@ pub trait Tool: Send + Sync { async fn execute( &self, args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), Box>; @@ -95,8 +95,8 @@ pub trait Tool: Send + Sync { async fn execute_synchronous( &self, _args: &serde_json::Value, - _id: &str, - _call_id: Option<&str>, + _id: &rap_protocol::ToolCallId, + _call_id: Option<&rap_protocol::ProviderCallId>, _context: &ToolContext, ) -> Option { None diff --git a/crates/infinity-agent-core/src/tools/rap_tool.rs b/crates/infinity-agent-core/src/tools/rap_tool.rs index 2a14d2fe..995c3e01 100644 --- a/crates/infinity-agent-core/src/tools/rap_tool.rs +++ b/crates/infinity-agent-core/src/tools/rap_tool.rs @@ -49,9 +49,9 @@ pub struct RapInvocationParams<'a> { /// The model-provided arguments. pub arguments: serde_json::Value, /// The tool-call ID from `execute`. - pub id: String, + pub id: rap_protocol::ToolCallId, /// The provider call ID from `execute`. - pub call_id: Option, + pub call_id: Option, /// Callback destination for the server's asynchronous results. `None` /// uses [`ToolContext::callback_url`]; embeddings that run their own /// callback listener pass its URL instead. @@ -163,8 +163,8 @@ impl Tool for RapTool { async fn execute( &self, args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), BoxError> { invoke_rap_tool( @@ -284,8 +284,8 @@ mod tests { endpoint: "http://server/invoke", operation: "lookup", arguments, - id: "tc-1".to_owned(), - call_id: Some("call-1".to_owned()), + id: "tc-1".into(), + call_id: Some("call-1".into()), callback_url: None, } } diff --git a/crates/infinity-agent-core/src/tools/sleep.rs b/crates/infinity-agent-core/src/tools/sleep.rs index db99659d..8dd719ae 100644 --- a/crates/infinity-agent-core/src/tools/sleep.rs +++ b/crates/infinity-agent-core/src/tools/sleep.rs @@ -39,8 +39,8 @@ impl Tool for SleepUntilEventOrInputTool { async fn execute( &self, _args: serde_json::Value, - _id: String, - _call_id: Option, + _id: rap_protocol::ToolCallId, + _call_id: Option, _context: &ToolContext, ) -> Result<(), Box> { tracing::info!("sleep_until_event_or_input invoked, agent will pause until next input"); @@ -87,8 +87,8 @@ impl Tool for TokioSleepTool { async fn execute( &self, args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), Box> { let seconds = args["seconds"].as_f64().unwrap_or(0.0); @@ -101,8 +101,8 @@ impl Tool for TokioSleepTool { } let msg = InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: id.clone(), - call_id, + id: id.clone().into_inner(), + call_id: call_id.map(|c| c.into_inner()), content: vec![ToolResultContent::Text(Text { text: format!("Slept for {} seconds", seconds), })], @@ -113,7 +113,7 @@ impl Tool for TokioSleepTool { display_as: None, subscription: false, }; - if let Err(e) = sender.send_to_input_queue(msg, &id).await { + if let Err(e) = sender.send_to_input_queue(msg, id.as_str()).await { tracing::error!("Failed to deliver sleep result: {}", e); } })); @@ -162,8 +162,8 @@ impl Tool for TokioSleepUntilTool { async fn execute( &self, args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), Box> { let date_str = args["date"].as_str().unwrap_or("").to_owned(); @@ -208,8 +208,8 @@ impl Tool for TokioSleepUntilTool { } let msg = InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: id.clone(), - call_id, + id: id.clone().into_inner(), + call_id: call_id.map(|c| c.into_inner()), content: vec![ToolResultContent::Text(Text { text: result_text })], })), group_id: group_id.clone(), @@ -218,7 +218,7 @@ impl Tool for TokioSleepUntilTool { display_as: None, subscription: false, }; - if let Err(e) = sender.send_to_input_queue(msg, &id).await { + if let Err(e) = sender.send_to_input_queue(msg, id.as_str()).await { tracing::error!("Failed to deliver sleep_until result: {}", e); } })); diff --git a/crates/infinity-agent-core/src/tools/thread.rs b/crates/infinity-agent-core/src/tools/thread.rs index dfd10309..365d1a3f 100644 --- a/crates/infinity-agent-core/src/tools/thread.rs +++ b/crates/infinity-agent-core/src/tools/thread.rs @@ -45,8 +45,8 @@ impl Tool for Spawn async fn execute( &self, _args: serde_json::Value, - _id: String, - _call_id: Option, + _id: rap_protocol::ToolCallId, + _call_id: Option, _context: &ToolContext, ) -> Result<(), Box> { Ok(()) @@ -59,8 +59,8 @@ impl Tool for Spawn async fn execute_synchronous( &self, args: &serde_json::Value, - id: &str, - call_id: Option<&str>, + id: &rap_protocol::ToolCallId, + call_id: Option<&rap_protocol::ProviderCallId>, context: &ToolContext, ) -> Option { // Validate child_of matches the actual thread stack @@ -71,8 +71,8 @@ impl Tool for Spawn if child_of != context.thread_stack { return Some(ToolResult { - id: id.to_owned(), - call_id: call_id.map(|c| c.to_owned()), + id: id.as_str().to_owned(), + call_id: call_id.map(|c| c.as_str().to_owned()), content: vec![ToolResultContent::Text(Text { text: format!( "Error: child_of {:?} does not match the actual thread stack {:?}. You may be confused and think you are in the parent thread, but you are not. You are in thread {}. Do NOT spawn threads — focus on your assigned task.", @@ -91,8 +91,8 @@ impl Tool for Spawn Err(e) => { tracing::error!("failed to spawn thread in conversation store: {e}"); return Some(ToolResult { - id: id.to_owned(), - call_id: call_id.map(|c| c.to_owned()), + id: id.as_str().to_owned(), + call_id: call_id.map(|c| c.as_str().to_owned()), content: vec![ToolResultContent::Text(Text { text: format!("Error: failed to spawn thread: {e}"), })], @@ -107,8 +107,8 @@ impl Tool for Spawn ); let parent_result = ToolResult { - id: id.to_owned(), - call_id: call_id.map(|c| c.to_owned()), + id: id.as_str().to_owned(), + call_id: call_id.map(|c| c.as_str().to_owned()), content: vec![ToolResultContent::Text(Text { text: format!( "Child thread is successfully spawned and has ID: {}. You will be notified automatically when the child has anything to report. Make sure that you **do not** do the task assigned to the child thread.", @@ -125,8 +125,8 @@ impl Tool for Spawn let child_result = InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: id.to_owned(), - call_id: call_id.map(|c| c.to_owned()), + id: id.as_str().to_owned(), + call_id: call_id.map(|c| c.as_str().to_owned()), content: vec![ToolResultContent::Text(Text { text: format!( "You are now INSIDE the thread that you requested to create. Your thread ID is {}. Your next task is to exactly follow these instructions: {}\n. Start by repeating to yourself the instructions, ignoring thinking from the parent context. Make sure to not be confused by the parent context. If the parent was planning to spawn more threads, you should not.", @@ -143,7 +143,7 @@ impl Tool for Spawn context .message_sender - .send_to_input_queue(child_result, id) + .send_to_input_queue(child_result, id.as_str()) .await .expect("failed to send child thread message to input queue"); @@ -182,8 +182,8 @@ impl Tool for Repor async fn execute( &self, args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), Box> { let Some(report_text) = args["report"].as_str() else { @@ -247,13 +247,13 @@ impl Tool for Repor context .message_sender - .send_to_input_queue(report_message, &id) + .send_to_input_queue(report_message, id.as_str()) .await?; let tool_result = InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: id.clone(), - call_id, + id: id.clone().into_inner(), + call_id: call_id.map(|c| c.into_inner()), content: vec![ToolResultContent::Text(Text { text: "Report sent to parent thread.".to_owned(), })], @@ -267,7 +267,7 @@ impl Tool for Repor context .message_sender - .send_to_input_queue(tool_result, &id) + .send_to_input_queue(tool_result, id.as_str()) .await?; Ok(()) @@ -314,8 +314,8 @@ impl, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), Box> { let Some(thread_id) = args["thread_id"].as_str() else { @@ -406,7 +406,7 @@ impl Tool async fn execute( &self, args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), Box> { let Some(child_thread_id) = args["thread_id"].as_str().map(ThreadId::from) else { @@ -576,13 +576,13 @@ impl Tool context .message_sender - .send_to_input_queue(child_message, &id) + .send_to_input_queue(child_message, id.as_str()) .await?; let tool_result = InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: id.clone(), - call_id, + id: id.clone().into_inner(), + call_id: call_id.map(|c| c.into_inner()), content: vec![ToolResultContent::Text(Text { text: "Message sent to child thread.".to_owned(), })], @@ -596,7 +596,7 @@ impl Tool context .message_sender - .send_to_input_queue(tool_result, &id) + .send_to_input_queue(tool_result, id.as_str()) .await?; Ok(()) diff --git a/crates/infinity-agent-core/src/traits.rs b/crates/infinity-agent-core/src/traits.rs index 3c0a5dd6..04627fac 100644 --- a/crates/infinity-agent-core/src/traits.rs +++ b/crates/infinity-agent-core/src/traits.rs @@ -1,7 +1,7 @@ use async_trait::async_trait; use infinity_provider_protocol::message::AssistantContent; -use rap_protocol::ThreadId; +use rap_protocol::{ChoiceId, ThreadId, ToolCallId}; use crate::message::{InfinityMessage, InputMessage}; use crate::system::UserChoice; @@ -114,7 +114,7 @@ pub trait ConversationStore: Send + Sync + Clone { async fn spawn_thread( &self, parent_thread_id: &ThreadId, - spawn_tool_call_id: &str, + spawn_tool_call_id: &ToolCallId, is_for_subscription_event: bool, spawn_order_override: Option, ) -> Result; @@ -131,7 +131,7 @@ pub trait ConversationStore: Send + Sync + Clone { async fn get_thread_parent_info( &self, thread_id: &ThreadId, - ) -> Result, Self::Error>; + ) -> Result, Self::Error>; async fn get_ancestor_chain( &self, @@ -201,20 +201,20 @@ pub trait StateStore: Send + Sync + Clone { async fn get_active_subscriptions( &self, thread_id: &ThreadId, - ) -> Result, Self::Error>; + ) -> Result, Self::Error>; /// Record a new active subscription (tool_call_id) for a specific thread. async fn add_active_subscription( &self, thread_id: &ThreadId, - tool_call_id: &str, + tool_call_id: &ToolCallId, ) -> Result<(), Self::Error>; /// Remove an active subscription (tool_call_id) from a specific thread. async fn remove_active_subscription( &self, thread_id: &ThreadId, - tool_call_id: &str, + tool_call_id: &ToolCallId, ) -> Result<(), Self::Error>; /// Add or replace a pending user choice for one thread. @@ -228,7 +228,7 @@ pub trait StateStore: Send + Sync + Clone { async fn remove_pending_user_choice( &self, thread_id: &ThreadId, - choice_id: &str, + choice_id: &ChoiceId, ) -> Result<(), Self::Error>; /// List choices awaiting a response for one thread. diff --git a/crates/infinity-agent-lambda/src/conversation_history.rs b/crates/infinity-agent-lambda/src/conversation_history.rs index d2151fb3..ba768599 100644 --- a/crates/infinity-agent-lambda/src/conversation_history.rs +++ b/crates/infinity-agent-lambda/src/conversation_history.rs @@ -280,7 +280,7 @@ impl ConversationStore for DsqlConversationStore { async fn spawn_thread( &self, parent_thread_id: &ThreadId, - spawn_tool_call_id: &str, + spawn_tool_call_id: &rap_protocol::ToolCallId, is_for_subscription_event: bool, spawn_order_override: Option, ) -> Result { @@ -316,7 +316,7 @@ impl ConversationStore for DsqlConversationStore { .bind(parent_thread_id.as_str()) .bind(&root_thread_id) .bind(spawn_message_order) - .bind(spawn_tool_call_id) + .bind(spawn_tool_call_id.as_str()) .bind(is_for_subscription_event) .execute(&self.pool) .await @@ -361,7 +361,7 @@ impl ConversationStore for DsqlConversationStore { async fn get_thread_parent_info( &self, thread_id: &ThreadId, - ) -> Result, DsqlError> { + ) -> Result, DsqlError> { let row = sqlx::query( r#"SELECT parent_thread_id, spawn_tool_call_id FROM thread_hierarchy WHERE thread_id = $1"#, ) @@ -373,7 +373,7 @@ impl ConversationStore for DsqlConversationStore { let parent: Option = row.get("parent_thread_id"); let tool_call_id: Option = row.get("spawn_tool_call_id"); match (parent, tool_call_id) { - (Some(p), Some(t)) => Ok(Some((ThreadId::from(p), t))), + (Some(p), Some(t)) => Ok(Some((ThreadId::from(p), rap_protocol::ToolCallId::from(t)))), _ => Ok(None), } } diff --git a/crates/infinity-agent-lambda/src/event_handler.rs b/crates/infinity-agent-lambda/src/event_handler.rs index f28a5ef8..f784e7c8 100644 --- a/crates/infinity-agent-lambda/src/event_handler.rs +++ b/crates/infinity-agent-lambda/src/event_handler.rs @@ -205,7 +205,7 @@ pub(crate) async fn function_handler(event: LambdaEvent) -> Result<(), entry.required_choices.push(choice); } AgentEvent::UserChoiceDismissed { choice_id } => { - entry.completed_choices.push(choice_id); + entry.completed_choices.push(choice_id.into_inner()); } _ => {} } @@ -250,7 +250,7 @@ pub(crate) async fn function_handler(event: LambdaEvent) -> Result<(), for choice in required_choices { let message = event_processor::UserChoiceOutputMessage { message_type: "user_choice_required".to_owned(), - id: choice.id, + id: choice.id.into_inner(), prompt: choice.prompt, choices: choice.choices, default: choice.default, diff --git a/crates/infinity-agent-lambda/src/state_store.rs b/crates/infinity-agent-lambda/src/state_store.rs index b4fa0d00..cd085169 100644 --- a/crates/infinity-agent-lambda/src/state_store.rs +++ b/crates/infinity-agent-lambda/src/state_store.rs @@ -128,7 +128,7 @@ impl StateStore for DynamoDbStateStore { async fn get_active_subscriptions( &self, thread_id: &ThreadId, - ) -> Result, DynamoError> { + ) -> Result, DynamoError> { let result = self .client .get_item() @@ -142,7 +142,12 @@ impl StateStore for DynamoDbStateStore { .item .and_then(|item| { if let Some(AttributeValue::Ss(ids)) = item.get("active_subscriptions") { - Some(ids.clone()) + Some( + ids.iter() + .cloned() + .map(rap_protocol::ToolCallId::from) + .collect(), + ) } else { None } @@ -153,14 +158,17 @@ impl StateStore for DynamoDbStateStore { async fn add_active_subscription( &self, thread_id: &ThreadId, - tool_call_id: &str, + tool_call_id: &rap_protocol::ToolCallId, ) -> Result<(), DynamoError> { self.client .update_item() .table_name(&self.table_name) .key("session", AttributeValue::S(thread_id.as_str().to_owned())) .update_expression("ADD active_subscriptions :id") - .expression_attribute_values(":id", AttributeValue::Ss(vec![tool_call_id.to_owned()])) + .expression_attribute_values( + ":id", + AttributeValue::Ss(vec![tool_call_id.as_str().to_owned()]), + ) .send() .await .map_err(|e| DynamoError(format!("Failed to add active subscription: {}", e)))?; @@ -208,13 +216,13 @@ impl StateStore for DynamoDbStateStore { async fn remove_pending_user_choice( &self, thread_id: &ThreadId, - choice_id: &str, + choice_id: &rap_protocol::ChoiceId, ) -> Result<(), DynamoError> { let Some(choice) = self .get_pending_user_choices(thread_id) .await? .into_iter() - .find(|choice| choice.id == choice_id) + .find(|choice| choice.id == *choice_id) else { return Ok(()); }; @@ -266,14 +274,17 @@ impl StateStore for DynamoDbStateStore { async fn remove_active_subscription( &self, thread_id: &ThreadId, - tool_call_id: &str, + tool_call_id: &rap_protocol::ToolCallId, ) -> Result<(), DynamoError> { self.client .update_item() .table_name(&self.table_name) .key("session", AttributeValue::S(thread_id.as_str().to_owned())) .update_expression("DELETE active_subscriptions :id") - .expression_attribute_values(":id", AttributeValue::Ss(vec![tool_call_id.to_owned()])) + .expression_attribute_values( + ":id", + AttributeValue::Ss(vec![tool_call_id.as_str().to_owned()]), + ) .send() .await .map_err(|e| DynamoError(format!("Failed to remove active subscription: {}", e)))?; diff --git a/crates/infinity-agent-lambda/src/tools/sleep.rs b/crates/infinity-agent-lambda/src/tools/sleep.rs index 56632874..e936afb1 100644 --- a/crates/infinity-agent-lambda/src/tools/sleep.rs +++ b/crates/infinity-agent-lambda/src/tools/sleep.rs @@ -20,15 +20,15 @@ type BoxError = Box; /// Build the tool-result message a sleep delivers when it wakes. fn wakeup_message( - id: &str, - call_id: Option, + id: &rap_protocol::ToolCallId, + call_id: Option, text: String, group_id: &ThreadId, ) -> InputMessage { InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: id.to_owned(), - call_id, + id: id.as_str().to_owned(), + call_id: call_id.map(|c| c.into_inner()), content: vec![ToolResultContent::Text(Text { text })], })), group_id: group_id.to_owned(), @@ -150,8 +150,8 @@ impl Tool for SleepTool { async fn execute( &self, args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), BoxError> { let seconds = args["seconds"].as_f64().unwrap_or(0.0) as i64; @@ -163,11 +163,14 @@ impl Tool for SleepTool { ); if seconds <= 0 { - context.message_sender.send_to_input_queue(msg, &id).await?; + context + .message_sender + .send_to_input_queue(msg, id.as_str()) + .await?; } else { let target = Utc::now() + Duration::seconds(seconds); self.scheduler - .deliver_at(context, &msg, seconds, target, "sleep", &id) + .deliver_at(context, &msg, seconds, target, "sleep", id.as_str()) .await?; } @@ -215,8 +218,8 @@ impl Tool for SleepUntilTool { async fn execute( &self, args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), BoxError> { let date_str = args["date"].as_str().unwrap_or(""); @@ -252,7 +255,10 @@ impl Tool for SleepUntilTool { let msg = wakeup_message(&id, call_id, text, &context.group_id); if target_utc <= now { - context.message_sender.send_to_input_queue(msg, &id).await?; + context + .message_sender + .send_to_input_queue(msg, id.as_str()) + .await?; return Ok(()); } diff --git a/crates/infinity-daemon/src/memory_store.rs b/crates/infinity-daemon/src/memory_store.rs index 819fc73e..968fbb6f 100644 --- a/crates/infinity-daemon/src/memory_store.rs +++ b/crates/infinity-daemon/src/memory_store.rs @@ -809,7 +809,7 @@ impl ConversationStore for PersistentConversationStore { async fn spawn_thread( &self, parent_thread_id: &ThreadId, - spawn_tool_call_id: &str, + spawn_tool_call_id: &rap_protocol::ToolCallId, is_for_subscription_event: bool, spawn_order_override: Option, ) -> Result { @@ -879,7 +879,7 @@ impl ConversationStore for PersistentConversationStore { async fn get_thread_parent_info( &self, thread_id: &ThreadId, - ) -> Result, MemoryError> { + ) -> Result, MemoryError> { self.ensure_thread_metadata_loaded(thread_id); Ok(self.core.get_thread_parent_info(thread_id).await?) } @@ -1006,13 +1006,17 @@ impl PersistentStateStore { .any(|thread_id| self.has_pending_choices(thread_id)) } - pub fn pending_choice(&self, thread_id: &ThreadId, choice_id: &str) -> Option { + pub fn pending_choice( + &self, + thread_id: &ThreadId, + choice_id: &rap_protocol::ChoiceId, + ) -> Option { self.ensure_loaded(thread_id); self.core .thread_state(thread_id) .pending_user_choices .into_iter() - .find(|choice| choice.id == choice_id) + .find(|choice| choice.id == *choice_id) } pub async fn clear_pending_choices( @@ -1114,7 +1118,7 @@ impl StateStore for PersistentStateStore { async fn get_active_subscriptions( &self, thread_id: &ThreadId, - ) -> Result, MemoryError> { + ) -> Result, MemoryError> { self.ensure_loaded(thread_id); Ok(self.core.get_active_subscriptions(thread_id).await?) } @@ -1122,7 +1126,7 @@ impl StateStore for PersistentStateStore { async fn add_active_subscription( &self, thread_id: &ThreadId, - tool_call_id: &str, + tool_call_id: &rap_protocol::ToolCallId, ) -> Result<(), MemoryError> { self.ensure_loaded(thread_id); self.core @@ -1135,7 +1139,7 @@ impl StateStore for PersistentStateStore { async fn remove_active_subscription( &self, thread_id: &ThreadId, - tool_call_id: &str, + tool_call_id: &rap_protocol::ToolCallId, ) -> Result<(), MemoryError> { self.ensure_loaded(thread_id); self.core @@ -1160,7 +1164,7 @@ impl StateStore for PersistentStateStore { async fn remove_pending_user_choice( &self, thread_id: &ThreadId, - choice_id: &str, + choice_id: &rap_protocol::ChoiceId, ) -> Result<(), MemoryError> { self.ensure_loaded(thread_id); self.core @@ -1264,7 +1268,12 @@ mod tests { .expect("append root messages"); let child = store - .spawn_thread(ThreadId::from_ref("root"), "tc-1", false, None) + .spawn_thread( + ThreadId::from_ref("root"), + rap_protocol::ToolCallId::from_ref("tc-1"), + false, + None, + ) .await .expect("spawn child thread"); @@ -1322,7 +1331,12 @@ mod tests { .expect("append root messages"); let child = store - .spawn_thread(ThreadId::from_ref("root"), "tc-1", false, None) + .spawn_thread( + ThreadId::from_ref("root"), + rap_protocol::ToolCallId::from_ref("tc-1"), + false, + None, + ) .await .expect("spawn child thread"); store @@ -1334,7 +1348,12 @@ mod tests { .expect("append child messages"); let grandchild = store - .spawn_thread(&child, "tc-2", false, None) + .spawn_thread( + &child, + rap_protocol::ToolCallId::from_ref("tc-2"), + false, + None, + ) .await .expect("spawn grandchild thread"); store @@ -1425,7 +1444,12 @@ mod tests { .expect("save compaction summary"); let child = store - .spawn_thread(ThreadId::from_ref("root"), "tc-1", false, None) + .spawn_thread( + ThreadId::from_ref("root"), + rap_protocol::ToolCallId::from_ref("tc-1"), + false, + None, + ) .await .expect("spawn child thread"); store @@ -1483,7 +1507,12 @@ mod tests { .expect("save later compaction summary"); let child = store - .spawn_thread(ThreadId::from_ref("root"), "tc-1", false, None) + .spawn_thread( + ThreadId::from_ref("root"), + rap_protocol::ToolCallId::from_ref("tc-1"), + false, + None, + ) .await .expect("spawn child thread"); store @@ -1531,7 +1560,12 @@ mod tests { .expect("save root compaction summary"); let child = store - .spawn_thread(ThreadId::from_ref("root"), "tc-1", false, None) + .spawn_thread( + ThreadId::from_ref("root"), + rap_protocol::ToolCallId::from_ref("tc-1"), + false, + None, + ) .await .expect("spawn child thread"); store diff --git a/crates/infinity-daemon/src/rap_servers.rs b/crates/infinity-daemon/src/rap_servers.rs index e269b37d..85312871 100644 --- a/crates/infinity-daemon/src/rap_servers.rs +++ b/crates/infinity-daemon/src/rap_servers.rs @@ -394,8 +394,8 @@ impl Tool for ManagedRapTool { async fn execute( &self, args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &infinity_agent_core::tools::ToolContext, ) -> Result<(), BoxError> { let endpoint = self.server.invoke_endpoint().await?; diff --git a/crates/infinity-daemon/src/session/display.rs b/crates/infinity-daemon/src/session/display.rs index ac9d6f06..f6ccc05d 100644 --- a/crates/infinity-daemon/src/session/display.rs +++ b/crates/infinity-daemon/src/session/display.rs @@ -99,7 +99,7 @@ pub(crate) fn history_message_to_daemon( .iter() .find_map(|m| { if let InfinityMessage::ToolCall { call, .. } = m - && call.id == *tool_call_id + && call.id == tool_call_id.as_str() { Some(format!( "{}({})", @@ -109,7 +109,7 @@ pub(crate) fn history_message_to_daemon( None } }) - .unwrap_or_else(|| tool_call_id.clone()) + .unwrap_or_else(|| tool_call_id.as_str().to_owned()) }; Some(DaemonMessage::SubscriptionEvent { name, diff --git a/crates/infinity-daemon/src/session/tests.rs b/crates/infinity-daemon/src/session/tests.rs index e6dc686a..56a8a2d8 100644 --- a/crates/infinity-daemon/src/session/tests.rs +++ b/crates/infinity-daemon/src/session/tests.rs @@ -199,8 +199,8 @@ impl Tool for AsyncStubTool { async fn execute( &self, _: serde_json::Value, - _: String, - _: Option, + _: rap_protocol::ToolCallId, + _: Option, _: &infinity_agent_core::tools::ToolContext, ) -> Result<(), Box> { Ok(()) @@ -500,7 +500,7 @@ async fn answered_user_choice_emits_complete_and_disappears_from_replay() { InputMessage { content: InputMessageContent::UserChoice(UserChoiceRequired { content_type: "user_choice_required".to_owned(), - id: "tc-choice".to_owned(), + id: "tc-choice".into(), call_id: None, prompt: "Pick one".to_owned(), choices: vec!["A".to_owned(), "B".to_owned()], @@ -518,7 +518,7 @@ async fn answered_user_choice_emits_complete_and_disappears_from_replay() { .await; assert!(matches!( display_rx.recv().await, - Some(DaemonMessage::UserChoiceRequired { id, .. }) if id == "tc-choice" + Some(DaemonMessage::UserChoiceRequired { id, .. }) if id.as_str() == "tc-choice" )); running @@ -535,7 +535,7 @@ async fn answered_user_choice_emits_complete_and_disappears_from_replay() { .expect("display channel closed") { DaemonMessage::UserChoiceComplete { choice_id } => { - assert_eq!(choice_id, "tc-choice"); + assert_eq!(choice_id.as_str(), "tc-choice"); saw_complete = true; } DaemonMessage::StartOutput { .. } => { @@ -618,11 +618,16 @@ async fn child_pending_choice_updates_root_session_status() { .expect("create session"); let child_id = manager .conversation_store - .spawn_thread(&session_id, "spawn-call", false, None) + .spawn_thread( + &session_id, + rap_protocol::ToolCallId::from_ref("spawn-call"), + false, + None, + ) .await .expect("spawn child thread"); let choice = UserChoice { - id: "choice-1".to_owned(), + id: "choice-1".into(), prompt: "Choose".to_owned(), choices: vec!["one".to_owned(), "two".to_owned()], default: 0, diff --git a/crates/infinity-daemon/src/set_title_tool.rs b/crates/infinity-daemon/src/set_title_tool.rs index e1a68384..63420b99 100644 --- a/crates/infinity-daemon/src/set_title_tool.rs +++ b/crates/infinity-daemon/src/set_title_tool.rs @@ -40,8 +40,8 @@ impl Tool for SetTitleTool { async fn execute( &self, args: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), Box> { let title = args["title"].as_str().unwrap_or("").to_owned(); @@ -51,8 +51,8 @@ impl Tool for SetTitleTool { let msg = InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: id.clone(), - call_id, + id: id.clone().into_inner(), + call_id: call_id.map(|c| c.into_inner()), content: vec![ToolResultContent::Text(Text { text: format!("Title set to: {}", title), })], @@ -64,7 +64,10 @@ impl Tool for SetTitleTool { subscription: false, }; - context.message_sender.send_to_input_queue(msg, &id).await?; + context + .message_sender + .send_to_input_queue(msg, id.as_str()) + .await?; Ok(()) } diff --git a/crates/infinity-daemon/tests/mcp_proxy_callback.rs b/crates/infinity-daemon/tests/mcp_proxy_callback.rs index 76f15780..5693d0f4 100644 --- a/crates/infinity-daemon/tests/mcp_proxy_callback.rs +++ b/crates/infinity-daemon/tests/mcp_proxy_callback.rs @@ -48,7 +48,7 @@ async fn list_tools_callback_deserializes() { let inv = RapInvocation { operation: "mock_list_tools".to_owned(), arguments: serde_json::json!({}), - id: "test-1".to_owned(), + id: "test-1".into(), call_id: None, callback_url, group_id: "g1".into(), diff --git a/crates/infinity-mcp-bridge/src/lib.rs b/crates/infinity-mcp-bridge/src/lib.rs index aee12290..b13ad642 100644 --- a/crates/infinity-mcp-bridge/src/lib.rs +++ b/crates/infinity-mcp-bridge/src/lib.rs @@ -650,8 +650,8 @@ impl Tool for McpTool { async fn execute( &self, arguments: serde_json::Value, - id: String, - call_id: Option, + id: rap_protocol::ToolCallId, + call_id: Option, context: &ToolContext, ) -> Result<(), BoxError> { let client = self.client.clone(); @@ -662,8 +662,8 @@ impl Tool for McpTool { let (text, display_as) = client.dispatch(&tool_name, &arguments).await; let message = InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: id.clone(), - call_id, + id: id.clone().into_inner(), + call_id: call_id.map(|c| c.into_inner()), content: vec![ToolResultContent::Text(Text { text })], })), group_id: group_id.clone(), @@ -672,7 +672,7 @@ impl Tool for McpTool { display_as, subscription: false, }; - if let Err(error) = sender.send_to_input_queue(message, &id).await { + if let Err(error) = sender.send_to_input_queue(message, id.as_str()).await { tracing::warn!(%error, "failed to deliver MCP tool result"); } }); diff --git a/crates/infinity-protocol/src/lib.rs b/crates/infinity-protocol/src/lib.rs index d2d665fa..f23fee74 100644 --- a/crates/infinity-protocol/src/lib.rs +++ b/crates/infinity-protocol/src/lib.rs @@ -3,9 +3,9 @@ use std::collections::HashMap; use std::path::PathBuf; use tokio_util::codec::LengthDelimitedCodec; -/// The thread identifier used throughout the runtime (re-exported from +/// The identifier kinds used throughout the runtime (re-exported from /// `rap-protocol`). -pub use rap_protocol::ThreadId; +pub use rap_protocol::{ChoiceId, ProviderCallId, ThreadId, ToolCallId}; strkind::strkind! { /// The name of a configured remote daemon (from `remotes.json`), e.g. @@ -288,7 +288,7 @@ pub enum ClientMessage { /// Notify the daemon that a user choice was answered so it can be /// removed from the pending replay list. UserChoiceAnswered { - choice_id: String, + choice_id: ChoiceId, selected: usize, }, /// Trigger compaction for the given session. @@ -388,13 +388,13 @@ pub enum DaemonMessage { }, UserChoiceRequired { thread_id: Option, - id: String, + id: ChoiceId, prompt: String, choices: Vec, default: usize, }, UserChoiceComplete { - choice_id: String, + choice_id: ChoiceId, }, ThinkingStart { thread_id: Option, diff --git a/crates/infinity-rap-bridge/src/callback.rs b/crates/infinity-rap-bridge/src/callback.rs index efe028cc..c71e7420 100644 --- a/crates/infinity-rap-bridge/src/callback.rs +++ b/crates/infinity-rap-bridge/src/callback.rs @@ -29,8 +29,8 @@ pub(crate) fn convert_callback(cb: RapCallback) -> Option { Some(match cb { RapCallback::ToolResult(tr) => InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: tr.id, - call_id: tr.call_id, + id: tr.id.into_inner(), + call_id: tr.call_id.map(|c| c.into_inner()), content: tool_result_content(tr.content, tr.text), })), group_id: tr.group_id, @@ -43,7 +43,7 @@ pub(crate) fn convert_callback(cb: RapCallback) -> Option { let is_final = se.r#final.unwrap_or(false); InputMessage { content: InputMessageContent::User(UserContent::ToolResult(ToolResult { - id: se.tool_call_id.clone(), + id: se.tool_call_id.clone().into_inner(), call_id: None, content: vec![ToolResultContent::Text(Text { text: se.text })], })), @@ -133,8 +133,8 @@ mod tests { ) -> RapCallback { RapCallback::ToolResult(RapToolResult { group_id: "t1".into(), - id: "call-1".to_owned(), - call_id: Some("prov-1".to_owned()), + id: "call-1".into(), + call_id: Some("prov-1".into()), text, content, display_as: None, @@ -166,7 +166,7 @@ mod tests { fn display_as_and_subscription_flag_pass_through() { let cb = RapCallback::ToolResult(RapToolResult { group_id: "t1".into(), - id: "call-1".to_owned(), + id: "call-1".into(), call_id: None, text: Some("done".to_owned()), content: None, @@ -184,7 +184,7 @@ mod tests { fn subscription_event_converts_with_flags() { let cb = RapCallback::SubscriptionEvent(RapSubscriptionEvent { group_id: "t1".into(), - tool_call_id: "sub-1".to_owned(), + tool_call_id: "sub-1".into(), text: "tick".to_owned(), associative: true, r#final: Some(true), @@ -196,7 +196,7 @@ mod tests { associative, r#final, })) => { - assert_eq!(tool_call_id, "sub-1"); + assert_eq!(tool_call_id.as_str(), "sub-1"); assert!(associative); assert!(r#final); } @@ -208,7 +208,7 @@ mod tests { fn oauth_converts() { let cb = RapCallback::OAuth(RapOAuth { group_id: "t1".into(), - id: "call-1".to_owned(), + id: "call-1".into(), call_id: None, auth_url: "https://auth".to_owned(), }); @@ -226,7 +226,7 @@ mod tests { fn user_choice_converts() { let cb = RapCallback::UserChoice(RapUserChoice { group_id: "t1".into(), - id: "choice-1".to_owned(), + id: "choice-1".into(), call_id: None, prompt: "pick one".to_owned(), choices: vec!["a".to_owned(), "b".to_owned()], diff --git a/crates/infinity-slack-bot/src/daemon_client.rs b/crates/infinity-slack-bot/src/daemon_client.rs index 8b5dc456..af1336ae 100644 --- a/crates/infinity-slack-bot/src/daemon_client.rs +++ b/crates/infinity-slack-bot/src/daemon_client.rs @@ -106,7 +106,11 @@ impl DaemonClient { Ok(()) } - pub async fn answer_choice(&self, choice_id: &str, selected: usize) -> Result<(), BoxError> { + pub async fn answer_choice( + &self, + choice_id: &infinity_protocol::ChoiceId, + selected: usize, + ) -> Result<(), BoxError> { self.tx .send(ClientMessage::UserChoiceAnswered { choice_id: choice_id.to_owned(), diff --git a/crates/infinity-slack-bot/src/daemon_sidecar.rs b/crates/infinity-slack-bot/src/daemon_sidecar.rs index c45fdff9..cf7eba43 100644 --- a/crates/infinity-slack-bot/src/daemon_sidecar.rs +++ b/crates/infinity-slack-bot/src/daemon_sidecar.rs @@ -36,7 +36,7 @@ pub enum DaemonCommand { /// Answer a choice prompt on the connection for this thread. AnswerChoice { thread_ts: String, - choice_id: String, + choice_id: infinity_protocol::ChoiceId, selected: usize, }, } diff --git a/crates/infinity-slack-bot/src/flow.rs b/crates/infinity-slack-bot/src/flow.rs index 93feebdc..39d57487 100644 --- a/crates/infinity-slack-bot/src/flow.rs +++ b/crates/infinity-slack-bot/src/flow.rs @@ -91,11 +91,11 @@ pub fn slack_dataflow<'a, P: 'a>( rt.choice_messages .lock() .expect("bug: lock poisoned") - .remove(&choice_id); + .remove(choice_id.as_str()); Some(crate::daemon_sidecar::DaemonCommand::AnswerChoice { thread_ts: event.thread_ts, - choice_id, + choice_id: choice_id.into(), selected, }) })); diff --git a/crates/infinity-slack-bot/src/session_store.rs b/crates/infinity-slack-bot/src/session_store.rs index 180d0e09..817c4de4 100644 --- a/crates/infinity-slack-bot/src/session_store.rs +++ b/crates/infinity-slack-bot/src/session_store.rs @@ -8,7 +8,7 @@ use crate::BoxError; /// A pending approval/choice waiting for user response. #[derive(Clone)] pub struct PendingChoice { - pub choice_id: String, + pub choice_id: infinity_protocol::ChoiceId, pub choices: Vec, } diff --git a/crates/infinity-slack-bot/src/sidecar.rs b/crates/infinity-slack-bot/src/sidecar.rs index 5090465c..cc3ca95d 100644 --- a/crates/infinity-slack-bot/src/sidecar.rs +++ b/crates/infinity-slack-bot/src/sidecar.rs @@ -71,12 +71,14 @@ pub enum SlackAction { /// If set, the sidecar will store the resulting message_ts under this /// choice_id so that a later `DismissChoiceButtons` can update it. #[serde(default, skip_serializing_if = "Option::is_none")] - choice_id: Option, + choice_id: Option, }, /// Dismiss interactive buttons for a completed choice (replace with a /// "resolved" indicator). The sidecar looks up the stored message_ts /// from the choice_id. - DismissChoiceButtons { choice_id: String }, + DismissChoiceButtons { + choice_id: infinity_protocol::ChoiceId, + }, /// Update an existing message's blocks (e.g. to replace buttons with a selection indicator). UpdateMessage { channel: String, @@ -546,7 +548,7 @@ pub fn create() -> (ReceiverStream, PollSender) { rt.choice_messages .lock() .expect("bug: lock poisoned") - .insert(cid, (channel, msg_ts)); + .insert(cid.into_inner(), (channel, msg_ts)); } } Ok(None) => { @@ -754,7 +756,7 @@ pub fn create() -> (ReceiverStream, PollSender) { rt.choice_messages .lock() .expect("bug: lock poisoned") - .remove(&choice_id) + .remove(choice_id.as_str()) }; if let Some((channel, msg_ts)) = info { let blocks = serde_json::json!([ diff --git a/crates/rap-client/src/notifier.rs b/crates/rap-client/src/notifier.rs index 4594294b..88d7d8a7 100644 --- a/crates/rap-client/src/notifier.rs +++ b/crates/rap-client/src/notifier.rs @@ -39,7 +39,7 @@ impl RapNotifier { pub async fn notify_tool_cancelled( &self, thread_id: &rap_protocol::ThreadId, - tool_call_id: &str, + tool_call_id: &rap_protocol::ToolCallId, ) { let payload = serde_json::json!({ "thread_id": thread_id, diff --git a/crates/rap-github-event-poller/src/lib.rs b/crates/rap-github-event-poller/src/lib.rs index 65027921..a25e0cce 100644 --- a/crates/rap-github-event-poller/src/lib.rs +++ b/crates/rap-github-event-poller/src/lib.rs @@ -24,8 +24,8 @@ pub struct SubscribeArgs { #[derive(Debug, Clone)] pub struct Subscription { - pub tool_call_id: String, - pub call_id: Option, + pub tool_call_id: rap_protocol::ToolCallId, + pub call_id: Option, pub callback_url: String, pub group_id: rap_protocol::ThreadId, pub filters: Filters, @@ -61,7 +61,7 @@ impl Filters { #[derive(Debug)] struct RepoState { /// Subscriptions keyed by tool_call_id - subscriptions: HashMap, + subscriptions: HashMap, /// ETag from last poll (for conditional requests) etag: Option, /// Poll interval from GitHub's X-Poll-Interval header (default 60s) @@ -167,7 +167,7 @@ impl Poller { } /// Remove a subscription by tool_call_id. - pub async fn cancel(&self, tool_call_id: &str) { + pub async fn cancel(&self, tool_call_id: &rap_protocol::ToolCallId) { let mut repos = self.repos.write().await; for state in repos.values_mut() { state.subscriptions.remove(tool_call_id); diff --git a/crates/rap-github-event-poller/src/main.rs b/crates/rap-github-event-poller/src/main.rs index 037d1c9a..7a04ef1f 100644 --- a/crates/rap-github-event-poller/src/main.rs +++ b/crates/rap-github-event-poller/src/main.rs @@ -122,7 +122,9 @@ async fn cancel_handler( State(state): State>, Json(req): Json, ) -> StatusCode { - state.cancel(&req.tool_call_id).await; + state + .cancel(rap_protocol::ToolCallId::from_ref(&req.tool_call_id)) + .await; StatusCode::OK } diff --git a/crates/rap-protocol/src/lib.rs b/crates/rap-protocol/src/lib.rs index 109f5fb5..8bfa5d73 100644 --- a/crates/rap-protocol/src/lib.rs +++ b/crates/rap-protocol/src/lib.rs @@ -14,8 +14,31 @@ strkind::strkind! { /// root thread ID is the caller-chosen conversation key (which doubles as /// the SQS FIFO `MessageGroupId`). Treat the contents as opaque. pub ThreadId; + + /// Identifies one tool call within a thread. + /// + /// This is the `id` field on RAP invocations and callbacks. It + /// originates from the model provider's tool-call stream (e.g. Bedrock + /// `tooluse_...` IDs) or, for runtime-synthesized calls (thread spawns, + /// subscription events), from a minted UUID. Treat the contents as + /// opaque; formats are provider-controlled. + pub ToolCallId; + + /// The provider-scoped companion ID of a tool call (`call_id` on RAP + /// invocations and callbacks). Some providers issue a second identifier + /// alongside the tool-call ID; it must be echoed back verbatim in tool + /// results. Absent for providers that do not use one. + pub ProviderCallId; } +/// Identifies a pending user choice. +/// +/// Currently an alias of [`ToolCallId`]: a choice is keyed by the tool call +/// that requested it (`RapUserChoice.id` is copied from the invocation). +/// Tracked as a candidate for becoming a distinct kind in +/// . +pub type ChoiceId = ToolCallId; + // ── RAP protocol types ── #[derive(Debug, Serialize, Deserialize)] @@ -23,8 +46,8 @@ pub struct RapInvocation { pub operation: String, #[serde(default)] pub arguments: serde_json::Value, - pub id: String, - pub call_id: Option, + pub id: ToolCallId, + pub call_id: Option, pub callback_url: String, pub group_id: ThreadId, #[serde(skip_serializing_if = "Option::is_none")] @@ -106,9 +129,9 @@ pub enum RapToolResultContent { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RapToolResult { pub group_id: ThreadId, - pub id: String, + pub id: ToolCallId, #[serde(default, skip_serializing_if = "Option::is_none")] - pub call_id: Option, + pub call_id: Option, /// Plain-text result — the shorthand for a text-only result. A tool MUST /// provide either `text` or `content`. When `content` is also present it /// supersedes `text` for the model-facing result. @@ -129,9 +152,9 @@ pub struct RapToolResult { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RapUserChoice { pub group_id: ThreadId, - pub id: String, + pub id: ChoiceId, #[serde(default, skip_serializing_if = "Option::is_none")] - pub call_id: Option, + pub call_id: Option, pub prompt: String, pub choices: Vec, #[serde(default)] @@ -142,7 +165,7 @@ pub struct RapUserChoice { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RapSubscriptionEvent { pub group_id: ThreadId, - pub tool_call_id: String, + pub tool_call_id: ToolCallId, pub text: String, #[serde(default)] pub associative: bool, @@ -164,9 +187,9 @@ pub struct RapViewUpdate { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RapOAuth { pub group_id: ThreadId, - pub id: String, + pub id: ToolCallId, #[serde(default, skip_serializing_if = "Option::is_none")] - pub call_id: Option, + pub call_id: Option, pub auth_url: String, } @@ -315,7 +338,7 @@ pub async fn send_subscription_event( client: &C, callback_url: &str, group_id: ThreadId, - tool_call_id: String, + tool_call_id: ToolCallId, text: &str, associative: bool, r#final: bool, diff --git a/crates/sandbox-core/src/server.rs b/crates/sandbox-core/src/server.rs index 2b3144c8..f57527d2 100644 --- a/crates/sandbox-core/src/server.rs +++ b/crates/sandbox-core/src/server.rs @@ -76,7 +76,7 @@ type PendingTasks = Arc>>>; /// is spawned) and the receiver is passed into the command handler. The /// `cancel_tool_call_handler` removes the sender and sends `()` to signal /// cancellation; the handler receives it and sends SIGTERM to the process. -type InFlightMap = Arc>>>; +type InFlightMap = Arc>>>; /// Send SIGTERM to a process group by PID. /// @@ -101,7 +101,8 @@ struct AppState { in_flight: InFlightMap, /// Pending user choice responses, keyed by tool call ID. /// The sender delivers the user's selected index. - pending_choices: Arc>>>, + pending_choices: + Arc>>>, /// Server base URL, set from the first request's Host header. server_base_url: std::sync::OnceLock, /// Whether this server advertises migration support. @@ -393,7 +394,11 @@ async fn cancel_tool_call_handler< "received cancel_tool_call notification" ); - let sender = state.in_flight.lock().await.remove(&request.tool_call_id); + let sender = state + .in_flight + .lock() + .await + .remove(rap_protocol::ToolCallId::from_ref(&request.tool_call_id)); if let Some(sender) = sender { // Signal the command handler to SIGTERM the process and clean up. @@ -434,7 +439,11 @@ async fn user_choice_response_handler< "received user_choice_response" ); - let sender = state.pending_choices.lock().await.remove(&response.id); + let sender = state + .pending_choices + .lock() + .await + .remove(rap_protocol::ToolCallId::from_ref(&response.id)); if let Some(sender) = sender { let _ = sender.send(response.selected); diff --git a/crates/sandbox-local/tests/common.rs b/crates/sandbox-local/tests/common.rs index 20635f73..322f4292 100644 --- a/crates/sandbox-local/tests/common.rs +++ b/crates/sandbox-local/tests/common.rs @@ -88,7 +88,8 @@ pub async fn invoke_raw( id: format!( "call-{operation}-{}", CALL_COUNTER.fetch_add(1, Ordering::Relaxed) - ), + ) + .into(), call_id: None, callback_url: callback_url.to_owned(), group_id: group_id.into(), @@ -173,7 +174,8 @@ pub async fn invoke_collecting_views( id: format!( "call-{operation}-{}", CALL_COUNTER.fetch_add(1, Ordering::Relaxed) - ), + ) + .into(), call_id: None, callback_url: callback_url.to_owned(), group_id: group_id.into(),