Skip to content

[codex] Add context-preserving invoke stream #21

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
senamakel merged 2 commits into from
Jul 5, 2026
Merged

[codex] Add context-preserving invoke stream #21

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
285dfe7
Add context-preserving invoke stream
senamakel Jul 5, 2026
a788b85
Fix invoke stream listener cleanup
senamakel Jul 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 92 additions & 12 deletions src/harness/agent_loop/stream.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<dyn EventListener>,
}

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<AgentLoopResult>) -> AgentStreamItem {
match result {
Expand All @@ -83,9 +96,15 @@ fn terminal_item(result: Result<AgentLoopResult>) -> AgentStreamItem {
/// Drive-phase for the streaming state machine.
enum Phase<'a> {
/// The run future is still executing.
Running(Pin<Box<dyn Future<Output = Result<AgentLoopResult>> + 'a>>),
Running {
run_fut: Pin<Box<dyn Future<Output = Result<AgentLoopResult>> + '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,
}
Expand Down Expand Up @@ -113,11 +132,32 @@ impl<State: Send + Sync, Ctx: Send + Sync + 'static> AgentHarness<State, Ctx> {
input: Vec<Message>,
) -> impl futures::Stream<Item = AgentStreamItem> + '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<Ctx>,
input: Vec<Message>,
) -> impl futures::Stream<Item = AgentStreamItem> + '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<dyn EventListener> = Arc::new(ChannelListener { tx });
ctx.events.subscribe(listener.clone());

@chatgpt-codex-connector chatgpt-codex-connector Bot Jul 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Avoid starving runs on shared event sinks

When ctx.events is a long-lived/shared sink, this subscribes the stream listener before this run future has been polled. The unfold below uses a biased select that drains rx.recv() before polling run_fut, so events emitted by another active run or host component on the same sink can keep this stream yielding unrelated records and prevent this run from starting or completing. Consider isolating/filtering the listener to this run or ensuring the run future is polled even while the shared sink is busy.

Useful? React with 👍 / 👎.

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`
Expand All @@ -127,24 +167,43 @@ impl<State: Send + Sync, Ctx: Send + Sync + 'static> AgentHarness<State, Ctx> {
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 —
// and every sub-agent clone of the sink —
// 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)))
}
},
Expand All @@ -156,19 +215,40 @@ impl<State: Send + Sync, Ctx: Send + Sync + 'static> AgentHarness<State, Ctx> {
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,
}
Expand Down
55 changes: 54 additions & 1 deletion src/harness/agent_loop/test.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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<AgentStreamItem> = 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<AgentStreamItem> = 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();
Expand Down
23 changes: 23 additions & 0 deletions src/harness/events/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<dyn EventListener>) -> 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.
///
Expand Down