From 86a56d5a7c9c9f1bef0ea7b0c2e04f40fc63b9c2 Mon Sep 17 00:00:00 2001 From: Shadaj Laddad Date: Thu, 18 Jun 2026 00:14:44 +0000 Subject: [PATCH] fix: subscription events inherit pretty tool call display_as MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In both the live processing path (batch_processor.rs) and the history replay path (display.rs), subscription events now use the `display_as` string from the originating `InfinityMessage::ToolCall` when available, falling back to `name(args)` only when no pretty display was computed. The history lookup iterates in reverse since the subscription tool call is typically near the end of the conversation. Added a unit test (`subscription_event_uses_display_as`) that seeds history with a tool call + result (subscription started), then sends a synthetic subscription event and asserts the display uses the pretty `display_as` string rather than raw `name(args)`. Co-authored-by: Infinity 🤖 PR: #40 --- .../src/batch_processor.rs | 121 ++++++++++++++++-- crates/infinity-daemon/src/session/display.rs | 10 +- 2 files changed, 114 insertions(+), 17 deletions(-) diff --git a/crates/infinity-agent-core/src/batch_processor.rs b/crates/infinity-agent-core/src/batch_processor.rs index f551b190..ee7addb1 100644 --- a/crates/infinity-agent-core/src/batch_processor.rs +++ b/crates/infinity-agent-core/src/batch_processor.rs @@ -12,11 +12,13 @@ use std::pin::Pin; use futures_util::StreamExt; use rig::completion::{GetTokenUsage, ToolDefinition}; -use rig::message::{AssistantContent, Message, ToolResultContent, UserContent}; +use rig::message::{ToolResultContent, UserContent}; use tokio::sync::{mpsc, oneshot}; use crate::event_processor::{self, CompletionAction, HistoryManager}; -use crate::message::{InputMessage, InputMessageContent, SyntheticKind, TaggedSyntheticKind}; +use crate::message::{ + InfinityMessage, InputMessage, InputMessageContent, SyntheticKind, TaggedSyntheticKind, +}; use crate::model_provider::{ModelProvider, ProviderStreamingResponse}; use crate::tools::{Tool, ToolContext}; use crate::traits::{ConversationStore, InputSender, StateStore}; @@ -131,20 +133,17 @@ where if let InputMessageContent::User(UserContent::ToolResult(res)) = &input_msg.content && let ToolResultContent::Text(text) = res.content.first() { - let orig_call = current_history.get_history().into_iter().find(|h| { - if let Message::Assistant { content, .. } = h - && let AssistantContent::ToolCall(c) = content.first() + let orig_call = current_history.history.borrow().iter().rev().find_map(|m| { + if let InfinityMessage::ToolCall { call, display_as } = m + && call.id == synth.tool_call_id() { - c.id == synth.tool_call_id() + Some((call.clone(), display_as.clone())) } else { - false + None } }); - if let Some(h) = orig_call - && let Message::Assistant { content, .. } = h - && let AssistantContent::ToolCall(c) = content.first() - { + if let Some((c, display_as)) = orig_call { let name = if let SyntheticKind::Tagged(TaggedSyntheticKind::ThreadReport { ref child_thread_id, @@ -153,7 +152,9 @@ where { format!("Report from child thread {}", child_thread_id) } else { - format!("{}({})", c.function.name, c.function.arguments) + display_as.unwrap_or_else(|| { + format!("{}({})", c.function.name, c.function.arguments) + }) }; let _ = display_tx.send(DisplayEvent::SubscriptionEvent { name, @@ -969,4 +970,100 @@ mod tests { assert!(ev.iter().any(|e| e.starts_with("OAuth:"))); assert!(ev.iter().any(|e| e == "UserInput:hello")); } + + #[tokio::test] + async fn subscription_event_uses_display_as() { + use crate::message::{SyntheticKind, TaggedSyntheticKind}; + + let s = StubConvo::new(); + let hm = HistoryManager::new_with_history(s.clone(), StubState, "t1".into()) + .await + .expect("create history manager"); + + // Seed history with a tool call + its initial result (subscription started). + *hm.history.borrow_mut() = vec![ + crate::message::InfinityMessage::from_rig_message(Message::User { + content: OneOrMany::one(UserContent::text("go")), + }), + crate::message::InfinityMessage::ToolCall { + call: rig::message::ToolCall { + id: "tc-sub".into(), + call_id: None, + function: rig::message::ToolFunction { + name: "sleep".into(), + arguments: serde_json::json!({"seconds": 30, "reason": "waiting"}), + }, + additional_params: None, + signature: None, + }, + display_as: Some("Sleeping 30s: waiting".into()), + }, + crate::message::InfinityMessage::ToolResult { + result: ToolResult { + id: "tc-sub".into(), + call_id: None, + content: OneOrMany::one(ToolResultContent::Text(rig::agent::Text { + text: "Subscription started".into(), + })), + }, + display_segments: None, + }, + ]; + + let (provider, _ctrl) = mock_provider(); + let (dtx, mut drx) = mpsc::unbounded_channel(); + let tn = HashSet::new(); + let td: Vec = vec![]; + let tr: HashMap> = HashMap::new(); + + // Simulate a subscription event arriving as a synthetic tool result. + let input = ( + InputMessage { + content: InputMessageContent::User(UserContent::ToolResult(ToolResult { + id: "tc-sub".into(), + call_id: None, + content: OneOrMany::one(ToolResultContent::Text(rig::agent::Text { + text: "woke up".into(), + })), + })), + group_id: "t1".into(), + metadata: None, + synthetic: Some(SyntheticKind::Tagged( + TaggedSyntheticKind::SubscriptionEvent { + tool_call_id: "tc-sub".into(), + associative: true, + r#final: false, + }, + )), + display_as: None, + subscription: false, + }, + "m-sub".into(), + ); + + let _r = super::process_batch( + vec![input].into_iter(), + &hm, + &s, + &dtx, + "t1", + &provider, + "mock", + &tn, + &td, + &tr, + ctx(), + &None, + NONE_NOTIFIER, + None, + ) + .await; + + let events = drain(&mut drx); + assert!( + events.contains(&"SubEvent:Sleeping 30s: waiting".to_owned()), + "Expected pretty display_as in subscription event, got: {:?}", + events + ); + } } diff --git a/crates/infinity-daemon/src/session/display.rs b/crates/infinity-daemon/src/session/display.rs index d3b1b159..bd59264e 100644 --- a/crates/infinity-daemon/src/session/display.rs +++ b/crates/infinity-daemon/src/session/display.rs @@ -105,14 +105,14 @@ pub(crate) fn history_message_to_daemon( } else { history .iter() + .rev() .find_map(|m| { - if let InfinityMessage::ToolCall { call, .. } = m + if let InfinityMessage::ToolCall { call, display_as } = m && call.id == *tool_call_id { - Some(format!( - "{}({})", - call.function.name, call.function.arguments - )) + Some(display_as.clone().unwrap_or_else(|| { + format!("{}({})", call.function.name, call.function.arguments) + })) } else { None }