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

ag_agent/channel/
app_server.rs
