mirror of
https://github.com/openai/codex.git
synced 2026-11-13 21:30:47 +00:00
Sean Huang
edd1ae37ec
## Why Code Mode sessions always used client-side output handling, leaving server truncation settings and observation budgets unused. ## What changed - Honor each session's `output_truncation` setting in the runtime and both host transports, and advertise the bounded-output capability. - Apply `max_output_tokens` independently to execute, wait, and terminate responses in server mode. Keep client mode as the default when the setting is omitted. - Use bounded output collection and rendering in server mode, preserving head and tail output with omission notices. Include script errors in the output budget while retaining failure status and a separate diagnostic capped at 64 KiB. - Estimate audio tokens without the Tokio-dependent cache when rendering bounded output. ## Testing Add runtime and host regression tests for session mode selection, per-request budgets, output overflow, oversized images, zero-token audio output, diagnostic preservation, and termination under output backpressure. GitOrigin-RevId: 3d9661edcdd2dde31526bb92eb83fa0614f0fa88
767 lines
24 KiB
Rust
767 lines
24 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::AtomicBool;
|
|
use std::sync::atomic::Ordering;
|
|
use std::sync::mpsc as std_mpsc;
|
|
use std::time::Duration;
|
|
|
|
use codex_code_mode_protocol::CodeModeOutputTruncation;
|
|
use codex_code_mode_protocol::ExecuteRequest;
|
|
use codex_code_mode_protocol::host::BOUNDED_OUTPUT_MAX_IMAGE_BYTES;
|
|
use pretty_assertions::assert_eq;
|
|
use serde_json::Value as JsonValue;
|
|
use tokio::sync::mpsc;
|
|
use tokio::sync::oneshot;
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
use super::*;
|
|
use crate::runtime::RuntimeEventSender;
|
|
use crate::session_runtime::OutputItem;
|
|
|
|
#[derive(Default)]
|
|
struct TestHost {
|
|
commit_started: CancellationToken,
|
|
resume_commit: Option<CancellationToken>,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct RecordingHost {
|
|
notified: AtomicBool,
|
|
}
|
|
|
|
impl CellHost for TestHost {
|
|
async fn invoke_tool(
|
|
&self,
|
|
_invocation: CellToolCall,
|
|
_cancellation_token: CancellationToken,
|
|
) -> Result<JsonValue, String> {
|
|
Err("unexpected tool call".to_string())
|
|
}
|
|
|
|
async fn notify(
|
|
&self,
|
|
_call_id: String,
|
|
_text: String,
|
|
_cancellation_token: CancellationToken,
|
|
) -> Result<(), String> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn commit_completion(
|
|
&self,
|
|
_stored_value_writes: HashMap<String, Arc<JsonValue>>,
|
|
event: impl FnOnce() -> CellEvent + Send,
|
|
cell_state: Arc<CellState>,
|
|
) -> CompletionCommit {
|
|
self.commit_started.cancel();
|
|
if let Some(resume) = &self.resume_commit {
|
|
resume.cancelled().await;
|
|
}
|
|
cell_state.commit_completion(event, || {})
|
|
}
|
|
|
|
async fn closed(&self) {}
|
|
}
|
|
|
|
impl CellHost for RecordingHost {
|
|
async fn invoke_tool(
|
|
&self,
|
|
_invocation: CellToolCall,
|
|
_cancellation_token: CancellationToken,
|
|
) -> Result<JsonValue, String> {
|
|
Err("unexpected tool call".to_string())
|
|
}
|
|
|
|
async fn notify(
|
|
&self,
|
|
_call_id: String,
|
|
_text: String,
|
|
_cancellation_token: CancellationToken,
|
|
) -> Result<(), String> {
|
|
self.notified.store(true, Ordering::Release);
|
|
Ok(())
|
|
}
|
|
|
|
async fn commit_completion(
|
|
&self,
|
|
_stored_value_writes: HashMap<String, Arc<JsonValue>>,
|
|
event: impl FnOnce() -> CellEvent + Send,
|
|
cell_state: Arc<CellState>,
|
|
) -> CompletionCommit {
|
|
cell_state.commit_completion(event, || {})
|
|
}
|
|
|
|
async fn closed(&self) {}
|
|
}
|
|
|
|
struct CellActorHarness {
|
|
event_tx: RuntimeEventSender,
|
|
handle: CellHandle,
|
|
yield_signal: CancellationToken,
|
|
initial_event_rx: oneshot::Receiver<Result<CellEvent, CellError>>,
|
|
task: tokio::task::JoinHandle<()>,
|
|
runtime_control_rx: std_mpsc::Receiver<RuntimeControlCommand>,
|
|
_runtime_event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
|
|
}
|
|
|
|
fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarness {
|
|
spawn_cell_actor_harness_with_host(initial_observe_mode, Arc::new(TestHost::default()))
|
|
}
|
|
|
|
fn spawn_cell_actor_harness_with_host<H: CellHost>(
|
|
initial_observe_mode: ObserveMode,
|
|
host: Arc<H>,
|
|
) -> CellActorHarness {
|
|
spawn_cell_actor_harness_with_host_and_failure_handler(
|
|
initial_observe_mode,
|
|
host,
|
|
CodeModeOutputTruncation::Client,
|
|
/*task_failure_handler*/ None,
|
|
)
|
|
}
|
|
|
|
fn spawn_cell_actor_harness_with_host_and_failure_handler<H: CellHost>(
|
|
initial_observe_mode: ObserveMode,
|
|
host: Arc<H>,
|
|
output_truncation: CodeModeOutputTruncation,
|
|
task_failure_handler: Option<TaskFailureHandler>,
|
|
) -> CellActorHarness {
|
|
let (event_tx, event_rx) = runtime_event_channel(output_truncation);
|
|
let output = OutputCollector::new(output_truncation);
|
|
let (command_tx, command_rx) = mpsc::unbounded_channel();
|
|
let (initial_event_tx, initial_event_rx) = oneshot::channel();
|
|
let (runtime_event_tx, runtime_event_rx) = mpsc::unbounded_channel();
|
|
let (runtime_tx, _runtime_control_tx, runtime_terminate_handle) = spawn_runtime(
|
|
HashMap::new(),
|
|
ExecuteRequest {
|
|
tool_call_id: "call-1".to_string(),
|
|
enabled_tools: Vec::new(),
|
|
source: "await new Promise(() => {});".to_string(),
|
|
yield_time_ms: None,
|
|
max_output_tokens: None,
|
|
},
|
|
runtime_event_tx,
|
|
PendingRuntimeMode::PauseUntilResumed,
|
|
/*task_failure_handler*/ None,
|
|
)
|
|
.unwrap();
|
|
let (runtime_control_tx, runtime_control_rx) = std_mpsc::channel();
|
|
let cell_state = Arc::new(CellState::new(CancellationToken::new()));
|
|
let handle = CellHandle::new(command_tx, Arc::clone(&cell_state));
|
|
let yield_signal = CancellationToken::new();
|
|
let task = tokio::spawn(run_cell(
|
|
host,
|
|
CellContext {
|
|
runtime_tx,
|
|
runtime_control_tx,
|
|
runtime_terminate_handle,
|
|
cell_state,
|
|
output,
|
|
},
|
|
event_rx,
|
|
command_rx,
|
|
Observer {
|
|
observation: initial_observe_mode.with_yield_signal(yield_signal.clone()),
|
|
response_tx: initial_event_tx,
|
|
},
|
|
task_failure_handler,
|
|
));
|
|
|
|
CellActorHarness {
|
|
event_tx,
|
|
handle,
|
|
yield_signal,
|
|
initial_event_rx,
|
|
task,
|
|
runtime_control_rx,
|
|
_runtime_event_rx: runtime_event_rx,
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn termination_drains_a_send_already_blocked_by_backpressure() {
|
|
let harness = spawn_cell_actor_harness_with_host_and_failure_handler(
|
|
ObserveMode::YieldAfter(Duration::from_secs(60)),
|
|
Arc::new(TestHost::default()),
|
|
CodeModeOutputTruncation::Server,
|
|
/*task_failure_handler*/ None,
|
|
);
|
|
let (termination, capacity) = {
|
|
let RuntimeEventSender::Bounded(tx) = &harness.event_tx else {
|
|
panic!("server runtime must use a bounded sender");
|
|
};
|
|
// This test owns the V8-side events. The current-thread actor cannot run
|
|
// before our first await; fill its actual FIFO and register a sender.
|
|
for _ in 0..tx.max_capacity() {
|
|
tx.try_send(RuntimeEvent::ContentItem(OutputItem::InputText {
|
|
text: "queued".to_string(),
|
|
}))
|
|
.unwrap();
|
|
}
|
|
let mut send = std::pin::pin!(tx.send(RuntimeEvent::ContentItem(OutputItem::InputText {
|
|
text: "blocked".to_string(),
|
|
})));
|
|
assert!(futures::poll!(send.as_mut()).is_pending());
|
|
// terminate() synchronously cancels. No event was consumed between the
|
|
// Pending poll and cancellation; run_cell must drain while terminating.
|
|
let termination = harness.handle.terminate();
|
|
tokio::time::timeout(Duration::from_secs(1), send)
|
|
.await
|
|
.expect("termination must release the registered sender")
|
|
.expect("actor must drain the blocked item");
|
|
(termination, tx.max_capacity())
|
|
};
|
|
drop(harness.event_tx);
|
|
let event = tokio::time::timeout(Duration::from_secs(1), termination)
|
|
.await
|
|
.expect("termination must finish after draining")
|
|
.expect("termination event");
|
|
let mut expected = OutputCollector::new(CodeModeOutputTruncation::Server);
|
|
for _ in 0..capacity {
|
|
expected.push(OutputItem::InputText {
|
|
text: "queued".to_string(),
|
|
});
|
|
}
|
|
expected.push(OutputItem::InputText {
|
|
text: "blocked".to_string(),
|
|
});
|
|
assert_eq!(
|
|
event,
|
|
CellEvent::Terminated {
|
|
content_items: expected,
|
|
}
|
|
);
|
|
assert_eq!(harness.initial_event_rx.await.unwrap(), Ok(event));
|
|
tokio::time::timeout(Duration::from_secs(1), harness.task)
|
|
.await
|
|
.expect("actor must close")
|
|
.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn unexpected_runtime_thread_exit_is_reported_to_the_session_owner() {
|
|
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
|
|
let harness = spawn_cell_actor_harness_with_host_and_failure_handler(
|
|
ObserveMode::YieldAfter(Duration::from_secs(60)),
|
|
Arc::new(TestHost::default()),
|
|
CodeModeOutputTruncation::Client,
|
|
Some(Arc::new(move |reason| {
|
|
let _ = failure_tx.send(reason);
|
|
})),
|
|
);
|
|
drop(harness.event_tx);
|
|
|
|
assert_eq!(
|
|
tokio::time::timeout(Duration::from_secs(1), failure_rx.recv())
|
|
.await
|
|
.expect("runtime failure timeout")
|
|
.expect("runtime failure"),
|
|
"code-mode V8 runtime thread ended unexpectedly"
|
|
);
|
|
assert!(
|
|
harness
|
|
.initial_event_rx
|
|
.await
|
|
.expect("initial event")
|
|
.is_ok()
|
|
);
|
|
harness.task.await.expect("cell task");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn termination_during_completion_commit_preserves_output_and_discards_losing_error() {
|
|
for mode in [
|
|
CodeModeOutputTruncation::Client,
|
|
CodeModeOutputTruncation::Server,
|
|
] {
|
|
// The actor has received Result but cannot commit until the host releases
|
|
// the session store. Termination must decide which output is published.
|
|
let resume_commit = CancellationToken::new();
|
|
let host = Arc::new(TestHost {
|
|
resume_commit: Some(resume_commit.clone()),
|
|
..TestHost::default()
|
|
});
|
|
let harness = spawn_cell_actor_harness_with_host_and_failure_handler(
|
|
ObserveMode::YieldAfter(Duration::from_secs(60)),
|
|
Arc::clone(&host),
|
|
mode,
|
|
/*task_failure_handler*/ None,
|
|
);
|
|
// Simulate V8's dedicated thread, where server-mode sends can block
|
|
// while the event queue is full.
|
|
let event_tx = harness.event_tx.clone();
|
|
tokio::task::spawn_blocking(move || {
|
|
event_tx
|
|
.blocking_send(RuntimeEvent::ContentItem(OutputItem::InputText {
|
|
text: "output that must survive termination".to_string(),
|
|
}))
|
|
.expect("output");
|
|
event_tx
|
|
.blocking_send(RuntimeEvent::Result {
|
|
stored_value_writes: HashMap::new(),
|
|
error_text: Some("uncommitted error".repeat(128 * 1024)),
|
|
})
|
|
.expect("completion");
|
|
})
|
|
.await
|
|
.expect("runtime producer");
|
|
tokio::time::timeout(Duration::from_secs(5), host.commit_started.cancelled())
|
|
.await
|
|
.expect("commit reached");
|
|
let termination = harness.handle.terminate();
|
|
resume_commit.cancel();
|
|
let mut expected_output = OutputCollector::new(mode);
|
|
expected_output.push(OutputItem::InputText {
|
|
text: "output that must survive termination".to_string(),
|
|
});
|
|
let expected = Ok(CellEvent::Terminated {
|
|
content_items: expected_output,
|
|
});
|
|
assert_eq!(termination.await, expected, "{mode:?}");
|
|
assert_eq!(harness.initial_event_rx.await.expect("observer"), expected);
|
|
harness.task.await.expect("cell task");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn runtime_thread_panic_remains_a_cell_error_without_owner_supervision() {
|
|
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(60)));
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::ThreadPanicked)
|
|
.expect("runtime panic event");
|
|
drop(harness.event_tx);
|
|
|
|
assert_eq!(
|
|
harness.initial_event_rx.await.expect("initial event"),
|
|
Ok(CellEvent::Completed {
|
|
content_items: OutputCollector::Client(Vec::new()),
|
|
result_buffered_for: Duration::ZERO,
|
|
error_text: Some("exec runtime ended unexpectedly".to_string()),
|
|
})
|
|
);
|
|
harness.task.await.expect("cell task");
|
|
}
|
|
|
|
async fn wait_for_notification(host: &RecordingHost) {
|
|
tokio::time::timeout(Duration::from_secs(1), async {
|
|
while !host.notified.load(Ordering::Acquire) {
|
|
tokio::task::yield_now().await;
|
|
}
|
|
})
|
|
.await
|
|
.expect("notification barrier timed out");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn yield_timer_preempts_buffered_runtime_output() {
|
|
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::ZERO));
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::Started)
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::ContentItem(OutputItem::InputText {
|
|
text: "queued output".to_string(),
|
|
}))
|
|
.unwrap();
|
|
|
|
assert_eq!(
|
|
harness.initial_event_rx.await.unwrap(),
|
|
Ok(CellEvent::Yielded {
|
|
content_items: OutputCollector::Client(Vec::new()),
|
|
})
|
|
);
|
|
|
|
let termination = harness.handle.terminate();
|
|
drop(harness.event_tx);
|
|
assert_eq!(
|
|
termination.await,
|
|
Ok(CellEvent::Terminated {
|
|
content_items: OutputCollector::Client(vec![OutputItem::InputText {
|
|
text: "queued output".to_string(),
|
|
}]),
|
|
})
|
|
);
|
|
harness.task.await.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn queued_termination_preempts_unobserved_runtime_completion() {
|
|
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::ZERO));
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::Started)
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::YieldRequested)
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::ContentItem(OutputItem::InputText {
|
|
text: "queued output".to_string(),
|
|
}))
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::Result {
|
|
stored_value_writes: HashMap::new(),
|
|
error_text: None,
|
|
})
|
|
.unwrap();
|
|
harness.yield_signal.cancel();
|
|
let termination = harness.handle.terminate();
|
|
|
|
let terminated = Ok(CellEvent::Terminated {
|
|
content_items: OutputCollector::Client(vec![OutputItem::InputText {
|
|
text: "queued output".to_string(),
|
|
}]),
|
|
});
|
|
assert_eq!(termination.await, terminated.clone());
|
|
assert_eq!(harness.initial_event_rx.await.unwrap(), terminated);
|
|
harness.task.await.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn observation_dropped_before_dequeue_does_not_consume_output() {
|
|
let host = Arc::new(RecordingHost::default());
|
|
let harness = spawn_cell_actor_harness_with_host(
|
|
ObserveMode::YieldAfter(Duration::from_secs(60)),
|
|
Arc::clone(&host),
|
|
);
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::YieldRequested)
|
|
.unwrap();
|
|
assert!(harness.initial_event_rx.await.unwrap().is_ok());
|
|
|
|
drop(
|
|
harness
|
|
.handle
|
|
.observe(ObserveMode::YieldAfter(Duration::from_secs(60))),
|
|
);
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::ContentItem(OutputItem::InputText {
|
|
text: "survives pre-dequeue cancellation".to_string(),
|
|
}))
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::YieldRequested)
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::Notify {
|
|
call_id: "after-dropped-command".to_string(),
|
|
text: "barrier".to_string(),
|
|
})
|
|
.unwrap();
|
|
wait_for_notification(&host).await;
|
|
|
|
assert_eq!(
|
|
harness
|
|
.handle
|
|
.observe(ObserveMode::YieldAfter(Duration::ZERO))
|
|
.await,
|
|
Ok(CellEvent::Yielded {
|
|
content_items: OutputCollector::Client(vec![OutputItem::InputText {
|
|
text: "survives pre-dequeue cancellation".to_string(),
|
|
}]),
|
|
})
|
|
);
|
|
|
|
let termination = harness.handle.terminate();
|
|
drop(harness.event_tx);
|
|
assert_eq!(
|
|
termination.await,
|
|
Ok(CellEvent::Terminated {
|
|
content_items: OutputCollector::Client(Vec::new()),
|
|
})
|
|
);
|
|
harness.task.await.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn dropped_yield_observer_preserves_output_for_the_next_observation() {
|
|
let host = Arc::new(RecordingHost::default());
|
|
let harness = spawn_cell_actor_harness_with_host(
|
|
ObserveMode::YieldAfter(Duration::from_secs(60)),
|
|
Arc::clone(&host),
|
|
);
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::YieldRequested)
|
|
.unwrap();
|
|
assert!(harness.initial_event_rx.await.unwrap().is_ok());
|
|
|
|
let dropped_observation = harness
|
|
.handle
|
|
.observe(ObserveMode::YieldAfter(Duration::from_secs(60)));
|
|
assert_eq!(
|
|
harness
|
|
.handle
|
|
.observe(ObserveMode::YieldAfter(Duration::ZERO))
|
|
.await,
|
|
Err(CellError::Busy)
|
|
);
|
|
drop(dropped_observation);
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::ContentItem(OutputItem::InputText {
|
|
text: "survives active cancellation".to_string(),
|
|
}))
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::YieldRequested)
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::Notify {
|
|
call_id: "after-dropped-observer".to_string(),
|
|
text: "barrier".to_string(),
|
|
})
|
|
.unwrap();
|
|
wait_for_notification(&host).await;
|
|
|
|
assert_eq!(
|
|
harness
|
|
.handle
|
|
.observe(ObserveMode::YieldAfter(Duration::ZERO))
|
|
.await,
|
|
Ok(CellEvent::Yielded {
|
|
content_items: OutputCollector::Client(vec![OutputItem::InputText {
|
|
text: "survives active cancellation".to_string(),
|
|
}]),
|
|
})
|
|
);
|
|
|
|
let termination = harness.handle.terminate();
|
|
drop(harness.event_tx);
|
|
assert_eq!(
|
|
termination.await,
|
|
Ok(CellEvent::Terminated {
|
|
content_items: OutputCollector::Client(Vec::new()),
|
|
})
|
|
);
|
|
harness.task.await.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observation() {
|
|
let host = Arc::new(RecordingHost::default());
|
|
let harness = spawn_cell_actor_harness_with_host(
|
|
ObserveMode::YieldAfter(Duration::from_secs(60)),
|
|
Arc::clone(&host),
|
|
);
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::YieldRequested)
|
|
.unwrap();
|
|
assert!(harness.initial_event_rx.await.unwrap().is_ok());
|
|
|
|
let dropped_observation = harness.handle.observe(ObserveMode::PendingFrontier);
|
|
assert_eq!(
|
|
harness.handle.observe(ObserveMode::PendingFrontier).await,
|
|
Err(CellError::Busy)
|
|
);
|
|
drop(dropped_observation);
|
|
let item = OutputItem::InputText {
|
|
text: "before frontier".to_owned(),
|
|
};
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::ContentItem(item.clone()))
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::ToolCall {
|
|
id: "tool-1".to_string(),
|
|
name: codex_protocol::ToolName {
|
|
name: "echo".to_string(),
|
|
namespace: None,
|
|
},
|
|
kind: codex_code_mode_protocol::CodeModeToolKind::Function,
|
|
input: Some(serde_json::json!({})),
|
|
})
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::Pending)
|
|
.unwrap();
|
|
harness
|
|
.event_tx
|
|
.blocking_send(RuntimeEvent::Notify {
|
|
call_id: "after-dropped-pending".to_string(),
|
|
text: "barrier".to_string(),
|
|
})
|
|
.unwrap();
|
|
wait_for_notification(&host).await;
|
|
|
|
assert_eq!(
|
|
harness.handle.observe(ObserveMode::PendingFrontier).await,
|
|
Ok(CellEvent::Pending {
|
|
content_items: OutputCollector::Client(vec![item]),
|
|
pending_tool_call_ids: vec!["tool-1".to_string()],
|
|
})
|
|
);
|
|
assert!(matches!(
|
|
harness.runtime_control_rx.try_recv(),
|
|
Err(std_mpsc::TryRecvError::Empty)
|
|
));
|
|
|
|
let termination = harness.handle.terminate();
|
|
drop(harness.event_tx);
|
|
assert_eq!(
|
|
termination.await,
|
|
Ok(CellEvent::Terminated {
|
|
content_items: OutputCollector::Client(Vec::new()),
|
|
})
|
|
);
|
|
harness.task.await.unwrap();
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
async fn only_the_first_termination_claims_a_buffered_completion() {
|
|
let cell_state = CellState::new(CancellationToken::new());
|
|
let completion = CellEvent::Completed {
|
|
content_items: OutputCollector::Client(Vec::new()),
|
|
result_buffered_for: Duration::ZERO,
|
|
error_text: None,
|
|
};
|
|
assert_eq!(
|
|
cell_state.commit_completion(|| completion.clone(), || {}),
|
|
CompletionCommit::Committed
|
|
);
|
|
assert!(matches!(
|
|
cell_state.deliver_completion(/*response_tx*/ None),
|
|
CompletionDelivery::Buffered
|
|
));
|
|
|
|
let first_termination = cell_state.request_termination();
|
|
assert_eq!(
|
|
cell_state.request_termination().await,
|
|
Err(CellError::AlreadyTerminating)
|
|
);
|
|
assert_eq!(first_termination.await, Ok(completion.clone()));
|
|
assert_eq!(
|
|
cell_state.finish_termination(CellEvent::Terminated {
|
|
content_items: OutputCollector::Client(Vec::new()),
|
|
}),
|
|
Some(completion)
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn termination_claim_prevents_stored_value_commit() {
|
|
let cell_state = CellState::new(CancellationToken::new());
|
|
let termination = cell_state.request_termination();
|
|
let mut commit_ran = false;
|
|
assert_eq!(
|
|
cell_state.commit_completion(
|
|
|| panic!("terminated cell must not construct completion output"),
|
|
|| commit_ran = true
|
|
),
|
|
CompletionCommit::Rejected
|
|
);
|
|
assert!(!commit_ran);
|
|
|
|
let terminated = CellEvent::Terminated {
|
|
content_items: OutputCollector::Client(Vec::new()),
|
|
};
|
|
assert_eq!(
|
|
cell_state.finish_termination(terminated.clone()),
|
|
Some(terminated.clone())
|
|
);
|
|
assert_eq!(termination.await, Ok(terminated));
|
|
}
|
|
|
|
#[tokio::test(start_paused = true)]
|
|
async fn failed_completion_delivery_rebuffers_the_event() {
|
|
let cell_state = CellState::new(CancellationToken::new());
|
|
let completion = |result_buffered_for| CellEvent::Completed {
|
|
content_items: OutputCollector::Client(vec![OutputItem::InputText {
|
|
text: "done".to_string(),
|
|
}]),
|
|
result_buffered_for,
|
|
error_text: None,
|
|
};
|
|
assert_eq!(
|
|
cell_state.commit_completion(|| completion(Duration::ZERO), || {}),
|
|
CompletionCommit::Committed
|
|
);
|
|
let (response_tx, response_rx) = oneshot::channel();
|
|
drop(response_rx);
|
|
assert!(matches!(
|
|
cell_state.deliver_completion(Some(response_tx)),
|
|
CompletionDelivery::Buffered
|
|
));
|
|
assert!(cell_state.accepting_observations());
|
|
|
|
let interval = Duration::from_secs(5);
|
|
tokio::time::advance(interval).await;
|
|
let (response_tx, response_rx) = oneshot::channel();
|
|
drop(response_rx);
|
|
assert!(matches!(
|
|
cell_state.route_observation(response_tx),
|
|
ObservationDelivery::Buffered
|
|
));
|
|
|
|
tokio::time::advance(interval).await;
|
|
let (response_tx, mut response_rx) = oneshot::channel();
|
|
assert!(matches!(
|
|
cell_state.route_observation(response_tx),
|
|
ObservationDelivery::Delivered
|
|
));
|
|
assert_eq!(response_rx.try_recv(), Ok(Ok(completion(2 * interval))));
|
|
}
|
|
|
|
#[test]
|
|
fn a_failed_yield_restores_image_allocation_and_omission_order() {
|
|
let mut output = OutputCollector::new(CodeModeOutputTruncation::Server);
|
|
let image_url = "i".repeat(BOUNDED_OUTPUT_MAX_IMAGE_BYTES);
|
|
let pointer = image_url.as_ptr();
|
|
output.push(OutputItem::InputImage {
|
|
image_url,
|
|
detail: None,
|
|
});
|
|
output.push(OutputItem::InputImage {
|
|
image_url: "excess".to_string(),
|
|
detail: None,
|
|
});
|
|
let (tx, rx) = oneshot::channel();
|
|
drop(rx);
|
|
send_yield(&mut output, tx);
|
|
|
|
output.push(OutputItem::InputText {
|
|
text: "after".to_string(),
|
|
});
|
|
let (tx, mut rx) = oneshot::channel();
|
|
send_yield(&mut output, tx);
|
|
let CellEvent::Yielded { content_items } = rx.try_recv().unwrap().unwrap() else {
|
|
panic!("expected a yield");
|
|
};
|
|
let items = crate::output::render_output(
|
|
content_items.into_collected_output(),
|
|
/*max_output_tokens*/ None,
|
|
);
|
|
let [
|
|
OutputItem::InputImage { image_url, .. },
|
|
OutputItem::InputText { text: notice },
|
|
OutputItem::InputText { text },
|
|
] = items.as_slice()
|
|
else {
|
|
panic!("expected image, omission notice, and subsequent text");
|
|
};
|
|
assert_eq!(image_url.as_ptr(), pointer);
|
|
assert!(notice.contains("1 images"), "{notice}");
|
|
assert_eq!(text, "after");
|
|
assert_eq!(
|
|
output.take(),
|
|
OutputCollector::new(CodeModeOutputTruncation::Server)
|
|
);
|
|
}
|