diff --git a/src/harness/agent_loop/stream.rs b/src/harness/agent_loop/stream.rs index cb11970..c52e88e 100644 --- a/src/harness/agent_loop/stream.rs +++ b/src/harness/agent_loop/stream.rs @@ -30,7 +30,7 @@ use std::sync::Arc; use crate::error::Result; use crate::harness::context::{RunConfig, RunContext}; -use crate::harness::events::{EventListener, EventRecord}; +use crate::harness::events::{EventListener, EventRecord, EventSink}; use crate::harness::message::Message; use crate::harness::middleware::AgentRun; use crate::harness::runtime::AgentHarness; @@ -72,6 +72,19 @@ impl EventListener for ChannelListener { } } +/// Removes the one-shot channel listener from the run's event sink when the +/// returned stream is exhausted or dropped early. +struct ChannelListenerGuard { + events: EventSink, + listener: Arc, +} + +impl Drop for ChannelListenerGuard { + fn drop(&mut self) { + let _ = self.events.unsubscribe(&self.listener); + } +} + /// Maps a finished run result onto its terminal [`AgentStreamItem`]. fn terminal_item(result: Result) -> AgentStreamItem { match result { @@ -83,9 +96,15 @@ fn terminal_item(result: Result) -> AgentStreamItem { /// Drive-phase for the streaming state machine. enum Phase<'a> { /// The run future is still executing. - Running(Pin> + 'a>>), + Running { + run_fut: Pin> + 'a>>, + listener_guard: ChannelListenerGuard, + }, /// The run has finished; drain any buffered events, then emit `terminal`. - Draining(AgentStreamItem), + Draining { + terminal: AgentStreamItem, + listener_guard: ChannelListenerGuard, + }, /// Terminal item already emitted; the stream is exhausted. Done, } @@ -113,11 +132,32 @@ impl AgentHarness { input: Vec, ) -> impl futures::Stream + 'a { let ctx = RunContext::new(config, ctx_data); + self.invoke_stream_in_context(state, ctx, input) + } + + /// Runs the agent loop while streaming every emitted event to the caller, + /// using a caller-supplied [`RunContext`]. + /// + /// This is the context-preserving counterpart of + /// [`AgentHarness::invoke_stream`]. Use it when the caller has already + /// attached cancellation, events, stores, workspace metadata, or steering to + /// the run context and still wants caller-consumable stream items. + pub fn invoke_stream_in_context<'a>( + &'a self, + state: &'a State, + ctx: RunContext, + input: Vec, + ) -> impl futures::Stream + 'a { let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); // Subscribe before driving so no event (starting with `RunStarted`) is // missed. The listener rides the run's `EventSink`, which sub-agents // clone, so their lifecycle events reach this stream too. - ctx.events.subscribe(Arc::new(ChannelListener { tx })); + let listener: Arc = Arc::new(ChannelListener { tx }); + ctx.events.subscribe(listener.clone()); + let listener_guard = ChannelListenerGuard { + events: ctx.events.clone(), + listener, + }; // `invoke_streaming_in_context_with_status` is the public wrapper over // the shared `drive(.., streaming = true)` path; it drives *our* `ctx` @@ -127,17 +167,35 @@ impl AgentHarness { Box::pin(self.invoke_streaming_in_context_with_status(state, ctx, input)); futures::stream::unfold( - (Phase::Running(run_fut), rx), + ( + Phase::Running { + run_fut, + listener_guard, + }, + rx, + ), |(phase, mut rx)| async move { match phase { - Phase::Running(mut run_fut) => { + Phase::Running { + mut run_fut, + listener_guard, + } => { tokio::select! { biased; // Prefer draining ready events so the consumer sees // fine-grained progress rather than a late burst. maybe = rx.recv() => match maybe { Some(record) => { - Some((AgentStreamItem::Event(record), (Phase::Running(run_fut), rx))) + Some(( + AgentStreamItem::Event(record), + ( + Phase::Running { + run_fut, + listener_guard, + }, + rx, + ), + )) } None => { // All senders dropped (the run's context — @@ -145,6 +203,7 @@ impl AgentHarness { // is gone): the run is finishing. Await it // for the terminal item. let terminal = terminal_item(run_fut.await); + drop(listener_guard); Some((terminal, (Phase::Done, rx))) } }, @@ -156,19 +215,40 @@ impl AgentHarness { match rx.try_recv() { Ok(record) => Some(( AgentStreamItem::Event(record), - (Phase::Draining(terminal), rx), + ( + Phase::Draining { + terminal, + listener_guard, + }, + rx, + ), )), - Err(_) => Some((terminal, (Phase::Done, rx))), + Err(_) => { + drop(listener_guard); + Some((terminal, (Phase::Done, rx))) + } } } } } - Phase::Draining(terminal) => match rx.try_recv() { + Phase::Draining { + terminal, + listener_guard, + } => match rx.try_recv() { Ok(record) => Some(( AgentStreamItem::Event(record), - (Phase::Draining(terminal), rx), + ( + Phase::Draining { + terminal, + listener_guard, + }, + rx, + ), )), - Err(_) => Some((terminal, (Phase::Done, rx))), + Err(_) => { + drop(listener_guard); + Some((terminal, (Phase::Done, rx))) + } }, Phase::Done => None, } diff --git a/src/harness/agent_loop/test.rs b/src/harness/agent_loop/test.rs index a179486..1035d81 100644 --- a/src/harness/agent_loop/test.rs +++ b/src/harness/agent_loop/test.rs @@ -15,7 +15,7 @@ use serde_json::json; use super::AgentStreamItem; use crate::error::{Result, TinyAgentsError}; use crate::harness::context::{RunConfig, RunContext}; -use crate::harness::events::AgentEvent; +use crate::harness::events::{AgentEvent, EventSink}; use crate::harness::limits::RunLimits; use crate::harness::message::{AssistantMessage, ContentBlock, Message, MessageDelta}; use crate::harness::middleware::{ @@ -2030,6 +2030,59 @@ async fn invoke_stream_yields_events_then_completed() { ); } +#[tokio::test] +async fn invoke_stream_in_context_preserves_caller_context() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", Arc::new(MockModel::constant("hello there"))); + let ctx = RunContext::new( + RunConfig::new("caller-run").with_thread("caller-thread"), + (), + ); + + let items: Vec = harness + .invoke_stream_in_context(&(), ctx, vec![Message::user("hi")]) + .collect() + .await; + let (events, terminal) = collect_stream(items).await; + + let started = events + .iter() + .find_map(|event| match event { + AgentEvent::RunStarted { run_id, thread_id } => Some((run_id, thread_id)), + _ => None, + }) + .expect("RunStarted event"); + assert_eq!(started.0.as_str(), "caller-run"); + assert_eq!( + started.1.as_ref().map(|thread| thread.as_str()), + Some("caller-thread") + ); + match terminal { + AgentStreamItem::Completed(run) => { + assert_eq!(run.text().as_deref(), Some("hello there")); + } + other => panic!("expected Completed terminal, got {other:?}"), + } +} + +#[tokio::test] +async fn invoke_stream_in_context_unsubscribes_channel_listener() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", Arc::new(MockModel::constant("hello there"))); + let events = EventSink::new(); + let ctx = RunContext::new(RunConfig::new("shared-events-run"), ()).with_events(events.clone()); + + assert_eq!(events.listener_count(), 0); + let items: Vec = harness + .invoke_stream_in_context(&(), ctx, vec![Message::user("hi")]) + .collect() + .await; + let (_events, terminal) = collect_stream(items).await; + + assert!(matches!(terminal, AgentStreamItem::Completed(_))); + assert_eq!(events.listener_count(), 0); +} + #[tokio::test] async fn invoke_stream_surfaces_tool_lifecycle() { let mut harness: AgentHarness<()> = AgentHarness::new(); diff --git a/src/harness/events/mod.rs b/src/harness/events/mod.rs index 828686b..b0086d5 100644 --- a/src/harness/events/mod.rs +++ b/src/harness/events/mod.rs @@ -82,6 +82,29 @@ impl EventSink { inner.listeners.push(listener); } + /// Removes a previously subscribed listener. + /// + /// This is primarily used by one-shot stream adapters that subscribe a + /// transient listener to a long-lived shared sink. Returns `true` when the + /// listener was present. + pub fn unsubscribe(&self, listener: &Arc) -> bool { + let mut inner = self.inner.lock().expect("EventSink lock poisoned"); + let before = inner.listeners.len(); + inner + .listeners + .retain(|candidate| !Arc::ptr_eq(candidate, listener)); + inner.listeners.len() != before + } + + #[cfg(test)] + pub(crate) fn listener_count(&self) -> usize { + self.inner + .lock() + .expect("EventSink lock poisoned") + .listeners + .len() + } + /// Emits an event, assigning a monotonic [`EventId`] and offset, then /// notifying all registered listeners in insertion order. ///