From 95b16816f522d357140696eccdc2bcfcc18a71df Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Mon, 9 Feb 2026 17:52:03 -0800 Subject: [PATCH] wip; make compact request async submit + poll to avoid timeouts --- codex-rs/codex-api/README.md | 5 + codex-rs/codex-api/src/endpoint/compact.rs | 199 ++++++++++++++++++-- codex-rs/codex-api/src/lib.rs | 1 + codex-rs/core/src/client.rs | 87 ++++++++- codex-rs/core/tests/common/responses.rs | 74 ++++++++ codex-rs/core/tests/suite/compact_remote.rs | 109 +++++++++++ 6 files changed, 452 insertions(+), 23 deletions(-) diff --git a/codex-rs/codex-api/README.md b/codex-rs/codex-api/README.md index 0570cf7570..aa5354bcbc 100644 --- a/codex-rs/codex-api/README.md +++ b/codex-rs/codex-api/README.md @@ -28,6 +28,11 @@ The public interface of this crate is intentionally small and uniform: - `instructions: &str` – fully-resolved compaction instructions. - Output: `Vec`. - `CompactClient::compact_input(&CompactionInput, extra_headers)` wraps the JSON encoding and retry/telemetry wiring. + - Async operation mode: + - `CompactClient::submit_compact_operation(body, extra_headers)` returns `operation_id` from `responses/compact/operations`. + - `CompactClient::poll_compact_operation(operation_id, extra_headers)` returns either: + - `CompactOperationPollResult::Pending { poll_after_ms }`, or + - `CompactOperationPollResult::Completed(Vec)`. - **Memory trace summarize endpoint** - Input: `MemoryTraceSummarizeInput` (re-exported as `codex_api::MemoryTraceSummarizeInput`): diff --git a/codex-rs/codex-api/src/endpoint/compact.rs b/codex-rs/codex-api/src/endpoint/compact.rs index 44a56a11a7..d2cb3eee89 100644 --- a/codex-rs/codex-api/src/endpoint/compact.rs +++ b/codex-rs/codex-api/src/endpoint/compact.rs @@ -5,9 +5,11 @@ use crate::error::ApiError; use crate::provider::Provider; use codex_client::HttpTransport; use codex_client::RequestTelemetry; +use codex_client::Response as HttpResponse; use codex_protocol::models::ResponseItem; use http::HeaderMap; use http::Method; +use http::StatusCode; use serde::Deserialize; use serde_json::to_value; use std::sync::Arc; @@ -16,6 +18,12 @@ pub struct CompactClient { session: EndpointSession, } +#[derive(Debug)] +pub enum CompactOperationPollResult { + Pending { poll_after_ms: Option }, + Completed(Vec), +} + impl CompactClient { pub fn new(transport: T, provider: Provider, auth: A) -> Self { Self { @@ -33,6 +41,14 @@ impl CompactClient { "responses/compact" } + fn operations_path() -> &'static str { + "responses/compact/operations" + } + + fn operation_path(operation_id: &str) -> String { + format!("responses/compact/operations/{operation_id}") + } + pub async fn compact( &self, body: serde_json::Value, @@ -56,6 +72,98 @@ impl CompactClient { .map_err(|e| ApiError::Stream(format!("failed to encode compaction input: {e}")))?; self.compact(body, extra_headers).await } + + pub async fn submit_compact_operation( + &self, + body: serde_json::Value, + extra_headers: HeaderMap, + ) -> Result { + let resp = self + .session + .execute( + Method::POST, + Self::operations_path(), + extra_headers, + Some(body), + ) + .await?; + + if resp.status != StatusCode::ACCEPTED { + return Err(response_to_api_error( + &resp, + "compact operation submit failed", + )); + } + + let parsed: CompactOperationPendingResponse = + serde_json::from_slice(&resp.body).map_err(|e| ApiError::Stream(e.to_string()))?; + if parsed.operation_id.is_empty() { + return Err(ApiError::Stream( + "compact operation submit returned empty operation_id".to_string(), + )); + } + if parsed.status != "pending" { + let status = parsed.status; + return Err(ApiError::Stream(format!( + "unexpected compact operation status: {status}" + ))); + } + + Ok(parsed.operation_id) + } + + pub async fn poll_compact_operation( + &self, + operation_id: &str, + extra_headers: HeaderMap, + ) -> Result { + let resp = self + .session + .execute( + Method::GET, + &Self::operation_path(operation_id), + extra_headers, + None, + ) + .await?; + + if resp.status == StatusCode::ACCEPTED { + let parsed: CompactOperationPendingResponse = + serde_json::from_slice(&resp.body).map_err(|e| ApiError::Stream(e.to_string()))?; + if parsed.status == "pending" { + return Ok(CompactOperationPollResult::Pending { + poll_after_ms: parsed.poll_after_ms, + }); + } + let status = parsed.status; + return Err(ApiError::Stream(format!( + "unexpected compact operation status: {status}" + ))); + } + + if resp.status != StatusCode::OK { + return Err(response_to_api_error( + &resp, + "compact operation poll failed", + )); + } + + let parsed: CompactHistoryResponse = + serde_json::from_slice(&resp.body).map_err(|e| ApiError::Stream(e.to_string()))?; + Ok(CompactOperationPollResult::Completed(parsed.output)) + } +} + +fn response_to_api_error(resp: &HttpResponse, fallback_message: &str) -> ApiError { + let message = String::from_utf8_lossy(&resp.body).trim().to_string(); + ApiError::Api { + status: resp.status, + message: if message.is_empty() { + fallback_message.to_string() + } else { + message + }, + } } #[derive(Debug, Deserialize)] @@ -63,28 +171,25 @@ struct CompactHistoryResponse { output: Vec, } +#[derive(Debug, Deserialize)] +struct CompactOperationPendingResponse { + operation_id: String, + status: String, + poll_after_ms: Option, +} + #[cfg(test)] mod tests { use super::*; + use crate::provider::RetryConfig; use async_trait::async_trait; use codex_client::Request; use codex_client::Response; use codex_client::StreamResponse; use codex_client::TransportError; - - #[derive(Clone, Default)] - struct DummyTransport; - - #[async_trait] - impl HttpTransport for DummyTransport { - async fn execute(&self, _req: Request) -> Result { - Err(TransportError::Build("execute should not run".to_string())) - } - - async fn stream(&self, _req: Request) -> Result { - Err(TransportError::Build("stream should not run".to_string())) - } - } + use http::HeaderMap; + use pretty_assertions::assert_eq; + use std::time::Duration; #[derive(Clone, Default)] struct DummyAuth; @@ -95,11 +200,67 @@ mod tests { } } - #[test] - fn path_is_responses_compact() { - assert_eq!( - CompactClient::::path(), - "responses/compact" + fn provider(base_url: &str) -> Provider { + Provider { + name: "test".to_string(), + base_url: base_url.to_string(), + query_params: None, + headers: HeaderMap::new(), + retry: RetryConfig { + max_attempts: 1, + base_delay: Duration::from_millis(1), + retry_429: false, + retry_5xx: true, + retry_transport: true, + }, + stream_idle_timeout: Duration::from_secs(1), + } + } + + #[derive(Clone)] + struct StaticResponseTransport { + response: Response, + } + + #[async_trait] + impl HttpTransport for StaticResponseTransport { + async fn execute(&self, _req: Request) -> Result { + Ok(self.response.clone()) + } + + async fn stream(&self, _req: Request) -> Result { + Err(TransportError::Build("stream should not run".to_string())) + } + } + + #[tokio::test] + async fn submit_compact_operation_returns_pending_status() { + let response_body = compact_operation_pending_payload("op_123"); + let client = CompactClient::new( + StaticResponseTransport { + response: Response { + status: StatusCode::ACCEPTED, + headers: HeaderMap::new(), + body: bytes::Bytes::from(response_body.to_string()), + }, + }, + provider("https://example.com/api/codex"), + DummyAuth, ); + + let submit_result = client + .submit_compact_operation(serde_json::json!({}), HeaderMap::new()) + .await + .expect("submit should succeed"); + + assert_eq!(submit_result, "op_123"); + } + + fn compact_operation_pending_payload(operation_id: &str) -> serde_json::Value { + serde_json::json!({ + "operation_id": operation_id, + "status": "pending", + "poll_after_ms": 1000, + }) } } diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index 70652d2d78..e0703b2d8e 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -28,6 +28,7 @@ pub use crate::common::ResponsesApiRequest; pub use crate::common::create_text_param_for_request; pub use crate::endpoint::aggregate::AggregateStreamExt; pub use crate::endpoint::compact::CompactClient; +pub use crate::endpoint::compact::CompactOperationPollResult; pub use crate::endpoint::memories::MemoriesClient; pub use crate::endpoint::models::ModelsClient; pub use crate::endpoint::responses::ResponsesClient; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index c3f56d5f86..05ce001b05 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -36,6 +36,7 @@ use crate::api_bridge::auth_provider_from_auth; use crate::api_bridge::map_api_error; use crate::auth::UnauthorizedRecovery; use codex_api::CompactClient as ApiCompactClient; +use codex_api::CompactOperationPollResult; use codex_api::CompactionInput as ApiCompactionInput; use codex_api::MemoriesClient as ApiMemoriesClient; use codex_api::MemoryTrace as ApiMemoryTrace; @@ -83,6 +84,7 @@ use tokio::sync::oneshot::error::TryRecvError; use tokio_tungstenite::tungstenite::Error; use tokio_tungstenite::tungstenite::Message; use tracing::warn; +use uuid::Uuid; use crate::AuthManager; use crate::auth::CodexAuth; @@ -105,6 +107,11 @@ pub const X_CODEX_TURN_METADATA_HEADER: &str = "x-codex-turn-metadata"; pub const X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER: &str = "x-responsesapi-include-timing-metrics"; const RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=2026-02-06"; +const X_CODEX_IDEMPOTENCY_KEY_HEADER: &str = "x-codex-idempotency-key"; +const COMPACT_OPERATION_INITIAL_POLL_INTERVAL: Duration = Duration::from_millis(250); +const COMPACT_OPERATION_MIN_POLL_INTERVAL: Duration = Duration::from_millis(100); +const COMPACT_OPERATION_MAX_POLL_INTERVAL: Duration = Duration::from_secs(2); +const COMPACT_OPERATION_POLL_DEADLINE: Duration = Duration::from_secs(15 * 60); /// Session-scoped state shared by all [`ModelClient`] clones. /// /// This is intentionally kept minimal so `ModelClient` does not need to hold a full `Config`. Most @@ -268,11 +275,83 @@ impl ModelClient { instructions: &instructions, }; + let body = serde_json::to_value(&payload).map_err(|err| { + map_api_error(ApiError::Stream(format!( + "failed to encode compaction input: {err}" + ))) + })?; + let extra_headers = self.build_subagent_headers(); - client - .compact_input(&payload, extra_headers) - .await - .map_err(map_api_error) + let mut submit_headers = extra_headers.clone(); + let compact_id = Uuid::new_v4(); + let idempotency_key = format!("compact-{compact_id}"); + let idempotency_header_value = HeaderValue::from_str(&idempotency_key).map_err(|err| { + map_api_error(ApiError::Stream(format!( + "failed to encode compact idempotency key header: {err}" + ))) + })?; + submit_headers.insert(X_CODEX_IDEMPOTENCY_KEY_HEADER, idempotency_header_value); + + let operation_id = match client.submit_compact_operation(body, submit_headers).await { + Ok(operation_id) => operation_id, + Err(err) if Self::is_compact_operations_endpoint_unsupported(&err) => { + return client + .compact_input(&payload, extra_headers) + .await + .map_err(map_api_error); + } + Err(err) => return Err(map_api_error(err)), + }; + + let mut poll_interval = COMPACT_OPERATION_INITIAL_POLL_INTERVAL; + let deadline = tokio::time::Instant::now() + COMPACT_OPERATION_POLL_DEADLINE; + + loop { + if tokio::time::Instant::now() >= deadline { + return Err(map_api_error(ApiError::Stream( + "timeout waiting for compact operation".to_string(), + ))); + } + + let poll_result = client + .poll_compact_operation(&operation_id, extra_headers.clone()) + .await + .map_err(map_api_error)?; + match poll_result { + CompactOperationPollResult::Completed(output) => return Ok(output), + CompactOperationPollResult::Pending { poll_after_ms } => { + let poll_after_duration = poll_after_ms + .map(Duration::from_millis) + .map(Self::clamp_compact_operation_poll_interval); + let sleep_duration = poll_after_duration.unwrap_or(poll_interval); + tokio::time::sleep(sleep_duration).await; + poll_interval = poll_after_duration.unwrap_or_else(|| { + Self::clamp_compact_operation_poll_interval(poll_interval.saturating_mul(2)) + }); + } + } + } + } + + fn is_compact_operations_endpoint_unsupported(err: &ApiError) -> bool { + match err { + ApiError::Api { status, .. } => { + *status == HttpStatusCode::NOT_FOUND + || *status == HttpStatusCode::METHOD_NOT_ALLOWED + } + ApiError::Transport(TransportError::Http { status, .. }) => { + *status == HttpStatusCode::NOT_FOUND + || *status == HttpStatusCode::METHOD_NOT_ALLOWED + } + _ => false, + } + } + + fn clamp_compact_operation_poll_interval(interval: Duration) -> Duration { + interval.clamp( + COMPACT_OPERATION_MIN_POLL_INTERVAL, + COMPACT_OPERATION_MAX_POLL_INTERVAL, + ) } /// Builds memory summaries for each provided normalized trace. diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs index 379e852407..7ad78d1b42 100644 --- a/codex-rs/core/tests/common/responses.rs +++ b/codex-rs/core/tests/common/responses.rs @@ -756,6 +756,31 @@ fn compact_mock() -> (MockBuilder, ResponseMock) { (mock, response_mock) } +fn compact_operation_submit_mock() -> (MockBuilder, ResponseMock) { + let response_mock = ResponseMock::new(); + let mock = Mock::given(method("POST")) + .and(path_regex(".*/responses/compact/operations$")) + .and(response_mock.clone()); + (mock, response_mock) +} + +fn compact_operation_poll_mock(operation_id: &str) -> (MockBuilder, ResponseMock) { + let response_mock = ResponseMock::new(); + let path = format!(".*/responses/compact/operations/{operation_id}$"); + let mock = Mock::given(method("GET")) + .and(path_regex(&path)) + .and(response_mock.clone()); + (mock, response_mock) +} + +fn compact_operation_pending_payload(operation_id: &str) -> serde_json::Value { + serde_json::json!({ + "operation_id": operation_id, + "status": "pending", + "poll_after_ms": 1000, + }) +} + fn models_mock() -> (MockBuilder, ModelsMock) { let models_mock = ModelsMock::new(); let mock = Mock::given(method("GET")) @@ -820,6 +845,55 @@ pub async fn mount_compact_json_once(server: &MockServer, body: serde_json::Valu response_mock } +pub async fn mount_compact_operation_submit_once( + server: &MockServer, + operation_id: &str, +) -> ResponseMock { + let (mock, response_mock) = compact_operation_submit_mock(); + mock.respond_with( + ResponseTemplate::new(202) + .insert_header("content-type", "application/json") + .set_body_json(compact_operation_pending_payload(operation_id)), + ) + .up_to_n_times(1) + .mount(server) + .await; + response_mock +} + +pub async fn mount_compact_operation_poll_pending_once( + server: &MockServer, + operation_id: &str, +) -> ResponseMock { + let (mock, response_mock) = compact_operation_poll_mock(operation_id); + mock.respond_with( + ResponseTemplate::new(202) + .insert_header("content-type", "application/json") + .set_body_json(compact_operation_pending_payload(operation_id)), + ) + .up_to_n_times(1) + .mount(server) + .await; + response_mock +} + +pub async fn mount_compact_operation_poll_done_once( + server: &MockServer, + operation_id: &str, + body: serde_json::Value, +) -> ResponseMock { + let (mock, response_mock) = compact_operation_poll_mock(operation_id); + mock.respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .set_body_json(body.clone()), + ) + .up_to_n_times(1) + .mount(server) + .await; + response_mock +} + pub async fn mount_models_once(server: &MockServer, body: ModelsResponse) -> ModelsMock { let (mock, models_mock) = models_mock(); mock.respond_with( diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 47a9a0a111..9f89f06589 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -156,6 +156,115 @@ async fn remote_compact_replaces_history_for_followups() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_compact_operations_happy_path_replaces_history_for_followups() -> Result<()> { + skip_if_no_network!(Ok(())); + + let harness = TestCodexHarness::with_builder( + test_codex().with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()), + ) + .await?; + let codex = harness.test().codex.clone(); + + let responses_mock = responses::mount_sse_sequence( + harness.server(), + vec![ + responses::sse(vec![ + responses::ev_assistant_message("m1", "FIRST_REMOTE_REPLY"), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("m2", "AFTER_COMPACT_REPLY"), + responses::ev_completed("resp-2"), + ]), + ], + ) + .await; + + let compacted_history = vec![ + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "REMOTE_COMPACTED_SUMMARY".to_string(), + }], + end_turn: None, + phase: None, + }, + ResponseItem::Compaction { + encrypted_content: "ENCRYPTED_COMPACTION_SUMMARY".to_string(), + }, + ]; + let operation_id = "op_compact_happy_1"; + let submit_mock = + responses::mount_compact_operation_submit_once(harness.server(), operation_id).await; + let poll_pending_mock = + responses::mount_compact_operation_poll_pending_once(harness.server(), operation_id).await; + let poll_done_mock = responses::mount_compact_operation_poll_done_once( + harness.server(), + operation_id, + serde_json::json!({ "output": compacted_history.clone() }), + ) + .await; + let legacy_compact_mock = responses::mount_compact_json_once( + harness.server(), + serde_json::json!({ "output": compacted_history.clone() }), + ) + .await; + + codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "hello remote compact".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + }) + .await?; + wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + codex.submit(Op::Compact).await?; + wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "after compact".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + }) + .await?; + wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + let submit_request = submit_mock.single_request(); + assert_eq!(submit_request.path(), "/v1/responses/compact/operations"); + assert_eq!(poll_pending_mock.requests().len(), 1); + assert_eq!(poll_done_mock.requests().len(), 1); + assert_eq!(legacy_compact_mock.requests().len(), 0); + + let follow_up_body = responses_mock + .requests() + .last() + .expect("follow-up request missing") + .body_json() + .to_string(); + assert!( + follow_up_body.contains("REMOTE_COMPACTED_SUMMARY"), + "expected follow-up request to use compacted history" + ); + assert!( + follow_up_body.contains("ENCRYPTED_COMPACTION_SUMMARY"), + "expected follow-up request to include compaction summary item" + ); + assert!( + !follow_up_body.contains("FIRST_REMOTE_REPLY"), + "expected follow-up request to drop pre-compaction assistant messages" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn remote_compact_runs_automatically() -> Result<()> { skip_if_no_network!(Ok(()));