From 4d7501a9684469998f2b518f6bd3da8bc95b266a Mon Sep 17 00:00:00 2001 From: Henry Park Date: Sun, 22 Mar 2026 20:33:52 -0700 Subject: [PATCH] Fix owner-scoped message routing fallbacks (#1574) * Fix owner-scoped message routing fallbacks * Address PR feedback on routing regressions * Address review notes on routing fallbacks --- src/testing/mod.rs | 71 ++++++++++++++- src/tools/builtin/message.rs | 167 ++++++++++++++++------------------- src/worker/job.rs | 65 ++++++++++++++ 3 files changed, 211 insertions(+), 92 deletions(-) diff --git a/src/testing/mod.rs b/src/testing/mod.rs index 953cbfcd..a633e91c 100644 --- a/src/testing/mod.rs +++ b/src/testing/mod.rs @@ -28,7 +28,7 @@ use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; use async_trait::async_trait; use rust_decimal::Decimal; -use tokio::sync::mpsc; +use tokio::sync::{Mutex as AsyncMutex, mpsc}; use crate::agent::AgentDeps; use crate::channels::{ @@ -361,6 +361,75 @@ impl Channel for StubChannel { } } +/// Captured broadcast deliveries keyed by the target user or chat identifier. +pub type BroadcastCapture = Arc>>; + +/// A lightweight channel double that only records `broadcast()` traffic. +/// +/// This is useful for unit tests that need to assert message routing without +/// spinning up a full interactive channel harness. +pub struct RecordingBroadcastChannel { + name: &'static str, + captures: BroadcastCapture, +} + +impl RecordingBroadcastChannel { + pub fn new(name: &'static str) -> (Self, BroadcastCapture) { + let captures = Arc::new(AsyncMutex::new(Vec::new())); + ( + Self { + name, + captures: Arc::clone(&captures), + }, + captures, + ) + } +} + +#[async_trait] +impl Channel for RecordingBroadcastChannel { + fn name(&self) -> &str { + self.name + } + + async fn start(&self) -> Result { + let (_tx, rx) = mpsc::channel::(1); + Ok(Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx))) + } + + async fn respond( + &self, + _msg: &IncomingMessage, + _response: OutgoingResponse, + ) -> Result<(), ChannelError> { + Ok(()) + } + + async fn send_status( + &self, + _status: StatusUpdate, + _metadata: &serde_json::Value, + ) -> Result<(), ChannelError> { + Ok(()) + } + + async fn broadcast( + &self, + user_id: &str, + response: OutgoingResponse, + ) -> Result<(), ChannelError> { + self.captures + .lock() + .await + .push((user_id.to_string(), response)); + Ok(()) + } + + async fn health_check(&self) -> Result<(), ChannelError> { + Ok(()) + } +} + /// Assembled test components. pub struct TestHarness { /// The agent dependencies, ready for use. diff --git a/src/tools/builtin/message.rs b/src/tools/builtin/message.rs index 83041b80..08029d6f 100644 --- a/src/tools/builtin/message.rs +++ b/src/tools/builtin/message.rs @@ -80,6 +80,12 @@ fn metadata_notify_user(metadata: &serde_json::Value) -> Option { metadata_string(metadata, "notify_user").filter(|value| value != "default") } +// Autonomous runs include `owner_id` when the job is executing on behalf of a +// durable owner scope instead of an interactive channel actor. +fn metadata_owner_id(metadata: &serde_json::Value) -> Option { + metadata_string(metadata, "owner_id") +} + fn channel_matches_source(resolved_channel: Option<&str>, source_channel: Option<&str>) -> bool { match (resolved_channel, source_channel) { (None, _) => true, @@ -91,11 +97,13 @@ fn channel_matches_source(resolved_channel: Option<&str>, source_channel: Option async fn resolve_channel_fallback_target( extension_manager: Option<&Arc>, channel: Option<&str>, + owner_scope_target: Option<&str>, ctx_user_id: &str, ) -> Option { - let channel_name = channel?; - - if let Some(extension_manager) = extension_manager + // Prefer an explicit channel binding when the extension manager knows the + // durable delivery target (for example, a bound Telegram chat ID). + if let Some(channel_name) = channel + && let Some(extension_manager) = extension_manager && let Some(target) = extension_manager .notification_target_for_channel(channel_name) .await @@ -103,13 +111,19 @@ async fn resolve_channel_fallback_target( return Some(target); } - Some(ctx_user_id.to_string()) + // `owner_id` is only present for autonomous owner-scoped executions. + // Interactive chat turns intentionally fall back to `ctx.user_id`, which is + // already the active conversation target for the current channel. + owner_scope_target + .map(ToOwned::to_owned) + .or_else(|| Some(ctx_user_id.to_string())) } struct MessageTargetResolution<'a> { extension_manager: Option<&'a Arc>, explicit_target: Option, metadata_target: Option, + owner_scope_target: Option, default_target: Option, channel: Option<&'a str>, metadata_channel: Option<&'a str>, @@ -133,6 +147,7 @@ async fn resolve_message_target(inputs: MessageTargetResolution<'_>) -> Option) -> Option>>; - - struct RecordingChannel { - name: &'static str, - captures: BroadcastCapture, - } - - impl RecordingChannel { - fn new(name: &'static str) -> (Self, BroadcastCapture) { - let captures = Arc::new(Mutex::new(Vec::new())); - ( - Self { - name, - captures: Arc::clone(&captures), - }, - captures, - ) - } - } - - #[async_trait] - impl Channel for RecordingChannel { - fn name(&self) -> &str { - self.name - } - - async fn start(&self) -> Result { - let (_tx, rx) = mpsc::channel::(1); - Ok(Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx))) - } - - async fn respond( - &self, - _msg: &IncomingMessage, - _response: OutgoingResponse, - ) -> Result<(), ChannelError> { - Ok(()) - } - - async fn send_status( - &self, - _status: StatusUpdate, - _metadata: &serde_json::Value, - ) -> Result<(), ChannelError> { - Ok(()) - } - - async fn broadcast( - &self, - user_id: &str, - response: OutgoingResponse, - ) -> Result<(), ChannelError> { - self.captures - .lock() - .await - .push((user_id.to_string(), response)); - Ok(()) - } - - async fn health_check(&self) -> Result<(), ChannelError> { - Ok(()) - } - } + use crate::testing::{BroadcastCapture, RecordingBroadcastChannel}; async fn message_tool_with_recording_channels() -> (MessageTool, BroadcastCapture, BroadcastCapture) { let channel_manager = ChannelManager::new(); - let (gateway, gateway_captures) = RecordingChannel::new("gateway"); - let (telegram, telegram_captures) = RecordingChannel::new("telegram"); + let (gateway, gateway_captures) = RecordingBroadcastChannel::new("gateway"); + let (telegram, telegram_captures) = RecordingBroadcastChannel::new("telegram"); channel_manager.add(Box::new(gateway)).await; channel_manager.add(Box::new(telegram)).await; @@ -870,28 +820,63 @@ mod tests { } #[tokio::test] - async fn message_tool_falls_back_to_ctx_user_when_channel_known() { - // Regression for owner-scoped notifications: a channel can be known - // even when the concrete delivery target is omitted, so the message - // tool should pass ctx.user_id through to the channel layer. - let tool = MessageTool::new(Arc::new(ChannelManager::new())); + async fn message_tool_falls_back_to_owner_scope_when_channel_known() { + let (tool, gateway_captures, telegram_captures) = + message_tool_with_recording_channels().await; let mut ctx = - crate::context::JobContext::with_user("owner-scope", "routine-job", "price alert"); + crate::context::JobContext::with_user("telegram", "routine-job", "price alert"); + ctx.metadata = serde_json::json!({ + "notify_channel": "telegram", + "owner_id": "owner-scope", + }); + + let result = tool + .execute(serde_json::json!({"content": "NEAR price is $5"}), &ctx) + .await + .expect("message tool should use owner scope before ctx.user_id"); + + assert_eq!( + result.result.as_str(), + Some("Sent message to telegram:owner-scope") + ); + assert!(gateway_captures.lock().await.is_empty()); + let telegram = telegram_captures.lock().await.clone(); + assert_eq!(telegram.len(), 1); + assert_eq!(telegram[0].0, "owner-scope"); + assert_eq!(telegram[0].1.content, "NEAR price is $5"); + } + + #[tokio::test] + async fn message_tool_falls_back_to_ctx_user_when_owner_scope_absent() { + let (tool, gateway_captures, telegram_captures) = + message_tool_with_recording_channels().await; + + let mut ctx = crate::context::JobContext::with_user( + "interactive-chat-user", + "routine-job", + "price alert", + ); ctx.metadata = serde_json::json!({ "notify_channel": "telegram", }); let result = tool .execute(serde_json::json!({"content": "NEAR price is $5"}), &ctx) - .await; + .await + .expect( + "message tool should fall back to ctx.user_id when owner scope metadata is absent", + ); - assert!(result.is_err()); // safety: test-only assertion - let err = result.unwrap_err().to_string(); - let mentions_missing_target = err.contains("No target specified"); - assert!(!mentions_missing_target); // safety: test-only assertion - let mentions_missing_channel = err.contains("No channel specified"); - assert!(!mentions_missing_channel); // safety: test-only assertion + assert_eq!( + result.result.as_str(), + Some("Sent message to telegram:interactive-chat-user") + ); + assert!(gateway_captures.lock().await.is_empty()); + let telegram = telegram_captures.lock().await.clone(); + assert_eq!(telegram.len(), 1); + assert_eq!(telegram[0].0, "interactive-chat-user"); + assert_eq!(telegram[0].1.content, "NEAR price is $5"); } #[tokio::test] diff --git a/src/worker/job.rs b/src/worker/job.rs index 436a23ce..ba5d47b9 100644 --- a/src/worker/job.rs +++ b/src/worker/job.rs @@ -1438,6 +1438,9 @@ impl From for Result { #[cfg(test)] mod tests { + use std::sync::Arc; + + use crate::channels::ChannelManager; use crate::llm::ToolSelection; use super::*; @@ -1448,6 +1451,8 @@ mod tests { ToolCompletionResponse, }; use crate::safety::SafetyLayer; + use crate::testing::{BroadcastCapture, RecordingBroadcastChannel}; + use crate::tools::builtin::MessageTool; use crate::tools::{Tool, ToolError as ToolExecError, ToolOutput}; /// A test tool that sleeps for a configurable duration before returning. @@ -1539,6 +1544,20 @@ mod tests { Worker::new(job_id, deps) } + async fn make_worker_with_message_tool() + -> (Worker, Arc, BroadcastCapture, BroadcastCapture) { + let channel_manager = ChannelManager::new(); + let (gateway, gateway_captures) = RecordingBroadcastChannel::new("gateway"); + let (telegram, telegram_captures) = RecordingBroadcastChannel::new("telegram"); + channel_manager.add(Box::new(gateway)).await; + channel_manager.add(Box::new(telegram)).await; + + let message_tool = Arc::new(MessageTool::new(Arc::new(channel_manager))); + let worker = make_worker(vec![message_tool.clone()]).await; + + (worker, message_tool, gateway_captures, telegram_captures) + } + #[test] fn test_tool_selection_preserves_call_id() { let selection = ToolSelection { @@ -2147,4 +2166,50 @@ mod tests { assert_eq!(ctx.metadata, original); // safety: test } + + #[tokio::test] + async fn autonomous_message_tool_ignores_stale_gateway_context_when_routine_metadata_targets_telegram() + { + let (worker, message_tool, gateway_captures, telegram_captures) = + make_worker_with_message_tool().await; + + message_tool + .set_context( + Some("gateway".to_string()), + Some("stale-gateway-target".to_string()), + ) + .await; + + worker + .context_manager() + .update_context(worker.job_id, |ctx| { + ctx.user_id = "telegram".to_string(); + ctx.metadata = serde_json::json!({ + "notify_channel": "telegram", + "owner_id": "owner-scope", + }); + Ok::<(), String>(()) + }) + .await + .unwrap() // safety: test + .unwrap(); // safety: test + + let result = worker + .execute_tool( + "message", + &serde_json::json!({"content": "hello from routine"}), + ) + .await + .unwrap(); // safety: test + assert!( + result.contains("telegram:owner-scope"), + "expected telegram owner-scope routing, got: {result}" + ); + + assert!(gateway_captures.lock().await.is_empty()); + let telegram = telegram_captures.lock().await.clone(); + assert_eq!(telegram.len(), 1); + assert_eq!(telegram[0].0, "owner-scope"); + assert_eq!(telegram[0].1.content, "hello from routine"); + } }