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_session_output: None,
123            model: AgentModel::Gpt55.as_str().to_string(),
124            prompt: "prompt".into(),
125            request_kind: crate::channel::AgentRequestKind::SessionStart,
126            provider_conversation_id: None,
127            persisted_instruction_conversation_id: None,
128            reasoning_level: ReasoningLevel::default(),
129            session_id: "session-1".to_string(),
130        };
131
132        // Act
133        let response = app_server_client
134            .run_turn(request, stream_tx)
135            .await
136            .expect("run_turn should succeed");
137
138        // Assert
139        assert_eq!(response.assistant_message, "codex");
140    }
141
142    #[tokio::test]
143    async fn run_turn_routes_gemini_models_to_gemini_client() {
144        // Arrange
145        let mut codex_client = MockAppServerClient::new();
146        codex_client.expect_run_turn().times(0);
147        let mut gemini_client = MockAppServerClient::new();
148        gemini_client.expect_run_turn().times(1).returning(|_, _| {
149            Box::pin(async {
150                Ok(AppServerTurnResponse {
151                    assistant_message: "gemini".to_string(),
152                    context_reset: false,
153                    input_tokens: 3,
154                    output_tokens: 4,
155                    pid: Some(222),
156                    provider_conversation_id: Some("thread-gemini".to_string()),
157                })
158            })
159        });
160        let app_server_client = RoutingAppServerClient::new_with_clients(
161            Arc::new(codex_client),
162            Arc::new(gemini_client),
163        );
164        let (stream_tx, _stream_rx) = mpsc::unbounded_channel();
165        let request = AppServerTurnRequest {
166            folder: std::env::temp_dir(),
167            live_session_output: None,
168            model: AgentModel::Gemini3FlashPreview.as_str().to_string(),
169            prompt: "prompt".into(),
170            request_kind: crate::channel::AgentRequestKind::SessionStart,
171            provider_conversation_id: None,
172            persisted_instruction_conversation_id: None,
173            reasoning_level: ReasoningLevel::default(),
174            session_id: "session-1".to_string(),
175        };
176
177        // Act
178        let response = app_server_client
179            .run_turn(request, stream_tx)
180            .await
181            .expect("run_turn should succeed");
182
183        // Assert
184        assert_eq!(response.assistant_message, "gemini");
185    }
186
187    #[tokio::test]
188    async fn run_turn_returns_error_for_unknown_model() {
189        // Arrange
190        let mut codex_client = MockAppServerClient::new();
191        codex_client.expect_run_turn().times(0);
192        let mut gemini_client = MockAppServerClient::new();
193        gemini_client.expect_run_turn().times(0);
194        let app_server_client = RoutingAppServerClient::new_with_clients(
195            Arc::new(codex_client),
196            Arc::new(gemini_client),
197        );
198        let (stream_tx, _stream_rx) = mpsc::unbounded_channel();
199        let request = AppServerTurnRequest {
200            folder: std::env::temp_dir(),
201            live_session_output: None,
202            model: "unknown-model".to_string(),
203            prompt: "prompt".into(),
204            request_kind: crate::channel::AgentRequestKind::SessionStart,
205            provider_conversation_id: None,
206            persisted_instruction_conversation_id: None,
207            reasoning_level: ReasoningLevel::default(),
208            session_id: "session-1".to_string(),
209        };
210
211        // Act
212        let result = app_server_client.run_turn(request, stream_tx).await;
213        let error = result.expect_err("should fail for unknown model");
214
215        // Assert
216        assert!(
217            error.to_string().contains("unknown model"),
218            "error should mention unknown model"
219        );
220    }
221
222    #[tokio::test]
223    async fn shutdown_session_propagates_to_codex_client() {
224        // Arrange
225        let mut codex_client = MockAppServerClient::new();
226        codex_client
227            .expect_shutdown_session()
228            .times(1)
229            .returning(|_| Box::pin(async {}));
230        let mut gemini_client = MockAppServerClient::new();
231        gemini_client
232            .expect_shutdown_session()
233            .times(1)
234            .returning(|_| Box::pin(async {}));
235        let app_server_client = RoutingAppServerClient::new_with_clients(
236            Arc::new(codex_client),
237            Arc::new(gemini_client),
238        );
239
240        // Act
241        app_server_client
242            .shutdown_session("session-1".to_string())
243            .await;
244
245        // Assert
246    }
247}
ag_agent/app_server_router.rs

ag_agent/
app_server_router.rs
