From a770e5b8470d3320eb53a56a286ea4a0a70a1f59 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Tue, 9 Jun 2026 08:29:32 -0700 Subject: [PATCH 01/17] [codex] Ignore pending PR review comments (#27080) ## Why The PR babysitter could surface inline comments from a GitHub review that was still in the `PENDING` state. That allowed Codex to start acting on feedback before the reviewer submitted it. ## What changed - Correlate inline comments with their parent review and ignore pending reviews and their comments. - Remove pending review IDs from saved watcher state so the feedback surfaces normally after publication. - Update the skill instructions and add regression coverage for the draft-to-published transition. ## Validation - `python3 -m pytest .codex/skills/babysit-pr/scripts/test_gh_pr_watch.py` - Skill package validation with `quick_validate.py` - Live verification on #26835: the draft comment stayed hidden and surfaced after the review was submitted. --- .codex/skills/babysit-pr/SKILL.md | 8 ++- .codex/skills/babysit-pr/agents/openai.yaml | 2 +- .../babysit-pr/references/github-api-notes.md | 3 + .../skills/babysit-pr/scripts/gh_pr_watch.py | 28 +++++++- .../babysit-pr/scripts/test_gh_pr_watch.py | 68 +++++++++++++++++++ 5 files changed, 103 insertions(+), 6 deletions(-) diff --git a/.codex/skills/babysit-pr/SKILL.md b/.codex/skills/babysit-pr/SKILL.md index 68fc7278ba..36c6bd093a 100644 --- a/.codex/skills/babysit-pr/SKILL.md +++ b/.codex/skills/babysit-pr/SKILL.md @@ -28,7 +28,7 @@ Accept any of the following: 3. Inspect the `actions` list in the JSON response. 4. If `diagnose_ci_failure` is present, inspect failed run logs and classify the failure. 5. If the failure is likely caused by the current branch, patch code locally, commit, and push. Do not patch random flaky tests, CI infrastructure, dependency outages, runner issues, or other failures that are unrelated to the branch. -6. If `process_review_comment` is present, inspect surfaced review items and decide whether to address them. +6. If `process_review_comment` is present, inspect surfaced published review items and decide whether to address them. 7. If a review item is actionable and correct, patch code locally, commit, push, and then resolve the associated review thread only when allowed by the GitHub state mutation policy below. 8. Do not post replies to human-authored review comments/threads unless the user explicitly confirms the exact response. If a human review item is non-actionable, already addressed, or not valid, surface the item and recommended response to the user instead of replying on GitHub. 9. If the failure is likely flaky/unrelated and `retry_failed_checks` is present, rerun failed jobs with `--retry-failed-now`. @@ -92,9 +92,13 @@ The watcher surfaces review items from: - Inline review comments - Review submissions (COMMENT / APPROVED / CHANGES_REQUESTED) +Only act on published feedback. Ignore review submissions in GitHub's `PENDING` state and inline +comments attached to those pending reviews. Do not mark pending review feedback as seen; it should +be eligible to surface after the reviewer submits the review. + It intentionally surfaces Codex reviewer bot feedback (for example comments/reviews from `chatgpt-codex-connector[bot]`) in addition to human reviewer feedback. Most unrelated bot noise should still be ignored. For safety, the watcher only auto-surfaces trusted human review authors (for example repo OWNER/MEMBER/COLLABORATOR, plus the authenticated operator) and approved review bots such as Codex. -On a fresh watcher state file, existing pending review feedback may be surfaced immediately (not only comments that arrive after monitoring starts). This is intentional so already-open review comments are not missed. +On a fresh watcher state file, existing unaddressed published review feedback may be surfaced immediately (not only comments that arrive after monitoring starts). This is intentional so already-open review comments are not missed. When you agree with a comment and it is actionable: diff --git a/.codex/skills/babysit-pr/agents/openai.yaml b/.codex/skills/babysit-pr/agents/openai.yaml index c6946cf8c0..e07637b903 100644 --- a/.codex/skills/babysit-pr/agents/openai.yaml +++ b/.codex/skills/babysit-pr/agents/openai.yaml @@ -1,4 +1,4 @@ interface: display_name: "PR Babysitter" short_description: "Watch PR review comments, CI, and merge conflicts" - default_prompt: "Babysit the current PR: monitor reviewer comments, CI, and merge-conflict status (prefer the watcher’s --watch mode for live monitoring); surface new review feedback before acting on CI or mergeability work, fix valid issues, push updates, and rerun flaky failures up to 3 times. Do not post replies to human-authored review comments unless the user explicitly confirms the exact response. Do not patch unrelated flaky tests, CI infrastructure, dependency outages, runner issues, or other failures that are not caused by the branch. Keep exactly one watcher session active for the PR (do not leave duplicate --watch terminals running). If you pause monitoring to patch review/CI feedback, restart --watch yourself immediately after the push in the same turn. If a watcher is still running and no strict stop condition has been reached, the task is still in progress: keep consuming watcher output and sending progress updates instead of ending the turn. Do not treat a green + mergeable PR as a terminal stop while it is still open; continue polling autonomously after any push/rerun so newly posted review comments are surfaced until a strict terminal stop condition is reached or the user interrupts." + default_prompt: "Babysit the current PR: monitor published reviewer comments, CI, and merge-conflict status (prefer the watcher’s --watch mode for live monitoring); ignore unpublished comments in pending GitHub reviews; surface new published review feedback before acting on CI or mergeability work, fix valid issues, push updates, and rerun flaky failures up to 3 times. Do not post replies to human-authored review comments unless the user explicitly confirms the exact response. Do not patch unrelated flaky tests, CI infrastructure, dependency outages, runner issues, or other failures that are not caused by the branch. Keep exactly one watcher session active for the PR (do not leave duplicate --watch terminals running). If you pause monitoring to patch review/CI feedback, restart --watch yourself immediately after the push in the same turn. If a watcher is still running and no strict stop condition has been reached, the task is still in progress: keep consuming watcher output and sending progress updates instead of ending the turn. Do not treat a green + mergeable PR as a terminal stop while it is still open; continue polling autonomously after any push/rerun so newly posted review comments are surfaced until a strict terminal stop condition is reached or the user interrupts." diff --git a/.codex/skills/babysit-pr/references/github-api-notes.md b/.codex/skills/babysit-pr/references/github-api-notes.md index 8c0a7c8a54..645e7453c1 100644 --- a/.codex/skills/babysit-pr/references/github-api-notes.md +++ b/.codex/skills/babysit-pr/references/github-api-notes.md @@ -44,6 +44,9 @@ Reruns only failed jobs (and dependencies) for a workflow run. - Review submissions: - `gh api repos/{owner}/{repo}/pulls//reviews?per_page=100` +Use each inline comment's `pull_request_review_id` to find its parent review. Ignore parent reviews +whose `state` is `PENDING`, along with their inline comments, until the review is submitted. + ## JSON fields consumed by the watcher ### `gh pr view` diff --git a/.codex/skills/babysit-pr/scripts/gh_pr_watch.py b/.codex/skills/babysit-pr/scripts/gh_pr_watch.py index face4e6981..f6398373b7 100755 --- a/.codex/skills/babysit-pr/scripts/gh_pr_watch.py +++ b/.codex/skills/babysit-pr/scripts/gh_pr_watch.py @@ -452,11 +452,14 @@ def normalize_issue_comments(items): return out -def normalize_review_comments(items): +def normalize_review_comments(items, review_states): out = [] for item in items: if not isinstance(item, dict): continue + review_id = str(item.get("pull_request_review_id") or "") + if review_states.get(review_id) == "PENDING": + continue line = item.get("line") if line is None: line = item.get("original_line") @@ -481,6 +484,8 @@ def normalize_reviews(items): for item in items: if not isinstance(item, dict): continue + if str(item.get("state") or "").upper() == "PENDING": + continue out.append( { "kind": "review", @@ -534,16 +539,33 @@ def fetch_new_review_items(pr, state, fresh_state, authenticated_login=None): review_payload = gh_api_list_paginated(endpoints["review"], repo=repo) issue_items = normalize_issue_comments(issue_payload) - review_comment_items = normalize_review_comments(review_comment_payload) + review_states = { + str(item.get("id")): str(item.get("state") or "").upper() + for item in review_payload + if isinstance(item, dict) and item.get("id") not in (None, "") + } + pending_review_ids = { + review_id for review_id, review_state in review_states.items() if review_state == "PENDING" + } + pending_review_comment_ids = { + str(item.get("id")) + for item in review_comment_payload + if isinstance(item, dict) + and item.get("id") not in (None, "") + and str(item.get("pull_request_review_id") or "") in pending_review_ids + } + review_comment_items = normalize_review_comments(review_comment_payload, review_states) review_items = normalize_reviews(review_payload) all_items = issue_items + review_comment_items + review_items seen_issue = {str(x) for x in state.get("seen_issue_comment_ids") or []} seen_review_comment = {str(x) for x in state.get("seen_review_comment_ids") or []} seen_review = {str(x) for x in state.get("seen_review_ids") or []} + seen_review_comment.difference_update(pending_review_comment_ids) + seen_review.difference_update(pending_review_ids) # On a brand-new state file, surface existing review activity instead of - # silently treating it as seen. This avoids missing already-pending review + # silently treating it as seen. This avoids missing already-published review # feedback when monitoring starts after comments were posted. new_items = [] diff --git a/.codex/skills/babysit-pr/scripts/test_gh_pr_watch.py b/.codex/skills/babysit-pr/scripts/test_gh_pr_watch.py index b636ee4c55..76293b56f8 100644 --- a/.codex/skills/babysit-pr/scripts/test_gh_pr_watch.py +++ b/.codex/skills/babysit-pr/scripts/test_gh_pr_watch.py @@ -118,6 +118,74 @@ def test_recommend_actions_prioritizes_review_comments(): ] +def test_pending_review_feedback_surfaces_only_after_publication(monkeypatch): + state = { + "seen_review_comment_ids": ["20"], + "seen_review_ids": ["10"], + } + review = { + "id": 10, + "user": {"login": "octocat"}, + "author_association": "MEMBER", + "state": "PENDING", + "body": "Please rename this.", + "created_at": "2026-06-08T10:00:00Z", + "submitted_at": None, + "html_url": "https://github.com/openai/codex/pull/123#pullrequestreview-10", + } + review_comment = { + "id": 20, + "pull_request_review_id": 10, + "user": {"login": "octocat"}, + "author_association": "MEMBER", + "body": "Please rename this.", + "created_at": "2026-06-08T10:00:00Z", + "path": "src/example.rs", + "line": 7, + "html_url": "https://github.com/openai/codex/pull/123#discussion_r20", + } + + def fake_list(endpoint, **kwargs): + if endpoint.endswith("/issues/123/comments"): + return [] + if endpoint.endswith("/pulls/123/comments"): + return [review_comment] + if endpoint.endswith("/pulls/123/reviews"): + return [review] + raise AssertionError(f"unexpected endpoint: {endpoint}") + + monkeypatch.setattr(gh_pr_watch, "gh_api_list_paginated", fake_list) + + assert ( + gh_pr_watch.fetch_new_review_items( + sample_pr(), + state, + fresh_state=True, + authenticated_login="octocat", + ) + == [] + ) + assert state["seen_review_comment_ids"] == [] + assert state["seen_review_ids"] == [] + + review["state"] = "COMMENTED" + review["submitted_at"] = "2026-06-08T10:05:00Z" + + published_items = gh_pr_watch.fetch_new_review_items( + sample_pr(), + state, + fresh_state=False, + authenticated_login="octocat", + ) + + assert {(item["kind"], item["id"]) for item in published_items} == { + ("review", "10"), + ("review_comment", "20"), + } + assert state["seen_review_comment_ids"] == ["20"] + assert state["seen_review_ids"] == ["10"] + + def test_run_watch_keeps_polling_open_ready_to_merge_pr(monkeypatch): sleeps = [] events = [] From 18ce671fed526be9033907bd88a3a63c6888bbf4 Mon Sep 17 00:00:00 2001 From: Zanie Blue Date: Tue, 9 Jun 2026 10:50:13 -0500 Subject: [PATCH 02/17] Avoid no-op backfill state writes (#26420) ## Summary - avoid acquiring SQLite's writer slot when the singleton backfill row already exists - preserve race-safe repair when the row is missing - add regressions for writer contention and missing-row repair ## Why State runtime initialization and backfill-state reads previously executed `INSERT ... ON CONFLICT DO NOTHING` even in the steady state. SQLite still enters the writer path for that statement, so TUI and app-server startup could wait behind another writer for up to the configured five-second busy timeout. ## Validation - `just test -p codex-state` (134 tests passed) - `just fix -p codex-state` - `just fmt` --- codex-rs/state/src/runtime.rs | 10 +++++ codex-rs/state/src/runtime/backfill.rs | 57 ++++++++++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index b08545d7a6..5ae41bcd8b 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -384,6 +384,16 @@ async fn open_sqlite( pub(super) async fn ensure_backfill_state_row_in_pool( pool: &sqlx::SqlitePool, ) -> anyhow::Result<()> { + // Eagerly check if the operation would have no effect to avoid blocking waiting for a SQLite + // writer for no reason in the hot startup path. + if sqlx::query_scalar::<_, i64>("SELECT 1 FROM backfill_state WHERE id = 1") + .fetch_optional(pool) + .await? + .is_some() + { + return Ok(()); + } + sqlx::query( r#" INSERT INTO backfill_state (id, status, last_watermark, last_success_at, updated_at) diff --git a/codex-rs/state/src/runtime/backfill.rs b/codex-rs/state/src/runtime/backfill.rs index 77ab57fced..3bfa74fb9a 100644 --- a/codex-rs/state/src/runtime/backfill.rs +++ b/codex-rs/state/src/runtime/backfill.rs @@ -110,9 +110,11 @@ WHERE id = 1 #[cfg(test)] mod tests { use super::StateRuntime; + use super::base_sqlite_options; use super::test_support::unique_temp_dir; use chrono::Utc; use pretty_assertions::assert_eq; + use sqlx::Connection; #[tokio::test] async fn backfill_state_persists_progress_and_completion() { @@ -167,6 +169,61 @@ mod tests { let _ = tokio::fs::remove_dir_all(codex_home).await; } + #[tokio::test] + async fn get_backfill_state_succeeds_while_another_connection_holds_writer_slot() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("initialize runtime"); + let mut write_connection = sqlx::SqliteConnection::connect_with(&base_sqlite_options( + &crate::state_db_path(codex_home.as_path()), + )) + .await + .expect("open write connection"); + let write_transaction = write_connection + .begin_with("BEGIN IMMEDIATE") + .await + .expect("acquire write lock"); + + let state = runtime + .get_backfill_state() + .await + .expect("get backfill state"); + assert_eq!(state, crate::BackfillState::default()); + + write_transaction + .rollback() + .await + .expect("release write lock"); + let _ = tokio::fs::remove_dir_all(codex_home).await; + } + + #[tokio::test] + async fn get_backfill_state_repairs_a_missing_singleton_row() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("initialize runtime"); + sqlx::query("DELETE FROM backfill_state WHERE id = 1") + .execute(runtime.pool.as_ref()) + .await + .expect("delete backfill state row"); + + let state = runtime + .get_backfill_state() + .await + .expect("get repaired backfill state"); + assert_eq!(state, crate::BackfillState::default()); + let row_count = + sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM backfill_state WHERE id = 1") + .fetch_one(runtime.pool.as_ref()) + .await + .expect("count repaired backfill state rows"); + assert_eq!(row_count, 1); + + let _ = tokio::fs::remove_dir_all(codex_home).await; + } + #[tokio::test] async fn backfill_claim_is_singleton_until_stale_and_blocked_when_complete() { let codex_home = unique_temp_dir(); From 6a9a49b334e8081756934b0ce7d909234b53aac7 Mon Sep 17 00:00:00 2001 From: Zanie Blue Date: Tue, 9 Jun 2026 11:16:27 -0500 Subject: [PATCH 03/17] Avoid rereading rollout history during cold resume (#27031) ## Summary - reuse the history-bearing `StoredThread` loaded while probing for a running thread - avoid rereading and reparsing the rollout when that probe finds no active process - reload after shutting down a loaded thread because shutdown may flush newer rollout items - add a regression test that verifies cold resume performs one history-bearing store read ## Problem `thread/resume` first reads the persisted thread with history while checking whether the thread is already running. When no running process exists, cold resume currently falls through to `resume_thread_from_rollout`, which reads and parses the same history again. That duplicate work grows with rollout size and remains on the synchronous resume path even when the caller requests `excludeTurns`. ## Background The duplicate read was introduced by #24528, which fixed resume overrides for idle cached threads. To support resumes specified by rollout path, `resume_running_thread` began loading the stored thread with history so it could resolve the canonical thread ID and determine whether a cached `CodexThread` was already loaded. That history is needed when the loaded-thread path handles the request. On a cold miss, however, the function's boolean result could only report that no loaded thread handled the request. It discarded the history-bearing `StoredThread`, and the normal cold-resume path immediately loaded and parsed the same rollout again. This change preserves the idle cached-thread behavior from #24528 while allowing the cold-resume path to reuse the probe result. ## Performance I benchmarked real retained rollouts using isolated `CODEX_HOME` directories, explicit rollout paths, debug builds of the commit and its exact parent, and alternating parent/patch order. The table below uses `thread/resume` with `excludeTurns: true`; response payload sizes were identical. | Rollout size | Records | Parent median | Patch median | Median paired saving | | ---: | ---: | ---: | ---: | ---: | | 6 MB | 3,574 | 541 ms | 441 ms | 132 ms | | 30 MB | 15,220 | 1.505 s | 1.041 s | 701 ms | | 60 MB | 31,453 | 2.644 s | 1.742 s | 970 ms | | 149 MB | 100,874 | 10.506 s | 7.156 s | 3.350 s | | 559 MB | 259,734 | 27.759 s | 16.725 s | 9.836 s | The absolute saving increases with thread size, as expected when removing one complete JSONL history read and parse. Total resume time is also content-dependent, so the relationship is not perfectly linear. I also tested full-history resume with `excludeTurns: false`. The response payload was byte-identical between variants, and the same size-dependent improvement remained visible: | Rollout size | Parent median | Patch median | Median paired saving | | ---: | ---: | ---: | ---: | | 6 MB | 1.052 s | 904 ms | 270 ms | | 30 MB | 2.667 s | 1.762 s | 924 ms | | 60 MB | 8.464 s | 6.272 s | 3.680 s | | 149 MB | 26.719 s | 12.118 s | 14.601 s | | 559 MB | 40.359 s | 25.475 s | 16.590 s | ## Validation - `just test -p codex-app-server cold_thread_resume_reuses_non_local_history_probe` - `just fix -p codex-app-server -p codex-thread-store` - `just fmt` --- .../request_processors/thread_processor.rs | 44 +++++-- .../tests/suite/v2/remote_thread_store.rs | 118 ++++++++++++++++++ codex-rs/thread-store/src/in_memory.rs | 4 + 3 files changed, 155 insertions(+), 11 deletions(-) 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 4cc169f80f..ec3bd57b07 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -327,6 +327,17 @@ pub(crate) struct ThreadRequestProcessor { pub(super) skills_watcher: Arc, } +/// Outcome of trying to satisfy a resume request from an already loaded thread. +enum RunningThreadResumeResult { + /// The request was delegated to the loaded thread. + Handled, + /// No loaded thread handled the request. + /// + /// The optional stored thread contains the history-bearing probe that cold + /// resume can reuse instead of reading the rollout again. + NotRunning(Option>), +} + impl ThreadRequestProcessor { #[allow(clippy::too_many_arguments)] pub(crate) fn new( @@ -2419,7 +2430,7 @@ impl ThreadRequestProcessor { return Ok(()); } }; - match self + let stored_thread_from_running_probe = match self .resume_running_thread( &request_id, ¶ms, @@ -2428,13 +2439,13 @@ impl ThreadRequestProcessor { ) .await { - Ok(true) => return Ok(()), - Ok(false) => {} + Ok(RunningThreadResumeResult::Handled) => return Ok(()), + Ok(RunningThreadResumeResult::NotRunning(stored_thread)) => stored_thread, Err(error) => { self.outgoing.send_error(request_id, error).await; return Ok(()); } - } + }; let ThreadResumeParams { thread_id, @@ -2458,15 +2469,20 @@ impl ThreadRequestProcessor { } = params; let include_turns = !exclude_turns; - let (thread_history, resume_source_thread) = match if let Some(history) = history { + let resume_result = if let Some(history) = history { self.resume_thread_from_history(history.as_slice()) .await .map(|thread_history| (thread_history, None)) + } else if let Some(stored_thread) = stored_thread_from_running_probe { + self.stored_thread_to_initial_history(&stored_thread) + .await + .map(|thread_history| (thread_history, Some(*stored_thread))) } else { self.resume_thread_from_rollout(&thread_id, path.as_ref()) .await .map(|(thread_history, stored_thread)| (thread_history, Some(stored_thread))) - } { + }; + let (thread_history, resume_source_thread) = match resume_result { Ok(value) => value, Err(error) => { self.outgoing.send_error(request_id, error).await; @@ -2707,7 +2723,7 @@ impl ThreadRequestProcessor { params: &ThreadResumeParams, app_server_client_name: Option, app_server_client_version: Option, - ) -> Result { + ) -> Result { let running_thread = if params.history.is_some() { if let Ok(existing_thread_id) = ThreadId::from_string(¶ms.thread_id) && self @@ -2743,7 +2759,11 @@ impl ThreadRequestProcessor { let existing_thread_id = source_thread.thread_id; match self.thread_manager.get_thread(existing_thread_id).await { Ok(existing_thread) => Some((existing_thread_id, existing_thread, source_thread)), - Err(_) => None, + Err(_) => { + return Ok(RunningThreadResumeResult::NotRunning(Some(Box::new( + source_thread, + )))); + } } }; @@ -2784,7 +2804,9 @@ impl ThreadRequestProcessor { ThreadShutdownResult::Complete => { self.thread_manager.remove_thread(&existing_thread_id).await; self.finalize_thread_teardown(existing_thread_id).await; - return Ok(false); + // Shutdown can flush newer rollout items, so reload the + // stored thread before starting the replacement session. + return Ok(RunningThreadResumeResult::NotRunning(None)); } ThreadShutdownResult::SubmitFailed => { warn!("failed to submit Shutdown to thread {existing_thread_id}"); @@ -2876,9 +2898,9 @@ impl ThreadRequestProcessor { "failed to enqueue running thread resume for thread {existing_thread_id}: thread listener command channel is closed" ))); } - return Ok(true); + return Ok(RunningThreadResumeResult::Handled); } - Ok(false) + Ok(RunningThreadResumeResult::NotRunning(None)) } async fn resume_thread_from_history( 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 016e4c3df8..ca2121dd4c 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 @@ -29,6 +29,7 @@ use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ThreadListParams; use codex_app_server_protocol::ThreadListResponse; +use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; @@ -37,6 +38,7 @@ use codex_arg0::Arg0DispatchPaths; use codex_config::CloudConfigBundleLoader; use codex_config::LoaderOverrides; use codex_config::NoopThreadConfigLoader; +use codex_core::config::Config; use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; @@ -183,6 +185,122 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste Ok(()) } +#[tokio::test] +async fn cold_thread_resume_reuses_non_local_history_probe() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + let store_id = Uuid::new_v4().to_string(); + create_config_toml_with_thread_store(codex_home.path(), &server.uri(), &store_id)?; + + let loader_overrides = LoaderOverrides::without_managed_config_for_tests(); + let config = Arc::new( + 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 = start_in_process_client(config.clone(), loader_overrides.clone()).await?; + let response = client + .request(ClientRequest::ThreadStart { + request_id: RequestId::Integer(1), + params: ThreadStartParams::default(), + }) + .await? + .expect("thread/start should succeed"); + let ThreadStartResponse { thread, .. } = serde_json::from_value(response)?; + + client + .request(ClientRequest::TurnStart { + request_id: RequestId::Integer(2), + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Materialize the thread".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, + }) + .await? + .expect("turn/start should succeed"); + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let Some(event) = client.next_event().await else { + anyhow::bail!("in-process app-server stopped before turn/completed"); + }; + if let InProcessServerEvent::ServerNotification(ServerNotification::TurnCompleted( + completed, + )) = event + && completed.thread_id == thread.id + { + return Ok::<(), anyhow::Error>(()); + } + } + }) + .await??; + client.shutdown().await?; + + let client = start_in_process_client(config, loader_overrides).await?; + let reads_before_resume = thread_store.calls().await.read_thread_with_history; + // The in-memory store is pathless, so resume currently fails later while + // assembling the response. The history-bearing probe must still be reused. + let _resume_result = client + .request(ClientRequest::ThreadResume { + request_id: RequestId::Integer(3), + params: ThreadResumeParams { + thread_id: thread.id.clone(), + ..Default::default() + }, + }) + .await?; + + assert_eq!( + thread_store.calls().await.read_thread_with_history, + reads_before_resume + 1 + ); + + client.shutdown().await?; + Ok(()) +} + +async fn start_in_process_client( + config: Arc, + loader_overrides: LoaderOverrides, +) -> std::io::Result { + in_process::start(InProcessStartArgs { + arg0_paths: Arg0DispatchPaths::default(), + 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 +} + 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/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index edc308a905..beed1d18fe 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -109,6 +109,7 @@ pub struct InMemoryThreadStoreCalls { pub discard_thread: usize, pub load_history: usize, pub read_thread: usize, + pub read_thread_with_history: usize, pub read_thread_by_rollout_path: usize, pub list_threads: usize, pub update_thread_metadata: usize, @@ -262,6 +263,9 @@ impl ThreadStore for InMemoryThreadStore { async fn read_thread(&self, params: ReadThreadParams) -> ThreadStoreResult { let mut state = self.state.lock().await; state.calls.read_thread += 1; + if params.include_history { + state.calls.read_thread_with_history += 1; + } stored_thread_from_state(&state, params.thread_id, params.include_history) } From 8a299fb7043b645cecb50b239beb7d0e9e4ac7e7 Mon Sep 17 00:00:00 2001 From: canvrno-oai Date: Tue, 9 Jun 2026 09:41:58 -0700 Subject: [PATCH 04/17] fix: Prevent /review crash when entering Esc on steer message (#22879) This changes the `/review` escape path so `Esc` no longer behaves like the normal queued-follow-up interrupt flow while a review is running. Steering is not currently supported in `/review` mode, without this change users are able to attempt a steer but it leads to a crash (see #22815). If the user has already tried to send additional guidance during `/review`, the TUI now keeps the review running and shows a warning that steer messages are not supported in that mode, while still pointing users to `Ctrl+C` if they actually want to cancel. It also adds regression coverage for the review-specific warning behavior. When users do cancel with Ctrl+C during /review, the TUI now tolerates the active-turn race that can happen during review handoff, and any queued steer messages are restored to the composer instead of being discarded. - Special-case `Esc` during an active `/review` when follow-up steer input is pending or has already been deferred. - Show a clear warning instead of interrupting the running review. - Make the Ctrl+C cancel path during /review resilient to active-turn races, while preserving any queued steer text by restoring it to the composer. - Add review-mode test coverage for the warning path. ## Testing 1. Start a `/review` with a diff large enough that the review stays active for more than a few seconds. 2. While the review is still running, type a follow-up / steer message, submit it, and then press `Esc`. Before: `Esc` causes the TUI to close abruptly. After: the review keeps running and the transcript shows a warning that steer messages are not supported during `/review`, with guidance to use `Ctrl+C` if you want to cancel. 3. Press `Ctrl+C` if you actually want to stop the review. Before: (after restarting the test since Pt. 2 crashed) this is the intentional cancellation path. After: this remains the intentional cancellation path, and any queued follow-up steer text is restored to the composer instead of being lost. ## Note: `/review` mode explicitly does not support steering at this time (as noted in `turn_processer.rs`, if we want to explore that in the future this code will need to be modified). This change keeps unsupported steer attempts from crashing the TUI and preserves queued follow-up text if the user cancels with Ctrl+C. --- codex-rs/tui/src/app.rs | 19 +++++++++ codex-rs/tui/src/app/tests.rs | 17 ++++++++ codex-rs/tui/src/app/thread_routing.rs | 41 +++++++++++++++++-- codex-rs/tui/src/app_server_session.rs | 10 +++-- codex-rs/tui/src/chatwidget/interaction.rs | 14 +++++++ ...s__review_submission_warning_snapshot.snap | 8 ++-- .../tui/src/chatwidget/tests/review_mode.rs | 23 +++++++++++ 7 files changed, 121 insertions(+), 11 deletions(-) diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 63a5d4a1c0..e5a3d07329 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -698,6 +698,25 @@ fn archived_session_guidance(err: &color_eyre::eyre::Report) -> Option { Some(message.to_string()) } +fn active_turn_interrupt_race(error: &TypedRequestError) -> Option { + let TypedRequestError::Server { method, source } = error else { + return None; + }; + if method != "turn/interrupt" { + return None; + } + let mismatch_prefix = "expected active turn id "; + let mismatch_separator = " but found "; + Some( + source + .message + .strip_prefix(mismatch_prefix)? + .split_once(mismatch_separator)? + .1 + .to_string(), + ) +} + impl App { pub fn chatwidget_init_for_forked_or_resumed_thread( &self, diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 99d4c35333..ed89c72878 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -4709,6 +4709,23 @@ fn active_turn_steer_race_extracts_actual_turn_id_from_mismatch() { ); } +#[test] +fn active_turn_interrupt_race_extracts_actual_turn_id_from_mismatch() { + let error = TypedRequestError::Server { + method: "turn/interrupt".to_string(), + source: JSONRPCErrorError { + code: -32602, + message: "expected active turn id turn-expected but found turn-actual".to_string(), + data: None, + }, + }; + + assert_eq!( + active_turn_interrupt_race(&error), + Some("turn-actual".to_string()) + ); +} + #[tokio::test] async fn fresh_session_config_uses_current_service_tier() { let mut app = make_test_app().await; diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 80b67b19d6..938174254a 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -500,9 +500,39 @@ impl App { match op { AppCommand::Interrupt { .. } => { if let Some(turn_id) = self.active_turn_id_for_thread(thread_id).await { - app_server.turn_interrupt(thread_id, turn_id).await?; + let mut interrupt_turn_id = turn_id; + for retried_after_turn_mismatch in [false, true] { + match app_server + .turn_interrupt(thread_id, interrupt_turn_id.clone()) + .await + { + Ok(()) => return Ok(true), + Err(error) if !retried_after_turn_mismatch => { + let Some(actual_turn_id) = active_turn_interrupt_race(&error) + else { + return Err(error).wrap_err("turn/interrupt failed in TUI"); + }; + if actual_turn_id == interrupt_turn_id { + return Err(error).wrap_err("turn/interrupt failed in TUI"); + } + // Review flows can swap the active turn before the TUI processes + // the corresponding notification. Retry once with the + // server-reported turn id so Ctrl+C/Esc do not fatally exit on that + // stale cache, but let lifecycle notifications own the cached + // active turn id. + interrupt_turn_id = actual_turn_id; + } + Err(error) => { + return Err(error).wrap_err("turn/interrupt failed in TUI"); + } + } + } + unreachable!("interrupt retry loop should return"); } else { - app_server.startup_interrupt(thread_id).await?; + app_server + .startup_interrupt(thread_id) + .await + .wrap_err("turn/interrupt failed in TUI")?; } Ok(true) } @@ -652,7 +682,12 @@ impl App { Ok(true) } AppCommand::Review { target } => { - app_server.review_start(thread_id, target.clone()).await?; + let response = app_server.review_start(thread_id, target.clone()).await?; + let review_thread_id = ThreadId::from_string(&response.review_thread_id) + .wrap_err("review/start returned invalid review thread id")?; + let store = Arc::clone(&self.ensure_thread_channel(review_thread_id).store); + let mut store = store.lock().await; + store.active_turn_id = Some(response.turn.id); Ok(true) } AppCommand::CleanBackgroundTerminals => { diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index 32db3d2cf7..e29b40111c 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -736,7 +736,7 @@ impl AppServerSession { &mut self, thread_id: ThreadId, turn_id: String, - ) -> Result<()> { + ) -> std::result::Result<(), TypedRequestError> { let request_id = self.next_request_id(); let _: TurnInterruptResponse = self .client @@ -747,12 +747,14 @@ impl AppServerSession { turn_id, }, }) - .await - .wrap_err("turn/interrupt failed in TUI")?; + .await?; Ok(()) } - pub(crate) async fn startup_interrupt(&mut self, thread_id: ThreadId) -> Result<()> { + pub(crate) async fn startup_interrupt( + &mut self, + thread_id: ThreadId, + ) -> std::result::Result<(), TypedRequestError> { self.turn_interrupt(thread_id, String::new()).await } diff --git a/codex-rs/tui/src/chatwidget/interaction.rs b/codex-rs/tui/src/chatwidget/interaction.rs index a94d483f0d..7a3385034d 100644 --- a/codex-rs/tui/src/chatwidget/interaction.rs +++ b/codex-rs/tui/src/chatwidget/interaction.rs @@ -112,6 +112,20 @@ impl ChatWidget { return; } + const REVIEW_STEER_UNAVAILABLE_MESSAGE: &str = "Steer messages aren't supported during /review. Press Ctrl+C now to cancel the review."; + + if self.chat_keymap.interrupt_turn.is_pressed(key_event) + && self.review.is_review_mode + && (!self.input_queue.pending_steers.is_empty() + || !self.input_queue.rejected_steers_queue.is_empty()) + && self.bottom_pane.is_task_running() + && self.bottom_pane.no_modal_or_popup_active() + && !self.should_handle_vim_insert_escape(key_event) + { + self.add_warning_message(REVIEW_STEER_UNAVAILABLE_MESSAGE.to_string()); + return; + } + if self.chat_keymap.interrupt_turn.is_pressed(key_event) && !self.input_queue.pending_steers.is_empty() && self.bottom_pane.is_task_running() diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__review_submission_warning_snapshot.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__review_submission_warning_snapshot.snap index c8bbb7c2a4..1a28a9222b 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__review_submission_warning_snapshot.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__review_submission_warning_snapshot.snap @@ -1,7 +1,7 @@ --- source: tui/src/chatwidget/tests/review_mode.rs -expression: lines_to_single_string(last) +assertion_line: 307 +expression: last --- -⚠ Steer messages aren't supported during /review. Send your message after the - review finishes, or press Ctrl+C now to cancel the review. - +⚠ Steer messages aren't supported during /review. Press Ctrl+C now to cancel the + review. diff --git a/codex-rs/tui/src/chatwidget/tests/review_mode.rs b/codex-rs/tui/src/chatwidget/tests/review_mode.rs index faf8636ad8..af5816cac2 100644 --- a/codex-rs/tui/src/chatwidget/tests/review_mode.rs +++ b/codex-rs/tui/src/chatwidget/tests/review_mode.rs @@ -285,6 +285,29 @@ async fn steer_rejection_queues_review_follow_up_before_existing_queued_messages } } +#[tokio::test] +async fn esc_with_review_queued_steers_shows_warning_and_does_not_interrupt() { + let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.thread_id = Some(ThreadId::new()); + handle_turn_started(&mut chat, "turn-1"); + handle_entered_review_mode(&mut chat, "feature branch"); + let _ = drain_insert_history(&mut rx); + chat.input_queue + .pending_steers + .push_back(pending_steer("review follow-up")); + chat.refresh_pending_input_preview(); + + chat.handle_key_event(KeyEvent::new(KeyCode::Esc, KeyModifiers::NONE)); + + assert!(!chat.input_queue.submit_pending_steers_after_interrupt); + assert_eq!(chat.input_queue.pending_steers.len(), 1); + assert_no_submit_op(&mut op_rx); + + let cells = drain_insert_history(&mut rx); + let last = lines_to_single_string(cells.last().expect("review warning")); + assert_chatwidget_snapshot!("review_submission_warning_snapshot", last); +} + #[tokio::test] async fn live_agent_message_renders_during_review_mode() { let (mut chat, mut rx, _ops) = make_chatwidget_manual(/*model_override*/ None).await; From 1026e9de1be292ebb01579dfcce5a34ab224917a Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 9 Jun 2026 19:40:40 +0200 Subject: [PATCH 05/17] app-server: reject direct input to multi-agent v2 sub-agents (#27173) ## Why Multi-agent v2 sub-agents are owned and coordinated by their parent agent. Allowing an app-server client to start or steer turns on a spawned child bypasses the multi-agent messaging path and creates a second, conflicting source of work for that sub-agent. ## What changed - Reject direct `turn/start` and `turn/steer` requests targeting multi-agent v2 thread-spawn sub-agents. - Identify these targets using both the thread's resolved multi-agent version and its `SubAgentSource::ThreadSpawn` session source, leaving root threads, v1 agents, and other sub-agent types unchanged. - Return a consistent invalid-request error before validating or applying the submitted input. ## Testing - Added an app-server integration test that spawns a real multi-agent v2 child and verifies that direct `turn/start` and `turn/steer` requests are rejected. --- .../src/request_processors/turn_processor.rs | 42 +++++- .../app-server/tests/suite/v2/turn_start.rs | 130 ++++++++++++++++++ 2 files changed, 166 insertions(+), 6 deletions(-) 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 7bdd9df0d8..3d44781946 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -1,6 +1,12 @@ use super::*; use codex_protocol::protocol::AdditionalContextEntry as CoreAdditionalContextEntry; use codex_protocol::protocol::AdditionalContextKind as CoreAdditionalContextKind; +use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; + +const DIRECT_INPUT_TO_MULTI_AGENT_V2_SUBAGENT_ERROR: &str = + "direct app-server input is not allowed for multi-agent v2 sub-agents"; #[derive(Clone)] pub(crate) struct TurnRequestProcessor { @@ -236,6 +242,26 @@ impl TurnRequestProcessor { Ok((thread_id, thread)) } + + async fn ensure_direct_input_allowed( + &self, + request_id: &ConnectionRequestId, + thread: &CodexThread, + ) -> Result<(), JSONRPCErrorError> { + if thread.multi_agent_version() == Some(MultiAgentVersion::V2) + && matches!( + thread.config_snapshot().await.session_source, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) + ) + { + let error = invalid_request(DIRECT_INPUT_TO_MULTI_AGENT_V2_SUBAGENT_ERROR); + self.track_error_response(request_id, &error, /*error_type*/ None); + return Err(error); + } + + Ok(()) + } + fn normalize_collaboration_mode( &self, mut collaboration_mode: CollaborationMode, @@ -370,6 +396,14 @@ impl TurnRequestProcessor { app_server_client_name: Option, app_server_client_version: Option, ) -> Result { + let (thread_id, thread) = + self.load_thread(¶ms.thread_id) + .await + .inspect_err(|error| { + self.track_error_response(&request_id, error, /*error_type*/ None); + })?; + self.ensure_direct_input_allowed(&request_id, thread.as_ref()) + .await?; if let Err(error) = Self::validate_v2_input_limit(¶ms.input) { self.track_error_response( &request_id, @@ -378,12 +412,6 @@ impl TurnRequestProcessor { ); return Err(error); } - let (thread_id, thread) = - self.load_thread(¶ms.thread_id) - .await - .inspect_err(|error| { - self.track_error_response(&request_id, error, /*error_type*/ None); - })?; Self::set_app_server_client_info( thread.as_ref(), app_server_client_name, @@ -762,6 +790,8 @@ impl TurnRequestProcessor { .inspect_err(|error| { self.track_error_response(request_id, error, /*error_type*/ None); })?; + self.ensure_direct_input_allowed(request_id, thread.as_ref()) + .await?; if params.expected_turn_id.is_empty() { return Err(invalid_request("expectedTurnId must not be empty")); 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 aed3c58697..5080937997 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -41,6 +41,7 @@ use codex_app_server_protocol::PatchChangeKind; use codex_app_server_protocol::RequestId; 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::ThreadItem; use codex_app_server_protocol::ThreadSource; @@ -53,6 +54,7 @@ use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStartedNotification; use codex_app_server_protocol::TurnStatus; +use codex_app_server_protocol::TurnSteerParams; use codex_app_server_protocol::UserInput as V2UserInput; use codex_app_server_protocol::WarningNotification; use codex_config::config_toml::ConfigToml; @@ -3363,6 +3365,134 @@ async fn turn_start_emits_spawn_agent_item_with_model_metadata_v2() -> Result<() Ok(()) } +#[tokio::test] +async fn direct_input_to_multi_agent_v2_subagent_is_rejected() -> Result<()> { + const CHILD_PROMPT: &str = "child: do work"; + const PARENT_PROMPT: &str = "spawn a child and continue"; + const SPAWN_CALL_ID: &str = "spawn-call-direct-input-rejection"; + const ERROR_MESSAGE: &str = + "direct app-server input is not allowed for multi-agent v2 sub-agents"; + + let server = responses::start_mock_server().await; + let spawn_args = serde_json::to_string(&json!({ + "message": CHILD_PROMPT, + "task_name": "worker", + }))?; + let _parent_turn = responses::mount_sse_once_match( + &server, + |req: &wiremock::Request| body_contains(req, PARENT_PROMPT), + responses::sse(vec![ + responses::ev_response_created("resp-parent-1"), + responses::ev_function_call(SPAWN_CALL_ID, "spawn_agent", &spawn_args), + responses::ev_completed("resp-parent-1"), + ]), + ) + .await; + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + &server.uri(), + "never", + &BTreeMap::from([(Feature::MultiAgentV2, true)]), + )?; + write_models_cache(codex_home.path())?; + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("gpt-5.3-codex".to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id, + input: vec![V2UserInput::Text { + text: PARENT_PROMPT.to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + let _: TurnStartResponse = to_response(turn_resp)?; + + let child_thread_id = timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let completed_notif = mcp + .read_stream_until_notification_message("item/completed") + .await?; + let completed: ItemCompletedNotification = + serde_json::from_value(completed_notif.params.expect("item/completed params"))?; + if let ThreadItem::SubAgentActivity { + id, + kind: SubAgentActivityKind::Started, + agent_thread_id, + .. + } = completed.item + && id == SPAWN_CALL_ID + { + return Ok::(agent_thread_id); + } + } + }) + .await??; + + let direct_turn_req = mcp + .send_turn_start_request(TurnStartParams { + thread_id: child_thread_id.clone(), + input: vec![V2UserInput::Text { + text: "direct app-server turn".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let direct_turn_error: JSONRPCError = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(direct_turn_req)), + ) + .await??; + assert_eq!(direct_turn_error.error.code, INVALID_REQUEST_ERROR_CODE); + assert_eq!(direct_turn_error.error.message, ERROR_MESSAGE); + + let direct_steer_req = mcp + .send_turn_steer_request(TurnSteerParams { + thread_id: child_thread_id, + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "direct app-server steer".to_string(), + text_elements: Vec::new(), + }], + responsesapi_client_metadata: None, + additional_context: None, + expected_turn_id: "any-active-turn".to_string(), + }) + .await?; + let direct_steer_error: JSONRPCError = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(direct_steer_req)), + ) + .await??; + assert_eq!(direct_steer_error.error.code, INVALID_REQUEST_ERROR_CODE); + assert_eq!(direct_steer_error.error.message, ERROR_MESSAGE); + + Ok(()) +} + #[tokio::test] async fn turn_start_emits_spawn_agent_item_with_effective_role_model_metadata_v2() -> Result<()> { skip_if_no_network!(Ok(())); From 89ac3ec27cae0000da14fbe4160a79ed465843fb Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 9 Jun 2026 19:51:54 +0200 Subject: [PATCH 06/17] Load selected executor skills through extensions (#27184) ## Why CCA is moving toward a split runtime where the orchestrator may not have a filesystem, while executors can expose preinstalled plugins and skills. A thread therefore needs to select capabilities without asking app-server or core to interpret executor-owned paths through the orchestrator's filesystem. The longer-term model is broader than executor skills: - A plugin is a bundle of skills, MCP servers, connectors/apps, and hooks. - A plugin root can be local, executor-owned, or hosted by a backend. - Components inside one plugin can use different access and execution mechanisms. A skill may be read from a filesystem or through backend tools; an HTTP MCP server can run without an executor; a stdio MCP server or hook needs an execution environment. - Core should carry generic extension initialization data. The extension that owns a component should discover it, expose it to the model, and invoke it through the appropriate runtime. This PR establishes that architecture through one complete vertical: selecting a root on an executor, discovering the skills beneath it, exposing those skills to the model, and reading an explicitly invoked `SKILL.md` through the same executor. ## Contract `thread/start` gains an experimental `selectedCapabilityRoots` field: ```json { "selectedCapabilityRoots": [ { "id": "deploy-plugin@1", "location": { "type": "environment", "environmentId": "workspace", "path": "/opt/codex/plugins/deploy" } } ] } ``` The root is intentionally not classified as a "plugin" or "skill" in the API. It can point at a standalone skill, a directory containing several skills, or a plugin containing skills and other components. This PR only teaches the skills extension how to consume it; later extensions can resolve MCP, connector, and hook components from the same selection. The platform-supplied `id` is stable selection identity. The location says which runtime owns the root and gives that runtime an opaque path. App-server does not inspect or canonicalize the path. ## What changed ### Generic thread extension initialization App-server converts selected roots into `ExtensionDataInit`. Core carries that generic initialization value until the final thread ID is known, then creates thread-scoped `ExtensionData` before lifecycle contributors run. This keeps `Session` and core independent of the capability-selection contract. The initialization value is consumed during construction; it is not retained as another long-lived `Session` field. ### Executor-backed skills The skills extension now owns an `ExecutorSkillProvider` that: - resolves the selected environment through `EnvironmentManager` - discovers, canonicalizes, and reads skills through that environment's `ExecutorFileSystem` - contributes the bounded selected-skill catalog as stable developer context - reads an explicitly invoked skill body through the authority that listed it - warns when an environment or root is unavailable - never falls back to the orchestrator filesystem for an executor-owned root Skill catalog and instruction fragments have hard byte bounds, which also bound them below the 10K-token per-item context limit. If a selected executor skill has the same name as a legacy local skill, the executor selection owns that invocation and the local body is not injected a second time. Existing local and bundled skill loading remains in place. Omitting `selectedCapabilityRoots` therefore preserves current local-only behavior. ## Current semantics - Only environment-owned locations are represented in this first contract. - Roots are resolved by the destination extension, not by app-server or core. - An unavailable executor or invalid root produces a warning and no capabilities from that root; it does not trigger a local-filesystem fallback. - Selection applies to a newly started active thread. - MCP servers, connectors, and hooks beneath a selected plugin root are not activated yet. - Selection is not yet persisted or inherited across resume, fork, or subagent creation. Existing local capabilities continue to behave as they do today in those flows. ## Planned vertical follow-ups 1. **Hosted HTTP MCP:** add an extension-backed HTTP MCP source that works without an executor, then replace the special-purpose MCP plugins loader with that implementation. 2. **Executor MCP:** register and execute stdio MCP servers through the environment that owns the selected plugin root. 3. **Backend skills:** add a hosted skill source whose catalog and bodies are accessed through extension tools rather than a filesystem. 4. **Connectors and hooks:** activate those components through their owning extensions, using the same selected-root boundary and component-specific runtime. 5. **Durable selection:** define the desired-selection lifecycle, persist it, and make resume, fork, and subagent inheritance explicit rather than accidental. 6. **Local convergence:** incrementally route existing local plugin, skill, and MCP loading through the same extension model while preserving current local behavior. Each follow-up remains reviewable as an end-to-end capability. The platform selects roots, generic thread extension data carries the selection, and the owning extension resolves and operates its component. ## Verification Coverage added for: - app-server end-to-end discovery and explicit invocation of a skill inside an executor-selected plugin root - exclusive invocation when a selected executor skill collides with a local skill name - executor filesystem authority for discovery, canonicalization, and reads - thread extension initialization before lifecycle contributors run - stable executor catalog context, explicit invocation, context rebuilding, hidden skills, and preserved host/remote catalog behavior Targeted protocol, core-skills, skills-extension, core lifecycle, and app-server executor-skill tests were run during development. --- codex-rs/Cargo.lock | 4 + codex-rs/Cargo.toml | 1 + .../schema/json/ClientRequest.json | 52 +++ .../codex_app_server_protocol.schemas.json | 52 +++ .../codex_app_server_protocol.v2.schemas.json | 52 +++ .../schema/json/v2/ThreadStartParams.json | 52 +++ .../typescript/v2/CapabilityRootLocation.ts | 8 + .../typescript/v2/SelectedCapabilityRoot.ts | 17 + .../schema/typescript/v2/index.ts | 2 + .../src/protocol/v2/thread.rs | 6 + codex-rs/app-server/Cargo.toml | 1 + codex-rs/app-server/README.md | 14 +- codex-rs/app-server/src/extensions.rs | 6 + codex-rs/app-server/src/mcp_refresh.rs | 10 +- codex-rs/app-server/src/message_processor.rs | 9 + .../request_processors/thread_processor.rs | 10 + .../tests/suite/v2/executor_skills.rs | 148 ++++++++ codex-rs/app-server/tests/suite/v2/mod.rs | 1 + .../app-server/tests/suite/v2/skills_list.rs | 1 + codex-rs/core-skills/src/loader.rs | 24 +- codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/session/mod.rs | 4 + codex-rs/core/src/session/session.rs | 8 +- codex-rs/core/src/session/tests.rs | 3 + .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/thread_manager.rs | 15 + codex-rs/core/src/thread_manager_tests.rs | 79 ++++ .../tests/suite/subagent_notifications.rs | 1 + codex-rs/ext/extension-api/src/lib.rs | 1 + codex-rs/ext/extension-api/src/state.rs | 33 +- codex-rs/ext/skills/Cargo.toml | 3 + codex-rs/ext/skills/src/catalog.rs | 35 +- codex-rs/ext/skills/src/extension.rs | 104 ++++-- codex-rs/ext/skills/src/lib.rs | 2 + codex-rs/ext/skills/src/provider.rs | 5 +- codex-rs/ext/skills/src/provider/executor.rs | 197 ++++++++++ codex-rs/ext/skills/src/provider/host.rs | 10 +- codex-rs/ext/skills/src/render.rs | 29 +- codex-rs/ext/skills/src/selection.rs | 4 +- codex-rs/ext/skills/src/sources.rs | 2 +- codex-rs/ext/skills/src/state.rs | 12 +- .../tests/executor_file_system_authority.rs | 337 ++++++++++++++++++ codex-rs/ext/skills/tests/skills_extension.rs | 199 +++++++---- codex-rs/memories/write/src/runtime.rs | 1 + codex-rs/protocol/src/capabilities.rs | 30 ++ codex-rs/protocol/src/lib.rs | 1 + 46 files changed, 1460 insertions(+), 127 deletions(-) create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/CapabilityRootLocation.ts create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/SelectedCapabilityRoot.ts create mode 100644 codex-rs/app-server/tests/suite/v2/executor_skills.rs create mode 100644 codex-rs/ext/skills/src/provider/executor.rs create mode 100644 codex-rs/ext/skills/tests/executor_file_system_authority.rs create mode 100644 codex-rs/protocol/src/capabilities.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 3d8086c27b..064d9ad944 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1994,6 +1994,7 @@ dependencies = [ "codex-rollout", "codex-sandboxing", "codex-shell-command", + "codex-skills-extension", "codex-state", "codex-thread-store", "codex-tools", @@ -3796,8 +3797,11 @@ dependencies = [ "async-trait", "codex-core", "codex-core-skills", + "codex-exec-server", "codex-extension-api", "codex-protocol", + "codex-utils-absolute-path", + "codex-utils-string", "pretty_assertions", "tokio", ] diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index c72b7c7844..955d8987f7 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -212,6 +212,7 @@ codex-sandboxing = { path = "sandboxing" } codex-secrets = { path = "secrets" } codex-shell-command = { path = "shell-command" } codex-shell-escalation = { path = "shell-escalation" } +codex-skills-extension = { path = "ext/skills" } codex-skills = { path = "skills" } codex-state = { path = "state" } codex-stdio-to-uds = { path = "stdio-to-uds" } diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 6d9be4ec02..c499376547 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -180,6 +180,36 @@ ], "type": "object" }, + "CapabilityRootLocation": { + "description": "Location used to resolve a selected capability root.", + "oneOf": [ + { + "description": "A path owned by an execution environment.", + "properties": { + "environmentId": { + "type": "string" + }, + "path": { + "type": "string" + }, + "type": { + "enum": [ + "environment" + ], + "title": "EnvironmentCapabilityRootLocationType", + "type": "string" + } + }, + "required": [ + "environmentId", + "path", + "type" + ], + "title": "EnvironmentCapabilityRootLocation", + "type": "object" + } + ] + }, "ClientInfo": { "properties": { "name": { @@ -2947,6 +2977,28 @@ } ] }, + "SelectedCapabilityRoot": { + "description": "A user-selected root that can expose one or more runtime capabilities.", + "properties": { + "id": { + "description": "Stable identifier supplied by the capability selection platform.", + "type": "string" + }, + "location": { + "allOf": [ + { + "$ref": "#/definitions/CapabilityRootLocation" + } + ], + "description": "Where the selected root can be resolved." + } + }, + "required": [ + "id", + "location" + ], + "type": "object" + }, "SendAddCreditsNudgeEmailParams": { "properties": { "creditType": { 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 7abc5af7d0..475154b38a 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 @@ -6765,6 +6765,36 @@ ], "type": "string" }, + "CapabilityRootLocation": { + "description": "Location used to resolve a selected capability root.", + "oneOf": [ + { + "description": "A path owned by an execution environment.", + "properties": { + "environmentId": { + "type": "string" + }, + "path": { + "type": "string" + }, + "type": { + "enum": [ + "environment" + ], + "title": "EnvironmentCapabilityRootLocationType", + "type": "string" + } + }, + "required": [ + "environmentId", + "path", + "type" + ], + "title": "EnvironmentCapabilityRootLocation", + "type": "object" + } + ] + }, "CodexErrorInfo": { "description": "This translation layer make sure that we expose codex error code in camel case.\n\nWhen an upstream HTTP status is available (for example, from the Responses API or a provider), it is forwarded in `httpStatusCode` on the relevant `codexErrorInfo` variant.", "oneOf": [ @@ -15061,6 +15091,28 @@ }, "type": "object" }, + "SelectedCapabilityRoot": { + "description": "A user-selected root that can expose one or more runtime capabilities.", + "properties": { + "id": { + "description": "Stable identifier supplied by the capability selection platform.", + "type": "string" + }, + "location": { + "allOf": [ + { + "$ref": "#/definitions/v2/CapabilityRootLocation" + } + ], + "description": "Where the selected root can be resolved." + } + }, + "required": [ + "id", + "location" + ], + "type": "object" + }, "SendAddCreditsNudgeEmailParams": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { 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 e5ac77c088..8ed40282a9 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 @@ -1087,6 +1087,36 @@ ], "type": "string" }, + "CapabilityRootLocation": { + "description": "Location used to resolve a selected capability root.", + "oneOf": [ + { + "description": "A path owned by an execution environment.", + "properties": { + "environmentId": { + "type": "string" + }, + "path": { + "type": "string" + }, + "type": { + "enum": [ + "environment" + ], + "title": "EnvironmentCapabilityRootLocationType", + "type": "string" + } + }, + "required": [ + "environmentId", + "path", + "type" + ], + "title": "EnvironmentCapabilityRootLocation", + "type": "object" + } + ] + }, "ClientInfo": { "properties": { "name": { @@ -11563,6 +11593,28 @@ }, "type": "object" }, + "SelectedCapabilityRoot": { + "description": "A user-selected root that can expose one or more runtime capabilities.", + "properties": { + "id": { + "description": "Stable identifier supplied by the capability selection platform.", + "type": "string" + }, + "location": { + "allOf": [ + { + "$ref": "#/definitions/CapabilityRootLocation" + } + ], + "description": "Where the selected root can be resolved." + } + }, + "required": [ + "id", + "location" + ], + "type": "object" + }, "SendAddCreditsNudgeEmailParams": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json index 99b25490ab..7f0d4ad9a4 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json @@ -64,6 +64,36 @@ } ] }, + "CapabilityRootLocation": { + "description": "Location used to resolve a selected capability root.", + "oneOf": [ + { + "description": "A path owned by an execution environment.", + "properties": { + "environmentId": { + "type": "string" + }, + "path": { + "type": "string" + }, + "type": { + "enum": [ + "environment" + ], + "title": "EnvironmentCapabilityRootLocationType", + "type": "string" + } + }, + "required": [ + "environmentId", + "path", + "type" + ], + "title": "EnvironmentCapabilityRootLocation", + "type": "object" + } + ] + }, "DynamicToolSpec": { "properties": { "deferLoading": { @@ -106,6 +136,28 @@ ], "type": "string" }, + "SelectedCapabilityRoot": { + "description": "A user-selected root that can expose one or more runtime capabilities.", + "properties": { + "id": { + "description": "Stable identifier supplied by the capability selection platform.", + "type": "string" + }, + "location": { + "allOf": [ + { + "$ref": "#/definitions/CapabilityRootLocation" + } + ], + "description": "Where the selected root can be resolved." + } + }, + "required": [ + "id", + "location" + ], + "type": "object" + }, "ThreadSource": { "enum": [ "user", diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/CapabilityRootLocation.ts b/codex-rs/app-server-protocol/schema/typescript/v2/CapabilityRootLocation.ts new file mode 100644 index 0000000000..6266101eb8 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/CapabilityRootLocation.ts @@ -0,0 +1,8 @@ +// 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. + +/** + * Location used to resolve a selected capability root. + */ +export type CapabilityRootLocation = { "type": "environment", environmentId: string, path: string, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/SelectedCapabilityRoot.ts b/codex-rs/app-server-protocol/schema/typescript/v2/SelectedCapabilityRoot.ts new file mode 100644 index 0000000000..849d5c7aae --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/SelectedCapabilityRoot.ts @@ -0,0 +1,17 @@ +// 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. +import type { CapabilityRootLocation } from "./CapabilityRootLocation"; + +/** + * A user-selected root that can expose one or more runtime capabilities. + */ +export type SelectedCapabilityRoot = { +/** + * Stable identifier supplied by the capability selection platform. + */ +id: string, +/** + * Where the selected root can be resolved. + */ +location: CapabilityRootLocation, }; 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 eab3494e1f..71f4c67cd0 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -40,6 +40,7 @@ export type { ByteRange } from "./ByteRange"; export type { CancelLoginAccountParams } from "./CancelLoginAccountParams"; export type { CancelLoginAccountResponse } from "./CancelLoginAccountResponse"; export type { CancelLoginAccountStatus } from "./CancelLoginAccountStatus"; +export type { CapabilityRootLocation } from "./CapabilityRootLocation"; export type { ChatgptAuthTokensRefreshParams } from "./ChatgptAuthTokensRefreshParams"; export type { ChatgptAuthTokensRefreshReason } from "./ChatgptAuthTokensRefreshReason"; export type { ChatgptAuthTokensRefreshResponse } from "./ChatgptAuthTokensRefreshResponse"; @@ -334,6 +335,7 @@ export type { ReviewTarget } from "./ReviewTarget"; export type { SandboxMode } from "./SandboxMode"; export type { SandboxPolicy } from "./SandboxPolicy"; export type { SandboxWorkspaceWrite } from "./SandboxWorkspaceWrite"; +export type { SelectedCapabilityRoot } from "./SelectedCapabilityRoot"; export type { SendAddCreditsNudgeEmailParams } from "./SendAddCreditsNudgeEmailParams"; export type { SendAddCreditsNudgeEmailResponse } from "./SendAddCreditsNudgeEmailResponse"; export type { ServerRequestResolvedNotification } from "./ServerRequestResolvedNotification"; 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 486e121cf7..0c7c362b52 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs @@ -11,6 +11,8 @@ use super::TurnEnvironmentParams; use super::TurnItemsView; use super::shared::v2_enum_from_core; use codex_experimental_api_macros::ExperimentalApi; +pub use codex_protocol::capabilities::CapabilityRootLocation; +pub use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::Personality; use codex_protocol::config_types::ReasoningSummary; @@ -153,6 +155,10 @@ pub struct ThreadStartParams { #[experimental("thread/start.dynamicTools")] #[ts(optional = nullable)] pub dynamic_tools: Option>, + /// Capability roots selected for this thread by the hosting platform. + #[experimental("thread/start.selectedCapabilityRoots")] + #[ts(optional = nullable)] + pub selected_capability_roots: Option>, /// Test-only experimental field used to validate experimental gating and /// schema filtering behavior in a stable way. #[experimental("thread/start.mockExperimentalField")] diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 4bdc430a1a..d6f71ab82f 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -49,6 +49,7 @@ codex-hooks = { workspace = true } codex-otel = { workspace = true } codex-plugin = { workspace = true } codex-shell-command = { workspace = true } +codex-skills-extension = { workspace = true } codex-utils-cli = { workspace = true } codex-utils-pty = { workspace = true } codex-backend-client = { workspace = true } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index a06adf03e7..e09866e196 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -130,7 +130,7 @@ Example with notification opt-out: ## API Overview -- `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. For permissions, prefer experimental `permissions` profile selection by id; the legacy `sandbox` shorthand is still accepted but cannot be combined with `permissions`. Experimental `environments` selects the sticky execution environments for turns on the thread; omit it to use the server default, pass `[]` to disable environments, or pass explicit environment ids with per-environment `cwd`. +- `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. For permissions, prefer experimental `permissions` profile selection by id; the legacy `sandbox` shorthand is still accepted but cannot be combined with `permissions`. Experimental `environments` selects the sticky execution environments for turns on the thread; omit it to use the server default, pass `[]` to disable environments, or pass explicit environment ids with per-environment `cwd`. Experimental `selectedCapabilityRoots` selects environment-owned plugin or standalone-skill roots. Skills found below those roots are listed and read through the owning environment; other plugin capabilities are not activated yet. - `thread/resume` — reopen an existing thread by id so subsequent `turn/start` calls append to it. Accepts the same permission override rules as `thread/start`. - `thread/fork` — fork an existing thread into a new thread id by copying the stored history; if the source thread is currently mid-turn, the fork records the same interruption marker as `turn/interrupt` instead of inheriting an unmarked partial turn suffix. The returned `thread.forkedFromId` points at the source thread when known. Accepts `ephemeral: true` for an in-memory temporary fork, emits `thread/started` (including the current `thread.status`), and auto-subscribes you to turn/item events for the new thread. Experimental clients can pass `excludeTurns: true` when they plan to page fork history via `thread/turns/list` instead of receiving the full turn array immediately. Accepts the same permission override rules as `thread/start`. - `thread/start`, `thread/resume`, and `thread/fork` responses include the legacy `sandbox` compatibility projection. Experimental clients can read `runtimeWorkspaceRoots` for the thread-scoped runtime roots and `activePermissionProfile` for the named or implicit built-in profile identity/provenance when known. @@ -249,6 +249,18 @@ Start a fresh thread when you need a new Codex conversation. // "permissions": ":workspace" // Experimental runtime roots for :workspace_roots materialization: // "runtimeWorkspaceRoots": ["/Users/me/project", "/Users/me/openai"], + // Experimental capability roots selected by the hosting platform: + "selectedCapabilityRoots": [ + { + "id": "github@openai", + "location": { + "type": "environment", + "environmentId": "workspace", + // Opaque to app-server; interpreted in the selected environment. + "path": "/opt/cca/plugins/github" + } + } + ], // Do not send both "sandbox" and "permissions". "personality": "friendly", "serviceName": "my_app_server_client", // optional metrics tag (`service_name`) diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index 53f19997cd..b18331064f 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -32,6 +32,7 @@ pub(crate) fn thread_extensions( state_db: Option, thread_manager: Weak, goal_service: Arc, + executor_skill_provider: Arc, ) -> Arc> where S: AgentSpawner + 'static, @@ -51,6 +52,11 @@ where codex_memories_extension::install(&mut builder, codex_otel::global()); codex_web_search_extension::install(&mut builder, auth_manager.clone()); codex_image_generation_extension::install(&mut builder, auth_manager); + codex_skills_extension::install_with_providers( + &mut builder, + codex_skills_extension::SkillProviders::new() + .with_executor_provider(executor_skill_provider), + ); Arc::new(builder.build()) } diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 8c72533ba5..821af862d5 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -181,12 +181,19 @@ mod tests { .await .expect("refresh tests require state db"); let thread_store = thread_store_from_config(&good_config, Some(state_db.clone())); + let environment_manager = Arc::new(EnvironmentManager::default_for_tests()); + let executor_skill_provider: Arc = Arc::new( + codex_skills_extension::ExecutorSkillProvider::new_with_restriction_product( + Arc::clone(&environment_manager), + SessionSource::Exec.restriction_product(), + ), + ); let thread_manager = Arc::new_cyclic(|thread_manager| { ThreadManager::new( &good_config, auth_manager.clone(), SessionSource::Exec, - Arc::new(EnvironmentManager::default_for_tests()), + Arc::clone(&environment_manager), thread_extensions( guardian_agent_spawner(thread_manager.clone()), Arc::new(NoopExtensionEventSink), @@ -194,6 +201,7 @@ mod tests { Some(state_db.clone()), thread_manager.clone(), Arc::new(codex_goal_extension::GoalService::new()), + Arc::clone(&executor_skill_provider), ), /*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 84d3d88dda..78ba2c12cf 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -307,6 +307,14 @@ impl MessageProcessor { // resumed, or forked threads to a different persistence backend/root. let thread_store = codex_core::thread_store_from_config(config.as_ref(), state_db.clone()); let environment_manager_for_requests = Arc::clone(&environment_manager); + let environment_manager_for_extensions = Arc::clone(&environment_manager); + let restriction_product = session_source.restriction_product(); + let executor_skill_provider: Arc = Arc::new( + codex_skills_extension::ExecutorSkillProvider::new_with_restriction_product( + environment_manager_for_extensions, + restriction_product, + ), + ); let goal_service = Arc::new(GoalService::new()); let thread_manager = Arc::new_cyclic(|thread_manager| { ThreadManager::new( @@ -321,6 +329,7 @@ impl MessageProcessor { state_db.clone(), thread_manager.clone(), Arc::clone(&goal_service), + Arc::clone(&executor_skill_provider), ), Some(analytics_events_client.clone()), Arc::clone(&thread_store), 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 ec3bd57b07..be5cffb49a 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1,5 +1,7 @@ use super::*; use crate::error_code::method_not_found; +use codex_app_server_protocol::SelectedCapabilityRoot; +use codex_extension_api::ExtensionDataInit; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; @@ -833,6 +835,7 @@ impl ThreadRequestProcessor { base_instructions, developer_instructions, dynamic_tools, + selected_capability_roots, mock_experimental_field: _mock_experimental_field, experimental_raw_events, personality, @@ -888,6 +891,7 @@ impl ThreadRequestProcessor { config, typesafe_overrides, dynamic_tools, + selected_capability_roots.unwrap_or_default(), session_start_source, thread_source.map(Into::into), environment_selections, @@ -960,6 +964,7 @@ impl ThreadRequestProcessor { config_overrides: Option>, typesafe_overrides: ConfigOverrides, dynamic_tools: Option>, + selected_capability_roots: Vec, session_start_source: Option, thread_source: Option, environments: Option>, @@ -1061,6 +1066,10 @@ impl ThreadRequestProcessor { .collect() }; let core_dynamic_tool_count = core_dynamic_tools.len(); + let mut thread_extension_init = ExtensionDataInit::new(); + if !selected_capability_roots.is_empty() { + thread_extension_init.insert(selected_capability_roots); + } let create_thread_started_at = std::time::Instant::now(); let NewThread { thread_id, @@ -1083,6 +1092,7 @@ impl ThreadRequestProcessor { metrics_service_name: service_name, parent_trace: request_trace, environments, + thread_extension_init, }) .instrument(tracing::info_span!( "app_server.thread_start.create_thread", diff --git a/codex-rs/app-server/tests/suite/v2/executor_skills.rs b/codex-rs/app-server/tests/suite/v2/executor_skills.rs new file mode 100644 index 0000000000..ede5b20fb8 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/executor_skills.rs @@ -0,0 +1,148 @@ +use std::time::Duration; + +use anyhow::Result; +use app_test_support::TestAppServer; +use app_test_support::to_response; +use codex_app_server_protocol::CapabilityRootLocation; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::SelectedCapabilityRoot; +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 core_test_support::responses; +use tempfile::TempDir; +use tokio::time::timeout; + +const READ_TIMEOUT: Duration = Duration::from_secs(10); +const SKILL_NAME: &str = "demo-plugin:deploy"; +const SKILL_MARKER: &str = "EXECUTOR_SKILL_BODY_MARKER"; +const LOCAL_SKILL_MARKER: &str = "LOCAL_SKILL_BODY_MARKER"; + +#[tokio::test] +async fn selected_executor_root_exposes_plugin_skill() -> Result<()> { + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-selected"), + responses::ev_assistant_message("msg-selected", "Done"), + responses::ev_completed("resp-selected"), + ]), + ) + .await; + + let codex_home = TempDir::new()?; + std::fs::write( + codex_home.path().join("config.toml"), + format!( + r#" +model = "mock-model" +approval_policy = "never" +sandbox_mode = "read-only" +model_provider = "mock_provider" + +[skills] +include_instructions = true + +[model_providers.mock_provider] +name = "Mock provider for test" +base_url = "{}/v1" +wire_api = "responses" +request_max_retries = 0 +stream_max_retries = 0 +"#, + server.uri() + ), + )?; + let local_skill_dir = codex_home.path().join("skills/local-deploy"); + std::fs::create_dir_all(&local_skill_dir)?; + std::fs::write( + local_skill_dir.join("SKILL.md"), + format!( + "---\nname: {SKILL_NAME}\ndescription: Colliding local skill.\n---\n\n# Local deploy\n\n{LOCAL_SKILL_MARKER}\n" + ), + )?; + let plugin_dir = TempDir::new()?; + let manifest_dir = plugin_dir.path().join(".codex-plugin"); + let skill_dir = plugin_dir.path().join("skills/deploy"); + std::fs::create_dir_all(&manifest_dir)?; + std::fs::create_dir_all(&skill_dir)?; + std::fs::write( + manifest_dir.join("plugin.json"), + r#"{"name":"demo-plugin"}"#, + )?; + std::fs::write( + skill_dir.join("SKILL.md"), + format!( + "---\nname: deploy\ndescription: Deploy through the executor.\n---\n\n# Deploy\n\n{SKILL_MARKER}\n" + ), + )?; + + let mut app_server = TestAppServer::new(codex_home.path()).await?; + timeout(READ_TIMEOUT, app_server.initialize()).await??; + + let request_id = app_server + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + selected_capability_roots: Some(vec![SelectedCapabilityRoot { + id: "demo-plugin@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "local".to_string(), + path: plugin_dir.path().to_string_lossy().into_owned(), + }, + }]), + ..Default::default() + }) + .await?; + let response: JSONRPCResponse = timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response(response)?; + + let request_id = app_server + .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( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + timeout( + READ_TIMEOUT, + app_server.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let request = response_mock.single_request(); + assert!( + request + .message_input_texts("developer") + .iter() + .any(|text| text.contains(SKILL_NAME)) + ); + let skill_fragments = request + .message_input_texts("user") + .into_iter() + .filter(|text| text.starts_with("")) + .collect::>(); + assert_eq!(1, skill_fragments.len()); + let skill_fragment = skill_fragments + .first() + .expect("executor skill instructions should be model-visible"); + assert!(skill_fragment.contains(&format!("{SKILL_NAME}"))); + assert!(skill_fragment.contains(SKILL_MARKER)); + assert!(!skill_fragment.contains(LOCAL_SKILL_MARKER)); + + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 9aef40aade..77d4b37589 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -12,6 +12,7 @@ mod connection_handling_websocket; #[cfg(unix)] mod connection_handling_websocket_unix; mod dynamic_tools; +mod executor_skills; mod experimental_api; mod experimental_feature_list; mod external_agent_config; diff --git a/codex-rs/app-server/tests/suite/v2/skills_list.rs b/codex-rs/app-server/tests/suite/v2/skills_list.rs index cba2f4ee3f..147d40c47a 100644 --- a/codex-rs/app-server/tests/suite/v2/skills_list.rs +++ b/codex-rs/app-server/tests/suite/v2/skills_list.rs @@ -793,6 +793,7 @@ async fn skills_changed_notification_is_emitted_after_skill_change() -> Result<( thread_source: None, dynamic_tools: None, environments: None, + selected_capability_roots: None, mock_experimental_field: None, experimental_raw_events: false, }) diff --git a/codex-rs/core-skills/src/loader.rs b/codex-rs/core-skills/src/loader.rs index 7dac256010..e2c3896cba 100644 --- a/codex-rs/core-skills/src/loader.rs +++ b/codex-rs/core-skills/src/loader.rs @@ -169,8 +169,8 @@ where let mut file_systems_by_skill_path: HashMap> = HashMap::new(); for root in roots { - let root_path = canonicalize_for_skill_identity(&root.path); let fs = root.file_system; + let root_path = canonicalize_for_skill_identity(fs.as_ref(), &root.path).await; let skills_before_root = outcome.skills.len(); discover_skills_under_root( fs.as_ref(), @@ -471,8 +471,13 @@ fn dedupe_skill_roots_by_path(roots: &mut Vec) { roots.retain(|root| seen.insert(root.path.clone())); } -fn canonicalize_for_skill_identity(path: &AbsolutePathBuf) -> AbsolutePathBuf { - path.canonicalize().unwrap_or_else(|_| path.clone()) +async fn canonicalize_for_skill_identity( + fs: &dyn ExecutorFileSystem, + path: &AbsolutePathBuf, +) -> AbsolutePathBuf { + fs.canonicalize(path, /*sandbox*/ None) + .await + .unwrap_or_else(|_| path.clone()) } async fn discover_skills_under_root( @@ -483,8 +488,11 @@ async fn discover_skills_under_root( plugin_root: Option<&AbsolutePathBuf>, outcome: &mut SkillLoadOutcome, ) { - let root = canonicalize_for_skill_identity(root); - let plugin_root = plugin_root.map(canonicalize_for_skill_identity); + let root = root.clone(); + let plugin_root = match plugin_root { + Some(plugin_root) => Some(canonicalize_for_skill_identity(fs, plugin_root).await), + None => None, + }; match fs.get_metadata(&root, /*sandbox*/ None).await { Ok(metadata) if metadata.is_directory => {} @@ -557,7 +565,7 @@ async fn discover_skills_under_root( } match fs.read_directory(&path, /*sandbox*/ None).await { Ok(_) => { - let resolved_dir = canonicalize_for_skill_identity(&path); + let resolved_dir = canonicalize_for_skill_identity(fs, &path).await; enqueue_dir( &mut queue, &mut visited_dirs, @@ -582,7 +590,7 @@ async fn discover_skills_under_root( } if metadata.is_directory { - let resolved_dir = canonicalize_for_skill_identity(&path); + let resolved_dir = canonicalize_for_skill_identity(fs, &path).await; enqueue_dir( &mut queue, &mut visited_dirs, @@ -672,7 +680,7 @@ async fn parse_skill_file( )?; } - let resolved_path = canonicalize_for_skill_identity(path); + let resolved_path = canonicalize_for_skill_identity(fs, path).await; Ok(SkillMetadata { name, diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 5f816bd075..5b68ecf917 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -103,6 +103,7 @@ pub(crate) async fn run_codex_thread_interactive( parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), parent_trace: None, environment_selections: parent_ctx.environments.clone(), + thread_extension_init: codex_extension_api::ExtensionDataInit::default(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), attestation_provider: parent_session.services.attestation_provider.clone(), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index bb66e8684e..23dbe1f45e 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -51,6 +51,7 @@ use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; use codex_exec_server::FileSystemSandboxContext; +use codex_extension_api::ExtensionDataInit; use codex_extension_api::PromptSlot; use codex_features::FEATURES; use codex_features::Feature; @@ -417,6 +418,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) user_shell_override: Option, pub(crate) parent_trace: Option, pub(crate) environment_selections: ResolvedTurnEnvironments, + pub(crate) thread_extension_init: ExtensionDataInit, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, @@ -497,6 +499,7 @@ impl Codex { parent_rollout_thread_trace, parent_trace: _, environment_selections, + thread_extension_init, analytics_events_client, thread_store, attestation_provider, @@ -641,6 +644,7 @@ impl Codex { plugins_manager, mcp_manager.clone(), extensions, + thread_extension_init, agent_control, environment_manager, analytics_events_client, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index c6d0d9767e..e51a83113b 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -4,6 +4,7 @@ use crate::agents_md::LoadedAgentsMd; use crate::config::ConstraintError; use crate::skills::SkillError; use crate::state::ActiveTurn; +use codex_extension_api::ExtensionDataInit; use codex_protocol::SessionId; use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; use codex_protocol::config_types::ServiceTier; @@ -487,6 +488,7 @@ impl Session { plugins_manager: Arc, mcp_manager: Arc, extensions: Arc>, + thread_extension_init: ExtensionDataInit, agent_control: AgentControl, environment_manager: Arc, analytics_events_client: Option, @@ -961,8 +963,10 @@ impl Session { ); let session_extension_data = codex_extension_api::ExtensionData::new(session_id.to_string()); - let thread_extension_data = - codex_extension_api::ExtensionData::new(thread_id.to_string()); + let thread_extension_data = codex_extension_api::ExtensionData::new_with_init( + thread_id.to_string(), + thread_extension_init, + ); for contributor in extensions.thread_lifecycle_contributors() { contributor.on_thread_start(codex_extension_api::ThreadStartInput { config: config.as_ref(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index ef3f6dbfba..f6166260a0 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4690,6 +4690,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { plugins_manager, mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), + codex_extension_api::ExtensionDataInit::default(), AgentControl::default(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, @@ -5030,6 +5031,7 @@ async fn make_session_with_config_and_rx( plugins_manager, mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), + codex_extension_api::ExtensionDataInit::default(), AgentControl::default(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, @@ -5131,6 +5133,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( plugins_manager, mcp_manager, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), + codex_extension_api::ExtensionDataInit::default(), agent_control, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), /*analytics_events_client*/ None, diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index d8b46c9999..c73add2920 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -731,6 +731,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { environment_selections: ResolvedTurnEnvironments { turn_environments: Vec::new(), }, + thread_extension_init: codex_extension_api::ExtensionDataInit::default(), analytics_events_client: None, thread_store, attestation_provider: None, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 1cde8c9320..a0c1d63d71 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -21,6 +21,7 @@ use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; +use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_extension_api::empty_extension_registry; use codex_features::Feature; @@ -182,6 +183,7 @@ pub struct StartThreadOptions { pub metrics_service_name: Option, pub parent_trace: Option, pub environments: Vec, + pub thread_extension_init: ExtensionDataInit, } pub(crate) struct ResumeThreadWithHistoryOptions { @@ -576,6 +578,7 @@ impl ThreadManager { metrics_service_name: None, parent_trace: None, environments, + thread_extension_init: ExtensionDataInit::default(), })) .await } @@ -614,6 +617,7 @@ impl ThreadManager { /*inherited_exec_policy*/ None, options.parent_trace, options.environments, + options.thread_extension_init, /*user_shell_override*/ None, )) .await @@ -702,6 +706,7 @@ impl ThreadManager { /*inherited_exec_policy*/ None, parent_trace, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ None, )) .await @@ -728,6 +733,7 @@ impl ThreadManager { /*metrics_service_name*/ None, /*parent_trace*/ None, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ Some(user_shell_override), )) .await @@ -763,6 +769,7 @@ impl ThreadManager { /*inherited_exec_policy*/ None, /*parent_trace*/ None, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ Some(user_shell_override), )) .await @@ -931,6 +938,7 @@ impl ThreadManager { /*metrics_service_name*/ None, parent_trace, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ None, )) .await @@ -1136,6 +1144,7 @@ impl ThreadManagerState { inherited_exec_policy, /*parent_trace*/ None, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ None, )) .await @@ -1172,6 +1181,7 @@ impl ThreadManagerState { inherited_exec_policy, /*parent_trace*/ None, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ None, )) .await @@ -1209,6 +1219,7 @@ impl ThreadManagerState { inherited_exec_policy, /*parent_trace*/ None, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ None, )) .await @@ -1229,6 +1240,7 @@ impl ThreadManagerState { metrics_service_name: Option, parent_trace: Option, environments: Vec, + thread_extension_init: ExtensionDataInit, user_shell_override: Option, ) -> CodexResult { Box::pin(self.spawn_thread_with_source( @@ -1246,6 +1258,7 @@ impl ThreadManagerState { /*inherited_exec_policy*/ None, parent_trace, environments, + thread_extension_init, user_shell_override, )) .await @@ -1268,6 +1281,7 @@ impl ThreadManagerState { inherited_exec_policy: Option>, parent_trace: Option, environments: Vec, + thread_extension_init: ExtensionDataInit, user_shell_override: Option, ) -> CodexResult { let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_)); @@ -1332,6 +1346,7 @@ impl ThreadManagerState { user_shell_override, parent_trace, environment_selections, + thread_extension_init, analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 116473258e..bb0c0af314 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -331,6 +331,7 @@ async fn start_thread_rejects_explicit_local_environment_when_default_provider_i environment_id: "local".to_string(), cwd: config.cwd.clone(), }], + thread_extension_init: Default::default(), }) .await; let err = match result { @@ -368,6 +369,7 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { metrics_service_name: None, parent_trace: None, environments: Vec::new(), + thread_extension_init: Default::default(), }) .await .expect("internal thread should start"); @@ -384,6 +386,81 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { assert!(manager.list_thread_ids().await.is_empty()); } +#[tokio::test] +async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() { + struct InitialMarker(&'static str); + + struct InitialDataRecorder { + 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(), + )); + } + } + + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let observed = Arc::new(std::sync::Mutex::new(None)); + let mut extensions = codex_extension_api::ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(Arc::new(InitialDataRecorder { + observed: Arc::clone(&observed), + })); + let manager = ThreadManager::new( + &config, + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()), + SessionSource::Exec, + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + Arc::new(extensions.build()), + /*analytics_events_client*/ None, + thread_store_from_config(&config, /*state_db*/ None), + /*state_db*/ None, + TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, + ); + let mut thread_extension_init = codex_extension_api::ExtensionDataInit::new(); + thread_extension_init.insert(InitialMarker("seeded")); + + let thread = manager + .start_thread_with_options(StartThreadOptions { + config, + initial_history: InitialHistory::New, + session_source: None, + thread_source: None, + dynamic_tools: Vec::new(), + metrics_service_name: None, + parent_trace: None, + environments: Vec::new(), + thread_extension_init, + }) + .await + .expect("start thread"); + + assert_eq!( + observed + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(), + Some((thread.thread_id.to_string(), "seeded".to_string())) + ); +} + #[tokio::test] async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { let temp_dir = tempdir().expect("tempdir"); @@ -423,6 +500,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { metrics_service_name: None, parent_trace: None, environments: environments.clone(), + thread_extension_init: Default::default(), }) .await .expect("start source thread"); @@ -693,6 +771,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { metrics_service_name: None, parent_trace: None, environments: Vec::new(), + thread_extension_init: Default::default(), }) .await .expect("start source thread"); diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 4d6cd17bb8..c3468e4653 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -757,6 +757,7 @@ async fn subagent_stop_replaces_stop_and_skips_internal_subagents() -> Result<() metrics_service_name: None, parent_trace: None, environments: Vec::new(), + thread_extension_init: Default::default(), }) .await?; diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 7fa60c0fe7..7d8d95d139 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -59,3 +59,4 @@ pub use registry::ExtensionRegistry; pub use registry::ExtensionRegistryBuilder; pub use registry::empty_extension_registry; pub use state::ExtensionData; +pub use state::ExtensionDataInit; diff --git a/codex-rs/ext/extension-api/src/state.rs b/codex-rs/ext/extension-api/src/state.rs index aab37f5059..97152df5a3 100644 --- a/codex-rs/ext/extension-api/src/state.rs +++ b/codex-rs/ext/extension-api/src/state.rs @@ -7,6 +7,32 @@ use std::sync::PoisonError; type ErasedData = Arc; +/// Typed values supplied before an [`ExtensionData`] scope is created. +/// +/// Hosts consume this value once to seed a scope before lifecycle contributors +/// run. It does not install extensions or provide persistence. +#[derive(Debug, Default)] +pub struct ExtensionDataInit { + entries: HashMap, +} + +impl ExtensionDataInit { + /// Creates an empty extension data initializer. + pub fn new() -> Self { + Self::default() + } + + /// Stores `value` as the initial attachment of type `T`. + pub fn insert(&mut self, value: T) -> Option> + where + T: Any + Send + Sync, + { + self.entries + .insert(TypeId::of::(), Arc::new(value)) + .map(downcast_data) + } +} + /// Typed extension-owned data attached to one host object. #[derive(Debug)] pub struct ExtensionData { @@ -17,9 +43,14 @@ pub struct ExtensionData { impl ExtensionData { /// Creates an empty attachment map for one host-owned scope. pub fn new(level_id: impl Into) -> Self { + Self::new_with_init(level_id, ExtensionDataInit::default()) + } + + /// Creates an attachment map seeded with host-supplied initial data. + pub fn new_with_init(level_id: impl Into, init: ExtensionDataInit) -> Self { Self { level_id: level_id.into(), - entries: Mutex::new(HashMap::new()), + entries: Mutex::new(init.entries), } } diff --git a/codex-rs/ext/skills/Cargo.toml b/codex-rs/ext/skills/Cargo.toml index 29767d58fe..9ff42cfafc 100644 --- a/codex-rs/ext/skills/Cargo.toml +++ b/codex-rs/ext/skills/Cargo.toml @@ -17,8 +17,11 @@ workspace = true async-trait = { workspace = true } codex-core = { workspace = true } codex-core-skills = { workspace = true } +codex-exec-server = { workspace = true } codex-extension-api = { workspace = true } codex-protocol = { workspace = true } +codex-utils-absolute-path = { workspace = true } +codex-utils-string = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/codex-rs/ext/skills/src/catalog.rs b/codex-rs/ext/skills/src/catalog.rs index a09a81ff79..c966cc09e9 100644 --- a/codex-rs/ext/skills/src/catalog.rs +++ b/codex-rs/ext/skills/src/catalog.rs @@ -1,4 +1,5 @@ use codex_core_skills::model::SkillDependencies; +use codex_exec_server::EnvironmentPathRef; /// Source authority that owns a skill package and must be used to read it. #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -56,9 +57,37 @@ impl SkillAuthority { #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct SkillPackageId(pub String); -/// Opaque resource id inside a skill package. +/// Opaque resource id inside a skill package, optionally bound to the +/// environment path that owns its contents. #[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct SkillResourceId(pub String); +pub struct SkillResourceId { + id: String, + environment_path: Option, +} + +impl SkillResourceId { + pub fn new(id: impl Into) -> Self { + Self { + id: id.into(), + environment_path: None, + } + } + + pub fn environment(id: impl Into, path: EnvironmentPathRef) -> Self { + Self { + id: id.into(), + environment_path: Some(path), + } + } + + pub fn as_str(&self) -> &str { + &self.id + } + + pub(crate) fn environment_path(&self) -> Option<&EnvironmentPathRef> { + self.environment_path.as_ref() + } +} /// Metadata shown in the always-visible skills catalog. #[derive(Clone, Debug, PartialEq, Eq)] @@ -125,7 +154,7 @@ impl SkillCatalogEntry { pub(crate) fn rendered_path(&self) -> &str { self.display_path .as_deref() - .unwrap_or(self.main_prompt.0.as_str()) + .unwrap_or_else(|| self.main_prompt.as_str()) } } diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index 3961d7d85d..45d7c95cb7 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -6,27 +6,32 @@ 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::ExtensionRegistryBuilder; +use codex_extension_api::PromptFragment; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadStartInput; use codex_extension_api::TurnInputContext; use codex_extension_api::TurnInputContributor; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::WarningEvent; -use crate::catalog::SkillAuthority; use crate::catalog::SkillCatalogEntry; use crate::catalog::SkillReadResult; use crate::catalog::SkillSourceKind; use crate::provider::HostSkillProvider; use crate::provider::SkillListQuery; use crate::provider::SkillReadRequest; +use crate::render::MAX_SKILL_NAME_BYTES; +use crate::render::MAX_SKILL_PATH_BYTES; use crate::render::available_skills_fragment; 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; @@ -42,11 +47,15 @@ struct SkillsExtension { #[async_trait::async_trait] impl ThreadLifecycleContributor for SkillsExtension { async fn on_thread_start(&self, input: ThreadStartInput<'_, Config>) { - input + let selected_roots = input .thread_store - .insert(SkillsThreadState::new(SkillsExtensionConfig::from_config( - input.config, - ))); + .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, + )); } } @@ -62,11 +71,47 @@ impl ConfigContributor for SkillsExtension { if let Some(state) = thread_store.get::() { state.set_config(next_config); } else { - thread_store.insert(SkillsThreadState::new(next_config)); + thread_store.insert(SkillsThreadState::new(next_config, Vec::new())); } } } +impl ContextContributor for SkillsExtension { + fn contribute<'a>( + &'a self, + _session_store: &'a ExtensionData, + thread_store: &'a ExtensionData, + ) -> std::pin::Pin> + Send + 'a>> { + Box::pin(async move { + let Some(thread_state) = thread_store.get::() else { + return Vec::new(); + }; + let config = thread_state.config(); + if !config.include_instructions || thread_state.selected_roots().is_empty() { + 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, + }) + .await; + for warning in &catalog.warnings { + self.emit_warning(thread_store.level_id(), warning.clone()); + } + available_skills_fragment(&catalog) + .map(|fragment| PromptFragment::developer_capability(fragment.render())) + .into_iter() + .collect() + }) + } +} + #[async_trait::async_trait] impl TurnInputContributor for SkillsExtension { async fn contribute( @@ -84,16 +129,7 @@ impl TurnInputContributor for SkillsExtension { let host_loaded_skills = turn_store.get::(); let query = SkillListQuery { turn_id: input.turn_id.clone(), - executor_authorities: input - .environments - .iter() - .map(|environment| { - SkillAuthority::new( - SkillSourceKind::Executor, - environment.environment_id.clone(), - ) - }) - .collect(), + executor_roots: thread_state.selected_roots().to_vec(), host: host_loaded_skills.clone(), include_host_skills: true, include_bundled_skills: config.bundled_skills_enabled, @@ -106,10 +142,14 @@ impl TurnInputContributor for SkillsExtension { let selected_entries = collect_explicit_skill_mentions(&input.user_input, &catalog); let mut fragments: Vec> = Vec::new(); - if config.include_instructions - && let Some(fragment) = available_skills_fragment(&catalog) - { - fragments.push(Box::new(fragment)); + 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(); @@ -132,14 +172,14 @@ impl TurnInputContributor for SkillsExtension { warnings.push(warning); } let injection = SkillInjection { - name: entry.name.clone(), - path: entry.rendered_path().to_string(), + 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.0.clone()); + injected_host_skill_prompts.insert_path(entry.main_prompt.as_str()); } } Err(message) => { @@ -150,6 +190,23 @@ impl TurnInputContributor for SkillsExtension { } } + 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, @@ -206,5 +263,6 @@ pub fn install_with_providers( }); registry.thread_lifecycle_contributor(extension.clone()); registry.config_contributor(extension.clone()); + registry.prompt_contributor(extension.clone()); registry.turn_input_contributor(extension); } diff --git a/codex-rs/ext/skills/src/lib.rs b/codex-rs/ext/skills/src/lib.rs index 3a28c5fbdf..5ab03ddd9e 100644 --- a/codex-rs/ext/skills/src/lib.rs +++ b/codex-rs/ext/skills/src/lib.rs @@ -8,6 +8,8 @@ mod state; pub use extension::install; pub use extension::install_with_providers; +pub use provider::ExecutorSkillProvider; pub use provider::HostSkillProvider; +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 a968774afd..8d13a2f981 100644 --- a/codex-rs/ext/skills/src/provider.rs +++ b/codex-rs/ext/skills/src/provider.rs @@ -2,9 +2,11 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +mod executor; mod host; use codex_core_skills::HostLoadedSkills; +use codex_protocol::capabilities::SelectedCapabilityRoot; use crate::catalog::SkillAuthority; use crate::catalog::SkillCatalog; @@ -14,12 +16,13 @@ use crate::catalog::SkillReadResult; use crate::catalog::SkillResourceId; use crate::catalog::SkillSearchResult; +pub use executor::ExecutorSkillProvider; pub use host::HostSkillProvider; #[derive(Clone, Debug)] pub struct SkillListQuery { pub turn_id: String, - pub executor_authorities: Vec, + pub executor_roots: Vec, pub host: Option>, pub include_host_skills: bool, pub include_bundled_skills: bool, diff --git a/codex-rs/ext/skills/src/provider/executor.rs b/codex-rs/ext/skills/src/provider/executor.rs new file mode 100644 index 0000000000..e23e5a95bd --- /dev/null +++ b/codex-rs/ext/skills/src/provider/executor.rs @@ -0,0 +1,197 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use codex_core_skills::SkillMetadata; +use codex_core_skills::filter_skill_load_outcome_for_product; +use codex_core_skills::loader::SkillRoot; +use codex_core_skills::loader::load_skills_from_roots; +use codex_exec_server::EnvironmentManager; +use codex_exec_server::EnvironmentPathRef; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::protocol::Product; +use codex_protocol::protocol::SkillScope; +use codex_utils_absolute_path::AbsolutePathBuf; + +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; + +/// Discovers and reads skills through the filesystem owned by an execution environment. +#[derive(Clone, Debug)] +pub struct ExecutorSkillProvider { + environment_manager: Arc, + restriction_product: Option, +} + +impl ExecutorSkillProvider { + pub fn new_with_restriction_product( + environment_manager: Arc, + restriction_product: Option, + ) -> Self { + Self { + environment_manager, + restriction_product, + } + } +} + +impl SkillProvider for ExecutorSkillProvider { + fn list(&self, query: SkillListQuery) -> SkillProviderFuture<'_, SkillCatalog> { + Box::pin(async move { + let mut catalog = SkillCatalog::default(); + for selected_root in query.executor_roots { + let selected_root_id = selected_root.id; + let CapabilityRootLocation::Environment { + environment_id, + path, + } = selected_root.location; + let authority = + SkillAuthority::new(SkillSourceKind::Executor, selected_root_id.clone()); + let Some(environment) = self.environment_manager.get_environment(&environment_id) + else { + catalog.warnings.push(format!( + "Selected capability root `{selected_root_id}` references unavailable environment `{environment_id}`." + )); + continue; + }; + let root_path = match executor_absolute_path(&path) { + Ok(root_path) => root_path, + Err(err) => { + catalog.warnings.push(format!( + "Selected capability root `{selected_root_id}` has invalid path `{path}`: {err}" + )); + continue; + } + }; + let file_system = environment.get_filesystem(); + let outcome = filter_skill_load_outcome_for_product( + load_skills_from_roots([SkillRoot { + path: root_path.clone(), + scope: SkillScope::User, + file_system: Arc::clone(&file_system), + plugin_id: None, + plugin_root: None, + }]) + .await, + self.restriction_product, + ); + catalog.warnings.extend(outcome.errors.iter().map(|err| { + format!( + "Failed to load executor skill at {}: {}", + err.path.display(), + err.message + ) + })); + for (skill, enabled) in outcome.skills_with_enabled() { + catalog.push_entry(catalog_entry_from_skill( + skill, + enabled, + authority.clone(), + &selected_root_id, + Arc::clone(&file_system), + )); + } + } + + Ok(catalog) + }) + } + + fn read(&self, request: SkillReadRequest) -> SkillProviderFuture<'_, SkillReadResult> { + Box::pin(async move { + if request.authority.kind != SkillSourceKind::Executor { + return Err(SkillProviderError::new(format!( + "executor skill provider cannot read {} resources", + request.authority.kind + ))); + } + if request.package.0 != request.resource.as_str() { + return Err(SkillProviderError::new( + "executor skill resource does not match its package", + )); + } + let Some(resource_path) = request.resource.environment_path() else { + return Err(SkillProviderError::new( + "executor skill resource is not bound to an environment", + )); + }; + let contents = resource_path + .read_to_string(/*sandbox*/ None) + .await + .map_err(|err| { + SkillProviderError::new(format!( + "failed to read executor skill resource {}: {err}", + 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_skill( + skill: &SkillMetadata, + enabled: bool, + authority: SkillAuthority, + selected_root_id: &str, + file_system: Arc, +) -> SkillCatalogEntry { + let skill_path = skill.path_to_skills_md.to_string_lossy().into_owned(); + let normalized_path = skill_path.replace('\\', "/"); + let display_path = format!( + "skill://{selected_root_id}/{}", + normalized_path.trim_start_matches('/') + ); + let mut entry = SkillCatalogEntry::new( + SkillPackageId(display_path.clone()), + authority, + skill.name.clone(), + skill.description.clone(), + SkillResourceId::environment( + display_path.clone(), + EnvironmentPathRef::new(file_system, skill.path_to_skills_md.clone()), + ), + ) + .with_short_description(skill.short_description.clone()) + .with_display_path(display_path) + .with_dependencies(skill.dependencies.clone()); + + if !enabled { + entry = entry.disabled(); + } + if !skill.allows_implicit_invocation() { + entry = entry.hidden_from_prompt(); + } + + entry +} + +fn executor_absolute_path(path: &str) -> std::io::Result { + let path = PathBuf::from(path); + if !path.is_absolute() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "executor path must be absolute", + )); + } + AbsolutePathBuf::from_absolute_path_checked(path) +} diff --git a/codex-rs/ext/skills/src/provider/host.rs b/codex-rs/ext/skills/src/provider/host.rs index 27293aa1d5..63d0d4588c 100644 --- a/codex-rs/ext/skills/src/provider/host.rs +++ b/codex-rs/ext/skills/src/provider/host.rs @@ -55,12 +55,12 @@ impl SkillProvider for HostSkillProvider { }; let Some(skill) = host_loaded_skills.outcome().skills.iter().find(|skill| { let skill_path = skill.path_to_skills_md.to_string_lossy(); - skill_path == request.resource.0.as_str() - || skill_path.replace('\\', "/") == request.resource.0 + skill_path == request.resource.as_str() + || skill_path.replace('\\', "/") == request.resource.as_str() }) else { return Err(SkillProviderError::new(format!( "host skill resource is not loaded: {}", - request.resource.0 + request.resource.as_str() ))); }; @@ -70,7 +70,7 @@ impl SkillProvider for HostSkillProvider { .map_err(|err| { SkillProviderError::new(format!( "failed to read host skill resource {}: {err}", - request.resource.0 + request.resource.as_str() )) })?; @@ -117,7 +117,7 @@ fn catalog_entry_from_skill(skill: &SkillMetadata, enabled: bool) -> SkillCatalo SkillAuthority::new(SkillSourceKind::Host, HOST_AUTHORITY_ID), skill.name.clone(), skill.description.clone(), - SkillResourceId(skill_path), + SkillResourceId::new(skill_path), ) .with_short_description(skill.short_description.clone()) .with_display_path(display_path) diff --git a/codex-rs/ext/skills/src/render.rs b/codex-rs/ext/skills/src/render.rs index 157ae4bff3..6b77d02140 100644 --- a/codex-rs/ext/skills/src/render.rs +++ b/codex-rs/ext/skills/src/render.rs @@ -2,11 +2,14 @@ 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; -const MAX_AVAILABLE_SKILLS_CHARS: usize = 8_000; -const MAX_MAIN_PROMPT_CHARS: usize = 40_000; +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 { @@ -32,7 +35,7 @@ impl ContextualUserFragment for AvailableSkillsFragment { } pub(crate) fn available_skills_fragment(catalog: &SkillCatalog) -> Option { - let mut total_chars = 0usize; + let mut total_bytes = 0usize; let mut omitted = 0usize; let mut skill_lines = Vec::new(); @@ -46,12 +49,12 @@ pub(crate) fn available_skills_fragment(catalog: &SkillCatalog) -> Option MAX_AVAILABLE_SKILLS_CHARS { + let next_bytes = total_bytes.saturating_add(line.len()); + if next_bytes > MAX_AVAILABLE_SKILLS_BYTES { omitted = omitted.saturating_add(1); continue; } - total_chars = next_chars; + total_bytes = next_bytes; skill_lines.push(line); } @@ -79,12 +82,10 @@ fn render_skill_line(name: &str, description: &str, path: &str) -> String { } pub(crate) fn truncate_main_prompt_contents(contents: &str) -> (String, bool) { - let mut chars = 0usize; - for (index, _) in contents.char_indices() { - if chars == MAX_MAIN_PROMPT_CHARS { - return (contents[..index].to_string(), true); - } - chars = chars.saturating_add(1); - } - (contents.to_string(), false) + truncate_utf8_to_bytes(contents, MAX_MAIN_PROMPT_BYTES) +} + +pub(crate) fn truncate_utf8_to_bytes(contents: &str, max_bytes: usize) -> (String, bool) { + let truncated = take_bytes_at_char_boundary(contents, max_bytes); + (truncated.to_string(), truncated.len() < contents.len()) } diff --git a/codex-rs/ext/skills/src/selection.rs b/codex-rs/ext/skills/src/selection.rs index c4405142e5..cfea2e3abd 100644 --- a/codex-rs/ext/skills/src/selection.rs +++ b/codex-rs/ext/skills/src/selection.rs @@ -93,12 +93,12 @@ fn push_selected( } fn entry_matches_path(entry: &SkillCatalogEntry, path: &str) -> bool { - entry.main_prompt.0 == path + entry.main_prompt.as_str() == path || entry.id.0 == path || entry .display_path .as_deref() - .is_some_and(|display_path| display_path == path) + .is_some_and(|display_path| normalize_skill_path(display_path) == path) } fn path_is_skill(path: &str) -> bool { diff --git a/codex-rs/ext/skills/src/sources.rs b/codex-rs/ext/skills/src/sources.rs index 9049fcb550..5ed51b878e 100644 --- a/codex-rs/ext/skills/src/sources.rs +++ b/codex-rs/ext/skills/src/sources.rs @@ -46,7 +46,7 @@ impl SkillProviderSource { fn should_list(&self, query: &SkillListQuery) -> bool { match &self.kind { SkillSourceKind::Host => query.include_host_skills, - SkillSourceKind::Executor => !query.executor_authorities.is_empty(), + SkillSourceKind::Executor => !query.executor_roots.is_empty(), SkillSourceKind::Remote => query.include_remote_skills, SkillSourceKind::Custom(_) => true, } diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index 4a639eb5f0..aa77cce822 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/codex-rs/ext/skills/src/state.rs @@ -1,4 +1,5 @@ use codex_core::config::Config; +use codex_protocol::capabilities::SelectedCapabilityRoot; use std::sync::Mutex; use crate::catalog::SkillCatalog; @@ -22,12 +23,17 @@ impl SkillsExtensionConfig { #[derive(Debug)] pub(crate) struct SkillsThreadState { config: Mutex, + selected_roots: Vec, } impl SkillsThreadState { - pub(crate) fn new(config: SkillsExtensionConfig) -> Self { + pub(crate) fn new( + config: SkillsExtensionConfig, + selected_roots: Vec, + ) -> Self { Self { config: Mutex::new(config), + selected_roots, } } @@ -44,6 +50,10 @@ impl SkillsThreadState { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = config; } + + pub(crate) fn selected_roots(&self) -> &[SelectedCapabilityRoot] { + &self.selected_roots + } } #[derive(Clone, Debug, Default, PartialEq, Eq)] diff --git a/codex-rs/ext/skills/tests/executor_file_system_authority.rs b/codex-rs/ext/skills/tests/executor_file_system_authority.rs new file mode 100644 index 0000000000..026c3f3348 --- /dev/null +++ b/codex-rs/ext/skills/tests/executor_file_system_authority.rs @@ -0,0 +1,337 @@ +use std::io; +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use async_trait::async_trait; +use codex_core_skills::HostLoadedSkills; +use codex_core_skills::loader::SkillRoot; +use codex_core_skills::loader::load_skills_from_roots; +use codex_exec_server::CopyOptions; +use codex_exec_server::CreateDirectoryOptions; +use codex_exec_server::EnvironmentManager; +use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::FileMetadata; +use codex_exec_server::FileSystemResult; +use codex_exec_server::FileSystemSandboxContext; +use codex_exec_server::ReadDirectoryEntry; +use codex_exec_server::RemoveOptions; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::protocol::SkillScope; +use codex_skills_extension::ExecutorSkillProvider; +use codex_skills_extension::catalog::SkillReadResult; +use codex_skills_extension::provider::SkillListQuery; +use codex_skills_extension::provider::SkillProvider; +use codex_skills_extension::provider::SkillReadRequest; +use codex_utils_absolute_path::AbsolutePathBuf; +use pretty_assertions::assert_eq; + +const SKILL_CONTENTS: &str = + "---\nname: synthetic\ndescription: Synthetic executor skill.\n---\n\nEXECUTOR_ONLY_BODY\n"; +static NEXT_TEST_ROOT_ID: AtomicUsize = AtomicUsize::new(0); + +struct SyntheticFileSystem { + alias_root: AbsolutePathBuf, + canonical_root: AbsolutePathBuf, +} + +impl SyntheticFileSystem { + fn metadata(&self, path: &AbsolutePathBuf) -> io::Result { + let skill_dir = self.canonical_root.join("skill"); + let skill_path = skill_dir.join("SKILL.md"); + let (is_directory, is_file) = if path == &self.canonical_root || path == &skill_dir { + (true, false) + } else if path == &skill_path { + (false, true) + } else { + return Err(io::Error::new(io::ErrorKind::NotFound, "not found")); + }; + Ok(FileMetadata { + is_directory, + is_file, + is_symlink: false, + created_at_ms: 0, + modified_at_ms: 0, + }) + } +} + +#[async_trait] +impl ExecutorFileSystem for SyntheticFileSystem { + async fn canonicalize( + &self, + path: &AbsolutePathBuf, + _sandbox: Option<&FileSystemSandboxContext>, + ) -> FileSystemResult { + if path == &self.alias_root { + return Ok(self.canonical_root.clone()); + } + self.metadata(path)?; + Ok(path.clone()) + } + + async fn join( + &self, + base_path: &AbsolutePathBuf, + path: &Path, + ) -> FileSystemResult { + Ok(base_path.join(path)) + } + + async fn parent(&self, path: &AbsolutePathBuf) -> FileSystemResult> { + Ok(path.parent()) + } + + async fn read_file( + &self, + path: &AbsolutePathBuf, + _sandbox: Option<&FileSystemSandboxContext>, + ) -> FileSystemResult> { + if path == &self.canonical_root.join("skill/SKILL.md") { + Ok(SKILL_CONTENTS.as_bytes().to_vec()) + } else { + Err(io::Error::new(io::ErrorKind::NotFound, "not found")) + } + } + + async fn write_file( + &self, + _path: &AbsolutePathBuf, + _contents: Vec, + _sandbox: Option<&FileSystemSandboxContext>, + ) -> FileSystemResult<()> { + Err(io::Error::new(io::ErrorKind::Unsupported, "read only")) + } + + async fn create_directory( + &self, + _path: &AbsolutePathBuf, + _options: CreateDirectoryOptions, + _sandbox: Option<&FileSystemSandboxContext>, + ) -> FileSystemResult<()> { + Err(io::Error::new(io::ErrorKind::Unsupported, "read only")) + } + + async fn get_metadata( + &self, + path: &AbsolutePathBuf, + _sandbox: Option<&FileSystemSandboxContext>, + ) -> FileSystemResult { + self.metadata(path) + } + + async fn read_directory( + &self, + path: &AbsolutePathBuf, + _sandbox: Option<&FileSystemSandboxContext>, + ) -> FileSystemResult> { + if path == &self.canonical_root { + Ok(vec![ReadDirectoryEntry { + file_name: "skill".to_string(), + is_directory: true, + is_file: false, + }]) + } else if path == &self.canonical_root.join("skill") { + Ok(vec![ReadDirectoryEntry { + file_name: "SKILL.md".to_string(), + is_directory: false, + is_file: true, + }]) + } else { + Err(io::Error::new(io::ErrorKind::NotFound, "not found")) + } + } + + async fn remove( + &self, + _path: &AbsolutePathBuf, + _options: RemoveOptions, + _sandbox: Option<&FileSystemSandboxContext>, + ) -> FileSystemResult<()> { + Err(io::Error::new(io::ErrorKind::Unsupported, "read only")) + } + + async fn copy( + &self, + _source_path: &AbsolutePathBuf, + _destination_path: &AbsolutePathBuf, + _options: CopyOptions, + _sandbox: Option<&FileSystemSandboxContext>, + ) -> FileSystemResult<()> { + Err(io::Error::new(io::ErrorKind::Unsupported, "read only")) + } +} + +#[tokio::test] +async fn skill_loading_and_reads_use_the_supplied_executor_file_system() { + let test_root = + std::env::temp_dir().join(format!("codex-executor-skill-fs-{}", std::process::id())); + let alias_root = AbsolutePathBuf::from_absolute_path_checked(test_root.join("alias")) + .expect("absolute path"); + let canonical_root = AbsolutePathBuf::from_absolute_path_checked(test_root.join("canonical")) + .expect("absolute path"); + assert!(!alias_root.as_path().exists()); + assert!(!canonical_root.as_path().exists()); + + let outcome = load_skills_from_roots([SkillRoot { + path: alias_root.clone(), + scope: SkillScope::User, + file_system: Arc::new(SyntheticFileSystem { + alias_root, + canonical_root: canonical_root.clone(), + }), + plugin_id: None, + plugin_root: None, + }]) + .await; + assert_eq!(outcome.errors, Vec::new()); + assert_eq!(outcome.skills.len(), 1); + + let skill = outcome.skills[0].clone(); + assert_eq!(skill.name, "synthetic"); + assert_eq!( + skill.path_to_skills_md, + canonical_root.join("skill/SKILL.md") + ); + let loaded = HostLoadedSkills::new(Arc::new(outcome)); + assert_eq!( + loaded.read_skill_text(&skill).await.expect("skill body"), + SKILL_CONTENTS + ); +} + +#[tokio::test] +async fn executor_provider_reads_from_the_environment_instance_used_for_listing() { + let test_root = create_local_skill_root("bound-instance").expect("create local skill root"); + let root_path = test_root.to_string_lossy().into_owned(); + let environment_manager = Arc::new(EnvironmentManager::default_for_tests()); + let provider = ExecutorSkillProvider::new_with_restriction_product( + Arc::clone(&environment_manager), + /*restriction_product*/ None, + ); + let catalog = provider + .list(SkillListQuery { + turn_id: "turn-1".to_string(), + executor_roots: vec![SelectedCapabilityRoot { + id: "root-a".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "local".to_string(), + path: root_path, + }, + }], + host: None, + include_host_skills: false, + include_bundled_skills: true, + include_remote_skills: false, + }) + .await + .expect("list executor skills"); + let entry = catalog + .entries + .into_iter() + .next() + .expect("listed executor skill"); + let resource = entry.main_prompt.clone(); + + environment_manager + .upsert_environment("local".to_string(), "http://127.0.0.1:1".to_string()) + .expect("replace environment"); + + assert_eq!( + provider + .read(SkillReadRequest { + authority: entry.authority, + package: entry.id, + resource: resource.clone(), + host: None, + }) + .await + .expect("read bound executor skill"), + SkillReadResult { + resource, + contents: SKILL_CONTENTS.to_string(), + } + ); + + std::fs::remove_dir_all(test_root).expect("remove skill directory"); +} + +#[tokio::test] +async fn selected_root_id_distinguishes_identical_executor_paths() { + let test_root = create_local_skill_root("root-identity").expect("create local skill root"); + let root_path = test_root.to_string_lossy().into_owned(); + let canonical_root = AbsolutePathBuf::from_absolute_path_checked(&test_root) + .expect("absolute skill root") + .canonicalize() + .expect("canonicalize skill root") + .to_string_lossy() + .replace('\\', "/"); + let provider = ExecutorSkillProvider::new_with_restriction_product( + Arc::new(EnvironmentManager::default_for_tests()), + /*restriction_product*/ None, + ); + + let catalog = provider + .list(SkillListQuery { + turn_id: "turn-1".to_string(), + executor_roots: ["root-a", "root-b"] + .into_iter() + .map(|id| SelectedCapabilityRoot { + id: id.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "local".to_string(), + path: root_path.clone(), + }, + }) + .collect(), + host: None, + include_host_skills: false, + include_bundled_skills: true, + include_remote_skills: false, + }) + .await + .expect("list executor skills"); + + assert_eq!( + catalog + .entries + .iter() + .map(|entry| ( + entry.authority.id.clone(), + entry.display_path.clone().expect("display path"), + )) + .collect::>(), + vec![ + ( + "root-a".to_string(), + format!( + "skill://root-a/{}/skill/SKILL.md", + canonical_root.trim_start_matches('/') + ), + ), + ( + "root-b".to_string(), + format!( + "skill://root-b/{}/skill/SKILL.md", + canonical_root.trim_start_matches('/') + ), + ), + ] + ); + + std::fs::remove_dir_all(test_root).expect("remove skill directory"); +} + +fn create_local_skill_root(label: &str) -> io::Result { + let id = NEXT_TEST_ROOT_ID.fetch_add(1, Ordering::Relaxed); + let test_root = std::env::temp_dir().join(format!( + "codex-executor-skill-{label}-{}-{id}", + std::process::id() + )); + let skill_dir = test_root.join("skill"); + std::fs::create_dir_all(&skill_dir)?; + std::fs::write(skill_dir.join("SKILL.md"), SKILL_CONTENTS)?; + Ok(test_root) +} diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index 1e2a2f0b60..dc1aa01407 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -14,7 +14,8 @@ use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ThreadStartInput; use codex_extension_api::TurnInputContext; -use codex_extension_api::TurnInputEnvironment; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use codex_protocol::protocol::SessionSource; use codex_protocol::user_input::UserInput; @@ -112,6 +113,7 @@ 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()); @@ -128,42 +130,35 @@ async fn installed_extension_loads_host_skills_from_legacy_roots() -> TestResult } #[tokio::test] -async fn installed_extension_injects_available_catalog_and_selected_entrypoint() -> TestResult { - let host_read_requests = Arc::new(Mutex::new(Vec::new())); - let remote_read_requests = Arc::new(Mutex::new(Vec::new())); - let host_provider = Arc::new(StaticSkillProvider { +async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_input() -> TestResult +{ + let read_requests = Arc::new(Mutex::new(Vec::new())); + let executor_provider = Arc::new(StaticSkillProvider { catalog: SkillCatalog { entries: vec![test_entry( - SkillSourceKind::Host, - "host", - "host/lint-fix", + SkillSourceKind::Executor, + "env-1", + "executor/lint-fix", "lint-fix/SKILL.md", )], warnings: Vec::new(), }, - read_requests: Arc::clone(&host_read_requests), + read_requests: Arc::clone(&read_requests), }); - let remote_provider = Arc::new(StaticSkillProvider { - catalog: SkillCatalog { - entries: vec![test_entry( - SkillSourceKind::Remote, - "remote", - "remote/lint-fix", - "lint-fix/SKILL.md", - )], - warnings: Vec::new(), - }, - read_requests: Arc::clone(&remote_read_requests), - }); - let providers = SkillProviders::new() - .with_host_provider(host_provider) - .with_remote_provider(remote_provider); + let providers = SkillProviders::new().with_executor_provider(executor_provider); let mut builder = ExtensionRegistryBuilder::new(); install_with_providers(&mut builder, providers); let registry = builder.build(); let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("thread"); + thread_store.insert(vec![SelectedCapabilityRoot { + id: "lint-fix".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "env-1".to_string(), + path: "/skills/lint-fix".to_string(), + }, + }]); let session_source = SessionSource::Cli; let config = default_config().await?; registry.thread_lifecycle_contributors()[0] @@ -176,6 +171,17 @@ async fn installed_extension_injects_available_catalog_and_selected_entrypoint() }) .await; + let prompt_fragments = registry.context_contributors()[0] + .contribute(&session_store, &thread_store) + .await; + assert_eq!(1, prompt_fragments.len()); + assert!( + prompt_fragments[0] + .text() + .starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) + ); + assert!(prompt_fragments[0].text().contains("lint-fix")); + let turn_store = ExtensionData::new("turn-1"); let fragments = registry.turn_input_contributors()[0] .contribute( @@ -185,11 +191,7 @@ async fn installed_extension_injects_available_catalog_and_selected_entrypoint() text: "$lint-fix please".to_string(), text_elements: Vec::new(), }], - environments: vec![TurnInputEnvironment { - environment_id: "env-1".to_string(), - cwd: std::env::temp_dir(), - is_primary: true, - }], + environments: Vec::new(), }, &session_store, &thread_store, @@ -197,31 +199,23 @@ async fn installed_extension_injects_available_catalog_and_selected_entrypoint() ) .await; - assert_eq!(2, fragments.len()); - assert_eq!("developer", fragments[0].role()); - assert!( - fragments[0] - .render() - .starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) - ); - assert!(fragments[0].render().contains("lint-fix")); - assert_eq!("user", fragments[1].role()); - assert!(fragments[1].render().contains("lint-fix")); - assert!(fragments[1].render().contains("# Lint Fix")); + assert_eq!(1, fragments.len()); + assert_eq!("user", fragments[0].role()); + assert!(fragments[0].render().contains("lint-fix")); + assert!(fragments[0].render().contains("# Lint Fix")); assert_eq!( vec![( - SkillAuthority::new(SkillSourceKind::Host, "host"), - SkillPackageId("host/lint-fix".to_string()), - SkillResourceId("lint-fix/SKILL.md".to_string()), + SkillAuthority::new(SkillSourceKind::Executor, "env-1"), + SkillPackageId("executor/lint-fix".to_string()), + SkillResourceId::new("lint-fix/SKILL.md"), )], - read_request_keys(&host_read_requests) - ); - assert!( - remote_read_requests - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .is_empty() + read_request_keys(&read_requests) ); + let rebuilt_prompt_fragments = registry.context_contributors()[0] + .contribute(&session_store, &thread_store) + .await; + assert_eq!(1, rebuilt_prompt_fragments.len()); + assert!(rebuilt_prompt_fragments[0].text().contains("lint-fix")); let next_turn_store = ExtensionData::new("turn-2"); let next_fragments = registry.turn_input_contributors()[0] @@ -240,9 +234,91 @@ async fn installed_extension_injects_available_catalog_and_selected_entrypoint() ) .await; - assert_eq!(1, next_fragments.len()); - assert_eq!("developer", next_fragments[0].role()); - assert!(next_fragments[0].render().contains("lint-fix")); + assert!(next_fragments.is_empty()); + + Ok(()) +} + +#[tokio::test] +async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> TestResult { + let read_requests = Arc::new(Mutex::new(Vec::new())); + let root_a_locator = "skill://root-a/shared/lint-fix/SKILL.md"; + let root_b_locator = "skill://root-b/shared/lint-fix/SKILL.md"; + let executor_provider = Arc::new(StaticSkillProvider { + catalog: SkillCatalog { + entries: [("root-a", root_a_locator), ("root-b", root_b_locator)] + .into_iter() + .map(|(root_id, locator)| { + SkillCatalogEntry::new( + SkillPackageId(locator.to_string()), + SkillAuthority::new(SkillSourceKind::Executor, root_id), + "lint-fix", + "Fix lint errors.", + SkillResourceId::new(locator), + ) + .with_display_path(locator) + }) + .collect(), + warnings: Vec::new(), + }, + read_requests: Arc::clone(&read_requests), + }); + let providers = SkillProviders::new().with_executor_provider(executor_provider); + let mut builder = ExtensionRegistryBuilder::new(); + install_with_providers(&mut builder, providers); + let registry = builder.build(); + let session_store = ExtensionData::new("session"); + let thread_store = ExtensionData::new("thread"); + thread_store.insert( + [("root-a", "/skills/root-a"), ("root-b", "/skills/root-b")] + .into_iter() + .map(|(id, path)| SelectedCapabilityRoot { + id: id.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "env-1".to_string(), + path: path.to_string(), + }, + }) + .collect::>(), + ); + let session_source = SessionSource::Cli; + let config = default_config().await?; + 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 fragments = registry.turn_input_contributors()[0] + .contribute( + TurnInputContext { + turn_id: "turn-1".to_string(), + user_input: vec![UserInput::Mention { + name: "lint-fix".to_string(), + path: root_b_locator.to_string(), + }], + environments: Vec::new(), + }, + &session_store, + &thread_store, + &ExtensionData::new("turn-1"), + ) + .await; + + assert_eq!(1, fragments.len()); + assert!(fragments[0].render().contains(root_b_locator)); + assert_eq!( + vec![( + SkillAuthority::new(SkillSourceKind::Executor, "root-b"), + SkillPackageId(root_b_locator.to_string()), + SkillResourceId::new(root_b_locator), + )], + read_request_keys(&read_requests) + ); Ok(()) } @@ -306,15 +382,14 @@ async fn prompt_hidden_skill_can_still_be_invoked() -> TestResult { .await; assert_eq!(2, fragments.len()); - let catalog_fragment = fragments[0].render(); - assert!(catalog_fragment.contains("visible-skill")); - assert!(!catalog_fragment.contains("hidden-skill")); + assert!(fragments[0].render().contains("visible-skill")); + assert!(!fragments[0].render().contains("hidden-skill")); assert!(fragments[1].render().contains("hidden-skill")); assert_eq!( vec![( SkillAuthority::new(SkillSourceKind::Host, "host"), SkillPackageId("host/hidden-skill".to_string()), - SkillResourceId("hidden-skill/SKILL.md".to_string()), + SkillResourceId::new("hidden-skill/SKILL.md"), )], read_request_keys(&read_requests) ); @@ -329,13 +404,9 @@ struct StaticSkillProvider { } impl SkillProvider for StaticSkillProvider { - fn list(&self, query: SkillListQuery) -> SkillProviderFuture<'_, SkillCatalog> { + fn list(&self, _query: SkillListQuery) -> SkillProviderFuture<'_, SkillCatalog> { let catalog = self.catalog.clone(); - Box::pin(async move { - assert!(query.include_host_skills); - assert!(query.include_bundled_skills); - Ok(catalog) - }) + Box::pin(async move { Ok(catalog) }) } fn read(&self, request: SkillReadRequest) -> SkillProviderFuture<'_, SkillReadResult> { @@ -369,7 +440,7 @@ fn test_entry( SkillAuthority::new(kind, authority_id), name, "Fix lint errors.", - SkillResourceId(main_prompt.to_string()), + SkillResourceId::new(main_prompt), ) .with_display_path(format!("skill://{package_id}/SKILL.md")) } diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 9c15ff04e2..d524a2e224 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -251,6 +251,7 @@ impl MemoryStartupContext { metrics_service_name: None, parent_trace: None, environments, + thread_extension_init: Default::default(), }) .await?; diff --git a/codex-rs/protocol/src/capabilities.rs b/codex-rs/protocol/src/capabilities.rs new file mode 100644 index 0000000000..cfdead7b28 --- /dev/null +++ b/codex-rs/protocol/src/capabilities.rs @@ -0,0 +1,30 @@ +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; +use ts_rs::TS; + +/// A user-selected root that can expose one or more runtime capabilities. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct SelectedCapabilityRoot { + /// Stable identifier supplied by the capability selection platform. + pub id: String, + /// Where the selected root can be resolved. + pub location: CapabilityRootLocation, +} + +/// Location used to resolve a selected capability root. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] +#[serde(tag = "type", rename_all = "camelCase")] +#[ts(tag = "type")] +#[ts(export_to = "v2/")] +pub enum CapabilityRootLocation { + /// A path owned by an execution environment. + Environment { + #[serde(rename = "environmentId")] + #[ts(rename = "environmentId")] + environment_id: String, + path: String, + }, +} diff --git a/codex-rs/protocol/src/lib.rs b/codex-rs/protocol/src/lib.rs index 63053159c6..3be16d08b7 100644 --- a/codex-rs/protocol/src/lib.rs +++ b/codex-rs/protocol/src/lib.rs @@ -9,6 +9,7 @@ pub use session_id::SessionId; pub use thread_id::ThreadId; pub use tool_name::ToolName; pub mod approvals; +pub mod capabilities; pub mod config_types; pub mod dynamic_tools; pub mod error; From 99da697e4c5c1fc908732a58b6548bf9cc227f83 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Tue, 9 Jun 2026 11:30:24 -0700 Subject: [PATCH 07/17] [codex] Test extension API contracts (#26835) ## Why `codex-extension-api` defines contracts shared by extension crates and their hosts, but it had no direct test suite. Host and feature tests cover downstream behavior, while regressions in the API crate's own typed state, registry ordering, and capability adapters could go unnoticed. ## What - Add public-surface integration tests for `ExtensionData`, including concurrent initialization and poison recovery. - Cover contributor registration order, approval short-circuiting, event sink retention, no-op response injection, and closure-based agent spawning. - Add the test-only dependencies used by the suite. ## Validation - `just test -p codex-extension-api` - `just argument-comment-lint -p codex-extension-api` - `just bazel-lock-check` --- codex-rs/Cargo.lock | 2 + codex-rs/ext/extension-api/Cargo.toml | 4 + .../ext/extension-api/tests/capabilities.rs | 56 +++ codex-rs/ext/extension-api/tests/registry.rs | 370 ++++++++++++++++++ codex-rs/ext/extension-api/tests/state.rs | 109 ++++++ 5 files changed, 541 insertions(+) create mode 100644 codex-rs/ext/extension-api/tests/capabilities.rs create mode 100644 codex-rs/ext/extension-api/tests/registry.rs create mode 100644 codex-rs/ext/extension-api/tests/state.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 064d9ad944..8783f0e16c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2923,6 +2923,8 @@ dependencies = [ "codex-context-fragments", "codex-protocol", "codex-tools", + "pretty_assertions", + "tokio", ] [[package]] diff --git a/codex-rs/ext/extension-api/Cargo.toml b/codex-rs/ext/extension-api/Cargo.toml index 85c7d8f98e..5497dc5439 100644 --- a/codex-rs/ext/extension-api/Cargo.toml +++ b/codex-rs/ext/extension-api/Cargo.toml @@ -18,3 +18,7 @@ async-trait = { workspace = true } codex-context-fragments = { workspace = true } codex-protocol = { workspace = true } codex-tools = { workspace = true } + +[dev-dependencies] +pretty_assertions = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/extension-api/tests/capabilities.rs b/codex-rs/ext/extension-api/tests/capabilities.rs new file mode 100644 index 0000000000..c01b8f5723 --- /dev/null +++ b/codex-rs/ext/extension-api/tests/capabilities.rs @@ -0,0 +1,56 @@ +use std::sync::Arc; +use std::sync::Mutex; + +use codex_extension_api::AgentSpawnFuture; +use codex_extension_api::AgentSpawner; +use codex_extension_api::NoopResponseItemInjector; +use codex_extension_api::ResponseItemInjector; +use codex_protocol::ThreadId; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseInputItem; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn noop_response_item_injector_returns_original_items() { + let items = vec![ResponseInputItem::Message { + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "keep this input".to_string(), + }], + phase: None, + }]; + + let returned_items = NoopResponseItemInjector + .inject_response_items(items.clone()) + .await + .expect_err("noop injector should reject same-turn injection"); + + assert_eq!(returned_items, items); +} + +#[tokio::test] +async fn closure_agent_spawner_forwards_arguments_and_result() { + let calls = Arc::new(Mutex::new(Vec::new())); + let recorded_calls = Arc::clone(&calls); + let spawner = move |thread_id: ThreadId, + request: String| + -> AgentSpawnFuture<'static, usize, &'static str> { + recorded_calls + .lock() + .expect("agent spawn calls lock") + .push((thread_id, request.clone())); + Box::pin(async move { Ok(request.len()) }) + }; + let thread_id = + ThreadId::from_string("11111111-1111-4111-8111-111111111111").expect("valid thread id"); + + let spawned = spawner + .spawn_subagent(thread_id, "delegate this".to_string()) + .await; + + assert_eq!(spawned, Ok(13)); + assert_eq!( + calls.lock().expect("agent spawn calls lock").as_slice(), + [(thread_id, "delegate this".to_string())] + ); +} diff --git a/codex-rs/ext/extension-api/tests/registry.rs b/codex-rs/ext/extension-api/tests/registry.rs new file mode 100644 index 0000000000..d26f1010cb --- /dev/null +++ b/codex-rs/ext/extension-api/tests/registry.rs @@ -0,0 +1,370 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::Mutex; + +use codex_extension_api::ApprovalReviewContributor; +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::ExtensionRegistryBuilder; +use codex_extension_api::PromptFragment; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::TokenUsageContributor; +use codex_extension_api::ToolCall; +use codex_extension_api::ToolContributor; +use codex_extension_api::ToolExecutor; +use codex_extension_api::ToolLifecycleContributor; +use codex_extension_api::TurnInputContext; +use codex_extension_api::TurnInputContributor; +use codex_extension_api::TurnItemContributor; +use codex_extension_api::TurnLifecycleContributor; +use codex_extension_api::empty_extension_registry; +use codex_protocol::items::HookPromptItem; +use codex_protocol::items::TurnItem; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ReviewDecision; +use codex_protocol::protocol::WarningEvent; +use pretty_assertions::assert_eq; + +struct AllContributors; + +impl ContextContributor for AllContributors { + fn contribute<'a>( + &'a self, + _session_store: &'a ExtensionData, + _thread_store: &'a ExtensionData, + ) -> Pin> + Send + 'a>> { + 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() + } +} + +impl ToolContributor for AllContributors { + fn tools( + &self, + _session_store: &ExtensionData, + _thread_store: &ExtensionData, + ) -> Vec>> { + Vec::new() + } +} + +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(()) + } +} + +#[async_trait::async_trait] +impl ApprovalReviewContributor for AllContributors { + async fn contribute( + &self, + _session_store: &ExtensionData, + _thread_store: &ExtensionData, + _prompt: &str, + ) -> Option { + Some(ReviewDecision::ApprovedForSession) + } +} + +#[tokio::test] +async fn build_round_trips_every_contributor_category() { + let contributor = Arc::new(AllContributors); + let mut builder = ExtensionRegistryBuilder::<()>::new(); + builder.thread_lifecycle_contributor(contributor.clone()); + builder.turn_lifecycle_contributor(contributor.clone()); + builder.config_contributor(contributor.clone()); + builder.token_usage_contributor(contributor.clone()); + builder.prompt_contributor(contributor.clone()); + builder.turn_input_contributor(contributor.clone()); + builder.tool_contributor(contributor.clone()); + builder.tool_lifecycle_contributor(contributor.clone()); + builder.turn_item_contributor(contributor.clone()); + builder.approval_review_contributor(contributor); + let registry = builder.build(); + + assert_eq!(registry.thread_lifecycle_contributors().len(), 1); + assert_eq!(registry.turn_lifecycle_contributors().len(), 1); + assert_eq!(registry.config_contributors().len(), 1); + assert_eq!(registry.token_usage_contributors().len(), 1); + assert_eq!(registry.context_contributors().len(), 1); + assert_eq!(registry.turn_input_contributors().len(), 1); + assert_eq!(registry.tool_contributors().len(), 1); + assert_eq!(registry.tool_lifecycle_contributors().len(), 1); + assert_eq!(registry.turn_item_contributors().len(), 1); + assert_eq!( + registry + .approval_review( + &ExtensionData::new("session"), + &ExtensionData::new("thread"), + "review this", + ) + .await, + Some(ReviewDecision::ApprovedForSession) + ); +} + +struct NamedContextContributor(&'static str); + +impl ContextContributor for NamedContextContributor { + fn contribute<'a>( + &'a self, + _session_store: &'a ExtensionData, + _thread_store: &'a ExtensionData, + ) -> Pin> + Send + 'a>> { + Box::pin(std::future::ready(vec![PromptFragment::developer_policy( + self.0, + )])) + } +} + +struct RecordingTurnItemContributor { + name: &'static str, + 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(()) + } +} + +#[tokio::test] +async fn contributors_preserve_registration_order() { + let turn_item_calls = Arc::new(Mutex::new(Vec::new())); + let mut builder = ExtensionRegistryBuilder::<()>::new(); + builder.prompt_contributor(Arc::new(NamedContextContributor("first"))); + builder.prompt_contributor(Arc::new(NamedContextContributor("second"))); + for name in ["first", "second"] { + builder.turn_item_contributor(Arc::new(RecordingTurnItemContributor { + name, + calls: Arc::clone(&turn_item_calls), + })); + } + let registry = builder.build(); + let session_store = ExtensionData::new("session"); + let thread_store = ExtensionData::new("thread"); + let turn_store = ExtensionData::new("turn"); + + let mut fragments = Vec::new(); + for contributor in registry.context_contributors() { + fragments.extend(contributor.contribute(&session_store, &thread_store).await); + } + let mut item = TurnItem::HookPrompt(HookPromptItem { + id: "item".to_string(), + fragments: Vec::new(), + }); + for contributor in registry.turn_item_contributors() { + contributor + .contribute(&thread_store, &turn_store, &mut item) + .await + .expect("turn item contribution should succeed"); + } + + assert_eq!( + fragments, + vec![ + PromptFragment::developer_policy("first"), + PromptFragment::developer_policy("second"), + ] + ); + assert_eq!( + turn_item_calls + .lock() + .expect("turn item calls lock") + .as_slice(), + ["first", "second"] + ); +} + +#[derive(Debug, PartialEq, Eq)] +struct ApprovalCall { + contributor: &'static str, + session_id: String, + thread_id: String, + prompt: String, +} + +struct RecordingApprovalContributor { + name: &'static str, + decision: Option, + 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() + } +} + +#[tokio::test] +async fn approval_review_returns_first_claim_and_short_circuits() { + let calls = Arc::new(Mutex::new(Vec::new())); + let mut builder = ExtensionRegistryBuilder::<()>::new(); + for (name, decision) in [ + ("first", None), + ("second", Some(ReviewDecision::Approved)), + ("third", Some(ReviewDecision::Denied)), + ] { + builder.approval_review_contributor(Arc::new(RecordingApprovalContributor { + name, + decision, + calls: Arc::clone(&calls), + })); + } + let registry = builder.build(); + + let decision = registry + .approval_review( + &ExtensionData::new("session-1"), + &ExtensionData::new("thread-1"), + "allow command?", + ) + .await; + + assert_eq!(decision, Some(ReviewDecision::Approved)); + assert_eq!( + calls.lock().expect("approval calls lock").as_slice(), + [ + ApprovalCall { + contributor: "first", + session_id: "session-1".to_string(), + thread_id: "thread-1".to_string(), + prompt: "allow command?".to_string(), + }, + ApprovalCall { + contributor: "second", + session_id: "session-1".to_string(), + thread_id: "thread-1".to_string(), + prompt: "allow command?".to_string(), + }, + ] + ); +} + +#[derive(Default)] +struct RecordingEventSink { + events: Mutex>, +} + +impl ExtensionEventSink for RecordingEventSink { + fn emit(&self, event: Event) { + let EventMsg::Warning(warning) = event.msg else { + panic!("test sink only accepts warning events"); + }; + self.events + .lock() + .unwrap_or_else(|error| panic!("recording event sink lock poisoned: {error}")) + .push((event.id, warning.message)); + } +} + +#[test] +fn custom_event_sink_survives_registry_build() { + let sink = Arc::new(RecordingEventSink::default()); + let builder = ExtensionRegistryBuilder::<()>::with_event_sink(sink.clone()); + + builder + .event_sink() + .emit(warning_event("builder", "before")); + let registry = builder.build(); + registry + .event_sink() + .emit(warning_event("registry", "after")); + + assert_eq!( + sink.events + .lock() + .expect("recording event sink lock") + .as_slice(), + [ + ("builder".to_string(), "before".to_string()), + ("registry".to_string(), "after".to_string()), + ] + ); +} + +#[tokio::test] +async fn empty_registry_does_not_claim_approval_review() { + let registry = empty_extension_registry::<()>(); + + assert_eq!( + registry + .approval_review( + &ExtensionData::new("session"), + &ExtensionData::new("thread"), + "unclaimed", + ) + .await, + None + ); +} + +fn warning_event(id: &str, message: &str) -> Event { + Event { + id: id.to_string(), + msg: EventMsg::Warning(WarningEvent { + message: message.to_string(), + }), + } +} diff --git a/codex-rs/ext/extension-api/tests/state.rs b/codex-rs/ext/extension-api/tests/state.rs new file mode 100644 index 0000000000..31b643c29b --- /dev/null +++ b/codex-rs/ext/extension-api/tests/state.rs @@ -0,0 +1,109 @@ +use std::panic::AssertUnwindSafe; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use codex_extension_api::ExtensionData; +use pretty_assertions::assert_eq; + +#[test] +fn typed_values_can_be_inserted_replaced_and_removed() { + let data = ExtensionData::new("thread-1"); + + assert_eq!(data.insert(/*value*/ 41_u64), None); + assert_eq!(data.insert("alpha".to_string()), None); + assert_eq!(data.get::().as_deref(), Some(&41)); + assert_eq!( + data.get::().map(|value| value.as_str().to_string()), + Some("alpha".to_string()) + ); + + assert_eq!(data.insert(/*value*/ 42_u64).as_deref(), Some(&41)); + assert_eq!(data.get::().as_deref(), Some(&42)); + assert_eq!( + data.remove::() + .map(|value| value.as_str().to_string()), + Some("alpha".to_string()) + ); + assert_eq!(data.get::(), None); + assert_eq!(data.get::().as_deref(), Some(&42)); +} + +#[test] +fn get_or_init_initializes_once_and_returns_shared_value() { + const CALLER_COUNT: usize = 8; + + #[derive(Debug, PartialEq, Eq)] + struct SharedValue(usize); + + let data = Arc::new(ExtensionData::new("session")); + let callers_started = Arc::new(AtomicUsize::new(0)); + let initialization_count = Arc::new(AtomicUsize::new(0)); + + let handles: [_; CALLER_COUNT] = std::array::from_fn(|_| { + let data = Arc::clone(&data); + let callers_started = Arc::clone(&callers_started); + let initialization_count = Arc::clone(&initialization_count); + std::thread::spawn(move || { + callers_started.fetch_add(1, Ordering::SeqCst); + data.get_or_init(|| { + initialization_count.fetch_add(1, Ordering::SeqCst); + // Keep the first initializer active until every worker has attempted + // get_or_init, forcing callers to overlap on the same missing entry. + while callers_started.load(Ordering::SeqCst) < CALLER_COUNT { + std::thread::yield_now(); + } + SharedValue(7) + }) + }) + }); + let values = handles + .into_iter() + .map(|handle| handle.join().expect("initializer thread should succeed")) + .collect::>(); + + assert_eq!(initialization_count.load(Ordering::SeqCst), 1); + assert_eq!( + values.iter().map(Arc::as_ref).collect::>(), + vec![&SharedValue(7); CALLER_COUNT] + ); + assert!( + values + .iter() + .skip(1) + .all(|value| Arc::ptr_eq(&values[0], value)) + ); +} + +#[test] +fn stores_are_isolated_and_preserve_level_id() { + let session_data = ExtensionData::new("root-1"); + let thread_data = ExtensionData::new("root-1"); + + session_data.insert(/*value*/ 17_u32); + thread_data.insert("thread value".to_string()); + + assert_eq!(session_data.level_id(), "root-1"); + assert_eq!(thread_data.level_id(), "root-1"); + assert_eq!(session_data.get::().as_deref(), Some(&17)); + assert_eq!(session_data.get::(), None); + assert_eq!(thread_data.get::(), None); + assert_eq!( + thread_data + .get::() + .map(|value| value.as_str().to_string()), + Some("thread value".to_string()) + ); +} + +#[test] +fn store_remains_usable_after_panicking_initializer() { + let data = ExtensionData::new("turn-1"); + + let result = std::panic::catch_unwind(AssertUnwindSafe(|| { + data.get_or_init::(|| panic!("initializer failed")); + })); + + assert!(result.is_err()); + assert_eq!(*data.get_or_init(|| 99_u64), 99); +} From a71e040df51ec0e1f5523bf0ff5f1bef39858128 Mon Sep 17 00:00:00 2001 From: marksteinbrick-oai Date: Tue, 9 Jun 2026 12:27:10 -0700 Subject: [PATCH 08/17] [codex-analytics] add extensible feature thread sources (#27063) ## Why - `ThreadSource` currently defines a closed set of core-owned values - Product features also create threads for background or scheduled work - Adding every product-specific value to the core enum would require repeated `codex-rs` protocol changes - Feature-backed values let product callers provide precise attribution while preserving the existing core classifications ## What Changed - Adds `ThreadSource::Feature(String)` for app-owned thread source values - Represents all app-server v2 thread sources as scalar strings, so a feature source is supplied as `"automation"` - Persists and emits the feature's plain string label, so `"automation"` produces `thread_source="automation"` in analytics - Keeps `user`, `subagent`, and `memory_consolidation` as explicit core-owned values and regenerates the app-server schemas and TypeScript bindings ## Verification - `just write-app-server-schema` - `cargo check --workspace` - `just test -p codex-protocol feature_thread_source_serializes_as_its_app_owned_label` - `just test -p codex-app-server-protocol thread_sources_round_trip_as_scalar_labels` - `cargo test -p codex-analytics thread_initialized_event_serializes_expected_shape` - `just fmt` --- .../analytics/src/analytics_client_tests.rs | 4 +- codex-rs/analytics/src/reducer.rs | 10 ++--- .../schema/json/ClientRequest.json | 5 --- .../schema/json/ServerNotification.json | 5 --- .../codex_app_server_protocol.schemas.json | 5 --- .../codex_app_server_protocol.v2.schemas.json | 5 --- .../schema/json/v2/ThreadForkParams.json | 5 --- .../schema/json/v2/ThreadForkResponse.json | 5 --- .../schema/json/v2/ThreadListResponse.json | 5 --- .../json/v2/ThreadMetadataUpdateResponse.json | 5 --- .../schema/json/v2/ThreadReadResponse.json | 5 --- .../schema/json/v2/ThreadResumeResponse.json | 5 --- .../json/v2/ThreadRollbackResponse.json | 5 --- .../schema/json/v2/ThreadStartParams.json | 5 --- .../schema/json/v2/ThreadStartResponse.json | 5 --- .../json/v2/ThreadStartedNotification.json | 5 --- .../json/v2/ThreadUnarchiveResponse.json | 5 --- .../schema/typescript/v2/ThreadSource.ts | 2 +- .../src/protocol/v2/tests.rs | 24 +++++++++++ .../src/protocol/v2/thread_data.rs | 36 ++++++++++++++-- .../request_processors/thread_processor.rs | 4 +- codex-rs/core/src/session/review.rs | 4 +- codex-rs/core/src/session/session.rs | 6 +-- codex-rs/core/src/session/turn_context.rs | 6 +-- codex-rs/exec/src/lib.rs | 4 +- codex-rs/protocol/src/protocol.rs | 43 ++++++++++++++++--- codex-rs/state/src/extract.rs | 2 +- codex-rs/state/src/model/thread_metadata.rs | 2 +- codex-rs/state/src/runtime/threads.rs | 2 + codex-rs/thread-store/src/in_memory.rs | 6 +-- .../src/local/update_thread_metadata.rs | 2 +- .../thread-store/src/thread_metadata_sync.rs | 4 +- 32 files changed, 123 insertions(+), 113 deletions(-) diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index 7172de6898..02779c6609 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -1374,7 +1374,7 @@ fn thread_initialized_event_serializes_expected_shape() { }, model: "gpt-5".to_string(), ephemeral: true, - thread_source: Some(ThreadSource::User), + thread_source: Some(ThreadSource::Feature("automation".to_string())), initialization_mode: ThreadInitializationMode::New, subagent_source: None, parent_thread_id: None, @@ -1407,7 +1407,7 @@ fn thread_initialized_event_serializes_expected_shape() { }, "model": "gpt-5", "ephemeral": true, - "thread_source": "user", + "thread_source": "automation", "initialization_mode": "new", "subagent_source": null, "parent_thread_id": null, diff --git a/codex-rs/analytics/src/reducer.rs b/codex-rs/analytics/src/reducer.rs index 0cceb2b46d..f8abe36e4e 100644 --- a/codex-rs/analytics/src/reducer.rs +++ b/codex-rs/analytics/src/reducer.rs @@ -1263,7 +1263,7 @@ impl AnalyticsReducer { thread_metadata.session_id.clone(), connection_state.app_server_client.clone(), connection_state.runtime.clone(), - thread_metadata.thread_source, + thread_metadata.thread_source.clone(), thread_metadata.subagent_source.clone(), thread_metadata.parent_thread_id.clone(), ), @@ -1368,7 +1368,7 @@ impl AnalyticsReducer { accepted_turn_id, app_server_client: connection_state.app_server_client.clone(), runtime: connection_state.runtime.clone(), - thread_source: thread_metadata.thread_source, + thread_source: thread_metadata.thread_source.clone(), subagent_source: thread_metadata.subagent_source.clone(), parent_thread_id: thread_metadata.parent_thread_id.clone(), num_input_images: pending_request.num_input_images, @@ -1411,7 +1411,7 @@ impl AnalyticsReducer { review_id: pending_review.review_id, app_server_client: connection_state.app_server_client.clone(), runtime: connection_state.runtime.clone(), - thread_source: thread_metadata.thread_source, + thread_source: thread_metadata.thread_source.clone(), subagent_source: thread_metadata.subagent_source.clone(), parent_thread_id: thread_metadata.parent_thread_id.clone(), subject_kind: pending_review.subject_kind, @@ -1979,7 +1979,7 @@ fn tool_item_base( item_id, app_server_client: context.connection_state.app_server_client.clone(), runtime: context.connection_state.runtime.clone(), - thread_source: thread_metadata.thread_source, + thread_source: thread_metadata.thread_source.clone(), subagent_source: thread_metadata.subagent_source.clone(), parent_thread_id: thread_metadata.parent_thread_id.clone(), tool_name, @@ -2459,7 +2459,7 @@ fn codex_turn_event_params( runtime, submission_type, ephemeral, - thread_source: thread_metadata.thread_source, + thread_source: thread_metadata.thread_source.clone(), initialization_mode: thread_metadata.initialization_mode, subagent_source: thread_metadata.subagent_source.clone(), parent_thread_id: thread_metadata.parent_thread_id.clone(), diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index c499376547..b360d4bddd 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -3847,11 +3847,6 @@ "type": "string" }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadSourceKind": { diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index 01f6f58b09..e0651d1727 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -4615,11 +4615,6 @@ "type": "object" }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStartedNotification": { 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 475154b38a..d65c043a7d 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 @@ -18109,11 +18109,6 @@ "type": "string" }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadSourceKind": { 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 8ed40282a9..8271dc485d 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 @@ -15926,11 +15926,6 @@ "type": "string" }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadSourceKind": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkParams.json index d9a543e939..e72522a26a 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkParams.json @@ -73,11 +73,6 @@ "type": "string" }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" } }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json index 0cea90749b..fa1277b0ef 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json @@ -1830,11 +1830,6 @@ ] }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStatus": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json index d0b3351a4e..1b53d4ca2e 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json @@ -1645,11 +1645,6 @@ ] }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStatus": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json index e2850b3dd9..c98fa1f2a5 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json @@ -1645,11 +1645,6 @@ ] }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStatus": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json index 01da3165b1..345d84d794 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json @@ -1645,11 +1645,6 @@ ] }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStatus": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json index d766d7a024..98ef79517b 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json @@ -1830,11 +1830,6 @@ ] }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStatus": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json index 6448ca6517..f63dfe6ffb 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json @@ -1645,11 +1645,6 @@ ] }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStatus": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json index 7f0d4ad9a4..57e6e0ab1b 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json @@ -159,11 +159,6 @@ "type": "object" }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStartSource": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json index a19a40f522..052023361a 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json @@ -1830,11 +1830,6 @@ ] }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStatus": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json index 0eb79075f6..cdeda64b7c 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json @@ -1645,11 +1645,6 @@ ] }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStatus": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json index 71064d5ea2..2a7281fcbb 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json @@ -1645,11 +1645,6 @@ ] }, "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" - ], "type": "string" }, "ThreadStatus": { diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSource.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSource.ts index 8f55524801..f27154ab6e 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSource.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSource.ts @@ -2,4 +2,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -export type ThreadSource = "user" | "subagent" | "memory_consolidation"; +export type ThreadSource = string; diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index 91bfbcc091..c570b87ce4 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -57,6 +57,30 @@ fn test_absolute_path() -> AbsolutePathBuf { absolute_path("readable") } +#[test] +fn thread_sources_round_trip_as_scalar_labels() { + for (source, label) in [ + (ThreadSource::User, "user"), + (ThreadSource::Subagent, "subagent"), + ( + ThreadSource::Feature("automation".to_string()), + "automation", + ), + (ThreadSource::MemoryConsolidation, "memory_consolidation"), + ] { + let value = serde_json::to_value(&source).expect("serialize thread source"); + + assert_eq!(value, json!(label)); + assert_eq!( + serde_json::from_value::(value).expect("deserialize thread source"), + source + ); + + let core_source: codex_protocol::protocol::ThreadSource = source.clone().into(); + assert_eq!(ThreadSource::from(core_source), source); + } +} + #[test] fn approvals_reviewer_serializes_auto_review_and_accepts_legacy_guardian_subagent() { assert_eq!( diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread_data.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread_data.rs index 35b6183871..fc0762361f 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread_data.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread_data.rs @@ -7,6 +7,8 @@ use codex_protocol::protocol::SubAgentSource as CoreSubAgentSource; use codex_protocol::protocol::ThreadSource as CoreThreadSource; use codex_utils_absolute_path::AbsolutePathBuf; use schemars::JsonSchema; +use schemars::r#gen::SchemaGenerator; +use schemars::schema::Schema; use serde::Deserialize; use serde::Serialize; use std::path::PathBuf; @@ -61,20 +63,47 @@ impl From for CoreSessionSource { } } -#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, JsonSchema, TS)] -#[serde(rename_all = "snake_case")] -#[ts(rename_all = "snake_case", export_to = "v2/")] +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, TS)] +#[serde(try_from = "String", into = "String")] +#[ts(type = "string")] +#[ts(export_to = "v2/")] pub enum ThreadSource { User, Subagent, + Feature(String), MemoryConsolidation, } +impl JsonSchema for ThreadSource { + fn schema_name() -> String { + "ThreadSource".to_string() + } + + fn json_schema(generator: &mut SchemaGenerator) -> Schema { + String::json_schema(generator) + } +} + +impl TryFrom for ThreadSource { + type Error = String; + + fn try_from(value: String) -> Result { + value.parse::().map(Into::into) + } +} + +impl From for String { + fn from(value: ThreadSource) -> Self { + CoreThreadSource::from(value).into() + } +} + impl From for ThreadSource { fn from(value: CoreThreadSource) -> Self { match value { CoreThreadSource::User => ThreadSource::User, CoreThreadSource::Subagent => ThreadSource::Subagent, + CoreThreadSource::Feature(feature) => ThreadSource::Feature(feature), CoreThreadSource::MemoryConsolidation => ThreadSource::MemoryConsolidation, } } @@ -85,6 +114,7 @@ impl From for CoreThreadSource { match value { ThreadSource::User => CoreThreadSource::User, ThreadSource::Subagent => CoreThreadSource::Subagent, + ThreadSource::Feature(feature) => CoreThreadSource::Feature(feature), ThreadSource::MemoryConsolidation => CoreThreadSource::MemoryConsolidation, } } 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 be5cffb49a..fd6432531d 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -4154,7 +4154,7 @@ fn summary_from_thread_metadata(metadata: &ThreadMetadata) -> ConversationSummar metadata.cwd.clone(), metadata.cli_version.clone(), metadata.source.clone(), - metadata.thread_source, + metadata.thread_source.clone(), metadata.agent_nickname.clone(), metadata.agent_role.clone(), metadata.git_sha.clone(), @@ -4243,7 +4243,7 @@ fn build_thread_from_snapshot( agent_nickname: config_snapshot.session_source.get_nickname(), agent_role: config_snapshot.session_source.get_agent_role(), source: config_snapshot.session_source.clone().into(), - thread_source: config_snapshot.thread_source.map(Into::into), + thread_source: config_snapshot.thread_source.clone().map(Into::into), git_info: None, name: None, turns: Vec::new(), diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index a3b23c6a78..00a89b2b5c 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -91,7 +91,7 @@ pub(super) async fn spawn_review_thread( forked_from_thread_id, parent_turn_context.parent_thread_id, &session_source, - parent_turn_context.thread_source, + parent_turn_context.thread_source.clone(), review_turn_id.clone(), #[allow(deprecated)] parent_turn_context.cwd.clone(), @@ -121,7 +121,7 @@ pub(super) async fn spawn_review_thread( reasoning_summary, session_source, parent_thread_id: parent_turn_context.parent_thread_id, - thread_source: parent_turn_context.thread_source, + thread_source: parent_turn_context.thread_source.clone(), environments: parent_turn_context.environments.clone(), available_models, unified_exec_shell_mode, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index e51a83113b..50ff984231 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -193,7 +193,7 @@ impl SessionConfiguration { session_source: self.session_source.clone(), forked_from_thread_id: self.forked_from_thread_id, parent_thread_id: self.parent_thread_id, - thread_source: self.thread_source, + thread_source: self.thread_source.clone(), } } @@ -549,7 +549,7 @@ impl Session { forked_from_id, parent_thread_id, source: session_source, - thread_source: session_configuration.thread_source, + thread_source: session_configuration.thread_source.clone(), base_instructions: BaseInstructions { text: session_configuration.base_instructions.clone(), }, @@ -1089,7 +1089,7 @@ impl Session { thread_id, forked_from_id, parent_thread_id, - thread_source: session_configuration.thread_source, + thread_source: session_configuration.thread_source.clone(), thread_name: session_configuration.thread_name.clone(), model: session_configuration.collaboration_mode.model().to_string(), model_provider_id: config.model_provider_id.clone(), diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index be6e316a00..e3bd01081a 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -239,7 +239,7 @@ impl TurnContext { reasoning_summary: self.reasoning_summary, session_source: self.session_source.clone(), parent_thread_id: self.parent_thread_id, - thread_source: self.thread_source, + thread_source: self.thread_source.clone(), environments: self.environments.clone(), #[allow(deprecated)] cwd: self.cwd.clone(), @@ -514,7 +514,7 @@ impl Session { session_configuration.forked_from_thread_id, session_configuration.parent_thread_id, &session_configuration.session_source, - session_configuration.thread_source, + session_configuration.thread_source.clone(), sub_id.clone(), cwd.clone(), &session_configuration.permission_profile(), @@ -538,7 +538,7 @@ impl Session { reasoning_summary, session_source, parent_thread_id: session_configuration.parent_thread_id, - thread_source: session_configuration.thread_source, + thread_source: session_configuration.thread_source.clone(), environments, #[allow(deprecated)] cwd, diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index b8b87a85ec..fbecfe5620 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -1146,7 +1146,7 @@ fn session_configured_from_thread_start_response( &response.thread.session_id, &response.thread.id, response.thread.parent_thread_id.as_deref(), - response.thread.thread_source.map(Into::into), + response.thread.thread_source.clone().map(Into::into), response.thread.name.clone(), response.thread.path.clone(), response.model.clone(), @@ -1169,7 +1169,7 @@ fn session_configured_from_thread_resume_response( &response.thread.session_id, &response.thread.id, response.thread.parent_thread_id.as_deref(), - response.thread.thread_source.map(Into::into), + response.thread.thread_source.clone().map(Into::into), response.thread.name.clone(), response.thread.path.clone(), response.model.clone(), diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index b141298288..5abc9c6399 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2480,12 +2480,12 @@ impl InitialHistory { pub fn get_resumed_session_sources(&self) -> Option<(SessionSource, Option)> { let meta = self.get_resumed_session_meta()?; - Some((meta.source.clone(), meta.thread_source)) + Some((meta.source.clone(), meta.thread_source.clone())) } pub fn get_resumed_thread_source(&self) -> Option { self.get_resumed_session_meta() - .and_then(|meta| meta.thread_source) + .and_then(|meta| meta.thread_source.clone()) } pub fn get_resumed_parent_thread_id(&self) -> Option { @@ -2529,20 +2529,23 @@ pub enum SessionSource { Unknown, } -#[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq, Eq, JsonSchema, TS)] -#[serde(rename_all = "snake_case")] -#[ts(rename_all = "snake_case")] +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS)] +#[serde(try_from = "String", into = "String")] +#[schemars(with = "String")] +#[ts(type = "string")] pub enum ThreadSource { User, Subagent, + Feature(String), MemoryConsolidation, } impl ThreadSource { - pub fn as_str(self) -> &'static str { + pub fn as_str(&self) -> &str { match self { ThreadSource::User => "user", ThreadSource::Subagent => "subagent", + ThreadSource::Feature(feature) => feature, ThreadSource::MemoryConsolidation => "memory_consolidation", } } @@ -2554,6 +2557,20 @@ impl fmt::Display for ThreadSource { } } +impl TryFrom for ThreadSource { + type Error = String; + + fn try_from(value: String) -> Result { + value.parse() + } +} + +impl From for String { + fn from(value: ThreadSource) -> Self { + value.to_string() + } +} + impl FromStr for ThreadSource { type Err = String; @@ -2562,7 +2579,7 @@ impl FromStr for ThreadSource { "user" => Ok(ThreadSource::User), "subagent" => Ok(ThreadSource::Subagent), "memory_consolidation" => Ok(ThreadSource::MemoryConsolidation), - other => Err(format!("unknown thread source: {other}")), + other => Ok(ThreadSource::Feature(other.to_string())), } } } @@ -4048,6 +4065,18 @@ mod tests { use tempfile::NamedTempFile; use tempfile::TempDir; + #[test] + fn feature_thread_source_serializes_as_its_app_owned_label() -> Result<()> { + let source = ThreadSource::Feature("automation".to_string()); + + assert_eq!(serde_json::to_value(&source)?, json!("automation")); + assert_eq!( + serde_json::from_value::(json!("automation"))?, + source + ); + Ok(()) + } + fn sorted_writable_roots(roots: Vec) -> Vec<(PathBuf, Vec)> { let mut sorted_roots: Vec<(PathBuf, Vec)> = roots .into_iter() diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 63cb06fea1..47285c5cf5 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -50,7 +50,7 @@ fn apply_session_meta_from_item(metadata: &mut ThreadMetadata, meta_line: &Sessi } metadata.id = meta_line.meta.id; metadata.source = enum_to_string(&meta_line.meta.source); - metadata.thread_source = meta_line.meta.thread_source; + metadata.thread_source = meta_line.meta.thread_source.clone(); metadata.agent_nickname = meta_line.meta.agent_nickname.clone(); metadata.agent_role = meta_line.meta.agent_role.clone(); metadata.agent_path = meta_line.meta.agent_path.clone(); diff --git a/codex-rs/state/src/model/thread_metadata.rs b/codex-rs/state/src/model/thread_metadata.rs index cfcd49d5f9..a30f7e479b 100644 --- a/codex-rs/state/src/model/thread_metadata.rs +++ b/codex-rs/state/src/model/thread_metadata.rs @@ -196,7 +196,7 @@ impl ThreadMetadataBuilder { created_at, updated_at, source, - thread_source: self.thread_source, + thread_source: self.thread_source.clone(), agent_nickname: self.agent_nickname.clone(), agent_role: self.agent_role.clone(), agent_path: self diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 4cc7ee4dc1..6f3c45f0e2 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -521,6 +521,7 @@ ON CONFLICT(id) DO NOTHING .bind( metadata .thread_source + .as_ref() .map(codex_protocol::protocol::ThreadSource::as_str), ) .bind(metadata.agent_nickname.as_deref()) @@ -753,6 +754,7 @@ ON CONFLICT(id) DO UPDATE SET .bind( metadata .thread_source + .as_ref() .map(codex_protocol::protocol::ThreadSource::as_str), ) .bind(metadata.agent_nickname.as_deref()) diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index beed1d18fe..2ab5103663 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -177,7 +177,7 @@ impl ThreadStore for InMemoryThreadStore { agent_role: params.source.get_agent_role(), agent_path: params.source.get_agent_path().map(Into::into), source: params.source.clone(), - thread_source: params.thread_source, + thread_source: params.thread_source.clone(), model_provider: Some(params.metadata.model_provider.clone()), base_instructions: Some(params.base_instructions.clone()), dynamic_tools: (!params.dynamic_tools.is_empty()).then(|| params.dynamic_tools.clone()), @@ -392,8 +392,8 @@ fn stored_thread_from_state( .and_then(|metadata| metadata.source.clone()) .unwrap_or_else(|| created.source.clone()), thread_source: metadata - .and_then(|metadata| metadata.thread_source) - .unwrap_or(created.thread_source), + .and_then(|metadata| metadata.thread_source.clone()) + .unwrap_or_else(|| created.thread_source.clone()), agent_nickname: metadata.and_then(|metadata| metadata.agent_nickname.clone().flatten()), agent_role: metadata.and_then(|metadata| metadata.agent_role.clone().flatten()), agent_path: metadata.and_then(|metadata| metadata.agent_path.clone().flatten()), diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index dc616e4da2..7754b1a841 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -221,7 +221,7 @@ async fn apply_metadata_update( patch.source.clone().unwrap_or(SessionSource::Unknown), ); builder.model_provider = patch.model_provider.clone(); - builder.thread_source = patch.thread_source.flatten(); + builder.thread_source = patch.thread_source.clone().flatten(); builder.agent_nickname = patch.agent_nickname.clone().flatten(); builder.agent_role = patch.agent_role.clone().flatten(); builder.agent_path = patch.agent_path.clone().flatten(); diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index 4fb9f398c1..3db827e3a2 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -66,7 +66,7 @@ impl ThreadMetadataSync { created_at: Some(created_at), updated_at: Some(created_at), source: Some(params.source.clone()), - thread_source: Some(params.thread_source), + thread_source: Some(params.thread_source.clone()), agent_nickname: Some(params.source.get_nickname()), agent_role: Some(params.source.get_agent_role()), agent_path: Some(params.source.get_agent_path().map(Into::into)), @@ -201,7 +201,7 @@ impl ThreadMetadataSync { RolloutItem::SessionMeta(meta_line) if meta_line.meta.id == self.thread_id => { update.created_at = parse_session_timestamp(meta_line.meta.timestamp.as_str()); update.source = Some(meta_line.meta.source.clone()); - update.thread_source = Some(meta_line.meta.thread_source); + update.thread_source = Some(meta_line.meta.thread_source.clone()); update.agent_nickname = Some(meta_line.meta.agent_nickname.clone()); update.agent_role = Some(meta_line.meta.agent_role.clone()); update.agent_path = Some(meta_line.meta.agent_path.clone()); From 6e7ab529297ac3e8ef9110718e1eab688d4391ef Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Tue, 9 Jun 2026 12:48:04 -0700 Subject: [PATCH 09/17] [codex] Speed up local nextest runs (#26479) ## Why `just test` currently uses the CI-oriented nextest profile, which serializes app-server integration tests even on developer machines that can run several safely. Bounded local parallelism substantially shortens this common iteration loop without changing CI behavior. Eight-worker experiments were faster, but keeping them reliable required relaxing several test deadlines. Four workers for integration tests is a solid tradeoff that speeds up local testing without needing to change test logic. ## What changed - Add a `local` nextest profile that inherits the existing defaults. - Allow up to four app-server integration tests to run concurrently under that profile. - Make `just test` select the local profile on Unix and Windows. - Keep the default CI profile serialized and leave all test deadlines unchanged. The tests use separate processes, randomized temporary `CODEX_HOME` directories, and ephemeral ports. The remaining shared constraints are system resources; each app-server also uses a multi-thread Tokio runtime, and fuzzy-search tests can create additional worker threads, so the local cap remains intentionally conservative. ## Performance and validation All measurements below are warm, execution-only app-server runs with nextest retries disabled. On the current rebased branch, an AMD EPYC 7763 machine with 16 logical CPUs and 62 GiB RAM completed three consecutive runs: | Run | Nextest time | Wall time | Result | | --- | ---: | ---: | --- | | 1 | 142.941s | 145.17s | 836/836 passed | | 2 | 143.402s | 145.59s | 836/836 passed | | 3 | 142.870s | 145.08s | 836/836 passed | The mean wall time was 145.28s. The slow-inventory, approval replay, and zsh-fork tests all passed with their original deadlines. Earlier measurements on the same Linux machine, before the suite grew, showed the scaling that motivated the change: | App-server concurrency | Nextest time | Result | | --- | ---: | --- | | 1 | 369.5s | 572/572 passed | | 2 | 194.5s | 572/572 passed | | 4 | 111.0s mean over 3 runs | 3/3 clean | Four workers reduced that execution time by about 70%, a roughly 3.3x speedup over serialization. --- codex-rs/.config/nextest.toml | 14 ++++++++++++++ justfile | 4 ++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/codex-rs/.config/nextest.toml b/codex-rs/.config/nextest.toml index 01f4a98dc8..f3f3ca2607 100644 --- a/codex-rs/.config/nextest.toml +++ b/codex-rs/.config/nextest.toml @@ -8,12 +8,20 @@ retries = 1 [profile.default.junit] path = "junit.xml" +[profile.local] +inherits = "default" + [test-groups.app_server_protocol_codegen] max-threads = 1 [test-groups.app_server_integration] max-threads = 1 +# Higher concurrency causes integration test timeouts under resource contention +# on common developer machines. +[test-groups.app_server_integration_local] +max-threads = 4 + [test-groups.core_apply_patch_cli_integration] max-threads = 1 @@ -42,6 +50,12 @@ test-group = 'app_server_protocol_codegen' filter = 'package(codex-app-server) & kind(test)' test-group = 'app_server_integration' +[[profile.local.overrides]] +# Use up to four app-server subprocesses locally. The global nextest pool still +# limits this to the machine's logical CPU count. +filter = 'package(codex-app-server) & kind(test)' +test-group = 'app_server_integration_local' + [[profile.default.overrides]] # These tests exercise full Codex turns and apply_patch execution, and they are # sensitive to Windows runner process-startup stalls when many cases launch at once. diff --git a/justfile b/justfile index 3b993b767d..fa32a19e7d 100644 --- a/justfile +++ b/justfile @@ -75,11 +75,11 @@ install: # there should be no need to add `--all-features`. [unix] test *args: - RUST_MIN_STACK={{ rust_min_stack }} cargo nextest run --no-fail-fast "$@" + RUST_MIN_STACK={{ rust_min_stack }} NEXTEST_PROFILE=local cargo nextest run --no-fail-fast "$@" [windows] test *args: - $env:RUST_MIN_STACK = "{{ rust_min_stack }}"; cargo nextest run --no-fail-fast @($args | Select-Object -Skip 1) + $env:RUST_MIN_STACK = "{{ rust_min_stack }}"; $env:NEXTEST_PROFILE = "local"; cargo nextest run --no-fail-fast @($args | Select-Object -Skip 1) # Run from the repository root so scripts that resolve paths from `cwd` see # the same layout they use in GitHub Actions. From 7a7cee1be4e82733f393c95180242d11b50064d5 Mon Sep 17 00:00:00 2001 From: Eric Ning Date: Tue, 9 Jun 2026 12:52:05 -0700 Subject: [PATCH 10/17] fix: use plugin service route for remote uninstall (#27223) --- .../tests/suite/v2/plugin_uninstall.rs | 18 +++++++++--------- codex-rs/core-plugins/src/remote.rs | 2 +- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/plugin_uninstall.rs b/codex-rs/app-server/tests/suite/v2/plugin_uninstall.rs index 46299d1186..5d31d8390f 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_uninstall.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_uninstall.rs @@ -206,7 +206,7 @@ async fn plugin_uninstall_writes_remote_plugin_to_cloud_when_remote_plugin_enabl Mock::given(method("POST")) .and(path(format!( - "/backend-api/plugins/{REMOTE_PLUGIN_ID}/uninstall" + "/backend-api/ps/plugins/{REMOTE_PLUGIN_ID}/uninstall" ))) .and(header("authorization", "Bearer chatgpt-token")) .and(header("chatgpt-account-id", "account-123")) @@ -249,7 +249,7 @@ async fn plugin_uninstall_writes_remote_plugin_to_cloud_when_remote_plugin_enabl wait_for_remote_plugin_request_count( &server, "POST", - &format!("/plugins/{REMOTE_PLUGIN_ID}/uninstall"), + &format!("/ps/plugins/{REMOTE_PLUGIN_ID}/uninstall"), /*expected_count*/ 1, ) .await?; @@ -278,7 +278,7 @@ async fn plugin_uninstall_uses_detail_scope_for_cache_namespace() -> Result<()> Mock::given(method("POST")) .and(path(format!( - "/backend-api/plugins/{REMOTE_PLUGIN_ID}/uninstall" + "/backend-api/ps/plugins/{REMOTE_PLUGIN_ID}/uninstall" ))) .and(header("authorization", "Bearer chatgpt-token")) .and(header("chatgpt-account-id", "account-123")) @@ -321,7 +321,7 @@ async fn plugin_uninstall_uses_detail_scope_for_cache_namespace() -> Result<()> wait_for_remote_plugin_request_count( &server, "POST", - &format!("/plugins/{REMOTE_PLUGIN_ID}/uninstall"), + &format!("/ps/plugins/{REMOTE_PLUGIN_ID}/uninstall"), /*expected_count*/ 1, ) .await?; @@ -357,7 +357,7 @@ async fn plugin_uninstall_accepts_workspace_remote_plugin_id_shape() -> Result<( Mock::given(method("POST")) .and(path(format!( - "/backend-api/plugins/{WORKSPACE_REMOTE_PLUGIN_ID}/uninstall" + "/backend-api/ps/plugins/{WORKSPACE_REMOTE_PLUGIN_ID}/uninstall" ))) .and(header("authorization", "Bearer chatgpt-token")) .and(header("chatgpt-account-id", "account-123")) @@ -395,7 +395,7 @@ async fn plugin_uninstall_accepts_workspace_remote_plugin_id_shape() -> Result<( wait_for_remote_plugin_request_count( &server, "POST", - &format!("/plugins/{WORKSPACE_REMOTE_PLUGIN_ID}/uninstall"), + &format!("/ps/plugins/{WORKSPACE_REMOTE_PLUGIN_ID}/uninstall"), /*expected_count*/ 1, ) .await?; @@ -451,7 +451,7 @@ async fn plugin_uninstall_rejects_before_post_when_remote_detail_fetch_fails() - wait_for_remote_plugin_request_count( &server, "POST", - &format!("/plugins/{REMOTE_PLUGIN_ID}/uninstall"), + &format!("/ps/plugins/{REMOTE_PLUGIN_ID}/uninstall"), /*expected_count*/ 0, ) .await?; @@ -487,7 +487,7 @@ async fn plugin_uninstall_rejects_remote_plugin_id_with_spaces_before_network_ca wait_for_remote_plugin_request_count( &server, "POST", - "/plugins/sample plugin/uninstall", + "/ps/plugins/sample plugin/uninstall", /*expected_count*/ 0, ) .await?; @@ -522,7 +522,7 @@ async fn plugin_uninstall_rejects_invalid_remote_plugin_id_before_network_call() wait_for_remote_plugin_request_count( &server, "POST", - "/plugins/linear/../../oops/uninstall", + "/ps/plugins/linear/../../oops/uninstall", /*expected_count*/ 0, ) .await?; diff --git a/codex-rs/core-plugins/src/remote.rs b/codex-rs/core-plugins/src/remote.rs index 1885318727..2beeed2a1b 100644 --- a/codex-rs/core-plugins/src/remote.rs +++ b/codex-rs/core-plugins/src/remote.rs @@ -1124,7 +1124,7 @@ pub async fn uninstall_remote_plugin( let plugin_name = plugin.name; let base_url = config.chatgpt_base_url.trim_end_matches('/'); - let url = format!("{base_url}/plugins/{plugin_id}/uninstall"); + let url = format!("{base_url}/ps/plugins/{plugin_id}/uninstall"); let client = build_reqwest_client(); let request = authenticated_request(client.post(&url), auth)?; let response: RemotePluginMutationResponse = send_and_decode(request, &url).await?; From 472619f9fc6739d5c4fa499b3ef10007655b031d Mon Sep 17 00:00:00 2001 From: Winston Howes Date: Tue, 9 Jun 2026 13:08:01 -0700 Subject: [PATCH 11/17] Add SOCKS5 TCP MITM coverage (#22685) ## Summary - reuse the MITM HTTPS serving path for raw SOCKS5 TCP streams - route limited-mode and hooked SOCKS5 TCP requests through MITM before dialing upstream - keep SOCKS5 UDP limited-mode behavior unchanged ## Validation - `just fmt` - `just test -p codex-network-proxy` - `just fix -p codex-network-proxy` - `git diff --check` --- codex-rs/network-proxy/README.md | 9 +- codex-rs/network-proxy/src/config.rs | 2 +- codex-rs/network-proxy/src/mitm.rs | 26 +- codex-rs/network-proxy/src/socks5.rs | 510 +++++++++++++++++++++------ 4 files changed, 419 insertions(+), 128 deletions(-) diff --git a/codex-rs/network-proxy/README.md b/codex-rs/network-proxy/README.md index 300cd3fcb0..d5d0730227 100644 --- a/codex-rs/network-proxy/README.md +++ b/codex-rs/network-proxy/README.md @@ -107,9 +107,9 @@ When a request is blocked, the proxy responds with `403` and includes: - `blocked-by-method-policy` - `blocked-by-policy` -In "limited" mode, only `GET`, `HEAD`, and `OPTIONS` are allowed. HTTPS `CONNECT` requests require -MITM to enforce limited-mode method policy; otherwise they are blocked. SOCKS5 remains blocked in -limited mode. +In "limited" mode, only `GET`, `HEAD`, and `OPTIONS` are allowed. HTTPS `CONNECT` requests and +HTTPS SOCKS5 TCP targets on `:443` require MITM to enforce limited-mode method policy; otherwise +they are blocked. SOCKS5 UDP and non-HTTPS SOCKS5 TCP remain blocked in limited mode. Websocket clients typically tunnel `wss://` through HTTPS `CONNECT`; those CONNECT targets still go through the same host allowlist/denylist checks. @@ -215,7 +215,8 @@ what it can reasonably guarantee. allowlisted (best-effort DNS lookup). - Limited mode enforcement: - only `GET`, `HEAD`, and `OPTIONS` are allowed - - HTTPS `CONNECT` remains a tunnel; limited-mode method enforcement does not apply to HTTPS + - HTTPS `CONNECT` requests and HTTPS SOCKS5 TCP targets on `:443` require MITM so the proxy can + enforce limited-mode method policy; SOCKS5 UDP and non-HTTPS SOCKS5 TCP remain blocked - Listener safety defaults: - the HTTP proxy listener clamps non-loopback binds unless explicitly enabled via `dangerously_allow_non_loopback_proxy` diff --git a/codex-rs/network-proxy/src/config.rs b/codex-rs/network-proxy/src/config.rs index 0b0a8233f5..d9cef46b15 100644 --- a/codex-rs/network-proxy/src/config.rs +++ b/codex-rs/network-proxy/src/config.rs @@ -276,7 +276,7 @@ impl NetworkProxySettings { pub enum NetworkMode { /// Limited (read-only) access: only GET/HEAD/OPTIONS are allowed for HTTP. HTTPS CONNECT is /// blocked unless MITM is enabled so the proxy can enforce method policy on inner requests. - /// SOCKS5 remains blocked in limited mode. + /// SOCKS5 UDP and non-HTTPS SOCKS5 TCP remain blocked in limited mode. Limited, /// Full network access: all HTTP methods are allowed. HTTPS CONNECTs are tunneled directly. /// MITM hooks do not currently make full mode enter MITM. diff --git a/codex-rs/network-proxy/src/mitm.rs b/codex-rs/network-proxy/src/mitm.rs index 9ee5bfc253..345c5b5032 100644 --- a/codex-rs/network-proxy/src/mitm.rs +++ b/codex-rs/network-proxy/src/mitm.rs @@ -21,10 +21,12 @@ use rama_core::Layer; use rama_core::Service; use rama_core::bytes::Bytes; use rama_core::error::BoxError; +use rama_core::extensions::ExtensionsMut; use rama_core::extensions::ExtensionsRef; -use rama_core::futures::stream::Stream; +use rama_core::futures::stream::Stream as FuturesStream; use rama_core::rt::Executor; use rama_core::service::service_fn; +use rama_core::stream::Stream; use rama_http::Body; use rama_http::BodyDataStream; use rama_http::HeaderMap; @@ -135,17 +137,25 @@ impl MitmState { /// Terminate the upgraded CONNECT stream with a generated leaf cert and proxy inner HTTPS traffic. pub(crate) async fn mitm_tunnel(upgraded: Upgraded) -> Result<()> { - let mitm = upgraded + mitm_stream(upgraded).await +} + +/// Terminate a raw client stream with a generated leaf cert and proxy inner HTTPS traffic. +pub(crate) async fn mitm_stream(stream: S) -> Result<()> +where + S: Stream + Unpin + ExtensionsMut, +{ + let mitm = stream .extensions() .get::>() .cloned() .context("missing MITM state")?; - let app_state = upgraded + let app_state = stream .extensions() .get::>() .cloned() .context("missing app state")?; - let target = upgraded + let target = stream .extensions() .get::() .context("missing proxy target")? @@ -154,7 +164,7 @@ pub(crate) async fn mitm_tunnel(upgraded: Upgraded) -> Result<()> { let target_host = normalize_host(&target.host.to_string()); let target_port = target.port; let acceptor_data = mitm.tls_acceptor_data_for_host(&target_host)?; - let mode = upgraded + let mode = stream .extensions() .get::() .copied() @@ -169,7 +179,7 @@ pub(crate) async fn mitm_tunnel(upgraded: Upgraded) -> Result<()> { mitm, }); - let executor = upgraded + let executor = stream .extensions() .get::() .cloned() @@ -194,7 +204,7 @@ pub(crate) async fn mitm_tunnel(upgraded: Upgraded) -> Result<()> { .into_layer(http_service); https_service - .serve(upgraded) + .serve(stream) .await .map_err(|err| anyhow!("MITM serve error: {err}"))?; Ok(()) @@ -456,7 +466,7 @@ struct InspectStream { max_body_bytes: usize, } -impl Stream for InspectStream { +impl FuturesStream for InspectStream { type Item = Result; fn poll_next(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { diff --git a/codex-rs/network-proxy/src/socks5.rs b/codex-rs/network-proxy/src/socks5.rs index 8369c4da80..d9a8a559fd 100644 --- a/codex-rs/network-proxy/src/socks5.rs +++ b/codex-rs/network-proxy/src/socks5.rs @@ -1,5 +1,6 @@ use crate::config::NetworkMode; use crate::connect_policy::TargetCheckedTcpConnector; +use crate::mitm; use crate::network_policy::BlockDecisionAuditEventArgs; use crate::network_policy::NetworkDecision; use crate::network_policy::NetworkDecisionSource; @@ -24,10 +25,17 @@ use anyhow::Result; use rama_core::Layer; use rama_core::Service; use rama_core::error::BoxError; +use rama_core::extensions::Extensions; +use rama_core::extensions::ExtensionsMut; use rama_core::extensions::ExtensionsRef; use rama_core::layer::AddInputExtensionLayer; use rama_core::service::service_fn; +use rama_net::address::HostWithPort; use rama_net::client::EstablishedClientConnection; +use rama_net::proxy::ProxyRequest; +use rama_net::proxy::ProxyTarget; +use rama_net::proxy::StreamForwardService; +use rama_net::stream::Socket; use rama_net::stream::SocketInfo; use rama_socks5::Socks5Acceptor; use rama_socks5::server::DefaultConnector; @@ -40,8 +48,14 @@ use rama_tcp::server::TcpListener; use std::io; use std::net::SocketAddr; use std::net::TcpListener as StdTcpListener; +use std::pin::Pin; use std::sync::Arc; +use std::task::Context as TaskContext; +use std::task::Poll; use std::time::Instant; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::io::ReadBuf; use tracing::error; use tracing::info; use tracing::warn; @@ -88,7 +102,9 @@ async fn run_socks5_with_listener( match state.network_mode().await { Ok(NetworkMode::Limited) => { - info!("SOCKS5 is blocked in limited mode; set mode=\"full\" to allow SOCKS5"); + info!( + "SOCKS5 UDP and non-HTTPS SOCKS5 TCP are blocked in limited mode; HTTPS SOCKS5 TCP requires MITM inspection" + ); } Ok(NetworkMode::Full) => {} Err(err) => { @@ -106,7 +122,10 @@ async fn run_socks5_with_listener( } }); - let socks_connector = DefaultConnector::default().with_connector(policy_tcp_connector); + let socks_proxy = service_fn(|request| async move { proxy_socks5_tcp(request).await }); + let socks_connector = DefaultConnector::default() + .with_connector(policy_tcp_connector) + .with_service(socks_proxy); let base = Socks5Acceptor::new().with_connector(socks_connector); if enable_socks5_udp { @@ -135,7 +154,7 @@ async fn handle_socks5_tcp( req: TcpRequest, tcp_connector: TargetCheckedTcpConnector, policy_decider: Option>, -) -> Result, BoxError> { +) -> Result, BoxError> { let app_state = req .extensions() .get::>() @@ -144,6 +163,7 @@ async fn handle_socks5_tcp( let host = normalize_host(&req.authority.host.to_string()); let port = req.authority.port; + let target = req.authority.clone(); if host.is_empty() { return Err(io::Error::new(io::ErrorKind::InvalidInput, "invalid host").into()); } @@ -196,94 +216,52 @@ async fn handle_socks5_tcp( } } - match app_state.network_mode().await { - Ok(NetworkMode::Limited) => { - emit_socks_block_decision_audit_event( - &app_state, - NetworkDecisionSource::ModeGuard, - REASON_METHOD_NOT_ALLOWED, - NetworkProtocol::Socks5Tcp, - host.as_str(), - port, - client.as_deref(), - ); - let details = PolicyDecisionDetails { - decision: NetworkPolicyDecision::Deny, - reason: REASON_METHOD_NOT_ALLOWED, - source: NetworkDecisionSource::ModeGuard, - protocol: NetworkProtocol::Socks5Tcp, - host: &host, - port, - }; - let _ = app_state - .record_blocked(BlockedRequest::new(BlockedRequestArgs { - host: host.clone(), - reason: REASON_METHOD_NOT_ALLOWED.to_string(), - client: client.clone(), - method: None, - mode: Some(NetworkMode::Limited), - protocol: "socks5".to_string(), - decision: Some(details.decision.as_str().to_string()), - source: Some(details.source.as_str().to_string()), - port: Some(port), - })) - .await; - let client = client.as_deref().unwrap_or_default(); - warn!( - "SOCKS blocked by method policy (client={client}, host={host}, mode=limited, allowed_methods=GET, HEAD, OPTIONS)" - ); - return Err(policy_denied_error(REASON_METHOD_NOT_ALLOWED, &details).into()); - } - Ok(NetworkMode::Full) => {} + let mode = match app_state.network_mode().await { + Ok(mode) => mode, Err(err) => { error!("failed to evaluate method policy: {err}"); return Err(io::Error::other("proxy error").into()); } - } - - match app_state.host_has_mitm_hooks(&host).await { - Ok(true) => { - emit_socks_block_decision_audit_event( - &app_state, - NetworkDecisionSource::ModeGuard, - REASON_MITM_REQUIRED, - NetworkProtocol::Socks5Tcp, - host.as_str(), - port, - client.as_deref(), - ); - let details = PolicyDecisionDetails { - decision: NetworkPolicyDecision::Deny, - reason: REASON_MITM_REQUIRED, - source: NetworkDecisionSource::ModeGuard, - protocol: NetworkProtocol::Socks5Tcp, - host: &host, - port, - }; - let _ = app_state - .record_blocked(BlockedRequest::new(BlockedRequestArgs { - host: host.clone(), - reason: REASON_MITM_REQUIRED.to_string(), - client: client.clone(), - method: None, - mode: Some(NetworkMode::Full), - protocol: "socks5".to_string(), - decision: Some(details.decision.as_str().to_string()), - source: Some(details.source.as_str().to_string()), - port: Some(port), - })) - .await; - let client = client.as_deref().unwrap_or_default(); - warn!( - "SOCKS blocked; MITM required to enforce HTTPS policy (client={client}, host={host}, mode=full)" - ); - return Err(policy_denied_error(REASON_MITM_REQUIRED, &details).into()); - } - Ok(false) => {} - Err(err) => { - error!("failed to inspect MITM hooks for {host}: {err}"); - return Err(io::Error::other("proxy error").into()); - } + }; + // SOCKS5 only exposes host and port, so only the default HTTPS port is identifiable as a + // TLS stream that the HTTPS MITM path can safely terminate. + let socks5_tcp_target_is_https = port == 443; + if mode == NetworkMode::Limited && !socks5_tcp_target_is_https { + emit_socks_block_decision_audit_event( + &app_state, + NetworkDecisionSource::ModeGuard, + REASON_METHOD_NOT_ALLOWED, + NetworkProtocol::Socks5Tcp, + host.as_str(), + port, + client.as_deref(), + ); + let details = PolicyDecisionDetails { + decision: NetworkPolicyDecision::Deny, + reason: REASON_METHOD_NOT_ALLOWED, + source: NetworkDecisionSource::ModeGuard, + protocol: NetworkProtocol::Socks5Tcp, + host: &host, + port, + }; + let _ = app_state + .record_blocked(BlockedRequest::new(BlockedRequestArgs { + host: host.clone(), + reason: REASON_METHOD_NOT_ALLOWED.to_string(), + client: client.clone(), + method: None, + mode: Some(NetworkMode::Limited), + protocol: "socks5".to_string(), + decision: Some(details.decision.as_str().to_string()), + source: Some(details.source.as_str().to_string()), + port: Some(port), + })) + .await; + let client = client.as_deref().unwrap_or_default(); + warn!( + "SOCKS blocked; limited mode only supports HTTPS MITM (client={client}, host={host}, port={port})" + ); + return Err(policy_denied_error(REASON_METHOD_NOT_ALLOWED, &details).into()); } let request = NetworkPolicyRequest::new(NetworkPolicyRequestArgs { @@ -337,9 +315,85 @@ async fn handle_socks5_tcp( } } + let host_has_mitm_hooks = match app_state.host_has_mitm_hooks(&host).await { + Ok(has_hooks) => has_hooks, + Err(err) => { + error!("failed to inspect MITM hooks for {host}: {err}"); + return Err(io::Error::other("proxy error").into()); + } + }; + let mitm_state = match app_state.mitm_state().await { + Ok(state) => state, + Err(err) => { + error!("failed to load MITM state: {err}"); + return Err(io::Error::other("proxy error").into()); + } + }; + let socks_needs_mitm = + socks5_tcp_target_is_https && (mode == NetworkMode::Limited || host_has_mitm_hooks); + if (host_has_mitm_hooks && !socks5_tcp_target_is_https) + || (socks_needs_mitm && mitm_state.is_none()) + { + emit_socks_block_decision_audit_event( + &app_state, + NetworkDecisionSource::ModeGuard, + REASON_MITM_REQUIRED, + NetworkProtocol::Socks5Tcp, + host.as_str(), + port, + client.as_deref(), + ); + let details = PolicyDecisionDetails { + decision: NetworkPolicyDecision::Deny, + reason: REASON_MITM_REQUIRED, + source: NetworkDecisionSource::ModeGuard, + protocol: NetworkProtocol::Socks5Tcp, + host: &host, + port, + }; + let _ = app_state + .record_blocked(BlockedRequest::new(BlockedRequestArgs { + host: host.clone(), + reason: REASON_MITM_REQUIRED.to_string(), + client: client.clone(), + method: None, + mode: Some(mode), + protocol: "socks5".to_string(), + decision: Some(details.decision.as_str().to_string()), + source: Some(details.source.as_str().to_string()), + port: Some(port), + })) + .await; + let client = client.as_deref().unwrap_or_default(); + warn!( + "SOCKS blocked; MITM required to enforce HTTPS policy (client={client}, host={host}, mode={mode:?}, hooked_host={host_has_mitm_hooks}, https_target={socks5_tcp_target_is_https})" + ); + return Err(policy_denied_error(REASON_MITM_REQUIRED, &details).into()); + } + + if socks_needs_mitm && let Some(mitm_state) = mitm_state { + let client = client.as_deref().unwrap_or_default(); + info!("SOCKS MITM enabled (client={client}, host={host}, port={port}, mode={mode:?})"); + return Ok(EstablishedClientConnection { + input: req, + conn: Socks5TcpConnection::Mitm { + target, + mode, + mitm: mitm_state, + extensions: Extensions::new(), + }, + }); + } + info!("SOCKS upstream dial started (host={host}, port={port})"); let connect_started_at = Instant::now(); - let result = tcp_connector.serve(req).await; + let result = tcp_connector.serve(req).await.map(|connection| { + let EstablishedClientConnection { input, conn } = connection; + EstablishedClientConnection { + input, + conn: Socks5TcpConnection::Direct(conn), + } + }); match &result { Ok(_) => info!( "SOCKS upstream dial established (host={host}, port={port}, elapsed_ms={})", @@ -353,6 +407,113 @@ async fn handle_socks5_tcp( result } +/// Internal connector output for SOCKS5 TCP. MITM requests do not dial upstream before the +/// inner HTTPS request is inspected, so they carry the target metadata instead of a socket. +#[derive(Debug)] +enum Socks5TcpConnection { + Direct(TcpStream), + Mitm { + target: HostWithPort, + mode: NetworkMode, + mitm: Arc, + extensions: Extensions, + }, +} + +impl AsyncRead for Socks5TcpConnection { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut TaskContext<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + match self.get_mut() { + Self::Direct(stream) => Pin::new(stream).poll_read(cx, buf), + Self::Mitm { .. } => Poll::Ready(Ok(())), + } + } +} + +impl AsyncWrite for Socks5TcpConnection { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut TaskContext<'_>, + buf: &[u8], + ) -> Poll> { + match self.get_mut() { + Self::Direct(stream) => Pin::new(stream).poll_write(cx, buf), + Self::Mitm { .. } => Poll::Ready(Ok(buf.len())), + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { + match self.get_mut() { + Self::Direct(stream) => Pin::new(stream).poll_flush(cx), + Self::Mitm { .. } => Poll::Ready(Ok(())), + } + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { + match self.get_mut() { + Self::Direct(stream) => Pin::new(stream).poll_shutdown(cx), + Self::Mitm { .. } => Poll::Ready(Ok(())), + } + } +} + +impl Socket for Socks5TcpConnection { + fn local_addr(&self) -> io::Result { + match self { + Self::Direct(stream) => stream.local_addr(), + Self::Mitm { .. } => Ok(SocketAddr::from(([0, 0, 0, 0], 0))), + } + } + + fn peer_addr(&self) -> io::Result { + match self { + Self::Direct(stream) => stream.peer_addr(), + Self::Mitm { .. } => Ok(SocketAddr::from(([0, 0, 0, 0], 0))), + } + } +} + +impl ExtensionsRef for Socks5TcpConnection { + fn extensions(&self) -> &Extensions { + match self { + Self::Direct(stream) => stream.extensions(), + Self::Mitm { extensions, .. } => extensions, + } + } +} + +impl ExtensionsMut for Socks5TcpConnection { + fn extensions_mut(&mut self) -> &mut Extensions { + match self { + Self::Direct(stream) => stream.extensions_mut(), + Self::Mitm { extensions, .. } => extensions, + } + } +} + +async fn proxy_socks5_tcp( + request: ProxyRequest, +) -> Result<(), BoxError> { + let ProxyRequest { mut source, target } = request; + match target { + Socks5TcpConnection::Direct(target) => StreamForwardService::default() + .serve(ProxyRequest { source, target }) + .await + .map_err(Into::into), + Socks5TcpConnection::Mitm { + target, mode, mitm, .. + } => { + source.extensions_mut().insert(ProxyTarget(target)); + source.extensions_mut().insert(mode); + source.extensions_mut().insert(mitm); + mitm::mitm_stream(source).await.map_err(Into::into) + } + } +} + async fn inspect_socks5_udp( request: RelayRequest, state: Arc, @@ -554,7 +715,6 @@ mod tests { use crate::network_policy::test_support::find_event_by_name; use crate::runtime::ConfigReloader; use crate::runtime::ConfigState; - use crate::runtime::network_proxy_state_for_policy; use crate::state::NetworkProxyConstraints; use crate::state::build_config_state; use async_trait::async_trait; @@ -567,6 +727,11 @@ mod tests { use std::net::IpAddr; use std::net::Ipv4Addr; use std::sync::Arc; + use std::sync::Mutex; + + // Managed MITM CA files live under the shared test CODEX_HOME, so MITM-enabled config state + // must be materialized one test at a time. + static MITM_CONFIG_STATE_LOCK: Mutex<()> = Mutex::new(()); #[derive(Clone)] struct StaticReloader { @@ -590,6 +755,10 @@ mod tests { fn state_for_settings(network: NetworkProxySettings) -> Arc { let config = NetworkProxyConfig { network }; + let _mitm_config_state_guard = config + .network + .mitm + .then(|| MITM_CONFIG_STATE_LOCK.lock().unwrap()); let state = build_config_state(config, NetworkProxyConstraints::default()).unwrap(); let reloader = Arc::new(StaticReloader { state: state.clone(), @@ -639,43 +808,56 @@ mod tests { } #[tokio::test(flavor = "current_thread")] - async fn handle_socks5_tcp_blocks_hooked_host_in_full_mode() { - let state = Arc::new(network_proxy_state_for_policy(NetworkProxySettings { + async fn handle_socks5_tcp_uses_mitm_in_limited_mode() { + let mut settings = NetworkProxySettings { enabled: true, - mode: NetworkMode::Full, + mode: NetworkMode::Limited, mitm: true, - mitm_hooks: vec![MitmHookConfig { - host: "api.github.com".to_string(), - matcher: MitmHookMatchConfig { - methods: vec!["GET".to_string()], - path_prefixes: vec!["/".to_string()], - ..MitmHookMatchConfig::default() - }, - ..MitmHookConfig::default() - }], ..NetworkProxySettings::default() - })); + }; + settings.set_allowed_domains(vec!["example.com".to_string()]); + let state = state_for_settings(settings); let mut request = - TcpRequest::new(HostWithPort::try_from("api.github.com:443").expect("valid authority")); + TcpRequest::new(HostWithPort::try_from("example.com:443").expect("valid authority")); + request.extensions_mut().insert(state.clone()); + + let result = handle_socks5_tcp( + request, + TargetCheckedTcpConnector::new(state), + /*policy_decider*/ None, + ) + .await + .expect("limited-mode HTTPS should use MITM"); + + assert!(matches!(result.conn, Socks5TcpConnection::Mitm { .. })); + } + + #[tokio::test(flavor = "current_thread")] + async fn handle_socks5_tcp_blocks_non_https_in_limited_mode() { + let mut settings = NetworkProxySettings { + enabled: true, + mode: NetworkMode::Limited, + ..NetworkProxySettings::default() + }; + settings.set_allowed_domains(vec!["example.com".to_string()]); + let state = state_for_settings(settings); + let mut request = + TcpRequest::new(HostWithPort::try_from("example.com:80").expect("valid authority")); request.extensions_mut().insert(state.clone()); let (result, events) = capture_events(|| async { handle_socks5_tcp( request, - TargetCheckedTcpConnector::new(state.clone()), + TargetCheckedTcpConnector::new(state), /*policy_decider*/ None, ) .await }) .await; - assert!(result.is_err(), "hooked host should require MITM"); - - let blocked = state.drain_blocked().await.unwrap(); - assert_eq!(blocked.len(), 1); - assert_eq!(blocked[0].reason, REASON_MITM_REQUIRED); - assert_eq!(blocked[0].host, "api.github.com"); - assert_eq!(blocked[0].port, Some(443)); - assert_eq!(blocked[0].protocol, "socks5"); + assert!( + result.is_err(), + "limited-mode non-HTTPS SOCKS should be denied" + ); let event = find_event_by_name(&events, POLICY_DECISION_EVENT_NAME) .expect("expected policy decision event"); @@ -684,18 +866,116 @@ mod tests { assert_eq!(event.field("network.policy.source"), Some("mode_guard")); assert_eq!( event.field("network.policy.reason"), - Some(REASON_MITM_REQUIRED) + Some(REASON_METHOD_NOT_ALLOWED) ); assert_eq!( event.field("network.transport.protocol"), Some("socks5_tcp") ); - assert_eq!(event.field("server.address"), Some("api.github.com")); - assert_eq!(event.field("server.port"), Some("443")); + assert_eq!(event.field("server.address"), Some("example.com")); + assert_eq!(event.field("server.port"), Some("80")); assert_eq!(event.field("http.request.method"), Some("none")); assert_eq!(event.field("client.address"), Some("unknown")); } + #[tokio::test(flavor = "current_thread")] + async fn handle_socks5_tcp_blocks_limited_mode_without_mitm_state() { + let mut settings = NetworkProxySettings { + enabled: true, + mode: NetworkMode::Limited, + ..NetworkProxySettings::default() + }; + settings.set_allowed_domains(vec!["example.com".to_string()]); + let state = state_for_settings(settings); + let mut request = + TcpRequest::new(HostWithPort::try_from("example.com:443").expect("valid authority")); + request.extensions_mut().insert(state.clone()); + + let err = handle_socks5_tcp( + request, + TargetCheckedTcpConnector::new(state), + /*policy_decider*/ None, + ) + .await + .expect_err("limited-mode HTTPS requires MITM"); + + assert!( + format!("{err:?}").contains("MITM required"), + "unexpected error: {err:?}" + ); + } + + #[tokio::test(flavor = "current_thread")] + async fn handle_socks5_tcp_uses_mitm_for_hooked_host_in_full_mode() { + let mut settings = NetworkProxySettings { + enabled: true, + mode: NetworkMode::Full, + mitm: true, + mitm_hooks: vec![MitmHookConfig { + host: "api.github.com".to_string(), + matcher: MitmHookMatchConfig { + methods: vec!["POST".to_string()], + path_prefixes: vec!["/repos/openai/".to_string()], + ..MitmHookMatchConfig::default() + }, + ..MitmHookConfig::default() + }], + ..NetworkProxySettings::default() + }; + settings.set_allowed_domains(vec!["api.github.com".to_string()]); + let state = state_for_settings(settings); + let mut request = + TcpRequest::new(HostWithPort::try_from("api.github.com:443").expect("valid authority")); + request.extensions_mut().insert(state.clone()); + + let result = handle_socks5_tcp( + request, + TargetCheckedTcpConnector::new(state), + /*policy_decider*/ None, + ) + .await + .expect("hooked HTTPS should use MITM"); + + assert!(matches!(result.conn, Socks5TcpConnection::Mitm { .. })); + } + + #[tokio::test(flavor = "current_thread")] + async fn handle_socks5_tcp_blocks_hooked_non_https_host_in_full_mode() { + let mut settings = NetworkProxySettings { + enabled: true, + mode: NetworkMode::Full, + mitm: true, + mitm_hooks: vec![MitmHookConfig { + host: "api.github.com".to_string(), + matcher: MitmHookMatchConfig { + methods: vec!["POST".to_string()], + path_prefixes: vec!["/repos/openai/".to_string()], + ..MitmHookMatchConfig::default() + }, + ..MitmHookConfig::default() + }], + ..NetworkProxySettings::default() + }; + settings.set_allowed_domains(vec!["api.github.com".to_string()]); + let state = state_for_settings(settings); + let mut request = + TcpRequest::new(HostWithPort::try_from("api.github.com:80").expect("valid authority")); + request.extensions_mut().insert(state.clone()); + + let err = handle_socks5_tcp( + request, + TargetCheckedTcpConnector::new(state), + /*policy_decider*/ None, + ) + .await + .expect_err("hooked non-HTTPS SOCKS should require MITM"); + + assert!( + format!("{err:?}").contains("MITM required"), + "unexpected error: {err:?}" + ); + } + #[tokio::test(flavor = "current_thread")] async fn inspect_socks5_udp_emits_block_decision_for_mode_guard_deny() { let state = state_for_settings(NetworkProxySettings { From 36bf63a5cf02dafb1574b2396597ed1ef16bb5d7 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Tue, 9 Jun 2026 13:23:31 -0700 Subject: [PATCH 12/17] Allow creating a new goal after completion (#26681) ## Why Users have indicated that they want an agent to be able to create a new goal for itself after completing the previous goal. Currently, that's not possible because agents cannot overwrite an existing goal even if it's complete. This PR removes this limitation and allows `create_goal` to overwrite an existing goal if it is in the `complete` state. ## What changed `create_goal` now replaces the existing goal only when its status is `complete`. The replacement is performed atomically in the goal store, creates a fresh active goal with reset usage, and continues to reject creation while any unfinished goal exists. App server clients see a single `thread/goal/updated` event when the previous goal is replaced with the new one. The tool description and error message now reflect these semantics. ## What didn't change Agents are not allowed to create a new goal (overwrite their existing goal) if an existing goal is still active, blocked, paused, or in any other state other than "completed". --- codex-rs/ext/goal/src/spec.rs | 4 +-- codex-rs/ext/goal/src/tool.rs | 2 +- .../ext/goal/tests/goal_extension_backend.rs | 25 +++++++++++++++++-- codex-rs/state/src/runtime/goals.rs | 11 +++++++- 4 files changed, 36 insertions(+), 6 deletions(-) diff --git a/codex-rs/ext/goal/src/spec.rs b/codex-rs/ext/goal/src/spec.rs index 2c92c03848..6b97a24e9c 100644 --- a/codex-rs/ext/goal/src/spec.rs +++ b/codex-rs/ext/goal/src/spec.rs @@ -27,7 +27,7 @@ pub fn create_create_goal_tool() -> ToolSpec { ( "objective".to_string(), JsonSchema::string(Some( - "Required. The concrete objective to start pursuing. This starts a new active goal only when no goal is currently defined; if a goal already exists, this tool fails." + "Required. The concrete objective to start pursuing. This starts a new active goal when no goal exists or replaces the current goal when it is complete." .to_string(), )), ), @@ -44,7 +44,7 @@ pub fn create_create_goal_tool() -> ToolSpec { name: CREATE_GOAL_TOOL_NAME.to_string(), description: format!( r#"Create a goal only when explicitly requested by the user or system/developer instructions; do not infer goals from ordinary tasks. -Set token_budget only when an explicit token budget is requested. Fails if a goal exists; use {UPDATE_GOAL_TOOL_NAME} only for status."# +Set token_budget only when an explicit token budget is requested. Fails if an unfinished goal exists; use {UPDATE_GOAL_TOOL_NAME} only for status."# ), strict: false, defer_loading: None, diff --git a/codex-rs/ext/goal/src/tool.rs b/codex-rs/ext/goal/src/tool.rs index 17947aa343..8c0f513004 100644 --- a/codex-rs/ext/goal/src/tool.rs +++ b/codex-rs/ext/goal/src/tool.rs @@ -191,7 +191,7 @@ impl GoalToolExecutor { .map_err(|err| FunctionCallError::RespondToModel(format!("failed to create goal: {err}")))? .ok_or_else(|| { FunctionCallError::RespondToModel( - "cannot create a new goal because this thread already has a goal; use update_goal only when the existing goal is complete" + "cannot create a new goal because this thread has an unfinished goal; complete the existing goal first" .to_string(), ) })?; diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index f89a08ebbb..b922f5332d 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -124,7 +124,7 @@ async fn goal_tools_hidden_for_review_subagents() -> anyhow::Result<()> { } #[tokio::test] -async fn installed_goal_tools_reject_duplicate_goal_creation() -> anyhow::Result<()> { +async fn installed_goal_tools_only_replace_complete_goal() -> anyhow::Result<()> { let runtime = test_runtime().await?; let thread_id = test_thread_id()?; seed_thread_metadata(runtime.as_ref(), thread_id).await?; @@ -152,10 +152,31 @@ async fn installed_goal_tools_reject_duplicate_goal_creation() -> anyhow::Result assert_eq!( err, FunctionCallError::RespondToModel( - "cannot create a new goal because this thread already has a goal; use update_goal only when the existing goal is complete" + "cannot create a new goal because this thread has an unfinished goal; complete the existing goal first" .to_string() ) ); + + let update_tool = tool_by_name(&tools, "update_goal"); + update_tool + .handle(tool_call( + "update_goal", + "call-complete-goal", + json!({ "status": "complete" }), + )) + .await?; + + let invocation = tool_call( + "create_goal", + "call-create-goal-3", + json!({ "objective": "replacement goal" }), + ); + let output = create_tool.handle(invocation.clone()).await?; + let result = output.code_mode_result(&invocation.payload); + + assert_eq!(json!("replacement goal"), result["goal"]["objective"]); + assert_eq!(json!("active"), result["goal"]["status"]); + assert_eq!(json!(0), result["goal"]["tokensUsed"]); Ok(()) } diff --git a/codex-rs/state/src/runtime/goals.rs b/codex-rs/state/src/runtime/goals.rs index aa0e90462b..50e8cf068e 100644 --- a/codex-rs/state/src/runtime/goals.rs +++ b/codex-rs/state/src/runtime/goals.rs @@ -141,7 +141,16 @@ INSERT INTO thread_goals ( created_at_ms, updated_at_ms ) VALUES (?, ?, ?, ?, ?, 0, 0, ?, ?) -ON CONFLICT(thread_id) DO NOTHING +ON CONFLICT(thread_id) DO UPDATE SET + goal_id = excluded.goal_id, + objective = excluded.objective, + status = excluded.status, + token_budget = excluded.token_budget, + tokens_used = 0, + time_used_seconds = 0, + created_at_ms = excluded.created_at_ms, + updated_at_ms = excluded.updated_at_ms +WHERE thread_goals.status = 'complete' RETURNING thread_id, goal_id, From 8e69d29521488506182d2b5814851cdbb90c1354 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Tue, 9 Jun 2026 13:26:00 -0700 Subject: [PATCH 13/17] Reduce TUI legacy core dependencies (#26711) ## Why The TUI still reached through `app-server-client::legacy_core` for thread-name normalization and project-instruction filename details. In particular, checking the TUI's local filesystem for `/init` is incorrect for remote app-server sessions, where the server owns the working directory and instruction discovery. ## What changed - use the instruction source paths supplied by the app server to decide whether `/init` should avoid overwriting project instructions - keep the small thread-name normalization helper local to the TUI - remove the now-unused instruction filename constants, utility module, and other unused `legacy_core` re-exports - make status helper tests independent of concrete instruction filenames ## Verification - `just test -p codex-app-server-client` - `just test -p codex-tui slash_init_skips_when_project_instructions_are_loaded` - `just test -p codex-tui` ran 2,799 tests; 2,797 passed and two unrelated guardian feature-flag tests failed reproducibly in untouched code ### Manual test Started an app server over WebSocket with a remote workspace containing `AGENTS.md`, then connected the TUI using `--remote`. After confirming `thread/start` returned the file in `instructionSources`, deleted `AGENTS.md` and ran `/init` in the existing session. The TUI still reported that project instructions already existed and skipped `/init`. The trace contained no `turn/start` request, confirming the decision came from app-server session state rather than a new client-local filesystem check. --- codex-rs/app-server-client/src/lib.rs | 20 ------------ codex-rs/tui/prompt_for_init_command.md | 1 + codex-rs/tui/src/chatwidget.rs | 6 +++- codex-rs/tui/src/chatwidget/interaction.rs | 2 +- codex-rs/tui/src/chatwidget/slash_dispatch.rs | 10 +----- .../src/chatwidget/tests/slash_commands.rs | 31 +++---------------- codex-rs/tui/src/status/helpers.rs | 12 +++---- 7 files changed, 18 insertions(+), 64 deletions(-) diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index c75915629c..11684c6360 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -70,13 +70,9 @@ pub use crate::remote::RemoteAppServerEndpoint; /// module exists so clients can remove a direct `codex-core` dependency /// while legacy startup/config paths are migrated to RPCs. pub mod legacy_core { - pub use codex_core::DEFAULT_AGENTS_MD_FILENAME; - pub use codex_core::LOCAL_AGENTS_MD_FILENAME; - pub use codex_core::McpManager; pub use codex_core::check_execpolicy_for_warnings; pub use codex_core::format_exec_policy_error_with_source; pub use codex_core::grant_read_root_non_elevated; - pub use codex_core::web_search_detail; pub mod config { pub use codex_core::config::*; @@ -86,10 +82,6 @@ pub mod legacy_core { } } - pub mod connectors { - pub use codex_core::connectors::*; - } - pub mod otel_init { pub use codex_core::otel_init::*; } @@ -98,22 +90,10 @@ pub mod legacy_core { pub use codex_core::personality_migration::*; } - pub mod review_format { - pub use codex_core::review_format::*; - } - - pub mod review_prompts { - pub use codex_core::review_prompts::*; - } - pub mod test_support { pub use codex_core::test_support::*; } - pub mod util { - pub use codex_core::util::*; - } - pub mod windows_sandbox { pub use codex_core::windows_sandbox::*; } diff --git a/codex-rs/tui/prompt_for_init_command.md b/codex-rs/tui/prompt_for_init_command.md index b8fd3886b3..e96b002df7 100644 --- a/codex-rs/tui/prompt_for_init_command.md +++ b/codex-rs/tui/prompt_for_init_command.md @@ -1,4 +1,5 @@ Generate a file named AGENTS.md that serves as a contributor guide for this repository. +Before writing, check whether AGENTS.md already exists in the current working directory. If it does, do not overwrite or modify it. Your goal is to produce a clear, concise, and well-structured document with descriptive headings and actionable explanations for each section. Follow the outline below, but adapt as needed — add sections if relevant, and omit those that do not apply to this project. diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index c9827a4e71..6fd95e8dd5 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -58,7 +58,6 @@ use crate::bottom_pane::TerminalTitleItem; use crate::bottom_pane::TerminalTitleSetupView; use crate::diff_model::FileChange; use crate::git_action_directives::parse_assistant_markdown; -use crate::legacy_core::DEFAULT_AGENTS_MD_FILENAME; use crate::legacy_core::config::Config; use crate::legacy_core::config::Constrained; use crate::legacy_core::config::ConstraintResult; @@ -257,6 +256,11 @@ fn queued_message_edit_hint_binding( .or_else(|| bindings.first().copied()) } +fn normalize_thread_name(name: &str) -> Option { + let trimmed = name.trim(); + (!trimmed.is_empty()).then(|| trimmed.to_string()) +} + use crate::app_event::AppEvent; use crate::app_event::ExitMode; use crate::app_event::PermissionProfileSelection; diff --git a/codex-rs/tui/src/chatwidget/interaction.rs b/codex-rs/tui/src/chatwidget/interaction.rs index 7a3385034d..49fdd438f8 100644 --- a/codex-rs/tui/src/chatwidget/interaction.rs +++ b/codex-rs/tui/src/chatwidget/interaction.rs @@ -302,7 +302,7 @@ impl ChatWidget { /*initial_text*/ existing_name.unwrap_or_default().to_string(), /*context_label*/ None, Box::new(move |name: String| { - let Some(name) = crate::legacy_core::util::normalize_thread_name(&name) else { + let Some(name) = normalize_thread_name(&name) else { tx.send(AppEvent::InsertHistoryCell(Box::new( history_cell::new_error_event("Thread name cannot be empty.".to_string()), ))); diff --git a/codex-rs/tui/src/chatwidget/slash_dispatch.rs b/codex-rs/tui/src/chatwidget/slash_dispatch.rs index 48caf8f54d..e7776a3007 100644 --- a/codex-rs/tui/src/chatwidget/slash_dispatch.rs +++ b/codex-rs/tui/src/chatwidget/slash_dispatch.rs @@ -213,14 +213,6 @@ impl ChatWidget { .send(AppEvent::OpenDesktopThread { thread_id }); } SlashCommand::Init => { - let init_target = self.config.cwd.join(DEFAULT_AGENTS_MD_FILENAME); - if init_target.exists() { - let message = format!( - "{DEFAULT_AGENTS_MD_FILENAME} already exists here. Skipping /init to avoid overwriting it." - ); - self.add_info_message(message, /*hint*/ None); - return; - } const INIT_PROMPT: &str = include_str!("../../prompt_for_init_command.md"); self.submit_user_message(INIT_PROMPT.to_string().into()); } @@ -662,7 +654,7 @@ impl ChatWidget { } self.session_telemetry .counter("codex.thread.rename", /*inc*/ 1, &[]); - let Some(name) = crate::legacy_core::util::normalize_thread_name(&args) else { + let Some(name) = normalize_thread_name(&args) else { self.add_error_message("Thread name cannot be empty.".to_string()); return; }; diff --git a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs index 9c2802fb60..02faef89f4 100644 --- a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs +++ b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs @@ -565,35 +565,14 @@ async fn ctrl_d_with_modal_open_does_not_quit() { } #[tokio::test] -async fn slash_init_skips_when_project_doc_exists() { - let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; - let tempdir = tempdir().unwrap(); - let existing_path = tempdir.path().join(DEFAULT_AGENTS_MD_FILENAME); - std::fs::write(&existing_path, "existing instructions").unwrap(); - chat.config.cwd = tempdir.path().to_path_buf().abs(); +async fn slash_init_does_not_depend_on_loaded_instruction_sources() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.instruction_source_paths = vec![chat.config.cwd.join("project-instructions.md")]; submit_composer_text(&mut chat, "/init"); - match op_rx.try_recv() { - Err(TryRecvError::Empty) => {} - other => panic!("expected no Codex op to be sent, got {other:?}"), - } - - let cells = drain_insert_history(&mut rx); - assert_eq!(cells.len(), 1, "expected one info message"); - let rendered = lines_to_single_string(&cells[0]); - assert!( - rendered.contains(DEFAULT_AGENTS_MD_FILENAME), - "info message should mention the existing file: {rendered:?}" - ); - assert!( - rendered.contains("Skipping /init"), - "info message should explain why /init was skipped: {rendered:?}" - ); - assert_eq!( - std::fs::read_to_string(existing_path).unwrap(), - "existing instructions" - ); + assert_eq!(chat.input_queue.queued_user_messages.len(), 1); + assert!(drain_insert_history(&mut rx).is_empty()); assert_eq!(recall_latest_after_clearing(&mut chat), "/init"); } diff --git a/codex-rs/tui/src/status/helpers.rs b/codex-rs/tui/src/status/helpers.rs index 27292e1364..f3ceeb435f 100644 --- a/codex-rs/tui/src/status/helpers.rs +++ b/codex-rs/tui/src/status/helpers.rs @@ -185,8 +185,6 @@ fn title_case(s: &str) -> String { #[cfg(test)] mod tests { use super::*; - use crate::legacy_core::DEFAULT_AGENTS_MD_FILENAME; - use crate::legacy_core::LOCAL_AGENTS_MD_FILENAME; use crate::legacy_core::config::ConfigBuilder; use codex_utils_absolute_path::test_support::PathBufExt; use pretty_assertions::assert_eq; @@ -227,7 +225,7 @@ mod tests { async fn compose_agents_summary_includes_global_agents_path() { let codex_home = TempDir::new().expect("temp codex home"); let cwd = TempDir::new().expect("temp cwd"); - let global_agents_path = codex_home.path().join(DEFAULT_AGENTS_MD_FILENAME); + let global_agents_path = codex_home.path().join("global.md"); let config = test_config(&codex_home, &cwd).await; assert_eq!( @@ -240,7 +238,7 @@ mod tests { async fn compose_agents_summary_names_global_agents_override() { let codex_home = TempDir::new().expect("temp codex home"); let cwd = TempDir::new().expect("temp cwd"); - let override_path = codex_home.path().join(LOCAL_AGENTS_MD_FILENAME); + let override_path = codex_home.path().join("override.md"); let config = test_config(&codex_home, &cwd).await; assert_eq!( @@ -253,8 +251,8 @@ mod tests { async fn compose_agents_summary_orders_global_before_project_agents() { let codex_home = TempDir::new().expect("temp codex home"); let cwd = TempDir::new().expect("temp cwd"); - let global_agents_path = codex_home.path().join(DEFAULT_AGENTS_MD_FILENAME); - let project_agents_path = cwd.path().join(DEFAULT_AGENTS_MD_FILENAME); + let global_agents_path = codex_home.path().join("global.md"); + let project_agents_path = cwd.path().join("project.md"); let config = test_config(&codex_home, &cwd).await; let summary = compose_agents_summary( @@ -270,7 +268,7 @@ mod tests { Some(format_directory_display(&global_agents_path, /*max_width*/ None).as_str()) ); let project_path = paths.next().expect("project agents path"); - assert!(project_path.ends_with(DEFAULT_AGENTS_MD_FILENAME)); + assert!(project_path.ends_with("project.md")); assert_eq!(paths.next(), None); } } From 5a0f9134267e0fed2406fce0ea0f14e30894513e Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Tue, 9 Jun 2026 13:35:29 -0700 Subject: [PATCH 14/17] [1/4] Add Python goal routing foundation (#27110) ## Why Goal continuation turns are emitted by the existing runtime as separate physical turns. The Python SDK needs private thread-scoped routing before it can present those notifications as one logical operation, without changing ordinary turn routing or the app-server protocol. ## What - add private goal operation state and thread-scoped notification routing - add internal wrappers for the existing `thread/goal/clear` and `thread/goal/set` RPCs - include existing goal notifications in the SDK notification union - preserve ordinary turn-ID routing unchanged - add focused routing coverage This PR does not expose a public goal API. It is the first PR in the Python goal operations stack. ## Test plan - online CI, including the Python SDK suite - focused typed-notification routing coverage --- sdk/python/src/openai_codex/_goal.py | 156 ++++++++++++++++++ .../src/openai_codex/_message_router.py | 71 ++++++-- sdk/python/src/openai_codex/async_client.py | 39 +++++ sdk/python/src/openai_codex/client.py | 47 ++++++ sdk/python/src/openai_codex/models.py | 4 + sdk/python/tests/test_client_rpc_methods.py | 26 +++ 6 files changed, 328 insertions(+), 15 deletions(-) create mode 100644 sdk/python/src/openai_codex/_goal.py diff --git a/sdk/python/src/openai_codex/_goal.py b/sdk/python/src/openai_codex/_goal.py new file mode 100644 index 0000000000..dc1fe358c2 --- /dev/null +++ b/sdk/python/src/openai_codex/_goal.py @@ -0,0 +1,156 @@ +import queue +import threading +import time +from dataclasses import dataclass, field + +from .generated.v2_all import ( + ThreadGoalClearedNotification, + ThreadGoalStatus, + ThreadGoalUpdatedNotification, + Turn, + TurnCompletedNotification, + TurnStartedNotification, +) +from .models import Notification + + +class _GoalStreamClosed(Exception): + """Wake a notification reader after its logical stream closes.""" + + +def _terminal_goal_status(status: ThreadGoalStatus | None) -> bool: + return status in { + ThreadGoalStatus.paused, + ThreadGoalStatus.blocked, + ThreadGoalStatus.usage_limited, + ThreadGoalStatus.budget_limited, + ThreadGoalStatus.complete, + } + + +@dataclass(slots=True) +class _GoalOperationState: + """Private state for one goal operation exposed as a logical turn.""" + + thread_id: str + logical_turn_id: str | None = None + current_turn_id: str | None = None + status: ThreadGoalStatus | None = None + started_turn: Turn | None = None + completed_turn: Turn | None = None + interrupted: bool = False + interrupt_requested: bool = False + cleared: bool = False + _condition: threading.Condition = field(default_factory=threading.Condition) + _notifications: queue.Queue[Notification | BaseException] = field(default_factory=queue.Queue) + _failure: BaseException | None = None + _finished: bool = False + + def observe(self, notification: Notification) -> None: + payload = notification.payload + with self._condition: + if isinstance(payload, TurnStartedNotification): + if self.logical_turn_id is None: + self.logical_turn_id = payload.turn.id + self.current_turn_id = payload.turn.id + if self.started_turn is None: + self.started_turn = payload.turn + elif isinstance(payload, TurnCompletedNotification): + self.completed_turn = payload.turn + if self.current_turn_id == payload.turn.id: + self.current_turn_id = None + elif isinstance(payload, ThreadGoalUpdatedNotification): + self.status = payload.goal.status + if self.status == ThreadGoalStatus.active: + self.cleared = False + elif isinstance(payload, ThreadGoalClearedNotification): + self.cleared = True + if ( + self.current_turn_id is None + and self.completed_turn is not None + and (self.cleared or _terminal_goal_status(self.status)) + ): + self._finished = True + self._condition.notify_all() + self._notifications.put(notification) + + def wait_for_start(self, timeout: float) -> str | None: + """Wait for the runtime-generated first turn without consuming its event.""" + deadline = time.monotonic() + timeout + with self._condition: + while self.started_turn is None or self.logical_turn_id is None: + if self._failure is not None: + raise self._failure + remaining = deadline - time.monotonic() + if remaining <= 0: + return None + self._condition.wait(remaining) + return self.logical_turn_id + + def fail(self, exc: BaseException) -> None: + with self._condition: + self._failure = exc + self._condition.notify_all() + self._notifications.put(exc) + + def next_notification(self) -> Notification: + item = self._notifications.get() + if isinstance(item, BaseException): + raise item + return item + + def finish(self) -> None: + """Mark the logical operation inactive and wake waiting controls.""" + with self._condition: + self._finished = True + self.current_turn_id = None + self._condition.notify_all() + + def is_finished(self) -> bool: + with self._condition: + return self._finished + + def begin_interrupt(self) -> bool: + with self._condition: + if self._finished: + return False + self.interrupt_requested = True + return True + + def confirm_interrupt(self) -> None: + with self._condition: + self.interrupted = True + self.interrupt_requested = False + + def cancel_interrupt(self) -> None: + with self._condition: + self.interrupt_requested = False + + def explicit_interrupt(self, status: ThreadGoalStatus | None) -> bool: + with self._condition: + return self.interrupted or ( + self.interrupt_requested and status == ThreadGoalStatus.paused + ) + + def active_turn(self, *, after: str | None = None) -> str | None: + """Wait for the current turn, or return None once the goal has ended.""" + with self._condition: + while True: + if self._failure is not None: + raise self._failure + if self._finished: + return None + if self.current_turn_id is not None and self.current_turn_id != after: + return self.current_turn_id + if self.cleared or _terminal_goal_status(self.status): + return None + self._condition.wait() + + def current_turn(self) -> str | None: + """Return the current physical turn without waiting for rollover.""" + with self._condition: + return self.current_turn_id + + def wake_notification_reader(self) -> None: + """Release a reader blocked after its stream has been closed.""" + self._notifications.put(_GoalStreamClosed()) diff --git a/sdk/python/src/openai_codex/_message_router.py b/sdk/python/src/openai_codex/_message_router.py index b500f16c56..e71544a0ef 100644 --- a/sdk/python/src/openai_codex/_message_router.py +++ b/sdk/python/src/openai_codex/_message_router.py @@ -4,6 +4,7 @@ import queue import threading from collections import deque +from ._goal import _GoalOperationState from .errors import CodexError, map_jsonrpc_error from .generated.notification_registry import notification_turn_id from .generated.v2_all import AccountLoginCompletedNotification @@ -30,6 +31,7 @@ class MessageRouter: self._pending_login_notifications: dict[str, deque[Notification]] = {} self._turn_notifications: dict[str, queue.Queue[NotificationQueueItem]] = {} self._pending_turn_notifications: dict[str, deque[Notification]] = {} + self._goal_operations: dict[str, _GoalOperationState] = {} self._global_notifications: queue.Queue[NotificationQueueItem] = queue.Queue() def create_response_waiter(self, request_id: str) -> queue.Queue[ResponseQueueItem]: @@ -116,6 +118,21 @@ class MessageRouter: raise item return item + def register_goal(self, thread_id: str) -> _GoalOperationState: + """Register one thread-scoped logical goal operation before it starts.""" + state = _GoalOperationState(thread_id=thread_id) + 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 + return state + + def unregister_goal(self, state: _GoalOperationState) -> None: + """Stop routing notifications to a completed logical goal operation.""" + with self._lock: + if self._goal_operations.get(state.thread_id) is state: + self._goal_operations.pop(state.thread_id) + def route_response(self, msg: dict[str, JsonValue]) -> None: """Deliver a JSON-RPC response or error to its request waiter.""" @@ -157,6 +174,17 @@ class MessageRouter: return turn_id = self._notification_turn_id(notification) + thread_id = self._notification_thread_id(notification) + if thread_id is not None: + with self._lock: + goal_state = self._goal_operations.get(thread_id) + 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 turn_id is None: self._global_notifications.put(notification) return @@ -182,6 +210,8 @@ class MessageRouter: self._pending_login_notifications.clear() turn_queues = list(self._turn_notifications.values()) self._pending_turn_notifications.clear() + goal_operations = list(self._goal_operations.values()) + self._goal_operations.clear() # Put the same transport failure into every queue so no SDK call blocks # forever waiting for a response that cannot arrive. for waiter in response_waiters: @@ -190,8 +220,34 @@ class MessageRouter: login_queue.put(exc) for turn_queue in turn_queues: turn_queue.put(exc) + for goal_operation in goal_operations: + goal_operation.fail(exc) self._global_notifications.put(exc) + def _notification_turn_id(self, notification: Notification) -> str | None: + """Extract routing ids from generated metadata or raw unknown payloads.""" + payload = notification.payload + if isinstance(payload, UnknownNotification): + raw_turn_id = payload.params.get("turnId") + if isinstance(raw_turn_id, str): + return raw_turn_id + raw_turn = payload.params.get("turn") + if isinstance(raw_turn, dict): + raw_nested_turn_id = raw_turn.get("id") + if isinstance(raw_nested_turn_id, str): + return raw_nested_turn_id + return None + return notification_turn_id(payload) + + def _notification_thread_id(self, notification: Notification) -> str | None: + """Extract thread ids from typed payloads or raw unknown payloads.""" + payload = notification.payload + if isinstance(payload, UnknownNotification): + raw_thread_id = payload.params.get("threadId") + return raw_thread_id if isinstance(raw_thread_id, str) else None + thread_id = getattr(payload, "thread_id", None) + return thread_id if isinstance(thread_id, str) else None + def _notification_login_id(self, notification: Notification) -> str | None: """Extract the login attempt id from completion notifications.""" if notification.method != "account/login/completed": @@ -205,18 +261,3 @@ class MessageRouter: if isinstance(raw_login_id, str): return raw_login_id return None - - def _notification_turn_id(self, notification: Notification) -> str | None: - """Extract routing ids from known generated payloads or raw unknown payloads.""" - payload = notification.payload - if isinstance(payload, UnknownNotification): - raw_turn_id = payload.params.get("turnId") - if isinstance(raw_turn_id, str): - return raw_turn_id - raw_turn = payload.params.get("turn") - if isinstance(raw_turn, dict): - raw_nested_turn_id = raw_turn.get("id") - if isinstance(raw_nested_turn_id, str): - return raw_nested_turn_id - return None - return notification_turn_id(payload) diff --git a/sdk/python/src/openai_codex/async_client.py b/sdk/python/src/openai_codex/async_client.py index f0b12b99ed..ff12f16884 100644 --- a/sdk/python/src/openai_codex/async_client.py +++ b/sdk/python/src/openai_codex/async_client.py @@ -6,6 +6,7 @@ from typing import AsyncIterator, Callable, ParamSpec, TypeVar from pydantic import BaseModel +from ._goal import _GoalOperationState from .client import CodexClient, CodexConfig from .generated.v2_all import ( AccountLoginCompletedNotification, @@ -21,6 +22,9 @@ from .generated.v2_all import ( ThreadCompactStartResponse, ThreadForkParams as V2ThreadForkParams, ThreadForkResponse, + ThreadGoalClearResponse, + ThreadGoalSetResponse, + ThreadGoalStatus, ThreadListParams as V2ThreadListParams, ThreadListResponse, ThreadReadResponse, @@ -107,6 +111,14 @@ class AsyncCodexClient: """Unregister a turn notification queue on the wrapped sync client.""" self._sync.unregister_turn_notifications(turn_id) + def register_goal_operation(self, thread_id: str) -> _GoalOperationState: + """Register a logical goal route on the wrapped sync client.""" + return self._sync.register_goal_operation(thread_id) + + def unregister_goal_operation(self, state: _GoalOperationState) -> None: + """Release one logical goal route.""" + self._sync.unregister_goal_operation(state) + async def request( self, method: str, @@ -192,6 +204,29 @@ class AsyncCodexClient: """Start thread compaction using the wrapped sync client.""" return await self._call_sync(self._sync.thread_compact, thread_id) + async def thread_goal_clear(self, thread_id: str) -> ThreadGoalClearResponse: + """Clear the persisted goal through the wrapped sync client.""" + return await self._call_sync(self._sync.thread_goal_clear, thread_id) + + async def thread_goal_set( + self, + thread_id: str, + *, + objective: str | None = None, + status: ThreadGoalStatus | None = None, + ) -> ThreadGoalSetResponse: + """Create or update a persisted goal through the wrapped sync client.""" + return await self._call_sync( + self._sync.thread_goal_set, + thread_id, + objective=objective, + status=status, + ) + + async def pause_goal(self, thread_id: str) -> ThreadGoalSetResponse: + """Pause the active goal through the wrapped sync client.""" + return await self._call_sync(self._sync.pause_goal, thread_id) + async def turn_start( self, thread_id: str, @@ -256,6 +291,10 @@ class AsyncCodexClient: """Wait for the next notification routed to one turn.""" return await self._call_sync(self._sync.next_turn_notification, turn_id) + async def next_goal_notification(self, state: _GoalOperationState) -> Notification: + """Wait for the next notification in a logical goal turn.""" + return await self._call_sync(self._sync.next_goal_notification, state) + async def wait_for_login_completed( self, login_id: str, diff --git a/sdk/python/src/openai_codex/client.py b/sdk/python/src/openai_codex/client.py index 9951d6d1f3..951f4f95f0 100644 --- a/sdk/python/src/openai_codex/client.py +++ b/sdk/python/src/openai_codex/client.py @@ -10,6 +10,7 @@ from typing import Callable, Iterator, TypeVar 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 @@ -30,6 +31,9 @@ from .generated.v2_all import ( ThreadCompactStartResponse, ThreadForkParams as V2ThreadForkParams, ThreadForkResponse, + ThreadGoalClearResponse, + ThreadGoalSetResponse, + ThreadGoalStatus, ThreadListParams as V2ThreadListParams, ThreadListResponse, ThreadReadResponse, @@ -352,6 +356,18 @@ class CodexClient: """Return the next routed notification for the requested turn id.""" return self._router.next_turn_notification(turn_id) + def register_goal_operation(self, thread_id: str) -> _GoalOperationState: + """Register a private thread-scoped route for a logical goal turn.""" + return self._router.register_goal(thread_id) + + def unregister_goal_operation(self, state: _GoalOperationState) -> None: + """Release routing state for one logical goal turn.""" + self._router.unregister_goal(state) + + def next_goal_notification(self, state: _GoalOperationState) -> Notification: + """Wait for the next notification in a logical goal turn.""" + return state.next_notification() + def account_login_start( self, params: V2LoginAccountParams | JsonObject, @@ -452,6 +468,37 @@ class CodexClient: response_model=ThreadCompactStartResponse, ) + def thread_goal_clear(self, thread_id: str) -> ThreadGoalClearResponse: + """Clear the persisted goal for a thread before replacing it.""" + return self.request( + "thread/goal/clear", + {"threadId": thread_id}, + response_model=ThreadGoalClearResponse, + ) + + def thread_goal_set( + self, + thread_id: str, + *, + objective: str | None = None, + status: ThreadGoalStatus | None = None, + ) -> ThreadGoalSetResponse: + """Create or update the persisted goal for a thread.""" + payload: JsonObject = {"threadId": thread_id} + if objective is not None: + payload["objective"] = objective + if status is not None: + payload["status"] = status.value + return self.request( + "thread/goal/set", + payload, + response_model=ThreadGoalSetResponse, + ) + + def pause_goal(self, thread_id: str) -> ThreadGoalSetResponse: + """Pause the active goal used by a logical goal turn.""" + return self.thread_goal_set(thread_id, status=ThreadGoalStatus.paused) + def turn_start( self, thread_id: str, diff --git a/sdk/python/src/openai_codex/models.py b/sdk/python/src/openai_codex/models.py index 70c61d44cd..d9d15dc684 100644 --- a/sdk/python/src/openai_codex/models.py +++ b/sdk/python/src/openai_codex/models.py @@ -27,6 +27,8 @@ from .generated.v2_all import ( ReasoningSummaryTextDeltaNotification, ReasoningTextDeltaNotification, TerminalInteractionNotification, + ThreadGoalClearedNotification, + ThreadGoalUpdatedNotification, ThreadNameUpdatedNotification, ThreadStartedNotification, ThreadTokenUsageUpdatedNotification, @@ -70,6 +72,8 @@ NotificationPayload: TypeAlias = ( | ReasoningTextDeltaNotification | TerminalInteractionNotification | ThreadNameUpdatedNotification + | ThreadGoalClearedNotification + | ThreadGoalUpdatedNotification | ThreadStartedNotification | ThreadTokenUsageUpdatedNotification | TurnCompletedNotification diff --git a/sdk/python/tests/test_client_rpc_methods.py b/sdk/python/tests/test_client_rpc_methods.py index 73fdd452ff..15be24f3b4 100644 --- a/sdk/python/tests/test_client_rpc_methods.py +++ b/sdk/python/tests/test_client_rpc_methods.py @@ -185,6 +185,32 @@ def test_turn_notification_router_demuxes_registered_turns() -> None: ] +def test_goal_notification_router_routes_by_thread_id() -> None: + """A goal operation should receive turn notifications across physical turn ids.""" + client = CodexClient() + state = client.register_goal_operation("thread-1") + + client._router.route_notification( + client._coerce_notification( + "item/agentMessage/delta", + { + "delta": "continued", + "itemId": "item-1", + "threadId": "thread-1", + "turnId": "turn-2", + }, + ) + ) + + event = client.next_goal_notification(state) + + assert isinstance(event.payload, AgentMessageDeltaNotification) + assert (event.method, event.payload.delta) == ( + "item/agentMessage/delta", + "continued", + ) + + def test_client_reader_routes_interleaved_turn_notifications_by_turn_id() -> None: """Reader-loop routing should preserve order within each interleaved turn stream.""" client = CodexClient() From 4ec3b8eeea38e9f693568e4f51bb607cc6cd9717 Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 9 Jun 2026 22:44:16 +0200 Subject: [PATCH 15/17] Route hosted Apps MCP through extensions (#27191) ## Stack - Base: #27184 - This PR is the second vertical and should be reviewed against `jif/external-plugins-1`, not `main`. ## Why CCA is moving toward a split runtime where the orchestrator may have no filesystem or executor, but it still needs to activate remotely hosted plugin components. HTTP MCP servers are the simplest complete example: they need configuration and host authentication, but they do not need an executor process. The Apps MCP endpoint is currently synthesized by a special-purpose loader inside the MCP runtime. That works locally, but it leaves hosted MCP activation outside the extension model being established in #27184. It also makes the Apps path a poor foundation for plugins whose skills, MCP servers, connectors, and hooks may come from different sources or execute in different places. This PR moves that one behavior behind an extension-owned contribution while preserving the existing local fallback. It deliberately does not introduce a generic plugin activation framework. ## What changed ### MCP extension contribution `codex-extension-api` gains an ordered `McpServerContributor` contract. A contributor returns typed `Set` or `Remove` overlays for MCP server configuration; later contributors win for the names they own. The contract stays at the existing MCP configuration boundary. Extensions do not create a second connection manager or transport abstraction. ### Hosted Apps MCP extension A new `codex-mcp-extension` contributes the reserved `codex_apps` server from the existing Apps feature, ChatGPT base URL, path override, and product SKU configuration. When `apps_mcp_path_override` is enabled for `https://chatgpt.com`, the resulting streamable HTTP endpoint is `https://chatgpt.com/backend-api/ps/mcp`. The existing ChatGPT-auth gate remains authoritative, so this server can run in an orchestrator-only process without being exposed for API-key sessions. ### One resolved runtime view `McpManager` now distinguishes three views: - **configured:** config- and plugin-backed servers before extension overlays; - **runtime:** configured servers plus host-installed extension contributions; - **effective:** runtime servers after auth gating and compatibility built-ins. App-server installs the hosted MCP extension and uses the runtime view for thread startup, refresh, status, threadless resource reads, connector discovery, and MCP OAuth lookup. This keeps `mcpServer/oauth/login` consistent with the servers exposed by the other MCP APIs. The hosted Apps server itself continues to use existing ChatGPT host authentication rather than MCP OAuth. ## Compatibility Hosts that do not install the MCP extension retain the existing Apps MCP synthesis path. This preserves current local-only, CLI, and standalone-host behavior while app-server exercises the extension path. Disabling Apps removes the reserved `codex_apps` entry, and losing ChatGPT auth removes it from the effective runtime view. Executor availability is not consulted for this HTTP transport. ## Follow-ups The next vertical will resolve a manifest-declared stdio MCP server from an executor-selected plugin root and execute it in the environment that owns that root. Later verticals can add backend-owned skills, connector metadata, hooks, durable selection semantics, and incremental local convergence without changing the component-specific runtime boundaries introduced here. ## Verification Focused coverage was added for: - contributing the hosted Apps MCP at `/backend-api/ps/mcp` without an executor; - requiring ChatGPT auth in the effective runtime view; - removing a reserved configured Apps server when the Apps feature is disabled. `cargo check -p codex-app-server -p codex-mcp-extension -p codex-extension-api -p codex-mcp` passed. Tests and Clippy were not run locally under the current development instruction; CI provides the full validation pass. --- codex-rs/Cargo.lock | 19 ++++ codex-rs/Cargo.toml | 2 + codex-rs/app-server/Cargo.toml | 1 + codex-rs/app-server/src/extensions.rs | 1 + codex-rs/app-server/src/mcp_refresh.rs | 5 +- codex-rs/app-server/src/request_processors.rs | 1 + .../src/request_processors/apps_processor.rs | 34 +++++-- .../src/request_processors/mcp_processor.rs | 22 +++-- .../src/request_processors/plugins.rs | 11 ++- codex-rs/chatgpt/src/connectors.rs | 1 + codex-rs/codex-mcp/src/lib.rs | 1 + codex-rs/codex-mcp/src/mcp/mod.rs | 36 ++++--- codex-rs/codex-mcp/src/mcp/mod_tests.rs | 11 +-- codex-rs/core/src/config/mod.rs | 1 + codex-rs/core/src/connectors.rs | 27 ++++-- codex-rs/core/src/mcp.rs | 70 +++++++++++++- codex-rs/core/src/mcp_skill_dependencies.rs | 22 +++-- codex-rs/core/src/session/mcp.rs | 6 +- codex-rs/core/src/thread_manager.rs | 5 +- codex-rs/ext/extension-api/Cargo.toml | 1 + .../ext/extension-api/src/contributors.rs | 14 +++ .../ext/extension-api/src/contributors/mcp.rs | 22 +++++ codex-rs/ext/extension-api/src/lib.rs | 2 + codex-rs/ext/extension-api/src/registry.rs | 15 +++ codex-rs/ext/mcp/BUILD.bazel | 6 ++ codex-rs/ext/mcp/Cargo.toml | 29 ++++++ codex-rs/ext/mcp/src/lib.rs | 31 +++++++ codex-rs/ext/mcp/tests/hosted_apps_mcp.rs | 93 +++++++++++++++++++ 28 files changed, 424 insertions(+), 65 deletions(-) create mode 100644 codex-rs/ext/extension-api/src/contributors/mcp.rs create mode 100644 codex-rs/ext/mcp/BUILD.bazel create mode 100644 codex-rs/ext/mcp/Cargo.toml create mode 100644 codex-rs/ext/mcp/src/lib.rs create mode 100644 codex-rs/ext/mcp/tests/hosted_apps_mcp.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8783f0e16c..cbc237a84c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1982,6 +1982,7 @@ dependencies = [ "codex-image-generation-extension", "codex-login", "codex-mcp", + "codex-mcp-extension", "codex-memories-extension", "codex-memories-write", "codex-model-provider", @@ -2920,6 +2921,7 @@ name = "codex-extension-api" version = "0.0.0" dependencies = [ "async-trait", + "codex-config", "codex-context-fragments", "codex-protocol", "codex-tools", @@ -3247,6 +3249,23 @@ dependencies = [ "url", ] +[[package]] +name = "codex-mcp-extension" +version = "0.0.0" +dependencies = [ + "async-trait", + "codex-config", + "codex-core", + "codex-core-plugins", + "codex-extension-api", + "codex-features", + "codex-login", + "codex-mcp", + "pretty_assertions", + "tempfile", + "tokio", +] + [[package]] name = "codex-mcp-server" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 955d8987f7..b87b0cf39f 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -49,6 +49,7 @@ members = [ "ext/guardian", "ext/image-generation", "ext/memories", + "ext/mcp", "ext/skills", "ext/web-search", "external-agent-migration", @@ -191,6 +192,7 @@ codex-web-search-extension = { path = "ext/web-search" } codex-memories-read = { path = "memories/read" } codex-memories-write = { path = "memories/write" } codex-mcp = { path = "codex-mcp" } +codex-mcp-extension = { path = "ext/mcp" } codex-mcp-server = { path = "mcp-server" } codex-model-provider-info = { path = "model-provider-info" } codex-models-manager = { path = "models-manager" } diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index d6f71ab82f..55ea78e617 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -61,6 +61,7 @@ codex-memories-extension = { workspace = true } codex-web-search-extension = { workspace = true } codex-memories-write = { workspace = true } codex-mcp = { workspace = true } +codex-mcp-extension = { workspace = true } codex-model-provider = { workspace = true } codex-models-manager = { workspace = true } codex-protocol = { workspace = true } diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index b18331064f..a0a44fd80d 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -50,6 +50,7 @@ where } codex_guardian::install(&mut builder, guardian_agent_spawner); codex_memories_extension::install(&mut builder, codex_otel::global()); + codex_mcp_extension::install(&mut builder); codex_web_search_extension::install(&mut builder, auth_manager.clone()); codex_image_generation_extension::install(&mut builder, auth_manager); codex_skills_extension::install_with_providers( diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 821af862d5..2cd1027754 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -67,10 +67,7 @@ async fn build_refresh_config( let config = config_manager .load_latest_config_for_thread(thread_config.as_ref()) .await?; - let mcp_servers = thread_manager - .mcp_manager() - .configured_servers(&config) - .await; + let mcp_servers = thread_manager.mcp_manager().runtime_servers(&config).await; Ok(McpServerRefreshConfig { mcp_servers: serde_json::to_value(mcp_servers).map_err(io::Error::other)?, mcp_oauth_credentials_store_mode: serde_json::to_value( diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index d4e9395627..00f17937ba 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -280,6 +280,7 @@ use codex_config::types::McpServerTransportConfig; use codex_core::CodexThread; use codex_core::CodexThreadSettingsOverrides; use codex_core::ForkSnapshot; +use codex_core::McpManager; use codex_core::NewThread; #[cfg(test)] use codex_core::SessionMeta; diff --git a/codex-rs/app-server/src/request_processors/apps_processor.rs b/codex-rs/app-server/src/request_processors/apps_processor.rs index 4a665d8ef3..95cf3f452a 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor.rs @@ -88,11 +88,19 @@ impl AppsRequestProcessor { let request = request_id.clone(); let outgoing = Arc::clone(&self.outgoing); let environment_manager = self.thread_manager.environment_manager(); + let mcp_manager = self.thread_manager.mcp_manager(); let shutdown_token = self.shutdown_token.child_token(); tokio::spawn(async move { tokio::select! { _ = shutdown_token.cancelled() => {} - _ = Self::apps_list_task(outgoing, request, params, config, environment_manager) => {} + _ = Self::apps_list_task( + outgoing, + request, + params, + config, + environment_manager, + mcp_manager, + ) => {} } }); Ok(None) @@ -108,11 +116,15 @@ impl AppsRequestProcessor { params: AppsListParams, config: Config, environment_manager: Arc, + mcp_manager: Arc, ) { let retry_params = params.clone(); let retry_config = config.clone(); let retry_environment_manager = Arc::clone(&environment_manager); - let result = Self::apps_list_response(&outgoing, params, config, environment_manager).await; + let retry_mcp_manager = Arc::clone(&mcp_manager); + let result = + Self::apps_list_response(&outgoing, params, config, environment_manager, mcp_manager) + .await; let should_retry = result .as_ref() .is_ok_and(|(_, codex_apps_ready)| !codex_apps_ready); @@ -128,6 +140,7 @@ impl AppsRequestProcessor { retry_params, retry_config, retry_environment_manager, + retry_mcp_manager, ) .await { @@ -141,6 +154,7 @@ impl AppsRequestProcessor { params: AppsListParams, config: Config, environment_manager: Arc, + mcp_manager: Arc, ) -> Result<(AppsListResponse, bool), JSONRPCErrorError> { let AppsListParams { cursor, @@ -167,14 +181,14 @@ impl AppsRequestProcessor { let accessible_config = config.clone(); let accessible_tx = tx.clone(); tokio::spawn(async move { - let result = - connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager( - &accessible_config, - force_refetch, - Arc::clone(&environment_manager), - ) - .await - .map_err(|err| format!("failed to load accessible apps: {err}")); + let result = connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( + &accessible_config, + force_refetch, + Arc::clone(&environment_manager), + mcp_manager, + ) + .await + .map_err(|err| format!("failed to load accessible apps: {err}")); let _ = accessible_tx.send(AppListLoadResult::Accessible(result)); }); diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index ae62e2e785..facd4facc6 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -120,12 +120,16 @@ impl McpRequestProcessor { timeout_secs, } = params; - let configured_servers = self + let auth = self.auth_manager.auth().await; + let effective_servers = self .thread_manager .mcp_manager() - .configured_servers(&config) + .effective_servers(&config, auth.as_ref()) .await; - let Some(server) = configured_servers.get(&name) else { + let Some(server) = effective_servers + .get(&name) + .and_then(codex_mcp::EffectiveMcpServer::configured_config) + else { return Err(invalid_request(format!( "No MCP server named '{name}' found." ))); @@ -210,8 +214,10 @@ impl McpRequestProcessor { } None => self.load_latest_config(/*fallback_cwd*/ None).await?, }; - let mcp_config = config - .to_mcp_config(self.thread_manager.plugins_manager().as_ref()) + let mcp_config = self + .thread_manager + .mcp_manager() + .runtime_config(&config) .await; let auth = self.auth_manager.auth().await; let environment_manager = self.thread_manager.environment_manager(); @@ -361,8 +367,10 @@ impl McpRequestProcessor { } let config = self.load_latest_config(/*fallback_cwd*/ None).await?; - let mcp_config = config - .to_mcp_config(self.thread_manager.plugins_manager().as_ref()) + let mcp_config = self + .thread_manager + .mcp_manager() + .runtime_config(&config) .await; let auth = self.auth_manager.auth().await; let environment_manager = self.thread_manager.environment_manager(); diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index bc2383c21f..54a5efbe7a 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -1036,6 +1036,7 @@ impl PluginRequestProcessor { &config, &outcome.plugin.apps, Arc::clone(&environment_manager), + self.thread_manager.mcp_manager(), ) .await; let visible_skills = outcome @@ -1118,6 +1119,7 @@ impl PluginRequestProcessor { &config, &plugin_apps, Arc::clone(&environment_manager), + self.thread_manager.mcp_manager(), ) .await; remote_plugin_detail_to_info(remote_detail, app_summaries) @@ -1611,10 +1613,11 @@ impl PluginRequestProcessor { let environment_manager = self.thread_manager.environment_manager(); let (all_connectors_result, accessible_connectors_result) = tokio::join!( connectors::list_all_connectors_with_options(config, /*force_refetch*/ false), - connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager( + connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( config, /*force_refetch*/ true, - Arc::clone(&environment_manager) + Arc::clone(&environment_manager), + self.thread_manager.mcp_manager(), ), ); @@ -1881,6 +1884,7 @@ 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(); @@ -1900,10 +1904,11 @@ 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_environment_manager( + match connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( config, /*force_refetch*/ false, environment_manager, + mcp_manager, ) .await { diff --git a/codex-rs/chatgpt/src/connectors.rs b/codex-rs/chatgpt/src/connectors.rs index 63b34e7355..4bbf40813e 100644 --- a/codex-rs/chatgpt/src/connectors.rs +++ b/codex-rs/chatgpt/src/connectors.rs @@ -14,6 +14,7 @@ use codex_connectors::merge::merge_plugin_connectors; use codex_core::config::Config; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager; +pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options_and_status; pub use codex_core::connectors::list_cached_accessible_connectors_from_mcp_tools; diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 83419175d4..a56365277a 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -24,6 +24,7 @@ pub use auth_elicitation::build_auth_elicitation_plan; pub use auth_elicitation::connector_auth_failure_from_tool_result; pub use codex_apps::CodexAppsToolsCacheKey; pub use codex_apps::codex_apps_tools_cache_key; +pub use mcp::codex_apps_mcp_server_config; pub use mcp::configured_mcp_servers; pub use mcp::effective_mcp_servers; pub use mcp::effective_mcp_servers_from_configured; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 9633655846..a74f95bae3 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -131,6 +131,11 @@ pub struct McpConfig { /// ChatGPT auth is checked separately at runtime before the host-owned apps /// MCP server is added. 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, @@ -218,10 +223,20 @@ pub fn with_codex_apps_mcp( 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)), + 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); @@ -381,13 +396,6 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( snapshot } -pub(crate) fn codex_apps_mcp_url(config: &McpConfig) -> String { - codex_apps_mcp_url_for_base_url( - &config.chatgpt_base_url, - config.apps_mcp_path_override.as_deref(), - ) -} - /// The Responses API requires tool names to match `^[a-zA-Z0-9_-]+$`. /// MCP server/tool names are user-controlled, so sanitize the fully-qualified /// name we expose to the model by replacing any disallowed character with `_`. @@ -443,10 +451,14 @@ fn codex_apps_mcp_url_for_base_url(base_url: &str, apps_mcp_path_override: Optio format!("{base_url}/{path}") } -fn codex_apps_mcp_server_config(config: &McpConfig) -> McpServerConfig { - let url = codex_apps_mcp_url(config); - let http_headers = config.apps_mcp_product_sku.as_ref().map(|product_sku| { - HashMap::from([("X-OpenAI-Product-Sku".to_string(), product_sku.clone())]) +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); + let http_headers = apps_mcp_product_sku.map(|product_sku| { + HashMap::from([("X-OpenAI-Product-Sku".to_string(), product_sku.to_string())]) }); McpServerConfig { diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index b29f7a9e58..fb1af362ca 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -27,6 +27,7 @@ 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(), @@ -206,16 +207,6 @@ fn codex_apps_mcp_url_for_base_url_keeps_existing_paths() { ); } -#[test] -fn codex_apps_mcp_url_uses_legacy_codex_apps_path() { - let config = test_mcp_config(PathBuf::from("/tmp")); - - assert_eq!( - codex_apps_mcp_url(&config), - "https://chatgpt.com/backend-api/wham/apps" - ); -} - #[test] fn codex_apps_server_config_uses_legacy_codex_apps_path() { let mut config = test_mcp_config(PathBuf::from("/tmp")); diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 93e203053f..5c2a93bc45 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1428,6 +1428,7 @@ 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 { diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 7a8344b3d7..623cb44216 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::collections::HashSet; use std::sync::Arc; use std::sync::LazyLock; @@ -41,8 +40,8 @@ use codex_mcp::ToolInfo; use codex_mcp::ToolPluginProvenance; use codex_mcp::codex_apps_tools_cache_key; use codex_mcp::compute_auth_statuses; +use codex_mcp::effective_mcp_servers; use codex_mcp::host_owned_codex_apps_enabled; -use codex_mcp::with_codex_apps_mcp; const CONNECTORS_READY_TIMEOUT_ON_EMPTY_TOOLS: Duration = Duration::from_secs(30); @@ -220,6 +219,23 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( config: &Config, force_refetch: bool, environment_manager: Arc, +) -> anyhow::Result { + let plugins_manager = Arc::new(PluginsManager::new(config.codex_home.to_path_buf())); + let mcp_manager = Arc::new(McpManager::new(plugins_manager)); + list_accessible_connectors_from_mcp_tools_with_mcp_manager( + config, + force_refetch, + environment_manager, + mcp_manager, + ) + .await +} + +pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( + config: &Config, + force_refetch: bool, + environment_manager: Arc, + mcp_manager: Arc, ) -> anyhow::Result { let auth_manager = AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; @@ -234,8 +250,6 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( }); } let cache_key = accessible_connectors_cache_key(config, auth.as_ref()); - let plugins_manager = Arc::new(PluginsManager::new(config.codex_home.to_path_buf())); - let mcp_manager = McpManager::new(Arc::clone(&plugins_manager)); let tool_plugin_provenance = mcp_manager.tool_plugin_provenance(config).await; if !force_refetch && let Some(cached_connectors) = read_cached_accessible_connectors(&cache_key) { @@ -250,8 +264,9 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( }); } - let mcp_config = config.to_mcp_config(plugins_manager.as_ref()).await; - let mcp_servers = with_codex_apps_mcp(HashMap::new(), auth.as_ref(), &mcp_config); + let mcp_config = mcp_manager.runtime_config(config).await; + let mut mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); + mcp_servers.retain(|name, _| name == CODEX_APPS_MCP_SERVER_NAME); let host_owned_codex_apps_enabled = host_owned_codex_apps_enabled(&mcp_config, auth.as_ref()); if mcp_servers.is_empty() { return Ok(AccessibleConnectorsStatus { diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 91ebce4ad8..1b693d70fd 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -4,8 +4,12 @@ use std::sync::Arc; use crate::config::Config; use codex_config::McpServerConfig; use codex_core_plugins::PluginsManager; +use codex_extension_api::ExtensionRegistry; +use codex_extension_api::McpServerContribution; use codex_login::CodexAuth; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::EffectiveMcpServer; +use codex_mcp::McpConfig; use codex_mcp::ToolPluginProvenance; use codex_mcp::configured_mcp_servers; use codex_mcp::effective_mcp_servers; @@ -14,29 +18,91 @@ use codex_mcp::tool_plugin_provenance as collect_tool_plugin_provenance; #[derive(Clone)] pub struct McpManager { plugins_manager: Arc, + extensions: Arc>, } impl McpManager { pub fn new(plugins_manager: Arc) -> Self { - Self { plugins_manager } + Self { + plugins_manager, + extensions: codex_extension_api::empty_extension_registry(), + } } + /// Creates a manager that resolves host-installed MCP contributions. + pub fn new_with_extensions( + plugins_manager: Arc, + extensions: Arc>, + ) -> Self { + Self { + plugins_manager, + extensions, + } + } + + /// Returns the MCP config after applying 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 + .iter() + .any(|contribution| contribution.name() == CODEX_APPS_MCP_SERVER_NAME) + { + mcp_config.legacy_apps_mcp_loader_enabled = false; + } + Self::apply_to_configured_servers(&contributions, &mut mcp_config.configured_mcp_servers); + mcp_config + } + + /// Returns config- and plugin-backed servers without runtime contributions. pub async fn configured_servers(&self, config: &Config) -> HashMap { let mcp_config = config.to_mcp_config(self.plugins_manager.as_ref()).await; configured_mcp_servers(&mcp_config) } + /// Returns configured and host-contributed servers before auth gating. + pub async fn runtime_servers(&self, config: &Config) -> HashMap { + let mcp_config = self.runtime_config(config).await; + configured_mcp_servers(&mcp_config) + } + + /// Returns runtime servers after auth gating and compatibility built-ins. pub async fn effective_servers( &self, config: &Config, auth: Option<&CodexAuth>, ) -> HashMap { - let mcp_config = config.to_mcp_config(self.plugins_manager.as_ref()).await; + let mcp_config = self.runtime_config(config).await; effective_mcp_servers(&mcp_config, auth) } + /// Returns provenance for plugin-owned servers in the configured view. pub async fn tool_plugin_provenance(&self, config: &Config) -> ToolPluginProvenance { let mcp_config = config.to_mcp_config(self.plugins_manager.as_ref()).await; collect_tool_plugin_provenance(&mcp_config) } + + async fn contributions(&self, config: &Config) -> Vec { + let mut contributions = Vec::new(); + for contributor in self.extensions.mcp_server_contributors() { + contributions.extend(contributor.contribute(config).await); + } + contributions + } + + fn apply_to_configured_servers( + contributions: &[McpServerContribution], + servers: &mut HashMap, + ) { + for contribution in contributions { + match contribution { + McpServerContribution::Set { name, config } => { + servers.insert(name.clone(), config.as_ref().clone()); + } + McpServerContribution::Remove { name } => { + servers.remove(name); + } + } + } + } } diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 936f1b5fb5..304deb500d 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -56,7 +56,7 @@ pub(crate) async fn maybe_prompt_and_install_mcp_dependencies( let installed = sess .services .mcp_manager - .configured_servers(config.as_ref()) + .runtime_servers(config.as_ref()) .await; let missing = collect_missing_mcp_dependencies(mentioned_skills, &installed); if missing.is_empty() { @@ -98,7 +98,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( } let codex_home = config.codex_home.clone(); - let installed = sess.services.mcp_manager.configured_servers(config).await; + let installed = sess.services.mcp_manager.runtime_servers(config).await; let missing = collect_missing_mcp_dependencies(mentioned_skills, &installed); if missing.is_empty() { return; @@ -190,16 +190,22 @@ pub(crate) async fn maybe_install_mcp_dependencies( } } - // Refresh from the config-backed merged MCP map (global + repo + managed) - // and overlay the updated global servers so we don't drop repo-scoped - // servers. Runtime additions such as built-ins are rebuilt by the refresh - // path from the current config. - let mut refresh_servers = sess.services.mcp_manager.configured_servers(config).await; + let mut refresh_config = config.clone(); + let mut configured_servers = config.mcp_servers.get().clone(); for (name, server_config) in &servers { - refresh_servers + configured_servers .entry(name.clone()) .or_insert_with(|| server_config.clone()); } + if let Err(err) = refresh_config.mcp_servers.set(configured_servers) { + warn!("failed to refresh MCP dependencies for mentioned skills: {err}"); + return; + } + let refresh_servers = sess + .services + .mcp_manager + .runtime_servers(&refresh_config) + .await; sess.refresh_mcp_servers_now( turn_context, refresh_servers, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index cc018343eb..f2643d0e5c 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -311,8 +311,10 @@ impl Session { ) { let auth = self.services.auth_manager.auth().await; let config = self.get_config().await; - let mcp_config = config - .to_mcp_config(self.services.plugins_manager.as_ref()) + let mcp_config = self + .services + .mcp_manager + .runtime_config(config.as_ref()) .await; let tool_plugin_provenance = self .services diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index a0c1d63d71..d5f54d555e 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -272,7 +272,10 @@ impl ThreadManager { codex_home.to_path_buf(), restriction_product, )); - let mcp_manager = Arc::new(McpManager::new(Arc::clone(&plugins_manager))); + let mcp_manager = Arc::new(McpManager::new_with_extensions( + Arc::clone(&plugins_manager), + Arc::clone(&extensions), + )); let skills_manager = Arc::new(SkillsManager::new_with_restriction_product( codex_home, config.bundled_skills_enabled(), diff --git a/codex-rs/ext/extension-api/Cargo.toml b/codex-rs/ext/extension-api/Cargo.toml index 5497dc5439..2e155b430e 100644 --- a/codex-rs/ext/extension-api/Cargo.toml +++ b/codex-rs/ext/extension-api/Cargo.toml @@ -15,6 +15,7 @@ workspace = true [dependencies] async-trait = { workspace = true } +codex-config = { workspace = true } codex-context-fragments = { workspace = true } codex-protocol = { workspace = true } codex-tools = { workspace = true } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 8706e8ee7a..0403294371 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -10,12 +10,14 @@ use codex_tools::ToolExecutor; use crate::ExtensionData; +mod mcp; mod prompt; mod thread_lifecycle; mod tool_lifecycle; mod turn_input; mod turn_lifecycle; +pub use mcp::McpServerContribution; pub use prompt::PromptFragment; pub use prompt::PromptSlot; pub use thread_lifecycle::ThreadIdleInput; @@ -34,6 +36,18 @@ pub use turn_lifecycle::TurnErrorInput; pub use turn_lifecycle::TurnStartInput; pub use turn_lifecycle::TurnStopInput; +/// Extension contribution that resolves runtime MCP servers from host config. +/// +/// Contributors run in registration order. Later contributions for the same +/// name replace earlier ones. Implementations must contribute only names they +/// 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; +} + /// Extension contribution that adds prompt fragments during prompt assembly. pub trait ContextContributor: Send + Sync { fn contribute<'a>( diff --git a/codex-rs/ext/extension-api/src/contributors/mcp.rs b/codex-rs/ext/extension-api/src/contributors/mcp.rs new file mode 100644 index 0000000000..a55715692b --- /dev/null +++ b/codex-rs/ext/extension-api/src/contributors/mcp.rs @@ -0,0 +1,22 @@ +use codex_config::McpServerConfig; + +/// One extension-owned overlay for the runtime MCP server configuration. +#[derive(Clone, Debug)] +pub enum McpServerContribution { + /// Adds or replaces a named MCP server. + Set { + name: String, + config: Box, + }, + /// 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 7d8d95d139..8b269f31d1 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -31,6 +31,8 @@ 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::McpServerContribution; +pub use contributors::McpServerContributor; pub use contributors::PromptFragment; pub use contributors::PromptSlot; pub use contributors::ThreadIdleInput; diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index 0849387110..e8cc989337 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/codex-rs/ext/extension-api/src/registry.rs @@ -7,6 +7,7 @@ use crate::ConfigContributor; use crate::ContextContributor; use crate::ExtensionData; use crate::ExtensionEventSink; +use crate::McpServerContributor; use crate::NoopExtensionEventSink; use crate::ThreadLifecycleContributor; use crate::TokenUsageContributor; @@ -24,6 +25,7 @@ pub struct ExtensionRegistryBuilder { config_contributors: Vec>>, token_usage_contributors: Vec>, context_contributors: Vec>, + mcp_server_contributors: Vec>>, turn_input_contributors: Vec>, tool_contributors: Vec>, tool_lifecycle_contributors: Vec>, @@ -41,6 +43,7 @@ impl Default for ExtensionRegistryBuilder { token_usage_contributors: Vec::new(), approval_review_contributors: Vec::new(), context_contributors: Vec::new(), + mcp_server_contributors: Vec::new(), turn_input_contributors: Vec::new(), tool_contributors: Vec::new(), tool_lifecycle_contributors: Vec::new(), @@ -101,6 +104,11 @@ impl ExtensionRegistryBuilder { self.context_contributors.push(contributor); } + /// Registers one runtime MCP server contributor. + pub fn mcp_server_contributor(&mut self, contributor: Arc>) { + self.mcp_server_contributors.push(contributor); + } + /// Registers one turn-input contributor. pub fn turn_input_contributor(&mut self, contributor: Arc) { self.turn_input_contributors.push(contributor); @@ -131,6 +139,7 @@ impl ExtensionRegistryBuilder { token_usage_contributors: self.token_usage_contributors, approval_review_contributors: self.approval_review_contributors, context_contributors: self.context_contributors, + mcp_server_contributors: self.mcp_server_contributors, turn_input_contributors: self.turn_input_contributors, tool_contributors: self.tool_contributors, tool_lifecycle_contributors: self.tool_lifecycle_contributors, @@ -147,6 +156,7 @@ pub struct ExtensionRegistry { config_contributors: Vec>>, token_usage_contributors: Vec>, context_contributors: Vec>, + mcp_server_contributors: Vec>>, turn_input_contributors: Vec>, tool_contributors: Vec>, tool_lifecycle_contributors: Vec>, @@ -205,6 +215,11 @@ impl ExtensionRegistry { &self.context_contributors } + /// Returns the registered runtime MCP server contributors. + pub fn mcp_server_contributors(&self) -> &[Arc>] { + &self.mcp_server_contributors + } + /// Returns the registered turn-input contributors. pub fn turn_input_contributors(&self) -> &[Arc] { &self.turn_input_contributors diff --git a/codex-rs/ext/mcp/BUILD.bazel b/codex-rs/ext/mcp/BUILD.bazel new file mode 100644 index 0000000000..30aefba692 --- /dev/null +++ b/codex-rs/ext/mcp/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "mcp", + crate_name = "codex_mcp_extension", +) diff --git a/codex-rs/ext/mcp/Cargo.toml b/codex-rs/ext/mcp/Cargo.toml new file mode 100644 index 0000000000..66c9ad1727 --- /dev/null +++ b/codex-rs/ext/mcp/Cargo.toml @@ -0,0 +1,29 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-mcp-extension" +version.workspace = true + +[lib] +name = "codex_mcp_extension" +path = "src/lib.rs" +test = false +doctest = false + +[lints] +workspace = true + +[dependencies] +async-trait = { workspace = true } +codex-core = { workspace = true } +codex-extension-api = { workspace = true } +codex-features = { workspace = true } +codex-mcp = { workspace = true } + +[dev-dependencies] +codex-config = { workspace = true } +codex-core-plugins = { workspace = true } +codex-login = { workspace = true } +pretty_assertions = { workspace = true } +tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/mcp/src/lib.rs b/codex-rs/ext/mcp/src/lib.rs new file mode 100644 index 0000000000..4a7838ccce --- /dev/null +++ b/codex-rs/ext/mcp/src/lib.rs @@ -0,0 +1,31 @@ +use codex_core::config::Config; +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; + +struct HostedAppsMcpExtension; + +#[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 }]; + } + + 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(), + )), + }] + } +} + +pub fn install(builder: &mut ExtensionRegistryBuilder) { + builder.mcp_server_contributor(std::sync::Arc::new(HostedAppsMcpExtension)); +} diff --git a/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs b/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs new file mode 100644 index 0000000000..0d025df954 --- /dev/null +++ b/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs @@ -0,0 +1,93 @@ +use std::sync::Arc; + +use codex_config::McpServerTransportConfig; +use codex_core::McpManager; +use codex_core::config::Config; +use codex_core::config::ConfigBuilder; +use codex_core_plugins::PluginsManager; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_login::CodexAuth; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn contributes_hosted_apps_mcp_without_an_executor() -> Result<(), Box> +{ + 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() + .await?; + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let manager = installed_manager(&config); + + let runtime_config = manager.runtime_config(&config).await; + assert!(!runtime_config.legacy_apps_mcp_loader_enabled); + 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("Apps MCP should be contributed as a configured server")?; + let McpServerTransportConfig::StreamableHttp { url, .. } = &server.transport else { + panic!("Apps MCP should use streamable HTTP"); + }; + assert_eq!(url, "https://chatgpt.com/backend-api/ps/mcp"); + + Ok(()) +} + +#[tokio::test] +async fn hosted_apps_mcp_requires_chatgpt_auth() -> Result<(), Box> { + 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::from_api_key("test"); + let manager = installed_manager(&config); + + let servers = manager.effective_servers(&config, Some(&auth)).await; + assert!(!servers.contains_key(CODEX_APPS_MCP_SERVER_NAME)); + + Ok(()) +} + +#[tokio::test] +async fn disabled_apps_remove_reserved_server_config() -> Result<(), Box> { + 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(), false.into()), + ( + "mcp_servers.codex_apps.url".to_string(), + "https://example.com/mcp".into(), + ), + ]) + .build() + .await?; + let manager = installed_manager(&config); + + let servers = manager.runtime_servers(&config).await; + + assert!(!servers.contains_key(CODEX_APPS_MCP_SERVER_NAME)); + Ok(()) +} + +fn installed_manager(config: &Config) -> McpManager { + let mut builder = ExtensionRegistryBuilder::new(); + codex_mcp_extension::install(&mut builder); + McpManager::new_with_extensions( + Arc::new(PluginsManager::new(config.codex_home.to_path_buf())), + Arc::new(builder.build()), + ) +} From 9e3081be9672c65f8a0cd958719065f49f47d839 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Tue, 9 Jun 2026 13:51:57 -0700 Subject: [PATCH 16/17] [codex] Characterize global instruction lifecycle (#26830) ## Why Global instruction behavior spans thread creation, resume, forks, subagents, and compaction. Characterization coverage is needed before changing those semantics so preserved history can be distinguished from newly loaded configuration. ## What changed - Extends the existing `agents_md` suite with fresh-thread, warning, resume, fork, and subagent lifecycle coverage. - Extends the existing `compact` suite with manual, mid-turn, and remote-v2 compaction coverage. - Asserts rendered instruction fragments, reported source paths, and structured request history before and after instruction-file mutations. --- codex-rs/core/tests/suite/agents_md.rs | 602 ++++++++++++++++++++++++- codex-rs/core/tests/suite/compact.rs | 384 ++++++++++++++++ 2 files changed, 985 insertions(+), 1 deletion(-) diff --git a/codex-rs/core/tests/suite/agents_md.rs b/codex-rs/core/tests/suite/agents_md.rs index 3fa2136ecb..88f014f757 100644 --- a/codex-rs/core/tests/suite/agents_md.rs +++ b/codex-rs/core/tests/suite/agents_md.rs @@ -1,18 +1,46 @@ use anyhow::Result; +use anyhow::anyhow; +use codex_core::ForkSnapshot; use codex_exec_server::CreateDirectoryOptions; +use codex_features::Feature; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; use codex_utils_absolute_path::AbsolutePathBuf; +use core_test_support::PathBufExt; use core_test_support::create_directory_symlink; +use core_test_support::load_default_config_for_test; +use core_test_support::responses; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_sse_once; 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::TestCodexBuilder; use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use core_test_support::wait_for_event_match; use pretty_assertions::assert_eq; +use serde_json::json; use std::sync::Arc; +use std::time::Duration; use tempfile::TempDir; +const GLOBAL_AGENTS_FILENAME: &str = "AGENTS.md"; +const GLOBAL_AGENTS_OVERRIDE_FILENAME: &str = "AGENTS.override.md"; +const GLOBAL_INSTRUCTIONS: &str = "global instructions"; +const NEW_GLOBAL_INSTRUCTIONS: &str = "new global instructions"; +const NEW_PROJECT_INSTRUCTIONS: &str = "new project instructions"; +const OLD_GLOBAL_INSTRUCTIONS: &str = "old global instructions"; +const PROJECT_INSTRUCTIONS: &str = "project instructions"; +const PROJECT_SEPARATOR: &str = "--- project-doc ---"; +const SPAWN_CALL_ID: &str = "spawn-global-instructions-child"; +const SPAWN_CHILD_PROMPT: &str = "inspect inherited global instructions"; +const SPAWN_FRESH_PARENT_PROMPT: &str = "spawn a child with fresh context"; +const SPAWN_PARENT_PROMPT: &str = "spawn a child with the parent context"; +const SPAWN_SEED_PROMPT: &str = "seed parent history"; + async fn agents_instructions(mut builder: TestCodexBuilder) -> Result { let server = start_mock_server().await; let resp_mock = mount_sse_once( @@ -32,6 +60,52 @@ async fn agents_instructions(mut builder: TestCodexBuilder) -> Result { .ok_or_else(|| anyhow::anyhow!("instructions message not found")) } +fn write_global_file( + home: &TempDir, + filename: &str, + contents: impl AsRef<[u8]>, +) -> Result { + let path = home.path().join(filename); + std::fs::write(&path, contents)?; + Ok(path.abs()) +} + +fn instruction_fragments(request: &responses::ResponsesRequest) -> Vec { + request + .message_input_texts("user") + .into_iter() + .filter(|text| text.starts_with("# AGENTS.md instructions for ")) + .collect() +} + +fn expected_instruction_fragment(cwd: &AbsolutePathBuf, contents: &str) -> String { + let cwd = cwd.as_path().display(); + format!("# AGENTS.md instructions for {cwd}\n\n\n{contents}\n") +} + +fn assert_single_instruction_fragment(request: &responses::ResponsesRequest, expected: &str) { + assert_eq!(instruction_fragments(request), vec![expected.to_string()]); +} + +fn request_body_contains(request: &wiremock::Request, text: &str) -> bool { + let is_zstd = request + .headers + .get("content-encoding") + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| { + value + .split(',') + .any(|entry| entry.trim().eq_ignore_ascii_case("zstd")) + }); + let body = if is_zstd { + zstd::stream::decode_all(std::io::Cursor::new(&request.body)).ok() + } else { + Some(request.body.clone()) + }; + body.and_then(|body| String::from_utf8(body).ok()) + .is_some_and(|body| body.contains(text)) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn agents_override_is_preferred_over_agents_md() -> Result<()> { let instructions = @@ -250,7 +324,7 @@ async fn selected_environment_sources_match_model_visible_instructions() -> Resu }); let test = builder.build_with_remote_env(&server).await?; let project_agents = test.config.cwd.join("AGENTS.md"); - let global_agents = AbsolutePathBuf::try_from(global_agents).expect("absolute path"); + let global_agents = global_agents.abs(); assert_eq!( test.codex.instruction_sources().await, @@ -268,3 +342,529 @@ async fn selected_environment_sources_match_model_visible_instructions() -> Resu Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fresh_thread_composes_global_before_project_and_reports_sources() -> Result<()> { + // Set up one global source, one project source, and two ordinary model turns. + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("response-1"), + responses::ev_completed("response-1"), + ]), + responses::sse(vec![ + responses::ev_response_created("response-2"), + responses::ev_completed("response-2"), + ]), + ], + ) + .await; + let home = Arc::new(TempDir::new()?); + let global_source = + write_global_file(home.as_ref(), GLOBAL_AGENTS_FILENAME, GLOBAL_INSTRUCTIONS)?; + + let mut builder = test_codex() + .with_home(Arc::clone(&home)) + .with_workspace_setup(|cwd, fs| async move { + fs.write_file( + &cwd.join("AGENTS.md"), + PROJECT_INSTRUCTIONS.as_bytes().to_vec(), + /*sandbox*/ None, + ) + .await?; + Ok(()) + }); + let test = builder.build_with_remote_env(&server).await?; + let project_source = test.config.cwd.join(GLOBAL_AGENTS_FILENAME); + let creation_sources = vec![global_source.clone(), project_source.clone()]; + + // Confirm the thread records both creation-time sources in composition order. + assert_eq!(test.codex.instruction_sources().await, creation_sources); + + // Materialize the initial snapshot, then rewrite both selected files in place before another + // ordinary turn. + test.submit_turn("first turn").await?; + let rewritten_global_source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_FILENAME, + NEW_GLOBAL_INSTRUCTIONS, + )?; + test.fs() + .write_file( + &project_source, + NEW_PROJECT_INSTRUCTIONS.as_bytes().to_vec(), + /*sandbox*/ None, + ) + .await?; + assert_eq!( + rewritten_global_source, global_source, + "same-path mutation should retain the selected global source path" + ); + test.submit_turn("second turn").await?; + + // Assert the running thread keeps its original rendering and structured prefix even though + // both files at the reported source paths now contain different text. + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + let expected_contents = + format!("{GLOBAL_INSTRUCTIONS}\n\n{PROJECT_SEPARATOR}\n\n{PROJECT_INSTRUCTIONS}"); + let expected_fragment = expected_instruction_fragment(&test.config.cwd, &expected_contents); + let fragments = instruction_fragments(&requests[0]); + assert_eq!(fragments, vec![expected_fragment.clone()]); + assert_single_instruction_fragment(&requests[1], &expected_fragment); + let rendered = fragments + .into_iter() + .next() + .ok_or_else(|| anyhow!("expected one rendered instruction fragment"))?; + let global_position = rendered.find(GLOBAL_INSTRUCTIONS).ok_or_else(|| { + anyhow!( + "expected rendered instructions to contain {GLOBAL_INSTRUCTIONS:?}; observed: {rendered}" + ) + })?; + let project_position = rendered.find(PROJECT_INSTRUCTIONS).ok_or_else(|| { + anyhow!( + "expected rendered instructions to contain {PROJECT_INSTRUCTIONS:?}; observed: {rendered}" + ) + })?; + assert!( + global_position < project_position, + "global instructions should precede project instructions: {rendered}" + ); + assert!( + rendered.contains(PROJECT_SEPARATOR), + "expected rendered instructions to contain {PROJECT_SEPARATOR:?}; observed: {rendered}" + ); + assert_eq!( + test.codex.instruction_sources().await, + creation_sources, + "ordinary turns retain the creation-time source list" + ); + let first_input = requests[0].input(); + let second_input = requests[1].input(); + assert_eq!( + second_input.get(..first_input.len()), + Some(first_input.as_slice()), + "the ordinary second turn should retain the cached prefix" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn global_loading_warning_surfaces_during_thread_creation() -> Result<()> { + // Set up a malformed global instruction file and one model response. + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("warning-response"), + responses::ev_completed("warning-response"), + ]), + ) + .await; + let home = Arc::new(TempDir::new()?); + let source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_FILENAME, + b"global\xFFinstructions", + )?; + + // Create the thread, capture its load warning, and submit one turn for rendered output. + let mut builder = test_codex().with_home(home); + let test = builder.build(&server).await?; + let warning = wait_for_event_match(&test.codex, |event| match event { + EventMsg::Warning(warning) + if warning + .message + .contains(source.as_path().display().to_string().as_str()) => + { + Some(warning.message.clone()) + } + _ => None, + }) + .await; + test.submit_turn("inspect lossy global instructions") + .await?; + + // Assert the source is reported, the warning is specific, and rendering is lossily decoded. + assert_eq!(test.codex.instruction_sources().await, vec![source.clone()]); + assert!( + warning.contains("invalid UTF-8"), + "expected warning to contain \"invalid UTF-8\"; observed: {warning}" + ); + let expected_fragment = + expected_instruction_fragment(&test.config.cwd, "global\u{FFFD}instructions"); + assert_single_instruction_fragment(&response_mock.single_request(), &expected_fragment); + + Ok(()) +} + +// TODO(anp): Align cold-resume instruction sources with the historical instructions replayed to +// the model so the API source list and model-visible context describe the same files. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cold_resume_replays_rendered_instructions_but_reports_current_config_sources() -> Result<()> +{ + // Set up an initial turn and a later cold-resumed turn against the same rollout. + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("initial-response"), + responses::ev_completed("initial-response"), + ]), + responses::sse(vec![ + responses::ev_response_created("resumed-response"), + responses::ev_completed("resumed-response"), + ]), + ], + ) + .await; + let home = Arc::new(TempDir::new()?); + let old_source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_FILENAME, + OLD_GLOBAL_INSTRUCTIONS, + )?; + + // Create the initial thread and persist its creation-time instruction snapshot. + let mut initial_builder = test_codex().with_home(Arc::clone(&home)); + let initial = initial_builder.build(&server).await?; + + // Assert the pre-resume thread reports the source used to create its snapshot. + assert_eq!( + initial.codex.instruction_sources().await, + vec![old_source.clone()], + "initial thread reports the creation-time global source" + ); + initial.submit_turn("persist instructions").await?; + let rollout_path = initial + .session_configured + .rollout_path + .clone() + .expect("rollout path"); + initial.codex.submit(Op::Shutdown).await?; + wait_for_event(&initial.codex, |event| { + matches!(event, EventMsg::ShutdownComplete) + }) + .await; + + // Add a preferred override source, then cold-resume with freshly loaded configuration. + let new_source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_OVERRIDE_FILENAME, + NEW_GLOBAL_INSTRUCTIONS, + )?; + assert_ne!(old_source, new_source); + let mut resume_builder = test_codex().with_home(Arc::clone(&home)); + let resumed = resume_builder + .resume(&server, Arc::clone(&home), rollout_path) + .await?; + + // Assert the API reports the new source while model history replays the old structured prefix. + assert_eq!( + resumed.codex.instruction_sources().await, + vec![new_source], + "resume reports sources from the newly loaded config" + ); + + resumed.submit_turn("continue resumed thread").await?; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + let initial_input = requests[0].input(); + let resumed_input = requests[1].input(); + assert_eq!( + resumed_input.get(..initial_input.len()), + Some(initial_input.as_slice()), + "cold resume should replay the original structured input prefix" + ); + let expected_fragment = + expected_instruction_fragment(&initial.config.cwd, OLD_GLOBAL_INSTRUCTIONS); + assert_single_instruction_fragment(&requests[0], &expected_fragment); + assert_single_instruction_fragment(&requests[1], &expected_fragment); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fork_replays_rendered_instructions_from_shared_history() -> Result<()> { + // Set up a parent turn and a later fork turn against the parent's rollout. + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("parent-response"), + responses::ev_completed("parent-response"), + ]), + responses::sse(vec![ + responses::ev_response_created("fork-response"), + responses::ev_completed("fork-response"), + ]), + ], + ) + .await; + let home = Arc::new(TempDir::new()?); + let source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_FILENAME, + OLD_GLOBAL_INSTRUCTIONS, + )?; + + // Create the parent and persist its creation-time instruction snapshot. + let mut builder = test_codex().with_home(Arc::clone(&home)); + let parent = builder.build(&server).await?; + + // Assert the parent reports the source used to create its snapshot. + assert_eq!( + parent.codex.instruction_sources().await, + vec![source.clone()], + "parent reports the creation-time global source" + ); + parent.submit_turn("persist instructions").await?; + parent.codex.ensure_rollout_materialized().await; + parent.codex.flush_rollout().await?; + let rollout_path = parent.codex.rollout_path().expect("rollout path"); + + // Add a preferred override source, then fork with freshly loaded configuration. + let new_source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_OVERRIDE_FILENAME, + NEW_GLOBAL_INSTRUCTIONS, + )?; + assert_ne!(source, new_source); + let mut fork_config = load_default_config_for_test(home.as_ref()).await; + fork_config.cwd = parent.config.cwd.clone(); + fork_config.model = parent.config.model.clone(); + fork_config.model_provider = parent.config.model_provider.clone(); + fork_config.model_catalog = parent.config.model_catalog.clone(); + fork_config.codex_self_exe = parent.config.codex_self_exe.clone(); + let forked = parent + .thread_manager + .fork_thread( + ForkSnapshot::Interrupted, + fork_config, + rollout_path, + /*thread_source*/ None, + /*parent_trace*/ None, + ) + .await?; + + // Assert the fork reports the new source before issuing its first turn. + assert_eq!( + forked.thread.instruction_sources().await, + vec![new_source], + "fork config should reflect the newly loaded global source" + ); + + forked + .thread + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "continue fork".to_string(), + 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(&forked.thread, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + // Assert the forked model request replays the parent's exact structured history. + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + let parent_input = requests[0].input(); + let fork_input = requests[1].input(); + assert_eq!( + fork_input.get(..parent_input.len()), + Some(parent_input.as_slice()), + "fork should replay the parent's original structured input prefix" + ); + let expected_fragment = + expected_instruction_fragment(&parent.config.cwd, OLD_GLOBAL_INSTRUCTIONS); + assert_single_instruction_fragment(&requests[0], &expected_fragment); + assert_single_instruction_fragment(&requests[1], &expected_fragment); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn forked_subagent_replays_one_creation_time_global_instruction_fragment() -> Result<()> { + skip_if_no_network!(Ok(())); + run_subagent_global_instruction_case(/*fork_context*/ true).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fresh_subagent_uses_creation_time_instructions_without_parent_history() -> Result<()> { + skip_if_no_network!(Ok(())); + run_subagent_global_instruction_case(/*fork_context*/ false).await +} + +async fn run_subagent_global_instruction_case(fork_context: bool) -> Result<()> { + // Set up matched responses for the parent seed, spawn call, child turn, and parent follow-up. + let server = responses::start_mock_server().await; + let parent_prompt = if fork_context { + SPAWN_PARENT_PROMPT + } else { + SPAWN_FRESH_PARENT_PROMPT + }; + let seed_mock = responses::mount_sse_once_match( + &server, + |request: &wiremock::Request| request_body_contains(request, SPAWN_SEED_PROMPT), + responses::sse(vec![ + responses::ev_response_created("seed-response"), + responses::ev_assistant_message("seed-message", "seeded"), + responses::ev_completed("seed-response"), + ]), + ) + .await; + let spawn_args = serde_json::to_string(&json!({ + "message": SPAWN_CHILD_PROMPT, + "fork_context": fork_context, + }))?; + let spawn_mock = responses::mount_sse_once_match( + &server, + move |request: &wiremock::Request| request_body_contains(request, parent_prompt), + responses::sse(vec![ + responses::ev_response_created("spawn-response"), + responses::ev_function_call_with_namespace( + SPAWN_CALL_ID, + "multi_agent_v1", + "spawn_agent", + &spawn_args, + ), + responses::ev_completed("spawn-response"), + ]), + ) + .await; + let child_mock = responses::mount_sse_once_match( + &server, + |request: &wiremock::Request| { + request_body_contains(request, SPAWN_CHILD_PROMPT) + && !request_body_contains(request, SPAWN_CALL_ID) + }, + responses::sse(vec![ + responses::ev_response_created("child-response"), + responses::ev_assistant_message("child-message", "done"), + responses::ev_completed("child-response"), + ]), + ) + .await; + responses::mount_sse_once_match( + &server, + |request: &wiremock::Request| request_body_contains(request, SPAWN_CALL_ID), + responses::sse(vec![ + responses::ev_response_created("spawn-follow-up-response"), + responses::ev_assistant_message("spawn-follow-up-message", "child started"), + responses::ev_completed("spawn-follow-up-response"), + ]), + ) + .await; + + // Create the parent thread, record its source, and seed the history inherited by the child. + let home = Arc::new(TempDir::new()?); + let source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_FILENAME, + OLD_GLOBAL_INSTRUCTIONS, + )?; + let mut builder = test_codex() + .with_home(Arc::clone(&home)) + .with_config(|config| { + let _ = config.features.enable(Feature::Collab); + let _ = config.features.disable(Feature::EnableRequestCompression); + }); + let test = builder.build(&server).await?; + + // Assert the parent reports the creation-time source before spawning. + assert_eq!( + test.codex.instruction_sources().await, + vec![source.clone()], + "parent reports the creation-time global source before spawning" + ); + test.submit_turn(SPAWN_SEED_PROMPT).await?; + let seed_request = seed_mock.single_request(); + + // Add a preferred override, then spawn a full-history child while observing its thread ID. + let new_source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_OVERRIDE_FILENAME, + NEW_GLOBAL_INSTRUCTIONS, + )?; + assert_ne!(source, new_source); + let mut created_threads = test.thread_manager.subscribe_thread_created(); + test.submit_turn(parent_prompt).await?; + let child_thread_id = tokio::time::timeout(Duration::from_secs(10), created_threads.recv()) + .await + .map_err(|_| anyhow!("timed out waiting for the subagent thread"))??; + let child_thread = test.thread_manager.get_thread(child_thread_id).await?; + let spawn_request = spawn_mock.single_request(); + let child_request = tokio::time::timeout(Duration::from_secs(10), async { + loop { + if let Some(request) = child_mock.requests().into_iter().find(|request| { + request + .message_input_texts("user") + .iter() + .any(|text| text == SPAWN_CHILD_PROMPT) + }) { + break request; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .map_err(|_| anyhow!("timed out waiting for the subagent request"))?; + + // Assert parent and child report and render the parent's creation-time snapshot exactly once. + let expected_fragment = + expected_instruction_fragment(&test.config.cwd, OLD_GLOBAL_INSTRUCTIONS); + assert_single_instruction_fragment(&seed_request, &expected_fragment); + assert_single_instruction_fragment(&spawn_request, &expected_fragment); + assert_single_instruction_fragment(&child_request, &expected_fragment); + assert_eq!( + test.codex.instruction_sources().await, + vec![source.clone()], + "running parent retains the creation-time global source after spawning" + ); + assert_eq!( + child_thread.instruction_sources().await, + vec![source], + "subagent reports the parent's creation-time source" + ); + if fork_context { + let seed_input = seed_request.input(); + let child_input = child_request.input(); + assert_eq!( + child_input.get(..seed_input.len()), + Some(seed_input.as_slice()), + "forked subagent should replay the parent's original structured input prefix" + ); + } else { + let child_user_texts = child_request.message_input_texts("user"); + assert_eq!( + child_user_texts + .iter() + .filter(|text| text.as_str() == SPAWN_SEED_PROMPT) + .count(), + 0, + "fresh-context subagent should omit parent user history; observed: {child_user_texts:?}" + ); + assert_eq!( + child_user_texts + .iter() + .filter(|text| text.as_str() == SPAWN_CHILD_PROMPT) + .count(), + 1, + "fresh-context subagent should contain its own prompt exactly once; observed: {child_user_texts:?}" + ); + } + + Ok(()) +} diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index ac497aff3b..339cadda4e 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -1,4 +1,6 @@ #![allow(clippy::expect_used)] +use anyhow::Result; +use anyhow::anyhow; use codex_core::compact::SUMMARIZATION_PROMPT; use codex_core::compact::SUMMARY_PREFIX; use codex_core::config::Config; @@ -23,11 +25,13 @@ use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; +use codex_utils_absolute_path::AbsolutePathBuf; use core_test_support::PathBufExt; use core_test_support::context_snapshot; use core_test_support::context_snapshot::ContextSnapshotOptions; use core_test_support::context_snapshot::ContextSnapshotRenderMode; use core_test_support::hooks::trust_discovered_hooks; +use core_test_support::responses; use core_test_support::responses::ev_reasoning_item; use core_test_support::responses::mount_models_once; use core_test_support::skip_if_no_network; @@ -57,6 +61,8 @@ use serde_json::Value; use serde_json::json; use std::fs; use std::path::Path; +use std::sync::Arc; +use tempfile::TempDir; use wiremock::MockServer; // --- Test helpers ----------------------------------------------------------- @@ -78,6 +84,11 @@ const DUMMY_CALL_ID: &str = "call-multi-auto"; const FUNCTION_CALL_LIMIT_MSG: &str = "function call limit push"; const POST_AUTO_USER_MSG: &str = "post auto follow-up"; const PRETURN_CONTEXT_DIFF_CWD: &str = "/tmp/PRETURN_CONTEXT_DIFF_CWD"; +const GLOBAL_AGENTS_FILENAME: &str = "AGENTS.md"; +const GLOBAL_AGENTS_OVERRIDE_FILENAME: &str = "AGENTS.override.md"; +const NEW_GLOBAL_INSTRUCTIONS: &str = "new global instructions"; +const OLD_GLOBAL_INSTRUCTIONS: &str = "old global instructions"; +const REMOTE_V2_SUMMARY: &str = "global-instructions-remote-v2-summary"; pub(super) const COMPACT_WARNING_MESSAGE: &str = "Heads up: Long threads and multiple compactions can cause the model to be less accurate. Start a new thread when possible to keep threads small and targeted."; @@ -267,6 +278,92 @@ fn non_openai_model_provider(server: &MockServer) -> ModelProviderInfo { provider } +fn write_global_file( + home: &TempDir, + filename: &str, + contents: impl AsRef<[u8]>, +) -> Result { + let path = home.path().join(filename); + std::fs::write(&path, contents)?; + Ok(path.abs()) +} + +fn instruction_fragments(request: &responses::ResponsesRequest) -> Vec { + request + .message_input_texts("user") + .into_iter() + .filter(|text| text.starts_with("# AGENTS.md instructions for ")) + .collect() +} + +fn instruction_fragments_in_items(items: &[Value]) -> Vec { + items + .iter() + .filter(|item| { + item.get("type").and_then(Value::as_str) == Some("message") + && item.get("role").and_then(Value::as_str) == Some("user") + }) + .filter_map(|item| item.get("content").and_then(Value::as_array)) + .flatten() + .filter_map(|span| span.get("text").and_then(Value::as_str)) + .filter(|text| text.starts_with("# AGENTS.md instructions for ")) + .map(str::to_string) + .collect() +} + +fn expected_instruction_fragment(cwd: &AbsolutePathBuf, contents: &str) -> String { + let cwd = cwd.as_path().display(); + format!("# AGENTS.md instructions for {cwd}\n\n\n{contents}\n") +} + +fn assert_single_instruction_fragment(request: &responses::ResponsesRequest, expected: &str) { + assert_eq!(instruction_fragments(request), vec![expected.to_string()]); +} + +fn replacement_history_from_rollout(path: &Path) -> Result> { + let rollout_text = fs::read_to_string(path)?; + let mut replacement_history = None; + for line in rollout_text + .lines() + .map(str::trim) + .filter(|line| !line.is_empty()) + { + let entry: RolloutLine = serde_json::from_str(line)?; + if let RolloutItem::Compacted(compacted) = entry.item + && let Some(items) = compacted.replacement_history + { + replacement_history = Some( + items + .into_iter() + .map(serde_json::to_value) + .collect::, _>>()?, + ); + } + } + replacement_history.ok_or_else(|| anyhow!("expected rollout replacement history")) +} + +fn remote_v2_compaction_response() -> String { + responses::sse(vec![ + json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": REMOTE_V2_SUMMARY, + } + }), + responses::ev_completed("remote-v2-compact-response"), + ]) +} + +fn local_compaction_provider(server: &wiremock::MockServer) -> ModelProviderInfo { + let mut provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone(); + provider.name = "OpenAI-compatible test provider".to_string(); + provider.base_url = Some(format!("{}/v1", server.uri())); + provider.supports_websockets = false; + provider +} + fn model_info_with_context_window(slug: &str, context_window: i64) -> ModelInfo { let models_response = bundled_models_response() .unwrap_or_else(|err| panic!("bundled models.json should parse: {err}")); @@ -3998,3 +4095,290 @@ async fn snapshot_request_shape_manual_compact_without_previous_user_messages() ) ); } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn manual_compaction_keeps_the_creation_time_global_instructions() -> Result<()> { + // Set up an initial turn, a manual compaction response, and a post-compaction turn. + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("first-response"), + responses::ev_completed("first-response"), + ]), + responses::sse(vec![ + responses::ev_response_created("compact-response"), + responses::ev_assistant_message("compact-message", "summary"), + responses::ev_completed("compact-response"), + ]), + responses::sse(vec![ + responses::ev_response_created("follow-up-response"), + responses::ev_completed("follow-up-response"), + ]), + ], + ) + .await; + let home = Arc::new(TempDir::new()?); + let source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_FILENAME, + OLD_GLOBAL_INSTRUCTIONS, + )?; + let provider = local_compaction_provider(&server); + + // Create the thread with the old global source loaded into its instruction snapshot. + let mut builder = test_codex() + .with_home(Arc::clone(&home)) + .with_config(move |config| { + config.model_provider = provider; + }); + let test = builder.build(&server).await?; + + // Assert the pre-compaction source list points at the creation-time file. + assert_eq!( + test.codex.instruction_sources().await, + vec![source.clone()], + "thread reports the creation-time global source before compaction" + ); + + // Materialize the old snapshot, rewrite the selected file in place, and manually compact. + test.submit_turn("first turn").await?; + let rewritten_source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_FILENAME, + NEW_GLOBAL_INSTRUCTIONS, + )?; + assert_eq!(source, rewritten_source); + + test.codex.submit(Op::Compact).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.submit_turn("after compact").await?; + + // Assert ordinary and compact turns keep the old rendering even though the reported source + // path now contains new text. + let requests = response_mock.requests(); + assert_eq!(requests.len(), 3); + let expected_fragment = + expected_instruction_fragment(&test.config.cwd, OLD_GLOBAL_INSTRUCTIONS); + assert_single_instruction_fragment(&requests[0], &expected_fragment); + assert_single_instruction_fragment(&requests[1], &expected_fragment); + assert_single_instruction_fragment(&requests[2], &expected_fragment); + assert_eq!( + test.codex.instruction_sources().await, + vec![source], + "thread retains the creation-time global source after compaction" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mid_turn_compaction_keeps_the_creation_time_global_instructions() -> Result<()> { + // Set up a turn that crosses the auto-compaction limit and a post-compaction response. + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_function_call("call-1", "unsupported_tool", "{}"), + responses::ev_completed_with_tokens("first-response", /*total_tokens*/ 96), + ]), + responses::sse(vec![ + responses::ev_assistant_message("compact-message", "summary"), + responses::ev_completed_with_tokens("compact-response", /*total_tokens*/ 10), + ]), + responses::sse(vec![ + responses::ev_assistant_message("final-message", "done"), + responses::ev_completed_with_tokens("follow-up-response", /*total_tokens*/ 10), + ]), + ], + ) + .await; + let home = Arc::new(TempDir::new()?); + let source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_FILENAME, + OLD_GLOBAL_INSTRUCTIONS, + )?; + let provider = local_compaction_provider(&server); + + // Create the thread with the old global source loaded into its instruction snapshot. + let mut builder = test_codex() + .with_home(Arc::clone(&home)) + .with_config(move |config| { + config.model_provider = provider; + config.model_context_window = Some(100); + config.model_auto_compact_token_limit = Some(90); + }); + let test = builder.build(&server).await?; + + // Assert the pre-compaction source list points at the creation-time file. + assert_eq!( + test.codex.instruction_sources().await, + vec![source.clone()], + "thread reports the creation-time global source before mid-turn compaction" + ); + + // Add a preferred override before the turn triggers automatic mid-turn compaction. + let new_source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_OVERRIDE_FILENAME, + NEW_GLOBAL_INSTRUCTIONS, + )?; + assert_ne!(source, new_source); + test.submit_turn("trigger mid-turn compaction").await?; + + // Assert the initial, compact, and resumed requests all keep the old snapshot and source. + let requests = response_mock.requests(); + assert_eq!(requests.len(), 3); + let expected_fragment = + expected_instruction_fragment(&test.config.cwd, OLD_GLOBAL_INSTRUCTIONS); + assert_single_instruction_fragment(&requests[0], &expected_fragment); + assert_single_instruction_fragment(&requests[1], &expected_fragment); + assert_single_instruction_fragment(&requests[2], &expected_fragment); + assert_eq!( + test.codex.instruction_sources().await, + vec![source], + "thread retains the creation-time global source after mid-turn compaction" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_v2_compaction_keeps_creation_time_instructions_after_same_path_mutation() +-> Result<()> { + skip_if_no_network!(Ok(())); + + // Set up an ordinary turn, a remote-v2 compact response, and a post-compaction turn. + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("remote-v2-initial-response"), + responses::ev_completed("remote-v2-initial-response"), + ]), + remote_v2_compaction_response(), + responses::sse(vec![ + responses::ev_response_created("remote-v2-follow-up-response"), + responses::ev_completed("remote-v2-follow-up-response"), + ]), + responses::sse(vec![ + responses::ev_response_created("remote-v2-resumed-response"), + responses::ev_completed("remote-v2-resumed-response"), + ]), + ], + ) + .await; + let home = Arc::new(TempDir::new()?); + let source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_FILENAME, + OLD_GLOBAL_INSTRUCTIONS, + )?; + let mut builder = test_codex() + .with_home(Arc::clone(&home)) + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + let _ = config.features.enable(Feature::RemoteCompactionV2); + }); + let test = builder.build(&server).await?; + + // Materialize the old snapshot, rewrite the selected file in place, and compact remotely. + test.submit_turn("before remote v2 compaction").await?; + let rewritten_source = write_global_file( + home.as_ref(), + GLOBAL_AGENTS_FILENAME, + NEW_GLOBAL_INSTRUCTIONS, + )?; + assert_eq!(source, rewritten_source); + test.codex.submit(Op::Compact).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.submit_turn("after remote v2 compaction").await?; + test.codex.flush_rollout().await?; + + // Assert the compact request, installed replacement history, and follow-up all keep the + // creation-time item despite the file-backed source now containing new text. + let requests = response_mock.requests(); + assert_eq!(requests.len(), 3); + let old_fragment = expected_instruction_fragment(&test.config.cwd, OLD_GLOBAL_INSTRUCTIONS); + assert_single_instruction_fragment(&requests[0], &old_fragment); + assert_single_instruction_fragment(&requests[1], &old_fragment); + assert_single_instruction_fragment(&requests[2], &old_fragment); + assert_eq!( + requests[1].input().last(), + Some(&json!({"type": "compaction_trigger"})), + "remote-v2 compact request should append exactly one compaction trigger" + ); + let rollout_path = test.codex.rollout_path().expect("rollout path"); + let replacement_history = replacement_history_from_rollout(&rollout_path)?; + assert_eq!( + instruction_fragments_in_items(&replacement_history), + Vec::::new(), + "remote-v2 replacement history currently omits the global-instruction fragment" + ); + assert_eq!( + test.codex.instruction_sources().await, + vec![source.clone()], + "running thread retains the selected same-path source" + ); + assert_eq!( + fs::read_to_string(source.as_path())?, + NEW_GLOBAL_INSTRUCTIONS, + "the selected source path should contain the rewritten text" + ); + + // Cold-resume the persisted replacement history with freshly loaded same-path configuration. + test.codex.submit(Op::Shutdown).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::ShutdownComplete) + }) + .await; + let resumed_cwd = test.config.cwd.clone(); + let mut resume_builder = test_codex() + .with_home(Arc::clone(&home)) + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config.cwd = resumed_cwd; + let _ = config.features.enable(Feature::RemoteCompactionV2); + }); + let resumed = resume_builder + .resume(&server, Arc::clone(&home), rollout_path) + .await?; + resumed + .submit_turn("after remote v2 compaction cold resume") + .await?; + + // Modern replacement-history resume replays the persisted checkpoint and its later old-context + // suffix even though the same source path now contains new text. + let requests = response_mock.requests(); + assert_eq!(requests.len(), 4); + assert_single_instruction_fragment(&requests[3], &old_fragment); + let resumed_input = requests[3].input(); + assert_eq!( + resumed_input.get(..replacement_history.len()), + Some(replacement_history.as_slice()), + "remote-v2 cold resume should replay persisted replacement history verbatim" + ); + let post_compact_input = requests[2].input(); + assert_eq!( + resumed_input.get(..post_compact_input.len()), + Some(post_compact_input.as_slice()), + "remote-v2 cold resume should replay the complete post-compaction structured prefix" + ); + assert_eq!( + resumed.codex.instruction_sources().await, + vec![source], + "cold-resumed thread reports the same rewritten source path" + ); + + Ok(()) +} From f574946960a4d21e884453ac9963a869153b847d Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Tue, 9 Jun 2026 14:18:24 -0700 Subject: [PATCH 17/17] [codex] Report unusable MCP OAuth credentials as logged out (#26713) ## Why Persisted MCP OAuth credentials were reported as authenticated whenever a credential record existed. An expired token without a usable refresh token could therefore appear as `OAuth` even though startup could not authenticate with it, leaving users with a misleading status instead of a login prompt. ## What changed - Classify stored OAuth credentials as missing, usable, or requiring authorization. - Reuse the existing refresh window so near-expiry credentials without a refresh path are also treated as logged out. - Validate required credential fields before reporting OAuth authentication. - Add unit coverage for credential usability and integration coverage for expired, unexpired, and refreshable persisted credentials. ## Validation - `just test -p codex-rmcp-client` --- codex-rs/rmcp-client/src/auth_status.rs | 11 +- codex-rs/rmcp-client/src/oauth.rs | 112 +++++++++++++++++- .../tests/streamable_http_oauth_startup.rs | 103 ++++++++++++++++ 3 files changed, 220 insertions(+), 6 deletions(-) diff --git a/codex-rs/rmcp-client/src/auth_status.rs b/codex-rs/rmcp-client/src/auth_status.rs index 0b3b3bf6a7..bef9edc6c9 100644 --- a/codex-rs/rmcp-client/src/auth_status.rs +++ b/codex-rs/rmcp-client/src/auth_status.rs @@ -12,7 +12,8 @@ use reqwest::header::HeaderMap; use serde::Deserialize; use tracing::debug; -use crate::oauth::has_oauth_tokens; +use crate::oauth::StoredOAuthTokenStatus; +use crate::oauth::oauth_token_status; use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; @@ -44,8 +45,12 @@ pub async fn determine_streamable_http_auth_status( return Ok(McpAuthStatus::BearerToken); } - if has_oauth_tokens(server_name, url, store_mode)? { - return Ok(McpAuthStatus::OAuth); + match oauth_token_status(server_name, url, store_mode)? { + StoredOAuthTokenStatus::Usable => return Ok(McpAuthStatus::OAuth), + StoredOAuthTokenStatus::AuthorizationRequired => { + return Ok(McpAuthStatus::NotLoggedIn); + } + StoredOAuthTokenStatus::Missing => {} } match discover_streamable_http_oauth_with_headers(url, &default_headers).await { diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index c348460795..5742660277 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -76,6 +76,13 @@ impl PartialEq for WrappedOAuthTokenResponse { } } +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum StoredOAuthTokenStatus { + Missing, + Usable, + AuthorizationRequired, +} + pub(crate) fn load_oauth_tokens( server_name: &str, url: &str, @@ -94,12 +101,33 @@ pub(crate) fn load_oauth_tokens( } } -pub(crate) fn has_oauth_tokens( +pub(crate) fn oauth_token_status( server_name: &str, url: &str, store_mode: OAuthCredentialsStoreMode, -) -> Result { - Ok(load_oauth_tokens(server_name, url, store_mode)?.is_some()) +) -> Result { + Ok( + match load_oauth_tokens(server_name, url, store_mode)?.as_ref() { + None => StoredOAuthTokenStatus::Missing, + Some(tokens) if oauth_tokens_are_usable(tokens) => StoredOAuthTokenStatus::Usable, + Some(_) => StoredOAuthTokenStatus::AuthorizationRequired, + }, + ) +} + +fn oauth_tokens_are_usable(tokens: &StoredOAuthTokens) -> bool { + if tokens.client_id.trim().is_empty() { + return false; + } + + let token_response = &tokens.token_response.0; + if token_needs_refresh(tokens.expires_at) { + return token_response + .refresh_token() + .is_some_and(|token| !token.secret().trim().is_empty()); + } + + !token_response.access_token().secret().trim().is_empty() } fn refresh_expires_in_from_timestamp(tokens: &mut StoredOAuthTokens) { @@ -852,6 +880,84 @@ mod tests { assert_eq!(tokens.token_response.0.expires_in(), Some(Duration::ZERO)); } + #[test] + fn oauth_tokens_are_usable_when_expiry_is_unknown() { + let mut tokens = sample_tokens(); + tokens.expires_at = None; + tokens.token_response.0.set_refresh_token(None); + + assert!(super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_usable_when_unexpired_without_refresh_token() { + let mut tokens = sample_tokens(); + tokens.token_response.0.set_refresh_token(None); + + assert!(super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_usable_when_expired_but_refreshable() { + let mut tokens = sample_tokens(); + tokens.expires_at = Some(0); + + assert!(super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_not_usable_when_expired_and_unrefreshable() { + let mut tokens = sample_tokens(); + tokens.expires_at = Some(0); + tokens.token_response.0.set_refresh_token(None); + + assert!(!super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_not_usable_when_near_expiry_and_unrefreshable() { + let mut tokens = sample_tokens(); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_millis() as u64; + tokens.expires_at = Some(now.saturating_add(REFRESH_SKEW_MILLIS - 1)); + tokens.token_response.0.set_refresh_token(None); + + assert!(!super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_not_usable_when_client_id_is_blank() { + let mut tokens = sample_tokens(); + tokens.client_id = " ".to_string(); + + assert!(!super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_not_usable_when_access_token_is_blank() { + let mut tokens = sample_tokens(); + tokens + .token_response + .0 + .set_access_token(AccessToken::new(" ".to_string())); + + assert!(!super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_not_usable_when_required_refresh_token_is_blank() { + let mut tokens = sample_tokens(); + tokens.expires_at = Some(0); + tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new(" ".to_string()))); + + assert!(!super::oauth_tokens_are_usable(&tokens)); + } + fn assert_tokens_match_without_expiry( actual: &StoredOAuthTokens, expected: &StoredOAuthTokens, diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs index 1c18f2c98b..d7d8cf4232 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,16 +1,21 @@ mod streamable_http_test_support; use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; +use codex_rmcp_client::McpAuthStatus; use codex_rmcp_client::RmcpClient; use codex_rmcp_client::StoredOAuthTokens; use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::determine_streamable_http_auth_status; use codex_rmcp_client::save_oauth_tokens; use oauth2::AccessToken; use oauth2::RefreshToken; use oauth2::basic::BasicTokenType; +use pretty_assertions::assert_eq; use rmcp::transport::auth::OAuthTokenResponse; use rmcp::transport::auth::VendorExtraTokenFields; use serde_json::Value; @@ -33,6 +38,9 @@ const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; +const UNREFRESHABLE_SERVER_URL: &str = "https://unrefreshable.example/mcp"; +const UNEXPIRED_SERVER_URL: &str = "https://unexpired.example/mcp"; +const REFRESHABLE_SERVER_URL: &str = "https://refreshable.example/mcp"; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { @@ -112,6 +120,101 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn reports_auth_status_for_persisted_credentials() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + + let status = Command::new(std::env::current_exe()?) + .args([ + "persisted_credentials_auth_status_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .status() + .await?; + + assert!( + status.success(), + "persisted credentials auth status child failed: {status}" + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by reports_auth_status_for_persisted_credentials"] +async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { + let response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: UNREFRESHABLE_SERVER_URL.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; + + let status = auth_status(UNREFRESHABLE_SERVER_URL).await?; + assert_eq!(status, McpAuthStatus::NotLoggedIn); + + let response = OAuthTokenResponse::new( + AccessToken::new("unexpired-access-token".to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_millis() as u64; + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: UNEXPIRED_SERVER_URL.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(now.saturating_add(/*rhs*/ 60_000)), + }; + save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; + + let status = auth_status(UNEXPIRED_SERVER_URL).await?; + assert_eq!(status, McpAuthStatus::OAuth); + + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string()))); + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: REFRESHABLE_SERVER_URL.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; + + let status = auth_status(REFRESHABLE_SERVER_URL).await?; + assert_eq!(status, McpAuthStatus::OAuth); + Ok(()) +} + +async fn auth_status(server_url: &str) -> anyhow::Result { + determine_streamable_http_auth_status( + SERVER_NAME, + server_url, + /*bearer_token_env_var*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + ) + .await +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"] async fn oauth_startup_child() -> anyhow::Result<()> {