ag_agent/
app_server_router.rs

1//! Provider router for app-server backed session execution.
2
3use std::sync::Arc;
4
5use tokio::sync::mpsc;
6
7use crate::agent;
8use crate::agent::app_server::{RealCodexAppServerClient, RealGeminiAcpClient};
9use crate::app_server::{
10    AppServerClient, AppServerError, AppServerFuture, AppServerStreamEvent, AppServerTurnRequest,
11    AppServerTurnResponse,
12};
13use crate::model::agent::AgentKind;
14
15/// Production router that dispatches app-server turns by model provider.
16pub struct RoutingAppServerClient {
17    codex_client: Arc<dyn AppServerClient>,
18    gemini_client: Arc<dyn AppServerClient>,
19}
20
21impl RoutingAppServerClient {
22    /// Creates a router backed by production Codex and Gemini clients.
23    pub fn new() -> Self {
24        Self::new_with_clients(
25            Arc::new(RealCodexAppServerClient::new()),
26            Arc::new(RealGeminiAcpClient::new()),
27        )
28    }
29
30    /// Creates a router with injected provider clients.
31    pub fn new_with_clients(
32        codex_client: Arc<dyn AppServerClient>,
33        gemini_client: Arc<dyn AppServerClient>,
34    ) -> Self {
35        Self {
36            codex_client,
37            gemini_client,
38        }
39    }
40}
41
42impl Default for RoutingAppServerClient {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl AppServerClient for RoutingAppServerClient {
49    fn run_turn(
50        &self,
51        request: AppServerTurnRequest,
52        stream_tx: mpsc::UnboundedSender<AppServerStreamEvent>,
53    ) -> AppServerFuture<Result<AppServerTurnResponse, AppServerError>> {
54        let codex_client = Arc::clone(&self.codex_client);
55        let gemini_client = Arc::clone(&self.gemini_client);
56        let model = request.model.clone();
57
58        Box::pin(async move {
59            let provider_kind = agent::provider_kind_for_model(&model).map_err(|error| {
60                AppServerError::Provider(format!("App-server routing failed for {error}"))
61            })?;
62
63            match provider_kind {
64                AgentKind::Codex => codex_client.run_turn(request, stream_tx).await,
65                AgentKind::Gemini => gemini_client.run_turn(request, stream_tx).await,
66                AgentKind::Antigravity => Err(AppServerError::Provider(
67                    "Antigravity does not support app-server session execution".to_string(),
68                )),
69                AgentKind::Claude => Err(AppServerError::Provider(
70                    "Claude does not support app-server session execution".to_string(),
71                )),
72            }
73        })
74    }
75
76    /// Shuts down the session on both provider clients concurrently so one
77    /// slow backend cannot block cleanup of the other.
78    fn shutdown_session(&self, session_id: String) -> AppServerFuture<()> {
79        let codex_client = Arc::clone(&self.codex_client);
80        let gemini_client = Arc::clone(&self.gemini_client);
81
82        Box::pin(async move {
83            tokio::join!(
84                codex_client.shutdown_session(session_id.clone()),
85                gemini_client.shutdown_session(session_id),
86            );
87        })
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use crate::app_server::MockAppServerClient;
95    use crate::model::agent::{AgentModel, ReasoningLevel};
96
97    #[tokio::test]
98    async fn run_turn_routes_codex_models_to_codex_client() {
99        // Arrange
100        let mut codex_client = MockAppServerClient::new();
101        codex_client.expect_run_turn().times(1).returning(|_, _| {
102            Box::pin(async {
103                Ok(AppServerTurnResponse {
104                    assistant_message: "codex".to_string(),
105                    context_reset: false,
106                    input_tokens: 1,
107                    output_tokens: 2,
108                    pid: Some(111),
109                    provider_conversation_id: Some("thread-codex".to_string()),
110                })
111            })
112        });
113        let mut gemini_client = MockAppServerClient::new();
114        gemini_client.expect_run_turn().times(0);
115        let app_server_client = RoutingAppServerClient::new_with_clients(
116            Arc::new(codex_client),
117            Arc::new(gemini_client),
118        );
119        let (stream_tx, _stream_rx) = mpsc::unbounded_channel();
120        let request = AppServerTurnRequest {
121            folder: std::env::temp_dir(),
122            live_transcript: None,
123            main_checkout_root: None,
124            model: AgentModel::Gpt55.as_str().to_string(),
125            prompt: "prompt".into(),
126            request_kind: crate::channel::AgentRequestKind::SessionStart,
127            replay_transcript: None,
128            provider_conversation_id: None,
129            persisted_instruction_conversation_id: None,
130            reasoning_level: ReasoningLevel::default(),
131            session_id: "session-1".to_string(),
132        };
133
134        // Act
135        let response = app_server_client
136            .run_turn(request, stream_tx)
137            .await
138            .expect("run_turn should succeed");
139
140        // Assert
141        assert_eq!(response.assistant_message, "codex");
142    }
143
144    #[tokio::test]
145    async fn run_turn_routes_gemini_models_to_gemini_client() {
146        // Arrange
147        let mut codex_client = MockAppServerClient::new();
148        codex_client.expect_run_turn().times(0);
149        let mut gemini_client = MockAppServerClient::new();
150        gemini_client.expect_run_turn().times(1).returning(|_, _| {
151            Box::pin(async {
152                Ok(AppServerTurnResponse {
153                    assistant_message: "gemini".to_string(),
154                    context_reset: false,
155                    input_tokens: 3,
156                    output_tokens: 4,
157                    pid: Some(222),
158                    provider_conversation_id: Some("thread-gemini".to_string()),
159                })
160            })
161        });
162        let app_server_client = RoutingAppServerClient::new_with_clients(
163            Arc::new(codex_client),
164            Arc::new(gemini_client),
165        );
166        let (stream_tx, _stream_rx) = mpsc::unbounded_channel();
167        let request = AppServerTurnRequest {
168            folder: std::env::temp_dir(),
169            live_transcript: None,
170            main_checkout_root: None,
171            model: AgentModel::Gemini3FlashPreview.as_str().to_string(),
172            prompt: "prompt".into(),
173            request_kind: crate::channel::AgentRequestKind::SessionStart,
174            replay_transcript: None,
175            provider_conversation_id: None,
176            persisted_instruction_conversation_id: None,
177            reasoning_level: ReasoningLevel::default(),
178            session_id: "session-1".to_string(),
179        };
180
181        // Act
182        let response = app_server_client
183            .run_turn(request, stream_tx)
184            .await
185            .expect("run_turn should succeed");
186
187        // Assert
188        assert_eq!(response.assistant_message, "gemini");
189    }
190
191    #[tokio::test]
192    async fn run_turn_returns_error_for_unknown_model() {
193        // Arrange
194        let mut codex_client = MockAppServerClient::new();
195        codex_client.expect_run_turn().times(0);
196        let mut gemini_client = MockAppServerClient::new();
197        gemini_client.expect_run_turn().times(0);
198        let app_server_client = RoutingAppServerClient::new_with_clients(
199            Arc::new(codex_client),
200            Arc::new(gemini_client),
201        );
202        let (stream_tx, _stream_rx) = mpsc::unbounded_channel();
203        let request = AppServerTurnRequest {
204            folder: std::env::temp_dir(),
205            live_transcript: None,
206            main_checkout_root: None,
207            model: "unknown-model".to_string(),
208            prompt: "prompt".into(),
209            request_kind: crate::channel::AgentRequestKind::SessionStart,
210            replay_transcript: None,
211            provider_conversation_id: None,
212            persisted_instruction_conversation_id: None,
213            reasoning_level: ReasoningLevel::default(),
214            session_id: "session-1".to_string(),
215        };
216
217        // Act
218        let result = app_server_client.run_turn(request, stream_tx).await;
219        let error = result.expect_err("should fail for unknown model");
220
221        // Assert
222        assert!(
223            error.to_string().contains("unknown model"),
224            "error should mention unknown model"
225        );
226    }
227
228    #[tokio::test]
229    async fn shutdown_session_propagates_to_codex_client() {
230        // Arrange
231        let mut codex_client = MockAppServerClient::new();
232        codex_client
233            .expect_shutdown_session()
234            .times(1)
235            .returning(|_| Box::pin(async {}));
236        let mut gemini_client = MockAppServerClient::new();
237        gemini_client
238            .expect_shutdown_session()
239            .times(1)
240            .returning(|_| Box::pin(async {}));
241        let app_server_client = RoutingAppServerClient::new_with_clients(
242            Arc::new(codex_client),
243            Arc::new(gemini_client),
244        );
245
246        // Act
247        app_server_client
248            .shutdown_session("session-1".to_string())
249            .await;
250
251        // Assert
252    }
253}
ag_agent/app_server_router.rs

ag_agent/
app_server_router.rs
