ag_agent/channel/
app_server.rs

1//! App-server RPC [`AgentChannel`] adapter.
2//!
3//! Delegates turn execution to [`AppServerClient`] and bridges
4//! [`AppServerStreamEvent`]s to the unified [`TurnEvent`] stream.
5
6use std::sync::Arc;
7
8use tokio::sync::mpsc;
9
10use crate::agent;
11use crate::app_server::{AppServerClient, AppServerStreamEvent, AppServerTurnRequest};
12use crate::channel::{
13    AgentChannel, AgentError, AgentFuture, SessionRef, StartSessionRequest, TurnEvent, TurnRequest,
14    TurnResult,
15};
16use crate::model::agent::AgentKind;
17
18/// [`AgentChannel`] adapter backed by a persistent app-server session.
19///
20/// Turn execution is delegated to [`AppServerClient::run_turn`].
21/// [`AppServerStreamEvent`]s emitted by the provider are bridged to
22/// [`TurnEvent::ThoughtDelta`] values when transient loader text should be
23/// updated.
24pub struct AppServerAgentChannel {
25    /// Provider-specific app-server client.
26    client: Arc<dyn AppServerClient>,
27    /// Provider kind routed through this channel instance.
28    kind: AgentKind,
29}
30
31impl AppServerAgentChannel {
32    /// Creates a new app-server channel backed by the given client.
33    pub fn new(client: Arc<dyn AppServerClient>, kind: AgentKind) -> Self {
34        Self { client, kind }
35    }
36}
37
38impl AgentChannel for AppServerAgentChannel {
39    /// Returns a [`SessionRef`] immediately; the app-server session is
40    /// initialised lazily on the first turn.
41    fn start_session(
42        &self,
43        req: StartSessionRequest,
44    ) -> AgentFuture<Result<SessionRef, AgentError>> {
45        let session_id = req.session_id;
46
47        Box::pin(async move { Ok(SessionRef { session_id }) })
48    }
49
50    /// Runs one app-server turn and bridges stream events to [`TurnEvent`]s.
51    ///
52    /// Assistant stream chunks are never appended directly to the transcript.
53    /// Instead, Codex thought-style deltas (`phase: thinking/plan`) and
54    /// provider progress updates are bridged to [`TurnEvent::ThoughtDelta`] so
55    /// the UI loader can reflect transient state while the final persisted
56    /// output still comes only from the parsed [`TurnResult`].
57    ///
58    /// # Errors
59    /// Returns [`AgentError`] when [`AppServerClient::run_turn`] fails.
60    fn run_turn(
61        &self,
62        session_id: String,
63        req: TurnRequest,
64        events: mpsc::UnboundedSender<TurnEvent>,
65    ) -> AgentFuture<Result<TurnResult, AgentError>> {
66        let client = Arc::clone(&self.client);
67        let kind = self.kind;
68        Box::pin(async move {
69            let request = AppServerTurnRequest {
70                folder: req.folder,
71                live_session_output: req.live_session_output,
72                model: req.model,
73                prompt: req.prompt,
74                request_kind: req.request_kind,
75                provider_conversation_id: req.provider_conversation_id,
76                persisted_instruction_conversation_id: req.persisted_instruction_conversation_id,
77                reasoning_level: req.reasoning_level,
78                session_id,
79            };
80            let protocol_profile = request.request_kind.protocol_profile();
81            let repair_request = request.clone();
82            let (stream_tx, mut stream_rx) = mpsc::unbounded_channel::<AppServerStreamEvent>();
83
84            let bridge_handle = {
85                let events = events.clone();
86
87                tokio::spawn(async move {
88                    while let Some(event) = stream_rx.recv().await {
89                        match event {
90                            AppServerStreamEvent::AssistantMessage {
91                                message,
92                                phase,
93                                is_delta,
94                            } => {
95                                let trimmed = message.trim_end();
96                                if trimmed.trim().is_empty() {
97                                    continue;
98                                }
99
100                                if agent::is_app_server_thought_chunk(
101                                    kind,
102                                    is_delta,
103                                    phase.as_deref(),
104                                ) {
105                                    // Fire-and-forget: receiver may be dropped during shutdown.
106                                    let _ =
107                                        events.send(TurnEvent::ThoughtDelta(trimmed.to_string()));
108                                }
109                            }
110                            AppServerStreamEvent::ProgressUpdate(progress) => {
111                                let trimmed = progress.trim();
112                                if trimmed.is_empty() {
113                                    continue;
114                                }
115
116                                // Fire-and-forget: receiver may be dropped during shutdown.
117                                let _ = events.send(TurnEvent::ThoughtDelta(trimmed.to_string()));
118                            }
119                        }
120                    }
121                })
122            };
123
124            let turn_result = client.run_turn(request, stream_tx).await;
125            // Task join: panic in the spawned task is not recoverable here.
126            let _ = bridge_handle.await;
127
128            match turn_result {
129                Ok(response) => {
130                    // Fire-and-forget: receiver may be dropped during shutdown.
131                    let _ = events.send(TurnEvent::PidUpdate(response.pid));
132                    let parsed = parse_or_repair_app_server_response(
133                        kind,
134                        &response,
135                        protocol_profile,
136                        repair_request,
137                        &client,
138                        &events,
139                    )
140                    .await?;
141
142                    Ok(TurnResult {
143                        assistant_message: parsed.assistant_message,
144                        context_reset: response.context_reset,
145                        input_tokens: response.input_tokens + parsed.repair_input_tokens,
146                        output_tokens: response.output_tokens + parsed.repair_output_tokens,
147                        provider_conversation_id: parsed.provider_conversation_id,
148                    })
149                }
150                Err(error) => Err(AgentError::AppServer(error)),
151            }
152        })
153    }
154
155    /// Shuts down the underlying app-server session.
156    fn shutdown_session(&self, session_id: String) -> AgentFuture<Result<(), AgentError>> {
157        let client = Arc::clone(&self.client);
158
159        Box::pin(async move {
160            client.shutdown_session(session_id).await;
161
162            Ok(())
163        })
164    }
165}
166
167/// Aggregated result from parsing an app-server turn response, including
168/// metadata from a repair turn when one was needed.
169struct AppServerParsedTurnResult {
170    /// Parsed agent response from the successful attempt.
171    assistant_message: agent::AgentResponse,
172    /// Provider conversation id from the latest successful attempt,
173    /// falling back to the original response when the repair turn does
174    /// not produce one.
175    provider_conversation_id: Option<String>,
176    /// Additional input tokens consumed by a repair turn (zero when no
177    /// repair was needed).
178    repair_input_tokens: u64,
179    /// Additional output tokens consumed by a repair turn (zero when no
180    /// repair was needed).
181    repair_output_tokens: u64,
182}
183
184/// Parses one app-server turn response strictly, falling back to a single
185/// protocol-repair retry when the initial parse fails.
186///
187/// The repair prompt is sent as a follow-up turn on the same session so the
188/// agent retains the original conversation context. When repair succeeds,
189/// the returned metadata reflects the repair turn's provider conversation id
190/// and token usage so the caller can propagate them correctly.
191///
192/// When repair is attempted, a [`TurnEvent::ThoughtDelta`] is emitted with
193/// the original parse error so the user can see what went wrong.
194async fn parse_or_repair_app_server_response(
195    kind: AgentKind,
196    response: &crate::app_server::AppServerTurnResponse,
197    protocol_profile: agent::ProtocolRequestProfile,
198    repair_request: AppServerTurnRequest,
199    client: &Arc<dyn AppServerClient>,
200    events: &mpsc::UnboundedSender<TurnEvent>,
201) -> Result<AppServerParsedTurnResult, AgentError> {
202    let parse_error =
203        match agent::parse_turn_response(kind, &response.assistant_message, protocol_profile) {
204            Ok(parsed) => {
205                return Ok(AppServerParsedTurnResult {
206                    assistant_message: parsed,
207                    provider_conversation_id: response.provider_conversation_id.clone(),
208                    repair_input_tokens: 0,
209                    repair_output_tokens: 0,
210                });
211            }
212            Err(error) => error,
213        };
214
215    let _ = events.send(TurnEvent::ThoughtDelta(format!(
216        "Protocol parse error — retrying: {parse_error}"
217    )));
218
219    let repair_prompt =
220        agent::protocol::build_protocol_repair_prompt(&parse_error, &response.assistant_message);
221
222    let repair_provider_conversation_id = response
223        .provider_conversation_id
224        .clone()
225        .or_else(|| repair_request.provider_conversation_id.clone());
226
227    let repair_turn_request = AppServerTurnRequest {
228        folder: repair_request.folder,
229        live_session_output: None,
230        model: repair_request.model,
231        prompt: crate::model::turn_prompt::TurnPrompt::from_agent_data(repair_prompt),
232        request_kind: repair_request.request_kind,
233        provider_conversation_id: repair_provider_conversation_id,
234        persisted_instruction_conversation_id: None,
235        reasoning_level: repair_request.reasoning_level,
236        session_id: repair_request.session_id,
237    };
238    let (repair_stream_tx, _repair_stream_rx) = mpsc::unbounded_channel();
239    let repair_result = client
240        .run_turn(repair_turn_request, repair_stream_tx)
241        .await
242        .map_err(|error| {
243            AgentError::Backend(format!(
244                "{parse_error}\nprotocol repair transport failed: {error}"
245            ))
246        })?;
247
248    let parsed =
249        agent::parse_turn_response(kind, &repair_result.assistant_message, protocol_profile)
250            .map_err(|error| {
251                AgentError::Backend(format!(
252                    "{parse_error}\nprotocol repair retry also failed: \
253                     {error}\nrepair_response:\n{}",
254                    repair_result.assistant_message
255                ))
256            })?;
257
258    Ok(AppServerParsedTurnResult {
259        assistant_message: parsed,
260        provider_conversation_id: repair_result
261            .provider_conversation_id
262            .or(response.provider_conversation_id.clone()),
263        repair_input_tokens: repair_result.input_tokens,
264        repair_output_tokens: repair_result.output_tokens,
265    })
266}
267
268#[cfg(test)]
269mod tests {
270    use std::path::PathBuf;
271    use std::sync::Arc;
272
273    use tokio::sync::mpsc;
274
275    use super::*;
276    use crate::app_server::{AppServerTurnResponse, MockAppServerClient};
277    use crate::channel::AgentRequestKind;
278    use crate::model::agent::ReasoningLevel;
279    use crate::model::turn_prompt::TurnPromptAttachment;
280
281    fn make_turn_request() -> TurnRequest {
282        TurnRequest {
283            folder: PathBuf::from("/tmp"),
284            live_session_output: None,
285            model: "gpt-5.5".to_string(),
286            request_kind: AgentRequestKind::SessionStart,
287            prompt: "Do something".into(),
288            provider_conversation_id: None,
289            persisted_instruction_conversation_id: None,
290            reasoning_level: ReasoningLevel::default(),
291        }
292    }
293
294    fn make_ok_response(assistant_message: &str) -> AppServerTurnResponse {
295        AppServerTurnResponse {
296            assistant_message: assistant_message.to_string(),
297            context_reset: false,
298            input_tokens: 10,
299            output_tokens: 5,
300            pid: None,
301            provider_conversation_id: None,
302        }
303    }
304
305    #[tokio::test]
306    /// Verifies non-thought assistant deltas are withheld from the unified
307    /// event stream so transcript output is only appended from the final turn
308    /// result.
309    async fn test_run_turn_suppresses_non_thought_assistant_delta_streaming() {
310        // Arrange
311        let mut mock_client = MockAppServerClient::new();
312        mock_client
313            .expect_run_turn()
314            .returning(|_request, stream_tx| {
315                let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
316                    message: "Hello world".to_string(),
317                    phase: None,
318                    is_delta: true,
319                });
320
321                Box::pin(async {
322                    Ok(make_ok_response(
323                        r#"{"answer":"Hello world","questions":[],"summary":null}"#,
324                    ))
325                })
326            });
327        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
328        let (events_tx, mut events_rx) = mpsc::unbounded_channel();
329
330        // Act
331        let result = channel
332            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
333            .await;
334
335        // Assert
336        assert!(result.is_ok());
337        let events = std::iter::from_fn(|| events_rx.try_recv().ok()).collect::<Vec<_>>();
338        assert!(!events.is_empty());
339        assert!(
340            events
341                .iter()
342                .all(|event| matches!(event, TurnEvent::PidUpdate(_))),
343            "only pid events should be emitted, got: {events:?}"
344        );
345    }
346
347    #[tokio::test]
348    /// Verifies completed assistant chunks are also withheld from the unified
349    /// event stream so the transcript only changes when the turn completes.
350    async fn test_run_turn_suppresses_non_delta_assistant_messages() {
351        // Arrange
352        let mut mock_client = MockAppServerClient::new();
353        mock_client
354            .expect_run_turn()
355            .returning(|_request, stream_tx| {
356                let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
357                    message: "Full paragraph   ".to_string(),
358                    phase: None,
359                    is_delta: false,
360                });
361
362                Box::pin(async {
363                    Ok(make_ok_response(
364                        r#"{"answer":"Full paragraph","questions":[],"summary":null}"#,
365                    ))
366                })
367            });
368        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
369        let (events_tx, mut events_rx) = mpsc::unbounded_channel();
370
371        // Act
372        let result = channel
373            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
374            .await;
375
376        // Assert
377        assert!(result.is_ok());
378        let events = std::iter::from_fn(|| events_rx.try_recv().ok()).collect::<Vec<_>>();
379        assert!(!events.is_empty());
380        assert!(
381            events
382                .iter()
383                .all(|event| matches!(event, TurnEvent::PidUpdate(_))),
384            "only pid events should be emitted, got: {events:?}"
385        );
386    }
387
388    #[tokio::test]
389    /// Verifies structured assistant payload chunks are not emitted as live
390    /// transcript output.
391    async fn test_run_turn_suppresses_non_delta_structured_json_streaming() {
392        // Arrange
393        let mut mock_client = MockAppServerClient::new();
394        mock_client
395            .expect_run_turn()
396            .returning(|_request, stream_tx| {
397                let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
398                    message: r#"{"answer":"Done.","questions":[{"text":"Need clarification.","options":[]}],"summary":null}"#.to_string(),
399                    phase: None,
400                    is_delta: false,
401                });
402
403                Box::pin(async {
404                    Ok(make_ok_response(
405                        r#"{"answer":"Done.","questions":[{"text":"Need clarification.","options":[]}],"summary":null}"#,
406                    ))
407                })
408            });
409        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
410        let (events_tx, mut events_rx) = mpsc::unbounded_channel();
411
412        // Act
413        let result = channel
414            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
415            .await;
416
417        // Assert
418        assert!(result.is_ok());
419        let events = std::iter::from_fn(|| events_rx.try_recv().ok()).collect::<Vec<_>>();
420        assert!(!events.is_empty());
421        assert!(
422            events
423                .iter()
424                .all(|event| matches!(event, TurnEvent::PidUpdate(_))),
425            "only pid events should be emitted, got: {events:?}"
426        );
427    }
428
429    #[tokio::test]
430    /// Verifies Codex thought-phase deltas are routed to `ThoughtDelta`.
431    async fn test_run_turn_routes_codex_thinking_delta_to_thought_event() {
432        // Arrange
433        let mut mock_client = MockAppServerClient::new();
434        mock_client
435            .expect_run_turn()
436            .returning(|_request, stream_tx| {
437                let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
438                    message: "Inspecting files".to_string(),
439                    phase: Some("thinking".to_string()),
440                    is_delta: true,
441                });
442
443                Box::pin(async {
444                    Ok(make_ok_response(
445                        r#"{"answer":"Done.","questions":[],"summary":null}"#,
446                    ))
447                })
448            });
449        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
450        let (events_tx, mut events_rx) = mpsc::unbounded_channel();
451
452        // Act
453        let result = channel
454            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
455            .await;
456
457        // Assert
458        assert!(result.is_ok());
459        let event = events_rx.try_recv().expect("should have received an event");
460        assert_eq!(
461            event,
462            TurnEvent::ThoughtDelta("Inspecting files".to_string())
463        );
464    }
465
466    #[tokio::test]
467    /// Verifies Codex thought-phase matching is case-insensitive for streamed
468    /// thought routing.
469    async fn test_run_turn_routes_uppercase_codex_thinking_delta_to_thought_event() {
470        // Arrange
471        let mut mock_client = MockAppServerClient::new();
472        mock_client
473            .expect_run_turn()
474            .returning(|_request, stream_tx| {
475                let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
476                    message: "Inspecting files".to_string(),
477                    phase: Some("Thinking".to_string()),
478                    is_delta: true,
479                });
480
481                Box::pin(async {
482                    Ok(make_ok_response(
483                        r#"{"answer":"Done.","questions":[],"summary":null}"#,
484                    ))
485                })
486            });
487        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
488        let (events_tx, mut events_rx) = mpsc::unbounded_channel();
489
490        // Act
491        let result = channel
492            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
493            .await;
494
495        // Assert
496        assert!(result.is_ok());
497        let event = events_rx.try_recv().expect("should have received an event");
498        assert_eq!(
499            event,
500            TurnEvent::ThoughtDelta("Inspecting files".to_string())
501        );
502    }
503
504    #[tokio::test]
505    /// Verifies `ProgressUpdate` events drive the transient thinking loader.
506    async fn test_run_turn_routes_progress_update_events_to_thought_delta() {
507        // Arrange
508        let mut mock_client = MockAppServerClient::new();
509        mock_client
510            .expect_run_turn()
511            .returning(|_request, stream_tx| {
512                let _ = stream_tx.send(AppServerStreamEvent::ProgressUpdate(
513                    "Running tool".to_string(),
514                ));
515
516                Box::pin(async {
517                    Ok(make_ok_response(
518                        r#"{"answer":"","questions":[],"summary":null}"#,
519                    ))
520                })
521            });
522        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
523        let (events_tx, mut events_rx) = mpsc::unbounded_channel();
524
525        // Act
526        let result = channel
527            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
528            .await;
529
530        // Assert
531        assert!(result.is_ok());
532        let event = events_rx
533            .try_recv()
534            .expect("should have received a progress event");
535        assert_eq!(event, TurnEvent::ThoughtDelta("Running tool".to_string()));
536    }
537
538    #[tokio::test]
539    /// Verifies strict session-turn responses synthesize an empty summary when
540    /// the provider returns `summary: null`.
541    async fn test_run_turn_fills_missing_summary_for_session_turn() {
542        // Arrange
543        let mut mock_client = MockAppServerClient::new();
544        mock_client
545            .expect_run_turn()
546            .returning(|_request, _stream_tx| {
547                Box::pin(async {
548                    Ok(make_ok_response(
549                        r#"{"answer":"Done.","questions":[],"summary":null}"#,
550                    ))
551                })
552            });
553        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Antigravity);
554        let (events_tx, _events_rx) = mpsc::unbounded_channel();
555
556        // Act
557        let result = channel
558            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
559            .await
560            .expect("turn should succeed");
561
562        // Assert
563        assert_eq!(
564            result.assistant_message.summary,
565            Some(agent::protocol::AgentResponseSummary {
566                turn: String::new(),
567                session: String::new(),
568            })
569        );
570    }
571
572    #[tokio::test]
573    /// Verifies whitespace-only `AssistantMessage` does not emit a thinking
574    /// update.
575    async fn test_run_turn_skips_whitespace_only_assistant_messages() {
576        // Arrange
577        let mut mock_client = MockAppServerClient::new();
578        mock_client
579            .expect_run_turn()
580            .returning(|_request, stream_tx| {
581                let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
582                    message: "   \n  ".to_string(),
583                    phase: None,
584                    is_delta: true,
585                });
586
587                Box::pin(async {
588                    Ok(make_ok_response(
589                        r#"{"answer":"","questions":[],"summary":null}"#,
590                    ))
591                })
592            });
593        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
594        let (events_tx, mut events_rx) = mpsc::unbounded_channel();
595
596        // Act
597        let result = channel
598            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
599            .await;
600
601        // Assert
602        assert!(result.is_ok());
603        while let Ok(event) = events_rx.try_recv() {
604            assert!(
605                !matches!(event, TurnEvent::ThoughtDelta(_)),
606                "no ThoughtDelta should be emitted for whitespace-only messages, got: {event:?}"
607            );
608        }
609    }
610
611    #[tokio::test]
612    /// Verifies delta protocol JSON fragments do not emit transient loader
613    /// updates.
614    async fn test_run_turn_skips_delta_protocol_json_fragments() {
615        // Arrange
616        let mut mock_client = MockAppServerClient::new();
617        mock_client
618            .expect_run_turn()
619            .returning(|_request, stream_tx| {
620                let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
621                    message: r#"{"answer":"#.to_string(),
622                    phase: None,
623                    is_delta: true,
624                });
625
626                Box::pin(async {
627                    Ok(make_ok_response(
628                        r#"{"answer":"Final answer.","questions":[],"summary":null}"#,
629                    ))
630                })
631            });
632        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
633        let (events_tx, mut events_rx) = mpsc::unbounded_channel();
634
635        // Act
636        let result = channel
637            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
638            .await
639            .expect("turn should succeed");
640
641        // Assert
642        assert_eq!(result.assistant_message.to_display_text(), "Final answer.");
643        while let Ok(event) = events_rx.try_recv() {
644            assert!(
645                !matches!(event, TurnEvent::ThoughtDelta(_)),
646                "no ThoughtDelta should be emitted for protocol fragments, got: {event:?}"
647            );
648        }
649    }
650
651    #[tokio::test]
652    /// Verifies app-server providers suppress streamed assistant chunks and
653    /// rely on the final parsed payload.
654    async fn test_run_turn_app_server_suppresses_streamed_assistant_messages() {
655        // Arrange
656        let mut mock_client = MockAppServerClient::new();
657        mock_client
658            .expect_run_turn()
659            .returning(|_request, stream_tx| {
660                let _ = stream_tx.send(AppServerStreamEvent::AssistantMessage {
661                    message: "streamed plain text".to_string(),
662                    phase: None,
663                    is_delta: true,
664                });
665
666                Box::pin(async {
667                    Ok(make_ok_response(
668                        r#"{"answer":"Final structured output.","questions":[],"summary":null}"#,
669                    ))
670                })
671            });
672        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
673        let (events_tx, mut events_rx) = mpsc::unbounded_channel();
674
675        // Act
676        let result = channel
677            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
678            .await
679            .expect("turn should succeed");
680
681        // Assert
682        assert_eq!(
683            result.assistant_message.to_display_text(),
684            "Final structured output."
685        );
686        while let Ok(event) = events_rx.try_recv() {
687            assert!(
688                !matches!(event, TurnEvent::ThoughtDelta(_)),
689                "no ThoughtDelta should be emitted for plain assistant deltas, got: {event:?}"
690            );
691        }
692    }
693
694    #[tokio::test]
695    /// Verifies app-server turns surface invalid structured output after both
696    /// the original parse and the protocol-repair retry fail.
697    async fn test_run_turn_returns_error_for_invalid_structured_output() {
698        // Arrange
699        let mut mock_client = MockAppServerClient::new();
700        mock_client
701            .expect_run_turn()
702            .times(2)
703            .returning(|request, _stream_tx| {
704                assert_eq!(request.request_kind, AgentRequestKind::SessionStart);
705
706                Box::pin(async { Ok(make_ok_response("plain non-json response")) })
707            });
708        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
709        let (events_tx, _events_rx) = mpsc::unbounded_channel();
710
711        // Act
712        let error = channel
713            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
714            .await
715            .expect_err("invalid structured output should fail");
716
717        // Assert
718        let error_message = error.to_string();
719        assert!(error_message.contains("did not match the required JSON schema"));
720        assert!(error_message.contains("response:\nplain non-json response"));
721    }
722
723    #[tokio::test]
724    /// Verifies app-server turns recover valid output when the initial parse
725    /// fails but the protocol-repair retry returns valid protocol JSON.
726    async fn test_run_turn_recovers_valid_output_via_protocol_repair() {
727        // Arrange
728        let call_counter = Arc::new(std::sync::atomic::AtomicUsize::new(0));
729        let mut mock_client = MockAppServerClient::new();
730        mock_client.expect_run_turn().times(2).returning({
731            let counter = Arc::clone(&call_counter);
732
733            move |_request, _stream_tx| {
734                let call_number = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
735
736                if call_number == 0 {
737                    Box::pin(async { Ok(make_ok_response("plain non-json response")) })
738                } else {
739                    Box::pin(async {
740                        Ok(make_ok_response(
741                            r#"{"answer":"Repaired response","questions":[],"summary":null}"#,
742                        ))
743                    })
744                }
745            }
746        });
747        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
748        let (events_tx, _events_rx) = mpsc::unbounded_channel();
749
750        // Act
751        let result = channel
752            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
753            .await
754            .expect("repair retry should succeed");
755
756        // Assert
757        assert_eq!(
758            result.assistant_message.to_display_text(),
759            "Repaired response"
760        );
761    }
762
763    #[tokio::test]
764    /// Verifies app-server turns pass pasted image prompt payloads through to
765    /// the underlying app-server client.
766    async fn test_run_turn_allows_image_attachments() {
767        // Arrange
768        let mut mock_client = MockAppServerClient::new();
769        mock_client
770            .expect_run_turn()
771            .times(1)
772            .returning(|request, _stream_tx| {
773                assert_eq!(request.prompt.attachments.len(), 1);
774
775                Box::pin(async {
776                    Ok(make_ok_response(
777                        r#"{"answer":"codex ok","questions":[],"summary":null}"#,
778                    ))
779                })
780            });
781        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
782        let (events_tx, _events_rx) = mpsc::unbounded_channel();
783        let mut request = make_turn_request();
784        request.prompt.attachments.push(TurnPromptAttachment {
785            placeholder: "[Image #1]".to_string(),
786            local_image_path: PathBuf::from("/tmp/image.png"),
787        });
788
789        // Act
790        let result = channel
791            .run_turn("sess-1".to_string(), request, events_tx)
792            .await
793            .expect("turn should succeed");
794
795        // Assert
796        assert_eq!(result.assistant_message.to_display_text(), "codex ok");
797    }
798
799    #[tokio::test]
800    /// Verifies Codex turns surface invalid plain-text output after both the
801    /// original parse and the protocol-repair retry fail.
802    async fn test_run_turn_codex_rejects_plain_text_after_repair_retry() {
803        // Arrange
804        let mut mock_client = MockAppServerClient::new();
805        mock_client
806            .expect_run_turn()
807            .times(2)
808            .returning(|_request, _stream_tx| Box::pin(async { Ok(make_ok_response("plain")) }));
809        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
810        let (events_tx, _events_rx) = mpsc::unbounded_channel();
811
812        // Act
813        let error = channel
814            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
815            .await
816            .expect_err("plain-text turn should fail");
817
818        // Assert
819        let error_message = error.to_string();
820        assert!(error_message.contains("did not match the required JSON schema"));
821        assert!(error_message.contains("response:\nplain"));
822    }
823
824    #[tokio::test]
825    /// Verifies client turn failure propagates as `Err(AgentError)`.
826    async fn test_run_turn_client_failure_returns_agent_error() {
827        // Arrange
828        let mut mock_client = MockAppServerClient::new();
829        mock_client
830            .expect_run_turn()
831            .returning(|_request, _stream_tx| {
832                Box::pin(async {
833                    Err(crate::app_server::AppServerError::Provider(
834                        "server timeout".to_string(),
835                    ))
836                })
837            });
838        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
839        let (events_tx, _events_rx) = mpsc::unbounded_channel();
840
841        // Act
842        let result = channel
843            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
844            .await;
845
846        // Assert
847        let error_message = result
848            .expect_err("expected Err on server timeout")
849            .to_string();
850        assert!(error_message.contains("server timeout"));
851    }
852
853    #[tokio::test]
854    /// Verifies `TurnResult` carries the correct token counts and context-reset
855    /// flag from the underlying `AppServerTurnResponse`.
856    async fn test_run_turn_returns_correct_token_counts_and_context_reset() {
857        // Arrange
858        let mut mock_client = MockAppServerClient::new();
859        mock_client
860            .expect_run_turn()
861            .returning(|_request, _stream_tx| {
862                Box::pin(async {
863                    Ok(AppServerTurnResponse {
864                        assistant_message: r#"{"answer":"Result","questions":[],"summary":null}"#
865                            .to_string(),
866                        context_reset: true,
867                        input_tokens: 100,
868                        output_tokens: 50,
869                        pid: Some(1234),
870                        provider_conversation_id: None,
871                    })
872                })
873            });
874        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
875        let (events_tx, _events_rx) = mpsc::unbounded_channel();
876
877        // Act
878        let result = channel
879            .run_turn("sess-1".to_string(), make_turn_request(), events_tx)
880            .await
881            .expect("turn should succeed");
882
883        // Assert
884        assert_eq!(result.assistant_message.to_display_text(), "Result");
885        assert!(result.context_reset);
886        assert_eq!(result.input_tokens, 100);
887        assert_eq!(result.output_tokens, 50);
888    }
889
890    #[tokio::test]
891    /// Verifies `provider_conversation_id` is forwarded from `TurnRequest` to
892    /// the underlying `AppServerTurnRequest` and propagated back from the
893    /// response into the returned `TurnResult`.
894    async fn test_run_turn_passes_and_returns_provider_conversation_id() {
895        // Arrange
896        let mut mock_client = MockAppServerClient::new();
897        mock_client
898            .expect_run_turn()
899            .returning(|request, _stream_tx| {
900                assert_eq!(
901                    request.provider_conversation_id,
902                    Some("thread-abc".to_string()),
903                    "request should carry the provider conversation id"
904                );
905                assert_eq!(
906                    request.reasoning_level,
907                    ReasoningLevel::Medium,
908                    "request should carry the codex reasoning level"
909                );
910
911                Box::pin(async {
912                    Ok(AppServerTurnResponse {
913                        assistant_message: r#"{"answer":"ok","questions":[],"summary":null}"#
914                            .to_string(),
915                        context_reset: false,
916                        input_tokens: 1,
917                        output_tokens: 1,
918                        pid: Some(42),
919                        provider_conversation_id: Some("thread-xyz".to_string()),
920                    })
921                })
922            });
923        let channel = AppServerAgentChannel::new(Arc::new(mock_client), AgentKind::Codex);
924        let (events_tx, mut events_rx) = mpsc::unbounded_channel();
925        let mut request = make_turn_request();
926        request.reasoning_level = ReasoningLevel::Medium;
927        request.provider_conversation_id = Some("thread-abc".to_string());
928
929        // Act
930        let result = channel
931            .run_turn("sess-1".to_string(), request, events_tx)
932            .await
933            .expect("turn should succeed");
934
935        // Assert
936        assert_eq!(
937            result.provider_conversation_id,
938            Some("thread-xyz".to_string()),
939            "result should carry the provider conversation id from the response"
940        );
941
942        // Verify PID event was emitted from the response.
943        let mut pid_event_seen = false;
944        while let Ok(event) = events_rx.try_recv() {
945            if matches!(event, TurnEvent::PidUpdate(Some(42))) {
946                pid_event_seen = true;
947            }
948        }
949        assert!(pid_event_seen, "should emit PidUpdate from response pid");
950    }
951}
ag_agent/channel/app_server.rs

ag_agent/channel/
app_server.rs
