#[cfg(test)]
use super::{
dispatcher::ServiceDispatcher,
protocol::{ASSISTANT_DELTA_EVENT, TURN_TERMINAL_EVENT, terminal_text_fits},
};
use super::{
dispatcher::{RequestIdGuard, ServiceOutbound},
protocol::{
ResponsePayload, ServiceErrorCode, ServiceEvent, ServiceMessage, ServiceRequest,
ServiceResponse, TurnCancelResult, TurnStartResult, TurnTerminalStatus,
},
runtime::ServiceRuntime,
};
use crate::{
agent::{self, AgentOutputSink},
cancellation::is_run_canceled,
output::OutputEvent,
sessions::{Session, SessionEventKind, TurnStatus, TurnStatusPayload},
};
use anyhow::Result;
use crossbeam_channel::Sender;
use std::{
collections::HashMap,
sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
},
thread::{self, JoinHandle},
};
use uuid::Uuid;
pub(crate) const MAX_ACTIVE_TURNS: usize = 4;
pub(crate) const INPUT_QUEUE_CAPACITY: usize = 32;
pub(crate) const TURN_EVENT_QUEUE_CAPACITY: usize = 128;
pub(crate) enum TurnWorkerMessage {
Configuration(super::dispatcher::ServiceOutbound),
Persisting {
turn_id: String,
},
Delta {
turn_id: String,
event: ServiceEvent,
},
Terminal {
turn_id: String,
event: ServiceEvent,
},
}
pub(super) struct TurnWorkerJob {
pub(super) runtime: Arc<ServiceRuntime>,
pub(super) settings_captured_at_admission: bool,
pub(super) maximum_sequence: u64,
pub(super) session: Session,
pub(super) prompt: String,
pub(super) turn_id: String,
pub(super) request_id: String,
pub(super) cancellation: Arc<AtomicBool>,
pub(super) sender: Sender<TurnWorkerMessage>,
pub(super) baseline_replay_generation: u64,
}
pub(super) type TurnWorker = Arc<dyn Fn(TurnWorkerJob) + Send + Sync + 'static>;
struct ActiveTurn {
request_id_guard: Option<RequestIdGuard>,
session_id: String,
cancellation: Arc<AtomicBool>,
join_handle: Option<JoinHandle<()>>,
terminal_seen: bool,
}
pub(crate) struct ServiceTurnManager {
runtime: Arc<ServiceRuntime>,
sender: Sender<TurnWorkerMessage>,
worker: TurnWorker,
active: HashMap<String, ActiveTurn>,
sessions: super::sessions::ServiceSessions,
}
impl ServiceTurnManager {
pub(crate) fn new(runtime: Arc<ServiceRuntime>, sender: Sender<TurnWorkerMessage>) -> Self {
Self {
sessions: super::sessions::ServiceSessions::new(Arc::clone(&runtime)),
runtime,
sender,
worker: Arc::new(run_turn_worker),
active: HashMap::new(),
}
}
#[cfg(test)]
fn new_with_worker(
runtime: Arc<ServiceRuntime>,
sender: Sender<TurnWorkerMessage>,
worker: TurnWorker,
) -> Self {
Self {
sessions: super::sessions::ServiceSessions::new(Arc::clone(&runtime)),
runtime,
sender,
worker,
active: HashMap::new(),
}
}
#[cfg(test)]
pub(super) fn set_worker(&mut self, worker: TurnWorker) {
self.worker = worker;
}
pub(super) fn claim_session(&mut self, id: &str) -> Result<(), ServiceErrorCode> {
self.sessions.claim(id)
}
pub(super) fn detach_session(&mut self, id: &str) {
self.sessions.detach(id, self.session_is_active(id));
}
#[cfg(test)]
pub(crate) fn start(
&mut self,
request: ServiceRequest,
guard: RequestIdGuard,
initialized: bool,
) -> ServiceOutbound {
self.start_with_settings_capture(request, guard, initialized, None, u64::MAX)
}
pub(crate) fn start_with_settings_capture(
&mut self,
request: ServiceRequest,
guard: RequestIdGuard,
initialized: bool,
prepared_runtime: Option<Arc<ServiceRuntime>>,
maximum_sequence: u64,
) -> ServiceOutbound {
let capture_settings_at_admission = prepared_runtime.is_some();
let identity = super::protocol::RequestIdentity::from_request(&request);
if !initialized {
return guarded_error(identity, guard, ServiceErrorCode::NotInitialized);
}
let params = match serde_json::from_value::<super::protocol::TurnStartParams>(
request.payload.clone(),
) {
Ok(params) => params,
Err(_) => return guarded_error(identity, guard, ServiceErrorCode::InvalidPayload),
};
if params.prompt.trim().is_empty() {
return guarded_error(identity, guard, ServiceErrorCode::InvalidPayload);
}
let max_active = if capture_settings_at_admission {
16
} else {
MAX_ACTIVE_TURNS
};
if self.active.len() >= max_active {
return guarded_error(identity, guard, ServiceErrorCode::TooManyActiveTurns);
}
let runtime = prepared_runtime.unwrap_or_else(|| Arc::clone(&self.runtime));
let (session, baseline_replay_generation) =
match self.session_for_request(request.session_id.as_deref()) {
Ok(session) => session,
Err(code) => return guarded_error(identity, guard, code),
};
let session_id = session.id().to_string();
if self.session_is_active(&session_id) {
return guarded_error(identity, guard, ServiceErrorCode::SessionBusy);
}
let turn_id = Uuid::new_v4().to_string();
let cancellation = Arc::new(AtomicBool::new(false));
let job = TurnWorkerJob {
runtime,
settings_captured_at_admission: capture_settings_at_admission,
session,
prompt: params.prompt,
turn_id: turn_id.clone(),
request_id: request.request_id.clone(),
cancellation: Arc::clone(&cancellation),
sender: self.sender.clone(),
baseline_replay_generation,
maximum_sequence,
};
let worker = Arc::clone(&self.worker);
let join_handle = match thread::Builder::new()
.name("magi-service-turn".to_string())
.spawn(move || worker(job))
{
Ok(handle) => handle,
Err(_) => {
self.sessions.finish_turn(&session_id);
return guarded_error(identity, guard, ServiceErrorCode::InternalError);
}
};
self.active.insert(
turn_id.clone(),
ActiveTurn {
request_id_guard: Some(guard),
session_id: session_id.clone(),
cancellation,
join_handle: Some(join_handle),
terminal_seen: false,
},
);
let response = ServiceResponse::success_with_identity(
request.request_id.clone(),
Some(session_id.clone()),
Some(request.method.clone()),
ResponsePayload::TurnStart(TurnStartResult {
turn_id: turn_id.clone(),
session_id: session_id.clone(),
status: "accepted".to_string(),
}),
);
let started = ServiceEvent::turn_started(request.request_id, session_id, turn_id);
ServiceOutbound::unguarded(vec![
ServiceMessage::response(response),
ServiceMessage::Event(started),
])
}
pub(crate) fn cancel(
&mut self,
request: ServiceRequest,
guard: RequestIdGuard,
initialized: bool,
) -> ServiceOutbound {
let identity = super::protocol::RequestIdentity::from_request(&request);
if !initialized {
return guarded_error(identity, guard, ServiceErrorCode::NotInitialized);
}
let params = match serde_json::from_value::<super::protocol::TurnCancelParams>(
request.payload.clone(),
) {
Ok(params) => params,
Err(_) => return guarded_error(identity, guard, ServiceErrorCode::InvalidPayload),
};
if let Err(code) = super::protocol::validate_service_id(¶ms.turn_id) {
return guarded_error(identity, guard, code);
}
let Some(active) = self.active.get(¶ms.turn_id) else {
return guarded_error(identity, guard, ServiceErrorCode::UnknownTurn);
};
active.cancellation.store(true, Ordering::SeqCst);
let response = ServiceResponse::success_with_identity(
request.request_id,
Some(active.session_id.clone()),
Some(request.method),
ResponsePayload::TurnCancel(TurnCancelResult {
turn_id: params.turn_id,
status: "cancellation_requested".to_string(),
}),
);
ServiceOutbound::guarded(vec![ServiceMessage::response(response)], guard)
}
pub(crate) fn handle_worker_message(
&mut self,
message: TurnWorkerMessage,
) -> Option<(String, ServiceEvent, bool)> {
let (turn_id, event, is_terminal) = match message {
TurnWorkerMessage::Delta { turn_id, event } => (turn_id, event, false),
TurnWorkerMessage::Terminal { turn_id, event } => (turn_id, event, true),
TurnWorkerMessage::Configuration(_) | TurnWorkerMessage::Persisting { .. } => {
return None;
}
};
let active = self.active.get_mut(&turn_id)?;
if active.terminal_seen {
return None;
}
if is_terminal {
active.terminal_seen = true;
}
Some((turn_id, event, is_terminal))
}
pub(super) fn worker_finished(&self, turn_id: &str) -> bool {
self.active.get(turn_id).is_some_and(|active| {
active
.join_handle
.as_ref()
.is_none_or(JoinHandle::is_finished)
})
}
pub(super) fn release_request_guard(&mut self, turn_id: &str) {
if let Some(active) = self.active.get_mut(turn_id) {
active.request_id_guard = None;
}
}
pub(crate) fn finish_output(&mut self, turn_id: &str) {
let Some(mut active) = self.active.remove(turn_id) else {
return;
};
if let Some(join_handle) = active.join_handle.take() {
let _ = join_handle.join();
}
drop(active.request_id_guard.take());
self.sessions.finish_turn(&active.session_id);
}
pub(crate) fn cancel_all(&self) {
for active in self.active.values() {
active.cancellation.store(true, Ordering::SeqCst);
}
}
pub(crate) fn join_all(&mut self) {
let active = std::mem::take(&mut self.active);
for (_, mut active) in active {
active.cancellation.store(true, Ordering::SeqCst);
drop(active.request_id_guard.take());
if let Some(join_handle) = active.join_handle.take() {
let _ = join_handle.join();
}
}
self.sessions.clear();
}
pub(crate) fn is_empty(&self) -> bool {
self.active.is_empty()
}
fn session_is_active(&self, session_id: &str) -> bool {
self.active
.values()
.any(|turn| turn.session_id == session_id)
}
pub(crate) fn dispatch_session(
&mut self,
request: ServiceRequest,
guard: RequestIdGuard,
initialized: bool,
) -> ServiceOutbound {
let identity = super::protocol::RequestIdentity::from_request(&request);
if !initialized {
return guarded_error(identity, guard, ServiceErrorCode::NotInitialized);
}
let active = request
.session_id
.as_ref()
.is_some_and(|id| self.session_is_active(id));
let payload = match self.sessions.operation(&request, active) {
Ok(payload) => payload,
Err(code) => return guarded_error(identity, guard, code),
};
if request.method == "session.close" && active {
for turn in self.active.values() {
if Some(turn.session_id.as_str()) == request.session_id.as_deref() {
turn.cancellation.store(true, Ordering::SeqCst);
}
}
}
let response = ServiceResponse::success_with_identity(
request.request_id,
request.session_id,
Some(request.method),
ResponsePayload::Session(payload),
);
ServiceOutbound::guarded(vec![ServiceMessage::response(response)], guard)
}
fn session_for_request(
&mut self,
requested_id: Option<&str>,
) -> std::result::Result<(Session, u64), ServiceErrorCode> {
if requested_id.is_some_and(|id| self.session_is_active(id)) {
return Err(ServiceErrorCode::SessionBusy);
}
let session = self.sessions.for_turn(requested_id)?;
let replay = session.read_events_tolerant_bounded(
crate::context::REPLAY_JSONL_MAX_LINES,
crate::context::REPLAY_JSONL_MAX_BYTES,
);
if !matches!(replay, Ok(ref replay) if replay.diagnostics.is_empty()) {
self.sessions.finish_turn(session.id());
return Err(ServiceErrorCode::SessionUnavailable);
}
let generation = session.replay_generation();
Ok((session, generation))
}
}
fn guarded_error(
identity: super::protocol::RequestIdentity,
guard: RequestIdGuard,
code: ServiceErrorCode,
) -> ServiceOutbound {
ServiceOutbound::guarded(
vec![ServiceMessage::response(ServiceResponse::error(
identity, code,
))],
guard,
)
}
#[derive(Default)]
struct TurnEventOrder {
sequence: u64,
activity_dropped: u64,
closed: bool,
maximum_sequence: Option<u64>,
exhausted: bool,
cancellation: Option<Arc<AtomicBool>>,
}
impl TurnEventOrder {
fn stamp(&mut self, event: &mut ServiceEvent) -> bool {
if self.sequence >= self.maximum_sequence.unwrap_or(u64::MAX).saturating_sub(1) {
self.exhausted = true;
if let Some(cancellation) = &self.cancellation {
cancellation.store(true, Ordering::SeqCst);
}
return false;
}
self.sequence += 1;
event.payload["sequence"] = serde_json::json!(self.sequence);
true
}
}
#[derive(Clone)]
struct ServiceActivitySender {
sender: Sender<TurnWorkerMessage>,
request_id: String,
session_id: String,
turn_id: String,
order: Arc<Mutex<TurnEventOrder>>,
}
impl ServiceActivitySender {
fn emit(&self, activity: super::activity::ActivityDto) {
let mut order = self.order.lock().unwrap_or_else(|e| e.into_inner());
if order.closed {
return;
}
let mut event = ServiceEvent::new(
self.request_id.clone(),
Some(self.session_id.clone()),
super::activity::ACTIVITY_EVENT,
serde_json::json!({"turn_id": self.turn_id, "activity": activity}),
);
if !order.stamp(&mut event) {
return;
}
if self
.sender
.try_send(TurnWorkerMessage::Delta {
turn_id: self.turn_id.clone(),
event,
})
.is_err()
{
order.activity_dropped = order.activity_dropped.saturating_add(1);
}
}
}
pub(super) struct ServiceTurnSink {
sender: Sender<TurnWorkerMessage>,
request_id: String,
session_id: String,
turn_id: String,
assistant_text: String,
activity: ServiceActivitySender,
persistent: bool,
output_limit_exceeded: bool,
}
impl ServiceTurnSink {
pub(super) fn new(
sender: Sender<TurnWorkerMessage>,
request_id: String,
session_id: String,
turn_id: String,
) -> Self {
Self {
activity: ServiceActivitySender {
sender: sender.clone(),
request_id: request_id.clone(),
session_id: session_id.clone(),
turn_id: turn_id.clone(),
order: Arc::new(Mutex::new(TurnEventOrder::default())),
},
sender,
request_id,
session_id,
turn_id,
assistant_text: String::new(),
persistent: false,
output_limit_exceeded: false,
}
}
pub(super) fn assistant_text(&self) -> &str {
&self.assistant_text
}
pub(super) fn finish(&self, event: &mut ServiceEvent) {
let mut order = self
.activity
.order
.lock()
.unwrap_or_else(|e| e.into_inner());
order.closed = true;
order.sequence = order.sequence.saturating_add(1);
event.payload["sequence"] = serde_json::json!(order.sequence);
event.payload["activity_dropped"] = serde_json::json!(order.activity_dropped);
}
}
impl AgentOutputSink for ServiceTurnSink {
fn wait_for_session_title_worker(&self) -> bool {
self.persistent
}
fn assistant_delta(&mut self, text: &str) -> Result<()> {
if text.len() > super::protocol::MAX_STRING_BYTES
|| self
.assistant_text
.len()
.checked_add(text.len())
.is_none_or(|length| length > super::protocol::MAX_STRING_BYTES)
{
self.output_limit_exceeded = true;
anyhow::bail!("assistant output exceeds the service text limit");
}
let candidate = format!("{}{}", self.assistant_text, text);
if !super::protocol::terminal_text_fits(
&self.request_id,
&self.session_id,
&self.turn_id,
&candidate,
) {
self.output_limit_exceeded = true;
anyhow::bail!("assistant output exceeds the service terminal envelope limit");
}
if self.persistent && serde_json::to_vec(&candidate)?.len() > 20_000 {
self.output_limit_exceeded = true;
anyhow::bail!("assistant output exceeds the service snapshot limit");
}
let mut event = ServiceEvent::assistant_delta(
self.request_id.clone(),
self.session_id.clone(),
self.turn_id.clone(),
text.to_string(),
);
let mut order = self
.activity
.order
.lock()
.unwrap_or_else(|e| e.into_inner());
if !order.stamp(&mut event) {
self.output_limit_exceeded = true;
anyhow::bail!("service event sequence limit reached");
}
self.sender
.send(TurnWorkerMessage::Delta {
turn_id: self.turn_id.clone(),
event,
})
.map_err(|_| anyhow::anyhow!("service output channel closed"))?;
self.assistant_text.push_str(text);
Ok(())
}
fn output_event(&mut self, event: OutputEvent) -> Result<()> {
match event {
OutputEvent::AssistantDelta { text } => self.assistant_delta(&text),
event => {
if let Some(activity) = super::activity::output(event) {
self.activity.emit(activity);
}
Ok(())
}
}
}
fn provider_reasoning_summary(&mut self, text: &str) -> Result<()> {
self.activity.emit(super::activity::provider_summary(text));
Ok(())
}
fn activity_event(&mut self, event: crate::output::ActivityEvent) -> Result<()> {
if let Some(activity) = super::activity::activity(event) {
self.activity.emit(activity);
}
Ok(())
}
fn activity_sender(&self) -> Option<crate::output::ActivitySender> {
let sender = self.activity.clone();
Some(Arc::new(move |event| {
if let Some(activity) = super::activity::activity(event) {
sender.emit(activity);
}
}))
}
fn tool_block(&mut self, _block: &str) -> Result<()> {
Ok(())
}
}
fn run_turn_worker(job: TurnWorkerJob) {
let turn_id = job.turn_id.clone();
let session_id = job.session.id().to_string();
let cancellation = Arc::clone(&job.cancellation);
let mut sink = ServiceTurnSink::new(
job.sender.clone(),
job.request_id.clone(),
session_id.clone(),
turn_id.clone(),
);
sink.persistent = job.settings_captured_at_admission;
{
let mut order = sink
.activity
.order
.lock()
.unwrap_or_else(|error| error.into_inner());
order.maximum_sequence = Some(job.maximum_sequence);
order.cancellation = Some(Arc::clone(&job.cancellation));
}
let run_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let reloaded;
let runtime = if job.settings_captured_at_admission {
&job.runtime
} else {
reloaded = Arc::new(job.runtime.reloaded()?);
&reloaded
};
agent::runner::run_provider_once_streaming(
&runtime.config,
&runtime.instructions,
&runtime.skills,
agent::runner::ProviderRunOptions {
settings: Some(runtime.settings.clone()),
prompt: &job.prompt,
session: Some(&job.session),
cwd: &job.runtime.cwd,
output_sink: Some(&mut sink),
selected_primary_agent: None,
cancellation: Some(Arc::clone(&cancellation)),
session_title_notifier: None,
herdr_reporter: None,
invocation_mode: crate::output::InvocationMode::Print,
disabled_tools: None,
disabled_subagent_profiles: None,
subagent_profile_discovery: None,
mcp: job.runtime.mcp.clone(),
},
)
}));
sink.activity
.order
.lock()
.unwrap_or_else(|error| error.into_inner())
.closed = true;
if job.settings_captured_at_admission {
let _ = job.sender.send(TurnWorkerMessage::Persisting {
turn_id: turn_id.clone(),
});
}
let status = match run_result {
Ok(Ok(_)) => TurnTerminalStatus::Completed,
Ok(Err(error)) if cancellation.load(Ordering::SeqCst) || is_run_canceled(&error) => {
TurnTerminalStatus::Cancelled
}
Ok(Err(_)) | Err(_) => TurnTerminalStatus::Failed,
};
persist_and_emit_terminal(job, &mut sink, status);
}
fn persist_and_emit_terminal(
job: TurnWorkerJob,
sink: &mut ServiceTurnSink,
mut status: TurnTerminalStatus,
) {
let turn_id = job.turn_id.clone();
let session_id = job.session.id().to_string();
sink.output_limit_exceeded |= sink
.activity
.order
.lock()
.unwrap_or_else(|error| error.into_inner())
.exhausted;
if sink.output_limit_exceeded {
status = TurnTerminalStatus::Failed;
}
let durable_status = match status {
TurnTerminalStatus::Completed => TurnStatus::Complete,
TurnTerminalStatus::Cancelled => TurnStatus::Cancelled,
TurnTerminalStatus::Failed => TurnStatus::Failed,
};
let persistence_failed = record_missing_terminal_status(
&job.session,
&job.runtime.cwd,
job.baseline_replay_generation,
durable_status,
sink.assistant_text(),
)
.is_err();
if persistence_failed {
status = TurnTerminalStatus::Failed;
}
let assistant_text = sink.assistant_text().to_string();
let mut terminal = ServiceEvent::turn_terminal(
job.request_id,
session_id,
turn_id.clone(),
status,
assistant_text,
);
sink.finish(&mut terminal);
if job.settings_captured_at_admission {
terminal.payload["persistence"] = serde_json::json!(if persistence_failed {
"failed"
} else {
"committed"
});
let code = if persistence_failed {
Some(ServiceErrorCode::PersistenceFailed)
} else if sink.output_limit_exceeded {
Some(ServiceErrorCode::OutputLimitExceeded)
} else if status == TurnTerminalStatus::Failed {
Some(ServiceErrorCode::InternalError)
} else {
None
};
terminal.payload["error"] = code.map_or(
serde_json::Value::Null,
|code| serde_json::json!({"code":code,"message":code.message()}),
);
}
let _ = job.sender.send(TurnWorkerMessage::Terminal {
turn_id,
event: terminal,
});
}
fn record_missing_terminal_status(
session: &Session,
cwd: &std::path::Path,
baseline_replay_generation: u64,
status: TurnStatus,
assistant_text: &str,
) -> Result<()> {
if session.terminal_status_recorded_since(baseline_replay_generation) {
return Ok(());
}
let payload = TurnStatusPayload::new(status, Some(assistant_text)).into_value()?;
crate::sessions::record_session_event(Some(session), cwd, SessionEventKind::TurnStatus, payload)
.map(|_| ())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::providers::{Provider, ProviderEvent, ToolCall};
use crate::{
config::{EffectiveConfig, McPaths, ProviderCredential, Settings},
sessions::SessionManager,
skills::SkillDiscovery,
};
use crossbeam_channel::{Receiver, bounded, unbounded};
use serde_json::{Value, json};
use std::{
collections::{BTreeMap, HashSet},
sync::Mutex,
};
fn service_dispatcher() -> ServiceDispatcher {
ServiceDispatcher::new(
super::super::dispatcher::ServiceSnapshot::new(|| Ok(false)),
super::super::dispatcher::ServiceTransportCapabilities::new(vec!["test".to_string()]),
)
}
fn response_value(outbound: &ServiceOutbound) -> Value {
serde_json::to_value(outbound.messages().first().unwrap()).unwrap()
}
fn error_code(outbound: &ServiceOutbound) -> ServiceErrorCode {
let ServiceMessage::Response(response) = outbound.messages().first().unwrap() else {
panic!("expected an error response");
};
response.error.as_ref().unwrap().code
}
fn accepted_turn_identity(outbound: &ServiceOutbound) -> (String, String) {
let value = response_value(outbound);
(
value["payload"]["turn_id"].as_str().unwrap().to_string(),
value["payload"]["session_id"].as_str().unwrap().to_string(),
)
}
fn released_terminal_worker(release: Receiver<()>) -> TurnWorker {
Arc::new(move |job| {
release.recv().expect("turn release");
let session_id = job.session.id().to_string();
let event = ServiceEvent::turn_terminal(
job.request_id.clone(),
session_id,
job.turn_id.clone(),
TurnTerminalStatus::Completed,
"answer".to_string(),
);
job.sender
.send(TurnWorkerMessage::Terminal {
turn_id: job.turn_id,
event,
})
.expect("turn event receiver");
})
}
struct SegmentedProvider {
first_segment: String,
request_count: Mutex<usize>,
}
impl SegmentedProvider {
fn new(first_segment: String) -> Self {
Self {
first_segment,
request_count: Mutex::new(0),
}
}
fn request_count(&self) -> usize {
*self.request_count.lock().unwrap()
}
}
impl Provider for SegmentedProvider {
fn stream_cancellable(
&self,
_request: crate::providers::ProviderRequest,
_cancellation: &crate::cancellation::AgentCancellation,
on_event: &mut dyn FnMut(ProviderEvent) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
let request_number = {
let mut count = self.request_count.lock().unwrap();
*count += 1;
*count
};
if request_number == 1 {
on_event(ProviderEvent::TextDelta(self.first_segment.clone()))?;
on_event(ProviderEvent::ToolCall(ToolCall {
id: "continuation-read".to_string(),
name: "read".to_string(),
arguments: json!({"path": "file.txt"}),
}))?;
} else {
on_event(ProviderEvent::TextDelta("second segment".to_string()))?;
}
on_event(ProviderEvent::Done)
}
}
fn largest_encoded_terminal_text(request_id: &str, session_id: &str, turn_id: &str) -> String {
let mut low = 0;
let mut high = super::super::protocol::MAX_STRING_BYTES;
while low < high {
let middle = (low + high).div_ceil(2);
if terminal_text_fits(request_id, session_id, turn_id, &"\"".repeat(middle)) {
low = middle;
} else {
high = middle - 1;
}
}
"\"".repeat(low)
}
use tempfile::TempDir;
fn runtime(temp: &TempDir) -> Arc<ServiceRuntime> {
let paths =
McPaths::from_root_and_project_dir(temp.path().join("mc"), temp.path().to_owned());
Arc::new(ServiceRuntime {
config: EffectiveConfig {
provider: Some("test".to_string()),
model: Some("test-model".to_string()),
no_color: true,
file_autocomplete_respects_gitignore: true,
custom_providers: BTreeMap::new(),
thinking_level: crate::thinking::ThinkingLevel::Default,
auth: Some(ProviderCredential::NoAuth),
paths: paths.clone(),
},
settings: Settings::default(),
instructions: Vec::new(),
skills: SkillDiscovery::default(),
session_manager: SessionManager::new(paths.sessions),
cwd: temp.path().to_path_buf(),
mcp: None,
})
}
#[test]
fn persistent_output_reserves_snapshot_and_terminal_sequence_capacity() {
let (sender, receiver) = bounded(4);
let mut sink =
ServiceTurnSink::new(sender, "request".into(), "session".into(), "turn".into());
sink.persistent = true;
let accepted = "\"".repeat(9_999);
sink.assistant_delta(&accepted).unwrap();
assert!(sink.assistant_delta("\"").is_err());
assert_eq!(sink.assistant_text(), accepted);
assert_eq!(receiver.len(), 1);
let (sender, _receiver) = bounded(4);
let mut sink =
ServiceTurnSink::new(sender, "request".into(), "session".into(), "turn".into());
sink.assistant_delta("accepted").unwrap();
let cancel = Arc::new(AtomicBool::new(false));
{
let mut order = sink.activity.order.lock().unwrap();
order.maximum_sequence = Some(2);
order.cancellation = Some(Arc::clone(&cancel));
}
assert!(sink.assistant_delta("x").is_err());
assert!(cancel.load(Ordering::SeqCst));
let mut terminal = ServiceEvent::turn_terminal(
"request".into(),
"session".into(),
"turn".into(),
TurnTerminalStatus::Failed,
"accepted".into(),
);
sink.finish(&mut terminal);
assert_eq!(terminal.payload["sequence"], 2);
}
#[test]
fn new_turn_snapshot_reloads_frontend_settings_without_mutating_running_snapshot() {
let env = crate::test_support::env::env_lock();
let _provider = env.save("MC_PROVIDER");
let _model = env.save("MC_MODEL");
env.remove_var("MC_PROVIDER");
env.remove_var("MC_MODEL");
let temp = TempDir::new().unwrap();
let original = runtime(&temp);
original.config.paths.ensure_runtime_dirs().unwrap();
let mut settings = Settings::default();
settings.selected_model.provider = Some("anthropic".into());
settings.selected_model.model = Some("claude-new".into());
settings.selected_model.thinking_level = Some(crate::thinking::ThinkingLevel::High);
settings.fast.enabled = true;
crate::config::write_settings(&original.config.paths, &settings).unwrap();
let updated = original.reloaded().unwrap();
assert_eq!(updated.config.provider.as_deref(), Some("anthropic"));
assert_eq!(updated.config.model.as_deref(), Some("claude-new"));
assert_eq!(
updated.config.thinking_level,
crate::thinking::ThinkingLevel::High
);
assert!(updated.settings.fast.enabled);
assert_eq!(original.config.model.as_deref(), Some("test-model"));
assert!(!original.settings.fast.enabled);
}
#[test]
fn admitted_worker_uses_settings_captured_before_later_commit() {
let env = crate::test_support::env::env_lock();
let _provider = env.save("MC_PROVIDER");
let _model = env.save("MC_MODEL");
env.remove_var("MC_PROVIDER");
env.remove_var("MC_MODEL");
let temp = TempDir::new().unwrap();
let runtime = runtime(&temp);
runtime.config.paths.ensure_runtime_dirs().unwrap();
let mut settings = Settings::default();
settings.selected_model.provider = Some("anthropic".into());
settings.selected_model.model = Some("before-admission".into());
crate::config::write_settings(&runtime.config.paths, &settings).unwrap();
let (sender, receiver) = bounded(8);
let (release, wait) = bounded(1);
let (observed, observation) = bounded(1);
let worker: TurnWorker = Arc::new(move |job| {
wait.recv().unwrap();
observed.send(job.runtime.config.model.clone()).unwrap();
let event = ServiceEvent::turn_terminal(
job.request_id,
job.session.id().into(),
job.turn_id.clone(),
TurnTerminalStatus::Completed,
String::new(),
);
job.sender
.send(TurnWorkerMessage::Terminal {
turn_id: job.turn_id,
event,
})
.unwrap();
});
let mut manager = ServiceTurnManager::new_with_worker(Arc::clone(&runtime), sender, worker);
let dispatcher = service_dispatcher();
let outbound = manager.start_with_settings_capture(
request(
"captured",
"turn.start",
json!({"prompt":"synthetic"}),
None,
),
dispatcher.begin_request("captured").unwrap(),
true,
Some(Arc::new(runtime.capture_turn_settings().unwrap())),
u64::MAX,
);
let (turn_id, _) = accepted_turn_identity(&outbound);
settings.selected_model.model = Some("after-admission".into());
crate::config::write_settings(&runtime.config.paths, &settings).unwrap();
release.send(()).unwrap();
assert_eq!(
observation.recv().unwrap().as_deref(),
Some("before-admission")
);
let (_, _, terminal) = manager
.handle_worker_message(receiver.recv().unwrap())
.unwrap();
assert!(terminal);
manager.finish_output(&turn_id);
assert_eq!(
runtime.reloaded().unwrap().config.model.as_deref(),
Some("after-admission")
);
}
#[test]
fn terminal_append_failure_is_reported_and_can_be_retried() {
let temp = TempDir::new().unwrap();
let runtime = runtime(&temp);
let session = runtime
.session_manager
.create()
.unwrap()
.activate()
.unwrap();
let baseline = session.replay_generation();
std::fs::create_dir(session.path()).unwrap();
assert!(
record_missing_terminal_status(
&session,
&runtime.cwd,
baseline,
TurnStatus::Complete,
"accepted"
)
.is_err()
);
assert!(!session.terminal_status_recorded_since(baseline));
std::fs::remove_dir(session.path()).unwrap();
record_missing_terminal_status(
&session,
&runtime.cwd,
baseline,
TurnStatus::Failed,
"accepted",
)
.unwrap();
let events = session.read_events().unwrap();
assert_eq!(events.len(), 1);
assert_eq!(events[0].payload["status"], "failed");
assert_eq!(events[0].payload["assistant_text"], "accepted");
}
#[test]
fn persistence_failure_projects_failed_terminal_with_authoritative_assistant_text() {
for (run_status, output_limit_exceeded) in [
(TurnTerminalStatus::Completed, false),
(TurnTerminalStatus::Cancelled, false),
(TurnTerminalStatus::Failed, true),
] {
let temp = TempDir::new().unwrap();
let runtime = runtime(&temp);
let session = runtime
.session_manager
.create()
.unwrap()
.activate()
.unwrap();
let baseline_replay_generation = session.replay_generation();
let session_id = session.id().to_string();
let (sender, receiver) = bounded(8);
let mut sink = ServiceTurnSink::new(
sender.clone(),
"request".into(),
session_id.clone(),
"turn".into(),
);
sink.persistent = true;
sink.assistant_delta("accepted ").unwrap();
sink.assistant_delta("answer λ").unwrap();
sink.output_limit_exceeded = output_limit_exceeded;
std::fs::create_dir(session.path()).unwrap();
let job = TurnWorkerJob {
runtime,
settings_captured_at_admission: true,
maximum_sequence: u64::MAX,
session: session.clone(),
prompt: "prompt".into(),
turn_id: "turn".into(),
request_id: "request".into(),
cancellation: Arc::new(AtomicBool::new(false)),
sender,
baseline_replay_generation,
};
persist_and_emit_terminal(job, &mut sink, run_status);
for expected_text in ["accepted ", "answer λ"] {
let TurnWorkerMessage::Delta { event, .. } = receiver.try_recv().unwrap() else {
panic!("accepted assistant delta expected");
};
assert_eq!(event.payload["text"], expected_text);
}
let TurnWorkerMessage::Terminal { turn_id, event } = receiver.try_recv().unwrap()
else {
panic!("production terminal event expected");
};
assert_eq!(turn_id, "turn");
assert_eq!(event.payload["status"], "failed");
assert_eq!(event.payload["persistence"], "failed");
assert_eq!(event.payload["error"]["code"], "persistence_failed");
assert_eq!(event.payload["assistant_text"], "accepted answer λ");
assert_eq!(event.payload["sequence"], 3);
assert!(receiver.try_recv().is_err());
assert!(!session.terminal_status_recorded_since(baseline_replay_generation));
}
}
fn request(id: &str, method: &str, payload: Value, session_id: Option<&str>) -> ServiceRequest {
ServiceRequest {
protocol_version: super::super::protocol::PROTOCOL_VERSION,
kind: super::super::protocol::MessageKind::Request,
request_id: id.to_string(),
session_id: session_id.map(str::to_string),
method: method.to_string(),
payload,
}
}
#[test]
fn activity_overflow_is_detectable_and_late_children_cannot_follow_terminal() {
use crate::output::{
ActivityEvent, ActivityId, ActivityKind, ActivityMetadata, ActivityStatus,
};
let (sender, receiver) = bounded(1);
let mut sink =
ServiceTurnSink::new(sender, "request".into(), "session".into(), "turn".into());
let child = sink.activity_sender().unwrap();
child(ActivityEvent::Started {
id: ActivityId::new("child"),
parent_id: None,
kind: ActivityKind::SubagentTask,
status: ActivityStatus::Running,
metadata: ActivityMetadata::new("private"),
});
child(ActivityEvent::Finished {
id: ActivityId::new("child"),
status: ActivityStatus::Success,
metadata: None,
});
let TurnWorkerMessage::Delta { event, .. } = receiver.recv().unwrap() else {
panic!("activity expected")
};
assert_eq!(event.payload["sequence"], 1);
sink.assistant_delta("answer").unwrap();
let TurnWorkerMessage::Delta { event, .. } = receiver.recv().unwrap() else {
panic!("delta expected")
};
assert_eq!(event.payload["sequence"], 3);
let mut terminal = ServiceEvent::turn_terminal(
"request".into(),
"session".into(),
"turn".into(),
TurnTerminalStatus::Completed,
sink.assistant_text().into(),
);
sink.finish(&mut terminal);
assert_eq!(terminal.payload["sequence"], 4);
assert_eq!(terminal.payload["activity_dropped"], 1);
assert_eq!(terminal.payload["activity_replay_available"], false);
assert_eq!(terminal.payload["replay_required"], true);
child(ActivityEvent::Finished {
id: ActivityId::new("late"),
status: ActivityStatus::Failed,
metadata: None,
});
assert!(receiver.try_recv().is_err());
}
#[test]
fn service_sink_orders_activity_with_accepted_deltas_and_bounds_text() {
let (sender, receiver) = crossbeam_channel::bounded(4);
let mut sink = ServiceTurnSink::new(
sender,
"request-1".to_string(),
"session-1".to_string(),
"turn-1".to_string(),
);
sink.output_event(OutputEvent::ToolStarted {
call: Box::new(crate::providers::ToolCall {
id: "call".to_string(),
name: "tool".to_string(),
arguments: json!({}),
}),
label: "tool".to_string(),
})
.unwrap();
sink.output_event(OutputEvent::AssistantDelta {
text: "one".to_string(),
})
.unwrap();
sink.output_event(OutputEvent::AssistantDelta {
text: "two".to_string(),
})
.unwrap();
assert_eq!(sink.assistant_text(), "onetwo");
let messages = receiver.try_iter().collect::<Vec<_>>();
assert_eq!(messages.len(), 3);
for (index, message) in messages.iter().enumerate() {
let TurnWorkerMessage::Delta { event, .. } = message else {
panic!("expected live event")
};
assert_eq!(event.payload["sequence"], index + 1);
assert_eq!(
event.event,
if index == 0 {
super::super::activity::ACTIVITY_EVENT
} else {
ASSISTANT_DELTA_EVENT
}
);
}
assert!(
sink.assistant_delta(&"x".repeat(super::super::protocol::MAX_STRING_BYTES))
.is_err()
);
}
#[test]
fn service_sink_rejects_text_that_only_fails_after_terminal_json_encoding() {
let (sender, receiver) = crossbeam_channel::bounded(4);
let mut sink = ServiceTurnSink::new(
sender,
"request-1".to_string(),
"session-1".to_string(),
"turn-1".to_string(),
);
let near_limit = "x".repeat(12_000);
sink.assistant_delta(&near_limit).unwrap();
let escaped_controls = "\u{0000}".repeat(4_000);
assert!(sink.assistant_delta(&escaped_controls).is_err());
assert_eq!(sink.assistant_text(), near_limit);
assert_eq!(receiver.try_iter().count(), 1);
}
#[test]
fn terminal_status_is_recorded_after_compaction_but_not_duplicated_without_it() {
let temp = TempDir::new().unwrap();
let runtime = runtime(&temp);
let session = runtime
.session_manager
.create()
.unwrap()
.activate()
.unwrap();
let baseline = session.replay_generation();
record_missing_terminal_status(
&session,
&runtime.cwd,
baseline,
TurnStatus::Cancelled,
"partial",
)
.unwrap();
assert_eq!(session.read_events().unwrap().len(), 1);
record_missing_terminal_status(
&session,
&runtime.cwd,
baseline,
TurnStatus::Cancelled,
"partial",
)
.unwrap();
assert_eq!(session.read_events().unwrap().len(), 1);
crate::sessions::record_session_compaction(
&session,
&runtime.cwd,
"checkpoint",
"provider",
"model",
0,
)
.unwrap();
record_missing_terminal_status(
&session,
&runtime.cwd,
baseline,
TurnStatus::Failed,
"partial",
)
.unwrap();
assert_eq!(
session
.read_events()
.unwrap()
.iter()
.filter(|event| event.kind() == Some(SessionEventKind::TurnStatus))
.count(),
2
);
}
#[test]
fn manager_keeps_malformed_duplicate_from_active_turn_through_terminal_output() {
let temp = TempDir::new().unwrap();
let runtime = runtime(&temp);
let (sender, receiver) = crossbeam_channel::bounded(8);
let (worker_ready_sender, worker_ready_receiver) = crossbeam_channel::bounded(1);
let (release_sender, release_receiver) = crossbeam_channel::bounded(1);
let worker: TurnWorker = Arc::new(move |job| {
worker_ready_sender.send(()).unwrap();
release_receiver.recv().unwrap();
let delta = ServiceEvent::assistant_delta(
job.request_id.clone(),
job.session.id().to_string(),
job.turn_id.clone(),
"answer".to_string(),
);
job.sender
.send(TurnWorkerMessage::Delta {
turn_id: job.turn_id.clone(),
event: delta,
})
.unwrap();
let terminal = ServiceEvent::turn_terminal(
job.request_id,
job.session.id().to_string(),
job.turn_id.clone(),
TurnTerminalStatus::Completed,
"answer".to_string(),
);
job.sender
.send(TurnWorkerMessage::Terminal {
turn_id: job.turn_id,
event: terminal,
})
.unwrap();
});
let mut manager = ServiceTurnManager::new_with_worker(runtime, sender, worker);
let mut dispatcher = service_dispatcher();
let guard = dispatcher.begin_request("request-1").unwrap();
let started = manager.start(
request(
"request-1",
super::super::protocol::TURN_START_METHOD,
json!({"prompt":"hello","options":{}}),
None,
),
guard,
true,
);
assert_eq!(started.messages().len(), 2);
let started_event =
serde_json::to_value(started.messages()[1].as_event().unwrap()).unwrap();
let turn_id = started_event["payload"]["turn_id"]
.as_str()
.unwrap()
.to_string();
drop(started);
worker_ready_receiver.recv().unwrap();
let malformed = br#"{"protocol_version":1,"kind":"request","request_id":"request-1","session_id":"session-1","method":"status","payload":{},"future":"malformed-secret"}"#;
let error = super::super::protocol::decode_request(malformed).unwrap_err();
assert_eq!(error.code, ServiceErrorCode::InvalidRequest);
let duplicate_outbound = dispatcher.error_outbound(error.identity, error.code);
let duplicate = response_value(&duplicate_outbound);
assert_eq!(duplicate["error"]["code"], "duplicate_request_id");
assert_eq!(duplicate["request_id"], "request-1");
assert_eq!(duplicate["method"], "status");
let encoded_duplicate =
super::super::protocol::encode_message(duplicate_outbound.messages().first().unwrap())
.unwrap();
assert!(
!String::from_utf8(encoded_duplicate)
.unwrap()
.contains("malformed-secret")
);
drop(duplicate_outbound);
release_sender.send(()).unwrap();
let delta = receiver.recv().unwrap();
let Some((received_turn_id, event, is_terminal)) = manager.handle_worker_message(delta)
else {
panic!("delta was not active");
};
assert_eq!(received_turn_id, turn_id);
assert_eq!(event.event, ASSISTANT_DELTA_EVENT);
assert_eq!(event.payload["text"], "answer");
assert!(!is_terminal);
let terminal = receiver.recv().unwrap();
let Some((received_turn_id, terminal_event, is_terminal)) =
manager.handle_worker_message(terminal)
else {
panic!("terminal was not active");
};
assert_eq!(received_turn_id, turn_id);
assert_eq!(terminal_event.event, TURN_TERMINAL_EVENT);
assert_eq!(terminal_event.payload["status"], "completed");
assert_eq!(terminal_event.payload["assistant_text"], "answer");
assert!(is_terminal);
let session_id = terminal_event.session_id.as_deref().unwrap();
assert!(matches!(
manager.session_for_request(Some(session_id)),
Err(ServiceErrorCode::SessionBusy)
));
let session_id = session_id.to_owned();
let terminal_outbound =
ServiceOutbound::unguarded(vec![ServiceMessage::Event(terminal_event)]);
assert!(
super::super::protocol::encode_message(terminal_outbound.messages().first().unwrap())
.is_ok()
);
assert!(matches!(
dispatcher.begin_request("request-1"),
Err(ServiceErrorCode::DuplicateRequestId)
));
drop(terminal_outbound);
manager.finish_output(&turn_id);
assert!(manager.is_empty());
assert!(dispatcher.begin_request("request-1").is_ok());
assert!(manager.session_for_request(Some(&session_id)).is_ok());
}
#[test]
fn manager_rejects_busy_sessions_allows_distinct_sessions_and_caps_active_turns() {
let temp = TempDir::new().unwrap();
let (sender, receiver) = bounded(16);
let (release_sender, release_receiver) = unbounded();
let worker = released_terminal_worker(release_receiver);
let mut manager = ServiceTurnManager::new_with_worker(runtime(&temp), sender, worker);
let dispatcher = service_dispatcher();
let first = manager.start(
request(
"turn-1",
super::super::protocol::TURN_START_METHOD,
json!({"prompt": "first", "options": {}}),
None,
),
dispatcher.begin_request("turn-1").unwrap(),
true,
);
let (first_turn_id, session_id) = accepted_turn_identity(&first);
drop(first);
let mut accepted_turn_ids = vec![first_turn_id];
let busy = manager.start(
request(
"turn-2",
super::super::protocol::TURN_START_METHOD,
json!({"prompt": "busy", "options": {}}),
Some(&session_id),
),
dispatcher.begin_request("turn-2").unwrap(),
true,
);
assert_eq!(error_code(&busy), ServiceErrorCode::SessionBusy);
drop(busy);
let distinct = manager.start(
request(
"turn-3",
super::super::protocol::TURN_START_METHOD,
json!({"prompt": "distinct", "options": {}}),
None,
),
dispatcher.begin_request("turn-3").unwrap(),
true,
);
let (distinct_turn_id, distinct_session_id) = accepted_turn_identity(&distinct);
assert_ne!(session_id, distinct_session_id);
accepted_turn_ids.push(distinct_turn_id);
drop(distinct);
for request_id in ["turn-4", "turn-5"] {
let accepted = manager.start(
request(
request_id,
super::super::protocol::TURN_START_METHOD,
json!({"prompt": request_id, "options": {}}),
None,
),
dispatcher.begin_request(request_id).unwrap(),
true,
);
let (turn_id, _) = accepted_turn_identity(&accepted);
accepted_turn_ids.push(turn_id);
drop(accepted);
}
assert_eq!(manager.active.len(), MAX_ACTIVE_TURNS);
let rejected = manager.start(
request(
"turn-6",
super::super::protocol::TURN_START_METHOD,
json!({"prompt": "over limit", "options": {}}),
None,
),
dispatcher.begin_request("turn-6").unwrap(),
true,
);
assert_eq!(error_code(&rejected), ServiceErrorCode::TooManyActiveTurns);
drop(rejected);
for _ in &accepted_turn_ids {
release_sender.send(()).unwrap();
}
let mut completed = HashSet::new();
for _ in &accepted_turn_ids {
let message = receiver.recv().unwrap();
let Some((turn_id, event, is_terminal)) = manager.handle_worker_message(message) else {
panic!("accepted turn output was filtered");
};
assert!(is_terminal);
assert_eq!(event.event, TURN_TERMINAL_EVENT);
assert!(accepted_turn_ids.contains(&turn_id));
assert!(completed.insert(turn_id.clone()));
manager.finish_output(&turn_id);
}
assert_eq!(completed.len(), MAX_ACTIVE_TURNS);
assert!(manager.is_empty());
}
#[test]
fn manager_keeps_one_terminal_when_cancellation_races_completion_or_failure() {
for (first_status, late_status) in [
(TurnTerminalStatus::Completed, TurnTerminalStatus::Failed),
(TurnTerminalStatus::Failed, TurnTerminalStatus::Completed),
(TurnTerminalStatus::Cancelled, TurnTerminalStatus::Completed),
] {
let temp = TempDir::new().unwrap();
let (sender, receiver) = bounded(8);
let (release_sender, release_receiver) = unbounded();
let worker: TurnWorker = Arc::new(move |job| {
let delta = ServiceEvent::assistant_delta(
job.request_id.clone(),
job.session.id().to_string(),
job.turn_id.clone(),
"before cancellation".to_string(),
);
job.sender
.send(TurnWorkerMessage::Delta {
turn_id: job.turn_id.clone(),
event: delta,
})
.unwrap();
release_receiver.recv().unwrap();
for status in [first_status, late_status] {
let terminal = ServiceEvent::turn_terminal(
job.request_id.clone(),
job.session.id().to_string(),
job.turn_id.clone(),
status,
"before cancellation".to_string(),
);
job.sender
.send(TurnWorkerMessage::Terminal {
turn_id: job.turn_id.clone(),
event: terminal,
})
.unwrap();
}
let late_delta = ServiceEvent::assistant_delta(
job.request_id,
job.session.id().to_string(),
job.turn_id.clone(),
"after terminal".to_string(),
);
job.sender
.send(TurnWorkerMessage::Delta {
turn_id: job.turn_id,
event: late_delta,
})
.unwrap();
});
let mut manager = ServiceTurnManager::new_with_worker(runtime(&temp), sender, worker);
let dispatcher = service_dispatcher();
let started = manager.start(
request(
"turn-1",
super::super::protocol::TURN_START_METHOD,
json!({"prompt": "race", "options": {}}),
None,
),
dispatcher.begin_request("turn-1").unwrap(),
true,
);
let (turn_id, session_id) = accepted_turn_identity(&started);
let contender = manager
.runtime
.session_manager
.open_existing(&session_id)
.unwrap();
assert!(contender.try_frontend_writer().unwrap().is_none());
drop(started);
let delta = receiver.recv().unwrap();
let Some((received_turn_id, event, is_terminal)) = manager.handle_worker_message(delta)
else {
panic!("initial delta was filtered");
};
assert_eq!(received_turn_id, turn_id);
assert_eq!(event.event, ASSISTANT_DELTA_EVENT);
assert!(!is_terminal);
let cancel = manager.cancel(
request(
"cancel-1",
super::super::protocol::TURN_CANCEL_METHOD,
json!({"turn_id": turn_id}),
None,
),
dispatcher.begin_request("cancel-1").unwrap(),
true,
);
assert_eq!(
response_value(&cancel)["payload"]["status"],
"cancellation_requested"
);
drop(cancel);
release_sender.send(()).unwrap();
let terminal = receiver.recv().unwrap();
let Some((received_turn_id, event, is_terminal)) =
manager.handle_worker_message(terminal)
else {
panic!("first terminal was filtered");
};
assert_eq!(received_turn_id, turn_id);
assert!(is_terminal);
assert_eq!(event.event, TURN_TERMINAL_EVENT);
assert_eq!(
event.payload["status"],
serde_json::to_value(first_status).unwrap()
);
let late_terminal = receiver.recv().unwrap();
assert!(manager.handle_worker_message(late_terminal).is_none());
let late_delta = receiver.recv().unwrap();
assert!(manager.handle_worker_message(late_delta).is_none());
manager.finish_output(&turn_id);
assert!(manager.is_empty());
assert!(contender.try_frontend_writer().unwrap().is_some());
}
}
#[test]
fn continuation_separator_is_rejected_before_exceeding_encoded_terminal_limit() {
let temp = TempDir::new().unwrap();
std::fs::write(temp.path().join("file.txt"), "tool result").unwrap();
let session = SessionManager::new(temp.path().join("sessions"))
.create()
.unwrap();
let session_id = session.id().to_string();
let segment = largest_encoded_terminal_text("request-1", &session_id, "turn-1");
assert!(!segment.is_empty());
assert!(terminal_text_fits(
"request-1",
&session_id,
"turn-1",
&segment,
));
assert!(!terminal_text_fits(
"request-1",
&session_id,
"turn-1",
&format!("{segment}\n\n"),
));
let provider = SegmentedProvider::new(segment.clone());
let tools = crate::tools::ToolRuntime::new(temp.path()).unwrap();
let agent = crate::agent::AgentSession::new("model", &[], &SkillDiscovery::default());
let (sender, receiver) = bounded(4);
let mut sink = ServiceTurnSink::new(
sender,
"request-1".to_string(),
session_id,
"turn-1".to_string(),
);
let error = agent
.run_print_with_tools_streaming_output(
&provider,
"read the file",
Some(&tools),
Some(&session),
temp.path(),
Some(&mut sink),
)
.unwrap_err()
.to_string();
assert!(error.contains("service terminal envelope limit"), "{error}");
assert_eq!(sink.assistant_text(), segment);
assert_eq!(provider.request_count(), 2);
let messages = receiver.try_iter().filter(|message| matches!(message, TurnWorkerMessage::Delta { event, .. } if event.event == ASSISTANT_DELTA_EVENT)).collect::<Vec<_>>();
assert_eq!(messages.len(), 1);
let TurnWorkerMessage::Delta { event, .. } = &messages[0] else {
panic!("separator failure emitted a non-delta message");
};
assert_eq!(event.payload["text"], segment);
}
#[test]
fn session_lifecycle_is_routed_bounded_and_private_across_coordinators() {
use super::super::{
coordinator::ServiceCoordinator, dispatcher::ServiceTransportCapabilities,
};
let temp = TempDir::new().unwrap();
let runtime = runtime(&temp);
let mut first = ServiceCoordinator::new(
Arc::clone(&runtime),
ServiceTransportCapabilities::new(vec!["test".into()]),
);
let mut second = ServiceCoordinator::new(
Arc::clone(&runtime),
ServiceTransportCapabilities::new(vec!["test".into()]),
);
for coordinator in [&mut first, &mut second] {
let rejected =
coordinator.dispatch_request(request("early", "session.list", json!({}), None));
assert_eq!(error_code(&rejected), ServiceErrorCode::NotInitialized);
drop(rejected);
let initialized = coordinator.dispatch_request(request(
"init",
"initialize",
json!({"supported_protocol_versions": [1]}),
None,
));
assert!(response_value(&initialized)["error"].is_null());
}
let created = response_value(&first.dispatch_request(request(
"create",
"session.create",
json!({}),
None,
)));
let id = created["payload"]["session_id"].as_str().unwrap();
assert!(Uuid::parse_str(id).is_ok());
let list = response_value(&second.dispatch_request(request(
"list",
"session.list",
json!({"limit": 1}),
None,
)));
assert_eq!(list["payload"]["sessions"][0]["session_id"], id);
assert!(
!serde_json::to_string(&list)
.unwrap()
.contains(temp.path().to_str().unwrap())
);
let conflict =
second.dispatch_request(request("open", "session.open", json!({}), Some(id)));
assert_eq!(error_code(&conflict), ServiceErrorCode::SessionBusy);
drop(conflict);
let replay = response_value(&second.dispatch_request(request(
"replay",
"session.replay",
json!({}),
Some(id),
)));
assert_eq!(replay["payload"]["events"], json!([]));
for invalid in ["../escape", "/tmp/session", "x/y"] {
let error = first.dispatch_request(request(
"invalid",
"session.open",
json!({}),
Some(invalid),
));
assert_eq!(error_code(&error), ServiceErrorCode::InvalidSessionId);
}
let error = first.dispatch_request(request(
"limit",
"session.replay",
json!({"limit": 33}),
Some(id),
));
assert_eq!(error_code(&error), ServiceErrorCode::LimitExceeded);
drop(error);
let closed = response_value(&first.dispatch_request(request(
"close",
"session.close",
json!({}),
Some(id),
)));
assert_eq!(closed["payload"]["status"], "closed");
let again = response_value(&first.dispatch_request(request(
"close",
"session.close",
json!({}),
Some(id),
)));
assert_eq!(again["payload"]["status"], "closed");
let opened = response_value(&second.dispatch_request(request(
"open",
"session.open",
json!({}),
Some(id),
)));
assert_eq!(opened["payload"]["status"], "open");
assert!(runtime.session_manager.open_existing(id).is_ok());
assert!(second.shutdown_next().unwrap().is_none());
let reopened = response_value(&first.dispatch_request(request(
"reopen",
"session.open",
json!({}),
Some(id),
)));
assert_eq!(reopened["payload"]["status"], "open");
}
#[test]
fn session_close_cancels_and_keeps_writer_until_terminal_output_then_replays_partial() {
let temp = TempDir::new().unwrap();
let runtime = runtime(&temp);
let dispatcher = service_dispatcher();
let (sender, receiver) = bounded(8);
let (release_sender, release_receiver) = bounded(1);
let worker: TurnWorker = Arc::new(move |job| {
release_receiver.recv().unwrap();
assert!(job.cancellation.load(Ordering::SeqCst));
crate::sessions::record_session_event(
Some(&job.session),
&job.runtime.cwd,
SessionEventKind::UserInput,
json!({"text": "question"}),
)
.unwrap();
record_missing_terminal_status(
&job.session,
&job.runtime.cwd,
job.baseline_replay_generation,
TurnStatus::Cancelled,
"partial before close",
)
.unwrap();
job.sender
.send(TurnWorkerMessage::Terminal {
turn_id: job.turn_id.clone(),
event: ServiceEvent::turn_terminal(
job.request_id,
job.session.id().to_owned(),
job.turn_id,
TurnTerminalStatus::Cancelled,
"partial before close".into(),
),
})
.unwrap();
});
let mut manager = ServiceTurnManager::new_with_worker(Arc::clone(&runtime), sender, worker);
let created = manager.dispatch_session(
request("create", "session.create", json!({}), None),
dispatcher.begin_request("create").unwrap(),
true,
);
let value = response_value(&created);
let id = value["payload"]["session_id"].as_str().unwrap();
drop(created);
let accepted = manager.start(
request(
"turn",
"turn.start",
json!({"prompt": "question"}),
Some(id),
),
dispatcher.begin_request("turn").unwrap(),
true,
);
let (turn_id, _) = accepted_turn_identity(&accepted);
drop(accepted);
let closed = manager.dispatch_session(
request("close", "session.close", json!({}), Some(id)),
dispatcher.begin_request("close").unwrap(),
true,
);
assert_eq!(response_value(&closed)["payload"]["status"], "closing");
let contender = runtime.session_manager.open_existing(id).unwrap();
assert!(contender.try_frontend_writer().unwrap().is_none());
release_sender.send(()).unwrap();
let (terminal_id, _, terminal) = manager
.handle_worker_message(receiver.recv().unwrap())
.unwrap();
assert!(terminal);
assert_eq!(terminal_id, turn_id);
assert!(contender.try_frontend_writer().unwrap().is_none());
manager.finish_output(&turn_id);
assert!(contender.try_frontend_writer().unwrap().is_some());
let replay = serde_json::to_value(contender.frontend_replay(None, 32).unwrap()).unwrap();
assert_eq!(replay["events"][0]["text"], "question");
assert_eq!(replay["events"][1]["text"], "partial before close");
assert_eq!(replay["events"][2]["status"], "cancelled");
}
#[test]
fn session_attachment_limit_rejects_excess_and_close_restores_capacity() {
let temp = TempDir::new().unwrap();
let runtime = runtime(&temp);
let mut sessions = super::super::sessions::ServiceSessions::new(runtime);
let mut first_id = String::new();
for index in 0..16 {
let result = sessions
.operation(&request("create", "session.create", json!({}), None), false)
.unwrap();
if index == 0 {
first_id = result["session_id"].as_str().unwrap().to_owned();
}
}
assert_eq!(
sessions
.operation(&request("create", "session.create", json!({}), None), false)
.unwrap_err(),
ServiceErrorCode::LimitExceeded
);
sessions
.operation(
&request("close", "session.close", json!({}), Some(&first_id)),
false,
)
.unwrap();
assert!(
sessions
.operation(&request("create", "session.create", json!({}), None), false)
.is_ok()
);
}
trait ServiceMessageExt {
fn as_event(&self) -> Option<&ServiceEvent>;
}
impl ServiceMessageExt for ServiceMessage {
fn as_event(&self) -> Option<&ServiceEvent> {
match self {
ServiceMessage::Event(event) => Some(event),
ServiceMessage::Response(_) => None,
}
}
}
}