use std::sync::Arc;
use tokio::sync::mpsc;
use crate::agent;
use crate::app_server::{AppServerClient, AppServerStreamEvent, AppServerTurnRequest};
use crate::channel::{
AgentChannel, AgentError, AgentFuture, SessionRef, StartSessionRequest, TurnEvent, TurnRequest,
TurnResult,
};
use crate::model::agent::AgentKind;
pub struct AppServerAgentChannel {
client: Arc<dyn AppServerClient>,
kind: AgentKind,
}
impl AppServerAgentChannel {
pub fn new(client: Arc<dyn AppServerClient>, kind: AgentKind) -> Self {
Self { client, kind }
}
}
impl AgentChannel for AppServerAgentChannel {
fn start_session(
&self,
req: StartSessionRequest,
) -> AgentFuture<Result<SessionRef, AgentError>> {
let session_id = req.session_id;
Box::pin(async move { Ok(SessionRef { session_id }) })
}
fn run_turn(
&self,
session_id: String,
req: TurnRequest,
events: mpsc::UnboundedSender<TurnEvent>,
) -> AgentFuture<Result<TurnResult, AgentError>> {
let client = Arc::clone(&self.client);
let kind = self.kind;
Box::pin(async move {
let request = AppServerTurnRequest {
folder: req.folder,
live_session_output: req.live_session_output,
model: req.model,
prompt: req.prompt,
request_kind: req.request_kind,
provider_conversation_id: req.provider_conversation_id,
persisted_instruction_conversation_id: req.persisted_instruction_conversation_id,
reasoning_level: req.reasoning_level,
session_id,
};
let protocol_profile = request.request_kind.protocol_profile();
let repair_request = request.clone();
let (stream_tx, mut stream_rx) = mpsc::unbounded_channel::<AppServerStreamEvent>();
let bridge_handle = {
let events = events.clone();
tokio::spawn(async move {
while let Some(event) = stream_rx.recv().await {
match event {
AppServerStreamEvent::AssistantMessage {
message,
phase,
is_delta,
} => {
let trimmed = message.trim_end();
if trimmed.trim().is_empty() {
continue;
}
if agent::is_app_server_thought_chunk(
kind,
is_delta,
phase.as_deref(),
) {
let _ =
events.send(TurnEvent::ThoughtDelta(trimmed.to_string()));
}
}
AppServerStreamEvent::ProgressUpdate(progress) => {
let trimmed = progress.trim();
if trimmed.is_empty() {
continue;
}
let _ = events.send(TurnEvent::ThoughtDelta(trimmed.to_string()));
}
}
}
})
};
let turn_result = client.run_turn(request, stream_tx).await;
let _ = bridge_handle.await;
match turn_result {
Ok(response) => {
let _ = events.send(TurnEvent::PidUpdate(response.pid));
let parsed = parse_or_repair_app_server_response(
kind,
&response,
protocol_profile,
repair_request,
&client,
&events,
)
.await?;
Ok(TurnResult {
assistant_message: parsed.assistant_message,
context_reset: response.context_reset,
input_tokens: response.input_tokens + parsed.repair_input_tokens,
output_tokens: response.output_tokens + parsed.repair_output_tokens,
provider_conversation_id: parsed.provider_conversation_id,
})
}
Err(error) => Err(AgentError::AppServer(error)),
}
})
}
fn shutdown_session(&self, session_id: String) -> AgentFuture<Result<(), AgentError>> {
let client = Arc::clone(&self.client);
Box::pin(async move {
client.shutdown_session(session_id).await;
Ok(())
})
}
}
struct AppServerParsedTurnResult {
assistant_message: agent::AgentResponse,
provider_conversation_id: Option<String>,
repair_input_tokens: u64,
repair_output_tokens: u64,
}
async fn parse_or_repair_app_server_response(
kind: AgentKind,
response: &crate::app_server::AppServerTurnResponse,
protocol_profile: agent::ProtocolRequestProfile,
repair_request: AppServerTurnRequest,
client: &Arc<dyn AppServerClient>,
events: &mpsc::UnboundedSender<TurnEvent>,
) -> Result<AppServerParsedTurnResult, AgentError> {
let parse_error =
match agent::parse_turn_response(kind, &response.assistant_message, protocol_profile) {
Ok(parsed) => {
return Ok(AppServerParsedTurnResult {
assistant_message: parsed,
provider_conversation_id: response.provider_conversation_id.clone(),
repair_input_tokens: 0,
repair_output_tokens: 0,
});
}
Err(error) => error,
};
let _ = events.send(TurnEvent::ThoughtDelta(format!(
"Protocol parse error — retrying: {parse_error}"
)));
let repair_prompt =
agent::protocol::build_protocol_repair_prompt(&parse_error, &response.assistant_message);
let repair_provider_conversation_id = response
.provider_conversation_id
.clone()
.or_else(|| repair_request.provider_conversation_id.clone());
let repair_turn_request = AppServerTurnRequest {
folder: repair_request.folder,
live_session_output: None,
model: repair_request.model,
prompt: crate::model::turn_prompt::TurnPrompt::from_agent_data(repair_prompt),
request_kind: repair_request.request_kind,
provider_conversation_id: repair_provider_conversation_id,
persisted_instruction_conversation_id: None,
reasoning_level: repair_request.reasoning_level,
session_id: repair_request.session_id,
};
let (repair_stream_tx, _repair_stream_rx) = mpsc::unbounded_channel();
let repair_result = client
.run_turn(repair_turn_request, repair_stream_tx)
.await
.map_err(|error| {
AgentError::Backend(format!(
"{parse_error}\nprotocol repair transport failed: {error}"
))
})?;
let parsed =
agent::parse_turn_response(kind, &repair_result.assistant_message, protocol_profile)
.map_err(|error| {
AgentError::Backend(format!(
"{parse_error}\nprotocol repair retry also failed: \
{error}\nrepair_response:\n{}",
repair_result.assistant_message
))
})?;
Ok(AppServerParsedTurnResult {
assistant_message: parsed,
provider_conversation_id: repair_result
.provider_conversation_id
.or(response.provider_conversation_id.clone()),
repair_input_tokens: repair_result.input_tokens,
repair_output_tokens: repair_result.output_tokens,
})
}
#[cfg(test)]
mod tests {
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::mpsc;
use super::*;
use crate::app_server::{AppServerTurnResponse, MockAppServerClient};
use crate::channel::AgentRequestKind;
use crate::model::agent::ReasoningLevel;
use crate::model::turn_prompt::TurnPromptAttachment;
fn make_turn_request() -> TurnRequest {
TurnRequest {
folder: PathBuf::from("/tmp"),
live_session_output: None,
model: "gpt-5.5".to_string(),
request_kind: AgentRequestKind::SessionStart,
prompt: "Do something".into(),
provider_conversation_id: None,
persisted_instruction_conversation_id: None,
reasoning_level: ReasoningLevel::default(),
}
}
fn make_ok_response(assistant_message: &str) -> AppServerTurnResponse {
AppServerTurnResponse {
assistant_message: assistant_message.to_string(),
context_reset: false,
input_tokens: 10,
output_tokens: 5,
pid: None,
provider_conversation_id: None,
}
}
#[tokio::test]
async fn test_run_turn_suppresses_non_thought_assistant_delta_streaming() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, stream_tx| {
let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
message: "Hello world".to_string(),
phase: None,
is_delta: true,
});
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"Hello world","questions":[],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, mut events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await;
assert!(result.is_ok());
let events = std::iter::from_fn(|| events_rx.try_recv().ok()).collect::<Vec<_>>();
assert!(!events.is_empty());
assert!(
events
.iter()
.all(|event| matches!(event, TurnEvent::PidUpdate(_))),
"only pid events should be emitted, got: {events:?}"
);
}
#[tokio::test]
async fn test_run_turn_suppresses_non_delta_assistant_messages() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, stream_tx| {
let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
message: "Full paragraph ".to_string(),
phase: None,
is_delta: false,
});
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"Full paragraph","questions":[],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, mut events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await;
assert!(result.is_ok());
let events = std::iter::from_fn(|| events_rx.try_recv().ok()).collect::<Vec<_>>();
assert!(!events.is_empty());
assert!(
events
.iter()
.all(|event| matches!(event, TurnEvent::PidUpdate(_))),
"only pid events should be emitted, got: {events:?}"
);
}
#[tokio::test]
async fn test_run_turn_suppresses_non_delta_structured_json_streaming() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, stream_tx| {
let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
message: r#"{"answer":"Done.","questions":[{"text":"Need clarification.","options":[]}],"summary":null}"#.to_string(),
phase: None,
is_delta: false,
});
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"Done.","questions":[{"text":"Need clarification.","options":[]}],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, mut events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await;
assert!(result.is_ok());
let events = std::iter::from_fn(|| events_rx.try_recv().ok()).collect::<Vec<_>>();
assert!(!events.is_empty());
assert!(
events
.iter()
.all(|event| matches!(event, TurnEvent::PidUpdate(_))),
"only pid events should be emitted, got: {events:?}"
);
}
#[tokio::test]
async fn test_run_turn_routes_codex_thinking_delta_to_thought_event() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, stream_tx| {
let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
message: "Inspecting files".to_string(),
phase: Some("thinking".to_string()),
is_delta: true,
});
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"Done.","questions":[],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, mut events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await;
assert!(result.is_ok());
let event = events_rx.try_recv().expect("should have received an event");
assert_eq!(
event,
TurnEvent::ThoughtDelta("Inspecting files".to_string())
);
}
#[tokio::test]
async fn test_run_turn_routes_uppercase_codex_thinking_delta_to_thought_event() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, stream_tx| {
let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
message: "Inspecting files".to_string(),
phase: Some("Thinking".to_string()),
is_delta: true,
});
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"Done.","questions":[],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, mut events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await;
assert!(result.is_ok());
let event = events_rx.try_recv().expect("should have received an event");
assert_eq!(
event,
TurnEvent::ThoughtDelta("Inspecting files".to_string())
);
}
#[tokio::test]
async fn test_run_turn_routes_progress_update_events_to_thought_delta() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, stream_tx| {
let _ = stream_tx.send(AppServerStreamEvent::ProgressUpdate(
"Running tool".to_string(),
));
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"","questions":[],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, mut events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await;
assert!(result.is_ok());
let event = events_rx
.try_recv()
.expect("should have received a progress event");
assert_eq!(event, TurnEvent::ThoughtDelta("Running tool".to_string()));
}
#[tokio::test]
async fn test_run_turn_fills_missing_summary_for_session_turn() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, _stream_tx| {
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"Done.","questions":[],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Antigravity);
let (events_tx, _events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await
.expect("turn should succeed");
assert_eq!(
result.assistant_message.summary,
Some(agent::protocol::AgentResponseSummary {
turn: String::new(),
session: String::new(),
})
);
}
#[tokio::test]
async fn test_run_turn_skips_whitespace_only_assistant_messages() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, stream_tx| {
let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
message: " \n ".to_string(),
phase: None,
is_delta: true,
});
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"","questions":[],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, mut events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await;
assert!(result.is_ok());
while let Ok(event) = events_rx.try_recv() {
assert!(
!matches!(event, TurnEvent::ThoughtDelta(_)),
"no ThoughtDelta should be emitted for whitespace-only messages, got: {event:?}"
);
}
}
#[tokio::test]
async fn test_run_turn_skips_delta_protocol_json_fragments() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, stream_tx| {
let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
message: r#"{"answer":"#.to_string(),
phase: None,
is_delta: true,
});
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"Final answer.","questions":[],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, mut events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await
.expect("turn should succeed");
assert_eq!(result.assistant_message.to_display_text(), "Final answer.");
while let Ok(event) = events_rx.try_recv() {
assert!(
!matches!(event, TurnEvent::ThoughtDelta(_)),
"no ThoughtDelta should be emitted for protocol fragments, got: {event:?}"
);
}
}
#[tokio::test]
async fn test_run_turn_app_server_suppresses_streamed_assistant_messages() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, stream_tx| {
let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
message: "streamed plain text".to_string(),
phase: None,
is_delta: true,
});
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"Final structured output.","questions":[],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, mut events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await
.expect("turn should succeed");
assert_eq!(
result.assistant_message.to_display_text(),
"Final structured output."
);
while let Ok(event) = events_rx.try_recv() {
assert!(
!matches!(event, TurnEvent::ThoughtDelta(_)),
"no ThoughtDelta should be emitted for plain assistant deltas, got: {event:?}"
);
}
}
#[tokio::test]
async fn test_run_turn_returns_error_for_invalid_structured_output() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.times(2)
.returning(|request, _stream_tx| {
assert_eq!(request.request_kind, AgentRequestKind::SessionStart);
Box::pin(async { Ok(make_ok_response("plain non-json response")) })
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, _events_rx) = mpsc::unbounded_channel();
let error = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await
.expect_err("invalid structured output should fail");
let error_message = error.to_string();
assert!(error_message.contains("did not match the required JSON schema"));
assert!(error_message.contains("response:\nplain non-json response"));
}
#[tokio::test]
async fn test_run_turn_recovers_valid_output_via_protocol_repair() {
let call_counter = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let mut mock_client = MockAppServerClient::new();
mock_client.expect_run_turn().times(2).returning({
let counter = Arc::clone(&call_counter);
move |_request, _stream_tx| {
let call_number = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if call_number == 0 {
Box::pin(async { Ok(make_ok_response("plain non-json response")) })
} else {
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"Repaired response","questions":[],"summary":null}"#,
))
})
}
}
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, _events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await
.expect("repair retry should succeed");
assert_eq!(
result.assistant_message.to_display_text(),
"Repaired response"
);
}
#[tokio::test]
async fn test_run_turn_allows_image_attachments() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.times(1)
.returning(|request, _stream_tx| {
assert_eq!(request.prompt.attachments.len(), 1);
Box::pin(async {
Ok(make_ok_response(
r#"{"answer":"codex ok","questions":[],"summary":null}"#,
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, _events_rx) = mpsc::unbounded_channel();
let mut request = make_turn_request();
request.prompt.attachments.push(TurnPromptAttachment {
placeholder: "[Image #1]".to_string(),
local_image_path: PathBuf::from("/tmp/image.png"),
});
let result = channel
.run_turn("sess-1".to_string(), request, events_tx)
.await
.expect("turn should succeed");
assert_eq!(result.assistant_message.to_display_text(), "codex ok");
}
#[tokio::test]
async fn test_run_turn_codex_rejects_plain_text_after_repair_retry() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.times(2)
.returning(|_request, _stream_tx| Box::pin(async { Ok(make_ok_response("plain")) }));
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, _events_rx) = mpsc::unbounded_channel();
let error = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await
.expect_err("plain-text turn should fail");
let error_message = error.to_string();
assert!(error_message.contains("did not match the required JSON schema"));
assert!(error_message.contains("response:\nplain"));
}
#[tokio::test]
async fn test_run_turn_client_failure_returns_agent_error() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, _stream_tx| {
Box::pin(async {
Err(crate::app_server::AppServerError::Provider(
"server timeout".to_string(),
))
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, _events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await;
let error_message = result
.expect_err("expected Err on server timeout")
.to_string();
assert!(error_message.contains("server timeout"));
}
#[tokio::test]
async fn test_run_turn_returns_correct_token_counts_and_context_reset() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|_request, _stream_tx| {
Box::pin(async {
Ok(AppServerTurnResponse {
assistant_message: r#"{"answer":"Result","questions":[],"summary":null}"#
.to_string(),
context_reset: true,
input_tokens: 100,
output_tokens: 50,
pid: Some(1234),
provider_conversation_id: None,
})
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, _events_rx) = mpsc::unbounded_channel();
let result = channel
.run_turn("sess-1".to_string(), make_turn_request(), events_tx)
.await
.expect("turn should succeed");
assert_eq!(result.assistant_message.to_display_text(), "Result");
assert!(result.context_reset);
assert_eq!(result.input_tokens, 100);
assert_eq!(result.output_tokens, 50);
}
#[tokio::test]
async fn test_run_turn_passes_and_returns_provider_conversation_id() {
let mut mock_client = MockAppServerClient::new();
mock_client
.expect_run_turn()
.returning(|request, _stream_tx| {
assert_eq!(
request.provider_conversation_id,
Some("thread-abc".to_string()),
"request should carry the provider conversation id"
);
assert_eq!(
request.reasoning_level,
ReasoningLevel::Medium,
"request should carry the codex reasoning level"
);
Box::pin(async {
Ok(AppServerTurnResponse {
assistant_message: r#"{"answer":"ok","questions":[],"summary":null}"#
.to_string(),
context_reset: false,
input_tokens: 1,
output_tokens: 1,
pid: Some(42),
provider_conversation_id: Some("thread-xyz".to_string()),
})
})
});
let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
let (events_tx, mut events_rx) = mpsc::unbounded_channel();
let mut request = make_turn_request();
request.reasoning_level = ReasoningLevel::Medium;
request.provider_conversation_id = Some("thread-abc".to_string());
let result = channel
.run_turn("sess-1".to_string(), request, events_tx)
.await
.expect("turn should succeed");
assert_eq!(
result.provider_conversation_id,
Some("thread-xyz".to_string()),
"result should carry the provider conversation id from the response"
);
let mut pid_event_seen = false;
while let Ok(event) = events_rx.try_recv() {
if matches!(event, TurnEvent::PidUpdate(Some(42))) {
pid_event_seen = true;
}
}
assert!(pid_event_seen, "should emit PidUpdate from response pid");
}
}