1use std::collections::BTreeMap;
2use std::collections::HashMap;
3use std::num::NonZeroU64;
4use std::time::Duration;
5
6use codex_model_provider_info::ModelProviderInfo;
7use codex_model_provider_info::WireApi;
8use codex_protocol::config_types::ModelProviderAuthInfo;
9use codex_utils_absolute_path::AbsolutePathBuf;
10
11use super::SessionThreadConfig;
12use super::ThreadConfigContext;
13use super::ThreadConfigLoadError;
14use super::ThreadConfigLoadErrorCode;
15use super::ThreadConfigLoader;
16use super::ThreadConfigLoaderFuture;
17use super::ThreadConfigSource;
18use super::UserThreadConfig;
19use proto::thread_config_loader_client::ThreadConfigLoaderClient;
20
21#[path = "proto/codex.thread_config.v1.rs"]
22mod proto;
23
24const REMOTE_THREAD_CONFIG_LOAD_TIMEOUT: Duration = Duration::from_secs(5);
25
26#[derive(Clone, Debug)]
28pub struct RemoteThreadConfigLoader {
29 endpoint: String,
30}
31
32impl RemoteThreadConfigLoader {
33 pub fn new(endpoint: impl Into<String>) -> Self {
34 Self {
35 endpoint: endpoint.into(),
36 }
37 }
38
39 async fn client(
40 &self,
41 ) -> Result<ThreadConfigLoaderClient<tonic::transport::Channel>, ThreadConfigLoadError> {
42 ThreadConfigLoaderClient::connect(self.endpoint.clone())
43 .await
44 .map_err(|err| {
45 ThreadConfigLoadError::new(
46 ThreadConfigLoadErrorCode::RequestFailed,
47 None,
48 format!("failed to connect to remote thread config loader: {err}"),
49 )
50 })
51 }
52
53 async fn load(
54 &self,
55 context: ThreadConfigContext,
56 ) -> Result<Vec<ThreadConfigSource>, ThreadConfigLoadError> {
57 let response = self
58 .client()
59 .await?
60 .load(load_thread_config_request(context))
61 .await
62 .map_err(remote_status_to_error)?
63 .into_inner();
64
65 response
66 .sources
67 .into_iter()
68 .map(thread_config_source_from_proto)
69 .collect()
70 }
71}
72
73impl ThreadConfigLoader for RemoteThreadConfigLoader {
74 fn load(
75 &self,
76 context: ThreadConfigContext,
77 ) -> ThreadConfigLoaderFuture<'_, Vec<ThreadConfigSource>> {
78 Box::pin(RemoteThreadConfigLoader::load(self, context))
79 }
80}
81
82fn load_thread_config_request(
83 context: ThreadConfigContext,
84) -> tonic::Request<proto::LoadThreadConfigRequest> {
85 let mut request = tonic::Request::new(proto::LoadThreadConfigRequest {
86 thread_id: context.thread_id,
87 cwd: context.cwd.map(|cwd| cwd.to_string_lossy().into_owned()),
88 });
89 request.set_timeout(REMOTE_THREAD_CONFIG_LOAD_TIMEOUT);
90 request
91}
92
93fn remote_status_to_error(status: tonic::Status) -> ThreadConfigLoadError {
94 let code = match status.code() {
95 tonic::Code::Unauthenticated | tonic::Code::PermissionDenied => {
96 ThreadConfigLoadErrorCode::Auth
97 }
98 tonic::Code::DeadlineExceeded => ThreadConfigLoadErrorCode::Timeout,
99 tonic::Code::Ok
100 | tonic::Code::Cancelled
101 | tonic::Code::Unknown
102 | tonic::Code::InvalidArgument
103 | tonic::Code::NotFound
104 | tonic::Code::AlreadyExists
105 | tonic::Code::ResourceExhausted
106 | tonic::Code::FailedPrecondition
107 | tonic::Code::Aborted
108 | tonic::Code::OutOfRange
109 | tonic::Code::Unimplemented
110 | tonic::Code::Internal
111 | tonic::Code::Unavailable
112 | tonic::Code::DataLoss => ThreadConfigLoadErrorCode::RequestFailed,
113 };
114 ThreadConfigLoadError::new(
115 code,
116 None,
117 format!("remote thread config request failed: {status}"),
118 )
119}
120
121fn thread_config_source_from_proto(
122 source: proto::ThreadConfigSource,
123) -> Result<ThreadConfigSource, ThreadConfigLoadError> {
124 match source.source {
125 Some(proto::thread_config_source::Source::Session(config)) => {
126 session_thread_config_from_proto(config).map(ThreadConfigSource::Session)
127 }
128 Some(proto::thread_config_source::Source::User(_)) => {
129 Ok(ThreadConfigSource::User(UserThreadConfig::default()))
130 }
131 None => Err(parse_error("remote thread config omitted source payload")),
132 }
133}
134
135fn session_thread_config_from_proto(
136 config: proto::SessionThreadConfig,
137) -> Result<SessionThreadConfig, ThreadConfigLoadError> {
138 let model_providers = config
139 .model_providers
140 .into_iter()
141 .map(model_provider_from_proto)
142 .collect::<Result<HashMap<_, _>, _>>()?;
143
144 Ok(SessionThreadConfig {
145 model_provider: config.model_provider,
146 model_providers,
147 features: config.features.into_iter().collect::<BTreeMap<_, _>>(),
148 })
149}
150
151fn model_provider_from_proto(
152 provider: proto::ModelProvider,
153) -> Result<(String, ModelProviderInfo), ThreadConfigLoadError> {
154 if provider.id.is_empty() {
155 return Err(parse_error(
156 "remote thread config returned model provider without an id",
157 ));
158 }
159 let id = provider.id;
160 let wire_api = match proto::WireApi::try_from(provider.wire_api) {
161 Ok(proto::WireApi::Responses) => WireApi::Responses,
162 Ok(proto::WireApi::Unspecified) => {
163 return Err(parse_error("remote thread config omitted wire_api"));
164 }
165 Err(_) => {
166 return Err(parse_error(format!(
167 "remote thread config returned unknown wire_api: {}",
168 provider.wire_api
169 )));
170 }
171 };
172 let info = ModelProviderInfo {
173 name: provider.name,
174 base_url: provider.base_url,
175 env_key: provider.env_key,
176 env_key_instructions: provider.env_key_instructions,
177 experimental_bearer_token: provider.experimental_bearer_token,
178 auth: provider
179 .auth
180 .map(model_provider_auth_from_proto)
181 .transpose()?,
182 aws: None,
183 wire_api,
184 query_params: provider.query_params.map(|map| map.values),
185 http_headers: provider.http_headers.map(|map| map.values),
186 env_http_headers: provider.env_http_headers.map(|map| map.values),
187 request_max_retries: provider.request_max_retries,
188 stream_max_retries: provider.stream_max_retries,
189 stream_idle_timeout_ms: provider.stream_idle_timeout_ms,
190 websocket_connect_timeout_ms: provider.websocket_connect_timeout_ms,
191 requires_openai_auth: provider.requires_openai_auth,
192 supports_websockets: provider.supports_websockets,
193 };
194 Ok((id, info))
195}
196
197#[cfg(test)]
198fn model_provider_to_proto(
199 id: impl Into<String>,
200 provider: ModelProviderInfo,
201) -> proto::ModelProvider {
202 let ModelProviderInfo {
203 name,
204 base_url,
205 env_key,
206 env_key_instructions,
207 experimental_bearer_token,
208 auth,
209 aws: _,
210 wire_api,
211 query_params,
212 http_headers,
213 env_http_headers,
214 request_max_retries,
215 stream_max_retries,
216 stream_idle_timeout_ms,
217 websocket_connect_timeout_ms,
218 requires_openai_auth,
219 supports_websockets,
220 } = provider;
221
222 proto::ModelProvider {
223 id: id.into(),
224 name,
225 base_url,
226 env_key,
227 env_key_instructions,
228 experimental_bearer_token,
229 auth: auth.map(model_provider_auth_to_proto),
230 wire_api: proto_wire_api(wire_api).into(),
231 query_params: query_params.map(proto_string_map),
232 http_headers: http_headers.map(proto_string_map),
233 env_http_headers: env_http_headers.map(proto_string_map),
234 request_max_retries,
235 stream_max_retries,
236 stream_idle_timeout_ms,
237 websocket_connect_timeout_ms,
238 requires_openai_auth,
239 supports_websockets,
240 }
241}
242
243fn model_provider_auth_from_proto(
244 auth: proto::ModelProviderAuthInfo,
245) -> Result<ModelProviderAuthInfo, ThreadConfigLoadError> {
246 let timeout_ms = NonZeroU64::new(auth.timeout_ms)
247 .ok_or_else(|| parse_error("remote thread config returned zero auth timeout_ms"))?;
248 let cwd = AbsolutePathBuf::from_absolute_path_checked(&auth.cwd).map_err(|err| {
249 parse_error(format!(
250 "remote thread config returned invalid auth cwd {:?}: {err}",
251 auth.cwd
252 ))
253 })?;
254
255 Ok(ModelProviderAuthInfo {
256 command: auth.command,
257 args: auth.args,
258 timeout_ms,
259 refresh_interval_ms: auth.refresh_interval_ms,
260 cwd,
261 })
262}
263
264#[cfg(test)]
265fn model_provider_auth_to_proto(auth: ModelProviderAuthInfo) -> proto::ModelProviderAuthInfo {
266 let ModelProviderAuthInfo {
267 command,
268 args,
269 timeout_ms,
270 refresh_interval_ms,
271 cwd,
272 } = auth;
273
274 proto::ModelProviderAuthInfo {
275 command,
276 args,
277 timeout_ms: timeout_ms.get(),
278 refresh_interval_ms,
279 cwd: cwd.to_string_lossy().into_owned(),
280 }
281}
282
283#[cfg(test)]
284fn proto_string_map(values: HashMap<String, String>) -> proto::StringMap {
285 proto::StringMap { values }
286}
287
288#[cfg(test)]
289fn proto_wire_api(wire_api: WireApi) -> proto::WireApi {
290 match wire_api {
291 WireApi::Responses => proto::WireApi::Responses,
292 }
293}
294
295fn parse_error(message: impl Into<String>) -> ThreadConfigLoadError {
296 ThreadConfigLoadError::new(
297 ThreadConfigLoadErrorCode::Parse,
298 None,
299 message.into(),
300 )
301}
302
303#[cfg(test)]
304mod tests {
305 use std::collections::BTreeMap;
306 use std::collections::HashMap;
307 use std::num::NonZeroU64;
308
309 use codex_model_provider_info::ModelProviderInfo;
310 use codex_model_provider_info::WireApi;
311 use codex_protocol::config_types::ModelProviderAuthInfo;
312 use codex_utils_absolute_path::AbsolutePathBuf;
313 use pretty_assertions::assert_eq;
314 use tonic::Request;
315 use tonic::Response;
316 use tonic::Status;
317 use tonic::transport::Server;
318
319 use super::proto::thread_config_loader_server;
320 use super::proto::thread_config_loader_server::ThreadConfigLoaderServer;
321 use super::*;
322 use crate::SessionThreadConfig;
323 use crate::UserThreadConfig;
324
325 struct TestServer {
326 sources: Vec<proto::ThreadConfigSource>,
327 expected_cwd: String,
328 }
329
330 impl TestServer {
331 async fn load(
332 &self,
333 request: Request<proto::LoadThreadConfigRequest>,
334 ) -> Result<Response<proto::LoadThreadConfigResponse>, Status> {
335 assert_eq!(
336 request.into_inner(),
337 proto::LoadThreadConfigRequest {
338 thread_id: Some("thread-1".to_string()),
339 cwd: Some(self.expected_cwd.clone()),
340 }
341 );
342
343 Ok(Response::new(proto::LoadThreadConfigResponse {
344 sources: self.sources.clone(),
345 }))
346 }
347 }
348
349 impl thread_config_loader_server::ThreadConfigLoader for TestServer {
350 fn load<'a, 'async_trait>(
351 &'a self,
352 request: Request<proto::LoadThreadConfigRequest>,
353 ) -> std::pin::Pin<
354 Box<
355 dyn std::future::Future<
356 Output = Result<Response<proto::LoadThreadConfigResponse>, Status>,
357 > + Send
358 + 'async_trait,
359 >,
360 >
361 where
362 'a: 'async_trait,
363 Self: 'async_trait,
364 {
365 Box::pin(TestServer::load(self, request))
366 }
367 }
368
369 #[tokio::test]
370 async fn load_thread_config_calls_remote_service() {
371 let cwd = workspace_dir().join("project");
372 let expected_cwd = cwd.to_string_lossy().into_owned();
373 let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
374 .await
375 .expect("bind test server");
376 let addr = listener.local_addr().expect("test server addr");
377 let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
378 let server = tokio::spawn(async move {
379 Server::builder()
380 .add_service(ThreadConfigLoaderServer::new(TestServer {
381 sources: proto_sources(),
382 expected_cwd,
383 }))
384 .serve_with_incoming_shutdown(
385 tokio_stream::wrappers::TcpListenerStream::new(listener),
386 async {
387 let _ = shutdown_rx.await;
388 },
389 )
390 .await
391 });
392
393 let loader = RemoteThreadConfigLoader::new(format!("http://{addr}"));
394 let loaded = loader
395 .load(ThreadConfigContext {
396 thread_id: Some("thread-1".to_string()),
397 cwd: Some(cwd),
398 })
399 .await;
400
401 let _ = shutdown_tx.send(());
402 server.await.expect("join server").expect("server");
403
404 assert_eq!(loaded.expect("load thread config"), expected_sources());
405 }
406
407 #[test]
408 fn load_thread_config_request_sets_timeout() {
409 let request = load_thread_config_request(ThreadConfigContext::default());
410
411 assert_eq!(
412 request
413 .metadata()
414 .get("grpc-timeout")
415 .and_then(|value| value.to_str().ok()),
416 Some("5000000u")
417 );
418 }
419
420 #[test]
421 fn model_provider_proto_roundtrips_through_domain_type() {
422 let expected = expected_provider();
423 let proto = model_provider_to_proto("local", expected.clone());
424 let (id, actual) = model_provider_from_proto(proto).expect("model provider from proto");
425
426 assert_eq!(id, "local");
427 assert_eq!(actual, expected);
428 }
429
430 fn proto_sources() -> Vec<proto::ThreadConfigSource> {
431 let workspace_cwd = workspace_dir().to_string_lossy().into_owned();
432 vec![
433 proto::ThreadConfigSource {
434 source: Some(proto::thread_config_source::Source::Session(
435 proto::SessionThreadConfig {
436 model_provider: Some("local".to_string()),
437 model_providers: vec![proto::ModelProvider {
438 id: "local".to_string(),
439 name: "Local".to_string(),
440 base_url: Some("http://127.0.0.1:8061/api/codex".to_string()),
441 env_key: None,
442 env_key_instructions: None,
443 experimental_bearer_token: None,
444 auth: Some(proto::ModelProviderAuthInfo {
445 command: "token-helper".to_string(),
446 args: vec!["--json".to_string()],
447 timeout_ms: 5_000,
448 refresh_interval_ms: 300_000,
449 cwd: workspace_cwd,
450 }),
451 wire_api: proto::WireApi::Responses.into(),
452 query_params: Some(proto::StringMap {
453 values: HashMap::from([(
454 "api-version".to_string(),
455 "2026-04-16".to_string(),
456 )]),
457 }),
458 http_headers: Some(proto::StringMap {
459 values: HashMap::from([(
460 "X-Test".to_string(),
461 "enabled".to_string(),
462 )]),
463 }),
464 env_http_headers: Some(proto::StringMap {
465 values: HashMap::from([(
466 "X-Env".to_string(),
467 "LOCAL_HEADER".to_string(),
468 )]),
469 }),
470 request_max_retries: Some(7),
471 stream_max_retries: Some(8),
472 stream_idle_timeout_ms: Some(9_000),
473 websocket_connect_timeout_ms: Some(10_000),
474 requires_openai_auth: false,
475 supports_websockets: true,
476 }],
477 features: HashMap::from([
478 ("plugins".to_string(), false),
479 ("tools".to_string(), true),
480 ]),
481 },
482 )),
483 },
484 proto::ThreadConfigSource {
485 source: Some(proto::thread_config_source::Source::User(
486 proto::UserThreadConfig {},
487 )),
488 },
489 ]
490 }
491
492 fn expected_sources() -> Vec<ThreadConfigSource> {
493 vec![
494 ThreadConfigSource::Session(SessionThreadConfig {
495 model_provider: Some("local".to_string()),
496 model_providers: HashMap::from([("local".to_string(), expected_provider())]),
497 features: BTreeMap::from([
498 ("plugins".to_string(), false),
499 ("tools".to_string(), true),
500 ]),
501 }),
502 ThreadConfigSource::User(UserThreadConfig::default()),
503 ]
504 }
505
506 fn expected_provider() -> ModelProviderInfo {
507 ModelProviderInfo {
508 name: "Local".to_string(),
509 base_url: Some("http://127.0.0.1:8061/api/codex".to_string()),
510 env_key: None,
511 env_key_instructions: None,
512 experimental_bearer_token: None,
513 auth: Some(ModelProviderAuthInfo {
514 command: "token-helper".to_string(),
515 args: vec!["--json".to_string()],
516 timeout_ms: NonZeroU64::new(5_000).expect("non-zero timeout"),
517 refresh_interval_ms: 300_000,
518 cwd: workspace_dir(),
519 }),
520 wire_api: WireApi::Responses,
521 query_params: Some(HashMap::from([(
522 "api-version".to_string(),
523 "2026-04-16".to_string(),
524 )])),
525 http_headers: Some(HashMap::from([(
526 "X-Test".to_string(),
527 "enabled".to_string(),
528 )])),
529 env_http_headers: Some(HashMap::from([(
530 "X-Env".to_string(),
531 "LOCAL_HEADER".to_string(),
532 )])),
533 request_max_retries: Some(7),
534 stream_max_retries: Some(8),
535 stream_idle_timeout_ms: Some(9_000),
536 websocket_connect_timeout_ms: Some(10_000),
537 requires_openai_auth: false,
538 supports_websockets: true,
539 aws: None,
540 }
541 }
542
543 fn workspace_dir() -> AbsolutePathBuf {
544 AbsolutePathBuf::current_dir()
545 .expect("current dir")
546 .join("workspace")
547 }
548}