From 010520de766c4f23f45a540b4c4d5c6c67fbf7ff Mon Sep 17 00:00:00 2001 From: Benjamin Carlsson Date: Wed, 4 Nov 2026 07:48:23 +0000 Subject: [PATCH] Add managed thread prewarming to the cloud client (#59351) ## What changed Add `Client::prewarm_default_thread` to prepare an ordinary coding thread in the default managed cloud environment, with optional model and provider selection. Return the thread identity and resolved settings as `ThreadActivation` without a separate activation RPC, preserving unknown thread protobuf fields. Send prewarm requests once without automatic retries or cleanup, and document that a failed response can leave a prepared thread whose identity the caller did not receive. ## Testing Add coverage for default and explicit model/provider requests, resolved settings and native thread metadata preservation, and rejection of responses missing sandbox policy. GitOrigin-RevId: 425ee5ee6bb30ec77ebdd0079955701a7f9704e4 --- codex-rs/cloud-client/README.md | 10 ++ codex-rs/cloud-client/src/activation.rs | 15 ++- codex-rs/cloud-client/src/client_tests.rs | 93 +++++++++++++++++++ codex-rs/cloud-client/src/lib.rs | 1 + codex-rs/cloud-client/src/prewarm.rs | 108 ++++++++++++++++++++++ 5 files changed, 219 insertions(+), 8 deletions(-) create mode 100644 codex-rs/cloud-client/src/prewarm.rs diff --git a/codex-rs/cloud-client/README.md b/codex-rs/cloud-client/README.md index 3953afac5e..389fd34ff7 100644 --- a/codex-rs/cloud-client/README.md +++ b/codex-rs/cloud-client/README.md @@ -85,3 +85,13 @@ one message after attaching: a persisted client message ID and the text. It uses ```sh cargo run -p codex-cloud-client --example attach -- "$CODEX_CLOUD_GRPC_ENDPOINT" "$CODEX_CLOUD_THREAD_ID" "$CODEX_CLOUD_MESSAGE_ID" 'Explain this project' ``` + +`prewarm_default_thread` prepares an ordinary coding thread in the default managed +cloud environment, with optional model/provider selection. It returns the native +thread identity and resolved settings without a separate activation RPC. The +service makes the thread discoverable when it admits the first user message. +The service owns executor selection and paths; the caller supplies no local cwd, +workspace roots, custom configuration, or Aeon mode. This API does not request an +ephemeral thread. Prewarm is sent once; a timeout, disconnect or malformed response +can leave a prepared thread whose identity the caller did not receive. No retry +or cleanup is automatic. diff --git a/codex-rs/cloud-client/src/activation.rs b/codex-rs/cloud-client/src/activation.rs index 00bde2fddc..327e6c6c57 100644 --- a/codex-rs/cloud-client/src/activation.rs +++ b/codex-rs/cloud-client/src/activation.rs @@ -255,23 +255,22 @@ impl Client { .await?; let response = wire::UpdateConfigResponse::decode(bytes) .map_err(|_| MutationError::OutcomeUnknown(Error::InvalidResponse))?; - ThreadActivation::try_from_wire(response, thread_id).map_err(MutationError::OutcomeUnknown) + let settings = + ThreadActivation::from_settings(response).map_err(MutationError::OutcomeUnknown)?; + if settings.thread.id != thread_id { + return Err(MutationError::OutcomeUnknown(Error::InvalidResponse)); + } + Ok(settings) } } impl ThreadActivation { - fn try_from_wire( - response: wire::UpdateConfigResponse, - requested_id: &str, - ) -> Result { + pub(crate) fn from_settings(response: wire::UpdateConfigResponse) -> Result { let raw_thread = required(response.thread)?; let thread = ThreadSnapshot::try_from_wire( wire::Thread::decode(raw_thread.clone()).map_err(|_| Error::InvalidResponse)?, raw_thread, )?; - if thread.id != requested_id { - return Err(Error::InvalidResponse); - } Ok(Self { thread, collaboration_mode: response diff --git a/codex-rs/cloud-client/src/client_tests.rs b/codex-rs/cloud-client/src/client_tests.rs index 236c5b1b41..5216fb02c2 100644 --- a/codex-rs/cloud-client/src/client_tests.rs +++ b/codex-rs/cloud-client/src/client_tests.rs @@ -213,6 +213,99 @@ async fn activation_preserves_native_settings_and_thread_metadata() { assert!(server.requests.try_recv().is_err()); } +#[tokio::test] +async fn prewarm_uses_managed_defaults_and_preserves_resolved_settings() { + // PrewarmThread uses Thread field numbers but omits turns entirely. + let thread = Bytes::from_static(&[ + 0x22, 1, b'v', // cli_version (4) + 0x28, 1, // created_at (5) + 0x32, 1, b'/', // cwd (6) + 0x48, 0, // ephemeral (9) + 0x72, 1, b't', // id (14) + 0x82, 1, 1, b'p', // model_provider (16) + 0xaa, 1, 0, // preview (21) + 0xda, 1, 1, b's', // session_id (27) + 0xe2, 1, 2, 8, 1, // source (28): cli + 0xea, 1, 2, 0x12, 0, // status (29): idle + 0x80, 2, 2, // updated_at (32) + 0xa0, 6, 7, // unknown field (100) + ]); + let native_thread = [b"\xfa\x01\x00".as_slice(), thread.as_ref()].concat(); // empty turns (31) + // The settings use different field numbers in Prewarm and UpdateConfig. + let fields = [ + (1, 1, b"\x0a\x01e\x12\x01i".as_slice()), // active permission profile + (2, 2, b"\x08\x02".as_slice()), // on-request approval + (4, 5, b"/".as_slice()), // cwd + (7, 10, b"m".as_slice()), // model + (8, 11, b"p".as_slice()), // provider + (10, 13, b"high".as_slice()), // reasoning effort + (11, 14, b"\x0a\x01/".as_slice()), // workspace roots + (12, 15, b"\x12\x02\x08\x01".as_slice()), // read-only sandbox + (13, 16, b"\x0a\x04fast".as_slice()), // service tier + ]; + let mut prewarm = vec![0x18, 1]; // approvals reviewer (3): user + let mut activation = prewarm.clone(); + for (prewarm_tag, activation_tag, bytes) in fields { + prost::encoding::bytes::encode(prewarm_tag, &bytes.to_vec(), &mut prewarm); + prost::encoding::bytes::encode(activation_tag, &bytes.to_vec(), &mut activation); + } + prost::encoding::bytes::encode(/* tag */ 14, &thread, &mut prewarm); + prost::encoding::bytes::encode(/* tag */ 17, &native_thread, &mut activation); + let mut server = fixture_responses(vec![ + response(vec![frame(&activation)], tonic::Status::ok("")), + response(vec![frame(&prewarm)], tonic::Status::ok("")), + response(vec![frame(&prewarm)], tonic::Status::ok("")), + ]) + .await; + let expected = server.client.activate("t").await.unwrap(); + assert_request(&mut server, "UpdateConfig", &[56, 1, 0x92, 1, 1, b't']).await; + for (model, model_provider, request) in [ + (None, None, vec![0x52, 4, 0x0a, 2, 0x0a, 0]), + ( + Some("m".to_owned()), + Some("p".to_owned()), + vec![0x52, 4, 0x0a, 2, 0x0a, 0, 0x7a, 1, b'm', 0x82, 1, 1, b'p'], + ), + ] { + let actual = server + .client + .prewarm_default_thread(model, model_provider) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(&actual).unwrap(), + serde_json::to_value(&expected).unwrap() + ); + assert_eq!(actual.thread.protobuf().as_ref(), native_thread.as_slice()); + // Exactly one default managed selection; no custom settings or Aeon mode. + assert_request(&mut server, "Prewarm", &request).await; + assert!(server.requests.try_recv().is_err()); + } + + // Missing authoritative policy must fail rather than supplying local defaults. + let mut missing_sandbox = vec![0x18, 1]; + for (tag, _, bytes) in fields { + if tag != 12 { + prost::encoding::bytes::encode(tag, &bytes.to_vec(), &mut missing_sandbox); + } + } + prost::encoding::bytes::encode(/* tag */ 14, &thread, &mut missing_sandbox); + let mut server = fixture(response( + vec![frame(&missing_sandbox)], + tonic::Status::ok(""), + )) + .await; + assert!(matches!( + server + .client + .prewarm_default_thread(/*model*/ None, /*model_provider*/ None) + .await, + Err(Error::InvalidResponse) + )); + assert_request(&mut server, "Prewarm", &[0x52, 4, 0x0a, 2, 0x0a, 0]).await; + assert!(server.requests.try_recv().is_err()); +} + #[tokio::test] async fn resume_sends_protobuf_and_preserves_admission_flags() { for wait in [false, true] { diff --git a/codex-rs/cloud-client/src/lib.rs b/codex-rs/cloud-client/src/lib.rs index 03c0429a34..367b1bef18 100644 --- a/codex-rs/cloud-client/src/lib.rs +++ b/codex-rs/cloud-client/src/lib.rs @@ -10,6 +10,7 @@ mod history; mod history_wire; mod item_wire; mod items; +mod prewarm; mod transcript; mod transcript_wire; mod turns; diff --git a/codex-rs/cloud-client/src/prewarm.rs b/codex-rs/cloud-client/src/prewarm.rs new file mode 100644 index 0000000000..4cdc9f8bfc --- /dev/null +++ b/codex-rs/cloud-client/src/prewarm.rs @@ -0,0 +1,108 @@ +//! Prepare an ordinary cloud coding thread with server-owned managed environment defaults. + +use crate::Client; +use crate::Error; +use bytes::Bytes; +use bytes::BytesMut; +use prost::Message; + +impl Client { + /// Prewarm an ordinary coding thread in the default managed environment. + /// The service makes it discoverable when it admits the first user message. + /// This mutation is never retried. A timeout or disconnect can leave a thread + /// prepared even if its identity and settings were not received. The caller + /// must not use this ordinary-thread API for an ephemeral session. + pub async fn prewarm_default_thread( + &self, + model: Option, + model_provider: Option, + ) -> Result { + let response = self + .raw_unary( + wire::PrewarmRequest { + model, + model_provider, + environments: Some(wire::EnvironmentList { + values: vec![wire::EnvironmentSelection { + managed: Some(wire::ManagedEnvironment {}), + }], + }), + }, + "/oaiproto.codex_app_server_backend.ThreadService/Prewarm", + ) + .await?; + let warm = wire::PrewarmResponse::decode(response).map_err(|_| Error::InvalidResponse)?; + // PrewarmThread omits turns, but the native Thread requires an empty + // ThreadTurns message. Prefix it so the original wire bytes, including + // unknown fields, are preserved and malformed trailing data stays invalid. + let mut thread = BytesMut::new(); + prost::encoding::bytes::encode(/* tag */ 31, &Bytes::new(), &mut thread); + thread.extend_from_slice(&warm.thread.ok_or(Error::InvalidResponse)?); + + crate::ThreadActivation::from_settings(crate::activation_wire::UpdateConfigResponse { + active_permission_profile: warm.active_permission_profile, + approval_policy: warm.approval_policy, + approvals_reviewer: warm.approvals_reviewer, + cwd: warm.cwd, + model: warm.model, + model_provider: warm.model_provider, + reasoning_effort: warm.reasoning_effort, + runtime_workspace_roots: warm.runtime_workspace_roots, + sandbox: warm.sandbox, + service_tier: warm.service_tier, + thread: Some(thread.freeze()), + ..Default::default() + }) + } +} + +mod wire { + use super::*; + + #[derive(Clone, Message)] + pub(super) struct PrewarmRequest { + #[prost(message, optional, tag = "10")] + pub environments: Option, + #[prost(string, optional, tag = "15")] + pub model: Option, + #[prost(string, optional, tag = "16")] + pub model_provider: Option, + } + #[derive(Clone, Message)] + pub(super) struct EnvironmentList { + #[prost(message, repeated, tag = "1")] + pub values: Vec, + } + #[derive(Clone, Message)] + pub(super) struct EnvironmentSelection { + #[prost(message, optional, tag = "1")] + pub managed: Option, + } + #[derive(Clone, Message)] + pub(super) struct ManagedEnvironment {} + #[derive(Clone, Message)] + pub(super) struct PrewarmResponse { + #[prost(message, optional, tag = "1")] + pub active_permission_profile: Option, + #[prost(message, optional, tag = "2")] + pub approval_policy: Option, + #[prost(int32, optional, tag = "3")] + pub approvals_reviewer: Option, + #[prost(string, optional, tag = "4")] + pub cwd: Option, + #[prost(string, optional, tag = "7")] + pub model: Option, + #[prost(string, optional, tag = "8")] + pub model_provider: Option, + #[prost(string, optional, tag = "10")] + pub reasoning_effort: Option, + #[prost(message, optional, tag = "11")] + pub runtime_workspace_roots: Option, + #[prost(message, optional, tag = "12")] + pub sandbox: Option, + #[prost(message, optional, tag = "13")] + pub service_tier: Option, + #[prost(bytes = "bytes", optional, tag = "14")] + pub thread: Option, + } +}