diff --git a/docs/sdk/rust/client/agent-trait.mdx b/docs/sdk/rust/client/agent-trait.mdx index a93ed9a30f..a89923d066 100644 --- a/docs/sdk/rust/client/agent-trait.mdx +++ b/docs/sdk/rust/client/agent-trait.mdx @@ -80,6 +80,8 @@ let params_typed = RunAgentParams::::new_typed() Builder helpers include: - with_run_id(run_id) +- with_parent_run_id(parent_run_id) +- with_resume(resume) - add_tool(tool) - add_context(ctx) - with_forwarded_props(props) diff --git a/docs/sdk/rust/client/subscriber.mdx b/docs/sdk/rust/client/subscriber.mdx index ef12cdb10d..a8b67defbf 100644 --- a/docs/sdk/rust/client/subscriber.mdx +++ b/docs/sdk/rust/client/subscriber.mdx @@ -44,7 +44,7 @@ impl AgentSubscriber for Logger { _buffer: &str, _params: ag_ui_client::subscriber::AgentSubscriberParams<'async_trait, serde_json::Value, serde_json::Value>, ) -> Result { - println!("chunk: {}", event.content); + println!("chunk: {}", event.delta); Ok(Default::default()) } } @@ -101,6 +101,9 @@ You can implement specific typed callbacks or the catch‑all on_event. Key call - on_tool_call_start_event / on_tool_call_args_event / on_tool_call_end_event / on_tool_call_result_event - on_state_snapshot_event / on_state_delta_event - on_messages_snapshot_event +- on_activity_snapshot_event / on_activity_delta_event +- on_reasoning_start_event / on_reasoning_message_start_event / on_reasoning_message_content_event / on_reasoning_message_end_event +- on_reasoning_message_chunk_event / on_reasoning_end_event / on_reasoning_encrypted_value_event - on_raw_event / on_custom_event All callbacks have sensible defaults that return Ok(Default::default()). Only implement what you need. diff --git a/docs/sdk/rust/core/events.mdx b/docs/sdk/rust/core/events.mdx index 3dfc80575b..3384210b88 100644 --- a/docs/sdk/rust/core/events.mdx +++ b/docs/sdk/rust/core/events.mdx @@ -25,28 +25,45 @@ pub enum Event { TextMessageEnd(TextMessageEndEvent), TextMessageChunk(TextMessageChunkEvent), ThinkingTextMessageStart(ThinkingTextMessageStartEvent), + ThinkingTextMessageContent(ThinkingTextMessageContentEvent), + ThinkingTextMessageEnd(ThinkingTextMessageEndEvent), ToolCallStart(ToolCallStartEvent), ToolCallArgs(ToolCallArgsEvent), ToolCallEnd(ToolCallEndEvent), + ToolCallChunk(ToolCallChunkEvent), ToolCallResult(ToolCallResultEvent), StateSnapshot(StateSnapshotEvent), StateDelta(StateDeltaEvent), MessagesSnapshot(MessagesSnapshotEvent), + ActivitySnapshot(ActivitySnapshotEvent), + ActivityDelta(ActivityDeltaEvent), Raw(RawEvent), Custom(CustomEvent), + ReasoningStart(ReasoningStartEvent), + ReasoningMessageStart(ReasoningMessageStartEvent), + ReasoningMessageContent(ReasoningMessageContentEvent), + ReasoningMessageEnd(ReasoningMessageEndEvent), + ReasoningMessageChunk(ReasoningMessageChunkEvent), + ReasoningEnd(ReasoningEndEvent), + ReasoningEncryptedValue(ReasoningEncryptedValueEvent), } ``` ## Frequently used events -- RunStartedEvent – signals the beginning of a run -- RunFinishedEvent – contains an optional result payload +- RunStartedEvent – signals the beginning of a run and can include parentRunId/input +- RunFinishedEvent – contains an optional result payload or interrupt-aware outcome - RunErrorEvent – reports an error with message/context - TextMessageStart/Content/End – streaming assistant message text +- TextMessageChunk – convenience chunk form matching the proto schema - ToolCall* – tool call lifecycle and argument streaming - StateSnapshotEvent – full state replacement - StateDeltaEvent – JSON-Patch style deltas (Vec) - MessagesSnapshotEvent – full replacement of the message history +- ActivitySnapshotEvent/ActivityDeltaEvent – structured in-progress activity updates +- Reasoning* – reasoning block/message events and encrypted reasoning values + +`Thinking*` events are still accepted for compatibility, but the protocol docs now mark them as deprecated in favor of the `Reasoning*` events. ## Handling in subscribers diff --git a/docs/sdk/rust/core/types.mdx b/docs/sdk/rust/core/types.mdx index 693b460b25..f42083d485 100644 --- a/docs/sdk/rust/core/types.mdx +++ b/docs/sdk/rust/core/types.mdx @@ -28,18 +28,19 @@ let tcid = ToolCallId::random(); ## Messages -Message is an enum covering different roles. Helpers exist to create messages. +Message is an enum covering protocol roles: developer, system, assistant, user, tool, activity, and reasoning. Helpers exist to create common message types. ```rust -use ag_ui_client::core::types::{Message, MessageId}; +use ag_ui_client::core::types::Message; let user = Message::new_user("Hello"); let system = Message::new_system("You are a helpful assistant"); let dev = Message::new_developer("internal instruction"); -let tool = Message::new_tool("result", ag_ui_client::core::types::ToolCallId::random()); +let tool = Message::new_tool("result"); ``` Each variant also has a builder form; see crate docs for full fields (name, error, tool call reference, etc.). +Activity messages carry structured JSON content for in-progress UI state. Reasoning messages carry reasoning content and optional encryptedValue. ## Tools and context @@ -60,6 +61,8 @@ let tool = Tool::new( let ctx = Context::new("trace_id".into(), "abc-123".into()); ``` +Tools can also carry optional metadata, and tool calls can carry encryptedValue for reasoning continuity workflows. + ## RunAgentInput The payload sent to an agent, typically over HTTP by HttpAgent. Constructed internally by Agent::run_agent but also available directly. @@ -79,7 +82,7 @@ let input = RunAgentInput::new( ); ``` -Type parameters allow strongly-typed state and forwarded_props when desired. +Type parameters allow strongly-typed state and forwarded_props when desired. Use `with_parent_run_id` to record run lineage and `with_resume` to resume interrupts with `ResumeEntry` values. ## Traits: AgentState and FwdProps diff --git a/sdks/community/rust/crates/ag-ui-client/src/agent.rs b/sdks/community/rust/crates/ag-ui-client/src/agent.rs index 98ba3086b9..3195a55d57 100644 --- a/sdks/community/rust/crates/ag-ui-client/src/agent.rs +++ b/sdks/community/rust/crates/ag-ui-client/src/agent.rs @@ -3,7 +3,7 @@ use std::collections::HashSet; use crate::core::JsonValue; use crate::core::types::{ - AgentId, Context, Message, MessageId, RunAgentInput, RunId, ThreadId, Tool, + AgentId, Context, Message, MessageId, ResumeEntry, RunAgentInput, RunId, ThreadId, Tool, }; use crate::core::{AgentState, FwdProps}; use crate::event_handler::EventHandler; @@ -41,11 +41,13 @@ where #[derive(Debug, Clone, Default)] pub struct RunAgentParams { pub run_id: Option, + pub parent_run_id: Option, pub tools: Vec, pub context: Vec, pub forwarded_props: FwdPropsT, pub messages: Vec, pub state: StateT, + pub resume: Option>, } impl RunAgentParams @@ -60,11 +62,13 @@ where pub fn new_typed() -> Self { Self { run_id: None, + parent_run_id: None, tools: Vec::new(), context: Vec::new(), forwarded_props: FwdPropsT::default(), messages: Vec::new(), state: StateT::default(), + resume: None, } } @@ -72,6 +76,14 @@ where self.run_id = Some(run_id); self } + pub fn with_parent_run_id(mut self, parent_run_id: RunId) -> Self { + self.parent_run_id = Some(parent_run_id); + self + } + pub fn with_resume(mut self, resume: Vec) -> Self { + self.resume = Some(resume); + self + } pub fn add_tool(mut self, tool: Tool) -> Self { self.tools.push(tool); self @@ -97,6 +109,7 @@ where id: MessageId::random(), content: content.into(), name: None, + encrypted_value: None, }); self } @@ -192,11 +205,13 @@ where let input = RunAgentInput { thread_id: ThreadId::random(), run_id: params.run_id.clone().unwrap_or_else(RunId::random), + parent_run_id: params.parent_run_id.clone(), state: params.state.clone(), messages: params.messages.clone(), tools: params.tools.clone(), context: params.context.clone(), forwarded_props: params.forwarded_props.clone(), + resume: params.resume.clone(), }; let current_message_ids: HashSet<&MessageId> = params.messages.iter().map(|m| m.id()).collect(); diff --git a/sdks/community/rust/crates/ag-ui-client/src/event_handler.rs b/sdks/community/rust/crates/ag-ui-client/src/event_handler.rs index c875abaa72..5aebe3a833 100644 --- a/sdks/community/rust/crates/ag-ui-client/src/event_handler.rs +++ b/sdks/community/rust/crates/ag-ui-client/src/event_handler.rs @@ -98,14 +98,42 @@ where match event { Event::TextMessageStart(e) => { // Default behavior - let new_message = Message::Assistant { - id: e.message_id.clone(), - content: Some(String::new()), - name: None, - tool_calls: None, - }; - self.messages.push(new_message); - current_mutation.messages = Some(self.messages.clone()); + if !self.messages.iter().any(|m| m.id() == &e.message_id) { + let new_message = match e.role { + Role::Developer => Message::Developer { + id: e.message_id.clone(), + content: String::new(), + name: e.name.clone(), + encrypted_value: None, + }, + Role::System => Message::System { + id: e.message_id.clone(), + content: String::new(), + name: e.name.clone(), + encrypted_value: None, + }, + Role::User => Message::User { + id: e.message_id.clone(), + content: String::new(), + name: e.name.clone(), + encrypted_value: None, + }, + Role::Reasoning => Message::Reasoning { + id: e.message_id.clone(), + content: String::new(), + encrypted_value: None, + }, + Role::Assistant | Role::Tool | Role::Activity => Message::Assistant { + id: e.message_id.clone(), + content: Some(String::new()), + name: e.name.clone(), + tool_calls: None, + encrypted_value: None, + }, + }; + self.messages.push(new_message); + current_mutation.messages = Some(self.messages.clone()); + } for subscriber in &self.subscribers { let params = self.to_subscriber_params(); @@ -115,7 +143,11 @@ where } Event::TextMessageContent(e) => { // Default behavior - if let Some(last_message) = self.messages.last_mut() { + if let Some(last_message) = self + .messages + .iter_mut() + .find(|message| message.id() == &e.message_id) + { let content = last_message.content_mut(); if let Some(s) = content { s.push_str(&e.delta) @@ -126,7 +158,8 @@ where // Get the current text message buffer let text_message_buffer = self .messages - .last() + .iter() + .find(|message| message.id() == &e.message_id) .and_then(|m| m.content()) .unwrap_or_default() .to_string(); // Clone to avoid borrowing issues @@ -143,7 +176,8 @@ where // Get the current text message buffer let text_message_buffer = self .messages - .last() + .iter() + .find(|message| message.id() == &e.message_id) .and_then(|m| m.content()) .unwrap_or_default() .to_string(); // Clone to avoid borrowing issues @@ -199,25 +233,37 @@ where name: e.tool_call_name.clone(), arguments: String::new(), }, + encrypted_value: None, }; - if let Some(last_message) = self.messages.last_mut() { - if Some(last_message.id()) == e.parent_message_id.clone().as_ref() { - let _ = last_message.tool_calls_mut().get_or_insert(&mut Vec::new()); + let parent_message_index = e.parent_message_id.as_ref().and_then(|parent_id| { + self.messages + .iter() + .position(|message| message.id() == parent_id) + }); + + let target_assistant_index = parent_message_index.and_then(|index| { + matches!(self.messages.get(index), Some(Message::Assistant { .. })) + .then_some(index) + }); - let _ = last_message - .tool_calls_mut() - .map(|tc| tc.push(new_tool_call)); + if let Some(index) = target_assistant_index { + if let Some(tool_calls) = self.messages[index].tool_calls_mut() { + tool_calls.push(new_tool_call); } } else { let new_message = Message::Assistant { - id: e - .parent_message_id - .clone() - .unwrap_or_else(MessageId::random), + id: if parent_message_index.is_some() { + MessageId::random() + } else { + e.parent_message_id + .clone() + .unwrap_or_else(MessageId::random) + }, content: None, name: None, - tool_calls: None, + tool_calls: Some(vec![new_tool_call]), + encrypted_value: None, }; self.messages.push(new_message); } @@ -231,39 +277,48 @@ where } Event::ToolCallArgs(e) => { // Default behavior - if let Some(last_message) = self.messages.last_mut() - && let Some(tool_calls) = last_message.tool_calls_mut() - && let Some(last_tool_call) = tool_calls.last_mut() + if let Some(tool_call) = self + .messages + .iter_mut() + .filter_map(|message| match message { + Message::Assistant { + tool_calls: Some(tool_calls), + .. + } => Some(tool_calls), + _ => None, + }) + .find_map(|tool_calls| { + tool_calls + .iter_mut() + .find(|tool_call| tool_call.id == e.tool_call_id) + }) { - last_tool_call.function.arguments.push_str(&e.delta); + tool_call.function.arguments.push_str(&e.delta); current_mutation.messages = Some(self.messages.clone()); } // Get the current tool call buffer and name - let (tool_call_buffer, tool_call_name, partial_args) = if let Some(last_message) = - self.messages.last() - { - if let Some(tool_calls) = last_message.tool_calls() { - if let Some(last_tool_call) = tool_calls.last() { - // Try to parse the arguments as JSON to get partial args - let partial_args = serde_json::from_str::>( - &last_tool_call.function.arguments, - ) - .unwrap_or_default(); - ( - last_tool_call.function.arguments.clone(), - last_tool_call.function.name.clone(), - partial_args, - ) - } else { - (String::new(), String::new(), HashMap::new()) - } - } else { - (String::new(), String::new(), HashMap::new()) - } - } else { - (String::new(), String::new(), HashMap::new()) - }; + let (tool_call_buffer, tool_call_name, partial_args) = self + .messages + .iter() + .filter_map(|message| message.tool_calls()) + .find_map(|tool_calls| { + tool_calls + .iter() + .find(|tool_call| tool_call.id == e.tool_call_id) + }) + .map(|tool_call| { + let partial_args = serde_json::from_str::>( + &tool_call.function.arguments, + ) + .unwrap_or_default(); + ( + tool_call.function.arguments.clone(), + tool_call.function.name.clone(), + partial_args, + ) + }) + .unwrap_or_else(|| (String::new(), String::new(), HashMap::new())); for subscriber in &self.subscribers { let params = self.to_subscriber_params(); @@ -281,25 +336,23 @@ where } Event::ToolCallEnd(e) => { // Get the current tool call buffer and name - let (tool_call_name, tool_call_args) = - if let Some(last_message) = self.messages.last() { - if let Some(tool_calls) = last_message.tool_calls() { - if let Some(last_tool_call) = tool_calls.last() { - // Try to parse the arguments as JSON - let args = serde_json::from_str::>( - &last_tool_call.function.arguments, - ) - .unwrap_or_default(); - (last_tool_call.function.name.clone(), args) - } else { - (String::new(), HashMap::new()) - } - } else { - (String::new(), HashMap::new()) - } - } else { - (String::new(), HashMap::new()) - }; + let (tool_call_name, tool_call_args) = self + .messages + .iter() + .filter_map(|message| message.tool_calls()) + .find_map(|tool_calls| { + tool_calls + .iter() + .find(|tool_call| tool_call.id == e.tool_call_id) + }) + .map(|tool_call| { + let args = serde_json::from_str::>( + &tool_call.function.arguments, + ) + .unwrap_or_default(); + (tool_call.function.name.clone(), args) + }) + .unwrap_or_else(|| (String::new(), HashMap::new())); for subscriber in &self.subscribers { let params = self.to_subscriber_params(); @@ -372,12 +425,96 @@ where } } Event::MessagesSnapshot(e) => { + // Default behavior + let snapshot_has_activity = e + .messages + .iter() + .any(|message| matches!(message, Message::Activity { .. })); + let snapshot_has_reasoning = e + .messages + .iter() + .any(|message| matches!(message, Message::Reasoning { .. })); + + self.messages.retain(|message| { + let in_snapshot = e + .messages + .iter() + .any(|snapshot_message| snapshot_message.id() == message.id()); + let preserved_client_only = matches!(message, Message::Activity { .. }) + && !snapshot_has_activity + || matches!(message, Message::Reasoning { .. }) && !snapshot_has_reasoning; + in_snapshot || preserved_client_only + }); + + for snapshot_message in &e.messages { + if let Some(existing_message) = self + .messages + .iter_mut() + .find(|message| message.id() == snapshot_message.id()) + { + *existing_message = snapshot_message.clone(); + } else { + self.messages.push(snapshot_message.clone()); + } + } + + current_mutation.messages = Some(self.messages.clone()); + for subscriber in &self.subscribers { let params = self.to_subscriber_params(); let mutation = subscriber.on_messages_snapshot_event(e, params).await?; mutations.push(mutation); } } + Event::ActivitySnapshot(e) => { + // Default behavior + let new_message = Message::Activity { + id: e.message_id.clone(), + activity_type: e.activity_type.clone(), + content: e.content.clone(), + }; + if let Some(index) = self.messages.iter().position(|m| m.id() == &e.message_id) { + let existing_is_activity = + matches!(self.messages.get(index), Some(Message::Activity { .. })); + if existing_is_activity || e.replace { + self.messages[index] = new_message; + current_mutation.messages = Some(self.messages.clone()); + } + } else { + self.messages.push(new_message); + current_mutation.messages = Some(self.messages.clone()); + } + + for subscriber in &self.subscribers { + let params = self.to_subscriber_params(); + let mutation = subscriber.on_activity_snapshot_event(e, params).await?; + mutations.push(mutation); + } + } + Event::ActivityDelta(e) => { + // Default behavior + if let Some(Message::Activity { + activity_type, + content, + .. + }) = self.messages.iter_mut().find(|m| m.id() == &e.message_id) + { + let patches: Vec = + serde_json::from_value(serde_json::to_value(e.patch.clone())?)?; + + json_patch::patch(content, &patches).map_err(|err| AgentError::Execution { + message: format!("Failed to apply activity patch: {err}"), + })?; + *activity_type = e.activity_type.clone(); + current_mutation.messages = Some(self.messages.clone()); + } + + for subscriber in &self.subscribers { + let params = self.to_subscriber_params(); + let mutation = subscriber.on_activity_delta_event(e, params).await?; + mutations.push(mutation); + } + } Event::Raw(e) => { for subscriber in &self.subscribers { let params = self.to_subscriber_params(); @@ -393,6 +530,15 @@ where } } Event::RunStarted(e) => { + if let Some(input) = &e.input { + for message in &input.messages { + if !self.messages.iter().any(|m| m.id() == message.id()) { + self.messages.push(message.clone()); + } + } + current_mutation.messages = Some(self.messages.clone()); + } + for subscriber in &self.subscribers { let params = self.to_subscriber_params(); let mutation = subscriber.on_run_started_event(e, params).await?; @@ -430,6 +576,132 @@ where mutations.push(mutation); } } + Event::ReasoningStart(e) => { + for subscriber in &self.subscribers { + let params = self.to_subscriber_params(); + let mutation = subscriber.on_reasoning_start_event(e, params).await?; + mutations.push(mutation); + } + } + Event::ReasoningMessageStart(e) => { + if !self.messages.iter().any(|m| m.id() == &e.message_id) { + self.messages.push(Message::Reasoning { + id: e.message_id.clone(), + content: String::new(), + encrypted_value: None, + }); + current_mutation.messages = Some(self.messages.clone()); + } + + for subscriber in &self.subscribers { + let params = self.to_subscriber_params(); + let mutation = subscriber + .on_reasoning_message_start_event(e, params) + .await?; + mutations.push(mutation); + } + } + Event::ReasoningMessageContent(e) => { + if let Some(message) = self.messages.iter_mut().find(|m| m.id() == &e.message_id) + && let Some(content) = message.content_mut() + { + content.push_str(&e.delta); + current_mutation.messages = Some(self.messages.clone()); + } + + let reasoning_message_buffer = self + .messages + .iter() + .find(|m| m.id() == &e.message_id) + .and_then(|m| m.content()) + .unwrap_or_default() + .to_string(); + + for subscriber in &self.subscribers { + let params = self.to_subscriber_params(); + let mutation = subscriber + .on_reasoning_message_content_event(e, &reasoning_message_buffer, params) + .await?; + mutations.push(mutation); + } + } + Event::ReasoningMessageEnd(e) => { + let reasoning_message_buffer = self + .messages + .iter() + .find(|m| m.id() == &e.message_id) + .and_then(|m| m.content()) + .unwrap_or_default() + .to_string(); + + for subscriber in &self.subscribers { + let params = self.to_subscriber_params(); + let mutation = subscriber + .on_reasoning_message_end_event(e, &reasoning_message_buffer, params) + .await?; + mutations.push(mutation); + } + } + Event::ReasoningMessageChunk(e) => { + for subscriber in &self.subscribers { + let params = self.to_subscriber_params(); + let mutation = subscriber + .on_reasoning_message_chunk_event(e, params) + .await?; + mutations.push(mutation); + } + } + Event::ReasoningEnd(e) => { + for subscriber in &self.subscribers { + let params = self.to_subscriber_params(); + let mutation = subscriber.on_reasoning_end_event(e, params).await?; + mutations.push(mutation); + } + } + Event::ReasoningEncryptedValue(e) => { + let mut entity_updated = false; + match e.subtype { + crate::core::event::ReasoningEncryptedValueSubtype::ToolCall => { + for message in &mut self.messages { + if let Message::Assistant { + tool_calls: Some(tool_calls), + .. + } = message + && let Some(tool_call) = tool_calls + .iter_mut() + .find(|tc| tc.id.to_string() == e.entity_id) + { + tool_call.encrypted_value = Some(e.encrypted_value.clone()); + entity_updated = true; + break; + } + } + } + crate::core::event::ReasoningEncryptedValueSubtype::Message => { + if let Some(encrypted_value) = self + .messages + .iter_mut() + .find(|message| message.id().to_string() == e.entity_id) + .and_then(|message| message.encrypted_value_mut()) + { + *encrypted_value = Some(e.encrypted_value.clone()); + entity_updated = true; + } + } + } + + if entity_updated { + current_mutation.messages = Some(self.messages.clone()); + } + + for subscriber in &self.subscribers { + let params = self.to_subscriber_params(); + let mutation = subscriber + .on_reasoning_encrypted_value_event(e, params) + .await?; + mutations.push(mutation); + } + } } for mutation in mutations { @@ -540,3 +812,134 @@ where Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::event::{ + BaseEvent, TextMessageContentEvent, TextMessageStartEvent, ToolCallArgsEvent, + ToolCallStartEvent, + }; + use crate::core::types::{RunId, ThreadId, ToolCallId}; + use crate::subscriber::Subscribers; + use serde_json::json; + + fn base_event() -> BaseEvent { + BaseEvent { + timestamp: None, + raw_event: None, + } + } + + fn input() -> RunAgentInput { + RunAgentInput::new( + ThreadId::random(), + RunId::random(), + json!({}), + vec![], + vec![], + vec![], + json!({}), + ) + } + + #[tokio::test] + async fn text_events_use_role_name_and_message_id() { + let input = input(); + let message_id = MessageId::random(); + let mut handler = EventHandler::new( + vec![Message::Assistant { + id: MessageId::random(), + content: Some("existing".to_string()), + name: None, + tool_calls: None, + encrypted_value: None, + }], + json!({}), + &input, + Subscribers::new(vec![]), + ); + + handler + .handle_event(&Event::TextMessageStart(TextMessageStartEvent { + base: base_event(), + message_id: message_id.clone(), + role: Role::User, + name: Some("Alice".to_string()), + })) + .await + .unwrap(); + + handler.messages.push(Message::Assistant { + id: MessageId::random(), + content: Some("tail".to_string()), + name: None, + tool_calls: None, + encrypted_value: None, + }); + + handler + .handle_event(&Event::TextMessageContent(TextMessageContentEvent { + base: base_event(), + message_id: message_id.clone(), + delta: "hello".to_string(), + })) + .await + .unwrap(); + + let target = handler + .messages + .iter() + .find(|message| message.id() == &message_id) + .unwrap(); + match target { + Message::User { content, name, .. } => { + assert_eq!(content, "hello"); + assert_eq!(name.as_deref(), Some("Alice")); + } + _ => panic!("expected user message"), + } + } + + #[tokio::test] + async fn tool_call_start_without_parent_keeps_tool_call() { + let input = input(); + let tool_call_id = ToolCallId::random(); + let mut handler = EventHandler::new( + vec![Message::new_user("hello")], + json!({}), + &input, + Subscribers::new(vec![]), + ); + + handler + .handle_event(&Event::ToolCallStart(ToolCallStartEvent { + base: base_event(), + tool_call_id: tool_call_id.clone(), + tool_call_name: "search".to_string(), + parent_message_id: None, + })) + .await + .unwrap(); + + handler + .handle_event(&Event::ToolCallArgs(ToolCallArgsEvent { + base: base_event(), + tool_call_id: tool_call_id.clone(), + delta: "{\"q\":\"ag-ui\"}".to_string(), + })) + .await + .unwrap(); + + let tool_call = handler + .messages + .iter() + .filter_map(|message| message.tool_calls()) + .flatten() + .find(|tool_call| tool_call.id == tool_call_id) + .unwrap(); + + assert_eq!(tool_call.function.name, "search"); + assert_eq!(tool_call.function.arguments, "{\"q\":\"ag-ui\"}"); + } +} diff --git a/sdks/community/rust/crates/ag-ui-client/src/subscriber.rs b/sdks/community/rust/crates/ag-ui-client/src/subscriber.rs index 9d880dcc8a..1bfc3c5e58 100644 --- a/sdks/community/rust/crates/ag-ui-client/src/subscriber.rs +++ b/sdks/community/rust/crates/ag-ui-client/src/subscriber.rs @@ -182,6 +182,22 @@ where Ok(AgentStateMutation::default()) } + async fn on_activity_snapshot_event( + &self, + event: &ActivitySnapshotEvent, + params: AgentSubscriberParams<'async_trait, StateT, FwdPropsT>, + ) -> Result, AgentError> { + Ok(AgentStateMutation::default()) + } + + async fn on_activity_delta_event( + &self, + event: &ActivityDeltaEvent, + params: AgentSubscriberParams<'async_trait, StateT, FwdPropsT>, + ) -> Result, AgentError> { + Ok(AgentStateMutation::default()) + } + async fn on_raw_event( &self, event: &RawEvent, @@ -254,6 +270,64 @@ where Ok(AgentStateMutation::default()) } + async fn on_reasoning_start_event( + &self, + event: &ReasoningStartEvent, + params: AgentSubscriberParams<'async_trait, StateT, FwdPropsT>, + ) -> Result, AgentError> { + Ok(AgentStateMutation::default()) + } + + async fn on_reasoning_message_start_event( + &self, + event: &ReasoningMessageStartEvent, + params: AgentSubscriberParams<'async_trait, StateT, FwdPropsT>, + ) -> Result, AgentError> { + Ok(AgentStateMutation::default()) + } + + async fn on_reasoning_message_content_event( + &self, + event: &ReasoningMessageContentEvent, + _reasoning_message_buffer: &str, + params: AgentSubscriberParams<'async_trait, StateT, FwdPropsT>, + ) -> Result, AgentError> { + Ok(AgentStateMutation::default()) + } + + async fn on_reasoning_message_end_event( + &self, + event: &ReasoningMessageEndEvent, + _reasoning_message_buffer: &str, + params: AgentSubscriberParams<'async_trait, StateT, FwdPropsT>, + ) -> Result, AgentError> { + Ok(AgentStateMutation::default()) + } + + async fn on_reasoning_message_chunk_event( + &self, + event: &ReasoningMessageChunkEvent, + params: AgentSubscriberParams<'async_trait, StateT, FwdPropsT>, + ) -> Result, AgentError> { + Ok(AgentStateMutation::default()) + } + + async fn on_reasoning_end_event( + &self, + event: &ReasoningEndEvent, + params: AgentSubscriberParams<'async_trait, StateT, FwdPropsT>, + ) -> Result, AgentError> { + Ok(AgentStateMutation::default()) + } + + async fn on_reasoning_encrypted_value_event( + &self, + event: &ReasoningEncryptedValueEvent, + params: AgentSubscriberParams<'async_trait, StateT, FwdPropsT>, + ) -> Result, AgentError> { + Ok(AgentStateMutation::default()) + } + // State changes async fn on_messages_changed( &self, diff --git a/sdks/community/rust/crates/ag-ui-core/src/event.rs b/sdks/community/rust/crates/ag-ui-core/src/event.rs index 64d42277be..c0d786dd54 100644 --- a/sdks/community/rust/crates/ag-ui-core/src/event.rs +++ b/sdks/community/rust/crates/ag-ui-core/src/event.rs @@ -1,6 +1,6 @@ use crate::JsonValue; use crate::state::AgentState; -use crate::types::{Message, Role}; +use crate::types::{Interrupt, Message, Role, RunAgentInput}; use crate::types::{MessageId, RunId, ThreadId, ToolCallId}; use serde::{Deserialize, Serialize}; @@ -42,6 +42,10 @@ pub enum EventType { StateDelta, /// Event containing a snapshot of the messages MessagesSnapshot, + /// Event containing a snapshot of structured activity + ActivitySnapshot, + /// Event containing a delta for structured activity + ActivityDelta, /// Event containing a raw event Raw, /// Event containing a custom event @@ -56,6 +60,20 @@ pub enum EventType { StepStarted, /// Event indicating that a step has finished StepFinished, + /// Event indicating the start of a reasoning block + ReasoningStart, + /// Event indicating the start of a reasoning message + ReasoningMessageStart, + /// Event containing a reasoning message delta + ReasoningMessageContent, + /// Event indicating the end of a reasoning message + ReasoningMessageEnd, + /// Event containing a chunk of reasoning message content + ReasoningMessageChunk, + /// Event indicating the end of a reasoning block + ReasoningEnd, + /// Event carrying an encrypted reasoning value + ReasoningEncryptedValue, } /// Base event for all events in the Agent User Interaction Protocol. @@ -68,6 +86,22 @@ pub struct BaseEvent { pub raw_event: Option, } +/// The interrupt-aware outcome of a completed run. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "lowercase")] +pub enum RunFinishedOutcome { + Success, + Interrupt { interrupts: Vec }, +} + +/// Entity subtype for encrypted reasoning values. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum ReasoningEncryptedValueSubtype { + Message, + ToolCall, +} + /// Event indicating the start of a text message. /// This event is sent when the agent begins generating a text message. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -76,7 +110,10 @@ pub struct TextMessageStartEvent { pub base: BaseEvent, #[serde(rename = "messageId")] pub message_id: MessageId, + #[serde(default = "Role::assistant")] pub role: Role, // "assistant" + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, } /// Event containing a piece of text message content. @@ -109,9 +146,12 @@ pub struct TextMessageChunkEvent { pub base: BaseEvent, #[serde(rename = "messageId", skip_serializing_if = "Option::is_none")] pub message_id: Option, - pub role: Role, + #[serde(skip_serializing_if = "Option::is_none")] + pub role: Option, #[serde(skip_serializing_if = "Option::is_none")] pub delta: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, } /// Event indicating the start of a thinking text message. @@ -252,6 +292,36 @@ pub struct MessagesSnapshotEvent { pub messages: Vec, } +/// Event containing a full structured activity snapshot. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ActivitySnapshotEvent { + #[serde(flatten)] + pub base: BaseEvent, + #[serde(rename = "messageId")] + pub message_id: MessageId, + #[serde(rename = "activityType")] + pub activity_type: String, + pub content: JsonValue, + #[serde(default = "default_activity_replace")] + pub replace: bool, +} + +fn default_activity_replace() -> bool { + true +} + +/// Event containing JSON Patch operations for an activity message. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ActivityDeltaEvent { + #[serde(flatten)] + pub base: BaseEvent, + #[serde(rename = "messageId")] + pub message_id: MessageId, + #[serde(rename = "activityType")] + pub activity_type: String, + pub patch: Vec, +} + /// Event containing a raw event. /// This event type allows wrapping arbitrary events from external sources. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -270,7 +340,8 @@ pub struct CustomEvent { #[serde(flatten)] pub base: BaseEvent, pub name: String, - pub value: JsonValue, + #[serde(skip_serializing_if = "Option::is_none")] + pub value: Option, } /// Event indicating that a run has started. @@ -283,6 +354,10 @@ pub struct RunStartedEvent { pub thread_id: ThreadId, #[serde(rename = "runId")] pub run_id: RunId, + #[serde(rename = "parentRunId", skip_serializing_if = "Option::is_none")] + pub parent_run_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub input: Option, } /// Event indicating that a run has finished. @@ -297,6 +372,8 @@ pub struct RunFinishedEvent { pub run_id: RunId, #[serde(skip_serializing_if = "Option::is_none")] pub result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub outcome: Option, } /// Event indicating that a run has encountered an error. @@ -330,6 +407,77 @@ pub struct StepFinishedEvent { pub step_name: String, } +/// Event indicating the start of a reasoning block. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ReasoningStartEvent { + #[serde(flatten)] + pub base: BaseEvent, + #[serde(rename = "messageId")] + pub message_id: MessageId, +} + +/// Event indicating the start of a reasoning message. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ReasoningMessageStartEvent { + #[serde(flatten)] + pub base: BaseEvent, + #[serde(rename = "messageId")] + pub message_id: MessageId, + #[serde(default = "Role::reasoning")] + pub role: Role, +} + +/// Event containing a piece of reasoning message content. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ReasoningMessageContentEvent { + #[serde(flatten)] + pub base: BaseEvent, + #[serde(rename = "messageId")] + pub message_id: MessageId, + pub delta: String, +} + +/// Event indicating the end of a reasoning message. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ReasoningMessageEndEvent { + #[serde(flatten)] + pub base: BaseEvent, + #[serde(rename = "messageId")] + pub message_id: MessageId, +} + +/// Event containing a chunk of reasoning message content. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ReasoningMessageChunkEvent { + #[serde(flatten)] + pub base: BaseEvent, + #[serde(rename = "messageId", skip_serializing_if = "Option::is_none")] + pub message_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub delta: Option, +} + +/// Event indicating the end of a reasoning block. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ReasoningEndEvent { + #[serde(flatten)] + pub base: BaseEvent, + #[serde(rename = "messageId")] + pub message_id: MessageId, +} + +/// Event carrying encrypted reasoning data for a message or tool call. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ReasoningEncryptedValueEvent { + #[serde(flatten)] + pub base: BaseEvent, + pub subtype: ReasoningEncryptedValueSubtype, + #[serde(rename = "entityId")] + pub entity_id: String, + #[serde(rename = "encryptedValue")] + pub encrypted_value: String, +} + /// Union of all possible events in the Agent User Interaction Protocol. /// This enum represents the full set of events that can be exchanged /// between the agent and the client. @@ -406,6 +554,12 @@ pub enum Event { /// Contains a vector of all current messages. MessagesSnapshot(MessagesSnapshotEvent), + /// Provides a complete structured activity snapshot. + ActivitySnapshot(ActivitySnapshotEvent), + + /// Provides incremental structured activity changes. + ActivityDelta(ActivityDeltaEvent), + /// Wraps a raw event from an external source. /// Contains the original event as a JSON value and an optional source identifier. Raw(RawEvent), @@ -433,9 +587,30 @@ pub enum Event { /// Signals the completion of a step within an agent run. /// Contains the name of the completed step. StepFinished(StepFinishedEvent), + + /// Signals the start of a reasoning block. + ReasoningStart(ReasoningStartEvent), + + /// Signals the start of a reasoning message. + ReasoningMessageStart(ReasoningMessageStartEvent), + + /// Represents content being added to an in-progress reasoning message. + ReasoningMessageContent(ReasoningMessageContentEvent), + + /// Signals the completion of a reasoning message. + ReasoningMessageEnd(ReasoningMessageEndEvent), + + /// Represents a complete or partial reasoning message chunk in a single event. + ReasoningMessageChunk(ReasoningMessageChunkEvent), + + /// Signals the end of a reasoning block. + ReasoningEnd(ReasoningEndEvent), + + /// Carries encrypted reasoning data for a message or tool call. + ReasoningEncryptedValue(ReasoningEncryptedValueEvent), } -impl Event { +impl Event { /// Get the event type pub fn event_type(&self) -> EventType { match self { @@ -456,6 +631,8 @@ impl Event { Event::StateSnapshot(_) => EventType::StateSnapshot, Event::StateDelta(_) => EventType::StateDelta, Event::MessagesSnapshot(_) => EventType::MessagesSnapshot, + Event::ActivitySnapshot(_) => EventType::ActivitySnapshot, + Event::ActivityDelta(_) => EventType::ActivityDelta, Event::Raw(_) => EventType::Raw, Event::Custom(_) => EventType::Custom, Event::RunStarted(_) => EventType::RunStarted, @@ -463,6 +640,13 @@ impl Event { Event::RunError(_) => EventType::RunError, Event::StepStarted(_) => EventType::StepStarted, Event::StepFinished(_) => EventType::StepFinished, + Event::ReasoningStart(_) => EventType::ReasoningStart, + Event::ReasoningMessageStart(_) => EventType::ReasoningMessageStart, + Event::ReasoningMessageContent(_) => EventType::ReasoningMessageContent, + Event::ReasoningMessageEnd(_) => EventType::ReasoningMessageEnd, + Event::ReasoningMessageChunk(_) => EventType::ReasoningMessageChunk, + Event::ReasoningEnd(_) => EventType::ReasoningEnd, + Event::ReasoningEncryptedValue(_) => EventType::ReasoningEncryptedValue, } } @@ -486,6 +670,8 @@ impl Event { Event::StateSnapshot(e) => e.base.timestamp, Event::StateDelta(e) => e.base.timestamp, Event::MessagesSnapshot(e) => e.base.timestamp, + Event::ActivitySnapshot(e) => e.base.timestamp, + Event::ActivityDelta(e) => e.base.timestamp, Event::Raw(e) => e.base.timestamp, Event::Custom(e) => e.base.timestamp, Event::RunStarted(e) => e.base.timestamp, @@ -493,6 +679,13 @@ impl Event { Event::RunError(e) => e.base.timestamp, Event::StepStarted(e) => e.base.timestamp, Event::StepFinished(e) => e.base.timestamp, + Event::ReasoningStart(e) => e.base.timestamp, + Event::ReasoningMessageStart(e) => e.base.timestamp, + Event::ReasoningMessageContent(e) => e.base.timestamp, + Event::ReasoningMessageEnd(e) => e.base.timestamp, + Event::ReasoningMessageChunk(e) => e.base.timestamp, + Event::ReasoningEnd(e) => e.base.timestamp, + Event::ReasoningEncryptedValue(e) => e.base.timestamp, } } } @@ -527,6 +720,7 @@ impl TextMessageStartEvent { }, message_id: message_id.into(), role: Role::Assistant, + name: None, } } @@ -539,6 +733,11 @@ impl TextMessageStartEvent { self.base.raw_event = Some(raw_event); self } + + pub fn with_name(mut self, name: String) -> Self { + self.name = Some(name); + self + } } impl TextMessageContentEvent { diff --git a/sdks/community/rust/crates/ag-ui-core/src/types/input.rs b/sdks/community/rust/crates/ag-ui-core/src/types/input.rs index 916a8b07b3..2bc5f320f9 100644 --- a/sdks/community/rust/crates/ag-ui-core/src/types/input.rs +++ b/sdks/community/rust/crates/ag-ui-core/src/types/input.rs @@ -5,6 +5,65 @@ use crate::types::message::Message; use crate::types::tool::Tool; use serde::{Deserialize, Serialize}; +/// An interrupt raised during an agent run for human-in-the-loop workflows. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct Interrupt { + pub id: String, + pub reason: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub message: Option, + #[serde(rename = "toolCallId", skip_serializing_if = "Option::is_none")] + pub tool_call_id: Option, + #[serde(rename = "responseSchema", skip_serializing_if = "Option::is_none")] + pub response_schema: Option, + #[serde(rename = "expiresAt", skip_serializing_if = "Option::is_none")] + pub expires_at: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub metadata: Option, +} + +impl Interrupt { + pub fn new(id: impl Into, reason: impl Into) -> Self { + Self { + id: id.into(), + reason: reason.into(), + message: None, + tool_call_id: None, + response_schema: None, + expires_at: None, + metadata: None, + } + } +} + +/// A resume response for a previously emitted interrupt. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ResumeEntry { + #[serde(rename = "interruptId")] + pub interrupt_id: String, + pub status: ResumeStatus, + #[serde(skip_serializing_if = "Option::is_none")] + pub payload: Option, +} + +impl ResumeEntry { + pub fn new(interrupt_id: impl Into, status: ResumeStatus) -> Self { + Self { + interrupt_id: interrupt_id.into(), + status, + payload: None, + } + } +} + +/// Status for a resume entry. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ResumeStatus { + Resolved, + Cancelled, +} + /// Input for running an agent. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct RunAgentInput { @@ -12,12 +71,16 @@ pub struct RunAgentInput { pub thread_id: ThreadId, #[serde(rename = "runId")] pub run_id: RunId, + #[serde(rename = "parentRunId", skip_serializing_if = "Option::is_none")] + pub parent_run_id: Option, pub state: StateT, pub messages: Vec, pub tools: Vec, pub context: Vec, #[serde(rename = "forwardedProps")] pub forwarded_props: FwdPropsT, + #[serde(skip_serializing_if = "Option::is_none")] + pub resume: Option>, } impl RunAgentInput { @@ -33,11 +96,23 @@ impl RunAgentInput { Self { thread_id: thread_id.into(), run_id: run_id.into(), + parent_run_id: None, state, messages, tools, context, forwarded_props, + resume: None, } } + + pub fn with_parent_run_id(mut self, parent_run_id: impl Into) -> Self { + self.parent_run_id = Some(parent_run_id.into()); + self + } + + pub fn with_resume(mut self, resume: Vec) -> Self { + self.resume = Some(resume); + self + } } diff --git a/sdks/community/rust/crates/ag-ui-core/src/types/message.rs b/sdks/community/rust/crates/ag-ui-core/src/types/message.rs index 66e63ac0e8..1e20ced07f 100644 --- a/sdks/community/rust/crates/ag-ui-core/src/types/message.rs +++ b/sdks/community/rust/crates/ag-ui-core/src/types/message.rs @@ -1,3 +1,4 @@ +use crate::JsonValue; use crate::types::ids::{MessageId, ToolCallId}; use crate::types::tool::ToolCall; use serde::{Deserialize, Serialize}; @@ -19,6 +20,8 @@ pub enum Role { Assistant, User, Tool, + Activity, + Reasoning, } // Utility methods for serde defaults @@ -38,6 +41,9 @@ impl Role { pub(crate) fn tool() -> Self { Self::Tool } + pub(crate) fn reasoning() -> Self { + Self::Reasoning + } } /// A basic message, where the only content should be an optional string. @@ -49,6 +55,8 @@ pub struct BaseMessage { pub content: Option, #[serde(skip_serializing_if = "Option::is_none")] pub name: Option, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + pub encrypted_value: Option, } /// A developer message. @@ -118,6 +126,8 @@ pub struct AssistantMessage { pub name: Option, #[serde(rename = "toolCalls", skip_serializing_if = "Option::is_none")] pub tool_calls: Option>, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + pub encrypted_value: Option, } impl AssistantMessage { @@ -128,6 +138,7 @@ impl AssistantMessage { content: None, name: None, tool_calls: None, + encrypted_value: None, } } @@ -145,6 +156,11 @@ impl AssistantMessage { self.tool_calls = Some(tool_calls); self } + + pub fn with_encrypted_value(mut self, encrypted_value: String) -> Self { + self.encrypted_value = Some(encrypted_value); + self + } } /// A user message. @@ -156,6 +172,8 @@ pub struct UserMessage { pub content: String, #[serde(skip_serializing_if = "Option::is_none")] pub name: Option, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + pub encrypted_value: Option, } impl UserMessage { @@ -165,6 +183,7 @@ impl UserMessage { role: Role::User, content, name: None, + encrypted_value: None, } } @@ -172,6 +191,11 @@ impl UserMessage { self.name = Some(name); self } + + pub fn with_encrypted_value(mut self, encrypted_value: String) -> Self { + self.encrypted_value = Some(encrypted_value); + self + } } /// A tool call result. @@ -185,6 +209,8 @@ pub struct ToolMessage { pub tool_call_id: ToolCallId, #[serde(skip_serializing_if = "Option::is_none")] pub error: Option, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + pub encrypted_value: Option, } impl ToolMessage { @@ -199,6 +225,7 @@ impl ToolMessage { role: Role::Tool, tool_call_id: tool_call_id.into(), error: None, + encrypted_value: None, } } @@ -206,6 +233,11 @@ impl ToolMessage { self.error = Some(error); self } + + pub fn with_encrypted_value(mut self, encrypted_value: String) -> Self { + self.encrypted_value = Some(encrypted_value); + self + } } /// Represents the different type of messages that you might receive, but as an enum. @@ -217,12 +249,16 @@ pub enum Message { content: String, #[serde(skip_serializing_if = "Option::is_none")] name: Option, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + encrypted_value: Option, }, System { id: MessageId, content: String, #[serde(skip_serializing_if = "Option::is_none")] name: Option, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + encrypted_value: Option, }, Assistant { id: MessageId, @@ -232,12 +268,16 @@ pub enum Message { name: Option, #[serde(rename = "toolCalls", skip_serializing_if = "Option::is_none")] tool_calls: Option>, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + encrypted_value: Option, }, User { id: MessageId, content: String, #[serde(skip_serializing_if = "Option::is_none")] name: Option, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + encrypted_value: Option, }, Tool { id: MessageId, @@ -246,6 +286,20 @@ pub enum Message { tool_call_id: ToolCallId, #[serde(skip_serializing_if = "Option::is_none")] error: Option, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + encrypted_value: Option, + }, + Activity { + id: MessageId, + #[serde(rename = "activityType")] + activity_type: String, + content: JsonValue, + }, + Reasoning { + id: MessageId, + content: String, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + encrypted_value: Option, }, } @@ -256,28 +310,43 @@ impl Message { id: id.into(), content: content.as_ref().to_string(), name: None, + encrypted_value: None, }, Role::System => Self::System { id: id.into(), content: content.as_ref().to_string(), name: None, + encrypted_value: None, }, Role::Assistant => Self::Assistant { id: id.into(), content: Some(content.as_ref().to_string()), name: None, tool_calls: None, + encrypted_value: None, }, Role::User => Self::User { id: id.into(), content: content.as_ref().to_string(), name: None, + encrypted_value: None, }, Role::Tool => Self::Tool { id: id.into(), content: content.as_ref().to_string(), tool_call_id: ToolCallId::random(), error: None, + encrypted_value: None, + }, + Role::Activity => Self::Activity { + id: id.into(), + activity_type: "activity".to_string(), + content: JsonValue::String(content.as_ref().to_string()), + }, + Role::Reasoning => Self::Reasoning { + id: id.into(), + content: content.as_ref().to_string(), + encrypted_value: None, }, } } @@ -314,6 +383,8 @@ impl Message { Message::Assistant { id, .. } => id, Message::User { id, .. } => id, Message::Tool { id, .. } => id, + Message::Activity { id, .. } => id, + Message::Reasoning { id, .. } => id, } } @@ -324,6 +395,8 @@ impl Message { Message::Assistant { id, .. } => id, Message::User { id, .. } => id, Message::Tool { id, .. } => id, + Message::Activity { id, .. } => id, + Message::Reasoning { id, .. } => id, } } @@ -334,6 +407,8 @@ impl Message { Message::Assistant { .. } => Role::Assistant, Message::User { .. } => Role::User, Message::Tool { .. } => Role::Tool, + Message::Activity { .. } => Role::Activity, + Message::Reasoning { .. } => Role::Reasoning, } } pub fn content(&self) -> Option<&str> { @@ -343,6 +418,8 @@ impl Message { Message::User { content, .. } => Some(content), Message::Tool { content, .. } => Some(content), Message::Assistant { content, .. } => content.as_deref(), + Message::Activity { .. } => None, + Message::Reasoning { content, .. } => Some(content), } } @@ -358,6 +435,8 @@ impl Message { } content.as_mut() } + Message::Activity { .. } => None, + Message::Reasoning { content, .. } => Some(content), } } @@ -379,4 +458,28 @@ impl Message { _ => None, } } + + pub fn encrypted_value_mut(&mut self) -> Option<&mut Option> { + match self { + Message::Developer { + encrypted_value, .. + } + | Message::System { + encrypted_value, .. + } + | Message::Assistant { + encrypted_value, .. + } + | Message::User { + encrypted_value, .. + } + | Message::Tool { + encrypted_value, .. + } + | Message::Reasoning { + encrypted_value, .. + } => Some(encrypted_value), + Message::Activity { .. } => None, + } + } } diff --git a/sdks/community/rust/crates/ag-ui-core/src/types/tool.rs b/sdks/community/rust/crates/ag-ui-core/src/types/tool.rs index c9195c2059..96d7e1bcf8 100644 --- a/sdks/community/rust/crates/ag-ui-core/src/types/tool.rs +++ b/sdks/community/rust/crates/ag-ui-core/src/types/tool.rs @@ -9,6 +9,8 @@ pub struct ToolCall { #[serde(rename = "type")] pub call_type: String, pub function: FunctionCall, + #[serde(rename = "encryptedValue", skip_serializing_if = "Option::is_none")] + pub encrypted_value: Option, } impl ToolCall { @@ -17,8 +19,14 @@ impl ToolCall { id: id.into(), call_type: "function".to_string(), function, + encrypted_value: None, } } + + pub fn with_encrypted_value(mut self, encrypted_value: impl Into) -> Self { + self.encrypted_value = Some(encrypted_value.into()); + self + } } /// A tool definition. @@ -30,6 +38,9 @@ pub struct Tool { pub description: String, /// The tool parameters pub parameters: serde_json::Value, + /// Arbitrary tool metadata. + #[serde(skip_serializing_if = "Option::is_none")] + pub metadata: Option, } impl Tool { @@ -38,6 +49,12 @@ impl Tool { name, description, parameters, + metadata: None, } } + + pub fn with_metadata(mut self, metadata: JsonValue) -> Self { + self.metadata = Some(metadata); + self + } } diff --git a/sdks/community/rust/crates/ag-ui-core/tests/unit.rs b/sdks/community/rust/crates/ag-ui-core/tests/unit.rs index 4890cfa66e..0e097c88bf 100644 --- a/sdks/community/rust/crates/ag-ui-core/tests/unit.rs +++ b/sdks/community/rust/crates/ag-ui-core/tests/unit.rs @@ -1,20 +1,36 @@ #[cfg(test)] mod tests { use ag_ui_core::error::AgUiError; + use ag_ui_core::event::{ + ActivityDeltaEvent, ActivitySnapshotEvent, BaseEvent, Event, EventType, + ReasoningEncryptedValueEvent, ReasoningEncryptedValueSubtype, ReasoningMessageContentEvent, + RunFinishedEvent, RunFinishedOutcome, TextMessageStartEvent, + }; use ag_ui_core::types::{ - AssistantMessage, Context, DeveloperMessage, FunctionCall, Message, MessageId, Role, - RunAgentInput, RunId, SystemMessage, ThreadId, Tool, ToolCall, ToolCallId, ToolMessage, - UserMessage, + AssistantMessage, Context, DeveloperMessage, FunctionCall, Interrupt, Message, MessageId, + ResumeEntry, ResumeStatus, Role, RunAgentInput, RunId, SystemMessage, ThreadId, Tool, + ToolCall, ToolCallId, ToolMessage, UserMessage, }; use serde::{Deserialize, Serialize}; use serde_json::json; use uuid::Uuid; + fn base_event() -> BaseEvent { + BaseEvent { + timestamp: None, + raw_event: None, + } + } + #[test] fn test_role_serialization() { let role = Role::Developer; let json = serde_json::to_string(&role).unwrap(); assert_eq!(json, r#""developer""#); + + let role = Role::Reasoning; + let json = serde_json::to_string(&role).unwrap(); + assert_eq!(json, r#""reasoning""#); } #[test] @@ -48,6 +64,7 @@ mod tests { id: MessageId::random(), content: "Hello".to_string(), name: None, + encrypted_value: None, }; let json = serde_json::to_string(&user_msg).unwrap(); @@ -65,6 +82,9 @@ mod tests { let tool_call = ToolCall::new(ToolCallId::random(), function_call); assert_eq!(tool_call.call_type, "function"); + + let tool_call = tool_call.with_encrypted_value("enc".to_string()); + assert_eq!(tool_call.encrypted_value, Some("enc".to_string())); } #[test] @@ -86,8 +106,10 @@ mod tests { "test_tool".to_string(), "tool desc".to_string(), json!({"type": "object"}), - ); + ) + .with_metadata(json!({"renderer": "a2ui"})); assert_eq!(tool.name, "test_tool"); + assert_eq!(tool.metadata, Some(json!({"renderer": "a2ui"}))); } #[test] @@ -172,6 +194,7 @@ mod tests { content, name, tool_calls, + .. } => { assert_eq!(id.to_string(), "00000000-0000-0000-0000-000000000000"); assert_eq!( @@ -215,7 +238,9 @@ mod tests { assert_eq!(messages.len(), 3); match &messages[0] { - Message::User { id, content, name } => { + Message::User { + id, content, name, .. + } => { assert_eq!(id.to_string(), "00000000-0000-0000-0000-000000000000"); assert_eq!(content, "Hello!"); assert_eq!(*name, Some("Alice".to_string())); @@ -224,6 +249,43 @@ mod tests { } } + #[test] + fn test_activity_and_reasoning_messages() { + let activity_id = MessageId::random(); + let reasoning_id = MessageId::random(); + let json_value = json!([ + { + "role": "activity", + "id": activity_id, + "activityType": "progress", + "content": {"pct": 10} + }, + { + "role": "reasoning", + "id": reasoning_id, + "content": "step 1", + "encryptedValue": "enc" + } + ]); + + let messages: Vec = serde_json::from_value(json_value).unwrap(); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].role(), Role::Activity); + assert_eq!(messages[1].role(), Role::Reasoning); + + match &messages[0] { + Message::Activity { + activity_type, + content, + .. + } => { + assert_eq!(activity_type, "progress"); + assert_eq!(content["pct"], 10); + } + _ => panic!("Wrong message type"), + } + } + #[test] fn test_complex_run_agent_input_deserialization() { let json_str = r#"{ @@ -317,4 +379,132 @@ mod tests { serde_json::from_str(json_str); assert!(wrong_input.is_err()) } + + #[test] + fn test_run_agent_input_parent_and_resume_serialization() { + let parent_run_id = RunId::random(); + let resume = ResumeEntry { + interrupt_id: "int1".to_string(), + status: ResumeStatus::Resolved, + payload: Some(json!({"answer": 42})), + }; + + let input = RunAgentInput::new( + ThreadId::random(), + RunId::random(), + json!({}), + vec![], + vec![], + vec![], + json!({}), + ) + .with_parent_run_id(parent_run_id.clone()) + .with_resume(vec![resume]); + + let payload = serde_json::to_value(&input).unwrap(); + assert_eq!(payload["parentRunId"], parent_run_id.to_string()); + assert_eq!(payload["resume"][0]["interruptId"], "int1"); + assert_eq!(payload["resume"][0]["status"], "resolved"); + + let round_trip: RunAgentInput = serde_json::from_value(payload).unwrap(); + assert_eq!(round_trip.parent_run_id, Some(parent_run_id)); + assert_eq!(round_trip.resume.unwrap()[0].status, ResumeStatus::Resolved); + } + + #[test] + fn test_new_event_shapes_serialization() { + let message_id = MessageId::random(); + let text_event: Event = Event::TextMessageStart(TextMessageStartEvent { + base: base_event(), + message_id: message_id.clone(), + role: Role::Assistant, + name: Some("research-agent".to_string()), + }); + let payload = serde_json::to_value(&text_event).unwrap(); + assert_eq!(payload["type"], "TEXT_MESSAGE_START"); + assert_eq!(payload["messageId"], message_id.to_string()); + assert_eq!(payload["name"], "research-agent"); + + let activity_event: Event = Event::ActivitySnapshot(ActivitySnapshotEvent { + base: base_event(), + message_id: message_id.clone(), + activity_type: "progress".to_string(), + content: json!({"pct": 20}), + replace: true, + }); + let payload = serde_json::to_value(&activity_event).unwrap(); + assert_eq!(payload["type"], "ACTIVITY_SNAPSHOT"); + assert_eq!(payload["activityType"], "progress"); + assert_eq!(payload["content"]["pct"], 20); + + let reasoning_event: Event = Event::ReasoningMessageContent(ReasoningMessageContentEvent { + base: base_event(), + message_id: message_id.clone(), + delta: "thinking".to_string(), + }); + let payload = serde_json::to_value(&reasoning_event).unwrap(); + assert_eq!(payload["type"], "REASONING_MESSAGE_CONTENT"); + assert_eq!(payload["delta"], "thinking"); + } + + #[test] + fn test_run_finished_outcome_serialization() { + let interrupt = Interrupt::new("int1", "input_required"); + let event: Event = Event::RunFinished(RunFinishedEvent { + base: base_event(), + thread_id: ThreadId::random(), + run_id: RunId::random(), + result: None, + outcome: Some(RunFinishedOutcome::Interrupt { + interrupts: vec![interrupt], + }), + }); + + let payload = serde_json::to_value(&event).unwrap(); + assert_eq!(payload["type"], "RUN_FINISHED"); + assert_eq!(payload["outcome"]["type"], "interrupt"); + assert_eq!(payload["outcome"]["interrupts"][0]["id"], "int1"); + + let round_trip: Event = serde_json::from_value(payload).unwrap(); + assert_eq!(round_trip.event_type(), EventType::RunFinished); + } + + #[test] + fn test_activity_delta_and_reasoning_encrypted_events_deserialize() { + let message_id = MessageId::random(); + let delta_payload = json!({ + "type": "ACTIVITY_DELTA", + "messageId": message_id, + "activityType": "progress", + "patch": [{"op": "replace", "path": "/pct", "value": 50}] + }); + let event: Event = serde_json::from_value(delta_payload).unwrap(); + match event { + Event::ActivityDelta(ActivityDeltaEvent { patch, .. }) => { + assert_eq!(patch[0]["op"], "replace"); + } + _ => panic!("Wrong event type"), + } + + let encrypted_payload = json!({ + "type": "REASONING_ENCRYPTED_VALUE", + "subtype": "tool-call", + "entityId": "call_123", + "encryptedValue": "enc" + }); + let event: Event = serde_json::from_value(encrypted_payload).unwrap(); + match event { + Event::ReasoningEncryptedValue(ReasoningEncryptedValueEvent { + subtype, + entity_id, + encrypted_value, + .. + }) => { + assert_eq!(subtype, ReasoningEncryptedValueSubtype::ToolCall); + assert_eq!(entity_id, "call_123"); + assert_eq!(encrypted_value, "enc"); + } + _ => panic!("Wrong event type"), + } + } }