diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 56a9157202..7ea7dce4d4 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2952,12 +2952,19 @@ dependencies = [ name = "codex-goal-extension" version = "0.0.0" dependencies = [ + "anyhow", "async-trait", + "chrono", "codex-extension-api", "codex-protocol", + "codex-state", "codex-tools", + "pretty_assertions", "serde", "serde_json", + "tempfile", + "tokio", + "tracing", ] [[package]] diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index ee68faf001..788e52e82b 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -984,6 +984,26 @@ ], "title": "InputImageFunctionCallOutputContentItem", "type": "object" + }, + { + "properties": { + "encrypted_content": { + "type": "string" + }, + "type": { + "enum": [ + "encrypted_content" + ], + "title": "EncryptedContentFunctionCallOutputContentItemType", + "type": "string" + } + }, + "required": [ + "encrypted_content", + "type" + ], + "title": "EncryptedContentFunctionCallOutputContentItem", + "type": "object" } ] }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 3db4a87ef2..f6e5a46c95 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -9242,6 +9242,26 @@ ], "title": "InputImageFunctionCallOutputContentItem", "type": "object" + }, + { + "properties": { + "encrypted_content": { + "type": "string" + }, + "type": { + "enum": [ + "encrypted_content" + ], + "title": "EncryptedContentFunctionCallOutputContentItemType", + "type": "string" + } + }, + "required": [ + "encrypted_content", + "type" + ], + "title": "EncryptedContentFunctionCallOutputContentItem", + "type": "object" } ] }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 7b62d6d4bf..c4e93338a1 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -5631,6 +5631,26 @@ ], "title": "InputImageFunctionCallOutputContentItem", "type": "object" + }, + { + "properties": { + "encrypted_content": { + "type": "string" + }, + "type": { + "enum": [ + "encrypted_content" + ], + "title": "EncryptedContentFunctionCallOutputContentItemType", + "type": "string" + } + }, + "required": [ + "encrypted_content", + "type" + ], + "title": "EncryptedContentFunctionCallOutputContentItem", + "type": "object" } ] }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/RawResponseItemCompletedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/RawResponseItemCompletedNotification.json index 74420ea57e..e69b342282 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/RawResponseItemCompletedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/RawResponseItemCompletedNotification.json @@ -140,6 +140,26 @@ ], "title": "InputImageFunctionCallOutputContentItem", "type": "object" + }, + { + "properties": { + "encrypted_content": { + "type": "string" + }, + "type": { + "enum": [ + "encrypted_content" + ], + "title": "EncryptedContentFunctionCallOutputContentItemType", + "type": "string" + } + }, + "required": [ + "encrypted_content", + "type" + ], + "title": "EncryptedContentFunctionCallOutputContentItem", + "type": "object" } ] }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeParams.json index 027989479a..d4bdeda0ed 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeParams.json @@ -199,6 +199,26 @@ ], "title": "InputImageFunctionCallOutputContentItem", "type": "object" + }, + { + "properties": { + "encrypted_content": { + "type": "string" + }, + "type": { + "enum": [ + "encrypted_content" + ], + "title": "EncryptedContentFunctionCallOutputContentItemType", + "type": "string" + } + }, + "required": [ + "encrypted_content", + "type" + ], + "title": "EncryptedContentFunctionCallOutputContentItem", + "type": "object" } ] }, diff --git a/codex-rs/app-server-protocol/schema/typescript/FunctionCallOutputContentItem.ts b/codex-rs/app-server-protocol/schema/typescript/FunctionCallOutputContentItem.ts index fb2996f1e5..cd18908145 100644 --- a/codex-rs/app-server-protocol/schema/typescript/FunctionCallOutputContentItem.ts +++ b/codex-rs/app-server-protocol/schema/typescript/FunctionCallOutputContentItem.ts @@ -7,4 +7,4 @@ import type { ImageDetail } from "./ImageDetail"; * Responses API compatible content items that can be returned by a tool call. * This is a subset of ContentItem with the types we support as function call outputs. */ -export type FunctionCallOutputContentItem = { "type": "input_text", text: string, } | { "type": "input_image", image_url: string, detail?: ImageDetail, }; +export type FunctionCallOutputContentItem = { "type": "input_text", text: string, } | { "type": "input_image", image_url: string, detail?: ImageDetail, } | { "type": "encrypted_content", encrypted_content: string, }; diff --git a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs index cf1e58345b..c18806ea03 100644 --- a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs @@ -161,7 +161,7 @@ impl ThreadGoalRequestProcessor { .thread_goals() .update_thread_goal( thread_id, - codex_state::ThreadGoalUpdate { + codex_state::GoalUpdate { objective: Some(objective.to_string()), status, token_budget: params.token_budget, @@ -206,7 +206,7 @@ impl ThreadGoalRequestProcessor { .thread_goals() .update_thread_goal( thread_id, - codex_state::ThreadGoalUpdate { + codex_state::GoalUpdate { objective: None, status, token_budget: params.token_budget, diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index 036f2368a0..d4baf7d773 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -1108,7 +1108,7 @@ async fn thread_goal_set_edits_objective_without_resetting_usage() -> Result<()> thread_id, /*time_delta_seconds*/ 12, /*token_delta*/ 50, - codex_state::ThreadGoalAccountingMode::ActiveOnly, + codex_state::GoalAccountingMode::ActiveOnly, Some(persisted_goal.goal_id.as_str()), ) .await?; diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 1b71ab15ae..f4127c437c 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -482,15 +482,15 @@ struct ExecServerCommand { #[arg(long = "listen", value_name = "URL", conflicts_with = "remote")] listen: Option, - /// Register this exec-server as a remote executor using the given base URL. - #[arg(long = "remote", value_name = "URL", requires = "executor_id")] + /// Register this exec-server as a remote environment using the given base URL. + #[arg(long = "remote", value_name = "URL", requires = "environment_id")] remote: Option, - /// Executor id to attach to when registering remotely. - #[arg(long = "executor-id", value_name = "ID")] - executor_id: Option, + /// Environment id to attach to when registering remotely. + #[arg(long = "environment-id", value_name = "ID")] + environment_id: Option, - /// Human-readable executor name. + /// Human-readable environment name. #[arg(long = "name", value_name = "NAME")] name: Option, @@ -1491,21 +1491,24 @@ async fn run_exec_server_command( arg0_paths.codex_linux_sandbox_exe.clone(), )?; if let Some(base_url) = cmd.remote { - let executor_id = cmd - .executor_id - .ok_or_else(|| anyhow::anyhow!("--executor-id is required when --remote is set"))?; + let environment_id = cmd + .environment_id + .ok_or_else(|| anyhow::anyhow!("--environment-id is required when --remote is set"))?; let auth_provider = load_exec_server_remote_auth_provider( root_config_overrides, config_profile, cmd.use_agent_identity_auth, ) .await?; - let mut remote_config = - codex_exec_server::RemoteExecutorConfig::new(base_url, executor_id, auth_provider)?; + let mut remote_config = codex_exec_server::RemoteEnvironmentConfig::new( + base_url, + environment_id, + auth_provider, + )?; if let Some(name) = cmd.name { remote_config.name = name; } - codex_exec_server::run_remote_executor(remote_config, runtime_paths).await?; + codex_exec_server::run_remote_environment(remote_config, runtime_paths).await?; return Ok(()); } let listen_url = cmd diff --git a/codex-rs/codex-api/src/api_bridge.rs b/codex-rs/codex-api/src/api_bridge.rs index 4f37c2ff9b..401dfab3a9 100644 --- a/codex-rs/codex-api/src/api_bridge.rs +++ b/codex-rs/codex-api/src/api_bridge.rs @@ -123,7 +123,7 @@ pub fn map_api_error(err: ApiError) -> CodexErr { status: http::StatusCode::INTERNAL_SERVER_ERROR, request_id: None, }), - TransportError::Timeout => CodexErr::Timeout, + TransportError::Timeout => CodexErr::RequestTimeout, TransportError::Network(msg) | TransportError::Build(msg) => { CodexErr::Stream(msg, None) } diff --git a/codex-rs/codex-api/src/endpoint/compact.rs b/codex-rs/codex-api/src/endpoint/compact.rs index 336e958d4e..a3da854844 100644 --- a/codex-rs/codex-api/src/endpoint/compact.rs +++ b/codex-rs/codex-api/src/endpoint/compact.rs @@ -11,6 +11,7 @@ use http::Method; use serde::Deserialize; use serde_json::to_value; use std::sync::Arc; +use std::time::Duration; pub struct CompactClient { session: EndpointSession, @@ -37,10 +38,19 @@ impl CompactClient { &self, body: serde_json::Value, extra_headers: HeaderMap, + request_timeout: Duration, ) -> Result, ApiError> { let resp = self .session - .execute(Method::POST, Self::path(), extra_headers, Some(body)) + .execute_with( + Method::POST, + Self::path(), + extra_headers, + Some(body), + |req| { + req.timeout = Some(request_timeout); + }, + ) .await?; let parsed: CompactHistoryResponse = serde_json::from_slice(&resp.body).map_err(|e| ApiError::Stream(e.to_string()))?; @@ -51,10 +61,11 @@ impl CompactClient { &self, input: &CompactionInput<'_>, extra_headers: HeaderMap, + request_timeout: Duration, ) -> Result, ApiError> { let body = to_value(input) .map_err(|e| ApiError::Stream(format!("failed to encode compaction input: {e}")))?; - self.compact(body, extra_headers).await + self.compact(body, extra_headers, request_timeout).await } } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index f604a63458..e3569e6d18 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -146,6 +146,9 @@ const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str = const RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=2026-02-06"; const RESPONSES_ENDPOINT: &str = "/responses"; const RESPONSES_COMPACT_ENDPOINT: &str = "/responses/compact"; +// `/responses/compact` is unary, so the timeout covers the full response rather than one idle +// period between stream events. +const COMPACT_REQUEST_TIMEOUT_IDLE_MULTIPLIER: u32 = 4; const MEMORIES_SUMMARIZE_ENDPOINT: &str = "/memories/trace_summarize"; #[cfg(test)] pub(crate) const WEBSOCKET_CONNECT_TIMEOUT: Duration = @@ -502,12 +505,16 @@ impl ModelClient { if let Some(header_value) = self.generate_attestation_header_for().await { extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } + let compact_request_timeout = client_setup + .api_provider + .stream_idle_timeout + .saturating_mul(COMPACT_REQUEST_TIMEOUT_IDLE_MULTIPLIER); let client = ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth) .with_telemetry(Some(request_telemetry)); let trace_attempt = compaction_trace.start_attempt(&payload); let result = client - .compact_input(&payload, extra_headers) + .compact_input(&payload, extra_headers, compact_request_timeout) .await .map_err(map_api_error); trace_attempt.record_result(result.as_deref()); diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index cfd9297382..6a05535488 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -504,6 +504,10 @@ fn estimate_reasoning_length(encoded_len: usize) -> usize { .saturating_sub(650) } +fn estimate_encrypted_function_output_length(encoded_len: usize) -> usize { + encoded_len.saturating_mul(9).div_ceil(16) +} + fn estimate_item_token_count(item: &ResponseItem) -> i64 { let model_visible_bytes = estimate_response_item_model_visible_bytes(item); approx_tokens_from_byte_count_i64(model_visible_bytes) @@ -547,16 +551,18 @@ pub(crate) fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> let raw = serde_json::to_string(item) .map(|serialized| i64::try_from(serialized.len()).unwrap_or(i64::MAX)) .unwrap_or_default(); - let (payload_bytes, replacement_bytes) = image_data_url_estimate_adjustment(item); - if payload_bytes == 0 || replacement_bytes == 0 { - raw - } else { - // Replace raw base64 payload bytes with a per-image estimate. - // We intentionally preserve the data URL prefix and JSON - // wrapper bytes already included in `raw`. - raw.saturating_sub(payload_bytes) - .saturating_add(replacement_bytes) - } + let (image_payload_bytes, image_replacement_bytes) = + image_data_url_estimate_adjustment(item); + let (encrypted_payload_bytes, encrypted_replacement_bytes) = + encrypted_function_output_estimate_adjustment(item); + // Replace raw base64 payload bytes with a per-image estimate. + // We intentionally preserve the data URL prefix and JSON + // wrapper bytes already included in `raw`. + let raw = raw + .saturating_sub(image_payload_bytes) + .saturating_add(image_replacement_bytes); + raw.saturating_sub(encrypted_payload_bytes) + .saturating_add(encrypted_replacement_bytes) } } } @@ -678,6 +684,31 @@ fn image_data_url_estimate_adjustment(item: &ResponseItem) -> (i64, i64) { (payload_bytes, replacement_bytes) } +fn encrypted_function_output_estimate_adjustment(item: &ResponseItem) -> (i64, i64) { + let ResponseItem::FunctionCallOutput { output, .. } = item else { + return (0, 0); + }; + let FunctionCallOutputBody::ContentItems(items) = &output.body else { + return (0, 0); + }; + + items.iter().fold((0i64, 0i64), |acc, item| { + let FunctionCallOutputContentItem::EncryptedContent { encrypted_content } = item else { + return acc; + }; + let payload_bytes = acc + .0 + .saturating_add(i64::try_from(encrypted_content.len()).unwrap_or(i64::MAX)); + let replacement_bytes = acc.1.saturating_add( + i64::try_from(estimate_encrypted_function_output_length( + encrypted_content.len(), + )) + .unwrap_or(i64::MAX), + ); + (payload_bytes, replacement_bytes) + }) +} + fn is_model_generated_item(item: &ResponseItem) -> bool { match item { ResponseItem::Message { role, .. } => role == "assistant", diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 8ba5d56e32..f63ae6ec1f 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -1754,6 +1754,26 @@ fn non_base64_image_urls_are_unchanged() { ); } +#[test] +fn encrypted_function_output_uses_plaintext_byte_estimate() { + let encrypted_content = "A".repeat(1_868); + let item = ResponseItem::FunctionCallOutput { + call_id: "call-encrypted".to_string(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: encrypted_content.clone(), + }, + ]), + }; + + let raw_len = serde_json::to_string(&item).unwrap().len() as i64; + let estimated = estimate_response_item_model_visible_bytes(&item); + let expected = raw_len - encrypted_content.len() as i64 + + estimate_encrypted_function_output_length(encrypted_content.len()) as i64; + + assert_eq!(estimated, expected); +} + #[test] fn data_url_without_base64_marker_is_unchanged() { let item = ResponseItem::Message { diff --git a/codex-rs/core/src/goals.rs b/codex-rs/core/src/goals.rs index 6843264e44..e58e82c36d 100644 --- a/codex-rs/core/src/goals.rs +++ b/codex-rs/core/src/goals.rs @@ -459,7 +459,7 @@ impl Session { self.account_thread_goal_wall_clock_usage( &state_db, - codex_state::ThreadGoalAccountingMode::ActiveOnly, + codex_state::GoalAccountingMode::ActiveOnly, TerminalMetricEmission::Emit, ) .await?; @@ -476,7 +476,7 @@ impl Session { .thread_goals() .update_thread_goal( self.conversation_id, - codex_state::ThreadGoalUpdate { + codex_state::GoalUpdate { objective: Some(objective.to_string()), status: status.map(state_goal_status_from_protocol), token_budget, @@ -516,7 +516,7 @@ impl Session { .thread_goals() .update_thread_goal( self.conversation_id, - codex_state::ThreadGoalUpdate { + codex_state::GoalUpdate { objective: None, status, token_budget, @@ -601,7 +601,7 @@ impl Session { let state_db = self.require_state_db_for_thread_goals().await?; self.account_thread_goal_wall_clock_usage( &state_db, - codex_state::ThreadGoalAccountingMode::ActiveOnly, + codex_state::GoalAccountingMode::ActiveOnly, TerminalMetricEmission::Emit, ) .await?; @@ -1028,7 +1028,7 @@ impl Session { self.conversation_id, time_delta_seconds, token_delta, - codex_state::ThreadGoalAccountingMode::ActiveOnly, + codex_state::GoalAccountingMode::ActiveOnly, expected_goal_id.as_deref(), ) .await?; @@ -1039,7 +1039,7 @@ impl Session { .is_some_and(|goal_id| reported_goal_id.as_deref() == Some(goal_id)) }; let goal = match outcome { - codex_state::ThreadGoalAccountingOutcome::Updated(goal) => { + codex_state::GoalAccountingOutcome::Updated(goal) => { let clear_active_goal = match goal.status { codex_state::ThreadGoalStatus::Active => false, codex_state::ThreadGoalStatus::BudgetLimited => { @@ -1072,7 +1072,7 @@ impl Session { } goal } - codex_state::ThreadGoalAccountingOutcome::Unchanged(_) => return Ok(()), + codex_state::GoalAccountingOutcome::Unchanged(_) => return Ok(()), }; let should_steer_budget_limit = matches!(budget_limit_steering, BudgetLimitSteering::Allowed) @@ -1119,7 +1119,7 @@ impl Session { }; self.account_thread_goal_wall_clock_usage( &state_db, - codex_state::ThreadGoalAccountingMode::ActiveOnly, + codex_state::GoalAccountingMode::ActiveOnly, TerminalMetricEmission::Suppress, ) .await?; @@ -1129,7 +1129,7 @@ impl Session { async fn account_thread_goal_wall_clock_usage( &self, state_db: &StateDbHandle, - mode: codex_state::ThreadGoalAccountingMode, + mode: codex_state::GoalAccountingMode, terminal_metric_emission: TerminalMetricEmission, ) -> anyhow::Result> { let _accounting_permit = self.goal_runtime.accounting_permit().await?; @@ -1158,7 +1158,7 @@ impl Session { ) .await? { - codex_state::ThreadGoalAccountingOutcome::Updated(goal) => { + codex_state::GoalAccountingOutcome::Updated(goal) => { if matches!(terminal_metric_emission, TerminalMetricEmission::Emit) { self.emit_goal_terminal_metrics_if_status_changed(previous_status, &goal); } @@ -1171,7 +1171,7 @@ impl Session { let goal = protocol_goal_from_state(goal); Ok(Some(goal)) } - codex_state::ThreadGoalAccountingOutcome::Unchanged(goal) => { + codex_state::GoalAccountingOutcome::Unchanged(goal) => { { let mut accounting = self.goal_runtime.accounting.lock().await; accounting.wall_clock.reset_baseline(); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index bfe3cfb3fe..5318b72151 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1911,6 +1911,108 @@ async fn record_token_usage_info_notifies_extension_contributors() { assert_eq!(expected, actual); } +#[tokio::test] +async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() { + struct SessionTurnStartMarker; + struct ThreadTurnStartMarker; + + #[derive(Debug, PartialEq, Eq)] + struct RecordedTurnStart { + turn_id: String, + collaboration_mode: CollaborationMode, + token_usage_at_turn_start: TokenUsage, + saw_session_store: bool, + saw_thread_store: bool, + } + + struct TurnStartRecorder { + records: Arc>>, + } + + #[async_trait::async_trait] + impl codex_extension_api::TurnLifecycleContributor for TurnStartRecorder { + async fn on_turn_start(&self, input: codex_extension_api::TurnStartInput<'_>) { + self.records + .lock() + .expect("turn start records lock") + .push(RecordedTurnStart { + turn_id: input.turn_id.to_string(), + collaboration_mode: input.collaboration_mode.clone(), + token_usage_at_turn_start: input.token_usage_at_turn_start.clone(), + saw_session_store: input + .session_store + .get::() + .is_some(), + saw_thread_store: input.thread_store.get::().is_some(), + }); + } + } + + let (mut session, turn_context) = make_session_and_context().await; + let records = Arc::new(std::sync::Mutex::new(Vec::new())); + let mut builder = codex_extension_api::ExtensionRegistryBuilder::::new(); + builder.turn_lifecycle_contributor(Arc::new(TurnStartRecorder { + records: Arc::clone(&records), + })); + session.services.extensions = Arc::new(builder.build()); + session + .services + .session_extension_data + .insert(SessionTurnStartMarker); + session + .services + .thread_extension_data + .insert(ThreadTurnStartMarker); + + let token_usage_at_turn_start = TokenUsage { + input_tokens: 120, + cached_input_tokens: 15, + output_tokens: 40, + reasoning_output_tokens: 9, + total_tokens: 169, + }; + session + .state + .lock() + .await + .set_token_info(Some(TokenUsageInfo { + total_token_usage: token_usage_at_turn_start.clone(), + last_token_usage: TokenUsage::default(), + model_context_window: turn_context.model_context_window(), + })); + + let turn_context = Arc::new(turn_context); + let session = Arc::new(session); + session + .spawn_task( + Arc::clone(&turn_context), + Vec::new(), + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: true, + }, + ) + .await; + + session.abort_all_tasks(TurnAbortReason::Interrupted).await; + + let actual = records + .lock() + .expect("turn start records lock") + .drain(..) + .collect::>(); + assert_eq!( + vec![RecordedTurnStart { + turn_id: turn_context.sub_id.clone(), + collaboration_mode: turn_context.collaboration_mode.clone(), + token_usage_at_turn_start, + saw_session_store: true, + saw_thread_store: true, + }], + actual + ); +} + #[tokio::test] async fn config_change_contributor_observes_effective_config_changes() { struct SessionConfigMarker; @@ -8620,7 +8722,7 @@ async fn external_goal_mutation_accounts_active_turn_before_status_change() -> a .thread_goals() .update_thread_goal( sess.conversation_id, - codex_state::ThreadGoalUpdate { + codex_state::GoalUpdate { objective: None, status: Some(codex_state::ThreadGoalStatus::Complete), token_budget: None, diff --git a/codex-rs/core/src/session/turn_tests.rs b/codex-rs/core/src/session/turn_tests.rs index f1cf86b46c..81a43ff5a9 100644 --- a/codex-rs/core/src/session/turn_tests.rs +++ b/codex-rs/core/src/session/turn_tests.rs @@ -1,6 +1,5 @@ use super::*; use codex_extension_api::ExtensionData; -use codex_extension_api::TurnItemContributionFuture; use codex_extension_api::TurnItemContributor; use codex_protocol::items::AgentMessageContent; use pretty_assertions::assert_eq; @@ -8,21 +7,20 @@ use std::sync::Arc; struct RewriteAgentMessageContributor; +#[async_trait::async_trait] impl TurnItemContributor for RewriteAgentMessageContributor { - fn contribute<'a>( - &'a self, - _thread_store: &'a ExtensionData, - _turn_store: &'a ExtensionData, - item: &'a mut TurnItem, - ) -> TurnItemContributionFuture<'a> { - Box::pin(async move { - if let TurnItem::AgentMessage(agent_message) = item { - agent_message.content = vec![AgentMessageContent::Text { - text: "plan contributed assistant text".to_string(), - }]; - } - Ok(()) - }) + async fn contribute( + &self, + _thread_store: &ExtensionData, + _turn_store: &ExtensionData, + item: &mut TurnItem, + ) -> Result<(), String> { + if let TurnItem::AgentMessage(agent_message) = item { + agent_message.content = vec![AgentMessageContent::Text { + text: "plan contributed assistant text".to_string(), + }]; + } + Ok(()) } } diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index 28ffbb7a6b..9c4be94576 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -13,7 +13,6 @@ use crate::tools::ToolRouter; use crate::tools::parallel::ToolCallRuntime; use crate::turn_diff_tracker::TurnDiffTracker; use codex_extension_api::ExtensionData; -use codex_extension_api::TurnItemContributionFuture; use codex_extension_api::TurnItemContributor; use codex_protocol::error::CodexErr; use codex_protocol::items::AgentMessageContent; @@ -168,43 +167,41 @@ struct TestTurnItemContributor; #[derive(Debug)] struct TurnItemContributorRan; +#[async_trait::async_trait] impl TurnItemContributor for TestTurnItemContributor { - fn contribute<'a>( - &'a self, - _thread_store: &'a ExtensionData, - turn_store: &'a ExtensionData, - item: &'a mut TurnItem, - ) -> TurnItemContributionFuture<'a> { - Box::pin(async move { - turn_store.insert(TurnItemContributorRan); - if let TurnItem::AgentMessage(agent_message) = item { - agent_message.memory_citation = Some(MemoryCitation { - entries: Vec::new(), - rollout_ids: Vec::new(), - }); - } - Ok(()) - }) + async fn contribute( + &self, + _thread_store: &ExtensionData, + turn_store: &ExtensionData, + item: &mut TurnItem, + ) -> Result<(), String> { + turn_store.insert(TurnItemContributorRan); + if let TurnItem::AgentMessage(agent_message) = item { + agent_message.memory_citation = Some(MemoryCitation { + entries: Vec::new(), + rollout_ids: Vec::new(), + }); + } + Ok(()) } } struct RewriteAgentMessageContributor; +#[async_trait::async_trait] impl TurnItemContributor for RewriteAgentMessageContributor { - fn contribute<'a>( - &'a self, - _thread_store: &'a ExtensionData, - _turn_store: &'a ExtensionData, - item: &'a mut TurnItem, - ) -> TurnItemContributionFuture<'a> { - Box::pin(async move { - if let TurnItem::AgentMessage(agent_message) = item { - agent_message.content = vec![AgentMessageContent::Text { - text: "contributed assistant text".to_string(), - }]; - } - Ok(()) - }) + async fn contribute( + &self, + _thread_store: &ExtensionData, + _turn_store: &ExtensionData, + item: &mut TurnItem, + ) -> Result<(), String> { + if let TurnItem::AgentMessage(agent_message) = item { + agent_message.content = vec![AgentMessageContent::Text { + text: "contributed assistant text".to_string(), + }]; + } + Ok(()) } } diff --git a/codex-rs/core/src/tasks/lifecycle.rs b/codex-rs/core/src/tasks/lifecycle.rs index 94b51647c9..8b934175cf 100644 --- a/codex-rs/core/src/tasks/lifecycle.rs +++ b/codex-rs/core/src/tasks/lifecycle.rs @@ -1,16 +1,25 @@ use codex_extension_api::ExtensionData; +use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; use crate::session::session::Session; +use crate::session::turn_context::TurnContext; impl Session { - pub(super) async fn emit_turn_start_lifecycle(&self, turn_store: &ExtensionData) { + pub(super) async fn emit_turn_start_lifecycle( + &self, + turn_context: &TurnContext, + token_usage_at_turn_start: &TokenUsage, + ) { for contributor in self.services.extensions.turn_lifecycle_contributors() { contributor .on_turn_start(codex_extension_api::TurnStartInput { + turn_id: turn_context.sub_id.as_str(), + collaboration_mode: &turn_context.collaboration_mode, + token_usage_at_turn_start, session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, - turn_store, + turn_store: turn_context.extension_data.as_ref(), }) .await; } diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 164e65f07c..80ddfee897 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -357,7 +357,7 @@ impl Session { debug_assert!(turn.tasks.is_empty()); Arc::clone(&turn.turn_state) }; - turn_state.lock().await.token_usage_at_turn_start = token_usage_at_turn_start; + turn_state.lock().await.token_usage_at_turn_start = token_usage_at_turn_start.clone(); let mut pending_items = queued_response_items .into_iter() .map(TurnInput::ResponseInputItem) @@ -366,7 +366,7 @@ impl Session { self.input_queue .extend_pending_input_for_turn_state(turn_state.as_ref(), pending_items) .await; - self.emit_turn_start_lifecycle(turn_context.extension_data.as_ref()) + self.emit_turn_start_lifecycle(turn_context.as_ref(), &token_usage_at_turn_start) .await; let turn_extension_data = Arc::clone(&turn_context.extension_data); diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index ab9d02f97c..a47ed9da34 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -367,19 +367,27 @@ fn build_code_mode_executors( return vec![]; } - let code_mode_nested_tool_specs = executors - .iter() - .filter_map(|executor| { - if executor.exposure() == ToolExposure::DirectModelOnly { - return None; - } + let mut code_mode_nested_tool_specs = Vec::new(); + let mut exec_prompt_tool_specs = Vec::new(); + for executor in executors { + let exposure = executor.exposure(); + if exposure == ToolExposure::DirectModelOnly { + continue; + } - executor.spec() - }) - .collect::>(); - let namespace_descriptions = code_mode_namespace_descriptions(&code_mode_nested_tool_specs); + let Some(spec) = executor.spec() else { + continue; + }; + + if exposure != ToolExposure::Deferred { + exec_prompt_tool_specs.push(spec.clone()); + } + code_mode_nested_tool_specs.push(spec); + } + + let namespace_descriptions = code_mode_namespace_descriptions(&exec_prompt_tool_specs); let mut enabled_tools = - collect_code_mode_exec_prompt_tool_definitions(code_mode_nested_tool_specs.iter()); + collect_code_mode_exec_prompt_tool_definitions(exec_prompt_tool_specs.iter()); enabled_tools .sort_by(|left, right| compare_code_mode_tools(left, right, &namespace_descriptions)); diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 7c811a88e0..210c287c2f 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -1808,7 +1808,7 @@ async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> { .await; // Phase 2: start the Streamable HTTP MCP test server in the active - // placement. In full CI this may be the remote executor container; locally + // placement. In full CI this may be the remote environment container; locally // it is a host process. let expected_env_value = "propagated-env-http"; let Some(http_server) = diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 01439108f8..468ac205f7 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -22,10 +22,10 @@ the wire. The CLI entrypoint supports: - `ws://IP:PORT` (default) -- `--remote URL --executor-id ID [--name NAME]` +- `--remote URL --environment-id ID [--name NAME]` -Remote mode registers the local exec-server with the executor registry, -then reconnects to the service-provided rendezvous websocket as the executor. +Remote mode registers the local exec-server with the environment registry, +then reconnects to the service-provided rendezvous websocket as the environment. It uses the standard Codex ChatGPT sign-in state; run `codex login` first when remote registration needs authentication. Containerized callers that receive an Agent Identity JWT in `CODEX_ACCESS_TOKEN` can opt into that auth path with @@ -39,7 +39,7 @@ Wire framing: ## Remote Relay Message Format -In remote mode, the harness and executor communicate through rendezvous using +In remote mode, the harness and environment communicate through rendezvous using `codex.exec_server.relay.v1.RelayMessageFrame`; the checked-in schema is in `src/proto/codex.exec_server.relay.v1.proto`. The relay frame carries stream identity plus endpoint-owned reliability metadata: @@ -58,8 +58,8 @@ next_seq // resume only: next sender seq reason // reset only: reset reason ``` -`stream_id` identifies one virtual harness/executor JSON-RPC session on the -executor websocket. The harness generates a UUIDv4 `stream_id`; the executor +`stream_id` identifies one virtual harness/environment JSON-RPC session on the +environment websocket. The harness generates a UUIDv4 `stream_id`; the environment demuxes frames by `stream_id` and runs an independent `ConnectionProcessor` per stream. @@ -375,12 +375,12 @@ The crate exports: - `DEFAULT_LISTEN_URL` and `ExecServerListenUrlParseError` - `ExecServerRuntimePaths` - `run_main()` for embedding the websocket server -- `RemoteExecutorConfig` and `run_remote_executor()` for embedding remote +- `RemoteEnvironmentConfig` and `run_remote_environment()` for embedding remote registration mode Callers must pass `ExecServerRuntimePaths` to `run_main()`. The top-level `codex exec-server` command builds these paths from the `codex` arg0 dispatch -state. `RemoteExecutorConfig::new(...)` also takes the auth provider that +state. `RemoteEnvironmentConfig::new(...)` also takes the auth provider that remote registration should use; the CLI builds that provider from Codex auth state before starting remote mode. diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 9261a59542..bb5a3855dc 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -162,7 +162,7 @@ struct Inner { // need serialization so concurrent register/remove operations do not // overwrite each other's copy-on-write updates. sessions_write_lock: Mutex<()>, - // Once the transport closes, every executor operation should fail quickly + // Once the transport closes, every environment operation should fail quickly // with the same canonical message. This client never reconnects, so the // latch only moves from unset to set once. disconnected: OnceLock, @@ -261,18 +261,18 @@ pub enum ExecServerError { Protocol(String), #[error("exec-server rejected request ({code}): {message}")] Server { code: i64, message: String }, - #[error("executor registry request failed ({status}{code_suffix}): {message}", code_suffix = .code.as_ref().map(|code| format!(", {code}")).unwrap_or_default())] - ExecutorRegistryHttp { + #[error("environment registry request failed ({status}{code_suffix}): {message}", code_suffix = .code.as_ref().map(|code| format!(", {code}")).unwrap_or_default())] + EnvironmentRegistryHttp { status: reqwest::StatusCode, code: Option, message: String, }, - #[error("executor registry configuration error: {0}")] - ExecutorRegistryConfig(String), - #[error("executor registry authentication error: {0}")] - ExecutorRegistryAuth(String), - #[error("executor registry request failed: {0}")] - ExecutorRegistryRequest(#[from] reqwest::Error), + #[error("environment registry configuration error: {0}")] + EnvironmentRegistryConfig(String), + #[error("environment registry authentication error: {0}")] + EnvironmentRegistryAuth(String), + #[error("environment registry request failed: {0}")] + EnvironmentRegistryRequest(#[from] reqwest::Error), } impl ExecServerClient { @@ -495,7 +495,7 @@ impl ExecServerClient { { // Reject new work before allocating a JSON-RPC request id. MCP tool // calls, process writes, and fs operations all pass through here, so - // this is the shared low-level failure path after executor disconnect. + // this is the shared low-level failure path after environment disconnect. if let Some(error) = self.inner.disconnected_error() { return Err(error); } @@ -720,7 +720,7 @@ impl Inner { session: Arc, ) -> Result<(), ExecServerError> { let _sessions_write_guard = self.sessions_write_lock.lock().await; - // Do not register a process session that can never receive executor + // Do not register a process session that can never receive environment // notifications. Without this check, remote MCP startup could create a // dead session and wait for process output that will never arrive. if let Some(error) = self.disconnected_error() { @@ -795,7 +795,7 @@ async fn fail_all_sessions(inner: &Arc, message: String) { for (_, session) in sessions { // Sessions synthesize a closed read response and emit a pushed Failed // event. That covers both polling consumers and streaming consumers - // such as executor-backed MCP stdio. + // such as environment-backed MCP stdio. session.set_failure(message.clone()).await; } } diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index bd556638ea..f2d16f8fac 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -91,8 +91,8 @@ pub use protocol::TerminateResponse; pub use protocol::WriteParams; pub use protocol::WriteResponse; pub use protocol::WriteStatus; -pub use remote::RemoteExecutorConfig; -pub use remote::run_remote_executor; +pub use remote::RemoteEnvironmentConfig; +pub use remote::run_remote_environment; pub use runtime_paths::ExecServerRuntimePaths; pub use server::DEFAULT_LISTEN_URL; pub use server::ExecServerListenUrlParseError; diff --git a/codex-rs/exec-server/src/process.rs b/codex-rs/exec-server/src/process.rs index 052e4aa643..cb6c832138 100644 --- a/codex-rs/exec-server/src/process.rs +++ b/codex-rs/exec-server/src/process.rs @@ -24,7 +24,7 @@ pub struct StartedExecProcess { /// stdout, stderr, or pty bytes. `Exited` reports the process exit status, while /// `Closed` means all output streams have ended and no more output events will /// arrive. `Failed` is used when the process session cannot continue, for -/// example because the remote executor connection disconnected. +/// example because the remote environment connection disconnected. #[derive(Debug, Clone, PartialEq, Eq)] pub enum ExecProcessEvent { Output(ProcessOutputChunk), diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index b3708f1d83..ae07d7652c 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -298,7 +298,7 @@ where } } -pub(crate) async fn run_multiplexed_executor( +pub(crate) async fn run_multiplexed_environment( stream: WebSocketStream, processor: ConnectionProcessor, ) where @@ -337,7 +337,7 @@ pub(crate) async fn run_multiplexed_executor( continue; } Some(Err(err)) => { - debug!("multiplexed executor websocket read failed: {err}"); + debug!("multiplexed environment websocket read failed: {err}"); break; } } diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 2493135c9a..de09d94e7c 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -11,28 +11,28 @@ use codex_utils_rustls_provider::ensure_rustls_crypto_provider; use crate::ExecServerError; use crate::ExecServerRuntimePaths; -use crate::relay::run_multiplexed_executor; +use crate::relay::run_multiplexed_environment; use crate::server::ConnectionProcessor; const ERROR_BODY_PREVIEW_BYTES: usize = 4096; #[derive(Clone)] -struct ExecutorRegistryClient { +struct EnvironmentRegistryClient { base_url: String, auth_provider: SharedAuthProvider, http: reqwest::Client, } -impl std::fmt::Debug for ExecutorRegistryClient { +impl std::fmt::Debug for EnvironmentRegistryClient { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ExecutorRegistryClient") + f.debug_struct("EnvironmentRegistryClient") .field("base_url", &self.base_url) .field("auth_provider", &"") .finish_non_exhaustive() } } -impl ExecutorRegistryClient { +impl EnvironmentRegistryClient { fn new(base_url: String, auth_provider: SharedAuthProvider) -> Result { let base_url = normalize_base_url(base_url)?; Ok(Self { @@ -42,15 +42,15 @@ impl ExecutorRegistryClient { }) } - async fn register_executor( + async fn register_environment( &self, - executor_id: &str, - ) -> Result { + environment_id: &str, + ) -> Result { let response = self .http .post(endpoint_url( &self.base_url, - &format!("/cloud/executor/{executor_id}/register"), + &format!("/cloud/environment/{environment_id}/register"), )) .headers(self.auth_provider.to_auth_headers()) .send() @@ -72,49 +72,49 @@ impl ExecutorRegistryClient { let status = response.status(); let body = response.text().await.unwrap_or_default(); if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) { - return Err(executor_registry_auth_error(status, &body)); + return Err(environment_registry_auth_error(status, &body)); } - Err(executor_registry_http_error(status, &body)) + Err(environment_registry_http_error(status, &body)) } } #[derive(Debug, Clone, Eq, PartialEq, Deserialize)] -struct ExecutorRegistryExecutorRegistrationResponse { - executor_id: String, +struct EnvironmentRegistryRegistrationResponse { + environment_id: String, url: String, } /// Configuration for registering an exec-server for remote use. #[derive(Clone)] -pub struct RemoteExecutorConfig { +pub struct RemoteEnvironmentConfig { pub base_url: String, - pub executor_id: String, + pub environment_id: String, pub name: String, auth_provider: SharedAuthProvider, } -impl std::fmt::Debug for RemoteExecutorConfig { +impl std::fmt::Debug for RemoteEnvironmentConfig { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("RemoteExecutorConfig") + f.debug_struct("RemoteEnvironmentConfig") .field("base_url", &self.base_url) - .field("executor_id", &self.executor_id) + .field("environment_id", &self.environment_id) .field("name", &self.name) .field("auth_provider", &"") .finish() } } -impl RemoteExecutorConfig { +impl RemoteEnvironmentConfig { pub fn new( base_url: String, - executor_id: String, + environment_id: String, auth_provider: SharedAuthProvider, ) -> Result { - let executor_id = normalize_executor_id(executor_id)?; + let environment_id = normalize_environment_id(environment_id)?; Ok(Self { base_url, - executor_id, + environment_id, name: "codex-exec-server".to_string(), auth_provider, }) @@ -123,27 +123,27 @@ impl RemoteExecutorConfig { /// Register an exec-server for remote use and serve requests over the returned /// rendezvous websocket. -pub async fn run_remote_executor( - config: RemoteExecutorConfig, +pub async fn run_remote_environment( + config: RemoteEnvironmentConfig, runtime_paths: ExecServerRuntimePaths, ) -> Result<(), ExecServerError> { ensure_rustls_crypto_provider(); let client = - ExecutorRegistryClient::new(config.base_url.clone(), config.auth_provider.clone())?; + EnvironmentRegistryClient::new(config.base_url.clone(), config.auth_provider.clone())?; let processor = ConnectionProcessor::new(runtime_paths); let mut backoff = Duration::from_secs(1); loop { - let response = client.register_executor(&config.executor_id).await?; + let response = client.register_environment(&config.environment_id).await?; eprintln!( - "codex exec-server remote executor registered with executor_id {}", - response.executor_id + "codex exec-server remote environment registered with environment_id {}", + response.environment_id ); match connect_async(response.url.as_str()).await { Ok((websocket, _)) => { backoff = Duration::from_secs(1); - run_multiplexed_executor(websocket, processor.clone()).await; + run_multiplexed_environment(websocket, processor.clone()).await; } Err(err) => { warn!("failed to connect remote exec-server websocket: {err}"); @@ -155,14 +155,14 @@ pub async fn run_remote_executor( } } -fn normalize_executor_id(executor_id: String) -> Result { - let executor_id = executor_id.trim().to_string(); - if executor_id.is_empty() { - return Err(ExecServerError::ExecutorRegistryConfig( - "executor id is required for remote exec-server registration".to_string(), +fn normalize_environment_id(environment_id: String) -> Result { + let environment_id = environment_id.trim().to_string(); + if environment_id.is_empty() { + return Err(ExecServerError::EnvironmentRegistryConfig( + "environment id is required for remote exec-server registration".to_string(), )); } - Ok(executor_id) + Ok(environment_id) } #[derive(Deserialize)] @@ -179,8 +179,8 @@ struct RegistryError { fn normalize_base_url(base_url: String) -> Result { let trimmed = base_url.trim().trim_end_matches('/').to_string(); if trimmed.is_empty() { - return Err(ExecServerError::ExecutorRegistryConfig( - "executor registry base URL is required".to_string(), + return Err(ExecServerError::EnvironmentRegistryConfig( + "environment registry base URL is required".to_string(), )); } Ok(trimmed) @@ -190,14 +190,14 @@ fn endpoint_url(base_url: &str, path: &str) -> String { format!("{base_url}/{}", path.trim_start_matches('/')) } -fn executor_registry_auth_error(status: StatusCode, body: &str) -> ExecServerError { +fn environment_registry_auth_error(status: StatusCode, body: &str) -> ExecServerError { let message = registry_error_message(body).unwrap_or_else(|| "empty error body".to_string()); - ExecServerError::ExecutorRegistryAuth(format!( - "executor registry authentication failed ({status}): {message}" + ExecServerError::EnvironmentRegistryAuth(format!( + "environment registry authentication failed ({status}): {message}" )) } -fn executor_registry_http_error(status: StatusCode, body: &str) -> ExecServerError { +fn environment_registry_http_error(status: StatusCode, body: &str) -> ExecServerError { let parsed = serde_json::from_str::(body).ok(); let (code, message) = parsed .and_then(|body| body.error) @@ -216,7 +216,7 @@ fn executor_registry_http_error(status: StatusCode, body: &str) -> ExecServerErr .unwrap_or_else(|| "empty or malformed error body".to_string()), ) }); - ExecServerError::ExecutorRegistryHttp { + ExecServerError::EnvironmentRegistryHttp { status, code, message, @@ -277,46 +277,46 @@ mod tests { } #[tokio::test] - async fn register_executor_posts_with_auth_provider_headers() { + async fn register_environment_posts_with_auth_provider_headers() { let server = MockServer::start().await; - let config = RemoteExecutorConfig::new( + let config = RemoteEnvironmentConfig::new( server.uri(), - "exec-requested".to_string(), + "environment-requested".to_string(), static_registry_auth_provider(), ) .expect("config"); Mock::given(method("POST")) - .and(path("/cloud/executor/exec-requested/register")) + .and(path("/cloud/environment/environment-requested/register")) .and(header("authorization", "Bearer registry-token")) .and(header("chatgpt-account-id", "workspace-123")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "executor_id": "exec-1", - "url": "wss://rendezvous.test/executor/exec-1?role=executor&sig=abc" + "environment_id": "env-1", + "url": "wss://rendezvous.test/cloud-agent/default/ws/environment/env-1?role=environment&sig=abc" }))) .mount(&server) .await; - let client = ExecutorRegistryClient::new(server.uri(), static_registry_auth_provider()) + let client = EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider()) .expect("client"); let response = client - .register_executor(&config.executor_id) + .register_environment(&config.environment_id) .await - .expect("register executor"); + .expect("register environment"); assert_eq!( response, - ExecutorRegistryExecutorRegistrationResponse { - executor_id: "exec-1".to_string(), - url: "wss://rendezvous.test/executor/exec-1?role=executor&sig=abc".to_string(), + EnvironmentRegistryRegistrationResponse { + environment_id: "env-1".to_string(), + url: "wss://rendezvous.test/cloud-agent/default/ws/environment/env-1?role=environment&sig=abc".to_string(), } ); } #[test] fn debug_output_redacts_auth_provider() { - let config = RemoteExecutorConfig::new( + let config = RemoteEnvironmentConfig::new( "https://registry.example".to_string(), - "exec-1".to_string(), + "env-1".to_string(), static_registry_auth_provider(), ) .expect("config"); diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index 9f985e6982..ff41dd525e 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -20,7 +20,7 @@ use codex_app_server_protocol::RequestId; use codex_exec_server::ExecServerRuntimePaths; use codex_exec_server::InitializeParams; use codex_exec_server::InitializeResponse; -use codex_exec_server::RemoteExecutorConfig; +use codex_exec_server::RemoteEnvironmentConfig; use futures::SinkExt; use futures::StreamExt; use http::HeaderMap; @@ -45,7 +45,7 @@ use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; -const EXECUTOR_ID: &str = "exec-mux-test"; +const ENVIRONMENT_ID: &str = "env-mux-test"; const REGISTRY_TOKEN: &str = "registry-token"; const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; const TEST_TIMEOUT: Duration = Duration::from_secs(5); @@ -67,15 +67,17 @@ fn static_registry_auth_provider() -> codex_api::SharedAuthProvider { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn multiplexed_remote_executor_routes_independent_virtual_streams() -> Result<()> { +async fn multiplexed_remote_environment_routes_independent_virtual_streams() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; let rendezvous_url = format!("ws://{}", listener.local_addr()?); let registry = MockServer::start().await; Mock::given(method("POST")) - .and(path(format!("/cloud/executor/{EXECUTOR_ID}/register"))) + .and(path(format!( + "/cloud/environment/{ENVIRONMENT_ID}/register" + ))) .and(header("authorization", format!("Bearer {REGISTRY_TOKEN}"))) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "executor_id": EXECUTOR_ID, + "environment_id": ENVIRONMENT_ID, "url": rendezvous_url, }))) .mount(®istry) @@ -83,22 +85,22 @@ async fn multiplexed_remote_executor_routes_independent_virtual_streams() -> Res let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?; let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?; - let config = RemoteExecutorConfig::new( + let config = RemoteEnvironmentConfig::new( registry.uri(), - EXECUTOR_ID.to_string(), + ENVIRONMENT_ID.to_string(), static_registry_auth_provider(), )?; - let remote_executor = tokio::spawn(codex_exec_server::run_remote_executor( + let remote_environment = tokio::spawn(codex_exec_server::run_remote_environment( config, runtime_paths, )); let (socket, _peer_addr) = timeout(TEST_TIMEOUT, listener.accept()) .await - .context("remote executor should connect to fake rendezvous")??; + .context("remote environment should connect to fake rendezvous")??; let mut websocket = timeout(TEST_TIMEOUT, accept_async(socket)) .await - .context("fake rendezvous should accept executor websocket")??; + .context("fake rendezvous should accept environment websocket")??; let stream_a = "stream-a"; let stream_b = "stream-b"; @@ -192,8 +194,8 @@ async fn multiplexed_remote_executor_routes_independent_virtual_streams() -> Res )?; websocket.close(None).await?; - remote_executor.abort(); - let _ = remote_executor.await; + remote_environment.abort(); + let _ = remote_environment.await; Ok(()) } @@ -285,7 +287,7 @@ where let frame = timeout(TEST_TIMEOUT, websocket.next()) .await .context("timed out waiting for relay frame")? - .ok_or_else(|| anyhow!("executor websocket closed"))??; + .ok_or_else(|| anyhow!("environment websocket closed"))??; match frame { Message::Binary(bytes) => { let frame = RelayMessageFrame::decode(bytes.as_ref())?; @@ -297,8 +299,8 @@ where return Ok((stream_id, message)); } Message::Ping(_) | Message::Pong(_) => {} - Message::Close(_) => bail!("executor websocket closed"), - Message::Text(_) => bail!("executor sent text frame on relay websocket"), + Message::Close(_) => bail!("environment websocket closed"), + Message::Text(_) => bail!("environment sent text frame on relay websocket"), Message::Frame(_) => {} } } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 8968e08573..2113fdef8e 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -134,34 +134,28 @@ pub trait ToolLifecycleContributor: Send + Sync { } } -/// Future returned by one claimed approval-review contribution. -pub type ApprovalReviewFuture<'a> = - std::pin::Pin + Send + 'a>>; - /// Extension contribution that can claim rendered approval-review prompts. +#[async_trait::async_trait] pub trait ApprovalReviewContributor: Send + Sync { - fn contribute<'a>( - &'a self, - session_store: &'a ExtensionData, - thread_store: &'a ExtensionData, - prompt: &'a str, - ) -> Option>; + async fn contribute( + &self, + session_store: &ExtensionData, + thread_store: &ExtensionData, + prompt: &str, + ) -> Option; } -/// Future returned by one ordered turn-item contribution. -pub type TurnItemContributionFuture<'a> = - std::pin::Pin> + Send + 'a>>; - /// Ordered post-processing contribution for one parsed turn item. /// /// Implementations may mutate the item before it is emitted and may use the /// explicitly exposed thread- and turn-lifetime stores when they need durable /// extension-private state. +#[async_trait::async_trait] pub trait TurnItemContributor: Send + Sync { - fn contribute<'a>( - &'a self, - thread_store: &'a ExtensionData, - turn_store: &'a ExtensionData, - item: &'a mut TurnItem, - ) -> TurnItemContributionFuture<'a>; + async fn contribute( + &self, + thread_store: &ExtensionData, + turn_store: &ExtensionData, + item: &mut TurnItem, + ) -> Result<(), String>; } diff --git a/codex-rs/ext/extension-api/src/contributors/turn_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/turn_lifecycle.rs index 0b53183e28..bbd3ae8f39 100644 --- a/codex-rs/ext/extension-api/src/contributors/turn_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/turn_lifecycle.rs @@ -1,9 +1,17 @@ +use codex_protocol::config_types::CollaborationMode; +use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; use crate::ExtensionData; /// Input supplied when the host starts a turn. pub struct TurnStartInput<'a> { + /// Stable host-owned turn identifier. + pub turn_id: &'a str, + /// Effective collaboration mode for this turn. + pub collaboration_mode: &'a CollaborationMode, + /// Total token usage snapshot captured when the turn started. + pub token_usage_at_turn_start: &'a TokenUsage, /// Store scoped to the host session runtime. pub session_store: &'a ExtensionData, /// Store scoped to this thread runtime. diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 373f3735a4..2ec132ef86 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -18,7 +18,6 @@ pub use codex_tools::ToolPayload; pub use codex_tools::ToolSpec; pub use codex_tools::parse_tool_input_schema; pub use contributors::ApprovalReviewContributor; -pub use contributors::ApprovalReviewFuture; pub use contributors::ConfigContributor; pub use contributors::ContextContributor; pub use contributors::PromptFragment; @@ -36,7 +35,6 @@ pub use contributors::ToolLifecycleContributor; pub use contributors::ToolLifecycleFuture; pub use contributors::ToolStartInput; pub use contributors::TurnAbortInput; -pub use contributors::TurnItemContributionFuture; pub use contributors::TurnItemContributor; pub use contributors::TurnLifecycleContributor; pub use contributors::TurnStartInput; diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index 4577ddc048..a359726257 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/codex-rs/ext/extension-api/src/registry.rs @@ -1,7 +1,8 @@ use std::sync::Arc; +use codex_protocol::protocol::ReviewDecision; + use crate::ApprovalReviewContributor; -use crate::ApprovalReviewFuture; use crate::ConfigContributor; use crate::ContextContributor; use crate::ExtensionData; @@ -171,15 +172,22 @@ impl ExtensionRegistry { /// Claims the first rendered approval-review prompt accepted by an /// installed contributor. - pub fn approval_review<'a>( - &'a self, - session_store: &'a ExtensionData, - thread_store: &'a ExtensionData, - prompt: &'a str, - ) -> Option> { - self.approval_review_contributors - .iter() - .find_map(|contributor| contributor.contribute(session_store, thread_store, prompt)) + pub async fn approval_review( + &self, + session_store: &ExtensionData, + thread_store: &ExtensionData, + prompt: &str, + ) -> Option { + for contributor in &self.approval_review_contributors { + if let Some(decision) = contributor + .contribute(session_store, thread_store, prompt) + .await + { + return Some(decision); + } + } + + None } /// Returns the registered prompt contributors. diff --git a/codex-rs/ext/goal/BUILD.bazel b/codex-rs/ext/goal/BUILD.bazel index 037313da37..05b80c3706 100644 --- a/codex-rs/ext/goal/BUILD.bazel +++ b/codex-rs/ext/goal/BUILD.bazel @@ -3,4 +3,5 @@ load("//:defs.bzl", "codex_rust_crate") codex_rust_crate( name = "goal", crate_name = "codex_goal_extension", + integration_compile_data_extra = ["src/accounting.rs"], ) diff --git a/codex-rs/ext/goal/Cargo.toml b/codex-rs/ext/goal/Cargo.toml index 7f8b9bd308..b6590f68e6 100644 --- a/codex-rs/ext/goal/Cargo.toml +++ b/codex-rs/ext/goal/Cargo.toml @@ -17,6 +17,15 @@ workspace = true async-trait = { workspace = true } codex-extension-api = { workspace = true } codex-protocol = { workspace = true } +codex-state = { workspace = true } codex-tools = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +tracing = { workspace = true } + +[dev-dependencies] +anyhow = { workspace = true } +chrono = { workspace = true } +pretty_assertions = { workspace = true } +tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/goal/src/accounting.rs b/codex-rs/ext/goal/src/accounting.rs index 712e325849..cc04684169 100644 --- a/codex-rs/ext/goal/src/accounting.rs +++ b/codex-rs/ext/goal/src/accounting.rs @@ -1,3 +1,4 @@ +use codex_protocol::config_types::ModeKind; use codex_protocol::protocol::TokenUsage; use std::collections::HashMap; use std::sync::Mutex; @@ -17,6 +18,8 @@ struct GoalAccountingInner { #[derive(Debug, Default)] struct GoalTurnAccounting { token_delta: i64, + last_accounted_token_usage: TokenUsage, + account_tokens: bool, stopped: bool, } @@ -27,24 +30,42 @@ pub(crate) struct RecordedTokenDelta { } impl GoalAccountingState { - pub(crate) fn start_turn(&self, turn_id: impl Into) { + pub(crate) fn start_turn( + &self, + turn_id: impl Into, + collaboration_mode: ModeKind, + token_usage_at_turn_start: &TokenUsage, + ) { let turn_id = turn_id.into(); - self.inner().turns.entry(turn_id).or_default().stopped = false; + self.inner().turns.insert( + turn_id, + GoalTurnAccounting { + token_delta: 0, + last_accounted_token_usage: token_usage_at_turn_start.clone(), + account_tokens: !matches!(collaboration_mode, ModeKind::Plan), + stopped: false, + }, + ); } pub(crate) fn record_token_usage( &self, turn_id: impl Into, - usage: &TokenUsage, + total_usage: &TokenUsage, ) -> Option { - let delta = goal_token_delta_for_usage(usage); - if delta <= 0 { + let turn_id = turn_id.into(); + let mut inner = self.inner(); + let turn = inner.turns.get_mut(&turn_id)?; + if turn.stopped || !turn.account_tokens { return None; } - let turn_id = turn_id.into(); - let mut inner = self.inner(); - let turn = inner.turns.entry(turn_id).or_default(); + let delta = + token_delta_since_last_accounting(&turn.last_accounted_token_usage, total_usage); + turn.last_accounted_token_usage = total_usage.clone(); + if delta <= 0 { + return None; + } turn.token_delta = turn.token_delta.saturating_add(delta); let turn_delta = turn.token_delta; inner.unflushed_token_delta = inner.unflushed_token_delta.saturating_add(delta); @@ -65,6 +86,21 @@ impl GoalAccountingState { } } +fn token_delta_since_last_accounting(last: &TokenUsage, current: &TokenUsage) -> i64 { + let delta = TokenUsage { + input_tokens: current.input_tokens.saturating_sub(last.input_tokens), + cached_input_tokens: current + .cached_input_tokens + .saturating_sub(last.cached_input_tokens), + output_tokens: current.output_tokens.saturating_sub(last.output_tokens), + reasoning_output_tokens: current + .reasoning_output_tokens + .saturating_sub(last.reasoning_output_tokens), + total_tokens: current.total_tokens.saturating_sub(last.total_tokens), + }; + goal_token_delta_for_usage(&delta) +} + pub(crate) fn goal_token_delta_for_usage(usage: &TokenUsage) -> i64 { usage .input_tokens diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index a8d4f5c289..6d8eb42f0e 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -5,7 +5,6 @@ use codex_extension_api::ConfigContributor; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistryBuilder; -use codex_extension_api::NoopExtensionEventSink; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadStartInput; use codex_extension_api::TokenUsageContributor; @@ -19,14 +18,12 @@ use codex_extension_api::TurnLifecycleContributor; use codex_extension_api::TurnStartInput; use codex_extension_api::TurnStopInput; use codex_protocol::ThreadId; -use codex_protocol::protocol::ThreadGoal; use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TurnAbortReason; use crate::accounting::GoalAccountingState; use crate::events::GoalEventEmitter; use crate::spec::UPDATE_GOAL_TOOL_NAME; -use crate::tool::CreateGoalRequest; use crate::tool::GoalToolExecutor; #[derive(Clone, Debug)] @@ -42,7 +39,7 @@ impl GoalExtensionConfig { #[derive(Clone)] pub struct GoalExtension { - backend: Arc, + state_dbs: Arc, event_emitter: GoalEventEmitter, goals_enabled: Arc bool + Send + Sync>, } @@ -54,70 +51,17 @@ impl std::fmt::Debug for GoalExtension { } impl GoalExtension { - pub fn new( - backend: Arc, - goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, - ) -> Self { - Self::new_with_event_sink(backend, Arc::new(NoopExtensionEventSink), goals_enabled) - } - - pub fn new_with_event_sink( - backend: Arc, + pub(crate) fn new_with_event_sink( + state_dbs: Arc, event_sink: Arc, goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, ) -> Self { Self { - backend, + state_dbs, event_emitter: GoalEventEmitter::new(event_sink), goals_enabled: Arc::new(goals_enabled), } } - - pub fn without_backend(goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static) -> Self { - Self::new(Arc::new(NoGoalToolBackend), goals_enabled) - } -} - -#[async_trait] -pub trait GoalToolBackend: Send + Sync { - async fn get_goal(&self, thread_id: ThreadId) -> Result, String>; - - async fn create_goal( - &self, - thread_id: ThreadId, - request: CreateGoalRequest, - ) -> Result; - - async fn complete_goal(&self, thread_id: ThreadId) -> Result; -} - -#[derive(Clone, Copy, Debug, Default)] -pub struct NoGoalToolBackend; - -#[async_trait] -impl GoalToolBackend for NoGoalToolBackend { - async fn get_goal(&self, _thread_id: ThreadId) -> Result, String> { - Err(missing_backend_message()) - } - - async fn create_goal( - &self, - _thread_id: ThreadId, - _request: CreateGoalRequest, - ) -> Result { - Err(missing_backend_message()) - } - - async fn complete_goal(&self, _thread_id: ThreadId) -> Result { - Err(missing_backend_message()) - } -} - -fn missing_backend_message() -> String { - // TODO: replace this fallback with a host-provided goal backend once - // ToolContributor invocations can reach thread-scoped goal persistence and - // the current turn context. - "goal tools are not connected to host goal persistence yet".to_string() } #[async_trait] @@ -164,10 +108,11 @@ where return; } - // TODO: TurnStartInput should expose collaboration mode and token usage - // at turn start. Goals need mode to suppress plan-mode accounting and - // the token baseline to account deltas exactly. - accounting_state(input.thread_store).start_turn(input.turn_store.level_id()); + accounting_state(input.thread_store).start_turn( + input.turn_id, + input.collaboration_mode.mode, + input.token_usage_at_turn_start, + ); } async fn on_turn_stop(&self, input: TurnStopInput<'_>) { @@ -214,7 +159,7 @@ where } let Some(_recorded) = accounting_state(thread_store) - .record_token_usage(turn_store.level_id(), &token_usage.last_token_usage) + .record_token_usage(turn_store.level_id(), &token_usage.total_token_usage) else { return; }; @@ -270,41 +215,32 @@ where vec![ Arc::new(GoalToolExecutor::get( thread_id, - Arc::clone(&self.backend), + Arc::clone(&self.state_dbs), self.event_emitter.clone(), )), Arc::new(GoalToolExecutor::create( thread_id, - Arc::clone(&self.backend), + Arc::clone(&self.state_dbs), self.event_emitter.clone(), )), Arc::new(GoalToolExecutor::update( thread_id, - Arc::clone(&self.backend), + Arc::clone(&self.state_dbs), self.event_emitter.clone(), )), ] } } -pub fn install( - registry: &mut ExtensionRegistryBuilder, - goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, -) where - C: Send + Sync + 'static, -{ - install_with_backend(registry, Arc::new(NoGoalToolBackend), goals_enabled); -} - pub fn install_with_backend( registry: &mut ExtensionRegistryBuilder, - backend: Arc, + state_dbs: Arc, goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, ) where C: Send + Sync + 'static, { let extension = Arc::new(GoalExtension::new_with_event_sink( - backend, + state_dbs, registry.event_sink(), goals_enabled, )); diff --git a/codex-rs/ext/goal/src/lib.rs b/codex-rs/ext/goal/src/lib.rs index 38332aee41..1625aeae14 100644 --- a/codex-rs/ext/goal/src/lib.rs +++ b/codex-rs/ext/goal/src/lib.rs @@ -12,9 +12,6 @@ mod tool; pub use extension::GoalExtension; pub use extension::GoalExtensionConfig; -pub use extension::GoalToolBackend; -pub use extension::NoGoalToolBackend; -pub use extension::install; pub use extension::install_with_backend; pub use spec::CREATE_GOAL_TOOL_NAME; pub use spec::GET_GOAL_TOOL_NAME; diff --git a/codex-rs/ext/goal/src/tool.rs b/codex-rs/ext/goal/src/tool.rs index 8ad4541a3c..96b160f938 100644 --- a/codex-rs/ext/goal/src/tool.rs +++ b/codex-rs/ext/goal/src/tool.rs @@ -16,7 +16,6 @@ use serde::Deserialize; use serde::Serialize; use crate::events::GoalEventEmitter; -use crate::extension::GoalToolBackend; use crate::spec::CREATE_GOAL_TOOL_NAME; use crate::spec::GET_GOAL_TOOL_NAME; use crate::spec::UPDATE_GOAL_TOOL_NAME; @@ -28,7 +27,7 @@ use crate::spec::create_update_goal_tool; pub(crate) struct GoalToolExecutor { kind: GoalToolKind, thread_id: ThreadId, - backend: Arc, + state_db: Arc, event_emitter: GoalEventEmitter, } @@ -69,39 +68,39 @@ enum CompletionBudgetReport { impl GoalToolExecutor { pub(crate) fn get( thread_id: ThreadId, - backend: Arc, + state_db: Arc, event_emitter: GoalEventEmitter, ) -> Self { Self { kind: GoalToolKind::Get, thread_id, - backend, + state_db, event_emitter, } } pub(crate) fn create( thread_id: ThreadId, - backend: Arc, + state_db: Arc, event_emitter: GoalEventEmitter, ) -> Self { Self { kind: GoalToolKind::Create, thread_id, - backend, + state_db, event_emitter, } } pub(crate) fn update( thread_id: ThreadId, - backend: Arc, + state_db: Arc, event_emitter: GoalEventEmitter, ) -> Self { Self { kind: GoalToolKind::Update, thread_id, - backend, + state_db, event_emitter, } } @@ -141,10 +140,14 @@ impl GoalToolExecutor { ) -> Result, FunctionCallError> { let _ = invocation.function_arguments()?; let goal = self - .backend - .get_goal(self.thread_id) + .state_db + .thread_goals() + .get_thread_goal(self.thread_id) .await - .map_err(FunctionCallError::RespondToModel)?; + .map(|goal| goal.map(protocol_goal_from_state)) + .map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to read goal: {err}")) + })?; goal_response(goal, CompletionBudgetReport::Omit) } @@ -159,10 +162,24 @@ impl GoalToolExecutor { validate_goal_budget(request.token_budget).map_err(FunctionCallError::RespondToModel)?; let goal = self - .backend - .create_goal(self.thread_id, request) + .state_db + .thread_goals() + .insert_thread_goal( + self.thread_id, + request.objective.as_str(), + codex_state::ThreadGoalStatus::Active, + request.token_budget, + ) .await - .map_err(FunctionCallError::RespondToModel)?; + .map_err(|err| FunctionCallError::RespondToModel(format!("failed to create goal: {err}")))? + .ok_or_else(|| { + FunctionCallError::RespondToModel( + "cannot create a new goal because this thread already has a goal; use update_goal only when the existing goal is complete" + .to_string(), + ) + })?; + fill_empty_thread_preview_if_possible(self.state_db.as_ref(), self.thread_id, &goal).await; + let goal = protocol_goal_from_state(goal); self.emit_goal_updated_from_tool_call(&invocation, goal.clone()); goal_response(Some(goal), CompletionBudgetReport::Omit) } @@ -182,10 +199,27 @@ impl GoalToolExecutor { // TODO: update_goal needs a host callback before completion to flush // final active-turn accounting with budget steering suppressed. let goal = self - .backend - .complete_goal(self.thread_id) + .state_db + .thread_goals() + .update_thread_goal( + self.thread_id, + codex_state::GoalUpdate { + objective: None, + status: Some(codex_state::ThreadGoalStatus::Complete), + token_budget: None, + expected_goal_id: None, + }, + ) .await - .map_err(FunctionCallError::RespondToModel)?; + .map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to complete goal: {err}")) + })? + .map(protocol_goal_from_state) + .ok_or_else(|| { + FunctionCallError::RespondToModel( + "cannot update goal because this thread has no goal".to_string(), + ) + })?; self.emit_goal_updated_from_tool_call(&invocation, goal.clone()); goal_response(Some(goal), CompletionBudgetReport::Include) } @@ -249,6 +283,45 @@ impl GoalToolResponse { } } +async fn fill_empty_thread_preview_if_possible( + state_db: &codex_state::StateRuntime, + thread_id: ThreadId, + goal: &codex_state::ThreadGoal, +) { + if let Err(err) = state_db + .set_thread_preview_if_empty(thread_id, goal.objective.as_str()) + .await + { + tracing::warn!( + "failed to set empty thread preview from goal objective for {thread_id}: {err}" + ); + } +} + +fn protocol_goal_from_state(goal: codex_state::ThreadGoal) -> ThreadGoal { + ThreadGoal { + thread_id: goal.thread_id, + objective: goal.objective, + status: protocol_status_from_state(goal.status), + token_budget: goal.token_budget, + tokens_used: goal.tokens_used, + time_used_seconds: goal.time_used_seconds, + created_at: goal.created_at.timestamp(), + updated_at: goal.updated_at.timestamp(), + } +} + +fn protocol_status_from_state(status: codex_state::ThreadGoalStatus) -> ThreadGoalStatus { + match status { + codex_state::ThreadGoalStatus::Active => ThreadGoalStatus::Active, + codex_state::ThreadGoalStatus::Paused => ThreadGoalStatus::Paused, + codex_state::ThreadGoalStatus::Blocked => ThreadGoalStatus::Blocked, + codex_state::ThreadGoalStatus::UsageLimited => ThreadGoalStatus::UsageLimited, + codex_state::ThreadGoalStatus::BudgetLimited => ThreadGoalStatus::BudgetLimited, + codex_state::ThreadGoalStatus::Complete => ThreadGoalStatus::Complete, + } +} + fn completion_budget_report(goal: &ThreadGoal) -> Option { if goal.token_budget.is_none() && goal.time_used_seconds <= 0 { None diff --git a/codex-rs/ext/goal/tests/accounting.rs b/codex-rs/ext/goal/tests/accounting.rs new file mode 100644 index 0000000000..99e9c93005 --- /dev/null +++ b/codex-rs/ext/goal/tests/accounting.rs @@ -0,0 +1,68 @@ +#![allow(dead_code)] + +#[path = "../src/accounting.rs"] +mod accounting; + +use accounting::GoalAccountingState; +use codex_protocol::config_types::ModeKind; +use codex_protocol::protocol::TokenUsage; +use pretty_assertions::assert_eq; + +#[test] +fn goal_accounting_uses_turn_start_baseline_for_exact_deltas() { + let state = GoalAccountingState::default(); + state.start_turn( + "turn-1", + ModeKind::Default, + &token_usage( + /*input_tokens*/ 100, /*cached_input_tokens*/ 10, /*output_tokens*/ 30, + /*reasoning_output_tokens*/ 5, /*total_tokens*/ 135, + ), + ); + + let recorded = state + .record_token_usage( + "turn-1", + &token_usage( + /*input_tokens*/ 120, /*cached_input_tokens*/ 14, + /*output_tokens*/ 42, /*reasoning_output_tokens*/ 8, + /*total_tokens*/ 162, + ), + ) + .expect("token delta should be recorded"); + + assert_eq!(28, recorded.turn_delta); + assert_eq!(28, recorded.thread_unflushed_delta); +} + +#[test] +fn goal_accounting_ignores_plan_mode_turns() { + let state = GoalAccountingState::default(); + state.start_turn("turn-1", ModeKind::Plan, &TokenUsage::default()); + + let recorded = state.record_token_usage( + "turn-1", + &token_usage( + /*input_tokens*/ 20, /*cached_input_tokens*/ 5, /*output_tokens*/ 8, + /*reasoning_output_tokens*/ 2, /*total_tokens*/ 30, + ), + ); + + assert_eq!(None, recorded); +} + +fn token_usage( + input_tokens: i64, + cached_input_tokens: i64, + output_tokens: i64, + reasoning_output_tokens: i64, + total_tokens: i64, +) -> TokenUsage { + TokenUsage { + input_tokens, + cached_input_tokens, + output_tokens, + reasoning_output_tokens, + total_tokens, + } +} diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs new file mode 100644 index 0000000000..a79c69c633 --- /dev/null +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -0,0 +1,171 @@ +use std::sync::Arc; + +use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::FunctionCallError; +use codex_extension_api::ThreadStartInput; +use codex_extension_api::ToolCall; +use codex_extension_api::ToolExecutor; +use codex_extension_api::ToolPayload; +use codex_goal_extension::install_with_backend; +use codex_protocol::ThreadId; +use codex_protocol::ToolName; +use codex_protocol::protocol::SessionSource; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; + +#[tokio::test] +async fn installed_goal_tools_create_goal_and_fill_empty_preview() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let tools = installed_tools(runtime.clone(), thread_id).await; + + let create_tool = tool_by_name(&tools, "create_goal"); + let invocation = ToolCall { + call_id: "call-create-goal".to_string(), + tool_name: ToolName::plain("create_goal"), + payload: ToolPayload::Function { + arguments: json!({ + "objective": "ship goal extension backend", + "token_budget": 123, + }) + .to_string(), + }, + }; + let output = create_tool.handle(invocation.clone()).await?; + let result = output.code_mode_result(&invocation.payload); + assert_eq!( + result, + json!({ + "goal": { + "threadId": thread_id, + "objective": "ship goal extension backend", + "status": "active", + "tokenBudget": 123, + "tokensUsed": 0, + "timeUsedSeconds": 0, + "createdAt": result["goal"]["createdAt"], + "updatedAt": result["goal"]["updatedAt"], + }, + "remainingTokens": 123, + "completionBudgetReport": serde_json::Value::Null, + }) + ); + + let metadata = runtime + .get_thread(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("seeded thread metadata should exist"))?; + assert_eq!( + metadata.preview.as_deref(), + Some("ship goal extension backend") + ); + Ok(()) +} + +#[tokio::test] +async fn installed_goal_tools_reject_duplicate_goal_creation() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let tools = installed_tools(runtime, thread_id).await; + + let create_tool = tool_by_name(&tools, "create_goal"); + let first = tool_call( + "create_goal", + "call-create-goal-1", + json!({ "objective": "first goal" }), + ); + create_tool.handle(first).await?; + + let second = tool_call( + "create_goal", + "call-create-goal-2", + json!({ "objective": "second goal" }), + ); + let err = match create_tool.handle(second).await { + Ok(_) => panic!("duplicate create should fail"), + Err(err) => err, + }; + + assert_eq!( + err, + FunctionCallError::RespondToModel( + "cannot create a new goal because this thread already has a goal; use update_goal only when the existing goal is complete" + .to_string() + ) + ); + Ok(()) +} + +async fn installed_tools( + runtime: Arc, + thread_id: ThreadId, +) -> Vec>> { + let mut builder = ExtensionRegistryBuilder::<()>::new(); + install_with_backend(&mut builder, runtime, |_| true); + let registry = builder.build(); + let session_store = ExtensionData::new("session-1"); + let thread_store = ExtensionData::new(thread_id.to_string()); + for contributor in registry.thread_lifecycle_contributors() { + contributor + .on_thread_start(ThreadStartInput { + config: &(), + session_store: &session_store, + thread_store: &thread_store, + }) + .await; + } + + registry + .tool_contributors() + .iter() + .flat_map(|contributor| contributor.tools(&session_store, &thread_store)) + .collect() +} + +fn tool_by_name<'a>( + tools: &'a [Arc>], + name: &str, +) -> &'a Arc> { + tools + .iter() + .find(|tool| tool.tool_name().namespace.is_none() && tool.tool_name().name == name) + .unwrap_or_else(|| panic!("missing tool {name}")) +} + +fn tool_call(tool_name: &str, call_id: &str, arguments: serde_json::Value) -> ToolCall { + ToolCall { + call_id: call_id.to_string(), + tool_name: ToolName::plain(tool_name), + payload: ToolPayload::Function { + arguments: arguments.to_string(), + }, + } +} + +async fn test_runtime() -> anyhow::Result> { + let tempdir = TempDir::new()?; + codex_state::StateRuntime::init(tempdir.keep(), "test-provider".to_string()).await +} + +fn test_thread_id() -> anyhow::Result { + ThreadId::from_string("11111111-1111-4111-8111-111111111111").map_err(anyhow::Error::msg) +} + +async fn seed_thread_metadata( + runtime: &codex_state::StateRuntime, + thread_id: ThreadId, +) -> anyhow::Result<()> { + let builder = codex_state::ThreadMetadataBuilder::new( + thread_id, + runtime + .codex_home() + .join(format!("rollout-{thread_id}.jsonl")), + chrono::Utc::now(), + SessionSource::Cli, + ); + runtime.upsert_thread(&builder.build("test-provider")).await +} diff --git a/codex-rs/file-system/src/lib.rs b/codex-rs/file-system/src/lib.rs index 026523eb2f..606f987cbd 100644 --- a/codex-rs/file-system/src/lib.rs +++ b/codex-rs/file-system/src/lib.rs @@ -130,7 +130,7 @@ fn file_system_policy_has_cwd_dependent_entries( pub type FileSystemResult = io::Result; /// Abstract filesystem access used by components that may operate locally or via -/// a remote executor. +/// a remote environment. #[async_trait] pub trait ExecutorFileSystem: Send + Sync { async fn read_file( diff --git a/codex-rs/protocol/src/error.rs b/codex-rs/protocol/src/error.rs index 207fd94ca2..ef9c86cadf 100644 --- a/codex-rs/protocol/src/error.rs +++ b/codex-rs/protocol/src/error.rs @@ -89,6 +89,8 @@ pub enum CodexErr { /// Returned by run_command_stream when the spawned child process timed out (10s). #[error("timeout waiting for child process to exit")] Timeout, + #[error("request timed out")] + RequestTimeout, /// Returned by run_command_stream when the child could not be spawned (its stdout/stderr pipes /// could not be captured). Analogous to the previous `CodexError::Spawn` variant. #[error("spawn failed: child stdout/stderr not captured")] @@ -192,6 +194,7 @@ impl CodexErr { | CodexErr::CyberPolicy { .. } => false, CodexErr::Stream(..) | CodexErr::Timeout + | CodexErr::RequestTimeout | CodexErr::UnexpectedStatus(_) | CodexErr::ResponseStreamFailed(_) | CodexErr::ConnectionFailed(_) diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index d1d2abe755..2927cb1cd1 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -1314,6 +1314,9 @@ pub enum FunctionCallOutputContentItem { #[ts(optional)] detail: Option, }, + EncryptedContent { + encrypted_content: String, + }, } /// Converts structured function-call output content into plain text for @@ -1337,7 +1340,8 @@ pub fn function_call_output_content_items_to_text( Some(text.as_str()) } FunctionCallOutputContentItem::InputText { .. } - | FunctionCallOutputContentItem::InputImage { .. } => None, + | FunctionCallOutputContentItem::InputImage { .. } + | FunctionCallOutputContentItem::EncryptedContent { .. } => None, }) .collect::>(); @@ -2063,6 +2067,9 @@ mod tests { image_url: "data:image/png;base64,AAA".to_string(), detail: Some(DEFAULT_IMAGE_DETAIL), }, + FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: "enc_opaque".to_string(), + }, ]; let text = function_call_output_content_items_to_text(&content_items); @@ -2270,6 +2277,35 @@ mod tests { Ok(()) } + #[test] + fn serializes_encrypted_function_output_content_as_array() -> Result<()> { + let item = ResponseInputItem::FunctionCallOutput { + call_id: "call1".into(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: "enc_opaque".into(), + }, + ]), + }; + + let json = serde_json::to_value(&item)?; + assert_eq!( + json, + serde_json::json!({ + "type": "function_call_output", + "call_id": "call1", + "output": [ + { + "type": "encrypted_content", + "encrypted_content": "enc_opaque", + } + ], + }) + ); + + Ok(()) + } + #[test] fn preserves_existing_image_data_urls() -> Result<()> { let call_tool_result = CallToolResult { @@ -2394,6 +2430,30 @@ mod tests { Ok(()) } + #[test] + fn deserializes_encrypted_array_payload_into_items() -> Result<()> { + let json = r#"[ + {"type": "encrypted_content", "encrypted_content": "enc_opaque"} + ]"#; + + let payload: FunctionCallOutputPayload = serde_json::from_str(json)?; + let expected_items = vec![FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: "enc_opaque".into(), + }]; + + assert_eq!(payload.success, None); + assert_eq!( + payload.body, + FunctionCallOutputBody::ContentItems(expected_items.clone()) + ); + assert_eq!( + serde_json::to_string(&payload)?, + serde_json::to_string(&expected_items)? + ); + + Ok(()) + } + #[test] fn deserializes_compaction_alias() -> Result<()> { let json = r#"{"type":"compaction_summary","encrypted_content":"abc"}"#; diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index 8746ff3a0d..f35b3420f6 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -48,13 +48,13 @@ pub use model::ThreadGoalStatus; pub use model::ThreadMetadata; pub use model::ThreadMetadataBuilder; pub use model::ThreadsPage; +pub use runtime::GoalAccountingMode; +pub use runtime::GoalAccountingOutcome; pub use runtime::GoalStore; +pub use runtime::GoalUpdate; pub use runtime::RemoteControlEnrollmentRecord; pub use runtime::RuntimeDbPath; pub use runtime::ThreadFilterOptions; -pub use runtime::ThreadGoalAccountingMode; -pub use runtime::ThreadGoalAccountingOutcome; -pub use runtime::ThreadGoalUpdate; pub use runtime::goals_db_filename; pub use runtime::goals_db_path; pub use runtime::logs_db_filename; diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index ac03f227da..8c9c6d0267 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -66,10 +66,10 @@ mod remote_control; mod test_support; mod threads; +pub use goals::GoalAccountingMode; +pub use goals::GoalAccountingOutcome; pub use goals::GoalStore; -pub use goals::ThreadGoalAccountingMode; -pub use goals::ThreadGoalAccountingOutcome; -pub use goals::ThreadGoalUpdate; +pub use goals::GoalUpdate; pub use remote_control::RemoteControlEnrollmentRecord; pub use threads::ThreadFilterOptions; diff --git a/codex-rs/state/src/runtime/goals.rs b/codex-rs/state/src/runtime/goals.rs index bf5a7bd825..38a64d0845 100644 --- a/codex-rs/state/src/runtime/goals.rs +++ b/codex-rs/state/src/runtime/goals.rs @@ -13,20 +13,20 @@ impl GoalStore { } } -pub struct ThreadGoalUpdate { +pub struct GoalUpdate { pub objective: Option, pub status: Option, pub token_budget: Option>, pub expected_goal_id: Option, } -pub enum ThreadGoalAccountingOutcome { +pub enum GoalAccountingOutcome { Unchanged(Option), Updated(crate::ThreadGoal), } #[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum ThreadGoalAccountingMode { +pub enum GoalAccountingMode { ActiveStatusOnly, ActiveOnly, ActiveOrComplete, @@ -170,9 +170,9 @@ RETURNING pub async fn update_thread_goal( &self, thread_id: ThreadId, - update: ThreadGoalUpdate, + update: GoalUpdate, ) -> anyhow::Result> { - let ThreadGoalUpdate { + let GoalUpdate { objective, status, token_budget, @@ -387,33 +387,33 @@ WHERE thread_id = ? thread_id: ThreadId, time_delta_seconds: i64, token_delta: i64, - mode: ThreadGoalAccountingMode, + mode: GoalAccountingMode, expected_goal_id: Option<&str>, - ) -> anyhow::Result { + ) -> anyhow::Result { let time_delta_seconds = time_delta_seconds.max(0); let token_delta = token_delta.max(0); if time_delta_seconds == 0 && token_delta == 0 { - return Ok(ThreadGoalAccountingOutcome::Unchanged( + return Ok(GoalAccountingOutcome::Unchanged( self.get_thread_goal(thread_id).await?, )); } let now_ms = datetime_to_epoch_millis(Utc::now()); let status_filter = match mode { - ThreadGoalAccountingMode::ActiveStatusOnly => "status = 'active'", - ThreadGoalAccountingMode::ActiveOnly => "status IN ('active', 'budget_limited')", - ThreadGoalAccountingMode::ActiveOrComplete => { + GoalAccountingMode::ActiveStatusOnly => "status = 'active'", + GoalAccountingMode::ActiveOnly => "status IN ('active', 'budget_limited')", + GoalAccountingMode::ActiveOrComplete => { "status IN ('active', 'budget_limited', 'complete')" } - ThreadGoalAccountingMode::ActiveOrStopped => { + GoalAccountingMode::ActiveOrStopped => { "status IN ('active', 'paused', 'blocked', 'usage_limited', 'budget_limited')" } }; let budget_limit_status_filter = match mode { - ThreadGoalAccountingMode::ActiveStatusOnly - | ThreadGoalAccountingMode::ActiveOnly - | ThreadGoalAccountingMode::ActiveOrComplete => "status = 'active'", - ThreadGoalAccountingMode::ActiveOrStopped => { + GoalAccountingMode::ActiveStatusOnly + | GoalAccountingMode::ActiveOnly + | GoalAccountingMode::ActiveOrComplete => "status = 'active'", + GoalAccountingMode::ActiveOrStopped => { "status IN ('active', 'paused', 'blocked', 'usage_limited', 'budget_limited')" } }; @@ -464,13 +464,13 @@ RETURNING let row = query.fetch_optional(self.pool.as_ref()).await?; let Some(row) = row else { - return Ok(ThreadGoalAccountingOutcome::Unchanged( + return Ok(GoalAccountingOutcome::Unchanged( self.get_thread_goal(thread_id).await?, )); }; let updated = thread_goal_from_row(&row)?; - Ok(ThreadGoalAccountingOutcome::Updated(updated)) + Ok(GoalAccountingOutcome::Updated(updated)) } } @@ -556,7 +556,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: None, status: Some(crate::ThreadGoalStatus::Paused), token_budget: Some(Some(200_000)), @@ -732,7 +732,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: None, status: Some(crate::ThreadGoalStatus::Complete), token_budget: None, @@ -756,7 +756,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: None, status: Some(crate::ThreadGoalStatus::Complete), token_budget: None, @@ -802,13 +802,13 @@ mod tests { thread_id, /*time_delta_seconds*/ 5, /*token_delta*/ 5, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, Some(original.goal_id.as_str()), ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Unchanged(Some(goal)) = outcome else { + let GoalAccountingOutcome::Unchanged(Some(goal)) = outcome else { panic!("stale goal version should not be updated"); }; assert_ne!(replacement.goal_id, original.goal_id); @@ -840,12 +840,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 12, /*token_delta*/ 30, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Updated(accounted) = outcome else { + let GoalAccountingOutcome::Updated(accounted) = outcome else { panic!("active goal should account usage"); }; @@ -853,7 +853,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: Some("draft the report clearly".to_string()), status: Some(crate::ThreadGoalStatus::Paused), token_budget: Some(Some(200)), @@ -891,7 +891,7 @@ mod tests { let status_update = runtime.thread_goals().update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: None, status: Some(crate::ThreadGoalStatus::Paused), token_budget: None, @@ -900,7 +900,7 @@ mod tests { ); let budget_update = runtime.thread_goals().update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: None, status: None, token_budget: Some(Some(200_000)), @@ -954,7 +954,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: None, status: Some(crate::ThreadGoalStatus::Complete), token_budget: None, @@ -1062,12 +1062,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 7, /*token_delta*/ 5, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Updated(goal) = outcome else { + let GoalAccountingOutcome::Updated(goal) = outcome else { panic!("active goal should be updated"); }; assert_eq!(crate::ThreadGoalStatus::Active, goal.status); @@ -1080,12 +1080,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 3, /*token_delta*/ 15, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Updated(goal) = outcome else { + let GoalAccountingOutcome::Updated(goal) = outcome else { panic!("budget crossing should update the goal"); }; assert_eq!(crate::ThreadGoalStatus::BudgetLimited, goal.status); @@ -1098,12 +1098,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 5, /*token_delta*/ 5, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Updated(goal) = outcome else { + let GoalAccountingOutcome::Updated(goal) = outcome else { panic!("budget-limited goal should still account in-flight active usage"); }; assert_eq!(crate::ThreadGoalStatus::BudgetLimited, goal.status); @@ -1133,12 +1133,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 5, /*token_delta*/ 5, - ThreadGoalAccountingMode::ActiveStatusOnly, + GoalAccountingMode::ActiveStatusOnly, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Unchanged(Some(goal)) = outcome else { + let GoalAccountingOutcome::Unchanged(Some(goal)) = outcome else { panic!("budget-limited goal should not be updated"); }; assert_eq!(crate::ThreadGoalStatus::BudgetLimited, goal.status); @@ -1165,7 +1165,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - crate::ThreadGoalUpdate { + crate::GoalUpdate { objective: None, status: Some(crate::ThreadGoalStatus::Paused), token_budget: None, @@ -1181,12 +1181,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 3, /*token_delta*/ 25, - ThreadGoalAccountingMode::ActiveOrStopped, + GoalAccountingMode::ActiveOrStopped, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Updated(goal) = outcome else { + let GoalAccountingOutcome::Updated(goal) = outcome else { panic!("stopped goal should account final usage"); }; assert_eq!(crate::ThreadGoalStatus::BudgetLimited, goal.status); @@ -1215,7 +1215,7 @@ mod tests { thread_id, /*time_delta_seconds*/ 1, /*token_delta*/ 50, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ) .await @@ -1225,7 +1225,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: None, status: None, token_budget: Some(Some(40)), @@ -1262,7 +1262,7 @@ mod tests { thread_id, /*time_delta_seconds*/ 1, /*token_delta*/ 50, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ) .await @@ -1272,7 +1272,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: Some("stay within budget, with clearer wording".to_string()), status: Some(crate::ThreadGoalStatus::Active), token_budget: None, @@ -1313,7 +1313,7 @@ mod tests { thread_id, /*time_delta_seconds*/ 1, /*token_delta*/ 50, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ) .await @@ -1323,7 +1323,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: None, status: Some(crate::ThreadGoalStatus::Paused), token_budget: None, @@ -1360,12 +1360,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 1, /*token_delta*/ 50, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Updated(budget_limited) = outcome else { + let GoalAccountingOutcome::Updated(budget_limited) = outcome else { panic!("budget crossing should update the goal"); }; @@ -1373,7 +1373,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: None, status: Some(crate::ThreadGoalStatus::Blocked), token_budget: None, @@ -1413,12 +1413,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 30, /*token_delta*/ 200, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Unchanged(Some(goal)) = active_only else { + let GoalAccountingOutcome::Unchanged(Some(goal)) = active_only else { panic!("completed goal should not be updated by active-only accounting"); }; assert_eq!(crate::ThreadGoalStatus::Complete, goal.status); @@ -1431,12 +1431,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 30, /*token_delta*/ 200, - ThreadGoalAccountingMode::ActiveOrComplete, + GoalAccountingMode::ActiveOrComplete, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Updated(goal) = completing_turn else { + let GoalAccountingOutcome::Updated(goal) = completing_turn else { panic!("completed goal should be updated for final accounting"); }; assert_eq!(crate::ThreadGoalStatus::Complete, goal.status); @@ -1463,7 +1463,7 @@ mod tests { .thread_goals() .update_thread_goal( thread_id, - ThreadGoalUpdate { + GoalUpdate { objective: None, status: Some(crate::ThreadGoalStatus::Paused), token_budget: None, @@ -1480,12 +1480,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 30, /*token_delta*/ 200, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Unchanged(Some(goal)) = active_only else { + let GoalAccountingOutcome::Unchanged(Some(goal)) = active_only else { panic!("paused goal should not be updated by active-only accounting"); }; assert_eq!(crate::ThreadGoalStatus::Paused, goal.status); @@ -1498,12 +1498,12 @@ mod tests { thread_id, /*time_delta_seconds*/ 30, /*token_delta*/ 200, - ThreadGoalAccountingMode::ActiveOrStopped, + GoalAccountingMode::ActiveOrStopped, /*expected_goal_id*/ None, ) .await .expect("usage accounting should succeed"); - let ThreadGoalAccountingOutcome::Updated(goal) = in_flight_turn else { + let GoalAccountingOutcome::Updated(goal) = in_flight_turn else { panic!("stopped goal should be updated for in-flight accounting"); }; assert_eq!(crate::ThreadGoalStatus::Paused, goal.status); @@ -1531,14 +1531,14 @@ mod tests { thread_id, /*time_delta_seconds*/ 4, /*token_delta*/ 40, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ); let second = runtime.thread_goals().account_thread_goal_usage( thread_id, /*time_delta_seconds*/ 6, /*token_delta*/ 60, - ThreadGoalAccountingMode::ActiveOnly, + GoalAccountingMode::ActiveOnly, /*expected_goal_id*/ None, ); let (first, second) = tokio::join!(first, second); diff --git a/codex-rs/tools/src/tool_output.rs b/codex-rs/tools/src/tool_output.rs index 64a8164da3..b61bb80de8 100644 --- a/codex-rs/tools/src/tool_output.rs +++ b/codex-rs/tools/src/tool_output.rs @@ -209,7 +209,8 @@ fn content_items_to_code_mode_result(items: &[FunctionCallOutputContentItem]) -> Some(image_url.clone()) } FunctionCallOutputContentItem::InputText { .. } - | FunctionCallOutputContentItem::InputImage { .. } => None, + | FunctionCallOutputContentItem::InputImage { .. } + | FunctionCallOutputContentItem::EncryptedContent { .. } => None, }) .collect::>() .join("\n"), diff --git a/codex-rs/utils/output-truncation/src/lib.rs b/codex-rs/utils/output-truncation/src/lib.rs index 24b1630da1..52cd741b7d 100644 --- a/codex-rs/utils/output-truncation/src/lib.rs +++ b/codex-rs/utils/output-truncation/src/lib.rs @@ -34,7 +34,8 @@ pub fn formatted_truncate_text_content_items_with_policy( .iter() .filter_map(|item| match item { FunctionCallOutputContentItem::InputText { text } => Some(text.as_str()), - FunctionCallOutputContentItem::InputImage { .. } => None, + FunctionCallOutputContentItem::InputImage { .. } + | FunctionCallOutputContentItem::EncryptedContent { .. } => None, }) .collect::>(); @@ -64,6 +65,11 @@ pub fn formatted_truncate_text_content_items_with_policy( detail: *detail, }) } + FunctionCallOutputContentItem::EncryptedContent { encrypted_content } => { + Some(FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: encrypted_content.clone(), + }) + } FunctionCallOutputContentItem::InputText { .. } => None, })); @@ -117,6 +123,11 @@ pub fn truncate_function_output_items_with_policy( detail: *detail, }); } + FunctionCallOutputContentItem::EncryptedContent { encrypted_content } => { + out.push(FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: encrypted_content.clone(), + }); + } } } diff --git a/codex-rs/utils/output-truncation/src/truncate_tests.rs b/codex-rs/utils/output-truncation/src/truncate_tests.rs index 74acb15ca3..baf26a058e 100644 --- a/codex-rs/utils/output-truncation/src/truncate_tests.rs +++ b/codex-rs/utils/output-truncation/src/truncate_tests.rs @@ -251,6 +251,60 @@ fn formatted_truncate_text_content_items_with_policy_merges_text_and_appends_ima assert_eq!(original_token_count, Some(4)); } +#[test] +fn formatted_truncate_text_content_items_with_policy_preserves_encrypted_content() { + let items = vec![ + FunctionCallOutputContentItem::InputText { + text: "abcdefgh".to_string(), + }, + FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: "enc_opaque".to_string(), + }, + ]; + + let (output, original_token_count) = + formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(2)); + + assert_eq!( + output, + vec![ + FunctionCallOutputContentItem::InputText { + text: "Total output lines: 1\n\na…6 chars truncated…h".to_string(), + }, + FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: "enc_opaque".to_string(), + }, + ] + ); + assert_eq!(original_token_count, Some(2)); +} + +#[test] +fn truncate_function_output_items_with_policy_preserves_encrypted_content() { + let items = vec![ + FunctionCallOutputContentItem::InputText { + text: "abcdefgh".to_string(), + }, + FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: "enc_opaque".to_string(), + }, + ]; + + let output = truncate_function_output_items_with_policy(&items, TruncationPolicy::Bytes(2)); + + assert_eq!( + output, + vec![ + FunctionCallOutputContentItem::InputText { + text: "a…6 chars truncated…h".to_string(), + }, + FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: "enc_opaque".to_string(), + }, + ] + ); +} + #[test] fn formatted_truncate_text_content_items_with_policy_merges_all_text_for_token_budget() { let items = vec![