diff --git a/.github/workflows/rust-release-windows.yml b/.github/workflows/rust-release-windows.yml index e7574d85d2..49e627c1be 100644 --- a/.github/workflows/rust-release-windows.yml +++ b/.github/workflows/rust-release-windows.yml @@ -89,6 +89,11 @@ jobs: with: targets: ${{ matrix.target }} + - name: Configure LLVM linker + uses: ./.github/actions/setup-msvc-env + with: + target: ${{ matrix.target }} + - name: Cargo build (Windows binaries) shell: bash run: | diff --git a/.github/workflows/rust-release.yml b/.github/workflows/rust-release.yml index b821700b49..22b873f7d7 100644 --- a/.github/workflows/rust-release.yml +++ b/.github/workflows/rust-release.yml @@ -1204,6 +1204,7 @@ jobs: ./scripts/stage_npm_packages.py \ --release-version "$RELEASE_VERSION" \ --workflow-url "$workflow_url" \ + --artifacts-dir "${GITHUB_WORKSPACE}/dist" \ --package codex \ --package codex-responses-api-proxy \ --package codex-sdk diff --git a/MODULE.bazel b/MODULE.bazel index c4ef9500a7..63c8df726c 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -94,6 +94,7 @@ single_version_override( module_name = "rules_rs", patch_strip = 1, patches = [ + "//patches:rules_rs_build_script_deps_annotation.patch", "//patches:rules_rs_windows_gnullvm_exec.patch", "//patches:rules_rs_windows_exec_linker.patch", ], @@ -295,6 +296,11 @@ bazel_dep(name = "openssl", version = "3.5.4.bcr.0") inject_repo(crate, "xz") crate.annotation( + # Build scripts compile in Bazel's exec configuration, so target-specific + # optional build deps are otherwise dropped for the musl release platforms. + build_script_deps = [ + "@crates//:openssl-src-300.5.5+3.5.5", + ], build_script_data = [ "@openssl//:gen_dir", ], diff --git a/bazel/platforms/BUILD.bazel b/bazel/platforms/BUILD.bazel new file mode 100644 index 0000000000..398d5d8083 --- /dev/null +++ b/bazel/platforms/BUILD.bazel @@ -0,0 +1 @@ +# Release platform rules live in this package. diff --git a/bazel/platforms/release_binaries.bzl b/bazel/platforms/release_binaries.bzl new file mode 100644 index 0000000000..2e2f445036 --- /dev/null +++ b/bazel/platforms/release_binaries.bzl @@ -0,0 +1,27 @@ +"""Rules for building release binaries across supported target platforms.""" + +load("@rules_platform//platform_data:defs.bzl", "platform_data") + +PLATFORMS = [ + "linux_arm64_musl", + "linux_amd64_musl", + "macos_amd64", + "macos_arm64", + "windows_amd64", + "windows_arm64", +] + +def multiplatform_binaries(name, platforms = PLATFORMS): + for platform in platforms: + platform_data( + name = name + "_" + platform, + platform = "@llvm//platforms:" + platform, + target = name, + tags = ["manual"], + ) + + native.filegroup( + name = "release_binaries", + srcs = [name + "_" + platform for platform in platforms], + tags = ["manual"], + ) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index cbc237a84c..90a4335dc3 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1892,6 +1892,7 @@ dependencies = [ "codex-model-provider", "codex-plugin", "codex-protocol", + "codex-state", "codex-utils-absolute-path", "os_info", "pretty_assertions", @@ -2920,7 +2921,6 @@ dependencies = [ name = "codex-extension-api" version = "0.0.0" dependencies = [ - "async-trait", "codex-config", "codex-context-fragments", "codex-protocol", @@ -3047,8 +3047,8 @@ name = "codex-goal-extension" version = "0.0.0" dependencies = [ "anyhow", - "async-trait", "chrono", + "codex-analytics", "codex-core", "codex-extension-api", "codex-otel", @@ -3068,7 +3068,6 @@ dependencies = [ name = "codex-guardian" version = "0.0.0" dependencies = [ - "async-trait", "codex-core", "codex-extension-api", "codex-protocol", @@ -3101,7 +3100,6 @@ dependencies = [ name = "codex-image-generation-extension" version = "0.0.0" dependencies = [ - "async-trait", "codex-api", "codex-core", "codex-extension-api", @@ -3222,6 +3220,7 @@ name = "codex-mcp" version = "0.0.0" dependencies = [ "anyhow", + "arc-swap", "async-channel", "codex-api", "codex-async-utils", @@ -3253,7 +3252,6 @@ dependencies = [ name = "codex-mcp-extension" version = "0.0.0" dependencies = [ - "async-trait", "codex-config", "codex-core", "codex-core-plugins", @@ -3302,7 +3300,6 @@ dependencies = [ name = "codex-memories-extension" version = "0.0.0" dependencies = [ - "async-trait", "codex-core", "codex-extension-api", "codex-features", @@ -3342,6 +3339,8 @@ dependencies = [ "codex-features", "codex-git-utils", "codex-login", + "codex-model-provider", + "codex-model-provider-info", "codex-models-manager", "codex-otel", "codex-protocol", @@ -3816,15 +3815,21 @@ name = "codex-skills-extension" version = "0.0.0" dependencies = [ "async-trait", - "codex-core", "codex-core-skills", "codex-exec-server", "codex-extension-api", + "codex-mcp", "codex-protocol", + "codex-tools", "codex-utils-absolute-path", "codex-utils-string", "pretty_assertions", + "schemars 0.8.22", + "serde", + "serde_json", "tokio", + "tracing", + "url", ] [[package]] @@ -3915,7 +3920,6 @@ dependencies = [ name = "codex-tools" version = "0.0.0" dependencies = [ - "async-trait", "codex-app-server-protocol", "codex-code-mode", "codex-features", @@ -4296,7 +4300,6 @@ dependencies = [ name = "codex-web-search-extension" version = "0.0.0" dependencies = [ - "async-trait", "codex-api", "codex-core", "codex-extension-api", diff --git a/codex-rs/analytics/Cargo.toml b/codex-rs/analytics/Cargo.toml index 918e7edc72..8546427469 100644 --- a/codex-rs/analytics/Cargo.toml +++ b/codex-rs/analytics/Cargo.toml @@ -19,6 +19,7 @@ codex-login = { workspace = true } codex-model-provider = { workspace = true } codex-plugin = { workspace = true } codex-protocol = { workspace = true } +codex-state = { workspace = true } os_info = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index 02779c6609..49cbf0a858 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -43,6 +43,7 @@ use crate::facts::AppInvocation; use crate::facts::AppMentionedInput; use crate::facts::AppUsedInput; use crate::facts::CodexCompactionEvent; +use crate::facts::CodexErrKind; use crate::facts::CompactionImplementation; use crate::facts::CompactionPhase; use crate::facts::CompactionReason; @@ -778,6 +779,33 @@ fn sample_initialize_fact(connection_id: u64) -> AnalyticsFact { } } +async fn ingest_complete_child_turn( + reducer: &mut AnalyticsReducer, + events: &mut Vec, + thread_id: &str, + turn_id: &str, +) { + for fact in [ + AnalyticsFact::Custom(CustomAnalyticsFact::TurnResolvedConfig(Box::new( + sample_turn_resolved_config(thread_id, turn_id), + ))), + AnalyticsFact::Custom(CustomAnalyticsFact::TurnProfile(Box::new( + TurnProfileFact { + turn_id: turn_id.to_string(), + profile: sample_turn_profile(), + }, + ))), + AnalyticsFact::Notification(Box::new(sample_turn_completed_notification( + thread_id, + turn_id, + AppServerTurnStatus::Completed, + /*codex_error_info*/ None, + ))), + ] { + reducer.ingest(fact, events).await; + } +} + fn sample_command_execution_item( status: CommandExecutionStatus, exit_code: Option, @@ -1251,11 +1279,13 @@ fn compaction_event_serializes_expected_shape() { phase: CompactionPhase::MidTurn, strategy: CompactionStrategy::Memento, status: CompactionStatus::Completed, - error: None, + codex_error_kind: None, + codex_error_http_status_code: None, active_context_tokens_before: 120_000, active_context_tokens_after: 18_000, retained_image_count: None, compaction_summary_tokens: None, + cached_input_tokens: None, started_at: 100, completed_at: 106, duration_ms: Some(6543), @@ -1301,11 +1331,13 @@ fn compaction_event_serializes_expected_shape() { "phase": "mid_turn", "strategy": "memento", "status": "completed", - "error": null, + "codex_error_kind": null, + "codex_error_http_status_code": null, "active_context_tokens_before": 120000, "active_context_tokens_after": 18000, "retained_image_count": null, "compaction_summary_tokens": null, + "cached_input_tokens": null, "started_at": 100, "completed_at": 106, "duration_ms": 6543 @@ -1823,11 +1855,13 @@ async fn compaction_event_ingests_custom_fact() { phase: CompactionPhase::StandaloneTurn, strategy: CompactionStrategy::Memento, status: CompactionStatus::Failed, - error: Some("context limit exceeded".to_string()), + codex_error_kind: Some(CodexErrKind::ContextWindowExceeded), + codex_error_http_status_code: None, active_context_tokens_before: 131_000, active_context_tokens_after: 131_000, retained_image_count: None, compaction_summary_tokens: None, + cached_input_tokens: None, started_at: 100, completed_at: 101, duration_ms: Some(1200), @@ -1843,6 +1877,14 @@ async fn compaction_event_ingests_custom_fact() { assert_eq!(payload[0]["event_params"]["session_id"], "session-thread-1"); assert_eq!(payload[0]["event_params"]["thread_id"], "thread-1"); assert_eq!(payload[0]["event_params"]["turn_id"], "turn-compact"); + assert_eq!( + payload[0]["event_params"]["codex_error_kind"], + json!("context_window_exceeded") + ); + assert_eq!( + payload[0]["event_params"]["codex_error_http_status_code"], + json!(null) + ); assert_eq!( payload[0]["event_params"]["app_server_client"]["product_client_id"], DEFAULT_ORIGINATOR @@ -1937,6 +1979,7 @@ async fn guardian_review_event_ingests_custom_fact_with_optional_target_item() { decision: GuardianReviewDecision::Denied, terminal_status: GuardianReviewTerminalStatus::TimedOut, failure_reason: Some(GuardianReviewFailureReason::Timeout), + attempt_count: 1, risk_level: None, user_authorization: None, outcome: None, @@ -2008,6 +2051,7 @@ async fn guardian_review_event_ingests_custom_fact_with_optional_target_item() { ); assert_eq!(payload[0]["event_params"]["terminal_status"], "timed_out"); assert_eq!(payload[0]["event_params"]["failure_reason"], "timeout"); + assert_eq!(payload[0]["event_params"]["attempt_count"], 1); assert_eq!(payload[0]["event_params"]["review_timeout_ms"], 90_000); } @@ -2665,7 +2709,7 @@ async fn subagent_thread_started_publishes_without_initialize() { } #[tokio::test] -async fn subagent_thread_started_inherits_parent_connection_for_new_thread() { +async fn subagent_events_use_inherited_connection_unless_turn_connection_is_explicit() { let mut reducer = AnalyticsReducer::default(); let mut events = Vec::new(); let parent_thread_id = @@ -2747,11 +2791,13 @@ async fn subagent_thread_started_inherits_parent_connection_for_new_thread() { phase: CompactionPhase::StandaloneTurn, strategy: CompactionStrategy::Memento, status: CompactionStatus::Completed, - error: None, + codex_error_kind: None, + codex_error_http_status_code: None, active_context_tokens_before: 131_000, active_context_tokens_after: 64_000, retained_image_count: None, compaction_summary_tokens: None, + cached_input_tokens: None, started_at: 100, completed_at: 101, duration_ms: Some(1200), @@ -2772,6 +2818,68 @@ async fn subagent_thread_started_inherits_parent_connection_for_new_thread() { payload[0]["event_params"]["parent_thread_id"], "44444444-4444-4444-4444-444444444444" ); + + events.clear(); + ingest_complete_child_turn(&mut reducer, &mut events, "thread-review", "turn-inherited").await; + let [TrackEventRequest::TurnEvent(event)] = events.as_slice() else { + panic!("expected one turn event"); + }; + let params = &event.event_params; + assert_eq!(params.session_id, "session-root"); + assert_eq!(params.thread_source, Some(ThreadSource::Subagent)); + assert_eq!(params.subagent_source.as_deref(), Some("thread_spawn")); + assert_eq!( + params.parent_thread_id.as_deref(), + Some("44444444-4444-4444-4444-444444444444") + ); + assert_eq!(params.app_server_client.product_client_id, "parent-client"); + assert_eq!(params.runtime.codex_rs_version, "0.1.0"); + + reducer + .ingest( + AnalyticsFact::Custom(CustomAnalyticsFact::TurnTokenUsage(Box::new( + sample_turn_token_usage_fact("thread-review", "turn-inherited"), + ))), + &mut events, + ) + .await; + assert_eq!(events.len(), 1); + + events.clear(); + reducer + .ingest(sample_initialize_fact(/*connection_id*/ 8), &mut events) + .await; + reducer + .ingest( + AnalyticsFact::ClientRequest { + connection_id: 8, + request_id: RequestId::Integer(3), + request: Box::new(sample_turn_start_request( + "thread-review", + /*request_id*/ 3, + )), + }, + &mut events, + ) + .await; + reducer + .ingest( + AnalyticsFact::ClientResponse { + connection_id: 8, + request_id: RequestId::Integer(3), + response: Box::new(sample_turn_start_response("turn-explicit")), + }, + &mut events, + ) + .await; + ingest_complete_child_turn(&mut reducer, &mut events, "thread-review", "turn-explicit").await; + let [TrackEventRequest::TurnEvent(event)] = events.as_slice() else { + panic!("expected one turn event"); + }; + assert_eq!( + event.event_params.app_server_client.product_client_id, + DEFAULT_ORIGINATOR + ); } #[tokio::test] diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index b99a2ec86f..193de9ef65 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -9,6 +9,7 @@ use crate::facts::AnalyticsJsonRpcError; use crate::facts::AppInvocation; use crate::facts::AppMentionedInput; use crate::facts::AppUsedInput; +use crate::facts::CodexGoalEvent; use crate::facts::CustomAnalyticsFact; use crate::facts::HookRunFact; use crate::facts::HookRunInput; @@ -246,6 +247,12 @@ impl AnalyticsEventsClient { ))); } + pub fn track_goal_event(&self, event: CodexGoalEvent) { + self.record_fact(AnalyticsFact::Custom(CustomAnalyticsFact::Goal(Box::new( + event, + )))); + } + pub fn track_turn_resolved_config(&self, fact: TurnResolvedConfigFact) { self.record_fact(AnalyticsFact::Custom( CustomAnalyticsFact::TurnResolvedConfig(Box::new(fact)), diff --git a/codex-rs/analytics/src/events.rs b/codex-rs/analytics/src/events.rs index 0b6c9a0122..997ee51c56 100644 --- a/codex-rs/analytics/src/events.rs +++ b/codex-rs/analytics/src/events.rs @@ -4,12 +4,14 @@ use crate::facts::AcceptedLineFingerprint; use crate::facts::AppInvocation; use crate::facts::CodexCompactionEvent; use crate::facts::CodexErrKind; +use crate::facts::CodexGoalEvent; use crate::facts::CompactionImplementation; use crate::facts::CompactionPhase; use crate::facts::CompactionReason; use crate::facts::CompactionStatus; use crate::facts::CompactionStrategy; use crate::facts::CompactionTrigger; +use crate::facts::GoalEventKind; use crate::facts::HookRunFact; use crate::facts::InvocationType; use crate::facts::PluginState; @@ -63,6 +65,7 @@ pub(crate) enum TrackEventRequest { AppUsed(CodexAppUsedEventRequest), HookRun(CodexHookRunEventRequest), Compaction(Box), + Goal(Box), TurnEvent(Box), TurnSteer(CodexTurnSteerEventRequest), CommandExecution(CodexCommandExecutionEventRequest), @@ -263,6 +266,7 @@ pub struct GuardianReviewEventParams { pub decision: GuardianReviewDecision, pub terminal_status: GuardianReviewTerminalStatus, pub failure_reason: Option, + pub attempt_count: i64, pub risk_level: Option, pub user_authorization: Option, pub outcome: Option, @@ -335,6 +339,7 @@ impl GuardianReviewTrackContext { decision: result.decision, terminal_status: result.terminal_status, failure_reason: result.failure_reason, + attempt_count: result.attempt_count, risk_level: result.risk_level, user_authorization: result.user_authorization, outcome: result.outcome, @@ -370,6 +375,7 @@ pub struct GuardianReviewAnalyticsResult { pub decision: GuardianReviewDecision, pub terminal_status: GuardianReviewTerminalStatus, pub failure_reason: Option, + pub attempt_count: i64, pub risk_level: Option, pub user_authorization: Option, pub outcome: Option, @@ -389,6 +395,7 @@ impl GuardianReviewAnalyticsResult { decision: GuardianReviewDecision::Denied, terminal_status: GuardianReviewTerminalStatus::FailedClosed, failure_reason: None, + attempt_count: 1, risk_level: None, user_authorization: None, outcome: None, @@ -755,11 +762,13 @@ pub(crate) struct CodexCompactionEventParams { pub(crate) phase: CompactionPhase, pub(crate) strategy: CompactionStrategy, pub(crate) status: CompactionStatus, - pub(crate) error: Option, + pub(crate) codex_error_kind: Option, + pub(crate) codex_error_http_status_code: Option, pub(crate) active_context_tokens_before: i64, pub(crate) active_context_tokens_after: i64, pub(crate) retained_image_count: Option, pub(crate) compaction_summary_tokens: Option, + pub(crate) cached_input_tokens: Option, pub(crate) started_at: u64, pub(crate) completed_at: u64, pub(crate) duration_ms: Option, @@ -771,6 +780,30 @@ pub(crate) struct CodexCompactionEventRequest { pub(crate) event_params: CodexCompactionEventParams, } +#[derive(Serialize)] +pub(crate) struct CodexGoalEventParams { + pub(crate) thread_id: String, + pub(crate) session_id: String, + pub(crate) turn_id: Option, + pub(crate) app_server_client: CodexAppServerClientMetadata, + pub(crate) runtime: CodexRuntimeMetadata, + pub(crate) thread_source: Option, + pub(crate) subagent_source: Option, + pub(crate) parent_thread_id: Option, + pub(crate) goal_id: String, + pub(crate) event_kind: GoalEventKind, + pub(crate) goal_status: codex_state::ThreadGoalStatus, + pub(crate) has_token_budget: bool, + pub(crate) cumulative_tokens_accounted: Option, + pub(crate) cumulative_time_accounted_seconds: Option, +} + +#[derive(Serialize)] +pub(crate) struct CodexGoalEventRequest { + pub(crate) event_type: &'static str, + pub(crate) event_params: CodexGoalEventParams, +} + #[derive(Serialize)] pub(crate) struct CodexTurnEventParams { pub(crate) thread_id: String, @@ -968,17 +1001,46 @@ pub(crate) fn codex_compaction_event_params( phase: input.phase, strategy: input.strategy, status: input.status, - error: input.error, + codex_error_kind: input.codex_error_kind, + codex_error_http_status_code: input.codex_error_http_status_code, active_context_tokens_before: input.active_context_tokens_before, active_context_tokens_after: input.active_context_tokens_after, retained_image_count: input.retained_image_count, compaction_summary_tokens: input.compaction_summary_tokens, + cached_input_tokens: input.cached_input_tokens, started_at: input.started_at, completed_at: input.completed_at, duration_ms: input.duration_ms, } } +pub(crate) fn codex_goal_event_params( + input: CodexGoalEvent, + session_id: String, + app_server_client: CodexAppServerClientMetadata, + runtime: CodexRuntimeMetadata, + thread_source: Option, + subagent_source: Option, + parent_thread_id: Option, +) -> CodexGoalEventParams { + CodexGoalEventParams { + thread_id: input.thread_id, + session_id, + turn_id: input.turn_id, + app_server_client, + runtime, + thread_source, + subagent_source, + parent_thread_id, + goal_id: input.goal_id, + event_kind: input.event_kind, + goal_status: input.goal_status, + has_token_budget: input.has_token_budget, + cumulative_tokens_accounted: input.cumulative_tokens_accounted, + cumulative_time_accounted_seconds: input.cumulative_time_accounted_seconds, + } +} + pub(crate) fn codex_plugin_used_metadata( tracking: &TrackEventsContext, plugin: PluginTelemetryMetadata, diff --git a/codex-rs/analytics/src/facts.rs b/codex-rs/analytics/src/facts.rs index a035bc98b9..ab278e223c 100644 --- a/codex-rs/analytics/src/facts.rs +++ b/codex-rs/analytics/src/facts.rs @@ -137,7 +137,7 @@ impl TurnCodexErrorFact { #[derive(Clone, Copy, Debug, Serialize)] #[serde(rename_all = "snake_case")] -pub(crate) enum CodexErrKind { +pub enum CodexErrKind { TurnAborted, Stream, ContextWindowExceeded, @@ -364,6 +364,7 @@ pub enum CompactionReason { UserRequested, ContextLimit, ModelDownshift, + CompHashChanged, } #[derive(Clone, Copy, Debug, Serialize)] @@ -407,16 +408,39 @@ pub struct CodexCompactionEvent { pub phase: CompactionPhase, pub strategy: CompactionStrategy, pub status: CompactionStatus, - pub error: Option, + pub codex_error_kind: Option, + pub codex_error_http_status_code: Option, pub active_context_tokens_before: i64, pub active_context_tokens_after: i64, pub retained_image_count: Option, pub compaction_summary_tokens: Option, + pub cached_input_tokens: Option, pub started_at: u64, pub completed_at: u64, pub duration_ms: Option, } +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum GoalEventKind { + Created, + UsageAccounted, + StatusChanged, + Cleared, +} + +#[derive(Clone)] +pub struct CodexGoalEvent { + pub thread_id: String, + pub turn_id: Option, + pub goal_id: String, + pub event_kind: GoalEventKind, + pub goal_status: codex_state::ThreadGoalStatus, + pub has_token_budget: bool, + pub cumulative_tokens_accounted: Option, + pub cumulative_time_accounted_seconds: Option, +} + #[allow(dead_code)] pub(crate) enum AnalyticsFact { Initialize { @@ -468,6 +492,7 @@ pub(crate) enum AnalyticsFact { pub(crate) enum CustomAnalyticsFact { SubAgentThreadStarted(SubAgentThreadStartedInput), Compaction(Box), + Goal(Box), GuardianReview(Box), TurnResolvedConfig(Box), TurnTokenUsage(Box), diff --git a/codex-rs/analytics/src/lib.rs b/codex-rs/analytics/src/lib.rs index e2e16dfacd..fd236cbce9 100644 --- a/codex-rs/analytics/src/lib.rs +++ b/codex-rs/analytics/src/lib.rs @@ -24,6 +24,8 @@ pub use facts::AcceptedLineFingerprint; pub use facts::AnalyticsJsonRpcError; pub use facts::AppInvocation; pub use facts::CodexCompactionEvent; +pub use facts::CodexErrKind; +pub use facts::CodexGoalEvent; pub use facts::CodexTurnSteerEvent; pub use facts::CompactionImplementation; pub use facts::CompactionPhase; @@ -31,6 +33,7 @@ pub use facts::CompactionReason; pub use facts::CompactionStatus; pub use facts::CompactionStrategy; pub use facts::CompactionTrigger; +pub use facts::GoalEventKind; pub use facts::HookRunFact; pub use facts::InputError; pub use facts::InvocationType; diff --git a/codex-rs/analytics/src/reducer.rs b/codex-rs/analytics/src/reducer.rs index f8abe36e4e..fef1621c91 100644 --- a/codex-rs/analytics/src/reducer.rs +++ b/codex-rs/analytics/src/reducer.rs @@ -15,6 +15,7 @@ use crate::events::CodexDynamicToolCallEventParams; use crate::events::CodexDynamicToolCallEventRequest; use crate::events::CodexFileChangeEventParams; use crate::events::CodexFileChangeEventRequest; +use crate::events::CodexGoalEventRequest; use crate::events::CodexHookRunEventRequest; use crate::events::CodexImageGenerationEventParams; use crate::events::CodexImageGenerationEventRequest; @@ -51,6 +52,7 @@ use crate::events::TrackEventRequest; use crate::events::WebSearchActionKind; use crate::events::codex_app_metadata; use crate::events::codex_compaction_event_params; +use crate::events::codex_goal_event_params; use crate::events::codex_hook_run_metadata; use crate::events::codex_plugin_metadata; use crate::events::codex_plugin_used_metadata; @@ -62,6 +64,7 @@ use crate::facts::AnalyticsJsonRpcError; use crate::facts::AppMentionedInput; use crate::facts::AppUsedInput; use crate::facts::CodexCompactionEvent; +use crate::facts::CodexGoalEvent; use crate::facts::CustomAnalyticsFact; use crate::facts::HookRunInput; use crate::facts::PluginState; @@ -192,6 +195,16 @@ impl<'a> AnalyticsDropSite<'a> { } } + fn goal(input: &'a CodexGoalEvent) -> Self { + Self { + event_name: "goal", + thread_id: &input.thread_id, + turn_id: input.turn_id.as_deref(), + review_id: None, + item_id: None, + } + } + fn tool_item( notification: &'a codex_app_server_protocol::ItemCompletedNotification, item_id: &'a str, @@ -461,6 +474,9 @@ impl AnalyticsReducer { CustomAnalyticsFact::Compaction(input) => { self.ingest_compaction(*input, out); } + CustomAnalyticsFact::Goal(input) => { + self.ingest_goal(*input, out); + } CustomAnalyticsFact::GuardianReview(input) => { self.ingest_guardian_review(*input, out); } @@ -1271,6 +1287,26 @@ impl AnalyticsReducer { ))); } + fn ingest_goal(&mut self, input: CodexGoalEvent, out: &mut Vec) { + let Some((connection_state, thread_metadata)) = + self.thread_context_or_warn(AnalyticsDropSite::goal(&input)) + else { + return; + }; + out.push(TrackEventRequest::Goal(Box::new(CodexGoalEventRequest { + event_type: "codex_goal_event", + event_params: codex_goal_event_params( + input, + thread_metadata.session_id.clone(), + connection_state.app_server_client.clone(), + connection_state.runtime.clone(), + thread_metadata.thread_source.clone(), + thread_metadata.subagent_source.clone(), + thread_metadata.parent_thread_id.clone(), + ), + }))); + } + fn ingest_guardian_review_completed( &mut self, notification: codex_app_server_protocol::ItemGuardianApprovalReviewCompletedNotification, @@ -1461,17 +1497,23 @@ impl AnalyticsReducer { let Some(thread_id) = turn_state.thread_id.as_ref() else { return; }; - let Some(connection_id) = turn_state.connection_id else { + let drop_site = AnalyticsDropSite::turn(thread_id, turn_id); + let connection_id = turn_state.connection_id.or_else(|| { + self.threads + .get(drop_site.thread_id) + .and_then(|thread| thread.connection_id) + }); + let Some(connection_id) = connection_id else { + warn_missing_analytics_context(&drop_site, MissingAnalyticsContext::ThreadConnection); return; }; let Some(connection_state) = self.connections.get(&connection_id) else { warn_missing_analytics_context( - &AnalyticsDropSite::turn(thread_id, turn_id), + &drop_site, MissingAnalyticsContext::Connection { connection_id }, ); return; }; - let drop_site = AnalyticsDropSite::turn(thread_id, turn_id); let Some(thread_metadata) = self .threads .get(drop_site.thread_id) diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index 11684c6360..004b7d1709 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -22,6 +22,7 @@ use std::fmt; use std::io::Error as IoError; use std::io::ErrorKind; use std::io::Result as IoResult; +use std::path::Path; use std::sync::Arc; use std::time::Duration; @@ -48,7 +49,11 @@ use codex_config::LoaderOverrides; use codex_config::NoopThreadConfigLoader; use codex_config::RemoteThreadConfigLoader; use codex_config::ThreadConfigLoader; +use codex_config::config_toml::ConfigToml; use codex_core::config::Config; +pub use codex_core::otel_init::build_provider as build_otel_provider; +use codex_core::personality_migration::PersonalityMigrationStatus; +use codex_core::personality_migration::maybe_migrate_personality; pub use codex_exec_server::EnvironmentManager; pub use codex_exec_server::ExecServerRuntimePaths; use codex_feedback::CodexFeedback; @@ -82,18 +87,6 @@ pub mod legacy_core { } } - pub mod otel_init { - pub use codex_core::otel_init::*; - } - - pub mod personality_migration { - pub use codex_core::personality_migration::*; - } - - pub mod test_support { - pub use codex_core::test_support::*; - } - pub mod windows_sandbox { pub use codex_core::windows_sandbox::*; } @@ -101,6 +94,23 @@ pub mod legacy_core { const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); +/// Runs the embedded app-server personality migration. +/// +/// Returns `true` when the migration changed config and the caller should reload it. +pub async fn migrate_personality_if_needed( + codex_home: &Path, + config_toml: &ConfigToml, + state_db: Option, +) -> IoResult { + let status = maybe_migrate_personality(codex_home, config_toml, state_db).await?; + match status { + PersonalityMigrationStatus::Applied => Ok(true), + PersonalityMigrationStatus::SkippedMarker + | PersonalityMigrationStatus::SkippedExplicitPersonality + | PersonalityMigrationStatus::SkippedNoSessions => Ok(false), + } +} + /// Raw app-server request result for typed in-process requests. /// /// Even on the in-process path, successful responses still travel back through @@ -158,6 +168,7 @@ pub(crate) fn server_notification_requires_delivery(notification: &ServerNotific ServerNotification::TurnCompleted(_) | ServerNotification::ThreadSettingsUpdated(_) | ServerNotification::ItemCompleted(_) + | ServerNotification::ExternalAgentConfigImportCompleted(_) | ServerNotification::AgentMessageDelta(_) | ServerNotification::PlanDelta(_) | ServerNotification::ReasoningSummaryTextDelta(_) @@ -2139,6 +2150,13 @@ mod tests { ) ) )); + assert!(event_requires_delivery( + &InProcessServerEvent::ServerNotification( + codex_app_server_protocol::ServerNotification::ExternalAgentConfigImportCompleted( + codex_app_server_protocol::ExternalAgentConfigImportCompletedNotification {}, + ) + ) + )); assert!(!event_requires_delivery(&InProcessServerEvent::Lagged { skipped: 1 })); diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index a91a5e7bf9..1e6420d804 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -1993,6 +1993,13 @@ ], "type": "object" }, + "RealtimeConversationVersion": { + "enum": [ + "v1", + "v2" + ], + "type": "string" + }, "RealtimeOutputModality": { "enum": [ "text", @@ -3196,6 +3203,17 @@ ], "type": "object" }, + "ThreadDeleteParams": { + "properties": { + "threadId": { + "type": "string" + } + }, + "required": [ + "threadId" + ], + "type": "object" + }, "ThreadForkParams": { "description": "There are two ways to fork a thread: 1. By thread_id: load the thread from disk by thread_id and fork it into a new thread. 2. By path: load the thread from disk by path and fork it into a new thread.\n\nIf using a non-empty path, the thread_id param will be ignored. Empty string path values are treated as absent.\n\nPrefer using thread_id whenever possible.", "properties": { @@ -4511,6 +4529,30 @@ "title": "Thread/archiveRequest", "type": "object" }, + { + "properties": { + "id": { + "$ref": "#/definitions/RequestId" + }, + "method": { + "enum": [ + "thread/delete" + ], + "title": "Thread/deleteRequestMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/ThreadDeleteParams" + } + }, + "required": [ + "id", + "method", + "params" + ], + "title": "Thread/deleteRequest", + "type": "object" + }, { "properties": { "id": { diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index e0651d1727..1c37ea6f33 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -532,6 +532,13 @@ "personalAccessToken" ], "type": "string" + }, + { + "description": "Amazon Bedrock bearer token managed by Codex.", + "enum": [ + "bedrockApiKey" + ], + "type": "string" } ] }, @@ -3523,6 +3530,17 @@ ], "type": "object" }, + "ThreadDeletedNotification": { + "properties": { + "threadId": { + "type": "string" + } + }, + "required": [ + "threadId" + ], + "type": "object" + }, "ThreadGoal": { "properties": { "createdAt": { @@ -5447,6 +5465,26 @@ "title": "Thread/archivedNotification", "type": "object" }, + { + "properties": { + "method": { + "enum": [ + "thread/deleted" + ], + "title": "Thread/deletedNotificationMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/ThreadDeletedNotification" + } + }, + "required": [ + "method", + "params" + ], + "title": "Thread/deletedNotification", + "type": "object" + }, { "properties": { "method": { 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 8b93be2558..c35f8e16dc 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 @@ -324,6 +324,30 @@ "title": "Thread/archiveRequest", "type": "object" }, + { + "properties": { + "id": { + "$ref": "#/definitions/v2/RequestId" + }, + "method": { + "enum": [ + "thread/delete" + ], + "title": "Thread/deleteRequestMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/v2/ThreadDeleteParams" + } + }, + "required": [ + "id", + "method", + "params" + ], + "title": "Thread/deleteRequest", + "type": "object" + }, { "properties": { "id": { @@ -4061,6 +4085,26 @@ "title": "Thread/archivedNotification", "type": "object" }, + { + "properties": { + "method": { + "enum": [ + "thread/deleted" + ], + "title": "Thread/deletedNotificationMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/v2/ThreadDeletedNotification" + } + }, + "required": [ + "method", + "params" + ], + "title": "Thread/deletedNotification", + "type": "object" + }, { "properties": { "method": { @@ -6389,15 +6433,11 @@ }, "name": { "type": "string" - }, - "needsAuth": { - "type": "boolean" } }, "required": [ "id", - "name", - "needsAuth" + "name" ], "type": "object" }, @@ -6684,6 +6724,13 @@ "personalAccessToken" ], "type": "string" + }, + { + "description": "Amazon Bedrock bearer token managed by Codex.", + "enum": [ + "bedrockApiKey" + ], + "type": "string" } ] }, @@ -16053,6 +16100,37 @@ "title": "ThreadCompactStartResponse", "type": "object" }, + "ThreadDeleteParams": { + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "threadId": { + "type": "string" + } + }, + "required": [ + "threadId" + ], + "title": "ThreadDeleteParams", + "type": "object" + }, + "ThreadDeleteResponse": { + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "ThreadDeleteResponse", + "type": "object" + }, + "ThreadDeletedNotification": { + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "threadId": { + "type": "string" + } + }, + "required": [ + "threadId" + ], + "title": "ThreadDeletedNotification", + "type": "object" + }, "ThreadForkParams": { "$schema": "http://json-schema.org/draft-07/schema#", "description": "There are two ways to fork a thread: 1. By thread_id: load the thread from disk by thread_id and fork it into a new thread. 2. By path: load the thread from disk by path and fork it into a new thread.\n\nIf using a non-empty path, the thread_id param will be ignored. Empty string path values are treated as absent.\n\nPrefer using thread_id whenever possible.", 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 f3691bd0df..6831ce3357 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 @@ -711,15 +711,11 @@ }, "name": { "type": "string" - }, - "needsAuth": { - "type": "boolean" } }, "required": [ "id", - "name", - "needsAuth" + "name" ], "type": "object" }, @@ -1006,6 +1002,13 @@ "personalAccessToken" ], "type": "string" + }, + { + "description": "Amazon Bedrock bearer token managed by Codex.", + "enum": [ + "bedrockApiKey" + ], + "type": "string" } ] }, @@ -1263,6 +1266,30 @@ "title": "Thread/archiveRequest", "type": "object" }, + { + "properties": { + "id": { + "$ref": "#/definitions/RequestId" + }, + "method": { + "enum": [ + "thread/delete" + ], + "title": "Thread/deleteRequestMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/ThreadDeleteParams" + } + }, + "required": [ + "id", + "method", + "params" + ], + "title": "Thread/deleteRequest", + "type": "object" + }, { "properties": { "id": { @@ -11726,6 +11753,26 @@ "title": "Thread/archivedNotification", "type": "object" }, + { + "properties": { + "method": { + "enum": [ + "thread/deleted" + ], + "title": "Thread/deletedNotificationMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/ThreadDeletedNotification" + } + }, + "required": [ + "method", + "params" + ], + "title": "Thread/deletedNotification", + "type": "object" + }, { "properties": { "method": { @@ -13870,6 +13917,37 @@ "title": "ThreadCompactStartResponse", "type": "object" }, + "ThreadDeleteParams": { + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "threadId": { + "type": "string" + } + }, + "required": [ + "threadId" + ], + "title": "ThreadDeleteParams", + "type": "object" + }, + "ThreadDeleteResponse": { + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "ThreadDeleteResponse", + "type": "object" + }, + "ThreadDeletedNotification": { + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "threadId": { + "type": "string" + } + }, + "required": [ + "threadId" + ], + "title": "ThreadDeletedNotification", + "type": "object" + }, "ThreadForkParams": { "$schema": "http://json-schema.org/draft-07/schema#", "description": "There are two ways to fork a thread: 1. By thread_id: load the thread from disk by thread_id and fork it into a new thread. 2. By path: load the thread from disk by path and fork it into a new thread.\n\nIf using a non-empty path, the thread_id param will be ignored. Empty string path values are treated as absent.\n\nPrefer using thread_id whenever possible.", diff --git a/codex-rs/app-server-protocol/schema/json/v2/AccountUpdatedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/AccountUpdatedNotification.json index 013592c9ed..5c1847c56c 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/AccountUpdatedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/AccountUpdatedNotification.json @@ -38,6 +38,13 @@ "personalAccessToken" ], "type": "string" + }, + { + "description": "Amazon Bedrock bearer token managed by Codex.", + "enum": [ + "bedrockApiKey" + ], + "type": "string" } ] }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/PluginInstallResponse.json b/codex-rs/app-server-protocol/schema/json/v2/PluginInstallResponse.json index 2ca7fda461..b02af0bf53 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/PluginInstallResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/PluginInstallResponse.json @@ -21,15 +21,11 @@ }, "name": { "type": "string" - }, - "needsAuth": { - "type": "boolean" } }, "required": [ "id", - "name", - "needsAuth" + "name" ], "type": "object" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/PluginReadResponse.json b/codex-rs/app-server-protocol/schema/json/v2/PluginReadResponse.json index 5a4c2e5b09..40720a5249 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/PluginReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/PluginReadResponse.json @@ -25,15 +25,11 @@ }, "name": { "type": "string" - }, - "needsAuth": { - "type": "boolean" } }, "required": [ "id", - "name", - "needsAuth" + "name" ], "type": "object" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadDeleteParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadDeleteParams.json new file mode 100644 index 0000000000..1711e11a2f --- /dev/null +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadDeleteParams.json @@ -0,0 +1,13 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "threadId": { + "type": "string" + } + }, + "required": [ + "threadId" + ], + "title": "ThreadDeleteParams", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadDeleteResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadDeleteResponse.json new file mode 100644 index 0000000000..ff9f485338 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadDeleteResponse.json @@ -0,0 +1,5 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "ThreadDeleteResponse", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadDeletedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadDeletedNotification.json new file mode 100644 index 0000000000..53011ea014 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadDeletedNotification.json @@ -0,0 +1,13 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "threadId": { + "type": "string" + } + }, + "required": [ + "threadId" + ], + "title": "ThreadDeletedNotification", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/app-server-protocol/schema/typescript/AuthMode.ts b/codex-rs/app-server-protocol/schema/typescript/AuthMode.ts index 1cb6ccb6b3..77bc037ebe 100644 --- a/codex-rs/app-server-protocol/schema/typescript/AuthMode.ts +++ b/codex-rs/app-server-protocol/schema/typescript/AuthMode.ts @@ -5,4 +5,4 @@ /** * Authentication mode for OpenAI-backed providers. */ -export type AuthMode = "apikey" | "chatgpt" | "chatgptAuthTokens" | "agentIdentity" | "personalAccessToken"; +export type AuthMode = "apikey" | "chatgpt" | "chatgptAuthTokens" | "agentIdentity" | "personalAccessToken" | "bedrockApiKey"; diff --git a/codex-rs/app-server-protocol/schema/typescript/ClientRequest.ts b/codex-rs/app-server-protocol/schema/typescript/ClientRequest.ts index c91792c726..b6a13fe307 100644 --- a/codex-rs/app-server-protocol/schema/typescript/ClientRequest.ts +++ b/codex-rs/app-server-protocol/schema/typescript/ClientRequest.ts @@ -62,6 +62,7 @@ import type { SkillsListParams } from "./v2/SkillsListParams"; import type { ThreadApproveGuardianDeniedActionParams } from "./v2/ThreadApproveGuardianDeniedActionParams"; import type { ThreadArchiveParams } from "./v2/ThreadArchiveParams"; import type { ThreadCompactStartParams } from "./v2/ThreadCompactStartParams"; +import type { ThreadDeleteParams } from "./v2/ThreadDeleteParams"; import type { ThreadForkParams } from "./v2/ThreadForkParams"; import type { ThreadGoalClearParams } from "./v2/ThreadGoalClearParams"; import type { ThreadGoalGetParams } from "./v2/ThreadGoalGetParams"; @@ -86,4 +87,4 @@ import type { WindowsSandboxSetupStartParams } from "./v2/WindowsSandboxSetupSta /** * Request from the client to the server. */ -export type ClientRequest ={ "method": "initialize", id: RequestId, params: InitializeParams, } | { "method": "thread/start", id: RequestId, params: ThreadStartParams, } | { "method": "thread/resume", id: RequestId, params: ThreadResumeParams, } | { "method": "thread/fork", id: RequestId, params: ThreadForkParams, } | { "method": "thread/archive", id: RequestId, params: ThreadArchiveParams, } | { "method": "thread/unsubscribe", id: RequestId, params: ThreadUnsubscribeParams, } | { "method": "thread/name/set", id: RequestId, params: ThreadSetNameParams, } | { "method": "thread/goal/set", id: RequestId, params: ThreadGoalSetParams, } | { "method": "thread/goal/get", id: RequestId, params: ThreadGoalGetParams, } | { "method": "thread/goal/clear", id: RequestId, params: ThreadGoalClearParams, } | { "method": "thread/metadata/update", id: RequestId, params: ThreadMetadataUpdateParams, } | { "method": "thread/unarchive", id: RequestId, params: ThreadUnarchiveParams, } | { "method": "thread/compact/start", id: RequestId, params: ThreadCompactStartParams, } | { "method": "thread/shellCommand", id: RequestId, params: ThreadShellCommandParams, } | { "method": "thread/approveGuardianDeniedAction", id: RequestId, params: ThreadApproveGuardianDeniedActionParams, } | { "method": "thread/rollback", id: RequestId, params: ThreadRollbackParams, } | { "method": "thread/list", id: RequestId, params: ThreadListParams, } | { "method": "thread/loaded/list", id: RequestId, params: ThreadLoadedListParams, } | { "method": "thread/read", id: RequestId, params: ThreadReadParams, } | { "method": "thread/inject_items", id: RequestId, params: ThreadInjectItemsParams, } | { "method": "skills/list", id: RequestId, params: SkillsListParams, } | { "method": "skills/extraRoots/set", id: RequestId, params: SkillsExtraRootsSetParams, } | { "method": "hooks/list", id: RequestId, params: HooksListParams, } | { "method": "marketplace/add", id: RequestId, params: MarketplaceAddParams, } | { "method": "marketplace/remove", id: RequestId, params: MarketplaceRemoveParams, } | { "method": "marketplace/upgrade", id: RequestId, params: MarketplaceUpgradeParams, } | { "method": "plugin/list", id: RequestId, params: PluginListParams, } | { "method": "plugin/installed", id: RequestId, params: PluginInstalledParams, } | { "method": "plugin/read", id: RequestId, params: PluginReadParams, } | { "method": "plugin/skill/read", id: RequestId, params: PluginSkillReadParams, } | { "method": "plugin/share/save", id: RequestId, params: PluginShareSaveParams, } | { "method": "plugin/share/updateTargets", id: RequestId, params: PluginShareUpdateTargetsParams, } | { "method": "plugin/share/list", id: RequestId, params: PluginShareListParams, } | { "method": "plugin/share/checkout", id: RequestId, params: PluginShareCheckoutParams, } | { "method": "plugin/share/delete", id: RequestId, params: PluginShareDeleteParams, } | { "method": "app/list", id: RequestId, params: AppsListParams, } | { "method": "fs/readFile", id: RequestId, params: FsReadFileParams, } | { "method": "fs/writeFile", id: RequestId, params: FsWriteFileParams, } | { "method": "fs/createDirectory", id: RequestId, params: FsCreateDirectoryParams, } | { "method": "fs/getMetadata", id: RequestId, params: FsGetMetadataParams, } | { "method": "fs/readDirectory", id: RequestId, params: FsReadDirectoryParams, } | { "method": "fs/remove", id: RequestId, params: FsRemoveParams, } | { "method": "fs/copy", id: RequestId, params: FsCopyParams, } | { "method": "fs/watch", id: RequestId, params: FsWatchParams, } | { "method": "fs/unwatch", id: RequestId, params: FsUnwatchParams, } | { "method": "skills/config/write", id: RequestId, params: SkillsConfigWriteParams, } | { "method": "plugin/install", id: RequestId, params: PluginInstallParams, } | { "method": "plugin/uninstall", id: RequestId, params: PluginUninstallParams, } | { "method": "turn/start", id: RequestId, params: TurnStartParams, } | { "method": "turn/steer", id: RequestId, params: TurnSteerParams, } | { "method": "turn/interrupt", id: RequestId, params: TurnInterruptParams, } | { "method": "review/start", id: RequestId, params: ReviewStartParams, } | { "method": "model/list", id: RequestId, params: ModelListParams, } | { "method": "modelProvider/capabilities/read", id: RequestId, params: ModelProviderCapabilitiesReadParams, } | { "method": "experimentalFeature/list", id: RequestId, params: ExperimentalFeatureListParams, } | { "method": "permissionProfile/list", id: RequestId, params: PermissionProfileListParams, } | { "method": "experimentalFeature/enablement/set", id: RequestId, params: ExperimentalFeatureEnablementSetParams, } | { "method": "mcpServer/oauth/login", id: RequestId, params: McpServerOauthLoginParams, } | { "method": "config/mcpServer/reload", id: RequestId, params: undefined, } | { "method": "mcpServerStatus/list", id: RequestId, params: ListMcpServerStatusParams, } | { "method": "mcpServer/resource/read", id: RequestId, params: McpResourceReadParams, } | { "method": "mcpServer/tool/call", id: RequestId, params: McpServerToolCallParams, } | { "method": "windowsSandbox/setupStart", id: RequestId, params: WindowsSandboxSetupStartParams, } | { "method": "windowsSandbox/readiness", id: RequestId, params: undefined, } | { "method": "account/login/start", id: RequestId, params: LoginAccountParams, } | { "method": "account/login/cancel", id: RequestId, params: CancelLoginAccountParams, } | { "method": "account/logout", id: RequestId, params: undefined, } | { "method": "account/rateLimits/read", id: RequestId, params: undefined, } | { "method": "account/usage/read", id: RequestId, params: undefined, } | { "method": "account/sendAddCreditsNudgeEmail", id: RequestId, params: SendAddCreditsNudgeEmailParams, } | { "method": "feedback/upload", id: RequestId, params: FeedbackUploadParams, } | { "method": "command/exec", id: RequestId, params: CommandExecParams, } | { "method": "command/exec/write", id: RequestId, params: CommandExecWriteParams, } | { "method": "command/exec/terminate", id: RequestId, params: CommandExecTerminateParams, } | { "method": "command/exec/resize", id: RequestId, params: CommandExecResizeParams, } | { "method": "config/read", id: RequestId, params: ConfigReadParams, } | { "method": "externalAgentConfig/detect", id: RequestId, params: ExternalAgentConfigDetectParams, } | { "method": "externalAgentConfig/import", id: RequestId, params: ExternalAgentConfigImportParams, } | { "method": "config/value/write", id: RequestId, params: ConfigValueWriteParams, } | { "method": "config/batchWrite", id: RequestId, params: ConfigBatchWriteParams, } | { "method": "configRequirements/read", id: RequestId, params: undefined, } | { "method": "account/read", id: RequestId, params: GetAccountParams, } | { "method": "getConversationSummary", id: RequestId, params: GetConversationSummaryParams, } | { "method": "gitDiffToRemote", id: RequestId, params: GitDiffToRemoteParams, } | { "method": "getAuthStatus", id: RequestId, params: GetAuthStatusParams, } | { "method": "fuzzyFileSearch", id: RequestId, params: FuzzyFileSearchParams, }; +export type ClientRequest ={ "method": "initialize", id: RequestId, params: InitializeParams, } | { "method": "thread/start", id: RequestId, params: ThreadStartParams, } | { "method": "thread/resume", id: RequestId, params: ThreadResumeParams, } | { "method": "thread/fork", id: RequestId, params: ThreadForkParams, } | { "method": "thread/archive", id: RequestId, params: ThreadArchiveParams, } | { "method": "thread/delete", id: RequestId, params: ThreadDeleteParams, } | { "method": "thread/unsubscribe", id: RequestId, params: ThreadUnsubscribeParams, } | { "method": "thread/name/set", id: RequestId, params: ThreadSetNameParams, } | { "method": "thread/goal/set", id: RequestId, params: ThreadGoalSetParams, } | { "method": "thread/goal/get", id: RequestId, params: ThreadGoalGetParams, } | { "method": "thread/goal/clear", id: RequestId, params: ThreadGoalClearParams, } | { "method": "thread/metadata/update", id: RequestId, params: ThreadMetadataUpdateParams, } | { "method": "thread/unarchive", id: RequestId, params: ThreadUnarchiveParams, } | { "method": "thread/compact/start", id: RequestId, params: ThreadCompactStartParams, } | { "method": "thread/shellCommand", id: RequestId, params: ThreadShellCommandParams, } | { "method": "thread/approveGuardianDeniedAction", id: RequestId, params: ThreadApproveGuardianDeniedActionParams, } | { "method": "thread/rollback", id: RequestId, params: ThreadRollbackParams, } | { "method": "thread/list", id: RequestId, params: ThreadListParams, } | { "method": "thread/loaded/list", id: RequestId, params: ThreadLoadedListParams, } | { "method": "thread/read", id: RequestId, params: ThreadReadParams, } | { "method": "thread/inject_items", id: RequestId, params: ThreadInjectItemsParams, } | { "method": "skills/list", id: RequestId, params: SkillsListParams, } | { "method": "skills/extraRoots/set", id: RequestId, params: SkillsExtraRootsSetParams, } | { "method": "hooks/list", id: RequestId, params: HooksListParams, } | { "method": "marketplace/add", id: RequestId, params: MarketplaceAddParams, } | { "method": "marketplace/remove", id: RequestId, params: MarketplaceRemoveParams, } | { "method": "marketplace/upgrade", id: RequestId, params: MarketplaceUpgradeParams, } | { "method": "plugin/list", id: RequestId, params: PluginListParams, } | { "method": "plugin/installed", id: RequestId, params: PluginInstalledParams, } | { "method": "plugin/read", id: RequestId, params: PluginReadParams, } | { "method": "plugin/skill/read", id: RequestId, params: PluginSkillReadParams, } | { "method": "plugin/share/save", id: RequestId, params: PluginShareSaveParams, } | { "method": "plugin/share/updateTargets", id: RequestId, params: PluginShareUpdateTargetsParams, } | { "method": "plugin/share/list", id: RequestId, params: PluginShareListParams, } | { "method": "plugin/share/checkout", id: RequestId, params: PluginShareCheckoutParams, } | { "method": "plugin/share/delete", id: RequestId, params: PluginShareDeleteParams, } | { "method": "app/list", id: RequestId, params: AppsListParams, } | { "method": "fs/readFile", id: RequestId, params: FsReadFileParams, } | { "method": "fs/writeFile", id: RequestId, params: FsWriteFileParams, } | { "method": "fs/createDirectory", id: RequestId, params: FsCreateDirectoryParams, } | { "method": "fs/getMetadata", id: RequestId, params: FsGetMetadataParams, } | { "method": "fs/readDirectory", id: RequestId, params: FsReadDirectoryParams, } | { "method": "fs/remove", id: RequestId, params: FsRemoveParams, } | { "method": "fs/copy", id: RequestId, params: FsCopyParams, } | { "method": "fs/watch", id: RequestId, params: FsWatchParams, } | { "method": "fs/unwatch", id: RequestId, params: FsUnwatchParams, } | { "method": "skills/config/write", id: RequestId, params: SkillsConfigWriteParams, } | { "method": "plugin/install", id: RequestId, params: PluginInstallParams, } | { "method": "plugin/uninstall", id: RequestId, params: PluginUninstallParams, } | { "method": "turn/start", id: RequestId, params: TurnStartParams, } | { "method": "turn/steer", id: RequestId, params: TurnSteerParams, } | { "method": "turn/interrupt", id: RequestId, params: TurnInterruptParams, } | { "method": "review/start", id: RequestId, params: ReviewStartParams, } | { "method": "model/list", id: RequestId, params: ModelListParams, } | { "method": "modelProvider/capabilities/read", id: RequestId, params: ModelProviderCapabilitiesReadParams, } | { "method": "experimentalFeature/list", id: RequestId, params: ExperimentalFeatureListParams, } | { "method": "permissionProfile/list", id: RequestId, params: PermissionProfileListParams, } | { "method": "experimentalFeature/enablement/set", id: RequestId, params: ExperimentalFeatureEnablementSetParams, } | { "method": "mcpServer/oauth/login", id: RequestId, params: McpServerOauthLoginParams, } | { "method": "config/mcpServer/reload", id: RequestId, params: undefined, } | { "method": "mcpServerStatus/list", id: RequestId, params: ListMcpServerStatusParams, } | { "method": "mcpServer/resource/read", id: RequestId, params: McpResourceReadParams, } | { "method": "mcpServer/tool/call", id: RequestId, params: McpServerToolCallParams, } | { "method": "windowsSandbox/setupStart", id: RequestId, params: WindowsSandboxSetupStartParams, } | { "method": "windowsSandbox/readiness", id: RequestId, params: undefined, } | { "method": "account/login/start", id: RequestId, params: LoginAccountParams, } | { "method": "account/login/cancel", id: RequestId, params: CancelLoginAccountParams, } | { "method": "account/logout", id: RequestId, params: undefined, } | { "method": "account/rateLimits/read", id: RequestId, params: undefined, } | { "method": "account/usage/read", id: RequestId, params: undefined, } | { "method": "account/sendAddCreditsNudgeEmail", id: RequestId, params: SendAddCreditsNudgeEmailParams, } | { "method": "feedback/upload", id: RequestId, params: FeedbackUploadParams, } | { "method": "command/exec", id: RequestId, params: CommandExecParams, } | { "method": "command/exec/write", id: RequestId, params: CommandExecWriteParams, } | { "method": "command/exec/terminate", id: RequestId, params: CommandExecTerminateParams, } | { "method": "command/exec/resize", id: RequestId, params: CommandExecResizeParams, } | { "method": "config/read", id: RequestId, params: ConfigReadParams, } | { "method": "externalAgentConfig/detect", id: RequestId, params: ExternalAgentConfigDetectParams, } | { "method": "externalAgentConfig/import", id: RequestId, params: ExternalAgentConfigImportParams, } | { "method": "config/value/write", id: RequestId, params: ConfigValueWriteParams, } | { "method": "config/batchWrite", id: RequestId, params: ConfigBatchWriteParams, } | { "method": "configRequirements/read", id: RequestId, params: undefined, } | { "method": "account/read", id: RequestId, params: GetAccountParams, } | { "method": "getConversationSummary", id: RequestId, params: GetConversationSummaryParams, } | { "method": "gitDiffToRemote", id: RequestId, params: GitDiffToRemoteParams, } | { "method": "getAuthStatus", id: RequestId, params: GetAuthStatusParams, } | { "method": "fuzzyFileSearch", id: RequestId, params: FuzzyFileSearchParams, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/ServerNotification.ts b/codex-rs/app-server-protocol/schema/typescript/ServerNotification.ts index 2520c71c3c..7cae63a553 100644 --- a/codex-rs/app-server-protocol/schema/typescript/ServerNotification.ts +++ b/codex-rs/app-server-protocol/schema/typescript/ServerNotification.ts @@ -43,6 +43,7 @@ import type { SkillsChangedNotification } from "./v2/SkillsChangedNotification"; import type { TerminalInteractionNotification } from "./v2/TerminalInteractionNotification"; import type { ThreadArchivedNotification } from "./v2/ThreadArchivedNotification"; import type { ThreadClosedNotification } from "./v2/ThreadClosedNotification"; +import type { ThreadDeletedNotification } from "./v2/ThreadDeletedNotification"; import type { ThreadGoalClearedNotification } from "./v2/ThreadGoalClearedNotification"; import type { ThreadGoalUpdatedNotification } from "./v2/ThreadGoalUpdatedNotification"; import type { ThreadNameUpdatedNotification } from "./v2/ThreadNameUpdatedNotification"; @@ -71,4 +72,4 @@ import type { WindowsWorldWritableWarningNotification } from "./v2/WindowsWorldW /** * Notification sent from the server to the client. */ -export type ServerNotification = { "method": "error", "params": ErrorNotification } | { "method": "thread/started", "params": ThreadStartedNotification } | { "method": "thread/status/changed", "params": ThreadStatusChangedNotification } | { "method": "thread/archived", "params": ThreadArchivedNotification } | { "method": "thread/unarchived", "params": ThreadUnarchivedNotification } | { "method": "thread/closed", "params": ThreadClosedNotification } | { "method": "skills/changed", "params": SkillsChangedNotification } | { "method": "thread/name/updated", "params": ThreadNameUpdatedNotification } | { "method": "thread/goal/updated", "params": ThreadGoalUpdatedNotification } | { "method": "thread/goal/cleared", "params": ThreadGoalClearedNotification } | { "method": "thread/settings/updated", "params": ThreadSettingsUpdatedNotification } | { "method": "thread/tokenUsage/updated", "params": ThreadTokenUsageUpdatedNotification } | { "method": "turn/started", "params": TurnStartedNotification } | { "method": "hook/started", "params": HookStartedNotification } | { "method": "turn/completed", "params": TurnCompletedNotification } | { "method": "hook/completed", "params": HookCompletedNotification } | { "method": "turn/diff/updated", "params": TurnDiffUpdatedNotification } | { "method": "turn/plan/updated", "params": TurnPlanUpdatedNotification } | { "method": "item/started", "params": ItemStartedNotification } | { "method": "item/autoApprovalReview/started", "params": ItemGuardianApprovalReviewStartedNotification } | { "method": "item/autoApprovalReview/completed", "params": ItemGuardianApprovalReviewCompletedNotification } | { "method": "item/completed", "params": ItemCompletedNotification } | { "method": "rawResponseItem/completed", "params": RawResponseItemCompletedNotification } | { "method": "item/agentMessage/delta", "params": AgentMessageDeltaNotification } | { "method": "item/plan/delta", "params": PlanDeltaNotification } | { "method": "command/exec/outputDelta", "params": CommandExecOutputDeltaNotification } | { "method": "process/outputDelta", "params": ProcessOutputDeltaNotification } | { "method": "process/exited", "params": ProcessExitedNotification } | { "method": "item/commandExecution/outputDelta", "params": CommandExecutionOutputDeltaNotification } | { "method": "item/commandExecution/terminalInteraction", "params": TerminalInteractionNotification } | { "method": "item/fileChange/outputDelta", "params": FileChangeOutputDeltaNotification } | { "method": "item/fileChange/patchUpdated", "params": FileChangePatchUpdatedNotification } | { "method": "serverRequest/resolved", "params": ServerRequestResolvedNotification } | { "method": "item/mcpToolCall/progress", "params": McpToolCallProgressNotification } | { "method": "mcpServer/oauthLogin/completed", "params": McpServerOauthLoginCompletedNotification } | { "method": "mcpServer/startupStatus/updated", "params": McpServerStatusUpdatedNotification } | { "method": "account/updated", "params": AccountUpdatedNotification } | { "method": "account/rateLimits/updated", "params": AccountRateLimitsUpdatedNotification } | { "method": "app/list/updated", "params": AppListUpdatedNotification } | { "method": "remoteControl/status/changed", "params": RemoteControlStatusChangedNotification } | { "method": "externalAgentConfig/import/completed", "params": ExternalAgentConfigImportCompletedNotification } | { "method": "fs/changed", "params": FsChangedNotification } | { "method": "item/reasoning/summaryTextDelta", "params": ReasoningSummaryTextDeltaNotification } | { "method": "item/reasoning/summaryPartAdded", "params": ReasoningSummaryPartAddedNotification } | { "method": "item/reasoning/textDelta", "params": ReasoningTextDeltaNotification } | { "method": "thread/compacted", "params": ContextCompactedNotification } | { "method": "model/rerouted", "params": ModelReroutedNotification } | { "method": "model/verification", "params": ModelVerificationNotification } | { "method": "turn/moderationMetadata", "params": TurnModerationMetadataNotification } | { "method": "warning", "params": WarningNotification } | { "method": "guardianWarning", "params": GuardianWarningNotification } | { "method": "deprecationNotice", "params": DeprecationNoticeNotification } | { "method": "configWarning", "params": ConfigWarningNotification } | { "method": "fuzzyFileSearch/sessionUpdated", "params": FuzzyFileSearchSessionUpdatedNotification } | { "method": "fuzzyFileSearch/sessionCompleted", "params": FuzzyFileSearchSessionCompletedNotification } | { "method": "thread/realtime/started", "params": ThreadRealtimeStartedNotification } | { "method": "thread/realtime/itemAdded", "params": ThreadRealtimeItemAddedNotification } | { "method": "thread/realtime/transcript/delta", "params": ThreadRealtimeTranscriptDeltaNotification } | { "method": "thread/realtime/transcript/done", "params": ThreadRealtimeTranscriptDoneNotification } | { "method": "thread/realtime/outputAudio/delta", "params": ThreadRealtimeOutputAudioDeltaNotification } | { "method": "thread/realtime/sdp", "params": ThreadRealtimeSdpNotification } | { "method": "thread/realtime/error", "params": ThreadRealtimeErrorNotification } | { "method": "thread/realtime/closed", "params": ThreadRealtimeClosedNotification } | { "method": "windows/worldWritableWarning", "params": WindowsWorldWritableWarningNotification } | { "method": "windowsSandbox/setupCompleted", "params": WindowsSandboxSetupCompletedNotification } | { "method": "account/login/completed", "params": AccountLoginCompletedNotification }; +export type ServerNotification = { "method": "error", "params": ErrorNotification } | { "method": "thread/started", "params": ThreadStartedNotification } | { "method": "thread/status/changed", "params": ThreadStatusChangedNotification } | { "method": "thread/archived", "params": ThreadArchivedNotification } | { "method": "thread/deleted", "params": ThreadDeletedNotification } | { "method": "thread/unarchived", "params": ThreadUnarchivedNotification } | { "method": "thread/closed", "params": ThreadClosedNotification } | { "method": "skills/changed", "params": SkillsChangedNotification } | { "method": "thread/name/updated", "params": ThreadNameUpdatedNotification } | { "method": "thread/goal/updated", "params": ThreadGoalUpdatedNotification } | { "method": "thread/goal/cleared", "params": ThreadGoalClearedNotification } | { "method": "thread/settings/updated", "params": ThreadSettingsUpdatedNotification } | { "method": "thread/tokenUsage/updated", "params": ThreadTokenUsageUpdatedNotification } | { "method": "turn/started", "params": TurnStartedNotification } | { "method": "hook/started", "params": HookStartedNotification } | { "method": "turn/completed", "params": TurnCompletedNotification } | { "method": "hook/completed", "params": HookCompletedNotification } | { "method": "turn/diff/updated", "params": TurnDiffUpdatedNotification } | { "method": "turn/plan/updated", "params": TurnPlanUpdatedNotification } | { "method": "item/started", "params": ItemStartedNotification } | { "method": "item/autoApprovalReview/started", "params": ItemGuardianApprovalReviewStartedNotification } | { "method": "item/autoApprovalReview/completed", "params": ItemGuardianApprovalReviewCompletedNotification } | { "method": "item/completed", "params": ItemCompletedNotification } | { "method": "rawResponseItem/completed", "params": RawResponseItemCompletedNotification } | { "method": "item/agentMessage/delta", "params": AgentMessageDeltaNotification } | { "method": "item/plan/delta", "params": PlanDeltaNotification } | { "method": "command/exec/outputDelta", "params": CommandExecOutputDeltaNotification } | { "method": "process/outputDelta", "params": ProcessOutputDeltaNotification } | { "method": "process/exited", "params": ProcessExitedNotification } | { "method": "item/commandExecution/outputDelta", "params": CommandExecutionOutputDeltaNotification } | { "method": "item/commandExecution/terminalInteraction", "params": TerminalInteractionNotification } | { "method": "item/fileChange/outputDelta", "params": FileChangeOutputDeltaNotification } | { "method": "item/fileChange/patchUpdated", "params": FileChangePatchUpdatedNotification } | { "method": "serverRequest/resolved", "params": ServerRequestResolvedNotification } | { "method": "item/mcpToolCall/progress", "params": McpToolCallProgressNotification } | { "method": "mcpServer/oauthLogin/completed", "params": McpServerOauthLoginCompletedNotification } | { "method": "mcpServer/startupStatus/updated", "params": McpServerStatusUpdatedNotification } | { "method": "account/updated", "params": AccountUpdatedNotification } | { "method": "account/rateLimits/updated", "params": AccountRateLimitsUpdatedNotification } | { "method": "app/list/updated", "params": AppListUpdatedNotification } | { "method": "remoteControl/status/changed", "params": RemoteControlStatusChangedNotification } | { "method": "externalAgentConfig/import/completed", "params": ExternalAgentConfigImportCompletedNotification } | { "method": "fs/changed", "params": FsChangedNotification } | { "method": "item/reasoning/summaryTextDelta", "params": ReasoningSummaryTextDeltaNotification } | { "method": "item/reasoning/summaryPartAdded", "params": ReasoningSummaryPartAddedNotification } | { "method": "item/reasoning/textDelta", "params": ReasoningTextDeltaNotification } | { "method": "thread/compacted", "params": ContextCompactedNotification } | { "method": "model/rerouted", "params": ModelReroutedNotification } | { "method": "model/verification", "params": ModelVerificationNotification } | { "method": "turn/moderationMetadata", "params": TurnModerationMetadataNotification } | { "method": "warning", "params": WarningNotification } | { "method": "guardianWarning", "params": GuardianWarningNotification } | { "method": "deprecationNotice", "params": DeprecationNoticeNotification } | { "method": "configWarning", "params": ConfigWarningNotification } | { "method": "fuzzyFileSearch/sessionUpdated", "params": FuzzyFileSearchSessionUpdatedNotification } | { "method": "fuzzyFileSearch/sessionCompleted", "params": FuzzyFileSearchSessionCompletedNotification } | { "method": "thread/realtime/started", "params": ThreadRealtimeStartedNotification } | { "method": "thread/realtime/itemAdded", "params": ThreadRealtimeItemAddedNotification } | { "method": "thread/realtime/transcript/delta", "params": ThreadRealtimeTranscriptDeltaNotification } | { "method": "thread/realtime/transcript/done", "params": ThreadRealtimeTranscriptDoneNotification } | { "method": "thread/realtime/outputAudio/delta", "params": ThreadRealtimeOutputAudioDeltaNotification } | { "method": "thread/realtime/sdp", "params": ThreadRealtimeSdpNotification } | { "method": "thread/realtime/error", "params": ThreadRealtimeErrorNotification } | { "method": "thread/realtime/closed", "params": ThreadRealtimeClosedNotification } | { "method": "windows/worldWritableWarning", "params": WindowsWorldWritableWarningNotification } | { "method": "windowsSandbox/setupCompleted", "params": WindowsSandboxSetupCompletedNotification } | { "method": "account/login/completed", "params": AccountLoginCompletedNotification }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/AppSummary.ts b/codex-rs/app-server-protocol/schema/typescript/v2/AppSummary.ts index 586c76f8f7..3cdb17d705 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/AppSummary.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/AppSummary.ts @@ -5,4 +5,4 @@ /** * EXPERIMENTAL - app metadata summary for plugin responses. */ -export type AppSummary = { id: string, name: string, description: string | null, installUrl: string | null, needsAuth: boolean, }; +export type AppSummary = { id: string, name: string, description: string | null, installUrl: string | null, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadDeleteParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadDeleteParams.ts new file mode 100644 index 0000000000..909ccda715 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadDeleteParams.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type ThreadDeleteParams = { threadId: string, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadDeleteResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadDeleteResponse.ts new file mode 100644 index 0000000000..1af1c3077f --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadDeleteResponse.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type ThreadDeleteResponse = Record; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadDeletedNotification.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadDeletedNotification.ts new file mode 100644 index 0000000000..5122a2229c --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadDeletedNotification.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type ThreadDeletedNotification = { threadId: string, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts index 71f4c67cd0..b4cd88aadf 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -374,6 +374,9 @@ export type { ThreadArchivedNotification } from "./ThreadArchivedNotification"; export type { ThreadClosedNotification } from "./ThreadClosedNotification"; export type { ThreadCompactStartParams } from "./ThreadCompactStartParams"; export type { ThreadCompactStartResponse } from "./ThreadCompactStartResponse"; +export type { ThreadDeleteParams } from "./ThreadDeleteParams"; +export type { ThreadDeleteResponse } from "./ThreadDeleteResponse"; +export type { ThreadDeletedNotification } from "./ThreadDeletedNotification"; export type { ThreadForkParams } from "./ThreadForkParams"; export type { ThreadForkResponse } from "./ThreadForkResponse"; export type { ThreadGoal } from "./ThreadGoal"; diff --git a/codex-rs/app-server-protocol/src/export.rs b/codex-rs/app-server-protocol/src/export.rs index cd6c7e47b6..5e6c2ad015 100644 --- a/codex-rs/app-server-protocol/src/export.rs +++ b/codex-rs/app-server-protocol/src/export.rs @@ -39,8 +39,11 @@ use ts_rs::TS; pub(crate) const GENERATED_TS_HEADER: &str = "// GENERATED CODE! DO NOT MODIFY BY HAND!\n\n"; const IGNORED_DEFINITIONS: &[&str] = &["Option<()>"]; const JSON_V1_ALLOWLIST: &[&str] = &["InitializeParams", "InitializeResponse"]; -const EXPERIMENTAL_CLIENT_METHOD_DEPENDENCY_TYPES: &[&str] = - &["RemoteControlClient", "RemoteControlClientsListOrder"]; +const EXPERIMENTAL_CLIENT_METHOD_DEPENDENCY_TYPES: &[&str] = &[ + "RemoteControlClient", + "RemoteControlClientsListOrder", + "ThreadBackgroundTerminal", +]; const SPECIAL_DEFINITIONS: &[&str] = &[ "ClientNotification", "ClientRequest", diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index d3c8eae22b..9311fe9058 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -41,6 +41,11 @@ pub enum AuthMode { #[ts(rename = "personalAccessToken")] #[strum(serialize = "personalAccessToken")] PersonalAccessToken, + /// Amazon Bedrock bearer token managed by Codex. + #[serde(rename = "bedrockApiKey")] + #[ts(rename = "bedrockApiKey")] + #[strum(serialize = "bedrockApiKey")] + BedrockApiKey, } impl AuthMode { @@ -48,7 +53,7 @@ impl AuthMode { pub fn has_chatgpt_account(self) -> bool { match self { Self::Chatgpt | Self::ChatgptAuthTokens | Self::PersonalAccessToken => true, - Self::ApiKey | Self::AgentIdentity => false, + Self::ApiKey | Self::AgentIdentity | Self::BedrockApiKey => false, } } } @@ -480,6 +485,11 @@ client_request_definitions! { serialization: thread_id(params.thread_id), response: v2::ThreadArchiveResponse, }, + ThreadDelete => "thread/delete" { + params: v2::ThreadDeleteParams, + serialization: thread_id(params.thread_id), + response: v2::ThreadDeleteResponse, + }, ThreadUnsubscribe => "thread/unsubscribe" { params: v2::ThreadUnsubscribeParams, serialization: thread_id(params.thread_id), @@ -574,6 +584,18 @@ client_request_definitions! { serialization: thread_id(params.thread_id), response: v2::ThreadBackgroundTerminalsCleanResponse, }, + #[experimental("thread/backgroundTerminals/list")] + ThreadBackgroundTerminalsList => "thread/backgroundTerminals/list" { + params: v2::ThreadBackgroundTerminalsListParams, + serialization: thread_id(params.thread_id), + response: v2::ThreadBackgroundTerminalsListResponse, + }, + #[experimental("thread/backgroundTerminals/terminate")] + ThreadBackgroundTerminalsTerminate => "thread/backgroundTerminals/terminate" { + params: v2::ThreadBackgroundTerminalsTerminateParams, + serialization: thread_id(params.thread_id), + response: v2::ThreadBackgroundTerminalsTerminateResponse, + }, ThreadRollback => "thread/rollback" { params: v2::ThreadRollbackParams, serialization: thread_id(params.thread_id), @@ -1522,6 +1544,7 @@ server_notification_definitions! { ThreadStarted => "thread/started" (v2::ThreadStartedNotification), ThreadStatusChanged => "thread/status/changed" (v2::ThreadStatusChangedNotification), ThreadArchived => "thread/archived" (v2::ThreadArchivedNotification), + ThreadDeleted => "thread/deleted" (v2::ThreadDeletedNotification), ThreadUnarchived => "thread/unarchived" (v2::ThreadUnarchivedNotification), ThreadClosed => "thread/closed" (v2::ThreadClosedNotification), SkillsChanged => "skills/changed" (v2::SkillsChangedNotification), @@ -2936,16 +2959,66 @@ mod tests { Ok(()) } + #[test] + fn serialize_thread_background_terminals_list() -> Result<()> { + let request = ClientRequest::ThreadBackgroundTerminalsList { + request_id: RequestId::Integer(8), + params: v2::ThreadBackgroundTerminalsListParams { + thread_id: "thr_123".to_string(), + cursor: None, + limit: None, + }, + }; + assert_eq!( + json!({ + "method": "thread/backgroundTerminals/list", + "id": 8, + "params": { + "threadId": "thr_123", + "cursor": null, + "limit": null + } + }), + serde_json::to_value(&request)?, + ); + Ok(()) + } + + #[test] + fn serialize_thread_background_terminals_terminate() -> Result<()> { + let request = ClientRequest::ThreadBackgroundTerminalsTerminate { + request_id: RequestId::Integer(8), + params: v2::ThreadBackgroundTerminalsTerminateParams { + thread_id: "thr_123".to_string(), + process_id: "42".to_string(), + }, + }; + assert_eq!( + json!({ + "method": "thread/backgroundTerminals/terminate", + "id": 8, + "params": { + "threadId": "thr_123", + "processId": "42" + } + }), + serde_json::to_value(&request)?, + ); + Ok(()) + } + #[test] fn serialize_thread_realtime_start() -> Result<()> { let request = ClientRequest::ThreadRealtimeStart { request_id: RequestId::Integer(9), params: v2::ThreadRealtimeStartParams { thread_id: "thr_123".to_string(), + model: Some("realtime-treatment-model".to_string()), output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("You are on a call".to_string())), realtime_session_id: Some("sess_456".to_string()), transport: None, + version: Some(RealtimeConversationVersion::V1), voice: Some(RealtimeVoice::Marin), }, }; @@ -2955,10 +3028,12 @@ mod tests { "id": 9, "params": { "threadId": "thr_123", + "model": "realtime-treatment-model", "outputModality": "audio", "prompt": "You are on a call", "realtimeSessionId": "sess_456", "transport": null, + "version": "v1", "voice": "marin" } }), @@ -2973,10 +3048,12 @@ mod tests { request_id: RequestId::Integer(9), params: v2::ThreadRealtimeStartParams { thread_id: "thr_123".to_string(), + model: None, output_modality: RealtimeOutputModality::Audio, prompt: None, realtime_session_id: None, transport: None, + version: None, voice: None, }, }; @@ -2986,9 +3063,11 @@ mod tests { "id": 9, "params": { "threadId": "thr_123", + "model": null, "outputModality": "audio", "realtimeSessionId": null, "transport": null, + "version": null, "voice": null } }), @@ -2999,10 +3078,12 @@ mod tests { request_id: RequestId::Integer(9), params: v2::ThreadRealtimeStartParams { thread_id: "thr_123".to_string(), + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(None), realtime_session_id: None, transport: None, + version: None, voice: None, }, }; @@ -3012,10 +3093,12 @@ mod tests { "id": 9, "params": { "threadId": "thr_123", + "model": null, "outputModality": "audio", "prompt": null, "realtimeSessionId": null, "transport": null, + "version": null, "voice": null } }), @@ -3168,10 +3251,12 @@ mod tests { request_id: RequestId::Integer(1), params: v2::ThreadRealtimeStartParams { thread_id: "thr_123".to_string(), + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("You are on a call".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, }, }; diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 075b3e043e..1305ec00fb 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -2939,6 +2939,7 @@ mod tests { RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: None, + window_id: None, }), RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { turn_id: "turn-compact".into(), diff --git a/codex-rs/app-server-protocol/src/protocol/v2/apps.rs b/codex-rs/app-server-protocol/src/protocol/v2/apps.rs index 9f46525e6c..9d4ccbd3af 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/apps.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/apps.rs @@ -111,7 +111,6 @@ pub struct AppSummary { pub name: String, pub description: Option, pub install_url: Option, - pub needs_auth: bool, } impl From for AppSummary { @@ -121,7 +120,6 @@ impl From for AppSummary { name: value.name, description: value.description, install_url: value.install_url, - needs_auth: false, } } } diff --git a/codex-rs/app-server-protocol/src/protocol/v2/realtime.rs b/codex-rs/app-server-protocol/src/protocol/v2/realtime.rs index c6ea0744de..a86cd8067d 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/realtime.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/realtime.rs @@ -65,6 +65,9 @@ impl From for CoreRealtimeAudioFrame { #[ts(export_to = "v2/")] pub struct ThreadRealtimeStartParams { pub thread_id: String, + /// Overrides the configured realtime model for this session only. + #[ts(optional = nullable)] + pub model: Option, /// Selects text or audio output for the realtime session. Transport and voice stay /// independent so clients can choose how they connect separately from what the model emits. pub output_modality: RealtimeOutputModality, @@ -80,6 +83,9 @@ pub struct ThreadRealtimeStartParams { pub realtime_session_id: Option, #[ts(optional = nullable)] pub transport: Option, + /// Overrides the configured realtime protocol version for this session only. + #[ts(optional = nullable)] + pub version: Option, #[ts(optional = nullable)] pub voice: Option, } diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs index a54479fdde..cdc8087553 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs @@ -601,6 +601,18 @@ pub struct ThreadArchiveParams { #[ts(export_to = "v2/")] pub struct ThreadArchiveResponse {} +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadDeleteParams { + pub thread_id: String, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadDeleteResponse {} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] @@ -942,6 +954,57 @@ pub struct ThreadBackgroundTerminalsCleanParams { #[ts(export_to = "v2/")] pub struct ThreadBackgroundTerminalsCleanResponse {} +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadBackgroundTerminalsListParams { + pub thread_id: String, + /// Opaque pagination cursor returned by a previous call. + #[ts(optional = nullable)] + pub cursor: Option, + /// Optional page size. + #[ts(optional = nullable)] + pub limit: Option, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadBackgroundTerminal { + pub item_id: String, + pub process_id: String, + pub command: String, + pub cwd: AbsolutePathBuf, + pub os_pid: Option, + pub cpu_percent: Option, + pub rss_kb: Option, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadBackgroundTerminalsListResponse { + pub data: Vec, + /// Opaque cursor to pass to the next call to continue after the last item. + /// If None, there are no more items to return. + pub next_cursor: Option, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadBackgroundTerminalsTerminateParams { + pub thread_id: String, + pub process_id: String, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadBackgroundTerminalsTerminateResponse { + pub terminated: bool, +} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] @@ -1337,6 +1400,13 @@ pub struct ThreadArchivedNotification { pub thread_id: String, } +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadDeletedNotification { + pub thread_id: String, +} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs index 1e653c2969..1d46404855 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs @@ -116,6 +116,7 @@ fn remote_control_auth_dot_json(account_id: Option<&str>) -> AuthDotJson { last_refresh: Some(chrono::Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, } } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index e0208b1f1c..ef74527e06 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -1842,6 +1842,7 @@ mod tests { last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, } } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 0ba359b64f..2e074dc5ba 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -151,18 +151,21 @@ Example with notification opt-out: - `thread/settings/updated` — experimental notification emitted to subscribed clients when a loaded thread’s effective next-turn settings change; includes `threadId` and the full `threadSettings`. - `thread/status/changed` — notification emitted when a loaded thread’s status changes (`threadId` + new `status`). - `thread/archive` — move a thread’s rollout file into the archived directory and attempt to move any spawned descendant thread rollout files; returns `{}` on success and emits `thread/archived` for each archived thread. +- `thread/delete` — hard-delete an active or archived thread and any spawned descendant threads; returns `{}` on success and emits `thread/deleted` for each deleted thread. - `thread/unsubscribe` — unsubscribe this connection from thread turn/item events. If this was the last subscriber, the server keeps the thread loaded and unloads it only after it has had no subscribers and no thread activity for 30 minutes, then emits `thread/closed`. - `thread/name/set` — set or update a thread’s user-facing name for either a loaded thread or a persisted rollout; returns `{}` on success and emits `thread/name/updated` to initialized, opted-in clients. Thread names are not required to be unique; name lookups resolve to the most recently updated thread. - `thread/unarchive` — move an archived rollout file back into the sessions directory; returns the restored `thread` on success and emits `thread/unarchived`. - `thread/compact/start` — trigger conversation history compaction for a thread; returns `{}` immediately while progress streams through standard turn/item notifications. - `thread/shellCommand` — run a user-initiated `!` shell command against a thread; this runs unsandboxed with full access rather than inheriting the thread sandbox policy. Returns `{}` immediately while progress streams through standard turn/item notifications and any active turn receives the formatted output in its message stream. - `thread/backgroundTerminals/clean` — terminate all running background terminals for a thread (experimental; requires `capabilities.experimentalApi`); returns `{}` when the cleanup request is accepted. +- `thread/backgroundTerminals/list` — list running background terminals for a loaded thread (experimental; requires `capabilities.experimentalApi`); returns `data` with the running terminal ids. +- `thread/backgroundTerminals/terminate` — terminate one running background terminal by app-server `processId` (experimental; requires `capabilities.experimentalApi`); returns whether a process was terminated. - `thread/rollback` — drop the last N turns from the agent’s in-memory context and persist a rollback marker in the rollout so future resumes see the pruned history; returns the updated `thread` (with `turns` populated) on success. - `turn/start` — add user input to a thread and begin Codex generation; responds with the initial `turn` object and streams `turn/started`, `item/*`, and `turn/completed` notifications. `clientUserMessageId` is optional; when supplied, the corresponding `userMessage` item echoes it as `clientId`. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. Prefer experimental `permissions` profile selection by id for permission overrides; the legacy `sandboxPolicy` field is still accepted but cannot be combined with `permissions`. For `collaborationMode`, `settings.developer_instructions: null` means "use built-in instructions for the selected mode". - `thread/inject_items` — append raw Responses API items to a loaded thread’s model-visible history without starting a user turn; returns `{}` on success. - `turn/steer` — add user input to an already in-flight regular turn without starting a new turn; returns the active `turnId` that accepted the input. `clientUserMessageId` is optional; when supplied, the corresponding `userMessage` item echoes it as `clientId`. Review and manual compaction turns reject `turn/steer`. - `turn/interrupt` — request cancellation of an in-flight turn by `(thread_id, turn_id)`; success is an empty `{}` response and the turn finishes with `status: "interrupted"`. -- `thread/realtime/start` — start a thread-scoped realtime session (experimental); pass `outputModality: "text"` or `outputModality: "audio"` to choose model output, returns `{}` and streams `thread/realtime/*` notifications. Omit `transport` for the websocket transport, or pass `{ "type": "webrtc", "sdp": "..." }` to create a WebRTC session from a browser-generated SDP offer; the remote answer SDP is emitted as `thread/realtime/sdp`. +- `thread/realtime/start` — start a thread-scoped realtime session (experimental); pass `outputModality: "text"` or `outputModality: "audio"` to choose model output, and optionally pass `model` and `version` to override configured realtime selection for this session only. Returns `{}` and streams `thread/realtime/*` notifications. Omit `transport` for the websocket transport, or pass `{ "type": "webrtc", "sdp": "..." }` to create a WebRTC session from a browser-generated SDP offer; the remote answer SDP is emitted as `thread/realtime/sdp`. - `thread/realtime/appendAudio` — append an input audio chunk to the active realtime session (experimental); returns `{}`. - `thread/realtime/appendText` — append text input to the active realtime session (experimental); returns `{}`. - `thread/realtime/stop` — stop the active realtime session for the thread (experimental); returns `{}`. @@ -203,7 +206,7 @@ Example with notification opt-out: - `marketplace/upgrade` — upgrade all configured Git plugin marketplaces, or one named marketplace when `marketplaceName` is provided. Returns selected marketplace names, upgraded roots, and per-marketplace errors. - `plugin/list` — list discovered plugin marketplaces and plugin state, including effective marketplace install/auth policy metadata, plugin `availability` (`AVAILABLE` by default or `DISABLED_BY_ADMIN` for remote plugins blocked upstream), fail-open `marketplaceLoadErrors` entries for marketplace files that could not be parsed or loaded, and best-effort `featuredPluginIds` for the official curated marketplace. `interface.category` uses the marketplace category when present; otherwise it falls back to the plugin manifest category (**under development; do not call from production clients yet**). - `plugin/installed` — list installed plugin rows plus any explicitly requested local install-suggestion plugin names, without fetching the broader remote catalog. Mention surfaces can use this narrower view when they need plugin mention payloads rather than plugin-page discovery data (**under development; do not call from production clients yet**). -- `plugin/read` — read one plugin by `marketplacePath` plus `pluginName`, returning marketplace info, a list-style `summary`, manifest descriptions/interface metadata, and bundled skills/hooks/apps/MCP server names. Returned plugin skills include their current `enabled` state after local config filtering; bundled hooks are returned as lightweight declaration summaries keyed for correlation with `hooks/list`. Plugin app summaries also include `needsAuth` when the server can determine connector accessibility (**under development; do not call from production clients yet**). +- `plugin/read` — read one plugin by `marketplacePath` plus `pluginName`, returning marketplace info, a list-style `summary`, manifest descriptions/interface metadata, and bundled skills/hooks/apps/MCP server names. Returned plugin skills include their current `enabled` state after local config filtering; bundled hooks are returned as lightweight declaration summaries keyed for correlation with `hooks/list`. Use `plugin/install`'s `appsNeedingAuth` to drive post-install authentication and `app/list`'s `isAccessible` to determine current connector accessibility (**under development; do not call from production clients yet**). - `plugin/skill/read` — read remote plugin skill markdown on demand by `remoteMarketplaceName`, `remotePluginId`, and `skillName`. This lets clients preview uninstalled remote plugin skills without downloading the plugin bundle. - `skills/changed` — notification emitted when watched local skill files change. - `app/list` — list available apps. @@ -610,6 +613,16 @@ Use `thread/archive` to move the persisted rollout (stored as a JSONL file on di An archived thread will not appear in `thread/list` unless `archived` is set to `true`. +### Example: Delete a thread + +Use `thread/delete` to hard-delete a thread and its spawned descendant threads. Existing rollout files and associated metadata must be removed before the request succeeds; missing rollout files are treated as already deleted. + +```json +{ "method": "thread/delete", "id": 23, "params": { "threadId": "thr_b" } } +{ "id": 23, "result": {} } +{ "method": "thread/deleted", "params": { "threadId": "thr_b" } } +``` + ### Example: Unarchive a thread Use `thread/unarchive` to move an archived rollout back into the sessions directory. @@ -837,6 +850,8 @@ Then send `offer.sdp` to app-server. Core uses `experimental_realtime_ws_backend Omit `prompt` to use Codex's default realtime backend prompt. Send `prompt: null` or `prompt: ""` when the session should start without that default backend prompt. +Clients may also pass `model` and `version` on `thread/realtime/start` to select a +different realtime session configuration without changing thread or user config. ```javascript await pc.setRemoteDescription({ @@ -870,6 +885,32 @@ Use `thread/backgroundTerminals/clean` to terminate all running background termi { "id": 35, "result": {} } ``` +### Example: List and terminate background terminals + +Use `thread/backgroundTerminals/list` to inspect running background terminals associated with a loaded thread. The `backgroundTerminals` segment intentionally follows the existing `thread/backgroundTerminals/clean` method. The returned `processId` is the app-server process id; host OS metadata is nullable. The request accepts the standard `cursor` and `limit` pagination fields. When `nextCursor` is non-null, pass it as `cursor` to fetch the next page. + +```json +{ "method": "thread/backgroundTerminals/list", "id": 36, "params": { "threadId": "thr_123" } } +{ "id": 36, "result": { "data": [ + { + "itemId": "item_456", + "processId": "42", + "command": "python3 -m http.server", + "cwd": "/workspace", + "osPid": null, + "cpuPercent": null, + "rssKb": null + } +], "nextCursor": null } } +``` + +Use `thread/backgroundTerminals/terminate` to terminate one running background terminal by that `processId`. + +```json +{ "method": "thread/backgroundTerminals/terminate", "id": 37, "params": { "threadId": "thr_123", "processId": "42" } } +{ "id": 37, "result": { "terminated": true } } +``` + ### Example: Steer an active turn Use `turn/steer` to append additional user input to the currently active regular turn. This does diff --git a/codex-rs/app-server/src/config/external_agent_config.rs b/codex-rs/app-server/src/config/external_agent_config.rs index eca927fc23..665588213f 100644 --- a/codex-rs/app-server/src/config/external_agent_config.rs +++ b/codex-rs/app-server/src/config/external_agent_config.rs @@ -1147,7 +1147,7 @@ fn configured_marketplace_plugins( ) -> io::Result>> { let plugins_input = config.plugins_config_input(); let marketplaces = plugins_manager - .list_marketplaces_for_config(&plugins_input, &[]) + .list_marketplaces_for_config(&plugins_input, &[], /*include_openai_curated*/ true) .map_err(|err| { invalid_data_error(format!("failed to list configured marketplaces: {err}")) })?; diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index a0a44fd80d..1c397b48cd 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use std::sync::Weak; +use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ThreadGoal; use codex_app_server_protocol::ThreadGoalUpdatedNotification; @@ -20,28 +21,47 @@ use codex_protocol::error::CodexErr; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_rollout::state_db::StateDbHandle; +use codex_thread_store::ThreadStore; use crate::outgoing_message::OutgoingMessageSender; use crate::thread_state::ThreadListenerCommand; use crate::thread_state::ThreadStateManager; +pub(crate) struct ThreadExtensionDependencies { + pub(crate) event_sink: Arc, + pub(crate) auth_manager: Arc, + pub(crate) state_db: Option, + pub(crate) analytics_events_client: AnalyticsEventsClient, + pub(crate) thread_manager: Weak, + pub(crate) goal_service: Arc, + pub(crate) executor_skill_provider: Arc, + /// Process-scoped persistence backend for extensions that need stored thread history. + pub(crate) thread_store: Arc, +} + pub(crate) fn thread_extensions( guardian_agent_spawner: S, - event_sink: Arc, - auth_manager: Arc, - state_db: Option, - thread_manager: Weak, - goal_service: Arc, - executor_skill_provider: Arc, + dependencies: ThreadExtensionDependencies, ) -> Arc> where S: AgentSpawner + 'static, { + let ThreadExtensionDependencies { + event_sink, + auth_manager, + state_db, + analytics_events_client, + thread_manager, + goal_service, + executor_skill_provider, + thread_store: _thread_store, + } = dependencies; let mut builder = ExtensionRegistryBuilder::::with_event_sink(event_sink); if let Some(state_db) = state_db { codex_goal_extension::install_with_backend( &mut builder, state_db, + analytics_events_client, codex_otel::global(), thread_manager, goal_service, @@ -56,7 +76,14 @@ where codex_skills_extension::install_with_providers( &mut builder, codex_skills_extension::SkillProviders::new() - .with_executor_provider(executor_skill_provider), + .with_executor_provider(executor_skill_provider) + .with_orchestrator_provider(Arc::new( + codex_skills_extension::OrchestratorSkillProvider::new(), + )), + |config: &Config| codex_skills_extension::SkillsExtensionConfig { + include_instructions: config.include_skill_instructions, + bundled_skills_enabled: config.bundled_skills_enabled(), + }, ); Arc::new(builder.build()) } @@ -140,7 +167,6 @@ pub(crate) fn guardian_agent_spawner( mod tests { use std::time::Duration; - use codex_analytics::AnalyticsEventsClient; use codex_protocol::protocol::ThreadGoal as CoreThreadGoal; use codex_protocol::protocol::ThreadGoalStatus; use codex_protocol::protocol::ThreadGoalUpdatedEvent; diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 448a389851..27191e7215 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -104,7 +104,9 @@ type PendingClientRequestResponse = std::result::Result bool { matches!( notification, - ServerNotification::TurnCompleted(_) | ServerNotification::ThreadSettingsUpdated(_) + ServerNotification::TurnCompleted(_) + | ServerNotification::ThreadSettingsUpdated(_) + | ServerNotification::ExternalAgentConfigImportCompleted(_) ) } @@ -729,6 +731,7 @@ mod tests { use super::*; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::ConfigRequirementsReadResponse; + use codex_app_server_protocol::ExternalAgentConfigImportCompletedNotification; use codex_app_server_protocol::SessionSource as ApiSessionSource; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -893,5 +896,10 @@ mod tests { }, }) )); + assert!(server_notification_requires_delivery( + &ServerNotification::ExternalAgentConfigImportCompleted( + ExternalAgentConfigImportCompletedNotification {}, + ) + )); } } diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index b6f0b8c859..d9c75a54e1 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -1,3 +1,4 @@ +#![recursion_limit = "256"] #![deny(clippy::print_stdout, clippy::print_stderr)] use codex_arg0::Arg0DispatchPaths; @@ -74,6 +75,8 @@ use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::registry::Registry; use tracing_subscriber::util::SubscriberInitExt; +const SQLITE_RECOVERY_CONFIG_WARNING_SUMMARY: &str = "Codex rebuilt its local database."; + mod analytics_utils; mod app_server_tracing; mod attestation; @@ -533,8 +536,8 @@ pub async fn run_main_with_transport_options( } _ => None, }; - let state_db = match rollout_state_db::try_init(&config).await { - Ok(state_db) => Some(state_db), + let state_db_init = match init_sqlite_state_db_with_fresh_start_on_corruption(&config).await { + Ok(state_db_init) => state_db_init, Err(err) => { return Err(std::io::Error::other(format!( "failed to initialize sqlite state runtime under {}: {err}", @@ -542,6 +545,15 @@ pub async fn run_main_with_transport_options( ))); } }; + let state_db = state_db_init.state_db; + if let Some(recovery_notice) = state_db_init.recovery_notice { + config_warnings.push(ConfigWarningNotification { + summary: SQLITE_RECOVERY_CONFIG_WARNING_SUMMARY.to_string(), + details: Some(recovery_notice.details), + path: None, + range: None, + }); + } if should_run_personality_migration { let effective_toml = config.config_layer_stack.effective_config(); @@ -1099,6 +1111,121 @@ pub async fn run_main_with_transport_options( Ok(()) } +struct SqliteRecoveryNotice { + details: String, +} + +struct RecoveredSqliteDatabase { + database_path: String, + backup_folder: String, +} + +struct StateDbInitResult { + state_db: Option, + recovery_notice: Option, +} + +async fn init_sqlite_state_db_with_fresh_start_on_corruption( + config: &Config, +) -> anyhow::Result { + let mut attempted_backups = HashSet::new(); + let mut recovered_databases = Vec::new(); + loop { + let err = match rollout_state_db::try_init(config).await { + Ok(state_db) => { + let recovery_notice = sqlite_recovery_notice(&recovered_databases); + if recovery_notice.is_some() { + emit_state_db_backup_warning(SQLITE_RECOVERY_CONFIG_WARNING_SUMMARY); + for recovered_database in &recovered_databases { + emit_state_db_backup_warning(&format!( + "Database path: {}", + recovered_database.database_path + )); + emit_state_db_backup_warning(&format!( + "Backup folder: {}", + recovered_database.backup_folder + )); + } + } + return Ok(StateDbInitResult { + state_db: Some(state_db), + recovery_notice, + }); + } + Err(err) => err, + }; + if !codex_state::is_sqlite_corruption_error(&err) { + return Err(err); + } + + let database_path = codex_state::runtime_db_path_for_corruption_error(&err) + .unwrap_or_else(|| codex_state::state_db_path(config.sqlite_home.as_path())); + if !attempted_backups.insert(database_path.clone()) { + return Err(anyhow::anyhow!( + "failed to initialize sqlite state runtime after moving damaged database file into a backup folder: {err}" + )); + } + + let original_error = err.to_string(); + emit_state_db_backup_warning(&format!( + "Codex local database at {} appears damaged. Moving it into a backup folder so the app server can rebuild it from saved data.", + database_path.display() + )); + let backups = codex_state::backup_runtime_db_for_fresh_start(database_path.as_path()) + .await + .map_err(|backup_err| { + anyhow::anyhow!( + "failed to move damaged sqlite state database files into a backup folder: {backup_err}; original error: {original_error}" + ) + })?; + for backup in &backups { + emit_state_db_backup_warning(&format!( + "Moved damaged Codex local database file {} to {}", + backup.original_path.display(), + backup.backup_path.display() + )); + } + if let Some(first_backup) = backups.first() + && let Some(backup_folder) = first_backup.backup_path.parent() + { + recovered_databases.push(RecoveredSqliteDatabase { + database_path: first_backup.original_path.display().to_string(), + backup_folder: backup_folder.display().to_string(), + }); + } + } +} + +fn sqlite_recovery_notice( + recovered_databases: &[RecoveredSqliteDatabase], +) -> Option { + if recovered_databases.is_empty() { + return None; + } + + let details = recovered_databases + .iter() + .map(|recovered_database| { + format!( + "Database path: {}\nBackup folder: {}", + recovered_database.database_path, recovered_database.backup_folder + ) + }) + .collect::>() + .join("\n\n"); + Some(SqliteRecoveryNotice { details }) +} + +fn emit_state_db_backup_warning(message: &str) { + warn!("{message}"); + if !tracing::dispatcher::has_been_set() { + #[allow(clippy::print_stderr)] + { + eprintln!("{message}"); + } + } +} + fn analytics_rpc_transport(transport: &AppServerTransport) -> AppServerRpcTransport { match transport { AppServerTransport::Stdio => AppServerRpcTransport::Stdio, diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 2cd1027754..d11678ff5a 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -96,6 +96,7 @@ async fn queue_refresh( #[cfg(test)] mod tests { use super::*; + use crate::extensions::ThreadExtensionDependencies; use crate::extensions::guardian_agent_spawner; use crate::extensions::thread_extensions; use async_trait::async_trait; @@ -193,12 +194,16 @@ mod tests { Arc::clone(&environment_manager), thread_extensions( guardian_agent_spawner(thread_manager.clone()), - Arc::new(NoopExtensionEventSink), - auth_manager.clone(), - Some(state_db.clone()), - thread_manager.clone(), - Arc::new(codex_goal_extension::GoalService::new()), - Arc::clone(&executor_skill_provider), + ThreadExtensionDependencies { + event_sink: Arc::new(NoopExtensionEventSink), + auth_manager: auth_manager.clone(), + state_db: Some(state_db.clone()), + analytics_events_client: codex_analytics::AnalyticsEventsClient::disabled(), + thread_manager: thread_manager.clone(), + goal_service: Arc::new(codex_goal_extension::GoalService::new()), + executor_skill_provider: Arc::clone(&executor_skill_provider), + thread_store: Arc::clone(&thread_store), + }, ), /*analytics_events_client*/ None, Arc::clone(&thread_store), diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 78ba2c12cf..6ba59e733f 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -8,6 +8,7 @@ use crate::attestation::app_server_attestation_provider; use crate::config_manager::ConfigManager; use crate::connection_rpc_gate::ConnectionRpcGate; use crate::error_code::invalid_request; +use crate::extensions::ThreadExtensionDependencies; use crate::extensions::app_server_extension_event_sink; use crate::extensions::guardian_agent_spawner; use crate::extensions::thread_extensions; @@ -324,12 +325,19 @@ impl MessageProcessor { environment_manager, thread_extensions( guardian_agent_spawner(thread_manager.clone()), - app_server_extension_event_sink(outgoing.clone(), thread_state_manager.clone()), - auth_manager.clone(), - state_db.clone(), - thread_manager.clone(), - Arc::clone(&goal_service), - Arc::clone(&executor_skill_provider), + ThreadExtensionDependencies { + event_sink: app_server_extension_event_sink( + outgoing.clone(), + thread_state_manager.clone(), + ), + auth_manager: auth_manager.clone(), + state_db: state_db.clone(), + analytics_events_client: analytics_events_client.clone(), + thread_manager: thread_manager.clone(), + goal_service: Arc::clone(&goal_service), + executor_skill_provider: Arc::clone(&executor_skill_provider), + thread_store: Arc::clone(&thread_store), + }, ), Some(analytics_events_client.clone()), Arc::clone(&thread_store), @@ -393,7 +401,7 @@ impl MessageProcessor { Arc::clone(&thread_manager), Arc::clone(&config), feedback, - log_db, + log_db.clone(), state_db.clone(), ); let git_processor = GitRequestProcessor::new(); @@ -447,6 +455,7 @@ impl MessageProcessor { Arc::clone(&thread_list_state_permit), thread_goal_processor.clone(), state_db, + log_db, Arc::clone(&skills_watcher), ); let turn_processor = TurnRequestProcessor::new( @@ -1064,6 +1073,11 @@ impl MessageProcessor { .thread_archive(request_id.clone(), params) .await } + ClientRequest::ThreadDelete { params, .. } => { + self.thread_processor + .thread_delete(request_id.clone(), params) + .await + } ClientRequest::ThreadIncrementElicitation { params, .. } => { self.thread_processor .thread_increment_elicitation(params) @@ -1119,6 +1133,16 @@ impl MessageProcessor { .thread_background_terminals_clean(&request_id, params) .await } + ClientRequest::ThreadBackgroundTerminalsList { params, .. } => { + self.thread_processor + .thread_background_terminals_list(params) + .await + } + ClientRequest::ThreadBackgroundTerminalsTerminate { params, .. } => { + self.thread_processor + .thread_background_terminals_terminate(params) + .await + } ClientRequest::ThreadRollback { params, .. } => { self.thread_processor .thread_rollback(&request_id, params) diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index afb47b00ae..94727c3538 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -171,13 +171,21 @@ use codex_app_server_protocol::ThreadApproveGuardianDeniedActionResponse; use codex_app_server_protocol::ThreadArchiveParams; use codex_app_server_protocol::ThreadArchiveResponse; use codex_app_server_protocol::ThreadArchivedNotification; +use codex_app_server_protocol::ThreadBackgroundTerminal; use codex_app_server_protocol::ThreadBackgroundTerminalsCleanParams; use codex_app_server_protocol::ThreadBackgroundTerminalsCleanResponse; +use codex_app_server_protocol::ThreadBackgroundTerminalsListParams; +use codex_app_server_protocol::ThreadBackgroundTerminalsListResponse; +use codex_app_server_protocol::ThreadBackgroundTerminalsTerminateParams; +use codex_app_server_protocol::ThreadBackgroundTerminalsTerminateResponse; use codex_app_server_protocol::ThreadClosedNotification; use codex_app_server_protocol::ThreadCompactStartParams; use codex_app_server_protocol::ThreadCompactStartResponse; use codex_app_server_protocol::ThreadDecrementElicitationParams; use codex_app_server_protocol::ThreadDecrementElicitationResponse; +use codex_app_server_protocol::ThreadDeleteParams; +use codex_app_server_protocol::ThreadDeleteResponse; +use codex_app_server_protocol::ThreadDeletedNotification; use codex_app_server_protocol::ThreadForkParams; use codex_app_server_protocol::ThreadForkResponse; use codex_app_server_protocol::ThreadGoal; @@ -416,6 +424,7 @@ use codex_rollout::state_db::reconcile_rollout; use codex_state::ThreadMetadata; use codex_state::log_db::LogDbLayer; use codex_thread_store::ArchiveThreadParams as StoreArchiveThreadParams; +use codex_thread_store::DeleteThreadParams as StoreDeleteThreadParams; use codex_thread_store::GitInfoPatch as StoreGitInfoPatch; use codex_thread_store::ListThreadsParams as StoreListThreadsParams; use codex_thread_store::LocalThreadStore; @@ -536,6 +545,7 @@ fn resolve_runtime_workspace_roots(workspace_roots: Vec) -> Vec mod config_errors; mod request_errors; +mod thread_delete; mod thread_goal_processor; mod thread_lifecycle; mod thread_resume_redaction; diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index a4bbe42d61..019533484c 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -158,6 +158,9 @@ impl AccountRequestProcessor { pub(crate) fn clear_external_auth(&self) { self.auth_manager.clear_external_auth(); + self.thread_manager + .plugins_manager() + .set_auth_mode(self.auth_manager.get_api_auth_mode()); } fn current_account_updated_notification(&self) -> AccountUpdatedNotification { @@ -173,6 +176,10 @@ impl AccountRequestProcessor { thread_manager: &Arc, auth: Option, ) { + thread_manager + .plugins_manager() + .set_auth_mode(auth.as_ref().map(CodexAuth::api_auth_mode)); + match config_manager .load_latest_config(/*fallback_cwd*/ None) .await diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index 54a5efbe7a..08fc99d01b 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -573,6 +573,7 @@ impl PluginRequestProcessor { .list_marketplaces_for_config( &config_for_marketplace_listing, &roots_for_marketplace_listing, + /*include_openai_curated*/ true, )?; Ok::< ( @@ -837,8 +838,11 @@ impl PluginRequestProcessor { let config_for_marketplace_listing = plugins_input.clone(); let shared_plugin_ids_by_local_path = load_shared_plugin_ids_by_local_path(config)?; match tokio::task::spawn_blocking(move || { - let outcome = plugins_manager - .list_marketplaces_for_config(&config_for_marketplace_listing, &roots)?; + let outcome = plugins_manager.list_marketplaces_for_config( + &config_for_marketplace_listing, + &roots, + /*include_openai_curated*/ true, + )?; Ok::< ( Vec, @@ -1031,14 +1035,7 @@ impl PluginRequestProcessor { } None => None, }; - let environment_manager = self.thread_manager.environment_manager(); - let app_summaries = load_plugin_app_summaries( - &config, - &outcome.plugin.apps, - Arc::clone(&environment_manager), - self.thread_manager.mcp_manager(), - ) - .await; + let app_summaries = load_plugin_app_summaries(&config, &outcome.plugin.apps).await; let visible_skills = outcome .plugin .skills @@ -1114,14 +1111,7 @@ impl PluginRequestProcessor { .cloned() .map(codex_plugin::AppConnectorId) .collect::>(); - let environment_manager = self.thread_manager.environment_manager(); - let app_summaries = load_plugin_app_summaries( - &config, - &plugin_apps, - Arc::clone(&environment_manager), - self.thread_manager.mcp_manager(), - ) - .await; + let app_summaries = load_plugin_app_summaries(&config, &plugin_apps).await; remote_plugin_detail_to_info(remote_detail, app_summaries) } }; @@ -1578,7 +1568,6 @@ impl PluginRequestProcessor { name: connector.name, description: connector.description, install_url: connector.install_url, - needs_auth: true, }) .collect() } @@ -1883,8 +1872,6 @@ impl PluginRequestProcessor { async fn load_plugin_app_summaries( config: &Config, plugin_apps: &[codex_plugin::AppConnectorId], - environment_manager: Arc, - mcp_manager: Arc, ) -> Vec { if plugin_apps.is_empty() { return Vec::new(); @@ -1902,49 +1889,9 @@ async fn load_plugin_app_summaries( }; let plugin_connectors = connectors::connectors_for_plugin_apps(connectors, plugin_apps); - - let accessible_connectors = - match connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( - config, - /*force_refetch*/ false, - environment_manager, - mcp_manager, - ) - .await - { - Ok(status) if status.codex_apps_ready => status.connectors, - Ok(_) => { - return plugin_connectors - .into_iter() - .map(AppSummary::from) - .collect(); - } - Err(err) => { - warn!("failed to load app auth state for plugin/read: {err:#}"); - return plugin_connectors - .into_iter() - .map(AppSummary::from) - .collect(); - } - }; - - let accessible_ids = accessible_connectors - .iter() - .map(|connector| connector.id.as_str()) - .collect::>(); - plugin_connectors .into_iter() - .map(|connector| { - let needs_auth = !accessible_ids.contains(connector.id.as_str()); - AppSummary { - id: connector.id, - name: connector.name, - description: connector.description, - install_url: connector.install_url, - needs_auth, - } - }) + .map(AppSummary::from) .collect() } @@ -1979,7 +1926,6 @@ fn plugin_apps_needing_auth( name: connector.name, description: connector.description, install_url: connector.install_url, - needs_auth: true, }) .collect() } diff --git a/codex-rs/app-server/src/request_processors/thread_delete.rs b/codex-rs/app-server/src/request_processors/thread_delete.rs new file mode 100644 index 0000000000..fc7e58573c --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_delete.rs @@ -0,0 +1,188 @@ +//! `thread/delete` request handling. + +use super::thread_processor::core_thread_write_error; +use super::thread_processor::unsupported_thread_store_operation; +use super::*; + +impl ThreadRequestProcessor { + pub(crate) async fn thread_delete( + &self, + request_id: ConnectionRequestId, + params: ThreadDeleteParams, + ) -> Result, JSONRPCErrorError> { + let mut deleted_thread_ids = Vec::new(); + let result = { + let _thread_list_state_permit = self.acquire_thread_list_state_permit().await?; + self.thread_delete_response(params, &mut deleted_thread_ids) + .await + }; + match result { + Ok(response) => { + self.outgoing + .send_response(request_id.clone(), response) + .await; + self.send_thread_deleted_notifications(deleted_thread_ids) + .await; + Ok(None) + } + Err(error) => Err(error), + } + } + + async fn thread_delete_response( + &self, + params: ThreadDeleteParams, + deleted_thread_ids: &mut Vec, + ) -> Result { + let thread_id = ThreadId::from_string(¶ms.thread_id) + .map_err(|err| invalid_request(format!("invalid thread id: {err}")))?; + + let mut thread_ids = self.state_db_spawn_subtree_thread_ids(thread_id).await?; + let mut seen = thread_ids.iter().copied().collect::>(); + + match self + .thread_manager + .list_agent_subtree_thread_ids(thread_id) + .await + { + Ok(live_thread_ids) => { + for live_thread_id in live_thread_ids { + if seen.insert(live_thread_id) { + thread_ids.push(live_thread_id); + } + } + } + Err(err) => return Err(core_thread_write_error("delete thread", err)), + } + + self.validate_root_thread_delete(thread_id, thread_ids.len() > 1) + .await?; + for thread_id_to_delete in thread_ids.iter().copied() { + self.prepare_thread_for_delete(thread_id_to_delete).await; + } + + let mut delete_order: Vec<_> = thread_ids.iter().skip(1).rev().copied().collect(); + delete_order.push(thread_id); + + for thread_id_to_delete in delete_order.iter().copied() { + match self + .thread_store + .delete_thread(StoreDeleteThreadParams { + thread_id: thread_id_to_delete, + }) + .await + { + Ok(()) => {} + Err(ThreadStoreError::ThreadNotFound { .. }) => { + warn!( + "thread {thread_id_to_delete} was already missing while deleting {thread_id}" + ); + } + Err(err) => { + return Err(thread_store_delete_error(err)); + } + } + } + + if let Some(state_db) = self.state_db.as_ref() { + state_db + .delete_threads_strict(thread_ids.as_slice()) + .await + .map_err(|err| { + internal_error(format!( + "failed to delete app-server state for {thread_id}: {err}" + )) + })?; + } + + deleted_thread_ids.extend( + delete_order + .into_iter() + .map(|thread_id| thread_id.to_string()), + ); + Ok(ThreadDeleteResponse {}) + } + + async fn send_thread_deleted_notifications(&self, deleted_thread_ids: Vec) { + for thread_id in deleted_thread_ids { + self.outgoing + .send_server_notification(ServerNotification::ThreadDeleted( + ThreadDeletedNotification { thread_id }, + )) + .await; + } + } + + async fn validate_root_thread_delete( + &self, + thread_id: ThreadId, + has_descendants: bool, + ) -> Result<(), JSONRPCErrorError> { + if let Ok(thread) = self.thread_manager.get_thread(thread_id).await { + if !thread.config_snapshot().await.ephemeral { + return Ok(()); + } + return Err(invalid_request(format!( + "thread is not persisted and cannot be deleted: {thread_id}" + ))); + } + match self + .thread_store + .read_thread(StoreReadThreadParams { + thread_id, + include_archived: true, + include_history: false, + }) + .await + { + Ok(_) => Ok(()), + Err(ThreadStoreError::ThreadNotFound { .. }) => { + if has_descendants { + return Ok(()); + } + let Some(state_db) = self.state_db.as_ref() else { + return Err(thread_store_delete_error( + ThreadStoreError::ThreadNotFound { thread_id }, + )); + }; + if state_db + .get_thread(thread_id) + .await + .map_err(|err| { + internal_error(format!( + "failed to read app-server state for {thread_id}: {err}" + )) + })? + .is_some() + { + Ok(()) + } else { + Err(thread_store_delete_error( + ThreadStoreError::ThreadNotFound { thread_id }, + )) + } + } + Err(err) => Err(thread_store_delete_error(err)), + } + } + + async fn prepare_thread_for_delete(&self, thread_id: ThreadId) { + self.prepare_thread_for_removal(thread_id, "delete").await; + if let Some(log_db) = self.log_db.as_ref() { + log_db.flush().await; + } + } +} + +fn thread_store_delete_error(err: ThreadStoreError) -> JSONRPCErrorError { + match err { + ThreadStoreError::ThreadNotFound { thread_id } => { + invalid_request(format!("thread not found: {thread_id}")) + } + ThreadStoreError::InvalidRequest { message } => invalid_request(message), + ThreadStoreError::Unsupported { operation } => { + unsupported_thread_store_operation(operation) + } + err => internal_error(format!("failed to delete thread: {err}")), + } +} diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index edc7c08b88..e211732720 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -325,6 +325,7 @@ pub(crate) struct ThreadRequestProcessor { pub(super) thread_list_state_permit: Arc, pub(super) thread_goal_processor: ThreadGoalRequestProcessor, pub(super) state_db: Option, + pub(super) log_db: Option, pub(super) background_tasks: TaskTracker, pub(super) skills_watcher: Arc, } @@ -356,6 +357,7 @@ impl ThreadRequestProcessor { thread_list_state_permit: Arc, thread_goal_processor: ThreadGoalRequestProcessor, state_db: Option, + log_db: Option, skills_watcher: Arc, ) -> Self { Self { @@ -372,6 +374,7 @@ impl ThreadRequestProcessor { thread_list_state_permit, thread_goal_processor, state_db, + log_db, background_tasks: TaskTracker::new(), skills_watcher, } @@ -569,6 +572,24 @@ impl ThreadRequestProcessor { .map(|response| Some(response.into())) } + pub(crate) async fn thread_background_terminals_list( + &self, + params: ThreadBackgroundTerminalsListParams, + ) -> Result, JSONRPCErrorError> { + self.thread_background_terminals_list_inner(params) + .await + .map(|response| Some(response.into())) + } + + pub(crate) async fn thread_background_terminals_terminate( + &self, + params: ThreadBackgroundTerminalsTerminateParams, + ) -> Result, JSONRPCErrorError> { + self.thread_background_terminals_terminate_inner(params) + .await + .map(|response| Some(response.into())) + } + pub(crate) async fn thread_rollback( &self, request_id: &ConnectionRequestId, @@ -678,7 +699,7 @@ impl ThreadRequestProcessor { Ok((thread_id, thread)) } - async fn acquire_thread_list_state_permit( + pub(super) async fn acquire_thread_list_state_permit( &self, ) -> Result, JSONRPCErrorError> { self.thread_list_state_permit @@ -750,6 +771,10 @@ impl ThreadRequestProcessor { } async fn prepare_thread_for_archive(&self, thread_id: ThreadId) { + self.prepare_thread_for_removal(thread_id, "archive").await; + } + + pub(super) async fn prepare_thread_for_removal(&self, thread_id: ThreadId, operation: &str) { let removed_conversation = self.thread_manager.remove_thread(&thread_id).await; if let Some(conversation) = removed_conversation { info!("thread {thread_id} was active; shutting down"); @@ -757,11 +782,11 @@ impl ThreadRequestProcessor { ThreadShutdownResult::Complete => {} ThreadShutdownResult::SubmitFailed => { error!( - "failed to submit Shutdown to thread {thread_id}; proceeding with archive" + "failed to submit Shutdown to thread {thread_id}; proceeding with {operation}" ); } ThreadShutdownResult::TimedOut => { - warn!("thread {thread_id} shutdown timed out; proceeding with archive"); + warn!("thread {thread_id} shutdown timed out; proceeding with {operation}"); } } } @@ -1294,23 +1319,7 @@ impl ThreadRequestProcessor { let thread_id = ThreadId::from_string(¶ms.thread_id) .map_err(|err| invalid_request(format!("invalid session id: {err}")))?; - let mut thread_ids = vec![thread_id]; - if let Some(state_db_ctx) = self.state_db.as_ref() { - let descendants = state_db_ctx - .list_thread_spawn_descendants(thread_id) - .await - .map_err(|err| { - internal_error(format!( - "failed to list spawned descendants for session {thread_id}: {err}" - )) - })?; - let mut seen = HashSet::from([thread_id]); - for descendant_id in descendants { - if seen.insert(descendant_id) { - thread_ids.push(descendant_id); - } - } - } + let thread_ids = self.state_db_spawn_subtree_thread_ids(thread_id).await?; let mut archive_thread_ids = Vec::new(); match self @@ -1395,6 +1404,31 @@ impl ThreadRequestProcessor { Ok((ThreadArchiveResponse {}, archived_thread_ids)) } + pub(super) async fn state_db_spawn_subtree_thread_ids( + &self, + thread_id: ThreadId, + ) -> Result, JSONRPCErrorError> { + let mut thread_ids = vec![thread_id]; + let Some(state_db_ctx) = self.state_db.as_ref() else { + return Ok(thread_ids); + }; + let mut seen = HashSet::from([thread_id]); + let descendants = state_db_ctx + .list_thread_spawn_descendants(thread_id) + .await + .map_err(|err| { + internal_error(format!( + "failed to list spawned descendants for thread id {thread_id}: {err}" + )) + })?; + for descendant_id in descendants { + if seen.insert(descendant_id) { + thread_ids.push(descendant_id); + } + } + Ok(thread_ids) + } + async fn thread_increment_elicitation_inner( &self, params: ThreadIncrementElicitationParams, @@ -1725,6 +1759,54 @@ impl ThreadRequestProcessor { Ok(ThreadBackgroundTerminalsCleanResponse {}) } + async fn thread_background_terminals_list_inner( + &self, + params: ThreadBackgroundTerminalsListParams, + ) -> Result { + let ThreadBackgroundTerminalsListParams { + thread_id, + cursor, + limit, + } = params; + + let (_, thread) = self.load_thread(&thread_id).await?; + let terminals = thread + .list_background_terminals() + .await + .into_iter() + .map(|terminal| ThreadBackgroundTerminal { + item_id: terminal.item_id, + process_id: terminal.process_id, + command: terminal.command, + cwd: terminal.cwd, + os_pid: None, + cpu_percent: None, + rss_kb: None, + }) + .collect::>(); + + let (data, next_cursor) = paginate_background_terminals(&terminals, cursor, limit)?; + + Ok(ThreadBackgroundTerminalsListResponse { data, next_cursor }) + } + + async fn thread_background_terminals_terminate_inner( + &self, + params: ThreadBackgroundTerminalsTerminateParams, + ) -> Result { + let ThreadBackgroundTerminalsTerminateParams { + thread_id, + process_id, + } = params; + let process_id = process_id.parse::().map_err(|err| { + invalid_request(format!("invalid background terminal process id: {err}")) + })?; + + let (_, thread) = self.load_thread(&thread_id).await?; + let terminated = thread.terminate_background_terminal(process_id).await; + Ok(ThreadBackgroundTerminalsTerminateResponse { terminated }) + } + async fn thread_shell_command_inner( &self, request_id: &ConnectionRequestId, @@ -3851,7 +3933,7 @@ fn thread_read_view_error(err: ThreadReadViewError) -> JSONRPCErrorError { } } -fn unsupported_thread_store_operation(operation: &'static str) -> JSONRPCErrorError { +pub(super) fn unsupported_thread_store_operation(operation: &'static str) -> JSONRPCErrorError { method_not_found(format!("{operation} is not supported yet")) } @@ -3972,7 +4054,7 @@ fn conversation_summary_rollout_path_read_error( } } -fn core_thread_write_error(operation: &str, err: CodexErr) -> JSONRPCErrorError { +pub(super) fn core_thread_write_error(operation: &str, err: CodexErr) -> JSONRPCErrorError { match err { CodexErr::ThreadNotFound(thread_id) => { invalid_request(format!("thread not found: {thread_id}")) @@ -4258,6 +4340,34 @@ fn build_thread_from_snapshot( } } +fn paginate_background_terminals( + terminals: &[ThreadBackgroundTerminal], + cursor: Option, + limit: Option, +) -> Result<(Vec, Option), JSONRPCErrorError> { + let start = match cursor { + Some(cursor) => { + let cursor = cursor + .parse::() + .map_err(|err| invalid_request(format!("invalid cursor: {err}")))?; + terminals + .iter() + .position(|terminal| { + terminal + .process_id + .parse::() + .is_ok_and(|process_id| process_id > cursor) + }) + .unwrap_or(terminals.len()) + } + None => 0, + }; + let effective_limit = limit.unwrap_or(terminals.len() as u32).max(1) as usize; + let end = start.saturating_add(effective_limit).min(terminals.len()); + let next_cursor = (end < terminals.len()).then(|| terminals[end - 1].process_id.clone()); + Ok((terminals[start..end].to_vec(), next_cursor)) +} + fn build_thread_from_loaded_snapshot( thread_id: ThreadId, config_snapshot: &ThreadConfigSnapshot, diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 54c0f0343f..7952517c92 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -36,6 +36,65 @@ mod thread_list_cwd_filter_tests { } } +mod background_terminal_pagination_tests { + use super::super::paginate_background_terminals; + use codex_app_server_protocol::ThreadBackgroundTerminal; + use codex_utils_absolute_path::AbsolutePathBuf; + use pretty_assertions::assert_eq; + + fn terminal(process_id: &str) -> ThreadBackgroundTerminal { + let cwd = if cfg!(windows) { r"C:\tmp" } else { "/tmp" }; + + ThreadBackgroundTerminal { + item_id: format!("item-{process_id}"), + process_id: process_id.to_string(), + command: format!("command-{process_id}"), + cwd: AbsolutePathBuf::from_absolute_path(cwd).expect("absolute cwd"), + os_pid: None, + cpu_percent: None, + rss_kb: None, + } + } + + #[test] + fn paginates_with_process_id_cursor() { + let terminals = vec![ + terminal("1"), + terminal("2"), + terminal("3"), + terminal("4"), + terminal("5"), + ]; + + let (data, next_cursor) = + paginate_background_terminals(&terminals, /*cursor*/ None, Some(2)) + .expect("valid page"); + + assert_eq!(data, vec![terminal("1"), terminal("2")]); + assert_eq!(next_cursor, Some("2".to_string())); + let first_cursor = next_cursor; + + let terminals_without_anchor = vec![terminal("1"), terminal("3"), terminal("4")]; + let (data, next_cursor) = + paginate_background_terminals(&terminals_without_anchor, first_cursor.clone(), Some(2)) + .expect("valid page"); + + assert_eq!(data, vec![terminal("3"), terminal("4")]); + assert_eq!(next_cursor, None); + + let (data, next_cursor) = + paginate_background_terminals(&terminals, first_cursor, Some(2)).expect("valid page"); + + assert_eq!(data, vec![terminal("3"), terminal("4")]); + assert_eq!(next_cursor, Some("4".to_string())); + + assert!( + paginate_background_terminals(&terminals, Some("missing".to_string()), Some(1)) + .is_err() + ); + } +} + mod thread_processor_behavior_tests { async fn forked_from_id_from_rollout(path: &Path) -> Option { codex_core::read_session_meta_line(path) diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index 355fdcd270..0f3f0f54fd 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -934,6 +934,7 @@ impl TurnRequestProcessor { request_id, thread.as_ref(), Op::RealtimeConversationStart(ConversationStartParams { + model: params.model, output_modality: params.output_modality, prompt: params.prompt, realtime_session_id: params.realtime_session_id, @@ -945,6 +946,7 @@ impl TurnRequestProcessor { ConversationStartTransport::Webrtc { sdp } } }), + version: params.version, voice: params.voice, }), ) diff --git a/codex-rs/app-server/tests/common/auth_fixtures.rs b/codex-rs/app-server/tests/common/auth_fixtures.rs index cf78e788de..18cfd7ed68 100644 --- a/codex-rs/app-server/tests/common/auth_fixtures.rs +++ b/codex-rs/app-server/tests/common/auth_fixtures.rs @@ -165,6 +165,7 @@ pub fn write_chatgpt_auth( last_refresh, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth(codex_home, &auth, cli_auth_credentials_store_mode).context("write auth.json") diff --git a/codex-rs/app-server/tests/common/models_cache.rs b/codex-rs/app-server/tests/common/models_cache.rs index 8233b1b296..2a0c10a27c 100644 --- a/codex-rs/app-server/tests/common/models_cache.rs +++ b/codex-rs/app-server/tests/common/models_cache.rs @@ -47,6 +47,7 @@ fn preset_to_info(preset: &ModelPreset, priority: i32) -> ModelInfo { context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), input_modalities: default_input_modalities(), diff --git a/codex-rs/app-server/tests/common/test_app_server.rs b/codex-rs/app-server/tests/common/test_app_server.rs index 646de6ad21..91c42e2fe4 100644 --- a/codex-rs/app-server/tests/common/test_app_server.rs +++ b/codex-rs/app-server/tests/common/test_app_server.rs @@ -79,6 +79,7 @@ use codex_app_server_protocol::SkillsExtraRootsSetParams; use codex_app_server_protocol::SkillsListParams; use codex_app_server_protocol::ThreadArchiveParams; use codex_app_server_protocol::ThreadCompactStartParams; +use codex_app_server_protocol::ThreadDeleteParams; use codex_app_server_protocol::ThreadForkParams; use codex_app_server_protocol::ThreadInjectItemsParams; use codex_app_server_protocol::ThreadListParams; @@ -456,6 +457,15 @@ impl TestAppServer { self.send_request("thread/archive", params).await } + /// Send a `thread/delete` JSON-RPC request. + pub async fn send_thread_delete_request( + &mut self, + params: ThreadDeleteParams, + ) -> anyhow::Result { + let params = Some(serde_json::to_value(params)?); + self.send_request("thread/delete", params).await + } + /// Send a `thread/name/set` JSON-RPC request. pub async fn send_thread_set_name_request( &mut self, diff --git a/codex-rs/app-server/tests/suite/v2/analytics.rs b/codex-rs/app-server/tests/suite/v2/analytics.rs index 8e8f2a7557..ea876940ab 100644 --- a/codex-rs/app-server/tests/suite/v2/analytics.rs +++ b/codex-rs/app-server/tests/suite/v2/analytics.rs @@ -124,6 +124,31 @@ pub(crate) async fn wait_for_analytics_event( server: &MockServer, read_timeout: Duration, event_type: &str, +) -> Result { + wait_for_matching_analytics_event(server, read_timeout, |event| { + event["event_type"] == event_type + }) + .await +} + +pub(crate) async fn wait_for_goal_event( + server: &MockServer, + read_timeout: Duration, + event_kind: &str, + goal_status: &str, +) -> Result { + wait_for_matching_analytics_event(server, read_timeout, |event| { + event["event_type"] == "codex_goal_event" + && event["event_params"]["event_kind"] == event_kind + && event["event_params"]["goal_status"] == goal_status + }) + .await +} + +async fn wait_for_matching_analytics_event( + server: &MockServer, + read_timeout: Duration, + matches: impl Fn(&Value) -> bool, ) -> Result { timeout(read_timeout, async { loop { @@ -142,10 +167,7 @@ pub(crate) async fn wait_for_analytics_event( let Some(events) = payload["events"].as_array() else { continue; }; - if let Some(event) = events - .iter() - .find(|event| event["event_type"] == event_type) - { + if let Some(event) = events.iter().find(|event| matches(event)) { return Ok::(event.clone()); } } diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index c9615d3a27..fa4c2042af 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -119,6 +119,7 @@ async fn list_apps_returns_empty_with_api_key_auth() -> Result<()> { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }, AuthCredentialsStoreMode::File, )?; @@ -1516,7 +1517,7 @@ async fn start_apps_server_with_delays_and_control_inner( get(workspace_settings_response), ) .with_state(state) - .nest_service("/api/codex/apps", mcp_service); + .nest_service("/api/codex/ps/mcp", mcp_service); let handle = tokio::spawn(async move { let _ = axum::serve(listener, router).await; diff --git a/codex-rs/app-server/tests/suite/v2/compaction.rs b/codex-rs/app-server/tests/suite/v2/compaction.rs index 690c6d8bff..08ddfa4c61 100644 --- a/codex-rs/app-server/tests/suite/v2/compaction.rs +++ b/codex-rs/app-server/tests/suite/v2/compaction.rs @@ -29,6 +29,7 @@ use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; use codex_config::types::AuthCredentialsStoreMode; +use codex_features::Feature; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use core_test_support::responses; @@ -150,7 +151,7 @@ async fn auto_compaction_remote_emits_started_and_completed_items() -> Result<() write_mock_responses_config_toml( codex_home.path(), &server.uri(), - &BTreeMap::default(), + &BTreeMap::from([(Feature::RemoteCompactionV2, false)]), REMOTE_AUTO_COMPACT_LIMIT, Some(true), "mock_provider", diff --git a/codex-rs/app-server/tests/suite/v2/experimental_api.rs b/codex-rs/app-server/tests/suite/v2/experimental_api.rs index fc1100d17b..03c9d455f0 100644 --- a/codex-rs/app-server/tests/suite/v2/experimental_api.rs +++ b/codex-rs/app-server/tests/suite/v2/experimental_api.rs @@ -80,10 +80,12 @@ async fn realtime_conversation_start_requires_experimental_api_capability() -> R let request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { thread_id: "thr_123".to_string(), + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("hello".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, }) .await?; @@ -186,12 +188,14 @@ async fn realtime_webrtc_start_requires_experimental_api_capability() -> Result< let request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { thread_id: "thr_123".to_string(), + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("hello".to_string())), realtime_session_id: None, transport: Some(ThreadRealtimeStartTransport::Webrtc { sdp: "v=offer\r\n".to_string(), }), + version: None, voice: None, }) .await?; diff --git a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs index d8ecb88996..59d8bece01 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs @@ -19,6 +19,8 @@ use codex_app_server_protocol::McpResourceReadResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::UserInput; use codex_arg0::Arg0DispatchPaths; use codex_config::CloudConfigBundleLoader; use codex_config::LoaderOverrides; @@ -30,9 +32,14 @@ use codex_protocol::protocol::SessionSource; use core_test_support::responses; use pretty_assertions::assert_eq; use rmcp::handler::server::ServerHandler; +use rmcp::model::ListResourcesResult; +use rmcp::model::Meta; +use rmcp::model::PaginatedRequestParams; use rmcp::model::ProtocolVersion; +use rmcp::model::RawResource; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; +use rmcp::model::Resource; use rmcp::model::ResourceContents; use rmcp::model::ServerCapabilities; use rmcp::model::ServerInfo; @@ -41,6 +48,7 @@ use rmcp::service::RoleServer; use rmcp::transport::StreamableHttpServerConfig; use rmcp::transport::StreamableHttpService; use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; +use serde_json::json; use tempfile::TempDir; use tokio::net::TcpListener; use tokio::task::JoinHandle; @@ -51,49 +59,34 @@ const TEST_RESOURCE_URI: &str = "test://codex/resource"; const TEST_BLOB_RESOURCE_URI: &str = "test://codex/resource.bin"; const TEST_RESOURCE_BLOB: &str = "YmluYXJ5LXJlc291cmNl"; const TEST_RESOURCE_TEXT: &str = "Resource body from the MCP server."; +const SKILL_NAME: &str = "demo-plugin:deploy"; +const RAW_SKILL_DESCRIPTION: &str = "Deploy\nthrough the orchestrator."; +const SKILL_DESCRIPTION: &str = "Deploy through the <hosted> orchestrator."; +const SKILL_RESOURCE_URI: &str = "skill://plugin_demo/deploy"; +const SKILL_MAIN_PROMPT_URI: &str = "skill://plugin_demo/deploy/SKILL.md"; +const SKILL_REFERENCE_URI: &str = "skill://plugin_demo/deploy/references/deploy.md"; +const SKILL_MARKER: &str = "ORCHESTRATOR_SKILL_BODY_MARKER"; +const SKILL_CONTENTS: &str = concat!( + "---\n", + "name: deploy\n", + "description: Deploy through the orchestrator.\n", + "---\n\n", + "# Deploy\n\n", + "ORCHESTRATOR_SKILL_BODY_MARKER\n\n", + "Read the [deployment reference](skill://plugin_demo/deploy/references/deploy.md).\n", +); +const SKILL_REFERENCE_CONTENTS: &str = + "# Deploy reference\n\nUse the orchestrator deployment API.\n"; +const SKILLS_LIST_CALL_ID: &str = "skills-list"; +const SKILLS_READ_CALL_ID: &str = "skills-read"; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn mcp_resource_read_returns_resource_contents() -> Result<()> { let responses_server = responses::start_mock_server().await; let (apps_server_url, apps_server_handle) = start_resource_apps_mcp_server().await?; - - let codex_home = TempDir::new()?; let responses_server_uri = responses_server.uri(); - std::fs::write( - codex_home.path().join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "untrusted" -sandbox_mode = "read-only" - -model_provider = "mock_provider" -chatgpt_base_url = "{apps_server_url}" -mcp_oauth_credentials_store = "file" - -[features] -apps = true - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{responses_server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - )?; - write_chatgpt_auth( - codex_home.path(), - ChatGptAuthFixture::new("chatgpt-token") - .account_id("account-123") - .chatgpt_user_id("user-123") - .chatgpt_account_id("account-123"), - AuthCredentialsStoreMode::File, - )?; - - let mut mcp = TestAppServer::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let (_codex_home, mut mcp) = + start_resource_test_app_server(&apps_server_url, &responses_server_uri).await?; let thread_start_id = mcp .send_thread_start_request(ThreadStartParams { @@ -120,7 +113,6 @@ stream_max_retries = 0 mcp.read_stream_until_response_message(RequestId::Integer(read_request_id)), ) .await??; - assert_eq!( to_response::(read_response)?, expected_resource_read_response() @@ -131,6 +123,164 @@ stream_max_retries = 0 Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn orchestrator_skill_can_read_referenced_resource_without_an_executor() -> Result<()> { + let responses_server = responses::start_mock_server().await; + let (apps_server_url, apps_server_handle) = start_resource_apps_mcp_server().await?; + let responses_server_uri = responses_server.uri(); + let (_codex_home, mut mcp) = + start_resource_test_app_server(&apps_server_url, &responses_server_uri).await?; + + let thread_start_id = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + environments: Some(Vec::new()), + ..Default::default() + }) + .await?; + let thread_start_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; + + let response_mock = responses::mount_sse_sequence( + &responses_server, + vec![ + responses::sse(vec![ + responses::ev_response_created("resp-skills-list"), + responses::ev_function_call_with_namespace( + SKILLS_LIST_CALL_ID, + "skills", + "list", + &json!({ + "authority": { + "kind": "orchestrator", + }, + }) + .to_string(), + ), + responses::ev_completed("resp-skills-list"), + ]), + responses::sse(vec![ + responses::ev_response_created("resp-skills-read"), + responses::ev_function_call_with_namespace( + SKILLS_READ_CALL_ID, + "skills", + "read", + &json!({ + "authority": { + "kind": "orchestrator", + }, + "package": SKILL_RESOURCE_URI, + "resource": SKILL_REFERENCE_URI, + }) + .to_string(), + ), + responses::ev_completed("resp-skills-read"), + ]), + responses::sse(vec![ + responses::ev_response_created("resp-orchestrator-skill"), + responses::ev_assistant_message("msg-orchestrator-skill", "Done"), + responses::ev_completed("resp-orchestrator-skill"), + ]), + ], + ) + .await; + let turn_start_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id, + input: vec![UserInput::Text { + text: format!("Use ${SKILL_NAME}"), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), + ) + .await??; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 3); + let first_request = &requests[0]; + assert!(first_request.tool_by_name("skills", "list").is_some()); + assert!(first_request.tool_by_name("skills", "read").is_some()); + assert!(first_request.tool_by_name("skills", "search").is_none()); + + let developer_messages = first_request.message_input_texts("developer"); + let catalog_line = format!( + "- {SKILL_NAME}: {SKILL_DESCRIPTION} (orchestrator resource: {SKILL_RESOURCE_URI})" + ); + assert_eq!( + 1, + developer_messages + .iter() + .filter(|text| text.contains(&catalog_line)) + .count() + ); + assert!( + developer_messages + .iter() + .all(|text| !text.contains("ignored-plugin:ignored")) + ); + assert!( + developer_messages + .iter() + .any(|text| text.contains("do not treat `skill://` identifiers as filesystem paths")) + ); + let skill_fragments = first_request + .message_input_texts("user") + .into_iter() + .filter(|text| text.starts_with("")) + .collect::>(); + assert_eq!(1, skill_fragments.len()); + assert!(skill_fragments[0].contains(&format!("{SKILL_NAME}"))); + assert!(skill_fragments[0].contains(SKILL_MARKER)); + assert!(skill_fragments[0].contains(SKILL_REFERENCE_URI)); + + let list_output = requests[1] + .function_call_output_text(SKILLS_LIST_CALL_ID) + .ok_or_else(|| anyhow::anyhow!("skills.list output should be sent to the model"))?; + assert_eq!( + serde_json::from_str::(&list_output)?, + json!({ + "skills": [{ + "authority": { + "kind": "orchestrator", + }, + "package": SKILL_RESOURCE_URI, + "name": SKILL_NAME, + "description": SKILL_DESCRIPTION, + "main_resource": SKILL_MAIN_PROMPT_URI, + }], + "warnings": ["Orchestrator skill discovery stopped after 2 resource pages: failed to list orchestrator skill resources: resources/list failed for `codex_apps`: Mcp error: -32603: simulated later-page failure"], + }) + ); + + let read_output = requests[2] + .function_call_output_text(SKILLS_READ_CALL_ID) + .ok_or_else(|| anyhow::anyhow!("skills.read output should be sent to the model"))?; + assert_eq!( + serde_json::from_str::(&read_output)?, + json!({ + "resource": SKILL_REFERENCE_URI, + "contents": SKILL_REFERENCE_CONTENTS, + }) + ); + apps_server_handle.abort(); + let _ = apps_server_handle.await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn mcp_resource_read_returns_resource_contents_without_thread() -> Result<()> { let (apps_server_url, apps_server_handle) = start_resource_apps_mcp_server().await?; @@ -246,6 +396,52 @@ async fn mcp_resource_read_returns_error_for_unknown_thread() -> Result<()> { Ok(()) } +async fn start_resource_test_app_server( + apps_server_url: &str, + responses_server_uri: &str, +) -> Result<(TempDir, TestAppServer)> { + let codex_home = TempDir::new()?; + std::fs::write( + codex_home.path().join("config.toml"), + format!( + r#" +model = "mock-model" +approval_policy = "untrusted" +sandbox_mode = "read-only" + +model_provider = "mock_provider" +chatgpt_base_url = "{apps_server_url}" +mcp_oauth_credentials_store = "file" + +[features] +apps = true + +[skills] +include_instructions = true + +[model_providers.mock_provider] +name = "Mock provider for test" +base_url = "{responses_server_uri}/v1" +wire_api = "responses" +request_max_retries = 0 +stream_max_retries = 0 +"# + ), + )?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("chatgpt-token") + .account_id("account-123") + .chatgpt_user_id("user-123") + .chatgpt_account_id("account-123"), + AuthCredentialsStoreMode::File, + )?; + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + Ok((codex_home, mcp)) +} + async fn start_resource_apps_mcp_server() -> Result<(String, JoinHandle<()>)> { let listener = TcpListener::bind("127.0.0.1:0").await?; let addr = listener.local_addr()?; @@ -256,7 +452,7 @@ async fn start_resource_apps_mcp_server() -> Result<(String, JoinHandle<()>)> { Arc::new(LocalSessionManager::default()), StreamableHttpServerConfig::default(), ); - let router = Router::new().nest_service("/api/codex/apps", mcp_service); + let router = Router::new().nest_service("/api/codex/ps/mcp", mcp_service); let apps_server_handle = tokio::spawn(async move { let _ = axum::serve(listener, router).await; }); @@ -292,12 +488,79 @@ impl ServerHandler for ResourceAppsMcpServer { .with_protocol_version(ProtocolVersion::V_2025_06_18) } + async fn list_resources( + &self, + request: Option, + _context: RequestContext, + ) -> Result { + let cursor = request.and_then(|request| request.cursor); + if cursor.is_none() { + return Ok(ListResourcesResult { + resources: vec![skill_resource( + "skill://plugin_ignored/ignored", + "plugin_ignored/ignored", + "Not an MCP skill resource.", + "text/plain", + "ignored-plugin", + "ignored", + )], + next_cursor: Some("skills-page".to_string()), + meta: None, + }); + } + if cursor.as_deref() == Some("failing-page") { + return Err(rmcp::ErrorData::internal_error( + "simulated later-page failure", + /*data*/ None, + )); + } + if cursor.as_deref() != Some("skills-page") { + return Err(rmcp::ErrorData::invalid_params( + "unexpected resources/list cursor", + /*data*/ None, + )); + } + + Ok(ListResourcesResult { + resources: vec![skill_resource( + SKILL_RESOURCE_URI, + "plugin_demo/deploy", + RAW_SKILL_DESCRIPTION, + "mcp/skill", + "demo-plugin", + "deploy", + )], + next_cursor: Some("failing-page".to_string()), + meta: None, + }) + } + async fn read_resource( &self, request: ReadResourceRequestParams, _context: RequestContext, ) -> Result { let uri = request.uri; + if uri == SKILL_MAIN_PROMPT_URI { + return Ok(ReadResourceResult::new(vec![ + ResourceContents::TextResourceContents { + uri: SKILL_MAIN_PROMPT_URI.to_string(), + mime_type: Some("text/markdown".to_string()), + text: SKILL_CONTENTS.to_string(), + meta: None, + }, + ])); + } + if uri == SKILL_REFERENCE_URI { + return Ok(ReadResourceResult::new(vec![ + ResourceContents::TextResourceContents { + uri: SKILL_REFERENCE_URI.to_string(), + mime_type: Some("text/markdown".to_string()), + text: SKILL_REFERENCE_CONTENTS.to_string(), + meta: None, + }, + ])); + } if uri != TEST_RESOURCE_URI { return Err(rmcp::ErrorData::resource_not_found( format!("resource not found: {uri}"), @@ -321,3 +584,27 @@ impl ServerHandler for ResourceAppsMcpServer { ])) } } + +fn skill_resource( + uri: &str, + name: &str, + description: &str, + mime_type: &str, + plugin_name: &str, + skill_name: &str, +) -> Resource { + Resource::new( + RawResource::new(uri, name) + .with_description(description) + .with_mime_type(mime_type) + .with_meta(skill_resource_meta(plugin_name, skill_name)), + /*annotations*/ None, + ) +} + +fn skill_resource_meta(plugin_name: &str, skill_name: &str) -> Meta { + Meta(serde_json::Map::from_iter([ + ("plugin_name".to_string(), json!(plugin_name)), + ("skill_name".to_string(), json!(skill_name)), + ])) +} diff --git a/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs b/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs index 8d0387ee97..b4bd1d377a 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_server_elicitation.rs @@ -406,7 +406,7 @@ async fn start_apps_server() -> Result<(String, JoinHandle<()>)> { get(list_directory_connectors), ) .with_state(state) - .nest_service("/api/codex/apps", mcp_service); + .nest_service("/api/codex/ps/mcp", mcp_service); let handle = tokio::spawn(async move { let _ = axum::serve(listener, router).await; diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 77d4b37589..e038a131e4 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -50,6 +50,7 @@ mod review; mod safety_check_downgrade; mod skills_list; mod thread_archive; +mod thread_delete; mod thread_fork; mod thread_inject_items; mod thread_list; diff --git a/codex-rs/app-server/tests/suite/v2/plugin_install.rs b/codex-rs/app-server/tests/suite/v2/plugin_install.rs index 5ebe4899d0..1a39c1f8bd 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_install.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_install.rs @@ -327,7 +327,6 @@ async fn plugin_install_uses_remote_apps_needing_auth_response() -> Result<()> { name: "alpha".to_string(), description: None, install_url: Some("https://chatgpt.com/apps/alpha/alpha".to_string()), - needs_auth: true, }], } ); @@ -1001,7 +1000,6 @@ async fn plugin_install_returns_apps_needing_auth() -> Result<()> { name: "Alpha".to_string(), description: Some("Alpha connector".to_string()), install_url: Some("https://chatgpt.com/apps/alpha/alpha".to_string()), - needs_auth: true, }], } ); @@ -1089,7 +1087,6 @@ async fn plugin_install_filters_disallowed_apps_needing_auth() -> Result<()> { name: "Alpha".to_string(), description: Some("Alpha connector".to_string()), install_url: Some("https://chatgpt.com/apps/alpha/alpha".to_string()), - needs_auth: true, }], } ); @@ -1289,7 +1286,7 @@ async fn start_apps_server( get(list_directory_connectors), ) .with_state(state) - .nest_service("/api/codex/apps", mcp_service); + .nest_service("/api/codex/ps/mcp", mcp_service); let handle = tokio::spawn(async move { let _ = axum::serve(listener, router).await; diff --git a/codex-rs/app-server/tests/suite/v2/plugin_read.rs b/codex-rs/app-server/tests/suite/v2/plugin_read.rs index e675a6172a..7f90bc1257 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_read.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_read.rs @@ -1,6 +1,3 @@ -use std::borrow::Cow; -use std::sync::Arc; -use std::sync::Mutex as StdMutex; use std::time::Duration; use anyhow::Result; @@ -8,15 +5,6 @@ use app_test_support::ChatGptAuthFixture; use app_test_support::TestAppServer; use app_test_support::to_response; use app_test_support::write_chatgpt_auth; -use axum::Json; -use axum::Router; -use axum::extract::State; -use axum::http::HeaderMap; -use axum::http::StatusCode; -use axum::http::Uri; -use axum::http::header::AUTHORIZATION; -use axum::routing::get; -use codex_app_server_protocol::AppInfo; use codex_app_server_protocol::AppTemplateSummary; use codex_app_server_protocol::AppTemplateUnavailableReason; use codex_app_server_protocol::HookEventName; @@ -37,21 +25,8 @@ use codex_app_server_protocol::RequestId; use codex_config::types::AuthCredentialsStoreMode; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; -use rmcp::handler::server::ServerHandler; -use rmcp::model::JsonObject; -use rmcp::model::ListToolsResult; -use rmcp::model::Meta; -use rmcp::model::ServerCapabilities; -use rmcp::model::ServerInfo; -use rmcp::model::Tool; -use rmcp::model::ToolAnnotations; -use rmcp::transport::StreamableHttpServerConfig; -use rmcp::transport::StreamableHttpService; -use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; use serde_json::json; use tempfile::TempDir; -use tokio::net::TcpListener; -use tokio::task::JoinHandle; use tokio::time::timeout; use wiremock::Mock; use wiremock::MockServer; @@ -1483,104 +1458,11 @@ enabled = false response.plugin.apps[0].install_url.as_deref(), Some("https://chatgpt.com/apps/gmail/gmail") ); - assert_eq!(response.plugin.apps[0].needs_auth, true); assert_eq!(response.plugin.mcp_servers.len(), 1); assert_eq!(response.plugin.mcp_servers[0], "demo"); Ok(()) } -#[tokio::test] -async fn plugin_read_returns_app_needs_auth() -> Result<()> { - let connectors = vec![ - AppInfo { - id: "alpha".to_string(), - name: "Alpha".to_string(), - description: Some("Alpha connector".to_string()), - logo_url: Some("https://example.com/alpha.png".to_string()), - logo_url_dark: None, - distribution_channel: Some("featured".to_string()), - branding: None, - app_metadata: None, - labels: None, - install_url: None, - is_accessible: false, - is_enabled: true, - plugin_display_names: Vec::new(), - }, - AppInfo { - id: "beta".to_string(), - name: "Beta".to_string(), - description: Some("Beta connector".to_string()), - logo_url: None, - logo_url_dark: None, - distribution_channel: None, - branding: None, - app_metadata: None, - labels: None, - install_url: None, - is_accessible: false, - is_enabled: true, - plugin_display_names: Vec::new(), - }, - ]; - let tools = vec![connector_tool("beta", "Beta App")?]; - let (server_url, server_handle) = start_apps_server(connectors, tools).await?; - - let codex_home = TempDir::new()?; - write_connectors_config(codex_home.path(), &server_url)?; - write_chatgpt_auth( - codex_home.path(), - ChatGptAuthFixture::new("chatgpt-token") - .account_id("account-123") - .chatgpt_user_id("user-123") - .chatgpt_account_id("account-123"), - AuthCredentialsStoreMode::File, - )?; - - let repo_root = TempDir::new()?; - write_plugin_marketplace( - repo_root.path(), - "debug", - "sample-plugin", - "./sample-plugin", - )?; - write_plugin_source(repo_root.path(), "sample-plugin", &["alpha", "beta"])?; - let marketplace_path = - AbsolutePathBuf::try_from(repo_root.path().join(".agents/plugins/marketplace.json"))?; - - let mut mcp = TestAppServer::new(codex_home.path()).await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - - let request_id = mcp - .send_plugin_read_request(PluginReadParams { - marketplace_path: Some(marketplace_path), - remote_marketplace_name: None, - plugin_name: "sample-plugin".to_string(), - }) - .await?; - - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; - - assert_eq!( - response - .plugin - .apps - .iter() - .map(|app| (app.id.as_str(), app.needs_auth)) - .collect::>(), - vec![("alpha", true), ("beta", false)] - ); - - server_handle.abort(); - let _ = server_handle.await; - Ok(()) -} - #[tokio::test] async fn plugin_read_accepts_legacy_string_default_prompt() -> Result<()> { let codex_home = TempDir::new()?; @@ -1846,151 +1728,6 @@ plugins = true Ok(()) } -#[derive(Clone)] -struct AppsServerState { - response: Arc>, -} - -#[derive(Clone)] -struct PluginReadMcpServer { - tools: Arc>>, -} - -impl ServerHandler for PluginReadMcpServer { - fn get_info(&self) -> ServerInfo { - ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) - } - - fn list_tools( - &self, - _request: Option, - _context: rmcp::service::RequestContext, - ) -> impl std::future::Future> + Send + '_ - { - let tools = self.tools.clone(); - async move { - let tools = tools - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone(); - Ok(ListToolsResult { - tools, - next_cursor: None, - meta: None, - }) - } - } -} - -async fn start_apps_server( - connectors: Vec, - tools: Vec, -) -> Result<(String, JoinHandle<()>)> { - let state = Arc::new(AppsServerState { - response: Arc::new(StdMutex::new( - json!({ "apps": connectors, "next_token": null }), - )), - }); - let tools = Arc::new(StdMutex::new(tools)); - - let listener = TcpListener::bind("127.0.0.1:0").await?; - let addr = listener.local_addr()?; - let mcp_service = StreamableHttpService::new( - { - let tools = tools.clone(); - move || { - Ok(PluginReadMcpServer { - tools: tools.clone(), - }) - } - }, - Arc::new(LocalSessionManager::default()), - StreamableHttpServerConfig::default(), - ); - let router = Router::new() - .route("/connectors/directory/list", get(list_directory_connectors)) - .route( - "/connectors/directory/list_workspace", - get(list_directory_connectors), - ) - .with_state(state) - .nest_service("/api/codex/apps", mcp_service); - - let handle = tokio::spawn(async move { - let _ = axum::serve(listener, router).await; - }); - - Ok((format!("http://{addr}"), handle)) -} - -async fn list_directory_connectors( - State(state): State>, - headers: HeaderMap, - uri: Uri, -) -> Result { - let bearer_ok = headers - .get(AUTHORIZATION) - .and_then(|value| value.to_str().ok()) - .is_some_and(|value| value == "Bearer chatgpt-token"); - let account_ok = headers - .get("chatgpt-account-id") - .and_then(|value| value.to_str().ok()) - .is_some_and(|value| value == "account-123"); - let external_logos_ok = uri - .query() - .is_some_and(|query| query.split('&').any(|pair| pair == "external_logos=true")); - - if !bearer_ok || !account_ok { - Err(StatusCode::UNAUTHORIZED) - } else if !external_logos_ok { - Err(StatusCode::BAD_REQUEST) - } else { - let response = state - .response - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone(); - Ok(Json(response)) - } -} - -fn connector_tool(connector_id: &str, connector_name: &str) -> Result { - let schema: JsonObject = serde_json::from_value(json!({ - "type": "object", - "additionalProperties": false - }))?; - let mut tool = Tool::new( - Cow::Owned(format!("connector_{connector_id}")), - Cow::Borrowed("Connector test tool"), - Arc::new(schema), - ); - tool.annotations = Some(ToolAnnotations::new().read_only(true)); - - let mut meta = Meta::new(); - meta.0 - .insert("connector_id".to_string(), json!(connector_id)); - meta.0 - .insert("connector_name".to_string(), json!(connector_name)); - tool.meta = Some(meta); - Ok(tool) -} - -fn write_connectors_config(codex_home: &std::path::Path, base_url: &str) -> std::io::Result<()> { - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -chatgpt_base_url = "{base_url}" -mcp_oauth_credentials_store = "file" - -[features] -plugins = true -connectors = true -"# - ), - ) -} - fn write_remote_plugin_catalog_config( codex_home: &std::path::Path, base_url: &str, diff --git a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs index f88fb54929..ddc0eae404 100644 --- a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs +++ b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs @@ -36,7 +36,10 @@ use codex_app_server_protocol::ThreadRealtimeTranscriptDoneNotification; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnCompletedNotification; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStartedNotification; +use codex_app_server_protocol::UserInput as V2UserInput; use codex_features::FEATURES; use codex_features::Feature; use codex_protocol::protocol::RealtimeConversationVersion; @@ -311,12 +314,14 @@ impl RealtimeE2eHarness { .mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { thread_id: self.thread_id.clone(), + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: Some(ThreadRealtimeStartTransport::Webrtc { sdp: offer_sdp.to_string(), }), + version: None, voice: None, }) .await?; @@ -556,10 +561,12 @@ async fn realtime_conversation_streams_v2_notifications() -> Result<()> { let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { thread_id: thread_start.thread.id.clone(), + model: Some("realtime-treatment-model".to_string()), output_modality: RealtimeOutputModality::Audio, prompt: None, realtime_session_id: None, transport: None, + version: None, voice: Some(RealtimeVoice::Cedar), }) .await?; @@ -588,6 +595,10 @@ async fn realtime_conversation_streams_v2_notifications() -> Result<()> { startup_context_request.body_json()["session"]["audio"]["output"]["voice"], "cedar" ); + assert_eq!( + realtime_server.single_handshake().uri(), + "/v1/realtime?model=realtime-treatment-model" + ); assert_eq!( startup_context_request.body_json()["session"]["output_modalities"], json!(["audio"]) @@ -805,10 +816,12 @@ async fn realtime_text_output_modality_requests_text_output_and_final_transcript let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { thread_id: thread_start.thread.id.clone(), + model: None, output_modality: RealtimeOutputModality::Text, prompt: None, realtime_session_id: None, transport: None, + version: None, voice: None, }) .await?; @@ -979,10 +992,12 @@ async fn realtime_conversation_stop_emits_closed_notification() -> Result<()> { let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { thread_id: thread_start.thread.id.clone(), + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, }) .await?; @@ -1076,12 +1091,14 @@ async fn realtime_webrtc_start_emits_sdp_notification() -> Result<()> { let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { thread_id: thread_id.clone(), + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: Some(ThreadRealtimeStartTransport::Webrtc { sdp: "v=offer\r\n".to_string(), }), + version: None, voice: None, }) .await?; @@ -1328,6 +1345,136 @@ async fn webrtc_v1_handoff_request_delegates_and_appends_result() -> Result<()> Ok(()) } +#[tokio::test] +async fn webrtc_assistant_output_without_handoff_reaches_realtime() -> Result<()> { + skip_if_no_network!(Ok(())); + + let final_answer = "long output ".repeat(1_000); + for (version, expected_version, preamble) in [ + ( + RealtimeTestVersion::V1, + RealtimeConversationVersion::V1, + "direct preamble from v1", + ), + ( + RealtimeTestVersion::V2, + RealtimeConversationVersion::V2, + "direct preamble from v2", + ), + ] { + let mut harness = RealtimeE2eHarness::new( + version, + main_loop_responses(vec![responses::sse(vec![ + responses::ev_response_created("resp-1"), + json!({ + "type": "response.output_item.done", + "item": { + "type": "message", + "role": "assistant", + "id": "msg-preamble", + "phase": "commentary", + "content": [{"type": "output_text", "text": preamble}] + } + }), + responses::ev_assistant_message("msg-final", &final_answer), + responses::ev_completed("resp-1"), + ])]), + realtime_sideband(vec![realtime_sideband_connection(vec![ + vec![session_updated("sess_standalone_output")], + vec![], + match version { + RealtimeTestVersion::V1 => vec![], + RealtimeTestVersion::V2 => vec![ + json!({ + "type": "response.created", + "response": { "id": "resp_preamble" } + }), + json!({ + "type": "response.done", + "response": { "id": "resp_preamble" } + }), + ], + }, + vec![], + vec![], + ])]), + ) + .await?; + + let started = harness.start_webrtc_realtime("v=offer\r\n").await?; + assert_eq!(started.started.version, expected_version); + + let request_id = harness + .mcp + .send_turn_start_request(TurnStartParams { + thread_id: harness.thread_id.clone(), + input: vec![V2UserInput::Text { + text: "direct text turn".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + harness + .mcp + .read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let _: TurnStartResponse = to_response(response)?; + let _ = harness + .read_notification::("turn/completed") + .await?; + + let preamble_request = harness.sideband_outbound_request(/*request_index*/ 1).await; + let output_text = match version { + RealtimeTestVersion::V1 => { + let final_request = harness.sideband_outbound_request(/*request_index*/ 2).await; + assert_eq!( + preamble_request, + json!({ + "type": "conversation.handoff.append", + "handoff_id": "codex", + "output_text": preamble, + }) + ); + assert_eq!(final_request["type"], "conversation.handoff.append"); + assert_eq!(final_request["handoff_id"], "codex"); + final_request["output_text"] + .as_str() + .expect("output text") + .to_string() + } + RealtimeTestVersion::V2 => { + assert_v2_progress_update(&preamble_request, preamble); + assert_v2_response_create( + &harness.sideband_outbound_request(/*request_index*/ 2).await, + ); + let final_request = harness.sideband_outbound_request(/*request_index*/ 3).await; + assert_eq!(final_request["type"], "conversation.item.create"); + assert_eq!(final_request["item"]["type"], "message"); + assert_eq!(final_request["item"]["role"], "user"); + assert_eq!(final_request["item"]["content"][0]["type"], "input_text"); + let output_text = final_request["item"]["content"][0]["text"] + .as_str() + .expect("output text"); + assert!(output_text.starts_with("[BACKEND] ")); + assert_v2_response_create( + &harness.sideband_outbound_request(/*request_index*/ 4).await, + ); + output_text.to_string() + } + }; + assert!(output_text.contains("tokens truncated")); + assert!(output_text.len() <= 4_000); + + harness.shutdown().await; + } + + Ok(()) +} + #[tokio::test] async fn webrtc_v2_forwards_audio_and_text_between_client_and_sideband() -> Result<()> { skip_if_no_network!(Ok(())); @@ -1991,12 +2138,14 @@ async fn realtime_webrtc_start_surfaces_backend_error() -> Result<()> { let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { thread_id: thread_start.thread.id, + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: Some(ThreadRealtimeStartTransport::Webrtc { sdp: "v=offer\r\n".to_string(), }), + version: None, voice: None, }) .await?; @@ -2050,10 +2199,12 @@ async fn realtime_conversation_requires_feature_flag() -> Result<()> { let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { thread_id: thread_start.thread.id.clone(), + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, }) .await?; diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index ca2121dd4c..18d1dcd43f 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -20,6 +20,7 @@ use std::sync::Arc; use anyhow::Result; use app_test_support::create_mock_responses_server_repeating_assistant; use codex_app_server::in_process; +use codex_app_server::in_process::InProcessClientHandle; use codex_app_server::in_process::InProcessServerEvent; use codex_app_server::in_process::InProcessStartArgs; use codex_app_server_protocol::ClientInfo; @@ -27,6 +28,8 @@ use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ThreadDeleteParams; +use codex_app_server_protocol::ThreadDeleteResponse; use codex_app_server_protocol::ThreadListParams; use codex_app_server_protocol::ThreadListResponse; use codex_app_server_protocol::ThreadResumeParams; @@ -42,8 +45,14 @@ use codex_core::config::Config; use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; +use codex_protocol::ThreadId; +use codex_protocol::models::BaseInstructions; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadMemoryMode; +use codex_thread_store::CreateThreadParams as StoreCreateThreadParams; use codex_thread_store::InMemoryThreadStore; +use codex_thread_store::ThreadPersistenceMetadata; +use codex_thread_store::ThreadStore; use pretty_assertions::assert_eq; use tempfile::TempDir; use tokio::time::timeout; @@ -52,7 +61,7 @@ use uuid::Uuid; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); #[tokio::test] -async fn thread_start_with_non_local_thread_store_does_not_create_local_persistence() -> Result<()> +async fn thread_delete_with_non_local_thread_store_does_not_create_local_persistence() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; @@ -61,43 +70,10 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste // here so this regression stays focused on thread persistence artifacts. create_config_toml_with_thread_store(codex_home.path(), &server.uri(), &store_id)?; - let loader_overrides = LoaderOverrides::without_managed_config_for_tests(); - let config = ConfigBuilder::default() - .codex_home(codex_home.path().to_path_buf()) - .fallback_cwd(Some(codex_home.path().to_path_buf())) - .loader_overrides(loader_overrides.clone()) - .build() - .await?; - let thread_store = InMemoryThreadStore::for_id(store_id.clone()); let _in_memory_store = InMemoryThreadStoreId { store_id }; - let mut client = in_process::start(InProcessStartArgs { - arg0_paths: Arg0DispatchPaths::default(), - config: Arc::new(config), - cli_overrides: Vec::new(), - loader_overrides, - strict_config: false, - cloud_config_bundle: CloudConfigBundleLoader::default(), - thread_config_loader: Arc::new(NoopThreadConfigLoader), - feedback: CodexFeedback::new(), - log_db: None, - state_db: None, - environment_manager: Arc::new(EnvironmentManager::default_for_tests()), - config_warnings: Vec::new(), - session_source: SessionSource::Cli, - enable_codex_api_key_env: false, - initialize: InitializeParams { - client_info: ClientInfo { - name: "codex-app-server-tests".to_string(), - title: None, - version: "0.1.0".to_string(), - }, - capabilities: None, - }, - channel_capacity: in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, - }) - .await?; + let mut client = start_in_process_server(codex_home.path()).await?; let response = client .request(ClientRequest::ThreadStart { @@ -166,11 +142,39 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste assert_eq!(data[0].id, thread.id); assert_eq!(data[0].path, None); + delete_thread(&client, /*request_id*/ 4, thread.id.clone()).await?; + let unloaded_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string())?; + thread_store + .create_thread(StoreCreateThreadParams { + thread_id: unloaded_thread_id, + extra_config: None, + forked_from_id: None, + parent_thread_id: None, + source: SessionSource::Cli, + thread_source: None, + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + multi_agent_version: None, + metadata: ThreadPersistenceMetadata { + cwd: Some(codex_home.path().to_path_buf()), + model_provider: "mock_provider".to_string(), + memory_mode: ThreadMemoryMode::Enabled, + }, + }) + .await?; + delete_thread( + &client, + /*request_id*/ 5, + unloaded_thread_id.to_string(), + ) + .await?; + client.shutdown().await?; let calls = thread_store.calls().await; - assert_eq!(calls.create_thread, 1); + assert_eq!(calls.create_thread, 2); assert_eq!(calls.list_threads, 1); + assert_eq!(calls.delete_thread, 2); assert!( calls.append_items > 0, "turn/start should append rollout items through the injected store" @@ -269,10 +273,24 @@ async fn cold_thread_resume_reuses_non_local_history_probe() -> Result<()> { Ok(()) } +async fn start_in_process_server(codex_home: &Path) -> Result { + let loader_overrides = LoaderOverrides::without_managed_config_for_tests(); + let config = Arc::new( + ConfigBuilder::default() + .codex_home(codex_home.to_path_buf()) + .fallback_cwd(Some(codex_home.to_path_buf())) + .loader_overrides(loader_overrides.clone()) + .build() + .await?, + ); + + Ok(start_in_process_client(config, loader_overrides).await?) +} + async fn start_in_process_client( config: Arc, loader_overrides: LoaderOverrides, -) -> std::io::Result { +) -> std::io::Result { in_process::start(InProcessStartArgs { arg0_paths: Arg0DispatchPaths::default(), config, @@ -301,6 +319,22 @@ async fn start_in_process_client( .await } +async fn delete_thread( + client: &InProcessClientHandle, + request_id: i64, + thread_id: String, +) -> Result<()> { + let response = client + .request(ClientRequest::ThreadDelete { + request_id: RequestId::Integer(request_id), + params: ThreadDeleteParams { thread_id }, + }) + .await? + .map_err(|error| anyhow::anyhow!("thread/delete failed: {}", error.message))?; + let _: ThreadDeleteResponse = serde_json::from_value(response)?; + Ok(()) +} + fn assert_no_local_persistence_artifacts(codex_home: &Path) -> Result<()> { // These are the observable tripwires for accidental local persistence. If a // future code path constructs a local rollout/session store or opens the diff --git a/codex-rs/app-server/tests/suite/v2/thread_delete.rs b/codex-rs/app-server/tests/suite/v2/thread_delete.rs new file mode 100644 index 0000000000..bef155f9ad --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/thread_delete.rs @@ -0,0 +1,200 @@ +use anyhow::Result; +use app_test_support::TestAppServer; +use app_test_support::create_fake_rollout; +use app_test_support::to_response; +use codex_app_server_protocol::JSONRPCError; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ThreadDeleteParams; +use codex_app_server_protocol::ThreadDeleteResponse; +use codex_app_server_protocol::ThreadDeletedNotification; +use codex_app_server_protocol::ThreadLoadedListParams; +use codex_app_server_protocol::ThreadLoadedListResponse; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_core::find_thread_path_by_id_str; +use codex_protocol::ThreadId; +use codex_state::DirectionalThreadSpawnEdgeStatus; +use codex_state::StateRuntime; +use pretty_assertions::assert_eq; +use std::path::Path; +use tempfile::TempDir; +use tokio::time::timeout; + +const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); + +#[tokio::test] +async fn thread_delete_deletes_spawned_descendants() -> Result<()> { + let codex_home = TempDir::new()?; + + let parent_id = create_delete_test_rollout(codex_home.path(), /*minute*/ 0, "parent")?; + let child_id = create_delete_test_rollout(codex_home.path(), /*minute*/ 1, "child")?; + let grandchild_id = + create_delete_test_rollout(codex_home.path(), /*minute*/ 2, "grandchild")?; + + let state_db = + StateRuntime::init(codex_home.path().to_path_buf(), "mock_provider".into()).await?; + let parent_thread_id = ThreadId::from_string(&parent_id)?; + let child_thread_id = ThreadId::from_string(&child_id)?; + let grandchild_thread_id = ThreadId::from_string(&grandchild_id)?; + + for (parent, child, status) in [ + ( + parent_thread_id, + child_thread_id, + DirectionalThreadSpawnEdgeStatus::Closed, + ), + ( + child_thread_id, + grandchild_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ), + ] { + state_db + .upsert_thread_spawn_edge(parent, child, status) + .await?; + } + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let delete_id = mcp + .send_thread_delete_request(ThreadDeleteParams { + thread_id: parent_id.clone(), + }) + .await?; + let delete_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(delete_id)), + ) + .await??; + let _: ThreadDeleteResponse = to_response::(delete_resp)?; + + let mut deleted_ids = Vec::new(); + for _ in 0..3 { + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("thread/deleted"), + ) + .await??; + let deleted_notification: ThreadDeletedNotification = serde_json::from_value( + notification + .params + .expect("thread/deleted notification params"), + )?; + deleted_ids.push(deleted_notification.thread_id); + } + assert_eq!(deleted_ids, vec![grandchild_id, child_id, parent_id]); + + for thread_id in [parent_thread_id, child_thread_id, grandchild_thread_id] { + let rollout_path = find_thread_path_by_id_str( + codex_home.path(), + &thread_id.to_string(), + /*state_db_ctx*/ None, + ) + .await?; + assert!( + rollout_path.is_none(), + "expected active rollout for {thread_id} to be deleted" + ); + } + assert_eq!( + state_db + .list_thread_spawn_descendants(parent_thread_id) + .await?, + Vec::::new() + ); + Ok(()) +} + +fn create_delete_test_rollout(codex_home: &Path, minute: u8, preview: &str) -> Result { + create_fake_rollout( + codex_home, + &format!("2025-01-01T00-{minute:02}-00"), + &format!("2025-01-01T00:{minute:02}:00Z"), + preview, + Some("mock_provider"), + /*git_info*/ None, + ) +} + +#[tokio::test] +async fn thread_delete_handles_live_threads_before_rollout_exists() -> Result<()> { + let codex_home = TempDir::new()?; + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let start_id = mcp + .send_thread_start_request(ThreadStartParams::default()) + .await?; + let start_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(start_id)), + ) + .await??; + let persisted_thread = to_response::(start_resp)?.thread; + let rollout_path = find_thread_path_by_id_str( + codex_home.path(), + &persisted_thread.id, + /*state_db_ctx*/ None, + ) + .await?; + assert_eq!(rollout_path, None); + + let delete_id = mcp + .send_thread_delete_request(ThreadDeleteParams { + thread_id: persisted_thread.id, + }) + .await?; + let delete_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(delete_id)), + ) + .await??; + let _: ThreadDeleteResponse = to_response::(delete_resp)?; + + let start_id = mcp + .send_thread_start_request(ThreadStartParams { + ephemeral: Some(true), + ..Default::default() + }) + .await?; + let start_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(start_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + + let delete_id = mcp + .send_thread_delete_request(ThreadDeleteParams { + thread_id: thread.id.clone(), + }) + .await?; + let delete_err: JSONRPCError = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(delete_id)), + ) + .await??; + let expected_message = format!( + "thread is not persisted and cannot be deleted: {}", + thread.id + ); + assert_eq!(delete_err.error.message, expected_message); + + let list_id = mcp + .send_thread_loaded_list_request(ThreadLoadedListParams::default()) + .await?; + let list_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(list_id)), + ) + .await??; + let ThreadLoadedListResponse { mut data, .. } = + to_response::(list_resp)?; + data.sort(); + assert_eq!(data, vec![thread.id]); + + Ok(()) +} 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 582177ee7d..cf797ff6be 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -103,6 +103,7 @@ use super::analytics::assert_basic_thread_initialized_event; use super::analytics::mount_analytics_capture; use super::analytics::thread_initialized_event; use super::analytics::wait_for_analytics_payload; +use super::analytics::wait_for_goal_event; #[cfg(windows)] const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(25); @@ -1316,19 +1317,30 @@ async fn thread_goal_set_edits_objective_without_resetting_usage() -> Result<()> } #[tokio::test] -async fn thread_goal_clear_deletes_goal_and_notifies() -> Result<()> { - let server = create_mock_responses_server_repeating_assistant("Done").await; +async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Result<()> { + let server = create_mock_responses_server_sequence_unchecked(vec![ + responses::sse(vec![ + responses::ev_response_created("materialize-thread"), + responses::ev_completed("materialize-thread"), + ]), + responses::sse(vec![ + responses::ev_response_created("goal-continuation"), + responses::ev_completed_with_tokens("goal-continuation", /*total_tokens*/ 200), + ]), + ]) + .await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + create_config_toml_with_chatgpt_base_url(codex_home.path(), &server.uri(), &server.uri())?; let config_path = codex_home.path().join("config.toml"); let config = std::fs::read_to_string(&config_path)?; std::fs::write( &config_path, config.replace("personality = true\n", "personality = true\ngoals = true\n"), )?; + mount_analytics_capture(&server, codex_home.path()).await?; let mut mcp = TestAppServer::new_without_managed_config(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + timeout(DEFAULT_READ_TIMEOUT.saturating_mul(2), mcp.initialize()).await??; let start_id = mcp .send_thread_start_request(ThreadStartParams { @@ -1370,7 +1382,8 @@ async fn thread_goal_clear_deletes_goal_and_notifies() -> Result<()> { "thread/goal/set", Some(json!({ "threadId": thread.id, - "objective": "keep polishing", + "objective": "do not serialize this objective", + "tokenBudget": 100, })), ) .await?; @@ -1386,6 +1399,55 @@ async fn thread_goal_clear_deletes_goal_and_notifies() -> Result<()> { ) .await??; + let created = wait_for_goal_event(&server, DEFAULT_READ_TIMEOUT, "created", "active").await?; + let persisted_goal_id = created["event_params"]["goal_id"] + .as_str() + .expect("created goal id"); + assert_eq!(created["event_params"]["thread_id"], thread.id); + assert_eq!(created["event_params"]["turn_id"], serde_json::Value::Null); + assert_eq!(created["event_params"]["has_token_budget"], true); + assert!(created["event_params"]["session_id"].is_string()); + assert!(created["event_params"]["app_server_client"].is_object()); + assert!(created["event_params"]["runtime"].is_object()); + assert!(created["event_params"].get("objective").is_none()); + assert!(created["event_params"].get("token_budget").is_none()); + + let usage = wait_for_goal_event( + &server, + DEFAULT_READ_TIMEOUT, + "usage_accounted", + "budget_limited", + ) + .await?; + let causal_turn_id = usage["event_params"]["turn_id"] + .as_str() + .expect("accounted usage turn id"); + assert_eq!(usage["event_params"]["goal_id"], persisted_goal_id); + assert_eq!(usage["event_params"]["cumulative_tokens_accounted"], 200); + assert!( + usage["event_params"]["cumulative_time_accounted_seconds"] + .as_i64() + .is_some() + ); + + let status = wait_for_goal_event( + &server, + DEFAULT_READ_TIMEOUT, + "status_changed", + "budget_limited", + ) + .await?; + assert_eq!(status["event_params"]["goal_id"], persisted_goal_id); + assert_eq!(status["event_params"]["turn_id"], causal_turn_id); + assert_eq!( + status["event_params"]["cumulative_tokens_accounted"], + serde_json::Value::Null + ); + assert_eq!( + status["event_params"]["cumulative_time_accounted_seconds"], + serde_json::Value::Null + ); + let clear_id = mcp .send_raw_request( "thread/goal/clear", @@ -1408,6 +1470,11 @@ async fn thread_goal_clear_deletes_goal_and_notifies() -> Result<()> { ) .await??; + let cleared = + wait_for_goal_event(&server, DEFAULT_READ_TIMEOUT, "cleared", "budget_limited").await?; + assert_eq!(cleared["event_params"]["goal_id"], persisted_goal_id); + assert_eq!(cleared["event_params"]["turn_id"], serde_json::Value::Null); + let get_id = mcp .send_raw_request( "thread/goal/get", diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index 5080937997..7fc0d8c71d 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -43,7 +43,12 @@ use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ServerRequestResolvedNotification; use codex_app_server_protocol::SubAgentActivityKind; use codex_app_server_protocol::TextElement; +use codex_app_server_protocol::ThreadDeleteParams; +use codex_app_server_protocol::ThreadDeleteResponse; +use codex_app_server_protocol::ThreadDeletedNotification; use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::ThreadLoadedListParams; +use codex_app_server_protocol::ThreadLoadedListResponse; use codex_app_server_protocol::ThreadSource; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -591,7 +596,15 @@ async fn turn_start_emits_thread_scoped_warning_notification_for_trimmed_skills( write_test_skill(codex_home.path(), "alpha-skill")?; write_test_skill(codex_home.path(), "beta-skill")?; - let mut mcp = TestAppServer::new(codex_home.path()).await?; + let isolated_home = codex_home.path().to_string_lossy(); + let mut mcp = TestAppServer::new_with_env( + codex_home.path(), + &[ + ("HOME", Some(isolated_home.as_ref())), + ("USERPROFILE", Some(isolated_home.as_ref())), + ], + ) + .await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_req = mcp @@ -3362,6 +3375,49 @@ async fn turn_start_emits_spawn_agent_item_with_model_metadata_v2() -> Result<() assert_eq!(turn_completed.thread_id, thread.id); assert_eq!(turn_completed.turn.id, turn.turn.id); + // Reuse this live spawn setup to cover thread/delete's ThreadManager descendant path. + let delete_req = mcp + .send_thread_delete_request(ThreadDeleteParams { + thread_id: thread.id.clone(), + }) + .await?; + let delete_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(delete_req)), + ) + .await??; + let _: ThreadDeleteResponse = to_response::(delete_resp)?; + + let mut deleted_thread_ids = Vec::new(); + for _ in 0..2 { + let deleted_notif = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("thread/deleted"), + ) + .await??; + let deleted: ThreadDeletedNotification = serde_json::from_value( + deleted_notif + .params + .expect("thread/deleted notification params"), + )?; + deleted_thread_ids.push(deleted.thread_id); + } + assert_eq!( + deleted_thread_ids, + vec![receiver_thread_id, thread.id.clone()] + ); + + let list_req = mcp + .send_thread_loaded_list_request(ThreadLoadedListParams::default()) + .await?; + let list_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(list_req)), + ) + .await??; + let ThreadLoadedListResponse { data, .. } = to_response::(list_resp)?; + assert_eq!(data, Vec::::new()); + Ok(()) } diff --git a/codex-rs/cli/BUILD.bazel b/codex-rs/cli/BUILD.bazel index faa8627477..97c868845b 100644 --- a/codex-rs/cli/BUILD.bazel +++ b/codex-rs/cli/BUILD.bazel @@ -1,4 +1,5 @@ -load("//:defs.bzl", "MACOS_WEBRTC_RUSTC_LINK_FLAGS", "codex_rust_crate", "multiplatform_binaries") +load("//:defs.bzl", "MACOS_WEBRTC_RUSTC_LINK_FLAGS", "codex_rust_crate") +load("//bazel/platforms:release_binaries.bzl", "multiplatform_binaries") codex_rust_crate( name = "cli", diff --git a/codex-rs/cli/src/doctor.rs b/codex-rs/cli/src/doctor.rs index d817f1a192..0b7e7cfe7b 100644 --- a/codex-rs/cli/src/doctor.rs +++ b/codex-rs/cli/src/doctor.rs @@ -1324,6 +1324,7 @@ fn stored_auth_mode(auth: &codex_login::AuthDotJson) -> &'static str { codex_app_server_protocol::AuthMode::ChatgptAuthTokens => "chatgpt_auth_tokens", codex_app_server_protocol::AuthMode::AgentIdentity => "agent_identity", codex_app_server_protocol::AuthMode::PersonalAccessToken => "personal_access_token", + codex_app_server_protocol::AuthMode::BedrockApiKey => "bedrock_api_key", } } @@ -1331,10 +1332,12 @@ fn stored_auth_mode_value(auth: &AuthDotJson) -> codex_app_server_protocol::Auth if let Some(mode) = auth.auth_mode { return mode; } - if auth.openai_api_key.is_some() { - codex_app_server_protocol::AuthMode::ApiKey - } else if auth.personal_access_token.is_some() { + if auth.personal_access_token.is_some() { codex_app_server_protocol::AuthMode::PersonalAccessToken + } else if auth.bedrock_api_key.is_some() { + codex_app_server_protocol::AuthMode::BedrockApiKey + } else if auth.openai_api_key.is_some() { + codex_app_server_protocol::AuthMode::ApiKey } else { codex_app_server_protocol::AuthMode::Chatgpt } @@ -1407,6 +1410,11 @@ fn stored_auth_issues( issues.push("personal access token auth is missing a personal access token"); } } + codex_app_server_protocol::AuthMode::BedrockApiKey => { + if auth.bedrock_api_key.is_none() { + issues.push("Bedrock API key auth is missing a Bedrock API key"); + } + } } issues } @@ -2141,8 +2149,9 @@ async fn state_check(config: &Config) -> DoctorCheck { }; let mut check = DoctorCheck::new("state.paths", "state", status, summary).details(details); if status == CheckStatus::Fail { - check = check - .remediation("Back up CODEX_HOME, then remove or repair the affected SQLite database."); + check = check.remediation( + "Move the damaged SQLite database aside, then restart the interactive CLI or app server so it can rebuild that runtime database from saved data. Other entry points may not rebuild automatically.", + ); } check } @@ -2436,6 +2445,7 @@ fn auth_mode_name(auth: &CodexAuth) -> &'static str { codex_app_server_protocol::AuthMode::ChatgptAuthTokens => "chatgpt_auth_tokens", codex_app_server_protocol::AuthMode::AgentIdentity => "agent_identity", codex_app_server_protocol::AuthMode::PersonalAccessToken => "personal_access_token", + codex_app_server_protocol::AuthMode::BedrockApiKey => "bedrock_api_key", } } @@ -2565,7 +2575,10 @@ fn provider_auth_reachability_mode_from_auth( return ProviderAuthReachabilityMode::Chatgpt; } match stored_auth.map(stored_auth_mode_value) { - Some(codex_app_server_protocol::AuthMode::ApiKey) => ProviderAuthReachabilityMode::ApiKey, + Some( + codex_app_server_protocol::AuthMode::ApiKey + | codex_app_server_protocol::AuthMode::BedrockApiKey, + ) => ProviderAuthReachabilityMode::ApiKey, Some( codex_app_server_protocol::AuthMode::Chatgpt | codex_app_server_protocol::AuthMode::ChatgptAuthTokens @@ -3470,6 +3483,7 @@ mod tests { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; assert_eq!( @@ -3488,6 +3502,7 @@ mod tests { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; assert_eq!( @@ -3508,6 +3523,7 @@ mod tests { last_refresh: None, agent_identity: None, personal_access_token: Some("at-test".to_string()), + bedrock_api_key: None, }; assert_eq!(stored_auth_mode(&auth), "personal_access_token"); @@ -3530,6 +3546,7 @@ mod tests { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; assert_eq!( diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index b3fc1bd544..388a9f4526 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -395,6 +395,10 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { eprintln!("Logged in using personal access token"); std::process::exit(0); } + AuthMode::BedrockApiKey => { + eprintln!("Logged in using Amazon Bedrock API key"); + std::process::exit(0); + } }, Ok(None) => { eprintln!("Not logged in"); diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index b194a5ab1b..c9fb185556 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -38,6 +38,7 @@ use codex_utils_cli::ProfileV2Name; use codex_utils_cli::SharedCliOptions; use codex_utils_cli::resume_hint; use owo_colors::OwoColorize; +use std::collections::HashSet; use std::io::IsTerminal; use std::path::PathBuf; use supports_color::Stream; @@ -179,6 +180,9 @@ enum Subcommand { /// Archive a saved session by id or session name. Archive(SessionArchiveCommand), + /// Permanently delete a saved session by id or session name. + Delete(DeleteCommand), + /// Unarchive a saved session by id or session name. Unarchive(SessionArchiveCommand), @@ -353,6 +357,16 @@ struct SessionArchiveConfigOverrides { config_overrides: CliConfigOverrides, } +#[derive(Debug, Args)] +struct DeleteCommand { + #[clap(flatten)] + session: SessionArchiveCommand, + + /// Delete without prompting. SESSION must be a UUID. + #[arg(long, default_value_t = false)] + force: bool, +} + #[derive(Debug, Parser)] struct ForkCommand { /// Conversation/session id (UUID). When provided, forks this session. @@ -828,6 +842,17 @@ async fn run_session_archive_cli_command( .map_err(|err| anyhow::anyhow!("{err}")) } +fn delete_action(target: &str, force: bool) -> anyhow::Result { + if force && codex_protocol::ThreadId::from_string(target).is_err() { + anyhow::bail!("--force requires a session UUID; names must be confirmed interactively"); + } + let confirmation = match force { + true => codex_tui::DeleteConfirmation::Skip, + false => codex_tui::DeleteConfirmation::Prompt, + }; + Ok(codex_tui::SessionArchiveAction::Delete(confirmation)) +} + async fn run_debug_app_server_command(cmd: DebugAppServerCommand) -> anyhow::Result<()> { match cmd.subcommand { DebugAppServerSubcommand::SendMessageV2(cmd) => { @@ -1232,6 +1257,20 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { .await?; println!("{output}"); } + Some(Subcommand::Delete(DeleteCommand { session, force })) => { + let action = delete_action(&session.target, force)?; + let output = run_session_archive_cli_command( + action, + session, + interactive, + root_config_overrides.clone(), + root_remote.clone(), + root_remote_auth_token_env.clone(), + arg0_paths.clone(), + ) + .await?; + println!("{output}"); + } Some(Subcommand::Unarchive(cmd)) => { let output = run_session_archive_cli_command( codex_tui::SessionArchiveAction::Unarchive, @@ -1596,6 +1635,7 @@ fn profile_v2_for_subcommand<'a>( | Subcommand::Review(_) | Subcommand::Resume(_) | Subcommand::Archive(_) + | Subcommand::Delete(_) | Subcommand::Unarchive(_) | Subcommand::Fork(_) | Subcommand::Mcp(_) @@ -1604,7 +1644,7 @@ fn profile_v2_for_subcommand<'a>( subcommand: DebugSubcommand::PromptInput(_), }) => Ok(Some(profile_v2)), _ => anyhow::bail!( - "--profile only applies to runtime commands and `codex mcp`: `codex`, `codex exec`, `codex review`, `codex resume`, `codex archive`, `codex unarchive`, `codex fork`, `codex mcp`, `codex sandbox`, and `codex debug prompt-input`." + "--profile only applies to runtime commands and `codex mcp`: `codex`, `codex exec`, `codex review`, `codex resume`, `codex archive`, `codex delete`, `codex unarchive`, `codex fork`, `codex mcp`, `codex sandbox`, and `codex debug prompt-input`." ), } } @@ -2025,6 +2065,7 @@ fn unsupported_subcommand_name_for_strict_config( | Some(Subcommand::ExecServer(_)) | Some(Subcommand::Resume(_)) | Some(Subcommand::Archive(_)) + | Some(Subcommand::Delete(_)) | Some(Subcommand::Unarchive(_)) | Some(Subcommand::Fork(_)) | Some(Subcommand::Doctor(_)) => None, @@ -2188,7 +2229,7 @@ async fn run_interactive_tui( remote_endpoint.clone(), ) }; - let mut attempted_repair = false; + let mut attempted_backups = HashSet::new(); loop { let err = match start_tui().await { Ok(exit_info) => return Ok(exit_info), @@ -2201,25 +2242,25 @@ async fn run_interactive_tui( local_state_db::print_locked_guidance(startup_error); return Ok(AppExitInfo::fatal(startup_error.to_string())); } - if attempted_repair { + if !local_state_db::is_auto_backup_recoverable(startup_error) { local_state_db::print_diagnostic_guidance(startup_error); return Ok(AppExitInfo::fatal(startup_error.to_string())); } - if !local_state_db::confirm_repair(startup_error)? { + if !attempted_backups.insert(startup_error.database_path().to_path_buf()) { local_state_db::print_diagnostic_guidance(startup_error); return Ok(AppExitInfo::fatal(startup_error.to_string())); } - match local_state_db::repair_files(startup_error).await { - Ok(backups) => local_state_db::print_repair_backups(&backups), - Err(repair_err) => { + local_state_db::print_auto_backup_start(startup_error); + match local_state_db::backup_files_for_fresh_start(startup_error).await { + Ok(backups) => local_state_db::confirm_fresh_start_rebuild(startup_error, &backups)?, + Err(backup_err) => { local_state_db::print_diagnostic_guidance(startup_error); return Ok(AppExitInfo::fatal(format!( - "failed to repair Codex local data automatically: {repair_err}" + "failed to move damaged Codex local database files into a backup folder automatically: {backup_err}" ))); } } - attempted_repair = true; } } @@ -2599,6 +2640,14 @@ mod tests { ); } + #[test] + fn import_remains_an_interactive_prompt() { + let cli = MultitoolCli::try_parse_from(["codex", "import"]).expect("parse"); + + assert!(cli.subcommand.is_none()); + assert_eq!(cli.interactive.prompt.as_deref(), Some("import")); + } + #[test] fn profile_v2_rejects_non_plain_names_at_parse_time() { assert!( @@ -2859,6 +2908,17 @@ mod tests { assert!(interactive.bypass_hook_trust); } + #[test] + fn delete_force_requires_uuid() { + assert!(delete_action("123e4567-e89b-12d3-a456-426614174000", true).is_ok()); + + let err = delete_action("my-thread", true).expect_err("name should require prompt"); + assert_eq!( + err.to_string(), + "--force requires a session UUID; names must be confirmed interactively" + ); + } + #[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))] #[test] fn sandbox_parses_permissions_profile() { diff --git a/codex-rs/cli/src/plugin_cmd.rs b/codex-rs/cli/src/plugin_cmd.rs index 94c388066d..d6e28b5da0 100644 --- a/codex-rs/cli/src/plugin_cmd.rs +++ b/codex-rs/cli/src/plugin_cmd.rs @@ -204,7 +204,7 @@ pub async fn run_plugin_list( .. } = load_plugin_command_context(overrides).await?; let outcome = manager - .list_marketplaces_for_config(&plugins_input, &[]) + .list_marketplaces_for_config(&plugins_input, &[], /*include_openai_curated*/ true) .context("failed to list marketplace plugins")?; ensure_configured_marketplace_snapshots_loaded( codex_home.as_path(), @@ -609,7 +609,7 @@ fn find_marketplace_for_plugin( plugin_name: &str, ) -> Result { let outcome = manager - .list_marketplaces_for_config(plugins_input, &[]) + .list_marketplaces_for_config(plugins_input, &[], /*include_openai_curated*/ true) .context("failed to list marketplace plugins")?; ensure_configured_marketplace_snapshots_loaded( codex_home, diff --git a/codex-rs/cli/src/state_db_recovery.rs b/codex-rs/cli/src/state_db_recovery.rs index 7aeffaca3a..89bb16643c 100644 --- a/codex-rs/cli/src/state_db_recovery.rs +++ b/codex-rs/cli/src/state_db_recovery.rs @@ -1,10 +1,12 @@ -//! CLI recovery for local state database startup failures. +//! CLI handling for local state database startup failures. //! -//! This keeps user-facing repair and lock-contention handling out of the main +//! This keeps user-facing backup and lock-contention handling out of the main //! CLI dispatch path while preserving the TUI startup error as the boundary type. +use codex_state::RuntimeDbBackup; use codex_tui::LocalStateDbStartupError; -use std::path::PathBuf; +use std::io::IsTerminal; +use std::path::Path; pub(crate) fn startup_error(err: &std::io::Error) -> Option<&LocalStateDbStartupError> { err.get_ref() @@ -12,66 +14,60 @@ pub(crate) fn startup_error(err: &std::io::Error) -> Option<&LocalStateDbStartup } pub(crate) fn is_locked(detail: &str) -> bool { - let detail = detail.to_ascii_lowercase(); - detail.contains("database is locked") || detail.contains("database is busy") + codex_state::sqlite_error_detail_is_lock(detail) } -pub(crate) fn confirm_repair(startup_error: &LocalStateDbStartupError) -> std::io::Result { +pub(crate) fn is_corruption(detail: &str) -> bool { + codex_state::sqlite_error_detail_is_corruption(detail) +} + +pub(crate) fn is_auto_backup_recoverable(startup_error: &LocalStateDbStartupError) -> bool { + is_corruption(startup_error.detail()) || sqlite_home_is_blocking_file(startup_error) +} + +fn sqlite_home_is_blocking_file(startup_error: &LocalStateDbStartupError) -> bool { + startup_error + .database_path() + .parent() + .and_then(|path| std::fs::metadata(path).ok()) + .is_some_and(|metadata| metadata.is_file()) +} + +pub(crate) fn print_auto_backup_start(startup_error: &LocalStateDbStartupError) { eprintln!("Codex couldn't start because its local database appears to be damaged."); - eprintln!("Codex can try a safe repair by backing up those files and rebuilding them."); + eprintln!("Moving the damaged local database aside so Codex can rebuild it from saved data."); print_technical_details(startup_error); - crate::confirm("Repair Codex local data now? [y/N]: ") } -pub(crate) async fn repair_files( +pub(crate) async fn backup_files_for_fresh_start( startup_error: &LocalStateDbStartupError, -) -> std::io::Result> { - let state_db_path = startup_error.state_db_path(); - let sqlite_home = state_db_path.parent().ok_or_else(|| { - std::io::Error::other("state database path does not have a parent directory") - })?; - let timestamp = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map_or(0, |duration| duration.as_secs()); - let repair_suffix = format!("codex-repair-{timestamp}"); - let mut backups = Vec::new(); - - match tokio::fs::metadata(sqlite_home).await { - Ok(metadata) if metadata.is_dir() => {} - Ok(_) => { - backups.push(backup_path(sqlite_home, &repair_suffix).await?); - tokio::fs::create_dir_all(sqlite_home).await?; - } - Err(err) if err.kind() == std::io::ErrorKind::NotFound => { - tokio::fs::create_dir_all(sqlite_home).await?; - } - Err(err) => return Err(err), - } - - for path in codex_state::runtime_db_paths(sqlite_home) - .into_iter() - .flat_map(|db| sqlite_paths(db.path.as_path())) - { - if tokio::fs::try_exists(path.as_path()).await? { - backups.push(backup_path(path.as_path(), &repair_suffix).await?); - } - } - - if backups.is_empty() { - return Err(std::io::Error::other( - "no repairable Codex local data files were found", - )); - } - - Ok(backups) +) -> std::io::Result> { + codex_state::backup_runtime_db_for_fresh_start(startup_error.database_path()).await } -pub(crate) fn print_repair_backups(backups: &[PathBuf]) { - eprintln!("Backed up Codex local data before repair:"); - for backup in backups { - eprintln!(" {}", backup.display()); +pub(crate) fn confirm_fresh_start_rebuild( + startup_error: &LocalStateDbStartupError, + backups: &[RuntimeDbBackup], +) -> std::io::Result<()> { + eprintln!("Codex rebuilt its local database."); + eprintln!( + "Codex detected a damaged local database, moved it into a backup folder, and will continue startup with a fresh database." + ); + eprintln!("Database path: {}", startup_error.database_path().display()); + if let Some(backup_folder) = backup_folder(backups) { + eprintln!("Backup folder: {}", backup_folder.display()); + } else { + eprintln!("Backup folder: unavailable"); } - eprintln!("Retrying startup with rebuilt local data..."); + + if std::io::stdin().is_terminal() && std::io::stderr().is_terminal() { + eprintln!("Press Enter to continue."); + let mut input = String::new(); + std::io::stdin().read_line(&mut input)?; + } else { + eprintln!("Continuing startup with a fresh local database..."); + } + Ok(()) } pub(crate) fn print_diagnostic_guidance(startup_error: &LocalStateDbStartupError) { @@ -87,79 +83,50 @@ pub(crate) fn print_locked_guidance(startup_error: &LocalStateDbStartupError) { print_technical_details(startup_error); } -fn sqlite_paths(db_path: &std::path::Path) -> Vec { - let mut wal_path = db_path.as_os_str().to_os_string(); - wal_path.push("-wal"); - let mut shm_path = db_path.as_os_str().to_os_string(); - shm_path.push("-shm"); - vec![ - db_path.to_path_buf(), - PathBuf::from(wal_path), - PathBuf::from(shm_path), - ] -} - -async fn backup_path(path: &std::path::Path, repair_suffix: &str) -> std::io::Result { - let file_name = path.file_name().ok_or_else(|| { - std::io::Error::other(format!( - "cannot create a repair backup name for {}", - path.display() - )) - })?; - let mut sequence = 0; - loop { - let mut backup_name = file_name.to_os_string(); - backup_name.push(format!(".{repair_suffix}.{sequence}.bak")); - let backup_path = path.with_file_name(backup_name); - if !tokio::fs::try_exists(backup_path.as_path()).await? { - tokio::fs::rename(path, backup_path.as_path()).await?; - return Ok(backup_path); - } - sequence += 1; - } -} - fn print_technical_details(startup_error: &LocalStateDbStartupError) { eprintln!("Technical details:"); - eprintln!(" Location: {}", startup_error.state_db_path().display()); + eprintln!(" Location: {}", startup_error.database_path().display()); eprintln!(" Cause: {}", startup_error.detail()); } +fn backup_folder(backups: &[RuntimeDbBackup]) -> Option<&Path> { + backups.first()?.backup_path.parent() +} + #[cfg(test)] mod tests { use super::*; use pretty_assertions::assert_eq; + use std::path::PathBuf; use tempfile::TempDir; #[tokio::test] - async fn repair_backs_up_owned_database_files() -> std::io::Result<()> { + async fn backup_backs_up_only_failed_database_file() -> std::io::Result<()> { let temp_dir = TempDir::new()?; let state_path = codex_state::state_db_path(temp_dir.path()); - let logs_path = codex_state::logs_db_path(temp_dir.path()); - let goals_path = codex_state::goals_db_path(temp_dir.path()); - let state_sidecars = sqlite_paths(state_path.as_path()); + let failed_db_path = codex_state::logs_db_path(temp_dir.path()); tokio::fs::write(state_path.as_path(), b"state").await?; - tokio::fs::write(state_sidecars[1].as_path(), b"state-wal").await?; - tokio::fs::write(logs_path.as_path(), b"logs").await?; - tokio::fs::write(goals_path.as_path(), b"goals").await?; + tokio::fs::write(failed_db_path.as_path(), b"logs").await?; let startup_error = - LocalStateDbStartupError::new(state_path.clone(), "corrupt".to_string()); - let backups = repair_files(&startup_error).await?; + LocalStateDbStartupError::new(failed_db_path.clone(), "corrupt".to_string()); + let backups = backup_files_for_fresh_start(&startup_error).await?; - assert_eq!(backups.len(), 4); - assert!(!tokio::fs::try_exists(state_path.as_path()).await?); - assert!(!tokio::fs::try_exists(state_sidecars[1].as_path()).await?); - assert!(!tokio::fs::try_exists(logs_path.as_path()).await?); - assert!(!tokio::fs::try_exists(goals_path.as_path()).await?); - for backup in backups { - assert!(tokio::fs::try_exists(backup.as_path()).await?); - } + assert_eq!( + backups + .iter() + .map(|backup| &backup.original_path) + .collect::>(), + vec![&failed_db_path] + ); + assert!(!tokio::fs::try_exists(failed_db_path.as_path()).await?); + assert!(tokio::fs::try_exists(state_path.as_path()).await?); + assert!(tokio::fs::try_exists(backups[0].backup_path.as_path()).await?); Ok(()) } #[tokio::test] - async fn repair_replaces_blocking_sqlite_home_file() -> std::io::Result<()> { + async fn backup_replaces_blocking_sqlite_home_file() -> std::io::Result<()> { let temp_dir = TempDir::new()?; let sqlite_home = temp_dir.path().join("sqlite-home"); tokio::fs::write(sqlite_home.as_path(), b"not-a-directory").await?; @@ -168,18 +135,25 @@ mod tests { "File exists".to_string(), ); - let backups = repair_files(&startup_error).await?; + assert!(is_auto_backup_recoverable(&startup_error)); + let backups = backup_files_for_fresh_start(&startup_error).await?; assert_eq!(backups.len(), 1); assert!(tokio::fs::metadata(sqlite_home.as_path()).await?.is_dir()); - assert!(tokio::fs::try_exists(backups[0].as_path()).await?); + assert!(tokio::fs::try_exists(backups[0].backup_path.as_path()).await?); Ok(()) } #[test] - fn lock_failures_skip_repair() { - assert!(is_locked("database is locked")); - assert!(is_locked("database is busy")); - assert!(!is_locked("database disk image is malformed")); + fn backup_folder_uses_parent_of_first_backup_path() { + let backups = vec![RuntimeDbBackup { + original_path: PathBuf::from("/tmp/state_5.sqlite"), + backup_path: PathBuf::from("/tmp/db-backups/sqlite-1-0/state_5.sqlite"), + }]; + + assert_eq!( + backup_folder(&backups), + Some(Path::new("/tmp/db-backups/sqlite-1-0")) + ); } } diff --git a/codex-rs/cli/tests/delete.rs b/codex-rs/cli/tests/delete.rs new file mode 100644 index 0000000000..854b87b56e --- /dev/null +++ b/codex-rs/cli/tests/delete.rs @@ -0,0 +1,17 @@ +use predicates::prelude::*; + +#[test] +fn missing_session_fails_before_delete_confirmation() -> anyhow::Result<()> { + let codex_home = tempfile::tempdir()?; + let mut cmd = assert_cmd::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?); + cmd.env("CODEX_HOME", codex_home.path()) + .args(["delete", "123e4567-e89b-12d3-a456-426614174000"]); + + cmd.assert() + .failure() + .stderr(predicate::str::contains( + "No active or archived session found matching", + )) + .stderr(predicate::str::contains("cannot confirm").not()); + Ok(()) +} diff --git a/codex-rs/cloud-config/src/cache.rs b/codex-rs/cloud-config/src/cache.rs index 1e68646159..e5d1d4dc84 100644 --- a/codex-rs/cloud-config/src/cache.rs +++ b/codex-rs/cloud-config/src/cache.rs @@ -22,7 +22,7 @@ use tokio::fs; const CLOUD_CONFIG_BUNDLE_CACHE_VERSION: u32 = 1; pub(super) const CLOUD_CONFIG_BUNDLE_CACHE_FILENAME: &str = "cloud-config-bundle-cache.json"; -const CLOUD_CONFIG_BUNDLE_CACHE_TTL: Duration = Duration::from_secs(30 * 60); +const CLOUD_CONFIG_BUNDLE_CACHE_TTL: Duration = Duration::from_secs(60 * 60); const CLOUD_CONFIG_BUNDLE_CACHE_WRITE_HMAC_KEY: &[u8] = b"codex-cloud-config-bundle-cache-v1-6160ae70-bcfd-4ca8-a99b-40f73b3b072e"; const CLOUD_CONFIG_BUNDLE_CACHE_READ_HMAC_KEYS: &[&[u8]] = diff --git a/codex-rs/cloud-config/src/cache_tests.rs b/codex-rs/cloud-config/src/cache_tests.rs index 0416ed7e77..28899f930d 100644 --- a/codex-rs/cloud-config/src/cache_tests.rs +++ b/codex-rs/cloud-config/src/cache_tests.rs @@ -81,7 +81,7 @@ async fn save_writes_signed_payload_and_loads_for_matching_identity() { .expect("parse cache"); assert!( cache_file.signed_payload.expires_at - <= cache_file.signed_payload.cached_at + ChronoDuration::minutes(30) + <= cache_file.signed_payload.cached_at + ChronoDuration::minutes(60) ); assert!(cache_file.signed_payload.expires_at > cache_file.signed_payload.cached_at); assert_eq!( diff --git a/codex-rs/cloud-config/src/service.rs b/codex-rs/cloud-config/src/service.rs index eed94c0a10..1a7104bf52 100644 --- a/codex-rs/cloud-config/src/service.rs +++ b/codex-rs/cloud-config/src/service.rs @@ -32,7 +32,7 @@ use tokio::time::timeout; pub(crate) const CLOUD_CONFIG_BUNDLE_TIMEOUT: Duration = Duration::from_secs(15); const CLOUD_CONFIG_BUNDLE_MAX_ATTEMPTS: usize = 5; -const CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_INTERVAL: Duration = Duration::from_secs(5 * 60); +const CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_INTERVAL: Duration = Duration::from_secs(15 * 60); const CLOUD_CONFIG_BUNDLE_LOAD_FAILED_MESSAGE: &str = "Failed to load cloud config bundle (workspace-managed policies)."; const CLOUD_CONFIG_BUNDLE_AUTH_RECOVERY_FAILED_MESSAGE: &str = concat!( diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs index ad26549a26..d9a67c7b97 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs @@ -291,6 +291,18 @@ impl RealtimeWebsocketWriter { .await } + pub async fn send_conversation_handoff_append( + &self, + handoff_id: String, + output_text: String, + ) -> Result<(), ApiError> { + self.send_json(&RealtimeOutboundMessage::ConversationHandoffAppend { + handoff_id, + output_text, + }) + .await + } + pub async fn send_conversation_function_call_output( &self, call_id: String, diff --git a/codex-rs/codex-api/tests/models_integration.rs b/codex-rs/codex-api/tests/models_integration.rs index 4ee6069898..5f421dc1ad 100644 --- a/codex-rs/codex-api/tests/models_integration.rs +++ b/codex-rs/codex-api/tests/models_integration.rs @@ -93,6 +93,7 @@ async fn models_client_hits_models_endpoint() { context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), input_modalities: default_input_modalities(), diff --git a/codex-rs/codex-mcp/Cargo.toml b/codex-rs/codex-mcp/Cargo.toml index d7634f0aae..433cbde8b8 100644 --- a/codex-rs/codex-mcp/Cargo.toml +++ b/codex-rs/codex-mcp/Cargo.toml @@ -14,6 +14,7 @@ workspace = true [dependencies] anyhow = { workspace = true } +arc-swap = { workspace = true } async-channel = { workspace = true } codex-async-utils = { workspace = true } codex-api = { workspace = true } diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index d15a2e1500..9d1be43d42 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -68,6 +68,7 @@ use serde_json::Value as JsonValue; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use tracing::Instrument; +use tracing::info_span; use tracing::instrument; use tracing::trace; use tracing::trace_span; @@ -105,6 +106,7 @@ pub fn tool_is_model_visible(tool: &ToolInfo) -> bool { pub struct McpConnectionManager { clients: HashMap, server_metadata: HashMap, + required_servers: Vec, tool_plugin_provenance: Arc, host_owned_codex_apps_enabled: bool, prefix_mcp_tool_names: bool, @@ -113,103 +115,7 @@ pub struct McpConnectionManager { } impl McpConnectionManager { - pub fn new_uninitialized( - approval_policy: &Constrained, - permission_profile: &Constrained, - prefix_mcp_tool_names: bool, - ) -> Self { - Self::new_uninitialized_with_permission_profile( - approval_policy, - permission_profile.get(), - prefix_mcp_tool_names, - ) - } - - pub fn new_uninitialized_with_permission_profile( - approval_policy: &Constrained, - permission_profile: &PermissionProfile, - prefix_mcp_tool_names: bool, - ) -> Self { - Self { - clients: HashMap::new(), - server_metadata: HashMap::new(), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - host_owned_codex_apps_enabled: false, - prefix_mcp_tool_names, - elicitation_requests: ElicitationRequestManager::new( - approval_policy.value(), - permission_profile.clone(), - /*reviewer*/ None, - ), - startup_cancellation_token: CancellationToken::new(), - } - } - - pub fn has_servers(&self) -> bool { - !self.clients.is_empty() - } - - /// Drain all MCP clients from this manager and return a future that stops - /// them and terminates their stdio server processes. - pub fn begin_shutdown(&mut self) -> impl std::future::Future + Send + 'static { - self.startup_cancellation_token.cancel(); - let clients = std::mem::take(&mut self.clients); - self.server_metadata.clear(); - async move { - for client in clients.into_values() { - client.shutdown().await; - } - } - } - - /// Stop all MCP clients owned by this manager and terminate stdio server processes. - pub async fn shutdown(&mut self) { - self.begin_shutdown().await; - } - - pub fn server_origin(&self, server_name: &str) -> Option<&str> { - self.server_metadata - .get(server_name) - .and_then(|metadata| metadata.origin.as_ref()) - .map(super::server::McpServerOrigin::as_str) - } - - pub fn server_pollutes_memory(&self, server_name: &str) -> bool { - self.server_metadata - .get(server_name) - .is_none_or(|metadata| metadata.pollutes_memory) - } - - pub fn plugin_id_for_mcp_server_name(&self, server_name: &str) -> Option<&str> { - self.tool_plugin_provenance - .plugin_id_for_mcp_server_name(server_name) - } - - pub fn is_host_owned_codex_apps_server(&self, server_name: &str) -> bool { - self.host_owned_codex_apps_enabled && server_name == CODEX_APPS_MCP_SERVER_NAME - } - - pub fn set_approval_policy(&self, approval_policy: &Constrained) { - if let Ok(mut policy) = self.elicitation_requests.approval_policy.lock() { - *policy = approval_policy.value(); - } - } - - pub fn set_permission_profile(&self, permission_profile: PermissionProfile) { - if let Ok(mut profile) = self.elicitation_requests.permission_profile.lock() { - *profile = permission_profile; - } - } - - pub fn elicitations_auto_deny(&self) -> bool { - self.elicitation_requests.auto_deny() - } - - pub fn set_elicitations_auto_deny(&self, auto_deny: bool) { - self.elicitation_requests.set_auto_deny(auto_deny); - } - - #[allow(clippy::new_ret_no_self, clippy::too_many_arguments)] + #[allow(clippy::too_many_arguments)] pub async fn new( mcp_servers: &HashMap, store_mode: OAuthCredentialsStoreMode, @@ -217,6 +123,7 @@ impl McpConnectionManager { approval_policy: &Constrained, submit_id: String, tx_event: Sender, + startup_cancellation_token: CancellationToken, initial_permission_profile: PermissionProfile, runtime_context: McpRuntimeContext, codex_home: PathBuf, @@ -227,8 +134,13 @@ impl McpConnectionManager { tool_plugin_provenance: ToolPluginProvenance, auth: Option<&CodexAuth>, elicitation_reviewer: Option, - ) -> (Self, CancellationToken) { - let cancel_token = CancellationToken::new(); + ) -> Self { + let mut required_servers = mcp_servers + .iter() + .filter(|(_, server)| server.enabled() && server.required()) + .map(|(name, _)| name.clone()) + .collect::>(); + required_servers.sort(); let mut clients = HashMap::new(); let mut server_metadata = HashMap::new(); let mut join_set = JoinSet::new(); @@ -248,7 +160,7 @@ impl McpConnectionManager { .filter(|(_, server)| server.enabled()) { server_metadata.insert(server_name.clone(), McpServerMetadata::from(&server)); - let cancel_token = cancel_token.child_token(); + let cancel_token = startup_cancellation_token.child_token(); let _ = emit_update( startup_submit_id.as_str(), &tx_event, @@ -333,11 +245,12 @@ impl McpConnectionManager { let manager = Self { clients, server_metadata, + required_servers, tool_plugin_provenance, host_owned_codex_apps_enabled, prefix_mcp_tool_names, elicitation_requests: elicitation_requests.clone(), - startup_cancellation_token: cancel_token.clone(), + startup_cancellation_token: startup_cancellation_token.clone(), }; tokio::spawn(async move { let outcomes = join_set.join_all().await; @@ -361,7 +274,132 @@ impl McpConnectionManager { }) .await; }); - (manager, cancel_token) + manager + } + + /// Waits for every required server and reports their startup failures together. + /// + /// Callers must make the manager reachable to request handlers before awaiting this method, + /// because server initialization may require client elicitation. + pub async fn validate_required_servers(&self) -> Result<()> { + let failures = async { + let mut failures = Vec::new(); + for server_name in &self.required_servers { + let Some(async_managed_client) = self.clients.get(server_name).cloned() else { + failures.push(McpStartupFailure { + server: server_name.clone(), + error: format!("required MCP server `{server_name}` was not initialized"), + }); + continue; + }; + + match async_managed_client.client().await { + Ok(_) => {} + Err(error) => failures.push(McpStartupFailure { + server: server_name.clone(), + error: startup_outcome_error_message(error), + }), + } + } + failures + } + .instrument(info_span!( + "session_init.required_mcp_wait", + otel.name = "session_init.required_mcp_wait", + session_init.required_mcp_server_count = self.required_servers.len(), + )) + .await; + if failures.is_empty() { + return Ok(()); + } + + let details = failures + .iter() + .map(|failure| format!("{}: {}", failure.server, failure.error)) + .collect::>() + .join("; "); + Err(anyhow!( + "required MCP servers failed to initialize: {details}" + )) + } + + pub fn new_uninitialized_with_permission_profile( + approval_policy: &Constrained, + permission_profile: &PermissionProfile, + prefix_mcp_tool_names: bool, + ) -> Self { + Self { + clients: HashMap::new(), + server_metadata: HashMap::new(), + required_servers: Vec::new(), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + host_owned_codex_apps_enabled: false, + prefix_mcp_tool_names, + elicitation_requests: ElicitationRequestManager::new( + approval_policy.value(), + permission_profile.clone(), + /*reviewer*/ None, + ), + startup_cancellation_token: CancellationToken::new(), + } + } + + pub fn has_servers(&self) -> bool { + !self.clients.is_empty() + } + + pub(crate) fn contains_server(&self, server_name: &str) -> bool { + self.clients.contains_key(server_name) + } + + /// Stop all MCP clients owned by this manager and terminate stdio server processes. + pub async fn shutdown(&self) { + self.startup_cancellation_token.cancel(); + for client in self.clients.values() { + client.shutdown().await; + } + } + + pub fn server_origin(&self, server_name: &str) -> Option<&str> { + self.server_metadata + .get(server_name) + .and_then(|metadata| metadata.origin.as_ref()) + .map(super::server::McpServerOrigin::as_str) + } + + pub fn server_pollutes_memory(&self, server_name: &str) -> bool { + self.server_metadata + .get(server_name) + .is_none_or(|metadata| metadata.pollutes_memory) + } + + pub fn plugin_id_for_mcp_server_name(&self, server_name: &str) -> Option<&str> { + self.tool_plugin_provenance + .plugin_id_for_mcp_server_name(server_name) + } + + pub fn is_host_owned_codex_apps_server(&self, server_name: &str) -> bool { + self.host_owned_codex_apps_enabled && server_name == CODEX_APPS_MCP_SERVER_NAME + } + + pub fn set_approval_policy(&self, approval_policy: &Constrained) { + if let Ok(mut policy) = self.elicitation_requests.approval_policy.lock() { + *policy = approval_policy.value(); + } + } + + pub fn set_permission_profile(&self, permission_profile: PermissionProfile) { + if let Ok(mut profile) = self.elicitation_requests.permission_profile.lock() { + *profile = permission_profile; + } + } + + pub fn elicitations_auto_deny(&self) -> bool { + self.elicitation_requests.auto_deny() + } + + pub fn set_elicitations_auto_deny(&self, auto_deny: bool) { + self.elicitation_requests.set_auto_deny(auto_deny); } pub async fn resolve_elicitation( @@ -386,31 +424,6 @@ impl McpConnectionManager { } } - pub async fn required_startup_failures( - &self, - required_servers: &[String], - ) -> Vec { - let mut failures = Vec::new(); - for server_name in required_servers { - let Some(async_managed_client) = self.clients.get(server_name).cloned() else { - failures.push(McpStartupFailure { - server: server_name.clone(), - error: format!("required MCP server `{server_name}` was not initialized"), - }); - continue; - }; - - match async_managed_client.client().await { - Ok(_) => {} - Err(error) => failures.push(McpStartupFailure { - server: server_name.clone(), - error: startup_outcome_error_message(error), - }), - } - } - failures - } - /// Returns all tools with model-visible names normalized. #[instrument(level = "trace", skip_all, fields(mcp_server_count = self.clients.len()))] pub async fn list_all_tools(&self) -> Vec { @@ -452,31 +465,6 @@ impl McpConnectionManager { normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names) } - /// Returns presentation metadata without waiting for uncached clients still initializing. - /// Cached values will be used if available and the server is still starting up. - pub async fn list_available_server_infos(&self) -> HashMap { - let mut server_infos = HashMap::new(); - for (server_name, client) in &self.clients { - if !client.startup_complete.load(Ordering::Acquire) { - if let Some(server_info) = client.cached_server_info.clone() { - server_infos.insert(server_name.clone(), server_info); - } - continue; - } - match client.client().await { - Ok(managed_client) => { - server_infos.insert(server_name.clone(), managed_client.server_info); - } - Err(_) => { - if let Some(server_info) = client.cached_server_info.clone() { - server_infos.insert(server_name.clone(), server_info); - } - } - } - } - server_infos - } - /// Force-refresh codex apps tools by bypassing the in-process cache. /// /// On success, the refreshed tools replace the cache contents and the @@ -532,21 +520,6 @@ impl McpConnectionManager { )) } - fn with_server_metadata(&self, mut tool: ToolInfo) -> ToolInfo { - let Some(metadata) = self.server_metadata.get(&tool.server_name) else { - tool.supports_parallel_tool_calls = false; - tool.server_origin = None; - return tool; - }; - - tool.supports_parallel_tool_calls = metadata.supports_parallel_tool_calls; - tool.server_origin = metadata - .origin - .as_ref() - .map(|origin| origin.as_str().to_string()); - tool - } - /// Returns a single map that contains all resources. Each key is the /// server name and the value is a vector of resources. pub async fn list_all_resources(&self) -> HashMap> { @@ -778,6 +751,46 @@ impl McpConnectionManager { .with_context(|| format!("resources/read failed for `{server}` ({uri})")) } + /// Returns presentation metadata without waiting for uncached clients still initializing. + /// Cached values will be used if available and the server is still starting up. + pub(crate) async fn list_available_server_infos(&self) -> HashMap { + let mut server_infos = HashMap::new(); + for (server_name, client) in &self.clients { + if !client.startup_complete.load(Ordering::Acquire) { + if let Some(server_info) = client.cached_server_info.clone() { + server_infos.insert(server_name.clone(), server_info); + } + continue; + } + match client.client().await { + Ok(managed_client) => { + server_infos.insert(server_name.clone(), managed_client.server_info); + } + Err(_) => { + if let Some(server_info) = client.cached_server_info.clone() { + server_infos.insert(server_name.clone(), server_info); + } + } + } + } + server_infos + } + + fn with_server_metadata(&self, mut tool: ToolInfo) -> ToolInfo { + let Some(metadata) = self.server_metadata.get(&tool.server_name) else { + tool.supports_parallel_tool_calls = false; + tool.server_origin = None; + return tool; + }; + + tool.supports_parallel_tool_calls = metadata.supports_parallel_tool_calls; + tool.server_origin = metadata + .origin + .as_ref() + .map(|origin| origin.as_str().to_string()); + tool + } + async fn client_by_name(&self, name: &str) -> Result { self.clients .get(name) @@ -786,6 +799,19 @@ impl McpConnectionManager { .await .context("failed to get client") } + + #[cfg(test)] + fn new_uninitialized( + approval_policy: &Constrained, + permission_profile: &Constrained, + prefix_mcp_tool_names: bool, + ) -> Self { + Self::new_uninitialized_with_permission_profile( + approval_policy, + permission_profile.get(), + prefix_mcp_tool_names, + ) + } } impl Drop for McpConnectionManager { diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 60a0026eeb..ea405ced9a 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -980,6 +980,48 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_ assert!(timeout_result.is_err()); } +#[tokio::test] +async fn shutdown_cancels_pending_tool_listing() { + let cancel_token = CancellationToken::new(); + let cancel_token_for_startup = cancel_token.clone(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let pending_client = async move { + let _ = started_tx.send(()); + cancel_token_for_startup.cancelled().await; + Err(StartupOutcomeError::Cancelled) + } + .boxed() + .shared(); + let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager.clients.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + AsyncManagedClient { + client: pending_client, + cached_tool_info_snapshot: None, + cached_server_info: None, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token, + }, + ); + let manager = Arc::new(manager); + let manager_for_list = Arc::clone(&manager); + let list_task = tokio::spawn(async move { manager_for_list.list_all_tools().await }); + + started_rx.await.expect("tool listing should start"); + tokio::time::timeout(Duration::from_secs(1), manager.shutdown()) + .await + .expect("shutdown should cancel speculative tool listing"); + let tools = list_task.await.expect("tool listing task should not panic"); + assert!(tools.is_empty()); +} + #[tokio::test] async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() { let pending_client = futures::future::pending::>() @@ -1167,13 +1209,15 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { ), ]); - let (manager, cancel_token) = McpConnectionManager::new( + let cancel_token = CancellationToken::new(); + let manager = McpConnectionManager::new( &mcp_servers, OAuthCredentialsStoreMode::default(), HashMap::new(), &approval_policy, String::new(), tx_event, + cancel_token.clone(), PermissionProfile::default(), McpRuntimeContext::new( Arc::new(EnvironmentManager::without_environments()), @@ -1201,13 +1245,18 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { .wait_for_server_ready("stdio", Duration::from_millis(10)) .await ); - let failures = manager - .required_startup_failures(&["stdio".to_string()]) - .await; - assert_eq!(failures.len(), 1); - assert_eq!(failures[0].server, "stdio"); + let error = match manager + .clients + .get("stdio") + .expect("stdio client") + .client() + .await + { + Ok(_) => panic!("local stdio MCP startup should fail"), + Err(error) => error, + }; assert_eq!( - failures[0].error, + startup_outcome_error_message(error), "local stdio MCP server `stdio` requires a local environment" ); cancel_token.cancel(); diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index a56365277a..195e8614e1 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -3,6 +3,9 @@ pub use connection_manager::tool_is_model_visible; pub use elicitation::ElicitationReviewRequest; pub use elicitation::ElicitationReviewer; pub use elicitation::ElicitationReviewerHandle; +pub use resource_client::McpResourceClient; +pub use resource_client::McpResourcePage; +pub use resource_client::McpResourceReadResult; pub use rmcp_client::MCP_SANDBOX_STATE_META_CAPABILITY; pub use runtime::McpRuntimeContext; pub use runtime::SandboxState; @@ -29,8 +32,8 @@ pub use mcp::configured_mcp_servers; pub use mcp::effective_mcp_servers; pub use mcp::effective_mcp_servers_from_configured; pub use mcp::host_owned_codex_apps_enabled; +pub use mcp::hosted_plugin_runtime_mcp_server_config; pub use mcp::tool_plugin_provenance; -pub use mcp::with_codex_apps_mcp; pub use mcp::McpServerStatusSnapshot; pub use mcp::McpSnapshotDetail; @@ -58,6 +61,7 @@ pub(crate) mod codex_apps; pub(crate) mod connection_manager; pub(crate) mod elicitation; pub(crate) mod mcp; +mod resource_client; pub(crate) mod rmcp_client; pub(crate) mod runtime; pub(crate) mod server; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index a74f95bae3..e6b17b14c3 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -35,6 +35,7 @@ use rmcp::model::ElicitationCapability; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; use serde_json::Value; +use tokio_util::sync::CancellationToken; use crate::codex_apps::codex_apps_tools_cache_key; use crate::connection_manager::McpConnectionManager; @@ -106,8 +107,6 @@ pub struct McpPermissionPromptAutoApproveContext { pub struct McpConfig { /// Base URL for ChatGPT-hosted app MCP servers, copied from the root config. pub chatgpt_base_url: String, - /// Optional path override for the host-owned apps MCP server. - pub apps_mcp_path_override: Option, /// Optional product SKU forwarded to the host-owned apps MCP server. pub apps_mcp_product_sku: Option, /// Codex home directory used for MCP OAuth state and app-tool cache files. @@ -128,22 +127,18 @@ pub struct McpConfig { pub use_legacy_landlock: bool, /// Whether the app MCP integration is enabled by config. /// - /// ChatGPT auth is checked separately at runtime before the host-owned apps - /// MCP server is added. + /// ChatGPT auth is checked separately before a materialized host-owned Apps + /// server can be used. pub apps_enabled: bool, - /// Whether to synthesize the legacy host-owned Apps MCP server. - /// - /// Hosts that install an MCP extension for this server disable the legacy - /// loader and contribute the server through the normal runtime overlay. - pub legacy_apps_mcp_loader_enabled: bool, /// Whether model-visible MCP tool namespaces should keep the legacy /// `mcp__` prefix. pub prefix_mcp_tool_names: bool, /// Client-side elicitation capabilities advertised during MCP initialization. pub client_elicitation_capability: ElicitationCapability, - /// Config-backed MCP servers keyed by server name. + /// Materialized MCP servers keyed by server name. /// - /// Runtime-only additions are merged later by [`effective_mcp_servers`]. + /// A host may add compatibility built-ins and extension overlays before + /// calling runtime entry points in this crate. pub configured_mcp_servers: HashMap, /// Winning plugin owner for plugin-provided MCP servers, keyed by server name. pub plugin_ids_by_mcp_server_name: HashMap, @@ -218,32 +213,6 @@ impl ToolPluginProvenance { } } -pub fn with_codex_apps_mcp( - mut servers: HashMap, - auth: Option<&CodexAuth>, - config: &McpConfig, -) -> HashMap { - if !config.legacy_apps_mcp_loader_enabled { - if !host_owned_codex_apps_enabled(config, auth) { - servers.remove(CODEX_APPS_MCP_SERVER_NAME); - } - return servers; - } - if host_owned_codex_apps_enabled(config, auth) { - servers.insert( - CODEX_APPS_MCP_SERVER_NAME.to_string(), - EffectiveMcpServer::configured(codex_apps_mcp_server_config( - &config.chatgpt_base_url, - config.apps_mcp_path_override.as_deref(), - config.apps_mcp_product_sku.as_deref(), - )), - ); - } else { - servers.remove(CODEX_APPS_MCP_SERVER_NAME); - } - servers -} - pub fn host_owned_codex_apps_enabled(config: &McpConfig, auth: Option<&CodexAuth>) -> bool { config.apps_enabled && auth.is_some_and(CodexAuth::uses_codex_backend) } @@ -259,16 +228,23 @@ pub fn effective_mcp_servers( effective_mcp_servers_from_configured(configured_mcp_servers(config), config, auth) } +/// Converts a materialized server map to its auth-gated runtime view. +/// +/// Compatibility built-ins and extension overlays must already be reflected in +/// `configured_servers`; this function does not synthesize missing servers. pub fn effective_mcp_servers_from_configured( configured_servers: HashMap, config: &McpConfig, auth: Option<&CodexAuth>, ) -> HashMap { - let servers = configured_servers + let mut servers = configured_servers .into_iter() .map(|(name, server)| (name, EffectiveMcpServer::configured(server))) .collect::>(); - with_codex_apps_mcp(servers, auth, config) + if !host_owned_codex_apps_enabled(config, auth) { + servers.remove(CODEX_APPS_MCP_SERVER_NAME); + } + servers } pub fn tool_plugin_provenance(config: &McpConfig) -> ToolPluginProvenance { @@ -293,13 +269,15 @@ pub async fn read_mcp_resource( .await; let (tx_event, rx_event) = unbounded(); drop(rx_event); - let (manager, cancel_token) = McpConnectionManager::new( + let cancel_token = CancellationToken::new(); + let manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, auth_statuses, &config.approval_policy, String::new(), tx_event, + cancel_token.clone(), PermissionProfile::default(), runtime_context, config.codex_home.clone(), @@ -363,13 +341,15 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( let (tx_event, rx_event) = unbounded(); drop(rx_event); - let (mcp_connection_manager, cancel_token) = McpConnectionManager::new( + let cancel_token = CancellationToken::new(); + let mcp_connection_manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, auth_status_entries.clone(), &config.approval_policy, submit_id, tx_event, + cancel_token.clone(), PermissionProfile::default(), runtime_context, config.codex_home.clone(), @@ -436,7 +416,7 @@ fn normalize_codex_apps_base_url(base_url: &str) -> String { base_url } -fn codex_apps_mcp_url_for_base_url(base_url: &str, apps_mcp_path_override: Option<&str>) -> String { +fn codex_apps_mcp_url_for_base_url(base_url: &str) -> String { let base_url = normalize_codex_apps_base_url(base_url); let (base_url, default_path) = if base_url.contains("/backend-api") { (base_url, "wham/apps") @@ -445,18 +425,34 @@ fn codex_apps_mcp_url_for_base_url(base_url: &str, apps_mcp_path_override: Optio } else { (format!("{base_url}/api/codex"), "apps") }; - let path = apps_mcp_path_override - .unwrap_or(default_path) - .trim_start_matches('/'); - format!("{base_url}/{path}") + format!("{base_url}/{default_path}") } pub fn codex_apps_mcp_server_config( chatgpt_base_url: &str, - apps_mcp_path_override: Option<&str>, apps_mcp_product_sku: Option<&str>, ) -> McpServerConfig { - let url = codex_apps_mcp_url_for_base_url(chatgpt_base_url, apps_mcp_path_override); + mcp_server_config_for_url( + codex_apps_mcp_url_for_base_url(chatgpt_base_url), + apps_mcp_product_sku, + ) +} + +/// Builds the ChatGPT-hosted plugin runtime served by plugin-service. +pub fn hosted_plugin_runtime_mcp_server_config( + chatgpt_base_url: &str, + apps_mcp_product_sku: Option<&str>, +) -> McpServerConfig { + let base_url = normalize_codex_apps_base_url(chatgpt_base_url); + let base_url = if base_url.contains("/backend-api") || base_url.contains("/api/codex") { + base_url + } else { + format!("{base_url}/api/codex") + }; + mcp_server_config_for_url(format!("{base_url}/ps/mcp"), apps_mcp_product_sku) +} + +fn mcp_server_config_for_url(url: String, apps_mcp_product_sku: Option<&str>) -> McpServerConfig { let http_headers = apps_mcp_product_sku.map(|product_sku| { HashMap::from([("X-OpenAI-Product-Sku".to_string(), product_sku.to_string())]) }); diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index fb1af362ca..767f20611f 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -16,7 +16,6 @@ use std::path::PathBuf; fn test_mcp_config(codex_home: PathBuf) -> McpConfig { McpConfig { chatgpt_base_url: "https://chatgpt.com".to_string(), - apps_mcp_path_override: None, apps_mcp_product_sku: None, codex_home, mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode::default(), @@ -27,7 +26,6 @@ fn test_mcp_config(codex_home: PathBuf) -> McpConfig { codex_linux_sandbox_exe: None, use_legacy_landlock: false, apps_enabled: false, - legacy_apps_mcp_loader_enabled: true, prefix_mcp_tool_names: true, client_elicitation_capability: ElicitationCapability::default(), configured_mcp_servers: HashMap::new(), @@ -178,52 +176,27 @@ fn tool_plugin_provenance_collects_app_and_mcp_sources() { #[test] fn codex_apps_mcp_url_for_base_url_keeps_existing_paths() { assert_eq!( - codex_apps_mcp_url_for_base_url( - "https://chatgpt.com/backend-api", - /*apps_mcp_path_override*/ None, - ), + codex_apps_mcp_url_for_base_url("https://chatgpt.com/backend-api"), "https://chatgpt.com/backend-api/wham/apps" ); assert_eq!( - codex_apps_mcp_url_for_base_url( - "https://chat.openai.com", - /*apps_mcp_path_override*/ None, - ), + codex_apps_mcp_url_for_base_url("https://chat.openai.com"), "https://chat.openai.com/backend-api/wham/apps" ); assert_eq!( - codex_apps_mcp_url_for_base_url( - "http://localhost:8080/api/codex", - /*apps_mcp_path_override*/ None, - ), + codex_apps_mcp_url_for_base_url("http://localhost:8080/api/codex"), "http://localhost:8080/api/codex/apps" ); assert_eq!( - codex_apps_mcp_url_for_base_url( - "http://localhost:8080", - /*apps_mcp_path_override*/ None, - ), + codex_apps_mcp_url_for_base_url("http://localhost:8080"), "http://localhost:8080/api/codex/apps" ); } #[test] fn codex_apps_server_config_uses_legacy_codex_apps_path() { - let mut config = test_mcp_config(PathBuf::from("/tmp")); - let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - - let mut servers = with_codex_apps_mcp(HashMap::new(), /*auth*/ None, &config); - assert!(!servers.contains_key(CODEX_APPS_MCP_SERVER_NAME)); - - config.apps_enabled = true; - - servers = with_codex_apps_mcp(servers, Some(&auth), &config); - let server = servers - .get(CODEX_APPS_MCP_SERVER_NAME) - .expect("codex apps should be present when apps is enabled"); - let config = server - .configured_config() - .expect("codex apps should use configured transport"); + let config = + codex_apps_mcp_server_config("https://chatgpt.com", /*apps_mcp_product_sku*/ None); let url = match &config.transport { McpServerTransportConfig::StreamableHttp { url, .. } => url, _ => panic!("expected streamable http transport for codex apps"), @@ -232,42 +205,9 @@ fn codex_apps_server_config_uses_legacy_codex_apps_path() { assert_eq!(url, "https://chatgpt.com/backend-api/wham/apps"); } -#[test] -fn codex_apps_server_config_uses_configured_apps_mcp_path_override() { - let mut config = test_mcp_config(PathBuf::from("/tmp")); - config.apps_mcp_path_override = Some("/custom/mcp".to_string()); - config.apps_enabled = true; - let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - - let servers = with_codex_apps_mcp(HashMap::new(), Some(&auth), &config); - let server = servers - .get(CODEX_APPS_MCP_SERVER_NAME) - .expect("codex apps should be present when apps is enabled"); - let config = server - .configured_config() - .expect("codex apps should use configured transport"); - let url = match &config.transport { - McpServerTransportConfig::StreamableHttp { url, .. } => url, - _ => panic!("expected streamable http transport for codex apps"), - }; - - assert_eq!(url, "https://chatgpt.com/backend-api/custom/mcp"); -} - #[test] fn codex_apps_server_config_forwards_configured_product_sku_header() { - let mut config = test_mcp_config(PathBuf::from("/tmp")); - config.apps_mcp_product_sku = Some("tpp".to_string()); - config.apps_enabled = true; - let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - - let servers = with_codex_apps_mcp(HashMap::new(), Some(&auth), &config); - let server = servers - .get(CODEX_APPS_MCP_SERVER_NAME) - .expect("codex apps should be present when apps is enabled"); - let config = server - .configured_config() - .expect("codex apps should use configured transport"); + let config = codex_apps_mcp_server_config("https://chatgpt.com", Some("tpp")); match &config.transport { McpServerTransportConfig::StreamableHttp { @@ -289,7 +229,7 @@ fn codex_apps_server_config_forwards_configured_product_sku_header() { } #[tokio::test] -async fn effective_mcp_servers_preserve_user_servers_and_add_codex_apps() { +async fn effective_mcp_servers_preserve_runtime_servers() { let codex_home = tempfile::tempdir().expect("tempdir"); let mut config = test_mcp_config(codex_home.path().to_path_buf()); config.apps_enabled = true; @@ -345,6 +285,13 @@ async fn effective_mcp_servers_preserve_user_servers_and_add_codex_apps() { tools: HashMap::new(), }, ); + config.configured_mcp_servers.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + codex_apps_mcp_server_config( + &config.chatgpt_base_url, + config.apps_mcp_product_sku.as_deref(), + ), + ); let effective = effective_mcp_servers(&config, Some(&auth)); diff --git a/codex-rs/codex-mcp/src/resource_client.rs b/codex-rs/codex-mcp/src/resource_client.rs new file mode 100644 index 0000000000..85dc974f20 --- /dev/null +++ b/codex-rs/codex-mcp/src/resource_client.rs @@ -0,0 +1,108 @@ +use std::sync::Arc; + +use anyhow::Context; +use anyhow::Result; +use arc_swap::ArcSwap; +use codex_protocol::mcp::Resource; +use codex_protocol::mcp::ResourceContent; +use rmcp::model::PaginatedRequestParams; +use rmcp::model::ReadResourceRequestParams; + +use crate::McpConnectionManager; + +/// One page of resources returned by an MCP server. +#[derive(Clone, Debug, PartialEq)] +pub struct McpResourcePage { + /// Resources advertised on this page. + pub resources: Vec, + /// Opaque cursor to supply when requesting the next page. + pub next_cursor: Option, +} + +/// Contents returned after reading one MCP resource. +#[derive(Clone, Debug, PartialEq)] +pub struct McpResourceReadResult { + /// Text or blob content returned for the requested resource. + pub contents: Vec, +} + +/// Session-scoped access to MCP resources through the currently installed manager. +/// +/// The client retains the manager's shared publication handle rather than a manager +/// snapshot, so calls automatically use replacements installed during startup and refresh. +#[derive(Clone)] +pub struct McpResourceClient { + manager: Arc>, +} + +impl std::fmt::Debug for McpResourceClient { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("McpResourceClient") + .finish_non_exhaustive() + } +} + +impl McpResourceClient { + /// Creates a resource client backed by the session's replaceable MCP manager. + pub fn new(manager: Arc>) -> Self { + Self { manager } + } + + /// Returns whether the current manager contains the named server. + /// + /// This does not wait for server startup or imply that startup succeeded. + pub async fn has_server(&self, server: &str) -> bool { + self.manager.load_full().contains_server(server) + } + + /// Lists one resource page from the named server. + pub async fn list_resources( + &self, + server: &str, + cursor: Option, + ) -> Result { + let params = + cursor.map(|cursor| PaginatedRequestParams::default().with_cursor(Some(cursor))); + let result = self + .manager + .load_full() + .list_resources(server, params) + .await?; + let resources = result + .resources + .into_iter() + .map(resource_from_rmcp) + .collect::>>()?; + Ok(McpResourcePage { + resources, + next_cursor: result.next_cursor, + }) + } + + /// Reads one resource from the named server. + pub async fn read_resource(&self, server: &str, uri: &str) -> Result { + let result = self + .manager + .load_full() + .read_resource(server, ReadResourceRequestParams::new(uri.to_string())) + .await?; + let contents = result + .contents + .into_iter() + .map(resource_content_from_rmcp) + .collect::>>()?; + Ok(McpResourceReadResult { contents }) + } +} + +fn resource_from_rmcp(resource: rmcp::model::Resource) -> Result { + let value = serde_json::to_value(resource).context("failed to serialize MCP resource")?; + Resource::from_mcp_value(value).context("failed to convert MCP resource") +} + +fn resource_content_from_rmcp(content: rmcp::model::ResourceContents) -> Result { + let value = + serde_json::to_value(content).context("failed to serialize MCP resource content")?; + serde_json::from_value(value).context("failed to convert MCP resource content") +} diff --git a/codex-rs/config/src/schema.rs b/codex-rs/config/src/schema.rs index 6b012494d9..c641f1703d 100644 --- a/codex-rs/config/src/schema.rs +++ b/codex-rs/config/src/schema.rs @@ -9,6 +9,7 @@ use schemars::schema::ObjectValidation; use schemars::schema::RootSchema; use schemars::schema::Schema; use schemars::schema::SchemaObject; +use schemars::schema::SubschemaValidation; use serde_json::Map; use serde_json::Value; use std::path::Path; @@ -46,9 +47,7 @@ pub fn features_schema(schema_gen: &mut SchemaGenerator) -> Schema { if feature.id == codex_features::Feature::AppsMcpPathOverride { validation.properties.insert( feature.key.to_string(), - schema_gen.subschema_for::>(), + removed_apps_mcp_path_override_schema(schema_gen), ); continue; } @@ -76,6 +75,30 @@ pub fn features_schema(schema_gen: &mut SchemaGenerator) -> Schema { Schema::Object(object) } +fn removed_apps_mcp_path_override_schema(schema_gen: &mut SchemaGenerator) -> Schema { + let mut config_validation = ObjectValidation::default(); + config_validation + .properties + .insert("enabled".to_string(), schema_gen.subschema_for::()); + config_validation + .properties + .insert("path".to_string(), schema_gen.subschema_for::()); + config_validation.additional_properties = Some(Box::new(Schema::Bool(false))); + + let config = Schema::Object(SchemaObject { + instance_type: Some(InstanceType::Object.into()), + object: Some(Box::new(config_validation)), + ..Default::default() + }); + Schema::Object(SchemaObject { + subschemas: Some(Box::new(SubschemaValidation { + any_of: Some(vec![schema_gen.subschema_for::(), config]), + ..Default::default() + })), + ..Default::default() + }) +} + /// Schema for the `[mcp_servers]` map using the raw input shape. pub fn mcp_servers_schema(schema_gen: &mut SchemaGenerator) -> Schema { let mut object = SchemaObject { diff --git a/codex-rs/core-plugins/src/discoverable.rs b/codex-rs/core-plugins/src/discoverable.rs index 28a90c1928..a630877f55 100644 --- a/codex-rs/core-plugins/src/discoverable.rs +++ b/codex-rs/core-plugins/src/discoverable.rs @@ -68,6 +68,7 @@ pub struct ToolSuggestPluginDiscoveryInput { #[derive(Clone, Debug, PartialEq, Eq)] pub struct ToolSuggestDiscoverablePlugin { pub id: String, + pub remote_plugin_id: Option, pub name: String, pub description: Option, pub has_skills: bool, @@ -86,7 +87,11 @@ impl PluginsManager { } let marketplaces = self - .list_marketplaces_for_config(&input.plugins, &[]) + .list_marketplaces_for_config( + &input.plugins, + &[], + /*include_openai_curated*/ !input.plugins.remote_plugin_enabled, + ) .context("failed to list plugin marketplaces for tool suggestions")? .marketplaces; let mut installed_app_connector_ids = self @@ -109,11 +114,6 @@ impl PluginsManager { let mut discoverable_plugins = Vec::::new(); for marketplace in marketplaces { let marketplace_name = marketplace.name; - if input.plugins.remote_plugin_enabled - && marketplace_name == OPENAI_CURATED_MARKETPLACE_NAME - { - continue; - } let use_legacy_local_curated_filter = should_use_legacy_local_curated_discovery_filter( &marketplace_name, marketplace.path.as_path(), @@ -162,6 +162,7 @@ impl PluginsManager { discoverable_plugins.push(ToolSuggestDiscoverablePlugin { id: plugin.config_name, + remote_plugin_id: None, name: plugin.display_name, description: plugin.description, has_skills: plugin.has_skills, @@ -217,6 +218,7 @@ impl PluginsManager { discoverable_plugins.push(ToolSuggestDiscoverablePlugin { id: plugin.config_id, + remote_plugin_id: Some(plugin.remote_plugin_id), name: plugin.name, description: plugin.description, has_skills: plugin.has_skills, diff --git a/codex-rs/core-plugins/src/manager.rs b/codex-rs/core-plugins/src/manager.rs index 095c5fd69d..1c5c8d78f0 100644 --- a/codex-rs/core-plugins/src/manager.rs +++ b/codex-rs/core-plugins/src/manager.rs @@ -48,6 +48,7 @@ use crate::store::PluginInstallResult as StorePluginInstallResult; use crate::store::PluginStore; use crate::store::PluginStoreError; use codex_analytics::AnalyticsEventsClient; +use codex_app_server_protocol::AuthMode; use codex_config::ConfigLayerStack; use codex_config::clear_user_plugin; use codex_config::set_user_plugin_enabled; @@ -203,6 +204,13 @@ fn featured_plugin_ids_cache_key( } } +fn project_plugin_load_outcome_for_auth( + outcome: PluginLoadOutcome, + _auth_mode: Option, +) -> PluginLoadOutcome { + outcome +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct PluginInstallRequest { pub plugin_name: String, @@ -319,6 +327,7 @@ pub struct PluginsManager { remote_installed_plugins_cache_refresh_state: RwLock, global_remote_catalog_cache_refresh_state: RwLock, restriction_product: Option, + auth_mode: RwLock>, analytics_events_client: RwLock>, } @@ -375,10 +384,30 @@ impl PluginsManager { GlobalRemoteCatalogCacheRefreshState::default(), ), restriction_product, + auth_mode: RwLock::new(None), analytics_events_client: RwLock::new(None), } } + pub fn set_auth_mode(&self, auth_mode: Option) -> bool { + let mut stored_auth_mode = match self.auth_mode.write() { + Ok(auth_mode_guard) => auth_mode_guard, + Err(err) => err.into_inner(), + }; + if *stored_auth_mode == auth_mode { + return false; + } + *stored_auth_mode = auth_mode; + true + } + + pub fn auth_mode(&self) -> Option { + match self.auth_mode.read() { + Ok(auth_mode_guard) => *auth_mode_guard, + Err(err) => *err.into_inner(), + } + } + pub fn set_analytics_events_client(&self, analytics_events_client: AnalyticsEventsClient) { let mut stored_client = match self.analytics_events_client.write() { Ok(client_guard) => client_guard, @@ -417,7 +446,7 @@ impl PluginsManager { remote_plugin_enabled: config.remote_plugin_enabled, }; if !force_reload && let Some(outcome) = self.cached_enabled_outcome(&cache_key) { - return outcome; + return self.project_plugins_for_auth(outcome); } let Ok(_load_permit) = self.enabled_outcome_load_semaphore.acquire().await else { @@ -425,7 +454,7 @@ impl PluginsManager { return PluginLoadOutcome::default(); }; if !force_reload && let Some(outcome) = self.cached_enabled_outcome(&cache_key) { - return outcome; + return self.project_plugins_for_auth(outcome); } let cache_generation = self.enabled_outcome_cache_generation(); let outcome = load_plugins_from_layer_stack( @@ -438,7 +467,11 @@ impl PluginsManager { .await; log_plugin_load_errors(&outcome); self.cache_enabled_outcome_if_current(cache_generation, cache_key, outcome.clone()); - outcome + self.project_plugins_for_auth(outcome) + } + + fn project_plugins_for_auth(&self, outcome: PluginLoadOutcome) -> PluginLoadOutcome { + project_plugin_load_outcome_for_auth(outcome, self.auth_mode()) } pub fn clear_cache(&self) { @@ -468,14 +501,15 @@ impl PluginsManager { if !config.plugins_enabled { return PluginLoadOutcome::default(); } - load_plugins_from_layer_stack( + let outcome = load_plugins_from_layer_stack( config_layer_stack, self.remote_installed_plugin_configs(), &self.store, self.restriction_product, config.remote_plugin_enabled, ) - .await + .await; + self.project_plugins_for_auth(outcome) } /// Resolve plugin hooks for a config layer stack without loading other plugin capabilities. @@ -992,14 +1026,19 @@ impl PluginsManager { &self, config: &PluginsConfigInput, additional_roots: &[AbsolutePathBuf], + include_openai_curated: bool, ) -> Result { if !config.plugins_enabled { return Ok(ConfiguredMarketplaceListOutcome::default()); } let (installed_plugins, enabled_plugins) = self.configured_plugin_states(config); - let marketplace_outcome = - self.discover_marketplaces_for_config(config, additional_roots)?; + let mut marketplace_roots = self.marketplace_roots(config, additional_roots); + if !include_openai_curated { + let curated_repo_root = curated_plugins_repo_path(self.codex_home.as_path()); + marketplace_roots.retain(|root| root.as_path() != curated_repo_root.as_path()); + } + let marketplace_outcome = list_marketplaces(&marketplace_roots)?; let mut seen_plugin_keys = HashSet::new(); let marketplaces = marketplace_outcome .marketplaces diff --git a/codex-rs/core-plugins/src/manager_tests.rs b/codex-rs/core-plugins/src/manager_tests.rs index c3961796cf..1a268a4888 100644 --- a/codex-rs/core-plugins/src/manager_tests.rs +++ b/codex-rs/core-plugins/src/manager_tests.rs @@ -17,6 +17,7 @@ use crate::test_support::load_plugins_config as load_plugins_config_input; use crate::test_support::write_curated_plugin_sha_with as write_curated_plugin_sha; use crate::test_support::write_file; use crate::test_support::write_openai_curated_marketplace; +use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::ConfigLayerSource; use codex_config::AppToolApproval; use codex_config::CONFIG_TOML_FILE; @@ -47,6 +48,21 @@ use wiremock::matchers::query_param; const MAX_CAPABILITY_SUMMARY_DESCRIPTION_LEN: usize = 1024; +#[test] +fn plugins_manager_tracks_auth_mode() { + let tmp = TempDir::new().unwrap(); + let manager = PluginsManager::new(tmp.path().to_path_buf()); + + assert_eq!(manager.auth_mode(), None); + assert!(manager.set_auth_mode(Some(AuthMode::ApiKey))); + assert_eq!(manager.auth_mode(), Some(AuthMode::ApiKey)); + assert!(!manager.set_auth_mode(Some(AuthMode::ApiKey))); + assert!(manager.set_auth_mode(Some(AuthMode::ChatgptAuthTokens))); + assert_eq!(manager.auth_mode(), Some(AuthMode::ChatgptAuthTokens)); + assert!(manager.set_auth_mode(/*auth_mode*/ None)); + assert_eq!(manager.auth_mode(), None); +} + fn write_plugin_with_version( root: &Path, dir_name: &str, @@ -1815,7 +1831,11 @@ enabled = false let config = load_config(tmp.path(), &repo_root).await; let marketplaces = PluginsManager::new(tmp.path().to_path_buf()) - .list_marketplaces_for_config(&config, &[AbsolutePathBuf::try_from(repo_root).unwrap()]) + .list_marketplaces_for_config( + &config, + &[AbsolutePathBuf::try_from(repo_root).unwrap()], + /*include_openai_curated*/ true, + ) .unwrap() .marketplaces; @@ -1917,7 +1937,11 @@ enabled = true let config = load_config(tmp.path(), &repo_root).await; let marketplaces = PluginsManager::new(tmp.path().to_path_buf()) - .list_marketplaces_for_config(&config, &[AbsolutePathBuf::try_from(repo_root).unwrap()]) + .list_marketplaces_for_config( + &config, + &[AbsolutePathBuf::try_from(repo_root).unwrap()], + /*include_openai_curated*/ true, + ) .unwrap() .marketplaces; @@ -1965,7 +1989,11 @@ plugins = true let config = load_config(tmp.path(), &repo_root).await; let marketplaces = PluginsManager::new(tmp.path().to_path_buf()) - .list_marketplaces_for_config(&config, &[AbsolutePathBuf::try_from(repo_root).unwrap()]) + .list_marketplaces_for_config( + &config, + &[AbsolutePathBuf::try_from(repo_root).unwrap()], + /*include_openai_curated*/ true, + ) .unwrap() .marketplaces; @@ -2402,7 +2430,11 @@ enabled = true let config = load_config(tmp.path(), &repo_root).await; let marketplaces = PluginsManager::new(tmp.path().to_path_buf()) - .list_marketplaces_for_config(&config, &[AbsolutePathBuf::try_from(repo_root).unwrap()]) + .list_marketplaces_for_config( + &config, + &[AbsolutePathBuf::try_from(repo_root).unwrap()], + /*include_openai_curated*/ true, + ) .unwrap() .marketplaces; @@ -2496,7 +2528,7 @@ plugins = true let config = load_config(tmp.path(), tmp.path()).await; let marketplaces = PluginsManager::new(tmp.path().to_path_buf()) - .list_marketplaces_for_config(&config, &[]) + .list_marketplaces_for_config(&config, &[], /*include_openai_curated*/ true) .unwrap() .marketplaces; @@ -2534,6 +2566,37 @@ plugins = true ); } +#[tokio::test] +async fn list_marketplaces_can_skip_openai_curated_before_loading() { + let tmp = tempfile::tempdir().unwrap(); + let curated_root = curated_plugins_repo_path(tmp.path()); + + write_file( + &tmp.path().join(CONFIG_TOML_FILE), + r#"[features] +plugins = true +"#, + ); + write_file( + &curated_root.join(".agents/plugins/marketplace.json"), + "{not valid json", + ); + + let config = load_config(tmp.path(), tmp.path()).await; + let outcome = PluginsManager::new(tmp.path().to_path_buf()) + .list_marketplaces_for_config(&config, &[], /*include_openai_curated*/ false) + .unwrap(); + + assert_eq!(outcome.errors, Vec::new()); + assert_eq!( + outcome + .marketplaces + .iter() + .any(|marketplace| marketplace.name == OPENAI_CURATED_MARKETPLACE_NAME), + false + ); +} + #[tokio::test] async fn list_marketplaces_includes_installed_marketplace_roots() { let tmp = tempfile::tempdir().unwrap(); @@ -2576,7 +2639,7 @@ source = "/tmp/debug" .unwrap(); let config = load_config(tmp.path(), tmp.path()).await; let marketplaces = PluginsManager::new(tmp.path().to_path_buf()) - .list_marketplaces_for_config(&config, &[]) + .list_marketplaces_for_config(&config, &[], /*include_openai_curated*/ true) .unwrap() .marketplaces; @@ -2652,7 +2715,7 @@ source = "/tmp/debug" let config = load_config(tmp.path(), tmp.path()).await; let marketplaces = PluginsManager::new(tmp.path().to_path_buf()) - .list_marketplaces_for_config(&config, &[]) + .list_marketplaces_for_config(&config, &[], /*include_openai_curated*/ true) .unwrap() .marketplaces; @@ -2707,7 +2770,7 @@ plugins = true .unwrap(); let config = load_config(tmp.path(), tmp.path()).await; let marketplaces = PluginsManager::new(tmp.path().to_path_buf()) - .list_marketplaces_for_config(&config, &[]) + .list_marketplaces_for_config(&config, &[], /*include_openai_curated*/ true) .unwrap() .marketplaces; @@ -2792,6 +2855,7 @@ enabled = false AbsolutePathBuf::try_from(repo_a_root).unwrap(), AbsolutePathBuf::try_from(repo_b_root).unwrap(), ], + /*include_openai_curated*/ true, ) .unwrap() .marketplaces; @@ -2902,7 +2966,11 @@ enabled = true let config = load_config(tmp.path(), &repo_root).await; let marketplaces = PluginsManager::new(tmp.path().to_path_buf()) - .list_marketplaces_for_config(&config, &[AbsolutePathBuf::try_from(repo_root).unwrap()]) + .list_marketplaces_for_config( + &config, + &[AbsolutePathBuf::try_from(repo_root).unwrap()], + /*include_openai_curated*/ true, + ) .unwrap() .marketplaces; diff --git a/codex-rs/core-skills/src/render.rs b/codex-rs/core-skills/src/render.rs index 1eff1f879f..1477a299b6 100644 --- a/codex-rs/core-skills/src/render.rs +++ b/codex-rs/core-skills/src/render.rs @@ -22,17 +22,17 @@ pub const SKILL_DESCRIPTION_TRUNCATED_WARNING: &str = "Skill descriptions were s pub const SKILL_DESCRIPTION_TRUNCATED_WARNING_WITH_PERCENT: &str = "Skill descriptions were shortened to fit the 2% skills context budget. Codex can still see every skill, but some descriptions are shorter. Disable unused skills or plugins to leave more room for the rest."; pub const SKILL_DESCRIPTIONS_REMOVED_WARNING_PREFIX: &str = "Exceeded skills context budget. All skill descriptions were removed and"; -pub const SKILLS_INTRO_WITH_ABSOLUTE_PATHS: &str = "A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill."; +pub const SKILLS_INTRO_WITH_ABSOLUTE_PATHS: &str = "A skill is a set of instructions provided through a `SKILL.md` source. Below is the list of skills that can be used. Each entry includes a name, description, and source locator. `file` locators are on the host filesystem, `environment resource` locators are owned by an execution environment, `orchestrator resource` locators are opaque non-filesystem resources, and `custom resource` locators use their provider's access mechanism."; pub const SKILLS_INTRO_WITH_ALIASES: &str = "A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and a short path that can be expanded into an absolute path using the skill roots table."; -pub const SKILLS_HOW_TO_USE_WITH_ABSOLUTE_PATHS: &str = r###"- Discovery: The list above is the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths. +pub const SKILLS_HOW_TO_USE_WITH_ABSOLUTE_PATHS: &str = r###"- Discovery: The list above is the skills available in this session (name + description + source locator). `file` entries live on the host filesystem, `environment resource` entries are owned by their execution environment, `orchestrator resource` entries must be accessed through `skills.list` and `skills.read`, and `custom resource` entries use their provider's access mechanism. - Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned. -- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback. +- Missing/blocked: If a named skill isn't in the list or its source can't be read, say so briefly and continue with the best fallback. - How to use a skill (progressive disclosure): - 1) After deciding to use a skill, the main agent must open and read its `SKILL.md` completely before taking task actions. If a read is truncated or paginated, continue until EOF. - 2) When `SKILL.md` references relative paths (e.g., `scripts/foo.py`), resolve them relative to the skill directory listed above first, and only consider other paths if needed. - 3) If `SKILL.md` points to extra folders such as `references/`, use its routing instructions to identify the files required for the task. The main agent must read each required instruction or reference file itself before acting on it. Do not delegate reading, summarizing, or interpreting skill instructions to a subagent. Subagents may still perform task work when the selected skill allows it. - 4) If `scripts/` exist, prefer running or patching them instead of retyping large code blocks. - 5) If `assets/` or templates exist, reuse them instead of recreating from scratch. + 1) After deciding to use a skill, the main agent must read its `SKILL.md` completely before taking task actions. For a `file` entry, open the listed path. For an `environment resource`, use the filesystem of the owning environment. For an `orchestrator resource`, call `skills.list` with `{"authority":{"kind":"orchestrator"}}`, select the matching package, and pass its `main_resource` to `skills.read`. If a read is truncated or paginated, continue until EOF. + 2) When `SKILL.md` references another resource, use the same access mechanism. Resolve relative paths against a filesystem-backed skill directory. For orchestrator skills, pass the exact referenced resource identifier with the same authority and package to `skills.read`; do not treat `skill://` identifiers as filesystem paths. + 3) If `SKILL.md` points to extra folders such as `references/`, use its routing instructions to identify the resources required for the task. The main agent must read each required instruction or reference file itself before acting on it. Do not delegate reading, summarizing, or interpreting skill instructions to a subagent. Subagents may still perform task work when the selected skill allows it. + 4) For filesystem-backed skills, prefer running or patching provided scripts instead of retyping large code blocks. For orchestrator skills, use `skills.read` and the available tools; do not invent a local path. + 5) Reuse provided assets or templates through the same source access mechanism instead of recreating them. - Coordination and sequencing: - If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them. - Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why. diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index b30b868bc7..5fcd8eb070 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -226,18 +226,6 @@ }, "type": "object" }, - "AppsMcpPathOverrideConfigToml": { - "additionalProperties": false, - "properties": { - "enabled": { - "type": "boolean" - }, - "path": { - "type": "string" - } - }, - "type": "object" - }, "AskForApproval": { "description": "Determines the conditions under which the user is consulted to approve running the command proposed by Codex.", "oneOf": [ @@ -408,7 +396,23 @@ "type": "boolean" }, "apps_mcp_path_override": { - "$ref": "#/definitions/FeatureToml_for_AppsMcpPathOverrideConfigToml" + "anyOf": [ + { + "type": "boolean" + }, + { + "additionalProperties": false, + "properties": { + "enabled": { + "type": "boolean" + }, + "path": { + "type": "string" + } + }, + "type": "object" + } + ] }, "auth_elicitation": { "type": "boolean" @@ -572,6 +576,9 @@ "request_rule": { "type": "boolean" }, + "resize_all_images": { + "type": "boolean" + }, "responses_websockets": { "type": "boolean" }, @@ -617,6 +624,9 @@ "terminal_visualization_instructions": { "type": "boolean" }, + "token_budget": { + "type": "boolean" + }, "tool_call_mcp_elicitation": { "type": "boolean" }, @@ -836,16 +846,6 @@ }, "type": "object" }, - "FeatureToml_for_AppsMcpPathOverrideConfigToml": { - "anyOf": [ - { - "type": "boolean" - }, - { - "$ref": "#/definitions/AppsMcpPathOverrideConfigToml" - } - ] - }, "FeatureToml_for_CodeModeConfigToml": { "anyOf": [ { @@ -4534,7 +4534,23 @@ "type": "boolean" }, "apps_mcp_path_override": { - "$ref": "#/definitions/FeatureToml_for_AppsMcpPathOverrideConfigToml" + "anyOf": [ + { + "type": "boolean" + }, + { + "additionalProperties": false, + "properties": { + "enabled": { + "type": "boolean" + }, + "path": { + "type": "string" + } + }, + "type": "object" + } + ] }, "auth_elicitation": { "type": "boolean" @@ -4698,6 +4714,9 @@ "request_rule": { "type": "boolean" }, + "resize_all_images": { + "type": "boolean" + }, "responses_websockets": { "type": "boolean" }, @@ -4743,6 +4762,9 @@ "terminal_visualization_instructions": { "type": "boolean" }, + "token_budget": { + "type": "boolean" + }, "tool_call_mcp_elicitation": { "type": "boolean" }, @@ -5224,4 +5246,4 @@ }, "title": "ConfigToml", "type": "object" -} +} \ No newline at end of file diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index c626dda4a5..ac42984b58 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1097,6 +1097,7 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(replacement_history), + window_id: None, }), RolloutItem::TurnContext(turn_context.to_turn_context_item()), RolloutItem::ResponseItem(spawn_agent_call(&parent_spawn_call_id)), @@ -2490,7 +2491,7 @@ async fn list_agent_subtree_thread_ids_includes_anonymous_and_closed_descendants } #[tokio::test] -async fn list_agent_subtree_thread_ids_includes_live_descendants_without_state_db() { +async fn list_agent_subtree_thread_ids_finds_live_descendants_of_unloaded_root() { let (_home, config) = test_config().await; let manager = ThreadManager::with_models_provider_home_and_state_for_tests( CodexAuth::from_api_key("dummy"), @@ -2535,6 +2536,8 @@ async fn list_agent_subtree_thread_ids_includes_live_descendants_without_state_d .await .expect("grandchild spawn should succeed"); + manager.remove_thread(&parent_thread_id).await; + let mut subtree_thread_ids = manager .list_agent_subtree_thread_ids(parent_thread_id) .await diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 02d231c962..a7220d20d8 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -28,7 +28,6 @@ use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::OnceLock; use std::sync::atomic::AtomicBool; -use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use codex_api::ApiError; @@ -172,7 +171,6 @@ pub(crate) struct CompactConversationRequestSettings { struct ModelClientState { session_id: SessionId, thread_id: ThreadId, - window_generation: AtomicU64, installation_id: String, provider: SharedModelProvider, auth_env_telemetry: AuthEnvTelemetry, @@ -345,7 +343,6 @@ impl ModelClient { state: Arc::new(ModelClientState { session_id, thread_id, - window_generation: AtomicU64::new(0), installation_id, provider: model_provider, auth_env_telemetry, @@ -394,24 +391,6 @@ impl ModelClient { self.state.provider.auth_manager() } - pub(crate) fn set_window_generation(&self, window_generation: u64) { - self.state - .window_generation - .store(window_generation, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); - } - - pub(crate) fn advance_window_generation(&self) { - self.state.window_generation.fetch_add(1, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); - } - - pub(crate) fn current_window_id(&self) -> String { - let thread_id = self.state.thread_id; - let window_generation = self.state.window_generation.load(Ordering::Relaxed); - format!("{thread_id}:{window_generation}") - } - fn take_cached_websocket_session(&self) -> WebsocketSession { let mut cached_websocket_session = self .state @@ -457,6 +436,7 @@ impl ModelClient { /// /// The model selection and telemetry context are passed explicitly to keep `ModelClient` /// session-scoped. + #[allow(clippy::too_many_arguments)] pub(crate) async fn compact_conversation_history( &self, prompt: &Prompt, @@ -464,6 +444,7 @@ impl ModelClient { settings: CompactConversationRequestSettings, session_telemetry: &SessionTelemetry, compaction_trace: &CompactionTraceContext, + window_id: &str, turn_metadata_header: Option<&str>, ) -> Result> { if prompt.input.is_empty() { @@ -488,6 +469,7 @@ impl ModelClient { settings.effort, settings.summary, settings.service_tier, + window_id, )?; let ResponsesApiRequest { model, @@ -522,7 +504,7 @@ impl ModelClient { /*turn_state*/ None, parse_turn_metadata_header(turn_metadata_header).as_ref(), )); - extra_headers.extend(self.build_responses_identity_headers()); + extra_headers.extend(self.build_responses_identity_headers(Some(window_id))); extra_headers.extend(build_session_headers( Some(self.state.session_id.to_string()), Some(self.state.thread_id.to_string()), @@ -644,14 +626,16 @@ impl ModelClient { extra_headers } - fn build_responses_identity_headers(&self) -> ApiHeaderMap { + fn build_responses_identity_headers(&self, window_id: Option<&str>) -> ApiHeaderMap { let mut extra_headers = self.build_subagent_headers(); if let Some(parent_thread_id) = parent_thread_id_header_value(self.state.parent_thread_id) && let Ok(val) = HeaderValue::from_str(&parent_thread_id) { extra_headers.insert(X_CODEX_PARENT_THREAD_ID_HEADER, val); } - if let Ok(val) = HeaderValue::from_str(&self.current_window_id()) { + if let Some(window_id) = window_id + && let Ok(val) = HeaderValue::from_str(window_id) + { extra_headers.insert(X_CODEX_WINDOW_ID_HEADER, val); } extra_headers @@ -659,6 +643,7 @@ impl ModelClient { fn build_ws_client_metadata( &self, + window_id: &str, turn_metadata_header: Option<&str>, use_responses_lite: bool, ) -> HashMap { @@ -667,10 +652,7 @@ impl ModelClient { X_CODEX_INSTALLATION_ID_HEADER.to_string(), self.state.installation_id.clone(), ); - client_metadata.insert( - X_CODEX_WINDOW_ID_HEADER.to_string(), - self.current_window_id(), - ); + client_metadata.insert(X_CODEX_WINDOW_ID_HEADER.to_string(), window_id.to_string()); if let Some(subagent) = subagent_header_value(&self.state.session_source) { client_metadata.insert(X_OPENAI_SUBAGENT_HEADER.to_string(), subagent); } @@ -752,6 +734,7 @@ impl ModelClient { } } + #[allow(clippy::too_many_arguments)] fn build_responses_request( &self, provider: &codex_api::Provider, @@ -760,9 +743,10 @@ impl ModelClient { effort: Option, summary: ReasoningSummaryConfig, service_tier: Option, + window_id: &str, ) -> Result { let instructions = &prompt.base_instructions.text; - let input = prompt.get_formatted_input(); + let input = prompt.get_formatted_input_for_request(model_info.use_responses_lite); let tools = create_tools_json_for_responses_api(&prompt.tools)?; let reasoning = Self::build_reasoning(model_info, effort, summary); let include = if reasoning.is_some() { @@ -807,10 +791,7 @@ impl ModelClient { X_CODEX_INSTALLATION_ID_HEADER.to_string(), self.state.installation_id.clone(), ), - ( - X_CODEX_WINDOW_ID_HEADER.to_string(), - self.current_window_id(), - ), + (X_CODEX_WINDOW_ID_HEADER.to_string(), window_id.to_string()), ])), }; Ok(request) @@ -955,7 +936,7 @@ impl ModelClient { headers.insert("x-client-request-id", header_value); } headers.extend(build_session_headers(Some(session_id), Some(thread_id))); - headers.extend(self.build_responses_identity_headers()); + headers.extend(self.build_responses_identity_headers(/*window_id*/ None)); if let Some(header_value) = self.generate_attestation_header_for().await { headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } @@ -998,6 +979,7 @@ impl ModelClientSession { /// regardless of transport choice. async fn build_responses_options( &self, + window_id: &str, turn_metadata_header: Option<&str>, compression: Compression, use_responses_lite: bool, @@ -1015,7 +997,10 @@ impl ModelClientSession { Some(&self.turn_state), turn_metadata_header.as_ref(), ); - headers.extend(self.client.build_responses_identity_headers()); + headers.extend( + self.client + .build_responses_identity_headers(Some(window_id)), + ); if let Some(header_value) = self.client.generate_attestation_header_for().await { headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } @@ -1115,7 +1100,6 @@ impl ModelClientSession { pub async fn preconnect_websocket( &mut self, session_telemetry: &SessionTelemetry, - _model_info: &ModelInfo, ) -> std::result::Result<(), ApiError> { if !self.client.responses_websocket_enabled() { return Ok(()); @@ -1257,6 +1241,7 @@ impl ModelClientSession { )] async fn stream_responses_api( &self, + window_id: &str, prompt: &Prompt, model_info: &ModelInfo, session_telemetry: &SessionTelemetry, @@ -1288,6 +1273,7 @@ impl ModelClientSession { let compression = self.responses_request_compression(client_setup.auth.as_ref()); let mut options = self .build_responses_options( + window_id, turn_metadata_header, compression, model_info.use_responses_lite, @@ -1301,6 +1287,7 @@ impl ModelClientSession { effort.clone(), summary, service_tier.clone(), + window_id, )?; let inference_trace_attempt = inference_trace.start_attempt(); inference_trace_attempt.add_request_headers(&mut options.extra_headers); @@ -1374,6 +1361,7 @@ impl ModelClientSession { )] async fn stream_responses_websocket( &mut self, + window_id: &str, prompt: &Prompt, model_info: &ModelInfo, session_telemetry: &SessionTelemetry, @@ -1402,6 +1390,7 @@ impl ModelClientSession { let options = self .build_responses_options( + window_id, turn_metadata_header, compression, model_info.use_responses_lite, @@ -1414,10 +1403,12 @@ impl ModelClientSession { effort.clone(), summary, service_tier.clone(), + window_id, )?; let mut ws_payload = ResponseCreateWsRequest { client_metadata: response_create_client_metadata( Some(self.client.build_ws_client_metadata( + window_id, turn_metadata_header, model_info.use_responses_lite, )), @@ -1553,6 +1544,7 @@ impl ModelClientSession { #[allow(clippy::too_many_arguments)] pub async fn prewarm_websocket( &mut self, + window_id: &str, prompt: &Prompt, model_info: &ModelInfo, session_telemetry: &SessionTelemetry, @@ -1571,6 +1563,7 @@ impl ModelClientSession { let disabled_trace = InferenceTraceContext::disabled(); match self .stream_responses_websocket( + window_id, prompt, model_info, session_telemetry, @@ -1614,6 +1607,7 @@ impl ModelClientSession { /// branches. pub async fn stream( &mut self, + window_id: &str, prompt: &Prompt, model_info: &ModelInfo, session_telemetry: &SessionTelemetry, @@ -1630,6 +1624,7 @@ impl ModelClientSession { let request_trace = current_span_w3c_trace_context(); match self .stream_responses_websocket( + window_id, prompt, model_info, session_telemetry, @@ -1651,6 +1646,7 @@ impl ModelClientSession { } self.stream_responses_api( + window_id, prompt, model_info, session_telemetry, @@ -1989,7 +1985,7 @@ impl AuthRequestTelemetryContext { let auth_telemetry = auth_header_telemetry(api_auth); Self { auth_mode: auth_mode.map(|mode| match mode { - AuthMode::ApiKey => "ApiKey", + AuthMode::ApiKey | AuthMode::BedrockApiKey => "ApiKey", AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens | AuthMode::AgentIdentity diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index 48d0523df4..4d7729d11b 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -2,6 +2,8 @@ pub use codex_api::ResponseEvent; use codex_config::types::Personality; use codex_protocol::error::Result; use codex_protocol::models::BaseInstructions; +use codex_protocol::models::ContentItem; +use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::InterAgentCommunication; use codex_tools::ToolSpec; @@ -71,6 +73,56 @@ impl Prompt { }) .collect() } + + pub(crate) fn get_formatted_input_for_request( + &self, + use_responses_lite: bool, + ) -> Vec { + let mut input = self.get_formatted_input(); + if use_responses_lite { + strip_image_details(&mut input); + } + input + } +} + +fn strip_image_details(items: &mut [ResponseItem]) { + for item in items { + match item { + ResponseItem::Message { content, .. } => { + for content_item in content { + if let ContentItem::InputImage { detail, .. } = content_item { + *detail = None; + } + } + } + ResponseItem::FunctionCallOutput { output, .. } + | ResponseItem::CustomToolCallOutput { output, .. } => { + if let Some(content) = output.content_items_mut() { + for content_item in content { + if let FunctionCallOutputContentItem::InputImage { detail, .. } = + content_item + { + *detail = None; + } + } + } + } + ResponseItem::Reasoning { .. } + | ResponseItem::AgentMessage { .. } + | ResponseItem::LocalShellCall { .. } + | ResponseItem::FunctionCall { .. } + | ResponseItem::ToolSearchCall { .. } + | ResponseItem::CustomToolCall { .. } + | ResponseItem::ToolSearchOutput { .. } + | ResponseItem::WebSearchCall { .. } + | ResponseItem::ImageGenerationCall { .. } + | ResponseItem::Compaction { .. } + | ResponseItem::CompactionTrigger + | ResponseItem::ContextCompaction { .. } + | ResponseItem::Other => {} + } + } } pub struct ResponseStream { diff --git a/codex-rs/core/src/client_common_tests.rs b/codex-rs/core/src/client_common_tests.rs index 45d4334222..4b218e490c 100644 --- a/codex-rs/core/src/client_common_tests.rs +++ b/codex-rs/core/src/client_common_tests.rs @@ -3,10 +3,95 @@ use codex_api::ResponsesApiRequest; use codex_api::TextControls; use codex_api::create_text_param_for_request; use codex_protocol::config_types::ServiceTier; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ImageDetail; use pretty_assertions::assert_eq; use super::*; +fn prompt_with_image_outputs() -> Prompt { + Prompt { + input: vec![ + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputImage { + image_url: "https://example.com/image.png".to_string(), + detail: Some(ImageDetail::Original), + }], + phase: None, + }, + ResponseItem::FunctionCallOutput { + call_id: "function-call".to_string(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,function".to_string(), + detail: Some(ImageDetail::High), + }, + ]), + }, + ResponseItem::CustomToolCallOutput { + call_id: "custom-call".to_string(), + name: None, + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,custom".to_string(), + detail: Some(ImageDetail::Auto), + }, + ]), + }, + ], + ..Default::default() + } +} + +#[test] +fn responses_lite_request_copies_strip_image_details() { + let prompt = prompt_with_image_outputs(); + let original = prompt.input.clone(); + + let stripped = prompt.get_formatted_input_for_request(/*use_responses_lite*/ true); + + assert_eq!( + stripped, + vec![ + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputImage { + image_url: "https://example.com/image.png".to_string(), + detail: None, + }], + phase: None, + }, + ResponseItem::FunctionCallOutput { + call_id: "function-call".to_string(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,function".to_string(), + detail: None, + }, + ]), + }, + ResponseItem::CustomToolCallOutput { + call_id: "custom-call".to_string(), + name: None, + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,custom".to_string(), + detail: None, + }, + ]), + }, + ] + ); + assert_eq!(prompt.input, original); + assert_eq!( + prompt.get_formatted_input_for_request(/*use_responses_lite*/ false), + original + ); +} + #[test] fn serializes_text_verbosity_when_set() { let input: Vec = vec![]; diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 2ebffe3922..d93bd8bf7a 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -291,13 +291,13 @@ fn build_ws_client_metadata_includes_window_lineage_and_turn_metadata() { Some(parent_thread_id), ); - client.advance_window_generation(); - + let thread_id = client.state.thread_id; + let window_id = format!("{thread_id}:1"); let client_metadata = client.build_ws_client_metadata( + &window_id, Some(r#"{"turn_id":"turn-123"}"#), /*use_responses_lite*/ false, ); - let thread_id = client.state.thread_id; assert_eq!( client_metadata, std::collections::HashMap::from([ diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 7dabb09d33..cb7767a16c 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -452,7 +452,7 @@ async fn delegated_mcp_guardian_abort_returns_synthetic_decline_answer() { } #[tokio::test] -async fn delegated_mcp_user_reviewer_waits_for_metadata_lookup() { +async fn delegated_mcp_user_reviewer_returns_none_without_metadata() { let (parent_session, parent_ctx, _rx_events) = crate::session::tests::make_session_and_context_with_rx().await; let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([( @@ -464,21 +464,6 @@ async fn delegated_mcp_user_reviewer_waits_for_metadata_lookup() { }, )]))); let cancel_token = CancellationToken::new(); - let manager = Arc::clone(&parent_session.services.mcp_connection_manager); - let (manager_locked_tx, manager_locked_rx) = std::sync::mpsc::sync_channel(0); - let (release_manager_tx, release_manager_rx) = std::sync::mpsc::sync_channel(0); - let manager_lock = tokio::task::spawn_blocking(move || { - let _manager_guard = manager.blocking_write(); - manager_locked_tx - .send(()) - .expect("manager lock receiver should remain open"); - release_manager_rx - .recv() - .expect("manager lock release sender should remain open"); - }); - manager_locked_rx - .recv_timeout(Duration::from_secs(1)) - .expect("manager write lock should be acquired"); let event = RequestUserInputEvent { call_id: "call-1".to_string(), @@ -498,24 +483,7 @@ async fn delegated_mcp_user_reviewer_waits_for_metadata_lookup() { &pending_mcp_invocations, &event, &cancel_token, - ); - tokio::pin!(response); - assert!( - timeout(Duration::from_millis(100), &mut response) - .await - .is_err(), - "manual reviewer should wait for MCP metadata" - ); - release_manager_tx - .send(()) - .expect("manager lock holder should remain open"); - manager_lock - .await - .expect("manager lock task should not panic"); - assert_eq!( - timeout(Duration::from_secs(1), response) - .await - .expect("manual reviewer should finish after MCP metadata lookup"), - None - ); + ) + .await; + assert_eq!(response, None); } diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index ab1f6df03e..c4cb1e8e12 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -159,6 +159,14 @@ pub struct CodexThread { out_of_band_elicitation_count: Mutex, } +#[derive(Debug, Eq, PartialEq)] +pub struct BackgroundTerminalInfo { + pub item_id: String, + pub process_id: String, + pub command: String, + pub cwd: AbsolutePathBuf, +} + /// Conduit for the bidirectional stream of messages that compose a thread /// (formerly called a conversation) in Codex. impl CodexThread { @@ -396,6 +404,17 @@ impl CodexThread { self.codex.agent_status().await } + pub async fn list_background_terminals(&self) -> Vec { + self.codex.session.list_background_terminals().await + } + + pub async fn terminate_background_terminal(&self, process_id: i32) -> bool { + self.codex + .session + .terminate_background_terminal(process_id) + .await + } + pub(crate) fn subscribe_status(&self) -> watch::Receiver { self.codex.agent_status.clone() } diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index e2573d6d6d..26b5e33bf6 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -143,17 +143,17 @@ async fn run_compact_task_inner( let pre_compact_outcome = run_pre_compact_hooks(&sess, &turn_context, trigger).await; match pre_compact_outcome { PreCompactHookOutcome::Continue => {} - PreCompactHookOutcome::Stopped { reason } => { - let error = reason.unwrap_or_else(|| "PreCompact hook stopped execution".to_string()); + PreCompactHookOutcome::Stopped => { + let error = CodexErr::TurnAborted; attempt .track( sess.as_ref(), CompactionStatus::Interrupted, - Some(error), + Some(&error), CompactionAnalyticsDetails::default(), ) .await; - return Err(CodexErr::TurnAborted); + return Err(error); } } let result = run_compact_task_inner_impl( @@ -165,7 +165,7 @@ async fn run_compact_task_inner( ) .await; let status = compaction_status_from_result(&result); - let error = result.as_ref().err().map(ToString::to_string); + let codex_error = result.as_ref().err(); if result.is_ok() { let post_compact_outcome = run_post_compact_hooks(&sess, &turn_context, trigger).await; if let PostCompactHookOutcome::Stopped = post_compact_outcome { @@ -173,7 +173,7 @@ async fn run_compact_task_inner( .track( sess.as_ref(), status, - error, + codex_error, CompactionAnalyticsDetails::default(), ) .await; @@ -184,7 +184,7 @@ async fn run_compact_task_inner( .track( sess.as_ref(), status, - error, + codex_error, CompactionAnalyticsDetails::default(), ) .await; @@ -228,7 +228,7 @@ async fn run_compact_task_inner_impl( personality: turn_context.personality, ..Default::default() }; - let window_id = sess.services.model_client.current_window_id(); + let window_id = sess.current_window_id().await; let turn_metadata_header = turn_context .turn_metadata_state .current_header_value_for_compaction(&window_id, compaction_metadata); @@ -236,6 +236,7 @@ async fn run_compact_task_inner_impl( &sess, turn_context.as_ref(), &mut client_session, + &window_id, turn_metadata_header.as_deref(), &prompt, ) @@ -293,6 +294,7 @@ async fn run_compact_task_inner_impl( let user_messages = collect_user_messages(history_items); let mut new_history = build_compacted_history(Vec::new(), &user_messages, &summary_text); + let window_id = sess.advance_auto_compact_window_id().await; if matches!( initial_context_injection, @@ -309,6 +311,7 @@ async fn run_compact_task_inner_impl( let compacted_item = CompactedItem { message: summary_text.clone(), replacement_history: Some(new_history.clone()), + window_id: Some(window_id), }; sess.replace_compacted_history(new_history, reference_context_item, compacted_item) .await; @@ -340,6 +343,7 @@ pub(crate) struct CompactionAnalyticsDetails { pub(crate) active_context_tokens_before: Option, pub(crate) retained_image_count: Option, pub(crate) compaction_summary_tokens: Option, + pub(crate) cached_input_tokens: Option, } impl CompactionAnalyticsAttempt { @@ -369,13 +373,14 @@ impl CompactionAnalyticsAttempt { self, sess: &Session, status: CompactionStatus, - error: Option, + codex_error: Option<&CodexErr>, details: CompactionAnalyticsDetails, ) { let CompactionAnalyticsDetails { active_context_tokens_before, retained_image_count, compaction_summary_tokens, + cached_input_tokens, } = details; let active_context_tokens_before = active_context_tokens_before.unwrap_or(self.active_context_tokens_before); @@ -391,11 +396,14 @@ impl CompactionAnalyticsAttempt { phase: self.phase, strategy: CompactionStrategy::Memento, status, - error, + codex_error_kind: codex_error.map(Into::into), + codex_error_http_status_code: codex_error + .and_then(CodexErr::http_status_code_value), active_context_tokens_before, active_context_tokens_after, retained_image_count, compaction_summary_tokens, + cached_input_tokens, started_at: self.started_at, completed_at: now_unix_seconds(), duration_ms: Some( @@ -579,11 +587,13 @@ async fn drain_to_completed( sess: &Session, turn_context: &TurnContext, client_session: &mut ModelClientSession, + window_id: &str, turn_metadata_header: Option<&str>, prompt: &Prompt, ) -> CodexResult<()> { let mut stream = client_session .stream( + window_id, prompt, &turn_context.model_info, &turn_context.session_telemetry, diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index bd852b4e79..93457053f5 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -113,17 +113,17 @@ async fn run_remote_compact_task_inner( let pre_compact_outcome = run_pre_compact_hooks(sess, turn_context, trigger).await; match pre_compact_outcome { PreCompactHookOutcome::Continue => {} - PreCompactHookOutcome::Stopped { reason } => { - let error = reason.unwrap_or_else(|| "PreCompact hook stopped execution".to_string()); + PreCompactHookOutcome::Stopped => { + let error = CodexErr::TurnAborted; attempt .track( sess.as_ref(), codex_analytics::CompactionStatus::Interrupted, - Some(error), + Some(&error), analytics_details, ) .await; - return Err(CodexErr::TurnAborted); + return Err(error); } } let result = run_remote_compact_task_inner_impl( @@ -135,18 +135,18 @@ async fn run_remote_compact_task_inner( ) .await; let status = compaction_status_from_result(&result); - let error = result.as_ref().err().map(ToString::to_string); + let codex_error = result.as_ref().err(); if result.is_ok() { let post_compact_outcome = run_post_compact_hooks(sess, turn_context, trigger).await; if let PostCompactHookOutcome::Stopped = post_compact_outcome { attempt - .track(sess.as_ref(), status, error, analytics_details) + .track(sess.as_ref(), status, codex_error, analytics_details) .await; return Err(CodexErr::TurnAborted); } } attempt - .track(sess.as_ref(), status, error.clone(), analytics_details) + .track(sess.as_ref(), status, codex_error, analytics_details) .await; if let Err(err) = result { sess.track_turn_codex_error(turn_context, &err); @@ -224,7 +224,7 @@ async fn run_remote_compact_task_inner_impl( output_schema: None, output_schema_strict: true, }; - let window_id = sess.services.model_client.current_window_id(); + let window_id = sess.current_window_id().await; let turn_metadata_header = turn_context .turn_metadata_state .current_header_value_for_compaction(&window_id, compaction_metadata); @@ -245,9 +245,11 @@ async fn run_remote_compact_task_inner_impl( }, &turn_context.session_telemetry, &compaction_trace, + &window_id, turn_metadata_header.as_deref(), ) .await?; + let new_window_id = sess.advance_auto_compact_window_id().await; new_history = process_compacted_history( sess.as_ref(), turn_context.as_ref(), @@ -263,6 +265,7 @@ async fn run_remote_compact_task_inner_impl( let compacted_item = CompactedItem { message: String::new(), replacement_history: Some(new_history.clone()), + window_id: Some(new_window_id), }; // Install is the semantic boundary where the compact endpoint's output becomes live // thread history. Keep it distinct from the later inference request so the reducer can diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index a80d4dd56f..5ab7a68a5b 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -127,17 +127,17 @@ async fn run_remote_compact_task_inner( let pre_compact_outcome = run_pre_compact_hooks(sess, turn_context, trigger).await; match pre_compact_outcome { PreCompactHookOutcome::Continue => {} - PreCompactHookOutcome::Stopped { reason } => { - let error = reason.unwrap_or_else(|| "PreCompact hook stopped execution".to_string()); + PreCompactHookOutcome::Stopped => { + let error = CodexErr::TurnAborted; attempt .track( sess.as_ref(), codex_analytics::CompactionStatus::Interrupted, - Some(error), + Some(&error), analytics_details, ) .await; - return Err(CodexErr::TurnAborted); + return Err(error); } } let result = run_remote_compact_task_inner_impl( @@ -150,18 +150,18 @@ async fn run_remote_compact_task_inner( ) .await; let status = compaction_status_from_result(&result); - let error = result.as_ref().err().map(ToString::to_string); + let codex_error = result.as_ref().err(); if result.is_ok() { let post_compact_outcome = run_post_compact_hooks(sess, turn_context, trigger).await; if let PostCompactHookOutcome::Stopped = post_compact_outcome { attempt - .track(sess.as_ref(), status, error, analytics_details) + .track(sess.as_ref(), status, codex_error, analytics_details) .await; return Err(CodexErr::TurnAborted); } } attempt - .track(sess.as_ref(), status, error.clone(), analytics_details) + .track(sess.as_ref(), status, codex_error, analytics_details) .await; if let Err(err) = result { sess.track_turn_codex_error(turn_context, &err); @@ -240,7 +240,7 @@ async fn run_remote_compact_task_inner_impl( output_schema_strict: true, }; - let window_id = sess.services.model_client.current_window_id(); + let window_id = sess.current_window_id().await; let turn_metadata_header = turn_context .turn_metadata_state .current_header_value_for_compaction(&window_id, compaction_metadata); @@ -264,6 +264,7 @@ async fn run_remote_compact_task_inner_impl( turn_context, client_session, &prompt, + &window_id, turn_metadata_header.as_deref(), ) .await; @@ -280,10 +281,12 @@ async fn run_remote_compact_task_inner_impl( if let Some(token_usage) = token_usage { analytics_details.active_context_tokens_before = Some(token_usage.input_tokens); analytics_details.compaction_summary_tokens = Some(token_usage.output_tokens); + analytics_details.cached_input_tokens = Some(token_usage.cached_input_tokens); } let (compacted_history, retained_images) = build_v2_compacted_history(&prompt_input, compaction_output); analytics_details.retained_image_count = Some(retained_images); + let new_window_id = sess.advance_auto_compact_window_id().await; let new_history = process_compacted_history( sess.as_ref(), turn_context.as_ref(), @@ -299,6 +302,7 @@ async fn run_remote_compact_task_inner_impl( let compacted_item = CompactedItem { message: String::new(), replacement_history: Some(new_history.clone()), + window_id: Some(new_window_id), }; compaction_trace.record_installed(&CompactionCheckpointTracePayload { input_history: &trace_input_history, @@ -323,6 +327,7 @@ async fn run_remote_compaction_request_v2( turn_context: &TurnContext, client_session: &mut ModelClientSession, prompt: &Prompt, + window_id: &str, turn_metadata_header: Option<&str>, ) -> CodexResult { let max_retries = turn_context @@ -334,6 +339,7 @@ async fn run_remote_compaction_request_v2( loop { let result = match client_session .stream( + window_id, prompt, &turn_context.model_info, &turn_context.session_telemetry, diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index 69bd95f2df..550a7ec846 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -483,6 +483,7 @@ async fn process_compacted_history_reinjects_model_switch_message() { }]; let previous_turn_settings = PreviousTurnSettings { model: "previous-regular-model".to_string(), + comp_hash: None, realtime_active: None, }; diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 3a3b490d69..58893f8c04 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -5500,14 +5500,9 @@ async fn to_mcp_config_preserves_apps_feature_from_config() -> std::io::Result<( .await?; let plugins_manager = PluginsManager::new(codex_home.path().to_path_buf()); - config.apps_mcp_path_override = Some("/custom/mcp".to_string()); config.apps_mcp_product_sku = Some("tpp".to_string()); let mcp_config = config.to_mcp_config(&plugins_manager).await; assert!(mcp_config.apps_enabled); - assert_eq!( - mcp_config.apps_mcp_path_override.as_deref(), - Some("/custom/mcp") - ); assert_eq!(mcp_config.apps_mcp_product_sku.as_deref(), Some("tpp")); let _ = config.features.disable(Feature::Apps); @@ -8862,84 +8857,6 @@ allow_login_shell = false Ok(()) } -#[tokio::test] -async fn config_loads_apps_mcp_path_override_from_feature_config() -> std::io::Result<()> { - let codex_home = TempDir::new()?; - let toml = r#" -model = "gpt-5.4" - -[features.apps_mcp_path_override] -path = "/custom/mcp" -"#; - let cfg: ConfigToml = - toml::from_str(toml).expect("TOML deserialization should succeed for apps MCP feature"); - - let config = Config::load_from_base_config_with_overrides( - cfg, - ConfigOverrides::default(), - codex_home.abs(), - ) - .await?; - - assert_eq!( - config.apps_mcp_path_override.as_deref(), - Some("/custom/mcp") - ); - Ok(()) -} - -#[tokio::test] -async fn config_defaults_enabled_apps_mcp_path_override_to_plugin_service() -> std::io::Result<()> { - let codex_home = TempDir::new()?; - let toml = r#" -model = "gpt-5.4" - -[features] -apps_mcp_path_override = true -"#; - let cfg: ConfigToml = - toml::from_str(toml).expect("TOML deserialization should succeed for apps MCP feature"); - - let config = Config::load_from_base_config_with_overrides( - cfg, - ConfigOverrides::default(), - codex_home.abs(), - ) - .await?; - - assert!(config.features.enabled(Feature::AppsMcpPathOverride)); - assert_eq!(config.apps_mcp_path_override.as_deref(), Some("/ps/mcp")); - Ok(()) -} - -#[tokio::test] -async fn config_preserves_explicit_apps_mcp_path_override_path() -> std::io::Result<()> { - let codex_home = TempDir::new()?; - let toml = r#" -model = "gpt-5.4" - -[features.apps_mcp_path_override] -enabled = true -path = "/custom/mcp" -"#; - let cfg: ConfigToml = - toml::from_str(toml).expect("TOML deserialization should succeed for apps MCP feature"); - - let config = Config::load_from_base_config_with_overrides( - cfg, - ConfigOverrides::default(), - codex_home.abs(), - ) - .await?; - - assert_eq!( - config.apps_mcp_path_override.as_deref(), - Some("/custom/mcp") - ); - assert!(config.features.enabled(Feature::AppsMcpPathOverride)); - Ok(()) -} - #[tokio::test] async fn config_loads_apps_mcp_product_sku_from_toml() -> std::io::Result<()> { let codex_home = TempDir::new()?; @@ -9914,10 +9831,7 @@ enabled = true .build() .await?; - assert_eq!(config.multi_agent_v2.max_concurrent_threads_per_session, 4); - assert_eq!(config.multi_agent_v2.min_wait_timeout_ms, 10_000); - assert_eq!(config.multi_agent_v2.max_wait_timeout_ms, 3_600_000); - assert_eq!(config.multi_agent_v2.default_wait_timeout_ms, 30_000); + assert_eq!(config.multi_agent_v2, MultiAgentV2Config::default()); assert_eq!( ( config.agent_max_threads, @@ -9925,36 +9839,32 @@ enabled = true ), (None, Some(3)) ); - assert_eq!( - config.multi_agent_v2.root_agent_usage_hint_text.as_deref(), - Some(DEFAULT_MULTI_AGENT_V2_ROOT_AGENT_USAGE_HINT_TEXT) - ); - assert!( - !config - .multi_agent_v2 - .root_agent_usage_hint_text - .as_deref() - .unwrap_or_default() - .contains("maximum concurrency"), - ); - assert_eq!( - config.multi_agent_v2.subagent_usage_hint_text.as_deref(), - Some(DEFAULT_MULTI_AGENT_V2_SUBAGENT_USAGE_HINT_TEXT) - ); - assert!( - !config - .multi_agent_v2 - .subagent_usage_hint_text - .as_deref() - .unwrap_or_default() - .contains("maximum concurrency"), - ); - assert!(config.multi_agent_v2.hide_spawn_agent_metadata); - assert!(config.multi_agent_v2.non_code_mode_only); Ok(()) } +#[test] +fn multi_agent_v2_default_usage_hints_use_configured_thread_cap() { + let config_toml = toml::from_str( + r#"[features.multi_agent_v2] +enabled = true +max_concurrent_threads_per_session = 17 +"#, + ) + .expect("multi-agent v2 config should parse"); + + let config = resolve_multi_agent_v2_config(&config_toml); + let concurrency_guidance = "There are 17 available concurrency slots, meaning that up to 17 agents can be active at once, including you."; + assert!( + [ + config.root_agent_usage_hint_text, + config.subagent_usage_hint_text, + ] + .into_iter() + .all(|hint| hint.is_some_and(|hint| hint.ends_with(concurrency_guidance))) + ); +} + #[tokio::test] async fn multi_agent_v2_empty_usage_hint_overrides_clear_default_hints() -> std::io::Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 5c2a93bc45..597033e4e1 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -58,7 +58,6 @@ use codex_config::types::WindowsSandboxModeToml; use codex_core_plugins::PluginsConfigInput; use codex_exec_server::ExecutorFileSystem; use codex_exec_server::LOCAL_FS; -use codex_features::AppsMcpPathOverrideConfigToml; use codex_features::CodeModeConfigToml; use codex_features::Feature; use codex_features::FeatureConfigSource; @@ -234,6 +233,13 @@ Payload: ``` You may also see them addressed as to=/root/..., which indicates your identity is /root/... "#; + +fn default_multi_agent_v2_usage_hint_text(usage_hint_text: &str, max_concurrency: usize) -> String { + format!( + "{usage_hint_text}\nThere are {max_concurrency} available concurrency slots, meaning that up to {max_concurrency} agents can be active at once, including you." + ) +} + pub(crate) const HARD_MIN_MULTI_AGENT_V2_TIMEOUT_MS: i64 = 0; pub(crate) const HARD_MAX_MULTI_AGENT_V2_TIMEOUT_MS: i64 = DEFAULT_MULTI_AGENT_V2_MAX_WAIT_TIMEOUT_MS; @@ -931,9 +937,6 @@ pub struct Config { /// Base URL for requests to ChatGPT (as opposed to the OpenAI API). pub chatgpt_base_url: String, - /// Optional path override for the host-owned apps MCP server. - pub apps_mcp_path_override: Option, - /// Optional product SKU forwarded to the host-owned apps MCP server. pub apps_mcp_product_sku: Option, @@ -1059,22 +1062,23 @@ pub struct MultiAgentV2Config { pub non_code_mode_only: bool, } -impl Default for MultiAgentV2Config { - fn default() -> Self { +impl MultiAgentV2Config { + fn defaults_for_max_concurrency(max_concurrent_threads_per_session: usize) -> Self { Self { - max_concurrent_threads_per_session: - DEFAULT_MULTI_AGENT_V2_MAX_CONCURRENT_THREADS_PER_SESSION, + max_concurrent_threads_per_session, min_wait_timeout_ms: DEFAULT_MULTI_AGENT_V2_MIN_WAIT_TIMEOUT_MS, max_wait_timeout_ms: DEFAULT_MULTI_AGENT_V2_MAX_WAIT_TIMEOUT_MS, default_wait_timeout_ms: DEFAULT_MULTI_AGENT_V2_DEFAULT_WAIT_TIMEOUT_MS, usage_hint_enabled: true, usage_hint_text: None, - root_agent_usage_hint_text: Some( - DEFAULT_MULTI_AGENT_V2_ROOT_AGENT_USAGE_HINT_TEXT.to_string(), - ), - subagent_usage_hint_text: Some( - DEFAULT_MULTI_AGENT_V2_SUBAGENT_USAGE_HINT_TEXT.to_string(), - ), + root_agent_usage_hint_text: Some(default_multi_agent_v2_usage_hint_text( + DEFAULT_MULTI_AGENT_V2_ROOT_AGENT_USAGE_HINT_TEXT, + max_concurrent_threads_per_session, + )), + subagent_usage_hint_text: Some(default_multi_agent_v2_usage_hint_text( + DEFAULT_MULTI_AGENT_V2_SUBAGENT_USAGE_HINT_TEXT, + max_concurrent_threads_per_session, + )), tool_namespace: None, hide_spawn_agent_metadata: true, non_code_mode_only: true, @@ -1082,6 +1086,14 @@ impl Default for MultiAgentV2Config { } } +impl Default for MultiAgentV2Config { + fn default() -> Self { + Self::defaults_for_max_concurrency( + DEFAULT_MULTI_AGENT_V2_MAX_CONCURRENT_THREADS_PER_SESSION, + ) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum TerminalResizeReflowMaxRows { /// Use the runtime terminal detector to choose a scrollback-sized cap. @@ -1415,7 +1427,6 @@ impl Config { McpConfig { chatgpt_base_url: self.chatgpt_base_url.clone(), - apps_mcp_path_override: self.apps_mcp_path_override.clone(), apps_mcp_product_sku: self.apps_mcp_product_sku.clone(), codex_home: self.codex_home.to_path_buf(), mcp_oauth_credentials_store_mode: self.mcp_oauth_credentials_store_mode, @@ -1428,7 +1439,6 @@ impl Config { codex_linux_sandbox_exe: self.codex_linux_sandbox_exe.clone(), use_legacy_landlock: self.features.use_legacy_landlock(), apps_enabled: self.features.enabled(Feature::Apps), - legacy_apps_mcp_loader_enabled: true, prefix_mcp_tool_names: self.prefix_mcp_tool_names(), client_elicitation_capability: if self.features.enabled(Feature::AuthElicitation) { ElicitationCapability { @@ -2319,11 +2329,11 @@ fn resolve_code_mode_config(config_toml: &ConfigToml) -> CodeModeConfig { fn resolve_multi_agent_v2_config(config_toml: &ConfigToml) -> MultiAgentV2Config { let base = multi_agent_v2_toml_config(config_toml.features.as_ref()); - let default = MultiAgentV2Config::default(); - let max_concurrent_threads_per_session = base .and_then(|config| config.max_concurrent_threads_per_session) - .unwrap_or(default.max_concurrent_threads_per_session); + .unwrap_or(DEFAULT_MULTI_AGENT_V2_MAX_CONCURRENT_THREADS_PER_SESSION); + let default = + MultiAgentV2Config::defaults_for_max_concurrency(max_concurrent_threads_per_session); let min_wait_timeout_ms = base .and_then(|config| config.min_wait_timeout_ms) .unwrap_or(default.min_wait_timeout_ms); @@ -2413,15 +2423,6 @@ fn multi_agent_v2_toml_config(features: Option<&FeaturesToml>) -> Option<&MultiA } } -fn apps_mcp_path_override_toml_config( - features: Option<&FeaturesToml>, -) -> Option<&AppsMcpPathOverrideConfigToml> { - match features?.apps_mcp_path_override.as_ref()? { - FeatureToml::Enabled(_) => None, - FeatureToml::Config(config) => Some(config), - } -} - fn network_proxy_toml_config(features: Option<&FeaturesToml>) -> Option<&NetworkProxyConfigToml> { match features?.network_proxy.as_ref()? { FeatureToml::Enabled(_) => None, @@ -3043,14 +3044,6 @@ impl Config { resolve_experimental_request_user_input_enabled(&cfg); let code_mode = resolve_code_mode_config(&cfg); let multi_agent_v2 = resolve_multi_agent_v2_config(&cfg); - let apps_mcp_path_override = if features.enabled(Feature::AppsMcpPathOverride) { - let base = apps_mcp_path_override_toml_config(cfg.features.as_ref()); - base.and_then(|config| config.path.as_ref()) - .cloned() - .or_else(|| Some("/ps/mcp".to_string())) - } else { - None - }; let terminal_resize_reflow = resolve_terminal_resize_reflow_config(&cfg); let agent_roles = @@ -3553,7 +3546,6 @@ impl Config { chatgpt_base_url: cfg .chatgpt_base_url .unwrap_or("https://chatgpt.com/backend-api/".to_string()), - apps_mcp_path_override, apps_mcp_product_sku: cfg.apps_mcp_product_sku.clone(), realtime_audio: cfg .audio diff --git a/codex-rs/core/src/config_lock.rs b/codex-rs/core/src/config_lock.rs index 14ee7e1105..f99ded0bf8 100644 --- a/codex-rs/core/src/config_lock.rs +++ b/codex-rs/core/src/config_lock.rs @@ -125,6 +125,9 @@ fn config_lock_for_comparison( ) -> ConfigLockfileToml { let mut lockfile = lockfile.clone(); clear_config_lock_debug_controls(&mut lockfile.config); + if let Some(features) = lockfile.config.features.as_mut() { + features.clear_removed_compatibility_entries(); + } if options.allow_codex_version_mismatch { lockfile.codex_version.clear(); } diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 623cb44216..d8df728554 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -17,6 +17,7 @@ use codex_protocol::models::PermissionProfile; use codex_tools::DiscoverableTool; use rmcp::model::ToolAnnotations; use serde::Deserialize; +use tokio_util::sync::CancellationToken; use tracing::warn; use crate::config::Config; @@ -285,13 +286,15 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( let (tx_event, rx_event) = unbounded(); drop(rx_event); - let (mut mcp_connection_manager, cancel_token) = McpConnectionManager::new( + let cancel_token = CancellationToken::new(); + let mcp_connection_manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, auth_status_entries, &config.permissions.approval_policy, INITIAL_SUBMIT_ID.to_owned(), tx_event, + cancel_token.clone(), PermissionProfile::default(), // Connector discovery is threadless. Use an actually configured env if // one exists, but do not reintroduce the old hidden-local fallback. diff --git a/codex-rs/core/src/context/available_skills_instructions.rs b/codex-rs/core/src/context/available_skills_instructions.rs index f7921072b2..4e166d436c 100644 --- a/codex-rs/core/src/context/available_skills_instructions.rs +++ b/codex-rs/core/src/context/available_skills_instructions.rs @@ -5,12 +5,23 @@ use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use super::ContextualUserFragment; +/// Model-context fragment describing the skills available to Codex. #[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct AvailableSkillsInstructions { +pub struct AvailableSkillsInstructions { skill_root_lines: Vec, skill_lines: Vec, } +impl AvailableSkillsInstructions { + /// Creates a skills context fragment from pre-rendered catalog lines. + pub fn from_skill_lines(skill_lines: Vec) -> Self { + Self { + skill_root_lines: Vec::new(), + skill_lines, + } + } +} + impl From for AvailableSkillsInstructions { fn from(available_skills: AvailableSkills) -> Self { Self { diff --git a/codex-rs/core/src/context/environment_context_tests.rs b/codex-rs/core/src/context/environment_context_tests.rs index 49c3fc0c58..c8e491b68c 100644 --- a/codex-rs/core/src/context/environment_context_tests.rs +++ b/codex-rs/core/src/context/environment_context_tests.rs @@ -178,6 +178,7 @@ fn turn_context_item_filesystem_uses_workspace_roots_instead_of_cwd() { network: None, file_system_sandbox_policy: None, model: "gpt-5".to_string(), + comp_hash: None, personality: None, collaboration_mode: None, multi_agent_version: None, diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 78fc70fd49..a470c3afb7 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -23,6 +23,7 @@ mod realtime_end_instructions; mod realtime_start_instructions; mod realtime_start_with_instructions; mod subagent_notification; +mod token_budget_context; mod turn_aborted; mod user_instructions; mod user_shell_command; @@ -30,7 +31,7 @@ mod user_shell_command; pub(crate) use approved_command_prefix_saved::ApprovedCommandPrefixSaved; pub(crate) use apps_instructions::AppsInstructions; pub(crate) use available_plugins_instructions::AvailablePluginsInstructions; -pub(crate) use available_skills_instructions::AvailableSkillsInstructions; +pub use available_skills_instructions::AvailableSkillsInstructions; pub(crate) use codex_context_fragments::AdditionalContextDeveloperFragment; pub(crate) use codex_context_fragments::AdditionalContextUserFragment; pub use codex_context_fragments::ContextualUserFragment; @@ -60,6 +61,8 @@ pub(crate) use realtime_end_instructions::RealtimeEndInstructions; pub(crate) use realtime_start_instructions::RealtimeStartInstructions; pub(crate) use realtime_start_with_instructions::RealtimeStartWithInstructions; pub(crate) use subagent_notification::SubagentNotification; +pub(crate) use token_budget_context::TokenBudgetContext; +pub(crate) use token_budget_context::TokenBudgetRemainingContext; pub(crate) use turn_aborted::TurnAborted; pub(crate) use user_instructions::UserInstructions; pub(crate) use user_shell_command::UserShellCommand; diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs new file mode 100644 index 0000000000..df32ae4d71 --- /dev/null +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -0,0 +1,78 @@ +use super::ContextualUserFragment; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct TokenBudgetContext { + window_id: u64, + tokens_left: i64, +} + +impl TokenBudgetContext { + pub(crate) fn new(window_id: u64, tokens_left: i64) -> Self { + Self { + window_id, + tokens_left, + } + } +} + +impl ContextualUserFragment for TokenBudgetContext { + fn role(&self) -> &'static str { + "developer" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("\n", "\n") + } + + fn body(&self) -> String { + let window_id = self.window_id; + let tokens_left = self.tokens_left; + format!( + "Current context window {window_id}.\nYou have {tokens_left} tokens left in this context window." + ) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct TokenBudgetRemainingContext { + tokens_left: Option, +} + +impl TokenBudgetRemainingContext { + pub(crate) fn new(tokens_left: i64) -> Self { + Self { + tokens_left: Some(tokens_left), + } + } + + pub(crate) fn unknown() -> Self { + Self { tokens_left: None } + } +} + +impl ContextualUserFragment for TokenBudgetRemainingContext { + fn role(&self) -> &'static str { + "developer" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("\n", "\n") + } + + fn body(&self) -> String { + match self.tokens_left { + Some(tokens_left) => { + format!("You have {tokens_left} tokens left in this context window.") + } + None => "You have unknown tokens left in this context window.".to_string(), + } + } +} diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index bb22749cb3..f4bc87f5b8 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -130,6 +130,7 @@ fn reference_context_item() -> TurnContextItem { network: None, file_system_sandbox_policy: None, model: "gpt-test".to_string(), + comp_hash: None, personality: None, collaboration_mode: None, multi_agent_version: None, diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index d3cae7feed..e9383aef41 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -16,6 +16,7 @@ use codex_protocol::models::is_local_image_close_tag_text; use codex_protocol::models::is_local_image_open_tag_text; use codex_protocol::protocol::COLLABORATION_MODE_OPEN_TAG; use codex_protocol::protocol::REALTIME_CONVERSATION_OPEN_TAG; +use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use codex_protocol::user_input::UserInput; use tracing::warn; use uuid::Uuid; @@ -29,7 +30,9 @@ const CONTEXTUAL_DEVELOPER_PREFIXES: &[&str] = &[ "", COLLABORATION_MODE_OPEN_TAG, REALTIME_CONVERSATION_OPEN_TAG, + SKILLS_INSTRUCTIONS_OPEN_TAG, "", + "", ]; pub(crate) fn is_contextual_user_message_content(message: &[ContentItem]) -> bool { diff --git a/codex-rs/core/src/event_mapping_tests.rs b/codex-rs/core/src/event_mapping_tests.rs index 5f2ed66830..6f9dad64bb 100644 --- a/codex-rs/core/src/event_mapping_tests.rs +++ b/codex-rs/core/src/event_mapping_tests.rs @@ -1,3 +1,5 @@ +use super::has_non_contextual_dev_message_content; +use super::is_contextual_dev_message_content; use super::parse_turn_item; use crate::context::ContextualUserFragment; use crate::context::InternalContextSource; @@ -13,9 +15,30 @@ use codex_protocol::models::ReasoningItemContent; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ResponseItem; use codex_protocol::models::WebSearchAction; +use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use codex_protocol::user_input::UserInput; use pretty_assertions::assert_eq; +#[test] +fn recognizes_skills_instructions_as_contextual_developer_content() { + assert!(is_contextual_dev_message_content(&[ + ContentItem::InputText { + text: format!("{SKILLS_INSTRUCTIONS_OPEN_TAG}\n## Skills"), + }, + ])); +} + +#[test] +fn recognizes_token_budget_as_contextual_developer_content() { + let content = vec![ContentItem::InputText { + text: "\nYou have 710 tokens left in this context window.\n" + .to_string(), + }]; + + assert!(is_contextual_dev_message_content(&content)); + assert!(!has_non_contextual_dev_message_content(&content)); +} + #[test] fn parses_user_message_with_text_and_two_images() { let img1 = "https://example.com/one.png".to_string(); diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index b4920f1ff6..b06a2daad3 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -168,7 +168,7 @@ use prompt::render_guardian_transcript_entries; #[cfg(test)] use review::GuardianReviewOutcome; #[cfg(test)] -use review::run_guardian_review_session as run_guardian_review_session_for_test; +use review::run_guardian_review_session_with_retry as run_guardian_review_session_for_test; #[cfg(test)] use review_session::build_guardian_review_session_config as build_guardian_review_session_config_for_test; diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 934e68042f..99f9c5d999 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -7,6 +7,7 @@ use codex_analytics::GuardianReviewTrackContext; use codex_analytics::GuardianReviewedAction; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::GuardianAssessmentDecisionSource; use codex_protocol::protocol::GuardianAssessmentEvent; @@ -19,11 +20,14 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::WarningEvent; use std::sync::Arc; use tokio::sync::oneshot; +use tokio::time::Instant; +use tokio::time::sleep_until; use tokio_util::sync::CancellationToken; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::turn_timing::now_unix_timestamp_ms; +use crate::util::backoff; use super::AUTO_REVIEW_DENIAL_WINDOW_SIZE; use super::GUARDIAN_REVIEW_TIMEOUT; @@ -58,6 +62,8 @@ const GUARDIAN_TIMEOUT_INSTRUCTIONS: &str = concat!( "You may retry once, or ask the user for guidance or explicit approval.", ); +const GUARDIAN_REVIEW_MAX_ATTEMPTS: i64 = 3; + pub(crate) fn new_guardian_review_id() -> String { uuid::Uuid::new_v4().to_string() } @@ -95,9 +101,16 @@ pub(super) enum GuardianReviewOutcome { #[derive(Debug)] pub(super) enum GuardianReviewError { - PromptBuild { message: String }, - Session { message: String }, - Parse { message: String }, + PromptBuild { + message: String, + }, + Session { + message: String, + error_info: Option, + }, + Parse { + message: String, + }, Timeout, Cancelled, } @@ -112,6 +125,14 @@ impl GuardianReviewError { fn session(err: anyhow::Error) -> Self { Self::Session { message: err.to_string(), + error_info: None, + } + } + + fn session_with_error_info(err: anyhow::Error, error_info: CodexErrorInfo) -> Self { + Self::Session { + message: err.to_string(), + error_info: Some(error_info), } } @@ -335,13 +356,14 @@ async fn run_guardian_review( let schema = guardian_output_schema(); let terminal_action = action_summary.clone(); - let (outcome, analytics_result) = Box::pin(run_guardian_review_session( + let (outcome, analytics_result) = Box::pin(run_guardian_review_session_with_retry( session.clone(), turn.clone(), request, retry_reason.clone(), schema, external_cancel, + GUARDIAN_REVIEW_MAX_ATTEMPTS, )) .await; @@ -464,7 +486,7 @@ async fn run_guardian_review( | GuardianReviewError::Parse { .. } => { let message = match &error { GuardianReviewError::PromptBuild { message } - | GuardianReviewError::Session { message } + | GuardianReviewError::Session { message, .. } | GuardianReviewError::Parse { message } => message, GuardianReviewError::Timeout | GuardianReviewError::Cancelled => { "guardian review failed" @@ -657,13 +679,14 @@ pub(crate) fn spawn_approval_request_review( /// context. It may still reuse the parent's managed-network allowlist for /// read-only checks, but it intentionally runs without inherited exec-policy /// rules. -pub(super) async fn run_guardian_review_session( +async fn run_guardian_review_session_before_deadline( session: Arc, turn: Arc, request: GuardianApprovalRequest, retry_reason: Option, schema: serde_json::Value, external_cancel: Option, + deadline: Instant, ) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) { let network_proxy = session.services.network_proxy.load_full(); let live_network_config = match network_proxy.as_ref() { @@ -753,6 +776,7 @@ pub(super) async fn run_guardian_review_session( reasoning_summary: turn.reasoning_summary, personality: turn.personality, external_cancel, + deadline, }), ) .await; @@ -787,10 +811,16 @@ pub(super) async fn run_guardian_review_session( GuardianReviewOutcome::Error(GuardianReviewError::prompt_build(err)), session_analytics_result, ), - GuardianReviewSessionOutcome::SessionFailed(err) => ( - GuardianReviewOutcome::Error(GuardianReviewError::session(err)), - session_analytics_result, - ), + GuardianReviewSessionOutcome::SessionFailed { error, error_info } => { + let error = match error_info { + Some(error_info) => GuardianReviewError::session_with_error_info(error, error_info), + None => GuardianReviewError::session(error), + }; + ( + GuardianReviewOutcome::Error(error), + session_analytics_result, + ) + } GuardianReviewSessionOutcome::TimedOut => ( GuardianReviewOutcome::Error(GuardianReviewError::Timeout), session_analytics_result, @@ -802,9 +832,85 @@ pub(super) async fn run_guardian_review_session( } } +pub(super) async fn run_guardian_review_session_with_retry( + session: Arc, + turn: Arc, + request: GuardianApprovalRequest, + retry_reason: Option, + schema: serde_json::Value, + external_cancel: Option, + max_attempts: i64, +) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) { + assert!(max_attempts > 0, "guardian review must run at least once"); + let deadline = Instant::now() + GUARDIAN_REVIEW_TIMEOUT; + let mut attempt_count = 1; + loop { + let (outcome, mut analytics_result) = run_guardian_review_session_before_deadline( + Arc::clone(&session), + Arc::clone(&turn), + request.clone(), + retry_reason.clone(), + schema.clone(), + external_cancel.clone(), + deadline, + ) + .await; + analytics_result.attempt_count = attempt_count; + if attempt_count >= max_attempts || !should_retry_guardian_review(&outcome) { + return (outcome, analytics_result); + } + if let Some(error) = + wait_before_guardian_retry(attempt_count, deadline, external_cancel.as_ref()).await + { + return (GuardianReviewOutcome::Error(error), analytics_result); + } + attempt_count += 1; + } +} + +async fn wait_before_guardian_retry( + attempt_count: i64, + deadline: Instant, + external_cancel: Option<&CancellationToken>, +) -> Option { + let retry_delay = backoff(attempt_count as u64); + let retry_at = (Instant::now() + retry_delay).min(deadline); + tokio::select! { + _ = sleep_until(retry_at) => { + (Instant::now() >= deadline).then_some(GuardianReviewError::Timeout) + } + _ = async { + if let Some(cancel_token) = external_cancel { + cancel_token.cancelled().await; + } else { + std::future::pending::<()>().await; + } + } => Some(GuardianReviewError::Cancelled), + } +} + +fn should_retry_guardian_review(outcome: &GuardianReviewOutcome) -> bool { + matches!( + outcome, + GuardianReviewOutcome::Error( + GuardianReviewError::Session { + error_info: Some( + CodexErrorInfo::ServerOverloaded + | CodexErrorInfo::HttpConnectionFailed { .. } + | CodexErrorInfo::ResponseStreamConnectionFailed { .. } + | CodexErrorInfo::InternalServerError + | CodexErrorInfo::ResponseStreamDisconnected { .. } + ), + .. + } | GuardianReviewError::Parse { .. } + ) + ) +} + #[cfg(test)] mod review_tests { use super::*; + use std::time::Duration; #[test] fn guardian_review_error_reason_distinguishes_error_kinds() { @@ -812,6 +918,10 @@ mod review_tests { let prompt_error = GuardianReviewError::prompt_build(anyhow::anyhow!("bad prompt/config")); let session_error = GuardianReviewError::session(anyhow::anyhow!("guardian runtime failed")); + let structured_session_error = GuardianReviewError::session_with_error_info( + anyhow::anyhow!("temporary guardian failure"), + CodexErrorInfo::ServerOverloaded, + ); assert!(matches!( parse_error.failure_reason(), @@ -825,5 +935,109 @@ mod review_tests { session_error.failure_reason(), GuardianReviewFailureReason::SessionError )); + assert!(matches!( + structured_session_error.failure_reason(), + GuardianReviewFailureReason::SessionError + )); + } + + #[test] + fn guardian_review_retry_only_retries_transient_session_and_parse_errors() { + let assessment = GuardianAssessment { + risk_level: GuardianRiskLevel::High, + user_authorization: GuardianUserAuthorization::Unknown, + outcome: GuardianAssessmentOutcome::Deny, + rationale: "deny".to_string(), + }; + let transient_error_info = [ + CodexErrorInfo::ServerOverloaded, + CodexErrorInfo::HttpConnectionFailed { + http_status_code: Some(502), + }, + CodexErrorInfo::ResponseStreamConnectionFailed { + http_status_code: Some(503), + }, + CodexErrorInfo::InternalServerError, + CodexErrorInfo::ResponseStreamDisconnected { + http_status_code: None, + }, + ]; + let mut outcomes = transient_error_info + .into_iter() + .map(|error_info| { + ( + GuardianReviewOutcome::Error(GuardianReviewError::session_with_error_info( + anyhow::anyhow!("transient session"), + error_info, + )), + true, + ) + }) + .collect::>(); + outcomes.extend([ + (GuardianReviewOutcome::Completed(assessment), false), + ( + GuardianReviewOutcome::Error(GuardianReviewError::prompt_build(anyhow::anyhow!( + "prompt" + ))), + false, + ), + ( + GuardianReviewOutcome::Error(GuardianReviewError::session(anyhow::anyhow!( + "session" + ))), + false, + ), + ( + GuardianReviewOutcome::Error(GuardianReviewError::session_with_error_info( + anyhow::anyhow!("bad request"), + CodexErrorInfo::BadRequest, + )), + false, + ), + ( + GuardianReviewOutcome::Error(GuardianReviewError::parse(anyhow::anyhow!("parse"))), + true, + ), + ( + GuardianReviewOutcome::Error(GuardianReviewError::Timeout), + false, + ), + ( + GuardianReviewOutcome::Error(GuardianReviewError::Cancelled), + false, + ), + ]); + + for (outcome, expected) in outcomes { + assert_eq!(should_retry_guardian_review(&outcome), expected); + } + } + + #[tokio::test] + async fn guardian_review_retry_wait_honors_cancellation() { + let cancel_token = CancellationToken::new(); + cancel_token.cancel(); + + let error = wait_before_guardian_retry( + /*attempt_count*/ 1, + Instant::now() + Duration::from_secs(/*secs*/ 1), + Some(&cancel_token), + ) + .await; + + assert!(matches!(error, Some(GuardianReviewError::Cancelled))); + } + + #[tokio::test] + async fn guardian_review_retry_wait_honors_deadline() { + let error = wait_before_guardian_retry( + /*attempt_count*/ 1, + Instant::now(), + /*external_cancel*/ None, + ) + .await; + + assert!(matches!(error, Some(GuardianReviewError::Timeout))); } } diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index fcda9a5a23..54282c1c0a 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -15,6 +15,8 @@ use codex_protocol::models::PermissionProfile; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::CodexErrorInfo; +use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InitialHistory; @@ -46,7 +48,6 @@ use codex_features::Feature; use codex_model_provider_info::ModelProviderInfo; use codex_utils_absolute_path::AbsolutePathBuf; -use super::GUARDIAN_REVIEW_TIMEOUT; use super::GUARDIAN_REVIEWER_NAME; use super::GuardianApprovalRequest; use super::prompt::GuardianPromptMode; @@ -60,7 +61,10 @@ const GUARDIAN_INTERRUPT_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); pub(crate) enum GuardianReviewSessionOutcome { Completed(anyhow::Result>), PromptBuildFailed(anyhow::Error), - SessionFailed(anyhow::Error), + SessionFailed { + error: anyhow::Error, + error_info: Option, + }, TimedOut, Aborted, } @@ -77,6 +81,7 @@ pub(crate) struct GuardianReviewSessionParams { pub(crate) reasoning_summary: ReasoningSummaryConfig, pub(crate) personality: Option, pub(crate) external_cancel: Option, + pub(crate) deadline: tokio::time::Instant, } #[derive(Default)] @@ -312,7 +317,7 @@ impl GuardianReviewSessionManager { &self, params: GuardianReviewSessionParams, ) -> (GuardianReviewSessionOutcome, GuardianReviewAnalyticsResult) { - let deadline = tokio::time::Instant::now() + GUARDIAN_REVIEW_TIMEOUT; + let deadline = params.deadline; let next_reuse_key = GuardianReviewSessionReuseKey::from_spawn_config(¶ms.spawn_config); let mut stale_trunk_to_shutdown = None; let mut spawned_trunk = false; @@ -364,7 +369,9 @@ impl GuardianReviewSessionManager { state.trunk.as_ref().cloned() } - Err(outcome) => return (outcome, GuardianReviewAnalyticsResult::without_session()), + Err(outcome) => { + return (outcome, GuardianReviewAnalyticsResult::without_session()); + } }; if let Some(review_session) = stale_trunk_to_shutdown { @@ -560,7 +567,9 @@ impl GuardianReviewSessionManager { GuardianReviewAnalyticsResult::without_session(), ); } - Err(outcome) => return (outcome, GuardianReviewAnalyticsResult::without_session()), + Err(outcome) => { + return (outcome, GuardianReviewAnalyticsResult::without_session()); + } }; self.register_active_ephemeral(Arc::clone(&review_session)) .await; @@ -764,7 +773,10 @@ async fn run_review_on_session( Ok(Ok(child_turn_id)) => child_turn_id, Ok(Err(err)) => { return ( - GuardianReviewSessionOutcome::SessionFailed(err.into()), + GuardianReviewSessionOutcome::SessionFailed { + error: err.into(), + error_info: None, + }, false, analytics_result, ); @@ -825,7 +837,7 @@ async fn wait_for_guardian_review( ) -> (GuardianReviewSessionOutcome, bool, bool) { let timeout = tokio::time::sleep_until(deadline); tokio::pin!(timeout); - let mut last_error_message: Option = None; + let mut last_error: Option = None; loop { tokio::select! { @@ -862,10 +874,13 @@ async fn wait_for_guardian_review( .time_to_first_token_ms .and_then(|ms| u64::try_from(ms).ok()); if turn_complete.last_agent_message.is_none() - && let Some(error_message) = last_error_message + && let Some(error) = last_error { return ( - GuardianReviewSessionOutcome::Completed(Err(anyhow!(error_message))), + GuardianReviewSessionOutcome::SessionFailed { + error: anyhow!(error.message), + error_info: error.codex_error_info, + }, true, true, ); @@ -877,7 +892,7 @@ async fn wait_for_guardian_review( ); } EventMsg::Error(error) => { - last_error_message = Some(error.message); + last_error = Some(error); } EventMsg::TurnAborted(_) => { return (GuardianReviewSessionOutcome::Aborted, true, false); @@ -920,6 +935,8 @@ pub(crate) fn build_guardian_review_session_config( let mut guardian_config = parent_config.clone(); guardian_config.model = Some(active_model.to_string()); guardian_config.model_reasoning_effort = reasoning_effort; + guardian_config.model_provider.request_max_retries = Some(1); + guardian_config.model_provider.stream_max_retries = Some(1); guardian_config.include_skill_instructions = false; guardian_config.base_instructions = Some( parent_config @@ -1148,6 +1165,7 @@ mod tests { reasoning_summary, personality, external_cancel: None, + deadline: tokio::time::Instant::now() + Duration::from_secs(30), } } @@ -1473,6 +1491,29 @@ mod tests { assert!(keep_review_session); } + #[tokio::test] + async fn run_review_removes_trunk_when_event_stream_is_broken() { + let (mut review_session, tx_event, _rx_sub) = test_review_session().await; + let params = test_review_params().await; + review_session.reuse_key = + GuardianReviewSessionReuseKey::from_spawn_config(¶ms.spawn_config); + let manager = GuardianReviewSessionManager { + state: Arc::new(Mutex::new(GuardianReviewSessionState { + trunk: Some(Arc::new(review_session)), + ephemeral_reviews: Vec::new(), + })), + }; + drop(tx_event); + + let (outcome, _) = manager.run_review(params).await; + + assert!(matches!( + outcome, + GuardianReviewSessionOutcome::Completed(Err(_)) + )); + assert!(manager.state.lock().await.trunk.is_none()); + } + #[tokio::test] async fn wait_for_guardian_review_ignores_prior_turn_completion() { let (review_session, tx_event, _rx_sub) = test_review_session().await; @@ -1545,6 +1586,47 @@ mod tests { assert!(capture_token_usage); } + #[tokio::test] + async fn wait_for_guardian_review_preserves_structured_session_error() { + let (review_session, tx_event, _rx_sub) = test_review_session().await; + tx_event + .send(Event { + id: "current-turn".to_string(), + msg: EventMsg::Error(ErrorEvent { + message: "temporary failure".to_string(), + codex_error_info: Some(CodexErrorInfo::ServerOverloaded), + }), + }) + .await + .expect("queue guardian error"); + tx_event + .send(turn_complete_event( + "current-turn", + /*last_agent_message*/ None, + Some(42), + )) + .await + .expect("queue current turn completion"); + + let mut analytics_result = GuardianReviewAnalyticsResult::without_session(); + let (outcome, keep_review_session, capture_token_usage) = wait_for_guardian_review( + &review_session, + "current-turn", + tokio::time::Instant::now() + Duration::from_secs(1), + /*external_cancel*/ None, + &mut analytics_result, + ) + .await; + + let GuardianReviewSessionOutcome::SessionFailed { error, error_info } = outcome else { + panic!("expected structured session failure"); + }; + assert_eq!(error.to_string(), "temporary failure"); + assert_eq!(error_info, Some(CodexErrorInfo::ServerOverloaded)); + assert!(keep_review_session); + assert!(capture_token_usage); + } + #[tokio::test] async fn wait_for_guardian_review_ignores_prior_turn_aborts() { let (review_session, tx_event, _rx_sub) = test_review_session().await; diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index ffadd801ca..8f6e01298c 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -55,10 +55,11 @@ use core_test_support::context_snapshot::ContextSnapshotOptions; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_response_created; -use core_test_support::responses::mount_response_once; +use core_test_support::responses::mount_response_sequence; use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; +use core_test_support::responses::sse_failed; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::streaming_sse::StreamingSseChunk; @@ -172,6 +173,51 @@ async fn guardian_test_session_and_turn( guardian_test_session_and_turn_with_base_url(server.uri().as_str()).await } +async fn guardian_test_session_turn_and_rx( + server: &wiremock::MockServer, +) -> ( + Arc, + Arc, + async_channel::Receiver, +) { + let (mut session, mut turn, rx) = + crate::session::tests::make_session_and_context_with_rx().await; + Arc::get_mut(&mut session) + .expect("session should be uniquely owned") + .thread_id = fixed_guardian_parent_session_id(); + let mut config = (*turn.config).clone(); + config.model_provider.base_url = Some(format!("{}/v1", server.uri())); + config.user_instructions = None; + let config = Arc::new(config); + let models_manager = test_support::models_manager_with_provider( + config.codex_home.to_path_buf(), + Arc::clone(&session.services.auth_manager), + config.model_provider.clone(), + ); + Arc::get_mut(&mut session) + .expect("session should be uniquely owned") + .services + .models_manager = models_manager; + let turn_mut = Arc::get_mut(&mut turn).expect("turn should be uniquely owned"); + turn_mut.config = Arc::clone(&config); + turn_mut.provider = + create_model_provider(config.model_provider.clone(), turn_mut.auth_manager.clone()); + turn_mut.user_instructions = None; + + (session, turn, rx) +} + +fn guardian_shell_request(id: &str) -> GuardianApprovalRequest { + GuardianApprovalRequest::Shell { + id: id.to_string(), + command: vec!["git".to_string(), "push".to_string()], + cwd: test_path_buf("/repo/codex-rs/core").abs(), + sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault, + additional_permissions: None, + justification: Some("Need to push the reviewed docs fix.".to_string()), + } +} + async fn guardian_test_session_and_turn_with_base_url( base_url: &str, ) -> (Arc, Arc) { @@ -1343,6 +1389,7 @@ async fn guardian_request_model_for_auto_review_override( Some("Sandbox denied outbound git push to github.com.".to_string()), guardian_output_schema(), /*external_cancel*/ None, + /*max_attempts*/ 1, ) .await; let (GuardianReviewOutcome::Completed(_), _) = outcome else { @@ -1454,6 +1501,7 @@ async fn guardian_review_request_layout_matches_model_visible_request_snapshot() Some("Sandbox denied outbound git push to github.com.".to_string()), guardian_output_schema(), /*external_cancel*/ None, + /*max_attempts*/ 1, ) .await; let (GuardianReviewOutcome::Completed(assessment), metadata) = outcome else { @@ -1632,6 +1680,7 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: Some("First retry reason".to_string()), guardian_output_schema(), /*external_cancel*/ None, + /*max_attempts*/ 1, ) .await; session @@ -1676,6 +1725,7 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: Some("Second retry reason".to_string()), guardian_output_schema(), /*external_cancel*/ None, + /*max_attempts*/ 1, ) .await; session @@ -1716,6 +1766,7 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: Some("Third retry reason".to_string()), guardian_output_schema(), /*external_cancel*/ None, + /*max_attempts*/ 1, ) .await; @@ -1907,6 +1958,7 @@ async fn guardian_reused_trunk_ignores_stale_prior_turn_completion() -> anyhow:: /*retry_reason*/ None, guardian_output_schema(), /*external_cancel*/ None, + /*max_attempts*/ 1, ) .await; let (GuardianReviewOutcome::Completed(first_assessment), first_metadata) = first_outcome else { @@ -1949,6 +2001,7 @@ async fn guardian_reused_trunk_ignores_stale_prior_turn_completion() -> anyhow:: /*retry_reason*/ None, guardian_output_schema(), /*external_cancel*/ None, + /*max_attempts*/ 1, ) .await; let (GuardianReviewOutcome::Completed(second_assessment), second_metadata) = second_outcome @@ -1978,15 +2031,17 @@ async fn guardian_review_surfaces_responses_api_errors_in_rejection_reason() -> let server = start_mock_server().await; let error_message = "Item 'rs_test' of type 'reasoning' was provided without its required following item."; - let _request_log = mount_response_once( + let request_log = mount_response_sequence( &server, - wiremock::ResponseTemplate::new(400).set_body_json(serde_json::json!({ - "error": { - "message": error_message, - "type": "invalid_request_error", - "param": "input" - } - })), + vec![ + wiremock::ResponseTemplate::new(400).set_body_json(serde_json::json!({ + "error": { + "message": error_message, + "type": "invalid_request_error", + "param": "input" + } + })), + ], ) .await; @@ -2030,6 +2085,7 @@ async fn guardian_review_surfaces_responses_api_errors_in_rejection_reason() -> .await; assert_eq!(decision, ReviewDecision::Denied); + assert_eq!(request_log.requests().len(), 1); let mut warnings = Vec::new(); let mut denial_rationales = Vec::new(); @@ -2080,13 +2136,255 @@ async fn guardian_review_surfaces_responses_api_errors_in_rejection_reason() -> Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn guardian_review_retries_transient_session_failure_then_approves() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let approval = serde_json::json!({ + "risk_level": "low", + "user_authorization": "high", + "outcome": "allow", + "rationale": "retry succeeded", + }) + .to_string(); + let request_log = mount_sse_sequence( + &server, + vec![ + sse_failed( + "resp-session-failure", + "server_is_overloaded", + "temporary reviewer overload", + ), + sse(vec![ + ev_response_created("resp-approved"), + ev_assistant_message("msg-approved", &approval), + ev_completed("resp-approved"), + ]), + ], + ) + .await; + let (session, turn) = guardian_test_session_and_turn(&server).await; + seed_guardian_parent_history(&session, &turn).await; + + let (outcome, metadata) = run_guardian_review_session_for_test( + Arc::clone(&session), + Arc::clone(&turn), + guardian_shell_request("shell-session-retry"), + /*retry_reason*/ None, + guardian_output_schema(), + /*external_cancel*/ None, + /*max_attempts*/ 3, + ) + .await; + + let GuardianReviewOutcome::Completed(assessment) = outcome else { + panic!("expected guardian assessment"); + }; + assert_eq!(assessment.outcome, GuardianAssessmentOutcome::Allow); + assert_eq!(assessment.rationale, "retry succeeded"); + assert_eq!(metadata.attempt_count, 2); + assert!(matches!( + metadata.guardian_session_kind, + Some(codex_analytics::GuardianReviewSessionKind::TrunkReused) + )); + assert_eq!(request_log.requests().len(), 2); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn guardian_review_does_not_retry_missing_assessment_payload() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let request_log = mount_sse_sequence( + &server, + vec![sse(vec![ + ev_response_created("resp-missing-assessment"), + ev_completed("resp-missing-assessment"), + ])], + ) + .await; + let (session, turn) = guardian_test_session_and_turn(&server).await; + seed_guardian_parent_history(&session, &turn).await; + + let decision = review_approval_request( + &session, + &turn, + "review-missing-assessment".to_string(), + guardian_shell_request("shell-missing-assessment"), + /*retry_reason*/ None, + ) + .await; + + assert_eq!(decision, ReviewDecision::Denied); + assert_eq!(request_log.requests().len(), 1); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn guardian_review_retries_two_parse_failures_then_approves() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let approval = serde_json::json!({ + "risk_level": "low", + "user_authorization": "high", + "outcome": "allow", + "rationale": "retry succeeded", + }) + .to_string(); + let request_log = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-parse-failure-1"), + ev_assistant_message("msg-parse-failure-1", "not valid guardian json"), + ev_completed("resp-parse-failure-1"), + ]), + sse(vec![ + ev_response_created("resp-parse-failure-2"), + ev_assistant_message("msg-parse-failure-2", "still not valid guardian json"), + ev_completed("resp-parse-failure-2"), + ]), + sse(vec![ + ev_response_created("resp-approved"), + ev_assistant_message("msg-approved", &approval), + ev_completed("resp-approved"), + ]), + ], + ) + .await; + let (session, turn) = guardian_test_session_and_turn(&server).await; + seed_guardian_parent_history(&session, &turn).await; + + let (outcome, metadata) = run_guardian_review_session_for_test( + Arc::clone(&session), + Arc::clone(&turn), + guardian_shell_request("shell-parse-retry"), + /*retry_reason*/ None, + guardian_output_schema(), + /*external_cancel*/ None, + /*max_attempts*/ 3, + ) + .await; + + let GuardianReviewOutcome::Completed(assessment) = outcome else { + panic!("expected guardian assessment"); + }; + assert_eq!(assessment.outcome, GuardianAssessmentOutcome::Allow); + assert_eq!(assessment.rationale, "retry succeeded"); + assert_eq!(metadata.attempt_count, 3); + assert!(matches!( + metadata.guardian_session_kind, + Some(codex_analytics::GuardianReviewSessionKind::TrunkReused) + )); + assert_eq!(request_log.requests().len(), 3); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn guardian_review_exhausts_three_failures_with_one_terminal_event() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let request_log = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-parse-failure-1"), + ev_assistant_message("msg-parse-failure-1", "invalid one"), + ev_completed("resp-parse-failure-1"), + ]), + sse(vec![ + ev_response_created("resp-parse-failure-2"), + ev_assistant_message("msg-parse-failure-2", "invalid two"), + ev_completed("resp-parse-failure-2"), + ]), + sse(vec![ + ev_response_created("resp-parse-failure-3"), + ev_assistant_message("msg-parse-failure-3", "invalid three"), + ev_completed("resp-parse-failure-3"), + ]), + ], + ) + .await; + let (session, turn, rx) = guardian_test_session_turn_and_rx(&server).await; + seed_guardian_parent_history(&session, &turn).await; + + let decision = review_approval_request( + &session, + &turn, + "review-exhausted-retry".to_string(), + guardian_shell_request("shell-exhausted-retry"), + /*retry_reason*/ None, + ) + .await; + + assert_eq!(decision, ReviewDecision::Denied); + assert_eq!(request_log.requests().len(), 3); + let mut statuses = Vec::new(); + while let Ok(event) = rx.try_recv() { + if let EventMsg::GuardianAssessment(event) = event.msg { + statuses.push(event.status); + } + } + assert_eq!( + statuses, + vec![ + GuardianAssessmentStatus::InProgress, + GuardianAssessmentStatus::Denied, + ] + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn guardian_review_does_not_retry_valid_denial() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let denial = serde_json::json!({ + "risk_level": "high", + "user_authorization": "unknown", + "outcome": "deny", + "rationale": "unsafe", + }) + .to_string(); + let request_log = mount_sse_sequence( + &server, + vec![sse(vec![ + ev_response_created("resp-denied"), + ev_assistant_message("msg-denied", &denial), + ev_completed("resp-denied"), + ])], + ) + .await; + let (session, turn) = guardian_test_session_and_turn(&server).await; + seed_guardian_parent_history(&session, &turn).await; + + let decision = review_approval_request( + &session, + &turn, + "review-valid-denial".to_string(), + guardian_shell_request("shell-valid-denial"), + /*retry_reason*/ None, + ) + .await; + + assert_eq!(decision, ReviewDecision::Denied); + assert_eq!(request_log.requests().len(), 1); + Ok(()) +} + #[tokio::test] -async fn guardian_parallel_reviews_fork_from_last_committed_trunk_history() -> anyhow::Result<()> { +async fn guardian_ephemeral_retry_preserves_parallel_trunk_and_fork_history() -> anyhow::Result<()> +{ const TEST_STACK_SIZE_BYTES: usize = 4 * 1024 * 1024; let handle = std::thread::Builder::new() - .name("guardian_parallel_reviews_fork_from_last_committed_trunk_history".to_string()) + .name("guardian_ephemeral_retry_preserves_parallel_trunk_and_fork_history".to_string()) .stack_size(TEST_STACK_SIZE_BYTES) .spawn(|| -> anyhow::Result<()> { let runtime = tokio::runtime::Builder::new_current_thread() @@ -2141,10 +2439,18 @@ async fn guardian_parallel_reviews_fork_from_last_committed_trunk_history() -> a gate: None, body: sse(vec![ ev_response_created("resp-guardian-3"), - ev_assistant_message("msg-guardian-3", &third_assessment), + ev_assistant_message("msg-guardian-3", "not valid guardian json"), ev_completed("resp-guardian-3"), ]), }], + vec![StreamingSseChunk { + gate: None, + body: sse(vec![ + ev_response_created("resp-guardian-4"), + ev_assistant_message("msg-guardian-4", &third_assessment), + ev_completed("resp-guardian-4"), + ]), + }], ]) .await; @@ -2271,20 +2577,28 @@ async fn guardian_parallel_reviews_fork_from_last_committed_trunk_history() -> a .await; assert_eq!(third_decision, ReviewDecision::Approved); let requests = server.requests().await; - assert_eq!(requests.len(), 3); + assert_eq!(requests.len(), 4); let second_request_body = serde_json::from_slice::(&requests[1])?; - let third_request_body = serde_json::from_slice::(&requests[2])?; + let failed_ephemeral_request_body = + serde_json::from_slice::(&requests[2])?; + let retried_ephemeral_request_body = + serde_json::from_slice::(&requests[3])?; assert_eq!( second_request_body["prompt_cache_key"], - third_request_body["prompt_cache_key"], + failed_ephemeral_request_body["prompt_cache_key"], "forked guardian review should reuse the trunk guardian prompt cache key" ); - let third_request_body_text = third_request_body.to_string(); + assert_eq!( + failed_ephemeral_request_body["prompt_cache_key"], + retried_ephemeral_request_body["prompt_cache_key"], + "retried ephemeral review should preserve the guardian prompt cache key" + ); + let third_request_body_text = retried_ephemeral_request_body.to_string(); assert!( third_request_body_text.contains("first guardian rationale"), "forked guardian review should include the last committed trunk assessment" ); - let third_user_message = last_user_message_text_from_body(&third_request_body); + let third_user_message = last_user_message_text_from_body(&retried_ephemeral_request_body); assert!(third_user_message.contains(">>> TRANSCRIPT DELTA START\n")); assert!( third_user_message.contains("[5] user: Please inspect pending changes before pushing.") @@ -2315,7 +2629,7 @@ async fn guardian_parallel_reviews_fork_from_last_committed_trunk_history() -> a match handle.join() { Ok(result) => result, Err(_) => Err(anyhow::anyhow!( - "guardian_parallel_reviews_fork_from_last_committed_trunk_history thread panicked" + "guardian_ephemeral_retry_preserves_parallel_trunk_and_fork_history thread panicked" )), } } @@ -2541,6 +2855,10 @@ async fn guardian_review_session_config_keeps_bedrock_provider_for_bedrock_gpt_5 ) .expect("guardian config"); + let mut expected_model_provider = + ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None); + expected_model_provider.request_max_retries = Some(1); + expected_model_provider.stream_max_retries = Some(1); assert_eq!( ( guardian_config.model, @@ -2550,7 +2868,7 @@ async fn guardian_review_session_config_keeps_bedrock_provider_for_bedrock_gpt_5 ( Some(AMAZON_BEDROCK_GPT_5_4_MODEL_ID.to_string()), AMAZON_BEDROCK_PROVIDER_ID.to_string(), - ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + expected_model_provider, ) ); } diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 6e2d4fe33c..1097f7bc78 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -36,6 +36,7 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_thread_store::ReadThreadParams; use serde_json::Value; +use tracing::instrument; use crate::context::ContextualUserFragment; use crate::context::HookAdditionalContext; @@ -97,6 +98,7 @@ impl From for ContextInjectingHookOutcome { } } +#[instrument(level = "trace", skip_all)] pub(crate) async fn run_pending_session_start_hooks( sess: &Arc, turn_context: &Arc, @@ -291,6 +293,7 @@ pub(crate) async fn run_post_tool_use_hooks( outcome } +#[instrument(level = "trace", skip_all)] pub(crate) async fn run_turn_stop_hooks( sess: &Arc, turn_context: &Arc, @@ -382,9 +385,7 @@ pub(crate) async fn run_pre_compact_hooks( let outcome = sess.hooks().run_pre_compact(request).await; emit_hook_completed_events(sess, turn_context, outcome.hook_events).await; if outcome.should_stop { - PreCompactHookOutcome::Stopped { - reason: outcome.stop_reason, - } + PreCompactHookOutcome::Stopped } else { PreCompactHookOutcome::Continue } @@ -392,7 +393,7 @@ pub(crate) async fn run_pre_compact_hooks( pub(crate) enum PreCompactHookOutcome { Continue, - Stopped { reason: Option }, + Stopped, } pub(crate) enum PostCompactHookOutcome { @@ -427,6 +428,7 @@ pub(crate) async fn run_post_compact_hooks( } } +#[instrument(level = "trace", skip_all)] pub(crate) async fn run_legacy_after_agent_hook( sess: &Arc, turn_context: &Arc, diff --git a/codex-rs/core/src/image_preparation.rs b/codex-rs/core/src/image_preparation.rs new file mode 100644 index 0000000000..d992adb809 --- /dev/null +++ b/codex-rs/core/src/image_preparation.rs @@ -0,0 +1,122 @@ +use codex_protocol::models::ContentItem; +use codex_protocol::models::FunctionCallOutputContentItem; +use codex_protocol::models::ImageDetail; +use codex_protocol::models::ResponseItem; +use codex_utils_image::ImageProcessingError; +use codex_utils_image::PromptImageMode; +use codex_utils_image::PromptImageResizeLimits; +use codex_utils_image::load_data_url_for_prompt; +use tracing::warn; + +pub(crate) const IMAGE_PROCESSING_ERROR_PLACEHOLDER: &str = + "image content omitted because it could not be processed"; +const IMAGE_TOO_LARGE_PLACEHOLDER: &str = + "image content omitted because it exceeded the supported size limit; use a smaller image"; +const UNSUPPORTED_LOW_DETAIL_PLACEHOLDER: &str = "image content omitted because detail 'low' is not supported; use 'high', 'original', or 'auto'"; + +const HIGH_DETAIL_LIMITS: PromptImageResizeLimits = PromptImageResizeLimits { + max_dimension: 2048, + max_patches: 2_500, +}; +const ORIGINAL_DETAIL_LIMITS: PromptImageResizeLimits = PromptImageResizeLimits { + max_dimension: 6000, + max_patches: 10_000, +}; +#[derive(Debug, thiserror::Error)] +enum ImagePreparationError { + #[error("image detail `low` is not supported")] + UnsupportedLowDetail, + #[error(transparent)] + Processing(#[from] ImageProcessingError), +} + +impl ImagePreparationError { + fn placeholder(&self) -> &'static str { + match self { + ImagePreparationError::UnsupportedLowDetail => UNSUPPORTED_LOW_DETAIL_PLACEHOLDER, + ImagePreparationError::Processing(ImageProcessingError::ImageTooLarge { .. }) => { + IMAGE_TOO_LARGE_PLACEHOLDER + } + ImagePreparationError::Processing(_) => IMAGE_PROCESSING_ERROR_PLACEHOLDER, + } + } +} + +pub(crate) fn prepare_response_items(items: &mut [ResponseItem]) { + for item in items { + match item { + ResponseItem::Message { content, .. } => prepare_message_content(content), + ResponseItem::FunctionCallOutput { output, .. } + | ResponseItem::CustomToolCallOutput { output, .. } => { + if let Some(content) = output.content_items_mut() { + prepare_tool_output_content(content); + } + } + ResponseItem::Reasoning { .. } + | ResponseItem::AgentMessage { .. } + | ResponseItem::LocalShellCall { .. } + | ResponseItem::FunctionCall { .. } + | ResponseItem::ToolSearchCall { .. } + | ResponseItem::CustomToolCall { .. } + | ResponseItem::ToolSearchOutput { .. } + | ResponseItem::WebSearchCall { .. } + | ResponseItem::ImageGenerationCall { .. } + | ResponseItem::Compaction { .. } + | ResponseItem::CompactionTrigger + | ResponseItem::ContextCompaction { .. } + | ResponseItem::Other => {} + } + } +} + +fn prepare_message_content(items: &mut [ContentItem]) { + for item in items { + if let ContentItem::InputImage { image_url, detail } = item + && is_data_url(image_url) + && let Err(error) = prepare_image(image_url, *detail) + { + warn!(%error, "failed to prepare message image"); + *item = ContentItem::InputText { + text: error.placeholder().to_string(), + }; + } + } +} + +fn prepare_tool_output_content(items: &mut [FunctionCallOutputContentItem]) { + for item in items { + if let FunctionCallOutputContentItem::InputImage { image_url, detail } = item + && is_data_url(image_url) + && let Err(error) = prepare_image(image_url, *detail) + { + warn!(%error, "failed to prepare tool output image"); + *item = FunctionCallOutputContentItem::InputText { + text: error.placeholder().to_string(), + }; + } + } +} + +fn is_data_url(image_url: &str) -> bool { + image_url + .get(.."data:".len()) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case("data:")) +} + +fn prepare_image( + image_url: &mut String, + detail: Option, +) -> Result<(), ImagePreparationError> { + let limits = match detail { + None | Some(ImageDetail::Auto | ImageDetail::High) => HIGH_DETAIL_LIMITS, + Some(ImageDetail::Original) => ORIGINAL_DETAIL_LIMITS, + Some(ImageDetail::Low) => return Err(ImagePreparationError::UnsupportedLowDetail), + }; + let image = load_data_url_for_prompt(image_url, PromptImageMode::ResizeWithLimits(limits))?; + *image_url = image.into_data_url(); + Ok(()) +} + +#[cfg(test)] +#[path = "image_preparation_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/image_preparation_tests.rs b/codex-rs/core/src/image_preparation_tests.rs new file mode 100644 index 0000000000..1d54973cc2 --- /dev/null +++ b/codex-rs/core/src/image_preparation_tests.rs @@ -0,0 +1,193 @@ +use std::io::Cursor; + +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_protocol::models::FunctionCallOutputBody; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_utils_image::data_url_from_bytes; +use image::DynamicImage; +use image::GenericImageView; +use image::ImageBuffer; +use image::ImageFormat; +use image::Rgba; +use pretty_assertions::assert_eq; + +use super::*; + +fn png_data_url(width: u32, height: u32) -> (String, Vec) { + let image = ImageBuffer::from_pixel(width, height, Rgba([10u8, 20, 30, 255])); + let mut encoded = Cursor::new(Vec::new()); + DynamicImage::ImageRgba8(image) + .write_to(&mut encoded, ImageFormat::Png) + .expect("encode PNG"); + let bytes = encoded.into_inner(); + (data_url_from_bytes("image/png", &bytes), bytes) +} + +fn decoded_image(image_url: &str) -> (Vec, DynamicImage) { + let (_, payload) = image_url.split_once(',').expect("data URL payload"); + let bytes = BASE64_STANDARD.decode(payload).expect("decode image URL"); + let image = image::load_from_memory(&bytes).expect("decode processed image"); + (bytes, image) +} + +#[test] +fn preparation_preserves_small_image_bytes_and_non_data_urls() { + let (data_url, original_bytes) = png_data_url(/*width*/ 64, /*height*/ 32); + let http_url = "https://example.com/image.png".to_string(); + let mut items = vec![ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ + ContentItem::InputImage { + image_url: data_url, + detail: Some(ImageDetail::High), + }, + ContentItem::InputImage { + image_url: http_url.clone(), + detail: Some(ImageDetail::Low), + }, + ], + phase: None, + }]; + + prepare_response_items(&mut items); + + let ResponseItem::Message { content, .. } = &items[0] else { + panic!("expected message"); + }; + let [ + ContentItem::InputImage { image_url, .. }, + ContentItem::InputImage { + image_url: preserved_http_url, + detail: Some(ImageDetail::Low), + }, + ] = content.as_slice() + else { + panic!("expected two images"); + }; + assert_eq!(decoded_image(image_url).0, original_bytes); + assert_eq!(preserved_http_url, &http_url); +} + +#[test] +fn detail_policies_apply_the_expected_budgets() { + for (detail, input_dimensions, expected_dimensions) in [ + (Some(ImageDetail::High), (2048, 2048), (1600, 1600)), + (Some(ImageDetail::Original), (6401, 100), (6000, 94)), + (Some(ImageDetail::Original), (3201, 3201), (3200, 3200)), + (Some(ImageDetail::Auto), (2048, 2048), (1600, 1600)), + (None, (2048, 2048), (1600, 1600)), + ] { + let (image_url, _) = png_data_url(input_dimensions.0, input_dimensions.1); + let mut items = vec![ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputImage { image_url, detail }], + phase: None, + }]; + + prepare_response_items(&mut items); + + let ResponseItem::Message { content, .. } = &items[0] else { + panic!("expected message"); + }; + let [ContentItem::InputImage { image_url, .. }] = content.as_slice() else { + panic!("expected image"); + }; + assert_eq!(decoded_image(image_url).1.dimensions(), expected_dimensions); + } +} + +#[test] +fn preparation_replaces_only_failed_tool_images_and_preserves_metadata() { + let (valid_image_url, _) = png_data_url(/*width*/ 64, /*height*/ 32); + let expected_valid_image_url = valid_image_url.clone(); + let mut items = vec![ResponseItem::CustomToolCallOutput { + call_id: "call-1".to_string(), + name: None, + output: FunctionCallOutputPayload { + body: FunctionCallOutputBody::ContentItems(vec![ + FunctionCallOutputContentItem::InputText { + text: "before".to_string(), + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,%%%".to_string(), + detail: Some(ImageDetail::High), + }, + FunctionCallOutputContentItem::InputImage { + image_url: data_url_from_bytes("image/png", b"not an image"), + detail: Some(ImageDetail::High), + }, + FunctionCallOutputContentItem::InputImage { + image_url: valid_image_url.clone(), + detail: Some(ImageDetail::Low), + }, + FunctionCallOutputContentItem::InputImage { + image_url: valid_image_url, + detail: Some(ImageDetail::High), + }, + ]), + success: Some(true), + }, + }]; + + prepare_response_items(&mut items); + + assert_eq!( + items, + vec![ResponseItem::CustomToolCallOutput { + call_id: "call-1".to_string(), + name: None, + output: FunctionCallOutputPayload { + body: FunctionCallOutputBody::ContentItems(vec![ + FunctionCallOutputContentItem::InputText { + text: "before".to_string(), + }, + FunctionCallOutputContentItem::InputText { + text: IMAGE_PROCESSING_ERROR_PLACEHOLDER.to_string(), + }, + FunctionCallOutputContentItem::InputText { + text: IMAGE_PROCESSING_ERROR_PLACEHOLDER.to_string(), + }, + FunctionCallOutputContentItem::InputText { + text: UNSUPPORTED_LOW_DETAIL_PLACEHOLDER.to_string(), + }, + FunctionCallOutputContentItem::InputImage { + image_url: expected_valid_image_url, + detail: Some(ImageDetail::High), + }, + ]), + success: Some(true), + }, + }] + ); +} + +#[test] +fn preparation_errors_use_bounded_actionable_placeholders() { + let cases = [ + ( + ImagePreparationError::UnsupportedLowDetail, + UNSUPPORTED_LOW_DETAIL_PLACEHOLDER, + ), + ( + ImagePreparationError::Processing(ImageProcessingError::ImageTooLarge { + representation: "decoded input", + size: 2, + max: 1, + }), + IMAGE_TOO_LARGE_PLACEHOLDER, + ), + ( + ImagePreparationError::Processing(ImageProcessingError::InvalidDataUrl { + reason: "details remain in logs".to_string(), + }), + IMAGE_PROCESSING_ERROR_PLACEHOLDER, + ), + ]; + + for (error, expected) in cases { + assert_eq!(error.placeholder(), expected); + } +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 5e452d8b81..c490360c8c 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -19,6 +19,7 @@ mod codex_thread; mod compact_remote; mod compact_remote_v2; mod config_lock; +pub use codex_thread::BackgroundTerminalInfo; pub use codex_thread::CodexThread; pub use codex_thread::CodexThreadSettingsOverrides; pub use codex_thread::ThreadConfigSnapshot; @@ -41,6 +42,7 @@ mod exec_policy; mod git_info_tests; mod guardian; mod hook_runtime; +mod image_preparation; mod installation_id; pub(crate) mod landlock; pub use landlock::spawn_command_under_linux_sandbox; diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 1b693d70fd..abd28ec619 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -11,6 +11,7 @@ use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::EffectiveMcpServer; use codex_mcp::McpConfig; use codex_mcp::ToolPluginProvenance; +use codex_mcp::codex_apps_mcp_server_config; use codex_mcp::configured_mcp_servers; use codex_mcp::effective_mcp_servers; use codex_mcp::tool_plugin_provenance as collect_tool_plugin_provenance; @@ -40,17 +41,36 @@ impl McpManager { } } - /// Returns the MCP config after applying runtime-only extension overlays. + /// Returns the MCP config after applying compatibility built-ins and + /// runtime-only extension overlays. pub async fn runtime_config(&self, config: &Config) -> McpConfig { let mut mcp_config = config.to_mcp_config(self.plugins_manager.as_ref()).await; - let contributions = self.contributions(config).await; - if contributions + let disabled_server_names = mcp_config + .configured_mcp_servers .iter() - .any(|contribution| contribution.name() == CODEX_APPS_MCP_SERVER_NAME) - { - mcp_config.legacy_apps_mcp_loader_enabled = false; + .filter(|(_, server)| !server.enabled) + .map(|(name, _)| name.clone()) + .collect::>(); + if mcp_config.apps_enabled { + mcp_config.configured_mcp_servers.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + codex_apps_mcp_server_config( + &mcp_config.chatgpt_base_url, + mcp_config.apps_mcp_product_sku.as_deref(), + ), + ); + } else { + mcp_config + .configured_mcp_servers + .remove(CODEX_APPS_MCP_SERVER_NAME); } + let contributions = self.contributions(config).await; Self::apply_to_configured_servers(&contributions, &mut mcp_config.configured_mcp_servers); + for name in disabled_server_names { + if let Some(server) = mcp_config.configured_mcp_servers.get_mut(&name) { + server.enabled = false; + } + } mcp_config } diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index b9a8d4ff08..e5273a6669 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -319,8 +319,7 @@ async fn handle_approved_mcp_tool_call( let server_origin = sess .services .mcp_connection_manager - .read() - .await + .load_full() .server_origin(&server) .map(str::to_string); @@ -606,8 +605,7 @@ async fn maybe_request_codex_apps_auth_elicitation( if !sess .services .mcp_connection_manager - .read() - .await + .load_full() .is_host_owned_codex_apps_server(server) { return result; @@ -669,13 +667,9 @@ async fn maybe_request_codex_apps_auth_elicitation( auth_elicitation_completed_result(&plan.auth_failure, result.meta) } -#[expect( - clippy::await_holding_invalid_type, - reason = "Codex Apps cache refresh reads through the session-owned manager guard" -)] async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: &TurnContext) { let mcp_tools_result = { - let manager = sess.services.mcp_connection_manager.read().await; + let manager = sess.services.mcp_connection_manager.load_full(); manager.hard_refresh_codex_apps_tools_cache().await }; @@ -694,10 +688,6 @@ async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: & } } -#[expect( - clippy::await_holding_invalid_type, - reason = "MCP sandbox metadata reads through the session-owned manager guard" -)] async fn augment_mcp_tool_request_meta_with_sandbox_state( sess: &Session, turn_context: &TurnContext, @@ -707,8 +697,7 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state( let supports_sandbox_state_meta = sess .services .mcp_connection_manager - .read() - .await + .load_full() .server_supports_sandbox_state_meta_capability(server) .await .unwrap_or(false); @@ -757,8 +746,7 @@ async fn maybe_mark_thread_memory_mode_polluted( let pollutes_memory = sess .services .mcp_connection_manager - .read() - .await + .load_full() .server_pollutes_memory(server); if !pollutes_memory { return; @@ -1414,17 +1402,13 @@ async fn mcp_tool_approval_decision_from_guardian( } } -#[expect( - clippy::await_holding_invalid_type, - reason = "MCP approval metadata reads through the session-owned manager guard" -)] pub(crate) async fn lookup_mcp_tool_metadata( sess: &Session, turn_context: &TurnContext, server: &str, tool_name: &str, ) -> Option { - let manager = sess.services.mcp_connection_manager.read().await; + let manager = sess.services.mcp_connection_manager.load_full(); let plugin_id = manager .plugin_id_for_mcp_server_name(server) .map(str::to_string); @@ -1509,10 +1493,6 @@ fn get_mcp_app_resource_uri( }) } -#[expect( - clippy::await_holding_invalid_type, - reason = "MCP app metadata reads through the session-owned manager guard" -)] async fn lookup_mcp_app_usage_metadata( sess: &Session, server: &str, @@ -1521,8 +1501,7 @@ async fn lookup_mcp_app_usage_metadata( let tools = sess .services .mcp_connection_manager - .read() - .await + .load_full() .list_all_tools() .await; diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index bd71c3a054..723948b7be 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -46,6 +46,7 @@ use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use tempfile::tempdir; +use tokio_util::sync::CancellationToken; use tracing::Instrument; use tracing::Level; use tracing_subscriber::fmt::format::FmtSpan; @@ -1256,13 +1257,14 @@ fn codex_apps_auth_failure_metadata() -> McpToolApprovalMetadata { async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: &TurnContext) { let auth = session.services.auth_manager.auth().await; - let (manager, _cancel_token) = codex_mcp::McpConnectionManager::new( + let manager = codex_mcp::McpConnectionManager::new( &HashMap::new(), turn_context.config.mcp_oauth_credentials_store_mode, HashMap::new(), &turn_context.approval_policy, turn_context.sub_id.clone(), session.get_tx_event(), + CancellationToken::new(), turn_context.permission_profile(), codex_mcp::McpRuntimeContext::new(Arc::clone(&session.services.environment_manager), { #[allow(deprecated)] @@ -1278,7 +1280,10 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: /*elicitation_reviewer*/ None, ) .await; - *session.services.mcp_connection_manager.write().await = manager; + session + .services + .mcp_connection_manager + .store(Arc::new(manager)); } #[tokio::test] diff --git a/codex-rs/core/src/mcp_tool_exposure.rs b/codex-rs/core/src/mcp_tool_exposure.rs index 5dc66af734..d430851570 100644 --- a/codex-rs/core/src/mcp_tool_exposure.rs +++ b/codex-rs/core/src/mcp_tool_exposure.rs @@ -4,6 +4,7 @@ use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::ToolInfo as McpToolInfo; use codex_mcp::tool_is_model_visible; +use tracing::instrument; use crate::config::Config; use crate::connectors; @@ -15,6 +16,7 @@ pub(crate) struct McpToolExposure { pub(crate) deferred_tools: Option>, } +#[instrument(level = "trace", skip_all)] pub(crate) fn build_mcp_tool_exposure( all_mcp_tools: &[McpToolInfo], connectors: Option<&[connectors::AppInfo]>, diff --git a/codex-rs/core/src/plugins/discoverable.rs b/codex-rs/core/src/plugins/discoverable.rs index 8ddec7056e..fc77ae21e6 100644 --- a/codex-rs/core/src/plugins/discoverable.rs +++ b/codex-rs/core/src/plugins/discoverable.rs @@ -41,6 +41,7 @@ pub(crate) async fn list_tool_suggest_discoverable_plugins( .into_iter() .map(|plugin| DiscoverablePluginInfo { id: plugin.id, + remote_plugin_id: plugin.remote_plugin_id, name: plugin.name, description: plugin.description, has_skills: plugin.has_skills, diff --git a/codex-rs/core/src/plugins/discoverable_tests.rs b/codex-rs/core/src/plugins/discoverable_tests.rs index c1dfe8e639..a6da04e77c 100644 --- a/codex-rs/core/src/plugins/discoverable_tests.rs +++ b/codex-rs/core/src/plugins/discoverable_tests.rs @@ -393,6 +393,7 @@ remote_plugin = true remote_plugins, vec![DiscoverablePluginInfo { id: "github@openai-curated-remote".to_string(), + remote_plugin_id: Some("plugins~Plugin_remote_github".to_string()), name: "Remote GitHub".to_string(), description: Some("Remote GitHub short".to_string()), has_skills: true, @@ -680,6 +681,7 @@ async fn list_tool_suggest_discoverable_plugins_normalizes_description() { discoverable_plugins, vec![DiscoverablePluginInfo { id: "slack@openai-curated".to_string(), + remote_plugin_id: None, name: "slack".to_string(), description: Some("Plugin with extra spacing".to_string()), has_skills: true, @@ -816,6 +818,7 @@ discoverables = [{ type = "plugin", id = "sample@openai-curated" }] discoverable_plugins, vec![DiscoverablePluginInfo { id: "sample@openai-curated".to_string(), + remote_plugin_id: None, name: "sample".to_string(), description: Some( "Plugin that includes skills, MCP servers, and app connectors".to_string(), diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index f6bf7d0022..617521608f 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -5,7 +5,6 @@ use codex_exec_server::ExecServerRuntimePaths; use codex_login::AuthManager; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; -use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::SessionSource; use codex_protocol::user_input::UserInput; @@ -78,8 +77,7 @@ pub(crate) async fn build_prompt_input_from_session( .await; if !input.is_empty() { - let input_item = ResponseInputItem::from(input); - let response_item = ResponseItem::from(input_item); + let response_item = sess.response_item_from_user_input(turn_context.as_ref(), input); sess.record_conversation_items(turn_context.as_ref(), std::slice::from_ref(&response_item)) .await; } diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index 9f35fb2fa4..33568f6f38 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -1,5 +1,6 @@ use crate::client::ModelClient; use crate::realtime_context::build_realtime_startup_context; +use crate::realtime_context::truncate_realtime_text_to_token_budget; use crate::realtime_prompt::prepare_realtime_backend_prompt; use crate::session::session::Session; use anyhow::Context; @@ -64,6 +65,8 @@ const USER_TEXT_IN_QUEUE_CAPACITY: usize = 64; const HANDOFF_OUT_QUEUE_CAPACITY: usize = 64; const OUTPUT_EVENTS_QUEUE_CAPACITY: usize = 256; const REALTIME_STARTUP_CONTEXT_TOKEN_BUDGET: usize = 5_300; +const REALTIME_ASSISTANT_OUTPUT_TOKEN_BUDGET: usize = 1_000; +const STANDALONE_HANDOFF_ID: &str = "codex"; const DEFAULT_REALTIME_MODEL: &str = "gpt-realtime-1.5"; pub(crate) const REALTIME_USER_TEXT_PREFIX: &str = "[USER] "; pub(crate) const REALTIME_BACKEND_TEXT_PREFIX: &str = "[BACKEND] "; @@ -106,6 +109,9 @@ struct RealtimeHandoffState { #[derive(Debug, PartialEq, Eq)] enum HandoffOutput { + StandaloneAssistantOutput { + output_text: String, + }, ProgressUpdate { handoff_id: String, output_text: String, @@ -460,22 +466,38 @@ impl RealtimeConversationManager { state.handoff.clone() }; - let Some(handoff_id) = handoff.active_handoff.lock().await.clone() else { - return Ok(()); + let active_handoff = handoff.active_handoff.lock().await.clone(); + let output = match active_handoff { + Some(handoff_id) => { + let output_text = prefix_realtime_text( + output_text, + REALTIME_BACKEND_TEXT_PREFIX, + handoff.session_kind, + ); + *handoff.last_output_text.lock().await = Some(output_text.clone()); + HandoffOutput::ProgressUpdate { + handoff_id, + output_text, + } + } + None if output_text.trim().is_empty() => return Ok(()), + None => { + let output_text = prefix_realtime_text( + output_text, + REALTIME_BACKEND_TEXT_PREFIX, + handoff.session_kind, + ); + HandoffOutput::StandaloneAssistantOutput { + output_text: truncate_realtime_text_to_token_budget( + &output_text, + REALTIME_ASSISTANT_OUTPUT_TOKEN_BUDGET, + ), + } + } }; - - let output_text = prefix_realtime_text( - output_text, - REALTIME_BACKEND_TEXT_PREFIX, - handoff.session_kind, - ); - *handoff.last_output_text.lock().await = Some(output_text.clone()); handoff .output_tx - .send(HandoffOutput::ProgressUpdate { - handoff_id, - output_text, - }) + .send(output) .await .map_err(|_| CodexErr::InvalidRequest("conversation is not running".to_string()))?; Ok(()) @@ -511,14 +533,6 @@ impl RealtimeConversationManager { .map_err(|_| CodexErr::InvalidRequest("conversation is not running".to_string())) } - pub(crate) async fn active_handoff_id(&self) -> Option { - let handoff = { - let guard = self.state.lock().await; - guard.as_ref().map(|state| state.handoff.clone()) - }?; - handoff.active_handoff.lock().await.clone() - } - pub(crate) async fn clear_active_handoff(&self) { let handoff = { let guard = self.state.lock().await; @@ -625,12 +639,14 @@ async fn prepare_realtime_start( if let Some(realtime_ws_base_url) = &config.experimental_realtime_ws_base_url { api_provider.base_url = realtime_ws_base_url.clone(); } - let version = config.realtime.version; + let version = params.version.unwrap_or(config.realtime.version); let session_config = build_realtime_session_config( sess, + params.model, params.prompt, params.realtime_session_id, params.output_modality, + version, params.voice, ) .await?; @@ -664,9 +680,11 @@ async fn prepare_realtime_start( pub(crate) async fn build_realtime_session_config( sess: &Arc, + model: Option, prompt: Option>, realtime_session_id: Option, output_modality: RealtimeOutputModality, + version: RealtimeWsVersion, voice: Option, ) -> CodexResult { let config = sess.get_config().await; @@ -689,18 +707,15 @@ pub(crate) async fn build_realtime_session_config( (false, false) => format!("{prompt}\n\n{startup_context}"), }; let model = Some( - config - .experimental_realtime_ws_model - .clone() + model + .or_else(|| config.experimental_realtime_ws_model.clone()) .unwrap_or_else(|| DEFAULT_REALTIME_MODEL.to_string()), ); - let event_parser = match config.realtime.version { + let event_parser = match version { RealtimeWsVersion::V1 => RealtimeEventParser::V1, RealtimeWsVersion::V2 => RealtimeEventParser::RealtimeV2, }; - if config.realtime.version == RealtimeWsVersion::V1 - && matches!(output_modality, RealtimeOutputModality::Text) - { + if version == RealtimeWsVersion::V1 && matches!(output_modality, RealtimeOutputModality::Text) { return Err(CodexErr::InvalidRequest( "text realtime output modality requires realtime v2".to_string(), )); @@ -711,8 +726,8 @@ pub(crate) async fn build_realtime_session_config( }; let voice = voice .or(config.realtime.voice) - .unwrap_or_else(|| default_realtime_voice(config.realtime.version)); - validate_realtime_voice(config.realtime.version, voice)?; + .unwrap_or_else(|| default_realtime_voice(version)); + validate_realtime_voice(version, voice)?; Ok(RealtimeSessionConfig { instructions: prompt, model, @@ -739,10 +754,6 @@ fn prefix_realtime_text(text: String, prefix: &str, session_kind: RealtimeSessio format!("{prefix}{text}") } -pub(crate) fn prefix_realtime_v2_text(text: String, prefix: &str) -> String { - prefix_realtime_text(text, prefix, RealtimeSessionKind::V2) -} - fn validate_realtime_voice(version: RealtimeWsVersion, voice: RealtimeVoice) -> CodexResult<()> { let voices = RealtimeVoicesList::builtin(); let allowed = match version { @@ -1191,6 +1202,15 @@ async fn handle_handoff_output( let result = match event_parser { RealtimeEventParser::V1 => match handoff_output { + HandoffOutput::StandaloneAssistantOutput { output_text } => { + // TODO(guinness): Use the new client event for standalone handoffs once the API changes are complete. + writer + .send_conversation_handoff_append( + STANDALONE_HANDOFF_ID.to_string(), + output_text, + ) + .await + } HandoffOutput::ProgressUpdate { handoff_id, output_text, @@ -1205,6 +1225,15 @@ async fn handle_handoff_output( } }, RealtimeEventParser::RealtimeV2 => match handoff_output { + HandoffOutput::StandaloneAssistantOutput { output_text } => { + if let Err(err) = writer.send_conversation_item_create(output_text).await { + Err(err) + } else { + return response_create_queue + .request_create(writer, events_tx, "standalone assistant output") + .await; + } + } HandoffOutput::ProgressUpdate { handoff_id, output_text, diff --git a/codex-rs/core/src/session/config_lock.rs b/codex-rs/core/src/session/config_lock.rs index c081a2aee6..ec2056581c 100644 --- a/codex-rs/core/src/session/config_lock.rs +++ b/codex-rs/core/src/session/config_lock.rs @@ -2,7 +2,6 @@ use anyhow::Context; use codex_config::config_toml::ConfigLockfileToml; use codex_config::config_toml::ConfigToml; use codex_config::types::MemoriesToml; -use codex_features::AppsMcpPathOverrideConfigToml; use codex_features::Feature; use codex_features::FeatureToml; use codex_features::FeaturesToml; @@ -149,10 +148,6 @@ fn save_config_resolved_fields( resolved_config_to_toml(&config.multi_agent_v2, "features.multi_agent_v2")?; multi_agent_v2.enabled = Some(config.features.enabled(Feature::MultiAgentV2)); features.multi_agent_v2 = Some(FeatureToml::Config(multi_agent_v2)); - features.apps_mcp_path_override = Some(FeatureToml::Config(AppsMcpPathOverrideConfigToml { - enabled: Some(config.features.enabled(Feature::AppsMcpPathOverride)), - path: config.apps_mcp_path_override.clone(), - })); lock_config.memories = Some(resolved_config_to_toml::( &config.memories, "memories", @@ -325,6 +320,32 @@ mod tests { assert!(message.contains("model = "), "{message}"); } + #[tokio::test] + async fn lock_validation_ignores_removed_apps_mcp_path_override() { + let sc = crate::session::tests::make_session_configuration_for_tests().await; + let actual = sc.to_config_lockfile_toml().expect("lock should serialize"); + let mut expected_value = toml::Value::try_from(&actual).expect("lock should become TOML"); + expected_value["config"]["features"] + .as_table_mut() + .expect("features should be a table") + .insert( + "apps_mcp_path_override".to_string(), + toml::Value::Table(toml::Table::from_iter([ + ("enabled".to_string(), toml::Value::Boolean(true)), + ( + "path".to_string(), + toml::Value::String("/custom/mcp".to_string()), + ), + ])), + ); + let expected: ConfigLockfileToml = expected_value + .try_into() + .expect("lock with removed input should deserialize"); + + validate_config_lock_replay(&expected, &actual, ConfigLockReplayOptions::default()) + .expect("removed compatibility input should not cause lock drift"); + } + #[tokio::test] async fn lock_validation_rejects_codex_version_mismatch_by_default() { let sc = crate::session::tests::make_session_configuration_for_tests().await; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index c0f62ec664..bfb87a0ff3 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -15,10 +15,6 @@ use crate::session::session::Session; use crate::session::session::SessionSettingsUpdate; use crate::config::Config; -use crate::realtime_context::REALTIME_TURN_TOKEN_BUDGET; -use crate::realtime_context::truncate_realtime_text_to_token_budget; -use crate::realtime_conversation::REALTIME_USER_TEXT_PREFIX; -use crate::realtime_conversation::prefix_realtime_v2_text; use crate::review_prompts::resolve_review_request; use crate::session::spawn_review_thread; use crate::tasks::CompactTask; @@ -54,9 +50,7 @@ use codex_protocol::request_user_input::RequestUserInputResponse; use crate::context_manager::is_user_turn_boundary; use codex_protocol::dynamic_tools::DynamicToolResponse; -use codex_protocol::items::UserMessageItem; use codex_protocol::mcp::RequestId as ProtocolRequestId; -use codex_protocol::user_input::UserInput; use codex_rmcp_client::ElicitationAction; use codex_rmcp_client::ElicitationResponse; use serde_json::Value; @@ -91,14 +85,7 @@ pub async fn user_input_or_turn( op: Op, client_user_message_id: Option, ) { - user_input_or_turn_inner( - sess, - sub_id, - op, - /*mirror_user_text_to_realtime*/ Some(()), - client_user_message_id, - ) - .await; + user_input_or_turn_inner(sess, sub_id, op, client_user_message_id).await; } pub async fn update_thread_settings( @@ -196,7 +183,6 @@ pub(super) async fn user_input_or_turn_inner( sess: &Arc, sub_id: String, op: Op, - mirror_user_text_to_realtime: Option<()>, client_user_message_id: Option, ) { let Op::UserInput { @@ -230,7 +216,7 @@ pub(super) async fn user_input_or_turn_inner( } sess.maybe_emit_unknown_model_warning_for_turn(current_context.as_ref()) .await; - let accepted_items = match sess + match sess .steer_input( items.clone(), additional_context.clone(), @@ -242,7 +228,6 @@ pub(super) async fn user_input_or_turn_inner( { Ok(_) => { current_context.session_telemetry.user_prompt(&items); - Some(items) } Err(SteerInputError::NoActiveTurn(items)) => { if let Some(responsesapi_client_metadata) = responsesapi_client_metadata { @@ -256,7 +241,6 @@ pub(super) async fn user_input_or_turn_inner( Some(sess.mcp_elicitation_reviewer()), ) .await; - let accepted_items = items.clone(); let additional_context_input = { let mut state = sess.state.lock().await; state.additional_context.merge(additional_context) @@ -278,7 +262,6 @@ pub(super) async fn user_input_or_turn_inner( crate::tasks::RegularTask::new(), ) .await; - Some(accepted_items) } Err(err) => { sess.send_event_raw(Event { @@ -286,33 +269,7 @@ pub(super) async fn user_input_or_turn_inner( msg: EventMsg::Error(err.to_error_event()), }) .await; - None } - }; - if let (Some(items), Some(())) = (accepted_items, mirror_user_text_to_realtime) { - self::mirror_user_text_to_realtime(sess, &items).await; - } -} - -async fn mirror_user_text_to_realtime(sess: &Arc, items: &[UserInput]) { - let text = UserMessageItem::new(items).message(); - if text.is_empty() { - return; - } - let text = if sess.conversation.is_running_v2().await { - prefix_realtime_v2_text(text, REALTIME_USER_TEXT_PREFIX) - } else { - text - }; - let text = truncate_realtime_text_to_token_budget(&text, REALTIME_TURN_TOKEN_BUDGET); - if text.is_empty() { - return; - } - if sess.conversation.running_state().await.is_none() { - return; - } - if let Err(err) = sess.conversation.text_in(text).await { - debug!("failed to mirror user text to realtime conversation: {err}"); } } @@ -623,6 +580,9 @@ pub async fn set_thread_memory_mode(sess: &Arc, sub_id: String, mode: T } async fn shutdown_session_runtime(sess: &Arc) { + if let Some(startup_prewarm) = sess.take_session_startup_prewarm().await { + startup_prewarm.abort().await; + } sess.abort_all_tasks(TurnAbortReason::Interrupted).await; let _ = sess.conversation.shutdown().await; sess.services @@ -632,11 +592,11 @@ async fn shutdown_session_runtime(sess: &Arc) { if let Err(err) = sess.services.code_mode_service.shutdown().await { warn!("failed to shutdown code mode session: {err}"); } - let mcp_shutdown = { - let mut manager = sess.services.mcp_connection_manager.write().await; - manager.begin_shutdown() - }; - mcp_shutdown.await; + sess.services + .mcp_connection_manager + .load_full() + .shutdown() + .await; sess.guardian_review_session.shutdown().await; } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index f2643d0e5c..39333801c0 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -91,8 +91,7 @@ impl Session { if self .services .mcp_connection_manager - .read() - .await + .load_full() .elicitations_auto_deny() { return McpServerElicitationOutcome { @@ -226,16 +225,11 @@ impl Session { self.services .mcp_connection_manager - .read() - .await + .load_full() .resolve_elicitation(server_name, id, response) .await } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP resource calls are serialized through the session-owned manager guard" - )] pub async fn list_resources( &self, server: &str, @@ -243,16 +237,11 @@ impl Session { ) -> anyhow::Result { self.services .mcp_connection_manager - .read() - .await + .load_full() .list_resources(server, params) .await } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP resource calls are serialized through the session-owned manager guard" - )] pub async fn list_resource_templates( &self, server: &str, @@ -260,16 +249,11 @@ impl Session { ) -> anyhow::Result { self.services .mcp_connection_manager - .read() - .await + .load_full() .list_resource_templates(server, params) .await } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP resource calls are serialized through the session-owned manager guard" - )] pub async fn read_resource( &self, server: &str, @@ -277,16 +261,11 @@ impl Session { ) -> anyhow::Result { self.services .mcp_connection_manager - .read() - .await + .load_full() .read_resource(server, params) .await } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP tool calls are serialized through the session-owned manager guard" - )] pub async fn call_tool( &self, server: &str, @@ -296,8 +275,7 @@ impl Session { ) -> anyhow::Result { self.services .mcp_connection_manager - .read() - .await + .load_full() .call_tool(server, tool, arguments, meta) .await } @@ -338,18 +316,21 @@ impl Session { turn_context.cwd.to_path_buf(), ), }; - { + let mcp_startup_cancellation_token = { let mut guard = self.services.mcp_startup_cancellation_token.lock().await; guard.cancel(); - *guard = CancellationToken::new(); - } - let (refreshed_manager, cancel_token) = McpConnectionManager::new( + let cancellation_token = CancellationToken::new(); + *guard = cancellation_token.clone(); + cancellation_token + }; + let refreshed_manager = McpConnectionManager::new( &mcp_servers, store_mode, auth_statuses, &turn_context.approval_policy, turn_context.sub_id.clone(), self.get_tx_event(), + mcp_startup_cancellation_token, turn_context.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), @@ -363,22 +344,12 @@ impl Session { ) .await; { - let current_manager = self.services.mcp_connection_manager.read().await; + let current_manager = self.services.mcp_connection_manager.load_full(); refreshed_manager.set_elicitations_auto_deny(current_manager.elicitations_auto_deny()); } - { - let mut guard = self.services.mcp_startup_cancellation_token.lock().await; - if guard.is_cancelled() { - cancel_token.cancel(); - } - *guard = cancel_token; - } - - let mut old_manager = { - let mut manager = self.services.mcp_connection_manager.write().await; - std::mem::replace(&mut *manager, refreshed_manager) - }; - old_manager.shutdown().await; + self.services + .mcp_connection_manager + .store(Arc::new(refreshed_manager)); } pub(crate) async fn refresh_mcp_servers_if_requested( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 23dbe1f45e..4da2b58ac5 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::collections::BTreeMap; use std::collections::HashMap; use std::collections::HashSet; @@ -31,6 +32,7 @@ use crate::context::PersonalitySpecInstructions; use crate::default_skill_metadata_budget; use crate::environment_selection::ResolvedTurnEnvironments; use crate::exec_policy::ExecPolicyManager; +use crate::image_preparation::prepare_response_items; use crate::parse_turn_item; use crate::realtime_conversation::RealtimeConversationManager; use crate::session_prefix::format_subagent_notification_message; @@ -63,6 +65,7 @@ use codex_login::CodexAuth; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_login::default_client::originator; use codex_mcp::McpConnectionManager; +use codex_mcp::McpResourceClient; use codex_mcp::McpRuntimeContext; use codex_mcp::codex_apps_tools_cache_key; use codex_models_manager::manager::RefreshStrategy; @@ -205,6 +208,7 @@ mod review; mod rollout_reconstruction; #[allow(clippy::module_inception)] pub(crate) mod session; +mod token_budget; pub(crate) mod turn; pub(crate) mod turn_context; use self::config_lock::export_config_lock_if_configured; @@ -273,10 +277,12 @@ impl SteerInputError { /// Conceptually this is the same role that `previous_model` used to fill, but /// it can carry other prior-turn settings that matter when constructing /// sensible state-change diffs or full-context reinjection, such as model -/// switches or detecting a prior `realtime_active -> false` transition. +/// switches, compaction compatibility, or detecting a prior +/// `realtime_active -> false` transition. #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct PreviousTurnSettings { pub(crate) model: String, + pub(crate) comp_hash: Option, pub(crate) realtime_active: Option, } @@ -326,6 +332,7 @@ use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::Personality; use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; use codex_protocol::config_types::WindowsSandboxLevel; +use codex_protocol::models::LocalImagePreparation; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; @@ -793,7 +800,7 @@ impl Codex { ..Default::default() }) .await?; - let mcp_connection_manager = self.session.services.mcp_connection_manager.read().await; + let mcp_connection_manager = self.session.services.mcp_connection_manager.load_full(); mcp_connection_manager.set_elicitations_auto_deny(mcp_elicitations_auto_deny); Ok(()) } @@ -1132,7 +1139,6 @@ impl Session { additional_context: Default::default(), thread_settings: Default::default(), }, - /*mirror_user_text_to_realtime*/ None, /*client_user_message_id*/ None, ) .await; @@ -1300,15 +1306,27 @@ impl Session { turn_context: &TurnContext, rollout_items: &[RolloutItem], ) -> Option { - let reconstructed_rollout = self + let rollout_reconstruction::RolloutReconstruction { + mut history, + previous_turn_settings, + reference_context_item, + window_id, + } = self .reconstruct_history_from_rollout(turn_context, rollout_items) .await; - let previous_turn_settings = reconstructed_rollout.previous_turn_settings.clone(); - self.replace_history( - reconstructed_rollout.history, - reconstructed_rollout.reference_context_item, - ) - .await; + if turn_context.features.enabled(Feature::ResizeAllImages) { + // Keep the recorded rollout unchanged. Prepare its reconstructed history before + // installing it, so legacy images are processed once for this resume or fork and + // will be processed again if the rollout is reconstructed in a future session. + // This meets image resizing requirements without modifying persisted rollouts. + prepare_response_items(&mut history); + } + { + let mut state = self.state.lock().await; + state.replace_history(history, reference_context_item); + state.set_auto_compact_window_id(window_id); + state.set_previous_turn_settings(previous_turn_settings.clone()); + } let prefix_tokens = if matches!( turn_context.config.model_auto_compact_token_limit_scope, AutoCompactTokenLimitScope::BodyAfterPrefix @@ -1323,8 +1341,6 @@ impl Session { self.set_auto_compact_window_estimated_prefill_for_scope(turn_context, prefix_tokens) .await; } - self.set_previous_turn_settings(previous_turn_settings.clone()) - .await; previous_turn_settings } @@ -1794,9 +1810,7 @@ impl Session { let Some(text) = realtime_text_for_event(msg) else { return; }; - if self.conversation.running_state().await.is_none() - || self.conversation.active_handoff_id().await.is_none() - { + if self.conversation.running_state().await.is_none() { return; } if let Err(err) = self.conversation.handoff_out(text).await { @@ -2583,11 +2597,43 @@ impl Session { /// Records conversation items: append to history, persist to rollout, and /// notify clients observing raw response items. + pub(crate) fn prepare_conversation_items_for_history<'a>( + &self, + turn_context: &TurnContext, + items: &'a [ResponseItem], + ) -> Cow<'a, [ResponseItem]> { + if !turn_context.features.enabled(Feature::ResizeAllImages) { + return Cow::Borrowed(items); + } + + let mut prepared_items = items.to_vec(); + prepare_response_items(&mut prepared_items); + Cow::Owned(prepared_items) + } + + pub(crate) fn response_item_from_user_input( + &self, + turn_context: &TurnContext, + input: Vec, + ) -> ResponseItem { + let local_image_preparation = if turn_context.features.enabled(Feature::ResizeAllImages) { + LocalImagePreparation::Defer + } else { + LocalImagePreparation::Process + }; + ResponseItem::from(ResponseInputItem::from_user_input( + input, + local_image_preparation, + )) + } + pub(crate) async fn record_conversation_items( &self, turn_context: &TurnContext, items: &[ResponseItem], ) { + let items = self.prepare_conversation_items_for_history(turn_context, items); + let items = items.as_ref(); { let mut state = self.state.lock().await; state.record_items(items.iter(), turn_context.truncation_policy); @@ -2656,6 +2702,7 @@ impl Session { .await; } + #[cfg(test)] pub(crate) async fn replace_history( &self, items: Vec, @@ -2674,7 +2721,6 @@ impl Session { { let mut state = self.state.lock().await; state.replace_history(items, reference_context_item.clone()); - state.start_next_auto_compact_window(); } self.persist_rollout_items(&[RolloutItem::Compacted(compacted_item)]) @@ -2687,7 +2733,6 @@ impl Session { let mut state = self.state.lock().await; state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact); } - self.services.model_client.advance_window_generation(); } async fn persist_rollout_response_items(&self, items: &[ResponseItem]) { @@ -2749,10 +2794,6 @@ impl Session { } } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP app context rendering reads through the session-owned manager guard" - )] pub(crate) async fn build_initial_context( &self, turn_context: &TurnContext, @@ -2766,6 +2807,7 @@ impl Session { collaboration_mode, base_instructions, session_source, + auto_compact_window_id, ) = { let state = self.state.lock().await; ( @@ -2774,6 +2816,7 @@ impl Session { state.session_configuration.collaboration_mode.clone(), state.session_configuration.base_instructions.clone(), state.session_configuration.session_source.clone(), + state.auto_compact_window_id(), ) }; if let Some(model_switch_message) = @@ -2845,7 +2888,7 @@ impl Session { } } if turn_context.config.include_apps_instructions && turn_context.apps_enabled() { - let mcp_connection_manager = self.services.mcp_connection_manager.read().await; + let mcp_connection_manager = self.services.mcp_connection_manager.load_full(); let accessible_and_enabled_connectors = connectors::list_accessible_and_enabled_connectors_from_manager( &mcp_connection_manager, @@ -2923,6 +2966,18 @@ impl Session { .render(), ); } + // This is full-context metadata. Steady-state context diffs should not re-emit it. + if turn_context.features.enabled(Feature::TokenBudget) + && let Some(model_context_window) = turn_context.model_context_window() + { + developer_sections.push( + crate::context::TokenBudgetContext::new( + auto_compact_window_id, + model_context_window, + ) + .render(), + ); + } if turn_context.config.include_environment_context { let shell = self.user_shell(); let subagents = self @@ -2994,6 +3049,56 @@ impl Session { state.clone_history() } + pub(crate) async fn current_window_id(&self) -> String { + let state = self.state.lock().await; + let thread_id = self.thread_id; + let window_id = state.auto_compact_window_id(); + format!("{thread_id}:{window_id}") + } + + pub(crate) async fn advance_auto_compact_window_id(&self) -> u64 { + let mut state = self.state.lock().await; + state.advance_auto_compact_window_id() + } + + pub(crate) async fn request_new_context_window(&self) { + let mut state = self.state.lock().await; + state.request_new_context_window(); + } + + pub(crate) async fn maybe_start_new_context_window( + &self, + turn_context: &TurnContext, + ) -> Option { + let window_id = { + let mut state = self.state.lock().await; + state.start_new_context_window_if_requested() + }; + let window_id = window_id?; + let context_items = self.build_initial_context(turn_context).await; + let turn_context_item = turn_context.to_turn_context_item(); + let replacement_history = context_items; + { + let mut state = self.state.lock().await; + state.replace_history(replacement_history.clone(), Some(turn_context_item.clone())); + }; + self.persist_rollout_items(&[ + RolloutItem::Compacted(CompactedItem { + message: String::new(), + replacement_history: Some(replacement_history), + window_id: Some(window_id), + }), + RolloutItem::TurnContext(turn_context_item), + ]) + .await; + { + let mut state = self.state.lock().await; + state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact); + } + self.recompute_token_usage(turn_context).await; + Some(window_id) + } + pub(crate) async fn reference_context_item(&self) -> Option { let state = self.state.lock().await; state.reference_context_item() @@ -3012,6 +3117,7 @@ impl Session { /// Mid-turn compaction is the other path that can re-establish that baseline when it /// reinjects full initial context into replacement history. Other non-regular tasks /// intentionally do not update the baseline. + #[instrument(level = "trace", skip_all)] pub(crate) async fn record_context_updates_and_set_reference_context_item( &self, turn_context: &TurnContext, @@ -3201,7 +3307,7 @@ impl Session { // Persist the user message to history, but emit the turn item from `UserInput` so // UI-only `text_elements` are preserved. `ResponseItem::Message` does not carry // those spans, and `record_response_item_and_emit_turn_item` would drop them. - let response_item = ResponseItem::from(ResponseInputItem::from(input.to_vec())); + let response_item = self.response_item_from_user_input(turn_context, input.to_vec()); self.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) .await; let mut user_message_item = UserMessageItem::new(input); diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 00a89b2b5c..60c0280365 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -114,6 +114,7 @@ pub(super) async fn spawn_review_thread( config: per_turn_config, auth_manager: auth_manager_for_context, model_info: model_info.clone(), + comp_hash: model_info.comp_hash.clone(), tool_mode, session_telemetry: session_telemetry_for_context, provider: provider_for_context, diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index a4c042af0c..e886648512 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -8,6 +8,7 @@ pub(super) struct RolloutReconstruction { pub(super) history: Vec, pub(super) previous_turn_settings: Option, pub(super) reference_context_item: Option, + pub(super) window_id: u64, } #[derive(Debug, Default)] @@ -33,6 +34,7 @@ struct ActiveReplaySegment<'a> { previous_turn_settings: Option, reference_context_item: TurnReferenceContextItem, base_replacement_history: Option<&'a [ResponseItem]>, + window_id: Option, } fn turn_ids_are_compatible(active_turn_id: Option<&str>, item_turn_id: Option<&str>) -> bool { @@ -45,6 +47,7 @@ fn finalize_active_segment<'a>( base_replacement_history: &mut Option<&'a [ResponseItem]>, previous_turn_settings: &mut Option, reference_context_item: &mut TurnReferenceContextItem, + window_id: &mut Option, pending_rollback_turns: &mut usize, ) { // Thread rollback drops the newest surviving real user-message boundaries. In replay, that @@ -65,6 +68,10 @@ fn finalize_active_segment<'a>( *base_replacement_history = Some(segment_base_replacement_history); } + if window_id.is_none() { + *window_id = active_segment.window_id; + } + // `previous_turn_settings` come from the newest surviving user turn that established them. if previous_turn_settings.is_none() && active_segment.counts_as_user_turn { *previous_turn_settings = active_segment.previous_turn_settings; @@ -97,6 +104,7 @@ impl Session { let mut base_replacement_history: Option<&[ResponseItem]> = None; let mut previous_turn_settings = None; let mut reference_context_item = TurnReferenceContextItem::NeverSet; + let mut window_id = None; // Rollback is "drop the newest N user turns". While scanning in reverse, that becomes // "skip the next N user-turn segments we finalize". let mut pending_rollback_turns = 0usize; @@ -112,6 +120,9 @@ impl Session { RolloutItem::Compacted(compacted) => { let active_segment = active_segment.get_or_insert_with(ActiveReplaySegment::default); + if active_segment.window_id.is_none() { + active_segment.window_id = compacted.window_id; + } // Looking backward, compaction clears any older baseline unless a newer // `TurnContextItem` in this same segment has already re-established it. if matches!( @@ -173,6 +184,7 @@ impl Session { ) { active_segment.previous_turn_settings = Some(PreviousTurnSettings { model: ctx.model.clone(), + comp_hash: ctx.comp_hash.clone(), realtime_active: ctx.realtime_active, }); if matches!( @@ -198,6 +210,7 @@ impl Session { &mut base_replacement_history, &mut previous_turn_settings, &mut reference_context_item, + &mut window_id, &mut pending_rollback_turns, ); } @@ -227,10 +240,19 @@ impl Session { &mut base_replacement_history, &mut previous_turn_settings, &mut reference_context_item, + &mut window_id, &mut pending_rollback_turns, ); } + let fallback_window_id = u64::try_from( + rollout_items + .iter() + .filter(|item| matches!(item, RolloutItem::Compacted(_))) + .count(), + ) + .unwrap_or(u64::MAX); + let mut history = ContextManager::new(); let mut saw_legacy_compaction_without_replacement_history = false; if let Some(base_replacement_history) = base_replacement_history { @@ -296,6 +318,7 @@ impl Session { history: history.raw_items().to_vec(), previous_turn_settings, reference_context_item, + window_id: window_id.unwrap_or(fallback_window_id), } } } diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 9ee70dada3..ca8494168c 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -70,6 +70,7 @@ async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previ network: None, file_system_sandbox_policy: None, model: previous_model.to_string(), + comp_hash: None, personality: turn_context.personality, collaboration_mode: Some(turn_context.collaboration_mode.clone()), multi_agent_version: None, @@ -109,6 +110,7 @@ async fn record_initial_history_resumed_hydrates_previous_turn_settings_from_lif network: None, file_system_sandbox_policy: None, model: previous_model.to_string(), + comp_hash: Some("comp-hash-a".to_string()), personality: turn_context.personality, collaboration_mode: Some(turn_context.collaboration_mode.clone()), multi_agent_version: None, @@ -166,6 +168,7 @@ async fn record_initial_history_resumed_hydrates_previous_turn_settings_from_lif session.previous_turn_settings().await, Some(PreviousTurnSettings { model: previous_model.to_string(), + comp_hash: Some("comp-hash-a".to_string()), realtime_active: Some(turn_context.realtime_active), }) ); @@ -271,6 +274,7 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_com reconstructed.previous_turn_settings, Some(PreviousTurnSettings { model: turn_context.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -364,6 +368,7 @@ async fn reconstruct_history_rollback_keeps_history_and_metadata_in_sync_for_inc reconstructed.previous_turn_settings, Some(PreviousTurnSettings { model: turn_context.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -489,6 +494,7 @@ async fn reconstruct_history_rollback_skips_non_user_turns_for_history_and_metad reconstructed.previous_turn_settings, Some(PreviousTurnSettings { model: turn_context.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -589,6 +595,7 @@ async fn reconstruct_history_rollback_counts_inter_agent_assistant_turns() { reconstructed.previous_turn_settings, Some(PreviousTurnSettings { model: turn_context.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -791,6 +798,7 @@ async fn record_initial_history_resumed_rollback_drops_incomplete_user_turn_comp RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_id: None, }), RolloutItem::EventMsg(EventMsg::ThreadRolledBack( codex_protocol::protocol::ThreadRolledBackEvent { num_turns: 1 }, @@ -809,6 +817,7 @@ async fn record_initial_history_resumed_rollback_drops_incomplete_user_turn_comp session.previous_turn_settings().await, Some(PreviousTurnSettings { model: turn_context.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -846,6 +855,7 @@ async fn record_initial_history_resumed_does_not_seed_reference_context_item_aft RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_id: None, }), ]; @@ -871,6 +881,7 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_does_ RolloutItem::Compacted(CompactedItem { message: "legacy summary".to_string(), replacement_history: None, + window_id: None, }), ]; @@ -902,6 +913,7 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_clear RolloutItem::Compacted(CompactedItem { message: "legacy summary".to_string(), replacement_history: None, + window_id: None, }), RolloutItem::EventMsg(EventMsg::TurnStarted( codex_protocol::protocol::TurnStartedEvent { @@ -959,6 +971,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis network: None, file_system_sandbox_policy: None, model: previous_model.to_string(), + comp_hash: None, personality: turn_context.personality, collaboration_mode: Some(turn_context.collaboration_mode.clone()), multi_agent_version: None, @@ -994,6 +1007,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_id: None, }), RolloutItem::TurnContext(previous_context_item), RolloutItem::EventMsg(EventMsg::TurnComplete( @@ -1019,6 +1033,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis session.previous_turn_settings().await, Some(PreviousTurnSettings { model: previous_model.to_string(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -1038,6 +1053,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis network: None, file_system_sandbox_policy: None, model: previous_model.to_string(), + comp_hash: None, personality: turn_context.personality, collaboration_mode: Some(turn_context.collaboration_mode.clone()), multi_agent_version: None, @@ -1067,6 +1083,7 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu network: None, file_system_sandbox_policy: None, model: previous_model.to_string(), + comp_hash: None, personality: turn_context.personality, collaboration_mode: Some(turn_context.collaboration_mode.clone()), multi_agent_version: None, @@ -1140,6 +1157,7 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_id: None, }), ]; @@ -1155,6 +1173,7 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu session.previous_turn_settings().await, Some(PreviousTurnSettings { model: previous_model.to_string(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -1186,6 +1205,7 @@ async fn record_initial_history_resumed_unmatched_abort_preserves_active_turn_fo network: None, file_system_sandbox_policy: None, model: current_model.to_string(), + comp_hash: None, personality: turn_context.personality, collaboration_mode: Some(turn_context.collaboration_mode.clone()), multi_agent_version: None, @@ -1275,6 +1295,7 @@ async fn record_initial_history_resumed_unmatched_abort_preserves_active_turn_fo session.previous_turn_settings().await, Some(PreviousTurnSettings { model: current_model.to_string(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -1304,6 +1325,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea network: None, file_system_sandbox_policy: None, model: previous_model.to_string(), + comp_hash: None, personality: turn_context.personality, collaboration_mode: Some(turn_context.collaboration_mode.clone()), multi_agent_version: None, @@ -1369,6 +1391,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_id: None, }), ]; @@ -1384,6 +1407,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea session.previous_turn_settings().await, Some(PreviousTurnSettings { model: previous_model.to_string(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -1434,6 +1458,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_preserves_turn_ session.previous_turn_settings().await, Some(PreviousTurnSettings { model: turn_context.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -1463,6 +1488,7 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear network: None, file_system_sandbox_policy: None, model: previous_model.to_string(), + comp_hash: None, personality: turn_context.personality, collaboration_mode: Some(turn_context.collaboration_mode.clone()), multi_agent_version: None, @@ -1529,6 +1555,7 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_id: None, }), // A newer TurnStarted replaces the incomplete compacted turn without a matching // completion/abort for the old one. @@ -1555,6 +1582,7 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear session.previous_turn_settings().await, Some(PreviousTurnSettings { model: previous_model.to_string(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 50ff984231..3da05ed0fd 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -469,10 +469,6 @@ impl Session { #[instrument(name = "session_init", level = "info", skip_all)] #[allow(clippy::too_many_arguments)] - #[expect( - clippy::await_holding_invalid_type, - reason = "session initialization must serialize access through session-owned manager guards" - )] pub(crate) async fn new( mut session_configuration: SessionConfiguration, config: Arc, @@ -519,17 +515,6 @@ impl Session { } InitialHistory::Resumed(resumed_history) => resumed_history.conversation_id, }; - let window_generation = match &initial_history { - InitialHistory::Resumed(resumed_history) => u64::try_from( - resumed_history - .history - .iter() - .filter(|item| matches!(item, RolloutItem::Compacted(_))) - .count(), - ) - .unwrap_or(u64::MAX), - InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => 0, - }; // Kick off independent async setup tasks in parallel to reduce startup latency. // // - initialize thread persistence with new or resumed session info @@ -961,8 +946,20 @@ impl Session { .effective_agent_max_threads(MultiAgentVersion::V2) .unwrap_or(usize::MAX), ); + // Keep one stable manager handle for the session so extension resource clients + // automatically observe the manager installed at startup and on later refreshes. + let mcp_connection_manager = Arc::new(arc_swap::ArcSwap::from_pointee( + McpConnectionManager::new_uninitialized_with_permission_profile( + &config.permissions.approval_policy, + config.permissions.permission_profile(), + config.prefix_mcp_tool_names(), + ), + )); let session_extension_data = codex_extension_api::ExtensionData::new(session_id.to_string()); + session_extension_data.insert(McpResourceClient::new(Arc::clone( + &mcp_connection_manager, + ))); let thread_extension_data = codex_extension_api::ExtensionData::new_with_init( thread_id.to_string(), thread_extension_init, @@ -985,13 +982,7 @@ impl Session { // before any MCP-related events. It is reasonable to consider // changing this to use Option or OnceCell, though the current // setup is straightforward enough and performs well. - mcp_connection_manager: Arc::new(RwLock::new( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), - )), + mcp_connection_manager, mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -1051,9 +1042,6 @@ impl Session { code_mode_service: crate::tools::code_mode::CodeModeService::new(), environment_manager, }; - services - .model_client - .set_window_generation(window_generation); let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) = watch::channel(false); @@ -1116,15 +1104,6 @@ impl Session { sess.send_event_raw(event).await; } - let mut required_mcp_servers: Vec = mcp_servers - .iter() - .filter(|(_, server)| server.enabled() && server.required()) - .map(|(name, _)| name.clone()) - .collect(); - required_mcp_servers.sort(); - let enabled_mcp_server_count = - mcp_servers.values().filter(|server| server.enabled()).count(); - let required_mcp_server_count = required_mcp_servers.len(); let tool_plugin_provenance = mcp_manager.tool_plugin_provenance(config.as_ref()).await; let host_owned_codex_apps_enabled = config .features @@ -1137,11 +1116,13 @@ impl Session { } else { ElicitationCapability::default() }; - { + let mcp_startup_cancellation_token = { let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; cancel_guard.cancel(); - *cancel_guard = CancellationToken::new(); - } + let cancel_token = CancellationToken::new(); + *cancel_guard = cancel_token.clone(); + cancel_token + }; let turn_environment = crate::environment_selection::resolve_environment_selections( sess.services.environment_manager.as_ref(), session_configuration.environment_selections(), @@ -1164,13 +1145,14 @@ impl Session { session_configuration.cwd().to_path_buf(), ), }; - let (mcp_connection_manager, cancel_token) = McpConnectionManager::new( + let mcp_connection_manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, - auth_statuses.clone(), + auth_statuses, &session_configuration.approval_policy, INITIAL_SUBMIT_ID.to_owned(), tx_event.clone(), + mcp_startup_cancellation_token, session_configuration.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), @@ -1185,43 +1167,11 @@ impl Session { .instrument(info_span!( "session_init.mcp_manager_init", otel.name = "session_init.mcp_manager_init", - session_init.enabled_mcp_server_count = enabled_mcp_server_count, - session_init.required_mcp_server_count = required_mcp_server_count, )) .await; - { - let mut manager_guard = sess.services.mcp_connection_manager.write().await; - *manager_guard = mcp_connection_manager; - } - { - let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; - if cancel_guard.is_cancelled() { - cancel_token.cancel(); - } - *cancel_guard = cancel_token; - } - if !required_mcp_servers.is_empty() { - let failures = sess - .services - .mcp_connection_manager - .read() - .await - .required_startup_failures(&required_mcp_servers) - .instrument(info_span!( - "session_init.required_mcp_wait", - otel.name = "session_init.required_mcp_wait", - session_init.required_mcp_server_count = required_mcp_server_count, - )) - .await; - if !failures.is_empty() { - let details = failures - .iter() - .map(|failure| format!("{}: {}", failure.server, failure.error)) - .collect::>() - .join("; "); - anyhow::bail!("required MCP servers failed to initialize: {details}"); - } - } + sess.services + .install_mcp_connection_manager(mcp_connection_manager) + .await?; sess.schedule_startup_prewarm(session_configuration.base_instructions.clone()) .await; let session_start_source = match &initial_history { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index f6166260a0..4079deb6b8 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -42,7 +42,9 @@ use codex_protocol::models::ActivePermissionProfile; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; use codex_protocol::models::FileSystemPermissions; use codex_protocol::models::FunctionCallOutputBody; +use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ImageDetail; use codex_protocol::models::PermissionProfile; use codex_protocol::models::SandboxEnforcement; use codex_protocol::openai_models::ModelServiceTier; @@ -319,8 +321,7 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled() session .services .mcp_connection_manager - .read() - .await + .load_full() .set_elicitations_auto_deny(/*auto_deny*/ true); let requested_schema: McpElicitationSchema = serde_json::from_value(json!({ @@ -1586,6 +1587,7 @@ async fn reconstruct_history_matches_live_compactions() { .await; assert_eq!(expected, reconstructed.history); + assert_eq!(2, reconstructed.window_id); } #[tokio::test] @@ -1613,6 +1615,7 @@ async fn reconstruct_history_uses_replacement_history_verbatim() { let rollout_items = vec![RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(replacement_history.clone()), + window_id: Some(42), })]; let reconstructed = session @@ -1620,6 +1623,7 @@ async fn reconstruct_history_uses_replacement_history_verbatim() { .await; assert_eq!(reconstructed.history, replacement_history); + assert_eq!(42, reconstructed.window_id); } #[tokio::test] @@ -1639,6 +1643,114 @@ async fn record_initial_history_reconstructs_resumed_transcript() { assert_eq!(expected, history.raw_items()); } +#[tokio::test] +async fn resize_all_images_prepares_failures_before_history_insertion() { + let (session, turn_context, _rx) = make_session_and_context_with_auth_and_config_and_rx( + CodexAuth::from_api_key("Test API Key"), + Vec::new(), + |config| { + let _ = config.features.enable(Feature::ResizeAllImages); + }, + ) + .await; + let item = ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload { + body: FunctionCallOutputBody::ContentItems(vec![ + FunctionCallOutputContentItem::InputText { + text: "before".to_string(), + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,%%%".to_string(), + detail: Some(ImageDetail::High), + }, + FunctionCallOutputContentItem::InputImage { + image_url: "https://example.com/image.png".to_string(), + detail: Some(ImageDetail::High), + }, + ]), + success: Some(true), + }, + }; + + session + .record_conversation_items(turn_context.as_ref(), std::slice::from_ref(&item)) + .await; + + assert_eq!( + session.state.lock().await.clone_history().raw_items(), + &[ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload { + body: FunctionCallOutputBody::ContentItems(vec![ + FunctionCallOutputContentItem::InputText { + text: "before".to_string(), + }, + FunctionCallOutputContentItem::InputText { + text: "image content omitted because it could not be processed".to_string(), + }, + FunctionCallOutputContentItem::InputImage { + image_url: "https://example.com/image.png".to_string(), + detail: Some(ImageDetail::High), + }, + ]), + success: Some(true), + }, + }] + ); +} + +#[tokio::test] +async fn resize_all_images_prepares_resumed_history_before_installing_it() { + let (session, _turn_context, _rx) = make_session_and_context_with_auth_and_config_and_rx( + CodexAuth::from_api_key("Test API Key"), + Vec::new(), + |config| { + let _ = config.features.enable(Feature::ResizeAllImages); + }, + ) + .await; + let resumed_item = ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ + ContentItem::InputImage { + image_url: "data:image/png;base64,%%%".to_string(), + detail: Some(ImageDetail::High), + }, + ContentItem::InputText { + text: "keep me".to_string(), + }, + ], + phase: None, + }; + + session + .record_initial_history(InitialHistory::Resumed(ResumedHistory { + conversation_id: ThreadId::default(), + history: vec![RolloutItem::ResponseItem(resumed_item)], + rollout_path: Some(PathBuf::from("/tmp/resume.jsonl")), + })) + .await; + + assert_eq!( + session.state.lock().await.clone_history().raw_items(), + &[ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ + ContentItem::InputText { + text: "image content omitted because it could not be processed".to_string(), + }, + ContentItem::InputText { + text: "keep me".to_string(), + }, + ], + phase: None, + }] + ); +} + #[test] fn resolve_multi_agent_version_handles_unset_and_legacy_history() { let thread_id = ThreadId::default(); @@ -1922,26 +2034,27 @@ async fn record_token_usage_info_notifies_extension_contributors() { records: Arc>>, } - #[async_trait::async_trait] impl codex_extension_api::TokenUsageContributor for TokenUsageRecorder { - async fn on_token_usage( - &self, - session_store: &codex_extension_api::ExtensionData, - thread_store: &codex_extension_api::ExtensionData, - turn_store: &codex_extension_api::ExtensionData, - token_usage: &TokenUsageInfo, - ) { - self.records - .lock() - .expect("token usage records lock") - .push(RecordedTokenUsage { - session_level_id: session_store.level_id().to_string(), - thread_level_id: thread_store.level_id().to_string(), - turn_level_id: turn_store.level_id().to_string(), - token_usage: token_usage.clone(), - saw_session_store: session_store.get::().is_some(), - saw_thread_store: thread_store.get::().is_some(), - }); + fn on_token_usage<'a>( + &'a self, + session_store: &'a codex_extension_api::ExtensionData, + thread_store: &'a codex_extension_api::ExtensionData, + turn_store: &'a codex_extension_api::ExtensionData, + token_usage: &'a TokenUsageInfo, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + self.records + .lock() + .expect("token usage records lock") + .push(RecordedTokenUsage { + session_level_id: session_store.level_id().to_string(), + thread_level_id: thread_store.level_id().to_string(), + turn_level_id: turn_store.level_id().to_string(), + token_usage: token_usage.clone(), + saw_session_store: session_store.get::().is_some(), + saw_thread_store: thread_store.get::().is_some(), + }); + }) } } @@ -2040,25 +2153,32 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() { 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 { - session_level_id: input.session_store.level_id().to_string(), - thread_level_id: input.thread_store.level_id().to_string(), - turn_level_id: input.turn_store.level_id().to_string(), - 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(), - }); + fn on_turn_start<'a>( + &'a self, + input: codex_extension_api::TurnStartInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + self.records + .lock() + .expect("turn start records lock") + .push(RecordedTurnStart { + session_level_id: input.session_store.level_id().to_string(), + thread_level_id: input.thread_store.level_id().to_string(), + turn_level_id: input.turn_store.level_id().to_string(), + 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(), + }); + }) } } @@ -2138,24 +2258,31 @@ async fn turn_error_lifecycle_exposes_error_and_stores() { records: Arc>>, } - #[async_trait::async_trait] impl codex_extension_api::TurnLifecycleContributor for TurnErrorRecorder { - async fn on_turn_error(&self, input: codex_extension_api::TurnErrorInput<'_>) { - self.records - .lock() - .expect("turn error records lock") - .push(RecordedTurnError { - session_level_id: input.session_store.level_id().to_string(), - thread_level_id: input.thread_store.level_id().to_string(), - turn_level_id: input.turn_store.level_id().to_string(), - turn_id: input.turn_id.to_string(), - error: input.error, - saw_session_store: input - .session_store - .get::() - .is_some(), - saw_thread_store: input.thread_store.get::().is_some(), - }); + fn on_turn_error<'a>( + &'a self, + input: codex_extension_api::TurnErrorInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + self.records + .lock() + .expect("turn error records lock") + .push(RecordedTurnError { + session_level_id: input.session_store.level_id().to_string(), + thread_level_id: input.thread_store.level_id().to_string(), + turn_level_id: input.turn_store.level_id().to_string(), + turn_id: input.turn_id.to_string(), + error: input.error, + saw_session_store: input + .session_store + .get::() + .is_some(), + saw_thread_store: input + .thread_store + .get::() + .is_some(), + }); + }) } } @@ -2542,6 +2669,7 @@ async fn record_initial_history_forked_hydrates_previous_turn_settings() { network: None, file_system_sandbox_policy: None, model: previous_model.to_string(), + comp_hash: None, personality: turn_context.personality, collaboration_mode: Some(turn_context.collaboration_mode.clone()), multi_agent_version: None, @@ -2594,6 +2722,7 @@ async fn record_initial_history_forked_hydrates_previous_turn_settings() { session.previous_turn_settings().await, Some(PreviousTurnSettings { model: previous_model.to_string(), + comp_hash: None, realtime_active: Some(turn_context.realtime_active), }) ); @@ -2636,6 +2765,7 @@ async fn thread_rollback_drops_last_turn_from_history() { sess.persist_rollout_items(&rollout_items).await; sess.set_previous_turn_settings(Some(PreviousTurnSettings { model: "stale-model".to_string(), + comp_hash: None, realtime_active: Some(tc.realtime_active), })) .await; @@ -2818,6 +2948,7 @@ async fn thread_rollback_recomputes_previous_turn_settings_and_reference_context .await; sess.set_previous_turn_settings(Some(PreviousTurnSettings { model: "stale-model".to_string(), + comp_hash: None, realtime_active: None, })) .await; @@ -2834,6 +2965,7 @@ async fn thread_rollback_recomputes_previous_turn_settings_and_reference_context sess.previous_turn_settings().await, Some(PreviousTurnSettings { model: tc.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(tc.realtime_active), }) ); @@ -2905,6 +3037,7 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio RolloutItem::Compacted(CompactedItem { message: "summary after compaction".to_string(), replacement_history: Some(compacted_history.clone()), + window_id: Some(7), }), RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { turn_id: compact_turn_id, @@ -2933,6 +3066,7 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio RolloutItem::TurnContext(TurnContextItem { turn_id: Some(rolled_back_turn_id.clone()), model: "rolled-back-model".to_string(), + comp_hash: None, ..first_context_item.clone() }), RolloutItem::ResponseItem(user_message("turn 2 user")), @@ -2951,6 +3085,10 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio Some(first_context_item), ) .await; + { + let mut state = sess.state.lock().await; + state.set_auto_compact_window_id(/*window_id*/ 99); + } handlers::thread_rollback(&sess, "sub-1".to_string(), /*num_turns*/ 1).await; let rollback_event = wait_for_thread_rolled_back(&rx).await; @@ -2958,6 +3096,7 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio assert_eq!(sess.clone_history().await.raw_items(), compacted_history); assert!(sess.reference_context_item().await.is_none()); + assert!(sess.current_window_id().await.ends_with(":7")); } #[tokio::test] @@ -4801,7 +4940,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { ); let services = SessionServices { - mcp_connection_manager: Arc::new(RwLock::new( + mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( McpConnectionManager::new_uninitialized_with_permission_profile( &config.permissions.approval_policy, config.permissions.permission_profile(), @@ -6411,16 +6550,20 @@ async fn submission_loop_channel_close_emits_thread_stop_lifecycle() { expected_thread_id: ThreadId, } - #[async_trait::async_trait] impl codex_extension_api::ThreadLifecycleContributor for ThreadStopRecorder { - async fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) { - assert_eq!( - self.expected_thread_id.to_string(), - input.thread_store.level_id() - ); - assert!(input.session_store.get::().is_some()); - assert!(input.thread_store.get::().is_some()); - self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + fn on_thread_stop<'a>( + &'a self, + input: codex_extension_api::ThreadStopInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + assert_eq!( + self.expected_thread_id.to_string(), + input.thread_store.level_id() + ); + assert!(input.session_store.get::().is_some()); + assert!(input.thread_store.get::().is_some()); + self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + }) } } @@ -6457,33 +6600,41 @@ async fn submission_loop_channel_close_aborts_active_turn_before_thread_stop_lif expected_turn_id: String, } - #[async_trait::async_trait] impl codex_extension_api::ThreadLifecycleContributor for LifecycleRecorder { - async fn on_thread_stop(&self, input: codex_extension_api::ThreadStopInput<'_>) { - assert_eq!( - self.expected_thread_id.to_string(), - input.thread_store.level_id() - ); - self.calls - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push("thread_stop"); + fn on_thread_stop<'a>( + &'a self, + input: codex_extension_api::ThreadStopInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + assert_eq!( + self.expected_thread_id.to_string(), + input.thread_store.level_id() + ); + self.calls + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push("thread_stop"); + }) } } - #[async_trait::async_trait] impl codex_extension_api::TurnLifecycleContributor for LifecycleRecorder { - async fn on_turn_abort(&self, input: codex_extension_api::TurnAbortInput<'_>) { - assert_eq!( - self.expected_thread_id.to_string(), - input.thread_store.level_id() - ); - assert_eq!(self.expected_turn_id, input.turn_store.level_id()); - assert_eq!(TurnAbortReason::Interrupted, input.reason); - self.calls - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .push("turn_abort"); + fn on_turn_abort<'a>( + &'a self, + input: codex_extension_api::TurnAbortInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + assert_eq!( + self.expected_thread_id.to_string(), + input.thread_store.level_id() + ); + assert_eq!(self.expected_turn_id, input.turn_store.level_id()); + assert_eq!(TurnAbortReason::Interrupted, input.reason); + self.calls + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push("turn_abort"); + }) } } @@ -6867,7 +7018,7 @@ where ); let services = SessionServices { - mcp_connection_manager: Arc::new(RwLock::new( + mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( McpConnectionManager::new_uninitialized_with_permission_profile( &config.permissions.approval_policy, config.permissions.permission_profile(), @@ -7312,6 +7463,7 @@ async fn build_settings_update_items_uses_previous_turn_settings_for_realtime_en previous_context_item.realtime_active = None; let previous_turn_settings = PreviousTurnSettings { model: previous_context.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(true), }; let mut current_context = previous_context @@ -7905,6 +8057,7 @@ async fn build_initial_context_uses_previous_turn_settings_for_realtime_end() { let (session, turn_context) = make_session_and_context().await; let previous_turn_settings = PreviousTurnSettings { model: turn_context.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(true), }; @@ -7927,6 +8080,7 @@ async fn build_initial_context_restates_realtime_start_when_reference_context_is turn_context.realtime_active = true; let previous_turn_settings = PreviousTurnSettings { model: turn_context.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(true), }; @@ -7960,6 +8114,18 @@ fn file_system_policy_with_unreadable_glob(turn_context: &TurnContext) -> FileSy policy } +#[tokio::test] +async fn turn_context_item_uses_turn_context_comp_hash_snapshot() { + let (_session, mut turn_context) = make_session_and_context().await; + turn_context.comp_hash = Some("turn-context-hash".to_string()); + turn_context.model_info.comp_hash = Some("model-info-hash".to_string()); + + assert_eq!( + turn_context.to_turn_context_item().comp_hash.as_deref(), + Some("turn-context-hash") + ); +} + #[tokio::test] async fn turn_context_item_omits_legacy_equivalent_file_system_sandbox_policy() { let (_session, turn_context) = make_session_and_context().await; @@ -8151,6 +8317,7 @@ async fn build_initial_context_prepends_model_switch_message() { let (session, turn_context) = make_session_and_context().await; let previous_turn_settings = PreviousTurnSettings { model: "previous-regular-model".to_string(), + comp_hash: None, realtime_active: None, }; @@ -8203,6 +8370,7 @@ async fn record_context_updates_and_set_reference_context_item_persists_full_rei session .set_previous_turn_settings(Some(PreviousTurnSettings { model: previous_context.model_info.slug.clone(), + comp_hash: None, realtime_active: Some(previous_context.realtime_active), })) .await; @@ -8698,15 +8866,19 @@ async fn task_finish_emits_thread_idle_lifecycle_after_active_turn_clears() { expected_thread_id: ThreadId, } - #[async_trait::async_trait] impl codex_extension_api::ThreadLifecycleContributor for ThreadIdleRecorder { - async fn on_thread_idle(&self, input: codex_extension_api::ThreadIdleInput<'_>) { - assert_eq!( - self.expected_thread_id.to_string(), - input.thread_store.level_id() - ); - self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); - self.idle_tx.send(()).await.expect("idle receiver open"); + fn on_thread_idle<'a>( + &'a self, + input: codex_extension_api::ThreadIdleInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + assert_eq!( + self.expected_thread_id.to_string(), + input.thread_store.level_id() + ); + self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + self.idle_tx.send(()).await.expect("idle receiver open"); + }) } } @@ -8740,10 +8912,14 @@ async fn thread_idle_lifecycle_waits_for_trigger_turn_mailbox_work() { calls: Arc, } - #[async_trait::async_trait] impl codex_extension_api::ThreadLifecycleContributor for ThreadIdleRecorder { - async fn on_thread_idle(&self, _input: codex_extension_api::ThreadIdleInput<'_>) { - self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + fn on_thread_idle<'a>( + &'a self, + _input: codex_extension_api::ThreadIdleInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + }) } } @@ -9405,18 +9581,13 @@ async fn abort_review_task_emits_exited_then_aborted_and_records_history() { } #[tokio::test] -#[expect( - clippy::await_holding_invalid_type, - reason = "test builds a router from session-owned MCP manager state" -)] async fn fatal_tool_error_stops_turn_and_reports_error() { let (session, turn_context, _rx) = make_session_and_context_with_rx().await; let tools = { session .services .mcp_connection_manager - .read() - .await + .load_full() .list_all_tools() .await }; @@ -9552,6 +9723,7 @@ async fn sample_rollout( rollout_items.push(RolloutItem::Compacted(CompactedItem { message: summary1.to_string(), replacement_history: None, + window_id: None, })); let user2 = ResponseItem::Message { @@ -9592,6 +9764,7 @@ async fn sample_rollout( rollout_items.push(RolloutItem::Compacted(CompactedItem { message: summary2.to_string(), replacement_history: None, + window_id: None, })); let user3 = ResponseItem::Message { diff --git a/codex-rs/core/src/session/token_budget.rs b/codex-rs/core/src/session/token_budget.rs new file mode 100644 index 0000000000..b8f9f82fbc --- /dev/null +++ b/codex-rs/core/src/session/token_budget.rs @@ -0,0 +1,44 @@ +use super::session::Session; +use super::turn_context::TurnContext; +use crate::context::ContextualUserFragment; +use codex_features::Feature; + +const TOKEN_BUDGET_USAGE_THRESHOLDS: [i64; 3] = [25, 50, 75]; + +pub(super) async fn maybe_record_token_budget_remaining_context( + sess: &Session, + turn_context: &TurnContext, + tokens_before_sampling: i64, + tokens_after_sampling: i64, +) { + if !turn_context.features.enabled(Feature::TokenBudget) { + return; + } + let Some(model_context_window) = turn_context.model_context_window() else { + return; + }; + if model_context_window <= 0 || tokens_after_sampling <= tokens_before_sampling { + return; + } + + let tokens_before_sampling = tokens_before_sampling.max(0); + let tokens_after_sampling = tokens_after_sampling.max(0); + let crossed_threshold = TOKEN_BUDGET_USAGE_THRESHOLDS.iter().any(|threshold| { + tokens_before_sampling.saturating_mul(100) < model_context_window.saturating_mul(*threshold) + && tokens_after_sampling.saturating_mul(100) + >= model_context_window.saturating_mul(*threshold) + }); + if !crossed_threshold { + return; + } + + let tokens_left = model_context_window + .saturating_sub(tokens_after_sampling) + .max(0); + + let response_item = ContextualUserFragment::into( + crate::context::TokenBudgetRemainingContext::new(tokens_left), + ); + sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) + .await; +} diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 959e28834a..7e429dcac5 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -172,6 +172,7 @@ pub(crate) async fn run_turn( .await; sess.set_previous_turn_settings(Some(PreviousTurnSettings { model: turn_context.model_info.slug.clone(), + comp_hash: turn_context.comp_hash.clone(), realtime_active: Some(turn_context.realtime_active), })) .await; @@ -186,10 +187,9 @@ pub(crate) async fn run_turn( let mut stop_hook_active = false; // Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains // many turns, from the perspective of the user, it is a single turn. + let display_roots = turn_diff_display_roots(turn_context.as_ref()).await; let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new( - TurnDiffTracker::with_environment_display_roots( - turn_diff_display_roots(turn_context.as_ref()).await, - ), + TurnDiffTracker::with_environment_display_roots(display_roots), )); // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse @@ -214,22 +214,26 @@ pub(crate) async fn run_turn( } // Construct the input that we will send to the model. - let sampling_request_input: Vec = { + let sampling_request_input: Vec = async { sess.clone_history() .await .for_prompt(&turn_context.model_info.input_modalities) - }; + } + .instrument(trace_span!("run_turn.prepare_sampling_request_input")) + .await; - let window_id = sess.services.model_client.current_window_id(); + let window_id = sess.current_window_id().await; let turn_metadata_header = turn_context .turn_metadata_state .current_header_value_for_model_request(&window_id); + let tokens_before_sampling = sess.get_total_token_usage().await; match run_sampling_request( Arc::clone(&sess), Arc::clone(&turn_context), Arc::clone(&turn_extension_data), Arc::clone(&turn_diff_tracker), &mut client_session, + &window_id, turn_metadata_header.as_deref(), sampling_request_input.clone(), cancellation_token.child_token(), @@ -242,15 +246,20 @@ pub(crate) async fn run_turn( last_agent_message: sampling_request_last_agent_message, } = sampling_request_output; can_drain_pending_input = true; - let has_pending_input = sess.input_queue.has_pending_input(&sess.active_turn).await; + let (has_pending_input, token_status, estimated_token_count) = async { + let has_pending_input = + sess.input_queue.has_pending_input(&sess.active_turn).await; + let token_status = + auto_compact_token_status(sess.as_ref(), turn_context.as_ref()).await; + let estimated_token_count = + sess.get_estimated_token_count(turn_context.as_ref()).await; + (has_pending_input, token_status, estimated_token_count) + } + .instrument(trace_span!("run_turn.collect_post_sampling_state")) + .await; let needs_follow_up = model_needs_follow_up || has_pending_input; - let token_status = - auto_compact_token_status(sess.as_ref(), turn_context.as_ref()).await; let token_limit_reached = token_status.token_limit_reached; - let estimated_token_count = - sess.get_estimated_token_count(turn_context.as_ref()).await; - trace!( turn_id = %turn_context.sub_id, total_usage_tokens = token_status.active_context_tokens, @@ -258,7 +267,6 @@ pub(crate) async fn run_turn( estimated_token_count = ?estimated_token_count, auto_compact_scope_limit = token_status.auto_compact_scope_limit, auto_compact_limit_scope = ?turn_context.config.model_auto_compact_token_limit_scope, - auto_compact_window_ordinal = ?token_status.auto_compact_window_ordinal, auto_compact_window_prefill_tokens = ?token_status.auto_compact_window_prefill_tokens, full_context_window_limit = ?token_status.full_context_window_limit, full_context_window_limit_reached = token_status.full_context_window_limit_reached, @@ -269,6 +277,24 @@ pub(crate) async fn run_turn( "post sampling token usage" ); + let tokens_after_sampling = token_status.active_context_tokens; + super::token_budget::maybe_record_token_budget_remaining_context( + sess.as_ref(), + turn_context.as_ref(), + tokens_before_sampling, + tokens_after_sampling, + ) + .await; + + let started_new_context_window = sess + .maybe_start_new_context_window(turn_context.as_ref()) + .await + .is_some(); + if started_new_context_window && needs_follow_up { + can_drain_pending_input = !model_needs_follow_up; + continue; + } + // as long as compaction works well in getting us way below the token limit, we shouldn't worry about being in an infinite loop. if token_limit_reached && needs_follow_up { if let Err(err) = run_auto_compact( @@ -381,6 +407,7 @@ pub(crate) async fn run_turn( last_agent_message } +#[instrument(level = "trace", skip_all)] async fn turn_diff_display_roots(turn_context: &TurnContext) -> Vec<(String, PathBuf)> { let mut display_roots = Vec::new(); for turn_environment in &turn_context.environments.turn_environments { @@ -396,6 +423,7 @@ async fn turn_diff_display_roots(turn_context: &TurnContext) -> Vec<(String, Pat display_roots } +#[instrument(level = "trace", skip_all)] async fn run_hooks_and_record_inputs( sess: &Arc, turn_context: &Arc, @@ -424,10 +452,7 @@ async fn run_hooks_and_record_inputs( blocked_input && !accepted_user_input } -#[expect( - clippy::await_holding_invalid_type, - reason = "MCP tool listing borrows the read guard across cancellation-aware await" -)] +#[instrument(level = "trace", skip_all)] async fn build_skills_and_plugins( sess: &Arc, turn_context: &TurnContext, @@ -464,8 +489,7 @@ async fn build_skills_and_plugins( match sess .services .mcp_connection_manager - .read() - .await + .load_full() .list_all_tools() .or_cancel(cancellation_token) .await @@ -698,7 +722,6 @@ struct AutoCompactTokenStatus { auto_compact_scope_tokens: i64, auto_compact_scope_limit: i64, full_context_window_limit: Option, - auto_compact_window_ordinal: Option, auto_compact_window_prefill_tokens: Option, full_context_window_limit_reached: bool, token_limit_reached: bool, @@ -709,7 +732,6 @@ async fn auto_compact_token_status( turn_context: &TurnContext, ) -> AutoCompactTokenStatus { let active_context_tokens = sess.get_total_token_usage().await; - let mut auto_compact_window_ordinal = None; let mut auto_compact_window_prefill_tokens = None; let (auto_compact_scope_tokens, auto_compact_scope_limit, full_context_window_limit) = match turn_context.config.model_auto_compact_token_limit_scope { @@ -723,7 +745,6 @@ async fn auto_compact_token_status( ), AutoCompactTokenLimitScope::BodyAfterPrefix => { let window = sess.auto_compact_window_snapshot().await; - auto_compact_window_ordinal = Some(window.ordinal); auto_compact_window_prefill_tokens = window.prefill_input_tokens; let baseline = window.prefill_input_tokens.unwrap_or(active_context_tokens); ( @@ -749,13 +770,13 @@ async fn auto_compact_token_status( auto_compact_scope_tokens, auto_compact_scope_limit, full_context_window_limit, - auto_compact_window_ordinal, auto_compact_window_prefill_tokens, full_context_window_limit_reached, token_limit_reached, } } +#[instrument(level = "trace", skip_all)] async fn run_pre_sampling_compact( sess: &Arc, turn_context: &Arc, @@ -778,8 +799,16 @@ async fn run_pre_sampling_compact( Ok(()) } -/// Runs pre-sampling compaction against the previous model when switching to a smaller -/// context-window model. +/// Returns true only when both turns declare compaction compatibility hashes and they differ. +/// A missing hash does not provide enough information to trigger compaction. +fn comp_hash_changed(previous: Option<&str>, current: Option<&str>) -> bool { + previous + .zip(current) + .is_some_and(|(previous, current)| previous != current) +} + +/// Runs pre-sampling compaction against the previous model when its compaction compatibility +/// hash changed or when switching to a smaller context-window model. /// /// Returns `Err(_)` only when compaction was attempted and failed. async fn maybe_run_previous_model_inline_compact( @@ -790,12 +819,29 @@ async fn maybe_run_previous_model_inline_compact( let Some(previous_turn_settings) = sess.previous_turn_settings().await else { return Ok(()); }; + let should_compact_for_comp_hash_change = comp_hash_changed( + previous_turn_settings.comp_hash.as_deref(), + turn_context.comp_hash.as_deref(), + ); let previous_model_turn_context = Arc::new( turn_context .with_model(previous_turn_settings.model, &sess.services.models_manager) .await, ); + if should_compact_for_comp_hash_change { + run_auto_compact( + sess, + &previous_model_turn_context, + client_session, + InitialContextInjection::DoNotInject, + CompactionReason::CompHashChanged, + CompactionPhase::PreTurn, + ) + .await?; + return Ok(()); + } + let Some(old_context_window) = previous_model_turn_context.model_context_window() else { return Ok(()); }; @@ -834,6 +880,11 @@ async fn maybe_run_previous_model_inline_compact( Ok(()) } +#[instrument( + level = "trace", + skip_all, + fields(reason = ?reason, phase = ?phase) +)] async fn run_auto_compact( sess: &Arc, turn_context: &Arc, @@ -942,6 +993,7 @@ pub(super) fn collect_explicit_app_ids_from_skill_items( connector_ids } +#[instrument(level = "trace", skip_all)] pub(crate) fn build_prompt( input: Vec, router: &ToolRouter, @@ -977,6 +1029,7 @@ async fn run_sampling_request( turn_store: Arc, turn_diff_tracker: SharedTurnDiffTracker, client_session: &mut ModelClientSession, + window_id: &str, turn_metadata_header: Option<&str>, input: Vec, cancellation_token: CancellationToken, @@ -1020,6 +1073,7 @@ async fn run_sampling_request( Arc::clone(&turn_context), Arc::clone(&turn_store), client_session, + window_id, turn_metadata_header, Arc::clone(&turn_diff_tracker), &prompt, @@ -1062,10 +1116,6 @@ async fn run_sampling_request( } } -#[expect( - clippy::await_holding_invalid_type, - reason = "tool router construction reads through the session-owned manager guard" -)] #[instrument(level = "trace", skip_all, fields( @@ -1079,22 +1129,17 @@ pub(crate) async fn built_tools( turn_context: &TurnContext, cancellation_token: &CancellationToken, ) -> CodexResult> { - let mcp_connection_manager = sess - .services - .mcp_connection_manager - .read() - .instrument(trace_span!("read_mcp_connection_manager")) - .await; + let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); let has_mcp_servers = mcp_connection_manager.has_servers(); let all_mcp_tools = mcp_connection_manager .list_all_tools() .or_cancel(cancellation_token) .await?; - drop(mcp_connection_manager); let loaded_plugins = sess .services .plugins_manager .plugins_for_config(&turn_context.config.plugins_config_input()) + .instrument(trace_span!("built_tools.load_plugins")) .await; let apps_enabled = turn_context.apps_enabled(); @@ -1125,35 +1170,39 @@ pub(crate) async fn built_tools( .into_iter() .map(|connector_id| connector_id.0) .collect::>(); - let discoverable_tools = if apps_enabled && tool_suggest_enabled(turn_context) { - if let Some(accessible_connectors) = accessible_connectors_with_enabled_state.as_ref() { - match connectors::list_tool_suggest_discoverable_tools_with_auth( - &turn_context.config, - sess.services.plugins_manager.as_ref(), - auth.as_ref(), - accessible_connectors.as_slice(), - &loaded_plugin_app_connector_ids, - ) - .await - .map(|discoverable_tools| { - filter_request_plugin_install_discoverable_tools_for_client( - discoverable_tools, - turn_context.app_server_client_name.as_deref(), + let discoverable_tools = async { + if apps_enabled && tool_suggest_enabled(turn_context) { + if let Some(accessible_connectors) = accessible_connectors_with_enabled_state.as_ref() { + match connectors::list_tool_suggest_discoverable_tools_with_auth( + &turn_context.config, + sess.services.plugins_manager.as_ref(), + auth.as_ref(), + accessible_connectors.as_slice(), + &loaded_plugin_app_connector_ids, ) - }) { - Ok(discoverable_tools) if discoverable_tools.is_empty() => None, - Ok(discoverable_tools) => Some(discoverable_tools), - Err(err) => { - warn!("failed to load discoverable tool suggestions: {err:#}"); - None + .await + .map(|discoverable_tools| { + filter_request_plugin_install_discoverable_tools_for_client( + discoverable_tools, + turn_context.app_server_client_name.as_deref(), + ) + }) { + Ok(discoverable_tools) if discoverable_tools.is_empty() => None, + Ok(discoverable_tools) => Some(discoverable_tools), + Err(err) => { + warn!("failed to load discoverable tool suggestions: {err:#}"); + None + } } + } else { + None } } else { None } - } else { - None - }; + } + .instrument(trace_span!("built_tools.load_discoverable_tools")) + .await; let mcp_tool_exposure = build_mcp_tool_exposure( &all_mcp_tools, @@ -1754,6 +1803,7 @@ async fn try_run_sampling_request( turn_context: Arc, turn_store: Arc, client_session: &mut ModelClientSession, + window_id: &str, turn_metadata_header: Option<&str>, turn_diff_tracker: SharedTurnDiffTracker, prompt: &Prompt, @@ -1775,6 +1825,7 @@ async fn try_run_sampling_request( let sampling_timing_guard = turn_context.turn_timing_state.begin_sampling(); let mut stream = client_session .stream( + window_id, prompt, &turn_context.model_info, &turn_context.session_telemetry, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index e3bd01081a..32c209c880 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -61,6 +61,7 @@ pub struct TurnContext { pub config: Arc, pub(crate) auth_manager: Option>, pub(crate) model_info: ModelInfo, + pub(crate) comp_hash: Option, pub(crate) tool_mode: ToolMode, pub(crate) session_telemetry: SessionTelemetry, pub(crate) provider: SharedModelProvider, @@ -229,6 +230,7 @@ impl TurnContext { config: Arc::new(config), auth_manager: self.auth_manager.clone(), model_info: model_info.clone(), + comp_hash: model_info.comp_hash.clone(), tool_mode, session_telemetry: self .session_telemetry @@ -355,6 +357,7 @@ impl TurnContext { network: self.turn_context_network_item(), file_system_sandbox_policy: self.non_legacy_file_system_sandbox_policy(), model: self.model_info.slug.clone(), + comp_hash: self.comp_hash.clone(), personality: self.personality, collaboration_mode: Some(self.collaboration_mode.clone()), multi_agent_version: Some(self.multi_agent_version), @@ -531,6 +534,7 @@ impl Session { config: per_turn_config.clone(), auth_manager: auth_manager_for_context, model_info: model_info.clone(), + comp_hash: model_info.comp_hash.clone(), tool_mode, session_telemetry: session_telemetry_for_context, provider: provider_for_context, @@ -724,7 +728,7 @@ impl Session { .unwrap_or_else(|| session_configuration.cwd().clone()); let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone()); { - let mcp_connection_manager = self.services.mcp_connection_manager.read().await; + let mcp_connection_manager = self.services.mcp_connection_manager.load_full(); mcp_connection_manager.set_approval_policy(&session_configuration.approval_policy); mcp_connection_manager .set_permission_profile(session_configuration.permission_profile()); diff --git a/codex-rs/core/src/session/turn_tests.rs b/codex-rs/core/src/session/turn_tests.rs index 81a43ff5a9..3511c03bda 100644 --- a/codex-rs/core/src/session/turn_tests.rs +++ b/codex-rs/core/src/session/turn_tests.rs @@ -7,20 +7,21 @@ use std::sync::Arc; struct RewriteAgentMessageContributor; -#[async_trait::async_trait] impl TurnItemContributor for RewriteAgentMessageContributor { - 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(()) + fn contribute<'a>( + &'a self, + _thread_store: &'a ExtensionData, + _turn_store: &'a ExtensionData, + item: &'a mut TurnItem, + ) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> { + 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(()) + }) } } diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index ba2cc42f29..fda4982545 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -4,6 +4,7 @@ use std::time::Instant; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; +use tokio_util::task::AbortOnDropHandle; use tracing::info; use tracing::warn; @@ -19,7 +20,7 @@ use codex_protocol::error::Result as CodexResult; use codex_protocol::models::BaseInstructions; pub(crate) struct SessionStartupPrewarmHandle { - task: JoinHandle>, + task: AbortOnDropHandle>, started_at: Instant, timeout: Duration, } @@ -40,12 +41,17 @@ impl SessionStartupPrewarmHandle { timeout: Duration, ) -> Self { Self { - task, + task: AbortOnDropHandle::new(task), started_at, timeout, } } + pub(crate) async fn abort(self) { + self.task.abort(); + let _ = self.task.await; + } + async fn resolve( self, session_telemetry: &SessionTelemetry, @@ -260,7 +266,7 @@ async fn schedule_startup_prewarm_inner( build_prompt_started_at.elapsed(), /*status*/ None, ); - let window_id = session.services.model_client.current_window_id(); + let window_id = session.current_window_id().await; let startup_turn_metadata_header = startup_turn_context .turn_metadata_state .current_header_value_for_prewarm(&window_id); @@ -268,6 +274,7 @@ async fn schedule_startup_prewarm_inner( let websocket_warmup_started_at = Instant::now(); client_session .prewarm_websocket( + &window_id, &startup_prompt, &startup_turn_context.model_info, &startup_turn_context.session_telemetry, diff --git a/codex-rs/core/src/state/auto_compact_window.rs b/codex-rs/core/src/state/auto_compact_window.rs index d1deb3c288..1e3eceadad 100644 --- a/codex-rs/core/src/state/auto_compact_window.rs +++ b/codex-rs/core/src/state/auto_compact_window.rs @@ -2,7 +2,6 @@ use codex_protocol::protocol::TokenUsage; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct AutoCompactWindowSnapshot { - pub(crate) ordinal: u64, pub(crate) prefill_input_tokens: Option, } @@ -14,7 +13,8 @@ enum AutoCompactWindowPrefill { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) struct AutoCompactWindow { - ordinal: u64, + window_id: u64, + new_context_window_requested: bool, /// Absolute input-token baseline for the current compaction window. /// /// `body_after_prefix` subtracts this from later active-context usage. It is @@ -26,7 +26,8 @@ pub(super) struct AutoCompactWindow { impl AutoCompactWindow { pub(super) fn new() -> Self { Self { - ordinal: 1, + window_id: 0, + new_context_window_requested: false, prefill_input_tokens: None, } } @@ -35,9 +36,28 @@ impl AutoCompactWindow { self.prefill_input_tokens = None; } - pub(super) fn start_next(&mut self) { - self.ordinal = self.ordinal.saturating_add(1); - self.clear_prefill(); + pub(super) fn window_id(&self) -> u64 { + self.window_id + } + + pub(super) fn set_window_id(&mut self, window_id: u64) { + self.window_id = window_id; + } + + pub(super) fn advance_window_id(&mut self) -> u64 { + self.window_id = self.window_id.saturating_add(1); + self.new_context_window_requested = false; + self.window_id + } + + pub(super) fn request_new_context_window(&mut self) { + self.new_context_window_requested = true; + } + + pub(super) fn take_new_context_window_request(&mut self) -> bool { + let requested = self.new_context_window_requested; + self.new_context_window_requested = false; + requested } /// Records the request-input side of the first server usage sample. The @@ -74,7 +94,6 @@ impl AutoCompactWindow { None => None, }; AutoCompactWindowSnapshot { - ordinal: self.ordinal, prefill_input_tokens, } } @@ -89,10 +108,20 @@ mod tests { fn tracks_prefill_and_window_boundaries() { let mut window = AutoCompactWindow::new(); + assert_eq!(window.window_id(), 0); + window.set_window_id(/*window_id*/ 3); + assert_eq!(window.window_id(), 3); + window.request_new_context_window(); + assert!(window.take_new_context_window_request()); + assert!(!window.take_new_context_window_request()); + window.request_new_context_window(); + assert_eq!(window.advance_window_id(), 4); + assert_eq!(window.window_id(), 4); + assert!(!window.take_new_context_window_request()); + assert_eq!( window.snapshot(), AutoCompactWindowSnapshot { - ordinal: 1, prefill_input_tokens: None, } ); @@ -101,7 +130,6 @@ mod tests { assert_eq!( window.snapshot(), AutoCompactWindowSnapshot { - ordinal: 1, prefill_input_tokens: Some(150), } ); @@ -114,7 +142,6 @@ mod tests { assert_eq!( window.snapshot(), AutoCompactWindowSnapshot { - ordinal: 1, prefill_input_tokens: Some(120), } ); @@ -128,18 +155,8 @@ mod tests { assert_eq!( window.snapshot(), AutoCompactWindowSnapshot { - ordinal: 1, prefill_input_tokens: Some(120), } ); - - window.start_next(); - assert_eq!( - window.snapshot(), - AutoCompactWindowSnapshot { - ordinal: 2, - prefill_input_tokens: None, - } - ); } } diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 59def9c7ab..d8d91e85a8 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -15,6 +15,7 @@ use crate::tools::code_mode::CodeModeService; use crate::tools::network_approval::NetworkApprovalService; use crate::tools::sandboxing::ApprovalStore; use crate::unified_exec::UnifiedExecProcessManager; +use anyhow::Result; use arc_swap::ArcSwap; use arc_swap::ArcSwapOption; use codex_analytics::AnalyticsEventsClient; @@ -34,12 +35,12 @@ use codex_thread_store::ThreadStore; use std::path::PathBuf; use tokio::runtime::Handle; use tokio::sync::Mutex; -use tokio::sync::RwLock; use tokio::sync::watch; use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { - pub(crate) mcp_connection_manager: Arc>, + /// The latest manager; callers retain an owned handle while performing MCP I/O. + pub(crate) mcp_connection_manager: Arc>, pub(crate) mcp_startup_cancellation_token: Mutex, pub(crate) unified_exec_manager: UnifiedExecProcessManager, #[cfg_attr(not(unix), allow(dead_code))] @@ -82,3 +83,18 @@ pub(crate) struct SessionServices { /// the same manager through child-thread spawn paths without reconstructing it. pub(crate) environment_manager: Arc, } + +impl SessionServices { + /// Installs the manager before validating required servers so startup-time elicitation can + /// resolve through the session's manager while validation waits. + pub(crate) async fn install_mcp_connection_manager( + &self, + manager: McpConnectionManager, + ) -> Result<()> { + self.mcp_connection_manager.store(Arc::new(manager)); + self.mcp_connection_manager + .load_full() + .validate_required_servers() + .await + } +} diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 024f1904cc..269d3e0f60 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -140,14 +140,36 @@ impl SessionState { self.auto_compact_window.set_estimated_prefill(tokens); } - pub(crate) fn start_next_auto_compact_window(&mut self) { - self.auto_compact_window.start_next(); - } - pub(crate) fn auto_compact_window_snapshot(&self) -> AutoCompactWindowSnapshot { self.auto_compact_window.snapshot() } + pub(crate) fn auto_compact_window_id(&self) -> u64 { + self.auto_compact_window.window_id() + } + + pub(crate) fn set_auto_compact_window_id(&mut self, window_id: u64) { + self.auto_compact_window.set_window_id(window_id); + } + + pub(crate) fn advance_auto_compact_window_id(&mut self) -> u64 { + self.auto_compact_window.advance_window_id() + } + + pub(crate) fn request_new_context_window(&mut self) { + self.auto_compact_window.request_new_context_window(); + } + + pub(crate) fn start_new_context_window_if_requested(&mut self) -> Option { + if !self.auto_compact_window.take_new_context_window_request() { + return None; + } + + let window_id = self.auto_compact_window.advance_window_id(); + self.auto_compact_window.clear_prefill(); + Some(window_id) + } + pub(crate) fn token_info(&self) -> Option { self.history.token_info() } diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index 5f0ae085e7..0fbb92b958 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -65,18 +65,16 @@ async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { } #[tokio::test] -async fn replace_history_clears_auto_compact_window_prefill_without_advancing() { +async fn replace_history_clears_auto_compact_window_prefill() { let session_configuration = make_session_configuration_for_tests().await; let mut state = SessionState::new(session_configuration); - state.start_next_auto_compact_window(); state.set_auto_compact_window_estimated_prefill(/*tokens*/ 100); state.replace_history(Vec::new(), /*reference_context_item*/ None); assert_eq!( state.auto_compact_window_snapshot(), AutoCompactWindowSnapshot { - ordinal: 2, prefill_input_tokens: None, } ); diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index 9c4be94576..c151a46fb3 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -167,41 +167,43 @@ struct TestTurnItemContributor; #[derive(Debug)] struct TurnItemContributorRan; -#[async_trait::async_trait] impl TurnItemContributor for TestTurnItemContributor { - 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(()) + fn contribute<'a>( + &'a self, + _thread_store: &'a ExtensionData, + turn_store: &'a ExtensionData, + item: &'a mut TurnItem, + ) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> { + 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(()) + }) } } struct RewriteAgentMessageContributor; -#[async_trait::async_trait] impl TurnItemContributor for RewriteAgentMessageContributor { - 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(()) + fn contribute<'a>( + &'a self, + _thread_store: &'a ExtensionData, + _turn_store: &'a ExtensionData, + item: &'a mut TurnItem, + ) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> { + Box::pin(async move { + 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/mod.rs b/codex-rs/core/src/tasks/mod.rs index 911f30b319..52af0d04db 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -21,6 +21,7 @@ use tracing::info_span; use tracing::trace; use tracing::warn; +use crate::codex_thread::BackgroundTerminalInfo; use crate::config::Config; use crate::context::ContextualUserFragment; use crate::hook_runtime::inspect_pending_input; @@ -782,6 +783,17 @@ impl Session { .await; } + pub(crate) async fn list_background_terminals(&self) -> Vec { + self.services.unified_exec_manager.list_processes().await + } + + pub(crate) async fn terminate_background_terminal(&self, process_id: i32) -> bool { + self.services + .unified_exec_manager + .terminate_process(process_id) + .await + } + async fn handle_task_abort(self: &Arc, task: RunningTask, reason: TurnAbortReason) { let sub_id = task.turn_context.sub_id.clone(); if task.cancellation_token.is_cancelled() { diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 5e9995ed39..90ba837bb2 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -282,6 +282,7 @@ impl ThreadManager { codex_home.to_path_buf(), restriction_product, )); + plugins_manager.set_auth_mode(auth_manager.get_api_auth_mode()); let mcp_manager = Arc::new(McpManager::new_with_extensions( Arc::clone(&plugins_manager), Arc::clone(&extensions), @@ -375,6 +376,7 @@ impl ThreadManager { codex_home.clone(), restriction_product, )); + plugins_manager.set_auth_mode(auth_manager.get_api_auth_mode()); let mcp_manager = Arc::new(McpManager::new(Arc::clone(&plugins_manager))); let skills_manager = Arc::new(SkillsManager::new_with_restriction_product( skills_codex_home, @@ -525,14 +527,12 @@ impl ThreadManager { &self, thread_id: ThreadId, ) -> CodexResult> { - let thread = self.state.get_thread(thread_id).await?; - let mut subtree_thread_ids = Vec::new(); let mut seen_thread_ids = HashSet::new(); subtree_thread_ids.push(thread_id); seen_thread_ids.insert(thread_id); - if let Some(state_db_ctx) = thread.state_db() { + if let Some(state_db_ctx) = self.state.state_db() { for status in [ DirectionalThreadSpawnEdgeStatus::Open, DirectionalThreadSpawnEdgeStatus::Closed, @@ -551,11 +551,8 @@ impl ThreadManager { } } - for descendant_id in thread - .codex - .session - .services - .agent_control + for descendant_id in self + .agent_control() .list_live_agent_subtree_thread_ids(thread_id) .await? { diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index bb0c0af314..48ca2598ab 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -394,20 +394,24 @@ async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() { observed: Arc>>, } - #[async_trait::async_trait] impl codex_extension_api::ThreadLifecycleContributor for InitialDataRecorder { - async fn on_thread_start(&self, input: codex_extension_api::ThreadStartInput<'_, Config>) { - let marker = input - .thread_store - .get::() - .expect("initial extension data should be available"); - *self - .observed - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(( - input.thread_store.level_id().to_string(), - marker.0.to_string(), - )); + fn on_thread_start<'a>( + &'a self, + input: codex_extension_api::ThreadStartInput<'a, Config>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + let marker = input + .thread_store + .get::() + .expect("initial extension data should be available"); + *self + .observed + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(( + input.thread_store.level_id().to_string(), + marker.0.to_string(), + )); + }) } } diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 1d69c8f4cc..2e46c1bd6b 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -91,7 +91,6 @@ impl CodeModeExecuteHandler { } } -#[async_trait::async_trait] impl ToolExecutor for CodeModeExecuteHandler { fn tool_name(&self) -> ToolName { ToolName::plain(PUBLIC_TOOL_NAME) @@ -101,7 +100,13 @@ impl ToolExecutor for CodeModeExecuteHandler { self.spec.clone() } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl CodeModeExecuteHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index 74cc314a7b..953391a573 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -44,7 +44,6 @@ where }) } -#[async_trait::async_trait] impl ToolExecutor for CodeModeWaitHandler { fn tool_name(&self) -> ToolName { ToolName::plain(WAIT_TOOL_NAME) @@ -54,7 +53,13 @@ impl ToolExecutor for CodeModeWaitHandler { create_wait_tool() } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl CodeModeWaitHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/events.rs b/codex-rs/core/src/tools/events.rs index e286af53d4..2751780d58 100644 --- a/codex-rs/core/src/tools/events.rs +++ b/codex-rs/core/src/tools/events.rs @@ -588,7 +588,7 @@ async fn emit_patch_end( if let Some(tracker) = ctx.turn_diff_tracker { let (should_emit_turn_diff, unified_diff) = { let mut guard = tracker.lock().await; - let previous_diff = guard.get_unified_diff(); + let had_unified_diff = guard.has_unified_diff(); let tracker_changed = match tracker_update { TurnDiffTrackerUpdate::Track { environment_id, @@ -605,7 +605,7 @@ async fn emit_patch_end( }; let unified_diff = guard.get_unified_diff(); ( - tracker_changed && (previous_diff.is_some() || unified_diff.is_some()), + tracker_changed && (had_unified_diff || unified_diff.is_some()), unified_diff.unwrap_or_default(), ) }; @@ -718,4 +718,98 @@ mod tests { ) .await; } + + #[tokio::test] + async fn net_zero_patch_emits_empty_turn_diff() { + let (session, turn, rx_event) = + make_session_and_context_with_dynamic_tools_and_rx(Vec::new()).await; + let tracker = Arc::new(Mutex::new(TurnDiffTracker::new())); + let dir = tempdir().expect("tempdir"); + let cwd = AbsolutePathBuf::from_absolute_path(dir.path()).expect("absolute cwd"); + + for patch in [ + "*** Begin Patch\n*** Add File: a.txt\n+one\n*** End Patch", + "*** Begin Patch\n*** Delete File: a.txt\n*** End Patch", + ] { + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + let delta = codex_apply_patch::apply_patch( + patch, + &cwd, + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .expect("apply patch"); + + emit_patch_end( + ToolEventCtx::new(session.as_ref(), turn.as_ref(), "call-id", Some(&tracker)), + HashMap::new(), + String::new(), + String::new(), + PatchApplyStatus::Completed, + TurnDiffTrackerUpdate::Track { + environment_id: None, + delta: &delta, + }, + ) + .await; + + rx_event.recv().await.expect("item completed event"); + let unified_diff = loop { + let event = rx_event.recv().await.expect("turn diff event"); + if let EventMsg::TurnDiff(TurnDiffEvent { unified_diff }) = event.msg { + break unified_diff; + } + }; + if patch.contains("Delete File") { + assert_eq!(unified_diff, ""); + } else { + assert!(unified_diff.contains("+one")); + } + } + } + + #[tokio::test] + async fn invalidation_emits_empty_turn_diff() { + let (session, turn, rx_event) = + make_session_and_context_with_dynamic_tools_and_rx(Vec::new()).await; + let tracker = Arc::new(Mutex::new(TurnDiffTracker::new())); + let dir = tempdir().expect("tempdir"); + let cwd = AbsolutePathBuf::from_absolute_path(dir.path()).expect("absolute cwd"); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + let delta = codex_apply_patch::apply_patch( + "*** Begin Patch\n*** Add File: a.txt\n+one\n*** End Patch", + &cwd, + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .expect("apply patch"); + tracker.lock().await.track_delta("", &delta); + + emit_patch_end( + ToolEventCtx::new(session.as_ref(), turn.as_ref(), "call-id", Some(&tracker)), + HashMap::new(), + String::new(), + String::new(), + PatchApplyStatus::Completed, + TurnDiffTrackerUpdate::Invalidate, + ) + .await; + + rx_event.recv().await.expect("item completed event"); + loop { + let event = rx_event.recv().await.expect("turn diff event"); + if let EventMsg::TurnDiff(TurnDiffEvent { unified_diff }) = event.msg { + assert_eq!(unified_diff, ""); + break; + } + } + } } diff --git a/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs b/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs index c7adb27610..579cc25ffb 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs @@ -13,7 +13,6 @@ use super::*; pub struct ReportAgentJobResultHandler; -#[async_trait::async_trait] impl ToolExecutor for ReportAgentJobResultHandler { fn tool_name(&self) -> ToolName { ToolName::plain("report_agent_job_result") @@ -23,7 +22,13 @@ impl ToolExecutor for ReportAgentJobResultHandler { create_report_agent_job_result_tool() } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl ReportAgentJobResultHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs b/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs index 26438ca7ba..6da517b9cc 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs @@ -14,7 +14,6 @@ use super::*; pub struct SpawnAgentsOnCsvHandler; -#[async_trait::async_trait] impl ToolExecutor for SpawnAgentsOnCsvHandler { fn tool_name(&self) -> ToolName { ToolName::plain("spawn_agents_on_csv") @@ -24,7 +23,13 @@ impl ToolExecutor for SpawnAgentsOnCsvHandler { create_spawn_agents_on_csv_tool() } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl SpawnAgentsOnCsvHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 66fb1d3080..70452b6657 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -305,7 +305,6 @@ async fn effective_patch_permissions( ) } -#[async_trait::async_trait] impl ToolExecutor for ApplyPatchHandler { fn tool_name(&self) -> ToolName { ToolName::plain("apply_patch") @@ -315,7 +314,13 @@ impl ToolExecutor for ApplyPatchHandler { create_apply_patch_freeform_tool(self.multi_environment) } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl ApplyPatchHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/dynamic.rs b/codex-rs/core/src/tools/handlers/dynamic.rs index a74f44fd6d..62eef88c30 100644 --- a/codex-rs/core/src/tools/handlers/dynamic.rs +++ b/codex-rs/core/src/tools/handlers/dynamic.rs @@ -33,7 +33,6 @@ pub struct DynamicToolHandler { tool_name: ToolName, spec: ToolSpec, exposure: ToolExposure, - search_text: String, } impl DynamicToolHandler { @@ -56,12 +55,10 @@ impl DynamicToolHandler { } else { ToolExposure::Direct }, - search_text: build_dynamic_search_text(tool), }) } } -#[async_trait::async_trait] impl ToolExecutor for DynamicToolHandler { fn tool_name(&self) -> ToolName { self.tool_name.clone() @@ -76,8 +73,7 @@ impl ToolExecutor for DynamicToolHandler { } fn search_info(&self) -> Option { - ToolSearchInfo::from_spec( - self.search_text.clone(), + ToolSearchInfo::from_tool_spec( self.spec(), Some(ToolSearchSourceInfo { name: "Dynamic tools".to_string(), @@ -86,7 +82,13 @@ impl ToolExecutor for DynamicToolHandler { ) } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl DynamicToolHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { @@ -212,27 +214,3 @@ async fn request_dynamic_tool( response } - -fn build_dynamic_search_text(tool: &DynamicToolSpec) -> String { - let mut schema_properties = tool - .input_schema - .get("properties") - .and_then(serde_json::Value::as_object) - .map(|map| map.keys().cloned().collect::>()) - .unwrap_or_default(); - schema_properties.sort(); - let mut parts = vec![ - tool.name.clone(), - tool.name.replace('_', " "), - tool.description.clone(), - ]; - if let Some(namespace) = &tool.namespace { - parts.push(namespace.clone()); - } - parts.extend(schema_properties); - parts.join(" ") -} - -#[cfg(test)] -#[path = "dynamic_tests.rs"] -mod tests; diff --git a/codex-rs/core/src/tools/handlers/dynamic_tests.rs b/codex-rs/core/src/tools/handlers/dynamic_tests.rs deleted file mode 100644 index 201ef45855..0000000000 --- a/codex-rs/core/src/tools/handlers/dynamic_tests.rs +++ /dev/null @@ -1,36 +0,0 @@ -use super::*; -use codex_tools::ToolSearchSourceInfo; -use pretty_assertions::assert_eq; -use serde_json::json; - -#[test] -fn search_info_uses_dynamic_tool_metadata_and_parameter_names() { - let handler = DynamicToolHandler::new(&DynamicToolSpec { - namespace: Some("codex_app".to_string()), - name: "automation_update".to_string(), - description: "Create or update automations.".to_string(), - input_schema: json!({ - "type": "object", - "properties": { - "timezone": { "type": "string" }, - "mode": { "type": "string" } - } - }), - defer_loading: true, - }) - .expect("dynamic handler should be created"); - - let search_info = handler.search_info().expect("dynamic search info"); - - assert_eq!( - search_info.entry.search_text, - "automation_update automation update Create or update automations. codex_app mode timezone" - ); - assert_eq!( - search_info.source_info, - Some(ToolSearchSourceInfo { - name: "Dynamic tools".to_string(), - description: Some("Tools provided by the current Codex thread.".to_string()), - }) - ); -} diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index ecbf772e90..d429b6c1fc 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -11,13 +11,11 @@ use codex_tools::ToolSpec; use codex_tools::TurnItemEmissionFuture; use codex_tools::TurnItemEmitter; -use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::stream_events_utils::TurnItemContributorPolicy; use crate::stream_events_utils::finalize_turn_item; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; @@ -30,7 +28,6 @@ impl ExtensionToolAdapter { } } -#[async_trait::async_trait] impl ToolExecutor for ExtensionToolAdapter { fn tool_name(&self) -> ToolName { self.0.tool_name() @@ -52,11 +49,8 @@ impl ToolExecutor for ExtensionToolAdapter { self.0.search_info() } - async fn handle( - &self, - invocation: ToolInvocation, - ) -> Result, FunctionCallError> { - self.0.handle(to_extension_call(&invocation).await).await + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async move { self.0.handle(to_extension_call(&invocation).await).await }) } } @@ -162,7 +156,6 @@ mod tests { struct StubExtensionExecutor; - #[async_trait::async_trait] impl codex_extension_api::ToolExecutor for StubExtensionExecutor { fn tool_name(&self) -> codex_tools::ToolName { codex_tools::ToolName::plain("extension_echo") @@ -187,13 +180,13 @@ mod tests { }) } - async fn handle( - &self, - _call: codex_tools::ToolCall, - ) -> Result, codex_tools::FunctionCallError> { - Ok(Box::new(codex_tools::JsonToolOutput::new( - json!({ "ok": true }), - ))) + fn handle(&self, _call: codex_tools::ToolCall) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async { + Ok( + Box::new(codex_tools::JsonToolOutput::new(json!({ "ok": true }))) + as Box, + ) + }) } } @@ -201,7 +194,6 @@ mod tests { captured_call: Arc>>, } - #[async_trait::async_trait] impl codex_extension_api::ToolExecutor for CapturingExtensionExecutor { fn tool_name(&self) -> codex_tools::ToolName { codex_tools::ToolName::plain("extension_echo") @@ -218,7 +210,13 @@ mod tests { }) } - async fn handle( + fn handle(&self, call: codex_tools::ToolCall) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(call)) + } + } + + impl CapturingExtensionExecutor { + async fn handle_call( &self, call: codex_tools::ToolCall, ) -> Result, codex_tools::FunctionCallError> { @@ -233,9 +231,10 @@ mod tests { call.turn_item_emitter.emit_started(item.clone()).await; call.turn_item_emitter.emit_completed(item).await; *self.captured_call.lock().await = Some(call); - Ok(Box::new(codex_tools::JsonToolOutput::new( - json!({ "ok": true }), - ))) + Ok( + Box::new(codex_tools::JsonToolOutput::new(json!({ "ok": true }))) + as Box, + ) } } @@ -388,16 +387,17 @@ mod tests { struct RecordExtensionTurnItemContributor; - #[async_trait::async_trait] impl TurnItemContributor for RecordExtensionTurnItemContributor { - async fn contribute( - &self, - _thread_store: &ExtensionData, - turn_store: &ExtensionData, - _item: &mut TurnItem, - ) -> Result<(), String> { - turn_store.insert(ExtensionTurnItemContributorRan); - Ok(()) + fn contribute<'a>( + &'a self, + _thread_store: &'a ExtensionData, + turn_store: &'a ExtensionData, + _item: &'a mut TurnItem, + ) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> { + Box::pin(async move { + turn_store.insert(ExtensionTurnItemContributorRan); + Ok(()) + }) } } @@ -431,7 +431,6 @@ mod tests { ); } - #[async_trait::async_trait] impl codex_extension_api::ToolExecutor for ImageGenerationExtensionExecutor { fn tool_name(&self) -> codex_tools::ToolName { codex_tools::ToolName::namespaced("image_gen", "imagegen") @@ -448,7 +447,13 @@ mod tests { }) } - async fn handle( + fn handle(&self, call: codex_tools::ToolCall) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(call)) + } + } + + impl ImageGenerationExtensionExecutor { + async fn handle_call( &self, call: codex_tools::ToolCall, ) -> Result, codex_tools::FunctionCallError> { @@ -474,9 +479,10 @@ mod tests { }, )) .await; - Ok(Box::new(codex_tools::JsonToolOutput::new( - json!({ "ok": true }), - ))) + Ok( + Box::new(codex_tools::JsonToolOutput::new(json!({ "ok": true }))) + as Box, + ) } } diff --git a/codex-rs/core/src/tools/handlers/get_context_remaining.rs b/codex-rs/core/src/tools/handlers/get_context_remaining.rs new file mode 100644 index 0000000000..a49ab65394 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/get_context_remaining.rs @@ -0,0 +1,54 @@ +use crate::context::ContextualUserFragment; +use crate::function_tool::FunctionCallError; +use crate::tools::context::FunctionToolOutput; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; +use crate::tools::context::boxed_tool_output; +use crate::tools::handlers::get_context_remaining_spec::GET_CONTEXT_REMAINING_TOOL_NAME; +use crate::tools::handlers::get_context_remaining_spec::create_get_context_remaining_tool; +use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::ToolExecutor; +use codex_tools::ToolName; +use codex_tools::ToolSpec; + +pub struct GetContextRemainingHandler; + +impl ToolExecutor for GetContextRemainingHandler { + fn tool_name(&self) -> ToolName { + ToolName::plain(GET_CONTEXT_REMAINING_TOOL_NAME) + } + + fn spec(&self) -> ToolSpec { + create_get_context_remaining_tool() + } + + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async move { + if !matches!(invocation.payload, ToolPayload::Function { .. }) { + return Err(FunctionCallError::RespondToModel( + "get_context_remaining handler received unsupported payload".to_string(), + )); + } + + let Some(model_context_window) = invocation.turn.model_context_window() else { + let fragment = crate::context::TokenBudgetRemainingContext::unknown().render(); + return Ok(boxed_tool_output(FunctionToolOutput::from_text( + fragment, + Some(true), + ))); + }; + let active_context_tokens = invocation.session.get_total_token_usage().await.max(0); + let tokens_left = model_context_window + .saturating_sub(active_context_tokens) + .max(0); + let fragment = crate::context::TokenBudgetRemainingContext::new(tokens_left).render(); + + Ok(boxed_tool_output(FunctionToolOutput::from_text( + fragment, + Some(true), + ))) + }) + } +} + +impl CoreToolRuntime for GetContextRemainingHandler {} diff --git a/codex-rs/core/src/tools/handlers/get_context_remaining_spec.rs b/codex-rs/core/src/tools/handlers/get_context_remaining_spec.rs new file mode 100644 index 0000000000..988f2492ff --- /dev/null +++ b/codex-rs/core/src/tools/handlers/get_context_remaining_spec.rs @@ -0,0 +1,17 @@ +use codex_tools::JsonSchema; +use codex_tools::ResponsesApiTool; +use codex_tools::ToolSpec; +use std::collections::BTreeMap; + +pub(crate) const GET_CONTEXT_REMAINING_TOOL_NAME: &str = "get_context_remaining"; + +pub fn create_get_context_remaining_tool() -> ToolSpec { + ToolSpec::Function(ResponsesApiTool { + name: GET_CONTEXT_REMAINING_TOOL_NAME.to_string(), + description: "Get the remaining tokens in the current context window.".to_string(), + strict: false, + defer_loading: None, + parameters: JsonSchema::object(BTreeMap::new(), /*required*/ None, Some(false.into())), + output_schema: None, + }) +} diff --git a/codex-rs/core/src/tools/handlers/list_available_plugins_to_install.rs b/codex-rs/core/src/tools/handlers/list_available_plugins_to_install.rs index d18a1ca36e..766b0e9b75 100644 --- a/codex-rs/core/src/tools/handlers/list_available_plugins_to_install.rs +++ b/codex-rs/core/src/tools/handlers/list_available_plugins_to_install.rs @@ -54,7 +54,6 @@ impl ListAvailablePluginsToInstallHandler { } } -#[async_trait::async_trait] impl ToolExecutor for ListAvailablePluginsToInstallHandler { fn tool_name(&self) -> ToolName { ToolName::plain(LIST_AVAILABLE_PLUGINS_TO_INSTALL_TOOL_NAME) @@ -68,7 +67,13 @@ impl ToolExecutor for ListAvailablePluginsToInstallHandler { false } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl ListAvailablePluginsToInstallHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 77ff3c14a3..efb32aab94 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -64,7 +64,6 @@ fn ensure_mcp_prefix(name: &str) -> String { } } -#[async_trait::async_trait] impl ToolExecutor for McpHandler { fn tool_name(&self) -> ToolName { self.tool_info.canonical_tool_name() @@ -113,7 +112,13 @@ impl ToolExecutor for McpHandler { ) } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl McpHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource.rs index 9357a9722d..e2f83859b6 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource.rs @@ -7,6 +7,8 @@ use codex_protocol::items::McpToolCallItem; use codex_protocol::items::McpToolCallStatus; use codex_protocol::items::TurnItem; use codex_protocol::mcp::CallToolResult; +use codex_protocol::protocol::TruncationPolicy; +use codex_utils_output_truncation::truncate_text; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; use rmcp::model::ReadResourceResult; @@ -270,7 +272,10 @@ fn normalize_required_string(field: &str, value: String) -> Result(payload: T) -> Result +fn serialize_function_output( + payload: T, + truncation_policy: TruncationPolicy, +) -> Result where T: Serialize, { @@ -279,6 +284,9 @@ where "failed to serialize MCP resource response: {err}" )) })?; + // Match regular MCP tool outputs by bounding the copy persisted to the + // rollout and injected into model context. + let content = truncate_text(&content, truncation_policy * 1.2); Ok(FunctionToolOutput::from_text(content, Some(true))) } diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs index b1e82aa8a6..80b183daba 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs @@ -26,7 +26,6 @@ use super::serialize_function_output; pub struct ListMcpResourceTemplatesHandler; -#[async_trait::async_trait] impl ToolExecutor for ListMcpResourceTemplatesHandler { fn tool_name(&self) -> ToolName { ToolName::plain("list_mcp_resource_templates") @@ -40,11 +39,13 @@ impl ToolExecutor for ListMcpResourceTemplatesHandler { true } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP resource template listing reads through the session-owned manager guard" - )] - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl ListMcpResourceTemplatesHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { @@ -107,8 +108,7 @@ impl ToolExecutor for ListMcpResourceTemplatesHandler { let templates = session .services .mcp_connection_manager - .read() - .await + .load_full() .list_all_resource_templates() .await; Ok(ListResourceTemplatesPayload::from_all_servers(templates)) @@ -117,7 +117,7 @@ impl ToolExecutor for ListMcpResourceTemplatesHandler { .await; match payload_result { - Ok(payload) => match serialize_function_output(payload) { + Ok(payload) => match serialize_function_output(payload, turn.truncation_policy) { Ok(output) => { let content = function_call_output_content_items_to_text(&output.body) .unwrap_or_default(); diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs index c2cdec8b15..a77d617b94 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs @@ -26,7 +26,6 @@ use super::serialize_function_output; pub struct ListMcpResourcesHandler; -#[async_trait::async_trait] impl ToolExecutor for ListMcpResourcesHandler { fn tool_name(&self) -> ToolName { ToolName::plain("list_mcp_resources") @@ -40,11 +39,13 @@ impl ToolExecutor for ListMcpResourcesHandler { true } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP resource listing reads through the session-owned manager guard" - )] - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl ListMcpResourcesHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { @@ -105,8 +106,7 @@ impl ToolExecutor for ListMcpResourcesHandler { let resources = session .services .mcp_connection_manager - .read() - .await + .load_full() .list_all_resources() .await; Ok(ListResourcesPayload::from_all_servers(resources)) @@ -115,7 +115,7 @@ impl ToolExecutor for ListMcpResourcesHandler { .await; match payload_result { - Ok(payload) => match serialize_function_output(payload) { + Ok(payload) => match serialize_function_output(payload, turn.truncation_policy) { Ok(output) => { let content = function_call_output_content_items_to_text(&output.body) .unwrap_or_default(); diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs index 126c5d85e8..9ca15718a0 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs @@ -26,7 +26,6 @@ use super::serialize_function_output; pub struct ReadMcpResourceHandler; -#[async_trait::async_trait] impl ToolExecutor for ReadMcpResourceHandler { fn tool_name(&self) -> ToolName { ToolName::plain("read_mcp_resource") @@ -40,7 +39,13 @@ impl ToolExecutor for ReadMcpResourceHandler { true } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl ReadMcpResourceHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { @@ -93,7 +98,7 @@ impl ToolExecutor for ReadMcpResourceHandler { .await; match payload_result { - Ok(payload) => match serialize_function_output(payload) { + Ok(payload) => match serialize_function_output(payload, turn.truncation_policy) { Ok(output) => { let content = function_call_output_content_items_to_text(&output.body) .unwrap_or_default(); diff --git a/codex-rs/core/src/tools/handlers/mcp_resource_tests.rs b/codex-rs/core/src/tools/handlers/mcp_resource_tests.rs index 8a8410b0bd..b52ab1681a 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource_tests.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource_tests.rs @@ -1,6 +1,7 @@ use super::*; use pretty_assertions::assert_eq; use rmcp::model::AnnotateAble; +use rmcp::model::ResourceContents; use serde_json::json; fn resource(uri: &str, name: &str) -> Resource { @@ -123,3 +124,39 @@ fn template_with_server_serializes_server_field() { }) ); } + +#[test] +fn serialize_function_output_preserves_small_payload() { + let payload = json!({"server": "hosted", "resources": []}); + let expected = serde_json::to_string(&payload).expect("serialize payload"); + + let output = serialize_function_output(payload, TruncationPolicy::Bytes(1_024)) + .expect("serialize function output") + .into_text(); + + assert_eq!(output, expected); +} + +#[test] +fn serialize_function_output_caps_read_resource_payload() { + let truncation_policy = TruncationPolicy::Bytes(8_000); + let payload = ReadResourcePayload { + server: "hosted".to_string(), + uri: "skill://large/SKILL.md".to_string(), + result: ReadResourceResult::new(vec![ResourceContents::TextResourceContents { + uri: "skill://large/SKILL.md".to_string(), + mime_type: Some("text/markdown".to_string()), + text: "x".repeat(16_000), + meta: None, + }]), + }; + let serialized = serde_json::to_string(&payload).expect("serialize payload"); + let expected = truncate_text(&serialized, truncation_policy * 1.2); + + let output = serialize_function_output(payload, truncation_policy) + .expect("serialize bounded function output") + .into_text(); + + assert_ne!(output, serialized); + assert_eq!(output, expected); +} diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index fdfb809965..258c9914dd 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -4,6 +4,8 @@ pub(crate) mod apply_patch; pub(crate) mod apply_patch_spec; mod dynamic; pub(crate) mod extension_tools; +mod get_context_remaining; +pub(crate) mod get_context_remaining_spec; mod list_available_plugins_to_install; pub(crate) mod list_available_plugins_to_install_spec; mod mcp; @@ -13,6 +15,8 @@ pub(crate) mod multi_agents; pub(crate) mod multi_agents_common; pub(crate) mod multi_agents_spec; pub(crate) mod multi_agents_v2; +mod new_context_window; +pub(crate) mod new_context_window_spec; mod plan; pub(crate) mod plan_spec; mod request_permissions; @@ -51,11 +55,13 @@ pub use apply_patch::ApplyPatchHandler; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::protocol::AskForApproval; pub use dynamic::DynamicToolHandler; +pub use get_context_remaining::GetContextRemainingHandler; pub use list_available_plugins_to_install::ListAvailablePluginsToInstallHandler; pub use mcp::McpHandler; pub use mcp_resource::ListMcpResourceTemplatesHandler; pub use mcp_resource::ListMcpResourcesHandler; pub use mcp_resource::ReadMcpResourceHandler; +pub use new_context_window::NewContextWindowHandler; pub use plan::PlanHandler; pub use request_permissions::RequestPermissionsHandler; pub use request_plugin_install::RequestPluginInstallHandler; diff --git a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs index 0d4e17447d..d1017adb65 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs @@ -6,7 +6,6 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::namespaced(MULTI_AGENT_V1_NAMESPACE, "close_agent") @@ -23,11 +22,8 @@ impl ToolExecutor for Handler { ) } - async fn handle( - &self, - invocation: ToolInvocation, - ) -> Result, FunctionCallError> { - handle_close_agent(invocation).await.map(boxed_tool_output) + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async move { handle_close_agent(invocation).await.map(boxed_tool_output) }) } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs index 8f95abe4e3..fa464ddbbb 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs @@ -7,7 +7,6 @@ use std::sync::Arc; pub(crate) struct Handler; -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::namespaced(MULTI_AGENT_V1_NAMESPACE, "resume_agent") @@ -24,11 +23,8 @@ impl ToolExecutor for Handler { ) } - async fn handle( - &self, - invocation: ToolInvocation, - ) -> Result, FunctionCallError> { - handle_resume_agent(invocation).await.map(boxed_tool_output) + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async move { handle_resume_agent(invocation).await.map(boxed_tool_output) }) } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs index 6e28d0a1d9..375bb50ac7 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs @@ -6,7 +6,6 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::namespaced(MULTI_AGENT_V1_NAMESPACE, "send_input") @@ -23,7 +22,13 @@ impl ToolExecutor for Handler { ) } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl Handler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index a50d17f34a..204385934e 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -22,7 +22,6 @@ impl Handler { } } -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::namespaced(MULTI_AGENT_V1_NAMESPACE, "spawn_agent") @@ -39,11 +38,8 @@ impl ToolExecutor for Handler { ) } - async fn handle( - &self, - invocation: ToolInvocation, - ) -> Result, FunctionCallError> { - handle_spawn_agent(invocation).await.map(boxed_tool_output) + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async move { handle_spawn_agent(invocation).await.map(boxed_tool_output) }) } } @@ -204,7 +200,7 @@ async fn handle_spawn_agent( turn.session_telemetry.counter( "codex.multi_agent.spawn", /*inc*/ 1, - &[("role", role_tag)], + &[("role", role_tag), ("version", "v1")], ); Ok(SpawnAgentResult { diff --git a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents/wait.rs index 525ee7f2ed..be22b36908 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/wait.rs @@ -27,7 +27,6 @@ impl Handler { } } -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::namespaced(MULTI_AGENT_V1_NAMESPACE, "wait_agent") @@ -44,7 +43,13 @@ impl ToolExecutor for Handler { ) } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl Handler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_spec.rs b/codex-rs/core/src/tools/handlers/multi_agents_spec.rs index 8f0d9d56b9..5fa2cea70a 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_spec.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_spec.rs @@ -26,7 +26,6 @@ pub struct SpawnAgentToolOptions { pub hide_agent_type_model_reasoning: bool, pub include_usage_hint: bool, pub usage_hint_text: Option, - pub max_concurrent_threads_per_session: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -102,7 +101,6 @@ pub fn create_spawn_agent_tool_v2(options: SpawnAgentToolOptions) -> ToolSpec { inherited_model_guidance, options.include_usage_hint, options.usage_hint_text, - options.max_concurrent_threads_per_session, ), strict: false, defer_loading: None, @@ -722,17 +720,9 @@ fn spawn_agent_tool_description_v2( inherited_model_guidance: Option<&str>, include_usage_hint: bool, usage_hint_text: Option, - max_concurrent_threads_per_session: Option, ) -> String { let agent_role_guidance = available_models_description.unwrap_or_default(); let inherited_model_guidance = inherited_model_guidance.unwrap_or_default(); - let concurrency_guidance = max_concurrent_threads_per_session - .map(|limit| { - format!( - "This session is configured with `max_concurrent_threads_per_session = {limit}` for concurrently open agent threads." - ) - }) - .unwrap_or_default(); let tool_description = format!( r#" @@ -743,8 +733,7 @@ The spawned agent will have the same tools as you and the ability to spawn its o {inherited_model_guidance} Only call this tool for a concrete, bounded subtask that can run independently alongside useful local work; otherwise continue locally. It will be able to send you and other running agents messages, and its final answer will be provided to you when it finishes. -The new agent's canonical task name will be provided to it along with the message. -{concurrency_guidance}"# +The new agent's canonical task name will be provided to it along with the message."# ); if !include_usage_hint { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_spec_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_spec_tests.rs index 512e318a82..1ef93b77f0 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_spec_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_spec_tests.rs @@ -47,7 +47,6 @@ fn spawn_agent_tool_v2_requires_task_name_and_lists_visible_models() { hide_agent_type_model_reasoning: false, include_usage_hint: true, usage_hint_text: None, - max_concurrent_threads_per_session: Some(4), }); let ToolSpec::Function(ResponsesApiTool { @@ -69,7 +68,7 @@ fn spawn_agent_tool_v2_requires_task_name_and_lists_visible_models() { .expect("spawn_agent should use object params"); assert!(description.contains("Spawns an agent to work on the specified task.")); assert!(description.contains("The spawned agent will have the same tools as you")); - assert!(description.contains("`max_concurrent_threads_per_session = 4`")); + assert!(!description.contains("max_concurrent_threads_per_session")); assert!(description.contains(SPAWN_AGENT_INHERITED_MODEL_GUIDANCE)); assert!( description @@ -124,7 +123,6 @@ fn spawn_agent_tool_v1_keeps_legacy_fork_context_field() { hide_agent_type_model_reasoning: false, include_usage_hint: true, usage_hint_text: None, - max_concurrent_threads_per_session: None, }); let ToolSpec::Namespace(namespace) = tool else { @@ -182,7 +180,6 @@ fn spawn_agent_tool_caps_visible_model_summaries() { hide_agent_type_model_reasoning: false, include_usage_hint: true, usage_hint_text: None, - max_concurrent_threads_per_session: Some(4), }); let ToolSpec::Function(ResponsesApiTool { description, .. }) = tool else { @@ -227,7 +224,6 @@ fn spawn_agent_tool_hides_service_tier_with_spawn_metadata() { hide_agent_type_model_reasoning: true, include_usage_hint: true, usage_hint_text: None, - max_concurrent_threads_per_session: Some(4), }); let ToolSpec::Function(ResponsesApiTool { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs index 2fe09f5f21..bd4dafe7c2 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs @@ -7,7 +7,6 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::plain("followup_task") @@ -17,7 +16,13 @@ impl ToolExecutor for Handler { create_followup_task_tool() } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl Handler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/interrupt_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/interrupt_agent.rs index 159c49127a..aeed055b4e 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/interrupt_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/interrupt_agent.rs @@ -6,7 +6,6 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::plain("interrupt_agent") @@ -16,13 +15,12 @@ impl ToolExecutor for Handler { create_interrupt_agent_tool_v2() } - async fn handle( - &self, - invocation: ToolInvocation, - ) -> Result, FunctionCallError> { - handle_interrupt_agent(invocation) - .await - .map(boxed_tool_output) + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async move { + handle_interrupt_agent(invocation) + .await + .map(boxed_tool_output) + }) } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs index 46b29bdf10..99e54e1ce9 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs @@ -5,7 +5,6 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::plain("list_agents") @@ -15,7 +14,13 @@ impl ToolExecutor for Handler { create_list_agents_tool() } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl Handler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs index a8e3605840..717e3bc920 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs @@ -7,7 +7,6 @@ use codex_tools::ToolSpec; pub(crate) struct Handler; -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::plain("send_message") @@ -17,7 +16,13 @@ impl ToolExecutor for Handler { create_send_message_tool() } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl Handler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index fc97df05dc..2186a4a6d6 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -22,7 +22,6 @@ impl Handler { } } -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::plain("spawn_agent") @@ -32,11 +31,8 @@ impl ToolExecutor for Handler { create_spawn_agent_tool_v2(self.options.clone()) } - async fn handle( - &self, - invocation: ToolInvocation, - ) -> Result, FunctionCallError> { - handle_spawn_agent(invocation).await.map(boxed_tool_output) + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async move { handle_spawn_agent(invocation).await.map(boxed_tool_output) }) } } @@ -165,7 +161,7 @@ async fn handle_spawn_agent( turn.session_telemetry.counter( "codex.multi_agent.spawn", /*inc*/ 1, - &[("role", role_tag)], + &[("role", role_tag), ("version", "v2")], ); let task_name = String::from(new_agent_path); diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs index 5e37725b77..36e5aa24f6 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs @@ -19,7 +19,6 @@ impl Handler { } } -#[async_trait::async_trait] impl ToolExecutor for Handler { fn tool_name(&self) -> ToolName { ToolName::plain("wait_agent") @@ -29,7 +28,13 @@ impl ToolExecutor for Handler { create_wait_agent_tool_v2(self.options) } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl Handler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/new_context_window.rs b/codex-rs/core/src/tools/handlers/new_context_window.rs new file mode 100644 index 0000000000..e1cd656399 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/new_context_window.rs @@ -0,0 +1,45 @@ +use crate::function_tool::FunctionCallError; +use crate::tools::context::FunctionToolOutput; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; +use crate::tools::context::boxed_tool_output; +use crate::tools::handlers::new_context_window_spec::NEW_CONTEXT_WINDOW_TOOL_NAME; +use crate::tools::handlers::new_context_window_spec::create_new_context_window_tool; +use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::ToolExecutor; +use codex_tools::ToolName; +use codex_tools::ToolSpec; + +pub(crate) const NEW_CONTEXT_WINDOW_MESSAGE: &str = + "A new context window will start without summarizing conversation history."; + +pub struct NewContextWindowHandler; + +impl ToolExecutor for NewContextWindowHandler { + fn tool_name(&self) -> ToolName { + ToolName::plain(NEW_CONTEXT_WINDOW_TOOL_NAME) + } + + fn spec(&self) -> ToolSpec { + create_new_context_window_tool() + } + + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async move { + if !matches!(invocation.payload, ToolPayload::Function { .. }) { + return Err(FunctionCallError::RespondToModel( + "new_context handler received unsupported payload".to_string(), + )); + } + + invocation.session.request_new_context_window().await; + + Ok(boxed_tool_output(FunctionToolOutput::from_text( + NEW_CONTEXT_WINDOW_MESSAGE.to_string(), + Some(true), + ))) + }) + } +} + +impl CoreToolRuntime for NewContextWindowHandler {} diff --git a/codex-rs/core/src/tools/handlers/new_context_window_spec.rs b/codex-rs/core/src/tools/handlers/new_context_window_spec.rs new file mode 100644 index 0000000000..f58b66e589 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/new_context_window_spec.rs @@ -0,0 +1,17 @@ +use codex_tools::JsonSchema; +use codex_tools::ResponsesApiTool; +use codex_tools::ToolSpec; +use std::collections::BTreeMap; + +pub(crate) const NEW_CONTEXT_WINDOW_TOOL_NAME: &str = "new_context"; + +pub fn create_new_context_window_tool() -> ToolSpec { + ToolSpec::Function(ResponsesApiTool { + name: NEW_CONTEXT_WINDOW_TOOL_NAME.to_string(), + description: "Start a new context window.".to_string(), + strict: false, + defer_loading: None, + parameters: JsonSchema::object(BTreeMap::new(), /*required*/ None, Some(false.into())), + output_schema: None, + }) +} diff --git a/codex-rs/core/src/tools/handlers/plan.rs b/codex-rs/core/src/tools/handlers/plan.rs index bd1060b3d7..1b30c281ce 100644 --- a/codex-rs/core/src/tools/handlers/plan.rs +++ b/codex-rs/core/src/tools/handlers/plan.rs @@ -45,7 +45,6 @@ impl ToolOutput for PlanToolOutput { } } -#[async_trait::async_trait] impl ToolExecutor for PlanHandler { fn tool_name(&self) -> ToolName { ToolName::plain("update_plan") @@ -55,7 +54,13 @@ impl ToolExecutor for PlanHandler { create_update_plan_tool() } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl PlanHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/request_permissions.rs b/codex-rs/core/src/tools/handlers/request_permissions.rs index a20ff2b665..4faf9d328b 100644 --- a/codex-rs/core/src/tools/handlers/request_permissions.rs +++ b/codex-rs/core/src/tools/handlers/request_permissions.rs @@ -25,7 +25,6 @@ struct RequestPermissionsEnvironmentArgs { environment_id: Option, } -#[async_trait::async_trait] impl ToolExecutor for RequestPermissionsHandler { fn tool_name(&self) -> ToolName { ToolName::plain("request_permissions") @@ -35,7 +34,13 @@ impl ToolExecutor for RequestPermissionsHandler { create_request_permissions_tool(request_permissions_tool_description()) } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl RequestPermissionsHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index 3f7bd92d25..55906185ab 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -48,7 +48,6 @@ impl RequestPluginInstallHandler { } } -#[async_trait::async_trait] impl ToolExecutor for RequestPluginInstallHandler { fn tool_name(&self) -> ToolName { ToolName::plain(REQUEST_PLUGIN_INSTALL_TOOL_NAME) @@ -62,7 +61,13 @@ impl ToolExecutor for RequestPluginInstallHandler { true } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl RequestPluginInstallHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { @@ -304,10 +309,6 @@ fn is_remote_plugin_install_suggestion(plugin_id: &str) -> bool { .is_some_and(|(_, marketplace_name)| marketplace_name == REMOTE_GLOBAL_MARKETPLACE_NAME) } -#[expect( - clippy::await_holding_invalid_type, - reason = "connector cache refresh reads through the session-owned manager guard" -)] async fn refresh_missing_requested_connectors( session: &crate::session::session::Session, turn: &crate::session::turn_context::TurnContext, @@ -319,7 +320,7 @@ async fn refresh_missing_requested_connectors( return Some(Vec::new()); } - let manager = session.services.mcp_connection_manager.read().await; + let manager = session.services.mcp_connection_manager.load_full(); let mcp_tools = manager.list_all_tools().await; let accessible_connectors = connectors::with_app_enabled_state( connectors::accessible_connectors_from_mcp_tools(&mcp_tools), @@ -358,7 +359,7 @@ fn verified_plugin_install_completed( ) -> bool { let plugins_input = config.plugins_config_input(); plugins_manager - .list_marketplaces_for_config(&plugins_input, &[]) + .list_marketplaces_for_config(&plugins_input, &[], /*include_openai_curated*/ true) .ok() .into_iter() .flat_map(|outcome| outcome.marketplaces) diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs b/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs index 830e9d408e..cacc25ff95 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs @@ -130,6 +130,7 @@ async fn persist_disabled_install_request_writes_plugin_config() { let codex_home = tempdir().expect("tempdir should succeed"); let tool = DiscoverableTool::Plugin(Box::new(DiscoverablePluginInfo { id: "slack@openai-curated".to_string(), + remote_plugin_id: None, name: "Slack".to_string(), description: None, has_skills: true, diff --git a/codex-rs/core/src/tools/handlers/request_user_input.rs b/codex-rs/core/src/tools/handlers/request_user_input.rs index 0aa4b4f6f8..4f26865a68 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input.rs @@ -20,7 +20,6 @@ pub struct RequestUserInputHandler { pub available_modes: Vec, } -#[async_trait::async_trait] impl ToolExecutor for RequestUserInputHandler { fn tool_name(&self) -> ToolName { ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME) @@ -30,7 +29,13 @@ impl ToolExecutor for RequestUserInputHandler { create_request_user_input_tool(request_user_input_tool_description(&self.available_modes)) } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl RequestUserInputHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/shell/shell_command.rs b/codex-rs/core/src/tools/handlers/shell/shell_command.rs index 7c492b3646..b754187d4a 100644 --- a/codex-rs/core/src/tools/handlers/shell/shell_command.rs +++ b/codex-rs/core/src/tools/handlers/shell/shell_command.rs @@ -124,7 +124,6 @@ impl From for ShellCommandHandler { } } -#[async_trait::async_trait] impl ToolExecutor for ShellCommandHandler { fn tool_name(&self) -> ToolName { ToolName::plain("shell_command") @@ -141,7 +140,13 @@ impl ToolExecutor for ShellCommandHandler { true } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl ShellCommandHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/test_sync.rs b/codex-rs/core/src/tools/handlers/test_sync.rs index 0e5fe894ad..08d983b616 100644 --- a/codex-rs/core/src/tools/handlers/test_sync.rs +++ b/codex-rs/core/src/tools/handlers/test_sync.rs @@ -57,7 +57,6 @@ fn barrier_map() -> &'static tokio::sync::Mutex> { BARRIERS.get_or_init(|| tokio::sync::Mutex::new(HashMap::new())) } -#[async_trait::async_trait] impl ToolExecutor for TestSyncHandler { fn tool_name(&self) -> ToolName { ToolName::plain("test_sync_tool") @@ -71,7 +70,13 @@ impl ToolExecutor for TestSyncHandler { true } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl TestSyncHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 2de1e20dd4..9254f4bee0 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -53,7 +53,6 @@ impl ToolSearchHandler { } } -#[async_trait::async_trait] impl ToolExecutor for ToolSearchHandler { fn tool_name(&self) -> ToolName { ToolName::plain(TOOL_SEARCH_TOOL_NAME) @@ -67,7 +66,13 @@ impl ToolExecutor for ToolSearchHandler { true } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl ToolSearchHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs index ba69f0d039..7c914deb4f 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs @@ -71,7 +71,6 @@ impl ExecCommandHandler { } } -#[async_trait::async_trait] impl ToolExecutor for ExecCommandHandler { fn tool_name(&self) -> ToolName { ToolName::plain("exec_command") @@ -92,7 +91,13 @@ impl ToolExecutor for ExecCommandHandler { true } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl ExecCommandHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs index a639ea0065..8ddce65840 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs @@ -31,7 +31,6 @@ struct WriteStdinArgs { pub struct WriteStdinHandler; -#[async_trait::async_trait] impl ToolExecutor for WriteStdinHandler { fn tool_name(&self) -> ToolName { ToolName::plain("write_stdin") @@ -41,7 +40,13 @@ impl ToolExecutor for WriteStdinHandler { create_write_stdin_tool() } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl WriteStdinHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index 5a67db1a66..c3023b4217 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -1,3 +1,4 @@ +use codex_features::Feature; use codex_protocol::items::ImageViewItem; use codex_protocol::items::TurnItem; use codex_protocol::models::DEFAULT_IMAGE_DETAIL; @@ -8,6 +9,7 @@ use codex_protocol::models::ImageDetail; use codex_protocol::models::ResponseInputItem; use codex_protocol::openai_models::InputModality; use codex_utils_image::PromptImageMode; +use codex_utils_image::data_url_from_bytes; use codex_utils_image::load_for_prompt_bytes; use serde::Deserialize; @@ -64,7 +66,6 @@ enum ViewImageDetail { Original, } -#[async_trait::async_trait] impl ToolExecutor for ViewImageHandler { fn tool_name(&self) -> ToolName { ToolName::plain("view_image") @@ -78,7 +79,13 @@ impl ToolExecutor for ViewImageHandler { true } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } +} + +impl ViewImageHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { @@ -170,25 +177,30 @@ impl ToolExecutor for ViewImageHandler { let can_request_original_detail = can_request_original_image_detail(&turn.model_info); let use_original_detail = can_request_original_detail && matches!(detail, Some(ViewImageDetail::Original)); - let image_mode = if use_original_detail { - PromptImageMode::Original - } else { - PromptImageMode::ResizeToFit - }; let image_detail = if use_original_detail { ImageDetail::Original } else { DEFAULT_IMAGE_DETAIL }; - let image = - load_for_prompt_bytes(abs_path.as_path(), file_bytes, image_mode).map_err(|error| { - FunctionCallError::RespondToModel(format!( - "unable to process image at `{}`: {error}", - abs_path.display() - )) - })?; - let image_url = image.into_data_url(); + let image_url = if turn.features.enabled(Feature::ResizeAllImages) { + // The history insertion path owns image decoding and resizing when this is enabled. + data_url_from_bytes("application/octet-stream", &file_bytes) + } else { + let image_mode = if use_original_detail { + PromptImageMode::Original + } else { + PromptImageMode::ResizeToFit + }; + load_for_prompt_bytes(abs_path.as_path(), file_bytes, image_mode) + .map_err(|error| { + FunctionCallError::RespondToModel(format!( + "unable to process image at `{}`: {error}", + abs_path.display() + )) + })? + .into_data_url() + }; let item = TurnItem::ImageView(ImageViewItem { id: call_id, @@ -213,7 +225,7 @@ pub struct ViewImageOutput { impl ToolOutput for ViewImageOutput { fn log_preview(&self) -> String { - self.image_url.clone() + format!("", self.image_url.len()) } fn success_for_logging(&self) -> bool { @@ -259,6 +271,16 @@ mod tests { use std::sync::Arc; use tokio::sync::Mutex; + #[test] + fn log_preview_omits_image_data() { + let output = ViewImageOutput { + image_url: "data:image/png;base64,AAA".to_string(), + image_detail: DEFAULT_IMAGE_DETAIL, + }; + + assert_eq!(output.log_preview(), ""); + } + #[test] fn code_mode_result_returns_image_url_object() { let output = ViewImageOutput { diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 7db1ec9643..c40886b56c 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -257,7 +257,6 @@ mod tests { tool_name: codex_tools::ToolName, } - #[async_trait::async_trait] impl ToolExecutor for ImmediateHandler { fn tool_name(&self) -> codex_tools::ToolName { self.tool_name.clone() @@ -274,14 +273,13 @@ mod tests { }) } - async fn handle( - &self, - _invocation: ToolInvocation, - ) -> Result, FunctionCallError> { - Ok(Box::new(FunctionToolOutput::from_text( - "ok".to_string(), - Some(true), - ))) + fn handle(&self, _invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async { + Ok( + Box::new(FunctionToolOutput::from_text("ok".to_string(), Some(true))) + as Box, + ) + }) } } @@ -294,7 +292,6 @@ mod tests { allow_cleanup: Arc, } - #[async_trait::async_trait] impl ToolExecutor for CancellationCleanupHandler { fn tool_name(&self) -> codex_tools::ToolName { self.tool_name.clone() @@ -311,7 +308,13 @@ mod tests { }) } - async fn handle( + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(invocation)) + } + } + + impl CancellationCleanupHandler { + async fn handle_call( &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { @@ -336,7 +339,7 @@ mod tests { Ok(Box::new(FunctionToolOutput::from_text( "cleanup complete".to_string(), Some(false), - ))) + )) as Box) } } diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 20e08e9f7d..76e3f5f341 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -250,7 +250,6 @@ struct ExposureOverride { exposure: ToolExposure, } -#[async_trait::async_trait] impl ToolExecutor for ExposureOverride { fn tool_name(&self) -> ToolName { self.handler.tool_name() @@ -272,11 +271,8 @@ impl ToolExecutor for ExposureOverride { self.handler.search_info() } - async fn handle( - &self, - invocation: ToolInvocation, - ) -> Result, FunctionCallError> { - self.handler.handle(invocation).await + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + self.handler.handle(invocation) } } diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index f2566251c5..67613f1d08 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -5,7 +5,6 @@ struct TestHandler { tool_name: codex_tools::ToolName, } -#[async_trait::async_trait] impl ToolExecutor for TestHandler { fn tool_name(&self) -> codex_tools::ToolName { self.tool_name.clone() @@ -15,13 +14,15 @@ impl ToolExecutor for TestHandler { test_spec(&self.tool_name) } - async fn handle( - &self, - _invocation: ToolInvocation, - ) -> Result, FunctionCallError> { - Ok(Box::new( - crate::tools::context::FunctionToolOutput::from_text("ok".to_string(), Some(true)), - )) + fn handle(&self, _invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async { + Ok( + Box::new(crate::tools::context::FunctionToolOutput::from_text( + "ok".to_string(), + Some(true), + )) as Box, + ) + }) } } @@ -38,7 +39,6 @@ struct LifecycleTestHandler { result: LifecycleTestResult, } -#[async_trait::async_trait] impl ToolExecutor for LifecycleTestHandler { fn tool_name(&self) -> codex_tools::ToolName { self.tool_name.clone() @@ -48,9 +48,14 @@ impl ToolExecutor for LifecycleTestHandler { test_spec(&self.tool_name) } - async fn handle( + fn handle(&self, _invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call()) + } +} + +impl LifecycleTestHandler { + async fn handle_call( &self, - _invocation: ToolInvocation, ) -> Result, FunctionCallError> { match self.result.clone() { LifecycleTestResult::Ok { success } => Ok(Box::new( @@ -58,7 +63,8 @@ impl ToolExecutor for LifecycleTestHandler { "ok".to_string(), Some(success), ), - )), + ) + as Box), LifecycleTestResult::Err => Err(FunctionCallError::RespondToModel( "handler failed".to_string(), )), diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index cece7fc171..4024cb61f7 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -44,7 +44,6 @@ impl codex_extension_api::ToolContributor for ExtensionEchoContributor { struct ExtensionEchoExecutor; -#[async_trait::async_trait] impl ToolExecutor for ExtensionEchoExecutor { fn tool_name(&self) -> ToolName { ToolName::namespaced("extension/", "echo") @@ -73,7 +72,13 @@ impl ToolExecutor for ExtensionEchoExecutor { }) } - async fn handle( + fn handle(&self, call: ExtensionToolCall) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(call)) + } +} + +impl ExtensionEchoExecutor { + async fn handle_call( &self, call: ExtensionToolCall, ) -> Result, codex_tools::FunctionCallError> { @@ -84,7 +89,7 @@ impl ToolExecutor for ExtensionEchoExecutor { "callId": call.call_id, "conversationHistory": call.conversation_history.items(), "ok": true, - })))) + }))) as Box) } } @@ -95,17 +100,12 @@ fn extension_tool_test_registry() -> Arc> { } #[tokio::test] -#[expect( - clippy::await_holding_invalid_type, - reason = "test builds a router from session-owned MCP manager state" -)] async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow::Result<()> { let (session, turn) = make_session_and_context().await; let mcp_tools = session .services .mcp_connection_manager - .read() - .await + .load_full() .list_all_tools() .await; let router = ToolRouter::from_turn_context( diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index fed62a0ee1..51a0da7ad8 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -9,10 +9,12 @@ use crate::tools::handlers::CodeModeWaitHandler; use crate::tools::handlers::DynamicToolHandler; use crate::tools::handlers::ExecCommandHandler; use crate::tools::handlers::ExecCommandHandlerOptions; +use crate::tools::handlers::GetContextRemainingHandler; use crate::tools::handlers::ListAvailablePluginsToInstallHandler; use crate::tools::handlers::ListMcpResourceTemplatesHandler; use crate::tools::handlers::ListMcpResourcesHandler; use crate::tools::handlers::McpHandler; +use crate::tools::handlers::NewContextWindowHandler; use crate::tools::handlers::PlanHandler; use crate::tools::handlers::ReadMcpResourceHandler; use crate::tools::handlers::RequestPermissionsHandler; @@ -72,7 +74,6 @@ use codex_tools::ToolCall as ExtensionToolCall; use codex_tools::ToolEnvironmentMode; use codex_tools::ToolExecutor; use codex_tools::ToolName; -use codex_tools::ToolOutput; use codex_tools::ToolSearchInfo; use codex_tools::ToolSpec; use codex_tools::UnifiedExecShellMode; @@ -86,6 +87,7 @@ use codex_tools::shell_type_for_model_and_features; use std::collections::BTreeMap; use std::collections::HashSet; use std::sync::Arc; +use tracing::instrument; use tracing::warn; const MULTI_AGENT_V2_NAMESPACE_DESCRIPTION: &str = "Tools for spawning and managing sub-agents."; @@ -148,6 +150,7 @@ struct CoreToolPlanContext<'a> { wait_agent_timeouts: WaitAgentTimeoutOptions, } +#[instrument(level = "trace", skip_all)] pub(crate) fn build_tool_router( turn_context: &TurnContext, params: ToolRouterParams<'_>, @@ -156,6 +159,7 @@ pub(crate) fn build_tool_router( ToolRouter::from_parts(registry, model_visible_specs) } +#[instrument(level = "trace", skip_all)] fn build_tool_specs_and_registry( turn_context: &TurnContext, params: ToolRouterParams<'_>, @@ -186,6 +190,7 @@ fn build_tool_specs_and_registry( build_model_visible_specs_and_registry(turn_context, planned_tools) } +#[instrument(level = "trace", skip_all)] fn build_model_visible_specs_and_registry( turn_context: &TurnContext, planned_tools: PlannedTools, @@ -383,15 +388,6 @@ fn wait_agent_timeout_options(turn_context: &TurnContext) -> WaitAgentTimeoutOpt } } -fn max_concurrent_threads_per_session(turn_context: &TurnContext) -> Option { - multi_agent_v2_enabled(turn_context).then_some( - turn_context - .config - .multi_agent_v2 - .max_concurrent_threads_per_session, - ) -} - fn agent_type_description( turn_context: &TurnContext, default_agent_type_description: &str, @@ -557,6 +553,7 @@ fn code_mode_namespace_descriptions( namespace_descriptions } +#[instrument(level = "trace", skip_all)] fn add_tool_sources(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { add_shell_tools(context, planned_tools); add_mcp_resource_tools(context, planned_tools); @@ -655,6 +652,11 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut planned_tools.add(RequestPermissionsHandler); } + if features.enabled(Feature::TokenBudget) { + planned_tools.add_with_exposure(NewContextWindowHandler, ToolExposure::DirectModelOnly); + planned_tools.add_with_exposure(GetContextRemainingHandler, ToolExposure::DirectModelOnly); + } + if tool_suggest_enabled(turn_context) && let Some(discoverable_tools) = context.discoverable_tools.filter(|tools| !tools.is_empty()) @@ -718,9 +720,6 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mu .hide_spawn_agent_metadata, include_usage_hint: turn_context.config.multi_agent_v2.usage_hint_enabled, usage_hint_text: turn_context.config.multi_agent_v2.usage_hint_text.clone(), - max_concurrent_threads_per_session: max_concurrent_threads_per_session( - turn_context, - ), }), tool_namespace, ), @@ -765,9 +764,6 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mu hide_agent_type_model_reasoning: false, include_usage_hint: turn_context.config.multi_agent_v2.usage_hint_enabled, usage_hint_text: turn_context.config.multi_agent_v2.usage_hint_text.clone(), - max_concurrent_threads_per_session: max_concurrent_threads_per_session( - turn_context, - ), }), exposure, ); @@ -837,6 +833,7 @@ fn add_extension_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut Pl ); } +#[instrument(level = "trace", skip_all)] fn append_tool_search_executor( context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools, @@ -940,7 +937,6 @@ struct MultiAgentV2NamespaceOverride { namespace: String, } -#[async_trait::async_trait] impl ToolExecutor for MultiAgentV2NamespaceOverride { fn tool_name(&self) -> ToolName { ToolName::namespaced(self.namespace.clone(), self.handler.tool_name().name) @@ -969,11 +965,8 @@ impl ToolExecutor for MultiAgentV2NamespaceOverride { self.handler.search_info() } - async fn handle( - &self, - invocation: ToolInvocation, - ) -> Result, codex_tools::FunctionCallError> { - self.handler.handle(invocation).await + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + self.handler.handle(invocation) } } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index f96f41f0f2..2f7f23947b 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -285,7 +285,6 @@ fn use_bedrock_provider(turn: &mut TurnContext) { struct WebRunExtensionTool; -#[async_trait::async_trait] impl ToolExecutor for WebRunExtensionTool { fn tool_name(&self) -> ToolName { ToolName::namespaced("web", "run") @@ -306,17 +305,15 @@ impl ToolExecutor for WebRunExtensionTool { }) } - async fn handle( - &self, - _call: ExtensionToolCall, - ) -> Result, codex_tools::FunctionCallError> { - Ok(Box::new(codex_tools::JsonToolOutput::new(json!({})))) + fn handle(&self, _call: ExtensionToolCall) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async { + Ok(Box::new(codex_tools::JsonToolOutput::new(json!({}))) as Box) + }) } } struct DeferredExtensionTool; -#[async_trait::async_trait] impl ToolExecutor for DeferredExtensionTool { fn tool_name(&self) -> ToolName { ToolName::plain("extension_echo") @@ -344,11 +341,8 @@ impl ToolExecutor for DeferredExtensionTool { ToolExposure::Deferred } - async fn handle( - &self, - _call: ExtensionToolCall, - ) -> Result, codex_tools::FunctionCallError> { - panic!("spec planning should not execute extension tools") + fn handle(&self, _call: ExtensionToolCall) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async { panic!("spec planning should not execute extension tools") }) } } @@ -406,6 +400,7 @@ fn dynamic_tool(namespace: Option<&str>, name: &str, defer_loading: bool) -> Dyn fn discoverable_plugin(id: &str, name: &str) -> DiscoverableTool { DiscoverablePluginInfo { id: id.to_string(), + remote_plugin_id: None, name: name.to_string(), description: Some(format!("{name} plugin")), has_skills: false, @@ -1020,7 +1015,7 @@ async fn multi_agent_feature_selects_one_agent_tool_family() { ToolSpec::Function(tool) => tool.description.as_str(), other => panic!("expected spawn_agent function spec, got {other:?}"), }; - assert!(spawn_agent_description.contains("max_concurrent_threads_per_session = 17")); + assert!(!spawn_agent_description.contains("max_concurrent_threads_per_session")); let direct_model_only = probe(|turn| { set_features( diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index 69dc8ab580..b09996872e 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -30,7 +30,6 @@ struct TestHandler { tool_name: codex_tools::ToolName, } -#[async_trait::async_trait] impl ToolExecutor for TestHandler { fn tool_name(&self) -> codex_tools::ToolName { self.tool_name.clone() @@ -47,14 +46,13 @@ impl ToolExecutor for TestHandler { }) } - async fn handle( - &self, - _invocation: ToolInvocation, - ) -> Result, FunctionCallError> { - Ok(Box::new(FunctionToolOutput::from_text( - "ok".to_string(), - Some(true), - ))) + fn handle(&self, _invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async { + Ok( + Box::new(FunctionToolOutput::from_text("ok".to_string(), Some(true))) + as Box, + ) + }) } } diff --git a/codex-rs/core/src/turn_diff_tracker.rs b/codex-rs/core/src/turn_diff_tracker.rs index 388c2210bd..ed8e9684ad 100644 --- a/codex-rs/core/src/turn_diff_tracker.rs +++ b/codex-rs/core/src/turn_diff_tracker.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::collections::HashSet; use std::path::Path; use std::path::PathBuf; +use std::time::Duration; use sha1::digest::Output; @@ -12,6 +13,14 @@ use codex_apply_patch::AppliedPatchFileChange; const ZERO_OID: &str = "0000000000000000000000000000000000000000"; const DEV_NULL: &str = "/dev/null"; const REGULAR_FILE_MODE: &str = "100644"; +// Normal edits finish well within 100 ms; pathological inputs fall back to a coarse, +// content-exact diff without stalling tool completion. +const DIFF_TIMEOUT: Duration = Duration::from_millis(100); + +struct TrackedContent { + content: String, + revision: u64, +} #[derive(Clone, Debug, Eq, Hash, PartialEq)] struct TrackedPath { @@ -28,14 +37,27 @@ impl TrackedPath { } } +#[derive(Eq, Hash, PartialEq)] +struct DiffCacheKey { + left_path: TrackedPath, + left_revision: Option, + right_path: TrackedPath, + right_revision: Option, +} + /// Tracks the net text diff for the current turn from committed apply_patch /// mutations, without rereading the workspace filesystem. pub struct TurnDiffTracker { valid: bool, display_roots_by_environment: HashMap, - baseline_by_path: HashMap, - current_by_path: HashMap, + baseline_by_path: HashMap, + current_by_path: HashMap, origin_by_current_path: HashMap, + next_revision: u64, + rendered_diffs: HashMap>, + unified_diff: Option, + #[cfg(test)] + rendered_diff_count: std::cell::Cell, } impl Default for TurnDiffTracker { @@ -46,6 +68,11 @@ impl Default for TurnDiffTracker { baseline_by_path: HashMap::new(), current_by_path: HashMap::new(), origin_by_current_path: HashMap::new(), + next_revision: 0, + rendered_diffs: HashMap::new(), + unified_diff: None, + #[cfg(test)] + rendered_diff_count: std::cell::Cell::new(0), } } } @@ -64,6 +91,10 @@ impl TurnDiffTracker { } pub fn track_delta(&mut self, environment_id: &str, delta: &AppliedPatchDelta) { + if !self.valid { + return; + } + if !delta.is_exact() { self.invalidate(); return; @@ -72,17 +103,24 @@ impl TurnDiffTracker { for change in delta.changes() { self.apply_change(environment_id, change); } + self.refresh_unified_diff(); } pub fn invalidate(&mut self) { self.valid = false; + self.rendered_diffs.clear(); + self.unified_diff = None; } pub fn get_unified_diff(&self) -> Option { - if !self.valid { - return None; - } + self.unified_diff.clone() + } + pub(crate) fn has_unified_diff(&self) -> bool { + self.unified_diff.is_some() + } + + fn refresh_unified_diff(&mut self) { let rename_pairs = self.rename_pairs(); let paired_destinations = rename_pairs.values().cloned().collect::>(); let mut handled = HashSet::new(); @@ -95,6 +133,8 @@ impl TurnDiffTracker { paths.sort_by_key(|path| self.display_path(path)); paths.dedup(); + let mut previous_diffs = std::mem::take(&mut self.rendered_diffs); + let mut rendered_diffs = HashMap::new(); let mut aggregated = String::new(); for path in paths { if !handled.insert(path.clone()) { @@ -105,22 +145,41 @@ impl TurnDiffTracker { continue; } - let diff = if let Some(dest) = rename_pairs.get(&path) { + let (left_path, right_path) = if let Some(dest) = rename_pairs.get(&path) { handled.insert(dest.clone()); - self.render_rename_diff(&path, dest) + (&path, dest) } else { - self.render_path_diff(&path) + (&path, &path) }; - if let Some(diff) = diff { - aggregated.push_str(&diff); + let left_content = self.baseline_by_path.get(left_path); + let right_content = self.current_by_path.get(right_path); + let key = DiffCacheKey { + left_path: left_path.clone(), + left_revision: left_content.map(|content| content.revision), + right_path: right_path.clone(), + right_revision: right_content.map(|content| content.revision), + }; + let rendered = previous_diffs.remove(&key).unwrap_or_else(|| { + self.render_diff( + left_path, + left_content.map(|content| content.content.as_str()), + right_path, + right_content.map(|content| content.content.as_str()), + ) + }); + + if let Some(diff) = rendered.as_deref() { + aggregated.push_str(diff); if !aggregated.ends_with('\n') { aggregated.push('\n'); } } + rendered_diffs.insert(key, rendered); } - (!aggregated.is_empty()).then_some(aggregated) + self.rendered_diffs = rendered_diffs; + self.unified_diff = (!aggregated.is_empty()).then_some(aggregated); } fn apply_change(&mut self, environment_id: &str, change: &AppliedPatchChange) { @@ -157,18 +216,20 @@ impl TurnDiffTracker { && !self.baseline_by_path.contains_key(&path) && let Some(overwritten_content) = overwritten_content { + let overwritten_content = self.tracked_content(overwritten_content); self.baseline_by_path - .insert(path.clone(), overwritten_content.to_string()); + .insert(path.clone(), overwritten_content); } - self.current_by_path.insert(path, content.to_string()); + let content = self.tracked_content(content); + self.current_by_path.insert(path, content); } fn apply_delete(&mut self, path: TrackedPath, content: &str) { if self.current_by_path.remove(&path).is_none() && !self.baseline_by_path.contains_key(&path) { - self.baseline_by_path - .insert(path.clone(), content.to_string()); + let content = self.tracked_content(content); + self.baseline_by_path.insert(path.clone(), content); } self.origin_by_current_path.remove(&path); } @@ -184,8 +245,9 @@ impl TurnDiffTracker { if !self.current_by_path.contains_key(&source_path) && !self.baseline_by_path.contains_key(&source_path) { + let old_content = self.tracked_content(old_content); self.baseline_by_path - .insert(source_path.clone(), old_content.to_string()); + .insert(source_path.clone(), old_content); } match move_path { @@ -194,28 +256,38 @@ impl TurnDiffTracker { && !self.baseline_by_path.contains_key(&dest_path) && let Some(overwritten_move_content) = overwritten_move_content { + let overwritten_move_content = self.tracked_content(overwritten_move_content); self.baseline_by_path - .insert(dest_path.clone(), overwritten_move_content.to_string()); + .insert(dest_path.clone(), overwritten_move_content); } let origin = self .origin_by_current_path .remove(&source_path) .unwrap_or_else(|| source_path.clone()); self.current_by_path.remove(&source_path); - self.current_by_path - .insert(dest_path.clone(), new_content.to_string()); + let new_content = self.tracked_content(new_content); + self.current_by_path.insert(dest_path.clone(), new_content); self.origin_by_current_path.remove(&dest_path); if dest_path != origin { self.origin_by_current_path.insert(dest_path, origin); } } None => { - self.current_by_path - .insert(source_path, new_content.to_string()); + let new_content = self.tracked_content(new_content); + self.current_by_path.insert(source_path, new_content); } } } + fn tracked_content(&mut self, content: &str) -> TrackedContent { + let revision = self.next_revision; + self.next_revision += 1; + TrackedContent { + content: content.to_string(), + revision, + } + } + fn rename_pairs(&self) -> HashMap { self.origin_by_current_path .iter() @@ -234,28 +306,6 @@ impl TurnDiffTracker { .collect() } - fn render_path_diff(&self, path: &TrackedPath) -> Option { - self.render_diff( - path, - self.baseline_by_path.get(path).map(String::as_str), - path, - self.current_by_path.get(path).map(String::as_str), - ) - } - - fn render_rename_diff( - &self, - source_path: &TrackedPath, - dest_path: &TrackedPath, - ) -> Option { - self.render_diff( - source_path, - self.baseline_by_path.get(source_path).map(String::as_str), - dest_path, - self.current_by_path.get(dest_path).map(String::as_str), - ) - } - fn render_diff( &self, left_path: &TrackedPath, @@ -267,6 +317,10 @@ impl TurnDiffTracker { return None; } + #[cfg(test)] + self.rendered_diff_count + .set(self.rendered_diff_count.get() + 1); + let left_display = self.display_path(left_path); let right_display = self.display_path(right_path); let left_oid = left_content.map_or_else( @@ -299,16 +353,23 @@ impl TurnDiffTracker { DEV_NULL.to_string() }; - let unified = - similar::TextDiff::from_lines(left_content.unwrap_or(""), right_content.unwrap_or("")) - .unified_diff() - .context_radius(3) - .header(&old_header, &new_header) - .to_string(); + let mut config = similar::TextDiff::configure(); + config.timeout(DIFF_TIMEOUT); + let unified = config + .diff_lines(left_content.unwrap_or(""), right_content.unwrap_or("")) + .unified_diff() + .context_radius(3) + .header(&old_header, &new_header) + .to_string(); diff.push_str(&unified); Some(diff) } + #[cfg(test)] + fn rendered_diff_count(&self) -> usize { + self.rendered_diff_count.get() + } + fn display_path(&self, path: &TrackedPath) -> String { let display = self .display_roots_by_environment diff --git a/codex-rs/core/src/turn_diff_tracker_tests.rs b/codex-rs/core/src/turn_diff_tracker_tests.rs index 6454430404..a7be3d2b0c 100644 --- a/codex-rs/core/src/turn_diff_tracker_tests.rs +++ b/codex-rs/core/src/turn_diff_tracker_tests.rs @@ -2,10 +2,15 @@ use super::*; use codex_apply_patch::AppliedPatchDelta; use codex_apply_patch::MaybeApplyPatchVerified; use codex_exec_server::LOCAL_FS; +use codex_git_utils::ApplyGitRequest; +use codex_git_utils::apply_git_patch; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use std::fs; use std::path::Path; +use std::process::Command; +use std::time::Duration; +use std::time::Instant; use tempfile::tempdir; fn git_blob_sha1_hex(data: &str) -> String { @@ -369,3 +374,122 @@ index {left_oid_b}..{right_oid_b} ); assert_eq!(tracker.get_unified_diff(), Some(expected)); } + +#[tokio::test] +async fn reuses_rendered_diffs_for_unchanged_paths() { + let dir = tempdir().expect("tempdir"); + let mut tracker = tracker_with_root(dir.path()); + + let add_a = apply_verified_patch( + dir.path(), + "*** Begin Patch\n*** Add File: a.txt\n+one\n*** End Patch", + ) + .await; + tracker.track_delta("", &add_a); + assert_eq!(tracker.rendered_diff_count(), 1); + + let add_b = apply_verified_patch( + dir.path(), + "*** Begin Patch\n*** Add File: b.txt\n+two\n*** End Patch", + ) + .await; + tracker.track_delta("", &add_b); + + assert_eq!(tracker.rendered_diff_count(), 2); + assert_eq!( + tracker.get_unified_diff(), + tracker.get_unified_diff(), + "reading the cached aggregate must not render file diffs", + ); + assert_eq!(tracker.rendered_diff_count(), 2); +} + +#[tokio::test] +async fn repeated_updates_only_rerender_the_touched_path() { + let dir = tempdir().expect("tempdir"); + let mut tracker = tracker_with_root(dir.path()); + + for patch in [ + "*** Begin Patch\n*** Add File: stable.txt\n+stable\n*** End Patch".to_string(), + "*** Begin Patch\n*** Add File: hot.txt\n+value 0\n*** End Patch".to_string(), + ] { + tracker.track_delta("", &apply_verified_patch(dir.path(), &patch).await); + } + + for value in 1..=40 { + let patch = format!( + "*** Begin Patch\n*** Update File: hot.txt\n@@\n-value {}\n+value {value}\n*** End Patch", + value - 1, + ); + tracker.track_delta("", &apply_verified_patch(dir.path(), &patch).await); + } + + assert_eq!(tracker.rendered_diff_count(), 42); +} + +#[test] +fn large_rewrite_returns_promptly_and_preserves_exact_content() { + let dir = tempdir().expect("tempdir"); + assert!( + Command::new("git") + .args(["init", "--quiet"]) + .current_dir(dir.path()) + .status() + .expect("run git init") + .success() + ); + assert!( + Command::new("git") + .args(["config", "core.autocrlf", "false"]) + .current_dir(dir.path()) + .status() + .expect("disable line ending conversion") + .success() + ); + let old_content = (0..48_000) + .map(|line| format!("old line {line:05}\n")) + .collect::(); + let new_content = (0..48_000) + .map(|line| format!("new line {line:05}\n")) + .collect::(); + let path = dir.path().join("large.txt"); + fs::write(&path, &old_content).expect("seed large file"); + assert!( + Command::new("git") + .args(["add", "large.txt"]) + .current_dir(dir.path()) + .status() + .expect("run git add") + .success() + ); + let tracker = tracker_with_root(dir.path()); + let tracked_path = TrackedPath::new("", &path); + + let started = Instant::now(); + let diff = tracker + .render_diff( + &tracked_path, + Some(&old_content), + &tracked_path, + Some(&new_content), + ) + .expect("complete rewrite should produce a diff"); + + assert!( + started.elapsed() < Duration::from_secs(2), + "large rewrite took {:?}", + started.elapsed(), + ); + let result = apply_git_patch(&ApplyGitRequest { + cwd: dir.path().to_path_buf(), + diff, + revert: false, + preflight: false, + }) + .expect("apply generated diff"); + assert_eq!(result.exit_code, 0, "{}", result.stderr); + assert_eq!( + fs::read_to_string(path).expect("read large file"), + new_content + ); +} diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 94a7ab7c73..65b18b064d 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -155,6 +155,8 @@ struct ProcessEntry { process: Arc, call_id: String, process_id: i32, + cwd: AbsolutePathBuf, + initial_exec_command_active: Arc, hook_command: String, tty: bool, network_approval: Option, diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index 9f24314da2..d994354130 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -1,5 +1,6 @@ use super::head_tail_buffer::HeadTailBuffer; use super::*; +use crate::codex_thread::BackgroundTerminalInfo; use crate::exec::ExecCapturePolicy; use crate::exec::ExecExpiration; use crate::sandboxing::ExecRequest; @@ -9,6 +10,16 @@ use crate::session::turn_context::TurnContext; use crate::tools::context::ExecCommandToolOutput; use crate::unified_exec::WriteStdinRequest; use crate::unified_exec::process::OutputHandles; +use async_trait::async_trait; +use codex_exec_server::ExecProcess; +use codex_exec_server::ExecProcessEventReceiver; +use codex_exec_server::ExecServerError; +use codex_exec_server::ProcessId; +use codex_exec_server::ProcessSignal; +use codex_exec_server::ReadResponse; +use codex_exec_server::StartedExecProcess; +use codex_exec_server::WriteResponse; +use codex_exec_server::WriteStatus; use codex_sandboxing::SandboxType; use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::approx_token_count; @@ -19,6 +30,8 @@ use pretty_assertions::assert_eq; use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; +use tokio::sync::Notify; +use tokio::sync::watch; use tokio::time::Duration; use tokio::time::Instant; @@ -116,6 +129,8 @@ async fn exec_command_with_tty( process: Arc::clone(&process), call_id: context.call_id.clone(), process_id, + cwd: cwd.clone(), + initial_exec_command_active: Arc::new(std::sync::atomic::AtomicBool::new(true)), hook_command: cmd.to_string(), tty, network_approval: None, @@ -158,6 +173,18 @@ async fn exec_command_with_tty( manager.release_process_id(process_id).await; None }; + if response_process_id.is_some() + && let Some(entry) = manager + .process_store + .lock() + .await + .processes + .get_mut(&process_id) + { + entry + .initial_exec_command_active + .store(false, std::sync::atomic::Ordering::Release); + } Ok(ExecCommandToolOutput { event_call_id: context.call_id, @@ -184,6 +211,82 @@ impl SpawnLifecycle for TestSpawnLifecycle { } } +struct BlockingTerminateExecProcess { + process_id: ProcessId, + terminate_started: watch::Sender, + allow_terminate: Arc, + wake_tx: watch::Sender, +} + +#[async_trait] +impl ExecProcess for BlockingTerminateExecProcess { + fn process_id(&self) -> &ProcessId { + &self.process_id + } + + fn subscribe_wake(&self) -> watch::Receiver { + self.wake_tx.subscribe() + } + + fn subscribe_events(&self) -> ExecProcessEventReceiver { + ExecProcessEventReceiver::empty() + } + + async fn read( + &self, + _after_seq: Option, + _max_bytes: Option, + _wait_ms: Option, + ) -> Result { + Ok(ReadResponse { + chunks: Vec::new(), + next_seq: 1, + exited: false, + exit_code: None, + closed: false, + failure: None, + }) + } + + async fn write(&self, _chunk: Vec) -> Result { + Ok(WriteResponse { + status: WriteStatus::Accepted, + }) + } + + async fn signal(&self, _signal: ProcessSignal) -> Result<(), ExecServerError> { + Ok(()) + } + + async fn terminate(&self) -> Result<(), ExecServerError> { + let _ = self.terminate_started.send(true); + self.allow_terminate.notified().await; + Ok(()) + } +} + +async fn blocking_terminate_unified_process( + process_id: i32, + terminate_started: watch::Sender, + allow_terminate: Arc, +) -> anyhow::Result> { + let (wake_tx, _wake_rx) = watch::channel(0); + Ok(Arc::new( + UnifiedExecProcess::from_exec_server_started( + StartedExecProcess { + process: Arc::new(BlockingTerminateExecProcess { + process_id: process_id.to_string().into(), + terminate_started, + allow_terminate, + wake_tx, + }), + }, + SandboxType::None, + ) + .await?, + )) +} + async fn write_stdin( session: &Arc, process_id: i32, @@ -241,12 +344,23 @@ async fn unified_exec_persists_across_requests() -> anyhow::Result<()> { skip_if_sandbox!(Ok(())); let (session, turn) = test_session_and_turn().await; + #[allow(deprecated)] + let cwd = turn.cwd.clone(); let open_shell = exec_command( &session, &turn, "bash -i", /*yield_time_ms*/ 2_500, /*workdir*/ None, ) .await?; let process_id = open_shell.process_id.expect("expected process_id"); + assert_eq!( + session.list_background_terminals().await, + vec![BackgroundTerminalInfo { + item_id: "call".to_string(), + process_id: process_id.to_string(), + command: "bash -i".to_string(), + cwd, + }] + ); write_stdin( &session, @@ -270,6 +384,10 @@ async fn unified_exec_persists_across_requests() -> anyhow::Result<()> { "expected environment variable output" ); + assert!(session.terminate_background_terminal(process_id).await); + assert!(!session.terminate_background_terminal(process_id).await); + assert!(session.list_background_terminals().await.is_empty()); + Ok(()) } @@ -523,6 +641,148 @@ async fn reusing_completed_process_returns_unknown_process() -> anyhow::Result<( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn terminating_initial_exec_command_rechecks_initial_response_state() -> anyhow::Result<()> { + let (session, turn) = test_session_and_turn().await; + let manager = &session.services.unified_exec_manager; + let process_id = manager.allocate_process_id().await; + let (terminate_started_tx, mut terminate_started_rx) = watch::channel(false); + let allow_terminate = Arc::new(Notify::new()); + let process = blocking_terminate_unified_process( + process_id, + terminate_started_tx, + Arc::clone(&allow_terminate), + ) + .await?; + #[allow(deprecated)] + let cwd = turn.cwd.clone(); + manager.process_store.lock().await.processes.insert( + process_id, + ProcessEntry { + process, + call_id: "call".to_string(), + process_id, + cwd, + initial_exec_command_active: Arc::new(std::sync::atomic::AtomicBool::new(true)), + hook_command: "sleep 60".to_string(), + tty: true, + network_approval: None, + session: Arc::downgrade(&session), + last_used: Instant::now(), + }, + ); + + let terminate_task = tokio::spawn({ + let session = Arc::clone(&session); + async move { session.terminate_background_terminal(process_id).await } + }); + tokio::time::timeout( + Duration::from_secs(2), + terminate_started_rx.wait_for(|started| *started), + ) + .await + .expect("terminate should start") + .expect("terminate signal sender should stay open"); + + { + let mut store = manager.process_store.lock().await; + let entry = store + .processes + .get_mut(&process_id) + .expect("process should remain stored until initial response returns"); + entry + .initial_exec_command_active + .store(false, std::sync::atomic::Ordering::Release); + } + + allow_terminate.notify_waiters(); + let terminated = tokio::time::timeout(Duration::from_secs(2), terminate_task) + .await + .expect("terminate should finish") + .expect("terminate task should not panic"); + assert!(terminated); + assert!( + !manager + .process_store + .lock() + .await + .processes + .contains_key(&process_id) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn terminating_during_stdin_poll_returns_exited_response() -> anyhow::Result<()> { + let (session, turn) = test_session_and_turn().await; + let manager = &session.services.unified_exec_manager; + let process_id = manager.allocate_process_id().await; + let (terminate_started_tx, _terminate_started_rx) = watch::channel(false); + let allow_terminate = Arc::new(Notify::new()); + let process = blocking_terminate_unified_process( + process_id, + terminate_started_tx, + Arc::clone(&allow_terminate), + ) + .await?; + #[allow(deprecated)] + let cwd = turn.cwd.clone(); + let last_used = Instant::now() - Duration::from_secs(1); + manager.process_store.lock().await.processes.insert( + process_id, + ProcessEntry { + process: Arc::clone(&process), + call_id: "call".to_string(), + process_id, + cwd, + initial_exec_command_active: Arc::new(std::sync::atomic::AtomicBool::new(false)), + hook_command: "sleep 60".to_string(), + tty: true, + network_approval: None, + session: Arc::downgrade(&session), + last_used, + }, + ); + + let poll_task = tokio::spawn({ + let session = Arc::clone(&session); + async move { + write_stdin(&session, process_id, "", /*yield_time_ms*/ 60_000).await + } + }); + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let poll_started = manager + .process_store + .lock() + .await + .processes + .get(&process_id) + .is_some_and(|entry| entry.last_used != last_used); + if poll_started { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("poll should clone process handles"); + + manager.release_process_id(process_id).await; + allow_terminate.notify_one(); + process.terminate_confirmed().await?; + + let output = tokio::time::timeout(Duration::from_secs(2), poll_task) + .await + .expect("poll should finish") + .expect("poll task should not panic")?; + assert_eq!(output.process_id, None); + assert!(manager.process_store.lock().await.processes.is_empty()); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn completed_pipe_commands_preserve_exit_code() -> anyhow::Result<()> { let (_, turn) = make_session_and_context().await; diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index 9671429d01..725be9eed2 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -14,6 +14,7 @@ use tokio_util::sync::CancellationToken; use crate::exec::is_likely_sandbox_denied; use codex_exec_server::ExecProcess; +use codex_exec_server::ProcessSignal as ExecServerProcessSignal; use codex_exec_server::ReadResponse as ExecReadResponse; use codex_exec_server::StartedExecProcess; use codex_exec_server::WriteStatus; @@ -23,6 +24,7 @@ use codex_protocol::protocol::TruncationPolicy; use codex_sandboxing::SandboxType; use codex_utils_output_truncation::formatted_truncate_text; use codex_utils_pty::ExecCommandSession; +use codex_utils_pty::ProcessSignal as PtyProcessSignal; use codex_utils_pty::SpawnedPty; use super::UNIFIED_EXEC_OUTPUT_MAX_TOKENS; @@ -31,7 +33,6 @@ use super::head_tail_buffer::HeadTailBuffer; use super::process_state::ProcessState; const EARLY_EXIT_GRACE_PERIOD: Duration = Duration::from_millis(150); - pub(crate) trait SpawnLifecycle: std::fmt::Debug + Send + Sync { /// Returns file descriptors that must stay open across the child `exec()`. /// @@ -194,9 +195,16 @@ impl UnifiedExecProcess { } } - pub(super) fn terminate(&self) { + fn finish_termination(&self) { self.output_closed.store(true, Ordering::Release); self.output_closed_notify.notify_waiters(); + self.cancellation_token.cancel(); + if let Some(output_task) = &self.output_task { + output_task.abort(); + } + } + + pub(super) fn terminate(&self) { match &self.process_handle { ProcessHandle::Local(process_handle) => process_handle.terminate(), ProcessHandle::ExecServer(process_handle) => { @@ -206,9 +214,33 @@ impl UnifiedExecProcess { }); } } - self.cancellation_token.cancel(); - if let Some(output_task) = &self.output_task { - output_task.abort(); + self.finish_termination(); + } + + pub(super) async fn terminate_confirmed(&self) -> Result<(), UnifiedExecError> { + match &self.process_handle { + ProcessHandle::Local(process_handle) => process_handle.terminate(), + ProcessHandle::ExecServer(process_handle) => { + process_handle + .terminate() + .await + .map_err(|err| UnifiedExecError::process_failed(err.to_string()))?; + } + } + self.signal_exit(self.exit_code()); + self.finish_termination(); + Ok(()) + } + + pub(super) async fn interrupt(&self) -> Result<(), UnifiedExecError> { + match &self.process_handle { + ProcessHandle::Local(process_handle) => process_handle + .signal(PtyProcessSignal::Interrupt) + .map_err(|err| UnifiedExecError::process_failed(err.to_string())), + ProcessHandle::ExecServer(process_handle) => process_handle + .signal(ExecServerProcessSignal::Interrupt) + .await + .map_err(|err| UnifiedExecError::process_failed(err.to_string())), } } diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index cb455b113b..d744a72c95 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -11,6 +11,7 @@ use tokio::time::Duration; use tokio::time::Instant; use tokio_util::sync::CancellationToken; +use crate::codex_thread::BackgroundTerminalInfo; use crate::exec_env::CODEX_THREAD_ID_ENV_VAR; use crate::exec_env::create_env; use crate::exec_policy::ExecApprovalRequest; @@ -72,6 +73,7 @@ const UNIFIED_EXEC_ENV: [(&str, &str); 10] = [ const NETWORK_ACCESS_DENIED_MESSAGE: &str = "Network access was denied by the Codex sandbox network proxy."; const LATE_NETWORK_DENIAL_GRACE_PERIOD: Duration = Duration::from_millis(100); +const INTERRUPT: &str = "\u{3}"; /// Test-only override for deterministic unified exec process IDs. /// @@ -174,11 +176,22 @@ struct PreparedProcessHandles { pause_state: Option>, session: Option>, network_approval: Option, + call_id: String, hook_command: String, process_id: i32, tty: bool, } +struct InitialExecCommandGuard { + active: Arc, +} + +impl Drop for InitialExecCommandGuard { + fn drop(&mut self) { + self.active.store(false, Ordering::Release); + } +} + fn exec_server_process_id(process_id: i32) -> String { process_id.to_string() } @@ -413,7 +426,8 @@ impl UnifiedExecProcessManager { // Persist live sessions before the initial yield wait so interrupting the // turn cannot drop the last Arc and terminate the background process. let process_started_alive = !process.has_exited() && process.exit_code().is_none(); - if process_started_alive { + let _initial_exec_command_guard = if process_started_alive { + let initial_exec_command_active = Arc::new(AtomicBool::new(true)); self.store_process( Arc::clone(&process), context, @@ -425,9 +439,15 @@ impl UnifiedExecProcessManager { request.tty, deferred_network_approval.clone(), Arc::clone(&transcript), + Arc::clone(&initial_exec_command_active), ) .await; - } + Some(InitialExecCommandGuard { + active: initial_exec_command_active, + }) + } else { + None + }; let yield_time_ms = clamp_yield_time(request.yield_time_ms); // For the initial exec_command call, we both stream output to events @@ -608,6 +628,7 @@ impl UnifiedExecProcessManager { pause_state, session, network_approval, + call_id, hook_command, process_id, tty, @@ -617,24 +638,29 @@ impl UnifiedExecProcessManager { if !request.input.is_empty() { if !tty { - return Err(UnifiedExecError::StdinClosed); - } - match process.write(request.input.as_bytes()).await { - Ok(()) => { - // Give the remote process a brief window to react so that we are - // more likely to capture its output in the poll below. - tokio::time::sleep(Duration::from_millis(100)).await; + if request.input == INTERRUPT { + process.interrupt().await?; + } else { + return Err(UnifiedExecError::StdinClosed); } - Err(err) => { - let status = self.refresh_process_state(process_id).await; - if matches!(status, ProcessStatus::Exited { .. }) { - status_after_write = Some(status); - } else if matches!(err, UnifiedExecError::ProcessFailed { .. }) { - process.terminate(); - self.release_process_id(process_id).await; - return Err(err); - } else { - return Err(err); + } else { + match process.write(request.input.as_bytes()).await { + Ok(()) => { + // Give the remote process a brief window to react so that we are + // more likely to capture its output in the poll below. + tokio::time::sleep(Duration::from_millis(100)).await; + } + Err(err) => { + let status = self.refresh_process_state(process_id).await; + if matches!(status, ProcessStatus::Exited { .. }) { + status_after_write = Some(status); + } else if matches!(err, UnifiedExecError::ProcessFailed { .. }) { + process.terminate(); + self.release_process_id(process_id).await; + return Err(err); + } else { + return Err(err); + } } } } @@ -715,9 +741,13 @@ impl UnifiedExecProcessManager { (None, exit_code, call_id) } ProcessStatus::Unknown => { - return Err(UnifiedExecError::UnknownProcessId { - process_id: request.process_id, - }); + if process.has_exited() { + (None, process.exit_code(), call_id) + } else { + return Err(UnifiedExecError::UnknownProcessId { + process_id: request.process_id, + }); + } } }; @@ -738,29 +768,27 @@ impl UnifiedExecProcessManager { } async fn refresh_process_state(&self, process_id: i32) -> ProcessStatus { - { - let mut store = self.process_store.lock().await; - let Some(entry) = store.processes.get(&process_id) else { + let mut store = self.process_store.lock().await; + let Some(entry) = store.processes.get_mut(&process_id) else { + return ProcessStatus::Unknown; + }; + + let exit_code = entry.process.exit_code(); + let process_id = entry.process_id; + + if entry.process.has_exited() { + let Some(entry) = store.remove(process_id) else { return ProcessStatus::Unknown; }; - - let exit_code = entry.process.exit_code(); - let process_id = entry.process_id; - - if entry.process.has_exited() { - let Some(entry) = store.remove(process_id) else { - return ProcessStatus::Unknown; - }; - ProcessStatus::Exited { - exit_code, - entry: Box::new(entry), - } - } else { - ProcessStatus::Alive { - exit_code, - call_id: entry.call_id.clone(), - process_id, - } + ProcessStatus::Exited { + exit_code, + entry: Box::new(entry), + } + } else { + ProcessStatus::Alive { + exit_code, + call_id: entry.call_id.clone(), + process_id, } } } @@ -798,6 +826,7 @@ impl UnifiedExecProcessManager { pause_state, session, network_approval: entry.network_approval.clone(), + call_id: entry.call_id.clone(), hook_command: entry.hook_command.clone(), process_id: entry.process_id, tty: entry.tty, @@ -817,11 +846,14 @@ impl UnifiedExecProcessManager { tty: bool, network_approval: Option, transcript: Arc>, + initial_exec_command_active: Arc, ) { let entry = ProcessEntry { process: Arc::clone(&process), call_id: context.call_id.clone(), process_id, + cwd: cwd.clone(), + initial_exec_command_active, hook_command, tty, network_approval, @@ -1268,6 +1300,59 @@ impl UnifiedExecProcessManager { entry.process.terminate(); } } + + pub(crate) async fn list_processes(&self) -> Vec { + let store = self.process_store.lock().await; + let mut entries = store + .processes + .values() + .filter(|entry| !entry.process.has_exited()) + .collect::>(); + entries.sort_by_key(|entry| entry.process_id); + entries + .into_iter() + .map(|entry| BackgroundTerminalInfo { + item_id: entry.call_id.clone(), + process_id: entry.process_id.to_string(), + command: entry.hook_command.clone(), + cwd: entry.cwd.clone(), + }) + .collect() + } + + pub(crate) async fn terminate_process(&self, process_id: i32) -> bool { + let (process, already_exited) = { + let store = self.process_store.lock().await; + let Some(entry) = store.processes.get(&process_id) else { + return false; + }; + (Arc::clone(&entry.process), entry.process.has_exited()) + }; + + if !already_exited && process.terminate_confirmed().await.is_err() { + return false; + } + + let entry = { + let mut store = self.process_store.lock().await; + let Some(entry) = store.processes.get(&process_id) else { + return true; + }; + if !Arc::ptr_eq(&entry.process, &process) { + return true; + } + if entry.initial_exec_command_active.load(Ordering::Acquire) { + return true; + } + let Some(entry) = store.remove(process_id) else { + return false; + }; + entry + }; + + unregister_network_approval_for_entry(&entry).await; + true + } } enum ProcessStatus { diff --git a/codex-rs/core/src/unified_exec/process_tests.rs b/codex-rs/core/src/unified_exec/process_tests.rs index ee0ea1cba3..814de95af4 100644 --- a/codex-rs/core/src/unified_exec/process_tests.rs +++ b/codex-rs/core/src/unified_exec/process_tests.rs @@ -5,6 +5,7 @@ use codex_exec_server::ExecProcess; use codex_exec_server::ExecProcessEventReceiver; use codex_exec_server::ExecServerError; use codex_exec_server::ProcessId; +use codex_exec_server::ProcessSignal; use codex_exec_server::ReadResponse; use codex_exec_server::StartedExecProcess; use codex_exec_server::WriteResponse; @@ -21,6 +22,7 @@ struct MockExecProcess { process_id: ProcessId, write_response: WriteResponse, read_responses: Mutex>, + terminate_error: Option, wake_tx: watch::Sender, } @@ -63,12 +65,22 @@ impl ExecProcess for MockExecProcess { Ok(self.write_response.clone()) } + async fn signal(&self, _signal: ProcessSignal) -> Result<(), ExecServerError> { + Ok(()) + } + async fn terminate(&self) -> Result<(), ExecServerError> { + if let Some(message) = &self.terminate_error { + return Err(ExecServerError::Protocol(message.clone())); + } Ok(()) } } -async fn remote_process(write_status: WriteStatus) -> UnifiedExecProcess { +async fn remote_process( + write_status: WriteStatus, + terminate_error: Option, +) -> UnifiedExecProcess { let (wake_tx, _wake_rx) = watch::channel(0); let started = StartedExecProcess { process: Arc::new(MockExecProcess { @@ -77,6 +89,7 @@ async fn remote_process(write_status: WriteStatus) -> UnifiedExecProcess { status: write_status, }, read_responses: Mutex::new(VecDeque::new()), + terminate_error, wake_tx, }), }; @@ -88,7 +101,7 @@ async fn remote_process(write_status: WriteStatus) -> UnifiedExecProcess { #[tokio::test] async fn remote_write_unknown_process_marks_process_exited() { - let process = remote_process(WriteStatus::UnknownProcess).await; + let process = remote_process(WriteStatus::UnknownProcess, /*terminate_error*/ None).await; let err = process .write(b"hello") @@ -101,7 +114,7 @@ async fn remote_write_unknown_process_marks_process_exited() { #[tokio::test] async fn remote_write_closed_stdin_marks_process_exited() { - let process = remote_process(WriteStatus::StdinClosed).await; + let process = remote_process(WriteStatus::StdinClosed, /*terminate_error*/ None).await; let err = process .write(b"hello") @@ -114,7 +127,7 @@ async fn remote_write_closed_stdin_marks_process_exited() { #[tokio::test] async fn fail_and_terminate_preserves_failure_message() { - let process = remote_process(WriteStatus::Accepted).await; + let process = remote_process(WriteStatus::Accepted, /*terminate_error*/ None).await; process.fail_and_terminate("network denied".to_string()); process.fail_and_terminate("second failure".to_string()); @@ -126,6 +139,32 @@ async fn fail_and_terminate_preserves_failure_message() { ); } +#[tokio::test] +async fn remote_terminate_confirmed_updates_state_on_success_only() { + let process = remote_process( + WriteStatus::Accepted, + Some("terminate unavailable".to_string()), + ) + .await; + + let err = process + .terminate_confirmed() + .await + .expect_err("expected terminate failure"); + + assert!(matches!(err, UnifiedExecError::ProcessFailed { .. })); + assert!(!process.has_exited()); + + let process = remote_process(WriteStatus::Accepted, /*terminate_error*/ None).await; + + process + .terminate_confirmed() + .await + .expect("terminate should succeed"); + + assert!(process.has_exited()); +} + #[tokio::test] async fn remote_process_waits_for_early_exit_event() { let (wake_tx, _wake_rx) = watch::channel(0); @@ -143,6 +182,7 @@ async fn remote_process_waits_for_early_exit_event() { closed: true, failure: None, }])), + terminate_error: None, wake_tx: wake_tx.clone(), }), }; diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index 1ba93aec10..af87762c04 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -85,6 +85,7 @@ async fn responses_stream_includes_subagent_header_on_review() { let session_source = SessionSource::SubAgent(SubAgentSource::Review); let model_info = codex_core::test_support::construct_model_info_offline(model.as_str(), &config); + let expected_window_id = format!("{thread_id}:0"); let session_telemetry = SessionTelemetry::new( thread_id, model.as_str(), @@ -126,6 +127,7 @@ async fn responses_stream_includes_subagent_header_on_review() { let mut stream = client_session .stream( + &expected_window_id, &prompt, &model_info, &session_telemetry, @@ -144,7 +146,6 @@ async fn responses_stream_includes_subagent_header_on_review() { } let request = request_recorder.single_request(); - let expected_window_id = format!("{thread_id}:0"); assert_eq!( request.header("x-openai-subagent").as_deref(), Some("review") @@ -217,6 +218,7 @@ async fn responses_stream_includes_subagent_header_on_other() { let session_source = SessionSource::SubAgent(SubAgentSource::Other("my-task".to_string())); let model_info = codex_core::test_support::construct_model_info_offline(model.as_str(), &config); + let window_id = format!("{thread_id}:0"); let session_telemetry = SessionTelemetry::new( thread_id, @@ -259,6 +261,7 @@ async fn responses_stream_includes_subagent_header_on_other() { let mut stream = client_session .stream( + &window_id, &prompt, &model_info, &session_telemetry, @@ -336,6 +339,7 @@ async fn responses_respects_model_info_overrides_from_config() { SessionSource::SubAgent(SubAgentSource::Other("override-check".to_string())); let model_info = codex_core::test_support::construct_model_info_offline(model.as_str(), &config); + let window_id = format!("{thread_id}:0"); let session_telemetry = SessionTelemetry::new( thread_id, model.as_str(), @@ -377,6 +381,7 @@ async fn responses_respects_model_info_overrides_from_config() { let mut stream = client_session .stream( + &window_id, &prompt, &model_info, &session_telemetry, diff --git a/codex-rs/core/tests/suite/auto_review.rs b/codex-rs/core/tests/suite/auto_review.rs index 5fe1cbc2f2..7dae7c4c62 100644 --- a/codex-rs/core/tests/suite/auto_review.rs +++ b/codex-rs/core/tests/suite/auto_review.rs @@ -256,6 +256,7 @@ fn remote_model_with_auto_review_override(slug: &str, review_model: &str) -> Mod context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), } diff --git a/codex-rs/core/tests/suite/cli_stream.rs b/codex-rs/core/tests/suite/cli_stream.rs index fc9b7d2569..1ef7c82c39 100644 --- a/codex-rs/core/tests/suite/cli_stream.rs +++ b/codex-rs/core/tests/suite/cli_stream.rs @@ -1,4 +1,3 @@ -use assert_cmd::Command as AssertCommand; use codex_git_utils::collect_git_info; use codex_login::CODEX_ACCESS_TOKEN_ENV_VAR; use codex_login::CODEX_API_KEY_ENV_VAR; @@ -7,6 +6,14 @@ use core_test_support::fs_wait; use core_test_support::responses; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; +use std::io; +#[cfg(unix)] +use std::os::unix::process::CommandExt as _; +use std::process::Command; +use std::process::Output; +use std::process::Stdio; +use std::sync::mpsc; +use std::thread; use std::time::Duration; use tempfile::TempDir; use uuid::Uuid; @@ -22,6 +29,7 @@ const PERSONAL_ACCESS_TOKEN_AUTHORIZATION: &str = "Bearer at-cli-test"; const PERSONAL_ACCESS_TOKEN_ACCOUNT_ID: &str = "account-pat"; const WHOAMI_PATH: &str = "/v1/user-auth-credential/whoami"; const CLOUD_CONFIG_BUNDLE_PATH: &str = "/backend-api/wham/config/bundle"; +const CLI_TIMEOUT: Duration = Duration::from_secs(30); fn repo_root() -> std::path::PathBuf { #[expect(clippy::expect_used)] @@ -60,10 +68,9 @@ async fn mount_personal_access_token_startup(server: &MockServer) { } #[expect(clippy::unwrap_used)] -fn personal_access_token_exec_command(server: &MockServer, home: &TempDir) -> AssertCommand { +fn personal_access_token_exec_command(server: &MockServer, home: &TempDir) -> Command { let bin = codex_utils_cargo_bin::cargo_bin("codex").unwrap(); - let mut cmd = AssertCommand::new(bin); - cmd.timeout(Duration::from_secs(30)); + let mut cmd = Command::new(bin); cmd.arg("exec") .arg("--skip-git-repo-check") .arg("-c") @@ -81,6 +88,62 @@ fn personal_access_token_exec_command(server: &MockServer, home: &TempDir) -> As cmd } +struct ChildProcessCleanupGuard(u32); + +impl Drop for ChildProcessCleanupGuard { + fn drop(&mut self) { + #[cfg(unix)] + { + let _ = codex_utils_pty::process_group::kill_process_group(self.0); + } + + #[cfg(windows)] + { + let _ = Command::new("taskkill") + .args(["/PID", &self.0.to_string(), "/T", "/F"]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + } + + #[cfg(not(any(unix, windows)))] + { + let _ = self.0; + } + } +} + +// Use this for new `codex exec` subprocess tests in this file. These commands +// can spawn shell/Python grandchildren, so the timeout path must reap the whole +// process group instead of only the direct CLI child. +fn run_cli_command(command: &mut Command) -> io::Result { + #[cfg(unix)] + command.process_group(0); + + command + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + + let child = command.spawn()?; + let _cleanup = ChildProcessCleanupGuard(child.id()); + let (sender, receiver) = mpsc::sync_channel(1); + let _waiter = thread::spawn(move || { + let _ = sender.send(child.wait_with_output()); + }); + + match receiver.recv_timeout(CLI_TIMEOUT) { + Ok(output) => output, + Err(mpsc::RecvTimeoutError::Timeout) => { + Err(io::Error::new(io::ErrorKind::TimedOut, "process timed out")) + } + Err(mpsc::RecvTimeoutError::Disconnected) => { + Err(io::Error::other("process output reader thread exited")) + } + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_mode_stream_cli_supports_personal_access_tokens() { skip_if_no_network!(); @@ -90,9 +153,8 @@ async fn responses_mode_stream_cli_supports_personal_access_tokens() { let resp_mock = responses::mount_sse_once(&server, cli_sse_response()).await; let home = TempDir::new().unwrap(); - let output = personal_access_token_exec_command(&server, &home) - .output() - .unwrap(); + let mut cmd = personal_access_token_exec_command(&server, &home); + let output = run_cli_command(&mut cmd).unwrap(); assert!( output.status.success(), @@ -140,9 +202,8 @@ async fn responses_mode_stream_cli_does_not_attempt_oauth_refresh_for_personal_a .await; let home = TempDir::new().unwrap(); - let output = personal_access_token_exec_command(&server, &home) - .output() - .unwrap(); + let mut cmd = personal_access_token_exec_command(&server, &home); + let output = run_cli_command(&mut cmd).unwrap(); assert!(!output.status.success()); server.verify().await; @@ -168,8 +229,7 @@ async fn responses_mode_stream_cli() { server.uri() ); let bin = codex_utils_cargo_bin::cargo_bin("codex").unwrap(); - let mut cmd = AssertCommand::new(bin); - cmd.timeout(Duration::from_secs(30)); + let mut cmd = Command::new(bin); cmd.arg("exec") .arg("--skip-git-repo-check") .arg("-c") @@ -182,7 +242,7 @@ async fn responses_mode_stream_cli() { cmd.env("CODEX_HOME", home.path()) .env("OPENAI_API_KEY", "dummy"); - let output = cmd.output().unwrap(); + let output = run_cli_command(&mut cmd).unwrap(); println!("Status: {}", output.status); println!("Stdout:\n{}", String::from_utf8_lossy(&output.stdout)); println!("Stderr:\n{}", String::from_utf8_lossy(&output.stderr)); @@ -211,8 +271,7 @@ async fn responses_mode_stream_cli_supports_openai_base_url_config_override() { let home = TempDir::new().unwrap(); let bin = codex_utils_cargo_bin::cargo_bin("codex").unwrap(); - let mut cmd = AssertCommand::new(bin); - cmd.timeout(Duration::from_secs(30)); + let mut cmd = Command::new(bin); cmd.arg("exec") .arg("--skip-git-repo-check") .arg("-c") @@ -223,7 +282,7 @@ async fn responses_mode_stream_cli_supports_openai_base_url_config_override() { cmd.env("CODEX_HOME", home.path()) .env("OPENAI_API_KEY", "dummy"); - let output = cmd.output().unwrap(); + let output = run_cli_command(&mut cmd).unwrap(); assert!(output.status.success()); let request = resp_mock.single_request(); @@ -264,7 +323,7 @@ async fn exec_cli_applies_model_instructions_file() { let home = TempDir::new().unwrap(); let repo_root = repo_root(); let bin = codex_utils_cargo_bin::cargo_bin("codex").unwrap(); - let mut cmd = AssertCommand::new(bin); + let mut cmd = Command::new(bin); cmd.arg("exec") .arg("--skip-git-repo-check") .arg("-c") @@ -279,7 +338,7 @@ async fn exec_cli_applies_model_instructions_file() { cmd.env("CODEX_HOME", home.path()) .env("OPENAI_API_KEY", "dummy"); - let output = cmd.output().unwrap(); + let output = run_cli_command(&mut cmd).unwrap(); println!("Status: {}", output.status); println!("Stdout:\n{}", String::from_utf8_lossy(&output.stdout)); println!("Stderr:\n{}", String::from_utf8_lossy(&output.stderr)); @@ -334,7 +393,7 @@ async fn exec_cli_profile_applies_model_instructions_file() { let repo_root = repo_root(); let bin = codex_utils_cargo_bin::cargo_bin("codex").unwrap(); - let mut cmd = AssertCommand::new(bin); + let mut cmd = Command::new(bin); cmd.arg("exec") .arg("--skip-git-repo-check") .arg("--profile") @@ -349,7 +408,7 @@ async fn exec_cli_profile_applies_model_instructions_file() { cmd.env("CODEX_HOME", home.path()) .env("OPENAI_API_KEY", "dummy"); - let output = cmd.output().unwrap(); + let output = run_cli_command(&mut cmd).unwrap(); println!("Status: {}", output.status); println!("Stdout:\n{}", String::from_utf8_lossy(&output.stdout)); println!("Stderr:\n{}", String::from_utf8_lossy(&output.stderr)); @@ -379,8 +438,7 @@ async fn responses_api_stream_cli() { let home = TempDir::new().unwrap(); let bin = codex_utils_cargo_bin::cargo_bin("codex").unwrap(); - let mut cmd = AssertCommand::new(bin); - cmd.timeout(Duration::from_secs(30)); + let mut cmd = Command::new(bin); cmd.arg("exec") .arg("--skip-git-repo-check") .arg("-c") @@ -391,7 +449,7 @@ async fn responses_api_stream_cli() { cmd.env("CODEX_HOME", home.path()) .env("OPENAI_API_KEY", "dummy"); - let output = cmd.output().unwrap(); + let output = run_cli_command(&mut cmd).unwrap(); assert!(output.status.success()); let stdout = String::from_utf8_lossy(&output.stdout); assert!(stdout.contains("fixture hello")); @@ -421,8 +479,7 @@ async fn integration_creates_and_checks_session_file() -> anyhow::Result<()> { // 4. Run the codex CLI and invoke `exec`, which is what records a session. let bin = codex_utils_cargo_bin::cargo_bin("codex").unwrap(); - let mut cmd = AssertCommand::new(bin); - cmd.timeout(Duration::from_secs(30)); + let mut cmd = Command::new(bin); cmd.arg("exec") .arg("--skip-git-repo-check") .arg("-c") @@ -433,7 +490,7 @@ async fn integration_creates_and_checks_session_file() -> anyhow::Result<()> { cmd.env("CODEX_HOME", home.path()) .env(CODEX_API_KEY_ENV_VAR, "dummy"); - let output = cmd.output().unwrap(); + let output = run_cli_command(&mut cmd).unwrap(); assert!( output.status.success(), "codex-cli exec failed: {}", @@ -542,8 +599,7 @@ async fn integration_creates_and_checks_session_file() -> anyhow::Result<()> { let marker2 = format!("integration-resume-{}", Uuid::new_v4()); let prompt2 = format!("echo {marker2}"); let bin2 = codex_utils_cargo_bin::cargo_bin("codex").unwrap(); - let mut cmd2 = AssertCommand::new(bin2); - cmd2.timeout(Duration::from_secs(30)); + let mut cmd2 = Command::new(bin2); cmd2.arg("exec") .arg("--skip-git-repo-check") .arg("-c") @@ -556,7 +612,7 @@ async fn integration_creates_and_checks_session_file() -> anyhow::Result<()> { cmd2.env("CODEX_HOME", home.path()) .env("OPENAI_API_KEY", "dummy"); - let output2 = cmd2.output().unwrap(); + let output2 = run_cli_command(&mut cmd2).unwrap(); assert!(output2.status.success(), "resume codex-cli run failed"); assert_eq!(resp_mock.requests().len(), 2); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 9a6f104d0c..a522e041c2 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -90,6 +90,7 @@ use wiremock::matchers::path; use wiremock::matchers::query_param; const INSTALLATION_ID_FILENAME: &str = "installation_id"; +const TEST_WINDOW_ID: &str = "test-thread:0"; #[expect(clippy::unwrap_used)] fn assert_message_role(request_body: &serde_json::Value, role: &str) { @@ -924,6 +925,7 @@ async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuth let mut stream = client_session .stream( + TEST_WINDOW_ID, &prompt, &model_info, &session_telemetry, @@ -2467,6 +2469,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { let mut stream = client_session .stream( + TEST_WINDOW_ID, &prompt, &model_info, &session_telemetry, diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 4dc674ec30..8b135b6b75 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -65,6 +65,7 @@ const X_CLIENT_REQUEST_ID_HEADER: &str = "x-client-request-id"; const WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY: &str = "ws_request_header_x_openai_internal_codex_responses_lite"; const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; +const TEST_WINDOW_ID: &str = "test-thread:0"; const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str = "x-codex-ws-stream-request-start-ms"; @@ -274,7 +275,7 @@ async fn responses_websocket_preconnect_does_not_replace_turn_trace_payload() { let harness = websocket_harness(&server).await; let mut client_session = harness.client.new_session(); client_session - .preconnect_websocket(&harness.session_telemetry, &harness.model_info) + .preconnect_websocket(&harness.session_telemetry) .await .expect("websocket preconnect failed"); let prompt = prompt_with_input(vec![message_item("hello")]); @@ -310,7 +311,7 @@ async fn responses_websocket_preconnect_reuses_connection() { let harness = websocket_harness(&server).await; let mut client_session = harness.client.new_session(); client_session - .preconnect_websocket(&harness.session_telemetry, &harness.model_info) + .preconnect_websocket(&harness.session_telemetry) .await .expect("websocket preconnect failed"); let prompt = prompt_with_input(vec![message_item("hello")]); @@ -321,6 +322,7 @@ async fn responses_websocket_preconnect_reuses_connection() { server.single_handshake().header(USER_AGENT_HEADER), Some(codex_login::default_client::get_codex_user_agent()) ); + assert_eq!(server.single_handshake().header("x-codex-window-id"), None); let connection = server.single_connection(); assert_eq!(connection.len(), 1); @@ -342,6 +344,7 @@ async fn responses_websocket_request_prewarm_reuses_connection() { let prompt = prompt_with_input(vec![message_item("hello")]); client_session .prewarm_websocket( + TEST_WINDOW_ID, &prompt, &harness.model_info, &harness.session_telemetry, @@ -396,6 +399,7 @@ async fn responses_websocket_request_prewarm_traces_logical_request() { client_session .prewarm_websocket( + TEST_WINDOW_ID, &prompt, &harness.model_info, &harness.session_telemetry, @@ -441,6 +445,7 @@ async fn responses_websocket_request_prewarm_traces_logical_request() { let mut stream = client_session .stream( + TEST_WINDOW_ID, &prompt, &harness.model_info, &harness.session_telemetry, @@ -606,12 +611,13 @@ async fn responses_websocket_preconnect_is_reused_even_with_header_changes() { let harness = websocket_harness(&server).await; let mut client_session = harness.client.new_session(); client_session - .preconnect_websocket(&harness.session_telemetry, &harness.model_info) + .preconnect_websocket(&harness.session_telemetry) .await .expect("websocket preconnect failed"); let prompt = prompt_with_input(vec![message_item("hello")]); let mut stream = client_session .stream( + TEST_WINDOW_ID, &prompt, &harness.model_info, &harness.session_telemetry, @@ -651,6 +657,7 @@ async fn responses_websocket_request_prewarm_is_reused_even_with_header_changes( let prompt = prompt_with_input(vec![message_item("hello")]); client_session .prewarm_websocket( + TEST_WINDOW_ID, &prompt, &harness.model_info, &harness.session_telemetry, @@ -663,6 +670,7 @@ async fn responses_websocket_request_prewarm_is_reused_even_with_header_changes( .expect("websocket prewarm failed"); let mut stream = client_session .stream( + TEST_WINDOW_ID, &prompt, &harness.model_info, &harness.session_telemetry, @@ -717,6 +725,7 @@ async fn responses_websocket_prewarm_uses_v2_when_provider_supports_websockets() let prompt = prompt_with_input(vec![message_item("hello")]); client_session .prewarm_websocket( + TEST_WINDOW_ID, &prompt, &harness.model_info, &harness.session_telemetry, @@ -772,7 +781,7 @@ async fn responses_websocket_preconnect_runs_when_only_v2_feature_enabled() { let harness = websocket_harness_with_options(&server, /*runtime_metrics_enabled*/ true).await; let mut client_session = harness.client.new_session(); client_session - .preconnect_websocket(&harness.session_telemetry, &harness.model_info) + .preconnect_websocket(&harness.session_telemetry) .await .expect("websocket preconnect failed"); @@ -1066,6 +1075,7 @@ async fn responses_websocket_emits_reasoning_included_event() { let mut stream = client_session .stream( + TEST_WINDOW_ID, &prompt, &harness.model_info, &harness.session_telemetry, @@ -1140,6 +1150,7 @@ async fn responses_websocket_emits_rate_limit_events() { let mut stream = client_session .stream( + TEST_WINDOW_ID, &prompt, &harness.model_info, &harness.session_telemetry, @@ -1794,6 +1805,7 @@ async fn responses_websocket_v2_after_error_uses_full_create_without_previous_re let mut second_stream = session .stream( + TEST_WINDOW_ID, &prompt_two, &harness.model_info, &harness.session_telemetry, @@ -1882,6 +1894,7 @@ async fn responses_websocket_v2_surfaces_terminal_error_without_close_handshake( let mut second_stream = session .stream( + TEST_WINDOW_ID, &prompt_two, &harness.model_info, &harness.session_telemetry, @@ -2116,6 +2129,7 @@ async fn stream_until_complete_with_model_info( ) { let mut stream = client_session .stream( + TEST_WINDOW_ID, prompt, model_info, &harness.session_telemetry, @@ -2183,6 +2197,7 @@ async fn stream_until_complete_with_request_metadata( ) { let mut stream = client_session .stream( + TEST_WINDOW_ID, prompt, &harness.model_info, &harness.session_telemetry, diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index e1b02292ef..fdd87274a7 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -42,11 +42,16 @@ use core_test_support::test_codex::turn_permission_fields; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use core_test_support::wait_for_mcp_server; +use image::DynamicImage; +use image::GenericImageView; +use image::ImageBuffer; +use image::Rgba; use pretty_assertions::assert_eq; use serde_json::Value; use std::collections::HashMap; use std::collections::HashSet; use std::fs; +use std::io::Cursor; use std::path::Path; use std::sync::Arc; use std::time::Duration; @@ -168,12 +173,21 @@ async fn run_code_mode_turn_with_config( code: &str, configure: impl FnOnce(&mut Config) + Send + 'static, ) -> Result<(TestCodex, ResponseMock)> { - let mut builder = test_codex() - .with_model("test-gpt-5.1-codex") - .with_config(move |config| { - let _ = config.features.enable(Feature::CodeMode); - configure(config); - }); + run_code_mode_turn_with_model_and_config(server, prompt, code, "test-gpt-5.1-codex", configure) + .await +} + +async fn run_code_mode_turn_with_model_and_config( + server: &MockServer, + prompt: &str, + code: &str, + model: &'static str, + configure: impl FnOnce(&mut Config) + Send + 'static, +) -> Result<(TestCodex, ResponseMock)> { + let mut builder = test_codex().with_model(model).with_config(move |config| { + let _ = config.features.enable(Feature::CodeMode); + configure(config); + }); let test = builder.build(server).await?; responses::mount_sse_once( @@ -2536,6 +2550,103 @@ image("data:image/png;base64,AAA"); Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn resize_all_images_replaces_malformed_code_mode_image_only() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (_test, second_mock) = run_code_mode_turn_with_config( + &server, + "use exec to return images", + r#" +image("https://example.com/image.jpg"); +image("data:image/png;base64,AAA"); +"#, + |config| { + let _ = config.features.enable(Feature::ResizeAllImages); + }, + ) + .await?; + + let req = second_mock.single_request(); + let items = custom_tool_output_items(&req, "call-1"); + let (_, success) = custom_tool_output_body_and_success(&req, "call-1"); + assert_ne!(success, Some(false)); + assert_eq!(items.len(), 3); + assert_eq!( + items[1], + serde_json::json!({ + "type": "input_image", + "image_url": "https://example.com/image.jpg", + "detail": "high" + }) + ); + assert_eq!( + items[2], + serde_json::json!({ + "type": "input_text", + "text": "image content omitted because it could not be processed" + }) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn resize_all_images_resizes_explicit_original_code_mode_image() -> Result<()> { + skip_if_no_network!(Ok(())); + + let original_dimensions = (6401, 100); + let image = ImageBuffer::from_pixel( + original_dimensions.0, + original_dimensions.1, + Rgba([20, 40, 60, 255]), + ); + let mut encoded = Cursor::new(Vec::new()); + DynamicImage::ImageRgba8(image).write_to(&mut encoded, image::ImageFormat::Png)?; + let image_data_url = format!( + "data:image/png;base64,{}", + BASE64_STANDARD.encode(encoded.into_inner()) + ); + let code = format!( + "image({}, \"original\");", + serde_json::to_string(&image_data_url)? + ); + + let server = responses::start_mock_server().await; + let (_test, second_mock) = run_code_mode_turn_with_model_and_config( + &server, + "use exec to return a large original-detail image", + &code, + "gpt-5.3-codex", + |config| { + config + .features + .enable(Feature::ResizeAllImages) + .expect("resize_all_images should be enabled"); + }, + ) + .await?; + + let req = second_mock.single_request(); + let items = custom_tool_output_items(&req, "call-1"); + let (_, success) = custom_tool_output_body_and_success(&req, "call-1"); + assert_ne!(success, Some(false)); + let resized_url = items[1]["image_url"] + .as_str() + .expect("code mode image output should contain a data URL"); + assert_eq!(items[1]["detail"], "original"); + let (_, resized_base64) = resized_url + .split_once(',') + .expect("resized image should contain a data URL prefix"); + let resized_bytes = BASE64_STANDARD.decode(resized_base64)?; + let resized = image::load_from_memory(&resized_bytes)?; + let resized_dimensions = resized.dimensions(); + assert_eq!(resized_dimensions, (6000, 94)); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_can_use_view_image_result_with_image_helper() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index 339cadda4e..73bc05df43 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -376,6 +376,12 @@ fn model_info_with_context_window(slug: &str, context_window: i64) -> ModelInfo model_info } +fn model_info_with_optional_comp_hash(slug: &str, comp_hash: Option<&str>) -> ModelInfo { + let mut model_info = model_info_with_context_window(slug, /*context_window*/ 273_000); + model_info.comp_hash = comp_hash.map(str::to_string); + model_info +} + fn assert_pre_sampling_switch_compaction_requests( first: &serde_json::Value, compact: &serde_json::Value, @@ -2009,6 +2015,7 @@ async fn auto_compact_runs_after_resume_when_token_usage_is_over_limit() { let mut builder = test_codex().with_config(move |config| { set_test_compact_prompt(config); config.model_auto_compact_token_limit = Some(limit); + let _ = config.features.disable(Feature::RemoteCompactionV2); }); let initial = builder.build(&server).await.unwrap(); let home = initial.home.clone(); @@ -2037,6 +2044,7 @@ async fn auto_compact_runs_after_resume_when_token_usage_is_over_limit() { let mut resume_builder = test_codex().with_config(move |config| { set_test_compact_prompt(config); config.model_auto_compact_token_limit = Some(limit); + let _ = config.features.disable(Feature::RemoteCompactionV2); }); let resumed = resume_builder .resume(&server, home, rollout_path) @@ -2189,6 +2197,203 @@ async fn pre_sampling_compact_runs_on_switch_to_smaller_context_model() { ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pre_sampling_compact_runs_when_comp_hash_changes() { + skip_if_no_network!(); + + let server = MockServer::start().await; + let previous_model = "gpt-5.3-codex"; + let next_model = "gpt-5.2"; + + let models_mock = mount_models_once( + &server, + ModelsResponse { + models: vec![ + model_info_with_optional_comp_hash(previous_model, Some("hash-a")), + model_info_with_optional_comp_hash(next_model, Some("hash-b")), + ], + }, + ) + .await; + + let request_log = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_assistant_message("m1", "before switch"), + ev_completed_with_tokens("r1", /*total_tokens*/ 100), + ]), + sse(vec![ + ev_assistant_message("m2", "COMP_HASH_SUMMARY"), + ev_completed_with_tokens("r2", /*total_tokens*/ 10), + ]), + sse(vec![ + ev_assistant_message("m3", "after switch"), + ev_completed_with_tokens("r3", /*total_tokens*/ 100), + ]), + ], + ) + .await; + + let model_provider = non_openai_model_provider(&server); + let mut builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model(previous_model) + .with_config(move |config| { + config.model_provider = model_provider; + set_test_compact_prompt(config); + }); + let test = builder.build(&server).await.expect("build test codex"); + + test.codex + .submit(disabled_permission_user_turn( + "before switch", + test.cwd.path().to_path_buf(), + previous_model.to_string(), + )) + .await + .expect("submit first user turn"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + test.codex + .submit(disabled_permission_user_turn( + "after switch", + test.cwd.path().to_path_buf(), + next_model.to_string(), + )) + .await + .expect("submit second user turn"); + assert_compaction_uses_turn_lifecycle_id(&test.codex).await; + + let requests = request_log.requests(); + assert_eq!(models_mock.requests().len(), 1); + assert_eq!( + requests.len(), + 3, + "a comp-hash change should compact before sampling the next turn" + ); + assert_pre_sampling_switch_compaction_requests( + &requests[0].body_json(), + &requests[1].body_json(), + &requests[2].body_json(), + previous_model, + next_model, + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pre_sampling_compact_skips_when_either_comp_hash_is_missing() { + skip_if_no_network!(); + + let server = MockServer::start().await; + let model_without_hash = "gpt-5.4"; + let model_with_hash = "gpt-5.3-codex"; + let next_model_without_hash = "gpt-5.2"; + + let models_mock = mount_models_once( + &server, + ModelsResponse { + models: vec![ + model_info_with_optional_comp_hash(model_without_hash, /*comp_hash*/ None), + model_info_with_optional_comp_hash(model_with_hash, Some("hash-a")), + model_info_with_optional_comp_hash( + next_model_without_hash, + /*comp_hash*/ None, + ), + ], + }, + ) + .await; + + let request_log = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_assistant_message("m1", "before hash"), + ev_completed_with_tokens("r1", /*total_tokens*/ 100), + ]), + sse(vec![ + ev_assistant_message("m2", "hash introduced"), + ev_completed_with_tokens("r2", /*total_tokens*/ 100), + ]), + sse(vec![ + ev_assistant_message("m3", "hash removed"), + ev_completed_with_tokens("r3", /*total_tokens*/ 100), + ]), + ], + ) + .await; + + let model_provider = non_openai_model_provider(&server); + let mut builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model(model_without_hash) + .with_config(move |config| { + config.model_provider = model_provider; + set_test_compact_prompt(config); + }); + let test = builder.build(&server).await.expect("build test codex"); + + test.codex + .submit(disabled_permission_user_turn( + "before hash", + test.cwd.path().to_path_buf(), + model_without_hash.to_string(), + )) + .await + .expect("submit first user turn"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + test.codex + .submit(disabled_permission_user_turn( + "hash introduced", + test.cwd.path().to_path_buf(), + model_with_hash.to_string(), + )) + .await + .expect("submit second user turn"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + test.codex + .submit(disabled_permission_user_turn( + "hash removed", + test.cwd.path().to_path_buf(), + next_model_without_hash.to_string(), + )) + .await + .expect("submit third user turn"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = request_log.requests(); + assert_eq!(models_mock.requests().len(), 1); + assert_eq!( + requests + .iter() + .map(|request| request.body_json()["model"].as_str().map(str::to_string)) + .collect::>(), + vec![ + Some(model_without_hash.to_string()), + Some(model_with_hash.to_string()), + Some(next_model_without_hash.to_string()), + ] + ); + assert!(requests.iter().all(|request| { + !body_contains_text(&request.body_json().to_string(), SUMMARIZATION_PROMPT) + })); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn body_after_prefix_model_switch_budget_compacts_with_next_model() { skip_if_no_network!(); @@ -2404,6 +2609,266 @@ async fn pre_sampling_compact_runs_after_resume_and_switch_to_smaller_model() { ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pre_sampling_compact_recovers_comp_hash_after_resume() { + skip_if_no_network!(); + + let server = MockServer::start().await; + let previous_model = "gpt-5.3-codex"; + let next_model = "gpt-5.2"; + + let models_mock = mount_models_once( + &server, + ModelsResponse { + models: vec![ + model_info_with_optional_comp_hash(previous_model, Some("hash-a")), + model_info_with_optional_comp_hash(next_model, Some("hash-b")), + ], + }, + ) + .await; + + let request_log = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_assistant_message("m1", "before resume"), + ev_completed_with_tokens("r1", /*total_tokens*/ 100), + ]), + sse(vec![ + ev_assistant_message("m2", "RESUMED_COMP_HASH_SUMMARY"), + ev_completed_with_tokens("r2", /*total_tokens*/ 10), + ]), + sse(vec![ + ev_assistant_message("m3", "after resume"), + ev_completed_with_tokens("r3", /*total_tokens*/ 100), + ]), + ], + ) + .await; + + let model_provider = non_openai_model_provider(&server); + let mut initial_builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model(previous_model) + .with_config(move |config| { + config.model_provider = model_provider; + set_test_compact_prompt(config); + }); + let initial = initial_builder + .build(&server) + .await + .expect("build initial test codex"); + let home = initial.home.clone(); + let rollout_path = initial + .session_configured + .rollout_path + .clone() + .expect("rollout path"); + + initial + .codex + .submit(disabled_permission_user_turn( + "before resume", + initial.cwd.path().to_path_buf(), + previous_model.to_string(), + )) + .await + .expect("submit pre-resume turn"); + wait_for_event(&initial.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + initial + .codex + .submit(Op::Shutdown) + .await + .expect("shutdown initial session"); + wait_for_event(&initial.codex, |event| { + matches!(event, EventMsg::ShutdownComplete) + }) + .await; + + let rollout = fs::read_to_string(&rollout_path).expect("read rollout"); + let persisted_comp_hash = rollout + .lines() + .filter_map(|line| serde_json::from_str::(line).ok()) + .find_map(|line| match line.item { + RolloutItem::TurnContext(context) => context.comp_hash, + _ => None, + }); + assert_eq!(persisted_comp_hash.as_deref(), Some("hash-a")); + + let model_provider = non_openai_model_provider(&server); + let mut resumed_builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model(previous_model) + .with_config(move |config| { + config.model_provider = model_provider; + set_test_compact_prompt(config); + }); + let resumed = resumed_builder + .resume(&server, home, rollout_path) + .await + .expect("resume codex"); + + resumed + .codex + .submit(disabled_permission_user_turn( + "after resume", + resumed.cwd.path().to_path_buf(), + next_model.to_string(), + )) + .await + .expect("submit resumed user turn"); + assert_compaction_uses_turn_lifecycle_id(&resumed.codex).await; + + let requests = request_log.requests(); + assert_eq!(models_mock.requests().len(), 1); + assert_eq!( + requests.len(), + 3, + "the resumed turn should compact using the comp hash recovered from rollout" + ); + assert_pre_sampling_switch_compaction_requests( + &requests[0].body_json(), + &requests[1].body_json(), + &requests[2].body_json(), + previous_model, + next_model, + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pre_sampling_compact_skips_missing_comp_hash_after_resume() { + skip_if_no_network!(); + + let server = MockServer::start().await; + let previous_model = "gpt-5.3-codex"; + let next_model = "gpt-5.2"; + + let models_mock = mount_models_once( + &server, + ModelsResponse { + models: vec![ + model_info_with_optional_comp_hash(previous_model, /*comp_hash*/ None), + model_info_with_optional_comp_hash(next_model, Some("hash-b")), + ], + }, + ) + .await; + + let request_log = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_assistant_message("m1", "before resume"), + ev_completed_with_tokens("r1", /*total_tokens*/ 100), + ]), + sse(vec![ + ev_assistant_message("m2", "after resume"), + ev_completed_with_tokens("r2", /*total_tokens*/ 100), + ]), + ], + ) + .await; + + let model_provider = non_openai_model_provider(&server); + let mut initial_builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model(previous_model) + .with_config(move |config| { + config.model_provider = model_provider; + set_test_compact_prompt(config); + }); + let initial = initial_builder + .build(&server) + .await + .expect("build initial test codex"); + let home = initial.home.clone(); + let rollout_path = initial + .session_configured + .rollout_path + .clone() + .expect("rollout path"); + + initial + .codex + .submit(disabled_permission_user_turn( + "before resume", + initial.cwd.path().to_path_buf(), + previous_model.to_string(), + )) + .await + .expect("submit pre-resume turn"); + wait_for_event(&initial.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + initial + .codex + .submit(Op::Shutdown) + .await + .expect("shutdown initial session"); + wait_for_event(&initial.codex, |event| { + matches!(event, EventMsg::ShutdownComplete) + }) + .await; + + let rollout = fs::read_to_string(&rollout_path).expect("read rollout"); + let persisted_turn_context = rollout + .lines() + .filter_map(|line| serde_json::from_str::(line).ok()) + .find(|line| line["type"] == "turn_context") + .expect("persisted turn context"); + assert!(persisted_turn_context["payload"].get("comp_hash").is_none()); + + let model_provider = non_openai_model_provider(&server); + let mut resumed_builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model(previous_model) + .with_config(move |config| { + config.model_provider = model_provider; + set_test_compact_prompt(config); + }); + let resumed = resumed_builder + .resume(&server, home, rollout_path) + .await + .expect("resume codex"); + + resumed + .codex + .submit(disabled_permission_user_turn( + "after resume", + resumed.cwd.path().to_path_buf(), + next_model.to_string(), + )) + .await + .expect("submit resumed user turn"); + wait_for_event(&resumed.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = request_log.requests(); + assert_eq!(models_mock.requests().len(), 1); + assert_eq!( + requests + .iter() + .map(|request| request.body_json()["model"].as_str().map(str::to_string)) + .collect::>(), + vec![ + Some(previous_model.to_string()), + Some(next_model.to_string()), + ] + ); + assert!(requests.iter().all(|request| { + !body_contains_text(&request.body_json().to_string(), SUMMARIZATION_PROMPT) + })); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn auto_compact_persists_rollout_entries() { skip_if_no_network!(); @@ -3560,6 +4025,7 @@ async fn auto_compact_counts_encrypted_reasoning_before_last_user() { config.chatgpt_base_url = chatgpt_base_url; set_test_compact_prompt(config); config.model_auto_compact_token_limit = Some(300); + let _ = config.features.disable(Feature::RemoteCompactionV2); }) .build(&server) .await @@ -3682,6 +4148,7 @@ async fn auto_compact_runs_when_reasoning_header_clears_between_turns() { .with_config(|config| { set_test_compact_prompt(config); config.model_auto_compact_token_limit = Some(300); + let _ = config.features.disable(Feature::RemoteCompactionV2); }) .build(&server) .await diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 73a500f04d..2a1d459f9c 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -36,7 +36,7 @@ use core_test_support::responses::start_websocket_server; use core_test_support::skip_if_no_network; use core_test_support::test_codex::TestCodexBuilder; use core_test_support::test_codex::TestCodexHarness; -use core_test_support::test_codex::test_codex; +use core_test_support::test_codex::test_codex as base_test_codex; use core_test_support::test_path_buf; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; @@ -160,6 +160,12 @@ fn compacted_summary_only_output(summary: &str) -> Vec { }] } +fn test_codex() -> TestCodexBuilder { + base_test_codex().with_config(|config| { + let _ = config.features.disable(Feature::RemoteCompactionV2); + }) +} + fn remote_realtime_test_codex_builder( realtime_server: &responses::WebSocketTestServer, ) -> TestCodexBuilder { @@ -194,10 +200,12 @@ async fn start_remote_realtime_server() -> responses::WebSocketTestServer { async fn start_realtime_conversation(codex: &codex_core::CodexThread) -> Result<()> { codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; diff --git a/codex-rs/core/tests/suite/compact_remote_parity.rs b/codex-rs/core/tests/suite/compact_remote_parity.rs index 1b7c50936c..81038337da 100644 --- a/codex-rs/core/tests/suite/compact_remote_parity.rs +++ b/codex-rs/core/tests/suite/compact_remote_parity.rs @@ -529,8 +529,8 @@ async fn build_harness_inner( if hooks { trust_discovered_hooks(config); } - if mode == Mode::V2 { - let _ = config.features.enable(Feature::RemoteCompactionV2); + if mode == Mode::Legacy { + let _ = config.features.disable(Feature::RemoteCompactionV2); } })) .await diff --git a/codex-rs/core/tests/suite/hooks.rs b/codex-rs/core/tests/suite/hooks.rs index 16412c2d27..31e54c5e7b 100644 --- a/codex-rs/core/tests/suite/hooks.rs +++ b/codex-rs/core/tests/suite/hooks.rs @@ -34,7 +34,6 @@ use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_message_item_added; use core_test_support::responses::ev_output_text_delta; use core_test_support::responses::ev_response_created; -use core_test_support::responses::mount_compact_json_once; use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; @@ -1441,21 +1440,32 @@ async fn resumed_thread_runs_resume_then_compact_session_start_hooks() -> Result let remote_summary = "remote compact summary"; let resume_context = "remember the resumed reef"; let compact_context = "remember the compacted reef"; - let compacted_history = vec![ - ResponseItem::Message { - id: None, - role: "assistant".to_string(), - content: vec![ContentItem::OutputText { - text: remote_summary.to_string(), - }], - phase: None, - }, - ResponseItem::Compaction { - encrypted_content: "encrypted compact summary".to_string(), - }, - ]; - let compact_mock = - mount_compact_json_once(&server, serde_json::json!({ "output": compacted_history })).await; + let responses_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "hello before resume"), + ev_completed_with_tokens("resp-1", over_limit_tokens), + ]), + sse(vec![ + serde_json::json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": remote_summary, + } + }), + ev_completed("resp-compact"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "hello after resume"), + ev_completed("resp-2"), + ]), + ], + ) + .await; let mut builder = test_codex() .with_pre_build_hook(move |home| { @@ -1479,37 +1489,19 @@ async fn resumed_thread_runs_resume_then_compact_session_start_hooks() -> Result .clone() .context("rollout path")?; - mount_sse_once( - &server, - sse(vec![ - ev_response_created("resp-1"), - ev_assistant_message("msg-1", "hello before resume"), - ev_completed_with_tokens("resp-1", over_limit_tokens), - ]), - ) - .await; initial.submit_turn("hello before resume").await?; - assert!(compact_mock.requests().is_empty()); + assert_eq!(responses_mock.requests().len(), 1); let mut resume_builder = test_codex().with_config(move |config| { config.model_auto_compact_token_limit = Some(limit); trust_discovered_hooks(config); }); let resumed = resume_builder.resume(&server, home, rollout_path).await?; - let follow_up = mount_sse_once( - &server, - sse(vec![ - ev_response_created("resp-2"), - ev_assistant_message("msg-2", "hello after resume"), - ev_completed("resp-2"), - ]), - ) - .await; - resumed.submit_turn("hello after resume").await?; - assert_eq!(compact_mock.requests().len(), 1); - let developer_messages = follow_up.single_request().message_input_texts("developer"); + let requests = responses_mock.requests(); + assert_eq!(requests.len(), 3); + let developer_messages = requests[2].message_input_texts("developer"); assert!( developer_messages .iter() diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 234fb8a2f0..5764068bd9 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -110,6 +110,7 @@ mod sqlite_state; mod stream_error_allows_next_turn; mod stream_no_completed; mod subagent_notifications; +mod token_budget; mod tool_harness; mod tool_parallelism; mod tools; diff --git a/codex-rs/core/tests/suite/model_switching.rs b/codex-rs/core/tests/suite/model_switching.rs index d176b5bc24..22363eabea 100644 --- a/codex-rs/core/tests/suite/model_switching.rs +++ b/codex-rs/core/tests/suite/model_switching.rs @@ -136,6 +136,7 @@ fn test_model_info( context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), } @@ -957,6 +958,7 @@ async fn model_switch_to_smaller_model_updates_token_context_window() -> Result< context_window: Some(large_context_window), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent, experimental_supported_tools: Vec::new(), }; diff --git a/codex-rs/core/tests/suite/models_cache_ttl.rs b/codex-rs/core/tests/suite/models_cache_ttl.rs index b343c4bbaf..5103f82123 100644 --- a/codex-rs/core/tests/suite/models_cache_ttl.rs +++ b/codex-rs/core/tests/suite/models_cache_ttl.rs @@ -365,6 +365,7 @@ fn test_remote_model(slug: &str, priority: i32) -> ModelInfo { context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), input_modalities: default_input_modalities(), diff --git a/codex-rs/core/tests/suite/personality.rs b/codex-rs/core/tests/suite/personality.rs index 06c2fc41f9..a58e5f20fe 100644 --- a/codex-rs/core/tests/suite/personality.rs +++ b/codex-rs/core/tests/suite/personality.rs @@ -587,6 +587,7 @@ async fn remote_model_friendly_personality_instructions_with_feature() -> anyhow context_window: Some(128_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), input_modalities: default_input_modalities(), @@ -701,6 +702,7 @@ async fn user_turn_personality_remote_model_template_includes_update_message() - context_window: Some(128_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), input_modalities: default_input_modalities(), diff --git a/codex-rs/core/tests/suite/realtime_conversation.rs b/codex-rs/core/tests/suite/realtime_conversation.rs index 1c6763aeba..25f0516b0d 100644 --- a/codex-rs/core/tests/suite/realtime_conversation.rs +++ b/codex-rs/core/tests/suite/realtime_conversation.rs @@ -27,7 +27,6 @@ use codex_protocol::protocol::RealtimeVoice; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; use codex_protocol::user_input::UserInput; -use codex_utils_output_truncation::approx_token_count; use core_test_support::responses; use core_test_support::responses::WebSocketConnectionConfig; use core_test_support::responses::start_mock_server; @@ -283,10 +282,12 @@ async fn conversation_start_audio_text_close_round_trip() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -421,10 +422,12 @@ async fn conversation_start_defaults_to_v2_and_gpt_realtime_1_5() -> Result<()> test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -503,18 +506,20 @@ async fn conversation_webrtc_start_posts_generated_session() -> Result<()> { config.experimental_realtime_ws_model = Some("realtime-test-model".to_string()); config.experimental_realtime_ws_startup_context = Some("startup context".to_string()); config.experimental_realtime_ws_base_url = Some(realtime_ws_base_url); - config.realtime.version = RealtimeWsVersion::V1; + config.realtime.version = RealtimeWsVersion::V2; }); let test = builder.build(&server).await?; test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: Some("session-override-model".to_string()), output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: Some(ConversationStartTransport::Webrtc { sdp: "v=offer\r\n".to_string(), }), + version: Some(RealtimeConversationVersion::V1), voice: None, })) .await?; @@ -573,7 +578,7 @@ async fn conversation_webrtc_start_posts_generated_session() -> Result<()> { Some("multipart/form-data; boundary=codex-realtime-call-boundary") ); let body = String::from_utf8(request.body).context("multipart body should be utf-8")?; - let session = r#"{"audio":{"input":{"format":{"type":"audio/pcm","rate":24000}},"output":{"voice":"cove"}},"type":"quicksilver","model":"realtime-test-model","instructions":"backend prompt\n\nstartup context"}"#; + let session = r#"{"audio":{"input":{"format":{"type":"audio/pcm","rate":24000}},"output":{"voice":"cove"}},"type":"quicksilver","model":"session-override-model","instructions":"backend prompt\n\nstartup context"}"#; let session = normalized_json_string(session)?; assert_eq!( body, @@ -679,12 +684,14 @@ async fn conversation_webrtc_close_while_sideband_connecting_drops_pending_join( test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: Some(ConversationStartTransport::Webrtc { sdp: "v=offer\r\n".to_string(), }), + version: None, voice: None, })) .await?; @@ -763,12 +770,14 @@ async fn conversation_webrtc_sideband_connect_failure_closes_with_error() -> Res test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: Some(ConversationStartTransport::Webrtc { sdp: "v=offer\r\n".to_string(), }), + version: None, voice: None, })) .await?; @@ -848,10 +857,12 @@ async fn conversation_start_uses_openai_env_key_fallback_with_chatgpt_auth() -> test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -914,10 +925,12 @@ async fn conversation_transport_close_emits_closed_event() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1004,10 +1017,12 @@ async fn conversation_start_preflight_failure_emits_realtime_error_only() -> Res test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1048,10 +1063,12 @@ async fn conversation_start_connect_failure_emits_realtime_error_only() -> Resul test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1139,10 +1156,12 @@ async fn conversation_second_start_replaces_runtime() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("old".to_string())), realtime_session_id: Some("conv_old".to_string()), transport: None, + version: None, voice: None, })) .await?; @@ -1162,10 +1181,12 @@ async fn conversation_second_start_replaces_runtime() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("new".to_string())), realtime_session_id: Some("conv_new".to_string()), transport: None, + version: None, voice: None, })) .await?; @@ -1256,10 +1277,12 @@ async fn conversation_uses_experimental_realtime_ws_base_url_override() -> Resul test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1318,10 +1341,12 @@ async fn conversation_uses_default_realtime_backend_prompt() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: None, realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1388,10 +1413,12 @@ async fn conversation_uses_empty_instructions_for_null_or_empty_prompt() -> Resu ] { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt, realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1451,10 +1478,12 @@ async fn conversation_uses_explicit_start_voice() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: Some(RealtimeVoice::Breeze), })) .await?; @@ -1506,10 +1535,12 @@ async fn conversation_uses_configured_realtime_voice() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1549,10 +1580,12 @@ async fn conversation_rejects_voice_for_wrong_realtime_version() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: Some(RealtimeVoice::Cove), })) .await?; @@ -1593,10 +1626,12 @@ async fn conversation_uses_experimental_realtime_ws_backend_prompt_override() -> test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("prompt from op".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1663,10 +1698,12 @@ async fn conversation_uses_experimental_realtime_ws_startup_context_override() - test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("prompt from op".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1727,10 +1764,12 @@ async fn conversation_disables_realtime_startup_context_with_empty_override() -> test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("prompt from op".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1784,10 +1823,12 @@ async fn conversation_start_injects_startup_context_from_thread_history() -> Res test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1893,10 +1934,12 @@ async fn conversation_startup_context_current_thread_selects_many_turns_by_budge codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -1998,10 +2041,12 @@ async fn conversation_startup_context_falls_back_to_workspace_map() -> Result<() test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -2055,10 +2100,12 @@ async fn conversation_startup_context_is_truncated_and_sent_once_per_start() -> test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -2099,7 +2146,7 @@ async fn conversation_startup_context_is_truncated_and_sent_once_per_start() -> } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn conversation_user_text_turn_is_sent_to_realtime_when_active() -> Result<()> { +async fn conversation_user_text_turn_is_not_sent_to_realtime() -> Result<()> { skip_if_no_network!(Ok(())); let api_server = start_mock_server().await; @@ -2107,7 +2154,6 @@ async fn conversation_user_text_turn_is_sent_to_realtime_when_active() -> Result &api_server, responses::sse(vec![ responses::ev_response_created("resp_user_text"), - responses::ev_assistant_message("msg_user_text", "ack"), responses::ev_completed("resp_user_text"), ]), ) @@ -2133,10 +2179,12 @@ async fn conversation_user_text_turn_is_sent_to_realtime_when_active() -> Result test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -2155,7 +2203,6 @@ async fn conversation_user_text_turn_is_sent_to_realtime_when_active() -> Result assert_eq!(session_updated, "sess_user_text"); let user_text = "typed follow-up for realtime"; - let prefixed_user_text = format!("[USER] {user_text}"); test.codex .submit(Op::UserInput { items: vec![UserInput::Text { @@ -2169,181 +2216,22 @@ async fn conversation_user_text_turn_is_sent_to_realtime_when_active() -> Result }) .await?; - wait_for_event(&test.codex, |event| { - matches!(event, EventMsg::TurnComplete(_)) - }) - .await; - - let realtime_text_request = wait_for_matching_websocket_request( - &realtime_server, - "normal user turn text mirrored to realtime", - |request| websocket_request_text(request).as_deref() == Some(prefixed_user_text.as_str()), - ) - .await; - let model_user_texts = response_mock.single_request().message_input_texts("user"); - assert_eq!( - ( - model_user_texts.iter().any(|text| text == user_text), - websocket_request_text(&realtime_text_request), - ), - (true, Some(prefixed_user_text)), - ); - let realtime_response_create = timeout(Duration::from_millis(200), async { - wait_for_matching_websocket_request( - &realtime_server, - "unexpected realtime response request for mirrored user text", - |request| request.body_json()["type"].as_str() == Some("response.create"), - ) - .await - }) - .await; - assert!( - realtime_response_create.is_err(), - "mirrored user text should not request a realtime response" - ); - - let realtime_request_body = realtime_text_request.body_json(); - let content = &realtime_request_body["item"]["content"][0]; - let snapshot = format!( - "type: {}\nitem.type: {}\nitem.role: {}\ncontent[0].type: {}\ncontent[0].text: {}\nresponse.create: {}", - realtime_request_body["type"].as_str().unwrap_or_default(), - realtime_request_body["item"]["type"] - .as_str() - .unwrap_or_default(), - realtime_request_body["item"]["role"] - .as_str() - .unwrap_or_default(), - content["type"].as_str().unwrap_or_default(), - content["text"].as_str().unwrap_or_default(), - realtime_response_create.is_ok(), - ); - insta::assert_snapshot!( - "conversation_user_text_turn_is_sent_to_realtime_when_active", - snapshot - ); - - realtime_server.shutdown().await; - Ok(()) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn conversation_user_text_turn_is_capped_when_mirrored_to_realtime() -> Result<()> { - skip_if_no_network!(Ok(())); - - let api_server = start_mock_server().await; - let response_mock = responses::mount_sse_once( - &api_server, - responses::sse(vec![ - responses::ev_response_created("resp_long_user_text"), - responses::ev_assistant_message("msg_long_user_text", "ack"), - responses::ev_completed("resp_long_user_text"), - ]), - ) - .await; - - let realtime_server = start_websocket_server(vec![vec![ - vec![json!({ - "type": "session.updated", - "session": { "id": "sess_long_user_text", "instructions": "backend prompt" } - })], - vec![], - ]]) - .await; - - let mut builder = test_codex().with_config({ - let realtime_base_url = realtime_server.uri().to_string(); - move |config| { - config.experimental_realtime_ws_base_url = Some(realtime_base_url); - config.experimental_realtime_ws_startup_context = Some(String::new()); - } - }); - let test = builder.build(&api_server).await?; - - // Phase 1: start realtime so the next normal user turn mirrors over the - // active WebSocket session. - test.codex - .submit(Op::RealtimeConversationStart(ConversationStartParams { - output_modality: RealtimeOutputModality::Audio, - prompt: Some(Some("backend prompt".to_string())), - realtime_session_id: None, - transport: None, - voice: None, - })) - .await?; - - let session_updated = wait_for_event_match(&test.codex, |msg| match msg { - EventMsg::RealtimeConversationRealtime(RealtimeConversationRealtimeEvent { - payload: - RealtimeEvent::SessionUpdated { - realtime_session_id: session_id, - .. - }, - }) => Some(session_id.clone()), + let turn_complete = wait_for_event_match(&test.codex, |event| match event { + EventMsg::TurnComplete(turn_complete) => Some(turn_complete.clone()), _ => None, }) .await; - assert_eq!(session_updated, "sess_long_user_text"); + assert_eq!(turn_complete.last_agent_message, None); - // Phase 2: submit one oversized text turn. The model request should keep - // the exact user text, while the realtime mirror should get the capped copy. - let user_text = format!( - "mirror-head {} mirror-middle {} mirror-tail", - "alpha ".repeat(900), - "omega ".repeat(900), - ); - test.codex - .submit(Op::UserInput { - items: vec![UserInput::Text { - text: user_text.clone(), - text_elements: Vec::new(), - }], - final_output_json_schema: None, - responsesapi_client_metadata: None, - additional_context: Default::default(), - thread_settings: Default::default(), - }) - .await?; - - wait_for_event(&test.codex, |event| { - matches!(event, EventMsg::TurnComplete(_)) - }) - .await; - - // Phase 3: capture the mirrored WebSocket item; the snapshot below records - // the capped payload shape. - let realtime_text_request = wait_for_matching_websocket_request( - &realtime_server, - "capped normal user turn text mirrored to realtime", - |request| websocket_request_text(request).is_some_and(|text| text.contains("mirror-head")), - ) - .await; - let realtime_text = - websocket_request_text(&realtime_text_request).expect("realtime request text"); let model_user_texts = response_mock.single_request().message_input_texts("user"); + assert!(model_user_texts.iter().any(|text| text == user_text)); - let realtime_request_body = realtime_text_request.body_json(); - let content = &realtime_request_body["item"]["content"][0]; - - // Snapshot the request envelope and capped text together so reviewers can - // see the preserved head/tail and truncation marker in one place. - let snapshot = format!( - "type: {}\nitem.type: {}\nitem.role: {}\ncontent[0].type: {}\nmodel_has_full_user_text: {}\nrealtime_text_equal_full_user_text: {}\nrealtime_text_approx_tokens: {}\ncontent[0].text: {}", - realtime_request_body["type"].as_str().unwrap_or_default(), - realtime_request_body["item"]["type"] - .as_str() - .unwrap_or_default(), - realtime_request_body["item"]["role"] - .as_str() - .unwrap_or_default(), - content["type"].as_str().unwrap_or_default(), - model_user_texts.iter().any(|text| text == &user_text), - realtime_text == user_text, - approx_token_count(&realtime_text), - realtime_text, - ); - insta::assert_snapshot!( - "conversation_user_text_turn_is_capped_when_mirrored_to_realtime", - snapshot + let realtime_connections = realtime_server.connections(); + assert_eq!(realtime_connections.len(), 1); + assert_eq!(realtime_connections[0].len(), 1); + assert_eq!( + realtime_connections[0][0].body_json()["type"].as_str(), + Some("session.update") ); realtime_server.shutdown().await; @@ -2387,10 +2275,12 @@ async fn realtime_v2_noop_tool_call_returns_empty_function_output_without_respon test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -2483,10 +2373,12 @@ async fn conversation_mirrors_assistant_message_text_to_realtime_handoff() -> Re test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -2617,10 +2509,12 @@ async fn conversation_handoff_persists_across_item_done_until_turn_complete() -> test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -2766,10 +2660,12 @@ async fn inbound_handoff_request_starts_turn() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -2865,10 +2761,12 @@ async fn inbound_handoff_request_uses_active_transcript() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -2965,10 +2863,12 @@ async fn inbound_handoff_request_sends_transcript_delta_after_each_handoff() -> test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -3063,10 +2963,12 @@ async fn inbound_conversation_item_does_not_start_turn_and_still_forwards_audio( test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -3183,10 +3085,12 @@ async fn delegated_turn_user_role_echo_does_not_redelegate_and_still_forwards_au test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -3333,10 +3237,12 @@ async fn inbound_handoff_request_does_not_block_realtime_event_forwarding() -> R test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -3441,7 +3347,6 @@ async fn inbound_handoff_request_steers_active_turn() -> Result<()> { "type": "session.updated", "session": { "id": "sess_steer", "instructions": "backend prompt" } })], - vec![], vec![ json!({ "type": "conversation.input_transcript.delta", @@ -3468,10 +3373,12 @@ async fn inbound_handoff_request_steers_active_turn() -> Result<()> { test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; @@ -3504,12 +3411,6 @@ async fn inbound_handoff_request_steers_active_turn() -> Result<()> { matches!(event, EventMsg::AgentMessageContentDelta(_)) }) .await; - let _ = wait_for_matching_websocket_request( - &realtime_server, - "first prompt mirrored to realtime", - |request| websocket_request_text(request).as_deref() == Some("first prompt"), - ) - .await; test.codex .submit(Op::RealtimeConversationAudio(ConversationAudioParams { @@ -3624,10 +3525,12 @@ async fn inbound_handoff_request_starts_turn_and_does_not_block_realtime_audio() test.codex .submit(Op::RealtimeConversationStart(ConversationStartParams { + model: None, output_modality: RealtimeOutputModality::Audio, prompt: Some(Some("backend prompt".to_string())), realtime_session_id: None, transport: None, + version: None, voice: None, })) .await?; diff --git a/codex-rs/core/tests/suite/remote_models.rs b/codex-rs/core/tests/suite/remote_models.rs index dc8df93c9b..b3fc9b9db7 100644 --- a/codex-rs/core/tests/suite/remote_models.rs +++ b/codex-rs/core/tests/suite/remote_models.rs @@ -499,6 +499,7 @@ async fn remote_models_remote_model_uses_unified_exec() -> Result<()> { context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), }; @@ -751,6 +752,7 @@ async fn remote_models_apply_remote_base_instructions() -> Result<()> { context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), }; @@ -1237,6 +1239,7 @@ fn test_remote_model_with_policy( context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), } diff --git a/codex-rs/core/tests/suite/responses_lite.rs b/codex-rs/core/tests/suite/responses_lite.rs index 60db4c064d..b137d5bc73 100644 --- a/codex-rs/core/tests/suite/responses_lite.rs +++ b/codex-rs/core/tests/suite/responses_lite.rs @@ -9,9 +9,11 @@ use codex_features::Feature; use codex_image_generation_extension::install as install_image_generation_extension; use codex_login::CodexAuth; use codex_protocol::config_types::WebSearchMode; +use codex_protocol::models::ImageDetail; use codex_protocol::openai_models::InputModality; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; use codex_web_search_extension::install as install_web_search_extension; use core_test_support::responses; use core_test_support::skip_if_no_network; @@ -52,6 +54,62 @@ fn has_hosted_tool(tools: &[Value], tool_type: &str) -> bool { .any(|tool| tool.get("type").and_then(Value::as_str) == Some(tool_type)) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_lite_strips_data_image_detail_without_resize_all_images() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + let image_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg=="; + let mut builder = test_codex().with_model_info_override("gpt-5.4", |model_info| { + model_info.use_responses_lite = true; + configure_image_capable_model(model_info); + }); + let test = builder.build(&server).await?; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Image { + image_url: image_url.to_string(), + detail: Some(ImageDetail::Original), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let request = response_mock.single_request(); + let input = request.input(); + let image = input + .iter() + .filter_map(|item| item.get("content").and_then(Value::as_array)) + .flatten() + .find(|item| item.get("type").and_then(Value::as_str) == Some("input_image")) + .context("request should contain an image")?; + assert_eq!( + image, + &serde_json::json!({ + "type": "input_image", + "image_url": image_url + }) + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_lite_uses_standalone_web_search_and_image_generation() -> Result<()> { skip_if_no_network!(Ok(())); @@ -119,10 +177,14 @@ async fn responses_lite_compact_request_uses_lite_transport_contract() -> Result let compact_mock = responses::mount_compact_json_once(&server, serde_json::json!({ "output": [] })).await; - let mut builder = test_codex().with_model_info_override("gpt-5.4", |model_info| { - model_info.use_responses_lite = true; - model_info.supports_parallel_tool_calls = true; - }); + let mut builder = test_codex() + .with_model_info_override("gpt-5.4", |model_info| { + model_info.use_responses_lite = true; + model_info.supports_parallel_tool_calls = true; + }) + .with_config(|config| { + let _ = config.features.disable(Feature::RemoteCompactionV2); + }); let test = builder.build(&server).await?; test.submit_turn("Compact this conversation").await?; diff --git a/codex-rs/core/tests/suite/resume_warning.rs b/codex-rs/core/tests/suite/resume_warning.rs index 422489a86a..2120559751 100644 --- a/codex-rs/core/tests/suite/resume_warning.rs +++ b/codex-rs/core/tests/suite/resume_warning.rs @@ -37,6 +37,7 @@ fn resume_history( network: None, file_system_sandbox_policy: None, model: previous_model.to_string(), + comp_hash: None, personality: None, collaboration_mode: None, multi_agent_version: None, diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 5cb61fa4a1..7f0493ea8d 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -2,6 +2,8 @@ use anyhow::Context as _; use anyhow::ensure; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use core_test_support::test_codex::local_selections; use std::collections::HashMap; use std::ffi::OsStr; @@ -25,6 +27,7 @@ use codex_core::config::Config; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::Environment; use codex_exec_server::HttpRequestParams; +use codex_features::Feature; use codex_login::CodexAuth; use codex_mcp::MCP_SANDBOX_STATE_META_CAPABILITY; use codex_models_manager::manager::RefreshStrategy; @@ -57,11 +60,16 @@ use core_test_support::test_codex::test_codex; use core_test_support::test_codex::turn_permission_fields; use core_test_support::wait_for_event; use core_test_support::wait_for_mcp_server; +use image::DynamicImage; +use image::GenericImageView; +use image::ImageBuffer; +use image::Rgba; use reqwest::Client; use reqwest::StatusCode; use serde_json::Value; use serde_json::json; use serial_test::serial; +use std::io::Cursor; use tempfile::tempdir; use tokio::process::Child; use tokio::process::Command; @@ -558,6 +566,52 @@ async fn stdio_server_round_trip() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shutdown_cancels_startup_prewarm_waiting_for_mcp_startup() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_websocket_server(vec![vec![vec![ + responses::ev_response_created("warm-1"), + responses::ev_completed("warm-1"), + ]]]) + .await; + let pending_mcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let pending_mcp_url = format!("http://{}/mcp", pending_mcp_listener.local_addr()?); + + let fixture = test_codex() + .with_config(move |config| { + insert_mcp_server( + config, + "shutdown_prewarm", + McpServerTransportConfig::StreamableHttp { + url: pending_mcp_url, + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + }, + TestMcpServerOptions::default(), + ); + }) + .build_with_websocket_server(&server) + .await?; + + let (_pending_mcp_connection, _) = + tokio::time::timeout(Duration::from_secs(5), pending_mcp_listener.accept()) + .await + .context("startup prewarm should start the MCP connection")??; + tokio::time::timeout(Duration::from_secs(2), fixture.codex.shutdown_and_wait()) + .await + .context("shutdown should not wait for startup prewarm MCP startup")??; + tokio::time::sleep(Duration::from_millis(100)).await; + assert!( + server.connections().is_empty(), + "startup prewarm should not send a websocket request after shutdown" + ); + + server.shutdown().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[serial(mcp_cwd)] async fn stdio_server_uses_configured_cwd_before_runtime_fallback() -> anyhow::Result<()> { @@ -1211,6 +1265,107 @@ async fn stdio_image_responses_round_trip() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial(mcp_test_value)] +async fn stdio_image_responses_resize_large_image() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let call_id = "img-resize-1"; + let server_name = "rmcp"; + let namespace = format!("mcp__{server_name}"); + + let original_dimensions = (3000, 2000); + let image = ImageBuffer::from_pixel( + original_dimensions.0, + original_dimensions.1, + Rgba([20, 40, 60, 255]), + ); + let mut encoded = Cursor::new(Vec::new()); + DynamicImage::ImageRgba8(image).write_to(&mut encoded, image::ImageFormat::Png)?; + let image_data_url = format!( + "data:image/png;base64,{}", + BASE64_STANDARD.encode(encoded.into_inner()) + ); + let tool_arguments = serde_json::to_string(&json!({ + "scenario": "image_only", + "data_url": image_data_url, + }))?; + + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call_with_namespace( + call_id, + &namespace, + "image_scenario", + &tool_arguments, + ), + responses::ev_completed("resp-1"), + ]), + ) + .await; + let final_mock = mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_assistant_message("msg-1", "done"), + responses::ev_completed("resp-2"), + ]), + ) + .await; + + let rmcp_test_server_bin = remote_aware_stdio_server_bin()?; + let fixture = test_codex() + .with_config(move |config| { + config + .features + .enable(Feature::ResizeAllImages) + .expect("resize_all_images should be enabled"); + insert_mcp_server( + config, + server_name, + stdio_transport(rmcp_test_server_bin, /*env*/ None, Vec::new()), + TestMcpServerOptions { + environment_id: remote_aware_environment_id(), + ..Default::default() + }, + ); + }) + .build_with_remote_env(&server) + .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; + + fixture + .codex + .submit(read_only_user_turn( + &fixture, + "call the rmcp image_scenario tool", + )) + .await?; + wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + let output_item = final_mock.single_request().function_call_output(call_id); + assert_eq!(output_item["call_id"], call_id); + let output = output_item["output"] + .as_array() + .expect("image MCP output should be content items"); + let resized_url = output[1]["image_url"] + .as_str() + .expect("MCP image output should contain a data URL"); + assert_eq!(output[1]["detail"], "high"); + let (_, resized_base64) = resized_url + .split_once(',') + .expect("resized image should contain a data URL prefix"); + let resized_bytes = BASE64_STANDARD.decode(resized_base64)?; + let resized = image::load_from_memory(&resized_bytes)?; + let resized_dimensions = resized.dimensions(); + assert_eq!(resized_dimensions, (1920, 1280)); + + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[serial(mcp_test_value)] async fn stdio_image_responses_preserve_original_detail_metadata() -> anyhow::Result<()> { @@ -1344,6 +1499,7 @@ async fn stdio_image_responses_are_sanitized_for_text_only_model() -> anyhow::Re context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), input_modalities: vec![InputModality::Text], diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__realtime_conversation__conversation_user_text_turn_is_capped_when_mirrored_to_realtime.snap b/codex-rs/core/tests/suite/snapshots/all__suite__realtime_conversation__conversation_user_text_turn_is_capped_when_mirrored_to_realtime.snap deleted file mode 100644 index 3daae5d378..0000000000 --- a/codex-rs/core/tests/suite/snapshots/all__suite__realtime_conversation__conversation_user_text_turn_is_capped_when_mirrored_to_realtime.snap +++ /dev/null @@ -1,12 +0,0 @@ ---- -source: core/tests/suite/realtime_conversation.rs -expression: snapshot ---- -type: conversation.item.create -item.type: message -item.role: user -content[0].type: input_text -model_has_full_user_text: true -realtime_text_equal_full_user_text: false -realtime_text_approx_tokens: 300 -content[0].text: [USER] mirror-head alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alpha alp…2419 tokens truncated…ega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega omega mirror-tail diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__realtime_conversation__conversation_user_text_turn_is_sent_to_realtime_when_active.snap b/codex-rs/core/tests/suite/snapshots/all__suite__realtime_conversation__conversation_user_text_turn_is_sent_to_realtime_when_active.snap deleted file mode 100644 index 3ca8233029..0000000000 --- a/codex-rs/core/tests/suite/snapshots/all__suite__realtime_conversation__conversation_user_text_turn_is_sent_to_realtime_when_active.snap +++ /dev/null @@ -1,10 +0,0 @@ ---- -source: core/tests/suite/realtime_conversation.rs -expression: snapshot ---- -type: conversation.item.create -item.type: message -item.role: user -content[0].type: input_text -content[0].text: [USER] typed follow-up for realtime -response.create: false diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap b/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap new file mode 100644 index 0000000000..d042616afc --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap @@ -0,0 +1,15 @@ +--- +source: core/tests/suite/token_budget.rs +assertion_line: 300 +expression: "context_snapshot::format_labeled_requests_snapshot(\"New context window tool installs fresh full context before the next follow-up request.\",\n&[(\"Final Follow-Up Request\", &requests[2])],\n&ContextSnapshotOptions::default(),)" +--- +Scenario: New context window tool installs fresh full context before the next follow-up request. + +## Final Follow-Up Request +00:message/developer[3]: + [01] + [02] + [03] \nCurrent context window 1.\nYou have 121600 tokens left in this context window.\n +01:message/user:> +02:function_call/update_plan +03:function_call_output:Plan updated diff --git a/codex-rs/core/tests/suite/spawn_agent_description.rs b/codex-rs/core/tests/suite/spawn_agent_description.rs index ba29b4f00c..30aab8926e 100644 --- a/codex-rs/core/tests/suite/spawn_agent_description.rs +++ b/codex-rs/core/tests/suite/spawn_agent_description.rs @@ -84,6 +84,7 @@ fn test_model_info( context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), } diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs new file mode 100644 index 0000000000..f50ec09894 --- /dev/null +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -0,0 +1,449 @@ +use anyhow::Result; +use codex_features::Feature; +use codex_model_provider_info::built_in_model_providers; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use core_test_support::PathBufExt; +use core_test_support::context_snapshot; +use core_test_support::context_snapshot::ContextSnapshotOptions; +use core_test_support::responses::ResponsesRequest; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_completed_with_tokens; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::local; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; + +const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000; +const EFFECTIVE_CONTEXT_WINDOW: i64 = CONFIGURED_CONTEXT_WINDOW * 95 / 100; + +fn token_budget_texts(request: &ResponsesRequest) -> Vec { + request + .message_input_texts("developer") + .into_iter() + .filter(|text| text.starts_with("")) + .collect() +} + +fn tool_names(request: &ResponsesRequest) -> Vec { + request + .body_json() + .get("tools") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(|tool| tool.get("name").and_then(Value::as_str).map(str::to_string)) + .collect() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + test.submit_turn("first turn").await?; + + let second_cwd = test.workspace_path("second-cwd"); + std::fs::create_dir_all(&second_cwd)?; + test.submit_turn_with_environments("second turn", Some(vec![local(second_cwd.abs())])) + .await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + + let expected = vec![format!( + "\nCurrent context window 0.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" + )]; + assert_eq!( + token_budget_texts(&requests[0]), + expected, + "initial full context should report context window 0" + ); + assert_eq!( + token_budget_texts(&requests[1]), + expected, + "steady-state context update should not advance the context window" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 2_500), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_completed_with_tokens("resp-2", /*total_tokens*/ 3_000), + ]), + sse(vec![ + ev_response_created("resp-3"), + ev_completed_with_tokens("resp-3", /*total_tokens*/ 5_000), + ]), + sse(vec![ + ev_response_created("resp-4"), + ev_completed_with_tokens("resp-4", /*total_tokens*/ 8_000), + ]), + sse(vec![ev_response_created("resp-5"), ev_completed("resp-5")]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + config.model_context_window = Some(10_000); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + for turn in 1..=5 { + test.submit_turn(&format!("turn {turn}")).await?; + } + + let requests = responses.requests(); + assert_eq!(requests.len(), 5); + + let full_context = "\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n" + .to_string(); + let threshold_25 = + "\nYou have 7000 tokens left in this context window.\n" + .to_string(); + let threshold_50 = + "\nYou have 4500 tokens left in this context window.\n" + .to_string(); + let threshold_75 = + "\nYou have 1500 tokens left in this context window.\n" + .to_string(); + + assert_eq!(token_budget_texts(&requests[0]), vec![full_context.clone()]); + assert_eq!( + token_budget_texts(&requests[1]), + vec![full_context.clone(), threshold_25.clone()] + ); + assert_eq!( + token_budget_texts(&requests[2]), + vec![full_context.clone(), threshold_25.clone()] + ); + assert_eq!( + token_budget_texts(&requests[3]), + vec![ + full_context.clone(), + threshold_25.clone(), + threshold_50.clone() + ] + ); + assert_eq!( + token_budget_texts(&requests[4]), + vec![full_context, threshold_25, threshold_50, threshold_75] + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "remaining-call"; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "noted"), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 2_500), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_function_call(call_id, "get_context_remaining", "{}"), + ev_completed_with_tokens("resp-2", /*total_tokens*/ 2_500), + ]), + sse(vec![ + ev_response_created("resp-3"), + ev_assistant_message("msg-3", "done"), + ev_completed("resp-3"), + ]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + config.model_context_window = Some(10_000); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + test.submit_turn("spend some tokens").await?; + test.submit_turn("check remaining context").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 3); + assert!( + tool_names(&requests[1]) + .iter() + .any(|name| name == "get_context_remaining"), + "get_context_remaining should be exposed when token budget is enabled" + ); + + let full_context = "\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n" + .to_string(); + let remaining_context = + "\nYou have 7000 tokens left in this context window.\n" + .to_string(); + assert_eq!( + token_budget_texts(&requests[1]), + vec![full_context, remaining_context.clone()] + ); + assert_eq!( + requests[2].function_call_output_content_and_success(call_id), + Some((Some(remaining_context), None)) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn get_context_remaining_returns_unknown_when_window_is_unavailable() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "remaining-call"; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "get_context_remaining", "{}"), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + let test = test_codex() + .with_model_info_override("gpt-5.2", |model_info| { + model_info.context_window = None; + model_info.max_context_window = None; + }) + .with_config(|config| { + config.model_context_window = None; + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + test.submit_turn("check remaining context").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + assert!( + tool_names(&requests[0]) + .iter() + .any(|name| name == "get_context_remaining"), + "get_context_remaining should be exposed when token budget is enabled" + ); + + assert_eq!(token_budget_texts(&requests[0]), Vec::::new()); + assert_eq!( + requests[1].function_call_output_content_and_success(call_id), + Some(( + Some( + "\nYou have unknown tokens left in this context window.\n" + .to_string() + ), + None, + )) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + sse(vec![ + ev_response_created("resp-compact"), + ev_assistant_message("msg-compact", "compact summary"), + ev_completed("resp-compact"), + ]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + + let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone(); + model_provider.name = "OpenAI-compatible test provider".to_string(); + model_provider.base_url = Some(format!("{}/v1", server.uri())); + model_provider.supports_websockets = false; + + let test = test_codex() + .with_config(move |config| { + config.model_provider = model_provider; + config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + test.submit_turn("before compact").await?; + test.codex.submit(Op::Compact).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.submit_turn("after compact").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 3); + + assert_eq!( + token_budget_texts(&requests[2]), + vec![format!( + "\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" + )], + "post-compaction full context should report context window 1" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "new-window-call"; + let continue_call_id = "continue-call"; + let continue_args = json!({ + "plan": [ + {"step": "Continue in the new context window", "status": "in_progress"} + ], + }) + .to_string(); + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "new_context", "{}"), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_function_call(continue_call_id, "update_plan", &continue_args), + ev_completed("resp-2"), + ]), + sse(vec![ + ev_response_created("resp-3"), + ev_assistant_message("msg-3", "done"), + ev_completed("resp-3"), + ]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + test.submit_turn("request new context window").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 3); + assert!( + tool_names(&requests[0]) + .iter() + .any(|name| name == "new_context"), + "new_context should be exposed when token budget is enabled" + ); + assert_eq!( + token_budget_texts(&requests[2]), + vec![format!( + "\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" + )] + ); + assert!( + !requests[2].body_contains_text("request new context window"), + "new_context should drop the prior window history before continuing the turn" + ); + assert_eq!( + requests[2].function_call_output_text(continue_call_id), + Some("Plan updated".to_string()) + ); + insta::assert_snapshot!( + "token_budget_new_context_window_tool_full_context", + context_snapshot::format_labeled_requests_snapshot( + "New context window tool installs fresh full context before the next follow-up request.", + &[("Final Follow-Up Request", &requests[2])], + &ContextSnapshotOptions::default(), + ) + ); + + Ok(()) +} diff --git a/codex-rs/core/tests/suite/unified_exec.rs b/codex-rs/core/tests/suite/unified_exec.rs index caea5ecf55..eb8bf29551 100644 --- a/codex-rs/core/tests/suite/unified_exec.rs +++ b/codex-rs/core/tests/suite/unified_exec.rs @@ -1995,6 +1995,249 @@ async fn write_stdin_returns_exit_metadata_and_clears_session() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn write_stdin_ctrl_c_interrupts_non_tty_session() -> Result<()> { + assert_write_stdin_ctrl_c_interrupts_non_tty_session( + "trap", + "trap 'echo INT-TRAP; exit 42' INT; echo READY; while true; do sleep 30; done", + /*expected_exit_code*/ 42, + Some("INT-TRAP"), + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn write_stdin_ctrl_c_default_interrupt_reports_130_for_non_tty_session() -> Result<()> { + assert_write_stdin_ctrl_c_interrupts_non_tty_session( + "default", + "echo READY; exec sleep 30", + /*expected_exit_code*/ 130, + /*expected_interrupt_output*/ None, + ) + .await +} + +async fn assert_write_stdin_ctrl_c_interrupts_non_tty_session( + test_name: &str, + command: &str, + expected_exit_code: i32, + expected_interrupt_output: Option<&str>, +) -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + skip_if_windows!(Ok(())); + + let server = start_mock_server().await; + + let mut builder = test_codex().with_config(|config| { + if let Err(err) = config.features.enable(Feature::UnifiedExec) { + panic!("test config should allow feature update: {err}"); + } + }); + let test = builder.build_with_remote_env(&server).await?; + + let start_call_id = format!("uexec-non-tty-interrupt-{test_name}-start"); + let interrupt_call_id = format!("uexec-non-tty-interrupt-{test_name}"); + + let start_args = serde_json::json!({ + "cmd": command, + "yield_time_ms": 250, + "tty": false, + }); + let interrupt_args = serde_json::json!({ + "chars": "\u{3}", + "session_id": 1000, + "yield_time_ms": 1000, + }); + + let responses = vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + &start_call_id, + "exec_command", + &serde_json::to_string(&start_args)?, + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_function_call( + &interrupt_call_id, + "write_stdin", + &serde_json::to_string(&interrupt_args)?, + ), + ev_completed("resp-2"), + ]), + sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-3"), + ]), + ]; + let request_log = mount_sse_sequence(&server, responses).await; + + submit_unified_exec_turn( + &test, + "interrupt non-tty unified exec", + PermissionProfile::Disabled, + ) + .await?; + + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = request_log.requests(); + assert!(!requests.is_empty(), "expected at least one POST request"); + let bodies = requests + .into_iter() + .map(|request| request.body_json()) + .collect::>(); + + let outputs = collect_tool_outputs(&bodies)?; + + let start_output = outputs + .get(&start_call_id) + .with_context(|| format!("missing start output for exec_command {start_call_id}"))?; + assert_eq!( + start_output.process_id.as_deref(), + Some("1000"), + "exec_command should leave a running non-TTY session" + ); + assert!( + start_output.exit_code.is_none(), + "initial exec_command should not include exit_code while session is running" + ); + assert!( + start_output.output.contains("READY"), + "start output should include command readiness marker, got {:?}", + start_output.output + ); + + let interrupt_output = outputs + .get(&interrupt_call_id) + .with_context(|| format!("missing interrupt output for write_stdin {interrupt_call_id}"))?; + assert!( + interrupt_output.process_id.is_none(), + "interrupted process should be cleared from the session map" + ); + assert_eq!( + interrupt_output.exit_code, + Some(expected_exit_code), + "interrupt should preserve the process-reported exit code" + ); + if let Some(expected_interrupt_output) = expected_interrupt_output { + assert!( + interrupt_output.output.contains(expected_interrupt_output), + "interrupt should drain output from the signal handler, got {:?}", + interrupt_output.output + ); + } + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[cfg_attr(not(windows), ignore = "Windows-only unified exec interrupt test")] +async fn write_stdin_ctrl_c_reports_unsupported_interrupt_to_model_on_windows() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + + let server = start_mock_server().await; + + let mut builder = test_codex().with_config(|config| { + config + .features + .enable(Feature::UnifiedExec) + .expect("test config should allow feature update"); + }); + let test = builder.build_with_remote_env(&server).await?; + + let start_call_id = "uexec-windows-interrupt-start"; + let interrupt_call_id = "uexec-windows-interrupt"; + + let start_args = serde_json::json!({ + "shell": "cmd", + "cmd": "echo READY && ping -n 30 127.0.0.1 >NUL", + "yield_time_ms": 250, + "tty": false, + }); + let interrupt_args = serde_json::json!({ + "chars": "\u{3}", + "session_id": 1000, + "yield_time_ms": 1000, + }); + + let responses = vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + start_call_id, + "exec_command", + &serde_json::to_string(&start_args)?, + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_function_call( + interrupt_call_id, + "write_stdin", + &serde_json::to_string(&interrupt_args)?, + ), + ev_completed("resp-2"), + ]), + sse(vec![ + ev_response_created("resp-3"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-3"), + ]), + ]; + let request_log = mount_sse_sequence(&server, responses).await; + + submit_unified_exec_turn( + &test, + "interrupt non-tty unified exec on Windows", + PermissionProfile::Disabled, + ) + .await?; + + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let start_output = request_log + .function_call_output_text(start_call_id) + .expect("missing start output for exec_command"); + let start_output = parse_unified_exec_output(&start_output)?; + assert_eq!( + start_output.process_id.as_deref(), + Some("1000"), + "exec_command should leave a running non-TTY session" + ); + assert!( + start_output.output.contains("READY"), + "start output should include command readiness marker, got {:?}", + start_output.output + ); + + let interrupt_output = request_log + .function_call_output_text(interrupt_call_id) + .expect("missing interrupt output for write_stdin"); + assert!( + interrupt_output.contains("write_stdin failed"), + "model-visible write_stdin output should report failure, got {interrupt_output:?}" + ); + assert!( + interrupt_output.contains("process interrupt is not supported by this process backend"), + "model-visible write_stdin output should explain unsupported interrupt, got {interrupt_output:?}" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn unified_exec_emits_end_event_when_session_dies_via_stdin() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs index 5820788566..240f386ec6 100644 --- a/codex-rs/core/tests/suite/view_image.rs +++ b/codex-rs/core/tests/suite/view_image.rs @@ -7,6 +7,7 @@ use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::LOCAL_ENVIRONMENT_ID; use codex_exec_server::REMOTE_ENVIRONMENT_ID; use codex_exec_server::RemoveOptions; +use codex_features::Feature; use codex_login::CodexAuth; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::models::PermissionProfile; @@ -177,10 +178,15 @@ async fn write_workspace_png( async fn assert_user_turn_local_image_resizes_to( original_dimensions: (u32, u32), expected_dimensions: (u32, u32), + resize_policy: TestImageResizePolicy, ) -> anyhow::Result<()> { let server = start_mock_server().await; - let mut builder = test_codex(); + let mut builder = test_codex().with_config(move |config| { + if resize_policy == TestImageResizePolicy::AllImages { + let _ = config.features.enable(Feature::ResizeAllImages); + } + }); let test = builder.build_with_remote_env(&server).await?; let TestCodex { codex, @@ -254,18 +260,42 @@ async fn assert_user_turn_local_image_resizes_to( Ok(()) } +#[derive(Clone, Copy, Eq, PartialEq)] +enum TestImageResizePolicy { + Legacy, + AllImages, +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn user_turn_with_local_image_attaches_image() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); - assert_user_turn_local_image_resizes_to((2304, 864), (2048, 768)).await + assert_user_turn_local_image_resizes_to((2304, 864), (2048, 768), TestImageResizePolicy::Legacy) + .await } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn user_turn_with_vertical_local_image_resizes_to_square_bounds() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); - assert_user_turn_local_image_resizes_to((1024, 4096), (512, 2048)).await + assert_user_turn_local_image_resizes_to( + (1024, 4096), + (512, 2048), + TestImageResizePolicy::Legacy, + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn resize_all_images_applies_patch_budget_to_local_user_image() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + assert_user_turn_local_image_resizes_to( + (2048, 2048), + (1600, 1600), + TestImageResizePolicy::AllImages, + ) + .await } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -1248,6 +1278,72 @@ async fn view_image_tool_errors_for_non_image_files() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn resize_all_images_turns_invalid_view_image_into_placeholder() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let mut builder = test_codex().with_config(|config| { + let _ = config.features.enable(Feature::ResizeAllImages); + }); + let test = builder.build_with_remote_env(&server).await?; + let TestCodex { + codex, + session_configured, + .. + } = &test; + + let rel_path = "assets/invalid-image.json"; + write_workspace_file(&test, rel_path, br#"{ "message": "hello" }"#.to_vec()).await?; + let call_id = "view-image-invalid-placeholder"; + let arguments = serde_json::json!({ "path": rel_path }).to_string(); + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "view_image", &arguments), + ev_completed("resp-1"), + ]), + ) + .await; + let second_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + ) + .await; + + codex + .submit(disabled_user_turn( + &test, + vec![UserInput::Text { + text: "please inspect the image".into(), + text_elements: Vec::new(), + }], + session_configured.model.clone(), + )) + .await?; + wait_for_event_with_timeout( + codex, + |event| matches!(event, EventMsg::TurnComplete(_)), + VIEW_IMAGE_TURN_COMPLETE_TIMEOUT, + ) + .await; + + let request = second_mock.single_request(); + assert_eq!( + request.function_call_output(call_id).get("output"), + Some(&serde_json::json!([{ + "type": "input_text", + "text": "image content omitted because it could not be processed" + }])) + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn view_image_tool_errors_when_file_missing() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); @@ -1374,6 +1470,7 @@ async fn view_image_tool_returns_unsupported_message_for_text_only_model() -> an context_window: Some(272_000), max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), }; diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 8cf18e7d1f..26b3e46e88 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -35,6 +35,7 @@ use crate::protocol::EXEC_EXITED_METHOD; use crate::protocol::EXEC_METHOD; use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; use crate::protocol::EXEC_READ_METHOD; +use crate::protocol::EXEC_SIGNAL_METHOD; use crate::protocol::EXEC_TERMINATE_METHOD; use crate::protocol::EXEC_WRITE_METHOD; use crate::protocol::EnvironmentInfo; @@ -80,8 +81,11 @@ use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; use crate::protocol::InitializeResponse; use crate::protocol::ProcessOutputChunk; +use crate::protocol::ProcessSignal; use crate::protocol::ReadParams; use crate::protocol::ReadResponse; +use crate::protocol::SignalParams; +use crate::protocol::SignalResponse; use crate::protocol::TerminateParams; use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; @@ -394,6 +398,23 @@ impl ExecServerClient { .await } + pub async fn signal( + &self, + process_id: &ProcessId, + signal: ProcessSignal, + ) -> Result<(), ExecServerError> { + let _response: SignalResponse = self + .call( + EXEC_SIGNAL_METHOD, + &SignalParams { + process_id: process_id.clone(), + signal, + }, + ) + .await?; + Ok(()) + } + pub async fn terminate( &self, process_id: &ProcessId, @@ -763,6 +784,10 @@ impl Session { self.client.write(&self.process_id, chunk).await } + pub(crate) async fn signal(&self, signal: ProcessSignal) -> Result<(), ExecServerError> { + self.client.signal(&self.process_id, signal).await + } + pub(crate) async fn terminate(&self) -> Result<(), ExecServerError> { self.client.terminate(&self.process_id).await?; Ok(()) diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index e306ec395e..378d12e632 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -93,9 +93,12 @@ pub use protocol::HttpRequestResponse; pub use protocol::InitializeParams; pub use protocol::InitializeResponse; pub use protocol::ProcessOutputChunk; +pub use protocol::ProcessSignal; pub use protocol::ReadParams; pub use protocol::ReadResponse; pub use protocol::ShellInfo; +pub use protocol::SignalParams; +pub use protocol::SignalResponse; pub use protocol::TerminateParams; pub use protocol::TerminateResponse; pub use protocol::WriteParams; diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index bc69ec6105..691f807032 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -10,6 +10,7 @@ use codex_protocol::config_types::EnvironmentVariablePattern; use codex_protocol::config_types::ShellEnvironmentPolicy; use codex_protocol::shell_environment; use codex_utils_pty::ExecCommandSession; +use codex_utils_pty::ProcessSignal as PtyProcessSignal; use codex_utils_pty::TerminalSize; use tokio::sync::Mutex; use tokio::sync::Notify; @@ -33,8 +34,11 @@ use crate::protocol::ExecOutputStream; use crate::protocol::ExecParams; use crate::protocol::ExecResponse; use crate::protocol::ProcessOutputChunk; +use crate::protocol::ProcessSignal; use crate::protocol::ReadParams; use crate::protocol::ReadResponse; +use crate::protocol::SignalParams; +use crate::protocol::SignalResponse; use crate::protocol::TerminateParams; use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; @@ -272,7 +276,6 @@ impl LocalProcess { &self, params: ReadParams, ) -> Result { - let _process_id = params.process_id.clone(); let after_seq = params.after_seq.unwrap_or(0); let max_bytes = params.max_bytes.unwrap_or(usize::MAX); let wait = Duration::from_millis(params.wait_ms.unwrap_or(0)); @@ -351,7 +354,6 @@ impl LocalProcess { &self, params: WriteParams, ) -> Result { - let _process_id = params.process_id.clone(); let _input_bytes = params.chunk.0.len(); let writer_tx = { let process_map = self.inner.processes.lock().await; @@ -383,11 +385,33 @@ impl LocalProcess { }) } + pub(crate) async fn signal_process( + &self, + params: SignalParams, + ) -> Result { + { + let process_map = self.inner.processes.lock().await; + match process_map.get(¶ms.process_id) { + Some(ProcessEntry::Running(process)) => { + if process.exit_code.is_some() { + return Ok(SignalResponse {}); + } + process + .session + .signal(pty_process_signal(params.signal)) + .map_err(|err| internal_error(format!("failed to signal process: {err}")))? + } + Some(ProcessEntry::Starting) | None => {} + } + } + + Ok(SignalResponse {}) + } + pub(crate) async fn terminate_process( &self, params: TerminateParams, ) -> Result { - let _process_id = params.process_id.clone(); let running = { let process_map = self.inner.processes.lock().await; match process_map.get(¶ms.process_id) { @@ -483,6 +507,10 @@ impl ExecProcess for LocalExecProcess { self.backend.write(&self.process_id, chunk).await } + async fn signal(&self, signal: ProcessSignal) -> Result<(), ExecServerError> { + self.backend.signal(&self.process_id, signal).await + } + async fn terminate(&self) -> Result<(), ExecServerError> { self.backend.terminate(&self.process_id).await } @@ -519,6 +547,20 @@ impl LocalProcess { .map_err(map_handler_error) } + async fn signal( + &self, + process_id: &ProcessId, + signal: ProcessSignal, + ) -> Result<(), ExecServerError> { + self.signal_process(SignalParams { + process_id: process_id.clone(), + signal, + }) + .await + .map_err(map_handler_error)?; + Ok(()) + } + async fn terminate(&self, process_id: &ProcessId) -> Result<(), ExecServerError> { self.terminate_process(TerminateParams { process_id: process_id.clone(), @@ -529,6 +571,12 @@ impl LocalProcess { } } +fn pty_process_signal(signal: ProcessSignal) -> PtyProcessSignal { + match signal { + ProcessSignal::Interrupt => PtyProcessSignal::Interrupt, + } +} + fn map_handler_error(error: JSONRPCErrorError) -> ExecServerError { ExecServerError::Server { code: error.code, diff --git a/codex-rs/exec-server/src/process.rs b/codex-rs/exec-server/src/process.rs index cb6c832138..97f53aaefb 100644 --- a/codex-rs/exec-server/src/process.rs +++ b/codex-rs/exec-server/src/process.rs @@ -10,6 +10,7 @@ use crate::ExecServerError; use crate::ProcessId; use crate::protocol::ExecParams; use crate::protocol::ProcessOutputChunk; +use crate::protocol::ProcessSignal; use crate::protocol::ReadResponse; use crate::protocol::WriteResponse; @@ -178,6 +179,8 @@ pub trait ExecProcess: Send + Sync { async fn write(&self, chunk: Vec) -> Result; + async fn signal(&self, signal: ProcessSignal) -> Result<(), ExecServerError>; + async fn terminate(&self) -> Result<(), ExecServerError>; } diff --git a/codex-rs/exec-server/src/protocol.rs b/codex-rs/exec-server/src/protocol.rs index 6eeeeefdd7..b85bff3577 100644 --- a/codex-rs/exec-server/src/protocol.rs +++ b/codex-rs/exec-server/src/protocol.rs @@ -15,6 +15,7 @@ pub const INITIALIZED_METHOD: &str = "initialized"; pub const EXEC_METHOD: &str = "process/start"; pub const EXEC_READ_METHOD: &str = "process/read"; pub const EXEC_WRITE_METHOD: &str = "process/write"; +pub const EXEC_SIGNAL_METHOD: &str = "process/signal"; pub const EXEC_TERMINATE_METHOD: &str = "process/terminate"; pub const EXEC_OUTPUT_DELTA_METHOD: &str = "process/output"; pub const EXEC_EXITED_METHOD: &str = "process/exited"; @@ -166,6 +167,23 @@ pub struct WriteResponse { pub status: WriteStatus, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum ProcessSignal { + Interrupt, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SignalParams { + pub process_id: ProcessId, + pub signal: ProcessSignal, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SignalResponse {} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct TerminateParams { diff --git a/codex-rs/exec-server/src/remote_process.rs b/codex-rs/exec-server/src/remote_process.rs index d8d06735cd..c7d16d4d41 100644 --- a/codex-rs/exec-server/src/remote_process.rs +++ b/codex-rs/exec-server/src/remote_process.rs @@ -12,6 +12,7 @@ use crate::StartedExecProcess; use crate::client::LazyRemoteExecServerClient; use crate::client::Session; use crate::protocol::ExecParams; +use crate::protocol::ProcessSignal; use crate::protocol::ReadResponse; use crate::protocol::WriteResponse; @@ -76,6 +77,11 @@ impl ExecProcess for RemoteExecProcess { self.session.write(chunk).await } + async fn signal(&self, signal: ProcessSignal) -> Result<(), ExecServerError> { + trace!("exec process signal"); + self.session.signal(signal).await + } + async fn terminate(&self) -> Result<(), ExecServerError> { trace!("exec process terminate"); self.session.terminate().await diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index b8934705d2..395ce67730 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -42,6 +42,8 @@ use crate::protocol::InitializeParams; use crate::protocol::InitializeResponse; use crate::protocol::ReadParams; use crate::protocol::ReadResponse; +use crate::protocol::SignalParams; +use crate::protocol::SignalResponse; use crate::protocol::TerminateParams; use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; @@ -171,6 +173,14 @@ impl ExecServerHandler { session.process().exec_write(params).await } + pub(crate) async fn signal( + &self, + params: SignalParams, + ) -> Result { + let session = self.require_initialized_for("exec")?; + session.process().signal(params).await + } + pub(crate) async fn terminate( &self, params: TerminateParams, diff --git a/codex-rs/exec-server/src/server/process_handler.rs b/codex-rs/exec-server/src/server/process_handler.rs index 38fbace1cd..9fced9c166 100644 --- a/codex-rs/exec-server/src/server/process_handler.rs +++ b/codex-rs/exec-server/src/server/process_handler.rs @@ -5,6 +5,8 @@ use crate::protocol::ExecParams; use crate::protocol::ExecResponse; use crate::protocol::ReadParams; use crate::protocol::ReadResponse; +use crate::protocol::SignalParams; +use crate::protocol::SignalResponse; use crate::protocol::TerminateParams; use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; @@ -49,6 +51,13 @@ impl ProcessHandler { self.process.exec_write(params).await } + pub(crate) async fn signal( + &self, + params: SignalParams, + ) -> Result { + self.process.signal_process(params).await + } + pub(crate) async fn terminate( &self, params: TerminateParams, diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index 26d2876c63..1f652c1c4c 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use crate::protocol::ENVIRONMENT_INFO_METHOD; use crate::protocol::EXEC_METHOD; use crate::protocol::EXEC_READ_METHOD; +use crate::protocol::EXEC_SIGNAL_METHOD; use crate::protocol::EXEC_TERMINATE_METHOD; use crate::protocol::EXEC_WRITE_METHOD; use crate::protocol::ExecParams; @@ -32,6 +33,7 @@ use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; use crate::protocol::ReadParams; +use crate::protocol::SignalParams; use crate::protocol::TerminateParams; use crate::protocol::WriteParams; use crate::rpc::RpcRouter; @@ -77,6 +79,12 @@ pub(crate) fn build_router() -> RpcRouter { handler.exec_write(params).await }, ); + router.request( + EXEC_SIGNAL_METHOD, + |handler: Arc, params: SignalParams| async move { + handler.signal(params).await + }, + ); router.request( EXEC_TERMINATE_METHOD, |handler: Arc, params: TerminateParams| async move { diff --git a/codex-rs/exec-server/tests/exec_process.rs b/codex-rs/exec-server/tests/exec_process.rs index e1f330fc4e..af8e722c66 100644 --- a/codex-rs/exec-server/tests/exec_process.rs +++ b/codex-rs/exec-server/tests/exec_process.rs @@ -1,5 +1,3 @@ -#![cfg(unix)] - mod common; use std::sync::Arc; @@ -13,6 +11,7 @@ use codex_exec_server::ExecParams; use codex_exec_server::ExecProcess; use codex_exec_server::ExecProcessEvent; use codex_exec_server::ProcessId; +use codex_exec_server::ProcessSignal; use codex_exec_server::ReadResponse; use codex_exec_server::StartedExecProcess; use codex_exec_server::WriteStatus; @@ -505,6 +504,98 @@ async fn assert_exec_process_rejects_write_without_pipe_stdin(use_remote: bool) Ok(()) } +async fn assert_exec_process_signal_interrupts_process(use_remote: bool) -> Result<()> { + let context = create_process_context(use_remote).await?; + let process_id = "proc-signal".to_string(); + let session = context + .backend + .start(ExecParams { + process_id: process_id.clone().into(), + argv: vec![ + "/bin/sh".to_string(), + "-c".to_string(), + "trap 'printf \"signal:2\\n\"; exit 7' INT; printf 'ready\\n'; while :; do :; done".to_string(), + ], + cwd: std::env::current_dir()?, + env_policy: /*env_policy*/ None, + env: Default::default(), + tty: false, + pipe_stdin: false, + arg0: None, + }) + .await?; + assert_eq!(session.process.process_id().as_str(), process_id); + + let StartedExecProcess { process } = session; + let mut wake_rx = process.subscribe_wake(); + let mut ready_output = String::new(); + let mut after_seq = None; + loop { + let response = + read_process_until_change(Arc::clone(&process), &mut wake_rx, after_seq).await?; + for chunk in response.chunks { + ready_output.push_str(&String::from_utf8_lossy(&chunk.chunk.into_inner())); + after_seq = Some(chunk.seq); + } + if ready_output.contains("ready\n") { + break; + } + if response.closed { + anyhow::bail!("process closed before readiness marker: {ready_output:?}"); + } + after_seq = response.next_seq.checked_sub(1).or(after_seq); + } + + process.signal(ProcessSignal::Interrupt).await?; + let (output, exit_code, closed) = collect_process_output_from_reads(process, wake_rx).await?; + + assert!( + output.contains("signal:2"), + "expected signal handler output, got {output:?}" + ); + assert_eq!(exit_code, Some(7)); + assert!(closed); + Ok(()) +} + +async fn assert_exec_process_signal_reports_unsupported_on_windows(use_remote: bool) -> Result<()> { + let context = create_process_context(use_remote).await?; + let session = context + .backend + .start(ExecParams { + process_id: ProcessId::from("proc-windows-signal"), + argv: vec![ + "cmd".to_string(), + "/C".to_string(), + "echo ready && ping -n 30 127.0.0.1 >NUL".to_string(), + ], + cwd: std::env::current_dir()?, + env_policy: /*env_policy*/ None, + env: Default::default(), + tty: false, + pipe_stdin: false, + arg0: None, + }) + .await?; + + let err = match session.process.signal(ProcessSignal::Interrupt).await { + Ok(()) => anyhow::bail!("Windows non-TTY signal should report unsupported"), + Err(err) => err, + }; + let message = err.to_string(); + assert!( + message.contains("failed to signal process"), + "unexpected signal error: {message}" + ); + assert!( + message.contains("process interrupt is not supported by this process backend"), + "unexpected signal error: {message}" + ); + + session.process.terminate().await?; + Ok(()) +} + async fn assert_exec_process_preserves_queued_events_before_subscribe( use_remote: bool, ) -> Result<()> { @@ -539,6 +630,7 @@ async fn assert_exec_process_preserves_queued_events_before_subscribe( } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] // Serialize tests that launch a real exec-server process through the full CLI. #[serial_test::serial(remote_exec_server)] async fn remote_exec_process_reports_transport_disconnect() -> Result<()> { @@ -630,6 +722,7 @@ async fn remote_exec_process_reports_transport_disconnect() -> Result<()> { #[test_case(false ; "local")] #[test_case(true ; "remote")] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] // Serialize tests that launch a real exec-server process through the full CLI. #[serial_test::serial(remote_exec_server)] @@ -639,6 +732,7 @@ async fn exec_process_starts_and_exits(use_remote: bool) -> Result<()> { #[test_case(false ; "local")] #[test_case(true ; "remote")] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] // Serialize tests that launch a real exec-server process through the full CLI. #[serial_test::serial(remote_exec_server)] @@ -648,6 +742,7 @@ async fn exec_process_streams_output(use_remote: bool) -> Result<()> { #[test_case(false ; "local")] #[test_case(true ; "remote")] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] // Serialize tests that launch a real exec-server process through the full CLI. #[serial_test::serial(remote_exec_server)] @@ -657,6 +752,7 @@ async fn exec_process_pushes_events(use_remote: bool) -> Result<()> { #[test_case(false ; "local")] #[test_case(true ; "remote")] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] // Serialize tests that launch a real exec-server process through the full CLI. #[serial_test::serial(remote_exec_server)] @@ -666,6 +762,7 @@ async fn exec_process_replays_events_after_close(use_remote: bool) -> Result<()> #[test_case(false ; "local")] #[test_case(true ; "remote")] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] // Serialize tests that launch a real exec-server process through the full CLI. #[serial_test::serial(remote_exec_server)] @@ -677,6 +774,7 @@ async fn exec_process_retains_output_after_exit_until_streams_close( #[test_case(false ; "local")] #[test_case(true ; "remote")] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] // Serialize tests that launch a real exec-server process through the full CLI. #[serial_test::serial(remote_exec_server)] @@ -686,6 +784,7 @@ async fn exec_process_write_then_read(use_remote: bool) -> Result<()> { #[test_case(false ; "local")] #[test_case(true ; "remote")] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] // Serialize tests that launch a real exec-server process through the full CLI. #[serial_test::serial(remote_exec_server)] @@ -695,6 +794,7 @@ async fn exec_process_write_then_read_without_tty(use_remote: bool) -> Result<() #[test_case(false ; "local")] #[test_case(true ; "remote")] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] // Serialize tests that launch a real exec-server process through the full CLI. #[serial_test::serial(remote_exec_server)] @@ -704,6 +804,27 @@ async fn exec_process_rejects_write_without_pipe_stdin(use_remote: bool) -> Resu #[test_case(false ; "local")] #[test_case(true ; "remote")] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +// Serialize tests that launch a real exec-server process through the full CLI. +#[serial_test::serial(remote_exec_server)] +async fn exec_process_signal_interrupts_process(use_remote: bool) -> Result<()> { + assert_exec_process_signal_interrupts_process(use_remote).await +} + +#[test_case(false ; "local")] +#[test_case(true ; "remote")] +#[cfg_attr(not(windows), ignore = "Windows-only exec-server process test")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +// Serialize tests that launch a real exec-server process through the full CLI. +#[serial_test::serial(remote_exec_server)] +async fn exec_process_signal_reports_unsupported_on_windows(use_remote: bool) -> Result<()> { + assert_exec_process_signal_reports_unsupported_on_windows(use_remote).await +} + +#[test_case(false ; "local")] +#[test_case(true ; "remote")] +#[cfg_attr(not(unix), ignore = "Unix-only exec-server process test")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] // Serialize tests that launch a real exec-server process through the full CLI. #[serial_test::serial(remote_exec_server)] diff --git a/codex-rs/ext/extension-api/Cargo.toml b/codex-rs/ext/extension-api/Cargo.toml index 2e155b430e..25ed8d0d71 100644 --- a/codex-rs/ext/extension-api/Cargo.toml +++ b/codex-rs/ext/extension-api/Cargo.toml @@ -14,7 +14,6 @@ doctest = false workspace = true [dependencies] -async-trait = { workspace = true } codex-config = { workspace = true } codex-context-fragments = { workspace = true } codex-protocol = { workspace = true } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 0403294371..e93088b5d0 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -1,4 +1,5 @@ use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use codex_context_fragments::ContextualUserFragment; @@ -36,6 +37,9 @@ pub use turn_lifecycle::TurnErrorInput; pub use turn_lifecycle::TurnStartInput; pub use turn_lifecycle::TurnStopInput; +/// Boxed, sendable future returned by asynchronous extension contributors. +pub type ExtensionFuture<'a, T> = Pin + Send + 'a>>; + /// Extension contribution that resolves runtime MCP servers from host config. /// /// Contributors run in registration order. Later contributions for the same @@ -43,9 +47,8 @@ pub use turn_lifecycle::TurnStopInput; /// own and must apply any source-specific policy before returning a server. /// Plugin-owned servers and their provenance continue to be resolved by the /// plugin manager until that ownership moves into an extension explicitly. -#[async_trait::async_trait] pub trait McpServerContributor: Send + Sync { - async fn contribute(&self, config: &C) -> Vec; + fn contribute<'a>(&'a self, config: &'a C) -> ExtensionFuture<'a, Vec>; } /// Extension contribution that adds prompt fragments during prompt assembly. @@ -54,7 +57,7 @@ pub trait ContextContributor: Send + Sync { &'a self, session_store: &'a ExtensionData, thread_store: &'a ExtensionData, - ) -> std::pin::Pin> + Send + 'a>>; + ) -> ExtensionFuture<'a, Vec>; } /// Contributor for host-owned thread lifecycle gates. @@ -62,24 +65,43 @@ pub trait ContextContributor: Send + Sync { /// Implementations should use these callbacks to seed, rehydrate, or flush /// extension-private thread state. Heavy dependencies belong on the extension /// value created by the host, not in these inputs. -#[async_trait::async_trait] pub trait ThreadLifecycleContributor: Send + Sync { /// Called after thread-scoped extension stores are created, before later /// contributors can read from them. - async fn on_thread_start(&self, _input: ThreadStartInput<'_, C>) {} + fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let _self = self; + let _input = input; + }) + } /// Called after the host constructs a runtime from persisted history. - async fn on_thread_resume(&self, _input: ThreadResumeInput<'_>) {} + fn on_thread_resume<'a>(&'a self, input: ThreadResumeInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let _self = self; + let _input = input; + }) + } /// Called after the host has drained immediately pending thread work. /// /// Implementations may use host capabilities captured by the extension to /// submit follow-up input. The host remains responsible for deciding /// whether that input starts a turn, is queued, or is ignored. - async fn on_thread_idle(&self, _input: ThreadIdleInput<'_>) {} + fn on_thread_idle<'a>(&'a self, input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let _self = self; + let _input = input; + }) + } /// Called before the host drops the thread runtime and thread-scoped store. - async fn on_thread_stop(&self, _input: ThreadStopInput<'_>) {} + fn on_thread_stop<'a>(&'a self, input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let _self = self; + let _input = input; + }) + } } /// Contributor for host-owned turn lifecycle gates. @@ -87,20 +109,39 @@ pub trait ThreadLifecycleContributor: Send + Sync { /// Implementations should use these callbacks to seed, observe, or clear /// extension-private turn state. The host exposes stable identifiers and /// extension stores instead of core runtime objects. -#[async_trait::async_trait] pub trait TurnLifecycleContributor: Send + Sync { /// Called after turn-scoped extension stores are created, before the task /// for the turn starts running. - async fn on_turn_start(&self, _input: TurnStartInput<'_>) {} + fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let _self = self; + let _input = input; + }) + } /// Called before the host drops the completed turn runtime and turn store. - async fn on_turn_stop(&self, _input: TurnStopInput<'_>) {} + fn on_turn_stop<'a>(&'a self, input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let _self = self; + let _input = input; + }) + } /// Called after the host aborts a running turn. - async fn on_turn_abort(&self, _input: TurnAbortInput<'_>) {} + fn on_turn_abort<'a>(&'a self, input: TurnAbortInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let _self = self; + let _input = input; + }) + } /// Called when the host observes an error for a running turn. - async fn on_turn_error(&self, _input: TurnErrorInput<'_>) {} + fn on_turn_error<'a>(&'a self, input: TurnErrorInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let _self = self; + let _input = input; + }) + } } /// Extension contribution that can add turn-local model input. @@ -109,16 +150,15 @@ pub trait TurnLifecycleContributor: Send + Sync { /// must preserve authority boundaries for external resources. Expensive or /// host-specific dependencies belong on the extension value installed by the /// host, not in this input. -#[async_trait::async_trait] pub trait TurnInputContributor: Send + Sync { /// Returns additional contextual fragments for one submitted turn. - async fn contribute( - &self, + fn contribute<'a>( + &'a self, input: TurnInputContext, - session_store: &ExtensionData, - thread_store: &ExtensionData, - turn_store: &ExtensionData, - ) -> Vec>; + session_store: &'a ExtensionData, + thread_store: &'a ExtensionData, + turn_store: &'a ExtensionData, + ) -> ExtensionFuture<'a, Vec>>; } /// Contributor for host-owned configuration changes. @@ -142,16 +182,19 @@ pub trait ConfigContributor: Send + Sync { /// Implementations should keep this callback cheap. The host calls it after /// updating cached token usage and before emitting the corresponding client /// token-count notification. -#[async_trait::async_trait] pub trait TokenUsageContributor: Send + Sync { /// Called each time the host records token usage from a model response. - async fn on_token_usage( - &self, - _session_store: &ExtensionData, - _thread_store: &ExtensionData, - _turn_store: &ExtensionData, - _token_usage: &TokenUsageInfo, - ) { + fn on_token_usage<'a>( + &'a self, + _session_store: &'a ExtensionData, + _thread_store: &'a ExtensionData, + _turn_store: &'a ExtensionData, + _token_usage: &'a TokenUsageInfo, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let _self = self; + let _inputs = (_session_store, _thread_store, _turn_store, _token_usage); + }) } } @@ -183,14 +226,13 @@ pub trait ToolLifecycleContributor: Send + Sync { } /// Extension contribution that can claim rendered approval-review prompts. -#[async_trait::async_trait] pub trait ApprovalReviewContributor: Send + Sync { - async fn contribute( - &self, - session_store: &ExtensionData, - thread_store: &ExtensionData, - prompt: &str, - ) -> Option; + fn contribute<'a>( + &'a self, + session_store: &'a ExtensionData, + thread_store: &'a ExtensionData, + prompt: &'a str, + ) -> ExtensionFuture<'a, Option>; } /// Ordered post-processing contribution for one parsed turn item. @@ -198,12 +240,11 @@ pub trait ApprovalReviewContributor: Send + Sync { /// 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 { - async fn contribute( - &self, - thread_store: &ExtensionData, - turn_store: &ExtensionData, - item: &mut TurnItem, - ) -> Result<(), String>; + fn contribute<'a>( + &'a self, + thread_store: &'a ExtensionData, + turn_store: &'a ExtensionData, + item: &'a mut TurnItem, + ) -> ExtensionFuture<'a, Result<(), String>>; } diff --git a/codex-rs/ext/extension-api/src/contributors/mcp.rs b/codex-rs/ext/extension-api/src/contributors/mcp.rs index a55715692b..7fdedcdda7 100644 --- a/codex-rs/ext/extension-api/src/contributors/mcp.rs +++ b/codex-rs/ext/extension-api/src/contributors/mcp.rs @@ -11,12 +11,3 @@ pub enum McpServerContribution { /// Removes a named MCP server. Remove { name: String }, } - -impl McpServerContribution { - /// Returns the stable server name owned by this contribution. - pub fn name(&self) -> &str { - match self { - Self::Set { name, .. } | Self::Remove { name } => name, - } - } -} diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 8b269f31d1..52f15c894c 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -20,6 +20,7 @@ pub use codex_tools::NoopTurnItemEmitter; pub use codex_tools::ResponsesApiTool; pub use codex_tools::ToolCall; pub use codex_tools::ToolExecutor; +pub use codex_tools::ToolExecutorFuture; pub use codex_tools::ToolName; pub use codex_tools::ToolOutput; pub use codex_tools::ToolPayload; @@ -31,6 +32,7 @@ pub use codex_tools::parse_tool_input_schema_without_compaction; pub use contributors::ApprovalReviewContributor; pub use contributors::ConfigContributor; pub use contributors::ContextContributor; +pub use contributors::ExtensionFuture; pub use contributors::McpServerContribution; pub use contributors::McpServerContributor; pub use contributors::PromptFragment; diff --git a/codex-rs/ext/extension-api/tests/registry.rs b/codex-rs/ext/extension-api/tests/registry.rs index d26f1010cb..42cb18af62 100644 --- a/codex-rs/ext/extension-api/tests/registry.rs +++ b/codex-rs/ext/extension-api/tests/registry.rs @@ -1,5 +1,3 @@ -use std::future::Future; -use std::pin::Pin; use std::sync::Arc; use std::sync::Mutex; @@ -9,6 +7,7 @@ use codex_extension_api::ContextContributor; use codex_extension_api::ContextualUserFragment; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::PromptFragment; use codex_extension_api::ThreadLifecycleContributor; @@ -37,32 +36,32 @@ impl ContextContributor for AllContributors { &'a self, _session_store: &'a ExtensionData, _thread_store: &'a ExtensionData, - ) -> Pin> + Send + 'a>> { + ) -> ExtensionFuture<'a, Vec> { Box::pin(std::future::ready(Vec::new())) } } -#[async_trait::async_trait] impl ThreadLifecycleContributor<()> for AllContributors {} -#[async_trait::async_trait] impl TurnLifecycleContributor for AllContributors {} impl ConfigContributor<()> for AllContributors {} -#[async_trait::async_trait] impl TokenUsageContributor for AllContributors {} -#[async_trait::async_trait] impl TurnInputContributor for AllContributors { - async fn contribute( - &self, - _input: TurnInputContext, - _session_store: &ExtensionData, - _thread_store: &ExtensionData, - _turn_store: &ExtensionData, - ) -> Vec> { - Vec::new() + fn contribute<'a>( + &'a self, + input: TurnInputContext, + _session_store: &'a ExtensionData, + _thread_store: &'a ExtensionData, + _turn_store: &'a ExtensionData, + ) -> ExtensionFuture<'a, Vec>> { + Box::pin(async move { + let _self = self; + let _input = input; + Vec::new() + }) } } @@ -78,27 +77,31 @@ impl ToolContributor for AllContributors { impl ToolLifecycleContributor for AllContributors {} -#[async_trait::async_trait] impl TurnItemContributor for AllContributors { - async fn contribute( - &self, - _thread_store: &ExtensionData, - _turn_store: &ExtensionData, - _item: &mut TurnItem, - ) -> Result<(), String> { - Ok(()) + fn contribute<'a>( + &'a self, + _thread_store: &'a ExtensionData, + _turn_store: &'a ExtensionData, + _item: &'a mut TurnItem, + ) -> ExtensionFuture<'a, Result<(), String>> { + Box::pin(async move { + let _self = self; + Ok(()) + }) } } -#[async_trait::async_trait] impl ApprovalReviewContributor for AllContributors { - async fn contribute( - &self, - _session_store: &ExtensionData, - _thread_store: &ExtensionData, - _prompt: &str, - ) -> Option { - Some(ReviewDecision::ApprovedForSession) + fn contribute<'a>( + &'a self, + _session_store: &'a ExtensionData, + _thread_store: &'a ExtensionData, + _prompt: &'a str, + ) -> ExtensionFuture<'a, Option> { + Box::pin(async move { + let _self = self; + Some(ReviewDecision::ApprovedForSession) + }) } } @@ -146,7 +149,7 @@ impl ContextContributor for NamedContextContributor { &'a self, _session_store: &'a ExtensionData, _thread_store: &'a ExtensionData, - ) -> Pin> + Send + 'a>> { + ) -> ExtensionFuture<'a, Vec> { Box::pin(std::future::ready(vec![PromptFragment::developer_policy( self.0, )])) @@ -158,19 +161,20 @@ struct RecordingTurnItemContributor { calls: Arc>>, } -#[async_trait::async_trait] impl TurnItemContributor for RecordingTurnItemContributor { - async fn contribute( - &self, - _thread_store: &ExtensionData, - _turn_store: &ExtensionData, - _item: &mut TurnItem, - ) -> Result<(), String> { - self.calls - .lock() - .unwrap_or_else(|error| panic!("turn item calls lock poisoned: {error}")) - .push(self.name); - Ok(()) + fn contribute<'a>( + &'a self, + _thread_store: &'a ExtensionData, + _turn_store: &'a ExtensionData, + _item: &'a mut TurnItem, + ) -> ExtensionFuture<'a, Result<(), String>> { + Box::pin(async move { + self.calls + .lock() + .unwrap_or_else(|error| panic!("turn item calls lock poisoned: {error}")) + .push(self.name); + Ok(()) + }) } } @@ -236,24 +240,25 @@ struct RecordingApprovalContributor { calls: Arc>>, } -#[async_trait::async_trait] impl ApprovalReviewContributor for RecordingApprovalContributor { - async fn contribute( - &self, - session_store: &ExtensionData, - thread_store: &ExtensionData, - prompt: &str, - ) -> Option { - self.calls - .lock() - .unwrap_or_else(|error| panic!("approval calls lock poisoned: {error}")) - .push(ApprovalCall { - contributor: self.name, - session_id: session_store.level_id().to_string(), - thread_id: thread_store.level_id().to_string(), - prompt: prompt.to_string(), - }); - self.decision.clone() + fn contribute<'a>( + &'a self, + session_store: &'a ExtensionData, + thread_store: &'a ExtensionData, + prompt: &'a str, + ) -> ExtensionFuture<'a, Option> { + Box::pin(async move { + self.calls + .lock() + .unwrap_or_else(|error| panic!("approval calls lock poisoned: {error}")) + .push(ApprovalCall { + contributor: self.name, + session_id: session_store.level_id().to_string(), + thread_id: thread_store.level_id().to_string(), + prompt: prompt.to_string(), + }); + self.decision.clone() + }) } } diff --git a/codex-rs/ext/goal/Cargo.toml b/codex-rs/ext/goal/Cargo.toml index 6116383243..4d75ba2f26 100644 --- a/codex-rs/ext/goal/Cargo.toml +++ b/codex-rs/ext/goal/Cargo.toml @@ -14,7 +14,7 @@ doctest = false workspace = true [dependencies] -async-trait = { workspace = true } +codex-analytics = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-otel = { workspace = true } diff --git a/codex-rs/ext/goal/src/analytics.rs b/codex-rs/ext/goal/src/analytics.rs new file mode 100644 index 0000000000..82d34962d1 --- /dev/null +++ b/codex-rs/ext/goal/src/analytics.rs @@ -0,0 +1,77 @@ +use codex_analytics::AnalyticsEventsClient; +use codex_analytics::CodexGoalEvent; +use codex_analytics::GoalEventKind; + +#[derive(Clone)] +pub(crate) struct GoalAnalytics { + client: AnalyticsEventsClient, +} + +pub(crate) enum GoalEventAttribution<'a> { + Turn(&'a str), + NoTurn, +} + +impl GoalAnalytics { + pub(crate) fn new(client: AnalyticsEventsClient) -> Self { + Self { client } + } + + pub(crate) fn created( + &self, + goal: &codex_state::ThreadGoal, + attribution: GoalEventAttribution<'_>, + ) { + self.track(goal, attribution, GoalEventKind::Created); + } + + pub(crate) fn usage_accounted( + &self, + goal: &codex_state::ThreadGoal, + attribution: GoalEventAttribution<'_>, + ) { + self.track(goal, attribution, GoalEventKind::UsageAccounted); + } + + pub(crate) fn status_changed( + &self, + goal: &codex_state::ThreadGoal, + previous_status: Option, + attribution: GoalEventAttribution<'_>, + ) { + if previous_status.is_some_and(|status| status != goal.status) { + self.track(goal, attribution, GoalEventKind::StatusChanged); + } + } + + pub(crate) fn cleared(&self, goal: &codex_state::ThreadGoal) { + self.track(goal, GoalEventAttribution::NoTurn, GoalEventKind::Cleared); + } + + fn track( + &self, + goal: &codex_state::ThreadGoal, + attribution: GoalEventAttribution<'_>, + event_kind: GoalEventKind, + ) { + let (cumulative_tokens_accounted, cumulative_time_accounted_seconds) = match event_kind { + GoalEventKind::UsageAccounted => (Some(goal.tokens_used), Some(goal.time_used_seconds)), + GoalEventKind::Created | GoalEventKind::StatusChanged | GoalEventKind::Cleared => { + (None, None) + } + }; + self.client.track_goal_event(CodexGoalEvent { + thread_id: goal.thread_id.to_string(), + turn_id: match attribution { + GoalEventAttribution::Turn(turn_id) => Some(turn_id.to_string()), + GoalEventAttribution::NoTurn => None, + }, + goal_id: goal.goal_id.clone(), + event_kind, + goal_status: goal.status, + has_token_budget: goal.token_budget.is_some(), + cumulative_tokens_accounted, + cumulative_time_accounted_seconds, + }); + } +} diff --git a/codex-rs/ext/goal/src/api.rs b/codex-rs/ext/goal/src/api.rs index 5123e6a5ce..a4dc93409e 100644 --- a/codex-rs/ext/goal/src/api.rs +++ b/codex-rs/ext/goal/src/api.rs @@ -259,19 +259,19 @@ impl GoalService { tracing::warn!("failed to prepare external goal mutation: {err}"); } - let cleared = state_db + let cleared_goal = state_db .thread_goals() .delete_thread_goal(thread_id) .await .map_err(|err| { GoalServiceError::Internal(format!("failed to clear thread goal: {err}")) })?; + let cleared = cleared_goal.is_some(); drop(goal_state_permit); drop(runtime); - if cleared - && let Some(runtime) = self.runtime_for_thread(thread_id) - && let Err(err) = runtime.apply_external_goal_clear().await + if let (Some(runtime), Some(goal)) = (self.runtime_for_thread(thread_id), cleared_goal) + && let Err(err) = runtime.apply_external_goal_clear(goal).await { tracing::warn!("failed to apply external goal clear runtime effects: {err}"); } diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 78ae0b5b44..8e5b20e7f5 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -1,11 +1,12 @@ use std::sync::Arc; use std::sync::Weak; -use async_trait::async_trait; +use codex_analytics::AnalyticsEventsClient; use codex_core::ThreadManager; use codex_extension_api::ConfigContributor; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ThreadIdleInput; use codex_extension_api::ThreadLifecycleContributor; @@ -33,6 +34,7 @@ use codex_protocol::protocol::TokenUsageInfo; use crate::accounting::BudgetLimitedGoalDisposition; use crate::accounting::GoalAccountingState; +use crate::analytics::GoalAnalytics; use crate::api::GoalService; use crate::events::GoalEventEmitter; use crate::metrics::GoalMetrics; @@ -57,6 +59,7 @@ impl GoalExtensionConfig { #[derive(Clone)] pub struct GoalExtension { state_dbs: Arc, + analytics: GoalAnalytics, event_emitter: GoalEventEmitter, metrics: GoalMetrics, thread_manager: Weak, @@ -73,6 +76,7 @@ impl std::fmt::Debug for GoalExtension { impl GoalExtension { pub(crate) fn new_with_host_capabilities( state_dbs: Arc, + analytics_events_client: AnalyticsEventsClient, event_sink: Arc, metrics_client: Option, thread_manager: Weak, @@ -81,6 +85,7 @@ impl GoalExtension { ) -> Self { Self { state_dbs, + analytics: GoalAnalytics::new(analytics_events_client), event_emitter: GoalEventEmitter::new(event_sink), metrics: GoalMetrics::new(metrics_client), thread_manager, @@ -90,75 +95,83 @@ impl GoalExtension { } } -#[async_trait] impl ThreadLifecycleContributor for GoalExtension where C: Send + Sync + 'static, { - async fn on_thread_start(&self, input: ThreadStartInput<'_, C>) { - let enabled = (self.goals_enabled)(input.config); - let tools_available_for_thread = input.persistent_thread_state_available - && !matches!( - input.session_source, - SessionSource::SubAgent(SubAgentSource::Review) - ); - input - .thread_store - .insert(GoalExtensionConfig::from_enabled(enabled)); - let accounting_state = input - .thread_store - .get_or_init::(GoalAccountingState::default); - let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) else { - return; - }; - let runtime = input.thread_store.get_or_init::(|| { - GoalRuntimeHandle::new( - thread_id, - Arc::clone(&self.state_dbs), - self.event_emitter.clone(), - self.metrics.clone(), - self.thread_manager.clone(), - accounting_state, - GoalRuntimeConfig { - enabled, - tools_available_for_thread, - }, - ) - }); - runtime.set_enabled(enabled); - self.goal_service.register_runtime(&runtime); + fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let enabled = (self.goals_enabled)(input.config); + let tools_available_for_thread = input.persistent_thread_state_available + && !matches!( + input.session_source, + SessionSource::SubAgent(SubAgentSource::Review) + ); + input + .thread_store + .insert(GoalExtensionConfig::from_enabled(enabled)); + let accounting_state = input + .thread_store + .get_or_init::(GoalAccountingState::default); + let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) else { + return; + }; + let runtime = input.thread_store.get_or_init::(|| { + GoalRuntimeHandle::new( + thread_id, + Arc::clone(&self.state_dbs), + self.event_emitter.clone(), + self.metrics.clone(), + self.thread_manager.clone(), + accounting_state, + GoalRuntimeConfig { + analytics: self.analytics.clone(), + enabled, + tools_available_for_thread, + }, + ) + }); + runtime.set_enabled(enabled); + self.goal_service.register_runtime(&runtime); + }) } - async fn on_thread_resume(&self, input: ThreadResumeInput<'_>) { - let Some(runtime) = goal_runtime_handle(input.thread_store) else { - return; - }; + fn on_thread_resume<'a>(&'a self, input: ThreadResumeInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let Some(runtime) = goal_runtime_handle(input.thread_store) else { + return; + }; - if let Err(err) = runtime.restore_after_resume().await { - tracing::warn!( - "failed to restore goal runtime after thread resume for {}: {err}", - runtime.thread_id() - ); - } + if let Err(err) = runtime.restore_after_resume().await { + tracing::warn!( + "failed to restore goal runtime after thread resume for {}: {err}", + runtime.thread_id() + ); + } + }) } - async fn on_thread_idle(&self, input: ThreadIdleInput<'_>) { - let Some(runtime) = goal_runtime_handle(input.thread_store) else { - return; - }; + fn on_thread_idle<'a>(&'a self, input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let Some(runtime) = goal_runtime_handle(input.thread_store) else { + return; + }; - if let Err(err) = runtime.continue_if_idle().await { - tracing::warn!( - "failed to continue active goal for idle thread {}: {err}", - runtime.thread_id() - ); - } + if let Err(err) = runtime.continue_if_idle().await { + tracing::warn!( + "failed to continue active goal for idle thread {}: {err}", + runtime.thread_id() + ); + } + }) } - async fn on_thread_stop(&self, input: ThreadStopInput<'_>) { - if let Some(runtime) = goal_runtime_handle(input.thread_store) { - self.goal_service.unregister_runtime(&runtime); - } + fn on_thread_stop<'a>(&'a self, input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + if let Some(runtime) = goal_runtime_handle(input.thread_store) { + self.goal_service.unregister_runtime(&runtime); + } + }) } } @@ -181,153 +194,161 @@ where } } -#[async_trait] impl TurnLifecycleContributor for GoalExtension where C: Send + Sync + 'static, { - async fn on_turn_start(&self, input: TurnStartInput<'_>) { - let Some(runtime) = goal_runtime_handle(input.thread_store) else { - return; - }; - if !runtime.is_enabled() { - return; - } + fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let Some(runtime) = goal_runtime_handle(input.thread_store) else { + return; + }; + if !runtime.is_enabled() { + return; + } - let accounting = runtime.accounting_state(); - accounting.start_turn( - input.turn_id, - input.collaboration_mode.mode, - input.token_usage_at_turn_start, - ); - if matches!( - input.collaboration_mode.mode, - codex_protocol::config_types::ModeKind::Plan - ) { - accounting.clear_current_turn_goal(); - return; - } - let Ok(goal) = self - .state_dbs - .thread_goals() - .get_thread_goal(runtime.thread_id()) - .await - else { - return; - }; - if let Some(goal) = goal - && matches!( - goal.status, - codex_state::ThreadGoalStatus::Active - | codex_state::ThreadGoalStatus::BudgetLimited - ) - { - accounting.mark_turn_goal_active(input.turn_id, goal.goal_id); - } + let accounting = runtime.accounting_state(); + accounting.start_turn( + input.turn_id, + input.collaboration_mode.mode, + input.token_usage_at_turn_start, + ); + if matches!( + input.collaboration_mode.mode, + codex_protocol::config_types::ModeKind::Plan + ) { + accounting.clear_current_turn_goal(); + return; + } + let Ok(goal) = self + .state_dbs + .thread_goals() + .get_thread_goal(runtime.thread_id()) + .await + else { + return; + }; + if let Some(goal) = goal + && matches!( + goal.status, + codex_state::ThreadGoalStatus::Active + | codex_state::ThreadGoalStatus::BudgetLimited + ) + { + accounting.mark_turn_goal_active(input.turn_id, goal.goal_id); + } + }) } - async fn on_turn_stop(&self, input: TurnStopInput<'_>) { - let Some(runtime) = goal_runtime_handle(input.thread_store) else { - return; - }; - if !runtime.is_enabled() { - return; - } + fn on_turn_stop<'a>(&'a self, input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let Some(runtime) = goal_runtime_handle(input.thread_store) else { + return; + }; + if !runtime.is_enabled() { + return; + } - let turn_id = input.turn_store.level_id(); - if let Err(err) = runtime - .account_active_goal_progress( - turn_id, - &format!("{turn_id}:turn-stop"), - codex_state::GoalAccountingMode::ActiveOnly, - BudgetLimitedGoalDisposition::ClearActive, - ) - .await - { - tracing::warn!( - "failed to account active goal progress at turn stop for {turn_id}: {err}" - ); - return; - } - runtime.accounting_state().finish_turn(turn_id); + let turn_id = input.turn_store.level_id(); + if let Err(err) = runtime + .account_active_goal_progress( + turn_id, + &format!("{turn_id}:turn-stop"), + codex_state::GoalAccountingMode::ActiveOnly, + BudgetLimitedGoalDisposition::ClearActive, + ) + .await + { + tracing::warn!( + "failed to account active goal progress at turn stop for {turn_id}: {err}" + ); + return; + } + runtime.accounting_state().finish_turn(turn_id); + }) } - async fn on_turn_abort(&self, input: TurnAbortInput<'_>) { - let Some(runtime) = goal_runtime_handle(input.thread_store) else { - return; - }; - if !runtime.is_enabled() { - return; - } + fn on_turn_abort<'a>(&'a self, input: TurnAbortInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let Some(runtime) = goal_runtime_handle(input.thread_store) else { + return; + }; + if !runtime.is_enabled() { + return; + } - let turn_id = input.turn_store.level_id(); - if let Err(err) = runtime - .account_active_goal_progress( - turn_id, - &format!("{turn_id}:turn-abort"), - codex_state::GoalAccountingMode::ActiveOnly, - BudgetLimitedGoalDisposition::ClearActive, - ) - .await - { - tracing::warn!( - "failed to account active goal progress after turn abort for {turn_id}: {err}" - ); - return; - } - runtime.accounting_state().finish_turn(turn_id); + let turn_id = input.turn_store.level_id(); + if let Err(err) = runtime + .account_active_goal_progress( + turn_id, + &format!("{turn_id}:turn-abort"), + codex_state::GoalAccountingMode::ActiveOnly, + BudgetLimitedGoalDisposition::ClearActive, + ) + .await + { + tracing::warn!( + "failed to account active goal progress after turn abort for {turn_id}: {err}" + ); + return; + } + runtime.accounting_state().finish_turn(turn_id); + }) } - async fn on_turn_error(&self, input: TurnErrorInput<'_>) { - let Some(runtime) = goal_runtime_handle(input.thread_store) else { - return; - }; + fn on_turn_error<'a>(&'a self, input: TurnErrorInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let Some(runtime) = goal_runtime_handle(input.thread_store) else { + return; + }; - let reason = match input.error { - CodexErrorInfo::UsageLimitExceeded => ActiveGoalStopReason::UsageLimit, - // The turn has ended because the error was non-retryable or its - // retries were exhausted. Block the goal to prevent automatic - // continuation from looping and consuming tokens, as can happen - // with compaction errors. - _ => ActiveGoalStopReason::TurnError, - }; - if let Err(err) = runtime - .stop_active_goal_for_turn(input.turn_id, reason) - .await - { - tracing::warn!( - error = ?input.error, - "failed to stop active goal after turn error: {err}" - ); - } + let reason = match input.error { + CodexErrorInfo::UsageLimitExceeded => ActiveGoalStopReason::UsageLimit, + // The turn has ended because the error was non-retryable or its + // retries were exhausted. Block the goal to prevent automatic + // continuation from looping and consuming tokens, as can happen + // with compaction errors. + _ => ActiveGoalStopReason::TurnError, + }; + if let Err(err) = runtime + .stop_active_goal_for_turn(input.turn_id, reason) + .await + { + tracing::warn!( + error = ?input.error, + "failed to stop active goal after turn error: {err}" + ); + } + }) } } -#[async_trait] impl TokenUsageContributor for GoalExtension where C: Send + Sync + 'static, { - async fn on_token_usage( - &self, - _session_store: &ExtensionData, - thread_store: &ExtensionData, - turn_store: &ExtensionData, - token_usage: &TokenUsageInfo, - ) { - let Some(runtime) = goal_runtime_handle(thread_store) else { - return; - }; - if !runtime.is_enabled() { - return; - } + fn on_token_usage<'a>( + &'a self, + _session_store: &'a ExtensionData, + thread_store: &'a ExtensionData, + turn_store: &'a ExtensionData, + token_usage: &'a TokenUsageInfo, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let Some(runtime) = goal_runtime_handle(thread_store) else { + return; + }; + if !runtime.is_enabled() { + return; + } - let Some(_recorded) = runtime - .accounting_state() - .record_token_usage(turn_store.level_id(), &token_usage.total_token_usage) - else { - return; - }; + let Some(_recorded) = runtime + .accounting_state() + .record_token_usage(turn_store.level_id(), &token_usage.total_token_usage) + else { + return; + }; + }) } } @@ -403,6 +424,7 @@ where runtime.thread_id(), Arc::clone(&self.state_dbs), runtime.accounting_state(), + self.analytics.clone(), self.event_emitter.clone(), self.metrics.clone(), )), @@ -410,6 +432,7 @@ where runtime.thread_id(), Arc::clone(&self.state_dbs), runtime.accounting_state(), + self.analytics.clone(), self.event_emitter.clone(), self.metrics.clone(), )), @@ -417,6 +440,7 @@ where runtime.thread_id(), Arc::clone(&self.state_dbs), runtime.accounting_state(), + self.analytics.clone(), self.event_emitter.clone(), self.metrics.clone(), )), @@ -427,6 +451,7 @@ where pub fn install_with_backend( registry: &mut ExtensionRegistryBuilder, state_dbs: Arc, + analytics_events_client: AnalyticsEventsClient, metrics_client: Option, thread_manager: Weak, goal_service: Arc, @@ -436,6 +461,7 @@ pub fn install_with_backend( { let extension = Arc::new(GoalExtension::new_with_host_capabilities( state_dbs, + analytics_events_client, registry.event_sink(), metrics_client, thread_manager, diff --git a/codex-rs/ext/goal/src/lib.rs b/codex-rs/ext/goal/src/lib.rs index bb640c6ce8..ccd091affc 100644 --- a/codex-rs/ext/goal/src/lib.rs +++ b/codex-rs/ext/goal/src/lib.rs @@ -1,6 +1,7 @@ //! Extension crate for the `/goal` feature. mod accounting; +mod analytics; mod api; mod events; mod extension; diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 2641dfb949..2f2202cbb8 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -10,6 +10,8 @@ use codex_protocol::protocol::ThreadGoal; use crate::accounting::BudgetLimitedGoalDisposition; use crate::accounting::GoalAccountingState; +use crate::analytics::GoalAnalytics; +use crate::analytics::GoalEventAttribution; use crate::events::GoalEventEmitter; use crate::metrics::GoalMetrics; use crate::steering::continuation_steering_item; @@ -24,6 +26,7 @@ pub struct GoalRuntimeHandle { } pub(crate) struct GoalRuntimeConfig { + pub(crate) analytics: GoalAnalytics, pub(crate) enabled: bool, pub(crate) tools_available_for_thread: bool, } @@ -36,6 +39,7 @@ pub(crate) enum ActiveGoalStopReason { struct GoalRuntimeInner { thread_id: ThreadId, state_dbs: Arc, + analytics: GoalAnalytics, event_emitter: GoalEventEmitter, metrics: GoalMetrics, thread_manager: Weak, @@ -87,6 +91,7 @@ impl GoalRuntimeHandle { inner: Arc::new(GoalRuntimeInner { thread_id, state_dbs, + analytics: config.analytics, event_emitter, metrics, thread_manager, @@ -165,6 +170,9 @@ impl GoalRuntimeHandle { .is_some_and(|previous_goal| previous_goal.goal_id != goal.goal_id); if previous_goal.is_none() || replaced_existing_goal { self.inner.metrics.record_created(); + self.inner + .analytics + .created(&goal, GoalEventAttribution::NoTurn); } let previous_status = previous_goal .as_ref() @@ -175,6 +183,9 @@ impl GoalRuntimeHandle { self.inner .metrics .record_terminal_if_status_changed(previous_status, &goal); + self.inner + .analytics + .status_changed(&goal, previous_status, GoalEventAttribution::NoTurn); let objective_changed = previous_goal.as_ref().is_some_and(|previous_goal| { !replaced_existing_goal && previous_goal.objective != goal.objective }); @@ -211,11 +222,15 @@ impl GoalRuntimeHandle { Ok(()) } - pub async fn apply_external_goal_clear(&self) -> Result<(), String> { + pub async fn apply_external_goal_clear( + &self, + goal: codex_state::ThreadGoal, + ) -> Result<(), String> { if !self.is_enabled() { return Ok(()); } + self.inner.analytics.cleared(&goal); self.inner.accounting_state.clear_active_goal(); Ok(()) } @@ -302,6 +317,11 @@ impl GoalRuntimeHandle { self.inner .metrics .record_terminal_if_status_changed(previous_status, &goal); + self.inner.analytics.status_changed( + &goal, + previous_status, + GoalEventAttribution::Turn(turn_id), + ); self.inner.accounting_state.clear_active_goal(); let goal = protocol_goal_from_state(goal); self.inner.event_emitter.thread_goal_updated( @@ -445,6 +465,14 @@ impl GoalRuntimeHandle { self.inner .metrics .record_terminal_if_status_changed(previous_status, &goal); + self.inner + .analytics + .usage_accounted(&goal, GoalEventAttribution::Turn(turn_id)); + self.inner.analytics.status_changed( + &goal, + previous_status, + GoalEventAttribution::Turn(turn_id), + ); accounting.mark_progress_accounted_for_status( turn_id, &snapshot, @@ -499,6 +527,14 @@ impl GoalRuntimeHandle { self.inner .metrics .record_terminal_if_status_changed(previous_status, &goal); + self.inner + .analytics + .usage_accounted(&goal, GoalEventAttribution::NoTurn); + self.inner.analytics.status_changed( + &goal, + previous_status, + GoalEventAttribution::NoTurn, + ); accounting.mark_idle_progress_accounted_for_status( &snapshot, goal.status, diff --git a/codex-rs/ext/goal/src/tool.rs b/codex-rs/ext/goal/src/tool.rs index 8c0f513004..d280f6e04b 100644 --- a/codex-rs/ext/goal/src/tool.rs +++ b/codex-rs/ext/goal/src/tool.rs @@ -1,6 +1,5 @@ use std::sync::Arc; -use async_trait::async_trait; use codex_extension_api::FunctionCallError; use codex_extension_api::JsonToolOutput; use codex_extension_api::ToolCall; @@ -17,6 +16,8 @@ use serde::Serialize; use crate::accounting::BudgetLimitedGoalDisposition; use crate::accounting::GoalAccountingState; +use crate::analytics::GoalAnalytics; +use crate::analytics::GoalEventAttribution; use crate::events::GoalEventEmitter; use crate::metrics::GoalMetrics; use crate::spec::CREATE_GOAL_TOOL_NAME; @@ -32,6 +33,7 @@ pub(crate) struct GoalToolExecutor { thread_id: ThreadId, state_db: Arc, accounting_state: Arc, + analytics: GoalAnalytics, event_emitter: GoalEventEmitter, metrics: GoalMetrics, } @@ -75,6 +77,7 @@ impl GoalToolExecutor { thread_id: ThreadId, state_db: Arc, accounting_state: Arc, + analytics: GoalAnalytics, event_emitter: GoalEventEmitter, metrics: GoalMetrics, ) -> Self { @@ -83,6 +86,7 @@ impl GoalToolExecutor { thread_id, state_db, accounting_state, + analytics, event_emitter, metrics, } @@ -92,6 +96,7 @@ impl GoalToolExecutor { thread_id: ThreadId, state_db: Arc, accounting_state: Arc, + analytics: GoalAnalytics, event_emitter: GoalEventEmitter, metrics: GoalMetrics, ) -> Self { @@ -100,6 +105,7 @@ impl GoalToolExecutor { thread_id, state_db, accounting_state, + analytics, event_emitter, metrics, } @@ -109,6 +115,7 @@ impl GoalToolExecutor { thread_id: ThreadId, state_db: Arc, accounting_state: Arc, + analytics: GoalAnalytics, event_emitter: GoalEventEmitter, metrics: GoalMetrics, ) -> Self { @@ -117,13 +124,13 @@ impl GoalToolExecutor { thread_id, state_db, accounting_state, + analytics, event_emitter, metrics, } } } -#[async_trait] impl ToolExecutor for GoalToolExecutor { fn tool_name(&self) -> ToolName { ToolName::plain(match self.kind { @@ -141,12 +148,14 @@ impl ToolExecutor for GoalToolExecutor { } } - async fn handle(&self, invocation: ToolCall) -> Result, FunctionCallError> { - match self.kind { - GoalToolKind::Get => self.handle_get(invocation).await, - GoalToolKind::Create => self.handle_create(invocation).await, - GoalToolKind::Update => self.handle_update(invocation).await, - } + fn handle(&self, invocation: ToolCall) -> codex_extension_api::ToolExecutorFuture<'_> { + Box::pin(async move { + match self.kind { + GoalToolKind::Get => self.handle_get(invocation).await, + GoalToolKind::Create => self.handle_create(invocation).await, + GoalToolKind::Update => self.handle_update(invocation).await, + } + }) } } @@ -200,6 +209,10 @@ impl GoalToolExecutor { .accounting_state .mark_current_turn_goal_active(goal.goal_id.clone()); self.metrics.record_created(); + self.analytics.created( + &goal, + GoalEventAttribution::Turn(invocation.turn_id.as_str()), + ); let goal = protocol_goal_from_state(goal); self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone()); goal_response(Some(goal), CompletionBudgetReport::Omit) @@ -259,6 +272,11 @@ impl GoalToolExecutor { })?; self.metrics .record_terminal_if_status_changed(previous_status, &goal); + self.analytics.status_changed( + &goal, + previous_status, + GoalEventAttribution::Turn(invocation.turn_id.as_str()), + ); let goal = protocol_goal_from_state(goal); let turn_id = self.accounting_state.clear_current_turn_goal(); self.emit_goal_updated_from_tool_call(&invocation, turn_id, goal.clone()); @@ -324,6 +342,13 @@ impl GoalToolExecutor { codex_state::GoalAccountingOutcome::Updated(goal) => { self.metrics .record_terminal_if_status_changed(previous_status, &goal); + self.analytics + .usage_accounted(&goal, GoalEventAttribution::Turn(turn_id.as_str())); + self.analytics.status_changed( + &goal, + previous_status, + GoalEventAttribution::Turn(turn_id.as_str()), + ); self.accounting_state.mark_progress_accounted_for_status( turn_id.as_str(), &snapshot, diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index b922f5332d..60c992ec78 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -4,6 +4,7 @@ use std::sync::PoisonError; use std::sync::Weak; use std::time::Duration; +use codex_analytics::AnalyticsEventsClient; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistryBuilder; @@ -1114,6 +1115,7 @@ async fn installed_tools_with_start( install_with_backend( &mut builder, runtime, + AnalyticsEventsClient::disabled(), /*metrics_client*/ None, Weak::new(), goal_service, @@ -1164,6 +1166,7 @@ impl GoalExtensionHarness { install_with_backend( &mut builder, runtime, + AnalyticsEventsClient::disabled(), /*metrics_client*/ None, Weak::new(), Arc::clone(&goal_service), diff --git a/codex-rs/ext/guardian/Cargo.toml b/codex-rs/ext/guardian/Cargo.toml index 53e553ec1c..513254b7cf 100644 --- a/codex-rs/ext/guardian/Cargo.toml +++ b/codex-rs/ext/guardian/Cargo.toml @@ -14,7 +14,6 @@ doctest = false workspace = true [dependencies] -async-trait = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-protocol = { workspace = true } diff --git a/codex-rs/ext/guardian/src/lib.rs b/codex-rs/ext/guardian/src/lib.rs index 0591887c25..a64cf4fed4 100644 --- a/codex-rs/ext/guardian/src/lib.rs +++ b/codex-rs/ext/guardian/src/lib.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use codex_core::config::Config; use codex_extension_api::AgentSpawnFuture; use codex_extension_api::AgentSpawner; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadStartInput; @@ -47,18 +48,23 @@ impl GuardianThreadContext { } } -#[async_trait::async_trait] impl ThreadLifecycleContributor for GuardianExtension where S: Send + Sync, { - async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) { - let Ok(forked_from_thread_id) = ThreadId::from_string(input.thread_store.level_id()) else { - return; - }; - input.thread_store.insert(GuardianThreadContext { - forked_from_thread_id, - }); + fn on_thread_start<'a>( + &'a self, + input: ThreadStartInput<'a, Config>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let Ok(forked_from_thread_id) = ThreadId::from_string(input.thread_store.level_id()) + else { + return; + }; + input.thread_store.insert(GuardianThreadContext { + forked_from_thread_id, + }); + }) } } diff --git a/codex-rs/ext/image-generation/Cargo.toml b/codex-rs/ext/image-generation/Cargo.toml index cb5ddda0a8..dfa85ee82e 100644 --- a/codex-rs/ext/image-generation/Cargo.toml +++ b/codex-rs/ext/image-generation/Cargo.toml @@ -13,7 +13,6 @@ doctest = false workspace = true [dependencies] -async-trait = { workspace = true } codex-api = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index 2c68f614f5..6a0016c037 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/codex-rs/ext/image-generation/src/extension.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use codex_core::config::Config; use codex_extension_api::ConfigContributor; use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadStartInput; @@ -41,13 +42,17 @@ impl From<&Config> for ImageGenerationExtensionConfig { } } -#[async_trait::async_trait] impl ThreadLifecycleContributor for ImageGenerationExtension { /// Seeds image-generation availability when a thread begins. - async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) { - input - .thread_store - .insert(ImageGenerationExtensionConfig::from(input.config)); + fn on_thread_start<'a>( + &'a self, + input: ThreadStartInput<'a, Config>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + input + .thread_store + .insert(ImageGenerationExtensionConfig::from(input.config)); + }) } } diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index 69b99bdbc4..07ed8f4bf6 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -78,7 +78,6 @@ struct ImagegenArgs { num_last_images_to_include: Option, } -#[async_trait::async_trait] impl ToolExecutor for ImageGenerationTool { /// Keeps the tool in the existing image-generation Responses namespace. fn tool_name(&self) -> ToolName { @@ -96,7 +95,13 @@ impl ToolExecutor for ImageGenerationTool { } /// Executes the selected image operation and returns the completed image result. - async fn handle(&self, call: ToolCall) -> Result, FunctionCallError> { + fn handle(&self, call: ToolCall) -> codex_extension_api::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(call)) + } +} + +impl ImageGenerationTool { + async fn handle_call(&self, call: ToolCall) -> Result, FunctionCallError> { let args = parse_args(&call)?; let request = request_for_args(&args, call.conversation_history.items())?; call.turn_item_emitter diff --git a/codex-rs/ext/mcp/Cargo.toml b/codex-rs/ext/mcp/Cargo.toml index 66c9ad1727..2d0e508b5e 100644 --- a/codex-rs/ext/mcp/Cargo.toml +++ b/codex-rs/ext/mcp/Cargo.toml @@ -14,7 +14,6 @@ doctest = false workspace = true [dependencies] -async-trait = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } diff --git a/codex-rs/ext/mcp/src/lib.rs b/codex-rs/ext/mcp/src/lib.rs index 4a7838ccce..8568d58150 100644 --- a/codex-rs/ext/mcp/src/lib.rs +++ b/codex-rs/ext/mcp/src/lib.rs @@ -1,31 +1,35 @@ use codex_core::config::Config; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributor; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; -use codex_mcp::codex_apps_mcp_server_config; +use codex_mcp::hosted_plugin_runtime_mcp_server_config; -struct HostedAppsMcpExtension; +struct HostedPluginRuntimeExtension; -#[async_trait::async_trait] -impl McpServerContributor for HostedAppsMcpExtension { - async fn contribute(&self, config: &Config) -> Vec { - let name = CODEX_APPS_MCP_SERVER_NAME.to_string(); - if !config.features.enabled(codex_features::Feature::Apps) { - return vec![McpServerContribution::Remove { name }]; - } +impl McpServerContributor for HostedPluginRuntimeExtension { + fn contribute<'a>( + &'a self, + config: &'a Config, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + let name = CODEX_APPS_MCP_SERVER_NAME.to_string(); + if !config.features.enabled(codex_features::Feature::Apps) { + return vec![McpServerContribution::Remove { name }]; + } - vec![McpServerContribution::Set { - name, - config: Box::new(codex_apps_mcp_server_config( - &config.chatgpt_base_url, - config.apps_mcp_path_override.as_deref(), - config.apps_mcp_product_sku.as_deref(), - )), - }] + vec![McpServerContribution::Set { + name, + config: Box::new(hosted_plugin_runtime_mcp_server_config( + &config.chatgpt_base_url, + config.apps_mcp_product_sku.as_deref(), + )), + }] + }) } } pub fn install(builder: &mut ExtensionRegistryBuilder) { - builder.mcp_server_contributor(std::sync::Arc::new(HostedAppsMcpExtension)); + builder.mcp_server_contributor(std::sync::Arc::new(HostedPluginRuntimeExtension)); } diff --git a/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs b/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs index 0d025df954..ea1b330caf 100644 --- a/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs +++ b/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs @@ -6,20 +6,22 @@ use codex_core::config::Config; use codex_core::config::ConfigBuilder; use codex_core_plugins::PluginsManager; use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::McpServerContribution; +use codex_extension_api::McpServerContributor; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use pretty_assertions::assert_eq; +type TestResult = Result<(), Box>; + #[tokio::test] -async fn contributes_hosted_apps_mcp_without_an_executor() -> Result<(), Box> -{ +async fn contributes_hosted_plugin_runtime_without_an_executor() -> TestResult { let codex_home = tempfile::tempdir()?; let config = ConfigBuilder::default() .codex_home(codex_home.path().to_path_buf()) .fallback_cwd(Some(codex_home.path().to_path_buf())) .cli_overrides(vec![ ("features.apps".to_string(), true.into()), - ("features.apps_mcp_path_override".to_string(), true.into()), ("chatgpt_base_url".to_string(), "https://chatgpt.com".into()), ]) .build() @@ -27,15 +29,13 @@ async fn contributes_hosted_apps_mcp_without_an_executor() -> Result<(), Box Result<(), Box Result<(), Box> { +async fn runtime_overlay_preserves_disabled_server() -> TestResult { + let codex_home = tempfile::tempdir()?; + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(codex_home.path().to_path_buf())) + .cli_overrides(vec![ + ("features.apps".to_string(), true.into()), + ( + "mcp_servers.codex_apps.url".to_string(), + "https://example.com/mcp".into(), + ), + ("mcp_servers.codex_apps.enabled".to_string(), false.into()), + ]) + .build() + .await?; + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let manager = installed_manager(&config); + + let servers = manager.effective_servers(&config, Some(&auth)).await; + let server = servers + .get(CODEX_APPS_MCP_SERVER_NAME) + .ok_or("hosted plugin runtime should remain configured")?; + + assert!(!server.enabled()); + Ok(()) +} + +#[tokio::test] +async fn legacy_fallback_overwrites_reserved_config_without_an_extension() -> TestResult { + let codex_home = tempfile::tempdir()?; + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(codex_home.path().to_path_buf())) + .cli_overrides(vec![ + ("features.apps".to_string(), true.into()), + ( + "mcp_servers.codex_apps.url".to_string(), + "https://example.com/mcp".into(), + ), + ]) + .build() + .await?; + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let manager = McpManager::new(Arc::new(PluginsManager::new( + config.codex_home.to_path_buf(), + ))); + + let servers = manager.effective_servers(&config, Some(&auth)).await; + let server = servers + .get(CODEX_APPS_MCP_SERVER_NAME) + .and_then(|server| server.configured_config()) + .ok_or("legacy Apps MCP should be present")?; + let McpServerTransportConfig::StreamableHttp { url, .. } = &server.transport else { + panic!("legacy Apps MCP should use streamable HTTP"); + }; + assert_eq!(url, "https://chatgpt.com/backend-api/wham/apps"); + + Ok(()) +} + +#[tokio::test] +async fn extension_can_remove_legacy_fallback_while_apps_are_enabled() -> TestResult { + let codex_home = tempfile::tempdir()?; + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(codex_home.path().to_path_buf())) + .cli_overrides(vec![("features.apps".to_string(), true.into())]) + .build() + .await?; + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let mut builder = ExtensionRegistryBuilder::new(); + builder.mcp_server_contributor(Arc::new(RemoveCodexApps)); + let manager = McpManager::new_with_extensions( + Arc::new(PluginsManager::new(config.codex_home.to_path_buf())), + Arc::new(builder.build()), + ); + + let servers = manager.effective_servers(&config, Some(&auth)).await; + + assert!(!servers.contains_key(CODEX_APPS_MCP_SERVER_NAME)); + Ok(()) +} + +#[tokio::test] +async fn hosted_apps_mcp_requires_chatgpt_auth() -> TestResult { let codex_home = tempfile::tempdir()?; let config = ConfigBuilder::default() .codex_home(codex_home.path().to_path_buf()) @@ -61,7 +145,7 @@ async fn hosted_apps_mcp_requires_chatgpt_auth() -> Result<(), Box Result<(), Box> { +async fn disabled_apps_remove_reserved_server_config() -> TestResult { let codex_home = tempfile::tempdir()?; let config = ConfigBuilder::default() .codex_home(codex_home.path().to_path_buf()) @@ -91,3 +175,18 @@ fn installed_manager(config: &Config) -> McpManager { Arc::new(builder.build()), ) } + +struct RemoveCodexApps; + +impl McpServerContributor for RemoveCodexApps { + fn contribute<'a>( + &'a self, + _config: &'a Config, + ) -> codex_extension_api::ExtensionFuture<'a, Vec> { + Box::pin(async move { + vec![McpServerContribution::Remove { + name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + }] + }) + } +} diff --git a/codex-rs/ext/memories/Cargo.toml b/codex-rs/ext/memories/Cargo.toml index 68eeaf8ec1..e283140b15 100644 --- a/codex-rs/ext/memories/Cargo.toml +++ b/codex-rs/ext/memories/Cargo.toml @@ -13,7 +13,6 @@ doctest = false workspace = true [dependencies] -async-trait = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } diff --git a/codex-rs/ext/memories/src/extension.rs b/codex-rs/ext/memories/src/extension.rs index 0be773b4c7..23c80a8779 100644 --- a/codex-rs/ext/memories/src/extension.rs +++ b/codex-rs/ext/memories/src/extension.rs @@ -4,6 +4,7 @@ use codex_core::config::Config; use codex_extension_api::ConfigContributor; use codex_extension_api::ContextContributor; use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::PromptFragment; use codex_extension_api::ThreadLifecycleContributor; @@ -69,12 +70,16 @@ impl ContextContributor for MemoriesExtension { } } -#[async_trait::async_trait] impl ThreadLifecycleContributor for MemoriesExtension { - async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) { - input - .thread_store - .insert(MemoriesExtensionConfig::from_config(input.config)); + fn on_thread_start<'a>( + &'a self, + input: ThreadStartInput<'a, Config>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + input + .thread_store + .insert(MemoriesExtensionConfig::from_config(input.config)); + }) } } diff --git a/codex-rs/ext/memories/src/tools/ad_hoc_note.rs b/codex-rs/ext/memories/src/tools/ad_hoc_note.rs index a6712a4022..0ce9cade67 100644 --- a/codex-rs/ext/memories/src/tools/ad_hoc_note.rs +++ b/codex-rs/ext/memories/src/tools/ad_hoc_note.rs @@ -41,7 +41,6 @@ pub(super) struct AddAdHocNoteTool { pub(super) metrics_client: Option, } -#[async_trait::async_trait] impl ToolExecutor for AddAdHocNoteTool where B: MemoriesBackend, @@ -57,7 +56,16 @@ where ) } - async fn handle( + fn handle(&self, call: ToolCall) -> codex_extension_api::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(call)) + } +} + +impl AddAdHocNoteTool +where + B: MemoriesBackend, +{ + async fn handle_call( &self, call: ToolCall, ) -> Result, codex_extension_api::FunctionCallError> diff --git a/codex-rs/ext/memories/src/tools/list.rs b/codex-rs/ext/memories/src/tools/list.rs index 301c7cab71..b7b9ad0586 100644 --- a/codex-rs/ext/memories/src/tools/list.rs +++ b/codex-rs/ext/memories/src/tools/list.rs @@ -39,7 +39,6 @@ pub(super) struct ListTool { pub(super) metrics_client: Option, } -#[async_trait::async_trait] impl ToolExecutor for ListTool where B: MemoriesBackend, @@ -55,7 +54,16 @@ where ) } - async fn handle( + fn handle(&self, call: ToolCall) -> codex_extension_api::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(call)) + } +} + +impl ListTool +where + B: MemoriesBackend, +{ + async fn handle_call( &self, call: ToolCall, ) -> Result, codex_extension_api::FunctionCallError> diff --git a/codex-rs/ext/memories/src/tools/read.rs b/codex-rs/ext/memories/src/tools/read.rs index 33ede3c600..6625cdaf8b 100644 --- a/codex-rs/ext/memories/src/tools/read.rs +++ b/codex-rs/ext/memories/src/tools/read.rs @@ -38,7 +38,6 @@ pub(super) struct ReadTool { pub(super) metrics_client: Option, } -#[async_trait::async_trait] impl ToolExecutor for ReadTool where B: MemoriesBackend, @@ -54,7 +53,16 @@ where ) } - async fn handle( + fn handle(&self, call: ToolCall) -> codex_extension_api::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(call)) + } +} + +impl ReadTool +where + B: MemoriesBackend, +{ + async fn handle_call( &self, call: ToolCall, ) -> Result, codex_extension_api::FunctionCallError> diff --git a/codex-rs/ext/memories/src/tools/search.rs b/codex-rs/ext/memories/src/tools/search.rs index 1d05072703..928c98b475 100644 --- a/codex-rs/ext/memories/src/tools/search.rs +++ b/codex-rs/ext/memories/src/tools/search.rs @@ -47,7 +47,6 @@ pub(super) struct SearchTool { pub(super) metrics_client: Option, } -#[async_trait::async_trait] impl ToolExecutor for SearchTool where B: MemoriesBackend, @@ -63,7 +62,16 @@ where ) } - async fn handle( + fn handle(&self, call: ToolCall) -> codex_extension_api::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(call)) + } +} + +impl SearchTool +where + B: MemoriesBackend, +{ + async fn handle_call( &self, call: ToolCall, ) -> Result, codex_extension_api::FunctionCallError> diff --git a/codex-rs/ext/skills/Cargo.toml b/codex-rs/ext/skills/Cargo.toml index 9ff42cfafc..0a30dddb99 100644 --- a/codex-rs/ext/skills/Cargo.toml +++ b/codex-rs/ext/skills/Cargo.toml @@ -14,15 +14,22 @@ doctest = false workspace = true [dependencies] -async-trait = { workspace = true } -codex-core = { workspace = true } codex-core-skills = { workspace = true } codex-exec-server = { workspace = true } codex-extension-api = { workspace = true } +codex-mcp = { workspace = true } codex-protocol = { workspace = true } +codex-tools = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-string = { workspace = true } +schemars = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["sync", "time"] } +tracing = { workspace = true } +url = { workspace = true } [dev-dependencies] +async-trait = { workspace = true } pretty_assertions = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/skills/src/catalog.rs b/codex-rs/ext/skills/src/catalog.rs index c966cc09e9..739c5320b5 100644 --- a/codex-rs/ext/skills/src/catalog.rs +++ b/codex-rs/ext/skills/src/catalog.rs @@ -9,8 +9,8 @@ pub enum SkillSourceKind { Host, /// Skills owned by an execution environment. Executor, - /// Skills read through an authenticated remote catalog/API. - Remote, + /// Skills owned by the orchestrator rather than an execution environment. + Orchestrator, /// Extension-private source kind for future providers that do not fit an /// existing transport category. Custom(String), @@ -25,7 +25,7 @@ impl SkillSourceKind { match self { Self::Host => "host", Self::Executor => "executor", - Self::Remote => "remote", + Self::Orchestrator => "orchestrator", Self::Custom(kind) => kind, } } diff --git a/codex-rs/ext/skills/src/config.rs b/codex-rs/ext/skills/src/config.rs new file mode 100644 index 0000000000..e883f0c5f5 --- /dev/null +++ b/codex-rs/ext/skills/src/config.rs @@ -0,0 +1,8 @@ +/// Host-supplied configuration used by the skills extension. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SkillsExtensionConfig { + /// Whether the available-skills catalog is included in model context. + pub include_instructions: bool, + /// Whether bundled skills are eligible for discovery. + pub bundled_skills_enabled: bool, +} diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index 45d7c95cb7..23dee27ebb 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -1,29 +1,34 @@ use std::sync::Arc; -use codex_core::config::Config; use codex_core_skills::HostLoadedSkills; -use codex_core_skills::SkillInstructions; use codex_core_skills::injection::InjectedHostSkillPrompts; -use codex_core_skills::injection::SkillInjection; use codex_extension_api::ConfigContributor; use codex_extension_api::ContextContributor; use codex_extension_api::ContextualUserFragment; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::PromptFragment; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadStartInput; +use codex_extension_api::ToolCall; +use codex_extension_api::ToolContributor; +use codex_extension_api::ToolExecutor; use codex_extension_api::TurnInputContext; use codex_extension_api::TurnInputContributor; +use codex_mcp::McpResourceClient; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::WarningEvent; +use crate::SkillsExtensionConfig; +use crate::catalog::SkillCatalog; use crate::catalog::SkillCatalogEntry; use crate::catalog::SkillReadResult; use crate::catalog::SkillSourceKind; +use crate::fragments::SkillInstructions; use crate::provider::HostSkillProvider; use crate::provider::SkillListQuery; use crate::provider::SkillReadRequest; @@ -34,40 +39,47 @@ use crate::render::truncate_main_prompt_contents; use crate::render::truncate_utf8_to_bytes; use crate::selection::collect_explicit_skill_mentions; use crate::sources::SkillProviders; -use crate::state::SkillsExtensionConfig; use crate::state::SkillsThreadState; use crate::state::SkillsTurnState; +use crate::tools::skill_tools; -#[derive(Clone)] -struct SkillsExtension { +struct SkillsExtension { providers: SkillProviders, event_sink: Arc, + config_from_host: Arc SkillsExtensionConfig + Send + Sync>, } -#[async_trait::async_trait] -impl ThreadLifecycleContributor for SkillsExtension { - async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) { - let selected_roots = input - .thread_store - .get::>() - .map(|selected_roots| selected_roots.as_ref().clone()) - .unwrap_or_default(); - input.thread_store.insert(SkillsThreadState::new( - SkillsExtensionConfig::from_config(input.config), - selected_roots, - )); +impl ThreadLifecycleContributor for SkillsExtension +where + C: Send + Sync + 'static, +{ + fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let selected_roots = input + .thread_store + .get::>() + .map(|selected_roots| selected_roots.as_ref().clone()) + .unwrap_or_default(); + input.thread_store.insert(SkillsThreadState::new( + (self.config_from_host)(input.config), + selected_roots, + )); + }) } } -impl ConfigContributor for SkillsExtension { +impl ConfigContributor for SkillsExtension +where + C: Send + Sync + 'static, +{ fn on_config_changed( &self, _session_store: &ExtensionData, thread_store: &ExtensionData, - _previous_config: &Config, - new_config: &Config, + _previous_config: &C, + new_config: &C, ) { - let next_config = SkillsExtensionConfig::from_config(new_config); + let next_config = (self.config_from_host)(new_config); if let Some(state) = thread_store.get::() { state.set_config(next_config); } else { @@ -76,10 +88,13 @@ impl ConfigContributor for SkillsExtension { } } -impl ContextContributor for SkillsExtension { +impl ContextContributor for SkillsExtension +where + C: Send + Sync + 'static, +{ fn contribute<'a>( &'a self, - _session_store: &'a ExtensionData, + session_store: &'a ExtensionData, thread_store: &'a ExtensionData, ) -> std::pin::Pin> + Send + 'a>> { Box::pin(async move { @@ -87,19 +102,22 @@ impl ContextContributor for SkillsExtension { return Vec::new(); }; let config = thread_state.config(); - if !config.include_instructions || thread_state.selected_roots().is_empty() { + if !config.include_instructions { return Vec::new(); } let catalog = self - .providers - .list_for_turn(SkillListQuery { - turn_id: thread_store.level_id().to_string(), - executor_roots: thread_state.selected_roots().to_vec(), - host: None, - include_host_skills: false, - include_bundled_skills: config.bundled_skills_enabled, - include_remote_skills: false, - }) + .list_skills( + SkillListQuery { + turn_id: thread_store.level_id().to_string(), + executor_roots: thread_state.selected_roots().to_vec(), + host: None, + include_host_skills: false, + include_bundled_skills: config.bundled_skills_enabled, + include_orchestrator_skills: true, + mcp_resources: session_store.get::(), + }, + &thread_state, + ) .await; for warning in &catalog.warnings { self.emit_warning(thread_store.level_id(), warning.clone()); @@ -112,120 +130,173 @@ impl ContextContributor for SkillsExtension { } } -#[async_trait::async_trait] -impl TurnInputContributor for SkillsExtension { - async fn contribute( +impl ToolContributor for SkillsExtension +where + C: Send + Sync + 'static, +{ + fn tools( &self, - input: TurnInputContext, - _session_store: &ExtensionData, - thread_store: &ExtensionData, - turn_store: &ExtensionData, - ) -> Vec> { - let Some(thread_state) = thread_store.get::() else { + session_store: &ExtensionData, + _thread_store: &ExtensionData, + ) -> Vec>> { + if !self.providers.has_orchestrator_provider() { return Vec::new(); - }; - - let config = thread_state.config(); - let host_loaded_skills = turn_store.get::(); - let query = SkillListQuery { - turn_id: input.turn_id.clone(), - executor_roots: thread_state.selected_roots().to_vec(), - host: host_loaded_skills.clone(), - include_host_skills: true, - include_bundled_skills: config.bundled_skills_enabled, - include_remote_skills: true, - }; - let catalog = self.providers.list_for_turn(query).await; - for warning in &catalog.warnings { - self.emit_warning(&input.turn_id, warning.clone()); } - let selected_entries = collect_explicit_skill_mentions(&input.user_input, &catalog); - let mut fragments: Vec> = Vec::new(); - if config.include_instructions { - let mut turn_catalog = catalog.clone(); - turn_catalog - .entries - .retain(|entry| entry.authority.kind != SkillSourceKind::Executor); - if let Some(fragment) = available_skills_fragment(&turn_catalog) { - fragments.push(Box::new(fragment)); - } - } - - let mut warnings = catalog.warnings.clone(); - let mut main_prompts_injected = false; - let mut injected_host_skill_prompts = InjectedHostSkillPrompts::default(); - for entry in &selected_entries { - match self - .read_main_prompt(entry, host_loaded_skills.clone()) - .await - { - Ok(read_result) => { - let (contents, truncated) = - truncate_main_prompt_contents(read_result.contents.as_str()); - if truncated { - let warning = format!( - "Skill `{}` exceeded the main prompt context limit and was truncated.", - entry.name - ); - self.emit_warning(&input.turn_id, warning.clone()); - warnings.push(warning); - } - let injection = SkillInjection { - name: truncate_utf8_to_bytes(&entry.name, MAX_SKILL_NAME_BYTES).0, - path: truncate_utf8_to_bytes(entry.rendered_path(), MAX_SKILL_PATH_BYTES).0, - contents, - }; - fragments.push(Box::new(SkillInstructions::from(&injection))); - main_prompts_injected = true; - if entry.authority.kind == SkillSourceKind::Host { - injected_host_skill_prompts.insert_path(entry.main_prompt.as_str()); - } - } - Err(message) => { - let warning = format!("Failed to load skill `{}`: {message}", entry.name); - self.emit_warning(&input.turn_id, warning.clone()); - warnings.push(warning); - } - } - } - - if let Some(host_loaded_skills) = &host_loaded_skills { - for entry in selected_entries - .iter() - .filter(|entry| entry.authority.kind != SkillSourceKind::Host) - { - for host_skill in host_loaded_skills - .outcome() - .skills - .iter() - .filter(|host_skill| host_skill.name == entry.name) - { - injected_host_skill_prompts - .insert_path(host_skill.path_to_skills_md.to_string_lossy()); - } - } - } - - turn_store.insert(SkillsTurnState { - catalog, - selected_entries, - warnings, - main_prompts_injected, - }); - if !injected_host_skill_prompts.is_empty() { - turn_store.insert(injected_host_skill_prompts); - } - - fragments + skill_tools( + self.providers.clone(), + session_store.get::(), + ) } } -impl SkillsExtension { +impl TurnInputContributor for SkillsExtension +where + C: Send + Sync + 'static, +{ + fn contribute<'a>( + &'a self, + input: TurnInputContext, + session_store: &'a ExtensionData, + thread_store: &'a ExtensionData, + turn_store: &'a ExtensionData, + ) -> ExtensionFuture<'a, Vec>> { + Box::pin(async move { + let Some(thread_state) = thread_store.get::() else { + return Vec::new(); + }; + + let config = thread_state.config(); + let host_loaded_skills = turn_store.get::(); + let query = SkillListQuery { + turn_id: input.turn_id.clone(), + executor_roots: thread_state.selected_roots().to_vec(), + host: host_loaded_skills.clone(), + include_host_skills: true, + include_bundled_skills: config.bundled_skills_enabled, + include_orchestrator_skills: true, + mcp_resources: session_store.get::(), + }; + let catalog = self.list_skills(query, &thread_state).await; + for warning in &catalog.warnings { + self.emit_warning(&input.turn_id, warning.clone()); + } + + let selected_entries = collect_explicit_skill_mentions(&input.user_input, &catalog); + let mut fragments: Vec> = Vec::new(); + if config.include_instructions { + let mut turn_catalog = catalog.clone(); + turn_catalog.entries.retain(|entry| { + entry.authority.kind != SkillSourceKind::Executor + && entry.authority.kind != SkillSourceKind::Orchestrator + }); + if let Some(fragment) = available_skills_fragment(&turn_catalog) { + fragments.push(Box::new(fragment)); + } + } + + let mut warnings = catalog.warnings.clone(); + let mut main_prompts_injected = false; + let mut injected_host_skill_prompts = InjectedHostSkillPrompts::default(); + for entry in &selected_entries { + match self + .read_main_prompt(entry, host_loaded_skills.clone(), session_store) + .await + { + Ok(read_result) => { + let (contents, truncated) = + truncate_main_prompt_contents(read_result.contents.as_str()); + if truncated { + let warning = format!( + "Skill `{}` exceeded the main prompt context limit and was truncated.", + entry.name + ); + self.emit_warning(&input.turn_id, warning.clone()); + warnings.push(warning); + } + let fragment = SkillInstructions { + name: truncate_utf8_to_bytes(&entry.name, MAX_SKILL_NAME_BYTES).0, + path: truncate_utf8_to_bytes( + entry.rendered_path(), + MAX_SKILL_PATH_BYTES, + ) + .0, + contents, + }; + fragments.push(Box::new(fragment)); + main_prompts_injected = true; + if entry.authority.kind == SkillSourceKind::Host { + injected_host_skill_prompts.insert_path(entry.main_prompt.as_str()); + } + } + Err(message) => { + let warning = format!("Failed to load skill `{}`: {message}", entry.name); + self.emit_warning(&input.turn_id, warning.clone()); + warnings.push(warning); + } + } + } + + if let Some(host_loaded_skills) = &host_loaded_skills { + for entry in selected_entries + .iter() + .filter(|entry| entry.authority.kind != SkillSourceKind::Host) + { + for host_skill in host_loaded_skills + .outcome() + .skills + .iter() + .filter(|host_skill| host_skill.name == entry.name) + { + injected_host_skill_prompts + .insert_path(host_skill.path_to_skills_md.to_string_lossy()); + } + } + } + + turn_store.insert(SkillsTurnState { + catalog, + selected_entries, + warnings, + main_prompts_injected, + }); + if !injected_host_skill_prompts.is_empty() { + turn_store.insert(injected_host_skill_prompts); + } + + fragments + }) + } +} + +impl SkillsExtension { + async fn list_skills( + &self, + mut query: SkillListQuery, + thread_state: &SkillsThreadState, + ) -> SkillCatalog { + let include_orchestrator_skills = query.include_orchestrator_skills; + let orchestrator_query = query.clone(); + query.include_orchestrator_skills = false; + + let mut catalog = self.providers.list_for_turn(query).await; + if include_orchestrator_skills { + let orchestrator_catalog = thread_state + .orchestrator_catalog_snapshot( + self.providers + .list_orchestrator_for_turn(orchestrator_query), + ) + .await; + catalog.extend(orchestrator_catalog); + } + catalog + } + async fn read_main_prompt( &self, entry: &SkillCatalogEntry, host_loaded_skills: Option>, + session_store: &ExtensionData, ) -> Result { self.providers .read(SkillReadRequest { @@ -233,6 +304,7 @@ impl SkillsExtension { package: entry.id.clone(), resource: entry.main_prompt.clone(), host: host_loaded_skills, + mcp_resources: session_store.get::(), }) .await .map_err(|err| err.message) @@ -246,23 +318,34 @@ impl SkillsExtension { } } -pub fn install(registry: &mut ExtensionRegistryBuilder) { +pub fn install( + registry: &mut ExtensionRegistryBuilder, + config_from_host: impl Fn(&C) -> SkillsExtensionConfig + Send + Sync + 'static, +) where + C: Send + Sync + 'static, +{ install_with_providers( registry, SkillProviders::new().with_host_provider(Arc::new(HostSkillProvider::new())), + config_from_host, ); } -pub fn install_with_providers( - registry: &mut ExtensionRegistryBuilder, +pub fn install_with_providers( + registry: &mut ExtensionRegistryBuilder, providers: SkillProviders, -) { + config_from_host: impl Fn(&C) -> SkillsExtensionConfig + Send + Sync + 'static, +) where + C: Send + Sync + 'static, +{ let extension = Arc::new(SkillsExtension { providers, event_sink: registry.event_sink(), + config_from_host: Arc::new(config_from_host), }); registry.thread_lifecycle_contributor(extension.clone()); registry.config_contributor(extension.clone()); registry.prompt_contributor(extension.clone()); - registry.turn_input_contributor(extension); + registry.turn_input_contributor(extension.clone()); + registry.tool_contributor(extension); } diff --git a/codex-rs/ext/skills/src/fragments.rs b/codex-rs/ext/skills/src/fragments.rs new file mode 100644 index 0000000000..0a24b48b5a --- /dev/null +++ b/codex-rs/ext/skills/src/fragments.rs @@ -0,0 +1,61 @@ +use codex_core_skills::render_available_skills_body; +use codex_extension_api::ContextualUserFragment; +use codex_protocol::protocol::SKILLS_INSTRUCTIONS_CLOSE_TAG; +use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct AvailableSkillsInstructions { + skill_lines: Vec, +} + +impl AvailableSkillsInstructions { + pub(crate) fn from_skill_lines(skill_lines: Vec) -> Self { + Self { skill_lines } + } +} + +impl ContextualUserFragment for AvailableSkillsInstructions { + fn role(&self) -> &'static str { + "developer" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + (SKILLS_INSTRUCTIONS_OPEN_TAG, SKILLS_INSTRUCTIONS_CLOSE_TAG) + } + + fn body(&self) -> String { + render_available_skills_body(&[], &self.skill_lines) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct SkillInstructions { + pub(crate) name: String, + pub(crate) path: String, + pub(crate) contents: String, +} + +impl ContextualUserFragment for SkillInstructions { + fn role(&self) -> &'static str { + "user" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("", "") + } + + fn body(&self) -> String { + let name = &self.name; + let path = &self.path; + let contents = &self.contents; + format!("\n{name}\n{path}\n{contents}\n") + } +} diff --git a/codex-rs/ext/skills/src/lib.rs b/codex-rs/ext/skills/src/lib.rs index 5ab03ddd9e..8abdf6a93e 100644 --- a/codex-rs/ext/skills/src/lib.rs +++ b/codex-rs/ext/skills/src/lib.rs @@ -1,15 +1,20 @@ pub mod catalog; +mod config; mod extension; +mod fragments; pub mod provider; mod render; mod selection; mod sources; mod state; +mod tools; +pub use config::SkillsExtensionConfig; pub use extension::install; pub use extension::install_with_providers; pub use provider::ExecutorSkillProvider; pub use provider::HostSkillProvider; +pub use provider::OrchestratorSkillProvider; pub use provider::SkillProvider; pub use sources::SkillProviderSource; pub use sources::SkillProviders; diff --git a/codex-rs/ext/skills/src/provider.rs b/codex-rs/ext/skills/src/provider.rs index 8d13a2f981..329c0c5813 100644 --- a/codex-rs/ext/skills/src/provider.rs +++ b/codex-rs/ext/skills/src/provider.rs @@ -4,8 +4,10 @@ use std::sync::Arc; mod executor; mod host; +mod orchestrator; use codex_core_skills::HostLoadedSkills; +use codex_mcp::McpResourceClient; use codex_protocol::capabilities::SelectedCapabilityRoot; use crate::catalog::SkillAuthority; @@ -18,6 +20,7 @@ use crate::catalog::SkillSearchResult; pub use executor::ExecutorSkillProvider; pub use host::HostSkillProvider; +pub use orchestrator::OrchestratorSkillProvider; #[derive(Clone, Debug)] pub struct SkillListQuery { @@ -26,7 +29,8 @@ pub struct SkillListQuery { pub host: Option>, pub include_host_skills: bool, pub include_bundled_skills: bool, - pub include_remote_skills: bool, + pub include_orchestrator_skills: bool, + pub mcp_resources: Option>, } #[derive(Clone, Debug)] @@ -35,6 +39,7 @@ pub struct SkillReadRequest { pub package: SkillPackageId, pub resource: SkillResourceId, pub host: Option>, + pub mcp_resources: Option>, } #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/codex-rs/ext/skills/src/provider/orchestrator.rs b/codex-rs/ext/skills/src/provider/orchestrator.rs new file mode 100644 index 0000000000..045b67de62 --- /dev/null +++ b/codex-rs/ext/skills/src/provider/orchestrator.rs @@ -0,0 +1,327 @@ +use std::collections::HashSet; +use std::time::Duration; + +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_protocol::mcp::Resource; +use codex_protocol::mcp::ResourceContent; +use url::Url; + +use crate::catalog::SkillAuthority; +use crate::catalog::SkillCatalog; +use crate::catalog::SkillCatalogEntry; +use crate::catalog::SkillPackageId; +use crate::catalog::SkillProviderError; +use crate::catalog::SkillReadResult; +use crate::catalog::SkillResourceId; +use crate::catalog::SkillSearchResult; +use crate::catalog::SkillSourceKind; +use crate::provider::SkillListQuery; +use crate::provider::SkillProvider; +use crate::provider::SkillProviderFuture; +use crate::provider::SkillReadRequest; +use crate::provider::SkillSearchRequest; + +const ORCHESTRATOR_SKILL_MIME_TYPE: &str = "mcp/skill"; +const ORCHESTRATOR_SKILL_DISCOVERY_TIMEOUT: Duration = Duration::from_secs(10); +const ORCHESTRATOR_SKILL_READ_TIMEOUT: Duration = Duration::from_secs(10); +const MAX_RESOURCE_PAGES: usize = 10; +const MAX_ORCHESTRATOR_SKILLS: usize = 100; +const MAX_SKILL_NAME_CHARS: usize = 64; +const MAX_QUALIFIED_SKILL_NAME_CHARS: usize = 128; +const MAX_SKILL_DESCRIPTION_CHARS: usize = 1_024; +const MAX_SKILL_PACKAGE_URI_CHARS: usize = 1_024; +const MAX_SKILL_RESOURCE_URI_CHARS: usize = 2_048; +const MAX_SKILL_RESOURCE_CONTENT_BYTES: usize = 1024 * 1024; + +/// Discovers and reads skills owned by the orchestrator. +/// +/// The provider uses session-scoped resources without exposing the transport or +/// resource server to callers that configure the skills extension. +#[derive(Clone, Debug, Default)] +pub struct OrchestratorSkillProvider; + +impl OrchestratorSkillProvider { + pub fn new() -> Self { + Self + } +} + +impl SkillProvider for OrchestratorSkillProvider { + fn list(&self, query: SkillListQuery) -> SkillProviderFuture<'_, SkillCatalog> { + Box::pin(async move { + let Some(client) = query.mcp_resources else { + return Ok(SkillCatalog::default()); + }; + if !client.has_server(CODEX_APPS_MCP_SERVER_NAME).await { + return Ok(SkillCatalog::default()); + } + + let discovery_deadline = + tokio::time::Instant::now() + ORCHESTRATOR_SKILL_DISCOVERY_TIMEOUT; + let mut catalog = SkillCatalog::default(); + let mut cursor = None; + let mut seen_cursors = HashSet::new(); + let mut skill_resources_seen = 0usize; + let mut skipped_resources = 0usize; + let mut truncated = false; + let mut completed_pages = 0usize; + + for _ in 0..MAX_RESOURCE_PAGES { + let page = match tokio::time::timeout_at( + discovery_deadline, + client.list_resources(CODEX_APPS_MCP_SERVER_NAME, cursor.clone()), + ) + .await + { + Ok(result) => result.map_err(|err| { + SkillProviderError::new(format!( + "failed to list orchestrator skill resources: {err:#}" + )) + }), + Err(_) => Err(SkillProviderError::new(format!( + "orchestrator skill discovery timed out after {ORCHESTRATOR_SKILL_DISCOVERY_TIMEOUT:?}" + ))), + }; + let result = match page { + Ok(result) => result, + Err(err) if completed_pages == 0 => return Err(err), + Err(err) => { + let page_word = if completed_pages == 1 { + "page" + } else { + "pages" + }; + catalog.warnings.push(format!( + "Orchestrator skill discovery stopped after {completed_pages} resource {page_word}: {}", + err.message + )); + cursor = None; + break; + } + }; + completed_pages = completed_pages.saturating_add(1); + + for resource in &result.resources { + if resource.mime_type.as_deref() != Some(ORCHESTRATOR_SKILL_MIME_TYPE) { + continue; + } + if skill_resources_seen >= MAX_ORCHESTRATOR_SKILLS { + truncated = true; + break; + } + skill_resources_seen = skill_resources_seen.saturating_add(1); + match catalog_entry_from_resource(resource) { + Some(entry) => catalog.push_entry(entry), + None => skipped_resources = skipped_resources.saturating_add(1), + } + } + + if truncated { + break; + } + let Some(next_cursor) = result.next_cursor else { + cursor = None; + break; + }; + if !seen_cursors.insert(next_cursor.clone()) { + catalog.warnings.push( + "Orchestrator skill resource pagination returned a duplicate cursor." + .to_string(), + ); + cursor = None; + break; + } + cursor = Some(next_cursor); + } + + if cursor.is_some() || truncated { + catalog.warnings.push(format!( + "Orchestrator skill discovery was truncated at {MAX_ORCHESTRATOR_SKILLS} skills or {MAX_RESOURCE_PAGES} resource pages." + )); + } + if skipped_resources > 0 { + catalog.warnings.push(format!( + "Skipped {skipped_resources} malformed orchestrator skill resources." + )); + } + + Ok(catalog) + }) + } + + fn read(&self, request: SkillReadRequest) -> SkillProviderFuture<'_, SkillReadResult> { + Box::pin(async move { + if request.authority + != SkillAuthority::new(SkillSourceKind::Orchestrator, CODEX_APPS_MCP_SERVER_NAME) + { + return Err(SkillProviderError::new(format!( + "orchestrator skill provider cannot read authority {}", + request.authority.id + ))); + } + if !resource_belongs_to_package(&request.package.0, request.resource.as_str()) { + return Err(SkillProviderError::new( + "orchestrator skill resource does not match its package", + )); + } + + let Some(client) = request.mcp_resources.as_ref() else { + return Err(SkillProviderError::new( + "session MCP resource client is not configured", + )); + }; + let result = tokio::time::timeout( + ORCHESTRATOR_SKILL_READ_TIMEOUT, + client.read_resource(CODEX_APPS_MCP_SERVER_NAME, request.resource.as_str()), + ) + .await + .map_err(|_| { + SkillProviderError::new(format!( + "orchestrator skill read timed out after {ORCHESTRATOR_SKILL_READ_TIMEOUT:?}" + )) + })? + .map_err(|err| { + SkillProviderError::new(format!( + "failed to read orchestrator skill resource {}: {err:#}", + request.resource.as_str() + )) + })?; + let contents = result + .contents + .into_iter() + .find_map(|contents| match contents { + ResourceContent::Text { uri, text, .. } if uri == request.resource.as_str() => { + Some(text) + } + ResourceContent::Text { .. } | ResourceContent::Blob { .. } => None, + }); + let Some(contents) = contents else { + return Err(SkillProviderError::new(format!( + "orchestrator skill resource {} did not return matching text contents", + request.resource.as_str() + ))); + }; + if contents.len() > MAX_SKILL_RESOURCE_CONTENT_BYTES { + return Err(SkillProviderError::new(format!( + "orchestrator skill resource {} exceeds the {MAX_SKILL_RESOURCE_CONTENT_BYTES}-byte read limit", + request.resource.as_str() + ))); + } + + Ok(SkillReadResult { + resource: request.resource, + contents, + }) + }) + } + + fn search(&self, _request: SkillSearchRequest) -> SkillProviderFuture<'_, SkillSearchResult> { + Box::pin(async { Ok(SkillSearchResult::default()) }) + } +} + +fn catalog_entry_from_resource(resource: &Resource) -> Option { + let uri = validated_skill_uri(resource.uri.as_str(), MAX_SKILL_PACKAGE_URI_CHARS)?; + let meta = resource.meta.as_ref()?.as_object()?; + let skill_name = normalized_label(meta.get("skill_name")?.as_str()?, MAX_SKILL_NAME_CHARS)?; + let name = if meta.get("source").and_then(|value| value.as_str()) == Some("user") { + skill_name + } else { + let plugin_name = + normalized_label(meta.get("plugin_name")?.as_str()?, MAX_SKILL_NAME_CHARS)?; + let qualified_name = format!("{plugin_name}:{skill_name}"); + (qualified_name.chars().count() <= MAX_QUALIFIED_SKILL_NAME_CHARS) + .then_some(qualified_name)? + }; + let description = normalized_description(resource.description.as_deref().unwrap_or_default())?; + let main_prompt = main_prompt_uri(uri); + + Some( + SkillCatalogEntry::new( + SkillPackageId(uri.to_string()), + SkillAuthority::new(SkillSourceKind::Orchestrator, CODEX_APPS_MCP_SERVER_NAME), + name, + description, + SkillResourceId::new(main_prompt), + ) + .with_display_path(uri), + ) +} + +fn validated_skill_uri(uri: &str, max_chars: usize) -> Option<&str> { + validated_skill_url(uri, max_chars).map(|_| uri) +} + +fn validated_skill_url(uri: &str, max_chars: usize) -> Option { + if uri.chars().count() > max_chars + || uri + .chars() + .any(|ch| ch.is_control() || ch.is_whitespace() || matches!(ch, '<' | '>')) + { + return None; + } + + let url = Url::parse(uri).ok()?; + let path_is_valid = url.path_segments().is_some_and(|segments| { + let segments = segments.collect::>(); + !segments.is_empty() && segments.iter().all(|segment| !segment.is_empty()) + }); + (url.scheme() == "skill" + && url.as_str() == uri + && url.host_str().is_some_and(|host| !host.is_empty()) + && url.username().is_empty() + && url.password().is_none() + && url.port().is_none() + && url.query().is_none() + && url.fragment().is_none() + && path_is_valid) + .then_some(url) +} + +fn resource_belongs_to_package(package: &str, resource: &str) -> bool { + let Some(package) = validated_skill_url(package, MAX_SKILL_PACKAGE_URI_CHARS) else { + return false; + }; + let Some(resource) = validated_skill_url(resource, MAX_SKILL_RESOURCE_URI_CHARS) else { + return false; + }; + + let Some(package_segments) = package.path_segments() else { + return false; + }; + let Some(resource_segments) = resource.path_segments() else { + return false; + }; + let package_segments = package_segments.collect::>(); + let resource_segments = resource_segments.collect::>(); + + package.scheme() == resource.scheme() + && package.host_str() == resource.host_str() + && resource_segments.len() > package_segments.len() + && resource_segments.starts_with(&package_segments) +} + +fn normalized_label(value: &str, max_chars: usize) -> Option { + let value = normalized_single_line(value, max_chars)?; + let invalid = value.is_empty() || value.chars().any(|ch| matches!(ch, '&' | '<' | '>')); + (!invalid).then_some(value) +} + +fn normalized_description(value: &str) -> Option { + normalized_single_line(value, MAX_SKILL_DESCRIPTION_CHARS).map(|value| { + value + .replace('&', "&") + .replace('<', "<") + .replace('>', ">") + }) +} + +fn normalized_single_line(value: &str, max_chars: usize) -> Option { + let value = value.split_whitespace().collect::>().join(" "); + let valid = value.chars().count() <= max_chars && !value.chars().any(char::is_control); + valid.then_some(value) +} + +fn main_prompt_uri(package_uri: &str) -> String { + format!("{}/SKILL.md", package_uri.trim_end_matches('/')) +} diff --git a/codex-rs/ext/skills/src/render.rs b/codex-rs/ext/skills/src/render.rs index 6b77d02140..d913e4161c 100644 --- a/codex-rs/ext/skills/src/render.rs +++ b/codex-rs/ext/skills/src/render.rs @@ -1,40 +1,18 @@ -use codex_core_skills::render_available_skills_body; -use codex_extension_api::ContextualUserFragment; -use codex_protocol::protocol::SKILLS_INSTRUCTIONS_CLOSE_TAG; -use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use codex_utils_string::take_bytes_at_char_boundary; use crate::catalog::SkillCatalog; +use crate::catalog::SkillCatalogEntry; +use crate::catalog::SkillSourceKind; +use crate::fragments::AvailableSkillsInstructions; const MAX_AVAILABLE_SKILLS_BYTES: usize = 8_000; const MAX_MAIN_PROMPT_BYTES: usize = 8_000; pub(crate) const MAX_SKILL_NAME_BYTES: usize = 256; pub(crate) const MAX_SKILL_PATH_BYTES: usize = 1_024; -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct AvailableSkillsFragment { - body: String, -} - -impl ContextualUserFragment for AvailableSkillsFragment { - fn role(&self) -> &'static str { - "developer" - } - - fn markers(&self) -> (&'static str, &'static str) { - Self::type_markers() - } - - fn body(&self) -> String { - self.body.clone() - } - - fn type_markers() -> (&'static str, &'static str) { - (SKILLS_INSTRUCTIONS_OPEN_TAG, SKILLS_INSTRUCTIONS_CLOSE_TAG) - } -} - -pub(crate) fn available_skills_fragment(catalog: &SkillCatalog) -> Option { +pub(crate) fn available_skills_fragment( + catalog: &SkillCatalog, +) -> Option { let mut total_bytes = 0usize; let mut omitted = 0usize; let mut skill_lines = Vec::new(); @@ -48,7 +26,7 @@ pub(crate) fn available_skills_fragment(catalog: &SkillCatalog) -> Option MAX_AVAILABLE_SKILLS_BYTES { omitted = omitted.saturating_add(1); @@ -68,16 +46,22 @@ pub(crate) fn available_skills_fragment(catalog: &SkillCatalog) -> Option String { +fn render_skill_line(entry: &SkillCatalogEntry, description: &str) -> String { + let locator_kind = match &entry.authority.kind { + SkillSourceKind::Host => "file", + SkillSourceKind::Executor => "environment resource", + SkillSourceKind::Orchestrator => "orchestrator resource", + SkillSourceKind::Custom(_) => "custom resource", + }; + let name = entry.name.as_str(); + let path = entry.rendered_path(); if description.is_empty() { - format!("- {name}: (file: {path})") + format!("- {name}: ({locator_kind}: {path})") } else { - format!("- {name}: {description} (file: {path})") + format!("- {name}: {description} ({locator_kind}: {path})") } } diff --git a/codex-rs/ext/skills/src/sources.rs b/codex-rs/ext/skills/src/sources.rs index 5ed51b878e..d558c13b3b 100644 --- a/codex-rs/ext/skills/src/sources.rs +++ b/codex-rs/ext/skills/src/sources.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use crate::catalog::SkillCatalog; use crate::catalog::SkillProviderError; +use crate::catalog::SkillProviderResult; use crate::catalog::SkillReadResult; use crate::catalog::SkillSearchResult; use crate::catalog::SkillSourceKind; @@ -39,15 +40,15 @@ impl SkillProviderSource { Self::new(SkillSourceKind::Executor, label, provider) } - pub fn remote(label: impl Into, provider: Arc) -> Self { - Self::new(SkillSourceKind::Remote, label, provider) + pub fn orchestrator(label: impl Into, provider: Arc) -> Self { + Self::new(SkillSourceKind::Orchestrator, label, provider) } fn should_list(&self, query: &SkillListQuery) -> bool { match &self.kind { SkillSourceKind::Host => query.include_host_skills, SkillSourceKind::Executor => !query.executor_roots.is_empty(), - SkillSourceKind::Remote => query.include_remote_skills, + SkillSourceKind::Orchestrator => query.include_orchestrator_skills, SkillSourceKind::Custom(_) => true, } } @@ -94,20 +95,54 @@ impl SkillProviders { self } - pub fn with_remote_provider(mut self, provider: Arc) -> Self { + pub fn with_orchestrator_provider(mut self, provider: Arc) -> Self { self.sources - .push(SkillProviderSource::remote("remote", provider)); + .push(SkillProviderSource::orchestrator("orchestrator", provider)); self } + pub(crate) fn has_orchestrator_provider(&self) -> bool { + self.sources + .iter() + .any(|source| source.kind == SkillSourceKind::Orchestrator) + } + pub(crate) async fn list_for_turn(&self, query: SkillListQuery) -> SkillCatalog { + self.list_matching(&query, |source| source.should_list(&query)) + .await + } + + pub(crate) async fn list_orchestrator_for_turn( + &self, + query: SkillListQuery, + ) -> SkillProviderResult { let mut catalog = SkillCatalog::default(); for source in self .sources .iter() - .filter(|source| source.should_list(&query)) + .filter(|source| source.kind == SkillSourceKind::Orchestrator) { + let source_catalog = source.provider.list(query.clone()).await.map_err(|err| { + SkillProviderError::new(format!( + "{} skills unavailable: {}", + source.label, err.message + )) + })?; + catalog.extend(source_catalog); + } + + Ok(catalog) + } + + async fn list_matching( + &self, + query: &SkillListQuery, + should_list: impl Fn(&SkillProviderSource) -> bool, + ) -> SkillCatalog { + let mut catalog = SkillCatalog::default(); + + for source in self.sources.iter().filter(|source| should_list(source)) { extend_catalog( &mut catalog, source.provider.list(query.clone()).await, diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index aa77cce822..0616ae4d71 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/codex-rs/ext/skills/src/state.rs @@ -1,29 +1,18 @@ -use codex_core::config::Config; use codex_protocol::capabilities::SelectedCapabilityRoot; +use std::future::Future; use std::sync::Mutex; +use tokio::sync::OnceCell; +use crate::SkillsExtensionConfig; use crate::catalog::SkillCatalog; use crate::catalog::SkillCatalogEntry; - -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) struct SkillsExtensionConfig { - pub(crate) include_instructions: bool, - pub(crate) bundled_skills_enabled: bool, -} - -impl SkillsExtensionConfig { - pub(crate) fn from_config(config: &Config) -> Self { - Self { - include_instructions: config.include_skill_instructions, - bundled_skills_enabled: config.bundled_skills_enabled(), - } - } -} +use crate::catalog::SkillProviderError; #[derive(Debug)] pub(crate) struct SkillsThreadState { config: Mutex, selected_roots: Vec, + orchestrator_catalog: OnceCell, } impl SkillsThreadState { @@ -34,6 +23,7 @@ impl SkillsThreadState { Self { config: Mutex::new(config), selected_roots, + orchestrator_catalog: OnceCell::new(), } } @@ -54,6 +44,21 @@ impl SkillsThreadState { pub(crate) fn selected_roots(&self) -> &[SelectedCapabilityRoot] { &self.selected_roots } + + pub(crate) async fn orchestrator_catalog_snapshot( + &self, + initialize: impl Future> + Send, + ) -> SkillCatalog { + self.orchestrator_catalog + .get_or_init(|| async { + initialize.await.unwrap_or_else(|err| SkillCatalog { + warnings: vec![err.message], + ..Default::default() + }) + }) + .await + .clone() + } } #[derive(Clone, Debug, Default, PartialEq, Eq)] diff --git a/codex-rs/ext/skills/src/tools/list.rs b/codex-rs/ext/skills/src/tools/list.rs new file mode 100644 index 0000000000..31cb68003d --- /dev/null +++ b/codex-rs/ext/skills/src/tools/list.rs @@ -0,0 +1,112 @@ +use codex_extension_api::ToolCall; +use codex_extension_api::ToolExecutor; +use codex_extension_api::ToolExecutorFuture; +use codex_extension_api::ToolName; +use codex_extension_api::ToolSpec; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; + +use crate::catalog::SkillCatalogEntry; +use crate::render::truncate_utf8_to_bytes; + +use super::MAX_HANDLE_BYTES; +use super::SkillToolAuthority; +use super::SkillToolContext; +use super::external_json_output; +use super::is_bounded_handle; +use super::parse_args; +use super::skill_function_tool; +use super::skill_tool_name; + +const TOOL_NAME: &str = "list"; +const MAX_WARNINGS: usize = 4; +const MAX_WARNING_BYTES: usize = 256; + +#[derive(Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +struct ListArgs { + authority: SkillToolAuthority, +} + +#[derive(Debug, Eq, JsonSchema, PartialEq, Serialize)] +#[schemars(deny_unknown_fields)] +struct ListedSkill { + authority: SkillToolAuthority, + package: String, + name: String, + description: String, + main_resource: String, +} + +#[derive(Debug, Eq, JsonSchema, PartialEq, Serialize)] +#[schemars(deny_unknown_fields)] +struct ListResponse { + skills: Vec, + warnings: Vec, +} + +#[derive(Clone)] +pub(super) struct ListTool { + pub(super) context: SkillToolContext, +} + +impl ToolExecutor for ListTool { + fn tool_name(&self) -> ToolName { + skill_tool_name(TOOL_NAME) + } + + fn spec(&self) -> ToolSpec { + skill_function_tool::( + TOOL_NAME, + "List enabled skills owned by the requested authority. Only orchestrator-owned skills are currently supported. Returns the opaque package and main-resource handles required by skills.read.", + ) + } + + fn handle(&self, call: ToolCall) -> ToolExecutorFuture<'_> { + Box::pin(async move { + let args: ListArgs = parse_args(&call)?; + let authority = args.authority.into_authority(); + let catalog = self.context.catalog(&call.turn_id, args.authority).await; + let response = ListResponse { + skills: catalog + .entries + .into_iter() + .filter(|entry| entry.enabled && entry.authority == authority) + .filter_map(listed_skill) + .collect(), + warnings: bounded_warnings(catalog.warnings), + }; + + external_json_output(&response) + }) + } +} + +fn listed_skill(entry: SkillCatalogEntry) -> Option { + let authority = SkillToolAuthority::from_authority(&entry.authority)?; + if !is_bounded_handle(&entry.id.0, MAX_HANDLE_BYTES) + || !is_bounded_handle(entry.main_prompt.as_str(), MAX_HANDLE_BYTES) + { + return None; + } + + Some(ListedSkill { + authority, + package: entry.id.0, + name: entry.name, + description: entry.description, + main_resource: entry.main_prompt.as_str().to_string(), + }) +} + +fn bounded_warnings(warnings: Vec) -> Vec { + warnings + .into_iter() + .take(MAX_WARNINGS) + .map(|warning| { + let (warning, _) = truncate_utf8_to_bytes(&warning, MAX_WARNING_BYTES); + warning + }) + .collect() +} diff --git a/codex-rs/ext/skills/src/tools/mod.rs b/codex-rs/ext/skills/src/tools/mod.rs new file mode 100644 index 0000000000..12b050338e --- /dev/null +++ b/codex-rs/ext/skills/src/tools/mod.rs @@ -0,0 +1,160 @@ +use std::sync::Arc; + +use codex_extension_api::FunctionCallError; +use codex_extension_api::JsonToolOutput; +use codex_extension_api::ResponsesApiTool; +use codex_extension_api::ToolCall; +use codex_extension_api::ToolExecutor; +use codex_extension_api::ToolName; +use codex_extension_api::ToolOutput; +use codex_extension_api::ToolSpec; +use codex_extension_api::parse_tool_input_schema; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::McpResourceClient; +use codex_tools::ResponsesApiNamespace; +use codex_tools::ResponsesApiNamespaceTool; +use codex_tools::default_namespace_description; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value; + +use crate::catalog::SkillAuthority; +use crate::catalog::SkillCatalog; +use crate::catalog::SkillSourceKind; +use crate::provider::SkillListQuery; +use crate::sources::SkillProviders; + +mod list; +mod read; +mod schema; + +const SKILLS_NAMESPACE: &str = "skills"; +const MAX_HANDLE_BYTES: usize = 2_048; + +pub(crate) fn skill_tools( + providers: SkillProviders, + mcp_resources: Option>, +) -> Vec>> { + let context = SkillToolContext { + providers, + mcp_resources, + }; + vec![ + Arc::new(list::ListTool { + context: context.clone(), + }), + Arc::new(read::ReadTool { context }), + ] +} + +#[derive(Clone)] +struct SkillToolContext { + providers: SkillProviders, + mcp_resources: Option>, +} + +impl SkillToolContext { + async fn catalog(&self, turn_id: &str, authority: SkillToolAuthority) -> SkillCatalog { + match authority { + SkillToolAuthority::Orchestrator => match self + .providers + .list_orchestrator_for_turn(SkillListQuery { + turn_id: turn_id.to_string(), + executor_roots: Vec::new(), + host: None, + include_host_skills: false, + include_bundled_skills: false, + include_orchestrator_skills: true, + mcp_resources: self.mcp_resources.clone(), + }) + .await + { + Ok(catalog) => catalog, + Err(err) => SkillCatalog { + warnings: vec![err.message], + ..Default::default() + }, + }, + } + } +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)] +#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] +enum SkillToolAuthority { + Orchestrator, +} + +impl SkillToolAuthority { + fn from_authority(authority: &SkillAuthority) -> Option { + if authority + != &SkillAuthority::new(SkillSourceKind::Orchestrator, CODEX_APPS_MCP_SERVER_NAME) + { + return None; + } + Some(Self::Orchestrator) + } + + fn into_authority(self) -> SkillAuthority { + match self { + Self::Orchestrator => { + SkillAuthority::new(SkillSourceKind::Orchestrator, CODEX_APPS_MCP_SERVER_NAME) + } + } + } +} + +fn skill_tool_name(name: &str) -> ToolName { + ToolName::namespaced(SKILLS_NAMESPACE, name) +} + +fn skill_function_tool(name: &str, description: &str) -> ToolSpec { + let tool = ResponsesApiTool { + name: name.to_string(), + description: description.to_string(), + strict: false, + defer_loading: None, + parameters: parse_tool_input_schema(&schema::input_schema_for::()) + .unwrap_or_else(|err| panic!("generated input schema for {name} should parse: {err}")), + output_schema: Some(schema::output_schema_for::()), + }; + + ToolSpec::Namespace(ResponsesApiNamespace { + name: SKILLS_NAMESPACE.to_string(), + description: default_namespace_description(SKILLS_NAMESPACE), + tools: vec![ResponsesApiNamespaceTool::Function(tool)], + }) +} + +fn parse_args Deserialize<'de>>(call: &ToolCall) -> Result { + let arguments = call.function_arguments()?; + let value = if arguments.trim().is_empty() { + Value::Object(serde_json::Map::new()) + } else { + serde_json::from_str(arguments) + .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))? + }; + serde_json::from_value(value).map_err(|err| FunctionCallError::RespondToModel(err.to_string())) +} + +fn validate_handle(name: &str, value: &str, max_bytes: usize) -> Result<(), FunctionCallError> { + if is_bounded_handle(value, max_bytes) { + return Ok(()); + } + + Err(FunctionCallError::RespondToModel(format!( + "{name} must be non-empty, contain no control characters, and be at most {max_bytes} bytes" + ))) +} + +fn is_bounded_handle(value: &str, max_bytes: usize) -> bool { + !value.is_empty() && value.len() <= max_bytes && !value.chars().any(char::is_control) +} + +fn external_json_output(value: &T) -> Result, FunctionCallError> { + let value = serde_json::to_value(value).map_err(|err| { + FunctionCallError::Fatal(format!("failed to serialize tool output: {err}")) + })?; + Ok(Box::new(JsonToolOutput::new(value).with_external_context())) +} diff --git a/codex-rs/ext/skills/src/tools/read.rs b/codex-rs/ext/skills/src/tools/read.rs new file mode 100644 index 0000000000..e8b29905b2 --- /dev/null +++ b/codex-rs/ext/skills/src/tools/read.rs @@ -0,0 +1,109 @@ +use codex_extension_api::FunctionCallError; +use codex_extension_api::ToolCall; +use codex_extension_api::ToolExecutor; +use codex_extension_api::ToolExecutorFuture; +use codex_extension_api::ToolName; +use codex_extension_api::ToolSpec; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; + +use crate::catalog::SkillPackageId; +use crate::catalog::SkillResourceId; +use crate::provider::SkillReadRequest; + +use super::MAX_HANDLE_BYTES; +use super::SkillToolAuthority; +use super::SkillToolContext; +use super::external_json_output; +use super::parse_args; +use super::skill_function_tool; +use super::skill_tool_name; +use super::validate_handle; + +const TOOL_NAME: &str = "read"; + +#[derive(Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +struct ReadArgs { + authority: SkillToolAuthority, + package: String, + resource: String, +} + +#[derive(Debug, Eq, JsonSchema, PartialEq, Serialize)] +#[schemars(deny_unknown_fields)] +struct ReadResponse { + resource: String, + contents: String, +} + +#[derive(Clone)] +pub(super) struct ReadTool { + pub(super) context: SkillToolContext, +} + +impl ToolExecutor for ReadTool { + fn tool_name(&self) -> ToolName { + skill_tool_name(TOOL_NAME) + } + + fn spec(&self) -> ToolSpec { + skill_function_tool::( + TOOL_NAME, + "Read one complete resource from an enabled skill. Pass the exact authority and package returned by skills.list; resource identifiers remain opaque and are routed to that authority.", + ) + } + + fn handle(&self, call: ToolCall) -> ToolExecutorFuture<'_> { + Box::pin(async move { + let args: ReadArgs = parse_args(&call)?; + let authority = args.authority.into_authority(); + validate_handle("package", &args.package, MAX_HANDLE_BYTES)?; + validate_handle("resource", &args.resource, MAX_HANDLE_BYTES)?; + + let catalog = self.context.catalog(&call.turn_id, args.authority).await; + let package_is_available = catalog.entries.iter().any(|entry| { + entry.enabled && entry.authority == authority && entry.id.0 == args.package + }); + if !package_is_available { + return Err(FunctionCallError::RespondToModel( + "skill package is not available from the requested authority".to_string(), + )); + } + + let requested_resource = SkillResourceId::new(args.resource); + let result = self + .context + .providers + .read(SkillReadRequest { + authority, + package: SkillPackageId(args.package), + resource: requested_resource.clone(), + host: None, + mcp_resources: self.context.mcp_resources.clone(), + }) + .await + .map_err(|err| { + tracing::warn!( + error = %err, + turn_id = %call.turn_id, + call_id = %call.call_id, + resource = requested_resource.as_str(), + "skills.read provider request failed" + ); + FunctionCallError::RespondToModel("failed to read skill resource".to_string()) + })?; + if result.resource != requested_resource { + return Err(FunctionCallError::Fatal( + "skill provider returned a different resource".to_string(), + )); + } + + external_json_output(&ReadResponse { + resource: result.resource.as_str().to_string(), + contents: result.contents, + }) + }) + } +} diff --git a/codex-rs/ext/skills/src/tools/schema.rs b/codex-rs/ext/skills/src/tools/schema.rs new file mode 100644 index 0000000000..d0e9425e47 --- /dev/null +++ b/codex-rs/ext/skills/src/tools/schema.rs @@ -0,0 +1,42 @@ +use schemars::JsonSchema; +use schemars::r#gen::SchemaSettings; +use serde_json::Map; +use serde_json::Value; + +pub(super) fn input_schema_for() -> Value { + schema_for::(/*option_add_null_type*/ false) +} + +pub(super) fn output_schema_for() -> Value { + schema_for::(/*option_add_null_type*/ true) +} + +fn schema_for(option_add_null_type: bool) -> Value { + let schema = SchemaSettings::draft2019_09() + .with(|settings| { + settings.inline_subschemas = true; + settings.option_add_null_type = option_add_null_type; + }) + .into_generator() + .into_root_schema_for::(); + let schema_value = serde_json::to_value(schema) + .unwrap_or_else(|err| panic!("generated skill tool schema should serialize: {err}")); + let Value::Object(mut schema_object) = schema_value else { + unreachable!("root tool schema must be an object"); + }; + + let mut tool_schema = Map::new(); + for key in [ + "properties", + "required", + "type", + "additionalProperties", + "$defs", + "definitions", + ] { + if let Some(value) = schema_object.remove(key) { + tool_schema.insert(key.to_string(), value); + } + } + Value::Object(tool_schema) +} diff --git a/codex-rs/ext/skills/tests/executor_file_system_authority.rs b/codex-rs/ext/skills/tests/executor_file_system_authority.rs index 026c3f3348..94c6f3821a 100644 --- a/codex-rs/ext/skills/tests/executor_file_system_authority.rs +++ b/codex-rs/ext/skills/tests/executor_file_system_authority.rs @@ -224,7 +224,8 @@ async fn executor_provider_reads_from_the_environment_instance_used_for_listing( host: None, include_host_skills: false, include_bundled_skills: true, - include_remote_skills: false, + include_orchestrator_skills: false, + mcp_resources: None, }) .await .expect("list executor skills"); @@ -246,6 +247,7 @@ async fn executor_provider_reads_from_the_environment_instance_used_for_listing( package: entry.id, resource: resource.clone(), host: None, + mcp_resources: None, }) .await .expect("read bound executor skill"), @@ -272,7 +274,6 @@ async fn selected_root_id_distinguishes_identical_executor_paths() { Arc::new(EnvironmentManager::default_for_tests()), /*restriction_product*/ None, ); - let catalog = provider .list(SkillListQuery { turn_id: "turn-1".to_string(), @@ -289,7 +290,8 @@ async fn selected_root_id_distinguishes_identical_executor_paths() { host: None, include_host_skills: false, include_bundled_skills: true, - include_remote_skills: false, + include_orchestrator_skills: false, + mcp_resources: None, }) .await .expect("list executor skills"); diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index dc1aa01407..dcbd828fce 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -4,26 +4,33 @@ use std::sync::Mutex; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; -use codex_core::config::Config; -use codex_core::config::ConfigBuilder; use codex_core_skills::HostLoadedSkills; -use codex_core_skills::SkillsLoadInput; -use codex_core_skills::SkillsManager; +use codex_core_skills::SKILLS_HOW_TO_USE_WITH_ABSOLUTE_PATHS; +use codex_core_skills::SKILLS_INTRO_WITH_ABSOLUTE_PATHS; +use codex_core_skills::SkillLoadOutcome; +use codex_core_skills::SkillMetadata; use codex_core_skills::injection::InjectedHostSkillPrompts; use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ThreadStartInput; use codex_extension_api::TurnInputContext; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::SKILLS_INSTRUCTIONS_CLOSE_TAG; use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SkillScope; use codex_protocol::user_input::UserInput; use codex_skills_extension::SkillProviders; +use codex_skills_extension::SkillsExtensionConfig; use codex_skills_extension::catalog::SkillAuthority; use codex_skills_extension::catalog::SkillCatalog; use codex_skills_extension::catalog::SkillCatalogEntry; use codex_skills_extension::catalog::SkillPackageId; +use codex_skills_extension::catalog::SkillProviderError; use codex_skills_extension::catalog::SkillReadResult; use codex_skills_extension::catalog::SkillResourceId; use codex_skills_extension::catalog::SkillSearchResult; @@ -35,14 +42,17 @@ use codex_skills_extension::provider::SkillProvider; use codex_skills_extension::provider::SkillProviderFuture; use codex_skills_extension::provider::SkillReadRequest; use codex_skills_extension::provider::SkillSearchRequest; +use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; type TestResult = Result<(), Box>; static NEXT_CODEX_HOME_ID: AtomicUsize = AtomicUsize::new(0); +const DEMO_SKILL_CONTENTS: &str = + "---\nname: demo\ndescription: Demo skill.\n---\n# Demo\n\nUse the demo skill.\n"; #[tokio::test] -async fn installed_extension_loads_host_skills_from_legacy_roots() -> TestResult { +async fn installed_extension_uses_host_loaded_skills() -> TestResult { let codex_home = test_codex_home(); let skill_path = codex_home.join("skills").join("demo").join("SKILL.md"); std::fs::create_dir_all( @@ -50,18 +60,11 @@ async fn installed_extension_loads_host_skills_from_legacy_roots() -> TestResult .parent() .ok_or("skill path should have a parent")?, )?; - std::fs::write( - &skill_path, - "---\nname: demo\ndescription: Demo skill.\n---\n# Demo\n\nUse the demo skill.\n", - )?; - let config = ConfigBuilder::default() - .codex_home(codex_home.clone()) - .fallback_cwd(Some(codex_home.clone())) - .build() - .await?; + std::fs::write(&skill_path, DEMO_SKILL_CONTENTS)?; + let config = default_config(); let mut builder = ExtensionRegistryBuilder::new(); - install(&mut builder); + install(&mut builder, skills_extension_config); let registry = builder.build(); let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("thread"); @@ -76,22 +79,21 @@ async fn installed_extension_loads_host_skills_from_legacy_roots() -> TestResult }) .await; - let manager = SkillsManager::new(config.codex_home.clone(), config.bundled_skills_enabled()); - let input = SkillsLoadInput::new( - config.cwd.clone(), - Vec::new(), - config.config_layer_stack.clone(), - config.bundled_skills_enabled(), - ); - let loaded_skills = Arc::new(manager.skills_for_config(&input, /*fs*/ None).await); - let skill_path_string = loaded_skills - .skills - .iter() - .find(|skill| skill.name == "demo") - .ok_or("demo skill should load")? - .path_to_skills_md - .to_string_lossy() - .into_owned(); + let skill_path = AbsolutePathBuf::try_from(skill_path)?; + let skill_path_string = skill_path.to_string_lossy().into_owned(); + let mut outcome = SkillLoadOutcome::default(); + outcome.skills.push(SkillMetadata { + name: "demo".to_string(), + description: "Demo skill.".to_string(), + short_description: None, + interface: None, + dependencies: None, + policy: None, + path_to_skills_md: skill_path, + scope: SkillScope::User, + plugin_id: None, + }); + let loaded_skills = Arc::new(outcome); let skill_prompt_path = skill_path_string.replace('\\', "/"); let turn_store = ExtensionData::new("turn-1"); turn_store.insert(HostLoadedSkills::new(Arc::clone(&loaded_skills))); @@ -112,14 +114,19 @@ async fn installed_extension_loads_host_skills_from_legacy_roots() -> TestResult ) .await; - assert_eq!(2, fragments.len()); - assert_eq!("developer", fragments[0].role()); - assert!(fragments[0].render().contains("demo")); - assert!(fragments[0].render().contains(&skill_prompt_path)); - assert_eq!("user", fragments[1].role()); - assert!(fragments[1].render().contains("demo")); - assert!(fragments[1].render().contains("# Demo")); - assert!(fragments[1].render().contains(&skill_prompt_path)); + let expected_catalog = format!( + "{SKILLS_INSTRUCTIONS_OPEN_TAG}\n## Skills\n{SKILLS_INTRO_WITH_ABSOLUTE_PATHS}\n### Available skills\n- demo: Demo skill. (file: {skill_prompt_path})\n### How to use skills\n{SKILLS_HOW_TO_USE_WITH_ABSOLUTE_PATHS}\n{SKILLS_INSTRUCTIONS_CLOSE_TAG}" + ); + let expected_skill = format!( + "\ndemo\n{skill_prompt_path}\n{DEMO_SKILL_CONTENTS}\n" + ); + assert_eq!( + vec![("developer", expected_catalog), ("user", expected_skill),], + fragments + .iter() + .map(|fragment| (fragment.role(), fragment.render())) + .collect::>() + ); let injected_host_skill_prompts = turn_store .get::() .ok_or("host skill prompt marker should be set")?; @@ -144,10 +151,12 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in warnings: Vec::new(), }, read_requests: Arc::clone(&read_requests), + list_calls: None, + fail_first_list: false, }); let providers = SkillProviders::new().with_executor_provider(executor_provider); let mut builder = ExtensionRegistryBuilder::new(); - install_with_providers(&mut builder, providers); + install_with_providers(&mut builder, providers, skills_extension_config); let registry = builder.build(); let session_store = ExtensionData::new("session"); @@ -160,7 +169,7 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in }, }]); let session_source = SessionSource::Cli; - let config = default_config().await?; + let config = default_config(); registry.thread_lifecycle_contributors()[0] .on_thread_start(ThreadStartInput { config: &config, @@ -181,6 +190,11 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in .starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) ); assert!(prompt_fragments[0].text().contains("lint-fix")); + assert!( + prompt_fragments[0] + .text() + .contains("(environment resource: skill://executor/lint-fix/SKILL.md)") + ); let turn_store = ExtensionData::new("turn-1"); let fragments = registry.turn_input_contributors()[0] @@ -239,6 +253,78 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in Ok(()) } +#[tokio::test] +async fn orchestrator_catalog_snapshot_caches_failure() -> TestResult { + let list_calls = Arc::new(AtomicUsize::new(0)); + let providers = + SkillProviders::new().with_orchestrator_provider(Arc::new(StaticSkillProvider { + catalog: SkillCatalog { + entries: vec![test_entry( + SkillSourceKind::Orchestrator, + "codex_apps", + "orchestrator/first", + "skill://orchestrator/first/SKILL.md", + )], + warnings: Vec::new(), + }, + read_requests: Arc::new(Mutex::new(Vec::new())), + list_calls: Some(Arc::clone(&list_calls)), + fail_first_list: true, + })); + let (event_tx, event_rx) = std::sync::mpsc::channel(); + let mut builder = + ExtensionRegistryBuilder::with_event_sink(Arc::new(ChannelEventSink(event_tx))); + install_with_providers(&mut builder, providers, skills_extension_config); + let registry = builder.build(); + let session_store = ExtensionData::new("session"); + let thread_store = ExtensionData::new("thread"); + let session_source = SessionSource::Cli; + let config = default_config(); + registry.thread_lifecycle_contributors()[0] + .on_thread_start(ThreadStartInput { + config: &config, + session_source: &session_source, + persistent_thread_state_available: true, + session_store: &session_store, + thread_store: &thread_store, + }) + .await; + + let initial_fragments = registry.context_contributors()[0] + .contribute(&session_store, &thread_store) + .await; + assert!(initial_fragments.is_empty()); + let EventMsg::Warning(warning) = event_rx.try_recv()?.msg else { + panic!("expected warning event"); + }; + assert_eq!( + warning.message, + "orchestrator skills unavailable: temporary orchestrator failure" + ); + + for turn_id in ["turn-1", "turn-2"] { + let fragments = registry.turn_input_contributors()[0] + .contribute( + TurnInputContext { + turn_id: turn_id.to_string(), + user_input: vec![UserInput::Text { + text: "$first".to_string(), + text_elements: Vec::new(), + }], + environments: Vec::new(), + }, + &session_store, + &thread_store, + &ExtensionData::new(turn_id), + ) + .await; + assert!(fragments.is_empty()); + } + assert_eq!(1, list_calls.load(Ordering::Relaxed)); + + Ok(()) +} + #[tokio::test] async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> TestResult { let read_requests = Arc::new(Mutex::new(Vec::new())); @@ -262,10 +348,12 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te warnings: Vec::new(), }, read_requests: Arc::clone(&read_requests), + list_calls: None, + fail_first_list: false, }); let providers = SkillProviders::new().with_executor_provider(executor_provider); let mut builder = ExtensionRegistryBuilder::new(); - install_with_providers(&mut builder, providers); + install_with_providers(&mut builder, providers, skills_extension_config); let registry = builder.build(); let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("thread"); @@ -282,7 +370,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te .collect::>(), ); let session_source = SessionSource::Cli; - let config = default_config().await?; + let config = default_config(); registry.thread_lifecycle_contributors()[0] .on_thread_start(ThreadStartInput { config: &config, @@ -346,15 +434,17 @@ async fn prompt_hidden_skill_can_still_be_invoked() -> TestResult { warnings: Vec::new(), }, read_requests: Arc::clone(&read_requests), + list_calls: None, + fail_first_list: false, }); let providers = SkillProviders::new().with_host_provider(provider); let mut builder = ExtensionRegistryBuilder::new(); - install_with_providers(&mut builder, providers); + install_with_providers(&mut builder, providers, skills_extension_config); let registry = builder.build(); let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("thread"); let session_source = SessionSource::Cli; - let config = default_config().await?; + let config = default_config(); registry.thread_lifecycle_contributors()[0] .on_thread_start(ThreadStartInput { config: &config, @@ -401,12 +491,33 @@ async fn prompt_hidden_skill_can_still_be_invoked() -> TestResult { struct StaticSkillProvider { catalog: SkillCatalog, read_requests: Arc>>, + list_calls: Option>, + fail_first_list: bool, +} + +struct ChannelEventSink(std::sync::mpsc::Sender); + +impl ExtensionEventSink for ChannelEventSink { + fn emit(&self, event: Event) { + let _ = self.0.send(event); + } } impl SkillProvider for StaticSkillProvider { fn list(&self, _query: SkillListQuery) -> SkillProviderFuture<'_, SkillCatalog> { + let list_call = self + .list_calls + .as_ref() + .map(|list_calls| list_calls.fetch_add(1, Ordering::Relaxed)); + let fail = self.fail_first_list && list_call == Some(0); let catalog = self.catalog.clone(); - Box::pin(async move { Ok(catalog) }) + Box::pin(async move { + if fail { + Err(SkillProviderError::new("temporary orchestrator failure")) + } else { + Ok(catalog) + } + }) } fn read(&self, request: SkillReadRequest) -> SkillProviderFuture<'_, SkillReadResult> { @@ -445,13 +556,24 @@ fn test_entry( .with_display_path(format!("skill://{package_id}/SKILL.md")) } -async fn default_config() -> std::io::Result { - let codex_home = test_codex_home(); - std::fs::create_dir_all(&codex_home)?; - let config = - Config::load_default_with_cli_overrides_for_codex_home(codex_home.clone(), vec![]).await?; - std::fs::remove_dir_all(codex_home)?; - Ok(config) +#[derive(Clone, Debug, Eq, PartialEq)] +struct TestConfig { + include_instructions: bool, + bundled_skills_enabled: bool, +} + +fn default_config() -> TestConfig { + TestConfig { + include_instructions: true, + bundled_skills_enabled: true, + } +} + +fn skills_extension_config(config: &TestConfig) -> SkillsExtensionConfig { + SkillsExtensionConfig { + include_instructions: config.include_instructions, + bundled_skills_enabled: config.bundled_skills_enabled, + } } fn test_codex_home() -> PathBuf { diff --git a/codex-rs/ext/web-search/Cargo.toml b/codex-rs/ext/web-search/Cargo.toml index 954ac3dac9..d007d6146d 100644 --- a/codex-rs/ext/web-search/Cargo.toml +++ b/codex-rs/ext/web-search/Cargo.toml @@ -13,7 +13,6 @@ doctest = false workspace = true [dependencies] -async-trait = { workspace = true } codex-api = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } diff --git a/codex-rs/ext/web-search/src/extension.rs b/codex-rs/ext/web-search/src/extension.rs index d081d4fb29..688b504dec 100644 --- a/codex-rs/ext/web-search/src/extension.rs +++ b/codex-rs/ext/web-search/src/extension.rs @@ -9,6 +9,7 @@ use codex_api::SearchSettings; use codex_core::config::Config; use codex_extension_api::ConfigContributor; use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadStartInput; @@ -80,12 +81,16 @@ fn search_settings(config: &Config, web_search_mode: WebSearchMode) -> SearchSet } } -#[async_trait::async_trait] impl ThreadLifecycleContributor for WebSearchExtension { - async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) { - input - .thread_store - .insert(WebSearchExtensionConfig::from(input.config)); + fn on_thread_start<'a>( + &'a self, + input: ThreadStartInput<'a, Config>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + input + .thread_store + .insert(WebSearchExtensionConfig::from(input.config)); + }) } } diff --git a/codex-rs/ext/web-search/src/tool.rs b/codex-rs/ext/web-search/src/tool.rs index 35f126dd6b..9b8dbf56d6 100644 --- a/codex-rs/ext/web-search/src/tool.rs +++ b/codex-rs/ext/web-search/src/tool.rs @@ -39,7 +39,6 @@ pub(crate) struct WebSearchTool { pub(crate) settings: SearchSettings, } -#[async_trait::async_trait] impl ToolExecutor for WebSearchTool { fn tool_name(&self) -> ToolName { ToolName::namespaced(WEB_NAMESPACE, RUN_TOOL_NAME) @@ -74,7 +73,13 @@ impl ToolExecutor for WebSearchTool { true } - async fn handle(&self, call: ToolCall) -> Result, FunctionCallError> { + fn handle(&self, call: ToolCall) -> codex_extension_api::ToolExecutorFuture<'_> { + Box::pin(self.handle_call(call)) + } +} + +impl WebSearchTool { + async fn handle_call(&self, call: ToolCall) -> Result, FunctionCallError> { let commands = parse_commands(&call)?; let command_action = command_action(&commands); let provider = self diff --git a/codex-rs/ext/web-search/web_run_description.md b/codex-rs/ext/web-search/web_run_description.md index 0f22f6df3c..77be9a0a03 100644 --- a/codex-rs/ext/web-search/web_run_description.md +++ b/codex-rs/ext/web-search/web_run_description.md @@ -48,6 +48,35 @@ Below is a list of scenarios where browsing the internet MUST be used. PAY CLOSE --- +## Citations + +Results from `web.run` include internal reference IDs such as `turn2search5`. Use +those reference IDs only in calls to `web.run`; do not expose them in the final +response. + +Cite sources in the final response using Markdown links: + +- Cite a single source as `[descriptive source title](https://example.com/page)`. +- Cite multiple sources with separate Markdown links, for example + `[first source](https://example.com/one), [second source](https://example.com/two)`. +- Link directly to the page that supports the claim. Do not link to search result + pages or use bare URLs. + +Formatting of citations: + +- Place each citation as near as possible to the claim it supports, normally at + the end of the sentence or paragraph and after punctuation. +- Do not place citations inside code fences. +- Do not put citations on a line by themselves or collect all citations at the + end of the response. + +If you browse the internet, cite statements supported by web sources. Each cited +source must directly support the associated claim. Prefer primary and +authoritative sources, and use sources from different domains when the response +benefits from multiple perspectives. + +--- + ## Special cases If these conflict with any other instructions, these should take precedence. @@ -74,7 +103,3 @@ Responses may not excessively quote or draw on a specific source. There are seve - You must avoid providing full articles, long verbatim passages, or extensive direct quotes due to copyright concerns. - If the user asked for a verbatim quote, the response should provide a short compliant excerpt and then answer with paraphrases and summaries. - Again, this limit does not apply to reddit content, as long as it's appropriately indicated that those are direct quotes and you link to the source. - ---- - -Cite your sources by using markdown links to the sites you used in your response. Do NOT use citations in the `turnX` style. diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index b7f666c21d..f59b89e9a8 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -70,21 +70,11 @@ impl FeatureConfig for MultiAgentV2ConfigToml { #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] #[serde(deny_unknown_fields)] -pub struct AppsMcpPathOverrideConfigToml { +pub(crate) struct RemovedAppsMcpPathOverrideConfigToml { #[serde(skip_serializing_if = "Option::is_none")] - pub enabled: Option, + enabled: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub path: Option, -} - -impl FeatureConfig for AppsMcpPathOverrideConfigToml { - fn enabled(&self) -> Option { - self.enabled.or(self.path.as_ref().map(|_| true)) - } - - fn set_enabled(&mut self, enabled: bool) { - self.enabled = Some(enabled); - } + path: Option, } #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 8ee05aed21..2f8e91d709 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -16,13 +16,13 @@ use toml::Table; mod feature_configs; mod legacy; -pub use feature_configs::AppsMcpPathOverrideConfigToml; pub use feature_configs::CodeModeConfigToml; pub use feature_configs::MultiAgentV2ConfigToml; pub use feature_configs::NetworkProxyConfigToml; pub use feature_configs::NetworkProxyDomainPermissionToml; pub use feature_configs::NetworkProxyModeToml; pub use feature_configs::NetworkProxyUnixSocketPermissionToml; +use feature_configs::RemovedAppsMcpPathOverrideConfigToml; use legacy::LegacyFeatureToggles; pub use legacy::legacy_feature_keys; @@ -143,7 +143,7 @@ pub enum Feature { Apps, /// Enable MCP apps. EnableMcpApps, - /// Use the new path for the host-owned apps MCP server. + /// Removed compatibility flag for the legacy Apps MCP path override. AppsMcpPathOverride, /// Removed compatibility flag retained as a no-op now that tool_search is always enabled. ToolSearch, @@ -177,12 +177,14 @@ pub enum Feature { RemotePlugin, /// Enable remote plugin sharing flows. PluginSharing, - /// Show the startup prompt for migrating external agent config into Codex. + /// Removed compatibility flag retained as a no-op. ExternalMigration, /// Allow the model to invoke the built-in image generation tool. ImageGeneration, /// Replace hosted image generation with the standalone image-generation extension. ImageGenExt, + /// Resize all inline data-URL images before recording them in history. + ResizeAllImages, /// Allow prompting and installing missing MCP dependencies. SkillMcpDependencyInstall, /// Removed compatibility flag for deleted skill env var dependency prompting. @@ -195,6 +197,8 @@ pub enum Feature { GuardianApproval, /// Enable persisted thread goals and automatic goal continuation. Goals, + /// Add current context-window metadata to model-visible context. + TokenBudget, /// Route MCP tool approval prompts through the MCP elicitation request path. ToolCallMcpElicitation, /// Prompt Codex Apps connector auth failures through MCP URL elicitations. @@ -443,7 +447,7 @@ impl Features { "apply_patch_freeform" => { continue; } - "tool_search" => { + "tool_search" | "apps_mcp_path_override" => { continue; } "image_detail_original" => { @@ -605,8 +609,9 @@ pub struct FeaturesToml { pub code_mode: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub multi_agent_v2: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub apps_mcp_path_override: Option>, + #[serde(default, rename = "apps_mcp_path_override", skip_serializing)] + #[schemars(skip)] + removed_apps_mcp_path_override: Option>, pub network_proxy: Option>, /// Boolean feature toggles keyed by canonical or legacy feature name. #[serde(flatten)] @@ -621,6 +626,13 @@ impl Features { } impl FeaturesToml { + /// Removes compatibility-only inputs that no longer affect runtime + /// behavior or belong in newly materialized config. + pub fn clear_removed_compatibility_entries(&mut self) { + self.removed_apps_mcp_path_override = None; + self.entries.remove("apps_mcp_path_override"); + } + pub fn entries(&self) -> BTreeMap { let mut entries = self.entries.clone(); if let Some(enabled) = self.code_mode.as_ref().and_then(FeatureToml::enabled) { @@ -629,13 +641,6 @@ impl FeaturesToml { if let Some(enabled) = self.multi_agent_v2.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::MultiAgentV2.key().to_string(), enabled); } - if let Some(enabled) = self - .apps_mcp_path_override - .as_ref() - .and_then(FeatureToml::enabled) - { - entries.insert(Feature::AppsMcpPathOverride.key().to_string(), enabled); - } if let Some(enabled) = self.network_proxy.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::NetworkProxy.key().to_string(), enabled); } @@ -643,10 +648,11 @@ impl FeaturesToml { } pub fn materialize_resolved_enabled(&mut self, features: &Features) { + self.clear_removed_compatibility_entries(); let Self { code_mode, multi_agent_v2, - apps_mcp_path_override, + removed_apps_mcp_path_override: _, network_proxy, entries, } = self; @@ -659,8 +665,6 @@ impl FeaturesToml { materialize_resolved_feature_enabled(code_mode, enabled); } else if spec.id == Feature::MultiAgentV2 { materialize_resolved_feature_enabled(multi_agent_v2, enabled); - } else if spec.id == Feature::AppsMcpPathOverride { - materialize_resolved_feature_enabled(apps_mcp_path_override, enabled); } else if spec.id == Feature::NetworkProxy { materialize_resolved_feature_enabled(network_proxy, enabled); } else { @@ -987,7 +991,7 @@ pub const FEATURES: &[FeatureSpec] = &[ FeatureSpec { id: Feature::AppsMcpPathOverride, key: "apps_mcp_path_override", - stage: Stage::UnderDevelopment, + stage: Stage::Removed, default_enabled: false, }, FeatureSpec { @@ -1071,11 +1075,7 @@ pub const FEATURES: &[FeatureSpec] = &[ FeatureSpec { id: Feature::ExternalMigration, key: "external_migration", - stage: Stage::Experimental { - name: "External migration", - menu_description: "Show a startup prompt when Codex detects migratable external agent config for this machine or project.", - announcement: "", - }, + stage: Stage::Removed, default_enabled: false, }, FeatureSpec { @@ -1090,6 +1090,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::ResizeAllImages, + key: "resize_all_images", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::SkillMcpDependencyInstall, key: "skill_mcp_dependency_install", @@ -1138,6 +1144,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::TokenBudget, + key: "token_budget", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::CollaborationModes, key: "collaboration_modes", @@ -1237,8 +1249,8 @@ pub const FEATURES: &[FeatureSpec] = &[ FeatureSpec { id: Feature::RemoteCompactionV2, key: "remote_compaction_v2", - stage: Stage::UnderDevelopment, - default_enabled: false, + stage: Stage::Stable, + default_enabled: true, }, FeatureSpec { id: Feature::WorkspaceDependencies, diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index aa5e0174ba..0a4f925f4f 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -83,6 +83,37 @@ fn plugin_hooks_is_removed_and_disabled_by_default() { assert_eq!(feature_for_key("plugin_hooks"), Some(Feature::PluginHooks)); } +#[test] +fn external_migration_is_removed_and_disabled_by_default() { + assert_eq!(Feature::ExternalMigration.stage(), Stage::Removed); + assert_eq!(Feature::ExternalMigration.default_enabled(), false); + assert_eq!( + feature_for_key("external_migration"), + Some(Feature::ExternalMigration) + ); +} + +#[test] +fn removed_apps_mcp_path_override_shapes_are_ignored() { + let features = [ + toml::from_str::("apps_mcp_path_override = true") + .expect("boolean compatibility form should deserialize"), + toml::from_str::( + r#" +[apps_mcp_path_override] +enabled = true +path = "/custom/mcp" +"#, + ) + .expect("structured compatibility form should deserialize"), + ]; + + assert_eq!( + features.map(|features| features.entries()), + [BTreeMap::new(), BTreeMap::new()] + ); +} + #[test] fn code_mode_only_requires_code_mode() { let mut features = Features::with_defaults(); @@ -101,23 +132,6 @@ fn guardian_approval_is_stable_and_enabled_by_default() { assert_eq!(Feature::GuardianApproval.default_enabled(), true); } -#[test] -fn external_migration_is_experimental_and_disabled_by_default() { - let spec = Feature::ExternalMigration.info(); - let stage = spec.stage; - - assert!(matches!(stage, Stage::Experimental { .. })); - assert_eq!(stage.experimental_menu_name(), Some("External migration")); - assert_eq!( - stage.experimental_menu_description(), - Some( - "Show a startup prompt when Codex detects migratable external agent config for this machine or project." - ) - ); - assert_eq!(stage.experimental_announcement(), None); - assert_eq!(Feature::ExternalMigration.default_enabled(), false); -} - #[test] fn request_permissions_is_under_development() { assert_eq!( @@ -136,16 +150,6 @@ fn request_permissions_tool_is_under_development() { assert_eq!(Feature::RequestPermissionsTool.default_enabled(), false); } -#[test] -fn remote_compaction_v2_is_under_development() { - assert_eq!(Feature::RemoteCompactionV2.stage(), Stage::UnderDevelopment); - assert_eq!(Feature::RemoteCompactionV2.default_enabled(), false); - assert_eq!( - feature_for_key("remote_compaction_v2"), - Some(Feature::RemoteCompactionV2) - ); -} - #[test] fn terminal_resize_reflow_is_experimental_and_enabled_by_default() { assert_eq!( diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index e7b5502ebc..ff12aaff61 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -166,6 +166,7 @@ async fn login_with_access_token_writes_only_personal_access_token() { last_refresh: None, agent_identity: None, personal_access_token: Some("at-login-test".to_string()), + bedrock_api_key: None, } ); assert_eq!(auth.resolved_mode(), AuthMode::PersonalAccessToken); @@ -325,6 +326,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { last_refresh: Some(last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }, auth_dot_json ); @@ -367,6 +369,7 @@ fn logout_removes_auth_file() -> Result<(), std::io::Error> { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; super::save_auth(dir.path(), &auth_dot_json, AuthCredentialsStoreMode::File)?; let auth_file = get_auth_file(dir.path()); @@ -1112,6 +1115,7 @@ async fn enforce_login_restrictions_logs_out_for_agent_identity_workspace_mismat last_refresh: None, agent_identity: Some(agent_identity), personal_access_token: None, + bedrock_api_key: None, }, AuthCredentialsStoreMode::File, ) diff --git a/codex-rs/login/src/auth/bedrock_api_key.rs b/codex-rs/login/src/auth/bedrock_api_key.rs new file mode 100644 index 0000000000..803b7a01ec --- /dev/null +++ b/codex-rs/login/src/auth/bedrock_api_key.rs @@ -0,0 +1,42 @@ +use std::path::Path; + +use codex_config::types::AuthCredentialsStoreMode; +use serde::Deserialize; +use serde::Serialize; + +use super::manager::save_auth; +use super::storage::AuthDotJson; +use codex_app_server_protocol::AuthMode; + +/// Managed Amazon Bedrock API key persisted in `auth.json`. +#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] +pub struct BedrockApiKeyAuth { + pub api_key: String, + pub region: String, +} + +/// Writes an `auth.json` that contains only the Amazon Bedrock API key auth. +pub fn login_with_bedrock_api_key( + codex_home: &Path, + api_key: &str, + region: &str, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result<()> { + let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::BedrockApiKey), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: Some(BedrockApiKeyAuth { + api_key: api_key.to_string(), + region: region.to_string(), + }), + }; + save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) +} + +#[cfg(test)] +#[path = "bedrock_api_key_tests.rs"] +mod tests; diff --git a/codex-rs/login/src/auth/bedrock_api_key_tests.rs b/codex-rs/login/src/auth/bedrock_api_key_tests.rs new file mode 100644 index 0000000000..372baadec7 --- /dev/null +++ b/codex-rs/login/src/auth/bedrock_api_key_tests.rs @@ -0,0 +1,162 @@ +use codex_app_server_protocol::AuthMode; +use codex_config::types::AuthCredentialsStoreMode; +use pretty_assertions::assert_eq; +use tempfile::tempdir; + +use super::*; +use crate::auth::AuthManager; +use crate::auth::CodexAuth; +use crate::auth::storage::AuthStorageBackend; +use crate::auth::storage::FileAuthStorage; + +fn api_key_auth() -> AuthDotJson { + AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), + openai_api_key: Some("sk-test-key".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: None, + } +} + +fn bedrock_only_auth() -> AuthDotJson { + AuthDotJson { + auth_mode: None, + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: Some(bedrock_auth()), + } +} + +fn bedrock_auth() -> BedrockApiKeyAuth { + BedrockApiKeyAuth { + api_key: "bedrock-api-key-test".to_string(), + region: "us-east-1".to_string(), + } +} + +#[tokio::test] +async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + storage.save(&api_key_auth())?; + login_with_bedrock_api_key( + codex_home.path(), + "bedrock-api-key-test", + "us-east-1", + AuthCredentialsStoreMode::File, + )?; + + let auth_manager = AuthManager::new( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + + let loaded = storage.load()?.expect("auth should be stored"); + let expected = AuthDotJson { + auth_mode: Some(AuthMode::BedrockApiKey), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: Some(bedrock_auth()), + }; + assert_eq!(loaded, expected); + assert_eq!(auth_manager.auth_mode(), Some(AuthMode::BedrockApiKey)); + assert_eq!( + auth_manager.auth_cached().and_then(|auth| match auth { + CodexAuth::BedrockApiKey(auth) => Some(auth), + CodexAuth::ApiKey(_) + | CodexAuth::Chatgpt(_) + | CodexAuth::ChatgptAuthTokens(_) + | CodexAuth::AgentIdentity(_) + | CodexAuth::PersonalAccessToken(_) => None, + }), + Some(bedrock_auth()) + ); + Ok(()) +} + +#[tokio::test] +async fn logout_removes_bedrock_auth() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + login_with_bedrock_api_key( + codex_home.path(), + "bedrock-api-key-test", + "us-east-1", + AuthCredentialsStoreMode::File, + )?; + let auth_manager = AuthManager::new( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + + assert!(auth_manager.logout().await?); + + assert_eq!(storage.load()?, None); + assert_eq!(auth_manager.auth_cached(), None); + Ok(()) +} + +#[tokio::test] +async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + storage.save(&bedrock_only_auth())?; + + let auth_manager = AuthManager::new( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + + assert_eq!(auth_manager.auth_mode(), Some(AuthMode::BedrockApiKey)); + assert_eq!( + auth_manager.auth_cached().and_then(|auth| match auth { + CodexAuth::BedrockApiKey(auth) => Some(auth), + CodexAuth::ApiKey(_) + | CodexAuth::Chatgpt(_) + | CodexAuth::ChatgptAuthTokens(_) + | CodexAuth::AgentIdentity(_) + | CodexAuth::PersonalAccessToken(_) => None, + }), + Some(bedrock_auth()) + ); + Ok(()) +} + +#[tokio::test] +async fn login_with_api_key_clears_bedrock_api_key() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + login_with_bedrock_api_key( + codex_home.path(), + "bedrock-api-key-test", + "us-east-1", + AuthCredentialsStoreMode::File, + )?; + + crate::auth::login_with_api_key( + codex_home.path(), + "sk-test-key", + AuthCredentialsStoreMode::File, + )?; + + assert_eq!(storage.load()?, Some(api_key_auth())); + Ok(()) +} diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 9e255a99a3..383011f246 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -29,6 +29,7 @@ use super::access_token::classify_codex_access_token; use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; pub use crate::auth::agent_identity::AgentIdentityAuth; +pub use crate::auth::bedrock_api_key::BedrockApiKeyAuth; pub use crate::auth::personal_access_token::PersonalAccessTokenAuth; pub use crate::auth::storage::AgentIdentityAuthRecord; pub use crate::auth::storage::AuthDotJson; @@ -57,12 +58,14 @@ pub enum CodexAuth { ChatgptAuthTokens(ChatgptAuthTokens), AgentIdentity(AgentIdentityAuth), PersonalAccessToken(PersonalAccessTokenAuth), + BedrockApiKey(BedrockApiKeyAuth), } impl PartialEq for CodexAuth { fn eq(&self, other: &Self) -> bool { match (self, other) { (Self::PersonalAccessToken(a), Self::PersonalAccessToken(b)) => a == b, + (Self::BedrockApiKey(a), Self::BedrockApiKey(b)) => a == b, _ => self.api_auth_mode() == other.api_auth_mode(), } } @@ -235,6 +238,14 @@ impl CodexAuth { }; return Self::from_personal_access_token(personal_access_token).await; } + if auth_mode == ApiAuthMode::BedrockApiKey { + let Some(auth) = auth_dot_json.bedrock_api_key else { + return Err(std::io::Error::other( + "Bedrock API key auth is missing a Bedrock API key.", + )); + }; + return Ok(Self::BedrockApiKey(auth)); + } let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); let state = ChatgptAuthState { @@ -255,6 +266,7 @@ impl CodexAuth { ApiAuthMode::PersonalAccessToken => { unreachable!("personal access token mode is handled above") } + ApiAuthMode::BedrockApiKey => unreachable!("bedrock api key mode is handled above"), } } @@ -296,6 +308,7 @@ impl CodexAuth { Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => AuthMode::Chatgpt, Self::AgentIdentity(_) => AuthMode::AgentIdentity, Self::PersonalAccessToken(_) => AuthMode::PersonalAccessToken, + Self::BedrockApiKey(_) => AuthMode::BedrockApiKey, } } @@ -306,6 +319,7 @@ impl CodexAuth { Self::ChatgptAuthTokens(_) => ApiAuthMode::ChatgptAuthTokens, Self::AgentIdentity(_) => ApiAuthMode::AgentIdentity, Self::PersonalAccessToken(_) => ApiAuthMode::PersonalAccessToken, + Self::BedrockApiKey(_) => ApiAuthMode::BedrockApiKey, } } @@ -346,7 +360,8 @@ impl CodexAuth { Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) | Self::AgentIdentity(_) - | Self::PersonalAccessToken(_) => None, + | Self::PersonalAccessToken(_) + | Self::BedrockApiKey(_) => None, } } @@ -375,6 +390,9 @@ impl CodexAuth { "agent identity auth does not expose a bearer token", )), Self::PersonalAccessToken(auth) => Ok(auth.access_token().to_string()), + Self::BedrockApiKey(_) => Err(std::io::Error::other( + "Bedrock API key auth does not expose a Codex bearer token", + )), } } @@ -447,7 +465,10 @@ impl CodexAuth { let state = match self { Self::Chatgpt(auth) => &auth.state, Self::ChatgptAuthTokens(auth) => &auth.state, - Self::ApiKey(_) | Self::AgentIdentity(_) | Self::PersonalAccessToken(_) => return None, + Self::ApiKey(_) + | Self::AgentIdentity(_) + | Self::PersonalAccessToken(_) + | Self::BedrockApiKey(_) => return None, }; #[expect(clippy::unwrap_used)] state.auth_dot_json.lock().unwrap().clone() @@ -472,6 +493,7 @@ impl CodexAuth { last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; let client = create_client(); @@ -589,6 +611,7 @@ pub fn login_with_api_key( last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) } @@ -612,6 +635,7 @@ pub async fn login_with_access_token( last_refresh: None, agent_identity: None, personal_access_token: Some(access_token.to_string()), + bedrock_api_key: None, } } CodexAccessToken::AgentIdentityJwt(jwt) => { @@ -627,6 +651,7 @@ pub async fn login_with_access_token( last_refresh: None, agent_identity: Some(jwt.to_string()), personal_access_token: None, + bedrock_api_key: None, } } }; @@ -698,7 +723,8 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result< if let Some(required_method) = config.forced_login_method { let method_violation = match (required_method, auth.auth_mode()) { - (ForcedLoginMethod::Api, AuthMode::ApiKey) => None, + (ForcedLoginMethod::Api, AuthMode::ApiKey) + | (ForcedLoginMethod::Api, AuthMode::BedrockApiKey) => None, (ForcedLoginMethod::Chatgpt, AuthMode::Chatgpt) | (ForcedLoginMethod::Chatgpt, AuthMode::ChatgptAuthTokens) | (ForcedLoginMethod::Chatgpt, AuthMode::AgentIdentity) @@ -710,7 +736,8 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result< "API key login is required, but ChatGPT is currently being used. Logging out." .to_string(), ), - (ForcedLoginMethod::Chatgpt, AuthMode::ApiKey) => Some( + (ForcedLoginMethod::Chatgpt, AuthMode::ApiKey) + | (ForcedLoginMethod::Chatgpt, AuthMode::BedrockApiKey) => Some( "ChatGPT login is required, but an API key is currently being used. Logging out." .to_string(), ), @@ -727,7 +754,9 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result< if let Some(expected_account_ids) = config.forced_chatgpt_workspace_id.as_deref() { let chatgpt_account_id = match &auth { - CodexAuth::ApiKey(_) | CodexAuth::PersonalAccessToken(_) => return Ok(()), + CodexAuth::ApiKey(_) + | CodexAuth::PersonalAccessToken(_) + | CodexAuth::BedrockApiKey(_) => return Ok(()), CodexAuth::AgentIdentity(_) => auth.get_account_id(), CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { let token_data = match auth.get_token_data() { @@ -1043,6 +1072,7 @@ impl AuthDotJson { last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }) } @@ -1059,13 +1089,16 @@ impl AuthDotJson { Self::from_external_tokens(&external) } - fn resolved_mode(&self) -> ApiAuthMode { + pub(super) fn resolved_mode(&self) -> ApiAuthMode { if let Some(mode) = self.auth_mode { return mode; } if self.personal_access_token.is_some() { return ApiAuthMode::PersonalAccessToken; } + if self.bedrock_api_key.is_some() { + return ApiAuthMode::BedrockApiKey; + } if self.openai_api_key.is_some() { return ApiAuthMode::ApiKey; } @@ -1590,6 +1623,7 @@ impl AuthManager { _ => false, }, (ApiAuthMode::PersonalAccessToken, ApiAuthMode::PersonalAccessToken) => a == b, + (ApiAuthMode::BedrockApiKey, ApiAuthMode::BedrockApiKey) => a == b, _ => false, }, _ => false, @@ -1851,7 +1885,8 @@ impl AuthManager { } CodexAuth::ApiKey(_) | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) => Ok(()), + | CodexAuth::PersonalAccessToken(_) + | CodexAuth::BedrockApiKey(_) => Ok(()), }; if let Err(RefreshTokenError::Permanent(error)) = &result { self.record_permanent_refresh_failure_if_unchanged(&attempted_auth, error); diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index 9bf6c350f7..2e18d530f7 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -1,5 +1,6 @@ mod access_token; mod agent_identity; +mod bedrock_api_key; pub mod default_client; pub mod error; mod personal_access_token; @@ -10,6 +11,8 @@ mod external_bearer; mod manager; mod revoke; +pub use bedrock_api_key::BedrockApiKeyAuth; +pub use bedrock_api_key::login_with_bedrock_api_key; pub use error::RefreshTokenFailedError; pub use error::RefreshTokenFailedReason; pub use manager::*; diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index d0e606a8cc..2edb4ca301 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -18,6 +18,7 @@ use std::sync::Arc; use std::sync::Mutex; use tracing::warn; +use super::BedrockApiKeyAuth; use crate::token_data::TokenData; use codex_agent_identity::AgentIdentityJwtClaims; use codex_agent_identity::decode_agent_identity_jwt; @@ -48,6 +49,9 @@ pub struct AuthDotJson { #[serde(default, skip_serializing_if = "Option::is_none")] pub personal_access_token: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub bedrock_api_key: Option, } #[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index debe0e18d3..c3491a8c4a 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -20,6 +20,7 @@ async fn file_storage_load_returns_auth_dot_json() -> anyhow::Result<()> { last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; storage @@ -42,6 +43,7 @@ async fn file_storage_save_persists_auth_dot_json() -> anyhow::Result<()> { last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; let file = get_auth_file(codex_home.path()); @@ -76,6 +78,7 @@ async fn file_storage_round_trips_agent_identity_auth() -> anyhow::Result<()> { last_refresh: None, agent_identity: Some(agent_identity), personal_access_token: None, + bedrock_api_key: None, }; storage.save(&auth_dot_json)?; @@ -96,6 +99,7 @@ async fn file_storage_round_trips_personal_access_token_auth() -> anyhow::Result last_refresh: None, agent_identity: None, personal_access_token: Some("at-example".to_string()), + bedrock_api_key: None, }; storage.save(&auth_dot_json)?; @@ -146,6 +150,7 @@ fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; let storage = create_auth_storage(dir.path().to_path_buf(), AuthCredentialsStoreMode::File); storage.save(&auth_dot_json)?; @@ -171,6 +176,7 @@ fn ephemeral_storage_save_load_delete_is_in_memory_only() -> anyhow::Result<()> last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; storage.save(&auth_dot_json)?; @@ -271,6 +277,7 @@ fn auth_with_prefix(prefix: &str) -> AuthDotJson { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, } } @@ -297,6 +304,7 @@ fn keyring_auth_storage_load_returns_deserialized_auth() -> anyhow::Result<()> { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; seed_keyring_with_auth( &mock_keyring, @@ -341,6 +349,7 @@ fn keyring_auth_storage_save_persists_and_removes_fallback_file() -> anyhow::Res last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; storage.save(&auth)?; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 990cf8b80e..b857592ade 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -40,6 +40,7 @@ pub use auth::enforce_login_restrictions; pub use auth::load_auth_dot_json; pub use auth::login_with_access_token; pub use auth::login_with_api_key; +pub use auth::login_with_bedrock_api_key; pub use auth::logout; pub use auth::logout_with_revoke; pub use auth::read_codex_access_token_from_env; diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 5a1833b58c..a91017bea3 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -823,6 +823,7 @@ pub(crate) async fn persist_tokens_async( last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth(&codex_home, &auth, auth_credentials_store_mode)?; Ok::<_, io::Error>((previous_auth, auth)) @@ -1327,6 +1328,7 @@ mod tests { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, } } diff --git a/codex-rs/login/tests/suite/auth_refresh.rs b/codex-rs/login/tests/suite/auth_refresh.rs index b30d738cd4..c86338e701 100644 --- a/codex-rs/login/tests/suite/auth_refresh.rs +++ b/codex-rs/login/tests/suite/auth_refresh.rs @@ -56,6 +56,7 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -121,6 +122,7 @@ async fn refresh_token_refreshes_when_auth_is_unchanged() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -187,6 +189,7 @@ async fn auth_refreshes_when_access_token_is_near_expiry() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -238,6 +241,7 @@ async fn auth_skips_access_token_outside_refresh_window() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -275,6 +279,7 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -286,6 +291,7 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth( ctx.codex_home.path(), @@ -342,6 +348,7 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -354,6 +361,7 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth( ctx.codex_home.path(), @@ -414,6 +422,7 @@ async fn returns_fresh_tokens_as_is() -> Result<()> { last_refresh: Some(stale_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -463,6 +472,7 @@ async fn refreshes_token_when_access_token_is_expired() -> Result<()> { last_refresh: Some(fresh_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -514,6 +524,7 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { last_refresh: Some(stale_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -526,6 +537,7 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { last_refresh: Some(fresh_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth( ctx.codex_home.path(), @@ -579,6 +591,7 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul last_refresh: Some(stale_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -591,6 +604,7 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul last_refresh: Some(fresh_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth( ctx.codex_home.path(), @@ -642,6 +656,7 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -696,6 +711,7 @@ async fn refresh_token_does_not_retry_after_permanent_failure() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -764,6 +780,7 @@ async fn refresh_token_does_not_retry_after_bad_request_reused_failure() -> Resu last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -832,6 +849,7 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -855,6 +873,7 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< last_refresh: Some(fresh_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth( ctx.codex_home.path(), @@ -915,6 +934,7 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -969,6 +989,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -980,6 +1001,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth( ctx.codex_home.path(), @@ -1065,6 +1087,7 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&initial_auth).await?; @@ -1077,6 +1100,7 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth( ctx.codex_home.path(), @@ -1136,6 +1160,7 @@ async fn unauthorized_recovery_requires_chatgpt_auth() -> Result<()> { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; ctx.write_auth(&auth).await?; diff --git a/codex-rs/login/tests/suite/logout.rs b/codex-rs/login/tests/suite/logout.rs index 2dd1bef83f..5d75803977 100644 --- a/codex-rs/login/tests/suite/logout.rs +++ b/codex-rs/login/tests/suite/logout.rs @@ -196,6 +196,7 @@ fn chatgpt_auth_with_refresh_token(refresh_token: &str) -> AuthDotJson { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, } } diff --git a/codex-rs/memories/write/Cargo.toml b/codex-rs/memories/write/Cargo.toml index ebe6b0f2a8..bf031a6a16 100644 --- a/codex-rs/memories/write/Cargo.toml +++ b/codex-rs/memories/write/Cargo.toml @@ -21,6 +21,7 @@ codex-config = { workspace = true } codex-features = { workspace = true } codex-git-utils = { workspace = true } codex-login = { workspace = true } +codex-model-provider = { workspace = true } codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-rollout = { workspace = true } @@ -39,6 +40,7 @@ tracing = { workspace = true, features = ["log"] } uuid = { workspace = true, features = ["v4", "v5"] } [dev-dependencies] +codex-model-provider-info = { workspace = true } codex-models-manager = { workspace = true } core_test_support = { workspace = true } pretty_assertions = { workspace = true } diff --git a/codex-rs/memories/write/src/lib.rs b/codex-rs/memories/write/src/lib.rs index 6764ebf5c6..ff86068ba6 100644 --- a/codex-rs/memories/write/src/lib.rs +++ b/codex-rs/memories/write/src/lib.rs @@ -76,7 +76,6 @@ signal to remove stale memories derived only from those resources. } mod stage_one { - pub(super) const MODEL: &str = "gpt-5.4-mini"; pub(super) const REASONING_EFFORT: codex_protocol::openai_models::ReasoningEffort = codex_protocol::openai_models::ReasoningEffort::Low; pub(super) const CONCURRENCY_LIMIT: usize = 8; @@ -101,7 +100,6 @@ mod stage_one { } mod stage_two { - pub(super) const MODEL: &str = "gpt-5.4"; pub(super) const REASONING_EFFORT: codex_protocol::openai_models::ReasoningEffort = codex_protocol::openai_models::ReasoningEffort::Medium; pub(super) const JOB_LEASE_SECONDS: i64 = 3_600; diff --git a/codex-rs/memories/write/src/phase1.rs b/codex-rs/memories/write/src/phase1.rs index 21b20759b1..693b7a1680 100644 --- a/codex-rs/memories/write/src/phase1.rs +++ b/codex-rs/memories/write/src/phase1.rs @@ -190,11 +190,12 @@ async fn build_request_context( context: &MemoryStartupContext, config: &Config, ) -> StageOneRequestContext { - let model_name = config - .memories - .extract_model - .clone() - .unwrap_or(crate::stage_one::MODEL.to_string()); + let model_name = config.memories.extract_model.clone().unwrap_or_else(|| { + context + .provider() + .memory_extraction_preferred_model() + .to_string() + }); context .stage_one_request_context(config, &model_name, crate::stage_one::REASONING_EFFORT) .await diff --git a/codex-rs/memories/write/src/phase2.rs b/codex-rs/memories/write/src/phase2.rs index c78032d9c2..0c588a6070 100644 --- a/codex-rs/memories/write/src/phase2.rs +++ b/codex-rs/memories/write/src/phase2.rs @@ -16,6 +16,7 @@ use crate::workspace::write_workspace_diff; use codex_config::Constrained; use codex_core::config::Config; use codex_features::Feature; +use codex_model_provider::ModelProvider; use codex_protocol::ThreadId; use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::AskForApproval; @@ -76,7 +77,7 @@ pub async fn run(context: Arc, config: Arc) { } // 3. Build the locked-down config used by the consolidation agent. - let Some(agent_config) = agent::get_config(config.as_ref()) else { + let Some(agent_config) = agent::get_config(config.as_ref(), context.provider()) else { // If we can't get the config, we can't consolidate. tracing::error!("failed to get agent config"); job::failed( @@ -297,7 +298,7 @@ mod agent { use super::*; use tracing::warn; - pub(super) fn get_config(config: &Config) -> Option { + pub(super) fn get_config(config: &Config, provider: &dyn ModelProvider) -> Option { let root = memory_root(&config.codex_home); let mut agent_config = config.clone(); @@ -339,7 +340,7 @@ mod agent { .memories .consolidation_model .clone() - .unwrap_or(crate::stage_two::MODEL.to_string()), + .unwrap_or_else(|| provider.memory_consolidation_preferred_model().to_string()), ); agent_config.model_reasoning_effort = Some(crate::stage_two::REASONING_EFFORT); diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index d524a2e224..0b6718f01e 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -13,6 +13,9 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_login::default_client::originator; +use codex_model_provider::ModelProvider; +use codex_model_provider::SharedModelProvider; +use codex_model_provider::create_model_provider; use codex_otel::SessionTelemetry; use codex_otel::TelemetryAuthMode; use codex_protocol::SessionId; @@ -68,6 +71,7 @@ pub(crate) struct MemoryStartupContext { thread: Arc, thread_manager: Arc, auth_manager: Arc, + provider: SharedModelProvider, session_telemetry: SessionTelemetry, } @@ -79,6 +83,51 @@ impl MemoryStartupContext { thread: Arc, config: &Config, source: SessionSource, + ) -> Self { + let provider = create_model_provider( + config.model_provider.clone(), + Some(Arc::clone(&auth_manager)), + ); + Self::new_with_provider( + thread_manager, + auth_manager, + thread_id, + thread, + config, + source, + provider, + ) + } + + #[cfg(test)] + pub(crate) fn new_for_testing( + thread_manager: Arc, + auth_manager: Arc, + thread_id: ThreadId, + thread: Arc, + config: &Config, + source: SessionSource, + provider: SharedModelProvider, + ) -> Self { + Self::new_with_provider( + thread_manager, + auth_manager, + thread_id, + thread, + config, + source, + provider, + ) + } + + fn new_with_provider( + thread_manager: Arc, + auth_manager: Arc, + thread_id: ThreadId, + thread: Arc, + config: &Config, + source: SessionSource, + provider: SharedModelProvider, ) -> Self { let auth = auth_manager.auth_cached(); let auth = auth.as_ref(); @@ -109,6 +158,7 @@ impl MemoryStartupContext { thread, thread_manager, auth_manager, + provider, session_telemetry, } } @@ -121,6 +171,10 @@ impl MemoryStartupContext { self.thread.state_db() } + pub(crate) fn provider(&self) -> &dyn ModelProvider { + self.provider.as_ref() + } + pub(crate) fn counter(&self, name: &str, inc: i64, tags: &[(&str, &str)]) { self.session_telemetry.counter(name, inc, tags); } @@ -189,8 +243,10 @@ impl MemoryStartupContext { ); let mut client_session = model_client.new_session(); + let window_id = format!("{}:0", self.thread_id); let mut stream = client_session .stream( + &window_id, prompt, &context.model_info, &context.session_telemetry, diff --git a/codex-rs/memories/write/src/startup_tests.rs b/codex-rs/memories/write/src/startup_tests.rs index bd0a9904e6..d5e46e09e0 100644 --- a/codex-rs/memories/write/src/startup_tests.rs +++ b/codex-rs/memories/write/src/startup_tests.rs @@ -1,12 +1,30 @@ +use crate::extensions::seed_extension_instructions; +use crate::memory_root; +use crate::phase1; +use crate::phase2; +use crate::runtime::MemoryStartupContext; use crate::start_memories_startup_task; +use codex_config::types::MemoriesConfig; use codex_features::Feature; use codex_git_utils::diff_since_latest_init; use codex_git_utils::reset_git_repository; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_model_provider::ModelProvider; +use codex_model_provider::ProviderAccountResult; +use codex_model_provider::SharedModelProvider; +use codex_model_provider::create_model_provider; +use codex_model_provider_info::ModelProviderInfo; use codex_protocol::ThreadId; use codex_protocol::config_types::ServiceTier; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::openai_models::ModelsResponse; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SessionSource; use core_test_support::responses::ResponseMock; use core_test_support::responses::ResponsesRequest; @@ -20,7 +38,10 @@ use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use pretty_assertions::assert_eq; +use std::future::Future; use std::path::Path; +use std::path::PathBuf; +use std::pin::Pin; use std::sync::Arc; use tempfile::TempDir; use tokio::time::Duration; @@ -329,18 +350,186 @@ async fn memories_startup_phase1_uses_live_thread_service_tier_and_detached_meta Ok(()) } +#[tokio::test] +async fn memories_startup_phase1_provider_default_drives_request_model() -> anyhow::Result<()> { + let server = start_mock_server().await; + let home = Arc::new(TempDir::new()?); + let request = + run_memory_phase_one_model_request_test(&server, home, startup_test_memories_config()) + .await?; + + assert_eq!( + request.body_json()["model"].as_str(), + Some(MOCK_PROVIDER_PHASE_ONE_MODEL) + ); + + Ok(()) +} + +#[tokio::test] +async fn memories_startup_phase2_provider_default_drives_request_model() -> anyhow::Result<()> { + let server = start_mock_server().await; + let home = Arc::new(TempDir::new()?); + let request = + run_memory_phase_two_model_request_test(&server, home, startup_test_memories_config()) + .await?; + + assert_eq!( + request.body_json()["model"].as_str(), + Some(MOCK_PROVIDER_PHASE_TWO_MODEL) + ); + + Ok(()) +} + +#[tokio::test] +async fn memories_startup_phase1_explicit_model_override_drives_request_model() -> anyhow::Result<()> +{ + let server = start_mock_server().await; + let home = Arc::new(TempDir::new()?); + let mut memories = startup_test_memories_config(); + memories.extract_model = Some("override.phase-one".to_string()); + let request = run_memory_phase_one_model_request_test(&server, home, memories).await?; + + assert_eq!( + request.body_json()["model"].as_str(), + Some("override.phase-one") + ); + + Ok(()) +} + +#[tokio::test] +async fn memories_startup_phase2_explicit_model_override_drives_request_model() -> anyhow::Result<()> +{ + let server = start_mock_server().await; + let home = Arc::new(TempDir::new()?); + let mut memories = startup_test_memories_config(); + memories.consolidation_model = Some("override.phase-two".to_string()); + let request = run_memory_phase_two_model_request_test(&server, home, memories).await?; + + assert_eq!( + request.body_json()["model"].as_str(), + Some("override.phase-two") + ); + + Ok(()) +} + +async fn run_memory_phase_one_model_request_test( + server: &wiremock::MockServer, + home: Arc, + memories: MemoriesConfig, +) -> anyhow::Result { + let test = build_test_codex_with_memories_config(server, Arc::clone(&home), memories).await?; + let provider = Arc::new(MockMemoryModelProvider::new( + test.config.model_provider.clone(), + Some(test.thread_manager.auth_manager()), + )); + let db = test + .codex + .state_db() + .ok_or_else(|| anyhow::anyhow!("state db should be enabled for memory startup test"))?; + seed_stage1_candidate( + db.as_ref(), + home.path(), + chrono::Utc::now() - chrono::Duration::hours(2), + "startup-models", + ) + .await?; + let response = mount_sse_once( + server, + sse(vec![ + ev_response_created("resp-phase1"), + ev_assistant_message( + "msg-phase1", + r#"{"raw_memory":"raw memory","rollout_summary":"rollout summary","rollout_slug":"startup-models"}"#, + ), + ev_completed("resp-phase1"), + ]), + ) + .await; + + let (context, config) = memory_startup_context_with_provider(&test, provider).await; + phase1::run(context, config).await; + let request = wait_for_single_request(&response).await; + shutdown_test_codex(&test).await?; + Ok(request) +} + +async fn run_memory_phase_two_model_request_test( + server: &wiremock::MockServer, + home: Arc, + memories: MemoriesConfig, +) -> anyhow::Result { + let test = build_test_codex_with_memories_config(server, home.clone(), memories).await?; + let provider = Arc::new(MockMemoryModelProvider::new( + test.config.model_provider.clone(), + Some(test.thread_manager.auth_manager()), + )); + let db = test + .codex + .state_db() + .ok_or_else(|| anyhow::anyhow!("state db should be enabled for memory startup test"))?; + seed_stage1_output( + db.as_ref(), + home.path(), + chrono::Utc::now(), + "raw memory for phase two", + "rollout summary for phase two", + "startup-models-phase-two", + ) + .await?; + + let response = mount_sse_once( + server, + sse(vec![ + ev_response_created("resp-phase2"), + ev_assistant_message("msg-phase2", "phase2 complete"), + ev_completed("resp-phase2"), + ]), + ) + .await; + + let (context, config) = memory_startup_context_with_provider(&test, provider).await; + let root = memory_root(&config.codex_home); + tokio::fs::create_dir_all(&root).await?; + seed_extension_instructions(&root).await?; + phase2::run(context, config).await; + let request = wait_for_single_request(&response).await; + wait_for_phase2_workspace_reset(&home.path().join("memories")).await?; + shutdown_test_codex(&test).await?; + Ok(request) +} + +fn startup_test_memories_config() -> MemoriesConfig { + MemoriesConfig { + max_raw_memories_for_consolidation: 1, + min_rollout_idle_hours: 0, + ..MemoriesConfig::default() + } +} + async fn build_test_codex( server: &wiremock::MockServer, home: Arc, +) -> anyhow::Result { + build_test_codex_with_memories_config(server, home, startup_test_memories_config()).await +} + +async fn build_test_codex_with_memories_config( + server: &wiremock::MockServer, + home: Arc, + memories: MemoriesConfig, ) -> anyhow::Result { test_codex() .with_home(home) - .with_config(|config| { + .with_config(move |config| { config .features .enable(Feature::Sqlite) .expect("test config should allow feature update"); - config.memories.max_raw_memories_for_consolidation = 1; + config.memories = memories; }) .build(server) .await @@ -370,6 +559,88 @@ async fn trigger_memories_startup(test: &TestCodex) { ); } +async fn memory_startup_context_with_provider( + test: &TestCodex, + provider: SharedModelProvider, +) -> (Arc, Arc) { + let config_snapshot = test.codex.config_snapshot().await; + let mut config = test.config.clone(); + config + .features + .enable(Feature::MemoryTool) + .expect("test config should allow feature update"); + let config = Arc::new(config); + let context = Arc::new(MemoryStartupContext::new_for_testing( + Arc::clone(&test.thread_manager), + test.thread_manager.auth_manager(), + test.session_configured.thread_id, + Arc::clone(&test.codex), + config.as_ref(), + config_snapshot.session_source, + provider, + )); + + (context, config) +} + +const MOCK_PROVIDER_PHASE_ONE_MODEL: &str = "mock.phase-one"; +const MOCK_PROVIDER_PHASE_TWO_MODEL: &str = "mock.phase-two"; + +#[derive(Debug)] +struct MockMemoryModelProvider { + delegate: SharedModelProvider, +} + +impl MockMemoryModelProvider { + fn new(info: ModelProviderInfo, auth_manager: Option>) -> Self { + Self { + delegate: create_model_provider(info, auth_manager), + } + } +} + +impl ModelProvider for MockMemoryModelProvider { + fn info(&self) -> &ModelProviderInfo { + self.delegate.info() + } + + fn memory_extraction_preferred_model(&self) -> &'static str { + MOCK_PROVIDER_PHASE_ONE_MODEL + } + + fn memory_consolidation_preferred_model(&self) -> &'static str { + MOCK_PROVIDER_PHASE_TWO_MODEL + } + + fn auth_manager(&self) -> Option> { + self.delegate.auth_manager() + } + + fn auth<'life0, 'async_trait>( + &'life0 self, + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + Self: 'async_trait, + { + let delegate = Arc::clone(&self.delegate); + Box::pin(async move { delegate.auth().await }) + } + + fn account_state(&self) -> ProviderAccountResult { + self.delegate.account_state() + } + + fn models_manager( + &self, + codex_home: PathBuf, + config_model_catalog: Option, + ) -> codex_models_manager::manager::SharedModelsManager { + self.delegate + .models_manager(codex_home, config_model_catalog) + } +} + async fn seed_stage1_output( db: &codex_state::StateRuntime, codex_home: &Path, @@ -404,6 +675,46 @@ async fn seed_stage1_output( Ok(thread_id) } +async fn seed_stage1_candidate( + db: &codex_state::StateRuntime, + codex_home: &Path, + updated_at: chrono::DateTime, + rollout_slug: &str, +) -> anyhow::Result { + let thread_id = ThreadId::new(); + let rollout_path = codex_home.join(format!("rollout-{thread_id}.jsonl")); + let line = RolloutLine { + timestamp: updated_at.to_rfc3339(), + item: RolloutItem::ResponseItem(ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "remember this startup test conversation".to_string(), + }], + phase: None, + }), + }; + let jsonl = serde_json::to_string(&line)?; + tokio::fs::write(&rollout_path, format!("{jsonl}\n")).await?; + + let mut metadata_builder = codex_state::ThreadMetadataBuilder::new( + thread_id, + rollout_path, + updated_at, + SessionSource::Cli, + ); + metadata_builder.cwd = codex_home.join(format!("workspace-{rollout_slug}")); + metadata_builder.model_provider = Some("test-provider".to_string()); + metadata_builder.git_branch = Some(format!("branch-{rollout_slug}")); + let mut metadata = metadata_builder.build("test-provider"); + metadata.preview = Some("remember this startup test conversation".to_string()); + metadata.first_user_message = metadata.preview.clone(); + db.upsert_thread(&metadata).await?; + db.set_thread_memory_mode(thread_id, "enabled").await?; + + Ok(thread_id) +} + async fn wait_for_single_request(mock: &ResponseMock) -> ResponsesRequest { wait_for_request(mock, /*expected_count*/ 1).await.remove(0) } @@ -447,7 +758,8 @@ async fn wait_for_request(mock: &ResponseMock, expected_count: usize) -> Vec &'static str { + AMAZON_BEDROCK_GPT_5_4_MODEL_ID + } + + fn memory_consolidation_preferred_model(&self) -> &'static str { + AMAZON_BEDROCK_GPT_5_4_MODEL_ID + } + fn auth_manager(&self) -> Option> { None } diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 4f87d46bae..8f5e02b175 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -8,11 +8,15 @@ use codex_api::SharedAuthProvider; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::error::CodexErr; use http::HeaderMap; use http::HeaderValue; use crate::bearer_auth_provider::BearerAuthProvider; +const BEDROCK_API_KEY_UNSUPPORTED_MESSAGE: &str = + "Bedrock API key auth is only supported by the Amazon Bedrock model provider"; + #[derive(Clone, Debug)] struct AgentIdentityAuthProvider { auth: codex_login::auth::AgentIdentityAuth, @@ -79,6 +83,12 @@ pub(crate) fn resolve_provider_auth( auth: Option<&CodexAuth>, provider: &ModelProviderInfo, ) -> codex_protocol::error::Result { + if matches!(auth, Some(CodexAuth::BedrockApiKey(_))) { + return Err(CodexErr::UnsupportedOperation( + BEDROCK_API_KEY_UNSUPPORTED_MESSAGE.to_string(), + )); + } + if let Some(auth) = bearer_auth_for_provider(provider)? { return Ok(Arc::new(auth)); } @@ -109,6 +119,7 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider { CodexAuth::AgentIdentity(auth) => { Arc::new(AgentIdentityAuthProvider { auth: auth.clone() }) } + CodexAuth::BedrockApiKey(_) => unreachable!("{BEDROCK_API_KEY_UNSUPPORTED_MESSAGE}"), CodexAuth::ApiKey(_) | CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) @@ -122,8 +133,10 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider { #[cfg(test)] mod tests { + use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; + use pretty_assertions::assert_eq; use super::*; @@ -135,4 +148,21 @@ mod tests { assert!(auth.to_auth_headers().is_empty()); } + + #[test] + fn openai_provider_rejects_bedrock_api_key_auth() { + let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + let auth = CodexAuth::BedrockApiKey(BedrockApiKeyAuth { + api_key: "bedrock-api-key-test".to_string(), + region: "us-east-1".to_string(), + }); + + match resolve_provider_auth(Some(&auth), &provider) { + Err(CodexErr::UnsupportedOperation(message)) => { + assert_eq!(message, BEDROCK_API_KEY_UNSUPPORTED_MESSAGE); + } + Err(err) => panic!("unexpected auth error: {err:?}"), + Ok(_) => panic!("Bedrock API key auth should be rejected"), + } + } } diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index afd07065d4..95427465eb 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -51,6 +51,7 @@ pub struct ProviderAccountState { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ProviderAccountError { MissingChatgptAccountDetails, + UnsupportedBedrockApiKeyAuth, } impl fmt::Display for ProviderAccountError { @@ -62,6 +63,12 @@ impl fmt::Display for ProviderAccountError { "email and plan type are required for chatgpt authentication" ) } + Self::UnsupportedBedrockApiKeyAuth => { + write!( + f, + "Bedrock API key auth is only supported by the Amazon Bedrock model provider" + ) + } } } } @@ -74,6 +81,14 @@ pub type ProviderAccountResult = std::result::Result &'static str { + DEFAULT_MEMORY_EXTRACTION_PREFERRED_MODEL + } + + /// Returns the preferred model used for memory consolidation. + /// + /// Providers that require backend-specific model IDs should override this. + fn memory_consolidation_preferred_model(&self) -> &'static str { + DEFAULT_MEMORY_CONSOLIDATION_PREFERRED_MODEL + } + /// Returns whether requests made through this provider should include attestation. fn supports_attestation(&self) -> bool { false @@ -210,6 +239,9 @@ impl ModelProvider for ConfiguredModelProvider { }) .map(|auth| match &auth { CodexAuth::ApiKey(_) => Ok(ProviderAccount::ApiKey), + CodexAuth::BedrockApiKey(_) => { + Err(ProviderAccountError::UnsupportedBedrockApiKeyAuth) + } CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) | CodexAuth::AgentIdentity(_) @@ -265,6 +297,7 @@ impl ModelProvider for ConfiguredModelProvider { mod tests { use std::num::NonZeroU64; + use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::WireApi; use codex_models_manager::manager::RefreshStrategy; @@ -352,6 +385,13 @@ mod tests { .expect("valid model") } + fn bedrock_api_key_auth() -> CodexAuth { + CodexAuth::BedrockApiKey(BedrockApiKeyAuth { + api_key: "bedrock-api-key-test".to_string(), + region: "us-east-1".to_string(), + }) + } + #[test] fn configured_provider_uses_default_capabilities() { let provider = create_model_provider( @@ -469,6 +509,19 @@ mod tests { ); } + #[test] + fn openai_provider_rejects_bedrock_api_key_account_state() { + let provider = create_model_provider( + ModelProviderInfo::create_openai_provider(/*base_url*/ None), + Some(AuthManager::from_auth_for_testing(bedrock_api_key_auth())), + ); + + assert_eq!( + provider.account_state(), + Err(ProviderAccountError::UnsupportedBedrockApiKeyAuth) + ); + } + #[test] fn custom_non_openai_provider_returns_no_account_state() { let provider = create_model_provider( diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 2e39ef3160..6f3a6a27f4 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -212,6 +212,7 @@ c2ln", last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; std::fs::create_dir_all(codex_home).expect("codex home should be created"); std::fs::write( diff --git a/codex-rs/models-manager/src/model_info.rs b/codex-rs/models-manager/src/model_info.rs index 81a5c6e5ed..566953dfc2 100644 --- a/codex-rs/models-manager/src/model_info.rs +++ b/codex-rs/models-manager/src/model_info.rs @@ -94,6 +94,7 @@ pub fn model_info_from_slug(slug: &str) -> ModelInfo { context_window: Some(272_000), max_context_window: Some(272_000), auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), input_modalities: default_input_modalities(), diff --git a/codex-rs/otel/src/lib.rs b/codex-rs/otel/src/lib.rs index 586dfa15a7..fe7a171ce9 100644 --- a/codex-rs/otel/src/lib.rs +++ b/codex-rs/otel/src/lib.rs @@ -54,7 +54,8 @@ pub enum TelemetryAuthMode { impl From for TelemetryAuthMode { fn from(mode: codex_app_server_protocol::AuthMode) -> Self { match mode { - codex_app_server_protocol::AuthMode::ApiKey => Self::ApiKey, + codex_app_server_protocol::AuthMode::ApiKey + | codex_app_server_protocol::AuthMode::BedrockApiKey => Self::ApiKey, codex_app_server_protocol::AuthMode::Chatgpt | codex_app_server_protocol::AuthMode::ChatgptAuthTokens | codex_app_server_protocol::AuthMode::AgentIdentity diff --git a/codex-rs/otel/src/metrics/client.rs b/codex-rs/otel/src/metrics/client.rs index 9d5ad5d166..a36157aae3 100644 --- a/codex-rs/otel/src/metrics/client.rs +++ b/codex-rs/otel/src/metrics/client.rs @@ -12,6 +12,7 @@ use crate::metrics::validation::validate_tags; use codex_utils_string::sanitize_metric_tag_value; use opentelemetry::KeyValue; use opentelemetry::metrics::Counter; +use opentelemetry::metrics::Gauge; use opentelemetry::metrics::Histogram; use opentelemetry::metrics::Meter; use opentelemetry::metrics::MeterProvider as _; @@ -89,6 +90,7 @@ struct MetricsClientInner { meter_provider: SdkMeterProvider, meter: Meter, counters: Mutex>>, + gauges: Mutex>>, histograms: Mutex>>, duration_histograms: Mutex>>, runtime_reader: Option>, @@ -146,6 +148,35 @@ impl MetricsClientInner { Ok(()) } + fn gauge( + &self, + name: &str, + description: Option<&str>, + value: i64, + tags: &[(&str, &str)], + ) -> Result<()> { + validate_metric_name(name)?; + let attributes = self.attributes(tags)?; + + let mut gauges = self + .gauges + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let key = InstrumentKey { + name: name.to_string(), + description: description.map(str::to_string), + }; + let gauge = gauges.entry(key).or_insert_with(|| { + let builder = self.meter.i64_gauge(name.to_string()); + match description { + Some(description) => builder.with_description(description.to_string()).build(), + None => builder.build(), + } + }); + gauge.record(value, &attributes); + Ok(()) + } + fn duration_histogram(&self, name: &str, value: i64, tags: &[(&str, &str)]) -> Result<()> { validate_metric_name(name)?; let attributes = self.attributes(tags)?; @@ -253,6 +284,7 @@ impl MetricsClient { meter_provider, meter, counters: Mutex::new(HashMap::new()), + gauges: Mutex::new(HashMap::new()), histograms: Mutex::new(HashMap::new()), duration_histograms: Mutex::new(HashMap::new()), runtime_reader, @@ -281,6 +313,22 @@ impl MetricsClient { self.0.histogram(name, value, tags) } + /// Send a single gauge measurement. + pub fn gauge(&self, name: &str, value: i64, tags: &[(&str, &str)]) -> Result<()> { + self.0.gauge(name, /*description*/ None, value, tags) + } + + /// Send a single gauge measurement with an instrument description. + pub fn gauge_with_description( + &self, + name: &str, + description: &str, + value: i64, + tags: &[(&str, &str)], + ) -> Result<()> { + self.0.gauge(name, Some(description), value, tags) + } + /// Record a duration in milliseconds using a histogram. pub fn record_duration( &self, diff --git a/codex-rs/otel/tests/suite/otlp_http_loopback.rs b/codex-rs/otel/tests/suite/otlp_http_loopback.rs index 4c2dd36f76..1f41336916 100644 --- a/codex-rs/otel/tests/suite/otlp_http_loopback.rs +++ b/codex-rs/otel/tests/suite/otlp_http_loopback.rs @@ -186,6 +186,12 @@ fn otlp_http_exporter_sends_metrics_to_collector() -> Result<()> { ))?; metrics.counter("codex.turns", /*inc*/ 1, &[("source", "test")])?; + metrics.gauge_with_description( + "codex.active", + "Number of active Codex operations.", + /*value*/ 1, + &[("component", "test")], + )?; metrics.shutdown()?; server.join().expect("server join"); @@ -220,6 +226,16 @@ fn otlp_http_exporter_sends_metrics_to_collector() -> Result<()> { "expected metric name not found; body prefix: {}", &body.chars().take(2000).collect::() ); + assert!( + body.contains("codex.active"), + "expected gauge not found; body prefix: {}", + &body.chars().take(2000).collect::() + ); + assert!( + body.contains("component") && body.contains("test"), + "expected gauge tag not found; body prefix: {}", + &body.chars().take(2000).collect::() + ); Ok(()) } diff --git a/codex-rs/otel/tests/suite/send.rs b/codex-rs/otel/tests/suite/send.rs index 3a268d3c7d..e2b8f2103e 100644 --- a/codex-rs/otel/tests/suite/send.rs +++ b/codex-rs/otel/tests/suite/send.rs @@ -24,6 +24,12 @@ fn send_builds_payload_with_tags_and_histograms() -> Result<()> { /*value*/ 25, &[("tool", "shell")], )?; + metrics.gauge_with_description( + "codex.active", + "Number of active Codex operations.", + /*value*/ 2, + &[("component", "test")], + )?; metrics.shutdown()?; let resource_metrics = latest_metrics(&exporter); @@ -80,6 +86,27 @@ fn send_builds_payload_with_tags_and_histograms() -> Result<()> { ]); assert_eq!(histogram_attrs, expected_histogram_attributes); + let gauge = find_metric(&resource_metrics, "codex.active").expect("gauge metric missing"); + assert_eq!(gauge.description(), "Number of active Codex operations."); + let gauge_point = match gauge.data() { + opentelemetry_sdk::metrics::data::AggregatedMetrics::I64(data) => match data { + opentelemetry_sdk::metrics::data::MetricData::Gauge(gauge) => { + gauge.data_points().next().expect("gauge point") + } + _ => panic!("unexpected gauge aggregation"), + }, + _ => panic!("unexpected gauge metric data type"), + }; + assert_eq!(gauge_point.value(), 2); + assert_eq!( + attributes_to_map(gauge_point.attributes()), + BTreeMap::from([ + ("component".to_string(), "test".to_string()), + ("env".to_string(), "prod".to_string()), + ("service".to_string(), "codex-cli".to_string()), + ]) + ); + Ok(()) } diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 517c33cac9..5df6d1e77f 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -4,6 +4,7 @@ use std::num::NonZeroUsize; use std::path::Path; use codex_utils_image::PromptImageMode; +use codex_utils_image::data_url_from_bytes; use codex_utils_image::load_for_prompt_bytes; use serde::Deserialize; use serde::Deserializer; @@ -1088,26 +1089,12 @@ pub fn local_image_content_items_with_label_number( }; match load_for_prompt_bytes(path, file_bytes, mode) { - Ok(image) => { - let mut items = Vec::with_capacity(3); - if let Some(label_number) = label_number { - items.push(ContentItem::InputText { - text: local_image_open_tag_text_with_path(label_number, path), - }); - } - items.push(ContentItem::InputImage { - image_url: image.into_data_url(), - detail: Some(detail), - }); - if label_number.is_some() { - items.push(ContentItem::InputText { - text: LOCAL_IMAGE_CLOSE_TAG.to_string(), - }); - } - items - } + Ok(image) => local_image_content_items(path, image.into_data_url(), label_number, detail), Err(err) => match &err { - ImageProcessingError::Read { .. } | ImageProcessingError::Encode { .. } => { + ImageProcessingError::Read { .. } + | ImageProcessingError::Encode { .. } + | ImageProcessingError::InvalidDataUrl { .. } + | ImageProcessingError::ImageTooLarge { .. } => { vec![local_image_error_placeholder(path, &err)] } ImageProcessingError::Decode { .. } if err.is_invalid_image() => { @@ -1123,6 +1110,36 @@ pub fn local_image_content_items_with_label_number( } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LocalImagePreparation { + Process, + Defer, +} + +fn local_image_content_items( + path: &std::path::Path, + image_url: String, + label_number: Option, + detail: ImageDetail, +) -> Vec { + let mut items = Vec::with_capacity(3); + if let Some(label_number) = label_number { + items.push(ContentItem::InputText { + text: local_image_open_tag_text_with_path(label_number, path), + }); + } + items.push(ContentItem::InputImage { + image_url, + detail: Some(detail), + }); + if label_number.is_some() { + items.push(ContentItem::InputText { + text: LOCAL_IMAGE_CLOSE_TAG.to_string(), + }); + } + items +} + impl From for ResponseItem { fn from(item: ResponseInputItem) -> Self { match item { @@ -1235,6 +1252,15 @@ pub enum ReasoningItemContent { impl From> for ResponseInputItem { fn from(items: Vec) -> Self { + Self::from_user_input(items, LocalImagePreparation::Process) + } +} + +impl ResponseInputItem { + pub fn from_user_input( + items: Vec, + local_image_preparation: LocalImagePreparation, + ) -> Self { let mut image_index = 0; Self::Message { role: "user".to_string(), @@ -1256,12 +1282,22 @@ impl From> for ResponseInputItem { image_index += 1; let detail = detail.unwrap_or(DEFAULT_IMAGE_DETAIL); match std::fs::read(&path) { - Ok(file_bytes) => local_image_content_items_with_label_number( - &path, - file_bytes, - Some(image_index), - detail, - ), + Ok(file_bytes) => match local_image_preparation { + LocalImagePreparation::Process => { + local_image_content_items_with_label_number( + &path, + file_bytes, + Some(image_index), + detail, + ) + } + LocalImagePreparation::Defer => local_image_content_items( + &path, + data_url_from_bytes("application/octet-stream", &file_bytes), + Some(image_index), + detail, + ), + }, Err(err) => vec![local_image_error_placeholder(&path, err)], } } diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index 1553ef3435..5b95fee29d 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -389,6 +389,9 @@ pub struct ModelInfo { /// context window when available. #[serde(default, skip_serializing_if = "Option::is_none")] pub auto_compact_token_limit: Option, + /// Opaque identifier for compaction-compatible model configurations. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub comp_hash: Option, /// Percentage of the context window considered usable for inputs, after /// reserving headroom for system prompts, tool overhead, and model output. #[serde(default = "default_effective_context_window_percent")] @@ -671,6 +674,7 @@ mod tests { context_window: None, max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: vec![], input_modalities: default_input_modalities(), @@ -940,6 +944,7 @@ mod tests { assert_eq!(model.web_search_tool_type, WebSearchToolType::Text); assert!(!model.supports_search_tool); assert!(!model.use_responses_lite); + assert_eq!(model.comp_hash, None); assert_eq!(model.auto_review_model_override, None); assert_eq!(model.tool_mode, None); } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 5abc9c6399..059c5692ce 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -183,11 +183,15 @@ pub struct McpServerRefreshConfig { #[derive(Debug, Clone, PartialEq)] pub struct ConversationStartParams { + /// Overrides the configured realtime model for this session only. + pub model: Option, /// Selects whether the realtime session should produce text or audio output. pub output_modality: RealtimeOutputModality, pub prompt: Option>, pub realtime_session_id: Option, pub transport: Option, + /// Overrides the configured realtime protocol version for this session only. + pub version: Option, pub voice: Option, } @@ -2888,6 +2892,8 @@ pub struct CompactedItem { pub message: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub replacement_history: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub window_id: Option, } impl From for ResponseItem { @@ -2935,6 +2941,8 @@ pub struct TurnContextItem { #[serde(default, skip_serializing_if = "Option::is_none")] pub file_system_sandbox_policy: Option, pub model: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub comp_hash: Option, #[serde(skip_serializing_if = "Option::is_none")] pub personality: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -5156,6 +5164,7 @@ mod tests { assert_eq!(item.network, None); assert_eq!(item.file_system_sandbox_policy, None); + assert_eq!(item.comp_hash, None); Ok(()) } @@ -5216,6 +5225,7 @@ mod tests { }, ])), model: "gpt-5".to_string(), + comp_hash: None, personality: None, collaboration_mode: None, multi_agent_version: None, diff --git a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs index 2384d39473..5232c04dd4 100644 --- a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs @@ -8,6 +8,7 @@ use std::time::Duration; use axum::Router; use axum::body::Body; +use axum::body::to_bytes; use axum::extract::Json; use axum::extract::State; use axum::http::HeaderMap; @@ -48,6 +49,7 @@ use rmcp::transport::StreamableHttpServerConfig; use rmcp::transport::StreamableHttpService; use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; use serde::Deserialize; +use serde_json::Value; use serde_json::json; use tokio::sync::Mutex; use tokio::task; @@ -64,18 +66,32 @@ const MEMO_URI: &str = "memo://codex/example-note"; const MEMO_CONTENT: &str = "This is a sample MCP resource served by the rmcp test server."; const MCP_SESSION_ID_HEADER: &str = "mcp-session-id"; const SESSION_POST_FAILURE_CONTROL_PATH: &str = "/test/control/session-post-failure"; +const INITIALIZE_POST_FAILURE_CONTROL_PATH: &str = "/test/control/initialize-post-failure"; +const INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH: &str = + "/test/control/initialized-notification-post-failure"; +const MAX_MCP_POST_BODY_BYTES: usize = 1024 * 1024; #[derive(Clone, Default)] -struct SessionFailureState { +struct PostFailureState { armed_failure: Arc>>, } +#[derive(Clone, Copy, Debug)] +enum ArmedFailureTarget { + Initialize, + InitializedNotification, + Session, +} + #[derive(Clone, Debug)] struct ArmedFailure { + target: ArmedFailureTarget, status: StatusCode, remaining: usize, /// Raw `WWW-Authenticate` challenge header field values returned with the failure. www_authenticate_headers: Vec, + content_type: Option, + body: Option, } #[derive(Debug, Deserialize)] @@ -85,6 +101,8 @@ struct ArmSessionPostFailureRequest { /// Raw `WWW-Authenticate` challenge header field values to add to the failure. #[serde(default)] www_authenticate_headers: Vec, + content_type: Option, + body: Option, } #[derive(Deserialize)] @@ -97,7 +115,7 @@ struct EchoArgs { #[tokio::main] async fn main() -> Result<(), Box> { let bind_addr = parse_bind_addr()?; - let session_failure_state = SessionFailureState::default(); + let post_failure_state = PostFailureState::default(); const MAX_BIND_RETRIES: u32 = 20; const BIND_RETRY_DELAY: Duration = Duration::from_millis(50); @@ -129,6 +147,14 @@ async fn main() -> Result<(), Box> { SESSION_POST_FAILURE_CONTROL_PATH, post(arm_session_post_failure), ) + .route( + INITIALIZE_POST_FAILURE_CONTROL_PATH, + post(arm_initialize_post_failure), + ) + .route( + INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH, + post(arm_initialized_notification_post_failure), + ) .route( "/.well-known/oauth-authorization-server/mcp", get({ @@ -162,10 +188,10 @@ async fn main() -> Result<(), Box> { ), ) .layer(middleware::from_fn_with_state( - session_failure_state.clone(), - fail_session_post_when_armed, + post_failure_state.clone(), + fail_mcp_post_when_armed, )) - .with_state(session_failure_state); + .with_state(post_failure_state); let router = if let Ok(token) = std::env::var("MCP_EXPECT_BEARER") { let expected = Arc::new(format!("Bearer {token}")); @@ -404,8 +430,30 @@ async fn require_bearer( } async fn arm_session_post_failure( - State(state): State, + State(state): State, Json(request): Json, +) -> Result { + arm_post_failure(state, request, ArmedFailureTarget::Session).await +} + +async fn arm_initialize_post_failure( + State(state): State, + Json(request): Json, +) -> Result { + arm_post_failure(state, request, ArmedFailureTarget::Initialize).await +} + +async fn arm_initialized_notification_post_failure( + State(state): State, + Json(request): Json, +) -> Result { + arm_post_failure(state, request, ArmedFailureTarget::InitializedNotification).await +} + +async fn arm_post_failure( + state: PostFailureState, + request: ArmSessionPostFailureRequest, + target: ArmedFailureTarget, ) -> Result { let status = StatusCode::from_u16(request.status).map_err(|_| StatusCode::BAD_REQUEST)?; let www_authenticate_headers = request @@ -413,46 +461,76 @@ async fn arm_session_post_failure( .into_iter() .map(|value| HeaderValue::from_str(&value).map_err(|_| StatusCode::BAD_REQUEST)) .collect::, _>>()?; + let content_type = request + .content_type + .map(|value| HeaderValue::from_str(&value).map_err(|_| StatusCode::BAD_REQUEST)) + .transpose()?; let armed_failure = if request.remaining == 0 { None } else { Some(ArmedFailure { + target, status, remaining: request.remaining, www_authenticate_headers, + content_type, + body: request.body, }) }; *state.armed_failure.lock().await = armed_failure; Ok(StatusCode::NO_CONTENT) } -async fn fail_session_post_when_armed( - State(state): State, +async fn fail_mcp_post_when_armed( + State(state): State, request: Request, next: Next, ) -> Response { - if request.uri().path() != "/mcp" - || request.method() != Method::POST - || !request.headers().contains_key(MCP_SESSION_ID_HEADER) - { + if request.uri().path() != "/mcp" || request.method() != Method::POST { return next.run(request).await; } + let (parts, body) = request.into_parts(); + let body_bytes = match to_bytes(body, MAX_MCP_POST_BODY_BYTES).await { + Ok(body_bytes) => body_bytes, + Err(_) => { + let mut response = Response::new(Body::from("failed to read request body")); + *response.status_mut() = StatusCode::BAD_REQUEST; + return response; + } + }; + let has_session_id = parts.headers.contains_key(MCP_SESSION_ID_HEADER); + let mcp_method = request_mcp_method(&body_bytes); { let mut armed_failure = state.armed_failure.lock().await; if let Some(failure) = armed_failure.as_mut() && failure.remaining > 0 + && match failure.target { + ArmedFailureTarget::Initialize => !has_session_id, + ArmedFailureTarget::InitializedNotification => { + has_session_id && mcp_method.as_deref() == Some("notifications/initialized") + } + ArmedFailureTarget::Session => { + has_session_id && mcp_method.as_deref() != Some("notifications/initialized") + } + } { failure.remaining -= 1; let status = failure.status; let www_authenticate_headers = failure.www_authenticate_headers.clone(); + let content_type = failure.content_type.clone(); + let body = failure + .body + .clone() + .unwrap_or_else(|| format!("forced session failure with status {status}")); if failure.remaining == 0 { *armed_failure = None; } - let mut response = Response::new(Body::from(format!( - "forced session failure with status {status}" - ))); + let mut response = Response::new(Body::from(body)); *response.status_mut() = status; + if let Some(content_type) = content_type { + response.headers_mut().insert(CONTENT_TYPE, content_type); + } for www_authenticate_header in www_authenticate_headers { response .headers_mut() @@ -462,5 +540,14 @@ async fn fail_session_post_when_armed( } } - next.run(request).await + next.run(Request::from_parts(parts, Body::from(body_bytes))) + .await +} + +fn request_mcp_method(body: &[u8]) -> Option { + serde_json::from_slice::(body) + .ok()? + .get("method")? + .as_str() + .map(ToString::to_string) } diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 6f98f78920..a9cdc7b88f 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -28,6 +28,8 @@ use reqwest::header::CONTENT_TYPE; use reqwest::header::HeaderMap; use reqwest::header::HeaderName; use rmcp::model::ClientJsonRpcMessage; +use rmcp::model::ClientNotification; +use rmcp::model::ConstString; use rmcp::model::JsonRpcMessage; use rmcp::model::ServerJsonRpcMessage; use rmcp::transport::streamable_http_client::AuthRequiredError; @@ -185,6 +187,25 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { let content_type = response_header(&response.headers, CONTENT_TYPE); let session_id = response_header(&response.headers, HEADER_SESSION_ID); + if !status_is_success(response.status) { + let body = collect_body(&mut body_stream).await?; + if !retryable_post_response_status(mcp_method.as_deref(), response.status) + && content_type + .as_deref() + .is_some_and(|content_type| content_type.starts_with(JSON_MIME_TYPE)) + && let Some(message) = parse_json_rpc_error(&body) + { + return Ok(StreamableHttpPostResponse::Json(message, session_id)); + } + return Err(StreamableHttpError::UnexpectedServerResponse( + format!( + "HTTP {}: {}", + response.status, + body_preview(String::from_utf8_lossy(&body).to_string()) + ) + .into(), + )); + } match content_type.as_deref() { Some(content_type) if content_type.starts_with(EVENT_STREAM_MIME_TYPE) => { let event_stream = sse_stream_from_body(body_stream); @@ -380,7 +401,26 @@ fn client_jsonrpc_message_fields( Some(request.id.to_string()), ), JsonRpcMessage::Response(response) => (None, Some(response.id.to_string())), - JsonRpcMessage::Notification(_) => (None, None), + JsonRpcMessage::Notification(notification) => { + let method = match ¬ification.notification { + ClientNotification::CancelledNotification(notification) => { + notification.method.as_str() + } + ClientNotification::ProgressNotification(notification) => { + notification.method.as_str() + } + ClientNotification::InitializedNotification(notification) => { + notification.method.as_str() + } + ClientNotification::RootsListChangedNotification(notification) => { + notification.method.as_str() + } + ClientNotification::CustomNotification(notification) => { + notification.method.as_str() + } + }; + (Some(method.to_string()), None) + } JsonRpcMessage::Error(error) => (None, error.id.as_ref().map(ToString::to_string)), } } @@ -463,6 +503,36 @@ fn status_is_success(status: u16) -> bool { StatusCode::from_u16(status).is_ok_and(|status| status.is_success()) } +fn retryable_post_response_status(mcp_method: Option<&str>, status: u16) -> bool { + let Ok(status) = StatusCode::from_u16(status) else { + return false; + }; + is_retryable_http_status(status) + && matches!( + mcp_method, + Some("initialize" | "notifications/initialized" | "tools/list") + ) +} + +fn is_retryable_http_status(status: StatusCode) -> bool { + matches!( + status, + StatusCode::REQUEST_TIMEOUT + | StatusCode::TOO_MANY_REQUESTS + | StatusCode::INTERNAL_SERVER_ERROR + | StatusCode::BAD_GATEWAY + | StatusCode::SERVICE_UNAVAILABLE + | StatusCode::GATEWAY_TIMEOUT + ) +} + +fn parse_json_rpc_error(body: &[u8]) -> Option { + match serde_json::from_slice::(body) { + Ok(message @ JsonRpcMessage::Error(_)) => Some(message), + _ => None, + } +} + async fn collect_body( body_stream: &mut HttpResponseBodyStream, ) -> std::result::Result, StreamableHttpError> { diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 90b09d724c..a28403e67c 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -74,6 +74,13 @@ use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; +#[path = "streamable_http_retry.rs"] +mod streamable_http_retry; + +use self::streamable_http_retry::HandshakeError; +use self::streamable_http_retry::STREAMABLE_HTTP_RETRY_DELAYS_MS; +use self::streamable_http_retry::sleep_with_retry_deadline; + enum PendingTransport { InProcess { transport: tokio::io::DuplexStream, @@ -223,6 +230,25 @@ enum ClientOperationError { Timeout { label: String, duration: Duration }, } +fn remaining_operation_timeout( + label: &str, + timeout: Option, + deadline: Option, +) -> std::result::Result, ClientOperationError> { + let Some(deadline) = deadline else { + return Ok(None); + }; + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + Err(ClientOperationError::Timeout { + label: label.to_string(), + duration: timeout.unwrap_or(remaining), + }) + } else { + Ok(Some(remaining)) + } +} + pub type Elicitation = CreateElicitationRequestParams; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -396,9 +422,13 @@ impl RmcpClient { } }; - let (service, oauth_persistor) = - Self::connect_pending_transport(pending_transport, client_service.clone(), timeout) - .await?; + let (service, oauth_persistor) = self + .connect_pending_transport_with_initialize_retries( + pending_transport, + client_service.clone(), + timeout, + ) + .await?; let initialize_result_rmcp = service .peer() @@ -845,14 +875,31 @@ impl RmcpClient { ), }; - let service = match timeout { - Some(duration) => time::timeout(duration, transport) - .await - .map_err(|_| anyhow!("timed out handshaking with MCP server after {duration:?}"))? - .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?, + let service_result = match timeout { + Some(duration) => match time::timeout(duration, transport).await { + Ok(result) => { + result.map_err(|source| anyhow::Error::from(HandshakeError { source })) + } + Err(_elapsed) => Err(anyhow!( + "timed out handshaking with MCP server after {duration:?}" + )), + }, None => transport .await - .map_err(|err| anyhow!("handshaking with MCP server failed: {err}"))?, + .map_err(|source| anyhow::Error::from(HandshakeError { source })), + }; + let service = match service_result { + Ok(service) => service, + Err(error) => { + if let Some(runtime) = oauth_persistor.as_ref() + && let Err(persist_error) = runtime.persist_if_needed().await + { + warn!( + "failed to persist OAuth tokens after failed initialize: {persist_error}" + ); + } + return Err(error); + } }; Ok((Arc::new(service), oauth_persistor)) @@ -869,7 +916,7 @@ impl RmcpClient { Fut: std::future::Future>, { let service = self.service().await?; - match Self::run_service_operation_once( + match Self::run_service_operation_with_transient_retries( Arc::clone(&service), label, timeout, @@ -882,7 +929,7 @@ impl RmcpClient { Err(error) if Self::is_session_expired_404(&error) => { self.reinitialize_after_session_expiry(&service).await?; let recovered_service = self.service().await?; - Self::run_service_operation_once( + Self::run_service_operation_with_transient_retries( recovered_service, label, timeout, @@ -896,6 +943,62 @@ impl RmcpClient { } } + async fn run_service_operation_with_transient_retries( + service: Arc>, + label: &str, + timeout: Option, + pause_state: ElicitationPauseState, + operation: &F, + ) -> std::result::Result + where + F: Fn(Arc>) -> Fut, + Fut: std::future::Future>, + { + let retry_deadline = timeout.map(|duration| Instant::now() + duration); + for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS + .iter() + .copied() + .map(Some) + .chain(std::iter::once(None)) + .enumerate() + { + let attempt_timeout = remaining_operation_timeout(label, timeout, retry_deadline)?; + match Self::run_service_operation_once( + Arc::clone(&service), + label, + attempt_timeout, + pause_state.clone(), + operation, + ) + .await + { + Ok(result) => return Ok(result), + Err(error) if Self::is_retryable_tools_list_error(label, &error) => { + let Some(retry_delay_ms) = retry_delay_ms else { + return Err(error); + }; + let delay = Duration::from_millis(retry_delay_ms); + warn!( + attempt = attempt + 1, + max_attempts = STREAMABLE_HTTP_RETRY_DELAYS_MS.len() + 1, + delay_ms = delay.as_millis(), + error = %error, + "streamable HTTP MCP tools/list failed with a retryable error; retrying" + ); + if !sleep_with_retry_deadline(delay, retry_deadline).await { + return Err(ClientOperationError::Timeout { + label: label.to_string(), + duration: timeout.unwrap_or(delay), + }); + } + } + Err(error) => return Err(error), + } + } + + unreachable!("service operation retry loop should return on success or final error") + } + async fn run_service_operation_once( service: Arc>, label: &str, @@ -921,6 +1024,22 @@ impl RmcpClient { } } + fn is_retryable_tools_list_error(label: &str, error: &ClientOperationError) -> bool { + if label != "tools/list" { + return false; + } + let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = + error + else { + return false; + }; + + error + .error + .downcast_ref::>() + .is_some_and(Self::is_retryable_streamable_http_error) + } + fn is_session_expired_404(error: &ClientOperationError) -> bool { let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = error @@ -974,12 +1093,13 @@ impl RmcpClient { .clone() .ok_or_else(|| anyhow!("MCP client cannot recover before initialize succeeds"))?; let pending_transport = Self::create_pending_transport(&self.transport_recipe).await?; - let (service, oauth_persistor) = Self::connect_pending_transport( - pending_transport, - initialize_context.client_service, - initialize_context.timeout, - ) - .await?; + let (service, oauth_persistor) = self + .connect_pending_transport_with_initialize_retries( + pending_transport, + initialize_context.client_service, + initialize_context.timeout, + ) + .await?; { let mut guard = self.state.lock().await; diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs new file mode 100644 index 0000000000..73da95de58 --- /dev/null +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -0,0 +1,239 @@ +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Result; +use anyhow::anyhow; +use codex_exec_server::ExecServerError; +use reqwest::StatusCode; +use rmcp::service::RoleClient; +use rmcp::service::RunningService; +use rmcp::transport::streamable_http_client::StreamableHttpError; +use tokio::time; +use tracing::warn; + +use crate::elicitation_client_service::ElicitationClientService; +use crate::http_client_adapter::StreamableHttpClientAdapterError; +use crate::oauth::OAuthPersistor; + +use super::PendingTransport; +use super::RmcpClient; + +const JSON_RPC_INTERNAL_ERROR_CODE: i64 = -32603; +pub(super) const STREAMABLE_HTTP_RETRY_DELAYS_MS: [u64; 2] = [250, 1_000]; + +impl RmcpClient { + pub(super) async fn connect_pending_transport_with_initialize_retries( + &self, + initial_transport: PendingTransport, + client_service: ElicitationClientService, + timeout: Option, + ) -> Result<( + Arc>, + Option, + )> { + let should_retry = match &initial_transport { + PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } => false, + PendingTransport::StreamableHttp { .. } + | PendingTransport::StreamableHttpWithOAuth { .. } => true, + }; + let retry_deadline = timeout.map(|duration| Instant::now() + duration); + let mut pending_transport = Some(initial_transport); + + for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS + .iter() + .copied() + .map(Some) + .chain(std::iter::once(None)) + .enumerate() + { + let transport = match pending_transport.take() { + Some(transport) => transport, + None => { + let remaining = remaining_initialize_timeout(timeout, retry_deadline)?; + match remaining { + Some(remaining) => time::timeout( + remaining, + Self::create_pending_transport(&self.transport_recipe), + ) + .await + .map_err(|_| initialize_timeout_error(timeout, remaining))??, + None => Self::create_pending_transport(&self.transport_recipe).await?, + } + } + }; + let attempt_timeout = remaining_initialize_timeout(timeout, retry_deadline)?; + + match Self::connect_pending_transport( + transport, + client_service.clone(), + attempt_timeout, + ) + .await + { + Ok(result) => return Ok(result), + Err(error) if should_retry && Self::is_retryable_initialize_error(&error) => { + let Some(retry_delay_ms) = retry_delay_ms else { + return Err(error); + }; + let delay = Duration::from_millis(retry_delay_ms); + warn!( + attempt = attempt + 1, + max_attempts = STREAMABLE_HTTP_RETRY_DELAYS_MS.len() + 1, + delay_ms = delay.as_millis(), + error = %error, + "streamable HTTP MCP initialize failed with a retryable error; retrying" + ); + if !sleep_with_retry_deadline(delay, retry_deadline).await { + let duration = timeout.unwrap_or(delay); + return Err(anyhow!( + "timed out handshaking with MCP server after {duration:?}" + )); + } + } + Err(error) => return Err(error), + } + } + + unreachable!("initialize retry loop should return on success or final error") + } + + fn is_retryable_initialize_error(error: &anyhow::Error) -> bool { + error.chain().any(|source| { + source + .downcast_ref::() + .is_some_and(|error| Self::is_retryable_client_initialize_error(&error.source)) + || source + .downcast_ref::() + .is_some_and(Self::is_retryable_client_initialize_error) + }) + } + + fn is_retryable_client_initialize_error(error: &rmcp::service::ClientInitializeError) -> bool { + match error { + rmcp::service::ClientInitializeError::TransportError { error, context } + if context.as_ref() == "send initialize request" => + { + error + .error + .downcast_ref::>() + .is_some_and(Self::is_retryable_streamable_http_error) + } + rmcp::service::ClientInitializeError::TransportError { error, context } + if context.as_ref() == "send initialized notification" => + { + error + .error + .downcast_ref::>() + .is_some_and(|error| { + matches!(error, StreamableHttpError::TransportChannelClosed) + || Self::is_retryable_streamable_http_error(error) + }) + } + _ => false, + } + } + + pub(super) fn is_retryable_streamable_http_error( + error: &StreamableHttpError, + ) -> bool { + match error { + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::HttpRequest(_), + )) => true, + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::Server { code, message }, + )) => { + *code == JSON_RPC_INTERNAL_ERROR_CODE && message.starts_with("http/request failed:") + } + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::Protocol(message), + )) => message.starts_with("http response stream `") && message.contains("` failed:"), + StreamableHttpError::UnexpectedServerResponse(message) => { + is_retryable_unexpected_server_response(message.as_ref()) + } + StreamableHttpError::AuthRequired(_) + | StreamableHttpError::InsufficientScope(_) + | StreamableHttpError::SessionExpired + | StreamableHttpError::UnexpectedContentType(_) + | StreamableHttpError::ServerDoesNotSupportSse + | StreamableHttpError::Deserialize(_) + | StreamableHttpError::Client(StreamableHttpClientAdapterError::SessionExpired404) + | StreamableHttpError::Client(StreamableHttpClientAdapterError::Header(_)) => false, + _ => false, + } + } +} + +fn is_retryable_unexpected_server_response(message: &str) -> bool { + let Some(message) = message.strip_prefix("HTTP ") else { + return false; + }; + let status_code = message + .chars() + .take_while(char::is_ascii_digit) + .collect::(); + let Ok(status) = status_code.parse::() else { + return false; + }; + let Ok(status) = StatusCode::from_u16(status) else { + return false; + }; + is_retryable_http_status(status) +} + +fn is_retryable_http_status(status: StatusCode) -> bool { + matches!( + status, + StatusCode::REQUEST_TIMEOUT + | StatusCode::TOO_MANY_REQUESTS + | StatusCode::INTERNAL_SERVER_ERROR + | StatusCode::BAD_GATEWAY + | StatusCode::SERVICE_UNAVAILABLE + | StatusCode::GATEWAY_TIMEOUT + ) +} + +fn remaining_initialize_timeout( + timeout: Option, + deadline: Option, +) -> Result> { + let Some(deadline) = deadline else { + return Ok(None); + }; + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + Err(initialize_timeout_error(timeout, remaining)) + } else { + Ok(Some(remaining)) + } +} + +fn initialize_timeout_error(timeout: Option, fallback: Duration) -> anyhow::Error { + let duration = timeout.unwrap_or(fallback); + anyhow!("timed out handshaking with MCP server after {duration:?}") +} + +pub(super) async fn sleep_with_retry_deadline(delay: Duration, deadline: Option) -> bool { + if let Some(deadline) = deadline { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return false; + } + time::timeout(remaining, time::sleep(delay)).await.is_ok() + } else { + time::sleep(delay).await; + true + } +} + +#[derive(Debug, thiserror::Error)] +#[error("handshaking with MCP server failed: {source}")] +pub(super) struct HandshakeError { + #[source] + pub(super) source: rmcp::service::ClientInitializeError, +} + +#[cfg(test)] +#[path = "streamable_http_retry_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs b/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs new file mode 100644 index 0000000000..ed0d8030e6 --- /dev/null +++ b/codex-rs/rmcp-client/src/streamable_http_retry_tests.rs @@ -0,0 +1,74 @@ +use std::any::TypeId; + +use codex_exec_server::ExecServerError; +use pretty_assertions::assert_eq; +use rmcp::transport::DynamicTransportError; +use rmcp::transport::streamable_http_client::StreamableHttpError; + +use crate::http_client_adapter::StreamableHttpClientAdapterError; + +use super::*; + +#[test] +fn retryable_initialize_error_includes_initialized_notification_context() { + let contexts = [ + "send initialize request", + "send initialized notification", + "receive initialize response", + ]; + + assert_eq!( + contexts.map(|context| { + RmcpClient::is_retryable_client_initialize_error(&retryable_initialize_error(context)) + }), + [true, true, false], + ); +} + +#[test] +fn retryable_streamable_http_error_includes_remote_body_stream_failure() { + let errors = [ + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::HttpRequest("error sending request for url".to_string()), + )), + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::Server { + code: JSON_RPC_INTERNAL_ERROR_CODE, + message: "http/request failed: error sending request for url".to_string(), + }, + )), + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::Protocol( + "http response stream `http-1` failed: exec-server transport disconnected" + .to_string(), + ), + )), + StreamableHttpError::Client(StreamableHttpClientAdapterError::HttpRequest( + ExecServerError::Protocol( + "http response stream `http-1` received seq 2, expected 1".to_string(), + ), + )), + StreamableHttpError::UnexpectedServerResponse("HTTP 502: upstream failure".into()), + StreamableHttpError::UnexpectedServerResponse("HTTP 400: bad request".into()), + ]; + + assert_eq!( + errors.map(|error| RmcpClient::is_retryable_streamable_http_error(&error)), + [true, true, true, false, true, false], + ); +} + +fn retryable_initialize_error(context: &'static str) -> rmcp::service::ClientInitializeError { + rmcp::service::ClientInitializeError::TransportError { + error: DynamicTransportError::from_parts( + "streamable_http", + TypeId::of::<()>(), + Box::new(StreamableHttpError::Client( + StreamableHttpClientAdapterError::HttpRequest(ExecServerError::HttpRequest( + "error sending request for url".to_string(), + )), + )), + ), + context: context.into(), + } +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs index 087d3d00df..29d2404ede 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs @@ -1,13 +1,229 @@ mod streamable_http_test_support; -use pretty_assertions::assert_eq; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; +use codex_exec_server::Environment; +use codex_exec_server::ExecServerError; +use codex_exec_server::HttpClient; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::HttpResponseBodyStream; +use futures::FutureExt as _; +use futures::future::BoxFuture; +use pretty_assertions::assert_eq; +use serde_json::Value; + +use streamable_http_test_support::arm_initialize_post_failure; +use streamable_http_test_support::arm_initialize_post_json_rpc_failure; +use streamable_http_test_support::arm_initialized_notification_post_json_rpc_failure; use streamable_http_test_support::arm_session_post_failure; +use streamable_http_test_support::arm_session_post_json_rpc_failure; use streamable_http_test_support::call_echo_tool; use streamable_http_test_support::create_client; +use streamable_http_test_support::create_client_with_http_client; use streamable_http_test_support::expected_echo_result; use streamable_http_test_support::spawn_streamable_http_server; +const JSON_RPC_INTERNAL_ERROR_CODE: i64 = -32603; +const SIMULATED_NO_RESPONSE_MESSAGE: &str = + "http/request failed: error sending request for url (simulated no response)"; + +#[derive(Clone)] +struct FailFirstInitializeHttpClient { + inner: Arc, + failures_remaining: Arc, + initialize_attempts: Arc, +} + +impl FailFirstInitializeHttpClient { + fn new(inner: Arc, failures_remaining: usize) -> Self { + Self { + inner, + failures_remaining: Arc::new(AtomicUsize::new(failures_remaining)), + initialize_attempts: Arc::new(AtomicUsize::new(0)), + } + } + + fn initialize_attempts(&self) -> usize { + self.initialize_attempts.load(Ordering::SeqCst) + } + + fn fail_next_initialize(&self) { + self.failures_remaining.store(1, Ordering::SeqCst); + } +} + +impl HttpClient for FailFirstInitializeHttpClient { + fn http_request( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + self.inner.http_request(params) + } + + fn http_request_stream( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + let inner = Arc::clone(&self.inner); + let failures_remaining = Arc::clone(&self.failures_remaining); + let initialize_attempts = Arc::clone(&self.initialize_attempts); + + async move { + if is_initialize_post(¶ms) { + initialize_attempts.fetch_add(1, Ordering::SeqCst); + if failures_remaining.swap(0, Ordering::SeqCst) > 0 { + return Err(ExecServerError::Server { + code: JSON_RPC_INTERNAL_ERROR_CODE, + message: SIMULATED_NO_RESPONSE_MESSAGE.to_string(), + }); + } + } + + inner.http_request_stream(params).await + } + .boxed() + } +} + +fn is_initialize_post(params: &HttpRequestParams) -> bool { + params.method.eq_ignore_ascii_case("POST") + && params + .body + .as_ref() + .and_then(|body| serde_json::from_slice::(&body.0).ok()) + .and_then(|body| { + body.get("method") + .and_then(Value::as_str) + .map(|method| method == "initialize") + }) + .unwrap_or(false) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_initialize_retries_remote_no_response_error() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let http_client = FailFirstInitializeHttpClient::new( + Environment::default_for_tests().get_http_client(), + /*failures_remaining*/ 1, + ); + + let client = create_client_with_http_client(&base_url, Arc::new(http_client.clone())).await?; + let result = call_echo_tool(&client, "after-init-retry").await?; + + assert_eq!(http_client.initialize_attempts(), 2); + assert_eq!(result, expected_echo_result("after-init-retry")); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_initialize_retries_transient_http_status() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + + arm_initialize_post_failure(&base_url, /*status*/ 502, /*remaining*/ 1).await?; + + let client = create_client(&base_url).await?; + let result = call_echo_tool(&client, "after-status-retry").await?; + + assert_eq!(result, expected_echo_result("after-status-retry")); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_initialize_retries_json_rpc_transient_status() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + + arm_initialize_post_json_rpc_failure(&base_url, /*status*/ 502, /*remaining*/ 1).await?; + + let client = create_client(&base_url).await?; + let result = call_echo_tool(&client, "after-json-status-retry").await?; + + assert_eq!(result, expected_echo_result("after-json-status-retry")); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_retries_initialized_notification_status() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + + arm_initialized_notification_post_json_rpc_failure( + &base_url, /*status*/ 502, /*remaining*/ 1, + ) + .await?; + + let client = create_client(&base_url).await?; + let result = call_echo_tool(&client, "after-notification-status-retry").await?; + + assert_eq!( + result, + expected_echo_result("after-notification-status-retry") + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_tools_list_retries_transient_http_status() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let client = create_client(&base_url).await?; + + let expected = client + .list_tools( + /*params*/ None, + /*timeout*/ Some(Duration::from_secs(5)), + ) + .await?; + arm_session_post_failure( + &base_url, + /*status*/ 502, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[], + ) + .await?; + + let result = client + .list_tools( + /*params*/ None, + /*timeout*/ Some(Duration::from_secs(5)), + ) + .await?; + + assert_eq!(result, expected); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_tools_list_retries_json_rpc_transient_status() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let client = create_client(&base_url).await?; + + let expected = client + .list_tools( + /*params*/ None, + /*timeout*/ Some(Duration::from_secs(5)), + ) + .await?; + arm_session_post_json_rpc_failure(&base_url, /*status*/ 502, /*remaining*/ 1).await?; + + let result = client + .list_tools( + /*params*/ None, + /*timeout*/ Some(Duration::from_secs(5)), + ) + .await?; + + assert_eq!(result, expected); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn streamable_http_404_session_expiry_recovers_and_retries_once() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; @@ -30,6 +246,34 @@ async fn streamable_http_404_session_expiry_recovers_and_retries_once() -> anyho Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_session_recovery_retries_initialize_failure() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let http_client = FailFirstInitializeHttpClient::new( + Environment::default_for_tests().get_http_client(), + /*failures_remaining*/ 0, + ); + let client = create_client_with_http_client(&base_url, Arc::new(http_client.clone())).await?; + + let warmup = call_echo_tool(&client, "warmup").await?; + assert_eq!(warmup, expected_echo_result("warmup")); + + arm_session_post_failure( + &base_url, + /*status*/ 404, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[], + ) + .await?; + http_client.fail_next_initialize(); + + let recovered = call_echo_tool(&client, "recovered-after-retry").await?; + assert_eq!(http_client.initialize_attempts(), 3); + assert_eq!(recovered, expected_echo_result("recovered-after-retry")); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn streamable_http_401_does_not_trigger_recovery() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; @@ -129,11 +373,10 @@ async fn streamable_http_404_recovery_only_retries_once() -> anyhow::Result<()> .await?; let error = call_echo_tool(&client, "double-404").await.unwrap_err(); + let error_message = error.to_string(); assert!( - error - .to_string() - .contains("handshaking with MCP server failed") - || error.to_string().contains("Transport channel closed") + error_message.contains("404") || error_message.contains("session expired"), + "expected session-expiry error, got: {error:#}" ); let recovered = call_echo_tool(&client, "after-double-404").await?; diff --git a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs index 03ee79392c..6b5df6e297 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -20,6 +20,7 @@ use anyhow::Context as _; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; use codex_exec_server::ExecServerClient; +use codex_exec_server::HttpClient; use codex_exec_server::RemoteExecServerConnectArgs; use codex_rmcp_client::ElicitationAction; use codex_rmcp_client::ElicitationResponse; @@ -44,6 +45,9 @@ use tokio::process::Command; use tokio::time::sleep; const SESSION_POST_FAILURE_CONTROL_PATH: &str = "/test/control/session-post-failure"; +const INITIALIZE_POST_FAILURE_CONTROL_PATH: &str = "/test/control/initialize-post-failure"; +const INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH: &str = + "/test/control/initialized-notification-post-failure"; fn streamable_http_server_bin() -> Result { codex_utils_cargo_bin::cargo_bin("test_streamable_http_server") @@ -74,6 +78,14 @@ pub(crate) fn expected_echo_result(message: &str) -> CallToolResult { } pub(crate) async fn create_client(base_url: &str) -> anyhow::Result { + create_client_with_http_client(base_url, Environment::default_for_tests().get_http_client()) + .await +} + +pub(crate) async fn create_client_with_http_client( + base_url: &str, + http_client: Arc, +) -> anyhow::Result { let client = RmcpClient::new_streamable_http_client( "test-streamable-http", &format!("{base_url}/mcp"), @@ -81,7 +93,7 @@ pub(crate) async fn create_client(base_url: &str) -> anyhow::Result /*http_headers*/ None, /*env_http_headers*/ None, OAuthCredentialsStoreMode::File, - Environment::default_for_tests().get_http_client(), + http_client, /*auth_provider*/ None, ) .await?; @@ -183,6 +195,107 @@ pub(crate) async fn arm_session_post_failure( Ok(()) } +pub(crate) async fn arm_session_post_json_rpc_failure( + base_url: &str, + status: u16, + remaining: usize, +) -> anyhow::Result<()> { + let response = reqwest::Client::new() + .post(format!("{base_url}{SESSION_POST_FAILURE_CONTROL_PATH}")) + .json(&json!({ + "status": status, + "remaining": remaining, + "content_type": "application/json", + "body": json!({ + "jsonrpc": "2.0", + "id": 1, + "error": { + "code": -32000, + "message": "transient session failure", + }, + }).to_string(), + })) + .send() + .await?; + + assert_eq!(response.status(), reqwest::StatusCode::NO_CONTENT); + Ok(()) +} + +pub(crate) async fn arm_initialized_notification_post_json_rpc_failure( + base_url: &str, + status: u16, + remaining: usize, +) -> anyhow::Result<()> { + let response = reqwest::Client::new() + .post(format!( + "{base_url}{INITIALIZED_NOTIFICATION_POST_FAILURE_CONTROL_PATH}" + )) + .json(&json!({ + "status": status, + "remaining": remaining, + "content_type": "application/json", + "body": json!({ + "jsonrpc": "2.0", + "id": 1, + "error": { + "code": -32000, + "message": "transient session failure", + }, + }).to_string(), + })) + .send() + .await?; + + assert_eq!(response.status(), reqwest::StatusCode::NO_CONTENT); + Ok(()) +} + +pub(crate) async fn arm_initialize_post_failure( + base_url: &str, + status: u16, + remaining: usize, +) -> anyhow::Result<()> { + let response = reqwest::Client::new() + .post(format!("{base_url}{INITIALIZE_POST_FAILURE_CONTROL_PATH}")) + .json(&json!({ + "status": status, + "remaining": remaining, + })) + .send() + .await?; + + assert_eq!(response.status(), reqwest::StatusCode::NO_CONTENT); + Ok(()) +} + +pub(crate) async fn arm_initialize_post_json_rpc_failure( + base_url: &str, + status: u16, + remaining: usize, +) -> anyhow::Result<()> { + let response = reqwest::Client::new() + .post(format!("{base_url}{INITIALIZE_POST_FAILURE_CONTROL_PATH}")) + .json(&json!({ + "status": status, + "remaining": remaining, + "content_type": "application/json", + "body": json!({ + "jsonrpc": "2.0", + "id": 1, + "error": { + "code": -32000, + "message": "transient initialize failure", + }, + }).to_string(), + })) + .send() + .await?; + + assert_eq!(response.status(), reqwest::StatusCode::NO_CONTENT); + Ok(()) +} + pub(crate) async fn spawn_streamable_http_server() -> anyhow::Result<(Child, String)> { let listener = TcpListener::bind("127.0.0.1:0")?; let port = listener.local_addr()?.port(); diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index eb81112965..682361ea9d 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -106,14 +106,13 @@ async fn append_rollout_item_materializes_compressed_rollout() -> anyhow::Result } #[tokio::test] -async fn search_rollout_matches_returns_compressed_snippet() -> anyhow::Result<()> { +async fn search_rollout_matches_uses_logical_path_for_compressed_rollout() -> anyhow::Result<()> { let home = TempDir::new()?; let uuid = Uuid::from_u128(15); let thread_id = ThreadId::from_string(&uuid.to_string())?; let rollout_path = rollout_path(home.path(), "2025-01-03T12-00-00", uuid); write_rollout(&rollout_path, thread_id, "targeted search term")?; compress_now(&rollout_path)?; - let compressed_path = compressed_rollout_path(&rollout_path); let matches = search_rollout_matches( std::path::Path::new("missing-rg-for-test"), @@ -124,7 +123,7 @@ async fn search_rollout_matches_returns_compressed_snippet() -> anyhow::Result<( .await?; assert_eq!( - matches.get(compressed_path.as_path()), + matches.get(rollout_path.as_path()), Some(&Some("targeted search term".to_string())) ); Ok(()) diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 9e6db0fb8d..23be07996c 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -73,6 +73,7 @@ pub use session_index::append_thread_name; pub use session_index::find_thread_meta_by_name_str; pub use session_index::find_thread_name_by_id; pub use session_index::find_thread_names_by_ids; +pub use session_index::remove_thread_name_entries; pub use state_db::StateDbHandle; pub use state_db::sqlite_telemetry_recorder; diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index fb740bb1df..af3785a975 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1348,32 +1348,56 @@ async fn find_thread_path_by_id_str_in_subdir( if !root.exists() { return Ok(unverified_db_path); } - // This is safe because we know the values are valid. - #[allow(clippy::unwrap_used)] - let limit = NonZero::new(1).unwrap(); - let options = file_search::FileSearchOptions { - limit, - compute_indices: false, - respect_gitignore: false, - ..Default::default() + let (filename_match, filename_scan_error) = match find_rollout_path_by_id_from_filenames( + root.as_path(), + id_str, + ) + .await + { + Ok(path) => (path, None), + Err(err) => { + tracing::warn!( + "rollout filename lookup failed during find_thread_path_by_id_str_in_subdir: {err}" + ); + (None, Some(err)) + } }; - let results = file_search::run( - id_str, - vec![root.clone()], - options, - /*cancel_flag*/ None, - ) - .map_err(|e| io::Error::other(format!("file search failed: {e}")))?; + let found = match filename_match { + Some(path) => Some(path), + None => { + // This is safe because we know the values are valid. + #[allow(clippy::unwrap_used)] + let limit = NonZero::new(1).unwrap(); + let options = file_search::FileSearchOptions { + limit, + compute_indices: false, + respect_gitignore: false, + ..Default::default() + }; - let found = match results - .matches - .into_iter() - .map(|m| m.full_path()) - .find_map(compression::RolloutFile::from_path) - { - Some(rollout_file) => Some(rollout_file.into_path()), - None => find_rollout_path_by_id_from_filenames(root.as_path(), id_str).await?, + let results = file_search::run( + id_str, + vec![root.clone()], + options, + /*cancel_flag*/ None, + ) + .map_err(|e| io::Error::other(format!("file search failed: {e}")))?; + + let found = results + .matches + .into_iter() + .map(|m| m.full_path()) + .find_map(compression::RolloutFile::from_path) + .map(compression::RolloutFile::into_path); + + if found.is_none() + && let Some(err) = filename_scan_error + { + return Err(err); + } + found + } }; if let Some(found_path) = found.as_ref() { tracing::debug!("state db missing rollout path for thread {id_str}"); diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index b53c958ae8..03d251e1b3 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -153,6 +153,7 @@ fn builder_from_items_falls_back_to_filename() { let items = vec![RolloutItem::Compacted(CompactedItem { message: "noop".to_string(), replacement_history: None, + window_id: None, })]; let builder = builder_from_items(items.as_slice(), path.as_path()).expect("builder"); diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 07029abe37..68896ad202 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -1144,6 +1144,7 @@ async fn resume_candidate_matches_cwd_reads_latest_turn_context() -> std::io::Re network: None, file_system_sandbox_policy: None, model: "test-model".to_string(), + comp_hash: None, personality: None, collaboration_mode: None, multi_agent_version: None, diff --git a/codex-rs/rollout/src/search.rs b/codex-rs/rollout/src/search.rs index 705d24e7b4..622fac8b0f 100644 --- a/codex-rs/rollout/src/search.rs +++ b/codex-rs/rollout/src/search.rs @@ -21,6 +21,7 @@ use super::compression; const MATCH_CONTEXT_BEFORE_CHARS: usize = 48; const MATCH_CONTEXT_AFTER_CHARS: usize = 96; +/// Search matches keyed by the canonical `.jsonl` path for each rollout. pub type RolloutSearchMatches = HashMap>; pub async fn search_rollout_paths( @@ -145,7 +146,10 @@ async fn scan_rollout_matches( if let Some(snippet) = first_rollout_content_match_snippet(rollout_file.path(), search_term).await? { - matches.insert(rollout_file.into_path(), Some(snippet)); + matches.insert( + compression::plain_rollout_path(rollout_file.path()), + Some(snippet), + ); } continue; } @@ -217,7 +221,10 @@ async fn scan_compressed_rollout_matches( if let Some(snippet) = first_rollout_content_match_snippet(rollout_file.path(), search_term).await? { - matches.insert(rollout_file.into_path(), Some(snippet)); + matches.insert( + compression::plain_rollout_path(rollout_file.path()), + Some(snippet), + ); } } } diff --git a/codex-rs/rollout/src/session_index.rs b/codex-rs/rollout/src/session_index.rs index e227515357..96ff9d825e 100644 --- a/codex-rs/rollout/src/session_index.rs +++ b/codex-rs/rollout/src/session_index.rs @@ -1,21 +1,25 @@ use std::collections::HashMap; use std::collections::HashSet; use std::fs::File; +use std::io::ErrorKind; use std::io::Read; use std::io::Seek; use std::io::SeekFrom; +use std::io::Write; use std::path::Path; use std::path::PathBuf; +use std::sync::LazyLock; +use std::sync::Mutex; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionMetaLine; use serde::Deserialize; use serde::Serialize; use tokio::io::AsyncBufReadExt; -use tokio::io::AsyncWriteExt; const SESSION_INDEX_FILE: &str = "session_index.jsonl"; const READ_CHUNK_SIZE: usize = 8192; +static SESSION_INDEX_LOCK: LazyLock> = LazyLock::new(|| Mutex::new(())); #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct SessionIndexEntry { @@ -25,7 +29,7 @@ pub struct SessionIndexEntry { } /// Append a thread name update to the session index. -/// The index is append-only; the most recent entry wins when resolving names or ids. +/// Name updates are append-only; the most recent entry wins when resolving names or ids. pub async fn append_thread_name( codex_home: &Path, thread_id: ThreadId, @@ -46,24 +50,60 @@ pub async fn append_thread_name( } /// Append a raw session index entry to `session_index.jsonl`. -/// The file is append-only; consumers scan from the end to find the newest match. +/// Consumers scan from the end to find the newest match. pub async fn append_session_index_entry( codex_home: &Path, entry: &SessionIndexEntry, ) -> std::io::Result<()> { + let _guard = SESSION_INDEX_LOCK + .lock() + .map_err(|err| std::io::Error::other(err.to_string()))?; let path = session_index_path(codex_home); - let mut file = tokio::fs::OpenOptions::new() + let mut file = std::fs::OpenOptions::new() .create(true) .append(true) - .open(&path) - .await?; + .open(&path)?; let mut line = serde_json::to_string(entry).map_err(std::io::Error::other)?; line.push('\n'); - file.write_all(line.as_bytes()).await?; - file.flush().await?; + file.write_all(line.as_bytes())?; + file.flush()?; Ok(()) } +/// Remove all recorded names for a thread from the session index. +pub async fn remove_thread_name_entries( + codex_home: &Path, + thread_id: ThreadId, +) -> std::io::Result<()> { + let _guard = SESSION_INDEX_LOCK + .lock() + .map_err(|err| std::io::Error::other(err.to_string()))?; + let path = session_index_path(codex_home); + let contents = match std::fs::read_to_string(&path) { + Ok(contents) => contents, + Err(err) if err.kind() == ErrorKind::NotFound => return Ok(()), + Err(err) => return Err(err), + }; + let mut removed = false; + let mut remaining = String::with_capacity(contents.len()); + for line in contents.lines() { + let should_remove = serde_json::from_str::(line.trim()) + .is_ok_and(|entry| entry.id == thread_id); + if should_remove { + removed = true; + } else { + remaining.push_str(line); + remaining.push('\n'); + } + } + if !removed { + return Ok(()); + } + let temp_path = path.with_extension("jsonl.tmp"); + std::fs::write(&temp_path, remaining)?; + std::fs::rename(temp_path, path) +} + /// Find the latest thread name for a thread id, if any. pub async fn find_thread_name_by_id( codex_home: &Path, diff --git a/codex-rs/rollout/src/state_db.rs b/codex-rs/rollout/src/state_db.rs index fa2905dd4f..c21d7002be 100644 --- a/codex-rs/rollout/src/state_db.rs +++ b/codex-rs/rollout/src/state_db.rs @@ -5,6 +5,7 @@ use crate::list::SortDirection; use crate::list::ThreadSortKey; use crate::metadata; use crate::sqlite_metrics; +use anyhow::Context; use chrono::DateTime; use chrono::Utc; use codex_protocol::ThreadId; @@ -50,7 +51,7 @@ pub async fn init(config: &impl RolloutConfigView) -> Option { { Ok(runtime) => Some(runtime), Err(err) => { - emit_startup_warning(&format!("failed to initialize state runtime: {err}")); + emit_startup_warning(&format!("failed to initialize state runtime: {err:#}")); None } } @@ -109,9 +110,9 @@ async fn try_init_with_roots_inner( let runtime = codex_state::StateRuntime::init(sqlite_home.clone(), default_model_provider_id.clone()) .await - .map_err(|err| { - anyhow::anyhow!( - "failed to initialize state runtime at {}: {err}", + .with_context(|| { + format!( + "failed to initialize state runtime at {}", sqlite_home.display() ) })?; @@ -128,7 +129,10 @@ async fn try_init_with_roots_inner( backfill_gate_started.elapsed(), &backfill_gate_result, ); - backfill_gate_result?; + if let Err(err) = backfill_gate_result { + runtime.close().await; + return Err(err); + } Ok(runtime) } diff --git a/codex-rs/state/migrations/0036_threads_visible_sort_indexes.sql b/codex-rs/state/migrations/0036_threads_visible_sort_indexes.sql new file mode 100644 index 0000000000..33a83a9d7b --- /dev/null +++ b/codex-rs/state/migrations/0036_threads_visible_sort_indexes.sql @@ -0,0 +1,7 @@ +CREATE INDEX idx_threads_visible_created_at_ms + ON threads(archived, created_at_ms DESC) + WHERE preview <> ''; + +CREATE INDEX idx_threads_visible_updated_at_ms + ON threads(archived, updated_at_ms DESC) + WHERE preview <> ''; diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 47285c5cf5..c1c2e56c93 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -356,6 +356,7 @@ mod tests { network: None, file_system_sandbox_policy: None, model: "gpt-5".to_string(), + comp_hash: None, personality: None, collaboration_mode: None, multi_agent_version: None, @@ -394,6 +395,7 @@ mod tests { network: None, file_system_sandbox_policy: None, model: "gpt-5".to_string(), + comp_hash: None, personality: None, collaboration_mode: None, multi_agent_version: None, @@ -429,6 +431,7 @@ mod tests { network: None, file_system_sandbox_policy: None, model: "gpt-5".to_string(), + comp_hash: None, personality: None, collaboration_mode: None, multi_agent_version: None, @@ -460,6 +463,7 @@ mod tests { network: None, file_system_sandbox_policy: None, model: "gpt-5".to_string(), + comp_hash: None, personality: None, collaboration_mode: None, multi_agent_version: None, diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index cb8c711e65..264b737e9f 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -57,15 +57,21 @@ pub use runtime::GoalStore; pub use runtime::GoalUpdate; pub use runtime::MemoryStore; pub use runtime::RemoteControlEnrollmentRecord; +pub use runtime::RuntimeDbBackup; pub use runtime::RuntimeDbPath; pub use runtime::ThreadFilterOptions; +pub use runtime::backup_runtime_db_for_fresh_start; pub use runtime::goals_db_filename; pub use runtime::goals_db_path; +pub use runtime::is_sqlite_corruption_error; pub use runtime::logs_db_filename; pub use runtime::logs_db_path; pub use runtime::memories_db_filename; pub use runtime::memories_db_path; +pub use runtime::runtime_db_path_for_corruption_error; pub use runtime::runtime_db_paths; +pub use runtime::sqlite_error_detail_is_corruption; +pub use runtime::sqlite_error_detail_is_lock; pub use runtime::sqlite_integrity_check; pub use runtime::state_db_filename; pub use runtime::state_db_path; diff --git a/codex-rs/state/src/model/thread_goal.rs b/codex-rs/state/src/model/thread_goal.rs index 4806250069..12d2a78dac 100644 --- a/codex-rs/state/src/model/thread_goal.rs +++ b/codex-rs/state/src/model/thread_goal.rs @@ -3,12 +3,14 @@ use anyhow::anyhow; use chrono::DateTime; use chrono::Utc; use codex_protocol::ThreadId; +use serde::Serialize; use sqlx::Row; use sqlx::sqlite::SqliteRow; use super::epoch_millis_to_datetime; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] pub enum ThreadGoalStatus { Active, Paused, diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 5ae41bcd8b..2937bf0c2e 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -62,6 +62,7 @@ mod backfill; mod goals; mod logs; mod memories; +mod recovery; mod remote_control; #[cfg(test)] mod test_support; @@ -72,6 +73,12 @@ pub use goals::GoalAccountingOutcome; pub use goals::GoalStore; pub use goals::GoalUpdate; pub use memories::MemoryStore; +pub use recovery::RuntimeDbBackup; +pub use recovery::backup_runtime_db_for_fresh_start; +pub use recovery::is_sqlite_corruption_error; +pub use recovery::runtime_db_path_for_corruption_error; +pub use recovery::sqlite_error_detail_is_corruption; +pub use recovery::sqlite_error_detail_is_lock; pub use remote_control::RemoteControlEnrollmentRecord; pub use threads::ThreadFilterOptions; @@ -200,6 +207,7 @@ impl StateRuntime { Ok(db) => Arc::new(db), Err(err) => { warn!("failed to open logs db at {}: {err}", logs_path.display()); + close_sqlite_pools(&[pool.as_ref()]).await; return Err(err); } }; @@ -208,6 +216,7 @@ impl StateRuntime { Ok(db) => Arc::new(db), Err(err) => { warn!("failed to open goals db at {}: {err}", goals_path.display()); + close_sqlite_pools(&[pool.as_ref(), logs_pool.as_ref()]).await; return Err(err); } }; @@ -224,6 +233,7 @@ impl StateRuntime { "failed to open memories db at {}: {err}", memories_path.display() ); + close_sqlite_pools(&[pool.as_ref(), logs_pool.as_ref(), goals_pool.as_ref()]).await; return Err(err); } }; @@ -236,7 +246,16 @@ impl StateRuntime { started.elapsed(), &backfill_state_result, ); - backfill_state_result?; + if let Err(err) = backfill_state_result { + close_sqlite_pools(&[ + pool.as_ref(), + logs_pool.as_ref(), + goals_pool.as_ref(), + memories_pool.as_ref(), + ]) + .await; + return Err(err); + } let started = Instant::now(); let thread_updated_at_millis_result: anyhow::Result> = sqlx::query_scalar("SELECT MAX(threads.updated_at_ms) FROM threads") @@ -250,7 +269,19 @@ impl StateRuntime { started.elapsed(), &thread_updated_at_millis_result, ); - let thread_updated_at_millis = thread_updated_at_millis_result?; + let thread_updated_at_millis = match thread_updated_at_millis_result { + Ok(value) => value, + Err(err) => { + close_sqlite_pools(&[ + pool.as_ref(), + logs_pool.as_ref(), + goals_pool.as_ref(), + memories_pool.as_ref(), + ]) + .await; + return Err(err); + } + }; let thread_updated_at_millis = thread_updated_at_millis.unwrap_or(0); let runtime = Arc::new(Self { thread_goals: GoalStore::new(Arc::clone(&goals_pool)), @@ -283,6 +314,14 @@ impl StateRuntime { &self.memories } + /// Close all SQLite pools and wait for outstanding pool workers to exit. + pub async fn close(&self) { + self.memories.close().await; + self.thread_goals.close().await; + self.logs_pool.close().await; + self.pool.close().await; + } + pub async fn clear_memory_data_in_sqlite_home(sqlite_home: &Path) -> anyhow::Result { let memories_path = MEMORIES_DB.path(sqlite_home); if !tokio::fs::try_exists(&memories_path).await? { @@ -302,6 +341,12 @@ impl StateRuntime { } } +async fn close_sqlite_pools(pools: &[&SqlitePool]) { + for pool in pools { + pool.close().await; + } +} + fn base_sqlite_options(path: &Path) -> SqliteConnectOptions { SqliteConnectOptions::new() .filename(path) @@ -367,7 +412,8 @@ async fn open_sqlite( started.elapsed(), &pool_result, ); - let pool = pool_result?; + let pool = pool_result + .map_err(|source| recovery::RuntimeDbInitError::new(spec.label, "open", path, source))?; let started = Instant::now(); let migrate_result = migrator.run(&pool).await.map_err(anyhow::Error::from); crate::telemetry::record_init_result( @@ -377,7 +423,10 @@ async fn open_sqlite( started.elapsed(), &migrate_result, ); - migrate_result?; + if let Err(source) = migrate_result { + pool.close().await; + return Err(recovery::RuntimeDbInitError::new(spec.label, "migrate", path, source).into()); + } Ok(pool) } diff --git a/codex-rs/state/src/runtime/goals.rs b/codex-rs/state/src/runtime/goals.rs index 50e8cf068e..e5769ccd1f 100644 --- a/codex-rs/state/src/runtime/goals.rs +++ b/codex-rs/state/src/runtime/goals.rs @@ -11,6 +11,10 @@ impl GoalStore { pub(crate) fn new(pool: Arc) -> Self { Self { pool } } + + pub(crate) async fn close(&self) { + self.pool.close().await; + } } pub struct GoalUpdate { @@ -377,18 +381,31 @@ WHERE thread_id = ? self.get_thread_goal(thread_id).await } - pub async fn delete_thread_goal(&self, thread_id: ThreadId) -> anyhow::Result { - let result = sqlx::query( + pub async fn delete_thread_goal( + &self, + thread_id: ThreadId, + ) -> anyhow::Result> { + let row = sqlx::query( r#" DELETE FROM thread_goals WHERE thread_id = ? +RETURNING + thread_id, + goal_id, + objective, + status, + token_budget, + tokens_used, + time_used_seconds, + created_at_ms, + updated_at_ms "#, ) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .fetch_optional(self.pool.as_ref()) .await?; - Ok(result.rows_affected() > 0) + row.map(|row| thread_goal_from_row(&row)).transpose() } pub async fn account_thread_goal_usage( @@ -622,7 +639,8 @@ mod tests { assert_eq!(0, replaced.tokens_used); assert_eq!(0, replaced.time_used_seconds); - assert!( + assert_eq!( + Some(replaced), runtime .thread_goals() .delete_thread_goal(thread_id) @@ -637,8 +655,9 @@ mod tests { .await .unwrap() ); - assert!( - !runtime + assert_eq!( + None, + runtime .thread_goals() .delete_thread_goal(thread_id) .await diff --git a/codex-rs/state/src/runtime/memories.rs b/codex-rs/state/src/runtime/memories.rs index 3101500796..895c254350 100644 --- a/codex-rs/state/src/runtime/memories.rs +++ b/codex-rs/state/src/runtime/memories.rs @@ -35,6 +35,10 @@ impl MemoryStore { Self { pool, state_pool } } + pub(crate) async fn close(&self) { + self.pool.close().await; + } + /// Deletes all persisted memory state in one transaction. /// /// This removes every `stage1_outputs` row and all `jobs` rows for the diff --git a/codex-rs/state/src/runtime/recovery.rs b/codex-rs/state/src/runtime/recovery.rs new file mode 100644 index 0000000000..d9cf0bcdc9 --- /dev/null +++ b/codex-rs/state/src/runtime/recovery.rs @@ -0,0 +1,243 @@ +//! Backup-and-rebuild support for Codex runtime SQLite databases. +//! +//! Codex keeps several independent runtime SQLite databases under one SQLite +//! home. When SQLite reports that one of them is corrupt, automatic recovery +//! moves only that database file and its sidecars into a backup folder so the +//! other databases keep their data. + +use std::borrow::Cow; +use std::path::Path; +use std::path::PathBuf; + +const BACKUP_DIR_NAME: &str = "db-backups"; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct RuntimeDbBackup { + /// Path where the runtime database or sidecar lived before it was moved. + pub original_path: PathBuf, + /// Path where the runtime database or sidecar was backed up. + pub backup_path: PathBuf, +} + +#[derive(Debug)] +pub(crate) struct RuntimeDbInitError { + label: &'static str, + operation: &'static str, + path: PathBuf, + source: anyhow::Error, +} + +impl RuntimeDbInitError { + pub(crate) fn new( + label: &'static str, + operation: &'static str, + path: &Path, + source: anyhow::Error, + ) -> Self { + Self { + label, + operation, + path: path.to_path_buf(), + source, + } + } + + fn path(&self) -> &Path { + self.path.as_path() + } +} + +impl std::fmt::Display for RuntimeDbInitError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "failed to {} {} at {}: {}", + self.operation, + self.label, + self.path.display(), + self.source + ) + } +} + +impl std::error::Error for RuntimeDbInitError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(self.source.as_ref()) + } +} + +/// Move one Codex runtime SQLite database out of the way so that database can +/// be recreated without discarding unrelated runtime databases. +pub async fn backup_runtime_db_for_fresh_start( + db_path: &Path, +) -> std::io::Result> { + let sqlite_home = db_path.parent().ok_or_else(|| { + std::io::Error::other(format!( + "database path does not have a parent directory: {}", + db_path.display() + )) + })?; + match tokio::fs::metadata(sqlite_home).await { + Ok(metadata) if metadata.is_dir() => backup_runtime_db_files(db_path).await, + Ok(_) => backup_blocking_sqlite_home(sqlite_home).await, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => { + tokio::fs::create_dir_all(sqlite_home).await?; + Err(std::io::Error::other(format!( + "no Codex runtime database files were found to back up for {}", + db_path.display() + ))) + } + Err(err) => Err(err), + } +} + +pub fn runtime_db_path_for_corruption_error(err: &anyhow::Error) -> Option { + if !is_sqlite_corruption_error(err) { + return None; + } + err.chain() + .find_map(|source| source.downcast_ref::()) + .map(|err| err.path().to_path_buf()) +} + +pub fn is_sqlite_corruption_error(err: &anyhow::Error) -> bool { + err.chain().any(sqlite_error_source_is_corruption) +} + +fn sqlite_error_source_is_corruption(source: &(dyn std::error::Error + 'static)) -> bool { + let Some(err) = source.downcast_ref::() else { + return false; + }; + let sqlx::Error::Database(database_error) = err else { + return false; + }; + sqlite_error_detail_is_corruption(database_error.message()) + || database_error + .code() + .is_some_and(sqlite_database_code_is_corruption) +} + +fn sqlite_database_code_is_corruption(code: Cow<'_, str>) -> bool { + matches!( + code.as_ref().to_ascii_lowercase().as_str(), + "11" | "26" | "sqlite_corrupt" | "sqlite_notadb" + ) +} + +pub fn sqlite_error_detail_is_corruption(detail: &str) -> bool { + let detail = detail.to_ascii_lowercase(); + detail.contains("database disk image is malformed") + || detail.contains("database schema is malformed") + || detail.contains("database is corrupt") + || detail.contains("file is not a database") + || detail.contains("sqlite_corrupt") + || detail.contains("sqlite_notadb") + || detail.contains("(code: 11)") + || detail.contains("(code: 26)") +} + +pub fn sqlite_error_detail_is_lock(detail: &str) -> bool { + let detail = detail.to_ascii_lowercase(); + detail.contains("database is locked") || detail.contains("database is busy") +} + +async fn backup_runtime_db_files(db_path: &Path) -> std::io::Result> { + let sqlite_home = db_path.parent().ok_or_else(|| { + std::io::Error::other(format!( + "database path does not have a parent directory: {}", + db_path.display() + )) + })?; + backup_sqlite_paths(sqlite_home, sqlite_paths(db_path)).await +} + +async fn backup_sqlite_paths( + sqlite_home: &Path, + paths: impl IntoIterator, +) -> std::io::Result> { + let backup_dir = create_unique_backup_dir(sqlite_home.join(BACKUP_DIR_NAME).as_path()).await?; + let mut backups = Vec::new(); + + for path in paths { + if tokio::fs::try_exists(path.as_path()).await? { + let backup_path = backup_dir.join(file_name(path.as_path())?); + tokio::fs::rename(path.as_path(), backup_path.as_path()).await?; + backups.push(RuntimeDbBackup { + original_path: path, + backup_path, + }); + } + } + + if backups.is_empty() { + let _ = tokio::fs::remove_dir(backup_dir).await; + return Err(std::io::Error::other( + "no Codex runtime database files were found to back up", + )); + } + + Ok(backups) +} + +async fn backup_blocking_sqlite_home(sqlite_home: &Path) -> std::io::Result> { + let parent = sqlite_home.parent().ok_or_else(|| { + std::io::Error::other(format!( + "cannot create a backup folder for {}", + sqlite_home.display() + )) + })?; + let mut backup_dir_name = file_name(sqlite_home)?.to_os_string(); + backup_dir_name.push(format!(".{BACKUP_DIR_NAME}")); + let backup_parent = parent.join(backup_dir_name); + let backup_dir = create_unique_backup_dir(backup_parent.as_path()).await?; + let backup_path = backup_dir.join(file_name(sqlite_home)?); + tokio::fs::rename(sqlite_home, backup_path.as_path()).await?; + tokio::fs::create_dir_all(sqlite_home).await?; + Ok(vec![RuntimeDbBackup { + original_path: sqlite_home.to_path_buf(), + backup_path, + }]) +} + +fn sqlite_paths(db_path: &Path) -> Vec { + let mut wal_path = db_path.as_os_str().to_os_string(); + wal_path.push("-wal"); + let mut shm_path = db_path.as_os_str().to_os_string(); + shm_path.push("-shm"); + vec![ + db_path.to_path_buf(), + PathBuf::from(wal_path), + PathBuf::from(shm_path), + ] +} + +async fn create_unique_backup_dir(backup_parent: &Path) -> std::io::Result { + tokio::fs::create_dir_all(backup_parent).await?; + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |duration| duration.as_secs()); + let mut sequence = 0_u32; + loop { + let backup_dir = backup_parent.join(format!("sqlite-{timestamp}-{sequence}")); + match tokio::fs::create_dir(backup_dir.as_path()).await { + Ok(()) => return Ok(backup_dir), + Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => { + sequence += 1; + } + Err(err) => return Err(err), + } + } +} + +fn file_name(path: &Path) -> std::io::Result<&std::ffi::OsStr> { + path.file_name().ok_or_else(|| { + std::io::Error::other(format!( + "cannot create a backup name for {}", + path.display() + )) + }) +} + +#[cfg(test)] +#[path = "recovery_tests.rs"] +mod tests; diff --git a/codex-rs/state/src/runtime/recovery_tests.rs b/codex-rs/state/src/runtime/recovery_tests.rs new file mode 100644 index 0000000000..7424dbcfe1 --- /dev/null +++ b/codex-rs/state/src/runtime/recovery_tests.rs @@ -0,0 +1,105 @@ +use super::*; +use crate::runtime::test_support::unique_temp_dir; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn backup_moves_only_requested_runtime_db_files_to_backup_folder() -> std::io::Result<()> { + let sqlite_home = unique_temp_dir(); + tokio::fs::create_dir_all(sqlite_home.as_path()).await?; + let runtime_paths = super::super::runtime_db_paths(sqlite_home.as_path()); + let mut expected_paths = Vec::new(); + for db_path in runtime_paths.iter().map(|db| db.path.as_path()) { + for path in sqlite_paths(db_path) { + tokio::fs::write(path.as_path(), path.display().to_string()).await?; + expected_paths.push(path); + } + } + let failed_db_path = super::super::logs_db_path(sqlite_home.as_path()); + let failed_paths = sqlite_paths(failed_db_path.as_path()); + + let backups = backup_runtime_db_for_fresh_start(failed_db_path.as_path()).await?; + + assert_eq!(backups.len(), failed_paths.len()); + for path in &failed_paths { + assert!(!tokio::fs::try_exists(path.as_path()).await?); + } + for path in expected_paths + .iter() + .filter(|path| !failed_paths.contains(path)) + { + assert!(tokio::fs::try_exists(path.as_path()).await?); + } + for backup in backups { + assert!( + backup + .backup_path + .starts_with(sqlite_home.join(BACKUP_DIR_NAME)) + ); + assert!(tokio::fs::try_exists(backup.backup_path.as_path()).await?); + } + Ok(()) +} + +#[tokio::test] +async fn backup_replaces_blocking_sqlite_home_file() -> std::io::Result<()> { + let temp_dir = unique_temp_dir(); + tokio::fs::create_dir_all(temp_dir.as_path()).await?; + let sqlite_home = temp_dir.join("sqlite-home"); + tokio::fs::write(sqlite_home.as_path(), b"not-a-directory").await?; + + let backups = backup_runtime_db_for_fresh_start( + super::super::state_db_path(sqlite_home.as_path()).as_path(), + ) + .await?; + + assert_eq!(backups.len(), 1); + assert!(tokio::fs::metadata(sqlite_home.as_path()).await?.is_dir()); + assert!( + backups[0] + .backup_path + .starts_with(temp_dir.join(format!("sqlite-home.{BACKUP_DIR_NAME}"))) + ); + assert!(tokio::fs::try_exists(backups[0].backup_path.as_path()).await?); + Ok(()) +} + +#[test] +fn sqlite_error_detail_classifies_corruption_and_lock_errors() { + assert!(sqlite_error_detail_is_corruption("file is not a database")); + assert!(sqlite_error_detail_is_corruption( + "error returned from database: (code: 11) database disk image is malformed" + )); + assert!(!sqlite_error_detail_is_corruption("database is locked")); + assert!(sqlite_error_detail_is_lock("database is locked")); + assert!(sqlite_error_detail_is_lock("database is busy")); +} + +#[tokio::test] +async fn runtime_db_path_for_corruption_error_returns_failed_database_path() -> std::io::Result<()> +{ + let sqlite_home = unique_temp_dir(); + tokio::fs::create_dir_all(sqlite_home.as_path()).await?; + let path = super::super::state_db_path(sqlite_home.as_path()); + tokio::fs::write(path.as_path(), b"not sqlite").await?; + + let err = match super::super::StateRuntime::init(sqlite_home, "openai".to_string()).await { + Ok(_) => panic!("malformed sqlite should fail to initialize"), + Err(err) => err, + }; + + assert_eq!(runtime_db_path_for_corruption_error(&err), Some(path)); + Ok(()) +} + +#[test] +fn runtime_db_path_for_corruption_error_ignores_corrupt_word_in_path() { + let path = PathBuf::from("/tmp/sqlite_corrupt/state_5.sqlite"); + let err = anyhow::Error::new(RuntimeDbInitError::new( + "state DB", + "open", + path.as_path(), + anyhow::anyhow!("permission denied"), + )); + + assert_eq!(runtime_db_path_for_corruption_error(&err), None); +} diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 6f3c45f0e2..98eff16c50 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -384,6 +384,7 @@ ON CONFLICT(child_thread_id) DO NOTHING &mut builder, crate::SortKey::UpdatedAt, SortDirection::Desc, + OrderByIndex::Enabled, /*limit*/ 1, ); @@ -399,14 +400,9 @@ ON CONFLICT(child_thread_id) DO NOTHING filters: ThreadFilterOptions<'_>, ) -> anyhow::Result { let limit = page_size.saturating_add(1); - let sort_key = filters.sort_key; - let sort_direction = filters.sort_direction; let mut builder = QueryBuilder::::new(""); - push_thread_select_columns(&mut builder); - builder.push(" FROM threads"); - push_thread_filters(&mut builder, filters); - push_thread_order_and_limit(&mut builder, sort_key, sort_direction, limit); + push_list_threads_query(&mut builder, filters, limit); let rows = builder.build().fetch_all(self.pool.as_ref()).await?; let mut items = rows @@ -418,7 +414,7 @@ ON CONFLICT(child_thread_id) DO NOTHING items.pop(); items .last() - .and_then(|item| anchor_from_item(item, sort_key)) + .and_then(|item| anchor_from_item(item, filters.sort_key)) } else { None }; @@ -453,7 +449,13 @@ ON CONFLICT(child_thread_id) DO NOTHING search_term: None, }, ); - push_thread_order_and_limit(&mut builder, sort_key, SortDirection::Desc, limit); + push_thread_order_and_limit( + &mut builder, + sort_key, + SortDirection::Desc, + OrderByIndex::Enabled, + limit, + ); let rows = builder.build().fetch_all(self.pool.as_ref()).await?; rows.into_iter() @@ -882,17 +884,117 @@ ON CONFLICT(id) DO UPDATE SET self.upsert_thread(&metadata).await } - /// Delete a thread metadata row by id. + /// Delete a thread and all associated state by id. pub async fn delete_thread(&self, thread_id: ThreadId) -> anyhow::Result { - let result = sqlx::query("DELETE FROM threads WHERE id = ?") - .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) - .await?; - let rows_affected = result.rows_affected(); - self.memories.delete_thread_memory(thread_id).await?; - if rows_affected > 0 { - self.thread_goals.delete_thread_goal(thread_id).await?; + self.delete_threads_strict(&[thread_id]).await + } + + /// Delete a set of threads and all associated state. + /// + /// Spawn edges and thread rows are deleted last so a failed delete can be retried with enough + /// state left to rediscover the same spawned subtree. + pub async fn delete_threads_strict(&self, thread_ids: &[ThreadId]) -> anyhow::Result { + if thread_ids.is_empty() { + return Ok(0); } + + let thread_id_strings = thread_ids + .iter() + .map(ThreadId::to_string) + .collect::>(); + for (thread_id, thread_id_string) in thread_ids.iter().zip(&thread_id_strings) { + sqlx::query("DELETE FROM logs WHERE thread_id = ?") + .bind(thread_id_string) + .execute(self.logs_pool.as_ref()) + .await?; + self.memories.delete_thread_memory(*thread_id).await?; + self.thread_goals.delete_thread_goal(*thread_id).await?; + } + + let now = Utc::now().timestamp(); + let mut tx = self.pool.begin().await?; + for thread_id_string in &thread_id_strings { + for parent_thread_id_string in &thread_id_strings { + // If both the job runner and worker are being deleted, requeueing + // the worker item would leave a running job with no loop to consume it. + sqlx::query( + r#" +UPDATE agent_jobs +SET status = ?, updated_at = ?, completed_at = ?, last_error = ? +WHERE status IN (?, ?) + AND id IN ( + SELECT item.job_id + FROM agent_job_items AS item + JOIN thread_spawn_edges AS edge ON edge.child_thread_id = item.assigned_thread_id + WHERE item.status = ? AND item.assigned_thread_id = ? AND edge.parent_thread_id = ? + ) + "#, + ) + .bind(AgentJobStatus::Cancelled.as_str()) + .bind(now) + .bind(now) + .bind("agent job runner thread was deleted") + .bind(AgentJobStatus::Pending.as_str()) + .bind(AgentJobStatus::Running.as_str()) + .bind(AgentJobItemStatus::Running.as_str()) + .bind(thread_id_string) + .bind(parent_thread_id_string) + .execute(&mut *tx) + .await?; + } + sqlx::query("DELETE FROM thread_dynamic_tools WHERE thread_id = ?") + .bind(thread_id_string) + .execute(&mut *tx) + .await?; + sqlx::query( + r#" +UPDATE agent_job_items +SET + status = ?, + assigned_thread_id = NULL, + updated_at = ?, + last_error = ? +WHERE assigned_thread_id = ? AND status = ? + "#, + ) + .bind(AgentJobItemStatus::Pending.as_str()) + .bind(now) + .bind("assigned thread was deleted") + .bind(thread_id_string) + .bind(AgentJobItemStatus::Running.as_str()) + .execute(&mut *tx) + .await?; + sqlx::query( + r#" +UPDATE agent_job_items +SET assigned_thread_id = NULL, updated_at = ? +WHERE assigned_thread_id = ? + "#, + ) + .bind(now) + .bind(thread_id_string) + .execute(&mut *tx) + .await?; + } + for thread_id_string in &thread_id_strings { + sqlx::query( + "DELETE FROM thread_spawn_edges WHERE parent_thread_id = ? OR child_thread_id = ?", + ) + .bind(thread_id_string) + .bind(thread_id_string) + .execute(&mut *tx) + .await?; + } + let mut rows_affected = 0; + for thread_id_string in &thread_id_strings { + rows_affected += sqlx::query("DELETE FROM threads WHERE id = ?") + .bind(thread_id_string) + .execute(&mut *tx) + .await? + .rows_affected(); + } + tx.commit().await?; + Ok(rows_affected) } } @@ -917,6 +1019,29 @@ fn one_thread_id_from_rows( } } +fn push_list_threads_query( + builder: &mut QueryBuilder, + filters: ThreadFilterOptions<'_>, + limit: usize, +) { + push_thread_select_columns(builder); + builder.push(" FROM threads"); + push_thread_filters(builder, filters); + let order_by_index = match filters.cwd_filters { + // Multi-cwd listing is supported but at the time of writing has no current use in production. + // Preserve its query plan so the global timestamp index does not regress cwd filtering into a scan. + Some(cwd_filters) if cwd_filters.len() > 1 => OrderByIndex::Disabled, + Some(_) | None => OrderByIndex::Enabled, + }; + push_thread_order_and_limit( + builder, + filters.sort_key, + filters.sort_direction, + order_by_index, + limit, + ); +} + pub(super) fn push_thread_select_columns(builder: &mut QueryBuilder) { builder.push( r#" @@ -1057,10 +1182,21 @@ pub(super) fn push_thread_filters<'a>( } } +/// Controls whether SQLite may use the ordered column to satisfy `ORDER BY` from an index. +/// +/// Disabling it adds a unary `+` to the ordered column. This preserves the sort semantics while +/// preventing a timestamp-only index from winning over a more selective filtering index. +#[derive(Clone, Copy)] +pub(super) enum OrderByIndex { + Enabled, + Disabled, +} + pub(super) fn push_thread_order_and_limit( builder: &mut QueryBuilder, sort_key: SortKey, sort_direction: SortDirection, + order_by_index: OrderByIndex, limit: usize, ) { let order_column = match sort_key { @@ -1072,6 +1208,12 @@ pub(super) fn push_thread_order_and_limit( SortDirection::Desc => "DESC", }; builder.push(" ORDER BY "); + match order_by_index { + OrderByIndex::Enabled => {} + OrderByIndex::Disabled => { + builder.push("+"); + } + } builder.push(order_column); builder.push(" "); builder.push(order_direction); @@ -1094,12 +1236,14 @@ mod tests { use crate::DirectionalThreadSpawnEdgeStatus; use crate::runtime::test_support::test_thread_metadata; use crate::runtime::test_support::unique_temp_dir; + use anyhow::Result; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::GitInfo; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use pretty_assertions::assert_eq; + use serde_json::json; use std::path::PathBuf; #[tokio::test] @@ -1140,6 +1284,180 @@ mod tests { assert_eq!(memory_mode, "disabled"); } + #[tokio::test] + async fn delete_thread_cleans_associated_state() -> Result<()> { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()).await?; + let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000401")?; + let child_thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000402")?; + runtime + .upsert_thread(&test_thread_metadata( + &codex_home, + thread_id, + codex_home.clone(), + )) + .await?; + seed_thread_cleanup_state(&runtime, thread_id, child_thread_id).await?; + sqlx::query("INSERT INTO thread_dynamic_tools (thread_id, position, name, description, input_schema) VALUES (?, ?, ?, ?, ?)") + .bind(thread_id.to_string()) + .bind(0_i64) + .bind("test_tool") + .bind("test dynamic tool") + .bind("{}") + .execute(runtime.pool.as_ref()) + .await?; + runtime + .create_agent_job( + &AgentJobCreateParams { + id: "job-1".to_string(), + name: "test-job".to_string(), + instruction: "Return a result".to_string(), + auto_export: true, + max_runtime_seconds: None, + output_schema_json: None, + input_headers: vec!["path".to_string()], + input_csv_path: "/tmp/in.csv".to_string(), + output_csv_path: "/tmp/out.csv".to_string(), + }, + &[AgentJobItemCreateParams { + item_id: "item-1".to_string(), + row_index: 0, + source_id: None, + row_json: json!({"path": "file-1"}), + }], + ) + .await?; + runtime.mark_agent_job_running("job-1").await?; + runtime + .mark_agent_job_item_running_with_thread( + "job-1", + "item-1", + &child_thread_id.to_string(), + ) + .await?; + + let rows = runtime + .delete_threads_strict(&[thread_id, child_thread_id]) + .await?; + + assert_eq!(rows, 1); + assert!(runtime.get_thread(thread_id).await?.is_none()); + let dynamic_tool_count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM thread_dynamic_tools WHERE thread_id = ?") + .bind(thread_id.to_string()) + .fetch_one(runtime.pool.as_ref()) + .await?; + assert_eq!(dynamic_tool_count, 0); + assert_thread_cleanup_state(&runtime, thread_id).await?; + let job_item = runtime + .get_agent_job_item("job-1", "item-1") + .await? + .expect("job item should exist"); + assert_eq!(job_item.status, AgentJobItemStatus::Pending); + assert_eq!(job_item.assigned_thread_id, None); + assert_eq!( + job_item.last_error, + Some("assigned thread was deleted".to_string()) + ); + let job = runtime + .get_agent_job("job-1") + .await? + .expect("job should exist"); + assert_eq!(job.status, AgentJobStatus::Cancelled); + + let missing_thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000403")?; + let missing_child_thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000404")?; + seed_thread_cleanup_state(&runtime, missing_thread_id, missing_child_thread_id).await?; + + assert_eq!(runtime.delete_thread(missing_thread_id).await?, 0); + assert_thread_cleanup_state(&runtime, missing_thread_id).await?; + Ok(()) + } + + #[tokio::test] + async fn delete_thread_keeps_retry_graph_on_cleanup_failure() -> Result<()> { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()).await?; + let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000405")?; + let child_thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000406")?; + runtime + .upsert_thread(&test_thread_metadata( + &codex_home, + thread_id, + codex_home.clone(), + )) + .await?; + seed_thread_cleanup_state(&runtime, thread_id, child_thread_id).await?; + + runtime.logs_pool.close().await; + runtime + .delete_thread(thread_id) + .await + .expect_err("closed log db should fail deletion"); + + assert!(runtime.get_thread(thread_id).await?.is_some()); + assert_eq!( + runtime.list_thread_spawn_descendants(thread_id).await?, + vec![child_thread_id] + ); + Ok(()) + } + + async fn seed_thread_cleanup_state( + runtime: &StateRuntime, + thread_id: ThreadId, + child_thread_id: ThreadId, + ) -> Result<()> { + runtime + .upsert_thread_spawn_edge( + thread_id, + child_thread_id, + DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await?; + runtime + .thread_goals() + .replace_thread_goal( + thread_id, + "test goal", + crate::ThreadGoalStatus::Active, + /*token_budget*/ None, + ) + .await?; + sqlx::query("INSERT INTO logs (ts, ts_nanos, level, target, feedback_log_body, thread_id) VALUES (1, 0, 'INFO', 'test', 'feedback log', ?)") + .bind(thread_id.to_string()) + .execute(runtime.logs_pool.as_ref()) + .await?; + Ok(()) + } + + async fn assert_thread_cleanup_state( + runtime: &StateRuntime, + thread_id: ThreadId, + ) -> Result<()> { + let spawn_edge_count: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM thread_spawn_edges WHERE parent_thread_id = ? OR child_thread_id = ?", + ) + .bind(thread_id.to_string()) + .bind(thread_id.to_string()) + .fetch_one(runtime.pool.as_ref()) + .await?; + assert_eq!(spawn_edge_count, 0); + assert_eq!( + runtime.thread_goals().get_thread_goal(thread_id).await?, + None + ); + let logs = runtime + .query_logs(&LogQuery { + thread_ids: vec![thread_id.to_string()], + ..Default::default() + }) + .await?; + assert!(logs.is_empty()); + Ok(()) + } + #[tokio::test] async fn list_threads_updated_after_returns_oldest_changes_first() { let codex_home = unique_temp_dir(); @@ -1255,9 +1573,9 @@ mod tests { } let cwd_filters = vec![first_cwd, second_cwd]; - let page = runtime + let first_page = runtime .list_threads( - /*page_size*/ 10, + /*page_size*/ 1, ThreadFilterOptions { archived_only: false, allowed_sources: &[], @@ -1272,8 +1590,44 @@ mod tests { .await .expect("list should succeed"); - let ids = page.items.iter().map(|item| item.id).collect::>(); - assert_eq!(ids, vec![second_id, first_id]); + let ids = first_page + .items + .iter() + .map(|item| item.id) + .collect::>(); + assert_eq!(ids, vec![second_id]); + assert_eq!( + first_page.next_anchor, + Some(Anchor { + ts: DateTime::::from_timestamp_millis(1_700_000_300_000) + .expect("valid timestamp"), + }) + ); + + let second_page = runtime + .list_threads( + /*page_size*/ 1, + ThreadFilterOptions { + archived_only: false, + allowed_sources: &[], + model_providers: None, + cwd_filters: Some(cwd_filters.as_slice()), + anchor: first_page.next_anchor.as_ref(), + sort_key: SortKey::UpdatedAt, + sort_direction: SortDirection::Desc, + search_term: None, + }, + ) + .await + .expect("second page should succeed"); + + let ids = second_page + .items + .iter() + .map(|item| item.id) + .collect::>(); + assert_eq!(ids, vec![first_id]); + assert_eq!(second_page.next_anchor, None); let page = runtime .list_threads( @@ -1295,6 +1649,85 @@ mod tests { assert_eq!(page.items, Vec::new()); } + #[tokio::test] + async fn list_threads_uses_indexes_matching_cwd_filters() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("state db should initialize"); + + let model_providers = ["test-provider".to_string()]; + let cwd_filters = [ + PathBuf::from("/workspace/one"), + PathBuf::from("/workspace/two"), + ]; + let anchor = Anchor { + ts: DateTime::::from_timestamp(1_700_000_000, 0).expect("valid timestamp"), + }; + for (sort_key, visible_index, cwd_index) in [ + ( + SortKey::CreatedAt, + "idx_threads_visible_created_at_ms", + "idx_threads_archived_cwd_created_at_ms", + ), + ( + SortKey::UpdatedAt, + "idx_threads_visible_updated_at_ms", + "idx_threads_archived_cwd_updated_at_ms", + ), + ] { + for (cwd_filters, anchor, expected_index, expect_temp_sort) in [ + (None, None, visible_index, false), + (Some(&cwd_filters[..1]), None, cwd_index, false), + ( + Some(&cwd_filters[..]), + None, + "idx_threads_archived_cwd_", + true, + ), + (Some(&cwd_filters[..]), Some(&anchor), cwd_index, true), + ] { + let mut builder = QueryBuilder::::new("EXPLAIN QUERY PLAN "); + push_list_threads_query( + &mut builder, + ThreadFilterOptions { + archived_only: false, + allowed_sources: &[], + model_providers: Some(&model_providers), + cwd_filters, + anchor, + sort_key, + sort_direction: SortDirection::Desc, + search_term: None, + }, + /*limit*/ 201, + ); + let plan_details = builder + .build() + .fetch_all(runtime.pool.as_ref()) + .await + .expect("query plan should load") + .into_iter() + .map(|row| row.get::("detail")) + .collect::>(); + + assert!( + plan_details + .iter() + .any(|detail| detail.contains(expected_index)), + "query plan did not use {expected_index}: {plan_details:?}" + ); + assert_eq!( + plan_details + .iter() + .any(|detail| detail.contains("TEMP B-TREE")), + expect_temp_sort, + "unexpected sorting plan: {plan_details:?}" + ); + } + } + } + #[tokio::test] async fn apply_rollout_items_restores_memory_mode_from_session_meta() { let codex_home = unique_temp_dir(); diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 97bfea8280..1bea0cae8b 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -251,7 +251,6 @@ fn new_config(model: Option, arg0_paths: Arg0DispatchPaths) -> anyhow::R model_catalog: None, model_verbosity: None, chatgpt_base_url: "https://chatgpt.com/backend-api/".to_string(), - apps_mcp_path_override: None, apps_mcp_product_sku: None, realtime_audio: RealtimeAudioConfig::default(), experimental_realtime_ws_base_url: None, diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index 2ab5103663..07156fe355 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -18,6 +18,7 @@ use codex_protocol::protocol::ThreadMemoryMode; use crate::AppendThreadItemsParams; use crate::ArchiveThreadParams; use crate::CreateThreadParams; +use crate::DeleteThreadParams; use crate::ListThreadsParams; use crate::LoadThreadHistoryParams; use crate::ReadThreadByRolloutPathParams; @@ -115,6 +116,7 @@ pub struct InMemoryThreadStoreCalls { pub update_thread_metadata: usize, pub archive_thread: usize, pub unarchive_thread: usize, + pub delete_thread: usize, } /// In-memory [`ThreadStore`] implementation for tests and debug configs. @@ -333,6 +335,25 @@ impl ThreadStore for InMemoryThreadStore { state.calls.unarchive_thread += 1; stored_thread_from_state(&state, params.thread_id, /*include_history*/ false) } + + async fn delete_thread(&self, params: DeleteThreadParams) -> ThreadStoreResult<()> { + let mut state = self.state.lock().await; + state.calls.delete_thread += 1; + let existed = state.histories.remove(¶ms.thread_id).is_some(); + state.created_threads.remove(¶ms.thread_id); + state.names.remove(¶ms.thread_id); + state.metadata_updates.remove(¶ms.thread_id); + state + .rollout_paths + .retain(|_, thread_id| *thread_id != params.thread_id); + if existed { + Ok(()) + } else { + Err(ThreadStoreError::ThreadNotFound { + thread_id: params.thread_id, + }) + } + } } fn stored_thread_from_state( diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 84b72f23f0..284614bf19 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -25,6 +25,7 @@ pub use types::AppendThreadItemsParams; pub use types::ArchiveThreadParams; pub use types::ClearableField; pub use types::CreateThreadParams; +pub use types::DeleteThreadParams; pub use types::ExtraConfig; pub use types::GitInfoPatch; pub use types::ItemPage; diff --git a/codex-rs/thread-store/src/local/delete_thread.rs b/codex-rs/thread-store/src/local/delete_thread.rs new file mode 100644 index 0000000000..ecd962f647 --- /dev/null +++ b/codex-rs/thread-store/src/local/delete_thread.rs @@ -0,0 +1,210 @@ +//! Local hard-delete support for persisted threads. +//! +//! Existing rollout files are deleted before this operation reports success. A rollout file that +//! vanishes after discovery counts as already deleted. SQLite cleanup happens at the app-server +//! layer after every associated rollout has been removed so failed deletes can be retried. + +use std::io::ErrorKind; +use std::path::Path; + +use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; +use codex_rollout::SESSIONS_SUBDIR; +use codex_rollout::find_archived_thread_path_by_id_str; +use codex_rollout::find_thread_path_by_id_str; +use codex_rollout::remove_thread_name_entries; + +use super::LocalThreadStore; +use super::helpers::matching_rollout_file_name; +use super::helpers::scoped_rollout_path; +use crate::DeleteThreadParams; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; + +pub(super) async fn delete_thread( + store: &LocalThreadStore, + params: DeleteThreadParams, +) -> ThreadStoreResult<()> { + let thread_id = params.thread_id; + let thread_id_str = thread_id.to_string(); + let state_db_ctx = store.state_db().await; + let mut rollout_paths = Vec::new(); + + match find_thread_path_by_id_str( + store.config.codex_home.as_path(), + thread_id_str.as_str(), + state_db_ctx.as_deref(), + ) + .await + { + Ok(Some(path)) => rollout_paths.push(path), + Ok(None) => {} + Err(err) => { + return Err(ThreadStoreError::InvalidRequest { + message: format!("failed to locate thread id {thread_id}: {err}"), + }); + } + } + + match find_archived_thread_path_by_id_str( + store.config.codex_home.as_path(), + thread_id_str.as_str(), + state_db_ctx.as_deref(), + ) + .await + { + Ok(Some(path)) => { + if !rollout_paths.contains(&path) { + rollout_paths.push(path); + } + } + Ok(None) => {} + Err(err) => { + return Err(ThreadStoreError::InvalidRequest { + message: format!("failed to locate archived thread id {thread_id}: {err}"), + }); + } + } + + let found_rollout_path = !rollout_paths.is_empty(); + for rollout_path in rollout_paths { + delete_rollout_file(store, rollout_path.as_path(), thread_id)?; + } + remove_thread_name_entries(store.config.codex_home.as_path(), thread_id) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to delete thread name index entries for {thread_id}: {err}"), + })?; + + if !found_rollout_path { + return Err(ThreadStoreError::ThreadNotFound { thread_id }); + } + + store.live_recorders.lock().await.remove(&thread_id); + + Ok(()) +} + +fn delete_rollout_file( + store: &LocalThreadStore, + rollout_path: &Path, + thread_id: codex_protocol::ThreadId, +) -> ThreadStoreResult { + let plain_path = codex_rollout::plain_rollout_path(rollout_path); + let compressed_path = plain_path.with_extension("jsonl.zst"); + let deleted_plain = delete_rollout_path(store, plain_path.as_path(), thread_id)?; + let deleted_compressed = delete_rollout_path(store, compressed_path.as_path(), thread_id)?; + Ok(deleted_plain || deleted_compressed) +} + +fn delete_rollout_path( + store: &LocalThreadStore, + rollout_path: &Path, + thread_id: codex_protocol::ThreadId, +) -> ThreadStoreResult { + let canonical_rollout_path = scoped_rollout_path( + store.config.codex_home.join(SESSIONS_SUBDIR), + rollout_path, + "sessions", + ) + .or_else(|_| { + scoped_rollout_path( + store.config.codex_home.join(ARCHIVED_SESSIONS_SUBDIR), + rollout_path, + "archived sessions", + ) + }) + .or_else(|err| match rollout_path.try_exists() { + Ok(false) => Ok(rollout_path.to_path_buf()), + Ok(true) | Err(_) => Err(err), + })?; + matching_rollout_file_name(&canonical_rollout_path, thread_id, rollout_path)?; + match std::fs::remove_file(&canonical_rollout_path) { + Ok(()) => Ok(true), + Err(err) if err.kind() == ErrorKind::NotFound => Ok(false), + Err(err) => Err(ThreadStoreError::Internal { + message: format!( + "failed to delete rollout file `{}`: {err}", + canonical_rollout_path.display() + ), + }), + } +} + +#[cfg(test)] +mod tests { + use codex_protocol::ThreadId; + use pretty_assertions::assert_eq; + use tempfile::TempDir; + use uuid::Uuid; + + use super::*; + use crate::ThreadStore; + use crate::local::LocalThreadStore; + use crate::local::test_support::test_config; + use crate::local::test_support::write_archived_session_file; + use crate::local::test_support::write_session_file; + + #[tokio::test] + async fn delete_thread_removes_active_and_archived_rollouts() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let active_path = + write_session_file(home.path(), "2025-01-03T12-00-00", Uuid::from_u128(301)) + .expect("session file"); + let compressed_path = active_path.with_extension("jsonl.zst"); + std::fs::write(&compressed_path, b"compressed sibling").expect("compressed sibling"); + let cases = [ + (Uuid::from_u128(301), active_path), + ( + Uuid::from_u128(302), + write_archived_session_file( + home.path(), + "2025-01-03T12-00-00", + Uuid::from_u128(302), + ) + .expect("archived session file"), + ), + ]; + + for (uuid, path) in cases { + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + store + .delete_thread(DeleteThreadParams { thread_id }) + .await + .expect("delete thread"); + + assert!(!path.exists()); + } + assert!(!compressed_path.exists()); + } + + #[tokio::test] + async fn delete_rollout_file_treats_vanished_path_as_already_deleted() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let uuid = Uuid::from_u128(305); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + let path = + write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); + std::fs::remove_file(&path).expect("remove session file"); + + assert!(!delete_rollout_file(&store, path.as_path(), thread_id).expect("delete rollout")); + } + + #[tokio::test] + async fn delete_thread_reports_missing_thread() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000304").expect("valid thread id"); + + let err = store + .delete_thread(DeleteThreadParams { thread_id }) + .await + .expect_err("missing thread should fail"); + assert_eq!( + err.to_string(), + "thread 00000000-0000-0000-0000-000000000304 not found" + ); + } +} diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index cef4496a02..3000fee9d9 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -1,5 +1,6 @@ mod archive_thread; mod create_thread; +mod delete_thread; mod helpers; mod list_threads; mod live_writer; @@ -24,6 +25,7 @@ use tokio::sync::Mutex; use crate::AppendThreadItemsParams; use crate::ArchiveThreadParams; use crate::CreateThreadParams; +use crate::DeleteThreadParams; use crate::ListThreadsParams; use crate::LoadThreadHistoryParams; use crate::ReadThreadByRolloutPathParams; @@ -288,6 +290,10 @@ impl ThreadStore for LocalThreadStore { ) -> ThreadStoreResult { unarchive_thread::unarchive_thread(self, params).await } + + async fn delete_thread(&self, params: DeleteThreadParams) -> ThreadStoreResult<()> { + delete_thread::delete_thread(self, params).await + } } #[cfg(test)] diff --git a/codex-rs/thread-store/src/local/search_threads.rs b/codex-rs/thread-store/src/local/search_threads.rs index 20f5da8243..c9394c3962 100644 --- a/codex-rs/thread-store/src/local/search_threads.rs +++ b/codex-rs/thread-store/src/local/search_threads.rs @@ -109,7 +109,8 @@ pub(super) async fn search_threads( ) .await?; for item in page.items { - let Some(snippet) = (match remaining_rollouts.remove(item.path.as_path()) { + let logical_path = codex_rollout::plain_rollout_path(item.path.as_path()); + let Some(snippet) = (match remaining_rollouts.remove(logical_path.as_path()) { Some(Some(snippet)) => Some(snippet), Some(None) => first_rollout_content_match_snippet(item.path.as_path(), search_term) .await diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 7276a980da..6373dd9b18 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -5,6 +5,7 @@ use std::any::Any; use crate::AppendThreadItemsParams; use crate::ArchiveThreadParams; use crate::CreateThreadParams; +use crate::DeleteThreadParams; use crate::ItemPage; use crate::ListItemsParams; use crate::ListThreadsParams; @@ -119,4 +120,7 @@ pub trait ThreadStore: Any + Send + Sync { &self, params: ArchiveThreadParams, ) -> ThreadStoreResult; + + /// Deletes a thread's persisted rollout data and associated metadata. + async fn delete_thread(&self, params: DeleteThreadParams) -> ThreadStoreResult<()>; } diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index 3db827e3a2..3a379a4815 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -474,6 +474,7 @@ mod tests { let item = RolloutItem::Compacted(CompactedItem { message: "compacted".to_string(), replacement_history: None, + window_id: None, }); let first = sync diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 1f18adf081..853d866bc8 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -653,6 +653,13 @@ pub struct ArchiveThreadParams { pub thread_id: ThreadId, } +/// Parameters for deleting a thread. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct DeleteThreadParams { + /// Thread id to delete. + pub thread_id: ThreadId, +} + #[cfg(test)] mod tests { use pretty_assertions::assert_eq; diff --git a/codex-rs/tools/Cargo.toml b/codex-rs/tools/Cargo.toml index 7cc0e34845..4b2a3257f3 100644 --- a/codex-rs/tools/Cargo.toml +++ b/codex-rs/tools/Cargo.toml @@ -8,7 +8,6 @@ version.workspace = true workspace = true [dependencies] -async-trait = { workspace = true } codex-app-server-protocol = { workspace = true } codex-code-mode = { workspace = true } codex-features = { workspace = true } diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 1cee3f7d2d..51c313eb0e 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -93,13 +93,13 @@ pub use tool_discovery::ToolSearchSourceInfo; pub use tool_discovery::collect_request_plugin_install_entries; pub use tool_discovery::filter_request_plugin_install_discoverable_tools_for_client; pub use tool_executor::ToolExecutor; +pub use tool_executor::ToolExecutorFuture; pub use tool_executor::ToolExposure; pub use tool_output::JsonToolOutput; pub use tool_output::ToolOutput; pub use tool_payload::ToolPayload; pub use tool_search::ToolSearchEntry; pub use tool_search::ToolSearchInfo; -pub use tool_search::default_tool_search_text; pub use tool_spec::ResponsesApiWebSearchFilters; pub use tool_spec::ResponsesApiWebSearchUserLocation; pub use tool_spec::ToolSpec; diff --git a/codex-rs/tools/src/request_plugin_install.rs b/codex-rs/tools/src/request_plugin_install.rs index 70e9cb093d..5d6352b7f8 100644 --- a/codex-rs/tools/src/request_plugin_install.rs +++ b/codex-rs/tools/src/request_plugin_install.rs @@ -47,6 +47,10 @@ pub struct RequestPluginInstallMeta<'a> { pub tool_name: &'a str, #[serde(skip_serializing_if = "Option::is_none")] pub install_url: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + pub remote_plugin_id: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + pub app_connector_ids: Option<&'a [String]>, } pub fn build_request_plugin_install_elicitation_request( @@ -57,8 +61,6 @@ pub fn build_request_plugin_install_elicitation_request( suggest_reason: &str, tool: &DiscoverableTool, ) -> McpServerElicitationRequestParams { - let tool_name = tool.name().to_string(); - let install_url = tool.install_url().map(ToString::to_string); let message = suggest_reason.to_string(); McpServerElicitationRequestParams { @@ -70,9 +72,7 @@ pub fn build_request_plugin_install_elicitation_request( args.tool_type, args.action_type, suggest_reason, - tool.id(), - tool_name.as_str(), - install_url.as_deref(), + tool, ))), message, requested_schema: McpElicitationSchema { @@ -108,19 +108,26 @@ fn build_request_plugin_install_meta<'a>( tool_type: DiscoverableToolType, action_type: DiscoverableToolAction, suggest_reason: &'a str, - tool_id: &'a str, - tool_name: &'a str, - install_url: Option<&'a str>, + tool: &'a DiscoverableTool, ) -> RequestPluginInstallMeta<'a> { + let (remote_plugin_id, app_connector_ids) = match tool { + DiscoverableTool::Connector(_) => (None, None), + DiscoverableTool::Plugin(plugin) => ( + plugin.remote_plugin_id.as_deref(), + Some(plugin.app_connector_ids.as_slice()), + ), + }; RequestPluginInstallMeta { codex_approval_kind: REQUEST_PLUGIN_INSTALL_APPROVAL_KIND_VALUE, persist: REQUEST_PLUGIN_INSTALL_PERSIST_ALWAYS_VALUE, tool_type, suggest_type: action_type, suggest_reason, - tool_id, - tool_name, - install_url, + tool_id: tool.id(), + tool_name: tool.name(), + install_url: tool.install_url(), + remote_plugin_id, + app_connector_ids, } } diff --git a/codex-rs/tools/src/request_plugin_install_tests.rs b/codex-rs/tools/src/request_plugin_install_tests.rs index ff2370ade7..ab8af9be67 100644 --- a/codex-rs/tools/src/request_plugin_install_tests.rs +++ b/codex-rs/tools/src/request_plugin_install_tests.rs @@ -57,6 +57,8 @@ fn build_request_plugin_install_elicitation_request_uses_expected_shape() { install_url: Some( "https://chatgpt.com/apps/google-calendar/connector_2128aebfecb84f64a069897515042a44" ), + remote_plugin_id: None, + app_connector_ids: None, })), message: "Plan and reference events from your calendar".to_string(), requested_schema: McpElicitationSchema { @@ -71,15 +73,16 @@ fn build_request_plugin_install_elicitation_request_uses_expected_shape() { } #[test] -fn build_request_plugin_install_elicitation_request_for_plugin_omits_install_url() { +fn build_request_plugin_install_elicitation_request_injects_plugin_metadata() { let args = RequestPluginInstallArgs { tool_type: DiscoverableToolType::Plugin, action_type: DiscoverableToolAction::Install, - tool_id: "sample@openai-curated".to_string(), + tool_id: "sample@openai-curated-remote".to_string(), suggest_reason: "Use the sample plugin's skills and MCP server".to_string(), }; let plugin = DiscoverableTool::Plugin(Box::new(DiscoverablePluginInfo { - id: "sample@openai-curated".to_string(), + id: "sample@openai-curated-remote".to_string(), + remote_plugin_id: Some("plugins~Plugin_sample".to_string()), name: "Sample Plugin".to_string(), description: Some("Includes skills, MCP servers, and apps.".to_string()), has_skills: true, @@ -109,9 +112,11 @@ fn build_request_plugin_install_elicitation_request_for_plugin_omits_install_url tool_type: DiscoverableToolType::Plugin, suggest_type: DiscoverableToolAction::Install, suggest_reason: "Use the sample plugin's skills and MCP server", - tool_id: "sample@openai-curated", + tool_id: "sample@openai-curated-remote", tool_name: "Sample Plugin", install_url: None, + remote_plugin_id: Some("plugins~Plugin_sample"), + app_connector_ids: Some(&["connector_calendar".to_string()]), })), message: "Use the sample plugin's skills and MCP server".to_string(), requested_schema: McpElicitationSchema { @@ -127,13 +132,28 @@ fn build_request_plugin_install_elicitation_request_for_plugin_omits_install_url #[test] fn build_request_plugin_install_meta_uses_expected_shape() { + let connector = DiscoverableTool::Connector(Box::new(AppInfo { + id: "connector_68df038e0ba48191908c8434991bbac2".to_string(), + name: "Gmail".to_string(), + description: None, + logo_url: None, + logo_url_dark: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: Some( + "https://chatgpt.com/apps/gmail/connector_68df038e0ba48191908c8434991bbac2".to_string(), + ), + is_accessible: false, + is_enabled: true, + plugin_display_names: Vec::new(), + })); let meta = build_request_plugin_install_meta( DiscoverableToolType::Connector, DiscoverableToolAction::Install, "Find and reference emails from your inbox", - "connector_68df038e0ba48191908c8434991bbac2", - "Gmail", - Some("https://chatgpt.com/apps/gmail/connector_68df038e0ba48191908c8434991bbac2"), + &connector, ); assert_eq!( @@ -149,6 +169,8 @@ fn build_request_plugin_install_meta_uses_expected_shape() { install_url: Some( "https://chatgpt.com/apps/gmail/connector_68df038e0ba48191908c8434991bbac2" ), + remote_plugin_id: None, + app_connector_ids: None, }, ); } diff --git a/codex-rs/tools/src/tool_config_tests.rs b/codex-rs/tools/src/tool_config_tests.rs index 0e3917d697..7a4738772c 100644 --- a/codex-rs/tools/src/tool_config_tests.rs +++ b/codex-rs/tools/src/tool_config_tests.rs @@ -39,6 +39,7 @@ fn model_with_shell_type(shell_type: ConfigShellToolType) -> ModelInfo { context_window: None, max_context_window: None, auto_compact_token_limit: None, + comp_hash: None, effective_context_window_percent: 95, experimental_supported_tools: Vec::new(), input_modalities: codex_protocol::openai_models::default_input_modalities(), diff --git a/codex-rs/tools/src/tool_discovery.rs b/codex-rs/tools/src/tool_discovery.rs index ac5aac904b..2584d7f17a 100644 --- a/codex-rs/tools/src/tool_discovery.rs +++ b/codex-rs/tools/src/tool_discovery.rs @@ -93,6 +93,7 @@ pub fn filter_request_plugin_install_discoverable_tools_for_client( #[derive(Clone, Debug, PartialEq, Eq)] pub struct DiscoverablePluginInfo { pub id: String, + pub remote_plugin_id: Option, pub name: String, pub description: Option, pub has_skills: bool, diff --git a/codex-rs/tools/src/tool_discovery_tests.rs b/codex-rs/tools/src/tool_discovery_tests.rs index 6e45260c0e..5a3310057c 100644 --- a/codex-rs/tools/src/tool_discovery_tests.rs +++ b/codex-rs/tools/src/tool_discovery_tests.rs @@ -37,6 +37,7 @@ fn filter_request_plugin_install_discoverable_tools_for_codex_tui_omits_plugins( })), DiscoverableTool::Plugin(Box::new(DiscoverablePluginInfo { id: "slack@openai-curated".to_string(), + remote_plugin_id: None, name: "Slack".to_string(), description: Some("Search Slack messages".to_string()), has_skills: true, diff --git a/codex-rs/tools/src/tool_executor.rs b/codex-rs/tools/src/tool_executor.rs index 541d2443e0..294b361f09 100644 --- a/codex-rs/tools/src/tool_executor.rs +++ b/codex-rs/tools/src/tool_executor.rs @@ -3,6 +3,12 @@ use crate::ToolName; use crate::ToolOutput; use crate::ToolSearchInfo; use crate::ToolSpec; +use std::future::Future; +use std::pin::Pin; + +/// The boxed future returned by [`ToolExecutor::handle`]. +pub type ToolExecutorFuture<'a> = + Pin, FunctionCallError>> + Send + 'a>>; /// Controls where a tool is exposed to the model. #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -40,7 +46,6 @@ impl ToolExposure { /// Implementations keep the model-visible spec tied to the executable runtime. /// Host crates can layer routing, hooks, telemetry, or other orchestration on /// top without reopening the spec/runtime split. -#[async_trait::async_trait] pub trait ToolExecutor: Send + Sync { /// The concrete tool name handled by this runtime instance. fn tool_name(&self) -> ToolName; @@ -53,15 +58,12 @@ pub trait ToolExecutor: Send + Sync { fn search_info(&self) -> Option { let spec = self.spec(); - ToolSearchInfo::from_tool_spec(&self.tool_name(), spec, /*source_info*/ None) + ToolSearchInfo::from_tool_spec(spec, /*source_info*/ None) } fn supports_parallel_tool_calls(&self) -> bool { false } - async fn handle( - &self, - invocation: Invocation, - ) -> Result, FunctionCallError>; + fn handle(&self, invocation: Invocation) -> ToolExecutorFuture<'_>; } diff --git a/codex-rs/tools/src/tool_output.rs b/codex-rs/tools/src/tool_output.rs index fe57f23ce9..f2eed18fac 100644 --- a/codex-rs/tools/src/tool_output.rs +++ b/codex-rs/tools/src/tool_output.rs @@ -93,6 +93,7 @@ where pub struct JsonToolOutput { value: JsonValue, success: Option, + contains_external_context: bool, } impl JsonToolOutput { @@ -100,11 +101,21 @@ impl JsonToolOutput { Self { value, success: Some(true), + contains_external_context: false, } } pub fn with_success(value: JsonValue, success: Option) -> Self { - Self { value, success } + Self { + value, + success, + contains_external_context: false, + } + } + + pub fn with_external_context(mut self) -> Self { + self.contains_external_context = true; + self } } @@ -117,6 +128,10 @@ impl ToolOutput for JsonToolOutput { self.success.unwrap_or(true) } + fn contains_external_context(&self) -> bool { + self.contains_external_context + } + fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem { let output = FunctionCallOutputPayload { body: FunctionCallOutputBody::Text(self.value.to_string()), diff --git a/codex-rs/tools/src/tool_search.rs b/codex-rs/tools/src/tool_search.rs index ecb5504d4f..1083d9ea92 100644 --- a/codex-rs/tools/src/tool_search.rs +++ b/codex-rs/tools/src/tool_search.rs @@ -2,7 +2,6 @@ use crate::JsonSchema; use crate::LoadableToolSpec; use crate::ResponsesApiNamespaceTool; use crate::ResponsesApiTool; -use crate::ToolName; use crate::ToolSearchSourceInfo; use crate::ToolSpec; use crate::default_namespace_description; @@ -21,11 +20,10 @@ pub struct ToolSearchInfo { impl ToolSearchInfo { pub fn from_tool_spec( - tool_name: &ToolName, spec: ToolSpec, source_info: Option, ) -> Option { - let search_text = default_tool_search_text(tool_name, &spec); + let search_text = default_tool_search_text(&spec); Self::from_spec(search_text, spec, source_info) } @@ -67,13 +65,8 @@ impl ToolSearchInfo { } } -pub fn default_tool_search_text(tool_name: &ToolName, spec: &ToolSpec) -> String { +fn default_tool_search_text(spec: &ToolSpec) -> String { let mut parts = Vec::new(); - push_search_part(&mut parts, tool_name.to_string()); - push_search_part(&mut parts, tool_name.name.replace('_', " ")); - if let Some(namespace) = &tool_name.namespace { - push_search_part(&mut parts, namespace.clone()); - } match spec { ToolSpec::Function(tool) => append_function_search_text(tool, &mut parts), @@ -137,3 +130,7 @@ fn push_search_part(parts: &mut Vec, part: String) { parts.push(part.to_string()); } } + +#[cfg(test)] +#[path = "tool_search_tests.rs"] +mod tests; diff --git a/codex-rs/tools/src/tool_search_tests.rs b/codex-rs/tools/src/tool_search_tests.rs new file mode 100644 index 0000000000..bf910c57f1 --- /dev/null +++ b/codex-rs/tools/src/tool_search_tests.rs @@ -0,0 +1,48 @@ +use super::*; +use pretty_assertions::assert_eq; +use std::collections::BTreeMap; + +#[test] +fn default_search_text_uses_model_visible_namespace_metadata_once() { + let mut schedule_schema = JsonSchema::object( + BTreeMap::from([( + "timezone".to_string(), + JsonSchema::string(Some("IANA timezone.".to_string())), + )]), + /*required*/ None, + /*additional_properties*/ None, + ); + schedule_schema.description = Some("Schedule settings.".to_string()); + let mut parameters = JsonSchema::object( + BTreeMap::from([ + ( + "mode".to_string(), + JsonSchema::string(Some("Update mode.".to_string())), + ), + ("schedule".to_string(), schedule_schema), + ]), + /*required*/ None, + /*additional_properties*/ None, + ); + parameters.description = Some("Automation options.".to_string()); + let spec = ToolSpec::Namespace(crate::ResponsesApiNamespace { + name: "codex_app".to_string(), + description: "Manage Codex automations.".to_string(), + tools: vec![ResponsesApiNamespaceTool::Function(ResponsesApiTool { + name: "automation_update".to_string(), + description: "Create or update automations.".to_string(), + strict: false, + defer_loading: None, + parameters, + output_schema: None, + })], + }); + + let search_info = ToolSearchInfo::from_tool_spec(spec, /*source_info*/ None) + .expect("namespace should be searchable"); + + assert_eq!( + search_info.entry.search_text, + "codex_app Manage Codex automations. automation_update automation update Create or update automations. Automation options. mode Update mode. schedule Schedule settings. timezone IANA timezone." + ); +} diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 5916456956..a58d5f2398 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -11,6 +11,7 @@ use crate::app_event::ExitMode; use crate::app_event::FeedbackCategory; use crate::app_event::HistoryLookupResponse; use crate::app_event::PermissionProfileSelection; +use crate::app_event::PluginLocation; use crate::app_event::RateLimitRefreshOrigin; use crate::app_event::RealtimeAudioDeviceKind; #[cfg(target_os = "windows")] @@ -37,8 +38,6 @@ use crate::cwd_prompt::CwdPromptAction; use crate::diff_render::DiffSummary; use crate::exec_command::split_command_string; use crate::exec_command::strip_bash_lc_and_escape; -use crate::external_agent_config_migration_startup::ExternalAgentConfigMigrationStartupOutcome; -use crate::external_agent_config_migration_startup::handle_external_agent_config_migration_prompt_if_needed; use crate::external_editor; use crate::file_search::FileSearchManager; use crate::history_cell; @@ -764,7 +763,6 @@ impl App { session_selection: SessionSelection, feedback: codex_feedback::CodexFeedback, is_first_run: bool, - entered_trust_nux: bool, should_prompt_windows_sandbox_nux_at_startup: bool, app_server_target: AppServerTarget, state_db: Option, @@ -786,39 +784,6 @@ impl App { let harness_overrides = normalize_harness_overrides_for_cwd(harness_overrides, &config.cwd)?; - let external_agent_config_migration_outcome = - handle_external_agent_config_migration_prompt_if_needed( - tui, - &mut app_server, - &mut config, - &cli_kv_overrides, - &harness_overrides, - &cloud_config_bundle, - entered_trust_nux, - ) - .await?; - let external_agent_config_migration_message = match external_agent_config_migration_outcome - { - ExternalAgentConfigMigrationStartupOutcome::Continue { success_message } => { - success_message - } - ExternalAgentConfigMigrationStartupOutcome::ExitRequested => { - app_server - .shutdown() - .await - .inspect_err(|err| { - tracing::warn!("app-server shutdown failed: {err}"); - }) - .ok(); - return Ok(AppExitInfo { - token_usage: TokenUsage::default(), - thread_id: None, - thread_name: None, - update_action: None, - exit_reason: ExitReason::UserRequested, - }); - } - }; let bootstrap = match startup_bootstrap { Some(bootstrap) => bootstrap, None => app_server.bootstrap(&config).await?, @@ -1018,10 +983,6 @@ impl App { }; chat_widget.remote_connection = remote_connection; let thread_and_widget_ms = thread_and_widget_started_at.elapsed().as_millis(); - if let Some(message) = external_agent_config_migration_message { - chat_widget.add_info_message(message, /*hint*/ None); - } - chat_widget .maybe_prompt_windows_sandbox_enable(should_prompt_windows_sandbox_nux_at_startup); diff --git a/codex-rs/tui/src/app/app_server_event_targets.rs b/codex-rs/tui/src/app/app_server_event_targets.rs index 800560d520..fa2bab0111 100644 --- a/codex-rs/tui/src/app/app_server_event_targets.rs +++ b/codex-rs/tui/src/app/app_server_event_targets.rs @@ -49,6 +49,7 @@ pub(super) fn server_notification_thread_target( Some(notification.thread_id.as_str()) } ServerNotification::ThreadArchived(notification) => Some(notification.thread_id.as_str()), + ServerNotification::ThreadDeleted(notification) => Some(notification.thread_id.as_str()), ServerNotification::ThreadUnarchived(notification) => Some(notification.thread_id.as_str()), ServerNotification::ThreadClosed(notification) => Some(notification.thread_id.as_str()), ServerNotification::ThreadNameUpdated(notification) => { diff --git a/codex-rs/tui/src/app/app_server_events.rs b/codex-rs/tui/src/app/app_server_events.rs index 269c4f6a43..2c9d620050 100644 --- a/codex-rs/tui/src/app/app_server_events.rs +++ b/codex-rs/tui/src/app/app_server_events.rs @@ -103,16 +103,25 @@ impl App { return; } ServerNotification::ExternalAgentConfigImportCompleted(_) => { - let cwd = self.chat_widget.config_ref().cwd.to_path_buf(); + let should_report_completion = + app_server_client.consume_external_agent_config_import_completion(); if let Err(err) = self.refresh_in_memory_config_from_disk().await { tracing::warn!( error = %err, "failed to refresh config after external agent config import" ); } + let cwd = self.chat_widget.config_ref().cwd.to_path_buf(); self.chat_widget.refresh_plugin_mentions(); self.chat_widget.submit_op(AppCommand::reload_user_config()); self.fetch_plugins_list(app_server_client, cwd); + if should_report_completion { + self.chat_widget.add_info_message( + crate::external_agent_config_migration_flow::EXTERNAL_AGENT_CONFIG_MIGRATION_FINISHED_MESSAGE + .to_string(), + /*hint*/ None, + ); + } return; } ServerNotification::AppListUpdated(notification) => { diff --git a/codex-rs/tui/src/app/background_requests.rs b/codex-rs/tui/src/app/background_requests.rs index 456676f796..a69d987447 100644 --- a/codex-rs/tui/src/app/background_requests.rs +++ b/codex-rs/tui/src/app/background_requests.rs @@ -241,7 +241,7 @@ impl App { &mut self, app_server: &AppServerSession, cwd: PathBuf, - marketplace_path: AbsolutePathBuf, + location: PluginLocation, plugin_name: String, plugin_display_name: String, ) { @@ -249,14 +249,14 @@ impl App { let app_event_tx = self.app_event_tx.clone(); tokio::spawn(async move { let cwd_for_event = cwd.clone(); - let marketplace_path_for_event = marketplace_path.clone(); + let location_for_event = location.clone(); let plugin_name_for_event = plugin_name.clone(); - let result = fetch_plugin_install(request_handle, marketplace_path, plugin_name) + let result = fetch_plugin_install(request_handle, location, plugin_name) .await .map_err(|err| format!("Failed to install plugin: {err}")); app_event_tx.send(AppEvent::PluginInstallLoaded { cwd: cwd_for_event, - marketplace_path: marketplace_path_for_event, + location: location_for_event, plugin_name: plugin_name_for_event, plugin_display_name, result, @@ -835,16 +835,17 @@ pub(super) async fn fetch_marketplace_upgrade( } pub(super) async fn fetch_plugin_install( request_handle: AppServerRequestHandle, - marketplace_path: AbsolutePathBuf, + location: PluginLocation, plugin_name: String, ) -> Result { let request_id = RequestId::String(format!("plugin-install-{}", Uuid::new_v4())); + let (marketplace_path, remote_marketplace_name) = location.into_request_params(); request_handle .request_typed(ClientRequest::PluginInstall { request_id, params: PluginInstallParams { - marketplace_path: Some(marketplace_path), - remote_marketplace_name: None, + marketplace_path, + remote_marketplace_name, plugin_name, }, }) @@ -1062,6 +1063,26 @@ mod tests { ); } + #[test] + fn plugin_location_request_params_select_exactly_one_location() { + let local_path = test_absolute_path("/marketplaces/local"); + + assert_eq!( + PluginLocation::Local { + marketplace_path: local_path.clone() + } + .into_request_params(), + (Some(local_path), None) + ); + assert_eq!( + PluginLocation::Remote { + marketplace_name: "workspace-directory".to_string() + } + .into_request_params(), + (None, Some("workspace-directory".to_string())) + ); + } + #[test] fn mcp_inventory_maps_prefix_tool_names_by_server() { let statuses = vec![ diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index c03bf18857..813f54f11a 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -6,6 +6,7 @@ use super::resize_reflow::trailing_run_start; use super::*; use crate::config_update::format_config_error; +use crate::external_agent_config_migration_flow::ExternalAgentConfigMigrationFlowOutcome; #[cfg(target_os = "windows")] use codex_config::types::WindowsSandboxModeToml; @@ -110,6 +111,31 @@ impl App { // Leaving alt-screen may blank the inline viewport; force a redraw either way. tui.frame_requester().schedule_frame(); } + AppEvent::OpenExternalAgentConfigMigration => { + match crate::external_agent_config_migration_flow::handle_external_agent_config_migration_prompt( + tui, + app_server, + &self.config, + ) + .await + { + Ok(ExternalAgentConfigMigrationFlowOutcome::Started(message)) => { + self.chat_widget.add_info_message(message, /*hint*/ None); + } + Ok(ExternalAgentConfigMigrationFlowOutcome::NoItems) => { + self.chat_widget.add_info_message( + crate::external_agent_config_migration_flow::EXTERNAL_AGENT_CONFIG_MIGRATION_NO_ITEMS_MESSAGE + .to_string(), + /*hint*/ None, + ); + } + Ok(ExternalAgentConfigMigrationFlowOutcome::Cancelled) => {} + Err(error_message) => { + self.chat_widget.add_error_message(error_message); + } + } + tui.frame_requester().schedule_frame(); + } AppEvent::ResumeSessionByIdOrName(id_or_name) => { match crate::lookup_session_target_with_app_server(app_server, &id_or_name).await? { Some(target_session) => { @@ -127,6 +153,9 @@ impl App { AppEvent::ArchiveCurrentThread => { return Ok(self.archive_current_thread(app_server).await); } + AppEvent::DeleteCurrentThread => { + return Ok(self.delete_current_thread(app_server).await); + } AppEvent::ForkCurrentSession => { self.session_telemetry.counter( "codex.thread.fork", @@ -580,14 +609,14 @@ impl App { } AppEvent::FetchPluginInstall { cwd, - marketplace_path, + location, plugin_name, plugin_display_name, } => { self.fetch_plugin_install( app_server, cwd, - marketplace_path, + location, plugin_name, plugin_display_name, ); @@ -608,7 +637,7 @@ impl App { } AppEvent::PluginInstallLoaded { cwd, - marketplace_path, + location, plugin_name, plugin_display_name, result, @@ -619,7 +648,7 @@ impl App { } let should_refresh_plugin_detail = self.chat_widget.on_plugin_install_loaded( cwd.clone(), - marketplace_path.clone(), + location.clone(), plugin_name.clone(), plugin_display_name, result, @@ -628,12 +657,14 @@ impl App { { self.fetch_plugins_list(app_server, cwd.clone()); if should_refresh_plugin_detail { + let (marketplace_path, remote_marketplace_name) = + location.into_request_params(); self.fetch_plugin_detail( app_server, cwd, PluginReadParams { - marketplace_path: Some(marketplace_path), - remote_marketplace_name: None, + marketplace_path, + remote_marketplace_name, plugin_name, }, ); @@ -2252,4 +2283,31 @@ impl App { } } } + + pub(super) async fn delete_current_thread( + &mut self, + app_server: &mut AppServerSession, + ) -> AppRunControl { + let Some(thread_id) = self.active_thread_id.or(self.chat_widget.thread_id()) else { + self.chat_widget + .add_error_message("A thread must start before it can be deleted.".to_string()); + return AppRunControl::Continue; + }; + if self.side_threads.contains_key(&thread_id) { + self.chat_widget.add_error_message( + "'/delete' is unavailable in side conversations. Press Ctrl+C to return to the main thread first." + .to_string(), + ); + return AppRunControl::Continue; + } + + match app_server.thread_delete(thread_id).await { + Ok(()) => AppRunControl::Exit(ExitReason::UserRequested), + Err(err) => { + self.chat_widget + .add_error_message(format!("Failed to delete current thread: {err}")); + AppRunControl::Continue + } + } + } } diff --git a/codex-rs/tui/src/app/test_support.rs b/codex-rs/tui/src/app/test_support.rs index 4e5aa8b10c..e0eafaa23b 100644 --- a/codex-rs/tui/src/app/test_support.rs +++ b/codex-rs/tui/src/app/test_support.rs @@ -6,12 +6,14 @@ use super::*; use crate::chatwidget::tests::make_chatwidget_manual_with_sender; +use codex_models_manager::test_support::construct_model_info_offline_for_tests; +use codex_models_manager::test_support::get_model_offline_for_tests; pub(super) async fn make_test_app() -> App { let (chat_widget, app_event_tx, _rx, _op_rx) = make_chatwidget_manual_with_sender().await; let config = chat_widget.config_ref().clone(); let file_search = FileSearchManager::new(config.cwd.to_path_buf(), app_event_tx.clone()); - let model = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model = get_model_offline_for_tests(config.model.as_deref()); let session_telemetry = test_session_telemetry(&config, model.as_str()); App { @@ -68,7 +70,8 @@ pub(super) async fn make_test_app() -> App { } fn test_session_telemetry(config: &Config, model: &str) -> SessionTelemetry { - let model_info = crate::legacy_core::test_support::construct_model_info_offline(model, config); + let model_info = + construct_model_info_offline_for_tests(model, &config.to_models_manager_config()); SessionTelemetry::new( ThreadId::new(), model, diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index fa2e5a90e4..0e0329addc 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -75,6 +75,8 @@ use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput; use codex_app_server_protocol::UserInput as AppServerUserInput; use codex_app_server_protocol::WarningNotification; +use codex_models_manager::test_support::construct_model_info_offline_for_tests; +use codex_models_manager::test_support::get_model_offline_for_tests; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; use codex_protocol::config_types::CollaborationMode; @@ -307,7 +309,7 @@ async fn enqueue_primary_thread_session_replays_turns_before_initial_prompt_subm let thread_id = ThreadId::new(); let initial_prompt = "follow-up after replay".to_string(); let config = app.config.clone(); - let model = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model = get_model_offline_for_tests(config.model.as_deref()); app.chat_widget = ChatWidget::new_with_app_event(ChatWidgetInit { config, frame_requester: crate::tui::FrameRequester::test_dummy(), @@ -4100,7 +4102,7 @@ async fn make_test_app() -> App { let (chat_widget, app_event_tx, _rx, _op_rx) = make_chatwidget_manual_with_sender().await; let config = chat_widget.config_ref().clone(); let file_search = FileSearchManager::new(config.cwd.to_path_buf(), app_event_tx.clone()); - let model = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model = get_model_offline_for_tests(config.model.as_deref()); let session_telemetry = test_session_telemetry(&config, model.as_str()); App { @@ -4164,7 +4166,7 @@ async fn make_test_app_with_channels() -> ( let (chat_widget, app_event_tx, rx, op_rx) = make_chatwidget_manual_with_sender().await; let config = chat_widget.config_ref().clone(); let file_search = FileSearchManager::new(config.cwd.to_path_buf(), app_event_tx.clone()); - let model = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model = get_model_offline_for_tests(config.model.as_deref()); let session_telemetry = test_session_telemetry(&config, model.as_str()); ( @@ -4683,7 +4685,8 @@ fn lines_to_single_string(lines: &[Line<'_>]) -> String { } fn test_session_telemetry(config: &Config, model: &str) -> SessionTelemetry { - let model_info = crate::legacy_core::test_support::construct_model_info_offline(model, config); + let model_info = + construct_model_info_offline_for_tests(model, &config.to_models_manager_config()); SessionTelemetry::new( ThreadId::new(), model, diff --git a/codex-rs/tui/src/app/tests/model_catalog.rs b/codex-rs/tui/src/app/tests/model_catalog.rs index 961c65159a..ac4a50af8f 100644 --- a/codex-rs/tui/src/app/tests/model_catalog.rs +++ b/codex-rs/tui/src/app/tests/model_catalog.rs @@ -6,7 +6,7 @@ use pretty_assertions::assert_eq; use tokio::sync::mpsc::unbounded_channel; fn all_model_presets() -> Vec { - crate::legacy_core::test_support::all_model_presets().clone() + crate::test_support::TEST_MODEL_PRESETS.clone() } fn model_availability_nux_config(shown_count: &[(&str, u32)]) -> ModelAvailabilityNuxConfig { diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index c2f9ee3f31..ee6364c0de 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -109,6 +109,21 @@ pub(crate) struct ConnectorsSnapshot { pub(crate) connectors: Vec, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum PluginLocation { + Local { marketplace_path: AbsolutePathBuf }, + Remote { marketplace_name: String }, +} + +impl PluginLocation { + pub(crate) fn into_request_params(self) -> (Option, Option) { + match self { + PluginLocation::Local { marketplace_path } => (Some(marketplace_path), None), + PluginLocation::Remote { marketplace_name } => (None, Some(marketplace_name)), + } + } +} + /// Distinguishes why a rate-limit refresh was requested so the completion /// handler can route the result correctly. /// @@ -206,12 +221,18 @@ pub(crate) enum AppEvent { /// Open the resume picker inside the running TUI session. OpenResumePicker, + /// Open the external agent migration picker inside the running TUI session. + OpenExternalAgentConfigMigration, + /// Resume a thread by UUID or thread name inside the running TUI session. ResumeSessionByIdOrName(String), /// Archive the current active main thread and exit after it succeeds. ArchiveCurrentThread, + /// Permanently delete the current active main thread and exit after it succeeds. + DeleteCurrentThread, + /// Fork the current session into a new thread. ForkCurrentSession, @@ -493,7 +514,7 @@ pub(crate) enum AppEvent { /// Install a specific plugin from a marketplace. FetchPluginInstall { cwd: PathBuf, - marketplace_path: AbsolutePathBuf, + location: PluginLocation, plugin_name: String, plugin_display_name: String, }, @@ -501,7 +522,7 @@ pub(crate) enum AppEvent { /// Result of installing a plugin. PluginInstallLoaded { cwd: PathBuf, - marketplace_path: AbsolutePathBuf, + location: PluginLocation, plugin_name: String, plugin_display_name: String, result: Result, diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index 6360bda354..b887f04db3 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -53,6 +53,8 @@ use codex_app_server_protocol::ThreadBackgroundTerminalsCleanParams; use codex_app_server_protocol::ThreadBackgroundTerminalsCleanResponse; use codex_app_server_protocol::ThreadCompactStartParams; use codex_app_server_protocol::ThreadCompactStartResponse; +use codex_app_server_protocol::ThreadDeleteParams; +use codex_app_server_protocol::ThreadDeleteResponse; use codex_app_server_protocol::ThreadForkParams; use codex_app_server_protocol::ThreadForkResponse; use codex_app_server_protocol::ThreadGoalClearParams; @@ -129,12 +131,16 @@ use color_eyre::eyre::WrapErr; use std::collections::HashMap; use std::collections::HashSet; use std::path::PathBuf; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; use uuid::Uuid; const JSONRPC_INVALID_REQUEST: i64 = -32600; const JSONRPC_METHOD_NOT_FOUND: i64 = -32601; +pub(crate) const EXTERNAL_AGENT_CONFIG_IMPORT_IN_PROGRESS_MESSAGE: &str = + "A previous agent import is still running. Wait for it to finish before importing again."; const THREAD_SETTINGS_UPDATE_METHOD: &str = "thread/settings/update"; const DISABLED_FORK_UNSUPPORTED_MESSAGE: &str = "connected app server did not confirm that multi-agent tools were disabled for the fork"; @@ -181,6 +187,7 @@ pub(crate) struct AppServerSession { available_models: Vec, rejected_fork_thread_ids: HashSet, disabled_forks_unsupported: bool, + external_agent_config_import_completion_pending: AtomicBool, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -232,6 +239,7 @@ impl AppServerSession { available_models: Vec::new(), rejected_fork_thread_ids: HashSet::new(), disabled_forks_unsupported: false, + external_agent_config_import_completion_pending: AtomicBool::new(false), } } @@ -248,6 +256,10 @@ impl AppServerSession { matches!(self.thread_params_mode, ThreadParamsMode::Remote) } + pub(crate) fn uses_embedded_app_server(&self) -> bool { + matches!(&self.client, AppServerClient::InProcess(_)) + } + pub(crate) fn server_version(&self) -> Option<&str> { let AppServerClient::Remote(client) = &self.client else { return None; @@ -370,21 +382,48 @@ impl AppServerSession { self.client .request_typed(ClientRequest::ExternalAgentConfigDetect { request_id, params }) .await - .wrap_err("externalAgentConfig/detect failed during TUI startup") + .wrap_err("externalAgentConfig/detect failed during agent import") } pub(crate) async fn external_agent_config_import( &mut self, migration_items: Vec, - ) -> Result { + ) -> Result<()> { + // Mark the import active before sending the request so a fast completion notification + // cannot arrive before the TUI records it. + if self + .external_agent_config_import_completion_pending + .swap(true, Ordering::Relaxed) + { + color_eyre::eyre::bail!(EXTERNAL_AGENT_CONFIG_IMPORT_IN_PROGRESS_MESSAGE); + } let request_id = self.next_request_id(); - self.client + let response: Result = self + .client .request_typed(ClientRequest::ExternalAgentConfigImport { request_id, params: ExternalAgentConfigImportParams { migration_items }, }) .await - .wrap_err("externalAgentConfig/import failed during TUI startup") + .wrap_err("externalAgentConfig/import failed during agent import"); + match response { + Ok(_) => Ok(()), + Err(err) => { + self.external_agent_config_import_completion_pending + .store(false, Ordering::Relaxed); + Err(err) + } + } + } + + pub(crate) fn external_agent_config_import_in_progress(&self) -> bool { + self.external_agent_config_import_completion_pending + .load(Ordering::Relaxed) + } + + pub(crate) fn consume_external_agent_config_import_completion(&self) -> bool { + self.external_agent_config_import_completion_pending + .swap(false, Ordering::Relaxed) } pub(crate) async fn next_event(&mut self) -> Option { @@ -626,6 +665,21 @@ impl AppServerSession { Ok(()) } + pub(crate) async fn thread_delete(&mut self, thread_id: ThreadId) -> Result<()> { + let request_id = self.next_request_id(); + let _: ThreadDeleteResponse = self + .client + .request_typed(ClientRequest::ThreadDelete { + request_id, + params: ThreadDeleteParams { + thread_id: thread_id.to_string(), + }, + }) + .await + .wrap_err("failed to delete session")?; + Ok(()) + } + pub(crate) async fn thread_unarchive(&mut self, thread_id: ThreadId) -> Result { let request_id = self.next_request_id(); let response: ThreadUnarchiveResponse = self @@ -1236,6 +1290,7 @@ pub(crate) fn status_account_display_from_auth_mode( email: None, plan: plan_type.map(plan_type_display_name), }), + Some(AuthMode::BedrockApiKey) => None, None => None, } } diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 6fd95e8dd5..832867d758 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -59,8 +59,6 @@ use crate::bottom_pane::TerminalTitleSetupView; use crate::diff_model::FileChange; use crate::git_action_directives::parse_assistant_markdown; use crate::legacy_core::config::Config; -use crate::legacy_core::config::Constrained; -use crate::legacy_core::config::ConstraintResult; use crate::legacy_core::config::PermissionProfileSnapshot; #[cfg(any(target_os = "windows", test))] use crate::legacy_core::windows_sandbox::WindowsSandboxLevelExt; @@ -126,6 +124,8 @@ use codex_app_server_protocol::TurnPlanStepStatus; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput; use codex_config::ConfigLayerStackOrdering; +use codex_config::Constrained; +use codex_config::ConstraintResult; use codex_config::types::ApprovalsReviewer; use codex_config::types::Notifications; use codex_config::types::WindowsSandboxModeToml; diff --git a/codex-rs/tui/src/chatwidget/plugins.rs b/codex-rs/tui/src/chatwidget/plugins.rs index 746f42fce9..9bc62fe261 100644 --- a/codex-rs/tui/src/chatwidget/plugins.rs +++ b/codex-rs/tui/src/chatwidget/plugins.rs @@ -5,6 +5,7 @@ use std::time::Instant; use super::ChatWidget; use crate::app_event::AppEvent; +use crate::app_event::PluginLocation; use crate::bottom_pane::ColumnWidthMode; use crate::bottom_pane::SelectionAction; use crate::bottom_pane::SelectionItem; @@ -25,17 +26,18 @@ use crate::tui::FrameRequester; use codex_app_server_protocol::MarketplaceAddResponse; use codex_app_server_protocol::MarketplaceRemoveResponse; use codex_app_server_protocol::MarketplaceUpgradeResponse; +use codex_app_server_protocol::PluginAvailability; use codex_app_server_protocol::PluginDetail; use codex_app_server_protocol::PluginInstallPolicy; use codex_app_server_protocol::PluginInstallResponse; use codex_app_server_protocol::PluginListResponse; use codex_app_server_protocol::PluginMarketplaceEntry; use codex_app_server_protocol::PluginReadResponse; +use codex_app_server_protocol::PluginSource; use codex_app_server_protocol::PluginSummary; use codex_app_server_protocol::PluginUninstallResponse; use codex_core_plugins::OPENAI_CURATED_MARKETPLACE_NAME; use codex_features::Feature; -use codex_utils_absolute_path::AbsolutePathBuf; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; use crossterm::event::KeyEventKind; @@ -453,7 +455,7 @@ impl ChatWidget { pub(crate) fn on_plugin_install_loaded( &mut self, cwd: PathBuf, - _marketplace_path: AbsolutePathBuf, + _location: PluginLocation, _plugin_name: String, plugin_display_name: String, result: Result, @@ -1625,19 +1627,22 @@ impl ChatWidget { ) -> SelectionViewParams { let marketplace_label = plugin.marketplace_name.clone(); let display_name = plugin_display_name(&plugin.summary); - let detail_status_label = if plugin.summary.installed { - if plugin.summary.enabled { - "Installed" + let detail_status_label = + if plugin.summary.availability == PluginAvailability::DisabledByAdmin { + "Disabled by admin" + } else if plugin.summary.installed { + if plugin.summary.enabled { + "Installed" + } else { + "Disabled" + } } else { - "Disabled" - } - } else { - match plugin.summary.install_policy { - PluginInstallPolicy::NotAvailable => "Not installable", - PluginInstallPolicy::Available => "Can be installed", - PluginInstallPolicy::InstalledByDefault => "Available by default", - } - }; + match plugin.summary.install_policy { + PluginInstallPolicy::NotAvailable => "Not installable", + PluginInstallPolicy::Available => "Can be installed", + PluginInstallPolicy::InstalledByDefault => "Available by default", + } + }; let mut header = ColumnRenderable::new(); header.push(Line::from("Plugins".bold())); header.push(Line::from( @@ -1676,23 +1681,40 @@ impl ChatWidget { }]; if plugin.summary.installed { - let uninstall_cwd = self.config.cwd.to_path_buf(); - let plugin_id = plugin.summary.id.clone(); - let plugin_display_name = display_name; + if let Some(plugin_id) = plugin_uninstall_id(&plugin.summary) { + let uninstall_cwd = self.config.cwd.to_path_buf(); + let plugin_display_name = display_name; + items.push(SelectionItem { + name: "Uninstall plugin".to_string(), + description: Some("Remove this plugin now.".to_string()), + selected_description: Some("Remove this plugin now.".to_string()), + actions: vec![Box::new(move |tx| { + tx.send(AppEvent::OpenPluginUninstallLoading { + plugin_display_name: plugin_display_name.clone(), + }); + tx.send(AppEvent::FetchPluginUninstall { + cwd: uninstall_cwd.clone(), + plugin_id: plugin_id.clone(), + plugin_display_name: plugin_display_name.clone(), + }); + })], + ..Default::default() + }); + } else { + items.push(SelectionItem { + name: "Uninstall plugin".to_string(), + description: Some( + "This remote plugin did not provide an uninstall identity.".to_string(), + ), + is_disabled: true, + ..Default::default() + }); + } + } else if plugin.summary.availability == PluginAvailability::DisabledByAdmin { items.push(SelectionItem { - name: "Uninstall plugin".to_string(), - description: Some("Remove this plugin now.".to_string()), - selected_description: Some("Remove this plugin now.".to_string()), - actions: vec![Box::new(move |tx| { - tx.send(AppEvent::OpenPluginUninstallLoading { - plugin_display_name: plugin_display_name.clone(), - }); - tx.send(AppEvent::FetchPluginUninstall { - cwd: uninstall_cwd.clone(), - plugin_id: plugin_id.clone(), - plugin_display_name: plugin_display_name.clone(), - }); - })], + name: "Install plugin".to_string(), + description: Some("This plugin is disabled by your workspace admin.".to_string()), + is_disabled: true, ..Default::default() }); } else if plugin.summary.install_policy == PluginInstallPolicy::NotAvailable { @@ -1704,9 +1726,9 @@ impl ChatWidget { is_disabled: true, ..Default::default() }); - } else if let Some(marketplace_path) = plugin.marketplace_path.clone() { + } else if let Some(location) = plugin_detail_location(plugin) { let install_cwd = self.config.cwd.to_path_buf(); - let plugin_name = plugin.summary.name.clone(); + let plugin_name = plugin_request_name(&plugin.summary); let plugin_display_name = display_name; items.push(SelectionItem { name: "Install plugin".to_string(), @@ -1718,7 +1740,7 @@ impl ChatWidget { }); tx.send(AppEvent::FetchPluginInstall { cwd: install_cwd.clone(), - marketplace_path: marketplace_path.clone(), + location: location.clone(), plugin_name: plugin_name.clone(), plugin_display_name: plugin_display_name.clone(), }); @@ -1728,7 +1750,7 @@ impl ChatWidget { } else { items.push(SelectionItem { name: "Install plugin".to_string(), - description: Some("Installing remote plugins is not supported yet.".to_string()), + description: Some("This plugin did not provide an install location.".to_string()), is_disabled: true, ..Default::default() }); @@ -1792,9 +1814,12 @@ impl ChatWidget { } else { plugin_brief_description_without_marketplace(plugin, status_label_width) }; - let can_view_details = marketplace.path.is_some(); + let plugin_detail_request = plugin_detail_request_for_entry(marketplace, plugin); + let can_view_details = plugin_detail_request.is_some(); + let disabled_by_admin = plugin.availability == PluginAvailability::DisabledByAdmin; + let can_toggle_plugin = plugin.installed && !disabled_by_admin; let selected_status_label = format!("{status_label: = if let Some(marketplace_path) = marketplace_path { - vec![Box::new(move |tx| { - tx.send(AppEvent::OpenPluginDetailLoading { - plugin_display_name: plugin_display_name.clone(), - }); - tx.send(AppEvent::FetchPluginDetail { - cwd: cwd.clone(), - params: codex_app_server_protocol::PluginReadParams { - marketplace_path: Some(marketplace_path.clone()), - remote_marketplace_name: None, - plugin_name: plugin_name.clone(), - }, - }); - })] - } else { - Vec::new() - }; + let actions: Vec = + if let Some((location, plugin_name)) = plugin_detail_request { + vec![Box::new(move |tx| { + tx.send(AppEvent::OpenPluginDetailLoading { + plugin_display_name: plugin_display_name.clone(), + }); + let (marketplace_path, remote_marketplace_name) = + location.clone().into_request_params(); + tx.send(AppEvent::FetchPluginDetail { + cwd: cwd.clone(), + params: codex_app_server_protocol::PluginReadParams { + marketplace_path, + remote_marketplace_name, + plugin_name: plugin_name.clone(), + }, + }); + })] + } else { + Vec::new() + }; let is_disabled = !can_view_details && !plugin.installed; - let disabled_reason = - is_disabled.then(|| "remote plugin details are not available yet".to_string()); + let disabled_reason = is_disabled.then(|| "plugin details are unavailable".to_string()); items.push(SelectionItem { name: display_name, toggle, - toggle_placeholder: (!plugin.installed).then_some("[-] "), + toggle_placeholder: (!can_toggle_plugin).then_some("[-] "), description: Some(description), selected_description: Some(selected_description), search_value: Some(search_value), @@ -2082,6 +2113,9 @@ fn plugin_brief_description_without_marketplace( } fn plugin_status_label(plugin: &PluginSummary) -> &'static str { + if plugin.availability == PluginAvailability::DisabledByAdmin { + return "Disabled by admin"; + } if plugin.installed { if plugin.enabled { "Installed" @@ -2097,6 +2131,59 @@ fn plugin_status_label(plugin: &PluginSummary) -> &'static str { } } +fn plugin_location_for_marketplace( + marketplace: &PluginMarketplaceEntry, + plugin: &PluginSummary, +) -> Option { + if let Some(marketplace_path) = marketplace.path.clone() { + return Some(PluginLocation::Local { marketplace_path }); + } + plugin_remote_identity(plugin).map(|_| PluginLocation::Remote { + marketplace_name: marketplace.name.clone(), + }) +} + +fn plugin_detail_location(plugin: &PluginDetail) -> Option { + if let Some(marketplace_path) = plugin.marketplace_path.clone() { + return Some(PluginLocation::Local { marketplace_path }); + } + plugin_remote_identity(&plugin.summary).map(|_| PluginLocation::Remote { + marketplace_name: plugin.marketplace_name.clone(), + }) +} + +fn plugin_detail_request_for_entry( + marketplace: &PluginMarketplaceEntry, + plugin: &PluginSummary, +) -> Option<(PluginLocation, String)> { + plugin_location_for_marketplace(marketplace, plugin) + .map(|location| (location, plugin_request_name(plugin))) +} + +fn plugin_request_name(plugin: &PluginSummary) -> String { + if matches!(&plugin.source, PluginSource::Remote) + && let Some(remote_plugin_id) = plugin_remote_identity(plugin) + { + return remote_plugin_id; + } + plugin.name.clone() +} + +fn plugin_remote_identity(plugin: &PluginSummary) -> Option { + plugin + .share_context + .as_ref() + .map(|context| context.remote_plugin_id.clone()) + .or_else(|| plugin.remote_plugin_id.clone()) +} + +fn plugin_uninstall_id(plugin: &PluginSummary) -> Option { + if matches!(&plugin.source, PluginSource::Remote) { + return plugin_remote_identity(plugin); + } + Some(plugin.id.clone()) +} + fn plugin_description(plugin: &PluginSummary) -> Option { plugin .interface diff --git a/codex-rs/tui/src/chatwidget/protocol.rs b/codex-rs/tui/src/chatwidget/protocol.rs index 9d29208075..51e4cedf85 100644 --- a/codex-rs/tui/src/chatwidget/protocol.rs +++ b/codex-rs/tui/src/chatwidget/protocol.rs @@ -6,6 +6,15 @@ impl ChatWidget { notification: ServerNotification, replay_kind: Option, ) { + // Reject misrouted child updates before shared notification handling mutates parent state. + if let ServerNotification::McpServerStatusUpdated(notification) = ¬ification + && let (Some(notification_thread_id), Some(thread_id)) = + (notification.thread_id.as_deref(), self.thread_id()) + && notification_thread_id != thread_id.to_string() + { + return; + } + let from_replay = replay_kind.is_some(); let is_resume_initial_replay = matches!(replay_kind, Some(ReplayKind::ResumeInitialMessages)); @@ -215,6 +224,7 @@ impl ChatWidget { | ServerNotification::ThreadStarted(_) | ServerNotification::ThreadStatusChanged(_) | ServerNotification::ThreadArchived(_) + | ServerNotification::ThreadDeleted(_) | ServerNotification::ThreadUnarchived(_) | ServerNotification::RawResponseItemCompleted(_) | ServerNotification::CommandExecOutputDelta(_) diff --git a/codex-rs/tui/src/chatwidget/slash_dispatch.rs b/codex-rs/tui/src/chatwidget/slash_dispatch.rs index e7776a3007..18820ce46a 100644 --- a/codex-rs/tui/src/chatwidget/slash_dispatch.rs +++ b/codex-rs/tui/src/chatwidget/slash_dispatch.rs @@ -193,6 +193,34 @@ impl ChatWidget { }); self.request_redraw(); } + SlashCommand::Delete => { + self.bottom_pane.show_selection_view(SelectionViewParams { + title: Some("Delete this session?".to_string()), + subtitle: Some( + "Cannot be undone. Subagent threads will also be deleted.".to_string(), + ), + footer_hint: Some(standard_popup_hint_line()), + items: vec![ + SelectionItem { + name: "No, keep this session".to_string(), + description: Some("Return to the current session".to_string()), + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Yes, delete and exit".to_string(), + description: Some("Permanently delete this session now".to_string()), + actions: vec![Box::new(|tx| { + tx.send(AppEvent::DeleteCurrentThread); + })], + dismiss_on_select: true, + ..Default::default() + }, + ], + ..Default::default() + }); + self.request_redraw(); + } SlashCommand::Clear => { self.app_event_tx.send(AppEvent::ClearUi); } @@ -398,6 +426,10 @@ impl ChatWidget { SlashCommand::Skills => { self.open_skills_menu(); } + SlashCommand::Import => { + self.app_event_tx + .send(AppEvent::OpenExternalAgentConfigMigration); + } SlashCommand::Hooks => { self.add_hooks_output(); } @@ -982,6 +1014,7 @@ impl ChatWidget { SlashCommand::Feedback | SlashCommand::New | SlashCommand::Archive + | SlashCommand::Delete | SlashCommand::Clear | SlashCommand::Resume | SlashCommand::Fork @@ -1010,6 +1043,7 @@ impl ChatWidget { | SlashCommand::Logout | SlashCommand::Mention | SlashCommand::Skills + | SlashCommand::Import | SlashCommand::Hooks | SlashCommand::Title | SlashCommand::Statusline diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__slash_delete_confirmation_popup.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__slash_delete_confirmation_popup.snap new file mode 100644 index 0000000000..de5c0610a0 --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__slash_delete_confirmation_popup.snap @@ -0,0 +1,11 @@ +--- +source: tui/src/chatwidget/tests/slash_commands.rs +expression: popup +--- + Delete this session? + Cannot be undone. Subagent threads will also be deleted. + +› 1. No, keep this session Return to the current session + 2. Yes, delete and exit Permanently delete this session now + + Press enter to confirm or esc to go back diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index ba89379e73..c20768e10e 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -21,8 +21,6 @@ pub(super) use crate::diff_model::FileChange; pub(super) use crate::history_cell::UserHistoryCell; pub(super) use crate::legacy_core::config::Config; pub(super) use crate::legacy_core::config::ConfigBuilder; -pub(super) use crate::legacy_core::config::Constrained; -pub(super) use crate::legacy_core::config::ConstraintError; pub(super) use crate::model_catalog::ModelCatalog; pub(super) use crate::test_backend::VT100Backend; pub(super) use crate::test_support::PathBufExt; @@ -120,6 +118,8 @@ pub(super) use codex_app_server_protocol::UserInput; pub(super) use codex_app_server_protocol::UserInput as AppServerUserInput; pub(super) use codex_app_server_protocol::WarningNotification; pub(super) use codex_config::ConfigLayerStack; +pub(super) use codex_config::Constrained; +pub(super) use codex_config::ConstraintError; pub(super) use codex_config::RequirementSource; pub(super) use codex_config::types::ApprovalsReviewer; pub(super) use codex_config::types::Notifications; @@ -129,6 +129,8 @@ pub(super) use codex_core_skills::model::SkillMetadata; pub(super) use codex_features::FEATURES; pub(super) use codex_features::Feature; pub(super) use codex_git_utils::CommitLogEntry; +pub(super) use codex_models_manager::test_support::construct_model_info_offline_for_tests; +pub(super) use codex_models_manager::test_support::get_model_offline_for_tests; pub(super) use codex_otel::RuntimeMetricsSummary; pub(super) use codex_otel::SessionTelemetry; pub(super) use codex_protocol::ThreadId; diff --git a/codex-rs/tui/src/chatwidget/tests/helpers.rs b/codex-rs/tui/src/chatwidget/tests/helpers.rs index b058b47cca..92a903b7de 100644 --- a/codex-rs/tui/src/chatwidget/tests/helpers.rs +++ b/codex-rs/tui/src/chatwidget/tests/helpers.rs @@ -119,7 +119,8 @@ pub(super) fn snapshot(percent: f64) -> RateLimitSnapshot { } pub(super) fn test_session_telemetry(config: &Config, model: &str) -> SessionTelemetry { - let model_info = crate::legacy_core::test_support::construct_model_info_offline(model, config); + let model_info = + construct_model_info_offline_for_tests(model, &config.to_models_manager_config()); SessionTelemetry::new( ThreadId::new(), model, @@ -136,7 +137,7 @@ pub(super) fn test_session_telemetry(config: &Config, model: &str) -> SessionTel pub(super) fn test_model_catalog(_config: &Config) -> Arc { Arc::new(ModelCatalog::new( - crate::legacy_core::test_support::all_model_presets().clone(), + crate::test_support::TEST_MODEL_PRESETS.clone(), )) } @@ -152,9 +153,9 @@ pub(super) async fn make_chatwidget_manual( let app_event_tx = AppEventSender::new(tx_raw); let (op_tx, op_rx) = unbounded_channel::(); let mut cfg = test_config().await; - let resolved_model = model_override.map(str::to_owned).unwrap_or_else(|| { - crate::legacy_core::test_support::get_model_offline(cfg.model.as_deref()) - }); + let resolved_model = model_override + .map(str::to_owned) + .unwrap_or_else(|| get_model_offline_for_tests(cfg.model.as_deref())); if let Some(model) = model_override { cfg.model = Some(model.to_string()); } @@ -1337,6 +1338,34 @@ pub(super) fn plugins_test_summary( } } +pub(super) fn plugins_test_remote_summary( + remote_plugin_id: &str, + name: &str, + display_name: Option<&str>, + description: Option<&str>, + installed: bool, +) -> PluginSummary { + PluginSummary { + id: remote_plugin_id.to_string(), + remote_plugin_id: Some(remote_plugin_id.to_string()), + local_version: None, + name: name.to_string(), + share_context: None, + source: PluginSource::Remote, + installed, + enabled: true, + install_policy: PluginInstallPolicy::Available, + auth_policy: PluginAuthPolicy::OnInstall, + availability: PluginAvailability::Available, + interface: Some(plugins_test_interface( + display_name, + description, + /*long_description*/ None, + )), + keywords: Vec::new(), + } +} + pub(super) fn plugins_test_curated_marketplace( plugins: Vec, ) -> PluginMarketplaceEntry { @@ -1386,7 +1415,7 @@ pub(super) fn plugins_test_detail( description: Option<&str>, skills: &[&str], hooks: &[(codex_app_server_protocol::HookEventName, usize)], - apps: &[(&str, bool)], + apps: &[&str], mcp_servers: &[&str], ) -> PluginDetail { PluginDetail { @@ -1421,12 +1450,11 @@ pub(super) fn plugins_test_detail( .collect(), apps: apps .iter() - .map(|(name, needs_auth)| AppSummary { + .map(|name| AppSummary { id: format!("{name}-id"), name: (*name).to_string(), description: Some(format!("{name} app")), install_url: Some(format!("https://example.test/{name}")), - needs_auth: *needs_auth, }) .collect(), app_templates: Vec::new(), diff --git a/codex-rs/tui/src/chatwidget/tests/mcp_startup.rs b/codex-rs/tui/src/chatwidget/tests/mcp_startup.rs index 8b0d5ca472..60efbe1353 100644 --- a/codex-rs/tui/src/chatwidget/tests/mcp_startup.rs +++ b/codex-rs/tui/src/chatwidget/tests/mcp_startup.rs @@ -25,6 +25,47 @@ fn notify_mcp_status_error(chat: &mut ChatWidget, name: &str, error: &str) { ); } +#[tokio::test] +async fn mcp_startup_ignores_status_for_other_thread() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.show_welcome_banner = false; + chat.set_mcp_startup_expected_servers(["sentry".to_string()]); + let parent_thread_id = ThreadId::new(); + let child_thread_id = ThreadId::new(); + chat.thread_id = Some(parent_thread_id); + chat.on_stream_error( + "Connection interrupted, retrying".to_string(), + /*additional_details*/ None, + ); + let status_before = chat.status_state.current_status.clone(); + let retry_status_header_before = chat.status_state.retry_status_header.clone(); + + for status in [ + McpServerStartupState::Starting, + McpServerStartupState::Failed, + ] { + chat.handle_server_notification( + ServerNotification::McpServerStatusUpdated(McpServerStatusUpdatedNotification { + thread_id: Some(child_thread_id.to_string()), + name: "sentry".to_string(), + status, + error: matches!(status, McpServerStartupState::Failed) + .then(|| "sentry is not logged in".to_string()), + }), + /*replay_kind*/ None, + ); + } + + assert!(drain_insert_history(&mut rx).is_empty()); + assert!(!chat.bottom_pane.is_task_running()); + assert!(chat.mcp_startup_status.is_none()); + assert_eq!(chat.status_state.current_status, status_before); + assert_eq!( + chat.status_state.retry_status_header, + retry_status_header_before + ); +} + #[tokio::test] async fn mcp_startup_dedupes_same_round_duplicate_failure_warning() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/tests/plan_mode.rs b/codex-rs/tui/src/chatwidget/tests/plan_mode.rs index 74cd527038..984c71d1ab 100644 --- a/codex-rs/tui/src/chatwidget/tests/plan_mode.rs +++ b/codex-rs/tui/src/chatwidget/tests/plan_mode.rs @@ -1447,7 +1447,7 @@ async fn collaboration_modes_defaults_to_code_on_startup() { assert_eq!(chat.active_collaboration_mode_kind(), ModeKind::Default); assert_eq!( chat.current_model(), - crate::legacy_core::test_support::get_model_offline(chat.config.model.as_deref()) + get_model_offline_for_tests(chat.config.model.as_deref()) ); } @@ -1482,7 +1482,7 @@ async fn make_startup_chat_with_cli_overrides( .build() .await .expect("config"); - let resolved_model = crate::legacy_core::test_support::get_model_offline(cfg.model.as_deref()); + let resolved_model = get_model_offline_for_tests(cfg.model.as_deref()); let session_telemetry = test_session_telemetry(&cfg, resolved_model.as_str()); let init = ChatWidgetInit { config: cfg.clone(), diff --git a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs index 89c332a4ef..736be9fa2a 100644 --- a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs +++ b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs @@ -6,6 +6,7 @@ use codex_app_server_protocol::HookErrorInfo; use codex_app_server_protocol::HooksListEntry; use codex_app_server_protocol::HooksListResponse; use codex_app_server_protocol::MarketplaceRemoveResponse; +use codex_app_server_protocol::PluginAvailability; use codex_features::Stage; use pretty_assertions::assert_eq; @@ -68,7 +69,7 @@ async fn experimental_mode_plan_is_ignored_on_startup() { .build() .await .expect("config"); - let resolved_model = crate::legacy_core::test_support::get_model_offline(cfg.model.as_deref()); + let resolved_model = get_model_offline_for_tests(cfg.model.as_deref()); let session_telemetry = test_session_telemetry(&cfg, resolved_model.as_str()); let init = ChatWidgetInit { config: cfg.clone(), @@ -653,7 +654,7 @@ async fn plugin_detail_popup_snapshot_shows_install_actions_and_capability_summa (codex_app_server_protocol::HookEventName::PreToolUse, 1), (codex_app_server_protocol::HookEventName::Stop, 2), ], - &[("Figma", true), ("Slack", false)], + &["Figma", "Slack"], &["figma-mcp", "docs-mcp"], ), }), @@ -697,7 +698,7 @@ async fn plugin_detail_popup_hides_disclosure_for_installed_plugins() { (codex_app_server_protocol::HookEventName::PreToolUse, 1), (codex_app_server_protocol::HookEventName::Stop, 2), ], - &[("Figma", true), ("Slack", false)], + &["Figma", "Slack"], &["figma-mcp", "docs-mcp"], ), }), @@ -714,6 +715,350 @@ async fn plugin_detail_popup_hides_disclosure_for_installed_plugins() { ); } +#[tokio::test] +async fn plugins_popup_remote_row_opens_remote_detail() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); + + let popup = render_loaded_plugins_popup( + &mut chat, + plugins_test_response(vec![PluginMarketplaceEntry { + name: "workspace-directory".to_string(), + path: None, + interface: Some(MarketplaceInterface { + display_name: Some("Workspace".to_string()), + }), + plugins: vec![plugins_test_remote_summary( + "plugins~Plugin_calendar", + "calendar", + Some("Calendar"), + Some("Workspace schedules."), + /*installed*/ false, + )], + }]), + ); + let remote_row = popup + .lines() + .find(|line| line.contains("Calendar")) + .expect("expected remote plugin row"); + assert!( + remote_row.contains("Available") + && remote_row.contains("Press Enter to install or view plugin details."), + "expected remote plugin row to be viewable, got:\n{remote_row}" + ); + + while rx.try_recv().is_ok() {} + chat.handle_key_event(KeyEvent::from(KeyCode::Enter)); + + match rx.try_recv() { + Ok(AppEvent::OpenPluginDetailLoading { + plugin_display_name, + }) => { + assert_eq!(plugin_display_name, "Calendar"); + } + other => panic!("expected OpenPluginDetailLoading event, got {other:?}"), + } + match rx.try_recv() { + Ok(AppEvent::FetchPluginDetail { cwd: _, params }) => { + assert_eq!(params.marketplace_path, None); + assert_eq!( + params.remote_marketplace_name, + Some("workspace-directory".to_string()) + ); + assert_eq!(params.plugin_name, "plugins~Plugin_calendar"); + } + other => panic!("expected FetchPluginDetail event, got {other:?}"), + } +} + +#[tokio::test] +async fn plugin_detail_remote_install_uses_remote_location() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); + + let summary = plugins_test_remote_summary( + "plugins~Plugin_linear", + "linear", + Some("Linear"), + Some("Issue tracking."), + /*installed*/ false, + ); + let cwd = chat.config.cwd.clone(); + chat.on_plugins_loaded( + cwd.to_path_buf(), + Ok(plugins_test_response(vec![PluginMarketplaceEntry { + name: "workspace-shared-with-me-private".to_string(), + path: None, + interface: Some(MarketplaceInterface { + display_name: Some("Shared with me".to_string()), + }), + plugins: vec![summary.clone()], + }])), + ); + chat.add_plugins_output(); + chat.on_plugin_detail_loaded( + cwd.to_path_buf(), + Ok(PluginReadResponse { + plugin: PluginDetail { + marketplace_name: "workspace-shared-with-me-private".to_string(), + marketplace_path: None, + summary, + description: Some("Install shared Linear plugin.".to_string()), + skills: Vec::new(), + hooks: Vec::new(), + apps: Vec::new(), + app_templates: Vec::new(), + mcp_servers: Vec::new(), + }, + }), + ); + let popup = render_bottom_popup(&chat, /*width*/ 100); + assert!( + popup.contains("Install plugin") && popup.contains("Install this plugin now."), + "expected remote detail to offer install, got:\n{popup}" + ); + + while rx.try_recv().is_ok() {} + chat.handle_key_event(KeyEvent::from(KeyCode::Down)); + chat.handle_key_event(KeyEvent::from(KeyCode::Enter)); + + match rx.try_recv() { + Ok(AppEvent::OpenPluginInstallLoading { + plugin_display_name, + }) => { + assert_eq!(plugin_display_name, "Linear"); + } + other => panic!("expected OpenPluginInstallLoading event, got {other:?}"), + } + match rx.try_recv() { + Ok(AppEvent::FetchPluginInstall { + cwd: _, + location: crate::app_event::PluginLocation::Remote { marketplace_name }, + plugin_name, + plugin_display_name, + }) => { + assert_eq!(marketplace_name, "workspace-shared-with-me-private"); + assert_eq!(plugin_name, "plugins~Plugin_linear"); + assert_eq!(plugin_display_name, "Linear"); + } + other => panic!("expected remote FetchPluginInstall event, got {other:?}"), + } +} + +#[tokio::test] +async fn plugin_detail_remote_uninstall_uses_remote_plugin_id() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); + + let summary = plugins_test_remote_summary( + "plugins~Plugin_linear", + "linear", + Some("Linear"), + Some("Issue tracking."), + /*installed*/ true, + ); + let cwd = chat.config.cwd.clone(); + chat.on_plugins_loaded( + cwd.to_path_buf(), + Ok(plugins_test_response(vec![PluginMarketplaceEntry { + name: "workspace-shared-with-me-private".to_string(), + path: None, + interface: Some(MarketplaceInterface { + display_name: Some("Shared with me".to_string()), + }), + plugins: vec![summary.clone()], + }])), + ); + chat.add_plugins_output(); + chat.on_plugin_detail_loaded( + cwd.to_path_buf(), + Ok(PluginReadResponse { + plugin: PluginDetail { + marketplace_name: "workspace-shared-with-me-private".to_string(), + marketplace_path: None, + summary, + description: Some("Installed shared Linear plugin.".to_string()), + skills: Vec::new(), + hooks: Vec::new(), + apps: Vec::new(), + app_templates: Vec::new(), + mcp_servers: Vec::new(), + }, + }), + ); + + while rx.try_recv().is_ok() {} + chat.handle_key_event(KeyEvent::from(KeyCode::Down)); + chat.handle_key_event(KeyEvent::from(KeyCode::Enter)); + + match rx.try_recv() { + Ok(AppEvent::OpenPluginUninstallLoading { + plugin_display_name, + }) => { + assert_eq!(plugin_display_name, "Linear"); + } + other => panic!("expected OpenPluginUninstallLoading event, got {other:?}"), + } + match rx.try_recv() { + Ok(AppEvent::FetchPluginUninstall { + plugin_id, + plugin_display_name, + .. + }) => { + assert_eq!(plugin_id, "plugins~Plugin_linear"); + assert_eq!(plugin_display_name, "Linear"); + } + other => panic!("expected remote FetchPluginUninstall event, got {other:?}"), + } +} + +#[tokio::test] +async fn plugin_detail_remote_without_remote_id_disables_uninstall_action() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); + + let summary = PluginSummary { + source: PluginSource::Remote, + ..plugins_test_summary( + "linear@workspace-shared-with-me-private", + "linear", + Some("Linear"), + Some("Issue tracking."), + /*installed*/ true, + /*enabled*/ true, + PluginInstallPolicy::Available, + ) + }; + let cwd = chat.config.cwd.clone(); + chat.on_plugins_loaded( + cwd.to_path_buf(), + Ok(plugins_test_response(vec![PluginMarketplaceEntry { + name: "workspace-shared-with-me-private".to_string(), + path: None, + interface: Some(MarketplaceInterface { + display_name: Some("Shared with me".to_string()), + }), + plugins: vec![summary.clone()], + }])), + ); + chat.add_plugins_output(); + chat.on_plugin_detail_loaded( + cwd.to_path_buf(), + Ok(PluginReadResponse { + plugin: PluginDetail { + marketplace_name: "workspace-shared-with-me-private".to_string(), + marketplace_path: None, + summary, + description: Some("Installed shared Linear plugin.".to_string()), + skills: Vec::new(), + hooks: Vec::new(), + apps: Vec::new(), + app_templates: Vec::new(), + mcp_servers: Vec::new(), + }, + }), + ); + + let popup = render_bottom_popup(&chat, /*width*/ 120); + assert!( + popup.contains("This remote plugin did not provide an uninstall identity.") + && !popup.contains("Remove this plugin now."), + "expected missing remote ID to disable uninstall, got:\n{popup}" + ); + + while rx.try_recv().is_ok() {} + assert!( + rx.try_recv().is_err(), + "expected no action after rendering disabled uninstall state" + ); +} + +#[tokio::test] +async fn plugin_detail_admin_disabled_plugin_blocks_install() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); + + let summary = PluginSummary { + availability: PluginAvailability::DisabledByAdmin, + ..plugins_test_summary( + "plugin-admin-blocked", + "admin-blocked", + Some("Admin Blocked"), + Some("Blocked by policy."), + /*installed*/ false, + /*enabled*/ true, + PluginInstallPolicy::Available, + ) + }; + let response = plugins_test_response(vec![plugins_test_curated_marketplace(vec![ + summary.clone(), + ])]); + let cwd = chat.config.cwd.clone(); + chat.on_plugins_loaded(cwd.to_path_buf(), Ok(response)); + chat.add_plugins_output(); + chat.on_plugin_detail_loaded( + cwd.to_path_buf(), + Ok(PluginReadResponse { + plugin: plugins_test_detail(summary, Some("Blocked by policy."), &[], &[], &[], &[]), + }), + ); + + let popup = render_bottom_popup(&chat, /*width*/ 100); + assert!( + popup.contains("Admin Blocked · Disabled by admin") + && popup.contains("This plugin is disabled by your workspace admin.") + && !popup.contains("Install this plugin now."), + "expected admin-disabled detail to block install, got:\n{popup}" + ); + + while rx.try_recv().is_ok() {} + assert!( + rx.try_recv().is_err(), + "expected no action after rendering disabled install state" + ); +} + +#[tokio::test] +async fn plugins_popup_admin_disabled_installed_plugin_has_no_toggle_hint() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); + + let summary = PluginSummary { + availability: PluginAvailability::DisabledByAdmin, + ..plugins_test_summary( + "plugin-admin-blocked", + "admin-blocked", + Some("Admin Blocked"), + Some("Blocked by policy."), + /*installed*/ true, + /*enabled*/ true, + PluginInstallPolicy::Available, + ) + }; + render_loaded_plugins_popup( + &mut chat, + plugins_test_response(vec![plugins_test_curated_marketplace(vec![summary])]), + ); + + let popup = render_bottom_popup(&chat, /*width*/ 100); + assert!( + popup.contains("Disabled by admin") + && popup.contains("Press Enter to view plugin details.") + && !popup.contains("Space to disable"), + "expected admin-disabled installed plugin to omit toggle hint, got:\n{popup}" + ); + + while rx.try_recv().is_ok() {} + let before = render_bottom_popup(&chat, /*width*/ 100); + chat.handle_key_event(KeyEvent::new(KeyCode::Char(' '), KeyModifiers::NONE)); + let after = render_bottom_popup(&chat, /*width*/ 100); + assert!( + rx.try_recv().is_err(), + "space should not toggle admin-disabled installed plugins" + ); + assert_eq!(after, before); +} + #[tokio::test] async fn plugin_detail_error_popup_skips_disabled_row_numbering() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs index 02faef89f4..e61ffd48fe 100644 --- a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs +++ b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs @@ -1861,6 +1861,18 @@ async fn slash_resume_opens_picker() { assert_matches!(rx.try_recv(), Ok(AppEvent::OpenResumePicker)); } +#[tokio::test] +async fn slash_import_opens_claude_code_import_picker() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + + chat.dispatch_command(SlashCommand::Import); + + assert_matches!( + rx.try_recv(), + Ok(AppEvent::OpenExternalAgentConfigMigration) + ); +} + #[tokio::test] async fn slash_archive_confirmation_requests_current_thread_archive() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; @@ -1879,6 +1891,24 @@ async fn slash_archive_confirmation_requests_current_thread_archive() { assert_matches!(rx.try_recv(), Ok(AppEvent::ArchiveCurrentThread)); } +#[tokio::test] +async fn slash_delete_confirmation_requests_current_thread_delete() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + + chat.dispatch_command(SlashCommand::Delete); + + assert!(chat.bottom_pane.has_active_view()); + assert_matches!(rx.try_recv(), Err(TryRecvError::Empty)); + + let popup = render_bottom_popup(&chat, /*width*/ 80); + assert_chatwidget_snapshot!("slash_delete_confirmation_popup", popup); + + chat.handle_key_event(KeyEvent::from(KeyCode::Down)); + chat.handle_key_event(KeyEvent::from(KeyCode::Enter)); + + assert_matches!(rx.try_recv(), Ok(AppEvent::DeleteCurrentThread)); +} + #[tokio::test] async fn slash_resume_with_arg_requests_named_session() { let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs index 3d08daa427..b7cd476834 100644 --- a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs +++ b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs @@ -411,7 +411,7 @@ async fn configured_pet_load_is_deferred_until_after_construction() { let mut cfg = test_config().await; cfg.tui_pet = Some(crate::pets::DEFAULT_PET_ID.to_string()); crate::pets::write_test_pack(&cfg.codex_home); - let resolved_model = crate::legacy_core::test_support::get_model_offline(cfg.model.as_deref()); + let resolved_model = get_model_offline_for_tests(cfg.model.as_deref()); let session_telemetry = test_session_telemetry(&cfg, resolved_model.as_str()); let init = ChatWidgetInit { config: cfg.clone(), diff --git a/codex-rs/tui/src/debug_config.rs b/codex-rs/tui/src/debug_config.rs index ebd2042a10..d1dde975f3 100644 --- a/codex-rs/tui/src/debug_config.rs +++ b/codex-rs/tui/src/debug_config.rs @@ -544,8 +544,6 @@ mod tests { use super::render_debug_config_lines; use super::sandbox_mode_is_allowed_by_permissions; use super::session_all_proxy_url; - use crate::legacy_core::config::Constrained; - use crate::legacy_core::config::ConstraintError; use crate::legacy_core::config::Permissions; use codex_app_server_protocol::AskForApproval; use codex_app_server_protocol::ConfigLayerSource; @@ -553,7 +551,9 @@ mod tests { use codex_config::ConfigLayerStack; use codex_config::ConfigRequirements; use codex_config::ConfigRequirementsToml; + use codex_config::Constrained; use codex_config::ConstrainedWithSource; + use codex_config::ConstraintError; use codex_config::FeatureRequirementsToml; use codex_config::FilesystemConstraints; use codex_config::HookEventsToml; diff --git a/codex-rs/tui/src/external_agent_config_migration.rs b/codex-rs/tui/src/external_agent_config_migration.rs index 1ce9d44751..d1c4582d68 100644 --- a/codex-rs/tui/src/external_agent_config_migration.rs +++ b/codex-rs/tui/src/external_agent_config_migration.rs @@ -1,10 +1,8 @@ use crate::diff_render::display_path_for; -use crate::key_hint; -use crate::line_truncation::truncate_line_with_ellipsis_if_overflow; -use crate::render::Insets; -use crate::render::RectExt as _; -use crate::selection_list::selection_option_row_with_dim; -use crate::style::accent_style; +use crate::external_agent_config_migration_model::ExternalAgentConfigMigrationGroupModel; +use crate::external_agent_config_migration_model::external_agent_config_migration_groups; +use crate::external_agent_config_migration_model::external_agent_config_migration_item_detail; +use crate::external_agent_config_migration_model::external_agent_config_migration_item_label; use crate::tui::FrameRequester; use crate::tui::Tui; use crate::tui::TuiEvent; @@ -14,25 +12,16 @@ use crossterm::event::KeyCode; use crossterm::event::KeyEvent; use crossterm::event::KeyEventKind; use crossterm::event::KeyModifiers; -use ratatui::buffer::Buffer; -use ratatui::layout::Constraint; -use ratatui::layout::Layout; -use ratatui::layout::Rect; use ratatui::prelude::Stylize as _; use ratatui::text::Line; -use ratatui::widgets::Clear; -use ratatui::widgets::Paragraph; -use ratatui::widgets::Widget; -use ratatui::widgets::WidgetRef; -use ratatui::widgets::Wrap; use tokio_stream::StreamExt; +mod render; + #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) enum ExternalAgentConfigMigrationOutcome { Proceed(Vec), Skip, - SkipForever, - Exit, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -44,34 +33,26 @@ enum FocusArea { #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum ActionMenuOption { Proceed, + Customize, Skip, - SkipForever, + Back, } impl ActionMenuOption { fn label(self) -> &'static str { match self { - Self::Proceed => "Proceed with selected", - Self::Skip => "Skip for now", - Self::SkipForever => "Don't ask again", + Self::Proceed => "Import selected", + Self::Customize => "Customize selection", + Self::Skip => "Cancel", + Self::Back => "Review selection", } } +} - fn previous(self) -> Option { - match self { - Self::Proceed => None, - Self::Skip => Some(Self::Proceed), - Self::SkipForever => Some(Self::Skip), - } - } - - fn next(self) -> Option { - match self { - Self::Proceed => Some(Self::Skip), - Self::Skip => Some(Self::SkipForever), - Self::SkipForever => None, - } - } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum MigrationView { + Summary, + Customize, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -136,6 +117,8 @@ pub(crate) async fn run_external_agent_config_migration_prompt( struct ExternalAgentConfigMigrationScreen { request_frame: FrameRequester, items: Vec, + groups: Vec, + view: MigrationView, selected_item_idx: Option, scroll_top: usize, focus: FocusArea, @@ -151,44 +134,67 @@ impl ExternalAgentConfigMigrationScreen { } fn first_available_action(&self) -> ActionMenuOption { - if self.proceed_enabled() { - ActionMenuOption::Proceed - } else { - ActionMenuOption::Skip + match self.available_actions().first() { + Some(action) => *action, + None => ActionMenuOption::Back, + } + } + + fn last_available_action(&self) -> ActionMenuOption { + match self.available_actions().last() { + Some(action) => *action, + None => ActionMenuOption::Back, } } fn previous_available_action(&self, action: ActionMenuOption) -> Option { - let mut candidate = action.previous(); - while let Some(option) = candidate { - if option != ActionMenuOption::Proceed || self.proceed_enabled() { - return Some(option); - } - candidate = option.previous(); - } - None + let actions = self.available_actions(); + actions + .iter() + .position(|candidate| *candidate == action) + .and_then(|idx| idx.checked_sub(1)) + .and_then(|idx| actions.get(idx)) + .copied() } fn next_available_action(&self, action: ActionMenuOption) -> Option { - let mut candidate = action.next(); - while let Some(option) = candidate { - if option != ActionMenuOption::Proceed || self.proceed_enabled() { - return Some(option); + let actions = self.available_actions(); + actions + .iter() + .position(|candidate| *candidate == action) + .and_then(|idx| actions.get(idx + 1)) + .copied() + } + + fn available_actions(&self) -> Vec { + match self.view { + MigrationView::Summary => { + let mut actions = Vec::new(); + if self.proceed_enabled() { + actions.push(ActionMenuOption::Proceed); + } + actions.extend([ActionMenuOption::Customize, ActionMenuOption::Skip]); + actions } - candidate = option.next(); + MigrationView::Customize => vec![ActionMenuOption::Back], } - None } fn normalize_highlighted_action(&mut self) { - if self.highlighted_action == ActionMenuOption::Proceed && !self.proceed_enabled() { + if !self.available_actions().contains(&self.highlighted_action) { self.highlighted_action = self.first_available_action(); } } fn display_description(item: &ExternalAgentConfigMigrationItem) -> String { + // App-server descriptions use migration vocabulary. Normalize that prefix so the TUI + // consistently uses the user-facing import vocabulary. + let description = item + .description + .strip_prefix("Migrate ") + .map_or_else(|| item.description.clone(), |rest| format!("Import {rest}")); let Some(cwd) = item.cwd.as_deref() else { - return item.description.clone(); + return description; }; fn reformat_description( @@ -207,33 +213,23 @@ impl ExternalAgentConfigMigrationScreen { )) } - if let Some(reformatted) = - reformat_description(&item.description, "Migrate ", " into ", cwd) - { + if let Some(reformatted) = reformat_description(&description, "Import ", " into ", cwd) { return reformatted; } if let Some(reformatted) = - reformat_description(&item.description, "Migrate skills from ", " to ", cwd) + reformat_description(&description, "Import skills from ", " to ", cwd) { return reformatted; } - if let Some(reformatted) = reformat_description(&item.description, "Migrate ", " to ", cwd) - { + if let Some(reformatted) = reformat_description(&description, "Import ", " to ", cwd) { return reformatted; } - if let Some(reformatted) = reformat_description(&item.description, "Import ", " to ", cwd) { - return reformatted; - } - - if let Some(source) = item - .description - .strip_prefix("Migrate enabled plugins from ") - { + if let Some(source) = description.strip_prefix("Import enabled plugins from ") { let description = format!( - "Migrate enabled plugins from {}", + "Import enabled plugins from {}", display_path_for(std::path::Path::new(source), cwd) ); if let Some(details) = &item.details { @@ -260,7 +256,7 @@ impl ExternalAgentConfigMigrationScreen { return description; } - item.description.clone() + description } fn new( @@ -269,6 +265,7 @@ impl ExternalAgentConfigMigrationScreen { selected_items: &[ExternalAgentConfigMigrationItem], error: Option, ) -> Self { + let groups = external_agent_config_migration_groups(items); let items = items .iter() .cloned() @@ -277,18 +274,22 @@ impl ExternalAgentConfigMigrationScreen { item, }) .collect::>(); - let selected_item_idx = (!items.is_empty()).then_some(0); - Self { + let selected_item_idx = (!groups.is_empty()).then_some(0); + let mut screen = Self { request_frame, items, + groups, + view: MigrationView::Summary, selected_item_idx, scroll_top: 0, - focus: FocusArea::Items, + focus: FocusArea::Actions, highlighted_action: ActionMenuOption::Proceed, done: false, outcome: ExternalAgentConfigMigrationOutcome::Skip, error, - } + }; + screen.normalize_highlighted_action(); + screen } fn plugin_detail_lines(plugin_groups: &[PluginsMigration]) -> Vec> { @@ -341,12 +342,6 @@ impl ExternalAgentConfigMigrationScreen { fn proceed(&mut self) { let selected = self.selected_items(); - if selected.is_empty() { - self.error = Some("Select at least one item or choose a skip option.".to_string()); - self.request_frame.schedule_frame(); - return; - } - self.finish_with(ExternalAgentConfigMigrationOutcome::Proceed(selected)); } @@ -354,14 +349,6 @@ impl ExternalAgentConfigMigrationScreen { self.finish_with(ExternalAgentConfigMigrationOutcome::Skip); } - fn skip_forever(&mut self) { - self.finish_with(ExternalAgentConfigMigrationOutcome::SkipForever); - } - - fn exit(&mut self) { - self.finish_with(ExternalAgentConfigMigrationOutcome::Exit); - } - fn selected_items(&self) -> Vec { self.items .iter() @@ -374,6 +361,22 @@ impl ExternalAgentConfigMigrationScreen { self.items.iter().filter(|item| item.enabled).count() } + fn group_selection_marker( + &self, + group: &ExternalAgentConfigMigrationGroupModel, + ) -> &'static str { + let enabled_count = group + .item_indices + .iter() + .filter(|idx| self.items[**idx].enabled) + .count(); + match enabled_count { + 0 => " ", + count if count == group.item_indices.len() => "x", + _ => "-", + } + } + fn set_all_enabled(&mut self, enabled: bool) { for item in &mut self.items { item.enabled = enabled; @@ -384,7 +387,7 @@ impl ExternalAgentConfigMigrationScreen { } fn toggle_selected_item(&mut self) { - if self.focus != FocusArea::Items { + if self.view != MigrationView::Customize || self.focus != FocusArea::Items { return; } let Some(selected_idx) = self.selected_item_idx else { @@ -393,26 +396,51 @@ impl ExternalAgentConfigMigrationScreen { let Some(item) = self.items.get_mut(selected_idx) else { return; }; - item.enabled = !item.enabled; self.error = None; self.normalize_highlighted_action(); self.request_frame.schedule_frame(); } + fn customize(&mut self) { + self.view = MigrationView::Customize; + self.selected_item_idx = (!self.items.is_empty()).then_some(0); + self.scroll_top = 0; + self.focus = FocusArea::Items; + self.highlighted_action = ActionMenuOption::Back; + self.request_frame.schedule_frame(); + } + + fn back_to_summary(&mut self) { + self.view = MigrationView::Summary; + self.selected_item_idx = (!self.groups.is_empty()).then_some(0); + self.scroll_top = 0; + self.focus = FocusArea::Actions; + self.highlighted_action = self.first_available_action(); + self.request_frame.schedule_frame(); + } + fn move_up(&mut self) { + if self.view == MigrationView::Summary { + self.focus = FocusArea::Actions; + self.highlighted_action = self + .previous_available_action(self.highlighted_action) + .unwrap_or_else(|| self.last_available_action()); + self.request_frame.schedule_frame(); + return; + } match self.focus { FocusArea::Items => match self.selected_item_idx { Some(0) => { self.focus = FocusArea::Actions; - self.highlighted_action = ActionMenuOption::SkipForever; + self.highlighted_action = self.last_available_action(); } Some(idx) => { self.selected_item_idx = Some(idx.saturating_sub(1)); } None => { self.focus = FocusArea::Actions; - self.highlighted_action = ActionMenuOption::SkipForever; + self.highlighted_action = self.last_available_action(); } }, FocusArea::Actions => { @@ -431,6 +459,14 @@ impl ExternalAgentConfigMigrationScreen { } fn move_down(&mut self) { + if self.view == MigrationView::Summary { + self.focus = FocusArea::Actions; + self.highlighted_action = self + .next_available_action(self.highlighted_action) + .unwrap_or_else(|| self.first_available_action()); + self.request_frame.schedule_frame(); + return; + } match self.focus { FocusArea::Items => match self.selected_item_idx { Some(idx) if idx + 1 < self.items.len() => { @@ -461,8 +497,9 @@ impl ExternalAgentConfigMigrationScreen { FocusArea::Items => self.toggle_selected_item(), FocusArea::Actions => match self.highlighted_action { ActionMenuOption::Proceed => self.proceed(), + ActionMenuOption::Customize => self.customize(), ActionMenuOption::Skip => self.skip(), - ActionMenuOption::SkipForever => self.skip_forever(), + ActionMenuOption::Back => self.back_to_summary(), }, } } @@ -471,39 +508,47 @@ impl ExternalAgentConfigMigrationScreen { if key_event.kind == KeyEventKind::Release { return; } - if is_ctrl_exit_combo(key_event) { - self.exit(); + self.skip(); return; } match key_event.code { KeyCode::Up | KeyCode::Char('k') => self.move_up(), KeyCode::Down | KeyCode::Char('j') => self.move_down(), - KeyCode::Char('1') => { - self.focus = FocusArea::Actions; - self.highlighted_action = ActionMenuOption::Proceed; - self.proceed(); + KeyCode::Char(number @ '1'..='9') => self.select_numbered_action(number), + KeyCode::Char('c') if self.view == MigrationView::Summary => self.customize(), + KeyCode::Char('b') if self.view == MigrationView::Customize => self.back_to_summary(), + KeyCode::Char(' ') if self.view == MigrationView::Customize => { + self.toggle_selected_item(); } - KeyCode::Char('2') => { - self.focus = FocusArea::Actions; - self.highlighted_action = ActionMenuOption::Skip; - self.skip(); + KeyCode::Char('a') if self.view == MigrationView::Customize => { + self.set_all_enabled(/*enabled*/ true); } - KeyCode::Char('3') => { - self.focus = FocusArea::Actions; - self.highlighted_action = ActionMenuOption::SkipForever; - self.skip_forever(); + KeyCode::Char('n') if self.view == MigrationView::Customize => { + self.set_all_enabled(/*enabled*/ false); } - KeyCode::Char(' ') => self.toggle_selected_item(), - KeyCode::Char('a') => self.set_all_enabled(/*enabled*/ true), - KeyCode::Char('n') => self.set_all_enabled(/*enabled*/ false), KeyCode::Enter => self.confirm_selection(), - KeyCode::Esc => self.skip(), + KeyCode::Esc => match self.view { + MigrationView::Summary => self.skip(), + MigrationView::Customize => self.back_to_summary(), + }, _ => {} } } + fn select_numbered_action(&mut self, number: char) { + let Some(index) = number.to_digit(10).and_then(|number| number.checked_sub(1)) else { + return; + }; + let Some(action) = self.available_actions().get(index as usize).copied() else { + return; + }; + self.focus = FocusArea::Actions; + self.highlighted_action = action; + self.confirm_selection(); + } + fn ensure_selected_item_visible(&mut self) { let Some(selected_idx) = self.selected_item_idx else { self.scroll_top = 0; @@ -534,12 +579,47 @@ impl ExternalAgentConfigMigrationScreen { fn section_title(cwd: Option<&std::path::Path>) -> Line<'static> { match cwd { - Some(cwd) => Line::from(vec!["Project: ".bold(), cwd.display().to_string().dim()]), + Some(cwd) => Line::from(vec![ + "Current project: ".bold(), + cwd.display().to_string().dim(), + ]), None => Line::from("Home".bold()), } } fn build_render_lines(&self) -> Vec { + match self.view { + MigrationView::Summary => self.build_summary_render_lines(), + MigrationView::Customize => self.build_customize_render_lines(), + } + } + + fn build_summary_render_lines(&self) -> Vec { + self.groups + .iter() + .enumerate() + .flat_map(|(idx, group)| { + [ + RenderLineEntry { + item_idx: Some(idx), + kind: RenderLineKind::Item, + line: Line::from(format!( + " [{}] {}", + self.group_selection_marker(group), + group.label + )), + }, + RenderLineEntry { + item_idx: None, + kind: RenderLineKind::ItemDetail, + line: Line::from(format!(" {}", group.description)), + }, + ] + }) + .collect() + } + + fn build_customize_render_lines(&self) -> Vec { let mut lines = Vec::new(); let mut current_scope: Option> = None; for (idx, item) in self.items.iter().enumerate() { @@ -562,12 +642,28 @@ impl ExternalAgentConfigMigrationScreen { lines.push(RenderLineEntry { item_idx: Some(idx), kind: RenderLineKind::Item, - line: Line::from(format!( - " [{}] {}", - if item.enabled { "x" } else { " " }, - Self::display_description(&item.item) - )), + line: Line::from(vec![ + " ".into(), + format!( + "[{}] {}", + if item.enabled { "x" } else { " " }, + external_agent_config_migration_item_label(&item.item) + ) + .into(), + ]), }); + lines.push(RenderLineEntry { + item_idx: None, + kind: RenderLineKind::ItemDetail, + line: Line::from(format!(" {}", Self::display_description(&item.item))), + }); + if let Some(details) = external_agent_config_migration_item_detail(&item.item) { + lines.push(RenderLineEntry { + item_idx: None, + kind: RenderLineKind::ItemDetail, + line: Line::from(format!(" {details}")), + }); + } if let Some(details) = &item.item.details { for line in Self::plugin_detail_lines(&details.plugins) { lines.push(RenderLineEntry { @@ -580,194 +676,25 @@ impl ExternalAgentConfigMigrationScreen { } lines } - - fn render_items(&self, area: Rect, buf: &mut Buffer) { - if area.height == 0 || area.width == 0 { - return; - } - let rows = self.build_render_lines(); - let visible_rows = area.height as usize; - let mut start_idx = self.scroll_top.min(rows.len().saturating_sub(1)); - if let Some(selected_item_idx) = self.selected_item_idx { - let selected_render_idx = self.selected_render_line_index(selected_item_idx); - if selected_render_idx < start_idx { - start_idx = selected_render_idx; - } else if visible_rows > 0 { - let bottom = start_idx + visible_rows - 1; - if selected_render_idx > bottom { - start_idx = selected_render_idx + 1 - visible_rows; - } - } - } - - let mut y = area.y; - for entry in rows.iter().skip(start_idx).take(visible_rows) { - if y >= area.y + area.height { - break; - } - - let selected = - self.focus == FocusArea::Items && self.selected_item_idx == entry.item_idx; - let mut line = entry.line.clone(); - if selected { - line.spans.iter_mut().for_each(|span| { - span.style = span.style.patch(accent_style()); - }); - } else if entry.kind != RenderLineKind::Item && !line.spans.is_empty() { - line.spans.iter_mut().for_each(|span| { - span.style = span.style.dim(); - }); - } - let line = truncate_line_with_ellipsis_if_overflow(line, area.width as usize); - line.render( - Rect { - x: area.x, - y, - width: area.width, - height: 1, - }, - buf, - ); - y = y.saturating_add(1); - } - } -} - -impl WidgetRef for &ExternalAgentConfigMigrationScreen { - fn render_ref(&self, area: Rect, buf: &mut Buffer) { - Clear.render(area, buf); - - let inner_area = area.inset(Insets::vh(/*v*/ 1, /*h*/ 2)); - let error_height = u16::from(self.error.is_some()); - let fixed_height = 1u16 + 2u16 + error_height + 1u16 + 4u16 + 1u16; - let list_height = - self.render_line_count() - .max(1) - .min(inner_area.height.saturating_sub(fixed_height) as usize) as u16; - let [ - header_area, - intro_area, - error_area, - list_area, - list_gap_area, - actions_area, - footer_area, - _spacer_area, - ] = Layout::vertical([ - Constraint::Length(1), - Constraint::Length(2), - Constraint::Length(error_height), - Constraint::Length(list_height), - Constraint::Length(1), - Constraint::Length(4), - Constraint::Length(1), - Constraint::Fill(1), - ]) - .areas(inner_area); - - let heading = Line::from(vec!["> ".into(), "External agent config detected".bold()]); - heading.render(header_area, buf); - - Paragraph::new(vec![ - Line::from("We found settings from another agent that you can add to this project."), - Line::from("Select what to import"), - ]) - .wrap(Wrap { trim: false }) - .render(intro_area, buf); - - if let Some(error) = &self.error { - Paragraph::new(error.clone().red().to_string()) - .wrap(Wrap { trim: false }) - .render(error_area, buf); - } - - self.render_items(list_area, buf); - Clear.render(list_gap_area, buf); - - let [ - actions_intro_area, - proceed_area, - skip_area, - skip_forever_area, - ] = Layout::vertical([ - Constraint::Length(1), - Constraint::Length(1), - Constraint::Length(1), - Constraint::Length(1), - ]) - .areas(actions_area); - let actions_intro = format!( - "Selected {} of {} item(s).", - self.selected_count(), - self.items.len() - ); - Paragraph::new(actions_intro) - .wrap(Wrap { trim: false }) - .render(actions_intro_area, buf); - selection_option_row_with_dim( - /*index*/ 0, - ActionMenuOption::Proceed.label().to_string(), - self.focus == FocusArea::Actions - && self.highlighted_action == ActionMenuOption::Proceed, - /*dim*/ self.focus != FocusArea::Actions || !self.proceed_enabled(), - ) - .render(proceed_area, buf); - selection_option_row_with_dim( - /*index*/ 1, - ActionMenuOption::Skip.label().to_string(), - self.focus == FocusArea::Actions && self.highlighted_action == ActionMenuOption::Skip, - /*dim*/ self.focus != FocusArea::Actions, - ) - .render(skip_area, buf); - selection_option_row_with_dim( - /*index*/ 2, - ActionMenuOption::SkipForever.label().to_string(), - self.focus == FocusArea::Actions - && self.highlighted_action == ActionMenuOption::SkipForever, - /*dim*/ self.focus != FocusArea::Actions, - ) - .render(skip_forever_area, buf); - - Line::from(vec![ - "Use ".dim(), - key_hint::plain(KeyCode::Up).into(), - "/".dim(), - key_hint::plain(KeyCode::Down).into(), - " to move, ".dim(), - key_hint::plain(KeyCode::Char(' ')).into(), - " to toggle, ".dim(), - "1".cyan(), - "/".dim(), - "2".cyan(), - "/".dim(), - "3".cyan(), - " to choose, ".dim(), - "a".cyan(), - "/".dim(), - "n".cyan(), - " for all/none".dim(), - ]) - .render(footer_area, buf); - } } fn is_ctrl_exit_combo(key_event: KeyEvent) -> bool { - key_event.modifiers.contains(KeyModifiers::CONTROL) - && matches!(key_event.code, KeyCode::Char('c') | KeyCode::Char('d')) + matches!(key_event.code, KeyCode::Char('c' | 'd')) + && key_event.modifiers.contains(KeyModifiers::CONTROL) } #[cfg(test)] mod tests { - use super::ActionMenuOption; use super::ExternalAgentConfigMigrationOutcome; use super::ExternalAgentConfigMigrationScreen; - use super::FocusArea; + use super::MigrationView; use crate::custom_terminal::Terminal; use crate::test_backend::VT100Backend; use crate::tui::FrameRequester; use codex_app_server_protocol::ExternalAgentConfigMigrationItem; use codex_app_server_protocol::ExternalAgentConfigMigrationItemType; use codex_app_server_protocol::PluginsMigration; + use codex_app_server_protocol::SessionMigration; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; use crossterm::event::KeyModifiers; @@ -829,6 +756,19 @@ mod tests { cwd: None, details: None, }, + ExternalAgentConfigMigrationItem { + item_type: ExternalAgentConfigMigrationItemType::Sessions, + description: "Migrate recent chat sessions".to_string(), + cwd: None, + details: Some(codex_app_server_protocol::MigrationDetails { + sessions: vec![SessionMigration { + path: PathBuf::from("/Users/alex/.claude/projects/project/session.jsonl"), + cwd: project_root.clone(), + title: Some("Investigate migration UX".to_string()), + }], + ..Default::default() + }), + }, ExternalAgentConfigMigrationItem { item_type: ExternalAgentConfigMigrationItemType::Plugins, description: format!( @@ -864,7 +804,13 @@ mod tests { frame.render_widget_ref(screen, frame.area()); } terminal.flush().expect("flush"); - terminal.backend().to_string() + terminal + .backend() + .to_string() + .lines() + .map(str::trim_end) + .collect::>() + .join("\n") } #[test] @@ -877,13 +823,56 @@ mod tests { /*error*/ None, ); - let rendered = render_screen(&screen, /*width*/ 80, /*height*/ 21); + let rendered = render_screen(&screen, /*width*/ 80, /*height*/ 24); #[cfg(windows)] assert_snapshot!("external_agent_config_migration_prompt_windows", rendered); #[cfg(not(windows))] assert_snapshot!("external_agent_config_migration_prompt", rendered); } + #[test] + fn customize_snapshot() { + let items = sample_items(); + let mut screen = ExternalAgentConfigMigrationScreen::new( + FrameRequester::test_dummy(), + &items, + &items, + /*error*/ None, + ); + screen.customize(); + + let rendered = render_screen(&screen, /*width*/ 80, /*height*/ 30); + #[cfg(windows)] + assert_snapshot!( + "external_agent_config_migration_customize_windows", + rendered + ); + #[cfg(not(windows))] + assert_snapshot!("external_agent_config_migration_customize", rendered); + } + + #[test] + fn customize_action_snapshot() { + let items = sample_items(); + let mut screen = ExternalAgentConfigMigrationScreen::new( + FrameRequester::test_dummy(), + &items, + &items, + /*error*/ None, + ); + screen.customize(); + screen.move_up(); + + let rendered = render_screen(&screen, /*width*/ 80, /*height*/ 30); + #[cfg(windows)] + assert_snapshot!( + "external_agent_config_migration_customize_action_windows", + rendered + ); + #[cfg(not(windows))] + assert_snapshot!("external_agent_config_migration_customize_action", rendered); + } + #[test] fn proceed_returns_selected_items() { let items = sample_items(); @@ -894,9 +883,6 @@ mod tests { /*error*/ None, ); - screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); - screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); - screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); screen.handle_key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); assert!(screen.is_done()); @@ -916,16 +902,19 @@ mod tests { /*error*/ None, ); + screen.handle_key(KeyEvent::new(KeyCode::Char('c'), KeyModifiers::NONE)); screen.handle_key(KeyEvent::new(KeyCode::Char(' '), KeyModifiers::NONE)); - screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); - screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); - screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); - screen.handle_key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + screen.handle_key(KeyEvent::new(KeyCode::Char('b'), KeyModifiers::NONE)); + screen.handle_key(KeyEvent::new(KeyCode::Char('1'), KeyModifiers::NONE)); assert!(screen.is_done()); assert_eq!( screen.outcome(), - ExternalAgentConfigMigrationOutcome::Proceed(vec![items[1].clone(), items[2].clone(),]) + ExternalAgentConfigMigrationOutcome::Proceed(vec![ + items[1].clone(), + items[2].clone(), + items[3].clone(), + ]) ); } @@ -946,30 +935,7 @@ mod tests { } #[test] - fn skip_forever_returns_skip_forever_outcome() { - let items = sample_items(); - let mut screen = ExternalAgentConfigMigrationScreen::new( - FrameRequester::test_dummy(), - &items, - &items, - /*error*/ None, - ); - - screen.move_down(); - screen.move_down(); - screen.move_down(); - screen.move_down(); - screen.move_down(); - screen.confirm_selection(); - - assert_eq!( - screen.outcome(), - ExternalAgentConfigMigrationOutcome::SkipForever - ); - } - - #[test] - fn proceed_requires_at_least_one_selected_item() { + fn numeric_shortcuts_follow_visible_actions_when_proceed_is_disabled() { let items = sample_items(); let mut screen = ExternalAgentConfigMigrationScreen::new( FrameRequester::test_dummy(), @@ -978,35 +944,46 @@ mod tests { /*error*/ None, ); + screen.handle_key(KeyEvent::new(KeyCode::Char('c'), KeyModifiers::NONE)); screen.handle_key(KeyEvent::new(KeyCode::Char('n'), KeyModifiers::NONE)); + screen.handle_key(KeyEvent::new(KeyCode::Char('b'), KeyModifiers::NONE)); screen.handle_key(KeyEvent::new(KeyCode::Char('1'), KeyModifiers::NONE)); - assert!(!screen.is_done()); - assert_eq!(screen.highlighted_action, ActionMenuOption::Proceed); - let rendered = render_screen(&screen, /*width*/ 80, /*height*/ 20); - assert!( - rendered.contains("Select at least one item or choose a skip option."), - "expected inline validation error, got:\n{rendered}" - ); + assert_eq!(screen.view, MigrationView::Customize); } #[test] - fn proceed_action_is_skipped_when_no_items_are_selected() { + fn empty_selection_enter_opens_customize_instead_of_proceeding() { let items = sample_items(); let mut screen = ExternalAgentConfigMigrationScreen::new( FrameRequester::test_dummy(), &items, - &items, + &[], /*error*/ None, ); - screen.handle_key(KeyEvent::new(KeyCode::Char('n'), KeyModifiers::NONE)); - screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); - screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); - screen.handle_key(KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); + screen.handle_key(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); - assert_eq!(screen.focus, FocusArea::Actions); - assert_eq!(screen.highlighted_action, ActionMenuOption::Skip); + assert!(!screen.is_done()); + assert_eq!(screen.view, MigrationView::Customize); + } + + #[test] + fn control_exit_shortcuts_cancel_prompt() { + let items = sample_items(); + for key_code in [KeyCode::Char('c'), KeyCode::Char('d')] { + let mut screen = ExternalAgentConfigMigrationScreen::new( + FrameRequester::test_dummy(), + &items, + &items, + /*error*/ None, + ); + + screen.handle_key(KeyEvent::new(key_code, KeyModifiers::CONTROL)); + + assert!(screen.is_done()); + assert_eq!(screen.outcome(), ExternalAgentConfigMigrationOutcome::Skip); + } } #[test] @@ -1025,28 +1002,42 @@ mod tests { ExternalAgentConfigMigrationOutcome::Proceed(items.clone()) ); + let mut customize_screen = ExternalAgentConfigMigrationScreen::new( + FrameRequester::test_dummy(), + &items, + &items, + /*error*/ None, + ); + customize_screen.handle_key(KeyEvent::new(KeyCode::Char('2'), KeyModifiers::NONE)); + assert_eq!(customize_screen.view, MigrationView::Customize); + customize_screen.handle_key(KeyEvent::new(KeyCode::Char('1'), KeyModifiers::NONE)); + assert_eq!(customize_screen.view, MigrationView::Summary); + let mut skip_screen = ExternalAgentConfigMigrationScreen::new( FrameRequester::test_dummy(), &items, &items, /*error*/ None, ); - skip_screen.handle_key(KeyEvent::new(KeyCode::Char('2'), KeyModifiers::NONE)); + skip_screen.handle_key(KeyEvent::new(KeyCode::Char('3'), KeyModifiers::NONE)); assert_eq!( skip_screen.outcome(), ExternalAgentConfigMigrationOutcome::Skip ); + } - let mut skip_forever_screen = ExternalAgentConfigMigrationScreen::new( + #[test] + fn summary_does_not_toggle_selection() { + let items = sample_items(); + let mut screen = ExternalAgentConfigMigrationScreen::new( FrameRequester::test_dummy(), &items, &items, /*error*/ None, ); - skip_forever_screen.handle_key(KeyEvent::new(KeyCode::Char('3'), KeyModifiers::NONE)); - assert_eq!( - skip_forever_screen.outcome(), - ExternalAgentConfigMigrationOutcome::SkipForever - ); + + screen.handle_key(KeyEvent::new(KeyCode::Char(' '), KeyModifiers::NONE)); + + assert_eq!(screen.selected_items(), items); } } diff --git a/codex-rs/tui/src/external_agent_config_migration/render.rs b/codex-rs/tui/src/external_agent_config_migration/render.rs new file mode 100644 index 0000000000..4d518906c4 --- /dev/null +++ b/codex-rs/tui/src/external_agent_config_migration/render.rs @@ -0,0 +1,206 @@ +use super::*; +use crate::key_hint; +use crate::line_truncation::truncate_line_with_ellipsis_if_overflow; +use crate::render::Insets; +use crate::render::RectExt as _; +use crate::selection_list::selection_option_row_with_dim; +use ratatui::buffer::Buffer; +use ratatui::layout::Constraint; +use ratatui::layout::Layout; +use ratatui::layout::Rect; +use ratatui::widgets::Clear; +use ratatui::widgets::Paragraph; +use ratatui::widgets::Widget; +use ratatui::widgets::WidgetRef; +use ratatui::widgets::Wrap; + +impl ExternalAgentConfigMigrationScreen { + fn render_items(&self, area: Rect, buf: &mut Buffer) { + if area.height == 0 || area.width == 0 { + return; + } + let rows = self.build_render_lines(); + let visible_rows = area.height as usize; + let mut start_idx = self.scroll_top.min(rows.len().saturating_sub(1)); + if let Some(selected_item_idx) = self.selected_item_idx { + let selected_render_idx = self.selected_render_line_index(selected_item_idx); + if selected_render_idx < start_idx { + start_idx = selected_render_idx; + } else if visible_rows > 0 { + let bottom = start_idx + visible_rows - 1; + if selected_render_idx > bottom { + start_idx = selected_render_idx + 1 - visible_rows; + } + } + } + + let mut y = area.y; + for entry in rows.iter().skip(start_idx).take(visible_rows) { + if y >= area.y + area.height { + break; + } + + let selected = + self.focus == FocusArea::Items && self.selected_item_idx == entry.item_idx; + let mut line = entry.line.clone(); + if selected { + if let Some(cursor) = line.spans.first_mut() { + cursor.content = "› ".into(); + } + line.spans.iter_mut().for_each(|span| { + span.style = span.style.cyan().bold(); + }); + } else if entry.kind != RenderLineKind::Item && !line.spans.is_empty() { + line.spans.iter_mut().for_each(|span| { + span.style = span.style.dim(); + }); + } + let line = truncate_line_with_ellipsis_if_overflow(line, area.width as usize); + line.render( + Rect { + x: area.x, + y, + width: area.width, + height: 1, + }, + buf, + ); + y = y.saturating_add(1); + } + } +} + +impl WidgetRef for &ExternalAgentConfigMigrationScreen { + fn render_ref(&self, area: Rect, buf: &mut Buffer) { + Clear.render(area, buf); + + let inner_area = area.inset(Insets::vh(/*v*/ 1, /*h*/ 2)); + let error_height = u16::from(self.error.is_some()); + let intro_lines = match self.view { + MigrationView::Summary => vec![ + Line::from("Bring over your setup, current project, and recent chats."), + Line::from("Codex may add files to your current project folder."), + Line::from("Your existing agent setup will not be changed."), + Line::from("Cloud-hosted chat data cannot be imported."), + ], + MigrationView::Customize => vec![ + Line::from("Choose the items to import."), + Line::from("Codex may add files to your current project folder."), + Line::from("Your existing agent setup will not be changed."), + ], + }; + let intro_height = intro_lines.len() as u16; + let actions = self.available_actions(); + let actions_height = actions.len() as u16 + 1; + let fixed_height = 1u16 + intro_height + error_height + 1u16 + actions_height + 1u16; + let list_height = + self.render_line_count() + .max(1) + .min(inner_area.height.saturating_sub(fixed_height) as usize) as u16; + let [ + header_area, + intro_area, + error_area, + list_area, + list_gap_area, + actions_area, + footer_area, + _spacer_area, + ] = Layout::vertical([ + Constraint::Length(1), + Constraint::Length(intro_height), + Constraint::Length(error_height), + Constraint::Length(list_height), + Constraint::Length(1), + Constraint::Length(actions_height), + Constraint::Length(1), + Constraint::Fill(1), + ]) + .areas(inner_area); + + let title = match self.view { + MigrationView::Summary => "Import from another coding agent", + MigrationView::Customize => "Choose what to import", + }; + let heading = Line::from(vec!["> ".into(), title.bold()]); + heading.render(header_area, buf); + + Paragraph::new(intro_lines) + .wrap(Wrap { trim: false }) + .render(intro_area, buf); + + if let Some(error) = &self.error { + Paragraph::new(error.clone().red().to_string()) + .wrap(Wrap { trim: false }) + .render(error_area, buf); + } + + self.render_items(list_area, buf); + Clear.render(list_gap_area, buf); + + let action_areas = Layout::vertical(std::iter::repeat_n( + Constraint::Length(1), + actions.len() + 1, + )) + .split(actions_area); + let item_label = if self.items.len() == 1 { + "item" + } else { + "items" + }; + let actions_intro = format!( + "Selected {} of {} {item_label}.", + self.selected_count(), + self.items.len() + ); + Paragraph::new(actions_intro) + .wrap(Wrap { trim: false }) + .render(action_areas[0], buf); + for (idx, action) in actions.iter().enumerate() { + selection_option_row_with_dim( + idx, + action.label().to_string(), + self.focus == FocusArea::Actions && self.highlighted_action == *action, + /*dim*/ self.focus != FocusArea::Actions, + ) + .render(action_areas[idx + 1], buf); + } + + let footer = match self.view { + MigrationView::Summary => Line::from(vec![ + "Use ".dim(), + key_hint::plain(KeyCode::Up).into(), + "/".dim(), + key_hint::plain(KeyCode::Down).into(), + " to move, ".dim(), + key_hint::plain(KeyCode::Enter).into(), + " to select, ".dim(), + "c".cyan(), + " to customize".dim(), + ]), + MigrationView::Customize if self.focus == FocusArea::Actions => Line::from(vec![ + "Press ".dim(), + key_hint::plain(KeyCode::Enter).into(), + " to continue, ".dim(), + key_hint::plain(KeyCode::Up).into(), + "/".dim(), + key_hint::plain(KeyCode::Down).into(), + " to move, ".dim(), + "b".cyan(), + " to go back".dim(), + ]), + MigrationView::Customize => Line::from(vec![ + "Use ".dim(), + key_hint::plain(KeyCode::Up).into(), + "/".dim(), + key_hint::plain(KeyCode::Down).into(), + " to move, ".dim(), + key_hint::plain(KeyCode::Char(' ')).into(), + " to toggle, ".dim(), + "b".cyan(), + " to go back".dim(), + ]), + }; + footer.render(footer_area, buf); + } +} diff --git a/codex-rs/tui/src/external_agent_config_migration_flow.rs b/codex-rs/tui/src/external_agent_config_migration_flow.rs new file mode 100644 index 0000000000..f0fabe3a01 --- /dev/null +++ b/codex-rs/tui/src/external_agent_config_migration_flow.rs @@ -0,0 +1,125 @@ +use crate::app_server_session::AppServerSession; +use crate::app_server_session::EXTERNAL_AGENT_CONFIG_IMPORT_IN_PROGRESS_MESSAGE; +use crate::external_agent_config_migration::ExternalAgentConfigMigrationOutcome; +use crate::external_agent_config_migration::run_external_agent_config_migration_prompt; +use crate::legacy_core::config::Config; +use crate::tui; +use codex_app_server_protocol::ExternalAgentConfigDetectParams; + +pub(crate) const EXTERNAL_AGENT_CONFIG_MIGRATION_FINISHED_MESSAGE: &str = + "Agent import finished. Run /import again to check for additional items."; +pub(crate) const EXTERNAL_AGENT_CONFIG_MIGRATION_NO_ITEMS_MESSAGE: &str = + "No supported agent setup was found to import."; +pub(crate) const EXTERNAL_AGENT_CONFIG_MIGRATION_REMOTE_UNAVAILABLE_MESSAGE: &str = + "Agent import is unavailable in remote sessions. Start Codex locally and run /import."; +pub(crate) const EXTERNAL_AGENT_CONFIG_MIGRATION_DAEMON_UNAVAILABLE_MESSAGE: &str = "Agent import is unavailable while Codex is connected to the local app-server daemon. Stop the daemon, restart Codex, and run /import."; + +pub(crate) enum ExternalAgentConfigMigrationFlowOutcome { + Started(String), + NoItems, + Cancelled, +} + +fn external_agent_config_migration_success_message(remaining_item_count: usize) -> String { + let message = "Agent import started. You can keep working while it finishes. Imported setup will apply to new chats."; + match remaining_items_handoff(remaining_item_count) { + Some(remaining_items_handoff) => format!("{message} {remaining_items_handoff}"), + None => message.to_string(), + } +} + +fn remaining_items_handoff(remaining_item_count: usize) -> Option { + match remaining_item_count { + 0 => None, + 1 => Some( + "1 additional item remains. After it finishes, run /import again to review it." + .to_string(), + ), + _ => Some(format!( + "{remaining_item_count} additional items remain. After it finishes, run /import again to review them." + )), + } +} + +pub(crate) async fn handle_external_agent_config_migration_prompt( + tui: &mut tui::Tui, + app_server: &mut AppServerSession, + config: &Config, +) -> Result { + if app_server.uses_remote_workspace() { + return Err(EXTERNAL_AGENT_CONFIG_MIGRATION_REMOTE_UNAVAILABLE_MESSAGE.to_string()); + } + if !app_server.uses_embedded_app_server() { + return Err(EXTERNAL_AGENT_CONFIG_MIGRATION_DAEMON_UNAVAILABLE_MESSAGE.to_string()); + } + if app_server.external_agent_config_import_in_progress() { + return Err(EXTERNAL_AGENT_CONFIG_IMPORT_IN_PROGRESS_MESSAGE.to_string()); + } + + let cwd = config.cwd.to_path_buf(); + let detected_items = match app_server + .external_agent_config_detect(ExternalAgentConfigDetectParams { + include_home: true, + cwds: Some(vec![cwd.clone()]), + }) + .await + { + Ok(response) => response.items, + Err(err) => { + tracing::warn!( + error = %err, + cwd = %cwd.display(), + "failed to detect external agent config migrations" + ); + return Err(format!("Could not check for agent setup: {err}")); + } + }; + + if detected_items.is_empty() { + return Ok(ExternalAgentConfigMigrationFlowOutcome::NoItems); + } + + let mut selected_items = detected_items.clone(); + let mut error: Option = None; + + loop { + match run_external_agent_config_migration_prompt( + tui, + &detected_items, + &selected_items, + error.as_deref(), + ) + .await + { + ExternalAgentConfigMigrationOutcome::Proceed(items) => { + selected_items = items.clone(); + match app_server.external_agent_config_import(items).await { + Ok(()) => { + let remaining_item_count = + detected_items.len().saturating_sub(selected_items.len()); + let success_message = + external_agent_config_migration_success_message(remaining_item_count); + return Ok(ExternalAgentConfigMigrationFlowOutcome::Started( + success_message, + )); + } + Err(err) => { + tracing::warn!( + error = %err, + cwd = %cwd.display(), + "failed to import external agent config migration items" + ); + error = Some(format!("Import failed: {err}")); + } + } + } + ExternalAgentConfigMigrationOutcome::Skip => { + return Ok(ExternalAgentConfigMigrationFlowOutcome::Cancelled); + } + } + } +} + +#[cfg(test)] +#[path = "external_agent_config_migration_flow_tests.rs"] +mod tests; diff --git a/codex-rs/tui/src/external_agent_config_migration_flow_tests.rs b/codex-rs/tui/src/external_agent_config_migration_flow_tests.rs new file mode 100644 index 0000000000..913b55607f --- /dev/null +++ b/codex-rs/tui/src/external_agent_config_migration_flow_tests.rs @@ -0,0 +1,21 @@ +use super::*; + +#[test] +fn external_agent_config_migration_messages_snapshot() { + let cases = [0, 1, 2]; + + let messages = cases + .map(external_agent_config_migration_success_message) + .into_iter() + .chain([ + EXTERNAL_AGENT_CONFIG_MIGRATION_FINISHED_MESSAGE.to_string(), + EXTERNAL_AGENT_CONFIG_MIGRATION_NO_ITEMS_MESSAGE.to_string(), + EXTERNAL_AGENT_CONFIG_MIGRATION_REMOTE_UNAVAILABLE_MESSAGE.to_string(), + EXTERNAL_AGENT_CONFIG_MIGRATION_DAEMON_UNAVAILABLE_MESSAGE.to_string(), + EXTERNAL_AGENT_CONFIG_IMPORT_IN_PROGRESS_MESSAGE.to_string(), + ]) + .collect::>() + .join("\n"); + + insta::assert_snapshot!("external_agent_config_migration_messages", messages); +} diff --git a/codex-rs/tui/src/external_agent_config_migration_model.rs b/codex-rs/tui/src/external_agent_config_migration_model.rs new file mode 100644 index 0000000000..d03e238da3 --- /dev/null +++ b/codex-rs/tui/src/external_agent_config_migration_model.rs @@ -0,0 +1,147 @@ +use codex_app_server_protocol::ExternalAgentConfigMigrationItem; +use codex_app_server_protocol::ExternalAgentConfigMigrationItemType; +use std::collections::BTreeSet; + +#[derive(Clone, Debug)] +pub(crate) struct ExternalAgentConfigMigrationGroupModel { + pub(crate) label: String, + pub(crate) description: &'static str, + pub(crate) item_indices: Vec, +} + +pub(crate) fn external_agent_config_migration_groups( + items: &[ExternalAgentConfigMigrationItem], +) -> Vec { + let tools_and_setup = items + .iter() + .enumerate() + .filter_map(|(idx, item)| { + (item.cwd.is_none() && item.item_type != ExternalAgentConfigMigrationItemType::Sessions) + .then_some(idx) + }) + .collect::>(); + let projects = items + .iter() + .enumerate() + .filter_map(|(idx, item)| { + (item.cwd.is_some() && item.item_type != ExternalAgentConfigMigrationItemType::Sessions) + .then_some(idx) + }) + .collect::>(); + let chat_sessions = items + .iter() + .enumerate() + .filter_map(|(idx, item)| { + (item.item_type == ExternalAgentConfigMigrationItemType::Sessions).then_some(idx) + }) + .collect::>(); + + let mut groups = Vec::new(); + if !tools_and_setup.is_empty() { + groups.push(ExternalAgentConfigMigrationGroupModel { + label: "Tools & setup".to_string(), + description: "Settings, instructions, integrations, agents, commands, and skills", + item_indices: tools_and_setup, + }); + } + if !projects.is_empty() { + let project_count = projects + .iter() + .filter_map(|idx| items[*idx].cwd.as_deref()) + .collect::>() + .len(); + groups.push(ExternalAgentConfigMigrationGroupModel { + label: if project_count == 1 { + "Current project".to_string() + } else { + format!("Projects ({project_count})") + }, + description: "Add Codex files alongside your existing project files", + item_indices: projects, + }); + } + if !chat_sessions.is_empty() { + let session_count = chat_sessions + .iter() + .filter_map(|idx| items[*idx].details.as_ref()) + .map(|details| details.sessions.len()) + .sum::(); + groups.push(ExternalAgentConfigMigrationGroupModel { + label: format!("Chat sessions ({session_count})"), + description: "Last 30 days of chats", + item_indices: chat_sessions, + }); + } + groups +} + +pub(crate) fn external_agent_config_migration_item_label( + item: &ExternalAgentConfigMigrationItem, +) -> &'static str { + match item.item_type { + ExternalAgentConfigMigrationItemType::AgentsMd => "Instructions (CLAUDE.md -> AGENTS.md)", + ExternalAgentConfigMigrationItemType::Config => "Settings (settings.json -> config.toml)", + ExternalAgentConfigMigrationItemType::Skills => "Skills", + ExternalAgentConfigMigrationItemType::Plugins => "Plugins", + ExternalAgentConfigMigrationItemType::McpServerConfig => "MCP servers", + ExternalAgentConfigMigrationItemType::Subagents => "Agents", + ExternalAgentConfigMigrationItemType::Hooks => "Hooks", + ExternalAgentConfigMigrationItemType::Commands => "Slash commands", + ExternalAgentConfigMigrationItemType::Sessions => "Recent chat sessions", + } +} + +pub(crate) fn external_agent_config_migration_item_detail( + item: &ExternalAgentConfigMigrationItem, +) -> Option { + let details = item.details.as_ref()?; + match item.item_type { + ExternalAgentConfigMigrationItemType::Plugins => None, + ExternalAgentConfigMigrationItemType::McpServerConfig => Some(format_counted_details( + "MCP server", + details.mcp_servers.len(), + details + .mcp_servers + .iter() + .map(|server| server.name.as_str()), + )), + ExternalAgentConfigMigrationItemType::Subagents => Some(format_counted_details( + "agent", + details.subagents.len(), + details.subagents.iter().map(|agent| agent.name.as_str()), + )), + ExternalAgentConfigMigrationItemType::Hooks => Some(format_counted_details( + "hook", + details.hooks.len(), + details.hooks.iter().map(|hook| hook.name.as_str()), + )), + ExternalAgentConfigMigrationItemType::Commands => Some(format_counted_details( + "slash command", + details.commands.len(), + details.commands.iter().map(|command| command.name.as_str()), + )), + ExternalAgentConfigMigrationItemType::Sessions => Some(format_counted_details( + "chat session", + details.sessions.len(), + details + .sessions + .iter() + .filter_map(|session| session.title.as_deref()), + )), + ExternalAgentConfigMigrationItemType::AgentsMd + | ExternalAgentConfigMigrationItemType::Config + | ExternalAgentConfigMigrationItemType::Skills => None, + } +} + +fn format_counted_details<'a>( + noun: &str, + count: usize, + names: impl Iterator, +) -> String { + let suffix = if count == 1 { "" } else { "s" }; + match names.take(4).collect::>() { + names if names.is_empty() => format!("{count} {noun}{suffix}"), + names => format!("{count} {noun}{suffix}: {}", names.join(", ")), + } +} diff --git a/codex-rs/tui/src/external_agent_config_migration_startup.rs b/codex-rs/tui/src/external_agent_config_migration_startup.rs deleted file mode 100644 index b2afbe7718..0000000000 --- a/codex-rs/tui/src/external_agent_config_migration_startup.rs +++ /dev/null @@ -1,570 +0,0 @@ -use crate::app_server_session::AppServerSession; -use crate::external_agent_config_migration::ExternalAgentConfigMigrationOutcome; -use crate::external_agent_config_migration::run_external_agent_config_migration_prompt; -use crate::legacy_core::config::Config; -use crate::legacy_core::config::ConfigBuilder; -use crate::legacy_core::config::ConfigOverrides; -use crate::legacy_core::config::edit::ConfigEdit; -use crate::legacy_core::config::edit::ConfigEditsBuilder; -use crate::tui; -use codex_app_server_protocol::ExternalAgentConfigDetectParams; -use codex_app_server_protocol::ExternalAgentConfigMigrationItem; -use codex_config::CloudConfigBundleLoader; -use codex_features::Feature; -use color_eyre::eyre::Result; -use color_eyre::eyre::WrapErr; -use std::collections::BTreeSet; -use std::path::Path; -use std::time::SystemTime; -use std::time::UNIX_EPOCH; -use toml::Value as TomlValue; - -const EXTERNAL_CONFIG_MIGRATION_PROMPT_COOLDOWN_SECS: i64 = 5 * 24 * 60 * 60; - -pub(crate) enum ExternalAgentConfigMigrationStartupOutcome { - Continue { success_message: Option }, - ExitRequested, -} - -pub(crate) fn should_show_external_agent_config_migration_prompt( - config: &Config, - entered_trust_nux: bool, -) -> bool { - entered_trust_nux && config.features.enabled(Feature::ExternalMigration) -} - -fn external_config_migration_project_key(path: &Path) -> String { - path.display().to_string() -} - -fn is_external_config_migration_scope_hidden(config: &Config, cwd: Option<&Path>) -> bool { - match cwd { - Some(cwd) => config - .notices - .external_config_migration_prompts - .projects - .get(&external_config_migration_project_key(cwd)) - .copied() - .unwrap_or(false), - None => config - .notices - .external_config_migration_prompts - .home - .unwrap_or(false), - } -} - -fn external_config_migration_last_prompted_at(config: &Config, cwd: Option<&Path>) -> Option { - match cwd { - Some(cwd) => config - .notices - .external_config_migration_prompts - .project_last_prompted_at - .get(&external_config_migration_project_key(cwd)) - .copied(), - None => { - config - .notices - .external_config_migration_prompts - .home_last_prompted_at - } - } -} - -fn is_external_config_migration_scope_cooling_down( - config: &Config, - cwd: Option<&Path>, - now_unix_seconds: i64, -) -> bool { - external_config_migration_last_prompted_at(config, cwd).is_some_and(|last_prompted_at| { - last_prompted_at.saturating_add(EXTERNAL_CONFIG_MIGRATION_PROMPT_COOLDOWN_SECS) - > now_unix_seconds - }) -} - -fn visible_external_agent_config_migration_items( - config: &Config, - items: Vec, - now_unix_seconds: i64, -) -> Vec { - items - .into_iter() - .filter(|item| { - !is_external_config_migration_scope_hidden(config, item.cwd.as_deref()) - && !is_external_config_migration_scope_cooling_down( - config, - item.cwd.as_deref(), - now_unix_seconds, - ) - }) - .collect() -} - -fn external_agent_config_migration_success_message( - items: &[ExternalAgentConfigMigrationItem], -) -> String { - if items.iter().any(|item| { - item.item_type == codex_app_server_protocol::ExternalAgentConfigMigrationItemType::Plugins - }) { - "External config migration completed. Plugin migration is still in progress and may take a few minutes." - .to_string() - } else { - "External config migration completed successfully.".to_string() - } -} - -fn unix_seconds_now() -> i64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs() as i64 -} - -async fn persist_external_agent_config_migration_prompt_shown( - config: &mut Config, - items: &[ExternalAgentConfigMigrationItem], - now_unix_seconds: i64, -) -> Result<()> { - let mut edits = Vec::new(); - if items.iter().any(|item| item.cwd.is_none()) { - edits.push( - ConfigEdit::SetNoticeExternalConfigMigrationPromptHomeLastPromptedAt(now_unix_seconds), - ); - } - - for project in items - .iter() - .filter_map(|item| item.cwd.as_deref()) - .map(external_config_migration_project_key) - { - edits.push( - ConfigEdit::SetNoticeExternalConfigMigrationPromptProjectLastPromptedAt( - project, - now_unix_seconds, - ), - ); - } - - if edits.is_empty() { - return Ok(()); - } - - ConfigEditsBuilder::for_config(config) - .with_edits(edits) - .apply() - .await - .map_err(|err| color_eyre::eyre::eyre!("{err}")) - .wrap_err("Failed to save external config migration prompt timestamp")?; - - if items.iter().any(|item| item.cwd.is_none()) { - config - .notices - .external_config_migration_prompts - .home_last_prompted_at = Some(now_unix_seconds); - } - for project in items - .iter() - .filter_map(|item| item.cwd.as_deref()) - .map(external_config_migration_project_key) - { - config - .notices - .external_config_migration_prompts - .project_last_prompted_at - .insert(project, now_unix_seconds); - } - - Ok(()) -} - -async fn persist_external_agent_config_migration_prompt_dismissal( - config: &mut Config, - items: &[ExternalAgentConfigMigrationItem], -) -> Result<()> { - let hide_home = items.iter().any(|item| item.cwd.is_none()); - let projects = items - .iter() - .filter_map(|item| item.cwd.as_deref()) - .map(external_config_migration_project_key) - .collect::>(); - - let mut edits = Vec::new(); - if hide_home - && !config - .notices - .external_config_migration_prompts - .home - .unwrap_or(false) - { - edits.push(ConfigEdit::SetNoticeHideExternalConfigMigrationPromptHome( - true, - )); - } - for project in &projects { - if !config - .notices - .external_config_migration_prompts - .projects - .get(project) - .copied() - .unwrap_or(false) - { - edits.push( - ConfigEdit::SetNoticeHideExternalConfigMigrationPromptProject( - project.clone(), - true, - ), - ); - } - } - - if edits.is_empty() { - return Ok(()); - } - - ConfigEditsBuilder::for_config(config) - .with_edits(edits) - .apply() - .await - .map_err(|err| color_eyre::eyre::eyre!("{err}")) - .wrap_err("Failed to save external config migration prompt preference")?; - - if hide_home { - config.notices.external_config_migration_prompts.home = Some(true); - } - for project in projects { - config - .notices - .external_config_migration_prompts - .projects - .insert(project, true); - } - - Ok(()) -} - -pub(crate) async fn handle_external_agent_config_migration_prompt_if_needed( - tui: &mut tui::Tui, - app_server: &mut AppServerSession, - config: &mut Config, - cli_kv_overrides: &[(String, TomlValue)], - harness_overrides: &ConfigOverrides, - cloud_config_bundle: &CloudConfigBundleLoader, - entered_trust_nux: bool, -) -> Result { - if !should_show_external_agent_config_migration_prompt(config, entered_trust_nux) { - return Ok(ExternalAgentConfigMigrationStartupOutcome::Continue { - success_message: None, - }); - } - - let now_unix_seconds = unix_seconds_now(); - let detected_items = match app_server - .external_agent_config_detect(ExternalAgentConfigDetectParams { - include_home: true, - cwds: Some(vec![config.cwd.to_path_buf()]), - }) - .await - { - Ok(response) => { - visible_external_agent_config_migration_items(config, response.items, now_unix_seconds) - } - Err(err) => { - tracing::warn!( - error = %err, - cwd = %config.cwd.display(), - "failed to detect external agent config migrations; continuing startup" - ); - return Ok(ExternalAgentConfigMigrationStartupOutcome::Continue { - success_message: None, - }); - } - }; - - if detected_items.is_empty() { - return Ok(ExternalAgentConfigMigrationStartupOutcome::Continue { - success_message: None, - }); - } - - if let Err(err) = persist_external_agent_config_migration_prompt_shown( - config, - &detected_items, - now_unix_seconds, - ) - .await - { - tracing::warn!( - error = %err, - cwd = %config.cwd.display(), - "failed to persist external config migration prompt timestamp" - ); - } - - let mut selected_items = detected_items.clone(); - let mut error: Option = None; - - loop { - match run_external_agent_config_migration_prompt( - tui, - &detected_items, - &selected_items, - error.as_deref(), - ) - .await - { - ExternalAgentConfigMigrationOutcome::Proceed(items) => { - selected_items = items.clone(); - match app_server.external_agent_config_import(items).await { - Ok(_) => { - let success_message = - external_agent_config_migration_success_message(&selected_items); - *config = ConfigBuilder::default() - .codex_home(config.codex_home.to_path_buf()) - .cli_overrides(cli_kv_overrides.to_vec()) - .harness_overrides(harness_overrides.clone()) - .cloud_config_bundle(cloud_config_bundle.clone()) - .build() - .await - .wrap_err("Failed to reload config after external agent migration")?; - return Ok(ExternalAgentConfigMigrationStartupOutcome::Continue { - success_message: Some(success_message), - }); - } - Err(err) => { - tracing::warn!( - error = %err, - cwd = %config.cwd.display(), - "failed to import external agent config migration items" - ); - error = Some(format!("Migration failed: {err}")); - } - } - } - ExternalAgentConfigMigrationOutcome::Skip => { - return Ok(ExternalAgentConfigMigrationStartupOutcome::Continue { - success_message: None, - }); - } - ExternalAgentConfigMigrationOutcome::SkipForever => { - match persist_external_agent_config_migration_prompt_dismissal( - config, - &detected_items, - ) - .await - { - Ok(()) => { - return Ok(ExternalAgentConfigMigrationStartupOutcome::Continue { - success_message: None, - }); - } - Err(err) => { - tracing::warn!( - error = %err, - cwd = %config.cwd.display(), - "failed to persist external config migration prompt dismissal" - ); - error = Some(format!("Failed to save preference: {err}")); - } - } - } - ExternalAgentConfigMigrationOutcome::Exit => { - return Ok(ExternalAgentConfigMigrationStartupOutcome::ExitRequested); - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use codex_app_server_protocol::ExternalAgentConfigMigrationItemType; - use pretty_assertions::assert_eq; - use std::path::PathBuf; - use tempfile::tempdir; - - #[tokio::test] - async fn visible_external_agent_config_migration_items_omits_hidden_scopes() { - let codex_home = tempdir().expect("temp codex home"); - let mut config = ConfigBuilder::default() - .codex_home(codex_home.path().to_path_buf()) - .build() - .await - .expect("config"); - config.notices.external_config_migration_prompts.home = Some(true); - config - .notices - .external_config_migration_prompts - .projects - .insert("/tmp/project".to_string(), true); - - let visible = visible_external_agent_config_migration_items( - &config, - vec![ - ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::Config, - description: "home".to_string(), - cwd: None, - details: None, - }, - ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::AgentsMd, - description: "project".to_string(), - cwd: Some(PathBuf::from("/tmp/project")), - details: None, - }, - ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::Skills, - description: "other project".to_string(), - cwd: Some(PathBuf::from("/tmp/other")), - details: None, - }, - ], - /*now_unix_seconds*/ 1_760_000_000, - ); - - assert_eq!( - visible, - vec![ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::Skills, - description: "other project".to_string(), - cwd: Some(PathBuf::from("/tmp/other")), - details: None, - }] - ); - } - - #[tokio::test] - async fn visible_external_agent_config_migration_items_omits_recently_prompted_scopes() { - let codex_home = tempdir().expect("temp codex home"); - let mut config = ConfigBuilder::default() - .codex_home(codex_home.path().to_path_buf()) - .build() - .await - .expect("config"); - config - .notices - .external_config_migration_prompts - .home_last_prompted_at = Some(1_760_000_000); - config - .notices - .external_config_migration_prompts - .project_last_prompted_at - .insert("/tmp/project".to_string(), 1_760_000_000); - - let visible = visible_external_agent_config_migration_items( - &config, - vec![ - ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::Config, - description: "home".to_string(), - cwd: None, - details: None, - }, - ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::AgentsMd, - description: "project".to_string(), - cwd: Some(PathBuf::from("/tmp/project")), - details: None, - }, - ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::Skills, - description: "other project".to_string(), - cwd: Some(PathBuf::from("/tmp/other")), - details: None, - }, - ], - /*now_unix_seconds*/ - 1_760_000_000 + EXTERNAL_CONFIG_MIGRATION_PROMPT_COOLDOWN_SECS - 1, - ); - - assert_eq!( - visible, - vec![ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::Skills, - description: "other project".to_string(), - cwd: Some(PathBuf::from("/tmp/other")), - details: None, - }] - ); - } - - #[tokio::test] - async fn external_config_migration_scope_cooldown_expires_after_five_days() { - let codex_home = tempdir().expect("temp codex home"); - let mut config = ConfigBuilder::default() - .codex_home(codex_home.path().to_path_buf()) - .build() - .await - .expect("config"); - config - .notices - .external_config_migration_prompts - .home_last_prompted_at = Some(1_760_000_000); - - assert!(is_external_config_migration_scope_cooling_down( - &config, - /*cwd*/ None, - 1_760_000_000 + EXTERNAL_CONFIG_MIGRATION_PROMPT_COOLDOWN_SECS - 1, - )); - assert!(!is_external_config_migration_scope_cooling_down( - &config, - /*cwd*/ None, - 1_760_000_000 + EXTERNAL_CONFIG_MIGRATION_PROMPT_COOLDOWN_SECS, - )); - } - - #[test] - fn external_agent_config_migration_success_message_mentions_plugins_when_present() { - let message = external_agent_config_migration_success_message(&[ - ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::Config, - description: String::new(), - cwd: None, - details: None, - }, - ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::Plugins, - description: String::new(), - cwd: None, - details: None, - }, - ]); - - assert_eq!( - message, - "External config migration completed. Plugin migration is still in progress and may take a few minutes." - ); - } - - #[test] - fn external_agent_config_migration_success_message_omits_plugins_copy_when_absent() { - let message = - external_agent_config_migration_success_message(&[ExternalAgentConfigMigrationItem { - item_type: ExternalAgentConfigMigrationItemType::AgentsMd, - description: String::new(), - cwd: None, - details: None, - }]); - - assert_eq!(message, "External config migration completed successfully."); - } - - #[tokio::test] - async fn external_agent_config_migration_prompt_requires_trust_nux_entry() { - let codex_home = tempdir().expect("temp codex home"); - let mut config = ConfigBuilder::default() - .codex_home(codex_home.path().to_path_buf()) - .build() - .await - .expect("config"); - let _ = config.features.enable(Feature::ExternalMigration); - - assert!(!should_show_external_agent_config_migration_prompt( - &config, /*entered_trust_nux*/ false, - )); - assert!(should_show_external_agent_config_migration_prompt( - &config, /*entered_trust_nux*/ true, - )); - } -} diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 281e4a52a0..c607bb11d6 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -64,6 +64,7 @@ use codex_utils_oss::ensure_oss_provider_ready; use codex_utils_oss::get_default_model_for_oss_provider; use color_eyre::eyre::WrapErr; use cwd_prompt::CwdPromptAction; +pub use session_archive_commands::DeleteConfirmation; pub use session_archive_commands::SessionArchiveAction; pub use session_archive_commands::SessionArchiveCommandOptions; pub use session_archive_commands::run_session_archive_command; @@ -131,7 +132,8 @@ mod diff_render; mod exec_cell; mod exec_command; mod external_agent_config_migration; -mod external_agent_config_migration_startup; +mod external_agent_config_migration_flow; +mod external_agent_config_migration_model; mod external_editor; mod file_search; mod frames; @@ -352,9 +354,11 @@ async fn init_state_db_for_app_server_target( ) -> std::io::Result> { match app_server_target { AppServerTarget::Embedded => state_db::try_init(config).await.map(Some).map_err(|err| { + let database_path = codex_state::runtime_db_path_for_corruption_error(&err) + .unwrap_or_else(|| codex_state::state_db_path(config.sqlite_home.as_path())); std::io::Error::other(LocalStateDbStartupError::new( - codex_state::state_db_path(config.sqlite_home.as_path()), - err.to_string(), + database_path, + format!("{err:#}"), )) }), AppServerTarget::LocalDaemon { .. } | AppServerTarget::Remote { .. } => { @@ -1121,7 +1125,7 @@ pub async fn run_main( let otel_originator = originator().value; let otel = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - crate::legacy_core::otel_init::build_provider( + codex_app_server_client::build_otel_provider( &config, env!("CARGO_PKG_VERSION"), /*service_name_override*/ None, @@ -1144,26 +1148,25 @@ pub async fn run_main( None } }; - crate::legacy_core::otel_init::record_process_start(otel.as_ref(), otel_originator.as_str()); - crate::legacy_core::otel_init::install_sqlite_telemetry( - otel.as_ref(), - otel_originator.as_str(), - ); + if let Some(metrics) = otel.as_ref().and_then(codex_otel::OtelProvider::metrics) { + let _ = codex_otel::record_process_start_once(metrics, otel_originator.as_str()); + let telemetry = + codex_rollout::sqlite_telemetry_recorder(metrics.clone(), otel_originator.as_str()); + let _ = codex_state::install_process_db_telemetry(telemetry); + } let state_db = init_state_db_for_app_server_target(&config, &app_server_target).await?; let effective_toml = config.config_layer_stack.effective_config(); match effective_toml.try_into() { Ok(config_toml) => { - match crate::legacy_core::personality_migration::maybe_migrate_personality( + match codex_app_server_client::migrate_personality_if_needed( &config.codex_home, &config_toml, state_db.clone(), ) .await { - Ok( - crate::legacy_core::personality_migration::PersonalityMigrationStatus::Applied, - ) => { + Ok(true) => { config = load_config_or_exit( cli_kv_overrides.clone(), overrides.clone(), @@ -1173,11 +1176,7 @@ pub async fn run_main( ) .await; } - Ok( - crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedMarker - | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedExplicitPersonality - | crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedNoSessions, - ) => {} + Ok(false) => {} Err(err) => { tracing::warn!(error = %err, "failed to run personality migration"); } @@ -1812,23 +1811,11 @@ async fn run_ratatui_app( let hooks_request_handle = app_server.request_handle(); let hooks_cwd = config.cwd.to_path_buf(); let startup_prefetch_started_at = Instant::now(); - let should_defer_bootstrap = - external_agent_config_migration_startup::should_show_external_agent_config_migration_prompt( - &config, - should_show_trust_screen_flag, - ); - let (startup_bootstrap, startup_hooks_entry) = if should_defer_bootstrap { - ( - None, - load_startup_hooks_review_entry(hooks_request_handle, hooks_cwd).await, - ) - } else { - let (bootstrap, entry) = tokio::join!( - app_server.bootstrap(&config), - load_startup_hooks_review_entry(hooks_request_handle, hooks_cwd), - ); - (Some(bootstrap?), entry) - }; + let (startup_bootstrap, startup_hooks_entry) = tokio::join!( + app_server.bootstrap(&config), + load_startup_hooks_review_entry(hooks_request_handle, hooks_cwd), + ); + let startup_bootstrap = Some(startup_bootstrap?); let startup_elapsed_before_app = startup_prefetch_started_at.elapsed(); let startup_hooks_browser = match maybe_run_startup_hooks_review( &mut app_server, @@ -1856,7 +1843,6 @@ async fn run_ratatui_app( session_selection, feedback, should_show_trust_screen, // Proxy to: is it a first run in this directory? - should_show_trust_screen_flag, // Preserve the startup-time trust NUX signal before onboarding should_prompt_windows_sandbox_nux_at_startup, app_server_target, state_db, @@ -2990,6 +2976,37 @@ mod tests { ); Ok(()) } + + #[tokio::test] + async fn embedded_state_db_corruption_preserves_failed_database_for_cli_recovery() + -> color_eyre::Result<()> { + let temp_dir = TempDir::new()?; + let mut config = build_config(&temp_dir).await?; + let sqlite_home = temp_dir.path().join("sqlite-home"); + std::fs::create_dir_all(&sqlite_home)?; + let logs_db_path = codex_state::logs_db_path(&sqlite_home); + std::fs::write(&logs_db_path, "not a sqlite database")?; + config.sqlite_home = sqlite_home; + + let err = + match init_state_db_for_app_server_target(&config, &AppServerTarget::Embedded).await { + Ok(_) => panic!("embedded startup should surface state db init failures"), + Err(err) => err, + }; + let startup_error = err + .get_ref() + .and_then(|err| err.downcast_ref::()) + .expect("state db startup failure should retain its typed context"); + + assert_eq!(startup_error.database_path(), logs_db_path.as_path()); + assert!( + codex_state::sqlite_error_detail_is_corruption(startup_error.detail()), + "startup error should preserve the SQLite corruption cause, got: {}", + startup_error.detail() + ); + Ok(()) + } + #[tokio::test] #[serial] async fn windows_shows_trust_prompt_with_sandbox() -> std::io::Result<()> { diff --git a/codex-rs/tui/src/local_chatgpt_auth.rs b/codex-rs/tui/src/local_chatgpt_auth.rs index d7ca01f618..8a9ba58fc9 100644 --- a/codex-rs/tui/src/local_chatgpt_auth.rs +++ b/codex-rs/tui/src/local_chatgpt_auth.rs @@ -110,6 +110,7 @@ mod tests { last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }; save_auth(codex_home, &auth, AuthCredentialsStoreMode::File) .expect("chatgpt auth should save"); @@ -158,6 +159,7 @@ mod tests { last_refresh: None, agent_identity: None, personal_access_token: None, + bedrock_api_key: None, }, AuthCredentialsStoreMode::File, ) diff --git a/codex-rs/tui/src/session_archive_commands.rs b/codex-rs/tui/src/session_archive_commands.rs index 4e5c420f63..7de1e577c7 100644 --- a/codex-rs/tui/src/session_archive_commands.rs +++ b/codex-rs/tui/src/session_archive_commands.rs @@ -1,8 +1,10 @@ -//! Shared implementation for `codex archive` and `codex unarchive`. +//! Shared implementation for `codex archive`, `codex delete`, and `codex unarchive`. //! //! The CLI commands are thin app-server clients: resolve a user-provided UUID or exact session -//! name, then call the existing `thread/archive` or `thread/unarchive` RPC. +//! name, then call the corresponding app-server RPC. +use std::io::IsTerminal; +use std::io::Write; use std::sync::Arc; use crate::Cli; @@ -26,16 +28,22 @@ use codex_protocol::ThreadId; use codex_utils_cli::CliConfigOverrides; use codex_utils_home_dir::find_codex_home; use codex_utils_oss::get_default_model_for_oss_provider; -use color_eyre::eyre::ContextCompat; use color_eyre::eyre::Result; use color_eyre::eyre::WrapErr; use color_eyre::eyre::eyre; use super::RemoteAppServerEndpoint; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeleteConfirmation { + Prompt, + Skip, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SessionArchiveAction { Archive, + Delete(DeleteConfirmation), Unarchive, } @@ -52,6 +60,7 @@ fn success_message( ) -> String { let action = match action { SessionArchiveAction::Archive => "Archived", + SessionArchiveAction::Delete(_) => "Deleted", SessionArchiveAction::Unarchive => "Unarchived", }; match session_name { @@ -80,25 +89,30 @@ async fn run_session_archive_action_with_app_server( target: &str, ) -> Result { let resolved = resolve_session_target(app_server, action, target).await?; - match action { + let session_name = match action { SessionArchiveAction::Archive => { app_server.thread_archive(resolved.session_id).await?; - Ok(success_message( - action, - resolved.session_id, - resolved.session_name.as_deref(), - )) + resolved.session_name + } + SessionArchiveAction::Delete(confirmation) => { + if matches!(confirmation, DeleteConfirmation::Prompt) + && !confirm_session_delete(&resolved)? + { + return Ok("Delete cancelled.".to_string()); + } + app_server.thread_delete(resolved.session_id).await?; + resolved.session_name } SessionArchiveAction::Unarchive => { let thread = app_server.thread_unarchive(resolved.session_id).await?; - let session_name = thread.name.or(resolved.session_name); - Ok(success_message( - action, - resolved.session_id, - session_name.as_deref(), - )) + thread.name.or(resolved.session_name) } - } + }; + Ok(success_message( + action, + resolved.session_id, + session_name.as_deref(), + )) } async fn resolve_session_target( @@ -107,61 +121,83 @@ async fn resolve_session_target( target: &str, ) -> Result { if let Ok(session_id) = ThreadId::from_string(target) { + if matches!( + action, + SessionArchiveAction::Delete(DeleteConfirmation::Prompt) + ) { + let thread = app_server + .thread_read(session_id, /*include_turns*/ false) + .await + .with_context(|| { + format!("No active or archived session found matching '{target}'.") + })?; + return Ok(ResolvedSessionTarget { + session_id, + session_name: thread.name, + }); + } return Ok(ResolvedSessionTarget { session_id, session_name: None, }); } - let search_scope = match action { - SessionArchiveAction::Archive => "active", - SessionArchiveAction::Unarchive => "archived", + let (search_scope, archived_values): (&str, &[bool]) = match action { + SessionArchiveAction::Archive => ("active", &[false]), + SessionArchiveAction::Delete(_) => ("active or archived", &[false, true]), + SessionArchiveAction::Unarchive => ("archived", &[true]), }; - let resolved = lookup_session_by_exact_name(app_server, action, target) - .await? - .map(session_target_from_app_server_thread) - .transpose()?; - - resolved.with_context(|| format!("No {search_scope} session found matching '{target}'.")) + for &archived in archived_values { + if let Some(thread) = lookup_session_by_exact_name(app_server, target, archived).await? { + return session_target_from_app_server_thread(thread); + } + } + Err(eyre!( + "No {search_scope} session found matching '{target}'." + )) } async fn lookup_session_by_exact_name( app_server: &mut AppServerSession, - action: SessionArchiveAction, name: &str, + archived: bool, ) -> Result> { - let mut cursor = None; - loop { - let response = app_server - .thread_list(ThreadListParams { - cursor: cursor.clone(), - limit: Some(100), - sort_key: Some(ThreadSortKey::UpdatedAt), - sort_direction: None, - model_providers: None, - source_kinds: Some(super::resume_source_kinds( - /*include_non_interactive*/ false, - )), - archived: Some(matches!(action, SessionArchiveAction::Unarchive)), - cwd: None, - use_state_db_only: false, - search_term: Some(name.to_string()), - }) - .await - .wrap_err("failed to list sessions while resolving session name")?; + // Search is the fast path, but some stores attach renamed titles after applying the filter. + for search_term in [Some(name), None] { + let mut cursor = None; + loop { + let response = app_server + .thread_list(ThreadListParams { + cursor: cursor.clone(), + limit: Some(100), + sort_key: Some(ThreadSortKey::UpdatedAt), + sort_direction: None, + model_providers: None, + source_kinds: Some(super::resume_source_kinds( + /*include_non_interactive*/ false, + )), + archived: Some(archived), + cwd: None, + use_state_db_only: false, + search_term: search_term.map(str::to_string), + }) + .await + .wrap_err("failed to list sessions while resolving session name")?; - if let Some(thread) = response - .data - .into_iter() - .find(|thread| thread.name.as_deref() == Some(name)) - { - return Ok(Some(thread)); + if let Some(thread) = response + .data + .into_iter() + .find(|thread| thread.name.as_deref() == Some(name)) + { + return Ok(Some(thread)); + } + let Some(next_cursor) = response.next_cursor else { + break; + }; + cursor = Some(next_cursor); } - if response.next_cursor.is_none() { - return Ok(None); - } - cursor = response.next_cursor; } + Ok(None) } fn session_target_from_app_server_thread(thread: AppServerThread) -> Result { @@ -173,6 +209,35 @@ fn session_target_from_app_server_thread(thread: AppServerThread) -> Result Result { + if !(std::io::stdin().is_terminal() && std::io::stderr().is_terminal()) { + return Err(eyre!( + "cannot confirm session deletion without an interactive terminal; rerun with --force and a session UUID" + )); + } + + let mut stderr = std::io::stderr().lock(); + match target.session_name.as_deref() { + Some(name) => writeln!( + stderr, + "Permanently delete session '{name}' ({})?", + target.session_id + ), + None => writeln!(stderr, "Permanently delete session {}?", target.session_id), + }?; + writeln!( + stderr, + "This cannot be undone. Subagent threads will also be deleted." + )?; + write!(stderr, "Continue? [y/N]: ")?; + stderr.flush()?; + + let mut input = String::new(); + std::io::stdin().read_line(&mut input)?; + let answer = input.trim(); + Ok(answer.eq_ignore_ascii_case("y") || answer.eq_ignore_ascii_case("yes")) +} + async fn start_app_server_for_archive_command( options: SessionArchiveCommandOptions, ) -> Result { diff --git a/codex-rs/tui/src/slash_command.rs b/codex-rs/tui/src/slash_command.rs index 805f266e0a..9ced21b26e 100644 --- a/codex-rs/tui/src/slash_command.rs +++ b/codex-rs/tui/src/slash_command.rs @@ -26,11 +26,13 @@ pub enum SlashCommand { AutoReview, Memories, Skills, + Import, Hooks, Review, Rename, New, Archive, + Delete, Resume, Fork, App, @@ -89,6 +91,7 @@ impl SlashCommand { SlashCommand::Rename => "rename the current thread", SlashCommand::Resume => "resume a saved chat", SlashCommand::Archive => "archive this session and exit", + SlashCommand::Delete => "permanently delete this session and exit", SlashCommand::Clear => "clear the terminal and start a new chat", SlashCommand::Fork => "fork the current chat", SlashCommand::App => "continue this session in Codex Desktop", @@ -98,6 +101,9 @@ impl SlashCommand { SlashCommand::Diff => "show git diff (including untracked files)", SlashCommand::Mention => "mention a file", SlashCommand::Skills => "use skills to improve how Codex performs specific tasks", + SlashCommand::Import => { + "import setup, this project, and recent chats from another coding agent" + } SlashCommand::Hooks => "view and manage lifecycle hooks", SlashCommand::Status => "show current session configuration and token usage", SlashCommand::DebugConfig => "show config layers and requirement sources for debugging", @@ -185,6 +191,7 @@ impl SlashCommand { match self { SlashCommand::New | SlashCommand::Archive + | SlashCommand::Delete | SlashCommand::Resume | SlashCommand::Fork | SlashCommand::Init @@ -198,6 +205,7 @@ impl SlashCommand { | SlashCommand::SandboxReadRoot | SlashCommand::Experimental | SlashCommand::Memories + | SlashCommand::Import | SlashCommand::Review | SlashCommand::Plan | SlashCommand::Clear diff --git a/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize.snap b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize.snap new file mode 100644 index 0000000000..f23691b6f3 --- /dev/null +++ b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize.snap @@ -0,0 +1,29 @@ +--- +source: tui/src/external_agent_config_migration.rs +expression: rendered +--- + + > Choose what to import + Choose the items to import. + Codex may add files to your current project folder. + Your existing agent setup will not be changed. + Home + › [x] Settings (settings.json -> config.toml) + Import /Users/alex/.claude/settings.json into /Users/alex/.codex/conf… + [x] Recent chat sessions + Import recent chat sessions + 1 chat session: Investigate migration UX + + Current project: /workspace/project + [x] Plugins + Import enabled plugins from .claude/settings.json (4 marketplaces, 6 … + • acme-tools: deployer, formatter, +1 more + • team-marketplace: asana + • debug: sample + • +1 more marketplaces + [x] Instructions (CLAUDE.md -> AGENTS.md) + Import CLAUDE.md to AGENTS.md + + Selected 4 of 4 items. + 1. Review selection + Use ↑/↓ to move, space to toggle, b to go back diff --git a/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize_action.snap b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize_action.snap new file mode 100644 index 0000000000..70df423cee --- /dev/null +++ b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize_action.snap @@ -0,0 +1,29 @@ +--- +source: tui/src/external_agent_config_migration.rs +expression: rendered +--- + + > Choose what to import + Choose the items to import. + Codex may add files to your current project folder. + Your existing agent setup will not be changed. + Home + [x] Settings (settings.json -> config.toml) + Import /Users/alex/.claude/settings.json into /Users/alex/.codex/conf… + [x] Recent chat sessions + Import recent chat sessions + 1 chat session: Investigate migration UX + + Current project: /workspace/project + [x] Plugins + Import enabled plugins from .claude/settings.json (4 marketplaces, 6 … + • acme-tools: deployer, formatter, +1 more + • team-marketplace: asana + • debug: sample + • +1 more marketplaces + [x] Instructions (CLAUDE.md -> AGENTS.md) + Import CLAUDE.md to AGENTS.md + + Selected 4 of 4 items. + › 1. Review selection + Press enter to continue, ↑/↓ to move, b to go back diff --git a/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize_action_windows.snap b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize_action_windows.snap new file mode 100644 index 0000000000..e6e35384ae --- /dev/null +++ b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize_action_windows.snap @@ -0,0 +1,29 @@ +--- +source: tui/src/external_agent_config_migration.rs +expression: rendered +--- + + > Choose what to import + Choose the items to import. + Codex may add files to your current project folder. + Your existing agent setup will not be changed. + Home + [x] Settings (settings.json -> config.toml) + Import /Users/alex/.claude/settings.json into /Users/alex/.codex/conf… + [x] Recent chat sessions + Import recent chat sessions + 1 chat session: Investigate migration UX + + Current project: C:\workspace\project + [x] Plugins + Import enabled plugins from .claude/settings.json (4 marketplaces, 6 … + • acme-tools: deployer, formatter, +1 more + • team-marketplace: asana + • debug: sample + • +1 more marketplaces + [x] Instructions (CLAUDE.md -> AGENTS.md) + Import CLAUDE.md to AGENTS.md + + Selected 4 of 4 items. + › 1. Review selection + Press enter to continue, ↑/↓ to move, b to go back diff --git a/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize_windows.snap b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize_windows.snap new file mode 100644 index 0000000000..12dff7a4e5 --- /dev/null +++ b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_customize_windows.snap @@ -0,0 +1,29 @@ +--- +source: tui/src/external_agent_config_migration.rs +expression: rendered +--- + + > Choose what to import + Choose the items to import. + Codex may add files to your current project folder. + Your existing agent setup will not be changed. + Home + › [x] Settings (settings.json -> config.toml) + Import /Users/alex/.claude/settings.json into /Users/alex/.codex/conf… + [x] Recent chat sessions + Import recent chat sessions + 1 chat session: Investigate migration UX + + Current project: C:\workspace\project + [x] Plugins + Import enabled plugins from .claude/settings.json (4 marketplaces, 6 … + • acme-tools: deployer, formatter, +1 more + • team-marketplace: asana + • debug: sample + • +1 more marketplaces + [x] Instructions (CLAUDE.md -> AGENTS.md) + Import CLAUDE.md to AGENTS.md + + Selected 4 of 4 items. + 1. Review selection + Use ↑/↓ to move, space to toggle, b to go back diff --git a/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_prompt.snap b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_prompt.snap index 5ce3c5ef09..64a5f0b1b0 100644 --- a/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_prompt.snap +++ b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_prompt.snap @@ -2,22 +2,21 @@ source: tui/src/external_agent_config_migration.rs expression: rendered --- - > External agent config detected - We found settings from another agent that you can add to this project. - Select what to import - Home - [x] Migrate /Users/alex/.claude/settings.json into /Users/alex/.codex/con… - Project: /workspace/project - [x] Migrate enabled plugins from .claude/settings.json (4 marketplaces, 6… - • acme-tools: deployer, formatter, +1 more - • team-marketplace: asana - • debug: sample - • +1 more marketplaces - [x] Migrate CLAUDE.md to AGENTS.md + > Import from another coding agent + Bring over your setup, current project, and recent chats. + Codex may add files to your current project folder. + Your existing agent setup will not be changed. + Cloud-hosted chat data cannot be imported. + [x] Tools & setup + Settings, instructions, integrations, agents, commands, and skills + [x] Current project + Add Codex files alongside your existing project files + [x] Chat sessions (1) + Last 30 days of chats - Selected 3 of 3 item(s). - 1. Proceed with selected - 2. Skip for now - 3. Don't ask again - Use ↑/↓ to move, space to toggle, 1/2/3 to choose, a/n for all/none + Selected 4 of 4 items. + › 1. Import selected + 2. Customize selection + 3. Cancel + Use ↑/↓ to move, enter to select, c to customize diff --git a/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_prompt_windows.snap b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_prompt_windows.snap index 44f41784e6..64a5f0b1b0 100644 --- a/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_prompt_windows.snap +++ b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration__tests__external_agent_config_migration_prompt_windows.snap @@ -3,22 +3,20 @@ source: tui/src/external_agent_config_migration.rs expression: rendered --- - > External agent config detected - We found settings from another agent that you can add to this project. - Select what to import - Home - [x] Migrate /Users/alex/.claude/settings.json into /Users/alex/.codex/con… + > Import from another coding agent + Bring over your setup, current project, and recent chats. + Codex may add files to your current project folder. + Your existing agent setup will not be changed. + Cloud-hosted chat data cannot be imported. + [x] Tools & setup + Settings, instructions, integrations, agents, commands, and skills + [x] Current project + Add Codex files alongside your existing project files + [x] Chat sessions (1) + Last 30 days of chats - Project: C:\workspace\project - [x] Migrate enabled plugins from .claude/settings.json (4 marketplaces, 6… - • acme-tools: deployer, formatter, +1 more - • team-marketplace: asana - • debug: sample - • +1 more marketplaces - [x] Migrate CLAUDE.md to AGENTS.md - - Selected 3 of 3 item(s). - 1. Proceed with selected - 2. Skip for now - 3. Don't ask again - Use ↑/↓ to move, space to toggle, 1/2/3 to choose, a/n for all/none + Selected 4 of 4 items. + › 1. Import selected + 2. Customize selection + 3. Cancel + Use ↑/↓ to move, enter to select, c to customize diff --git a/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration_flow__tests__external_agent_config_migration_messages.snap b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration_flow__tests__external_agent_config_migration_messages.snap new file mode 100644 index 0000000000..5e80f8ac29 --- /dev/null +++ b/codex-rs/tui/src/snapshots/codex_tui__external_agent_config_migration_flow__tests__external_agent_config_migration_messages.snap @@ -0,0 +1,12 @@ +--- +source: tui/src/external_agent_config_migration_flow_tests.rs +expression: messages +--- +Agent import started. You can keep working while it finishes. Imported setup will apply to new chats. +Agent import started. You can keep working while it finishes. Imported setup will apply to new chats. 1 additional item remains. After it finishes, run /import again to review it. +Agent import started. You can keep working while it finishes. Imported setup will apply to new chats. 2 additional items remain. After it finishes, run /import again to review them. +Agent import finished. Run /import again to check for additional items. +No supported agent setup was found to import. +Agent import is unavailable in remote sessions. Start Codex locally and run /import. +Agent import is unavailable while Codex is connected to the local app-server daemon. Stop the daemon, restart Codex, and run /import. +A previous agent import is still running. Wait for it to finish before importing again. diff --git a/codex-rs/tui/src/startup_error.rs b/codex-rs/tui/src/startup_error.rs index 55b691d836..585cec4b23 100644 --- a/codex-rs/tui/src/startup_error.rs +++ b/codex-rs/tui/src/startup_error.rs @@ -3,24 +3,28 @@ use std::path::PathBuf; #[derive(Debug, thiserror::Error)] #[error( - "failed to initialize sqlite state db at {}: {detail}", - state_db_path.display() + "failed to initialize sqlite local db at {}: {detail}", + database_path.display() )] pub struct LocalStateDbStartupError { - state_db_path: PathBuf, + database_path: PathBuf, detail: String, } impl LocalStateDbStartupError { - pub fn new(state_db_path: PathBuf, detail: String) -> Self { + pub fn new(database_path: PathBuf, detail: String) -> Self { Self { - state_db_path, + database_path, detail, } } + pub fn database_path(&self) -> &Path { + self.database_path.as_path() + } + pub fn state_db_path(&self) -> &Path { - self.state_db_path.as_path() + self.database_path() } pub fn detail(&self) -> &str { diff --git a/codex-rs/tui/src/status/tests.rs b/codex-rs/tui/src/status/tests.rs index 1fc4c4feba..4ab7a2cfa2 100644 --- a/codex-rs/tui/src/status/tests.rs +++ b/codex-rs/tui/src/status/tests.rs @@ -29,6 +29,8 @@ use codex_app_server_protocol::SpendControlLimitSnapshot; use codex_config::LoaderOverrides; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::ModelProviderInfo; +use codex_models_manager::test_support::construct_model_info_offline_for_tests; +use codex_models_manager::test_support::get_model_offline_for_tests; use codex_protocol::ThreadId; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::ReasoningSummary; @@ -147,7 +149,7 @@ fn test_status_account_display() -> Option { fn token_info_for(model_slug: &str, config: &Config, usage: &TokenUsage) -> TokenUsageInfo { let context_window = - crate::legacy_core::test_support::construct_model_info_offline(model_slug, config) + construct_model_info_offline_for_tests(model_slug, &config.to_models_manager_config()) .context_window; TokenUsageInfo { total_token_usage: usage.clone(), @@ -212,7 +214,7 @@ fn permissions_text_for(config: &Config) -> Option { .with_ymd_and_hms(2024, 1, 2, 3, 4, 5) .single() .expect("timestamp"); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let composite = new_status_output( config, test_status_account_display().as_ref(), @@ -287,7 +289,7 @@ async fn status_snapshot_includes_reasoning_details() { }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let reasoning_effort_override = Some(Some(ReasoningEffort::High)); @@ -603,7 +605,7 @@ async fn status_snapshot_shows_active_user_defined_profile() { .with_ymd_and_hms(2024, 1, 2, 3, 4, 5) .single() .expect("timestamp"); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( @@ -648,7 +650,7 @@ async fn status_model_provider_uses_bedrock_runtime_base_url_and_gates_usage_lin .with_ymd_and_hms(2024, 1, 2, 3, 4, 5) .single() .expect("timestamp"); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let runtime_base_url = "https://bedrock-mantle.eu-west-1.api.aws/openai/v1"; let (composite, _handle) = new_status_output_with_rate_limits_handle( @@ -758,7 +760,7 @@ async fn status_snapshot_shows_auto_review_permissions() { .with_ymd_and_hms(2024, 1, 2, 3, 4, 5) .single() .expect("timestamp"); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( @@ -854,7 +856,7 @@ async fn status_snapshot_includes_forked_from() { .single() .expect("valid time"); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let session_id = ThreadId::from_string("0f0f3c13-6cf9-4aa4-8b80-7d49c2f1be2e").expect("session id"); @@ -923,7 +925,7 @@ async fn status_snapshot_includes_monthly_limit() { }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -987,7 +989,7 @@ async fn status_snapshot_includes_enterprise_monthly_credit_limit() { }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1067,7 +1069,7 @@ async fn status_snapshot_uses_generic_limit_labels_for_unsupported_windows() { }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1119,7 +1121,7 @@ async fn status_snapshot_shows_unlimited_credits() { rate_limit_reached_type: None, }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1170,7 +1172,7 @@ async fn status_snapshot_shows_positive_credits() { rate_limit_reached_type: None, }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1221,7 +1223,7 @@ async fn status_snapshot_hides_zero_credits() { rate_limit_reached_type: None, }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1270,7 +1272,7 @@ async fn status_snapshot_hides_when_has_no_credits_flag() { rate_limit_reached_type: None, }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1315,7 +1317,7 @@ async fn status_card_token_usage_excludes_cached_tokens() { .single() .expect("timestamp"); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1378,7 +1380,7 @@ async fn status_snapshot_truncates_in_narrow_terminal() { }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let reasoning_effort_override = Some(Some(ReasoningEffort::High)); let composite = new_status_output( @@ -1428,7 +1430,7 @@ async fn status_snapshot_shows_missing_limits_message() { .single() .expect("timestamp"); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1480,7 +1482,7 @@ async fn status_snapshot_uses_default_reasoning_when_config_empty() { version: "v0.133.0".to_string(), }; - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let (composite, _) = new_status_output_with_rate_limits_handle( &config, @@ -1549,7 +1551,7 @@ async fn status_snapshot_shows_refreshing_limits_notice() { }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output_with_rate_limits( &config, @@ -1621,7 +1623,7 @@ async fn status_snapshot_includes_credits_and_limits() { }; let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1680,7 +1682,7 @@ async fn status_snapshot_shows_unavailable_limits_message() { .expect("timestamp"); let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1738,7 +1740,7 @@ async fn status_snapshot_treats_refreshing_empty_limits_as_unavailable() { .expect("timestamp"); let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output_with_rate_limits( &config, @@ -1807,7 +1809,7 @@ async fn status_snapshot_shows_stale_limits_message() { let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); let now = captured_at + ChronoDuration::minutes(20); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1879,7 +1881,7 @@ async fn status_snapshot_cached_limits_hide_credits_without_flag() { let rate_display = rate_limit_snapshot_display(&snapshot, captured_at); let now = captured_at + ChronoDuration::minutes(20); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = token_info_for(&model_slug, &config, &usage); let composite = new_status_output( &config, @@ -1933,7 +1935,7 @@ async fn status_context_window_uses_last_usage() { .single() .expect("timestamp"); - let model_slug = crate::legacy_core::test_support::get_model_offline(config.model.as_deref()); + let model_slug = get_model_offline_for_tests(config.model.as_deref()); let token_info = TokenUsageInfo { total_token_usage: total_usage.clone(), last_token_usage: last_usage, diff --git a/codex-rs/tui/src/test_support.rs b/codex-rs/tui/src/test_support.rs index 53fd8adf62..44e9899379 100644 --- a/codex-rs/tui/src/test_support.rs +++ b/codex-rs/tui/src/test_support.rs @@ -1,8 +1,23 @@ +//! Test-only helpers shared across the TUI crate. + +use std::sync::LazyLock; + +use codex_models_manager::bundled_models_response; +use codex_protocol::openai_models::ModelPreset; pub(crate) use codex_utils_absolute_path::test_support::PathBufExt; pub(crate) use codex_utils_absolute_path::test_support::test_path_buf; use serde::Serialize; use serde::de::DeserializeOwned; +pub(crate) static TEST_MODEL_PRESETS: LazyLock> = LazyLock::new(|| { + let mut response = bundled_models_response() + .unwrap_or_else(|err| panic!("bundled models.json should parse: {err}")); + response.models.sort_by_key(|model| model.priority); + let mut presets: Vec = response.models.into_iter().map(Into::into).collect(); + ModelPreset::mark_default_by_picker_visibility(&mut presets); + presets +}); + pub(crate) fn test_path_display(path: &str) -> String { test_path_buf(path).display().to_string() } diff --git a/codex-rs/utils/image/src/error.rs b/codex-rs/utils/image/src/error.rs index 28b73f4a7c..aaa3d4f2f3 100644 --- a/codex-rs/utils/image/src/error.rs +++ b/codex-rs/utils/image/src/error.rs @@ -25,6 +25,14 @@ pub enum ImageProcessingError { }, #[error("unsupported image `{mime}`")] UnsupportedImageFormat { mime: String }, + #[error("invalid image data URL: {reason}")] + InvalidDataUrl { reason: String }, + #[error("image {representation} is too large ({size} bytes; max {max} bytes)")] + ImageTooLarge { + representation: &'static str, + size: usize, + max: usize, + }, } impl ImageProcessingError { diff --git a/codex-rs/utils/image/src/image_tests.rs b/codex-rs/utils/image/src/image_tests.rs new file mode 100644 index 0000000000..e60c8c6be0 --- /dev/null +++ b/codex-rs/utils/image/src/image_tests.rs @@ -0,0 +1,341 @@ +use std::io::Cursor; + +use super::*; +use image::GenericImageView; +use image::ImageBuffer; +use image::ImageDecoder; +use image::Rgba; +use image::metadata::Orientation; + +const TEST_RGB_ICC_PROFILE: &[u8] = b"0123456789abcdefRGB "; +const TEST_CMYK_ICC_PROFILE: &[u8] = b"0123456789abcdefCMYK"; +const ROTATE_90_EXIF: &[u8] = &[ + 0x49, 0x49, 0x2a, 0x00, 0x08, 0x00, 0x00, 0x00, 0x01, 0x00, 0x12, 0x01, 0x03, 0x00, 0x01, 0x00, + 0x00, 0x00, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, +]; + +fn image_bytes(image: &ImageBuffer, Vec>, format: ImageFormat) -> Vec { + let mut encoded = Cursor::new(Vec::new()); + DynamicImage::ImageRgba8(image.clone()) + .write_to(&mut encoded, format) + .expect("encode image to bytes"); + encoded.into_inner() +} + +fn image_bytes_with_metadata( + image: &ImageBuffer, Vec>, + format: ImageFormat, + icc_profile: &[u8], +) -> Vec { + let mut encoded = Vec::new(); + match format { + ImageFormat::Png => { + let mut encoder = PngEncoder::new(&mut encoded); + encoder + .set_icc_profile(icc_profile.to_vec()) + .expect("set PNG ICC profile"); + encoder + .set_exif_metadata(ROTATE_90_EXIF.to_vec()) + .expect("set PNG EXIF metadata"); + encoder + .write_image( + image.as_raw(), + image.width(), + image.height(), + ColorType::Rgba8.into(), + ) + .expect("encode PNG with metadata"); + } + ImageFormat::Jpeg => { + let mut encoder = JpegEncoder::new_with_quality(&mut encoded, 90); + encoder + .set_icc_profile(icc_profile.to_vec()) + .expect("set JPEG ICC profile"); + encoder + .set_exif_metadata(ROTATE_90_EXIF.to_vec()) + .expect("set JPEG EXIF metadata"); + encoder + .encode_image(&DynamicImage::ImageRgba8(image.clone())) + .expect("encode JPEG with metadata"); + } + ImageFormat::WebP => { + let mut encoder = WebPEncoder::new_lossless(&mut encoded); + encoder + .set_icc_profile(icc_profile.to_vec()) + .expect("set WebP ICC profile"); + encoder + .set_exif_metadata(ROTATE_90_EXIF.to_vec()) + .expect("set WebP EXIF metadata"); + encoder + .write_image( + image.as_raw(), + image.width(), + image.height(), + ColorType::Rgba8.into(), + ) + .expect("encode WebP with metadata"); + } + _ => panic!("unsupported test format"), + } + encoded +} + +#[tokio::test(flavor = "multi_thread")] +async fn returns_original_image_when_within_bounds() { + for (format, mime) in [ + (ImageFormat::Png, "image/png"), + (ImageFormat::WebP, "image/webp"), + ] { + let image = ImageBuffer::from_pixel(64, 32, Rgba([10u8, 20, 30, 255])); + let original_bytes = image_bytes(&image, format); + + let encoded = load_for_prompt_bytes( + Path::new("in-memory-image"), + original_bytes.clone(), + PromptImageMode::ResizeToFit, + ) + .expect("process image"); + + assert_eq!(encoded.width, 64); + assert_eq!(encoded.height, 32); + assert_eq!(encoded.mime, mime); + assert_eq!(encoded.bytes, original_bytes); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn downscales_large_image() { + for (format, mime) in [ + (ImageFormat::Png, "image/png"), + (ImageFormat::WebP, "image/webp"), + ] { + let image = ImageBuffer::from_pixel(4096, 2048, Rgba([200u8, 10, 10, 255])); + let original_bytes = image_bytes(&image, format); + + let processed = load_for_prompt_bytes( + Path::new("in-memory-image"), + original_bytes, + PromptImageMode::ResizeToFit, + ) + .expect("process image"); + + assert!(processed.width <= MAX_DIMENSION); + assert!(processed.height <= MAX_DIMENSION); + assert_eq!(processed.mime, mime); + + let detected_format = + image::guess_format(&processed.bytes).expect("detect resized output format"); + assert_eq!(detected_format, format); + + let loaded = + image::load_from_memory(&processed.bytes).expect("read resized bytes back into image"); + assert_eq!(loaded.dimensions(), (processed.width, processed.height)); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn downscales_tall_image_to_fit_square_bounds() { + let image = ImageBuffer::from_pixel(1024, 4096, Rgba([200u8, 10, 10, 255])); + let original_bytes = image_bytes(&image, ImageFormat::Png); + + let processed = load_for_prompt_bytes( + Path::new("in-memory-image"), + original_bytes, + PromptImageMode::ResizeToFit, + ) + .expect("process image"); + + assert_eq!(processed.width, 512); + assert_eq!(processed.height, MAX_DIMENSION); + assert_eq!(processed.mime, "image/png"); +} + +#[tokio::test(flavor = "multi_thread")] +async fn resizing_preserves_supported_metadata() { + for format in [ImageFormat::Png, ImageFormat::Jpeg, ImageFormat::WebP] { + let image = ImageBuffer::from_pixel(2050, 2, Rgba([200u8, 10, 10, 255])); + let original_bytes = image_bytes_with_metadata(&image, format, TEST_RGB_ICC_PROFILE); + + let processed = load_for_prompt_bytes( + Path::new("in-memory-image"), + original_bytes, + PromptImageMode::ResizeToFit, + ) + .expect("process image"); + + assert_eq!((processed.width, processed.height), (2048, 2)); + + let mut decoder = ImageReader::with_format(Cursor::new(&processed.bytes), format) + .into_decoder() + .expect("create decoder"); + assert_eq!( + ( + decoder.dimensions(), + decoder.orientation().expect("read orientation"), + decoder.icc_profile().expect("read ICC profile"), + decoder.exif_metadata().expect("read EXIF metadata"), + ), + ( + (2048, 2), + Orientation::Rotate90, + Some(TEST_RGB_ICC_PROFILE.to_vec()), + Some(ROTATE_90_EXIF.to_vec()), + ) + ); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn resizing_drops_non_rgb_icc_profile() { + let image = ImageBuffer::from_pixel(2050, 2, Rgba([200u8, 10, 10, 255])); + let original_bytes = + image_bytes_with_metadata(&image, ImageFormat::Jpeg, TEST_CMYK_ICC_PROFILE); + + let processed = load_for_prompt_bytes( + Path::new("in-memory-image"), + original_bytes, + PromptImageMode::ResizeToFit, + ) + .expect("process image"); + + let mut decoder = ImageReader::with_format(Cursor::new(&processed.bytes), ImageFormat::Jpeg) + .into_decoder() + .expect("create decoder"); + assert_eq!( + ( + decoder.icc_profile().expect("read ICC profile"), + decoder.exif_metadata().expect("read EXIF metadata"), + ), + (None, Some(ROTATE_90_EXIF.to_vec())) + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn preserves_large_image_in_original_mode() { + let image = ImageBuffer::from_pixel(4096, 2048, Rgba([180u8, 30, 30, 255])); + let original_bytes = image_bytes(&image, ImageFormat::Png); + + let processed = load_for_prompt_bytes( + Path::new("in-memory-image"), + original_bytes.clone(), + PromptImageMode::Original, + ) + .expect("process image"); + + assert_eq!(processed.width, 4096); + assert_eq!(processed.height, 2048); + assert_eq!(processed.mime, "image/png"); + assert_eq!(processed.bytes, original_bytes); +} + +#[tokio::test(flavor = "multi_thread")] +async fn data_url_processing_preserves_supported_source_bytes() { + let image = ImageBuffer::from_pixel(64, 32, Rgba([10u8, 20, 30, 255])); + let original_bytes = image_bytes(&image, ImageFormat::Png); + let image_url = data_url_from_bytes("image/png", &original_bytes) + .replacen("data:", "DATA:", 1) + .replacen(";base64,", ";BASE64,", 1); + + let processed = load_data_url_for_prompt(&image_url, PromptImageMode::ResizeToFit) + .expect("process data URL image"); + + assert_eq!(processed.width, 64); + assert_eq!(processed.height, 32); + assert_eq!(processed.mime, "image/png"); + assert_eq!(processed.bytes, original_bytes); +} + +#[tokio::test(flavor = "multi_thread")] +async fn data_url_processing_converts_gif_to_png() { + let image = ImageBuffer::from_pixel(64, 32, Rgba([10u8, 20, 30, 255])); + let gif_bytes = image_bytes(&image, ImageFormat::Gif); + let image_url = data_url_from_bytes("image/gif", &gif_bytes); + + let processed = load_data_url_for_prompt(&image_url, PromptImageMode::ResizeToFit) + .expect("process GIF data URL"); + + assert_eq!(processed.mime, "image/png"); + assert_eq!( + image::guess_format(&processed.bytes).expect("detect processed format"), + ImageFormat::Png + ); +} + +#[test] +fn data_url_processing_rejects_malformed_input() { + for image_url in [ + "image/png;base64,AAAA", + "data:image/png;base64", + "data:image/png,AAAA", + "data:image/png;base64,not base64", + ] { + assert!(matches!( + load_data_url_for_prompt(image_url, PromptImageMode::ResizeToFit), + Err(ImageProcessingError::InvalidDataUrl { .. }) + )); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn resize_with_limits_respects_dimension_and_patch_budgets() { + let image = ImageBuffer::from_pixel(2048, 2048, Rgba([200u8, 10, 10, 255])); + let original_bytes = image_bytes(&image, ImageFormat::Png); + let limits = PromptImageResizeLimits { + max_dimension: 2048, + max_patches: 2_500, + }; + + let processed = load_for_prompt_bytes( + Path::new("in-memory-image"), + original_bytes, + PromptImageMode::ResizeWithLimits(limits), + ) + .expect("process image with explicit limits"); + + assert_eq!((processed.width, processed.height), (1600, 1600)); +} + +#[tokio::test(flavor = "multi_thread")] +async fn fails_cleanly_for_invalid_images() { + let err = load_for_prompt_bytes( + Path::new("in-memory-image"), + b"not an image".to_vec(), + PromptImageMode::ResizeToFit, + ) + .expect_err("invalid image should fail"); + assert!(matches!( + err, + ImageProcessingError::Decode { .. } | ImageProcessingError::UnsupportedImageFormat { .. } + )); +} + +#[tokio::test(flavor = "multi_thread")] +async fn reprocesses_updated_file_contents() { + IMAGE_CACHE.clear(); + + let first_image = ImageBuffer::from_pixel(32, 16, Rgba([20u8, 120, 220, 255])); + let first_bytes = image_bytes(&first_image, ImageFormat::Png); + + let first = load_for_prompt_bytes( + Path::new("in-memory-image"), + first_bytes, + PromptImageMode::ResizeToFit, + ) + .expect("process first image"); + + let second_image = ImageBuffer::from_pixel(96, 48, Rgba([50u8, 60, 70, 255])); + let second_bytes = image_bytes(&second_image, ImageFormat::Png); + + let second = load_for_prompt_bytes( + Path::new("in-memory-image"), + second_bytes, + PromptImageMode::ResizeToFit, + ) + .expect("process updated image"); + + assert_eq!(first.width, 32); + assert_eq!(first.height, 16); + assert_eq!(second.width, 96); + assert_eq!(second.height, 48); + assert_ne!(second.bytes, first.bytes); +} diff --git a/codex-rs/utils/image/src/lib.rs b/codex-rs/utils/image/src/lib.rs index a3f595a183..93ae427ea5 100644 --- a/codex-rs/utils/image/src/lib.rs +++ b/codex-rs/utils/image/src/lib.rs @@ -1,3 +1,4 @@ +use std::io::Cursor; use std::num::NonZeroUsize; use std::path::Path; use std::sync::LazyLock; @@ -9,14 +10,24 @@ use codex_utils_cache::sha1_digest; use image::ColorType; use image::DynamicImage; use image::GenericImageView; +use image::ImageDecoder; use image::ImageEncoder; use image::ImageFormat; +use image::ImageReader; use image::codecs::jpeg::JpegEncoder; use image::codecs::png::PngEncoder; use image::codecs::webp::WebPEncoder; use image::imageops::FilterType; + +const DATA_URL_PREFIX: &str = "data:"; +pub const PROMPT_IMAGE_PATCH_SIZE: u32 = 32; /// Maximum width or height used when resizing images before uploading. pub const MAX_DIMENSION: u32 = 2048; +/// Maximum accepted byte length for prompt image input representations. +/// +/// This is a high sanity guard against pathological inputs, not a protocol +/// requirement or target upload size. +pub const MAX_PROMPT_IMAGE_INPUT_BYTES: usize = 1024 * 1024 * 1024; pub mod error; @@ -32,15 +43,32 @@ pub struct EncodedImage { impl EncodedImage { pub fn into_data_url(self) -> String { - let encoded = BASE64_STANDARD.encode(&self.bytes); - format!("data:{};base64,{encoded}", self.mime) + data_url_from_bytes(&self.mime, &self.bytes) } } +/// Wraps image bytes in a data URL without decoding or validating them. +pub fn data_url_from_bytes(mime: &str, bytes: &[u8]) -> String { + let encoded = BASE64_STANDARD.encode(bytes); + format!("data:{mime};base64,{encoded}") +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum PromptImageMode { ResizeToFit, Original, + ResizeWithLimits(PromptImageResizeLimits), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct PromptImageResizeLimits { + pub max_dimension: u32, + pub max_patches: usize, +} + +struct ImageMetadata { + icc_profile: Option>, + exif: Option>, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -65,22 +93,70 @@ pub fn load_for_prompt_bytes( }; IMAGE_CACHE.get_or_try_insert_with(key, move || { - let format = match image::guess_format(&file_bytes) { - Ok(ImageFormat::Png) => Some(ImageFormat::Png), - Ok(ImageFormat::Jpeg) => Some(ImageFormat::Jpeg), - Ok(ImageFormat::Gif) => Some(ImageFormat::Gif), - Ok(ImageFormat::WebP) => Some(ImageFormat::WebP), + let guessed_format = image::guess_format(&file_bytes) + .map_err(|source| ImageProcessingError::decode_error(&path_buf, source))?; + let format = match guessed_format { + ImageFormat::Png => Some(ImageFormat::Png), + ImageFormat::Jpeg => Some(ImageFormat::Jpeg), + ImageFormat::Gif => Some(ImageFormat::Gif), + ImageFormat::WebP => Some(ImageFormat::WebP), _ => None, }; - let dynamic = image::load_from_memory(&file_bytes) + let mut decoder = ImageReader::with_format(Cursor::new(&file_bytes), guessed_format) + .into_decoder() + .map_err(|source| ImageProcessingError::decode_error(&path_buf, source))?; + // Preserve the metadata most important for rendering prompt images faithfully: the color + // profile and EXIF data, including orientation. Other format-specific metadata is + // intentionally not copied. + let metadata = ImageMetadata { + // Only RGB profiles are safe across every re-encoding path. For example, JPEG decoding + // can convert CMYK/YCCK pixels to RGB while retaining the source profile; copying it + // would mislabel the output. Bytes 16..20 are the ICC data color space signature. + icc_profile: decoder + .icc_profile() + .ok() + .flatten() + .filter(|profile| profile.get(16..20) == Some(b"RGB ")), + exif: decoder.exif_metadata().ok().flatten(), + }; + let dynamic = DynamicImage::from_decoder(decoder) .map_err(|source| ImageProcessingError::decode_error(&path_buf, source))?; let (width, height) = dynamic.dimensions(); - let encoded = if mode == PromptImageMode::Original - || (width <= MAX_DIMENSION && height <= MAX_DIMENSION) - { + let target_dimensions = match mode { + PromptImageMode::ResizeToFit if width > MAX_DIMENSION || height > MAX_DIMENSION => { + let resized = dynamic.resize(MAX_DIMENSION, MAX_DIMENSION, FilterType::Triangle); + Some((resized.width(), resized.height(), resized)) + } + PromptImageMode::ResizeWithLimits(limits) => { + let (target_width, target_height) = + prompt_image_output_dimensions_for_limits(width, height, limits); + if (target_width, target_height) == (width, height) { + None + } else { + let resized = + dynamic.resize_exact(target_width, target_height, FilterType::Triangle); + Some((target_width, target_height, resized)) + } + } + PromptImageMode::ResizeToFit | PromptImageMode::Original => None, + }; + + let encoded = if let Some((width, height, resized)) = target_dimensions { + let target_format = format + .filter(|format| can_preserve_source_bytes(*format)) + .unwrap_or(ImageFormat::Png); + let (bytes, output_format) = encode_image(&resized, target_format, metadata)?; + let mime = format_to_mime(output_format); + EncodedImage { + bytes, + mime, + width, + height, + } + } else { if let Some(format) = format.filter(|format| can_preserve_source_bytes(*format)) { let mime = format_to_mime(format); EncodedImage { @@ -90,7 +166,7 @@ pub fn load_for_prompt_bytes( height, } } else { - let (bytes, output_format) = encode_image(&dynamic, ImageFormat::Png)?; + let (bytes, output_format) = encode_image(&dynamic, ImageFormat::Png, metadata)?; let mime = format_to_mime(output_format); EncodedImage { bytes, @@ -99,25 +175,107 @@ pub fn load_for_prompt_bytes( height, } } - } else { - let resized = dynamic.resize(MAX_DIMENSION, MAX_DIMENSION, FilterType::Triangle); - let target_format = format - .filter(|format| can_preserve_source_bytes(*format)) - .unwrap_or(ImageFormat::Png); - let (bytes, output_format) = encode_image(&resized, target_format)?; - let mime = format_to_mime(output_format); - EncodedImage { - bytes, - mime, - width: resized.width(), - height: resized.height(), - } }; Ok(encoded) }) } +pub fn load_data_url_for_prompt( + image_url: &str, + mode: PromptImageMode, +) -> Result { + let rest = image_url + .get(..DATA_URL_PREFIX.len()) + .filter(|prefix| prefix.eq_ignore_ascii_case(DATA_URL_PREFIX)) + .and_then(|_| image_url.get(DATA_URL_PREFIX.len()..)) + .ok_or_else(|| ImageProcessingError::InvalidDataUrl { + reason: "missing data: prefix".to_string(), + })?; + let (metadata, encoded) = + rest.split_once(',') + .ok_or_else(|| ImageProcessingError::InvalidDataUrl { + reason: "missing comma separator".to_string(), + })?; + if !metadata + .split(';') + .any(|part| part.eq_ignore_ascii_case("base64")) + { + return Err(ImageProcessingError::InvalidDataUrl { + reason: "only base64 data URLs are supported".to_string(), + }); + } + + if encoded.len() > MAX_PROMPT_IMAGE_INPUT_BYTES { + return Err(ImageProcessingError::ImageTooLarge { + representation: "base64 payload", + size: encoded.len(), + max: MAX_PROMPT_IMAGE_INPUT_BYTES, + }); + } + let file_bytes = + BASE64_STANDARD + .decode(encoded) + .map_err(|source| ImageProcessingError::InvalidDataUrl { + reason: format!("invalid base64 payload: {source}"), + })?; + if file_bytes.len() > MAX_PROMPT_IMAGE_INPUT_BYTES { + return Err(ImageProcessingError::ImageTooLarge { + representation: "decoded input", + size: file_bytes.len(), + max: MAX_PROMPT_IMAGE_INPUT_BYTES, + }); + } + + load_for_prompt_bytes(Path::new(""), file_bytes, mode) +} + +fn prompt_image_output_dimensions_for_limits( + width: u32, + height: u32, + limits: PromptImageResizeLimits, +) -> (u32, u32) { + let width = width.max(1); + let height = height.max(1); + if prompt_image_dimensions_fit(width, height, limits) { + return (width, height); + } + + let max_dimension_scale = + (f64::from(limits.max_dimension) / f64::from(width.max(height))).min(1.0); + let width = ((f64::from(width) * max_dimension_scale).round() as u32).max(1); + let height = ((f64::from(height) * max_dimension_scale).round() as u32).max(1); + if prompt_image_dimensions_fit(width, height, limits) { + return (width, height); + } + + let width_f64 = f64::from(width); + let height_f64 = f64::from(height); + let patch_size = f64::from(PROMPT_IMAGE_PATCH_SIZE); + let mut scale = + (patch_size * patch_size * limits.max_patches as f64 / width_f64 / height_f64).sqrt(); + // Match Responses patch-budget math: shrink by area, then round the scaled + // patch grid down so integer output dimensions remain within the budget. + let scaled_patches_wide = width_f64 * scale / patch_size; + let scaled_patches_high = height_f64 * scale / patch_size; + scale *= (scaled_patches_wide.floor() / scaled_patches_wide) + .min(scaled_patches_high.floor() / scaled_patches_high); + + ( + ((width_f64 * scale).floor() as u32).max(1), + ((height_f64 * scale).floor() as u32).max(1), + ) +} + +fn prompt_image_dimensions_fit(width: u32, height: u32, limits: PromptImageResizeLimits) -> bool { + let patches_wide = width.div_ceil(PROMPT_IMAGE_PATCH_SIZE); + let patches_high = height.div_ceil(PROMPT_IMAGE_PATCH_SIZE); + let patch_count = u64::from(patches_wide) * u64::from(patches_high); + width <= limits.max_dimension + && height <= limits.max_dimension + && patch_count <= limits.max_patches as u64 +} + fn can_preserve_source_bytes(format: ImageFormat) -> bool { // Public API docs explicitly call out non-animated GIF support only. // Preserve byte-for-byte only for formats we can safely pass through. @@ -130,6 +288,7 @@ fn can_preserve_source_bytes(format: ImageFormat) -> bool { fn encode_image( image: &DynamicImage, preferred_format: ImageFormat, + metadata: ImageMetadata, ) -> Result<(Vec, ImageFormat), ImageProcessingError> { let target_format = match preferred_format { ImageFormat::Jpeg => ImageFormat::Jpeg, @@ -138,11 +297,13 @@ fn encode_image( }; let mut buffer = Vec::new(); + let ImageMetadata { icc_profile, exif } = metadata; match target_format { ImageFormat::Png => { let rgba = image.to_rgba8(); - let encoder = PngEncoder::new(&mut buffer); + let mut encoder = PngEncoder::new(&mut buffer); + apply_image_metadata(&mut encoder, icc_profile, exif, target_format)?; encoder .write_image( rgba.as_raw(), @@ -157,6 +318,7 @@ fn encode_image( } ImageFormat::Jpeg => { let mut encoder = JpegEncoder::new_with_quality(&mut buffer, 85); + apply_image_metadata(&mut encoder, icc_profile, exif, target_format)?; encoder .encode_image(image) .map_err(|source| ImageProcessingError::Encode { @@ -166,7 +328,8 @@ fn encode_image( } ImageFormat::WebP => { let rgba = image.to_rgba8(); - let encoder = WebPEncoder::new_lossless(&mut buffer); + let mut encoder = WebPEncoder::new_lossless(&mut buffer); + apply_image_metadata(&mut encoder, icc_profile, exif, target_format)?; encoder .write_image( rgba.as_raw(), @@ -185,6 +348,31 @@ fn encode_image( Ok((buffer, target_format)) } +fn apply_image_metadata( + encoder: &mut impl ImageEncoder, + icc_profile: Option>, + exif: Option>, + format: ImageFormat, +) -> Result<(), ImageProcessingError> { + if let Some(icc_profile) = icc_profile { + encoder + .set_icc_profile(icc_profile) + .map_err(|source| ImageProcessingError::Encode { + format, + source: image::ImageError::Unsupported(source), + })?; + } + if let Some(exif) = exif { + encoder + .set_exif_metadata(exif) + .map_err(|source| ImageProcessingError::Encode { + format, + source: image::ImageError::Unsupported(source), + })?; + } + Ok(()) +} + fn format_to_mime(format: ImageFormat) -> String { match format { ImageFormat::Jpeg => "image/jpeg".to_string(), @@ -195,155 +383,5 @@ fn format_to_mime(format: ImageFormat) -> String { } #[cfg(test)] -mod tests { - use std::io::Cursor; - - use super::*; - use image::GenericImageView; - use image::ImageBuffer; - use image::Rgba; - - fn image_bytes(image: &ImageBuffer, Vec>, format: ImageFormat) -> Vec { - let mut encoded = Cursor::new(Vec::new()); - DynamicImage::ImageRgba8(image.clone()) - .write_to(&mut encoded, format) - .expect("encode image to bytes"); - encoded.into_inner() - } - - #[tokio::test(flavor = "multi_thread")] - async fn returns_original_image_when_within_bounds() { - for (format, mime) in [ - (ImageFormat::Png, "image/png"), - (ImageFormat::WebP, "image/webp"), - ] { - let image = ImageBuffer::from_pixel(64, 32, Rgba([10u8, 20, 30, 255])); - let original_bytes = image_bytes(&image, format); - - let encoded = load_for_prompt_bytes( - Path::new("in-memory-image"), - original_bytes.clone(), - PromptImageMode::ResizeToFit, - ) - .expect("process image"); - - assert_eq!(encoded.width, 64); - assert_eq!(encoded.height, 32); - assert_eq!(encoded.mime, mime); - assert_eq!(encoded.bytes, original_bytes); - } - } - - #[tokio::test(flavor = "multi_thread")] - async fn downscales_large_image() { - for (format, mime) in [ - (ImageFormat::Png, "image/png"), - (ImageFormat::WebP, "image/webp"), - ] { - let image = ImageBuffer::from_pixel(4096, 2048, Rgba([200u8, 10, 10, 255])); - let original_bytes = image_bytes(&image, format); - - let processed = load_for_prompt_bytes( - Path::new("in-memory-image"), - original_bytes, - PromptImageMode::ResizeToFit, - ) - .expect("process image"); - - assert!(processed.width <= MAX_DIMENSION); - assert!(processed.height <= MAX_DIMENSION); - assert_eq!(processed.mime, mime); - - let detected_format = - image::guess_format(&processed.bytes).expect("detect resized output format"); - assert_eq!(detected_format, format); - - let loaded = image::load_from_memory(&processed.bytes) - .expect("read resized bytes back into image"); - assert_eq!(loaded.dimensions(), (processed.width, processed.height)); - } - } - - #[tokio::test(flavor = "multi_thread")] - async fn downscales_tall_image_to_fit_square_bounds() { - let image = ImageBuffer::from_pixel(1024, 4096, Rgba([200u8, 10, 10, 255])); - let original_bytes = image_bytes(&image, ImageFormat::Png); - - let processed = load_for_prompt_bytes( - Path::new("in-memory-image"), - original_bytes, - PromptImageMode::ResizeToFit, - ) - .expect("process image"); - - assert_eq!(processed.width, 512); - assert_eq!(processed.height, MAX_DIMENSION); - assert_eq!(processed.mime, "image/png"); - } - - #[tokio::test(flavor = "multi_thread")] - async fn preserves_large_image_in_original_mode() { - let image = ImageBuffer::from_pixel(4096, 2048, Rgba([180u8, 30, 30, 255])); - let original_bytes = image_bytes(&image, ImageFormat::Png); - - let processed = load_for_prompt_bytes( - Path::new("in-memory-image"), - original_bytes.clone(), - PromptImageMode::Original, - ) - .expect("process image"); - - assert_eq!(processed.width, 4096); - assert_eq!(processed.height, 2048); - assert_eq!(processed.mime, "image/png"); - assert_eq!(processed.bytes, original_bytes); - } - - #[tokio::test(flavor = "multi_thread")] - async fn fails_cleanly_for_invalid_images() { - let err = load_for_prompt_bytes( - Path::new("in-memory-image"), - b"not an image".to_vec(), - PromptImageMode::ResizeToFit, - ) - .expect_err("invalid image should fail"); - assert!(matches!( - err, - ImageProcessingError::Decode { .. } - | ImageProcessingError::UnsupportedImageFormat { .. } - )); - } - - #[tokio::test(flavor = "multi_thread")] - async fn reprocesses_updated_file_contents() { - { - IMAGE_CACHE.clear(); - } - - let first_image = ImageBuffer::from_pixel(32, 16, Rgba([20u8, 120, 220, 255])); - let first_bytes = image_bytes(&first_image, ImageFormat::Png); - - let first = load_for_prompt_bytes( - Path::new("in-memory-image"), - first_bytes, - PromptImageMode::ResizeToFit, - ) - .expect("process first image"); - - let second_image = ImageBuffer::from_pixel(96, 48, Rgba([50u8, 60, 70, 255])); - let second_bytes = image_bytes(&second_image, ImageFormat::Png); - - let second = load_for_prompt_bytes( - Path::new("in-memory-image"), - second_bytes, - PromptImageMode::ResizeToFit, - ) - .expect("process updated image"); - - assert_eq!(first.width, 32); - assert_eq!(first.height, 16); - assert_eq!(second.width, 96); - assert_eq!(second.height, 48); - assert_ne!(second.bytes, first.bytes); - } -} +#[path = "image_tests.rs"] +mod tests; diff --git a/codex-rs/utils/path-uri/src/lib.rs b/codex-rs/utils/path-uri/src/lib.rs index a788fdc9b4..444af1ea28 100644 --- a/codex-rs/utils/path-uri/src/lib.rs +++ b/codex-rs/utils/path-uri/src/lib.rs @@ -9,6 +9,8 @@ use serde::Deserializer; use serde::Serialize; use serde::Serializer; use std::fmt; +use std::io; +use std::path::Path; use std::str::FromStr; use thiserror::Error; use ts_rs::TS; @@ -25,7 +27,7 @@ pub const FILE_SCHEME: &str = "file"; /// /// `file:` paths retain their URI spelling so they can be parsed independently /// of the current host. In particular, `/C:/src` remains ambiguous between a -/// Windows drive path and a valid POSIX path until [`Self::to_native_path`] +/// Windows drive path and a valid POSIX path until [`Self::to_abs_path`] /// applies the current host's rules. A local POSIX `file:` URI can also retain /// percent-encoded non-UTF-8 bytes for lossless native round trips. /// @@ -55,15 +57,27 @@ impl PathUri { /// Converts an absolute path on the current host to a `file:` URI. /// - /// On Unix, [`AbsolutePathBuf`]'s absolute-path invariant is sufficient for - /// `url` to represent the path. On Windows, conversion can still fail for - /// absolute paths whose prefix has no `file:` URI representation, including - /// `\\.\` device paths and generic `\\?\` verbatim namespaces. Those cases - /// return [`PathUriParseError::InvalidFileUriPath`]. - pub fn from_file_path(path: &AbsolutePathBuf) -> Result { - let url = Url::from_file_path(path.as_path()) - .map_err(|()| PathUriParseError::InvalidFileUriPath)?; - Self::try_from(url) + /// On Windows, paths without a URI representation, including `\\.\` device + /// paths and generic `\\?\` verbatim namespaces, are reported as invalid + /// input. + pub fn from_abs_path(path: &AbsolutePathBuf) -> io::Result { + let url = Url::from_file_path(path.as_path()).map_err(|()| { + io::Error::new( + io::ErrorKind::InvalidInput, + PathUriParseError::InvalidFileUriPath, + ) + })?; + Self::try_from(url).map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err)) + } + + /// Converts a path on the current host to a `file:` URI. + /// + /// Relative paths and paths without a URI representation are reported as + /// invalid input. + pub fn from_path(path: impl AsRef) -> io::Result { + let path = AbsolutePathBuf::from_absolute_path_checked(path) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?; + Self::from_abs_path(&path) } /// Returns the percent-encoded URI path. @@ -149,13 +163,15 @@ impl PathUri { /// on POSIX or a POSIX root on Windows. Because a `file:` URI does not record /// its source operating system, callers should only use this method when the /// URI is known to identify a path on the current host. - pub fn to_native_path(&self) -> Result { - let path = self - .0 - .to_file_path() - .map_err(|()| PathUriParseError::InvalidFileUriPath)?; + pub fn to_abs_path(&self) -> io::Result { + let path = self.0.to_file_path().map_err(|()| { + io::Error::new( + io::ErrorKind::InvalidInput, + PathUriParseError::InvalidFileUriPath, + ) + })?; AbsolutePathBuf::from_absolute_path_checked(path) - .map_err(|_| PathUriParseError::InvalidFileUriPath) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err)) } /// Returns a clone of the canonical URL. @@ -216,7 +232,7 @@ impl<'de> Deserialize<'de> for PathUri { .map_or_else(|| path_error.to_string(), |error| error.to_string()), ) })?; - Self::from_file_path(&path).map_err(serde::de::Error::custom) + Self::from_abs_path(&path).map_err(serde::de::Error::custom) } } diff --git a/codex-rs/utils/path-uri/src/tests.rs b/codex-rs/utils/path-uri/src/tests.rs index 6a237a8c1e..4fb2ba5b45 100644 --- a/codex-rs/utils/path-uri/src/tests.rs +++ b/codex-rs/utils/path-uri/src/tests.rs @@ -12,7 +12,7 @@ fn file_uri_round_trips_an_absolute_path() { .expect("current directory") .join("a path/file.rs"); - let uri = PathUri::from_file_path(&path).expect("path should convert to a file URI"); + let uri = PathUri::from_abs_path(&path).expect("path should convert to a file URI"); let uri_string = uri.to_string(); assert!(uri_string.starts_with("file:")); @@ -22,12 +22,26 @@ fn file_uri_round_trips_an_absolute_path() { uri ); assert_eq!( - uri.to_native_path() + uri.to_abs_path() .expect("local file URI should convert to a native path"), path ); } +#[test] +fn non_native_uri_io_conversion_is_invalid_input() { + #[cfg(unix)] + let uri = PathUri::parse("file://server/share/file.txt").expect("valid file URI"); + #[cfg(windows)] + let uri = PathUri::parse("file:///usr/local/file.txt").expect("valid file URI"); + + let error = uri + .to_abs_path() + .expect_err("URI should not be host-native"); + + assert_eq!(error.kind(), io::ErrorKind::InvalidInput); +} + #[test] fn file_uri_parses_a_windows_path_on_any_host() { let uri = PathUri::parse("file:///C:/Users/Alice%20Smith/src/main.rs") @@ -52,8 +66,10 @@ fn file_uri_rejects_windows_prefixes_without_a_uri_representation() { .expect("Windows namespace path should be absolute"); assert_eq!( - PathUri::from_file_path(&path), - Err(PathUriParseError::InvalidFileUriPath), + PathUri::from_abs_path(&path) + .expect_err("Windows namespace path should not convert") + .kind(), + io::ErrorKind::InvalidInput, "converting {native_path}" ); } @@ -85,9 +101,9 @@ fn file_uri_accepts_non_utf8_posix_paths() { let path = PathBuf::from(std::ffi::OsString::from_vec(b"/tmp/non-utf8-\xff".to_vec())); let path = AbsolutePathBuf::from_absolute_path_checked(path).expect("absolute POSIX path"); - let uri = PathUri::from_file_path(&path).expect("non-UTF-8 path should convert to a file URI"); + let uri = PathUri::from_abs_path(&path).expect("non-UTF-8 path should convert to a file URI"); assert_eq!( - uri.to_native_path() + uri.to_abs_path() .expect("URI should convert to native path"), path ); @@ -111,10 +127,10 @@ fn file_uri_round_trips_literal_percent_characters() { fn file_uri_round_trips_windows_unc_paths() { let path = AbsolutePathBuf::from_absolute_path_checked(r"\\server\share\src\main.rs") .expect("absolute UNC path"); - let uri = PathUri::from_file_path(&path).expect("UNC path should convert to a file URI"); + let uri = PathUri::from_abs_path(&path).expect("UNC path should convert to a file URI"); assert_eq!(uri.encoded_path(), "/share/src/main.rs"); - assert_eq!(uri.to_native_path().expect("UNC URI should convert"), path); + assert_eq!(uri.to_abs_path().expect("UNC URI should convert"), path); } #[test] @@ -184,10 +200,17 @@ fn path_uri_deserializes_legacy_absolute_paths() { assert_eq!( uri, - PathUri::from_file_path(&path).expect("expected file URI") + PathUri::from_abs_path(&path).expect("expected file URI") ); } +#[test] +fn path_uri_rejects_relative_native_paths() { + let error = PathUri::from_path("src/lib.rs").expect_err("relative path should be rejected"); + + assert_eq!(error.kind(), io::ErrorKind::InvalidInput); +} + #[test] fn path_uri_rejects_legacy_relative_paths_with_absolute_path_guard() { let base = AbsolutePathBuf::current_dir().expect("current directory"); diff --git a/codex-rs/utils/pty/src/lib.rs b/codex-rs/utils/pty/src/lib.rs index 39fc9b5522..9f0cb442a2 100644 --- a/codex-rs/utils/pty/src/lib.rs +++ b/codex-rs/utils/pty/src/lib.rs @@ -17,6 +17,8 @@ pub use pipe::spawn_process_no_stdin as spawn_pipe_process_no_stdin; pub use process::ProcessDriver; /// Handle for interacting with a spawned process (PTY or pipe). pub use process::ProcessHandle; +/// Process signal supported by spawned-process handles. +pub use process::ProcessSignal; /// Bundle of process handles plus split output and exit receivers returned by spawn helpers. pub use process::SpawnedProcess; /// Terminal size in character cells used for PTY spawn and resize operations. diff --git a/codex-rs/utils/pty/src/pipe.rs b/codex-rs/utils/pty/src/pipe.rs index 541a2ecf2f..3a9b62d9b7 100644 --- a/codex-rs/utils/pty/src/pipe.rs +++ b/codex-rs/utils/pty/src/pipe.rs @@ -19,7 +19,9 @@ use tokio::task::JoinHandle; use crate::process::ChildTerminator; use crate::process::ProcessHandle; +use crate::process::ProcessSignal; use crate::process::SpawnedProcess; +use crate::process::exit_code_from_status; #[cfg(target_os = "linux")] use libc; @@ -32,6 +34,22 @@ struct PipeChildTerminator { } impl ChildTerminator for PipeChildTerminator { + fn signal(&mut self, signal: ProcessSignal) -> io::Result<()> { + match signal { + ProcessSignal::Interrupt => { + #[cfg(unix)] + { + crate::process_group::interrupt_process_group(self.process_group_id) + } + + #[cfg(not(unix))] + { + Err(crate::process::unsupported_signal(signal)) + } + } + } + } + fn kill(&mut self) -> io::Result<()> { #[cfg(unix)] { @@ -209,7 +227,7 @@ async fn spawn_process_with_stdin_mode( let wait_exit_code = Arc::clone(&exit_code); let wait_handle: JoinHandle<()> = tokio::spawn(async move { let code = match child.wait().await { - Ok(status) => status.code().unwrap_or(-1), + Ok(status) => exit_code_from_status(status), Err(_) => -1, }; wait_exit_status.store(true, std::sync::atomic::Ordering::SeqCst); diff --git a/codex-rs/utils/pty/src/process.rs b/codex-rs/utils/pty/src/process.rs index 898a3d90ff..fdd693f43d 100644 --- a/codex-rs/utils/pty/src/process.rs +++ b/codex-rs/utils/pty/src/process.rs @@ -2,6 +2,7 @@ use core::fmt; use std::io; #[cfg(unix)] use std::os::fd::RawFd; +use std::process::ExitStatus; use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::atomic::AtomicBool; @@ -17,7 +18,39 @@ use tokio::sync::watch; use tokio::task::AbortHandle; use tokio::task::JoinHandle; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ProcessSignal { + Interrupt, +} + +pub(crate) fn unsupported_signal(signal: ProcessSignal) -> io::Error { + match signal { + ProcessSignal::Interrupt => io::Error::new( + io::ErrorKind::Unsupported, + "process interrupt is not supported by this process backend", + ), + } +} + +pub(crate) fn exit_code_from_status(status: ExitStatus) -> i32 { + if let Some(code) = status.code() { + return code; + } + + #[cfg(unix)] + { + use std::os::unix::process::ExitStatusExt; + if let Some(signal) = status.signal() { + return 128 + signal; + } + } + + -1 +} + pub(crate) trait ChildTerminator: Send + Sync { + fn signal(&mut self, signal: ProcessSignal) -> io::Result<()>; + fn kill(&mut self) -> io::Result<()>; } @@ -193,6 +226,17 @@ impl ProcessHandle { } } + pub fn signal(&self, signal: ProcessSignal) -> io::Result<()> { + let Ok(mut killer_opt) = self.killer.lock() else { + return Ok(()); + }; + let Some(killer) = killer_opt.as_mut() else { + return Ok(()); + }; + + killer.signal(signal) + } + /// Attempts to kill the child and abort helper tasks. pub fn terminate(&self) { self.request_terminate(); @@ -232,6 +276,10 @@ struct ClosureTerminator { } impl ChildTerminator for ClosureTerminator { + fn signal(&mut self, signal: ProcessSignal) -> io::Result<()> { + Err(unsupported_signal(signal)) + } + fn kill(&mut self) -> io::Result<()> { if let Some(inner) = self.inner.as_mut() { (inner)(); diff --git a/codex-rs/utils/pty/src/process_group.rs b/codex-rs/utils/pty/src/process_group.rs index 22a934fab7..9d10d7c368 100644 --- a/codex-rs/utils/pty/src/process_group.rs +++ b/codex-rs/utils/pty/src/process_group.rs @@ -118,15 +118,10 @@ pub fn kill_process_group_by_pid(_pid: u32) -> io::Result<()> { } #[cfg(unix)] -/// Send SIGTERM to a specific process group ID (best-effort). -/// -/// Returns `Ok(true)` when SIGTERM was delivered to an existing group and -/// `Ok(false)` when the group no longer exists. -pub fn terminate_process_group(process_group_id: u32) -> io::Result { +fn signal_process_group_id(pgid: libc::pid_t, signal: libc::c_int) -> io::Result { use std::io::ErrorKind; - let pgid = process_group_id as libc::pid_t; - let result = unsafe { libc::killpg(pgid, libc::SIGTERM) }; + let result = unsafe { libc::killpg(pgid, signal) }; if result == -1 { let err = io::Error::last_os_error(); if err.kind() == ErrorKind::NotFound || err.raw_os_error() == Some(libc::ESRCH) { @@ -138,27 +133,37 @@ pub fn terminate_process_group(process_group_id: u32) -> io::Result { Ok(true) } +#[cfg(unix)] +/// Send SIGTERM to a specific process group ID (best-effort). +/// +/// Returns `Ok(true)` when SIGTERM was delivered to an existing group and +/// `Ok(false)` when the group no longer exists. +pub fn terminate_process_group(process_group_id: u32) -> io::Result { + signal_process_group_id(process_group_id as libc::pid_t, libc::SIGTERM) +} + #[cfg(not(unix))] /// No-op on non-Unix platforms. pub fn terminate_process_group(_process_group_id: u32) -> io::Result { Ok(false) } +#[cfg(unix)] +/// Send SIGINT to a specific process group ID (best-effort). +pub fn interrupt_process_group(process_group_id: u32) -> io::Result<()> { + signal_process_group_id(process_group_id as libc::pid_t, libc::SIGINT).map(|_| ()) +} + +#[cfg(not(unix))] +/// No-op on non-Unix platforms. +pub fn interrupt_process_group(_process_group_id: u32) -> io::Result<()> { + Ok(()) +} + #[cfg(unix)] /// Kill a specific process group ID (best-effort). pub fn kill_process_group(process_group_id: u32) -> io::Result<()> { - use std::io::ErrorKind; - - let pgid = process_group_id as libc::pid_t; - let result = unsafe { libc::killpg(pgid, libc::SIGKILL) }; - if result == -1 { - let err = io::Error::last_os_error(); - if err.kind() != ErrorKind::NotFound && err.raw_os_error() != Some(libc::ESRCH) { - return Err(err); - } - } - - Ok(()) + signal_process_group_id(process_group_id as libc::pid_t, libc::SIGKILL).map(|_| ()) } #[cfg(not(unix))] diff --git a/codex-rs/utils/pty/src/pty.rs b/codex-rs/utils/pty/src/pty.rs index 45c587b328..951818a80d 100644 --- a/codex-rs/utils/pty/src/pty.rs +++ b/codex-rs/utils/pty/src/pty.rs @@ -30,10 +30,13 @@ use tokio::task::JoinHandle; use crate::process::ChildTerminator; use crate::process::ProcessHandle; +use crate::process::ProcessSignal; use crate::process::PtyHandles; use crate::process::PtyMasterHandle; use crate::process::SpawnedProcess; use crate::process::TerminalSize; +#[cfg(unix)] +use crate::process::exit_code_from_status; /// Returns true when ConPTY support is available (Windows only). #[cfg(windows)] @@ -54,6 +57,19 @@ struct PtyChildTerminator { } impl ChildTerminator for PtyChildTerminator { + fn signal(&mut self, signal: ProcessSignal) -> std::io::Result<()> { + match signal { + ProcessSignal::Interrupt => { + #[cfg(unix)] + if let Some(process_group_id) = self.process_group_id { + return crate::process_group::interrupt_process_group(process_group_id); + } + + Err(crate::process::unsupported_signal(signal)) + } + } + } + fn kill(&mut self) -> std::io::Result<()> { #[cfg(unix)] if let Some(process_group_id) = self.process_group_id { @@ -81,6 +97,14 @@ struct RawPidTerminator { #[cfg(unix)] impl ChildTerminator for RawPidTerminator { + fn signal(&mut self, signal: ProcessSignal) -> std::io::Result<()> { + match signal { + ProcessSignal::Interrupt => { + crate::process_group::interrupt_process_group(self.process_group_id) + } + } + } + fn kill(&mut self) -> std::io::Result<()> { crate::process_group::kill_process_group(self.process_group_id) } @@ -368,7 +392,7 @@ async fn spawn_process_preserving_fds( let wait_exit_code = Arc::clone(&exit_code); let wait_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || { let code = match child.wait() { - Ok(status) => status.code().unwrap_or(-1), + Ok(status) => exit_code_from_status(status), Err(_) => -1, }; wait_exit_status.store(true, std::sync::atomic::Ordering::SeqCst); diff --git a/defs.bzl b/defs.bzl index 4af81ecef1..b255160b09 100644 --- a/defs.bzl +++ b/defs.bzl @@ -1,18 +1,8 @@ load("@crates//:data.bzl", "DEP_DATA") load("@crates//:defs.bzl", "all_crate_deps") -load("@rules_platform//platform_data:defs.bzl", "platform_data") load("@rules_rust//cargo/private:cargo_build_script_wrapper.bzl", "cargo_build_script") load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_library", "rust_proc_macro", "rust_test") -PLATFORMS = [ - "linux_arm64_musl", - "linux_amd64_musl", - "macos_amd64", - "macos_arm64", - "windows_amd64", - "windows_arm64", -] - # Match Cargo's Windows linker behavior so Bazel-built binaries and tests use # the same stack reserve on both Windows ABIs and resolve UCRT imports on MSVC. WINDOWS_RUSTC_LINK_FLAGS = select({ @@ -53,21 +43,6 @@ MACOS_WEBRTC_RUSTC_LINK_FLAGS = select({ "//conditions:default": [], }) -def multiplatform_binaries(name, platforms = PLATFORMS): - for platform in platforms: - platform_data( - name = name + "_" + platform, - platform = "@llvm//platforms:" + platform, - target = name, - tags = ["manual"], - ) - - native.filegroup( - name = "release_binaries", - srcs = [name + "_" + platform for platform in platforms], - tags = ["manual"], - ) - def _workspace_root_test_impl(ctx): is_windows = ctx.target_platform_has_constraint(ctx.attr._windows_constraint[platform_common.ConstraintValueInfo]) launcher = ctx.actions.declare_file(ctx.label.name + ".bat" if is_windows else ctx.label.name) diff --git a/patches/BUILD.bazel b/patches/BUILD.bazel index 72d6990790..347b0f7ae4 100644 --- a/patches/BUILD.bazel +++ b/patches/BUILD.bazel @@ -14,6 +14,7 @@ exports_files([ "rules_rust_windows_msvc_direct_link_args.patch", "rules_rust_windows_gnullvm_build_script.patch", "rules_cc_rusty_v8_custom_libcxx.patch", + "rules_rs_build_script_deps_annotation.patch", "rules_rs_windows_gnullvm_exec.patch", "rules_rs_windows_exec_linker.patch", "rusty_v8_prebuilt_out_dir.patch", diff --git a/patches/rules_rs_build_script_deps_annotation.patch b/patches/rules_rs_build_script_deps_annotation.patch new file mode 100644 index 0000000000..ef1c442bc4 --- /dev/null +++ b/patches/rules_rs_build_script_deps_annotation.patch @@ -0,0 +1,60 @@ +# What: expose the existing unconditional build-script dependency input through +# the `rules_rs` crate annotation API. +# Scope: annotation plumbing only; generated dependency resolution is unchanged. + +diff --git a/rs/extensions.bzl b/rs/extensions.bzl +--- a/rs/extensions.bzl ++++ b/rs/extensions.bzl +@@ -737,7 +737,7 @@ def _generate_hub_and_spokes( + additive_build_file = annotation.additive_build_file, + additive_build_file_content = annotation.additive_build_file_content, + gen_build_script = annotation.gen_build_script, +- build_script_deps = [], ++ build_script_deps = annotation.build_script_deps, + build_script_deps_select = _select(feature_resolutions.build_deps), + build_script_data = annotation.build_script_data, + build_script_data_select = annotation.build_script_data_select, +@@ -1417,9 +1417,9 @@ _annotation = tag_class( + "build_script_data_select": attr.string_list_dict( + doc = "A list of labels to add to a crate's `cargo_build_script::data` attribute. Keys should be the platform triplet. Value should be a list of labels.", + ), +- # "build_script_deps": _relative_label_list( +- # doc = "A list of labels to add to a crate's `cargo_build_script::deps` attribute.", +- # ), ++ "build_script_deps": _relative_label_list( ++ doc = "A list of labels to add to a crate's `cargo_build_script::deps` attribute.", ++ ), + "build_script_env": attr.string_dict( + doc = "Additional environment variables to set on a crate's `cargo_build_script::env` attribute.", + ), +diff --git a/rs/private/annotations.bzl b/rs/private/annotations.bzl +--- a/rs/private/annotations.bzl ++++ b/rs/private/annotations.bzl +@@ -5,6 +5,7 @@ def _crate_annotation( + gen_build_script = "auto", + build_script_data = [], + build_script_data_select = {}, ++ build_script_deps = [], + build_script_env = {}, + build_script_env_select = {}, + build_script_tools = [], +@@ -28,6 +29,7 @@ def _crate_annotation( + gen_build_script = gen_build_script, + build_script_data = build_script_data, + build_script_data_select = build_script_data_select, ++ build_script_deps = build_script_deps, + build_script_env = build_script_env, + build_script_env_select = build_script_env_select, + build_script_tools = build_script_tools, +diff --git a/rs/private/repository_utils.bzl b/rs/private/repository_utils.bzl +--- a/rs/private/repository_utils.bzl ++++ b/rs/private/repository_utils.bzl +@@ -216,7 +216,7 @@ common_attrs = { + "additive_build_file": attr.label(), + "additive_build_file_content": attr.string(), + "gen_build_script": attr.string(), +- "build_script_deps": attr.label_list(default = []), ++ "build_script_deps": attr.string_list(default = []), + "build_script_deps_select": attr.string_list_dict(), + "build_script_data": attr.label_list(default = []), + "build_script_data_select": attr.string_list_dict(), diff --git a/scripts/stage_npm_packages.py b/scripts/stage_npm_packages.py index 2329d2ba62..38374f523d 100755 --- a/scripts/stage_npm_packages.py +++ b/scripts/stage_npm_packages.py @@ -101,6 +101,11 @@ def parse_args() -> argparse.Namespace: "--workflow-url", help="Optional workflow URL to reuse for native artifacts.", ) + parser.add_argument( + "--artifacts-dir", + type=Path, + help="Directory containing previously downloaded workflow artifacts.", + ) parser.add_argument( "--output-dir", type=Path, @@ -496,6 +501,7 @@ def main() -> int: vendor_temp_roots: list[Path] = [] vendor_src_by_components: dict[tuple[str, ...], Path] = {} artifacts_temp_root: Path | None = None + remove_artifacts_temp_root = False resolved_head_sha: str | None = None final_messages = [] @@ -506,10 +512,15 @@ def main() -> int: args.release_version, args.workflow_url ) print(f"Using native artifacts from {workflow_url}", flush=True) - artifacts_temp_root = Path( - tempfile.mkdtemp(prefix="npm-native-artifacts-", dir=runner_temp) - ) - print(f"Caching downloaded artifacts in {artifacts_temp_root}", flush=True) + if args.artifacts_dir is not None: + artifacts_temp_root = args.artifacts_dir.resolve() + artifacts_temp_root.mkdir(parents=True, exist_ok=True) + else: + artifacts_temp_root = Path( + tempfile.mkdtemp(prefix="npm-native-artifacts-", dir=runner_temp) + ) + remove_artifacts_temp_root = True + print(f"Using artifact cache at {artifacts_temp_root}", flush=True) for components in native_component_sets: vendor_temp_root = Path( tempfile.mkdtemp(prefix="npm-native-", dir=runner_temp) @@ -570,7 +581,7 @@ def main() -> int: if not args.keep_staging_dirs: for vendor_temp_root in vendor_temp_roots: shutil.rmtree(vendor_temp_root, ignore_errors=True) - if artifacts_temp_root is not None: + if remove_artifacts_temp_root and artifacts_temp_root is not None: shutil.rmtree(artifacts_temp_root, ignore_errors=True) for msg in final_messages: diff --git a/sdk/python/src/openai_codex/_goal.py b/sdk/python/src/openai_codex/_goal.py index dc1fe358c2..37a63fc28f 100644 --- a/sdk/python/src/openai_codex/_goal.py +++ b/sdk/python/src/openai_codex/_goal.py @@ -1,8 +1,12 @@ +import asyncio import queue import threading import time +from collections import deque from dataclasses import dataclass, field +from typing import AsyncIterator, Awaitable, Callable, Iterator +from .generated.notification_registry import notification_turn_id from .generated.v2_all import ( ThreadGoalClearedNotification, ThreadGoalStatus, @@ -10,8 +14,9 @@ from .generated.v2_all import ( Turn, TurnCompletedNotification, TurnStartedNotification, + TurnStatus, ) -from .models import Notification +from .models import Notification, UnknownNotification class _GoalStreamClosed(Exception): @@ -45,10 +50,16 @@ class _GoalOperationState: _notifications: queue.Queue[Notification | BaseException] = field(default_factory=queue.Queue) _failure: BaseException | None = None _finished: bool = False + _turn_routing_active: bool = False - def observe(self, notification: Notification) -> None: + def observe(self, notification: Notification) -> bool: payload = notification.payload with self._condition: + if not self._turn_routing_active and not isinstance( + payload, + ThreadGoalClearedNotification | ThreadGoalUpdatedNotification, + ): + return False if isinstance(payload, TurnStartedNotification): if self.logical_turn_id is None: self.logical_turn_id = payload.turn.id @@ -73,6 +84,12 @@ class _GoalOperationState: self._finished = True self._condition.notify_all() self._notifications.put(notification) + return True + + def activate_turn_routing(self) -> None: + """Accept physical turns after the previous stored goal is cleared.""" + with self._condition: + self._turn_routing_active = True def wait_for_start(self, timeout: float) -> str | None: """Wait for the runtime-generated first turn without consuming its event.""" @@ -121,16 +138,18 @@ class _GoalOperationState: with self._condition: self.interrupted = True self.interrupt_requested = False + self._condition.notify_all() def cancel_interrupt(self) -> None: with self._condition: self.interrupt_requested = False + self._condition.notify_all() - def explicit_interrupt(self, status: ThreadGoalStatus | None) -> bool: + def explicit_interrupt(self) -> bool: with self._condition: - return self.interrupted or ( - self.interrupt_requested and status == ThreadGoalStatus.paused - ) + while self.interrupt_requested: + self._condition.wait() + return self.interrupted def active_turn(self, *, after: str | None = None) -> str | None: """Wait for the current turn, or return None once the goal has ended.""" @@ -151,6 +170,279 @@ class _GoalOperationState: with self._condition: return self.current_turn_id + def resolve_active_turn(self, expected: str, active: str) -> None: + """Adopt a server-reported active id when routed state is still stale.""" + with self._condition: + if self.current_turn_id in {None, expected}: + self.current_turn_id = active + self._condition.notify_all() + + def turn_for_interrupt(self) -> str | None: + """Return an active or stale turn id that can resolve rollover races.""" + with self._condition: + if self.current_turn_id is not None: + return self.current_turn_id + if self.completed_turn is not None: + return self.completed_turn.id + if self.started_turn is not None: + return self.started_turn.id + return None + def wake_notification_reader(self) -> None: """Release a reader blocked after its stream has been closed.""" self._notifications.put(_GoalStreamClosed()) + + +def _logical_notification(notification: Notification, logical_turn_id: str) -> Notification: + """Return a copy whose turn metadata uses the logical operation id.""" + payload = notification.payload + if isinstance(payload, UnknownNotification): + params = dict(payload.params) + if isinstance(params.get("turnId"), str): + params["turnId"] = logical_turn_id + turn = params.get("turn") + if isinstance(turn, dict) and isinstance(turn.get("id"), str): + params["turn"] = {**turn, "id": logical_turn_id} + return Notification(notification.method, UnknownNotification(params)) + + turn_id = notification_turn_id(payload) + if turn_id is None: + return notification + if hasattr(payload, "turn_id"): + return Notification( + notification.method, + payload.model_copy(update={"turn_id": logical_turn_id}), + ) + if hasattr(payload, "turn"): + logical_turn = payload.turn.model_copy(update={"id": logical_turn_id}) + return Notification( + notification.method, + payload.model_copy(update={"turn": logical_turn}), + ) + return notification + + +def _logical_completion( + completed: TurnCompletedNotification, + *, + logical_turn_id: str, + started: Turn | None, + interrupted: bool, +) -> TurnCompletedNotification: + """Coalesce the final physical completion into one logical completion.""" + final_turn = completed.turn + started_at = started.started_at if started is not None else final_turn.started_at + duration_ms = final_turn.duration_ms + if started_at is not None and final_turn.completed_at is not None: + duration_ms = max(0, final_turn.completed_at - started_at) * 1000 + updates: dict[str, object] = { + "id": logical_turn_id, + "started_at": started_at, + "duration_ms": duration_ms, + } + if interrupted: + updates["status"] = TurnStatus.interrupted + return completed.model_copy(update={"turn": final_turn.model_copy(update=updates)}) + + +@dataclass(slots=True) +class _GoalStreamCursor: + """Consume physical goal events as one ordered logical turn stream.""" + + state: _GoalOperationState + started: Turn | None = None + last_completed: TurnCompletedNotification | None = None + failed_completion: TurnCompletedNotification | None = None + status: ThreadGoalStatus | None = None + active: bool = False + cleared: bool = False + + def process(self, notification: Notification) -> tuple[list[Notification], bool]: + logical_turn_id = self.state.logical_turn_id + if logical_turn_id is None: + raise RuntimeError("goal operation has not been bound to a logical turn id") + + payload = notification.payload + if isinstance(payload, TurnStartedNotification): + self.active = True + if self.started is not None: + return [], False + self.started = payload.turn + return [_logical_notification(notification, logical_turn_id)], False + + if isinstance(payload, TurnCompletedNotification): + self.active = False + self.last_completed = payload + if payload.turn.status == TurnStatus.interrupted: + return [ + self._completion( + notification.method, + self.failed_completion or payload, + ) + ], True + if payload.turn.status == TurnStatus.failed: + self.failed_completion = payload + if self.cleared or _terminal_goal_status(self.status): + self.state.finish() + return [self._completion(notification.method, payload)], True + return [], False + if self.status is None and not self.cleared: + raise RuntimeError( + "the connected Codex runtime did not activate goal mode for this turn" + ) + if self.cleared or _terminal_goal_status(self.status): + self.state.finish() + return [ + self._completion( + notification.method, + self.failed_completion or payload, + ) + ], True + return [], False + + events = [_logical_notification(notification, logical_turn_id)] + if isinstance(payload, ThreadGoalUpdatedNotification): + self.status = payload.goal.status + if self.status == ThreadGoalStatus.active: + self.cleared = False + events = [] + elif isinstance(payload, ThreadGoalClearedNotification): + self.cleared = True + events = [] + + if ( + not self.active + and self.last_completed is not None + and (self.cleared or _terminal_goal_status(self.status)) + ): + self.state.finish() + events.append( + self._completion( + "turn/completed", + self.failed_completion or self.last_completed, + ) + ) + return events, True + return events, False + + def _completion( + self, + method: str, + payload: TurnCompletedNotification, + ) -> Notification: + logical_turn_id = self.state.logical_turn_id + if logical_turn_id is None: + raise RuntimeError("goal operation has not been bound to a logical turn id") + return Notification( + method, + _logical_completion( + payload, + logical_turn_id=logical_turn_id, + started=self.started, + interrupted=self.state.explicit_interrupt(), + ), + ) + + +@dataclass(slots=True) +class _GoalNotificationStream(Iterator[Notification]): + """Closeable synchronous view of one logical goal operation.""" + + state: _GoalOperationState + next_notification: Callable[[], Notification] + unregister: Callable[[], None] + cancel_goal: Callable[[], None] + _cursor: _GoalStreamCursor = field(init=False) + _pending: deque[Notification] = field(default_factory=deque) + _closed: bool = False + + def __post_init__(self) -> None: + self._cursor = _GoalStreamCursor(self.state) + + def __iter__(self) -> "_GoalNotificationStream": + return self + + def __next__(self) -> Notification: + if self._closed: + raise StopIteration + try: + while not self._pending: + notification = self.next_notification() + events, completed = self._cursor.process(notification) + self._pending.extend(events) + if completed: + self._finish() + return self._pending.popleft() + except _GoalStreamClosed: + self.close() + raise StopIteration from None + except KeyboardInterrupt: + self.cancel_goal() + self.close() + raise + except BaseException: + self.close() + raise + + def _finish(self) -> None: + if self._closed: + return + self.state.finish() + self.state.wake_notification_reader() + self.unregister() + self._closed = True + + def close(self) -> None: + self._finish() + + +@dataclass(slots=True) +class _AsyncGoalNotificationStream(AsyncIterator[Notification]): + """Closeable asynchronous view of one logical goal operation.""" + + state: _GoalOperationState + next_notification: Callable[[], Awaitable[Notification]] + unregister: Callable[[], None] + cancel_goal: Callable[[], Awaitable[None]] + _cursor: _GoalStreamCursor = field(init=False) + _pending: deque[Notification] = field(default_factory=deque) + _closed: bool = False + + def __post_init__(self) -> None: + self._cursor = _GoalStreamCursor(self.state) + + def __aiter__(self) -> "_AsyncGoalNotificationStream": + return self + + async def __anext__(self) -> Notification: + if self._closed: + raise StopAsyncIteration + try: + while not self._pending: + notification = await self.next_notification() + events, completed = self._cursor.process(notification) + self._pending.extend(events) + if completed: + self._finish() + return self._pending.popleft() + except _GoalStreamClosed: + await self.aclose() + raise StopAsyncIteration from None + except asyncio.CancelledError: + await self.cancel_goal() + await self.aclose() + raise + except BaseException: + await self.aclose() + raise + + def _finish(self) -> None: + if self._closed: + return + self.state.finish() + self.state.wake_notification_reader() + self.unregister() + self._closed = True + + async def aclose(self) -> None: + self._finish() diff --git a/sdk/python/src/openai_codex/_message_router.py b/sdk/python/src/openai_codex/_message_router.py index e71544a0ef..c979c8c8dd 100644 --- a/sdk/python/src/openai_codex/_message_router.py +++ b/sdk/python/src/openai_codex/_message_router.py @@ -121,10 +121,20 @@ class MessageRouter: def register_goal(self, thread_id: str) -> _GoalOperationState: """Register one thread-scoped logical goal operation before it starts.""" state = _GoalOperationState(thread_id=thread_id) + state.activate_turn_routing() + return self._register_goal(state) + + def reserve_goal(self, thread_id: str) -> _GoalOperationState: + """Reserve a thread route without accepting physical turns yet.""" + return self._register_goal(_GoalOperationState(thread_id=thread_id)) + + def _register_goal(self, state: _GoalOperationState) -> _GoalOperationState: with self._lock: - if thread_id in self._goal_operations: - raise RuntimeError(f"thread {thread_id!r} already has an active goal operation") - self._goal_operations[thread_id] = state + if state.thread_id in self._goal_operations: + raise RuntimeError( + f"thread {state.thread_id!r} already has an active goal operation" + ) + self._goal_operations[state.thread_id] = state return state def unregister_goal(self, state: _GoalOperationState) -> None: @@ -133,6 +143,11 @@ class MessageRouter: if self._goal_operations.get(state.thread_id) is state: self._goal_operations.pop(state.thread_id) + def has_goal(self, thread_id: str) -> bool: + """Return whether a logical goal operation owns this thread route.""" + with self._lock: + return thread_id in self._goal_operations + def route_response(self, msg: dict[str, JsonValue]) -> None: """Deliver a JSON-RPC response or error to its request waiter.""" @@ -181,10 +196,10 @@ class MessageRouter: if goal_state is not None and ( turn_id is not None or notification.method.startswith("thread/goal/") ): - goal_state.observe(notification) - if goal_state.is_finished(): - self.unregister_goal(goal_state) - return + if goal_state.observe(notification): + if goal_state.is_finished(): + self.unregister_goal(goal_state) + return if turn_id is None: self._global_notifications.put(notification) return diff --git a/sdk/python/src/openai_codex/async_client.py b/sdk/python/src/openai_codex/async_client.py index ff12f16884..07e7d0d405 100644 --- a/sdk/python/src/openai_codex/async_client.py +++ b/sdk/python/src/openai_codex/async_client.py @@ -1,7 +1,9 @@ from __future__ import annotations import asyncio +import threading from collections.abc import Iterator +from concurrent.futures import Future from typing import AsyncIterator, Callable, ParamSpec, TypeVar from pydantic import BaseModel @@ -227,6 +229,58 @@ class AsyncCodexClient: """Pause the active goal through the wrapped sync client.""" return await self._call_sync(self._sync.pause_goal, thread_id) + async def cancel_goal_operation(self, state: _GoalOperationState) -> None: + """Stop continuation work after a logical goal operation is cancelled.""" + await self._call_sync(self._sync.cancel_goal_operation, state) + + async def start_goal_operation( + self, + thread_id: str, + objective: str, + ) -> tuple[_GoalOperationState, str]: + """Start a logical goal through the wrapped sync client.""" + operation: Future[tuple[_GoalOperationState, str]] = Future() + + def start_operation() -> None: + try: + operation.set_result(self._sync.start_goal_operation(thread_id, objective)) + except BaseException as exc: + operation.set_exception(exc) + + worker = threading.Thread( + target=start_operation, + name="codex-goal-start", + daemon=True, + ) + worker.start() + try: + return await asyncio.shield(asyncio.wrap_future(operation)) + except asyncio.CancelledError: + + def cleanup_cancelled_start( + completed: Future[tuple[_GoalOperationState, str]], + ) -> None: + try: + state, _ = completed.result() + except BaseException: + return + + def stop_cancelled_goal() -> None: + try: + self._sync.cancel_goal_operation(state) + finally: + state.finish() + self._sync.unregister_goal_operation(state) + + threading.Thread( + target=stop_cancelled_goal, + name="codex-goal-start-cleanup", + daemon=True, + ).start() + + operation.add_done_callback(cleanup_cancelled_start) + raise + async def turn_start( self, thread_id: str, diff --git a/sdk/python/src/openai_codex/client.py b/sdk/python/src/openai_codex/client.py index 951f4f95f0..ab5390ae52 100644 --- a/sdk/python/src/openai_codex/client.py +++ b/sdk/python/src/openai_codex/client.py @@ -1,10 +1,13 @@ import json import os +import re import subprocess import threading import uuid +from _thread import LockType from collections import deque -from dataclasses import dataclass +from contextlib import contextmanager +from dataclasses import dataclass, field from pathlib import Path from typing import Callable, Iterator, TypeVar @@ -13,7 +16,7 @@ from pydantic import BaseModel from ._goal import _GoalOperationState from ._message_router import MessageRouter from ._version import __version__ as SDK_VERSION -from .errors import CodexError, TransportClosedError +from .errors import CodexError, InvalidRequestError, TransportClosedError from .generated.notification_registry import NOTIFICATION_MODELS from .generated.v2_all import ( AccountLoginCompletedNotification, @@ -23,6 +26,7 @@ from .generated.v2_all import ( ChatgptLoginAccountResponse, GetAccountParams as V2GetAccountParams, GetAccountResponse, + IdleThreadStatus, LoginAccountParams as V2LoginAccountParams, LoginAccountResponse, LogoutAccountResponse, @@ -61,6 +65,18 @@ from .retry import retry_on_overload ModelT = TypeVar("ModelT", bound=BaseModel) ApprovalHandler = Callable[[str, JsonObject | None], JsonObject] RUNTIME_PKG_NAME = "openai-codex-cli-bin" +_GOAL_START_TIMEOUT_S = 30.0 + + +@dataclass(slots=True) +class _ThreadStartLock: + lock: LockType = field(default_factory=threading.Lock) + users: int = 0 + + +def _active_turn_id_from_error(exc: InvalidRequestError) -> str | None: + match = re.search(r" but found `?([^`]+)`?$", exc.message) + return match.group(1) if match is not None else None def _params_dict( @@ -205,6 +221,8 @@ class CodexClient: self._approval_handler = approval_handler or self._default_approval_handler self._proc: subprocess.Popen[str] | None = None self._lock = threading.Lock() + self._thread_start_locks_guard = threading.Lock() + self._thread_start_locks: dict[str, _ThreadStartLock] = {} self._router = MessageRouter() self._stderr_lines: deque[str] = deque(maxlen=400) self._stderr_thread: threading.Thread | None = None @@ -360,6 +378,10 @@ class CodexClient: """Register a private thread-scoped route for a logical goal turn.""" return self._router.register_goal(thread_id) + def reserve_goal_operation(self, thread_id: str) -> _GoalOperationState: + """Reserve a private thread route before replacing its stored goal.""" + return self._router.reserve_goal(thread_id) + def unregister_goal_operation(self, state: _GoalOperationState) -> None: """Release routing state for one logical goal turn.""" self._router.unregister_goal(state) @@ -499,6 +521,84 @@ class CodexClient: """Pause the active goal used by a logical goal turn.""" return self.thread_goal_set(thread_id, status=ThreadGoalStatus.paused) + def cancel_goal_operation(self, state: _GoalOperationState) -> None: + """Best-effort cleanup after a logical goal operation is cancelled.""" + try: + self.pause_goal(state.thread_id) + except Exception: + pass + self._interrupt_goal_operation(state) + + def _interrupt_goal_operation(self, state: _GoalOperationState) -> None: + turn_id = state.turn_for_interrupt() + if turn_id is None: + return + try: + self.turn_interrupt(state.thread_id, turn_id) + except InvalidRequestError as exc: + if not exc.message.startswith("expected active turn id"): + return + next_turn_id = _active_turn_id_from_error(exc) or state.current_turn() + if next_turn_id is None or next_turn_id == turn_id: + return + try: + self.turn_interrupt(state.thread_id, next_turn_id) + except Exception: + pass + except Exception: + pass + + def start_goal_operation( + self, + thread_id: str, + objective: str, + ) -> tuple[_GoalOperationState, str]: + """Start a logical goal and wait for its runtime-generated first turn.""" + with self._thread_start_lock(thread_id): + return self._start_goal_operation(thread_id, objective) + + def _start_goal_operation( + self, + thread_id: str, + objective: str, + ) -> tuple[_GoalOperationState, str]: + thread = self.thread_read(thread_id).thread + if not isinstance(thread.status.root, IdleThreadStatus): + raise InvalidRequestError( + -32600, + f"thread must be idle before starting a goal: {thread_id}", + ) + if thread.ephemeral or thread.path is None: + raise InvalidRequestError( + -32600, + f"thread must be persisted before starting a goal: {thread_id}", + ) + + state = self.reserve_goal_operation(thread_id) + activated = False + try: + self.thread_goal_clear(thread_id) + state.activate_turn_routing() + self.thread_goal_set( + thread_id, + objective=objective, + status=ThreadGoalStatus.active, + ) + activated = True + turn_id = state.wait_for_start(_GOAL_START_TIMEOUT_S) + if turn_id is None: + raise CodexError( + "timed out waiting for goal turn to start after " + f"{int(_GOAL_START_TIMEOUT_S)} seconds" + ) + return state, turn_id + except BaseException as exc: + if activated or not isinstance(exc, InvalidRequestError): + self.cancel_goal_operation(state) + state.finish() + self.unregister_goal_operation(state) + raise + def turn_start( self, thread_id: str, @@ -506,14 +606,37 @@ class CodexClient: params: V2TurnStartParams | JsonObject | None = None, ) -> TurnStartResponse: """Start a turn and register its notification queue as early as possible.""" - payload = { - **_params_dict(params), - "threadId": thread_id, - "input": self._normalize_input_items(input_items), - } - started = self.request("turn/start", payload, response_model=TurnStartResponse) - self.register_turn_notifications(started.turn.id) - return started + with self._thread_start_lock(thread_id): + if self._router.has_goal(thread_id): + raise InvalidRequestError( + -32600, + f"thread has an active goal operation: {thread_id}", + ) + payload = { + **_params_dict(params), + "threadId": thread_id, + "input": self._normalize_input_items(input_items), + } + started = self.request("turn/start", payload, response_model=TurnStartResponse) + self.register_turn_notifications(started.turn.id) + return started + + @contextmanager + def _thread_start_lock(self, thread_id: str) -> Iterator[None]: + with self._thread_start_locks_guard: + entry = self._thread_start_locks.get(thread_id) + if entry is None: + entry = _ThreadStartLock() + self._thread_start_locks[thread_id] = entry + entry.users += 1 + try: + with entry.lock: + yield + finally: + with self._thread_start_locks_guard: + entry.users -= 1 + if entry.users == 0: + self._thread_start_locks.pop(thread_id, None) def turn_interrupt(self, thread_id: str, turn_id: str) -> TurnInterruptResponse: return self.request( diff --git a/sdk/python/tests/test_app_server_goal_operations.py b/sdk/python/tests/test_app_server_goal_operations.py new file mode 100644 index 0000000000..c280c7522c --- /dev/null +++ b/sdk/python/tests/test_app_server_goal_operations.py @@ -0,0 +1,90 @@ +from app_server_harness import ( + AppServerHarness, + ev_assistant_message, + ev_completed, + ev_function_call, + ev_response_created, + sse, +) +from app_server_helpers import agent_message_texts + +from openai_codex import Codex +from openai_codex._goal import _GoalNotificationStream +from openai_codex._run import _collect_turn_result +from openai_codex.generated.notification_registry import notification_turn_id +from openai_codex.generated.v2_all import TurnStatus + + +def test_private_goal_operation_coalesces_runtime_continuations(tmp_path) -> None: + """The private engine should expose automatic continuations as one turn.""" + with AppServerHarness(tmp_path) as harness: + harness.responses.enqueue_assistant_message( + "Initial pass complete.", + response_id="goal-initial", + ) + harness.responses.enqueue_sse( + sse( + [ + ev_response_created("goal-complete-tool"), + ev_function_call( + "call-goal-complete", + "update_goal", + '{"status":"complete"}', + ), + ev_completed("goal-complete-tool"), + ] + ) + ) + harness.responses.enqueue_sse( + sse( + [ + ev_response_created("goal-final"), + ev_assistant_message("msg-goal-final", "Goal complete."), + ev_completed("goal-final"), + ] + ) + ) + + with Codex(config=harness.app_server_config()) as codex: + thread = codex.thread_start() + state, turn_id = codex._client.start_goal_operation( # noqa: SLF001 + thread.id, + "Improve benchmark coverage", + ) + stream = _GoalNotificationStream( + state, + state.next_notification, + lambda: codex._client.unregister_goal_operation(state), # noqa: SLF001 + lambda: codex._client.cancel_goal_operation(state), # noqa: SLF001 + ) + events = list(stream) + result = _collect_turn_result(iter(events), turn_id=turn_id) + routes = codex._client._router._goal_operations.copy() # noqa: SLF001 + requests = harness.responses.wait_for_requests(3) + + lifecycle = [event.method for event in events if event.method.startswith("turn/")] + routed_ids = [ + routed_id + for event in events + if (routed_id := notification_turn_id(event.payload)) is not None + ] + assert { + "lifecycle": lifecycle, + "routed_ids": routed_ids, + "result": (result.id, result.status, result.final_response), + "messages": agent_message_texts(events), + "request_count": len(requests), + "objective_reached_model": ( + "\nImprove benchmark coverage\n" + in "\n".join(requests[0].message_input_texts("user")) + ), + "routes_after_completion": routes, + } == { + "lifecycle": ["turn/started", "turn/completed"], + "routed_ids": [turn_id] * len(routed_ids), + "result": (turn_id, TurnStatus.completed, "Goal complete."), + "messages": ["Initial pass complete.", "Goal complete."], + "request_count": 3, + "objective_reached_model": True, + "routes_after_completion": {}, + }