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 = [] 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/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 3d8086c27b..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", @@ -1994,6 +1995,7 @@ dependencies = [ "codex-rollout", "codex-sandboxing", "codex-shell-command", + "codex-skills-extension", "codex-state", "codex-thread-store", "codex-tools", @@ -2919,9 +2921,12 @@ name = "codex-extension-api" version = "0.0.0" dependencies = [ "async-trait", + "codex-config", "codex-context-fragments", "codex-protocol", "codex-tools", + "pretty_assertions", + "tokio", ] [[package]] @@ -3244,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" @@ -3796,8 +3818,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..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" } @@ -212,6 +214,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/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-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/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 2c3632d5cf..d1d2d7efe4 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": { @@ -3799,11 +3851,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 6b13ad73fe..4fe6277a61 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": { @@ -18061,11 +18113,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 0ed2dd6887..60e6078a83 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": { @@ -15878,11 +15930,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 405c2c921e..8971bda1a6 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 4fa4d0501a..d4376ca877 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,12 +136,29 @@ ], "type": "string" }, - "ThreadSource": { - "enum": [ - "user", - "subagent", - "memory_consolidation" + "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": { "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/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/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/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/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.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs index 7e406eeff1..754cbddc48 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; @@ -156,6 +158,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-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/Cargo.toml b/codex-rs/app-server/Cargo.toml index 4bdc430a1a..55ea78e617 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 } @@ -60,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/README.md b/codex-rs/app-server/README.md index 9f96277652..36483f8c45 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -130,7 +130,8 @@ 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. Set `disableMultiAgentTools: true` to prevent the thread from receiving multi-agent collaboration tools. 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`. +<<<<<<< HEAD +- `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. Set `disableMultiAgentTools: true` to prevent the thread from receiving multi-agent collaboration tools. 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. Set `disableMultiAgentTools: true` to prevent the fork from inheriting or enabling multi-agent collaboration tools. 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 +250,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..a0a44fd80d 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, @@ -49,8 +50,14 @@ 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( + &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..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( @@ -181,12 +178,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 +198,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.rs b/codex-rs/app-server/src/request_processors.rs index bac92c6b4e..731f2e1eb4 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/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 8bda39773c..a785343a50 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; @@ -331,6 +333,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( @@ -826,6 +839,7 @@ impl ThreadRequestProcessor { base_instructions, developer_instructions, dynamic_tools, + selected_capability_roots, mock_experimental_field: _mock_experimental_field, experimental_raw_events, personality, @@ -882,6 +896,7 @@ impl ThreadRequestProcessor { config, typesafe_overrides, dynamic_tools, + selected_capability_roots.unwrap_or_default(), session_start_source, thread_source.map(Into::into), disable_multi_agent_tools, @@ -955,6 +970,7 @@ impl ThreadRequestProcessor { config_overrides: Option>, typesafe_overrides: ConfigOverrides, dynamic_tools: Option>, + selected_capability_roots: Vec, session_start_source: Option, thread_source: Option, disable_multi_agent_tools: bool, @@ -1057,6 +1073,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, @@ -1080,6 +1100,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", @@ -2427,7 +2448,7 @@ impl ThreadRequestProcessor { return Ok(()); } }; - match self + let stored_thread_from_running_probe = match self .resume_running_thread( &request_id, ¶ms, @@ -2436,13 +2457,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, @@ -2466,15 +2487,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; @@ -2715,7 +2741,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 @@ -2751,7 +2777,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, + )))); + } } }; @@ -2792,7 +2822,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}"); @@ -2884,9 +2916,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( @@ -4132,7 +4164,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(), @@ -4221,7 +4253,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/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index f7c609fa22..c9dbe900d4 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/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/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/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/app-server/tests/suite/v2/skills_list.rs b/codex-rs/app-server/tests/suite/v2/skills_list.rs index cba2f4ee3f..2824a725a0 100644 --- a/codex-rs/app-server/tests/suite/v2/skills_list.rs +++ b/codex-rs/app-server/tests/suite/v2/skills_list.rs @@ -791,8 +791,10 @@ async fn skills_changed_notification_is_emitted_after_skill_change() -> Result<( ephemeral: None, session_start_source: None, thread_source: None, + disable_multi_agent_tools: false, dynamic_tools: None, environments: None, + selected_capability_roots: None, mock_experimental_field: None, experimental_raw_events: false, }) 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(())); 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-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?; 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/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/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/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 c6d0d9767e..50ff984231 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; @@ -192,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(), } } @@ -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, @@ -547,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(), }, @@ -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(), @@ -1085,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/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/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/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index c3f7355836..125087be15 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; @@ -183,6 +184,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 { @@ -271,7 +273,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(), @@ -578,6 +583,7 @@ impl ThreadManager { metrics_service_name: None, parent_trace: None, environments, + thread_extension_init: ExtensionDataInit::default(), })) .await } @@ -617,6 +623,7 @@ impl ThreadManager { /*inherited_exec_policy*/ None, options.parent_trace, options.environments, + options.thread_extension_init, /*user_shell_override*/ None, )) .await @@ -706,6 +713,7 @@ impl ThreadManager { /*inherited_exec_policy*/ None, parent_trace, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ None, )) .await @@ -733,6 +741,7 @@ impl ThreadManager { /*metrics_service_name*/ None, /*parent_trace*/ None, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ Some(user_shell_override), )) .await @@ -769,6 +778,7 @@ impl ThreadManager { /*inherited_exec_policy*/ None, /*parent_trace*/ None, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ Some(user_shell_override), )) .await @@ -952,6 +962,7 @@ impl ThreadManager { /*metrics_service_name*/ None, parent_trace, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ None, )) .await @@ -1158,6 +1169,7 @@ impl ThreadManagerState { inherited_exec_policy, /*parent_trace*/ None, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ None, )) .await @@ -1195,6 +1207,7 @@ impl ThreadManagerState { inherited_exec_policy, /*parent_trace*/ None, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ None, )) .await @@ -1233,6 +1246,7 @@ impl ThreadManagerState { inherited_exec_policy, /*parent_trace*/ None, environments, + /*thread_extension_init*/ ExtensionDataInit::default(), /*user_shell_override*/ None, )) .await @@ -1254,6 +1268,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( @@ -1272,6 +1287,7 @@ impl ThreadManagerState { /*inherited_exec_policy*/ None, parent_trace, environments, + thread_extension_init, user_shell_override, )) .await @@ -1295,6 +1311,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(_)); @@ -1363,6 +1380,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 a6ce4f4874..89f8f83951 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -332,6 +332,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 { @@ -370,6 +371,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"); @@ -386,6 +388,82 @@ 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, + multi_agent_version: 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"); @@ -426,6 +504,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"); @@ -697,6 +776,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/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(()) +} diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 8b54727d28..9c32ccac0c 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -758,6 +758,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/core/tests/suite/tools.rs b/codex-rs/core/tests/suite/tools.rs index afe172cdf9..17e2d95f98 100644 --- a/codex-rs/core/tests/suite/tools.rs +++ b/codex-rs/core/tests/suite/tools.rs @@ -100,6 +100,7 @@ async fn multi_agent_version_override_omits_multi_agent_tools_from_requests() -> metrics_service_name: None, parent_trace: None, environments: Vec::new(), + thread_extension_init: Default::default(), }) .await?; assert_eq!( 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/ext/extension-api/Cargo.toml b/codex-rs/ext/extension-api/Cargo.toml index 85c7d8f98e..2e155b430e 100644 --- a/codex-rs/ext/extension-api/Cargo.toml +++ b/codex-rs/ext/extension-api/Cargo.toml @@ -15,6 +15,11 @@ workspace = true [dependencies] async-trait = { workspace = true } +codex-config = { 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/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 7fa60c0fe7..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; @@ -59,3 +61,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/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/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/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); +} 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/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()), + ) +} 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..a0f64673a0 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -247,10 +247,12 @@ impl MemoryStartupContext { InternalSessionSource::MemoryConsolidation, )), thread_source: Some(ThreadSource::MemoryConsolidation), + multi_agent_version: None, dynamic_tools: Vec::new(), metrics_service_name: None, parent_trace: None, environments, + thread_extension_init: Default::default(), }) .await?; 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 { 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; 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/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<()> { 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.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(); 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, 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 edc308a905..2ab5103663 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, @@ -176,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()), @@ -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) } @@ -388,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()); 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/app.rs b/codex-rs/tui/src/app.rs index c64e323fe9..5916456956 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -127,7 +127,6 @@ use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadLoadedListParams; use codex_app_server_protocol::ThreadMemoryMode; use codex_app_server_protocol::ThreadRollbackResponse; -use codex_app_server_protocol::ThreadSource; use codex_app_server_protocol::ThreadStartSource; use codex_app_server_protocol::Turn; use codex_app_server_protocol::TurnError as AppServerTurnError; @@ -700,6 +699,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 164ea2c718..d3b596e97f 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -744,7 +744,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 @@ -755,12 +755,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.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 a94d483f0d..49fdd438f8 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() @@ -288,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/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; 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); } } 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. 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()