Merge branch 'main' into codex/refresh-turn-diff-environments

This commit is contained in:
sayan-oai
2026-06-25 18:06:18 -07:00
committed by GitHub
240 changed files with 11013 additions and 1695 deletions
+9 -7
View File
@@ -107,8 +107,9 @@ common:ci-bazel --config=ci
common:ci-bazel --build_metadata=TAG_workflow=bazel
# Keep code-mode integration cases out of ordinary Bazel legs. The
# Windows-cross config below re-enables them after generating its Windows V8
# snapshot on the Windows runner.
common:ci-bazel --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=suite::code_mode::
# snapshot on the Windows runner. Also skip the ConPTY Ctrl-C integration test
# until Windows Bazel CI can reliably interrupt foreground processes.
common:ci-bazel --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=suite::code_mode::,tests::windows_tests::conpty_ctrl_c_interrupts_powershell_foreground_child
# Shared config for Bazel-backed Rust linting.
build:clippy --aspects=@rules_rust//rust:defs.bzl%rust_clippy_aspect
@@ -191,11 +192,12 @@ common:ci-windows-cross --strategy=TestRunner=local
common:ci-windows-cross --strategy=V8Mksnapshot=local
common:ci-windows-cross --local_test_jobs=4
common:ci-windows-cross --test_env=RUST_TEST_THREADS=1
# Native Windows CI still covers the PowerShell tests. The cross-built gnullvm
# binaries currently hang in PowerShell AST parser tests when those binaries are
# run on the Windows runner. Keep V8-backed code-mode tests enabled except for
# the hidden dynamic-tool callback test, which currently times out on Windows.
common:ci-windows-cross --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=powershell,suite::code_mode::code_mode_can_call_hidden_dynamic_tools
# Native Windows CI still covers the PowerShell parser-process tests. The
# cross-built gnullvm binaries currently hang in those tests when run on the
# Windows runner. Keep V8-backed code-mode tests enabled except for the hidden
# dynamic-tool callback test, which currently times out on Windows. This config
# replaces the base skip list, so repeat the exact ConPTY Ctrl-C test exclusion.
common:ci-windows-cross --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=command_safety::powershell_parser::tests::,suite::code_mode::code_mode_can_call_hidden_dynamic_tools,tests::windows_tests::conpty_ctrl_c_interrupts_powershell_foreground_child
common:ci-windows-cross --platforms=//:windows_x86_64_gnullvm
common:ci-windows-cross --extra_execution_platforms=//:rbe,//:windows_x86_64_msvc
common:ci-windows-cross --extra_toolchains=//:windows_gnullvm_tests_on_msvc_host_toolchain
+4
View File
@@ -5,21 +5,25 @@
"macos-aarch64": {
"name": "codex-zsh-aarch64-apple-darwin.tar.gz",
"format": "tar.gz",
"hash": "sha256",
"path": "codex-zsh/bin/zsh"
},
"macos-x86_64": {
"name": "codex-zsh-x86_64-apple-darwin.tar.gz",
"format": "tar.gz",
"hash": "sha256",
"path": "codex-zsh/bin/zsh"
},
"linux-x86_64": {
"name": "codex-zsh-x86_64-unknown-linux-musl.tar.gz",
"format": "tar.gz",
"hash": "sha256",
"path": "codex-zsh/bin/zsh"
},
"linux-aarch64": {
"name": "codex-zsh-aarch64-unknown-linux-musl.tar.gz",
"format": "tar.gz",
"hash": "sha256",
"path": "codex-zsh/bin/zsh"
}
}
@@ -9,6 +9,7 @@ Usage: build-codex-package-archive.sh \
--entrypoint-dir <dir> \
--archive-dir <dir> \
[--bwrap-bin <path>] \
[--zsh-manifest <path>] \
[--codex-command-runner-bin <path>] \
[--codex-windows-sandbox-setup-bin <path>] \
[--target-suffixed-entrypoint]
@@ -48,6 +49,10 @@ while [[ $# -gt 0 ]]; do
bwrap_bin_provided="true"
shift 2
;;
--zsh-manifest)
resource_args+=(--zsh-manifest "${2:?--zsh-manifest requires a value}")
shift 2
;;
--codex-command-runner-bin)
resource_args+=(
--codex-command-runner-bin
+80 -1
View File
@@ -1,15 +1,56 @@
name: rust-release-zsh
on:
workflow_call:
push:
tags:
- "codex-zsh-v*.*.*"
env:
ZSH_COMMIT: 77045ef899e53b9598bebc5a41db93a548a40ca6
ZSH_PATCH: codex-rs/shell-escalation/patches/zsh-exec-wrapper.patch
concurrency:
group: ${{ github.workflow }}::${{ github.ref_name }}
cancel-in-progress: false
jobs:
metadata:
runs-on: ubuntu-latest
outputs:
release_tag: ${{ steps.release_tag.outputs.release_tag }}
steps:
- name: Validate release tag
id: release_tag
env:
RELEASE_TAG: ${{ github.ref_name }}
shell: bash
run: |
set -euo pipefail
if [[ ! "${RELEASE_TAG}" =~ ^codex-zsh-v[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
echo "Tag ${RELEASE_TAG} does not match codex-zsh-vX.Y.Z." >&2
exit 1
fi
echo "release_tag=${RELEASE_TAG}" >> "${GITHUB_OUTPUT}"
- name: Ensure release does not exist
env:
GH_TOKEN: ${{ github.token }}
RELEASE_TAG: ${{ steps.release_tag.outputs.release_tag }}
shell: bash
run: |
set -euo pipefail
if gh release view "${RELEASE_TAG}" --repo "${GITHUB_REPOSITORY}" > /dev/null 2>&1; then
echo "Release ${RELEASE_TAG} already exists; publish changed artifacts under a new tag." >&2
exit 1
fi
linux:
name: Build zsh (Linux) - ${{ matrix.variant }} - ${{ matrix.target }}
needs: metadata
runs-on: ${{ matrix.runner }}
timeout-minutes: 30
container:
@@ -62,6 +103,7 @@ jobs:
darwin:
name: Build zsh (macOS) - ${{ matrix.variant }} - ${{ matrix.target }}
needs: metadata
runs-on: ${{ matrix.runner }}
timeout-minutes: 30
@@ -101,3 +143,40 @@ jobs:
with:
name: codex-zsh-${{ matrix.target }}
path: dist/zsh/${{ matrix.target }}/*
publish-release:
needs:
- metadata
- linux
- darwin
runs-on: ubuntu-latest
permissions:
contents: write
actions: read
steps:
- name: Checkout repository
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
path: dist
- name: Create GitHub Release
uses: softprops/action-gh-release@153bb8e04406b158c6c84fc1615b65b24149a1fe # v2.6.1
with:
tag_name: ${{ needs.metadata.outputs.release_tag }}
name: ${{ needs.metadata.outputs.release_tag }}
files: dist/**
# Keep zsh artifact releases out of Codex's normal "latest release" channel.
prerelease: true
- name: Publish DotSlash manifest
uses: facebook/dotslash-publish-release@9c9ec027515c34db9282a09a25a9cab5880b2c52 # v2
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
tag: ${{ needs.metadata.outputs.release_tag }}
config: .github/dotslash-zsh-config.json
+26 -18
View File
@@ -14,6 +14,9 @@ on:
tags:
- "rust-v*.*.*"
env:
CODEX_ZSH_RELEASE_TAG: codex-zsh-v0.1.0
concurrency:
group: ${{ github.workflow }}
cancel-in-progress: true
@@ -354,6 +357,15 @@ jobs:
cp target/${{ matrix.target }}/release/codex-${{ matrix.target }}.dmg "$dest/codex-${{ matrix.target }}.dmg"
fi
- name: Download packaged zsh manifest
if: ${{ runner.os != 'macOS' }}
shell: bash
run: |
set -euo pipefail
curl -fsSL \
"https://github.com/${GITHUB_REPOSITORY}/releases/download/${CODEX_ZSH_RELEASE_TAG}/codex-zsh" \
-o "${RUNNER_TEMP}/codex-zsh"
- name: Build Codex package archive
if: ${{ runner.os != 'macOS' }}
shell: bash
@@ -366,7 +378,8 @@ jobs:
--target "$TARGET" \
--bundle "$BUNDLE" \
--entrypoint-dir "target/${TARGET}/release" \
--archive-dir "dist/${TARGET}"
--archive-dir "dist/${TARGET}" \
--zsh-manifest "${RUNNER_TEMP}/codex-zsh"
- name: Build Python runtime wheel
if: ${{ matrix.bundle == 'primary' && runner.os != 'macOS' }}
@@ -697,6 +710,14 @@ jobs:
cp "target/${{ matrix.target }}/release/${binary}" "$dest/${binary}-${{ matrix.target }}"
done
- name: Download packaged zsh manifest
shell: bash
run: |
set -euo pipefail
curl -fsSL \
"https://github.com/${GITHUB_REPOSITORY}/releases/download/${CODEX_ZSH_RELEASE_TAG}/codex-zsh" \
-o "${RUNNER_TEMP}/codex-zsh"
- name: Build Codex package archive
shell: bash
env:
@@ -708,7 +729,8 @@ jobs:
--target "$TARGET" \
--bundle "$BUNDLE" \
--entrypoint-dir "target/${TARGET}/release" \
--archive-dir "dist/${TARGET}"
--archive-dir "dist/${TARGET}" \
--zsh-manifest "${RUNNER_TEMP}/codex-zsh"
- name: Build Python runtime wheel
if: ${{ matrix.bundle == 'primary' }}
@@ -1043,11 +1065,6 @@ jobs:
with:
publish: true
zsh-release-assets:
name: zsh release assets
needs: tag-check
uses: ./.github/workflows/rust-release-zsh.yml
release:
needs:
- tag-check
@@ -1055,7 +1072,6 @@ jobs:
- finalize-macos
- build-windows
- argument-comment-lint-release-assets
- zsh-release-assets
if: >-
${{
always() &&
@@ -1063,8 +1079,7 @@ jobs:
needs.build.result == 'success' &&
needs.finalize-macos.result == 'success' &&
needs.build-windows.result == 'success' &&
needs.argument-comment-lint-release-assets.result == 'success' &&
needs.zsh-release-assets.result == 'success'
needs.argument-comment-lint-release-assets.result == 'success'
}}
name: release
runs-on: ubuntu-latest
@@ -1113,7 +1128,7 @@ jobs:
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
path: dist
pattern: "{*-symbols,argument-comment-lint-*,codex-zsh-*,python-runtime-wheel-*}"
pattern: "{*-symbols,argument-comment-lint-*,python-runtime-wheel-*}"
- name: List
run: ls -R dist/
@@ -1236,13 +1251,6 @@ jobs:
tag: ${{ github.ref_name }}
config: .github/dotslash-config.json
- uses: facebook/dotslash-publish-release@9c9ec027515c34db9282a09a25a9cab5880b2c52 # v2
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
tag: ${{ github.ref_name }}
config: .github/dotslash-zsh-config.json
- uses: facebook/dotslash-publish-release@9c9ec027515c34db9282a09a25a9cab5880b2c52 # v2
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+12 -1
View File
@@ -2503,6 +2503,7 @@ dependencies = [
"codex-code-mode-protocol",
"codex-protocol",
"deno_core_icudata",
"futures",
"pretty_assertions",
"serde_json",
"tokio",
@@ -2514,6 +2515,14 @@ dependencies = [
[[package]]
name = "codex-code-mode-host"
version = "0.0.0"
dependencies = [
"anyhow",
"codex-code-mode",
"codex-code-mode-protocol",
"pretty_assertions",
"tokio",
"tokio-util",
]
[[package]]
name = "codex-code-mode-protocol"
@@ -2557,6 +2566,7 @@ dependencies = [
"multimap",
"pretty_assertions",
"prost 0.14.3",
"regex-lite",
"schemars 0.8.22",
"serde",
"serde_ignored",
@@ -3009,6 +3019,7 @@ dependencies = [
"codex-tools",
"codex-utils-absolute-path",
"pretty_assertions",
"serde_json",
"tokio",
]
@@ -3356,6 +3367,7 @@ name = "codex-mcp-extension"
version = "0.0.0"
dependencies = [
"codex-config",
"codex-connectors-extension",
"codex-core",
"codex-core-plugins",
"codex-exec-server",
@@ -3757,7 +3769,6 @@ dependencies = [
"base64 0.22.1",
"bytes",
"codex-api",
"codex-client",
"codex-config",
"codex-exec-server",
"codex-keyring-store",
+1
View File
@@ -166,6 +166,7 @@ codex-code-mode-protocol = { path = "code-mode-protocol" }
codex-home = { path = "codex-home" }
codex-config = { path = "config" }
codex-connectors = { path = "connectors" }
codex-connectors-extension = { path = "ext/connectors" }
codex-context-fragments = { path = "context-fragments" }
codex-core = { path = "core" }
codex-core-api = { path = "core-api" }
@@ -1619,6 +1619,12 @@
"null"
]
},
"threadId": {
"type": [
"string",
"null"
]
},
"timeoutSecs": {
"format": "int64",
"type": [
@@ -4168,6 +4174,7 @@
"type": "object"
},
"ThreadRollbackParams": {
"description": "DEPRECATED: `thread/rollback` will be removed soon.",
"properties": {
"numTurns": {
"description": "The number of turns to drop from the end of the thread. Must be >= 1.\n\nThis only modifies the thread's history and does not revert local file changes that have been made by the agent. Clients are responsible for reverting these changes.",
@@ -2370,6 +2370,12 @@
},
"success": {
"type": "boolean"
},
"threadId": {
"type": [
"string",
"null"
]
}
},
"required": [
@@ -2378,6 +2384,12 @@
],
"type": "object"
},
"McpServerStartupFailureReason": {
"enum": [
"reauthenticationRequired"
],
"type": "string"
},
"McpServerStartupState": {
"enum": [
"starting",
@@ -2395,6 +2407,16 @@
"null"
]
},
"failureReason": {
"anyOf": [
{
"$ref": "#/definitions/McpServerStartupFailureReason"
},
{
"type": "null"
}
]
},
"name": {
"type": "string"
},
@@ -12026,6 +12026,12 @@
},
"success": {
"type": "boolean"
},
"threadId": {
"type": [
"string",
"null"
]
}
},
"required": [
@@ -12050,6 +12056,12 @@
"null"
]
},
"threadId": {
"type": [
"string",
"null"
]
},
"timeoutSecs": {
"format": "int64",
"type": [
@@ -12082,6 +12094,12 @@
"title": "McpServerRefreshResponse",
"type": "object"
},
"McpServerStartupFailureReason": {
"enum": [
"reauthenticationRequired"
],
"type": "string"
},
"McpServerStartupState": {
"enum": [
"starting",
@@ -12153,6 +12171,16 @@
"null"
]
},
"failureReason": {
"anyOf": [
{
"$ref": "#/definitions/v2/McpServerStartupFailureReason"
},
{
"type": "null"
}
]
},
"name": {
"type": "string"
},
@@ -18985,6 +19013,7 @@
},
"ThreadRollbackParams": {
"$schema": "http://json-schema.org/draft-07/schema#",
"description": "DEPRECATED: `thread/rollback` will be removed soon.",
"properties": {
"numTurns": {
"description": "The number of turns to drop from the end of the thread. Must be >= 1.\n\nThis only modifies the thread's history and does not revert local file changes that have been made by the agent. Clients are responsible for reverting these changes.",
@@ -8430,6 +8430,12 @@
},
"success": {
"type": "boolean"
},
"threadId": {
"type": [
"string",
"null"
]
}
},
"required": [
@@ -8454,6 +8460,12 @@
"null"
]
},
"threadId": {
"type": [
"string",
"null"
]
},
"timeoutSecs": {
"format": "int64",
"type": [
@@ -8486,6 +8498,12 @@
"title": "McpServerRefreshResponse",
"type": "object"
},
"McpServerStartupFailureReason": {
"enum": [
"reauthenticationRequired"
],
"type": "string"
},
"McpServerStartupState": {
"enum": [
"starting",
@@ -8557,6 +8575,16 @@
"null"
]
},
"failureReason": {
"anyOf": [
{
"$ref": "#/definitions/McpServerStartupFailureReason"
},
{
"type": "null"
}
]
},
"name": {
"type": "string"
},
@@ -16764,6 +16792,7 @@
},
"ThreadRollbackParams": {
"$schema": "http://json-schema.org/draft-07/schema#",
"description": "DEPRECATED: `thread/rollback` will be removed soon.",
"properties": {
"numTurns": {
"description": "The number of turns to drop from the end of the thread. Must be >= 1.\n\nThis only modifies the thread's history and does not revert local file changes that have been made by the agent. Clients are responsible for reverting these changes.",
@@ -12,6 +12,12 @@
},
"success": {
"type": "boolean"
},
"threadId": {
"type": [
"string",
"null"
]
}
},
"required": [
@@ -13,6 +13,12 @@
"null"
]
},
"threadId": {
"type": [
"string",
"null"
]
},
"timeoutSecs": {
"format": "int64",
"type": [
@@ -1,6 +1,12 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"definitions": {
"McpServerStartupFailureReason": {
"enum": [
"reauthenticationRequired"
],
"type": "string"
},
"McpServerStartupState": {
"enum": [
"starting",
@@ -18,6 +24,16 @@
"null"
]
},
"failureReason": {
"anyOf": [
{
"$ref": "#/definitions/McpServerStartupFailureReason"
},
{
"type": "null"
}
]
},
"name": {
"type": "string"
},
@@ -1,5 +1,6 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"description": "DEPRECATED: `thread/rollback` will be removed soon.",
"properties": {
"numTurns": {
"description": "The number of turns to drop from the end of the thread. Must be >= 1.\n\nThis only modifies the thread's history and does not revert local file changes that have been made by the agent. Clients are responsible for reverting these changes.",
@@ -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 McpServerOauthLoginCompletedNotification = { name: string, success: boolean, error?: string, };
export type McpServerOauthLoginCompletedNotification = { name: string, threadId: string | null, success: boolean, error?: string, };
@@ -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 McpServerOauthLoginParams = { name: string, scopes?: Array<string> | null, timeoutSecs?: bigint | null, };
export type McpServerOauthLoginParams = { name: string, threadId?: string | null, scopes?: Array<string> | null, timeoutSecs?: bigint | null, };
@@ -0,0 +1,5 @@
// GENERATED CODE! DO NOT MODIFY BY HAND!
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
export type McpServerStartupFailureReason = "reauthenticationRequired";
@@ -1,6 +1,7 @@
// 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 { McpServerStartupFailureReason } from "./McpServerStartupFailureReason";
import type { McpServerStartupState } from "./McpServerStartupState";
export type McpServerStatusUpdatedNotification = { threadId: string | null, name: string, status: McpServerStartupState, error: string | null, };
export type McpServerStatusUpdatedNotification = { threadId: string | null, name: string, status: McpServerStartupState, error: string | null, failureReason: McpServerStartupFailureReason | null, };
@@ -2,6 +2,9 @@
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
/**
* DEPRECATED: `thread/rollback` will be removed soon.
*/
export type ThreadRollbackParams = { threadId: string,
/**
* The number of turns to drop from the end of the thread. Must be >= 1.
@@ -238,6 +238,7 @@ export type { McpServerOauthLoginCompletedNotification } from "./McpServerOauthL
export type { McpServerOauthLoginParams } from "./McpServerOauthLoginParams";
export type { McpServerOauthLoginResponse } from "./McpServerOauthLoginResponse";
export type { McpServerRefreshResponse } from "./McpServerRefreshResponse";
export type { McpServerStartupFailureReason } from "./McpServerStartupFailureReason";
export type { McpServerStartupState } from "./McpServerStartupState";
export type { McpServerStatus } from "./McpServerStatus";
export type { McpServerStatusDetail } from "./McpServerStatusDetail";
@@ -1913,6 +1913,7 @@ mod tests {
request_id: request_id(),
params: v2::McpServerOauthLoginParams {
name: "server-a".to_string(),
thread_id: None,
scopes: None,
timeout_secs: None,
},
@@ -23,6 +23,12 @@ v2_enum_from_core!(
}
);
v2_enum_from_core!(
pub enum McpServerStartupFailureReason from codex_protocol::protocol::McpStartupFailureReason {
ReauthenticationRequired
}
);
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
@@ -186,6 +192,8 @@ pub struct McpServerRefreshResponse {}
#[ts(export_to = "v2/")]
pub struct McpServerOauthLoginParams {
pub name: String,
#[ts(optional = nullable)]
pub thread_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[ts(optional = nullable)]
pub scopes: Option<Vec<String>>,
@@ -215,6 +223,7 @@ pub struct McpToolCallProgressNotification {
#[ts(export_to = "v2/")]
pub struct McpServerOauthLoginCompletedNotification {
pub name: String,
pub thread_id: Option<String>,
pub success: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[ts(optional)]
@@ -239,6 +248,7 @@ pub struct McpServerStatusUpdatedNotification {
pub name: String,
pub status: McpServerStartupState,
pub error: Option<String>,
pub failure_reason: Option<McpServerStartupFailureReason>,
}
#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, JsonSchema, TS)]
@@ -1,4 +1,5 @@
use super::*;
use crate::ServerNotification;
use codex_protocol::approvals::ElicitationRequest as CoreElicitationRequest;
use codex_protocol::config_types::MultiAgentMode;
use codex_protocol::items::AgentMessageContent;
@@ -2155,6 +2156,7 @@ fn mcp_server_status_updated_accepts_missing_thread_id() {
name: "optional_broken".to_string(),
status: McpServerStartupState::Failed,
error: Some("handshake failed".to_string()),
failure_reason: None,
};
assert_eq!(notification, expected);
assert_eq!(
@@ -2164,6 +2166,33 @@ fn mcp_server_status_updated_accepts_missing_thread_id() {
"name": "optional_broken",
"status": "failed",
"error": "handshake failed",
"failureReason": null,
})
);
}
#[test]
fn mcp_server_status_updated_serializes_failure_reason() {
let notification =
ServerNotification::McpServerStatusUpdated(McpServerStatusUpdatedNotification {
thread_id: Some("thread-1".to_string()),
name: "expired-oauth".to_string(),
status: McpServerStartupState::Failed,
error: Some("OAuth credentials expired".to_string()),
failure_reason: Some(McpServerStartupFailureReason::ReauthenticationRequired),
});
assert_eq!(
serde_json::to_value(notification).expect("notification should serialize"),
json!({
"method": "mcpServer/startupStatus/updated",
"params": {
"threadId": "thread-1",
"name": "expired-oauth",
"status": "failed",
"error": "OAuth credentials expired",
"failureReason": "reauthenticationRequired",
},
})
);
}
@@ -57,6 +57,11 @@ pub struct ThreadStartParams {
pub model: Option<String>,
#[ts(optional = nullable)]
pub model_provider: Option<String>,
/// Allow a provider with an authoritative static model catalog to replace an unavailable
/// requested model with its default.
#[experimental("thread/start.allowProviderModelFallback")]
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
pub allow_provider_model_fallback: bool,
#[serde(
default,
deserialize_with = "crate::protocol::serde_helpers::deserialize_double_option",
@@ -1025,6 +1030,7 @@ pub struct ThreadBackgroundTerminalsTerminateResponse {
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
/// DEPRECATED: `thread/rollback` will be removed soon.
pub struct ThreadRollbackParams {
pub thread_id: String,
/// The number of turns to drop from the end of the thread. Must be >= 1.
+6 -6
View File
@@ -137,7 +137,7 @@ Example with notification opt-out:
## API Overview
- `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. For permissions, prefer experimental `permissions` profile selection by id; the legacy `sandbox` shorthand is still accepted but cannot be combined with `permissions`. Deprecated experimental `multiAgentMode` is ignored; use Ultra reasoning effort for proactive multi-agent behavior. 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 using environment-native absolute paths. Skills found below those roots are listed and read through the owning environment. Stdio MCP servers declared by selected plugins are also started in that environment; HTTP MCP declarations remain inactive.
- `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. Experimental `allowProviderModelFallback` lets providers backed by an authoritative static model catalog replace an unavailable requested `model` with the catalog default; dynamic or cached catalogs preserve the requested model. 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`. Deprecated experimental `multiAgentMode` is ignored; use Ultra reasoning effort for proactive multi-agent behavior. 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 using environment-native absolute paths. Skills found below those roots are listed and read through the owning environment. Stdio MCP servers declared by selected plugins are started in that environment, and HTTP MCP connections use that environment's HTTP client.
- `thread/resume` — reopen an existing thread by id so subsequent `turn/start` calls append to it. Accepts the same permission override rules as `thread/start`.
- `thread/fork` — fork an existing thread into a new thread id by copying the stored history; if the source thread is currently mid-turn, the fork records the same interruption marker as `turn/interrupt` instead of inheriting an unmarked partial turn suffix. The returned `thread.forkedFromId` points at the source thread when known. Accepts `ephemeral: true` for an in-memory temporary fork, emits `thread/started` (including the current `thread.status`), and auto-subscribes you to turn/item events for the new thread. Experimental clients can pass `excludeTurns: true` when they plan to page fork history via `thread/turns/list` instead of receiving the full turn array immediately. Accepts the same permission override rules as `thread/start`.
- `thread/start`, `thread/resume`, and `thread/fork` responses include the legacy `sandbox` compatibility projection. `instructionSources` lists loaded instruction files using each source environment's native absolute path syntax, including files loaded from remote environments. 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. Their deprecated experimental `multiAgentMode` field, and the corresponding thread setting, always report `explicitRequestOnly`; Ultra reasoning effort is the source of proactive multi-agent behavior.
@@ -167,7 +167,7 @@ Example with notification opt-out:
- `thread/backgroundTerminals/clean` — terminate all running background terminals for a thread (experimental; requires `capabilities.experimentalApi`); returns `{}` when the cleanup request is accepted.
- `thread/backgroundTerminals/list` — list running background terminals for a loaded thread (experimental; requires `capabilities.experimentalApi`); returns `data` with the running terminal ids.
- `thread/backgroundTerminals/terminate` — terminate one running background terminal by app-server `processId` (experimental; requires `capabilities.experimentalApi`); returns whether a process was terminated.
- `thread/rollback` — drop the last N turns from the agent’s in-memory context and persist a rollback marker in the rollout so future resumes see the pruned history; returns the updated `thread` (with `turns` populated) on success.
- `thread/rollback` — deprecated and will be removed soon. Drop the last N turns from the agent’s in-memory context and persist a rollback marker in the rollout so future resumes see the pruned history; returns the updated `thread` (with `turns` populated) on success.
- `turn/start` — add user input to a thread and begin Codex generation; responds with the initial `turn` object and streams `turn/started`, `item/*`, and `turn/completed` notifications. `clientUserMessageId` is optional; when supplied, the corresponding `userMessage` item echoes it as `clientId`. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. Prefer experimental `permissions` profile selection by id for permission overrides; the legacy `sandboxPolicy` field is still accepted but cannot be combined with `permissions`. For `collaborationMode`, `settings.developer_instructions: null` means "use built-in instructions for the selected mode". Deprecated experimental `multiAgentMode` is ignored; Ultra reasoning effort selects proactive behavior.
- `thread/inject_items` — append raw Responses API items to a loaded thread’s model-visible history without starting a user turn; returns `{}` on success.
- `turn/steer` — add user input to an already in-flight regular turn without starting a new turn; returns the active `turnId` that accepted the input. `clientUserMessageId` is optional; when supplied, the corresponding `userMessage` item echoes it as `clientId`. Review and manual compaction turns reject `turn/steer`.
@@ -229,7 +229,7 @@ Example with notification opt-out:
- `skills/config/write` — write user-level skill config by name or absolute path.
- `plugin/install` — install a plugin from a discovered marketplace entry, rejecting marketplace entries marked unavailable for install, install MCPs if any, and return the effective plugin auth policy plus any apps that still need auth (**under development; do not call from production clients yet**).
- `plugin/uninstall` — uninstall a local plugin by `pluginId` in `<plugin>@<marketplace>` form by removing its cached files and clearing its user-level config entry, or uninstall a remote ChatGPT plugin by backend `pluginId` by forwarding the uninstall to the ChatGPT plugin backend and removing any downloaded remote-plugin cache (**under development; do not call from production clients yet**).
- `mcpServer/oauth/login` — start an OAuth login for a configured MCP server; returns an `authorization_url` and later emits `mcpServer/oauthLogin/completed` once the browser flow finishes.
- `mcpServer/oauth/login` — start an OAuth login for a configured MCP server; pass `threadId` to resolve servers from that thread's selected plugins and executor, and receive an `authorization_url` followed by `mcpServer/oauthLogin/completed` once the browser flow finishes.
- `tool/requestUserInput` — prompt the user with 1–3 short questions for a tool call and return their answers (experimental).
- `config/mcpServer/reload` — reload MCP server config from disk and queue a refresh for loaded threads (applied on each thread's next active turn); returns `{}`. Use this after editing `config.toml` without restarting the server.
- `mcpServerStatus/list` — enumerate configured MCP servers with their tools, auth status, server info, plus resources/resource templates for `full` detail; supports optional `threadId` and cursor+limit pagination. If `threadId` is omitted, the server reads from the latest global config directly. If `detail` is omitted, the server defaults to `full`.
@@ -1345,7 +1345,7 @@ Because audio is intentionally separate from `ThreadItem`, clients can opt out o
### MCP server startup events
- `mcpServer/startupStatus/updated` — `{ threadId, name, status, error }` when app-server observes an MCP server startup transition. `threadId` identifies the owning thread when startup is thread-scoped and is `null` when startup is app-scoped. `status` is one of `starting`, `ready`, `failed`, or `cancelled`. `error` is `null` except for `failed`.
- `mcpServer/startupStatus/updated` — `{ threadId, name, status, error, failureReason }` when app-server observes an MCP server startup transition. `threadId` identifies the owning thread when startup is thread-scoped and is `null` when startup is app-scoped. `status` is one of `starting`, `ready`, `failed`, or `cancelled`. `error` and `failureReason` are `null` except for `failed`; `failureReason` is `reauthenticationRequired` when stored OAuth credentials have expired and cannot be refreshed, so clients can prompt the user to reconnect the named server.
### Turn events
@@ -1934,8 +1934,8 @@ Codex supports these authentication modes. The current mode is surfaced in `acco
- `account/workspaceMessages/read` — fetch active workspace messages, including workspace notification headlines when available.
- `account/rateLimits/updated` (notify) — emitted whenever a user's ChatGPT rate limits change. This is a sparse rolling update; merge available values into the most recent `account/rateLimits/read` response or refetch that snapshot.
- `account/sendAddCreditsNudgeEmail` — ask ChatGPT to email the workspace owner about depleted credits or a reached usage limit.
- `mcpServer/oauthLogin/completed` (notify) — emitted after a `mcpServer/oauth/login` flow finishes for a server; payload includes `{ name, success, error? }`.
- `mcpServer/startupStatus/updated` (notify) — emitted when a configured MCP server's startup status changes; payload includes `{ threadId, name, status, error }`, where `threadId` is the owning thread when startup is thread-scoped and `null` when it is app-scoped, and `status` is `starting`, `ready`, `failed`, or `cancelled`.
- `mcpServer/oauthLogin/completed` (notify) — emitted after a `mcpServer/oauth/login` flow finishes for a server; payload includes `{ name, threadId, success, error? }`.
- `mcpServer/startupStatus/updated` (notify) — emitted when a configured MCP server's startup status changes; payload includes `{ threadId, name, status, error, failureReason }`, where `threadId` is the owning thread when startup is thread-scoped and `null` when it is app-scoped, and `status` is `starting`, `ready`, `failed`, or `cancelled`. `failureReason` is `reauthenticationRequired` when stored OAuth credentials have expired and cannot be refreshed, so clients can prompt the user to reconnect the named server.
### 1) Check auth state
@@ -200,18 +200,20 @@ pub(crate) async fn apply_bespoke_event_handling(
.await;
}
EventMsg::McpStartupUpdate(update) => {
let (status, error) = match update.status {
let (status, error, failure_reason) = match update.status {
codex_protocol::protocol::McpStartupStatus::Starting => {
(McpServerStartupState::Starting, None)
(McpServerStartupState::Starting, None, None)
}
codex_protocol::protocol::McpStartupStatus::Ready => {
(McpServerStartupState::Ready, None)
}
codex_protocol::protocol::McpStartupStatus::Failed { error } => {
(McpServerStartupState::Failed, Some(error))
(McpServerStartupState::Ready, None, None)
}
codex_protocol::protocol::McpStartupStatus::Failed { error, reason } => (
McpServerStartupState::Failed,
Some(error),
reason.map(Into::into),
),
codex_protocol::protocol::McpStartupStatus::Cancelled => {
(McpServerStartupState::Cancelled, None)
(McpServerStartupState::Cancelled, None, None)
}
};
let notification = McpServerStatusUpdatedNotification {
@@ -219,6 +221,7 @@ pub(crate) async fn apply_bespoke_event_handling(
name: update.server,
status,
error,
failure_reason,
};
outgoing
.send_server_notification(ServerNotification::McpServerStatusUpdated(notification))
+26 -3
View File
@@ -23,6 +23,7 @@ use crate::outgoing_message::OutgoingMessageSender;
use crate::thread_state::ThreadStateManager;
const CURRENT_TIME_REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
const CURRENT_TIME_POLL_INTERVAL: Duration = Duration::from_secs(1);
pub(crate) fn app_server_time_provider(
outgoing: Arc<OutgoingMessageSender>,
@@ -51,10 +52,32 @@ impl TimeProvider for AppServerTimeProvider {
})
}
fn sleep(&self, _thread_id: ThreadId, duration: Duration) -> SleepFuture<'_> {
fn sleep(&self, thread_id: ThreadId, duration: Duration) -> SleepFuture<'_> {
let outgoing = self.outgoing.clone();
let thread_state_manager = self.thread_state_manager.clone();
Box::pin(async move {
tokio::time::sleep(duration).await;
Ok(())
let outgoing = outgoing
.upgrade()
.context("app-server current-time provider is unavailable")?;
let started_at =
request_current_time(outgoing.clone(), thread_state_manager.clone(), thread_id)
.await?;
let wake_at = started_at
.checked_add_signed(
chrono::Duration::from_std(duration)
.context("external sleep duration is outside the supported range")?,
)
.context("external sleep deadline is outside the supported range")?;
loop {
tokio::time::sleep(CURRENT_TIME_POLL_INTERVAL).await;
if request_current_time(outgoing.clone(), thread_state_manager.clone(), thread_id)
.await?
>= wake_at
{
return Ok(());
}
}
})
}
}
+1 -1
View File
@@ -1205,7 +1205,7 @@ impl MessageProcessor {
}
ClientRequest::ThreadRollback { params, .. } => {
self.thread_processor
.thread_rollback(&request_id, params)
.thread_rollback(&request_id, params, app_server_client_name.as_deref())
.await
}
ClientRequest::ThreadList { params, .. } => {
@@ -56,6 +56,7 @@ use codex_app_server_protocol::ConsumeAccountRateLimitResetCreditParams;
use codex_app_server_protocol::ConsumeAccountRateLimitResetCreditResponse;
use codex_app_server_protocol::ConversationGitInfo;
use codex_app_server_protocol::ConversationSummary;
use codex_app_server_protocol::DeprecationNoticeNotification;
use codex_app_server_protocol::DynamicToolFunctionSpec;
use codex_app_server_protocol::DynamicToolNamespaceTool;
use codex_app_server_protocol::DynamicToolSpec;
@@ -380,7 +381,7 @@ use codex_mcp::McpRuntimeContext;
use codex_mcp::McpServerStatusSnapshot;
use codex_mcp::McpSnapshotDetail;
use codex_mcp::collect_mcp_server_status_snapshot_with_detail;
use codex_mcp::discover_supported_scopes;
use codex_mcp::discover_supported_scopes_with_http_client;
use codex_mcp::read_mcp_resource as read_mcp_resource_without_thread;
use codex_mcp::resolve_oauth_scopes;
use codex_memories_write::clear_memory_roots_contents;
@@ -428,7 +429,7 @@ use codex_protocol::protocol::USER_MESSAGE_BEGIN;
use codex_protocol::protocol::W3cTraceContext;
use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS;
use codex_protocol::user_input::UserInput as CoreInputItem;
use codex_rmcp_client::perform_oauth_login_return_url;
use codex_rmcp_client::perform_oauth_login_return_url_with_http_client;
use codex_rollout::is_persisted_rollout_item;
use codex_rollout::state_db::StateDbHandle;
use codex_rollout::state_db::reconcile_rollout;
@@ -184,7 +184,11 @@ impl ExternalAgentSessionImporter {
.map_err(|err| format!("failed to load imported session config: {err}"))?;
let models_manager = self.thread_manager.get_models_manager();
let model = models_manager
.get_default_model(&config.model, RefreshStrategy::Offline)
.get_default_model(
&config.model,
/*allow_provider_model_fallback*/ false,
RefreshStrategy::Offline,
)
.await;
let model_info = models_manager
.get_model_info(model.as_str(), &config.to_models_manager_config())
@@ -215,6 +219,7 @@ impl ExternalAgentSessionImporter {
.unwrap_or_else(|| model_info.get_model_instructions(config.personality)),
},
dynamic_tools: Vec::new(),
selected_capability_roots: Vec::new(),
multi_agent_version: Some(MultiAgentVersion::V1),
initial_window_id: uuid::Uuid::now_v7().to_string(),
metadata: ThreadPersistenceMetadata {
@@ -113,19 +113,35 @@ impl McpRequestProcessor {
&self,
params: McpServerOauthLoginParams,
) -> Result<McpServerOauthLoginResponse, JSONRPCErrorError> {
let config = self.load_latest_config(/*fallback_cwd*/ None).await?;
let McpServerOauthLoginParams {
name,
thread_id,
scopes,
timeout_secs,
} = params;
let auth = self.auth_manager.auth().await;
let effective_servers = self
.thread_manager
.mcp_manager()
.effective_servers(&config, auth.as_ref())
.await;
let (mcp_config, runtime_context) = match thread_id.as_deref() {
Some(thread_id) => {
let (_, thread) = self.load_thread(thread_id).await?;
let runtime = thread.current_mcp_runtime().await;
(runtime.config().clone(), runtime.runtime_context().clone())
}
None => {
let config = self.load_latest_config(/*fallback_cwd*/ None).await?;
let mcp_config = self
.thread_manager
.mcp_manager()
.runtime_config(&config)
.await;
let runtime_context = McpRuntimeContext::new(
self.thread_manager.environment_manager(),
config.cwd.to_path_buf(),
);
(mcp_config, runtime_context)
}
};
let effective_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref());
let Some(server) = effective_servers
.get(&name)
.and_then(codex_mcp::EffectiveMcpServer::configured_config)
@@ -149,32 +165,41 @@ impl McpRequestProcessor {
}
};
let http_client = runtime_context
.resolve_http_client(&name, server)
.map_err(|err| {
internal_error(format!("failed to resolve MCP server runtime: {err}"))
})?;
let discovered_scopes = if scopes.is_none() && server.scopes.is_none() {
discover_supported_scopes(&server.transport).await
discover_supported_scopes_with_http_client(&server.transport, Arc::clone(&http_client))
.await
} else {
None
};
let resolved_scopes =
resolve_oauth_scopes(scopes, server.scopes.clone(), discovered_scopes);
let handle = perform_oauth_login_return_url(
let handle = perform_oauth_login_return_url_with_http_client(
&name,
&url,
config.mcp_oauth_credentials_store_mode,
config.auth_keyring_backend_kind(),
mcp_config.mcp_oauth_credentials_store_mode,
mcp_config.auth_keyring_backend_kind,
http_headers,
env_http_headers,
&resolved_scopes.scopes,
server.oauth_client_id(),
server.oauth_resource.as_deref(),
timeout_secs,
config.mcp_oauth_callback_port,
config.mcp_oauth_callback_url.as_deref(),
mcp_config.mcp_oauth_callback_port,
mcp_config.mcp_oauth_callback_url.as_deref(),
http_client,
)
.await
.map_err(|err| internal_error(format!("failed to login to MCP server '{name}': {err}")))?;
let authorization_url = handle.authorization_url().to_string();
let notification_name = name.clone();
let notification_thread_id = thread_id;
let outgoing = Arc::clone(&self.outgoing);
tokio::spawn(async move {
@@ -186,6 +211,7 @@ impl McpRequestProcessor {
let notification = ServerNotification::McpServerOauthLoginCompleted(
McpServerOauthLoginCompletedNotification {
name: notification_name,
thread_id: notification_thread_id,
success,
error,
},
@@ -217,22 +243,24 @@ impl McpRequestProcessor {
}
None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None),
};
let mcp_config = match thread {
Some(thread) => thread.runtime_mcp_config(&config).await,
let mcp_manager = self.thread_manager.mcp_manager();
let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache();
let auth = self.auth_manager.auth().await;
let (mcp_config, runtime_context) = match thread {
Some(thread) => {
let mcp_config = thread.runtime_mcp_config(&config).await;
let runtime = thread.current_mcp_runtime().await;
(mcp_config, runtime.runtime_context().clone())
}
None => {
self.thread_manager
.mcp_manager()
.runtime_config(&config)
.await
let mcp_config = mcp_manager.runtime_config(&config).await;
let runtime_context = McpRuntimeContext::new(
self.thread_manager.environment_manager(),
config.cwd.to_path_buf(),
);
(mcp_config, runtime_context)
}
};
let auth = self.auth_manager.auth().await;
let environment_manager = self.thread_manager.environment_manager();
// This status path has no turn-selected environment. Use config cwd
// as the local stdio fallback; named environment stdio MCPs must
// declare their own absolute cwd.
let runtime_context =
McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf());
tokio::spawn(async move {
Self::list_mcp_server_status_task(
@@ -242,6 +270,7 @@ impl McpRequestProcessor {
mcp_config,
auth,
runtime_context,
codex_apps_tools_cache,
)
.await;
});
@@ -255,6 +284,7 @@ impl McpRequestProcessor {
mcp_config: codex_mcp::McpConfig,
auth: Option<CodexAuth>,
runtime_context: McpRuntimeContext,
codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache,
) {
let result = Self::list_mcp_server_status_response(
request_id.request_id.to_string(),
@@ -262,6 +292,7 @@ impl McpRequestProcessor {
mcp_config,
auth,
runtime_context,
codex_apps_tools_cache,
)
.await;
outgoing.send_result(request_id, result).await;
@@ -273,6 +304,7 @@ impl McpRequestProcessor {
mcp_config: codex_mcp::McpConfig,
auth: Option<CodexAuth>,
runtime_context: McpRuntimeContext,
codex_apps_tools_cache: codex_mcp::CodexAppsToolsCache,
) -> Result<ListMcpServerStatusResponse, JSONRPCErrorError> {
let detail = match params.detail.unwrap_or(McpServerStatusDetail::Full) {
McpServerStatusDetail::Full => McpSnapshotDetail::Full,
@@ -284,6 +316,7 @@ impl McpRequestProcessor {
auth.as_ref(),
request_id,
runtime_context,
codex_apps_tools_cache,
detail,
)
.await;
@@ -374,11 +407,9 @@ impl McpRequestProcessor {
}
let config = self.load_latest_config(/*fallback_cwd*/ None).await?;
let mcp_config = self
.thread_manager
.mcp_manager()
.runtime_config(&config)
.await;
let mcp_manager = self.thread_manager.mcp_manager();
let mcp_config = mcp_manager.runtime_config(&config).await;
let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache();
let auth = self.auth_manager.auth().await;
let environment_manager = self.thread_manager.environment_manager();
// This threadless resource-read path has no turn cwd or turn-selected
@@ -393,6 +424,7 @@ impl McpRequestProcessor {
&mcp_config,
auth.as_ref(),
runtime_context,
codex_apps_tools_cache,
&server,
&uri,
)
@@ -1883,6 +1883,7 @@ impl PluginRequestProcessor {
let notification = ServerNotification::McpServerOauthLoginCompleted(
McpServerOauthLoginCompletedNotification {
name: notification_name,
thread_id: None,
success,
error,
},
@@ -8,6 +8,9 @@ use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE;
const THREAD_LIST_DEFAULT_LIMIT: usize = 25;
const THREAD_LIST_MAX_LIMIT: usize = 100;
const CODEX_TUI_CLIENT_NAME: &str = "codex-tui";
const THREAD_ROLLBACK_DEPRECATION_SUMMARY: &str =
"thread/rollback is deprecated and will be removed soon";
struct ThreadListFilters {
model_providers: Option<Vec<String>>,
@@ -632,12 +635,29 @@ impl ThreadRequestProcessor {
&self,
request_id: &ConnectionRequestId,
params: ThreadRollbackParams,
app_server_client_name: Option<&str>,
) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
if app_server_client_name != Some(CODEX_TUI_CLIENT_NAME) {
self.send_thread_rollback_deprecation_notice(request_id.connection_id)
.await;
}
self.thread_rollback_inner(request_id, params)
.await
.map(|()| None)
}
async fn send_thread_rollback_deprecation_notice(&self, connection_id: ConnectionId) {
self.outgoing
.send_server_notification_to_connections(
&[connection_id],
ServerNotification::DeprecationNotice(DeprecationNoticeNotification {
summary: THREAD_ROLLBACK_DEPRECATION_SUMMARY.to_string(),
details: None,
}),
)
.await;
}
pub(crate) async fn thread_list(
&self,
params: ThreadListParams,
@@ -887,6 +907,7 @@ impl ThreadRequestProcessor {
let ThreadStartParams {
model,
model_provider,
allow_provider_model_fallback,
service_tier,
cwd,
runtime_workspace_roots,
@@ -963,6 +984,7 @@ impl ThreadRequestProcessor {
thread_source.map(Into::into),
environment_selections,
service_name,
allow_provider_model_fallback,
experimental_raw_events,
request_trace,
)
@@ -1037,6 +1059,7 @@ impl ThreadRequestProcessor {
thread_source: Option<codex_protocol::protocol::ThreadSource>,
environments: Option<Vec<TurnEnvironmentSelection>>,
service_name: Option<String>,
allow_provider_model_fallback: bool,
experimental_raw_events: bool,
request_trace: Option<W3cTraceContext>,
) -> Result<(), JSONRPCErrorError> {
@@ -1132,7 +1155,6 @@ impl ThreadRequestProcessor {
let mut thread_extension_init = ExtensionDataInit::new();
if !selected_capability_roots.is_empty() {
thread_extension_init.insert(selected_capability_roots);
codex_mcp_extension::initialize_executor_plugin_thread_data(&mut thread_extension_init);
}
let create_thread_started_at = std::time::Instant::now();
let NewThread {
@@ -1144,6 +1166,7 @@ impl ThreadRequestProcessor {
.thread_manager
.start_thread_with_options(StartThreadOptions {
config,
allow_provider_model_fallback,
initial_history: match session_start_source
.unwrap_or(codex_app_server_protocol::ThreadStartSource::Startup)
{
@@ -200,6 +200,7 @@ fn create_fake_rollout_with_source_and_parent_thread_id(
model_provider: model_provider.map(str::to_string),
base_instructions: None,
dynamic_tools: None,
selected_capability_roots: Vec::new(),
memory_mode: None,
multi_agent_version: None,
context_window: None,
@@ -288,6 +289,7 @@ pub fn create_fake_rollout_with_text_elements(
model_provider: model_provider.map(str::to_string),
base_instructions: None,
dynamic_tools: None,
selected_capability_roots: Vec::new(),
memory_mode: None,
multi_agent_version: None,
context_window: None,
@@ -131,6 +131,7 @@ async fn get_conversation_summary_by_thread_id_reads_pathless_store_thread() ->
originator: "test_originator".to_string(),
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
selected_capability_roots: Vec::new(),
multi_agent_version: None,
initial_window_id: Uuid::now_v7().to_string(),
metadata: ThreadPersistenceMetadata {
@@ -2,9 +2,16 @@ use anyhow::Result;
use app_test_support::TestAppServer;
use app_test_support::to_response;
use app_test_support::write_mock_responses_config_toml;
use axum::Json;
use axum::Router;
use axum::body::Bytes;
use axum::routing::get;
use axum::routing::post;
use codex_app_server_protocol::CapabilityRootLocation;
use codex_app_server_protocol::ListMcpServerStatusParams;
use codex_app_server_protocol::ListMcpServerStatusResponse;
use codex_app_server_protocol::McpServerOauthLoginCompletedNotification;
use codex_app_server_protocol::McpServerOauthLoginResponse;
use codex_app_server_protocol::McpServerToolCallParams;
use codex_app_server_protocol::McpServerToolCallResponse;
use codex_app_server_protocol::RequestId;
@@ -17,14 +24,36 @@ use codex_utils_path_uri::PathUri;
use core_test_support::responses;
use core_test_support::stdio_server_bin;
use pretty_assertions::assert_eq;
use rmcp::handler::server::ServerHandler;
use rmcp::model::CallToolRequestParams;
use rmcp::model::CallToolResult;
use rmcp::model::JsonObject;
use rmcp::model::ListToolsResult;
use rmcp::model::ServerCapabilities;
use rmcp::model::ServerInfo;
use rmcp::model::Tool;
use rmcp::model::ToolAnnotations;
use rmcp::service::RequestContext;
use rmcp::service::RoleServer;
use rmcp::transport::StreamableHttpServerConfig;
use rmcp::transport::StreamableHttpService;
use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
use serde_json::json;
use std::borrow::Cow;
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::Duration;
use tempfile::TempDir;
use tokio::net::TcpListener;
use tokio::sync::mpsc;
use tokio::time::timeout;
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(20);
const EXECUTOR_HTTP_MCP_URL: &str = "http://executor-only.invalid/mcp";
const HTTP_MCP_SERVER_NAME: &str = "executor_http";
const MCP_SERVER_NAME: &str = "executor_demo";
const OAUTH_MCP_SERVER_NAME: &str = "executor_oauth";
const EXECUTOR_OAUTH_MCP_URL: &str = "http://oauth-only.invalid/oauth-mcp";
const EXECUTOR_ENV_NAME: &str = "MCP_EXECUTOR_MARKER";
const EXECUTOR_ENV_VALUE: &str = "executor-only";
const EXECUTOR_ID: &str = "executor-1";
@@ -32,8 +61,58 @@ const REFRESH_PROBE_SERVER_NAME: &str = "refresh_probe";
const TOOL_CALL_ID: &str = "executor-mcp-call";
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn selected_executor_plugin_exposes_its_stdio_mcp_only_to_that_thread() -> Result<()> {
async fn selected_executor_plugin_exposes_its_mcps_only_to_that_thread() -> Result<()> {
let responses_server = responses::start_mock_server().await;
let http_listener = TcpListener::bind("127.0.0.1:0").await?;
let http_addr = http_listener.local_addr()?;
let http_server_config = StreamableHttpServerConfig::default()
.with_allowed_hosts(["executor-only.invalid", "oauth-only.invalid"]);
let http_mcp_service = StreamableHttpService::new(
|| Ok(ExecutorHttpMcpServer),
Arc::new(LocalSessionManager::default()),
http_server_config.clone(),
);
let oauth_mcp_service = StreamableHttpService::new(
|| Ok(ExecutorHttpMcpServer),
Arc::new(LocalSessionManager::default()),
http_server_config,
);
let (token_request_tx, mut token_request_rx) = mpsc::unbounded_channel();
let oauth_metadata = json!({
"authorization_endpoint": "https://oauth-only.invalid/authorize",
"token_endpoint": "http://oauth-only.invalid/token",
"scopes_supported": ["read", "write"],
"response_types_supported": ["code"],
"code_challenge_methods_supported": ["S256"],
});
let http_router = Router::new()
.route(
"/.well-known/oauth-authorization-server/oauth-mcp",
get(move || {
let metadata = oauth_metadata.clone();
async move { Json(metadata) }
}),
)
.route(
"/token",
post(move |body: Bytes| {
let token_request_tx = token_request_tx.clone();
async move {
let _ = token_request_tx.send(String::from_utf8_lossy(&body).into_owned());
Json(json!({
"access_token": "executor-access-token",
"token_type": "Bearer",
"expires_in": 3600,
"refresh_token": "executor-refresh-token",
}))
}
}),
)
.nest_service("/mcp", http_mcp_service)
.nest_service("/oauth-mcp", oauth_mcp_service);
let http_server_handle = tokio::spawn(async move {
let _ = axum::serve(http_listener, http_router).await;
});
let codex_home = TempDir::new()?;
write_mock_responses_config_toml(
codex_home.path(),
@@ -44,6 +123,12 @@ async fn selected_executor_plugin_exposes_its_stdio_mcp_only_to_that_thread() ->
"mock_provider",
"compact",
)?;
let codex_bin = toml::Value::String(
codex_utils_cargo_bin::cargo_bin("codex")?
.to_string_lossy()
.into_owned(),
);
let http_proxy = toml::Value::String(format!("http://{http_addr}"));
std::fs::write(
codex_home.path().join("environments.toml"),
format!(
@@ -52,16 +137,12 @@ include_local = true
[[environments]]
id = "{EXECUTOR_ID}"
program = {}
program = {codex_bin}
args = ["exec-server", "--listen", "stdio"]
[environments.env]
{EXECUTOR_ENV_NAME} = "{EXECUTOR_ENV_VALUE}"
"#,
toml::Value::String(
codex_utils_cargo_bin::cargo_bin("codex")?
.to_string_lossy()
.into_owned()
)
HTTP_PROXY = {http_proxy}
"#
),
)?;
@@ -79,6 +160,17 @@ args = ["exec-server", "--listen", "stdio"]
"command": stdio_server_bin()?,
"env_vars": [EXECUTOR_ENV_NAME],
"startup_timeout_sec": 10,
},
(HTTP_MCP_SERVER_NAME): {
"url": EXECUTOR_HTTP_MCP_URL,
"environment_id": "local",
"startup_timeout_sec": 10,
},
(OAUTH_MCP_SERVER_NAME): {
"url": EXECUTOR_OAUTH_MCP_URL,
"environment_id": "local",
"oauth": {"clientId": "executor-oauth-client"},
"startup_timeout_sec": 10,
}
}
}))?,
@@ -99,7 +191,6 @@ args = ["exec-server", "--listen", "stdio"]
)
.await?;
std::fs::write(plugin.path().join(".mcp.json"), r#"{"mcpServers":{}}"#)?;
let config_path = codex_home.path().join("config.toml");
let mut config = std::fs::read_to_string(&config_path)?;
config.push_str(&format!(
@@ -120,6 +211,76 @@ startup_timeout_sec = 10
)
.await??;
let request_id = app_server
.send_raw_request(
"mcpServer/oauth/login",
Some(json!({
"name": OAUTH_MCP_SERVER_NAME,
"threadId": selected_thread.clone(),
"timeoutSecs": 10,
})),
)
.await?;
let response = timeout(
DEFAULT_READ_TIMEOUT,
app_server.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
let response: McpServerOauthLoginResponse = to_response(response)?;
assert!(
response
.authorization_url
.starts_with("https://oauth-only.invalid/authorize?")
);
assert!(
response
.authorization_url
.contains("client_id=executor-oauth-client")
);
let authorization_url = reqwest::Url::parse(&response.authorization_url)?;
let state = authorization_url
.query_pairs()
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
.expect("authorization URL should include state");
let redirect_uri = authorization_url
.query_pairs()
.find_map(|(key, value)| (key == "redirect_uri").then(|| value.into_owned()))
.expect("authorization URL should include redirect_uri");
let mut callback_url = reqwest::Url::parse(&redirect_uri)?;
callback_url
.query_pairs_mut()
.append_pair("code", "executor-test-code")
.append_pair("state", &state);
reqwest::Client::builder()
.no_proxy()
.build()?
.get(callback_url)
.send()
.await?
.error_for_status()?;
let token_request = timeout(DEFAULT_READ_TIMEOUT, token_request_rx.recv())
.await?
.expect("executor token endpoint should receive a request");
assert!(token_request.contains("grant_type=authorization_code"));
assert!(token_request.contains("code=executor-test-code"));
assert!(token_request.contains("code_verifier="));
let notification = timeout(
DEFAULT_READ_TIMEOUT,
app_server.read_stream_until_notification_message("mcpServer/oauthLogin/completed"),
)
.await??;
let completed: McpServerOauthLoginCompletedNotification =
serde_json::from_value(notification.params.expect("notification params"))?;
assert_eq!(
completed,
McpServerOauthLoginCompletedNotification {
name: OAUTH_MCP_SERVER_NAME.to_string(),
thread_id: Some(selected_thread.clone()),
success: true,
error: None,
}
);
let namespace = format!("mcp__{MCP_SERVER_NAME}");
let response_mock = responses::mount_sse_sequence(
&responses_server,
@@ -178,6 +339,26 @@ startup_timeout_sec = 10
assert!(output.contains("ECHOING: hello from executor"));
assert!(output.contains(EXECUTOR_ENV_VALUE));
let request_id = app_server
.send_mcp_server_tool_call_request(McpServerToolCallParams {
thread_id: selected_thread.clone(),
server: HTTP_MCP_SERVER_NAME.to_string(),
tool: "echo".to_string(),
arguments: Some(json!({"message": "hello over executor HTTP"})),
meta: None,
})
.await?;
let response = timeout(
DEFAULT_READ_TIMEOUT,
app_server.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
let response: McpServerToolCallResponse = to_response(response)?;
assert_eq!(
response.structured_content,
Some(json!({"echo": "ECHOING: hello over executor HTTP"}))
);
let request_id = app_server
.send_mcp_server_tool_call_request(McpServerToolCallParams {
thread_id: selected_thread.clone(),
@@ -200,25 +381,87 @@ startup_timeout_sec = 10
Some(json!("ECHOING: refresh applied"))
);
let selected_server_names = mcp_server_names(&mut app_server, selected_thread).await?;
assert!(
mcp_server_names(&mut app_server, selected_thread)
.await?
selected_server_names
.iter()
.any(|name| name == MCP_SERVER_NAME)
);
assert!(
selected_server_names
.iter()
.any(|name| name == HTTP_MCP_SERVER_NAME)
);
assert!(
selected_server_names
.iter()
.any(|name| name == OAUTH_MCP_SERVER_NAME)
);
let unselected_thread =
start_thread(&mut app_server, /*selected_capability_roots*/ None).await?;
assert!(
mcp_server_names(&mut app_server, unselected_thread)
.await?
.iter()
.all(|name| name != MCP_SERVER_NAME)
);
let unselected_server_names = mcp_server_names(&mut app_server, unselected_thread).await?;
assert!(unselected_server_names.iter().all(|name| {
name != MCP_SERVER_NAME && name != HTTP_MCP_SERVER_NAME && name != OAUTH_MCP_SERVER_NAME
}));
http_server_handle.abort();
let _ = http_server_handle.await;
Ok(())
}
#[derive(Clone, Copy)]
struct ExecutorHttpMcpServer;
impl ServerHandler for ExecutorHttpMcpServer {
fn get_info(&self) -> ServerInfo {
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
}
async fn list_tools(
&self,
_request: Option<rmcp::model::PaginatedRequestParams>,
_context: RequestContext<RoleServer>,
) -> Result<ListToolsResult, rmcp::ErrorData> {
let input_schema: JsonObject = serde_json::from_value(json!({
"type": "object",
"properties": {"message": {"type": "string"}},
"required": ["message"],
"additionalProperties": false
}))
.map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))?;
let mut tool = Tool::new(
Cow::Borrowed("echo"),
Cow::Borrowed("Echo a message."),
Arc::new(input_schema),
);
tool.annotations = Some(ToolAnnotations::new().read_only(true));
Ok(ListToolsResult {
tools: vec![tool],
next_cursor: None,
meta: None,
})
}
async fn call_tool(
&self,
request: CallToolRequestParams,
_context: RequestContext<RoleServer>,
) -> Result<CallToolResult, rmcp::ErrorData> {
let message = request
.arguments
.as_ref()
.and_then(|arguments| arguments.get("message"))
.and_then(serde_json::Value::as_str)
.unwrap_or_default();
Ok(CallToolResult::structured(json!({
"echo": format!("ECHOING: {message}")
})))
}
}
async fn mcp_server_names(
app_server: &mut TestAppServer,
thread_id: String,
@@ -158,6 +158,7 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist
originator: "test_originator".to_string(),
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
selected_capability_roots: Vec::new(),
multi_agent_version: None,
initial_window_id: Uuid::now_v7().to_string(),
metadata: ThreadPersistenceMetadata {
@@ -882,6 +882,7 @@ async fn skills_changed_notification_is_emitted_after_skill_change() -> Result<(
.send_thread_start_request(ThreadStartParams {
model: None,
model_provider: None,
allow_provider_model_fallback: false,
service_tier: None,
cwd: None,
runtime_workspace_roots: None,
+80 -35
View File
@@ -1,11 +1,12 @@
use anyhow::Result;
use app_test_support::TestAppServer;
use app_test_support::to_response;
use codex_app_server_protocol::CurrentTimeReadResponse;
use codex_app_server_protocol::ItemCompletedNotification;
use codex_app_server_protocol::ItemStartedNotification;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ServerRequest;
use codex_app_server_protocol::ThreadItem;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse;
@@ -19,12 +20,16 @@ use std::time::Duration;
use tempfile::TempDir;
use tokio::time::timeout;
#[cfg(windows)]
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(25);
#[cfg(not(windows))]
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10);
const CURRENT_TIME_AT: i64 = 1_781_717_655;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sleep_emits_started_and_completed_items() -> Result<()> {
async fn external_sleep_polls_current_time_and_emits_items() -> Result<()> {
const CALL_ID: &str = "sleep-1";
const DURATION_MS: u64 = 1;
const DURATION_MS: u64 = 2_000;
let server = responses::start_mock_server().await;
responses::mount_sse_sequence(
@@ -85,38 +90,19 @@ async fn sleep_emits_started_and_completed_items() -> Result<()> {
.await??;
let TurnStartResponse { turn, .. } = to_response(turn_start_response)?;
let (started, completed) = timeout(DEFAULT_READ_TIMEOUT, async {
let mut started = None;
let mut completed = None;
while started.is_none() || completed.is_none() {
let JSONRPCMessage::Notification(notification) = mcp.read_next_message().await? else {
continue;
};
match notification.method.as_str() {
"item/started" => {
let payload: ItemStartedNotification =
serde_json::from_value(notification.params.expect("item/started params"))?;
if matches!(&payload.item, ThreadItem::Sleep { .. }) {
started = Some(payload);
}
}
"item/completed" => {
let payload: ItemCompletedNotification = serde_json::from_value(
notification.params.expect("item/completed params"),
)?;
if matches!(&payload.item, ThreadItem::Sleep { .. }) {
completed = Some(payload);
}
}
_ => {}
}
}
Ok::<_, anyhow::Error>((
started.expect("sleep started"),
completed.expect("sleep completed"),
))
})
.await??;
// Read once for the initial reminder, then once to establish the sleep deadline.
respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT).await?;
let started = wait_for_sleep_started(&mut mcp, CALL_ID).await?;
respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT).await?;
// The first poll remains below the deadline, so the provider must request time again.
respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 1).await?;
respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 2).await?;
let completed = wait_for_sleep_completed(&mut mcp, CALL_ID).await?;
// The next inference boundary reads the same external clock after the sleep completes.
respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 2).await?;
timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("turn/completed"),
@@ -150,6 +136,64 @@ async fn sleep_emits_started_and_completed_items() -> Result<()> {
Ok(())
}
async fn wait_for_sleep_started(
mcp: &mut TestAppServer,
call_id: &str,
) -> Result<ItemStartedNotification> {
loop {
let notification = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("item/started"),
)
.await??;
let started: ItemStartedNotification =
serde_json::from_value(notification.params.expect("item/started params"))?;
if matches!(&started.item, ThreadItem::Sleep { id, .. } if id == call_id) {
return Ok(started);
}
}
}
async fn wait_for_sleep_completed(
mcp: &mut TestAppServer,
call_id: &str,
) -> Result<ItemCompletedNotification> {
loop {
let notification = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_notification_message("item/completed"),
)
.await??;
let completed: ItemCompletedNotification =
serde_json::from_value(notification.params.expect("item/completed params"))?;
if matches!(&completed.item, ThreadItem::Sleep { id, .. } if id == call_id) {
return Ok(completed);
}
}
}
async fn respond_to_current_time_read(
mcp: &mut TestAppServer,
thread_id: &str,
current_time_at: i64,
) -> Result<()> {
let request = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_request_message(),
)
.await??;
let ServerRequest::CurrentTimeRead { request_id, params } = request else {
panic!("expected CurrentTimeRead request, got: {request:?}");
};
assert_eq!(params.thread_id, thread_id);
mcp.send_response(
request_id,
serde_json::to_value(CurrentTimeReadResponse { current_time_at })?,
)
.await?;
Ok(())
}
fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> {
std::fs::write(
codex_home.join("config.toml"),
@@ -170,6 +214,7 @@ stream_max_retries = 0
[features.current_time_reminder]
enabled = true
sleep_tool = true
clock_source = "external"
"#
),
)
@@ -1369,6 +1369,7 @@ async fn seed_pathless_store_thread(
originator: "test_originator".to_string(),
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
selected_capability_roots: Vec::new(),
multi_agent_version: None,
initial_window_id: Uuid::now_v7().to_string(),
metadata: ThreadPersistenceMetadata {
@@ -2070,6 +2070,7 @@ stream_max_retries = 0
model_provider: Some("mock_provider".to_string()),
base_instructions: None,
dynamic_tools: None,
selected_capability_roots: Vec::new(),
memory_mode: None,
multi_agent_version: None,
context_window: None,
@@ -3,6 +3,9 @@ use app_test_support::TestAppServer;
use app_test_support::create_final_assistant_message_sse_response;
use app_test_support::create_mock_responses_server_sequence_unchecked;
use app_test_support::to_response;
use codex_app_server_protocol::ClientInfo;
use codex_app_server_protocol::DeprecationNoticeNotification;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ThreadItem;
@@ -22,6 +25,48 @@ use tokio::time::timeout;
const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
#[tokio::test]
async fn thread_rollback_does_not_emit_deprecation_notice_to_codex_tui() -> Result<()> {
let codex_home = TempDir::new()?;
let mut mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?;
let initialized = timeout(
DEFAULT_READ_TIMEOUT,
mcp.initialize_with_client_info(ClientInfo {
name: "codex-tui".to_string(),
title: None,
version: "0.1.0".to_string(),
}),
)
.await??;
let JSONRPCMessage::Response(_) = initialized else {
panic!("expected initialize response, got {initialized:?}");
};
mcp.clear_message_buffer();
let rollback_id = mcp
.send_thread_rollback_request(ThreadRollbackParams {
thread_id: "00000000-0000-0000-0000-000000000001".to_string(),
num_turns: 1,
})
.await?;
loop {
let message = timeout(DEFAULT_READ_TIMEOUT, mcp.read_next_message()).await??;
match message {
JSONRPCMessage::Notification(notification) => {
assert_ne!(notification.method, "deprecationNotice");
}
JSONRPCMessage::Error(error) if error.id == RequestId::Integer(rollback_id) => {
break;
}
message => {
panic!("expected rollback error response, got {message:?}");
}
}
}
Ok(())
}
#[tokio::test]
async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<()> {
// Three Codex turns hit the mock model (session start + two turn/start calls).
@@ -97,6 +142,7 @@ async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<()
mcp.read_stream_until_notification_message("turn/completed"),
)
.await??;
mcp.clear_message_buffer();
// Roll back the last turn.
let rollback_id = mcp
@@ -105,6 +151,23 @@ async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<()
num_turns: 1,
})
.await?;
let deprecation_notice = timeout(DEFAULT_READ_TIMEOUT, mcp.read_next_message()).await??;
let JSONRPCMessage::Notification(deprecation_notice) = deprecation_notice else {
panic!("thread/rollback should emit deprecationNotice before its response");
};
assert_eq!(deprecation_notice.method, "deprecationNotice");
let deprecation_notice: DeprecationNoticeNotification = serde_json::from_value(
deprecation_notice
.params
.expect("deprecationNotice params should be present"),
)?;
assert_eq!(
deprecation_notice,
DeprecationNoticeNotification {
summary: "thread/rollback is deprecated and will be removed soon".to_string(),
details: None,
}
);
let rollback_resp: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(rollback_id)),
@@ -54,6 +54,115 @@ use super::analytics::wait_for_analytics_payload;
const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
const INVALID_REQUEST_ERROR_CODE: i64 = -32600;
async fn start_thread_with_model(
mcp: &mut TestAppServer,
model: &str,
allow_provider_model_fallback: bool,
) -> Result<ThreadStartResponse> {
let request_id = mcp
.send_thread_start_request_with_auto_env(ThreadStartParams {
model: Some(model.to_string()),
allow_provider_model_fallback,
..Default::default()
})
.await?;
let response: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
to_response(response)
}
#[tokio::test]
async fn thread_start_provider_model_fallback_applies_to_configured_model() -> Result<()> {
let codex_home = TempDir::new()?;
std::fs::write(
codex_home.path().join("config.toml"),
r#"model_provider = "amazon-bedrock"
model = "gpt-5.4-mini"
"#,
)?;
let mut mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
let request_id = mcp
.send_thread_start_request_with_auto_env(ThreadStartParams {
allow_provider_model_fallback: true,
..Default::default()
})
.await?;
let response: JSONRPCResponse = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
let response: ThreadStartResponse = to_response(response)?;
assert_eq!(response.model, "openai.gpt-5.5");
Ok(())
}
#[tokio::test]
async fn thread_start_provider_model_fallback_uses_bedrock_static_catalog() -> Result<()> {
let codex_home = TempDir::new()?;
std::fs::write(
codex_home.path().join("config.toml"),
r#"model_provider = "amazon-bedrock"
"#,
)?;
let mut mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
let unsupported_with_fallback = start_thread_with_model(
&mut mcp,
"gpt-5.4-mini",
/*allow_provider_model_fallback*/ true,
)
.await?;
let supported_with_fallback = start_thread_with_model(
&mut mcp,
"openai.gpt-5.4",
/*allow_provider_model_fallback*/ true,
)
.await?;
let unsupported_without_fallback = start_thread_with_model(
&mut mcp,
"gpt-5.4-mini",
/*allow_provider_model_fallback*/ false,
)
.await?;
assert_eq!(
vec![
unsupported_with_fallback.model,
supported_with_fallback.model,
unsupported_without_fallback.model,
],
vec!["openai.gpt-5.5", "openai.gpt-5.4", "gpt-5.4-mini"]
);
Ok(())
}
#[tokio::test]
async fn thread_start_provider_model_fallback_ignores_dynamic_catalog() -> Result<()> {
let server = create_mock_responses_server_repeating_assistant("Done").await;
let codex_home = TempDir::new()?;
create_config_toml_without_approval_policy(codex_home.path(), &server.uri())?;
let mut mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
let response = start_thread_with_model(
&mut mcp,
"unlisted-dynamic-model",
/*allow_provider_model_fallback*/ true,
)
.await?;
assert_eq!(response.model, "unlisted-dynamic-model");
Ok(())
}
#[tokio::test]
async fn thread_start_creates_thread_and_emits_started() -> Result<()> {
// Provide a mock server and config so model wiring is valid.
@@ -836,6 +945,7 @@ async fn thread_start_emits_mcp_server_status_updated_notifications() -> Result<
name: "optional_broken".to_string(),
status: McpServerStartupState::Starting,
error: None,
failure_reason: None,
}
);
@@ -868,6 +978,7 @@ async fn thread_start_emits_mcp_server_status_updated_notifications() -> Result<
assert_eq!(failed.thread_id, Some(start_response.thread.id));
assert_eq!(failed.name, "optional_broken");
assert_eq!(failed.status, McpServerStartupState::Failed);
assert_eq!(failed.failure_reason, None);
assert!(
failed
.error
@@ -218,6 +218,7 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> {
originator: "test_originator".to_string(),
base_instructions: BaseInstructions::default(),
dynamic_tools: Vec::new(),
selected_capability_roots: Vec::new(),
multi_agent_version: None,
initial_window_id: Uuid::now_v7().to_string(),
metadata: ThreadPersistenceMetadata {
+80 -5
View File
@@ -1,3 +1,5 @@
use std::future::Future;
use tracing_subscriber::EnvFilter;
use tracing_subscriber::prelude::*;
@@ -7,7 +9,7 @@ const OTEL_SERVICE_NAME: &str = "codex-exec-server";
pub(crate) fn init(
config: Option<&codex_core::config::Config>,
) -> Result<impl Send + Sync, Box<dyn std::error::Error>> {
) -> (impl Send + Sync, codex_exec_server::ExecServerTelemetry) {
let fmt_layer = tracing_subscriber::fmt::layer()
.with_writer(std::io::stderr)
.with_filter(stderr_env_filter());
@@ -17,21 +19,94 @@ pub(crate) fn init(
env!("CARGO_PKG_VERSION"),
Some(OTEL_SERVICE_NAME),
DEFAULT_ANALYTICS_ENABLED,
),
None => Ok(None),
)
.unwrap_or_else(|error| {
eprintln!("Could not create otel exporter: {error}");
None
}),
None => None,
};
let provider = otel.as_ref().ok().and_then(Option::as_ref);
let provider = otel.as_ref();
codex_core::otel_init::record_process_start(provider, OTEL_SERVICE_NAME);
let otel_logger_layer = provider.and_then(|otel| otel.logger_layer());
let otel_tracing_layer = provider.and_then(|otel| otel.tracing_layer());
let telemetry = provider
.and_then(|otel| otel.metrics())
.cloned()
.map(codex_exec_server::ExecServerTelemetry::new)
.unwrap_or_default();
let _ = tracing_subscriber::registry()
.with(fmt_layer)
.with(otel_tracing_layer)
.with(otel_logger_layer)
.try_init();
tracing::callsite::rebuild_interest_cache();
otel
(otel, telemetry)
}
pub(crate) async fn run_until_shutdown<F, E>(run: F) -> Result<(), E>
where
F: Future<Output = Result<(), E>>,
{
let shutdown_signal = match shutdown_signal() {
Ok(signal) => Some(signal),
Err(error) => {
eprintln!("Could not listen for exec-server shutdown signal: {error}");
None
}
};
tokio::pin!(run);
if let Some(shutdown_signal) = shutdown_signal {
tokio::select! {
result = &mut run => result,
signal = wait_for_shutdown_signal(shutdown_signal) => {
match signal {
Ok(()) => Ok(()),
Err(error) => {
eprintln!("Could not listen for exec-server shutdown signal: {error}");
run.await
}
}
}
}
} else {
run.await
}
}
#[cfg(unix)]
struct ShutdownSignal {
terminate: tokio::signal::unix::Signal,
}
#[cfg(unix)]
fn shutdown_signal() -> std::io::Result<ShutdownSignal> {
Ok(ShutdownSignal {
terminate: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?,
})
}
#[cfg(unix)]
async fn wait_for_shutdown_signal(mut shutdown_signal: ShutdownSignal) -> std::io::Result<()> {
tokio::select! {
result = tokio::signal::ctrl_c() => result,
_ = shutdown_signal.terminate.recv() => Ok(()),
}
}
#[cfg(not(unix))]
struct ShutdownSignal;
#[cfg(not(unix))]
fn shutdown_signal() -> std::io::Result<ShutdownSignal> {
Ok(ShutdownSignal)
}
#[cfg(not(unix))]
async fn wait_for_shutdown_signal(_: ShutdownSignal) -> std::io::Result<()> {
tokio::signal::ctrl_c().await
}
fn stderr_env_filter() -> EnvFilter {
+18 -3
View File
@@ -53,7 +53,12 @@ pub struct SeatbeltCommand {
pub sandbox_state: SandboxStateArgs,
/// Named permissions profile to apply from the active configuration stack.
#[arg(long = "permissions-profile", short = 'P', value_name = "NAME")]
#[arg(
long = "permission-profile",
alias = "permissions-profile",
short = 'P',
value_name = "NAME"
)]
pub permissions_profile: Option<String>,
/// Layer $CODEX_HOME/<name>.config.toml on top of the base user config.
@@ -104,7 +109,12 @@ pub struct LandlockCommand {
pub sandbox_state: SandboxStateArgs,
/// Named permissions profile to apply from the active configuration stack.
#[arg(long = "permissions-profile", short = 'P', value_name = "NAME")]
#[arg(
long = "permission-profile",
alias = "permissions-profile",
short = 'P',
value_name = "NAME"
)]
pub permissions_profile: Option<String>,
/// Layer $CODEX_HOME/<name>.config.toml on top of the base user config.
@@ -142,7 +152,12 @@ pub struct WindowsCommand {
pub sandbox_state: SandboxStateArgs,
/// Named permissions profile to apply from the active configuration stack.
#[arg(long = "permissions-profile", short = 'P', value_name = "NAME")]
#[arg(
long = "permission-profile",
alias = "permissions-profile",
short = 'P',
value_name = "NAME"
)]
pub permissions_profile: Option<String>,
/// Layer $CODEX_HOME/<name>.config.toml on top of the base user config.
+43 -13
View File
@@ -1693,9 +1693,7 @@ async fn run_exec_server_command(
.environment_id
.ok_or_else(|| anyhow::anyhow!("--environment-id is required when --remote is set"))?;
let config = load_exec_server_config(root_config_overrides, strict_config).await?;
let _otel = exec_server_telemetry::init(Some(&config))
.inspect_err(|err| eprintln!("Could not create otel exporter: {err}"))
.ok();
let (_otel, telemetry) = exec_server_telemetry::init(Some(&config));
let auth_provider =
load_exec_server_remote_auth_provider(&config, &base_url, cmd.use_agent_identity_auth)
.await?;
@@ -1707,7 +1705,11 @@ async fn run_exec_server_command(
if let Some(name) = cmd.name {
remote_config.name = name;
}
codex_exec_server::run_remote_environment(remote_config, runtime_paths).await?;
let remote_config = remote_config.with_telemetry(telemetry);
exec_server_telemetry::run_until_shutdown(async move {
codex_exec_server::run_remote_environment(remote_config, runtime_paths).await
})
.await?;
Ok(())
} else {
let config_result = load_exec_server_config(root_config_overrides, strict_config).await;
@@ -1716,16 +1718,15 @@ async fn run_exec_server_command(
} else {
config_result.ok()
};
let _otel = exec_server_telemetry::init(config.as_ref())
.inspect_err(|err| eprintln!("Could not create otel exporter: {err}"))
.ok();
let (_otel, telemetry) = exec_server_telemetry::init(config.as_ref());
let listen_url = cmd
.listen
.as_deref()
.unwrap_or(codex_exec_server::DEFAULT_LISTEN_URL);
codex_exec_server::run_main(listen_url, runtime_paths)
.await
.map_err(anyhow::Error::from_boxed)
.unwrap_or_else(|| codex_exec_server::DEFAULT_LISTEN_URL.to_string());
exec_server_telemetry::run_until_shutdown(async move {
codex_exec_server::run_main_with_telemetry(&listen_url, runtime_paths, telemetry).await
})
.await
.map_err(anyhow::Error::from_boxed)
}
}
@@ -2965,7 +2966,28 @@ mod tests {
#[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))]
#[test]
fn sandbox_parses_permissions_profile() {
fn sandbox_parses_permission_profile() {
let cli = MultitoolCli::try_parse_from([
"codex",
"sandbox",
"--permission-profile",
":workspace",
"--",
"echo",
])
.expect("parse");
let Some(Subcommand::Sandbox(command)) = cli.subcommand else {
panic!("expected sandbox command");
};
assert_eq!(command.permissions_profile.as_deref(), Some(":workspace"));
assert_eq!(command.command, vec!["echo"]);
}
#[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))]
#[test]
fn sandbox_parses_legacy_permissions_profile_alias() {
let cli = MultitoolCli::try_parse_from([
"codex",
"sandbox",
@@ -2984,6 +3006,14 @@ mod tests {
assert_eq!(command.command, vec!["echo"]);
}
#[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))]
#[test]
fn sandbox_help_only_shows_singular_permission_profile() {
let help = help_from_args(&["codex", "sandbox", "--help"]);
assert!(help.contains("--permission-profile"), "{help}");
assert!(!help.contains("--permissions-profile"), "{help}");
}
#[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))]
#[test]
fn sandbox_parses_permissions_profile_short_alias() {
+9 -3
View File
@@ -18,8 +18,10 @@ use codex_core::config::edit::ConfigEditsBuilder;
use codex_core::config::find_codex_home;
use codex_core::config::load_global_mcp_servers;
use codex_core_plugins::PluginsManager;
use codex_exec_server::EnvironmentManager;
use codex_login::AuthManager;
use codex_mcp::McpOAuthLoginSupport;
use codex_mcp::McpRuntimeContext;
use codex_mcp::ResolvedMcpOAuthScopes;
use codex_mcp::compute_auth_statuses;
use codex_mcp::discover_supported_scopes;
@@ -558,6 +560,10 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) ->
config.mcp_oauth_credentials_store_mode,
config.auth_keyring_backend_kind(),
auth.as_ref(),
&McpRuntimeContext::new(
Arc::new(EnvironmentManager::without_environments()),
config.cwd.to_path_buf(),
),
)
.await;
@@ -567,7 +573,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) ->
.map(|(name, cfg)| {
let auth_status = auth_statuses
.get(name.as_str())
.map(|entry| entry.auth_status)
.map(|entry| McpAuthStatus::from(entry.auth_state))
.unwrap_or(McpAuthStatus::Unsupported);
let transport = match &cfg.transport {
McpServerTransportConfig::Stdio {
@@ -651,7 +657,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) ->
let status = format_mcp_status(cfg);
let auth_status = auth_statuses
.get(name.as_str())
.map(|entry| entry.auth_status)
.map(|entry| McpAuthStatus::from(entry.auth_state))
.unwrap_or(McpAuthStatus::Unsupported)
.to_string();
stdio_rows.push([
@@ -672,7 +678,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) ->
let status = format_mcp_status(cfg);
let auth_status = auth_statuses
.get(name.as_str())
.map(|entry| entry.auth_status)
.map(|entry| McpAuthStatus::from(entry.auth_state))
.unwrap_or(McpAuthStatus::Unsupported)
.to_string();
let bearer_token_display =
+295
View File
@@ -1,9 +1,34 @@
#[cfg(unix)]
use std::io::BufRead as _;
#[cfg(unix)]
use std::io::BufReader as StdBufReader;
#[cfg(unix)]
use std::io::Read as _;
#[cfg(unix)]
use std::io::Write as _;
#[cfg(unix)]
use std::net::TcpStream;
use std::path::Path;
use std::process::Stdio;
#[cfg(unix)]
use std::thread;
use std::time::Duration;
#[cfg(unix)]
use std::time::Instant;
use anyhow::Result;
use predicates::prelude::PredicateBooleanExt;
use predicates::str::contains;
use tempfile::TempDir;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
fn codex_command(codex_home: &Path) -> Result<assert_cmd::Command> {
let mut cmd = assert_cmd::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?);
@@ -48,3 +73,273 @@ fn local_exec_server_ignores_invalid_config_without_strict_config() -> Result<()
Ok(())
}
#[tokio::test]
async fn local_exec_server_flushes_telemetry_on_stdio_disconnect() -> Result<()> {
let collector = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/metrics"))
.respond_with(ResponseTemplate::new(202))
.mount(&collector)
.await;
let codex_home = TempDir::new()?;
let base_url = collector.uri();
std::fs::write(
codex_home.path().join("config.toml"),
format!(
r#"
[analytics]
enabled = true
[otel]
environment = "test"
metrics_exporter = {{ otlp-http = {{ endpoint = "{base_url}/v1/metrics", protocol = "json" }} }}
"#
),
)?;
let cwd = url::Url::from_directory_path(std::env::current_dir()?)
.map_err(|()| anyhow::anyhow!("could not convert cwd to file URL"))?;
#[cfg(windows)]
let argv = vec!["ping.exe", "-n", "61", "127.0.0.1"];
#[cfg(not(windows))]
let argv = vec!["/bin/sleep", "60"];
let codex_bin = codex_utils_cargo_bin::cargo_bin("codex")?;
let codex_home = codex_home.path().to_path_buf();
let subprocess = async move {
let mut command = tokio::process::Command::new(codex_bin);
command
.env("CODEX_HOME", codex_home)
.env("NO_PROXY", "127.0.0.1,localhost")
.env("no_proxy", "127.0.0.1,localhost")
.args(["exec-server", "--listen", "stdio"])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.kill_on_drop(true);
let mut child = command.spawn()?;
let mut stdin = child
.stdin
.take()
.ok_or_else(|| anyhow::anyhow!("exec-server stdin was not piped"))?;
let stdout = child
.stdout
.take()
.ok_or_else(|| anyhow::anyhow!("exec-server stdout was not piped"))?;
let mut stdout = BufReader::new(stdout);
send_json_line(
&mut stdin,
&serde_json::json!({
"id": 1,
"method": "initialize",
"params": {"clientName": "otel-test", "resumeSessionId": null}
}),
)
.await?;
wait_for_response(&mut stdout, /*expected_id*/ 1).await?;
send_json_line(
&mut stdin,
&serde_json::json!({"method": "initialized", "params": {}}),
)
.await?;
send_json_line(
&mut stdin,
&serde_json::json!({
"id": 2,
"method": "process/start",
"params": {
"processId": "otel-process",
"argv": argv,
"cwd": cwd,
"env": {},
"tty": false,
"pipeStdin": false,
"arg0": null
}
}),
)
.await?;
wait_for_response(&mut stdout, /*expected_id*/ 2).await?;
drop(stdin);
let mut remaining_stdout = String::new();
stdout.read_to_string(&mut remaining_stdout).await?;
let status = child.wait().await?;
anyhow::ensure!(
status.success(),
"exec-server exited with {status}; remaining stdout: {remaining_stdout}"
);
Ok::<(), anyhow::Error>(())
};
let subprocess_result = tokio::time::timeout(Duration::from_secs(30), subprocess)
.await
.map_err(|_| anyhow::anyhow!("exec-server subprocess timed out"))?;
subprocess_result?;
let requests = collector
.received_requests()
.await
.ok_or_else(|| anyhow::anyhow!("failed to read OTLP collector requests"))?;
let metrics = requests
.iter()
.filter(|request| request.url.path() == "/v1/metrics")
.map(|request| serde_json::from_slice::<serde_json::Value>(&request.body))
.collect::<serde_json::Result<Vec<_>>>()?;
assert_metric_point(
&metrics,
"exec_server_connections_active",
&[("transport", "stdio")],
Some(0),
);
assert_metric_point(
&metrics,
"exec_server_connections_total",
&[("transport", "stdio")],
Some(1),
);
assert_metric_point(
&metrics,
"exec_server_requests_total",
&[("method", "process/start"), ("result", "success")],
Some(1),
);
assert_metric_point(&metrics, "exec_server_processes_active", &[], Some(0));
assert_metric_point(
&metrics,
"exec_server_processes_finished_total",
&[("result", "terminated")],
Some(1),
);
assert_metric_point(
&metrics,
"exec_server_request_duration_seconds",
&[("method", "process/start"), ("result", "success")],
/*value*/ None,
);
assert_metric_point(
&metrics,
"exec_server_process_duration_seconds",
&[("result", "terminated")],
/*value*/ None,
);
Ok(())
}
async fn send_json_line(
stdin: &mut (impl tokio::io::AsyncWrite + Unpin),
message: &serde_json::Value,
) -> Result<()> {
let mut encoded = serde_json::to_vec(message)?;
encoded.push(b'\n');
stdin.write_all(&encoded).await?;
stdin.flush().await?;
Ok(())
}
#[cfg(unix)]
#[test]
fn local_exec_server_exits_successfully_on_sigterm() -> Result<()> {
let codex_home = TempDir::new()?;
let mut child = std::process::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?)
.env("CODEX_HOME", codex_home.path())
.args(["exec-server", "--listen", "ws://127.0.0.1:0"])
.stdout(Stdio::piped())
.spawn()?;
let mut listen_url = String::new();
StdBufReader::new(child.stdout.take().expect("child stdout")).read_line(&mut listen_url)?;
assert!(listen_url.starts_with("ws://127.0.0.1:"), "{listen_url}");
let listen_addr = listen_url
.trim()
.strip_prefix("ws://")
.expect("listen URL should use ws://")
.parse()?;
let deadline = Instant::now() + Duration::from_secs(5);
let mut ready = false;
while let Some(remaining) = deadline.checked_duration_since(Instant::now()) {
if let Ok(mut stream) =
TcpStream::connect_timeout(&listen_addr, remaining.min(Duration::from_millis(100)))
{
let _ = stream.set_read_timeout(Some(Duration::from_secs(1)));
let request =
format!("GET /readyz HTTP/1.1\r\nHost: {listen_addr}\r\nConnection: close\r\n\r\n");
let mut response = String::new();
if stream.write_all(request.as_bytes()).is_ok()
&& stream.read_to_string(&mut response).is_ok()
&& response.starts_with("HTTP/1.1 200")
{
ready = true;
break;
}
}
thread::sleep(Duration::from_millis(10));
}
assert!(ready, "exec-server did not become ready at {listen_url}");
// SAFETY: `child.id()` is the live process spawned above.
let result = unsafe { libc::kill(child.id() as libc::pid_t, libc::SIGTERM) };
assert_eq!(result, 0);
let status = child.wait()?;
assert!(status.success(), "{status}");
Ok(())
}
async fn wait_for_response(
stdout: &mut (impl tokio::io::AsyncBufRead + Unpin),
expected_id: i64,
) -> Result<()> {
loop {
let mut line = String::new();
if stdout.read_line(&mut line).await? == 0 {
anyhow::bail!("exec-server stdout closed before response {expected_id}");
}
let message: serde_json::Value = serde_json::from_str(&line)?;
if message["id"].as_i64() == Some(expected_id) {
anyhow::ensure!(
message.get("error").is_none(),
"exec-server request {expected_id} failed: {message}"
);
return Ok(());
}
}
}
fn assert_metric_point(
payloads: &[serde_json::Value],
name: &str,
attributes: &[(&str, &str)],
value: Option<i64>,
) {
let found = payloads
.iter()
.flat_map(|payload| payload["resourceMetrics"].as_array().into_iter().flatten())
.flat_map(|resource| resource["scopeMetrics"].as_array().into_iter().flatten())
.flat_map(|scope| scope["metrics"].as_array().into_iter().flatten())
.filter(|metric| metric["name"].as_str() == Some(name))
.flat_map(|metric| {
["gauge", "sum", "histogram"]
.into_iter()
.find_map(|kind| metric[kind]["dataPoints"].as_array())
.into_iter()
.flatten()
})
.any(|point| {
let actual_attributes = point["attributes"]
.as_array()
.map(Vec::as_slice)
.unwrap_or_default();
let attributes_match = actual_attributes.len() == attributes.len()
&& attributes.iter().all(|(expected_key, expected_value)| {
actual_attributes.iter().any(|actual| {
actual["key"].as_str() == Some(*expected_key)
&& actual["value"]["stringValue"].as_str() == Some(*expected_value)
})
});
let actual_value = point["asInt"]
.as_i64()
.or_else(|| point["asInt"].as_str()?.parse().ok());
attributes_match && value.is_none_or(|expected| actual_value == Some(expected))
});
assert!(
found,
"metric {name} with attributes {attributes:?} and value {value:?} missing"
);
}
+1 -1
View File
@@ -35,7 +35,7 @@ mode = "full"
.env("CODEX_HOME", codex_home.path())
.args([
"sandbox",
"--permissions-profile",
"--permission-profile",
"network-test",
"--",
"curl",
+15
View File
@@ -8,5 +8,20 @@ license.workspace = true
name = "codex-code-mode-host"
path = "src/main.rs"
[lib]
doctest = false
name = "codex_code_mode_host"
path = "src/lib.rs"
[lints]
workspace = true
[dependencies]
anyhow = { workspace = true }
codex-code-mode = { workspace = true }
codex-code-mode-protocol = { workspace = true }
tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt", "sync", "time"] }
tokio-util = { workspace = true, features = ["rt"] }
[dev-dependencies]
pretty_assertions = { workspace = true }
+85
View File
@@ -0,0 +1,85 @@
use std::sync::Arc;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeNestedToolCall;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::NotificationFuture;
use codex_code_mode_protocol::ToolInvocationFuture;
use codex_code_mode_protocol::host::DelegateRequest;
use codex_code_mode_protocol::host::DelegateResponse;
use codex_code_mode_protocol::host::SessionId;
use tokio_util::sync::CancellationToken;
use crate::peer::HostPeer;
pub(super) struct RemoteDelegate {
session_id: SessionId,
peer: Arc<HostPeer>,
}
impl RemoteDelegate {
pub(super) fn new(session_id: SessionId, peer: Arc<HostPeer>) -> Self {
Self { session_id, peer }
}
}
impl CodeModeSessionDelegate for RemoteDelegate {
fn invoke_tool<'a>(
&'a self,
invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async move {
match self
.peer
.call(
self.session_id.clone(),
DelegateRequest::InvokeTool {
invocation: invocation.into(),
},
cancellation_token,
)
.await?
{
DelegateResponse::ToolResult { result } => Ok(result),
DelegateResponse::NotificationDelivered => {
Err("code-mode client returned an invalid tool result".to_string())
}
}
})
}
fn notify<'a>(
&'a self,
call_id: String,
cell_id: CellId,
text: String,
cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async move {
match self
.peer
.call(
self.session_id.clone(),
DelegateRequest::Notify {
call_id,
cell_id: cell_id.into(),
text,
},
cancellation_token,
)
.await?
{
DelegateResponse::NotificationDelivered => Ok(()),
DelegateResponse::ToolResult { .. } => {
Err("code-mode client returned an invalid notification result".to_string())
}
}
})
}
fn cell_closed(&self, cell_id: &CellId) {
self.peer
.close_cell(self.session_id.clone(), cell_id.clone());
}
}
+486
View File
@@ -0,0 +1,486 @@
use codex_code_mode_protocol::host::Capability;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientHello;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::FramedReader;
use codex_code_mode_protocol::host::FramedWriter;
use codex_code_mode_protocol::host::HandshakeRejectReason;
use codex_code_mode_protocol::host::HostHello;
use codex_code_mode_protocol::host::HostRequest;
use codex_code_mode_protocol::host::HostResponse;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::ProtocolVersion;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SessionId;
use codex_code_mode_protocol::host::SupportedProtocolVersions;
use codex_code_mode_protocol::host::WireExecuteRequest;
use codex_code_mode_protocol::host::WireResult;
use pretty_assertions::assert_eq;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use super::HostState;
use super::MAX_ACTIVE_CELLS;
use super::MAX_IN_FLIGHT_REQUESTS;
use super::MAX_RECENT_REQUEST_IDS;
use super::RequestKind;
use super::RequestRegistry;
use super::SeenSessionIds;
use super::peer::HostPeer;
use super::run;
fn client_hello(
versions: impl IntoIterator<Item = ProtocolVersion>,
required_capabilities: CapabilitySet,
) -> ClientToHost {
ClientToHost::ClientHello(
ClientHello::new(
SupportedProtocolVersions::try_new(versions).expect("supported versions"),
required_capabilities,
CapabilitySet::empty(),
)
.expect("client hello"),
)
}
fn session_id(value: &str) -> SessionId {
SessionId::new(value).expect("session ID")
}
fn request_id(value: i64) -> RequestId {
RequestId::new(value)
}
async fn decode_frame(frame: EncodedFrame) -> HostToClient {
let (reader, writer) = tokio::io::duplex(/*max_buf_size*/ 4096);
let writer = tokio::spawn(async move {
FramedWriter::new(writer)
.write_frame(&frame)
.await
.expect("write encoded frame");
});
let message = FramedReader::new(reader)
.read()
.await
.expect("read encoded frame")
.expect("encoded frame message");
writer.await.expect("frame writer task");
message
}
fn execute_request(source: &str) -> WireExecuteRequest {
WireExecuteRequest {
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
yield_time_ms: Some(60_000),
max_output_tokens: Some(1_000),
}
}
#[tokio::test]
async fn handshake_and_multiple_session_lifecycles_are_ordered() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 4096);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
writer
.write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty()))
.await
.expect("write hello");
assert_eq!(
reader.read::<HostToClient>().await.expect("read hello"),
Some(HostToClient::HostHello(HostHello::new(
ProtocolVersion::V1,
CapabilitySet::empty(),
)))
);
for (request_id, id) in [
(request_id(/*value*/ 1), "session-1"),
(request_id(/*value*/ 2), "session-2"),
] {
writer
.write(&ClientToHost::Request {
id: request_id,
request: HostRequest::OpenSession {
session_id: session_id(id),
},
})
.await
.expect("open session");
assert_eq!(
reader.read::<HostToClient>().await.expect("session ready"),
Some(HostToClient::Response {
id: request_id,
result: WireResult::Ok {
value: HostResponse::SessionReady {
session_id: session_id(id),
},
},
})
);
}
for (request_id, id) in [
(request_id(/*value*/ 3), "session-1"),
(request_id(/*value*/ 4), "session-2"),
] {
writer
.write(&ClientToHost::Request {
id: request_id,
request: HostRequest::ShutdownSession {
session_id: session_id(id),
},
})
.await
.expect("shutdown session");
assert_eq!(
reader.read::<HostToClient>().await.expect("session closed"),
Some(HostToClient::Response {
id: request_id,
result: WireResult::Ok {
value: HostResponse::SessionClosed {
session_id: session_id(id),
},
},
})
);
}
drop(writer);
drop(reader);
host.await.expect("host task").expect("host connection");
}
#[tokio::test]
async fn incompatible_or_invalid_handshake_is_rejected() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
let version_two = ProtocolVersion::new(/*value*/ 2).expect("protocol version");
writer
.write(&client_hello([version_two], CapabilitySet::empty()))
.await
.expect("write hello");
assert_eq!(
reader.read::<HostToClient>().await.expect("rejection"),
Some(HostToClient::HandshakeRejected {
reason: HandshakeRejectReason::NoCompatibleVersion {
supported_versions: SupportedProtocolVersions::try_new([ProtocolVersion::V1])
.expect("host versions"),
},
})
);
host.await.expect("host task").expect("host connection");
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
writer
.write(&ClientToHost::Request {
id: request_id(/*value*/ 1),
request: HostRequest::OpenSession {
session_id: session_id("session-1"),
},
})
.await
.expect("write invalid first message");
assert_eq!(
reader.read::<HostToClient>().await.expect("rejection"),
Some(HostToClient::HandshakeRejected {
reason: HandshakeRejectReason::InvalidHello {
message: "first message must be connection/hello".to_string(),
},
})
);
host.await.expect("host task").expect("host connection");
}
#[tokio::test]
async fn unsupported_required_capability_is_rejected() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
let capability = Capability::new("required").expect("capability");
writer
.write(&client_hello(
[ProtocolVersion::V1],
CapabilitySet::try_new([capability.clone()]).expect("capabilities"),
))
.await
.expect("write hello");
assert_eq!(
reader.read::<HostToClient>().await.expect("rejection"),
Some(HostToClient::HandshakeRejected {
reason: HandshakeRejectReason::MissingRequiredCapability { capability },
})
);
host.await.expect("host task").expect("host connection");
}
#[tokio::test]
async fn session_id_cannot_be_reused_after_shutdown() {
let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 2048);
let (host_reader, host_writer) = tokio::io::split(host_stream);
let (client_reader, client_writer) = tokio::io::split(client_stream);
let host = tokio::spawn(run(host_reader, host_writer));
let mut reader = FramedReader::new(client_reader);
let mut writer = FramedWriter::new(client_writer);
writer
.write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty()))
.await
.expect("write hello");
reader
.read::<HostToClient>()
.await
.expect("read hello")
.expect("host hello");
let id = session_id("session-1");
for (request_id, request) in [
(
request_id(/*value*/ 1),
HostRequest::OpenSession {
session_id: id.clone(),
},
),
(
request_id(/*value*/ 2),
HostRequest::ShutdownSession {
session_id: id.clone(),
},
),
] {
writer
.write(&ClientToHost::Request {
id: request_id,
request,
})
.await
.expect("session request");
reader
.read::<HostToClient>()
.await
.expect("session response")
.expect("session response message");
}
writer
.write(&ClientToHost::Request {
id: request_id(/*value*/ 3),
request: HostRequest::OpenSession { session_id: id },
})
.await
.expect("reuse session ID");
assert_eq!(
reader.read::<HostToClient>().await.expect("reuse response"),
Some(HostToClient::Response {
id: request_id(/*value*/ 3),
result: WireResult::Err {
message: "code-mode session ID `session-1` was reused".to_string(),
},
})
);
drop(writer);
drop(reader);
host.await.expect("host task").expect("host connection");
}
#[test]
fn request_cancellation_tombstones_are_bounded() {
let mut requests = RequestRegistry::default();
let duplicate = request_id(/*value*/ -1);
requests
.start(duplicate, RequestKind::OpenSession)
.expect("start duplicate probe");
assert!(requests.start(duplicate, RequestKind::OpenSession).is_err());
requests.finish(duplicate);
for value in 1..=MAX_RECENT_REQUEST_IDS as i64 + 100 {
let id = request_id(value);
requests
.start(id, RequestKind::Wait)
.expect("start request");
requests.cancel(id);
requests.finish(id);
}
for value in 10_000..20_000 {
requests.cancel(request_id(value));
}
assert!(requests.active.is_empty());
assert_eq!(requests.recent.len(), MAX_RECENT_REQUEST_IDS);
assert_eq!(requests.recent_order.len(), MAX_RECENT_REQUEST_IDS);
}
#[tokio::test]
async fn request_task_panic_disconnects_host() {
let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let state = HostState {
sessions: Mutex::new(HashMap::new()),
seen_session_ids: Mutex::new(SeenSessionIds::default()),
requests: Mutex::new(RequestRegistry::default()),
request_tasks: TaskTracker::new(),
request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)),
active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)),
closing: AtomicBool::new(false),
peer: Arc::clone(&peer),
};
let task = state.request_tasks.spawn(async {
panic!("request panic probe");
});
state.supervise_request_task(task);
tokio::time::timeout(Duration::from_secs(1), peer.disconnected())
.await
.expect("request panic should disconnect host");
assert!(
peer.failure()
.expect("request failure")
.contains("request task failed")
);
}
#[tokio::test]
async fn execute_request_id_remains_active_until_initial_response() {
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 4);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let state = Arc::new(HostState {
sessions: Mutex::new(HashMap::new()),
seen_session_ids: Mutex::new(SeenSessionIds::default()),
requests: Mutex::new(RequestRegistry::default()),
request_tasks: TaskTracker::new(),
request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)),
active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)),
closing: AtomicBool::new(false),
peer,
});
let session_id = session_id("session-1");
state
.open_session(session_id.clone())
.expect("open session");
let request_id = request_id(/*value*/ 1);
state
.spawn_request(
request_id,
HostRequest::Execute {
session_id: session_id.clone(),
request: execute_request("await new Promise(() => {});"),
},
)
.expect("spawn execute request");
let started = decode_frame(outgoing_rx.recv().await.expect("execution started frame")).await;
let HostToClient::Response {
id,
result:
WireResult::Ok {
value: HostResponse::ExecutionStarted { cell_id },
},
} = started
else {
panic!("expected execution started response");
};
assert_eq!(id, request_id);
assert!(
state
.requests
.lock()
.unwrap_or_else(PoisonError::into_inner)
.active
.contains_key(&request_id)
);
state
.session(&session_id)
.expect("session")
.terminate(cell_id.into())
.await
.expect("terminate cell");
state.disconnect().await;
}
#[tokio::test]
async fn active_cell_limit_rejects_execute_without_disconnecting() {
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 1);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let state = HostState {
sessions: Mutex::new(HashMap::new()),
seen_session_ids: Mutex::new(SeenSessionIds::default()),
requests: Mutex::new(RequestRegistry::default()),
request_tasks: TaskTracker::new(),
request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)),
active_cell_permits: Arc::new(Semaphore::new(/*permits*/ 0)),
closing: AtomicBool::new(false),
peer: Arc::clone(&peer),
};
let session_id = session_id("session-1");
state
.open_session(session_id.clone())
.expect("open session");
let request_id = request_id(/*value*/ 1);
state
.handle_request(
request_id,
HostRequest::Execute {
session_id,
request: execute_request("text(\"hello\");"),
},
CancellationToken::new(),
)
.await;
assert_eq!(
decode_frame(outgoing_rx.recv().await.expect("execute response frame")).await,
HostToClient::Response {
id: request_id,
result: WireResult::Err {
message: "code-mode host has too many active cells".to_string(),
},
}
);
assert!(!peer.is_disconnected());
state.disconnect().await;
}
#[tokio::test]
async fn cell_forwarding_panic_disconnects_host() {
let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1);
let peer = Arc::new(HostPeer::new(outgoing_tx));
peer.spawn_critical("cell forwarding", async {
panic!("cell forwarding panic probe");
});
tokio::time::timeout(Duration::from_secs(1), peer.disconnected())
.await
.expect("cell panic should disconnect host");
assert!(
peer.failure()
.expect("cell failure")
.contains("cell forwarding task failed")
);
}
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::PoisonError;
use std::sync::atomic::AtomicBool;
use std::time::Duration;
+605
View File
@@ -0,0 +1,605 @@
use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::PoisonError;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
use codex_code_mode::InProcessCodeModeSession;
use codex_code_mode_protocol::host::CapabilitySet;
use codex_code_mode_protocol::host::ClientToHost;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::FramedReader;
use codex_code_mode_protocol::host::FramedWriter;
use codex_code_mode_protocol::host::HandshakeRejectReason;
use codex_code_mode_protocol::host::HostHello;
use codex_code_mode_protocol::host::HostRequest;
use codex_code_mode_protocol::host::HostResponse;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::ProtocolVersion;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SessionId;
use codex_code_mode_protocol::host::SupportedProtocolVersions;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use self::delegate::RemoteDelegate;
use self::peer::HostPeer;
mod delegate;
mod peer;
const MAX_IN_FLIGHT_REQUESTS: usize = 256;
const MAX_ACTIVE_CELLS: usize = 128;
const MAX_RECENT_REQUEST_IDS: usize = 4096;
const MAX_RECENT_SESSION_IDS: usize = 4096;
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
/// Runs one code-mode host connection over the process standard streams.
pub async fn run_stdio() -> Result<()> {
run(tokio::io::stdin(), tokio::io::stdout()).await
}
/// Runs one code-mode host connection over an ordered input/output pair.
async fn run<R, W>(reader: R, writer: W) -> Result<()>
where
R: AsyncRead + Send + Unpin + 'static,
W: AsyncWrite + Send + Unpin + 'static,
{
let mut reader = FramedReader::new(reader);
let mut writer = FramedWriter::new(writer);
if !negotiate(&mut reader, &mut writer).await? {
return Ok(());
}
let (outgoing_tx, mut outgoing_rx) = mpsc::channel::<EncodedFrame>(/*max_capacity*/ 128);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let state = Arc::new(HostState {
sessions: Mutex::new(HashMap::new()),
seen_session_ids: Mutex::new(SeenSessionIds::default()),
requests: Mutex::new(RequestRegistry::default()),
request_tasks: TaskTracker::new(),
request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)),
active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)),
closing: AtomicBool::new(false),
peer: Arc::clone(&peer),
});
let writer_disconnected = peer.disconnection_token();
let writer_task = tokio::spawn(async move {
loop {
tokio::select! {
_ = writer_disconnected.cancelled() => return Ok::<(), anyhow::Error>(()),
frame = outgoing_rx.recv() => {
let Some(frame) = frame else {
return Ok(());
};
if let Err(err) = writer.write_frame(&frame).await {
return Err(
anyhow::Error::new(err)
.context("failed to write code-mode host message")
);
}
}
}
}
});
let writer_peer = Arc::clone(&peer);
let writer_supervisor = tokio::spawn(async move {
match writer_task.await {
Ok(Ok(())) if !writer_peer.is_disconnected() => {
writer_peer.fail("code-mode writer task exited unexpectedly".to_string());
}
Ok(Ok(())) => {}
Ok(Err(err)) => {
writer_peer.fail(format!("code-mode writer task failed: {err:#}"));
}
Err(err) => {
writer_peer.fail(format!("code-mode writer task failed: {err}"));
}
}
});
let input_result = async {
loop {
let message = tokio::select! {
_ = peer.disconnected() => break,
message = reader.read::<ClientToHost>() => message
.context("failed to read code-mode client message")?,
};
let Some(message) = message else {
break;
};
match message {
ClientToHost::ClientHello(_) => {
anyhow::bail!("received a second code-mode client hello");
}
ClientToHost::Request { id, request } => {
state.spawn_request(id, request)?;
}
ClientToHost::CancelRequest { id } => {
state.cancel_request(id);
}
ClientToHost::DelegateResponse { id, result } => {
peer.complete(id, result.into_result()).await;
}
}
}
Ok::<(), anyhow::Error>(())
}
.await;
peer.disconnect();
if tokio::time::timeout(SHUTDOWN_TIMEOUT, state.disconnect())
.await
.is_err()
{
peer.fail("timed out shutting down code-mode host state".to_string());
}
drop(state);
tokio::time::timeout(SHUTDOWN_TIMEOUT, writer_supervisor)
.await
.context("timed out supervising code-mode writer task")?
.context("code-mode writer supervisor task failed")?;
let failure = peer.failure();
drop(peer);
input_result?;
if let Some(failure) = failure {
anyhow::bail!(failure);
}
Ok(())
}
async fn negotiate<R, W>(reader: &mut FramedReader<R>, writer: &mut FramedWriter<W>) -> Result<bool>
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
let Some(first_message) = reader
.read::<ClientToHost>()
.await
.context("failed to read code-mode client hello")?
else {
return Ok(false);
};
let ClientToHost::ClientHello(client_hello) = first_message else {
writer
.write(&HostToClient::HandshakeRejected {
reason: HandshakeRejectReason::InvalidHello {
message: "first message must be connection/hello".to_string(),
},
})
.await
.context("failed to reject invalid code-mode client hello")?;
return Ok(false);
};
let supported_versions = SupportedProtocolVersions::try_new([ProtocolVersion::V1])?;
if !client_hello
.supported_versions()
.contains(ProtocolVersion::V1)
{
writer
.write(&HostToClient::HandshakeRejected {
reason: HandshakeRejectReason::NoCompatibleVersion { supported_versions },
})
.await
.context("failed to reject incompatible code-mode client")?;
return Ok(false);
}
let host_capabilities = CapabilitySet::empty();
if let Some(capability) = client_hello
.required_capabilities()
.iter()
.find(|capability| !host_capabilities.contains(capability))
{
writer
.write(&HostToClient::HandshakeRejected {
reason: HandshakeRejectReason::MissingRequiredCapability {
capability: capability.clone(),
},
})
.await
.context("failed to reject unsupported code-mode capability")?;
return Ok(false);
}
writer
.write(&HostToClient::HostHello(HostHello::new(
ProtocolVersion::V1,
host_capabilities,
)))
.await
.context("failed to write code-mode host hello")?;
Ok(true)
}
struct HostState {
sessions: Mutex<HashMap<SessionId, Arc<InProcessCodeModeSession>>>,
seen_session_ids: Mutex<SeenSessionIds>,
requests: Mutex<RequestRegistry>,
request_tasks: TaskTracker,
request_permits: Arc<Semaphore>,
active_cell_permits: Arc<Semaphore>,
closing: AtomicBool,
peer: Arc<HostPeer>,
}
impl HostState {
fn spawn_request(
self: &Arc<Self>,
request_id: RequestId,
request: HostRequest,
) -> Result<(), anyhow::Error> {
let cancellation = self
.requests
.lock()
.unwrap_or_else(PoisonError::into_inner)
.start(request_id, RequestKind::from(&request))?;
let Ok(permit) = Arc::clone(&self.request_permits).try_acquire_owned() else {
self.respond(
request_id,
Err("code-mode host has too many in-flight requests".to_string()),
);
self.finish_request(request_id);
return Ok(());
};
let state = Arc::clone(self);
let request_task = self.request_tasks.spawn(async move {
let _permit = permit;
state
.handle_request(request_id, request, cancellation)
.await;
state.finish_request(request_id);
});
self.supervise_request_task(request_task);
Ok(())
}
fn supervise_request_task(&self, task: tokio::task::JoinHandle<()>) {
let peer = Arc::clone(&self.peer);
tokio::spawn(async move {
if let Err(err) = task.await {
peer.fail(format!("code-mode request task failed: {err}"));
}
});
}
async fn handle_request(
&self,
request_id: RequestId,
request: HostRequest,
cancellation: CancellationToken,
) {
if self.closing.load(Ordering::Acquire) {
self.respond(
request_id,
Err("code-mode host is shutting down".to_string()),
);
return;
}
match request {
HostRequest::OpenSession { session_id } => {
let result = self
.open_session(session_id.clone())
.map(|()| HostResponse::SessionReady { session_id });
self.respond(request_id, result);
}
HostRequest::Execute {
session_id,
request,
} => {
if cancellation.is_cancelled() {
self.respond(request_id, Err("code-mode request cancelled".to_string()));
return;
}
let request = match request.try_into() {
Ok(request) => request,
Err(err) => {
self.respond(
request_id,
Err(format!("invalid code-mode execute request: {err}")),
);
return;
}
};
let session = match self.session(&session_id) {
Ok(session) => session,
Err(err) => {
self.respond(request_id, Err(err));
return;
}
};
let Ok(active_cell_permit) =
Arc::clone(&self.active_cell_permits).try_acquire_owned()
else {
self.respond(
request_id,
Err("code-mode host has too many active cells".to_string()),
);
return;
};
let result = session.execute(request).await;
match result {
Ok(started) => {
let cell_id = started.cell_id.clone();
self.respond(
request_id,
Ok(HostResponse::ExecutionStarted {
cell_id: cell_id.into(),
}),
);
let initial_response_sent = self.peer.start_cell(
session_id,
request_id,
started,
active_cell_permit,
);
let _ = initial_response_sent.await;
}
Err(err) => self.respond(request_id, Err(err)),
}
}
HostRequest::Wait {
session_id,
request,
} => {
let result = match self.session(&session_id) {
Ok(session) => {
tokio::select! {
biased;
_ = cancellation.cancelled() => {
Err("code-mode request cancelled".to_string())
}
result = session.wait(request.into()) => result.map(|outcome| {
HostResponse::WaitCompleted {
outcome: outcome.into(),
}
}),
}
}
Err(err) => Err(err),
};
self.respond(request_id, result);
}
HostRequest::Terminate {
session_id,
cell_id,
} => {
let result = match self.session(&session_id) {
Ok(session) => session.terminate(cell_id.into()).await.map(|outcome| {
HostResponse::WaitCompleted {
outcome: outcome.into(),
}
}),
Err(err) => Err(err),
};
self.respond(request_id, result);
}
HostRequest::ShutdownSession { session_id } => {
let session = self
.sessions
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(&session_id);
let result = match session {
Some(session) => match session.shutdown().await {
Ok(()) => {
self.peer.wait_for_session_cells(&session_id).await;
Ok(HostResponse::SessionClosed { session_id })
}
Err(err) => Err(err),
},
None => Err(format!("unknown code-mode session {session_id}")),
};
self.respond(request_id, result);
}
}
}
fn open_session(&self, session_id: SessionId) -> Result<(), String> {
let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner);
if sessions.contains_key(&session_id) {
return Err(format!(
"code-mode session ID `{session_id}` is already open"
));
}
if self.closing.load(Ordering::Acquire) {
return Err("code-mode host is shutting down".to_string());
}
if !self
.seen_session_ids
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remember(session_id.clone())
{
return Err(format!("code-mode session ID `{session_id}` was reused"));
}
let delegate = Arc::new(RemoteDelegate::new(
session_id.clone(),
Arc::clone(&self.peer),
));
let peer = Arc::downgrade(&self.peer);
let task_failure_handler = Arc::new(move |reason| {
if let Some(peer) = peer.upgrade() {
peer.fail(reason);
}
});
sessions.insert(
session_id,
Arc::new(
InProcessCodeModeSession::with_delegate_and_task_failure_handler(
delegate,
task_failure_handler,
),
),
);
Ok(())
}
fn session(&self, session_id: &SessionId) -> Result<Arc<InProcessCodeModeSession>, String> {
self.sessions
.lock()
.unwrap_or_else(PoisonError::into_inner)
.get(session_id)
.cloned()
.ok_or_else(|| format!("unknown code-mode session {session_id}"))
}
fn respond(&self, id: RequestId, result: Result<HostResponse, String>) {
self.peer.respond(id, result);
}
fn cancel_request(&self, request_id: RequestId) {
self.requests
.lock()
.unwrap_or_else(PoisonError::into_inner)
.cancel(request_id);
}
fn finish_request(&self, request_id: RequestId) {
self.requests
.lock()
.unwrap_or_else(PoisonError::into_inner)
.finish(request_id);
}
async fn disconnect(&self) {
self.closing.store(true, Ordering::Release);
self.requests
.lock()
.unwrap_or_else(PoisonError::into_inner)
.cancel_all();
self.request_tasks.close();
self.request_tasks.wait().await;
let sessions = self
.sessions
.lock()
.unwrap_or_else(PoisonError::into_inner)
.drain()
.map(|(_, session)| session)
.collect::<Vec<_>>();
for session in sessions {
let _ = session.shutdown().await;
}
}
}
#[derive(Clone, Copy)]
enum RequestKind {
OpenSession,
Execute,
Wait,
Terminate,
ShutdownSession,
}
impl RequestKind {
fn from(request: &HostRequest) -> Self {
match request {
HostRequest::OpenSession { .. } => Self::OpenSession,
HostRequest::Execute { .. } => Self::Execute,
HostRequest::Wait { .. } => Self::Wait,
HostRequest::Terminate { .. } => Self::Terminate,
HostRequest::ShutdownSession { .. } => Self::ShutdownSession,
}
}
fn is_cancellable(self) -> bool {
matches!(self, Self::Execute | Self::Wait)
}
}
struct ActiveRequest {
kind: RequestKind,
cancellation: CancellationToken,
}
#[derive(Default)]
struct RequestRegistry {
active: HashMap<RequestId, ActiveRequest>,
recent: HashSet<RequestId>,
recent_order: VecDeque<RequestId>,
}
impl RequestRegistry {
fn start(
&mut self,
request_id: RequestId,
kind: RequestKind,
) -> Result<CancellationToken, anyhow::Error> {
if self.active.contains_key(&request_id) || self.recent.contains(&request_id) {
anyhow::bail!("duplicate code-mode request ID {request_id:?}");
}
let cancellation = CancellationToken::new();
self.active.insert(
request_id,
ActiveRequest {
kind,
cancellation: cancellation.clone(),
},
);
Ok(cancellation)
}
fn cancel(&self, request_id: RequestId) {
if let Some(request) = self.active.get(&request_id)
&& request.kind.is_cancellable()
{
request.cancellation.cancel();
}
}
fn finish(&mut self, request_id: RequestId) {
if self.active.remove(&request_id).is_none() {
return;
}
self.recent.insert(request_id);
self.recent_order.push_back(request_id);
while self.recent_order.len() > MAX_RECENT_REQUEST_IDS {
if let Some(expired) = self.recent_order.pop_front() {
self.recent.remove(&expired);
}
}
}
fn cancel_all(&self) {
for request in self.active.values() {
request.cancellation.cancel();
}
}
}
#[derive(Default)]
struct SeenSessionIds {
ids: HashSet<SessionId>,
order: VecDeque<SessionId>,
}
impl SeenSessionIds {
fn remember(&mut self, session_id: SessionId) -> bool {
if !self.ids.insert(session_id.clone()) {
return false;
}
self.order.push_back(session_id);
while self.order.len() > MAX_RECENT_SESSION_IDS {
if let Some(expired) = self.order.pop_front() {
self.ids.remove(&expired);
}
}
true
}
}
#[cfg(test)]
#[path = "host_tests.rs"]
mod tests;
+4 -1
View File
@@ -1 +1,4 @@
fn main() {}
#[tokio::main(flavor = "current_thread")]
async fn main() -> anyhow::Result<()> {
codex_code_mode_host::run_stdio().await
}
+541
View File
@@ -0,0 +1,541 @@
use std::collections::HashMap;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::PoisonError;
use std::sync::atomic::AtomicI64;
use std::sync::atomic::Ordering;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::host::DelegateRequest;
use codex_code_mode_protocol::host::DelegateRequestId;
use codex_code_mode_protocol::host::DelegateResponse;
use codex_code_mode_protocol::host::EncodedFrame;
use codex_code_mode_protocol::host::HostToClient;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SessionId;
use codex_code_mode_protocol::host::WireResult;
use tokio::sync::Mutex;
use tokio::sync::Notify;
use tokio::sync::OwnedSemaphorePermit;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
const CELL_MESSAGE_CAPACITY: usize = 128;
const MAX_PENDING_DELEGATE_CALLS: usize = 256;
pub(super) struct HostPeer {
outgoing_tx: mpsc::Sender<EncodedFrame>,
pending: Mutex<HashMap<DelegateRequestId, PendingDelegate>>,
delegate_permits: Arc<Semaphore>,
cell_routes: StdMutex<HashMap<(SessionId, CellId), CellRoute>>,
cell_routes_changed: Notify,
next_request_id: AtomicI64,
disconnected: CancellationToken,
failure: StdMutex<Option<String>>,
}
struct PendingDelegate {
response_tx: oneshot::Sender<Result<DelegateResponse, String>>,
dispatched: bool,
_permit: OwnedSemaphorePermit,
}
enum CellRoute {
Pending(VecDeque<CellMessage>),
Active(mpsc::Sender<CellMessage>),
}
enum CellMessage {
Delegate {
id: DelegateRequestId,
request: DelegateRequest,
dispatched_tx: oneshot::Sender<Result<(), String>>,
},
Closed,
}
impl HostPeer {
pub(super) fn new(outgoing_tx: mpsc::Sender<EncodedFrame>) -> Self {
Self {
outgoing_tx,
pending: Mutex::new(HashMap::new()),
delegate_permits: Arc::new(Semaphore::new(MAX_PENDING_DELEGATE_CALLS)),
cell_routes: StdMutex::new(HashMap::new()),
cell_routes_changed: Notify::new(),
next_request_id: AtomicI64::new(1),
disconnected: CancellationToken::new(),
failure: StdMutex::new(None),
}
}
pub(super) fn send(&self, message: HostToClient) -> Result<(), PeerSendError> {
let frame = EncodedFrame::encode(&message)
.map_err(|err| PeerSendError::Payload(err.to_string()))?;
self.send_frame(frame)
}
pub(super) fn respond(
&self,
id: RequestId,
result: Result<codex_code_mode_protocol::host::HostResponse, String>,
) {
let message = HostToClient::Response {
id,
result: WireResult::from_result(result),
};
if let Err(PeerSendError::Payload(err)) = self.send(message) {
let _ = self.send(HostToClient::Response {
id,
result: WireResult::Err {
message: format!("code-mode host response exceeds the IPC frame limit: {err}"),
},
});
}
}
fn initial_response(
&self,
id: RequestId,
result: Result<codex_code_mode_protocol::host::WireRuntimeResponse, String>,
) {
let message = HostToClient::InitialResponse {
id,
result: WireResult::from_result(result),
};
if let Err(PeerSendError::Payload(err)) = self.send(message) {
let _ = self.send(HostToClient::InitialResponse {
id,
result: WireResult::Err {
message: format!(
"code-mode initial response exceeds the IPC frame limit: {err}"
),
},
});
}
}
pub(super) async fn call(
self: &Arc<Self>,
session_id: SessionId,
request: DelegateRequest,
cancellation_token: CancellationToken,
) -> Result<DelegateResponse, String> {
if self.disconnected.is_cancelled() {
return Err("code-mode client connection closed".to_string());
}
let Ok(permit) = Arc::clone(&self.delegate_permits).try_acquire_owned() else {
return Err("code-mode host has too many pending delegate calls".to_string());
};
let id = DelegateRequestId::new(self.next_request_id.fetch_add(1, Ordering::Relaxed));
let (response_tx, response_rx) = oneshot::channel();
self.pending.lock().await.insert(
id,
PendingDelegate {
response_tx,
dispatched: false,
_permit: permit,
},
);
let mut pending = PendingDelegateRequest::new(Arc::clone(self), id);
let cell_id = match &request {
DelegateRequest::InvokeTool { invocation } => invocation.cell_id.clone().into(),
DelegateRequest::Notify { cell_id, .. } => cell_id.clone().into(),
};
let (dispatched_tx, dispatched_rx) = oneshot::channel();
if let Err(err) = self.route_cell_message(
(session_id, cell_id),
CellMessage::Delegate {
id,
request,
dispatched_tx,
},
) {
self.pending.lock().await.remove(&id);
pending.disarm();
return Err(err);
}
let dispatched = tokio::select! {
dispatched = dispatched_rx => dispatched.map_err(|_| {
"code-mode cell route closed before dispatching delegate request".to_string()
})?,
_ = self.disconnected.cancelled() => {
self.pending.lock().await.remove(&id);
pending.disarm();
return Err("code-mode client connection closed".to_string());
}
};
if let Err(err) = dispatched {
self.pending.lock().await.remove(&id);
pending.disarm();
return Err(err);
}
tokio::select! {
response = response_rx => {
pending.disarm();
response.map_err(|_| {
"code-mode client closed before returning delegate output".to_string()
})?
}
_ = cancellation_token.cancelled() => {
if self.remove_pending(id).await.is_some() {
let _ = self.send(HostToClient::CancelDelegateRequest { id });
}
pending.disarm();
Err("code mode delegate request cancelled".to_string())
}
_ = self.disconnected.cancelled() => {
self.pending.lock().await.remove(&id);
pending.disarm();
Err("code-mode client connection closed".to_string())
}
}
}
pub(super) async fn complete(
&self,
id: DelegateRequestId,
response: Result<DelegateResponse, String>,
) {
if let Some(pending) = self.remove_pending(id).await {
let _ = pending.response_tx.send(response);
}
}
pub(super) fn start_cell(
self: &Arc<Self>,
session_id: SessionId,
request_id: RequestId,
started: StartedCell,
active_cell_permit: OwnedSemaphorePermit,
) -> oneshot::Receiver<()> {
let (initial_response_sent_tx, initial_response_sent_rx) = oneshot::channel();
let key = (session_id, started.cell_id.clone());
let (messages_tx, messages_rx) = mpsc::channel(CELL_MESSAGE_CAPACITY);
let previous = self
.cell_routes
.lock()
.unwrap_or_else(PoisonError::into_inner)
.insert(key.clone(), CellRoute::Active(messages_tx.clone()));
match previous {
Some(CellRoute::Pending(messages)) => {
for message in messages {
if messages_tx.try_send(message).is_err() {
self.disconnect();
return initial_response_sent_rx;
}
}
}
Some(CellRoute::Active(_)) => {
self.disconnect();
return initial_response_sent_rx;
}
None => {}
}
let peer = Arc::clone(self);
self.spawn_critical("cell forwarding", async move {
drive_cell(
peer,
key,
request_id,
started,
messages_rx,
initial_response_sent_tx,
active_cell_permit,
)
.await;
});
initial_response_sent_rx
}
pub(super) fn close_cell(&self, session_id: SessionId, cell_id: CellId) {
let _ = self.route_cell_message((session_id, cell_id), CellMessage::Closed);
}
pub(super) fn disconnect(&self) {
self.disconnected.cancel();
}
pub(super) fn fail(&self, reason: String) {
let mut failure = self.failure.lock().unwrap_or_else(PoisonError::into_inner);
if failure.is_none() {
*failure = Some(reason);
}
drop(failure);
self.disconnect();
}
pub(super) fn failure(&self) -> Option<String> {
self.failure
.lock()
.unwrap_or_else(PoisonError::into_inner)
.clone()
}
pub(super) fn is_disconnected(&self) -> bool {
self.disconnected.is_cancelled()
}
pub(super) async fn disconnected(&self) {
self.disconnected.cancelled().await;
}
pub(super) fn disconnection_token(&self) -> CancellationToken {
self.disconnected.clone()
}
pub(super) async fn wait_for_session_cells(&self, session_id: &SessionId) {
loop {
let changed = self.cell_routes_changed.notified();
if !self
.cell_routes
.lock()
.unwrap_or_else(PoisonError::into_inner)
.keys()
.any(|(route_session_id, _)| route_session_id == session_id)
{
return;
}
tokio::select! {
_ = changed => {}
_ = self.disconnected.cancelled() => return,
}
}
}
async fn send_delegate_if_pending(
&self,
id: DelegateRequestId,
session_id: SessionId,
request: DelegateRequest,
dispatched_tx: oneshot::Sender<Result<(), String>>,
) {
let result = {
let mut pending = self.pending.lock().await;
let Some(pending) = pending.get_mut(&id) else {
let _ = dispatched_tx.send(Err(
"code-mode delegate request was cancelled before dispatch".to_string(),
));
return;
};
match self.send(HostToClient::DelegateRequest {
id,
session_id,
request,
}) {
Ok(()) => {
pending.dispatched = true;
Ok(())
}
Err(err) => Err(err.to_string()),
}
};
let _ = dispatched_tx.send(result);
}
fn route_cell_message(
&self,
key: (SessionId, CellId),
message: CellMessage,
) -> Result<(), String> {
use std::collections::hash_map::Entry;
let result = match self
.cell_routes
.lock()
.unwrap_or_else(PoisonError::into_inner)
.entry(key)
{
Entry::Occupied(mut entry) => match entry.get_mut() {
CellRoute::Pending(messages) if messages.len() < CELL_MESSAGE_CAPACITY => {
messages.push_back(message);
Ok(())
}
CellRoute::Pending(_) => Err("code-mode cell message queue is full".to_string()),
CellRoute::Active(sender) => sender
.try_send(message)
.map_err(|_| "code-mode cell message queue is unavailable".to_string()),
},
Entry::Vacant(entry) => {
entry.insert(CellRoute::Pending(VecDeque::from([message])));
Ok(())
}
};
if result.is_err() {
self.disconnect();
}
result
}
async fn remove_pending(&self, id: DelegateRequestId) -> Option<PendingDelegate> {
self.pending.lock().await.remove(&id)
}
pub(super) fn spawn_critical<F>(self: &Arc<Self>, task_name: &'static str, future: F)
where
F: std::future::Future<Output = ()> + Send + 'static,
{
let task = tokio::spawn(future);
let peer = Arc::clone(self);
tokio::spawn(async move {
if let Err(err) = task.await {
peer.fail(format!("code-mode {task_name} task failed: {err}"));
}
});
}
fn send_frame(&self, frame: EncodedFrame) -> Result<(), PeerSendError> {
match self.outgoing_tx.try_send(frame) {
Ok(()) => Ok(()),
Err(mpsc::error::TrySendError::Full(_)) => {
self.disconnect();
Err(PeerSendError::Unavailable(
"code-mode host outgoing queue is full".to_string(),
))
}
Err(mpsc::error::TrySendError::Closed(_)) => {
self.disconnect();
Err(PeerSendError::Unavailable(
"code-mode client connection closed".to_string(),
))
}
}
}
}
async fn drive_cell(
peer: Arc<HostPeer>,
key: (SessionId, CellId),
request_id: RequestId,
started: StartedCell,
mut messages_rx: mpsc::Receiver<CellMessage>,
initial_response_sent_tx: oneshot::Sender<()>,
_active_cell_permit: OwnedSemaphorePermit,
) {
let mut initial_response_sent_tx = Some(initial_response_sent_tx);
let initial_response = started.initial_response();
tokio::pin!(initial_response);
let closed = loop {
tokio::select! {
biased;
result = &mut initial_response => {
peer.initial_response(request_id, result.map(Into::into));
if let Some(initial_response_sent_tx) = initial_response_sent_tx.take() {
let _ = initial_response_sent_tx.send(());
}
break false;
}
message = messages_rx.recv() => match message {
Some(CellMessage::Delegate {
id,
request,
dispatched_tx,
}) => {
peer.send_delegate_if_pending(id, key.0.clone(), request, dispatched_tx).await;
}
Some(CellMessage::Closed) | None => break true,
},
_ = peer.disconnected.cancelled() => {
peer.remove_cell_route(&key);
return;
}
}
};
if closed {
peer.initial_response(request_id, initial_response.await.map(Into::into));
if let Some(initial_response_sent_tx) = initial_response_sent_tx.take() {
let _ = initial_response_sent_tx.send(());
}
} else {
loop {
tokio::select! {
message = messages_rx.recv() => match message {
Some(CellMessage::Delegate {
id,
request,
dispatched_tx,
}) => {
peer.send_delegate_if_pending(id, key.0.clone(), request, dispatched_tx).await;
}
Some(CellMessage::Closed) | None => break,
},
_ = peer.disconnected.cancelled() => {
peer.remove_cell_route(&key);
return;
}
}
}
}
let _ = peer.send(HostToClient::CellClosed {
session_id: key.0.clone(),
cell_id: (&key.1).into(),
});
peer.remove_cell_route(&key);
}
impl HostPeer {
fn remove_cell_route(&self, key: &(SessionId, CellId)) {
let removed = self
.cell_routes
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(key);
if removed.is_some() {
self.cell_routes_changed.notify_waiters();
}
}
}
pub(super) enum PeerSendError {
Payload(String),
Unavailable(String),
}
impl std::fmt::Display for PeerSendError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Payload(message) | Self::Unavailable(message) => formatter.write_str(message),
}
}
}
struct PendingDelegateRequest {
peer: Arc<HostPeer>,
id: Option<DelegateRequestId>,
}
impl PendingDelegateRequest {
fn new(peer: Arc<HostPeer>, id: DelegateRequestId) -> Self {
Self { peer, id: Some(id) }
}
fn disarm(&mut self) {
self.id = None;
}
}
impl Drop for PendingDelegateRequest {
fn drop(&mut self) {
let Some(id) = self.id.take() else {
return;
};
let peer = Arc::clone(&self.peer);
tokio::spawn(async move {
if let Some(pending) = peer.remove_pending(id).await
&& pending.dispatched
{
let _ = peer.send(HostToClient::CancelDelegateRequest { id });
}
});
}
}
#[cfg(test)]
#[path = "peer_tests.rs"]
mod tests;
+95
View File
@@ -0,0 +1,95 @@
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::RuntimeResponse;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::host::DelegateRequest;
use codex_code_mode_protocol::host::RequestId;
use codex_code_mode_protocol::host::SessionId;
use pretty_assertions::assert_eq;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::oneshot::error::TryRecvError;
use tokio_util::sync::CancellationToken;
use super::HostPeer;
use super::MAX_PENDING_DELEGATE_CALLS;
fn session_id(value: &str) -> SessionId {
SessionId::new(value).expect("session ID")
}
#[tokio::test]
async fn start_cell_reports_when_initial_response_is_enqueued() {
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 4);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let cell_id = CellId::new("cell-1".to_string());
let (response_tx, response_rx) = oneshot::channel();
let started = StartedCell::new(cell_id.clone(), response_rx);
let active_cell_permits = Arc::new(Semaphore::new(/*permits*/ 1));
let active_cell_permit = Arc::clone(&active_cell_permits)
.try_acquire_owned()
.expect("active cell permit");
let mut initial_response_sent = peer.start_cell(
session_id("session-1"),
RequestId::new(/*value*/ 1),
started,
active_cell_permit,
);
assert_eq!(initial_response_sent.try_recv(), Err(TryRecvError::Empty));
response_tx
.send(RuntimeResponse::Result {
cell_id: cell_id.clone(),
content_items: Vec::new(),
error_text: None,
})
.expect("initial response receiver");
initial_response_sent
.await
.expect("initial response completion");
outgoing_rx.recv().await.expect("initial response frame");
assert_eq!(active_cell_permits.available_permits(), 0);
peer.close_cell(session_id("session-1"), cell_id);
let permit = tokio::time::timeout(
Duration::from_secs(1),
Arc::clone(&active_cell_permits).acquire_owned(),
)
.await
.expect("cell permit should be released")
.expect("cell permit semaphore should remain open");
drop(permit);
}
#[tokio::test]
async fn pending_delegate_limit_rejects_call_without_disconnecting() {
let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1);
let peer = Arc::new(HostPeer::new(outgoing_tx));
let permits = Arc::clone(&peer.delegate_permits)
.acquire_many_owned(MAX_PENDING_DELEGATE_CALLS as u32)
.await
.expect("delegate permits");
let result = peer
.call(
session_id("session-1"),
DelegateRequest::Notify {
call_id: "call-1".to_string(),
cell_id: CellId::new("cell-1".to_string()).into(),
text: "hello".to_string(),
},
CancellationToken::new(),
)
.await;
assert_eq!(
result,
Err("code-mode host has too many pending delegate calls".to_string())
);
assert!(!peer.is_disconnected());
drop(permits);
}
+37 -17
View File
@@ -11,6 +11,36 @@ use tokio::io::AsyncWriteExt;
/// Maximum JSON payload size accepted for one IPC frame.
pub const MAX_FRAME_BYTES: usize = 64 * 1024 * 1024;
/// A serialized IPC frame that has already passed the payload size limit.
#[derive(Clone, Debug)]
pub struct EncodedFrame {
payload: Vec<u8>,
}
impl EncodedFrame {
pub fn encode<T>(message: &T) -> io::Result<Self>
where
T: Serialize,
{
let payload = serde_json::to_vec(message).map_err(|err| {
io::Error::new(
io::ErrorKind::InvalidData,
format!("failed to encode code-mode IPC frame: {err}"),
)
})?;
if payload.len() > MAX_FRAME_BYTES {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"code-mode IPC frame length {} exceeds {MAX_FRAME_BYTES} bytes",
payload.len()
),
));
}
Ok(Self { payload })
}
}
/// Decodes JSON messages prefixed by a four-byte little-endian payload length.
pub struct FramedReader<R> {
reader: R,
@@ -72,22 +102,12 @@ where
where
T: Serialize,
{
let payload = serde_json::to_vec(message).map_err(|err| {
io::Error::new(
io::ErrorKind::InvalidData,
format!("failed to encode code-mode IPC frame: {err}"),
)
})?;
if payload.len() > MAX_FRAME_BYTES {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"code-mode IPC frame length {} exceeds {MAX_FRAME_BYTES} bytes",
payload.len()
),
));
}
let length = u32::try_from(payload.len()).map_err(|_| {
self.write_frame(&EncodedFrame::encode(message)?).await
}
/// Writes and flushes a frame encoded before it entered an I/O queue.
pub async fn write_frame(&mut self, frame: &EncodedFrame) -> io::Result<()> {
let length = u32::try_from(frame.payload.len()).map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidData,
"code-mode IPC frame length exceeds u32",
@@ -95,7 +115,7 @@ where
})?;
self.writer.write_all(&length.to_le_bytes()).await?;
self.writer.write_all(&payload).await?;
self.writer.write_all(&frame.payload).await?;
self.writer.flush().await
}
}
@@ -365,6 +365,16 @@ fn client_to_host_v1_variants_are_pinned() {
}),
);
}
assert_wire_round_trip(
ClientToHost::CancelRequest {
id: request_id(/*value*/ 9),
},
json!({
"type": "operation/cancel",
"id": 9,
}),
);
}
#[test]
@@ -132,6 +132,8 @@ pub enum ClientToHost {
ClientHello(ClientHello),
#[serde(rename = "operation/request")]
Request { id: RequestId, request: HostRequest },
#[serde(rename = "operation/cancel")]
CancelRequest { id: RequestId },
#[serde(rename = "delegate/response")]
DelegateResponse {
id: DelegateRequestId,
@@ -11,6 +11,7 @@ mod message;
mod payload;
mod types;
pub use codec::EncodedFrame;
pub use codec::FramedReader;
pub use codec::FramedWriter;
pub use codec::MAX_FRAME_BYTES;
+1
View File
@@ -19,6 +19,7 @@ workspace = true
codex-code-mode-protocol = { workspace = true }
codex-protocol = { workspace = true }
deno_core_icudata = { workspace = true }
futures = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] }
tokio-util = { workspace = true, features = ["rt"] }
+63 -12
View File
@@ -1,11 +1,14 @@
use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use futures::FutureExt;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use tracing::warn;
use super::CellHost;
use super::CellToolCall;
use crate::TaskFailureHandler;
use crate::runtime::RuntimeCommand;
#[derive(Clone, Copy)]
@@ -20,10 +23,20 @@ pub(super) fn spawn_notification<H: CellHost>(
call_id: String,
text: String,
cancellation_token: CancellationToken,
task_failure_handler: Option<TaskFailureHandler>,
) {
tasks.spawn(async move {
if let Err(err) = host.notify(call_id, text, cancellation_token).await {
warn!("failed to deliver code mode notification: {err}");
let callback =
AssertUnwindSafe(async move { host.notify(call_id, text, cancellation_token).await })
.catch_unwind()
.await;
match callback {
Ok(Ok(())) => {}
Ok(Err(err)) => warn!("failed to deliver code mode notification: {err}"),
Err(_) => report_task_failure(
task_failure_handler.as_ref(),
"code mode notification task panicked".to_string(),
),
}
});
}
@@ -34,14 +47,32 @@ pub(super) fn spawn_tool<H: CellHost>(
invocation: CellToolCall,
runtime_tx: std::sync::mpsc::Sender<RuntimeCommand>,
cancellation_token: CancellationToken,
task_failure_handler: Option<TaskFailureHandler>,
) {
tasks.spawn(async move {
let id = invocation.id.clone();
let command = match host.invoke_tool(invocation, cancellation_token).await {
Ok(result) => RuntimeCommand::ToolResponse { id, result },
Err(error_text) => RuntimeCommand::ToolError { id, error_text },
let callback =
AssertUnwindSafe(async move { host.invoke_tool(invocation, cancellation_token).await })
.catch_unwind()
.await;
let (command, failure_reason) = match callback {
Ok(Ok(result)) => (RuntimeCommand::ToolResponse { id, result }, None),
Ok(Err(error_text)) => (RuntimeCommand::ToolError { id, error_text }, None),
Err(_) => {
let failure_reason = "code mode tool task panicked".to_string();
(
RuntimeCommand::ToolError {
id,
error_text: failure_reason.clone(),
},
Some(failure_reason),
)
}
};
let _ = runtime_tx.send(command);
if let Some(failure_reason) = failure_reason {
report_task_failure(task_failure_handler.as_ref(), failure_reason);
}
});
}
@@ -50,28 +81,48 @@ pub(super) async fn finish_callbacks(
notification_tasks: &mut JoinSet<()>,
tool_tasks: &mut JoinSet<()>,
completion: CallbackCompletion,
task_failure_handler: Option<&TaskFailureHandler>,
) {
if matches!(completion, CallbackCompletion::Cancel) {
cancellation_token.cancel();
}
drain_tasks(notification_tasks, "notification").await;
drain_tasks(notification_tasks, "notification", task_failure_handler).await;
cancellation_token.cancel();
drain_tasks(tool_tasks, "tool").await;
drain_tasks(tool_tasks, "tool", task_failure_handler).await;
}
pub(super) fn log_task_result(
pub(super) fn report_task_result(
task_result: Option<Result<(), tokio::task::JoinError>>,
description: &str,
task_failure_handler: Option<&TaskFailureHandler>,
) {
if let Some(Err(err)) = task_result
&& !err.is_cancelled()
{
warn!("code mode {description} task failed: {err}");
report_task_failure(
task_failure_handler,
format!("code mode {description} task failed: {err}"),
);
}
}
async fn drain_tasks(tasks: &mut JoinSet<()>, description: &str) {
while let Some(result) = tasks.join_next().await {
log_task_result(Some(result), description);
fn report_task_failure(task_failure_handler: Option<&TaskFailureHandler>, failure_reason: String) {
warn!("{failure_reason}");
if let Some(task_failure_handler) = task_failure_handler {
task_failure_handler(failure_reason);
}
}
async fn drain_tasks(
tasks: &mut JoinSet<()>,
description: &str,
task_failure_handler: Option<&TaskFailureHandler>,
) {
while let Some(result) = tasks.join_next().await {
report_task_result(Some(result), description, task_failure_handler);
}
}
#[cfg(test)]
#[path = "callbacks_tests.rs"]
mod tests;
@@ -0,0 +1,132 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::mpsc as std_mpsc;
use std::time::Duration;
use pretty_assertions::assert_eq;
use serde_json::Value as JsonValue;
use tokio::sync::mpsc;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use super::*;
use crate::cell_actor::CellState;
use crate::cell_actor::CompletionCommit;
use crate::runtime::RuntimeCommand;
use crate::session_runtime::CellEvent;
use crate::session_runtime::ToolKind;
use crate::session_runtime::ToolName;
struct PanickingCallbackHost;
impl CellHost for PanickingCallbackHost {
async fn invoke_tool(
&self,
_invocation: CellToolCall,
_cancellation_token: CancellationToken,
) -> Result<JsonValue, String> {
panic!("tool callback panic probe");
}
async fn notify(
&self,
_call_id: String,
_text: String,
_cancellation_token: CancellationToken,
) -> Result<(), String> {
panic!("notification callback panic probe");
}
async fn commit_completion(
&self,
_stored_value_writes: HashMap<String, JsonValue>,
_event: CellEvent,
_pending_initial_yield_items: Option<Vec<crate::session_runtime::OutputItem>>,
_cell_state: Arc<CellState>,
) -> CompletionCommit {
panic!("unexpected completion commit");
}
async fn closed(&self) {}
}
#[tokio::test]
async fn tool_callback_panic_rejects_the_js_promise_and_reports_failure() {
let mut tasks = JoinSet::new();
let (runtime_tx, runtime_rx) = std_mpsc::channel();
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
spawn_tool(
&mut tasks,
Arc::new(PanickingCallbackHost),
CellToolCall {
id: "tool-1".to_string(),
name: ToolName {
name: "panic".to_string(),
namespace: None,
},
kind: ToolKind::Function,
input: None,
},
runtime_tx,
CancellationToken::new(),
Some(Arc::new(move |reason| {
let _ = failure_tx.send(reason);
})),
);
tasks
.join_next()
.await
.expect("tool callback task")
.expect("tool callback wrapper");
let command = runtime_rx
.recv_timeout(Duration::from_secs(1))
.expect("tool error command");
let RuntimeCommand::ToolError { id, error_text } = command else {
panic!("expected a tool error command");
};
assert_eq!(id, "tool-1");
assert_eq!(error_text, "code mode tool task panicked");
assert_eq!(failure_rx.recv().await, Some(error_text));
}
#[tokio::test]
async fn notification_callback_panic_reports_failure() {
let mut tasks = JoinSet::new();
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
spawn_notification(
&mut tasks,
Arc::new(PanickingCallbackHost),
"notify-1".to_string(),
"hello".to_string(),
CancellationToken::new(),
Some(Arc::new(move |reason| {
let _ = failure_tx.send(reason);
})),
);
tasks
.join_next()
.await
.expect("notification callback task")
.expect("notification callback wrapper");
let failure_reason = failure_rx.recv().await.expect("notification failure");
assert_eq!(failure_reason, "code mode notification task panicked");
}
#[tokio::test]
async fn callback_wrapper_join_error_reports_failure() {
let task_result = tokio::spawn(async {
panic!("callback wrapper panic probe");
})
.await;
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
let task_failure_handler: TaskFailureHandler = Arc::new(move |reason| {
let _ = failure_tx.send(reason);
});
report_task_result(Some(task_result), "tool", Some(&task_failure_handler));
let failure_reason = failure_rx.recv().await.expect("wrapper failure");
assert!(failure_reason.contains("code mode tool task failed"));
}
+32 -3
View File
@@ -14,7 +14,7 @@ use tokio_util::sync::CancellationToken;
use self::callbacks::CallbackCompletion;
use self::callbacks::finish_callbacks;
use self::callbacks::log_task_result;
use self::callbacks::report_task_result;
use self::callbacks::spawn_notification;
use self::callbacks::spawn_tool;
use self::conversions::cell_tool_kind;
@@ -30,6 +30,7 @@ pub(crate) use self::types::CellToolCall;
pub(crate) use self::types::CompletionCommit;
use self::types::CompletionDelivery;
use self::types::ObservationDelivery;
use crate::TaskFailureHandler;
use crate::runtime::PendingRuntimeMode;
use crate::runtime::RuntimeCommand;
use crate::runtime::RuntimeControlCommand;
@@ -50,6 +51,7 @@ impl CellActor {
host: Arc<H>,
initial_observe_mode: ObserveMode,
cell_state: Arc<CellState>,
task_failure_handler: Option<TaskFailureHandler>,
) -> Result<
(
CellHandle,
@@ -66,6 +68,7 @@ impl CellActor {
runtime_request(request),
event_tx,
PendingRuntimeMode::PauseUntilResumed,
task_failure_handler.clone(),
)?;
let handle = CellHandle::new(command_tx, Arc::clone(&cell_state));
let task = run_cell(
@@ -82,6 +85,7 @@ impl CellActor {
mode: initial_observe_mode,
response_tx: initial_response_tx,
},
task_failure_handler,
);
let initial_response =
Box::pin(async move { initial_response_rx.await.unwrap_or(Err(CellError::Closed)) });
@@ -107,6 +111,7 @@ async fn run_cell<H: CellHost>(
mut event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
command_rx: mpsc::UnboundedReceiver<CellCommand>,
initial_observer: Observer,
task_failure_handler: Option<TaskFailureHandler>,
) {
let CellContext {
runtime_tx,
@@ -123,6 +128,7 @@ async fn run_cell<H: CellHost>(
let mut termination = false;
let mut runtime_closed = false;
let mut runtime_paused = false;
let mut runtime_failure_reported = false;
let mut yield_timer: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;
let mut notification_tasks = JoinSet::new();
let mut tool_tasks = JoinSet::new();
@@ -149,6 +155,7 @@ async fn run_cell<H: CellHost>(
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::Cancel,
task_failure_handler.as_ref(),
).await;
finish_termination(
&cell_state,
@@ -259,6 +266,7 @@ async fn run_cell<H: CellHost>(
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::Cancel,
task_failure_handler.as_ref(),
).await;
finish_termination(
&cell_state,
@@ -269,11 +277,20 @@ async fn run_cell<H: CellHost>(
);
break;
}
if !runtime_failure_reported
&& let Some(task_failure_handler) = &task_failure_handler
{
runtime_failure_reported = true;
task_failure_handler(
"code-mode V8 runtime thread ended unexpectedly".to_string(),
);
}
finish_callbacks(
&callback_cancellation_token,
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::DrainNotifications,
task_failure_handler.as_ref(),
)
.await;
let event = CellEvent::Completed {
@@ -376,6 +393,7 @@ async fn run_cell<H: CellHost>(
call_id,
text,
callback_cancellation_token.child_token(),
task_failure_handler.clone(),
);
}
RuntimeEvent::ToolCall { id, name, kind, input } => {
@@ -394,6 +412,7 @@ async fn run_cell<H: CellHost>(
},
runtime_tx.clone(),
callback_cancellation_token.child_token(),
task_failure_handler.clone(),
);
}
RuntimeEvent::Result { stored_value_writes, error_text } => {
@@ -405,6 +424,7 @@ async fn run_cell<H: CellHost>(
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::Cancel,
task_failure_handler.as_ref(),
).await;
finish_termination(
&cell_state,
@@ -420,6 +440,7 @@ async fn run_cell<H: CellHost>(
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::DrainNotifications,
task_failure_handler.as_ref(),
)
.await;
let event = CellEvent::Completed {
@@ -455,13 +476,20 @@ async fn run_cell<H: CellHost>(
}
}
}
RuntimeEvent::ThreadPanicked => {
runtime_failure_reported = true;
}
}
}
task_result = notification_tasks.join_next(), if !notification_tasks.is_empty() => {
log_task_result(task_result, "notification");
report_task_result(
task_result,
"notification",
task_failure_handler.as_ref(),
);
}
task_result = tool_tasks.join_next(), if !tool_tasks.is_empty() => {
log_task_result(task_result, "tool");
report_task_result(task_result, "tool", task_failure_handler.as_ref());
}
}
}
@@ -479,6 +507,7 @@ async fn run_cell<H: CellHost>(
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::Cancel,
task_failure_handler.as_ref(),
)
.await;
host.closed().await;
@@ -102,6 +102,18 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne
fn spawn_cell_actor_harness_with_host<H: CellHost>(
initial_observe_mode: ObserveMode,
host: Arc<H>,
) -> CellActorHarness {
spawn_cell_actor_harness_with_host_and_failure_handler(
initial_observe_mode,
host,
/*task_failure_handler*/ None,
)
}
fn spawn_cell_actor_harness_with_host_and_failure_handler<H: CellHost>(
initial_observe_mode: ObserveMode,
host: Arc<H>,
task_failure_handler: Option<TaskFailureHandler>,
) -> CellActorHarness {
let (event_tx, event_rx) = mpsc::unbounded_channel();
let (command_tx, command_rx) = mpsc::unbounded_channel();
@@ -118,6 +130,7 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
},
runtime_event_tx,
PendingRuntimeMode::PauseUntilResumed,
/*task_failure_handler*/ None,
)
.unwrap();
let (runtime_control_tx, runtime_control_rx) = std_mpsc::channel();
@@ -137,6 +150,7 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
mode: initial_observe_mode,
response_tx: initial_event_tx,
},
task_failure_handler,
));
CellActorHarness {
@@ -149,6 +163,54 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
}
}
#[tokio::test]
async fn unexpected_runtime_thread_exit_is_reported_to_the_session_owner() {
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
let harness = spawn_cell_actor_harness_with_host_and_failure_handler(
ObserveMode::YieldAfter(Duration::from_secs(60)),
Arc::new(TestHost),
Some(Arc::new(move |reason| {
let _ = failure_tx.send(reason);
})),
);
drop(harness.event_tx);
assert_eq!(
tokio::time::timeout(Duration::from_secs(1), failure_rx.recv())
.await
.expect("runtime failure timeout")
.expect("runtime failure"),
"code-mode V8 runtime thread ended unexpectedly"
);
assert!(
harness
.initial_event_rx
.await
.expect("initial event")
.is_ok()
);
harness.task.await.expect("cell task");
}
#[tokio::test]
async fn runtime_thread_panic_remains_a_cell_error_without_owner_supervision() {
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(60)));
harness
.event_tx
.send(RuntimeEvent::ThreadPanicked)
.expect("runtime panic event");
drop(harness.event_tx);
assert_eq!(
harness.initial_event_rx.await.expect("initial event"),
Ok(CellEvent::Completed {
content_items: Vec::new(),
error_text: Some("exec runtime ended unexpectedly".to_string()),
})
);
harness.task.await.expect("cell task");
}
async fn wait_for_notification(host: &RecordingHost) {
tokio::time::timeout(Duration::from_secs(1), async {
while !host.notified.load(Ordering::Acquire) {
+2
View File
@@ -3,6 +3,8 @@ mod runtime;
mod service;
mod session_runtime;
pub(crate) type TaskFailureHandler = std::sync::Arc<dyn Fn(String) + Send + Sync>;
pub use codex_code_mode_protocol::*;
pub use service::InProcessCodeModeSession;
pub use service::InProcessCodeModeSessionProvider;
+64 -1
View File
@@ -5,6 +5,8 @@ mod timers;
mod value;
use std::collections::HashMap;
use std::panic::AssertUnwindSafe;
use std::panic::catch_unwind;
use std::sync::OnceLock;
use std::sync::mpsc as std_mpsc;
use std::thread;
@@ -18,6 +20,8 @@ use codex_protocol::ToolName;
use serde_json::Value as JsonValue;
use tokio::sync::mpsc;
use crate::TaskFailureHandler;
const EXIT_SENTINEL: &str = "__codex_code_mode_exit__";
#[derive(Debug)]
@@ -63,6 +67,7 @@ pub(crate) enum RuntimeEvent {
stored_value_writes: HashMap<String, JsonValue>,
error_text: Option<String>,
},
ThreadPanicked,
}
pub(crate) fn spawn_runtime(
@@ -70,6 +75,7 @@ pub(crate) fn spawn_runtime(
request: ExecuteRequest,
event_tx: mpsc::UnboundedSender<RuntimeEvent>,
pending_mode: PendingRuntimeMode,
task_failure_handler: Option<TaskFailureHandler>,
) -> Result<
(
std_mpsc::Sender<RuntimeCommand>,
@@ -96,7 +102,7 @@ pub(crate) fn spawn_runtime(
stored_values,
};
thread::spawn(move || {
spawn_supervised_runtime_thread(event_tx.clone(), task_failure_handler, move || {
run_runtime(
config,
event_tx,
@@ -114,6 +120,21 @@ pub(crate) fn spawn_runtime(
Ok((command_tx, control_tx, isolate_handle))
}
fn spawn_supervised_runtime_thread(
event_tx: mpsc::UnboundedSender<RuntimeEvent>,
task_failure_handler: Option<TaskFailureHandler>,
runtime: impl FnOnce() + Send + 'static,
) {
thread::spawn(move || {
if catch_unwind(AssertUnwindSafe(runtime)).is_err() {
if let Some(task_failure_handler) = task_failure_handler {
task_failure_handler("code-mode V8 runtime thread panicked".to_string());
}
let _ = event_tx.send(RuntimeEvent::ThreadPanicked);
}
});
}
#[derive(Clone)]
struct RuntimeConfig {
tool_call_id: String,
@@ -336,6 +357,7 @@ mod tests {
use super::RuntimeControlCommand;
use super::RuntimeEvent;
use super::spawn_runtime;
use super::spawn_supervised_runtime_thread;
use crate::FunctionCallOutputContentItem;
fn execute_request(source: &str) -> ExecuteRequest {
@@ -348,6 +370,45 @@ mod tests {
}
}
#[tokio::test]
async fn runtime_thread_panic_before_initialization_is_reported_directly() {
let (event_tx, event_rx) = mpsc::unbounded_channel();
drop(event_rx);
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
spawn_supervised_runtime_thread(
event_tx,
Some(std::sync::Arc::new(move |reason| {
let _ = failure_tx.send(reason);
})),
|| panic!("runtime thread panic probe"),
);
assert_eq!(
tokio::time::timeout(Duration::from_secs(1), failure_rx.recv())
.await
.expect("runtime failure timeout")
.expect("runtime failure"),
"code-mode V8 runtime thread panicked"
);
}
#[tokio::test]
async fn runtime_thread_panic_is_forwarded_without_owner_supervision() {
let (event_tx, mut event_rx) = mpsc::unbounded_channel();
spawn_supervised_runtime_thread(
event_tx,
/*task_failure_handler*/ None,
|| panic!("runtime thread panic probe"),
);
assert!(matches!(
tokio::time::timeout(Duration::from_secs(1), event_rx.recv())
.await
.expect("runtime panic event timeout"),
Some(RuntimeEvent::ThreadPanicked)
));
}
#[tokio::test]
async fn terminate_execution_stops_cpu_bound_module() {
let (event_tx, mut event_rx) = mpsc::unbounded_channel();
@@ -356,6 +417,7 @@ mod tests {
execute_request("while (true) {}"),
event_tx,
PendingRuntimeMode::Continue,
/*task_failure_handler*/ None,
)
.unwrap();
@@ -398,6 +460,7 @@ await new Promise(() => {});
),
event_tx,
PendingRuntimeMode::PauseUntilResumed,
/*task_failure_handler*/ None,
)
.unwrap();
+12
View File
@@ -87,6 +87,18 @@ impl InProcessCodeModeSession {
}
}
pub fn with_delegate_and_task_failure_handler(
delegate: Arc<dyn CodeModeSessionDelegate>,
task_failure_handler: Arc<dyn Fn(String) + Send + Sync>,
) -> Self {
Self {
runtime: SessionRuntime::new_with_task_failure_handler(
Arc::new(ProtocolDelegate { delegate }),
Some(task_failure_handler),
),
}
}
pub async fn execute(&self, request: ExecuteRequest) -> Result<StartedCell, String> {
let yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS);
let started = self
+31 -9
View File
@@ -24,6 +24,7 @@ pub(crate) use self::types::SessionRuntimeDelegate;
pub(crate) use self::types::ToolDefinition;
pub(crate) use self::types::ToolKind;
pub(crate) use self::types::ToolName;
use crate::TaskFailureHandler;
use crate::cell_actor::CellActor;
use crate::cell_actor::CellError;
use crate::cell_actor::CellEventFuture;
@@ -46,11 +47,19 @@ struct Inner<D: SessionRuntimeDelegate> {
cell_tasks: TaskTracker,
shutdown_token: CancellationToken,
delegate: Arc<D>,
task_failure_handler: Option<TaskFailureHandler>,
next_cell_id: AtomicU64,
}
impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
pub(crate) fn new(delegate: Arc<D>) -> Self {
Self::new_with_task_failure_handler(delegate, /*task_failure_handler*/ None)
}
pub(crate) fn new_with_task_failure_handler(
delegate: Arc<D>,
task_failure_handler: Option<TaskFailureHandler>,
) -> Self {
Self {
inner: Arc::new(Inner {
stored_values: Mutex::new(HashMap::new()),
@@ -58,6 +67,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
cell_tasks: TaskTracker::new(),
shutdown_token: CancellationToken::new(),
delegate,
task_failure_handler,
next_cell_id: AtomicU64::new(1),
}),
}
@@ -71,7 +81,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
if self.inner.shutdown_token.is_cancelled() {
return Err(Error::ShuttingDown);
}
let cell_id = self.allocate_cell_id();
let cell_id = self.allocate_cell_id()?;
let initial_event = self
.start_cell(cell_id.clone(), request, initial_observe_mode)
.await?;
@@ -133,13 +143,14 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
Ok(())
}
fn allocate_cell_id(&self) -> CellId {
CellId::new(
self.inner
.next_cell_id
.fetch_add(1, Ordering::Relaxed)
.to_string(),
)
fn allocate_cell_id(&self) -> Result<CellId, Error> {
self.inner
.next_cell_id
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |next_cell_id| {
next_cell_id.checked_add(1)
})
.map(|cell_id| CellId::new(cell_id.to_string()))
.map_err(|_| Error::CellIdSpaceExhausted)
}
async fn start_cell(
@@ -167,10 +178,21 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
host,
initial_observe_mode,
cell_state,
self.inner.task_failure_handler.clone(),
)
.map_err(Error::Runtime)?;
cells.insert(cell_id.clone(), handle);
self.inner.cell_tasks.spawn(task);
let task = self.inner.cell_tasks.spawn(task);
if let Some(task_failure_handler) = self.inner.task_failure_handler.clone() {
let failed_cell_id = cell_id.clone();
let _failure_watcher = self.inner.cell_tasks.spawn(async move {
if let Err(err) = task.await {
task_failure_handler(format!(
"code-mode cell {failed_cell_id} task failed: {err}"
));
}
});
}
drop(cells);
Ok(map_actor_event(cell_id, initial_event))
}
@@ -15,6 +15,8 @@ use crate::cell_actor::CompletionCommit;
struct RecordingDelegate;
struct PanickingClosedDelegate;
impl SessionRuntimeDelegate for RecordingDelegate {
async fn invoke_tool(
&self,
@@ -37,6 +39,62 @@ impl SessionRuntimeDelegate for RecordingDelegate {
fn cell_closed(&self, _cell_id: &CellId) {}
}
impl SessionRuntimeDelegate for PanickingClosedDelegate {
async fn invoke_tool(
&self,
_invocation: NestedToolCall,
_cancellation_token: CancellationToken,
) -> Result<JsonValue, String> {
Ok(JsonValue::Null)
}
async fn notify(
&self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation_token: CancellationToken,
) -> Result<(), String> {
Ok(())
}
fn cell_closed(&self, _cell_id: &CellId) {
panic!("cell close panic probe");
}
}
#[tokio::test]
async fn reports_cell_actor_panics_to_the_owner() {
let (failure_tx, mut failure_rx) = tokio::sync::mpsc::unbounded_channel();
let runtime = SessionRuntime::new_with_task_failure_handler(
Arc::new(PanickingClosedDelegate),
Some(Arc::new(move |reason| {
let _ = failure_tx.send(reason);
})),
);
let started = runtime
.execute(
execute_request(r#"text("done");"#),
ObserveMode::YieldAfter(Duration::from_secs(1)),
)
.await
.expect("start cell");
assert_eq!(
started.initial_event().await,
Ok(CellEvent::Completed {
content_items: vec![OutputItem::Text {
text: "done".to_string(),
}],
error_text: None,
})
);
runtime.shutdown().await.expect("shutdown runtime");
let failure = failure_rx
.try_recv()
.expect("shutdown should wait for the cell failure watcher");
assert!(failure.contains("code-mode cell 1 task failed"));
}
#[tokio::test]
async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_load_it() {
let runtime = SessionRuntime::new(Arc::new(RecordingDelegate));
@@ -118,6 +176,26 @@ fn execute_request(source: &str) -> CreateCellRequest {
}
}
#[tokio::test]
async fn cell_id_allocation_fails_before_wrapping() {
let runtime = SessionRuntime::new(Arc::new(RecordingDelegate));
runtime
.inner
.next_cell_id
.store(u64::MAX, Ordering::Relaxed);
assert_eq!(
runtime
.execute(
execute_request(r#"text("unreachable");"#),
ObserveMode::YieldAfter(Duration::from_secs(1)),
)
.await
.err(),
Some(Error::CellIdSpaceExhausted)
);
}
#[tokio::test]
#[expect(
clippy::await_holding_invalid_type,
@@ -138,6 +138,7 @@ pub(crate) trait SessionRuntimeDelegate: Send + Sync + 'static {
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum Error {
ShuttingDown,
CellIdSpaceExhausted,
DuplicateCell(CellId),
MissingCell(CellId),
BusyObserver(CellId),
@@ -150,6 +151,9 @@ impl fmt::Display for Error {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::ShuttingDown => formatter.write_str("code mode session is shutting down"),
Self::CellIdSpaceExhausted => {
formatter.write_str("code mode session exhausted its cell ID space")
}
Self::DuplicateCell(cell_id) => write!(formatter, "exec cell {cell_id} already exists"),
Self::MissingCell(cell_id) => write!(formatter, "exec cell {cell_id} not found"),
Self::BusyObserver(cell_id) => {
+24
View File
@@ -371,6 +371,30 @@ impl ResolvedMcpCatalog {
.collect()
}
/// Replaces the resolved server set while preserving known server sources.
///
/// Names not present in the existing catalog are treated as config-owned.
pub fn with_materialized_servers(&self, servers: HashMap<String, McpServerConfig>) -> Self {
let mut builder = Self::builder();
for (name, config) in servers {
let source = self
.server(&name)
.map(|server| server.source.clone())
.unwrap_or(McpServerSource::Config);
let precedence = match &source {
McpServerSource::Plugin(_) => RegistrationPrecedence::Plugin(Reverse(0)),
McpServerSource::SelectedPlugin(_) => {
RegistrationPrecedence::SelectedPlugin(Reverse(0))
}
McpServerSource::Config => RegistrationPrecedence::Config,
McpServerSource::Compatibility { .. } => RegistrationPrecedence::Compatibility,
McpServerSource::Extension { .. } => RegistrationPrecedence::Extension(0),
};
builder.register(McpServerRegistration::new(name, source, config, precedence));
}
builder.build()
}
/// Returns package attribution for each winning plugin-owned server.
pub fn plugin_attributions_by_server_name(&self) -> HashMap<String, McpPluginAttribution> {
self.servers
+11
View File
@@ -304,6 +304,17 @@ fn selected_plugins_override_discovered_plugins_but_not_config() {
}]
);
let refreshed = server("https://refreshed.example/mcp");
let catalog =
catalog.with_materialized_servers(HashMap::from([("docs".to_string(), refreshed.clone())]));
assert_eq!(
catalog.server("docs"),
Some(&super::ResolvedMcpServer {
source: selected_plugin_source("selected-alpha"),
config: refreshed,
})
);
let mut builder = catalog.to_builder();
let configured = server("https://config.example/mcp");
builder.register(McpServerRegistration::from_config(
+1 -211
View File
@@ -1,70 +1,9 @@
//! Codex Apps support for the host-owned apps MCP server.
//!
//! This module owns the pieces that are unique to ChatGPT-hosted app
//! connectors: cache scoping by authenticated user, disk cache reads/writes,
//! connector allow-list filtering, and the normalization that turns app
//! This module owns the normalization that turns ChatGPT-hosted app
//! connector/tool metadata into model-visible MCP callable names.
use std::path::PathBuf;
use std::time::Instant;
use crate::runtime::emit_duration;
use crate::tools::MCP_TOOLS_CACHE_WRITE_DURATION_METRIC;
use crate::tools::ToolInfo;
use anyhow::Context;
use codex_login::CodexAuth;
use codex_protocol::mcp::McpServerInfo;
use codex_utils_plugins::mcp_connector::sanitize_name;
use serde::Deserialize;
use serde::Serialize;
use sha1::Digest;
use sha1::Sha1;
use tracing::instrument;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CodexAppsToolsCacheKey {
pub(crate) account_id: Option<String>,
pub(crate) chatgpt_user_id: Option<String>,
pub(crate) is_workspace_account: bool,
}
pub fn codex_apps_tools_cache_key(auth: Option<&CodexAuth>) -> CodexAppsToolsCacheKey {
CodexAppsToolsCacheKey {
account_id: auth.and_then(CodexAuth::get_account_id),
chatgpt_user_id: auth.and_then(CodexAuth::get_chatgpt_user_id),
is_workspace_account: auth.is_some_and(CodexAuth::is_workspace_account),
}
}
#[derive(Clone)]
pub(crate) struct CodexAppsToolsCacheContext {
pub(crate) codex_home: PathBuf,
pub(crate) user_key: CodexAppsToolsCacheKey,
}
impl CodexAppsToolsCacheContext {
pub(crate) fn tools_cache_path(&self) -> PathBuf {
self.cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR)
}
pub(crate) fn server_info_cache_path(&self) -> PathBuf {
self.cache_path_in(CODEX_APPS_SERVER_INFO_CACHE_DIR)
}
fn cache_path_in(&self, cache_dir: &str) -> PathBuf {
let user_key_json = serde_json::to_string(&self.user_key).unwrap_or_default();
let user_key_hash = sha1_hex(&user_key_json);
self.codex_home
.join(cache_dir)
.join(format!("{user_key_hash}.json"))
}
}
pub(crate) enum CachedCodexAppsToolsLoad {
Hit(Vec<ToolInfo>),
Missing,
Invalid,
}
pub(crate) fn normalize_codex_apps_tool_title(connector_name: Option<&str>, value: &str) -> String {
let Some(connector_name) = connector_name
@@ -124,152 +63,3 @@ pub(crate) fn normalize_codex_apps_callable_namespace(
server_name.to_string()
}
}
pub(crate) fn write_codex_apps_tools_cache(
cache_context: Option<&CodexAppsToolsCacheContext>,
server_info: &McpServerInfo,
tools: &[ToolInfo],
) {
if let Some(cache_context) = cache_context {
let cache_write_start = Instant::now();
write_cached_codex_apps_tools(cache_context, tools);
if let Err(err) = write_cached_codex_apps_server_info(cache_context, server_info) {
tracing::warn!("failed to write Codex Apps server info cache: {err:#}");
}
emit_duration(
MCP_TOOLS_CACHE_WRITE_DURATION_METRIC,
cache_write_start.elapsed(),
&[],
);
}
}
pub(crate) fn load_startup_cached_codex_apps_tools_snapshot(
cache_context: Option<&CodexAppsToolsCacheContext>,
) -> Option<Vec<ToolInfo>> {
let cache_context = cache_context?;
match load_cached_codex_apps_tools(cache_context) {
CachedCodexAppsToolsLoad::Hit(tools) => Some(tools),
CachedCodexAppsToolsLoad::Missing | CachedCodexAppsToolsLoad::Invalid => None,
}
}
pub(crate) fn load_startup_cached_codex_apps_server_info(
cache_context: Option<&CodexAppsToolsCacheContext>,
) -> Option<McpServerInfo> {
load_cached_codex_apps_server_info(cache_context?)
}
#[cfg(test)]
pub(crate) fn read_cached_codex_apps_tools(
cache_context: &CodexAppsToolsCacheContext,
) -> Option<Vec<ToolInfo>> {
match load_cached_codex_apps_tools(cache_context) {
CachedCodexAppsToolsLoad::Hit(tools) => Some(tools),
CachedCodexAppsToolsLoad::Missing | CachedCodexAppsToolsLoad::Invalid => None,
}
}
#[instrument(level = "trace", skip_all)]
pub(crate) fn load_cached_codex_apps_tools(
cache_context: &CodexAppsToolsCacheContext,
) -> CachedCodexAppsToolsLoad {
let cache_path = cache_context.tools_cache_path();
let bytes = match std::fs::read(cache_path) {
Ok(bytes) => bytes,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
return CachedCodexAppsToolsLoad::Missing;
}
Err(_) => return CachedCodexAppsToolsLoad::Invalid,
};
let cache: CodexAppsToolsDiskCache = match serde_json::from_slice(&bytes) {
Ok(cache) => cache,
Err(_) => return CachedCodexAppsToolsLoad::Invalid,
};
if cache.schema_version != CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION {
return CachedCodexAppsToolsLoad::Invalid;
}
CachedCodexAppsToolsLoad::Hit(cache.tools)
}
pub(crate) fn write_cached_codex_apps_tools(
cache_context: &CodexAppsToolsCacheContext,
tools: &[ToolInfo],
) {
let cache_path = cache_context.tools_cache_path();
if let Some(parent) = cache_path.parent()
&& std::fs::create_dir_all(parent).is_err()
{
return;
}
let Ok(bytes) = serde_json::to_vec_pretty(&CodexAppsToolsDiskCache {
schema_version: CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION,
tools: tools.to_vec(),
}) else {
return;
};
let _ = std::fs::write(cache_path, bytes);
}
#[instrument(level = "trace", skip_all)]
pub(crate) fn load_cached_codex_apps_server_info(
cache_context: &CodexAppsToolsCacheContext,
) -> Option<McpServerInfo> {
let bytes = std::fs::read(cache_context.server_info_cache_path()).ok()?;
let cache: CodexAppsServerInfoDiskCache = serde_json::from_slice(&bytes).ok()?;
(cache.schema_version == CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION)
.then_some(cache.server_info)
}
fn write_cached_codex_apps_server_info(
cache_context: &CodexAppsToolsCacheContext,
server_info: &McpServerInfo,
) -> anyhow::Result<()> {
let cache_path = cache_context.server_info_cache_path();
if let Some(parent) = cache_path.parent() {
std::fs::create_dir_all(parent).with_context(|| {
format!(
"failed to create Codex Apps server info cache directory `{}`",
parent.display()
)
})?;
}
let bytes = serde_json::to_vec_pretty(&CodexAppsServerInfoDiskCache {
schema_version: CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION,
server_info: server_info.clone(),
})
.context("failed to serialize Codex Apps server info cache")?;
std::fs::write(&cache_path, bytes).with_context(|| {
format!(
"failed to write Codex Apps server info cache `{}`",
cache_path.display()
)
})?;
Ok(())
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CodexAppsToolsDiskCache {
schema_version: u8,
tools: Vec<ToolInfo>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CodexAppsServerInfoDiskCache {
schema_version: u8,
server_info: McpServerInfo,
}
const CODEX_APPS_TOOLS_CACHE_DIR: &str = "cache/codex_apps_tools";
pub(crate) const CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION: u8 = 4;
const CODEX_APPS_SERVER_INFO_CACHE_DIR: &str = "cache/codex_apps_server_info";
const CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION: u8 = 1;
fn sha1_hex(s: &str) -> String {
let mut hasher = Sha1::new();
hasher.update(s.as_bytes());
let sha1 = hasher.finalize();
format!("{sha1:x}")
}
+378
View File
@@ -0,0 +1,378 @@
//! Shared raw tool cache for the host-owned Codex Apps MCP server.
//!
//! Cache entries are process-local live state scoped by the active Codex auth
//! key. Disk is best-effort cold-start persistence; entries do not reread disk
//! after creation.
use std::collections::HashMap;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::time::Instant;
use anyhow::Context;
use arc_swap::ArcSwapOption;
use codex_login::CodexAuth;
use codex_protocol::mcp::McpServerInfo;
use serde::Deserialize;
use serde::Serialize;
use sha1::Digest;
use sha1::Sha1;
use tracing::instrument;
use crate::runtime::emit_duration;
use crate::tools::MCP_TOOLS_CACHE_WRITE_DURATION_METRIC;
use crate::tools::ToolInfo;
const MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC: &str = "codex.mcp.tools.cache_publish.duration_ms";
/// The CodexAuth bits that identify a Codex Apps catalog.
///
/// Debug bearer-token overrides bypass the shared cache, so shared entries only
/// need the CodexAuth-backed identity.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct CodexAppsToolsCacheKey {
pub(crate) account_id: Option<String>,
pub(crate) chatgpt_user_id: Option<String>,
pub(crate) is_workspace_account: bool,
}
/// Builds the CodexAuth-backed Codex Apps cache key.
pub fn codex_apps_tools_cache_key(auth: Option<&CodexAuth>) -> CodexAppsToolsCacheKey {
CodexAppsToolsCacheKey {
account_id: auth.and_then(CodexAuth::get_account_id),
chatgpt_user_id: auth.and_then(CodexAuth::get_chatgpt_user_id),
is_workspace_account: auth.is_some_and(CodexAuth::is_workspace_account),
}
}
/// Process-scoped registry for shared Codex Apps raw tool snapshots.
///
/// Two clients share an entry only when they would read the same Codex Apps
/// catalog. New entries may seed from disk; live entries read from memory only.
#[derive(Clone, Default)]
pub struct CodexAppsToolsCache {
entries: Arc<Mutex<HashMap<CodexAppsToolsCacheIdentity, Arc<CodexAppsToolsCacheEntry>>>>,
}
/// Handle to one shared Codex Apps tools cache entry.
///
/// The connection manager creates this from the auth key, then tool
/// reads and refreshes for that managed client use the same entry.
#[derive(Clone)]
pub(crate) struct CodexAppsToolsCacheContext {
entry: Arc<CodexAppsToolsCacheEntry>,
}
impl CodexAppsToolsCacheContext {
pub(crate) fn tools_cache_path(&self) -> PathBuf {
self.entry
.identity
.cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR)
}
pub(crate) fn server_info_cache_path(&self) -> PathBuf {
self.entry
.identity
.cache_path_in(CODEX_APPS_SERVER_INFO_CACHE_DIR)
}
pub(crate) fn current_tools(&self) -> Option<Vec<ToolInfo>> {
self.entry
.current_tools
.load_full()
.map(|tools| tools.as_ref().clone())
}
pub(crate) fn has_current_tools(&self) -> bool {
self.entry.current_tools.load_full().is_some()
}
pub(crate) fn begin_fetch(
&self,
source: CodexAppsToolsFetchSource,
) -> CodexAppsToolsFetchTicket {
CodexAppsToolsFetchTicket {
generation: self
.entry
.next_fetch_generation
.fetch_add(1, Ordering::Relaxed)
+ 1,
source,
}
}
pub(crate) fn publish_if_newest_accepted(
&self,
ticket: CodexAppsToolsFetchTicket,
server_info: &McpServerInfo,
tools: Vec<ToolInfo>,
) -> Vec<ToolInfo> {
let publish_start = Instant::now();
let mut last_accepted_generation = lock_unpoisoned(&self.entry.last_accepted_generation);
if ticket.generation <= *last_accepted_generation {
emit_duration(
MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC,
publish_start.elapsed(),
&[("source", ticket.source.as_str()), ("result", "stale")],
);
return self.current_tools().unwrap_or(tools);
}
*last_accepted_generation = ticket.generation;
self.entry
.current_tools
.store(Some(Arc::new(tools.clone())));
persist_codex_apps_cache(self, server_info, &tools);
emit_duration(
MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC,
publish_start.elapsed(),
&[("source", ticket.source.as_str()), ("result", "published")],
);
tools
}
#[cfg(test)]
pub(crate) fn store_current_tools_for_test(&self, tools: Vec<ToolInfo>) {
self.entry.current_tools.store(Some(Arc::new(tools)));
}
}
impl CodexAppsToolsCache {
pub(crate) fn context(
&self,
codex_home: PathBuf,
auth_key: CodexAppsToolsCacheKey,
) -> CodexAppsToolsCacheContext {
let identity = CodexAppsToolsCacheIdentity {
codex_home,
auth_key,
};
let mut entries = lock_unpoisoned(&self.entries);
let entry = entries
.entry(identity.clone())
.or_insert_with(|| Arc::new(CodexAppsToolsCacheEntry::new(identity)))
.clone();
CodexAppsToolsCacheContext { entry }
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) enum CodexAppsToolsFetchSource {
Startup,
HardRefresh,
}
impl CodexAppsToolsFetchSource {
fn as_str(self) -> &'static str {
match self {
Self::Startup => "startup",
Self::HardRefresh => "hard_refresh",
}
}
}
pub(crate) struct CodexAppsToolsFetchTicket {
generation: u64,
source: CodexAppsToolsFetchSource,
}
struct CodexAppsToolsCacheEntry {
identity: CodexAppsToolsCacheIdentity,
current_tools: ArcSwapOption<Vec<ToolInfo>>,
next_fetch_generation: AtomicU64,
last_accepted_generation: Mutex<u64>,
}
impl CodexAppsToolsCacheEntry {
fn new(identity: CodexAppsToolsCacheIdentity) -> Self {
let current_tools = load_cached_codex_apps_tools_for_identity(&identity).map(Arc::new);
Self {
identity,
current_tools: ArcSwapOption::from(current_tools),
next_fetch_generation: AtomicU64::new(0),
last_accepted_generation: Mutex::new(0),
}
}
}
/// Everything that decides whether two Codex Apps clients can share tools.
///
/// The auth key says whose catalog we are reading. `codex_home` keeps the
/// persisted cache under the right home directory.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct CodexAppsToolsCacheIdentity {
codex_home: PathBuf,
auth_key: CodexAppsToolsCacheKey,
}
impl CodexAppsToolsCacheIdentity {
fn cache_path_in(&self, cache_dir: &str) -> PathBuf {
// `codex_home` is already the parent directory. Keep it out of the
// filename hash so non-UTF-8 Unix paths cannot collapse distinct auth
// keys onto the same disk cache file.
let identity_json = serde_json::to_string(&self.auth_key).unwrap_or_default();
let identity_hash = sha1_hex(&identity_json);
self.codex_home
.join(cache_dir)
.join(format!("{identity_hash}.json"))
}
}
#[cfg(test)]
fn write_cached_codex_apps_tools_for_test(
cache_context: &CodexAppsToolsCacheContext,
server_info: &McpServerInfo,
tools: &[ToolInfo],
) {
cache_context
.entry
.current_tools
.store(Some(Arc::new(tools.to_vec())));
persist_codex_apps_cache(cache_context, server_info, tools);
}
pub(crate) fn load_startup_cached_codex_apps_server_info(
cache_context: &CodexAppsToolsCacheContext,
) -> Option<McpServerInfo> {
load_cached_codex_apps_server_info(cache_context)
}
#[cfg(test)]
fn read_cached_codex_apps_tools(
cache_context: &CodexAppsToolsCacheContext,
) -> Option<Vec<ToolInfo>> {
load_cached_codex_apps_tools_for_identity(&cache_context.entry.identity)
}
#[instrument(level = "trace", skip_all)]
fn load_cached_codex_apps_tools_for_identity(
identity: &CodexAppsToolsCacheIdentity,
) -> Option<Vec<ToolInfo>> {
let cache_path = identity.cache_path_in(CODEX_APPS_TOOLS_CACHE_DIR);
let bytes = std::fs::read(cache_path).ok()?;
let cache: CodexAppsToolsDiskCache = serde_json::from_slice(&bytes).ok()?;
(cache.schema_version == CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION).then_some(cache.tools)
}
fn write_cached_codex_apps_tools(
cache_context: &CodexAppsToolsCacheContext,
tools: &[ToolInfo],
) -> anyhow::Result<()> {
let cache_path = cache_context.tools_cache_path();
let bytes = serde_json::to_vec_pretty(&CodexAppsToolsDiskCache {
schema_version: CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION,
tools: tools.to_vec(),
})
.context("failed to serialize Codex Apps tools cache")?;
write_codex_apps_cache_file(&cache_path, "tools", bytes)
}
#[instrument(level = "trace", skip_all)]
fn load_cached_codex_apps_server_info(
cache_context: &CodexAppsToolsCacheContext,
) -> Option<McpServerInfo> {
let bytes = std::fs::read(cache_context.server_info_cache_path()).ok()?;
let cache: CodexAppsServerInfoDiskCache = serde_json::from_slice(&bytes).ok()?;
(cache.schema_version == CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION)
.then_some(cache.server_info)
}
fn write_cached_codex_apps_server_info(
cache_context: &CodexAppsToolsCacheContext,
server_info: &McpServerInfo,
) -> anyhow::Result<()> {
let cache_path = cache_context.server_info_cache_path();
let bytes = serde_json::to_vec_pretty(&CodexAppsServerInfoDiskCache {
schema_version: CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION,
server_info: server_info.clone(),
})
.context("failed to serialize Codex Apps server info cache")?;
write_codex_apps_cache_file(&cache_path, "server info", bytes)
}
fn write_codex_apps_cache_file(
cache_path: &Path,
cache_name: &str,
bytes: Vec<u8>,
) -> anyhow::Result<()> {
if let Some(parent) = cache_path.parent() {
std::fs::create_dir_all(parent).with_context(|| {
format!(
"failed to create Codex Apps {cache_name} cache directory `{}`",
parent.display()
)
})?;
}
std::fs::write(cache_path, bytes).with_context(|| {
format!(
"failed to write Codex Apps {cache_name} cache `{}`",
cache_path.display()
)
})?;
Ok(())
}
fn persist_codex_apps_cache(
cache_context: &CodexAppsToolsCacheContext,
server_info: &McpServerInfo,
tools: &[ToolInfo],
) {
let cache_write_start = Instant::now();
let tools_result = write_cached_codex_apps_tools(cache_context, tools);
if let Err(err) = &tools_result {
tracing::warn!("failed to write Codex Apps tools cache: {err:#}");
}
let server_info_result = write_cached_codex_apps_server_info(cache_context, server_info);
if let Err(err) = &server_info_result {
tracing::warn!("failed to write Codex Apps server info cache: {err:#}");
}
let status = if tools_result.is_ok() && server_info_result.is_ok() {
"success"
} else {
"failure"
};
emit_duration(
MCP_TOOLS_CACHE_WRITE_DURATION_METRIC,
cache_write_start.elapsed(),
&[("status", status)],
);
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CodexAppsToolsDiskCache {
schema_version: u8,
tools: Vec<ToolInfo>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CodexAppsServerInfoDiskCache {
schema_version: u8,
server_info: McpServerInfo,
}
const CODEX_APPS_TOOLS_CACHE_DIR: &str = "cache/codex_apps_tools";
const CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION: u8 = 4;
const CODEX_APPS_SERVER_INFO_CACHE_DIR: &str = "cache/codex_apps_server_info";
const CODEX_APPS_SERVER_INFO_CACHE_SCHEMA_VERSION: u8 = 1;
fn sha1_hex(s: &str) -> String {
let mut hasher = Sha1::new();
hasher.update(s.as_bytes());
let sha1 = hasher.finalize();
format!("{sha1:x}")
}
fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
mutex
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[cfg(test)]
#[path = "codex_apps_cache_tests.rs"]
mod tests;
@@ -0,0 +1,464 @@
use super::*;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
use crate::tools::ToolInfo;
use codex_protocol::ToolName;
use codex_protocol::mcp::McpServerInfo;
use pretty_assertions::assert_eq;
use rmcp::model::JsonObject;
use rmcp::model::Tool;
use std::collections::HashSet;
#[cfg(unix)]
use std::os::unix::ffi::OsStringExt;
use std::path::PathBuf;
use std::sync::Arc;
use tempfile::tempdir;
fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo {
ToolInfo {
server_name: server_name.to_string(),
supports_parallel_tool_calls: false,
server_origin: None,
callable_name: tool_name.to_string(),
callable_namespace: server_name.to_string(),
namespace_description: None,
tool: Tool::new(
tool_name.to_string(),
format!("Test tool: {tool_name}"),
Arc::new(JsonObject::default()),
),
connector_id: None,
connector_name: None,
plugin_display_names: Vec::new(),
}
}
fn create_test_tool_with_connector(
server_name: &str,
tool_name: &str,
connector_id: &str,
connector_name: Option<&str>,
) -> ToolInfo {
let mut tool = create_test_tool(server_name, tool_name);
tool.connector_id = Some(connector_id.to_string());
tool.connector_name = connector_name.map(ToOwned::to_owned);
tool
}
fn create_codex_apps_tools_cache_context(
codex_home: PathBuf,
account_id: Option<&str>,
chatgpt_user_id: Option<&str>,
) -> CodexAppsToolsCacheContext {
CodexAppsToolsCache::default().context(
codex_home,
CodexAppsToolsCacheKey {
account_id: account_id.map(ToOwned::to_owned),
chatgpt_user_id: chatgpt_user_id.map(ToOwned::to_owned),
is_workspace_account: false,
},
)
}
fn create_test_server_info(title: &str) -> McpServerInfo {
McpServerInfo {
name: "codex-apps".to_string(),
title: Some(title.to_string()),
version: "1.0.0".to_string(),
description: None,
icons: None,
website_url: None,
}
}
fn model_tool_names(tools: &[ToolInfo]) -> HashSet<ToolName> {
tools
.iter()
.map(ToolInfo::canonical_tool_name)
.collect::<HashSet<_>>()
}
#[test]
fn codex_apps_tools_cache_is_overwritten_by_last_write() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let tools_gateway_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")];
let tools_gateway_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")];
write_cached_codex_apps_tools(&cache_context, &tools_gateway_1).expect("write first cache");
let cached_gateway_1 =
read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for first write");
assert_eq!(cached_gateway_1[0].callable_name, "one");
write_cached_codex_apps_tools(&cache_context, &tools_gateway_2).expect("write second cache");
let cached_gateway_2 =
read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for second write");
assert_eq!(cached_gateway_2[0].callable_name, "two");
}
#[test]
fn codex_apps_tools_cache_is_scoped_per_user() {
let codex_home = tempdir().expect("tempdir");
let cache_context_user_1 = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cache_context_user_2 = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-two"),
Some("user-two"),
);
let tools_user_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")];
let tools_user_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")];
write_cached_codex_apps_tools(&cache_context_user_1, &tools_user_1)
.expect("write user one cache");
write_cached_codex_apps_tools(&cache_context_user_2, &tools_user_2)
.expect("write user two cache");
let read_user_1 =
read_cached_codex_apps_tools(&cache_context_user_1).expect("cache entry for user one");
let read_user_2 =
read_cached_codex_apps_tools(&cache_context_user_2).expect("cache entry for user two");
assert_eq!(read_user_1[0].callable_name, "one");
assert_eq!(read_user_2[0].callable_name, "two");
assert_ne!(
cache_context_user_1.tools_cache_path(),
cache_context_user_2.tools_cache_path(),
"each user should get an isolated cache file"
);
}
#[test]
fn codex_apps_tools_cache_preserves_formerly_disallowed_connectors() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let tools = vec![
create_test_tool_with_connector(
CODEX_APPS_MCP_SERVER_NAME,
"formerly_blocked_tool",
"connector_2b0a9009c9c64bf9933a3dae3f2b1254",
Some("Formerly Blocked"),
),
create_test_tool_with_connector(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_tool",
"calendar",
Some("Calendar"),
),
];
write_cached_codex_apps_tools(&cache_context, &tools).expect("write cache");
let cached = read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for user");
assert_eq!(
cached
.iter()
.map(|tool| (tool.callable_name.as_str(), tool.connector_id.as_deref()))
.collect::<Vec<_>>(),
vec![
(
"formerly_blocked_tool",
Some("connector_2b0a9009c9c64bf9933a3dae3f2b1254")
),
("calendar_tool", Some("calendar")),
]
);
}
#[test]
fn codex_apps_tools_cache_is_ignored_when_schema_version_mismatches() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cache_path = cache_context.tools_cache_path();
if let Some(parent) = cache_path.parent() {
std::fs::create_dir_all(parent).expect("create parent");
}
let bytes = serde_json::to_vec_pretty(&serde_json::json!({
"schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION + 1,
"tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")],
}))
.expect("serialize");
std::fs::write(cache_path, bytes).expect("write");
assert!(read_cached_codex_apps_tools(&cache_context).is_none());
}
#[test]
fn codex_apps_tools_cache_is_ignored_when_json_is_invalid() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cache_path = cache_context.tools_cache_path();
if let Some(parent) = cache_path.parent() {
std::fs::create_dir_all(parent).expect("create parent");
}
std::fs::write(cache_path, b"{not json").expect("write");
assert!(read_cached_codex_apps_tools(&cache_context).is_none());
}
#[test]
fn startup_cached_codex_apps_tools_loads_from_disk_cache() {
let codex_home = tempdir().expect("tempdir");
let writer_cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cached_tools = vec![create_test_tool(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_search",
)];
let server_info = create_test_server_info("Codex Apps");
write_cached_codex_apps_tools_for_test(&writer_cache_context, &server_info, &cached_tools);
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let startup_tools = cache_context
.current_tools()
.expect("expected startup snapshot to load from cache");
let cached_server_info = load_startup_cached_codex_apps_server_info(&cache_context);
assert_eq!(startup_tools.len(), 1);
assert_eq!(startup_tools[0].server_name, CODEX_APPS_MCP_SERVER_NAME);
assert_eq!(startup_tools[0].callable_name, "calendar_search");
assert_eq!(cached_server_info, Some(server_info));
}
#[test]
fn startup_cached_codex_apps_tools_loads_without_server_info_cache() {
let codex_home = tempdir().expect("tempdir");
let writer_cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cache_path = writer_cache_context.tools_cache_path();
if let Some(parent) = cache_path.parent() {
std::fs::create_dir_all(parent).expect("create parent");
}
let bytes = serde_json::to_vec_pretty(&serde_json::json!({
"schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION,
"tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")],
}))
.expect("serialize");
std::fs::write(cache_path, bytes).expect("write");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let startup_tools = cache_context
.current_tools()
.expect("legacy startup snapshot should remain available");
let cached_server_info = load_startup_cached_codex_apps_server_info(&cache_context);
assert_eq!(startup_tools.len(), 1);
assert_eq!(startup_tools[0].callable_name, "calendar_search");
assert_eq!(cached_server_info, None);
}
#[test]
fn codex_apps_server_info_cache_survives_legacy_tools_cache_write() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let server_info = create_test_server_info("Codex Apps");
write_cached_codex_apps_tools_for_test(
&cache_context,
&server_info,
&[create_test_tool(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_search",
)],
);
let cache_path = cache_context.tools_cache_path();
if let Some(parent) = cache_path.parent() {
std::fs::create_dir_all(parent).expect("create parent");
}
let bytes = serde_json::to_vec_pretty(&serde_json::json!({
"schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION - 1,
"tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")],
}))
.expect("serialize");
std::fs::write(cache_path, bytes).expect("write legacy tools cache");
let startup_cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
assert_eq!(
load_startup_cached_codex_apps_server_info(&startup_cache_context),
Some(server_info)
);
assert!(startup_cache_context.current_tools().is_none());
}
#[test]
fn codex_apps_tools_cache_context_does_not_reread_disk_after_creation() {
let codex_home = tempdir().expect("tempdir");
let writer_cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cached_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "cached")];
write_cached_codex_apps_tools(&writer_cache_context, &cached_tools).expect("write cache");
let reader_cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let updated_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "updated")];
write_cached_codex_apps_tools(&writer_cache_context, &updated_tools).expect("rewrite cache");
assert_eq!(
reader_cache_context
.current_tools()
.expect("in-memory tools")[0]
.callable_name,
"cached"
);
assert_eq!(
read_cached_codex_apps_tools(&writer_cache_context).expect("disk tools")[0].callable_name,
"updated"
);
}
#[test]
fn codex_apps_tools_cache_publishes_newest_shared_snapshot() {
let codex_home = tempdir().expect("tempdir");
let cache = CodexAppsToolsCache::default();
let cache_context_1 = cache.context(
codex_home.path().to_path_buf(),
CodexAppsToolsCacheKey {
account_id: Some("account-one".to_string()),
chatgpt_user_id: Some("user-one".to_string()),
is_workspace_account: false,
},
);
let cache_context_2 = cache.context(
codex_home.path().to_path_buf(),
CodexAppsToolsCacheKey {
account_id: Some("account-one".to_string()),
chatgpt_user_id: Some("user-one".to_string()),
is_workspace_account: false,
},
);
let older_ticket = cache_context_1.begin_fetch(CodexAppsToolsFetchSource::Startup);
let newer_ticket = cache_context_2.begin_fetch(CodexAppsToolsFetchSource::HardRefresh);
let server_info = create_test_server_info("Codex Apps");
let newer_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newer")];
let older_tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "older")];
let published_tools =
cache_context_2.publish_if_newest_accepted(newer_ticket, &server_info, newer_tools);
assert_eq!(
model_tool_names(&published_tools),
model_tool_names(
&cache_context_1
.current_tools()
.expect("new snapshot should publish")
)
);
let current_tools =
cache_context_1.publish_if_newest_accepted(older_ticket, &server_info, older_tools);
assert_eq!(current_tools[0].callable_name, "newer");
assert_eq!(
cache_context_2.current_tools().expect("shared snapshot")[0].callable_name,
"newer"
);
assert_eq!(
read_cached_codex_apps_tools(&cache_context_1).expect("persisted snapshot")[0]
.callable_name,
"newer"
);
}
#[test]
fn codex_apps_tools_cache_keeps_live_publish_when_disk_persistence_fails() {
let codex_home = tempdir().expect("tempdir");
let codex_home_file = codex_home.path().join("not-a-directory");
std::fs::write(&codex_home_file, b"occupied").expect("create codex home file");
let cache_context = CodexAppsToolsCache::default().context(
codex_home_file,
CodexAppsToolsCacheKey {
account_id: Some("account-one".to_string()),
chatgpt_user_id: Some("user-one".to_string()),
is_workspace_account: false,
},
);
let tools = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "live")];
let published_tools = cache_context.publish_if_newest_accepted(
cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh),
&create_test_server_info("Codex Apps"),
tools.clone(),
);
assert_eq!(model_tool_names(&published_tools), model_tool_names(&tools));
assert_eq!(
model_tool_names(&cache_context.current_tools().expect("live snapshot")),
model_tool_names(&tools)
);
}
#[cfg(unix)]
#[test]
fn codex_apps_tools_cache_scopes_non_utf8_home_disk_paths() {
let codex_home = PathBuf::from(std::ffi::OsString::from_vec(
b"/tmp/codex-home-\xff".to_vec(),
));
let cache = CodexAppsToolsCache::default();
let user_one_context = cache.context(
codex_home.clone(),
CodexAppsToolsCacheKey {
account_id: Some("account-one".to_string()),
chatgpt_user_id: Some("user-one".to_string()),
is_workspace_account: false,
},
);
let user_two_context = cache.context(
codex_home,
CodexAppsToolsCacheKey {
account_id: Some("account-two".to_string()),
chatgpt_user_id: Some("user-two".to_string()),
is_workspace_account: false,
},
);
let cache_paths = [
user_one_context.tools_cache_path(),
user_two_context.tools_cache_path(),
];
assert_eq!(
cache_paths.iter().collect::<HashSet<_>>().len(),
cache_paths.len()
);
}
+87 -42
View File
@@ -7,7 +7,6 @@
//! `codex-core`.
use std::collections::HashMap;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::Ordering;
@@ -15,9 +14,9 @@ use std::time::Duration;
use std::time::Instant;
use crate::McpAuthStatusEntry;
use crate::codex_apps::CodexAppsToolsCacheContext;
use crate::codex_apps::CodexAppsToolsCacheKey;
use crate::codex_apps::write_codex_apps_tools_cache;
use crate::codex_apps_cache::CodexAppsToolsCache;
use crate::codex_apps_cache::CodexAppsToolsCacheKey;
use crate::codex_apps_cache::CodexAppsToolsFetchSource;
use crate::elicitation::ElicitationRequestManager;
use crate::elicitation::ElicitationReviewerHandle;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
@@ -56,9 +55,12 @@ use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::McpStartupCompleteEvent;
use codex_protocol::protocol::McpStartupFailure;
use codex_protocol::protocol::McpStartupFailureReason;
use codex_protocol::protocol::McpStartupStatus;
use codex_protocol::protocol::McpStartupUpdateEvent;
use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::McpAuthState;
use codex_rmcp_client::McpLoginRequirement;
use rmcp::model::ElicitationCapability;
use rmcp::model::ListResourceTemplatesResult;
use rmcp::model::ListResourcesResult;
@@ -131,6 +133,7 @@ impl McpConnectionManager {
initial_permission_profile: PermissionProfile,
runtime_context: McpRuntimeContext,
codex_home: PathBuf,
codex_apps_tools_cache: CodexAppsToolsCache,
codex_apps_tools_cache_key: CodexAppsToolsCacheKey,
prefix_mcp_tool_names: bool,
client_elicitation_capability: ElicitationCapability,
@@ -174,13 +177,32 @@ impl McpConnectionManager {
},
)
.await;
let codex_apps_tools_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME {
codex_apps_tools_cache_context(&codex_home, &codex_apps_tools_cache_key)
} else {
regular_mcp_tools_cache_context()
};
let configured_config = server.configured_config();
// For built-in Codex Apps, `CODEX_CONNECTORS_TOKEN` is a debug
// override: it supplies runtime auth but bypasses the shared tools
// cache.
let uses_env_bearer_token =
configured_config.is_some_and(|config| match &config.transport {
McpServerTransportConfig::StreamableHttp {
bearer_token_env_var,
..
} => bearer_token_env_var.is_some(),
McpServerTransportConfig::Stdio { .. } => false,
});
let shares_codex_apps_tools_cache =
should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token);
let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| {
codex_apps_tools_cache
.context(codex_home.clone(), codex_apps_tools_cache_key.clone())
});
// If Codex Apps has an env bearer token, that is its auth path. Do
// not also attach the ambient CodexAuth provider.
let runtime_auth_provider =
chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider.clone());
if server_name == CODEX_APPS_MCP_SERVER_NAME && uses_env_bearer_token {
None
} else {
chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider.clone())
};
let async_managed_client = AsyncManagedClient::new(
server_name.clone(),
server,
@@ -209,12 +231,16 @@ impl McpConnectionManager {
Ok(_) => McpStartupStatus::Ready,
Err(StartupOutcomeError::Cancelled) => McpStartupStatus::Cancelled,
Err(error) => {
let reason = mcp_startup_failure_reason(auth_entry.as_ref(), error);
let error_str = mcp_init_error_display(
server_name.as_str(),
auth_entry.as_ref(),
error,
);
McpStartupStatus::Failed { error: error_str }
McpStartupStatus::Failed {
error: error_str,
reason,
}
}
};
@@ -247,7 +273,7 @@ impl McpConnectionManager {
match outcome {
Ok(_) => summary.ready.push(server_name),
Err(StartupOutcomeError::Cancelled) => summary.cancelled.push(server_name),
Err(StartupOutcomeError::Failed { error }) => {
Err(StartupOutcomeError::Failed { error, .. }) => {
summary.failed.push(McpStartupFailure {
server: server_name,
error,
@@ -445,13 +471,13 @@ impl McpConnectionManager {
pub async fn list_all_tools(&self) -> Vec<ToolInfo> {
let mut tools = Vec::new();
for (server_name, managed_client) in &self.clients {
let has_cached_tool_info_snapshot = managed_client.cached_tool_info_snapshot.is_some();
let has_cached_tools = managed_client.has_cached_tools();
let startup_complete = managed_client
.startup_complete
.load(std::sync::atomic::Ordering::Acquire);
trace!(
server_name = %server_name,
has_cached_tool_info_snapshot,
has_cached_tools,
startup_complete,
"waiting for MCP server tools while building tool list"
);
@@ -460,7 +486,7 @@ impl McpConnectionManager {
.instrument(trace_span!(
"list_tools_for_server",
server_name = %server_name,
has_cached_tool_info_snapshot,
has_cached_tools,
startup_complete
))
.await
@@ -483,9 +509,9 @@ impl McpConnectionManager {
/// Force-refresh codex apps tools by bypassing the in-process cache.
///
/// On success, the refreshed tools replace the cache contents and the
/// latest filtered tools are returned directly to the caller. On
/// failure, the existing cache remains unchanged.
/// On success, the refreshed tools replace shared cache contents when the
/// cache is enabled and the latest filtered tools are returned directly to
/// the caller. On failure, existing shared cache contents remain unchanged.
pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result<Vec<ToolInfo>> {
let managed_client = self
.clients
@@ -497,6 +523,10 @@ impl McpConnectionManager {
let list_start = Instant::now();
let fetch_start = Instant::now();
let fetch_ticket = managed_client
.codex_apps_tools_cache_context
.as_ref()
.map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh));
let tools = list_tools_for_client_uncached(
CODEX_APPS_MCP_SERVER_NAME,
/*is_codex_apps_mcp_server*/ true,
@@ -514,11 +544,16 @@ impl McpConnectionManager {
&[],
);
write_codex_apps_tools_cache(
managed_client.codex_apps_tools_cache_context.as_ref(),
&managed_client.server_info,
&tools,
);
let tools =
match (
managed_client.codex_apps_tools_cache_context.as_ref(),
fetch_ticket,
) {
(Some(cache_context), Some(fetch_ticket)) => cache_context
.publish_if_newest_accepted(fetch_ticket, &managed_client.server_info, tools),
(None, None) => tools,
_ => unreachable!("Codex Apps fetch ticket requires cache context"),
};
emit_duration(
MCP_TOOLS_LIST_DURATION_METRIC,
list_start.elapsed(),
@@ -849,22 +884,6 @@ impl Drop for McpConnectionManager {
}
}
/// Creates the per-user tools cache context used only by the Codex Apps server.
fn codex_apps_tools_cache_context(
codex_home: &Path,
codex_apps_tools_cache_key: &CodexAppsToolsCacheKey,
) -> Option<CodexAppsToolsCacheContext> {
Some(CodexAppsToolsCacheContext {
codex_home: codex_home.to_path_buf(),
user_key: codex_apps_tools_cache_key.clone(),
})
}
/// Keeps regular MCP servers isolated from the Codex Apps tools cache.
fn regular_mcp_tools_cache_context() -> Option<CodexAppsToolsCacheContext> {
None
}
/// Makes ChatGPT authentication available to servers that explicitly opt in.
/// The HTTP transport applies it only when no configured authorization resolves.
fn chatgpt_auth_provider_for_server(
@@ -880,6 +899,10 @@ fn chatgpt_auth_provider_for_server(
chatgpt_auth_provider
}
fn should_share_codex_apps_tools_cache(server_name: &str, uses_env_bearer_token: bool) -> bool {
server_name == CODEX_APPS_MCP_SERVER_NAME && !uses_env_bearer_token
}
async fn emit_update(
submit_id: &str,
tx_event: &Sender<Event>,
@@ -893,6 +916,28 @@ async fn emit_update(
.await
}
fn mcp_startup_failure_reason(
entry: Option<&McpAuthStatusEntry>,
error: &StartupOutcomeError,
) -> Option<McpStartupFailureReason> {
if !error.is_authentication_required() {
return None;
}
match entry.map(|entry| entry.auth_state) {
Some(McpAuthState::LoggedOut(McpLoginRequirement::Reauthentication)) => {
Some(McpStartupFailureReason::ReauthenticationRequired)
}
Some(
McpAuthState::Unsupported
| McpAuthState::LoggedOut(McpLoginRequirement::Login)
| McpAuthState::BearerToken
| McpAuthState::OAuth,
)
| None => None,
}
}
fn mcp_init_error_display(
server_name: &str,
entry: Option<&McpAuthStatusEntry>,
@@ -939,20 +984,20 @@ fn mcp_init_error_display(
fn startup_outcome_error_message(error: StartupOutcomeError) -> String {
match error {
StartupOutcomeError::Cancelled => "MCP startup cancelled".to_string(),
StartupOutcomeError::Failed { error } => error,
StartupOutcomeError::Failed { error, .. } => error,
}
}
fn is_mcp_client_auth_required_error(error: &StartupOutcomeError) -> bool {
match error {
StartupOutcomeError::Failed { error } => error.contains("Auth required"),
StartupOutcomeError::Failed { error, .. } => error.contains("Auth required"),
_ => false,
}
}
fn is_mcp_client_startup_timeout_error(error: &StartupOutcomeError) -> bool {
match error {
StartupOutcomeError::Failed { error } => {
StartupOutcomeError::Failed { error, .. } => {
error.contains("request timed out")
|| error.contains("timed out handshaking with MCP server")
}
+163 -312
View File
@@ -1,11 +1,6 @@
use super::*;
use crate::codex_apps::CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION;
use crate::codex_apps::CodexAppsToolsCacheContext;
use crate::codex_apps::load_startup_cached_codex_apps_server_info;
use crate::codex_apps::load_startup_cached_codex_apps_tools_snapshot;
use crate::codex_apps::read_cached_codex_apps_tools;
use crate::codex_apps::write_cached_codex_apps_tools;
use crate::codex_apps::write_codex_apps_tools_cache;
use crate::codex_apps_cache::CodexAppsToolsCache;
use crate::codex_apps_cache::CodexAppsToolsCacheContext;
use crate::declared_openai_file_input_param_names;
use crate::elicitation::ElicitationRequestManager;
use crate::elicitation::elicitation_is_rejected_by_policy;
@@ -30,8 +25,10 @@ use codex_protocol::ToolName;
use codex_protocol::mcp::McpServerInfo;
use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::GranularApprovalConfig;
use codex_protocol::protocol::McpAuthStatus;
use codex_rmcp_client::InProcessTransportFactory;
use codex_rmcp_client::RmcpClient;
use futures::FutureExt;
use futures::future::BoxFuture;
use pretty_assertions::assert_eq;
use rmcp::model::CreateElicitationRequestParams;
use rmcp::model::ElicitationAction;
@@ -41,8 +38,10 @@ use rmcp::model::Meta;
use rmcp::model::NumberOrString;
use rmcp::model::Tool;
use std::collections::HashSet;
use std::io;
use std::sync::Arc;
use tempfile::tempdir;
use tokio::io::DuplexStream;
fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo {
ToolInfo {
@@ -63,31 +62,19 @@ fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo {
}
}
fn create_test_tool_with_connector(
server_name: &str,
tool_name: &str,
connector_id: &str,
connector_name: Option<&str>,
) -> ToolInfo {
let mut tool = create_test_tool(server_name, tool_name);
tool.connector_id = Some(connector_id.to_string());
tool.connector_name = connector_name.map(ToOwned::to_owned);
tool
}
fn create_codex_apps_tools_cache_context(
codex_home: PathBuf,
account_id: Option<&str>,
chatgpt_user_id: Option<&str>,
) -> CodexAppsToolsCacheContext {
CodexAppsToolsCacheContext {
CodexAppsToolsCache::default().context(
codex_home,
user_key: CodexAppsToolsCacheKey {
CodexAppsToolsCacheKey {
account_id: account_id.map(ToOwned::to_owned),
chatgpt_user_id: chatgpt_user_id.map(ToOwned::to_owned),
is_workspace_account: false,
},
}
)
}
fn create_test_server_info(title: &str) -> McpServerInfo {
@@ -101,6 +88,51 @@ fn create_test_server_info(title: &str) -> McpServerInfo {
}
}
struct TestInProcessTransportFactory;
impl InProcessTransportFactory for TestInProcessTransportFactory {
fn open(&self) -> BoxFuture<'static, io::Result<DuplexStream>> {
async {
let (client_stream, _server_stream) = tokio::io::duplex(1);
Ok(client_stream)
}
.boxed()
}
}
async fn create_ready_async_managed_client(tools: Vec<ToolInfo>) -> AsyncManagedClient {
let tool_filter = ToolFilter::default();
let managed_client = ManagedClient {
client: Arc::new(
RmcpClient::new_in_process_client(Arc::new(TestInProcessTransportFactory))
.await
.expect("create in-process RMCP client"),
),
server_info: create_test_server_info("Ready"),
tools,
tool_filter: tool_filter.clone(),
tool_timeout: None,
server_instructions: None,
server_supports_sandbox_state_meta_capability: false,
codex_apps_tools_cache_context: None,
};
AsyncManagedClient {
client: futures::future::ready::<Result<ManagedClient, StartupOutcomeError>>(Ok(
managed_client,
))
.boxed()
.shared(),
is_codex_apps_mcp_server: false,
cached_server_info: None,
codex_apps_tools_cache_context: None,
tool_filter,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
}
}
fn model_tool_names(tools: &[ToolInfo]) -> HashSet<ToolName> {
tools
.iter()
@@ -554,235 +586,25 @@ fn filter_tools_applies_per_server_filters() {
}
#[test]
fn codex_apps_tools_cache_is_overwritten_by_last_write() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let tools_gateway_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")];
let tools_gateway_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")];
write_cached_codex_apps_tools(&cache_context, &tools_gateway_1);
let cached_gateway_1 =
read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for first write");
assert_eq!(cached_gateway_1[0].callable_name, "one");
write_cached_codex_apps_tools(&cache_context, &tools_gateway_2);
let cached_gateway_2 =
read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for second write");
assert_eq!(cached_gateway_2[0].callable_name, "two");
}
#[test]
fn codex_apps_tools_cache_is_scoped_per_user() {
let codex_home = tempdir().expect("tempdir");
let cache_context_user_1 = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cache_context_user_2 = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-two"),
Some("user-two"),
);
let tools_user_1 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")];
let tools_user_2 = vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "two")];
write_cached_codex_apps_tools(&cache_context_user_1, &tools_user_1);
write_cached_codex_apps_tools(&cache_context_user_2, &tools_user_2);
let read_user_1 =
read_cached_codex_apps_tools(&cache_context_user_1).expect("cache entry for user one");
let read_user_2 =
read_cached_codex_apps_tools(&cache_context_user_2).expect("cache entry for user two");
assert_eq!(read_user_1[0].callable_name, "one");
assert_eq!(read_user_2[0].callable_name, "two");
assert_ne!(
cache_context_user_1.tools_cache_path(),
cache_context_user_2.tools_cache_path(),
"each user should get an isolated cache file"
);
}
#[test]
fn codex_apps_tools_cache_preserves_formerly_disallowed_connectors() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let tools = vec![
create_test_tool_with_connector(
CODEX_APPS_MCP_SERVER_NAME,
"formerly_blocked_tool",
"connector_2b0a9009c9c64bf9933a3dae3f2b1254",
Some("Formerly Blocked"),
),
create_test_tool_with_connector(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_tool",
"calendar",
Some("Calendar"),
),
];
write_cached_codex_apps_tools(&cache_context, &tools);
let cached = read_cached_codex_apps_tools(&cache_context).expect("cache entry exists for user");
assert_eq!(
cached
.iter()
.map(|tool| (tool.callable_name.as_str(), tool.connector_id.as_deref()))
.collect::<Vec<_>>(),
vec![
(
"formerly_blocked_tool",
Some("connector_2b0a9009c9c64bf9933a3dae3f2b1254")
),
("calendar_tool", Some("calendar")),
]
);
}
#[test]
fn codex_apps_tools_cache_is_ignored_when_schema_version_mismatches() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cache_path = cache_context.tools_cache_path();
if let Some(parent) = cache_path.parent() {
std::fs::create_dir_all(parent).expect("create parent");
}
let bytes = serde_json::to_vec_pretty(&serde_json::json!({
"schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION + 1,
"tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "one")],
}))
.expect("serialize");
std::fs::write(cache_path, bytes).expect("write");
assert!(read_cached_codex_apps_tools(&cache_context).is_none());
}
#[test]
fn codex_apps_tools_cache_is_ignored_when_json_is_invalid() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cache_path = cache_context.tools_cache_path();
if let Some(parent) = cache_path.parent() {
std::fs::create_dir_all(parent).expect("create parent");
}
std::fs::write(cache_path, b"{not json").expect("write");
assert!(read_cached_codex_apps_tools(&cache_context).is_none());
}
#[test]
fn startup_cached_codex_apps_tools_loads_from_disk_cache() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cached_tools = vec![create_test_tool(
fn codex_apps_env_bearer_token_bypasses_shared_tools_cache() {
assert!(!should_share_codex_apps_tools_cache(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_search",
)];
let server_info = create_test_server_info("Codex Apps");
write_codex_apps_tools_cache(Some(&cache_context), &server_info, &cached_tools);
let startup_tools = load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context))
.expect("expected startup snapshot to load from cache");
let cached_server_info = load_startup_cached_codex_apps_server_info(Some(&cache_context));
assert_eq!(startup_tools.len(), 1);
assert_eq!(startup_tools[0].server_name, CODEX_APPS_MCP_SERVER_NAME);
assert_eq!(startup_tools[0].callable_name, "calendar_search");
assert_eq!(cached_server_info, Some(server_info));
}
#[test]
fn startup_cached_codex_apps_tools_loads_without_server_info_cache() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let cache_path = cache_context.tools_cache_path();
if let Some(parent) = cache_path.parent() {
std::fs::create_dir_all(parent).expect("create parent");
}
let bytes = serde_json::to_vec_pretty(&serde_json::json!({
"schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION,
"tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")],
}))
.expect("serialize");
std::fs::write(cache_path, bytes).expect("write");
let startup_tools = load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context))
.expect("legacy startup snapshot should remain available");
let cached_server_info = load_startup_cached_codex_apps_server_info(Some(&cache_context));
assert_eq!(startup_tools.len(), 1);
assert_eq!(startup_tools[0].callable_name, "calendar_search");
assert_eq!(cached_server_info, None);
}
#[test]
fn codex_apps_server_info_cache_survives_legacy_tools_cache_write() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
let server_info = create_test_server_info("Codex Apps");
write_codex_apps_tools_cache(
Some(&cache_context),
&server_info,
&[create_test_tool(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_search",
)],
);
let cache_path = cache_context.tools_cache_path();
if let Some(parent) = cache_path.parent() {
std::fs::create_dir_all(parent).expect("create parent");
}
let bytes = serde_json::to_vec_pretty(&serde_json::json!({
"schema_version": CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION - 1,
"tools": [create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "calendar_search")],
}))
.expect("serialize");
std::fs::write(cache_path, bytes).expect("write legacy tools cache");
assert_eq!(
load_startup_cached_codex_apps_server_info(Some(&cache_context)),
Some(server_info)
);
assert!(load_startup_cached_codex_apps_tools_snapshot(Some(&cache_context)).is_none());
/*uses_env_bearer_token*/ true,
));
}
#[tokio::test]
async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() {
let startup_tools = vec![create_test_tool(
async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
cache_context.store_current_tools_for_test(vec![create_test_tool(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_create_event",
)];
)]);
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
@@ -798,8 +620,9 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending()
AsyncManagedClient {
client: pending_client,
is_codex_apps_mcp_server: true,
cached_tool_info_snapshot: Some(startup_tools),
cached_server_info: None,
codex_apps_tools_cache_context: Some(cache_context),
tool_filter: ToolFilter::default(),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
@@ -813,7 +636,7 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending()
tool.canonical_tool_name()
== ToolName::namespaced("mcp__codex_apps", "calendar_create_event")
})
.expect("tool from startup cache");
.expect("tool from shared cache");
assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME);
assert_eq!(tool.callable_name, "calendar_create_event");
}
@@ -836,8 +659,9 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() {
AsyncManagedClient {
client: pending_client,
is_codex_apps_mcp_server: true,
cached_tool_info_snapshot: Some(Vec::new()),
cached_server_info: Some(server_info.clone()),
codex_apps_tools_cache_context: None,
tool_filter: ToolFilter::default(),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
@@ -858,10 +682,8 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() {
#[tokio::test]
async fn list_all_tools_accepts_canonical_namespaced_tool_names() {
let startup_tools = vec![create_test_tool("rmcp", "echo")];
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
let managed_client =
create_ready_async_managed_client(vec![create_test_tool("rmcp", "echo")]).await;
let approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
let permission_profile = Constrained::allow_any(PermissionProfile::default());
let mut manager = McpConnectionManager::new_uninitialized(
@@ -869,18 +691,7 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() {
&permission_profile,
/*prefix_mcp_tool_names*/ false,
);
manager.clients.insert(
"rmcp".to_string(),
AsyncManagedClient {
client: pending_client,
is_codex_apps_mcp_server: false,
cached_tool_info_snapshot: Some(startup_tools),
cached_server_info: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);
manager.clients.insert("rmcp".to_string(), managed_client);
let tools = manager.list_all_tools().await;
let tool = tools
@@ -902,10 +713,8 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() {
#[tokio::test]
async fn list_all_tools_applies_legacy_mcp_prefix_by_default() {
let startup_tools = vec![create_test_tool("rmcp", "echo")];
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
let managed_client =
create_ready_async_managed_client(vec![create_test_tool("rmcp", "echo")]).await;
let approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
let permission_profile = Constrained::allow_any(PermissionProfile::default());
let mut manager = McpConnectionManager::new_uninitialized(
@@ -913,18 +722,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() {
&permission_profile,
/*prefix_mcp_tool_names*/ true,
);
manager.clients.insert(
"rmcp".to_string(),
AsyncManagedClient {
client: pending_client,
is_codex_apps_mcp_server: false,
cached_tool_info_snapshot: Some(startup_tools),
cached_server_info: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);
manager.clients.insert("rmcp".to_string(), managed_client);
let tools = manager.list_all_tools().await;
let tool = tools
@@ -945,7 +743,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() {
}
#[tokio::test]
async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_snapshot() {
async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() {
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
@@ -961,8 +759,9 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_
AsyncManagedClient {
client: pending_client,
is_codex_apps_mcp_server: true,
cached_tool_info_snapshot: None,
cached_server_info: None,
codex_apps_tools_cache_context: None,
tool_filter: ToolFilter::default(),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
@@ -998,8 +797,9 @@ async fn shutdown_cancels_pending_tool_listing() {
AsyncManagedClient {
client: pending_client,
is_codex_apps_mcp_server: true,
cached_tool_info_snapshot: None,
cached_server_info: None,
codex_apps_tools_cache_context: None,
tool_filter: ToolFilter::default(),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token,
@@ -1043,8 +843,9 @@ async fn shutdown_continues_after_caller_is_aborted() {
AsyncManagedClient {
client: blocking_client,
is_codex_apps_mcp_server: true,
cached_tool_info_snapshot: None,
cached_server_info: None,
codex_apps_tools_cache_context: None,
tool_filter: ToolFilter::default(),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
@@ -1071,7 +872,14 @@ async fn shutdown_continues_after_caller_is_aborted() {
}
#[tokio::test]
async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() {
async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
cache_context.store_current_tools_for_test(Vec::new());
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
@@ -1087,8 +895,9 @@ async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty()
AsyncManagedClient {
client: pending_client,
is_codex_apps_mcp_server: true,
cached_tool_info_snapshot: Some(Vec::new()),
cached_server_info: None,
codex_apps_tools_cache_context: Some(cache_context),
tool_filter: ToolFilter::default(),
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
@@ -1097,20 +906,27 @@ async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty()
let timeout_result =
tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()).await;
let tools = timeout_result.expect("cache-hit startup snapshot should not block");
let tools = timeout_result.expect("shared empty cache should not block");
assert!(tools.is_empty());
}
#[tokio::test]
async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails() {
let startup_tools = vec![create_test_tool(
async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
cache_context.store_current_tools_for_test(vec![create_test_tool(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_create_event",
)];
)]);
let server_info = create_test_server_info("Codex Apps");
let failed_client = futures::future::ready::<Result<ManagedClient, StartupOutcomeError>>(Err(
StartupOutcomeError::Failed {
error: "startup failed".to_string(),
is_authentication_required: false,
},
))
.boxed()
@@ -1128,8 +944,9 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails
AsyncManagedClient {
client: failed_client,
is_codex_apps_mcp_server: true,
cached_tool_info_snapshot: Some(startup_tools),
cached_server_info: Some(server_info.clone()),
codex_apps_tools_cache_context: Some(cache_context),
tool_filter: ToolFilter::default(),
startup_complete,
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
@@ -1143,7 +960,7 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails
tool.canonical_tool_name()
== ToolName::namespaced("mcp__codex_apps", "calendar_create_event")
})
.expect("tool from startup cache");
.expect("tool from shared cache");
assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME);
assert_eq!(tool.callable_name, "calendar_create_event");
assert_eq!(
@@ -1156,12 +973,10 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails
}
#[tokio::test]
async fn list_all_tools_adds_server_metadata_to_cached_tools() {
async fn list_all_tools_adds_server_metadata_to_tools() {
let server_name = "docs";
let startup_tools = vec![create_test_tool(server_name, "search")];
let pending_client = futures::future::pending::<Result<ManagedClient, StartupOutcomeError>>()
.boxed()
.shared();
let managed_client =
create_ready_async_managed_client(vec![create_test_tool(server_name, "search")]).await;
let approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
let permission_profile = Constrained::allow_any(PermissionProfile::default());
let mut manager = McpConnectionManager::new_uninitialized(
@@ -1182,18 +997,9 @@ async fn list_all_tools_adds_server_metadata_to_cached_tools() {
tool_approval_modes: HashMap::new(),
},
);
manager.clients.insert(
server_name.to_string(),
AsyncManagedClient {
client: pending_client,
is_codex_apps_mcp_server: false,
cached_tool_info_snapshot: Some(startup_tools),
cached_server_info: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()),
cancel_token: CancellationToken::new(),
},
);
manager
.clients
.insert(server_name.to_string(), managed_client);
let tools = manager.list_all_tools().await;
assert_eq!(tools.len(), 1);
@@ -1340,6 +1146,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() {
PathBuf::from("/tmp"),
),
codex_home.path().to_path_buf(),
CodexAppsToolsCache::default(),
CodexAppsToolsCacheKey {
account_id: None,
chatgpt_user_id: None,
@@ -1429,7 +1236,7 @@ fn mcp_init_error_display_prompts_for_github_pat() {
oauth_resource: None,
tools: HashMap::new(),
}),
auth_status: McpAuthStatus::Unsupported,
auth_state: McpAuthState::Unsupported,
};
let err: StartupOutcomeError = anyhow::anyhow!("OAuth is unsupported").into();
@@ -1456,6 +1263,50 @@ fn mcp_init_error_display_prompts_for_login_when_auth_required() {
assert_eq!(expected, display);
}
#[test]
fn mcp_startup_failure_reason_requires_existing_oauth_and_auth_failure() {
for (auth_state, is_authentication_required, expected) in [
(
Some(McpAuthState::LoggedOut(
McpLoginRequirement::Reauthentication,
)),
true,
Some(McpStartupFailureReason::ReauthenticationRequired),
),
(
Some(McpAuthState::LoggedOut(
McpLoginRequirement::Reauthentication,
)),
false,
None,
),
(
Some(McpAuthState::LoggedOut(McpLoginRequirement::Login)),
true,
None,
),
(Some(McpAuthState::Unsupported), true, None),
(Some(McpAuthState::BearerToken), true, None),
(Some(McpAuthState::OAuth), true, None),
(None, true, None),
] {
let entry = auth_state.map(|auth_state| McpAuthStatusEntry {
config: None,
auth_state,
});
let error = StartupOutcomeError::Failed {
error: "startup failed".to_string(),
is_authentication_required,
};
assert_eq!(
mcp_startup_failure_reason(entry.as_ref(), &error),
expected,
"auth_state={auth_state:?}, is_authentication_required={is_authentication_required}"
);
}
}
#[test]
fn mcp_init_error_display_reports_generic_errors() {
let server_name = "custom";
@@ -1483,7 +1334,7 @@ fn mcp_init_error_display_reports_generic_errors() {
oauth_resource: None,
tools: HashMap::new(),
}),
auth_status: McpAuthStatus::Unsupported,
auth_state: McpAuthState::Unsupported,
};
let err: StartupOutcomeError = anyhow::anyhow!("boom").into();
+6 -2
View File
@@ -35,8 +35,9 @@ pub use auth_elicitation::auth_elicitation_id;
pub use auth_elicitation::build_auth_elicitation;
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 codex_apps_cache::CodexAppsToolsCache;
pub use codex_apps_cache::CodexAppsToolsCacheKey;
pub use codex_apps_cache::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;
@@ -61,7 +62,9 @@ pub use mcp::McpOAuthScopesSource;
pub use mcp::ResolvedMcpOAuthScopes;
pub use mcp::compute_auth_statuses;
pub use mcp::discover_supported_scopes;
pub use mcp::discover_supported_scopes_with_http_client;
pub use mcp::oauth_login_support;
pub use mcp::oauth_login_support_with_http_client;
pub use mcp::resolve_oauth_scopes;
pub use mcp::should_retry_without_scopes;
@@ -73,6 +76,7 @@ pub use tools::declared_openai_file_input_param_names;
pub(crate) mod auth_elicitation;
mod catalog;
pub(crate) mod codex_apps;
pub(crate) mod codex_apps_cache;
pub(crate) mod connection_manager;
pub(crate) mod elicitation;
pub(crate) mod mcp;
+110 -39
View File
@@ -1,4 +1,5 @@
use std::collections::HashMap;
use std::sync::Arc;
use anyhow::Result;
use codex_config::McpServerAuth;
@@ -6,15 +7,19 @@ use codex_config::McpServerConfig;
use codex_config::McpServerTransportConfig;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::HttpClient;
use codex_login::CodexAuth;
use codex_protocol::protocol::McpAuthStatus;
use codex_rmcp_client::McpAuthState;
use codex_rmcp_client::OAuthProviderError;
use codex_rmcp_client::determine_streamable_http_auth_status;
use codex_rmcp_client::determine_streamable_http_auth_status_with_http_client;
use codex_rmcp_client::discover_streamable_http_oauth;
use codex_rmcp_client::discover_streamable_http_oauth_with_http_client;
use futures::FutureExt;
use futures::future::join_all;
use tracing::warn;
use crate::runtime::McpRuntimeContext;
use crate::server::EffectiveMcpServer;
#[derive(Debug, Clone)]
@@ -49,10 +54,54 @@ pub struct ResolvedMcpOAuthScopes {
#[derive(Debug, Clone)]
pub struct McpAuthStatusEntry {
pub config: Option<McpServerConfig>,
pub auth_status: McpAuthStatus,
pub auth_state: McpAuthState,
}
pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAuthLoginSupport {
let Some(mut config) = oauth_login_candidate(transport) else {
return McpOAuthLoginSupport::Unsupported;
};
match discover_streamable_http_oauth(
&config.url,
config.http_headers.clone(),
config.env_http_headers.clone(),
)
.await
{
Ok(Some(discovery)) => {
config.discovered_scopes = discovery.scopes_supported;
McpOAuthLoginSupport::Supported(config)
}
Ok(None) => McpOAuthLoginSupport::Unsupported,
Err(err) => McpOAuthLoginSupport::Unknown(err),
}
}
pub async fn oauth_login_support_with_http_client(
transport: &McpServerTransportConfig,
http_client: Arc<dyn HttpClient>,
) -> McpOAuthLoginSupport {
let Some(mut config) = oauth_login_candidate(transport) else {
return McpOAuthLoginSupport::Unsupported;
};
match discover_streamable_http_oauth_with_http_client(
&config.url,
config.http_headers.clone(),
config.env_http_headers.clone(),
http_client,
)
.await
{
Ok(Some(discovery)) => {
config.discovered_scopes = discovery.scopes_supported;
McpOAuthLoginSupport::Supported(config)
}
Ok(None) => McpOAuthLoginSupport::Unsupported,
Err(err) => McpOAuthLoginSupport::Unknown(err),
}
}
fn oauth_login_candidate(transport: &McpServerTransportConfig) -> Option<McpOAuthLoginConfig> {
let McpServerTransportConfig::StreamableHttp {
url,
bearer_token_env_var,
@@ -60,24 +109,17 @@ pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAu
env_http_headers,
} = transport
else {
return McpOAuthLoginSupport::Unsupported;
return None;
};
if bearer_token_env_var.is_some() {
return McpOAuthLoginSupport::Unsupported;
}
match discover_streamable_http_oauth(url, http_headers.clone(), env_http_headers.clone()).await
{
Ok(Some(discovery)) => McpOAuthLoginSupport::Supported(McpOAuthLoginConfig {
url: url.clone(),
http_headers: http_headers.clone(),
env_http_headers: env_http_headers.clone(),
discovered_scopes: discovery.scopes_supported,
}),
Ok(None) => McpOAuthLoginSupport::Unsupported,
Err(err) => McpOAuthLoginSupport::Unknown(err),
return None;
}
Some(McpOAuthLoginConfig {
url: url.clone(),
http_headers: http_headers.clone(),
env_http_headers: env_http_headers.clone(),
discovered_scopes: None,
})
}
pub async fn discover_supported_scopes(
@@ -89,6 +131,16 @@ pub async fn discover_supported_scopes(
}
}
pub async fn discover_supported_scopes_with_http_client(
transport: &McpServerTransportConfig,
http_client: Arc<dyn HttpClient>,
) -> Option<Vec<String>> {
match oauth_login_support_with_http_client(transport, http_client).await {
McpOAuthLoginSupport::Supported(config) => config.discovered_scopes,
McpOAuthLoginSupport::Unsupported | McpOAuthLoginSupport::Unknown(_) => None,
}
}
pub fn resolve_oauth_scopes(
explicit_scopes: Option<Vec<String>>,
configured_scopes: Option<Vec<String>>,
@@ -133,6 +185,7 @@ pub async fn compute_auth_statuses<'a, I>(
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
auth: Option<&CodexAuth>,
runtime_context: &McpRuntimeContext,
) -> HashMap<String, McpAuthStatusEntry>
where
I: IntoIterator<Item = (&'a String, &'a EffectiveMcpServer)>,
@@ -140,6 +193,7 @@ where
let futures = servers.into_iter().map(|(name, server)| {
let name = name.clone();
let config = server.configured_config().cloned();
let runtime_context = runtime_context.clone();
let has_runtime_auth = config
.as_ref()
.is_some_and(|config| matches!(&config.auth, McpServerAuth::ChatGpt))
@@ -154,7 +208,7 @@ where
)
});
async move {
let auth_status = match config.as_ref() {
let auth_state = match config.as_ref() {
Some(config) => {
match compute_auth_status(
&name,
@@ -162,6 +216,7 @@ where
store_mode,
keyring_backend_kind,
has_runtime_auth,
&runtime_context,
)
.await
{
@@ -170,16 +225,13 @@ where
warn!(
"failed to determine auth status for MCP server `{name}`: {error:?}"
);
McpAuthStatus::Unsupported
McpAuthState::Unsupported
}
}
}
None => McpAuthStatus::Unsupported,
};
let entry = McpAuthStatusEntry {
config,
auth_status,
None => McpAuthState::Unsupported,
};
let entry = McpAuthStatusEntry { config, auth_state };
(name, entry)
}
});
@@ -193,34 +245,53 @@ async fn compute_auth_status(
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
has_runtime_auth: bool,
) -> Result<McpAuthStatus> {
runtime_context: &McpRuntimeContext,
) -> Result<McpAuthState> {
if !config.enabled {
return Ok(McpAuthStatus::Unsupported);
return Ok(McpAuthState::Unsupported);
}
if has_runtime_auth {
return Ok(McpAuthStatus::BearerToken);
return Ok(McpAuthState::BearerToken);
}
match &config.transport {
McpServerTransportConfig::Stdio { .. } => Ok(McpAuthStatus::Unsupported),
McpServerTransportConfig::Stdio { .. } => Ok(McpAuthState::Unsupported),
McpServerTransportConfig::StreamableHttp {
url,
bearer_token_env_var,
http_headers,
env_http_headers,
} => {
determine_streamable_http_auth_status(
server_name,
url,
bearer_token_env_var.as_deref(),
http_headers.clone(),
env_http_headers.clone(),
store_mode,
keyring_backend_kind,
)
.boxed()
.await
if config.is_local_environment() {
determine_streamable_http_auth_status(
server_name,
url,
bearer_token_env_var.as_deref(),
http_headers.clone(),
env_http_headers.clone(),
store_mode,
keyring_backend_kind,
)
.boxed()
.await
} else {
let http_client = runtime_context
.resolve_http_client(server_name, config)
.map_err(anyhow::Error::msg)?;
determine_streamable_http_auth_status_with_http_client(
server_name,
url,
bearer_token_env_var.as_deref(),
http_headers.clone(),
env_http_headers.clone(),
store_mode,
keyring_backend_kind,
http_client,
)
.boxed()
.await
}
}
}
}
+11 -2
View File
@@ -5,7 +5,9 @@ pub use auth::McpOAuthScopesSource;
pub use auth::ResolvedMcpOAuthScopes;
pub use auth::compute_auth_statuses;
pub use auth::discover_supported_scopes;
pub use auth::discover_supported_scopes_with_http_client;
pub use auth::oauth_login_support;
pub use auth::oauth_login_support_with_http_client;
pub use auth::resolve_oauth_scopes;
pub use auth::should_retry_without_scopes;
@@ -42,7 +44,8 @@ use serde_json::Value;
use tokio_util::sync::CancellationToken;
use crate::ResolvedMcpCatalog;
use crate::codex_apps::codex_apps_tools_cache_key;
use crate::codex_apps_cache::CodexAppsToolsCache;
use crate::codex_apps_cache::codex_apps_tools_cache_key;
use crate::connection_manager::McpConnectionManager;
use crate::runtime::McpRuntimeContext;
use crate::server::EffectiveMcpServer;
@@ -299,6 +302,7 @@ pub async fn read_mcp_resource(
config: &McpConfig,
auth: Option<&CodexAuth>,
runtime_context: McpRuntimeContext,
codex_apps_tools_cache: CodexAppsToolsCache,
server: &str,
uri: &str,
) -> anyhow::Result<ReadResourceResult> {
@@ -309,6 +313,7 @@ pub async fn read_mcp_resource(
config.mcp_oauth_credentials_store_mode,
config.auth_keyring_backend_kind,
auth,
&runtime_context,
)
.await;
let (tx_event, rx_event) = unbounded();
@@ -326,6 +331,7 @@ pub async fn read_mcp_resource(
PermissionProfile::default(),
runtime_context,
config.codex_home.clone(),
codex_apps_tools_cache,
codex_apps_tools_cache_key(auth),
config.prefix_mcp_tool_names,
config.client_elicitation_capability.clone(),
@@ -358,6 +364,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
auth: Option<&CodexAuth>,
submit_id: String,
runtime_context: McpRuntimeContext,
codex_apps_tools_cache: CodexAppsToolsCache,
detail: McpSnapshotDetail,
) -> McpServerStatusSnapshot {
let mcp_servers = effective_mcp_servers(config, auth);
@@ -378,6 +385,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
config.mcp_oauth_credentials_store_mode,
config.auth_keyring_backend_kind,
auth,
&runtime_context,
)
.await;
@@ -399,6 +407,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
PermissionProfile::default(),
runtime_context,
config.codex_home.clone(),
codex_apps_tools_cache,
codex_apps_tools_cache_key(auth),
config.prefix_mcp_tool_names,
config.client_elicitation_capability.clone(),
@@ -558,7 +567,7 @@ fn auth_statuses_from_entries(
) -> HashMap<String, McpAuthStatus> {
auth_status_entries
.iter()
.map(|(name, entry)| (name.clone(), entry.auth_status))
.map(|(name, entry)| (name.clone(), McpAuthStatus::from(entry.auth_state)))
.collect::<HashMap<_, _>>()
}
+29 -2
View File
@@ -173,8 +173,35 @@ fn environment_cwd(
fn bind_environment_env_vars(config: &mut McpServerConfig) -> Result<(), String> {
let is_local_environment = config.is_local_environment();
let McpServerTransportConfig::Stdio { env_vars, .. } = &mut config.transport else {
return Ok(());
let env_vars = match &mut config.transport {
McpServerTransportConfig::Stdio { env_vars, .. } => env_vars,
// Never resolve executor-owned environment references in the host process.
// Remove this rejection once the owning executor resolves these fields.
McpServerTransportConfig::StreamableHttp {
bearer_token_env_var,
env_http_headers,
..
} => {
if is_local_environment {
return Ok(());
}
if bearer_token_env_var.is_some() {
return Err(
"`bearer_token_env_var` requires executor-side environment resolution for an executor-owned HTTP MCP"
.to_string(),
);
}
if env_http_headers
.as_ref()
.is_some_and(|headers| !headers.is_empty())
{
return Err(
"`env_http_headers` requires executor-side environment resolution for an executor-owned HTTP MCP"
.to_string(),
);
}
return Ok(());
}
};
for env_var in env_vars {
match env_var {
@@ -275,6 +275,98 @@ fn environment_placement_rejects_orchestrator_env_vars() {
);
}
#[test]
fn remote_environment_placement_rejects_http_env_references() {
let plugin_root = plugin_root();
let outcome = parse_executor_plugin_mcp_config(
&plugin_root_uri(&plugin_root),
r#"{
"bearer": {
"url": "https://example.com/bearer",
"bearer_token_env_var": "TOKEN"
},
"headers": {
"url": "https://example.com/headers",
"env_http_headers": {"Authorization": "TOKEN"}
}
}"#,
"executor-1",
)
.expect("parse plugin MCP config");
assert_eq!(
outcome,
PluginMcpConfigParseOutcome {
servers: BTreeMap::new(),
errors: vec![
PluginMcpServerParseError {
name: "bearer".to_string(),
message: "`bearer_token_env_var` requires executor-side environment resolution for an executor-owned HTTP MCP"
.to_string(),
},
PluginMcpServerParseError {
name: "headers".to_string(),
message: "`env_http_headers` requires executor-side environment resolution for an executor-owned HTTP MCP"
.to_string(),
},
],
}
);
}
#[test]
fn local_environment_placement_preserves_http_env_references() {
let plugin_root = plugin_root();
let outcome = parse_executor_plugin_mcp_config(
&plugin_root_uri(&plugin_root),
r#"{
"demo": {
"url": "https://example.com/mcp",
"bearer_token_env_var": "TOKEN",
"env_http_headers": {"X-Account": "ACCOUNT_ID"}
}
}"#,
DEFAULT_MCP_SERVER_ENVIRONMENT_ID,
)
.expect("parse plugin MCP config");
assert_eq!(
outcome,
PluginMcpConfigParseOutcome {
servers: BTreeMap::from([(
"demo".to_string(),
McpServerConfig {
auth: Default::default(),
transport: McpServerTransportConfig::StreamableHttp {
url: "https://example.com/mcp".to_string(),
bearer_token_env_var: Some("TOKEN".to_string()),
http_headers: None,
env_http_headers: Some(HashMap::from([(
"X-Account".to_string(),
"ACCOUNT_ID".to_string(),
)])),
},
environment_id: DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(),
enabled: true,
required: false,
supports_parallel_tool_calls: false,
disabled_reason: None,
startup_timeout_sec: None,
tool_timeout_sec: None,
default_tools_approval_mode: None,
enabled_tools: None,
disabled_tools: None,
scopes: None,
oauth: None,
oauth_resource: None,
tools: HashMap::new(),
},
)]),
errors: Vec::new(),
}
);
}
#[test]
fn local_environment_placement_preserves_local_env_vars() {
let plugin_root = plugin_root();
+91 -45
View File
@@ -2,7 +2,7 @@
//!
//! This module owns startup of individual RMCP clients: building the transport,
//! initializing the server, listing raw tools, applying per-server tool filters,
//! and exposing cached startup snapshots while a client is still connecting.
//! and exposing cached Codex Apps tools while a client is still connecting.
//! Higher-level aggregation and resource/tool APIs live in
//! [`crate::connection_manager`].
@@ -17,15 +17,12 @@ use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant;
use crate::codex_apps::CachedCodexAppsToolsLoad;
use crate::codex_apps::CodexAppsToolsCacheContext;
use crate::codex_apps::load_cached_codex_apps_tools;
use crate::codex_apps::load_startup_cached_codex_apps_server_info;
use crate::codex_apps::load_startup_cached_codex_apps_tools_snapshot;
use crate::codex_apps::normalize_codex_apps_callable_name;
use crate::codex_apps::normalize_codex_apps_callable_namespace;
use crate::codex_apps::normalize_codex_apps_tool_title;
use crate::codex_apps::write_codex_apps_tools_cache;
use crate::codex_apps_cache::CodexAppsToolsCacheContext;
use crate::codex_apps_cache::CodexAppsToolsFetchSource;
use crate::codex_apps_cache::load_startup_cached_codex_apps_server_info;
use crate::elicitation::ElicitationRequestManager;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
use crate::mcp::ToolPluginProvenance;
@@ -66,6 +63,7 @@ use rmcp::model::InitializeRequestParams;
use rmcp::model::JsonObject;
use rmcp::model::ProtocolVersion;
use rmcp::model::Tool as RmcpTool;
use rmcp::transport::auth::AuthError;
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use tracing::instrument;
@@ -105,9 +103,10 @@ pub(crate) struct ManagedClient {
impl ManagedClient {
fn listed_tools(&self) -> Vec<ToolInfo> {
let total_start = Instant::now();
if let Some(cache_context) = self.codex_apps_tools_cache_context.as_ref()
&& let CachedCodexAppsToolsLoad::Hit(tools) =
load_cached_codex_apps_tools(cache_context)
if let Some(tools) = self
.codex_apps_tools_cache_context
.as_ref()
.and_then(CodexAppsToolsCacheContext::current_tools)
{
emit_duration(
MCP_TOOLS_LIST_DURATION_METRIC,
@@ -133,8 +132,9 @@ impl ManagedClient {
pub(crate) struct AsyncManagedClient {
pub(crate) client: Shared<BoxFuture<'static, Result<ManagedClient, StartupOutcomeError>>>,
pub(crate) is_codex_apps_mcp_server: bool,
pub(crate) cached_tool_info_snapshot: Option<Vec<ToolInfo>>,
pub(crate) cached_server_info: Option<McpServerInfo>,
pub(crate) codex_apps_tools_cache_context: Option<CodexAppsToolsCacheContext>,
pub(crate) tool_filter: ToolFilter,
pub(crate) startup_complete: Arc<AtomicBool>,
pub(crate) tool_plugin_provenance: Arc<ToolPluginProvenance>,
pub(crate) cancel_token: CancellationToken,
@@ -165,19 +165,15 @@ impl AsyncManagedClient {
.configured_config()
.map(ToolFilter::from_config)
.unwrap_or_default();
let (cached_tool_info_snapshot, cached_server_info) = if is_codex_apps_mcp_server {
(
load_startup_cached_codex_apps_tools_snapshot(
codex_apps_tools_cache_context.as_ref(),
),
load_startup_cached_codex_apps_server_info(codex_apps_tools_cache_context.as_ref()),
)
let cached_server_info = if is_codex_apps_mcp_server {
codex_apps_tools_cache_context
.as_ref()
.and_then(load_startup_cached_codex_apps_server_info)
} else {
(None, None)
None
};
let cached_tool_info_snapshot =
cached_tool_info_snapshot.map(|tools| filter_tools(tools, &tool_filter));
let startup_tool_filter = tool_filter;
let startup_tool_filter = tool_filter.clone();
let codex_apps_tools_cache_context_for_fut = codex_apps_tools_cache_context.clone();
let startup_complete = Arc::new(AtomicBool::new(false));
let startup_complete_for_fut = Arc::clone(&startup_complete);
let cancel_token_for_fut = cancel_token.clone();
@@ -214,7 +210,7 @@ impl AsyncManagedClient {
tool_filter: startup_tool_filter,
tx_event,
elicitation_requests,
codex_apps_tools_cache_context,
codex_apps_tools_cache_context: codex_apps_tools_cache_context_for_fut,
client_elicitation_capability,
supports_openai_form_elicitation,
},
@@ -232,7 +228,10 @@ impl AsyncManagedClient {
outcome
};
let client = fut.in_current_span().boxed().shared();
if cached_tool_info_snapshot.is_some() {
if codex_apps_tools_cache_context
.as_ref()
.is_some_and(CodexAppsToolsCacheContext::has_current_tools)
{
let startup_task = client.clone();
tokio::spawn(async move {
let _ = startup_task.await;
@@ -242,8 +241,9 @@ impl AsyncManagedClient {
Self {
client,
is_codex_apps_mcp_server,
cached_tool_info_snapshot,
cached_server_info,
codex_apps_tools_cache_context,
tool_filter,
startup_complete,
tool_plugin_provenance,
cancel_token,
@@ -265,15 +265,29 @@ impl AsyncManagedClient {
}
}
pub(crate) fn has_cached_tools(&self) -> bool {
self.codex_apps_tools_cache_context
.as_ref()
.is_some_and(CodexAppsToolsCacheContext::has_current_tools)
}
fn cached_tools(&self) -> Option<Vec<ToolInfo>> {
self.codex_apps_tools_cache_context
.as_ref()
.and_then(CodexAppsToolsCacheContext::current_tools)
.map(|tools| filter_tools(tools, &self.tool_filter))
}
pub(crate) async fn listed_tools(&self) -> Option<Vec<ToolInfo>> {
// Keep cache payloads raw; plugin provenance is resolved per-session at read time.
let tools = if let Some(startup_tools) = self.cached_tool_info_snapshot_while_initializing()
let tools = if !self.startup_complete.load(Ordering::Acquire)
&& let Some(startup_tools) = self.cached_tools()
{
Some(startup_tools)
} else {
match self.client().await {
Ok(client) => Some(client.listed_tools()),
Err(_) => self.cached_tool_info_snapshot.clone(),
Err(_) => self.cached_tools(),
}
}?;
Some(if self.is_codex_apps_mcp_server {
@@ -282,13 +296,6 @@ impl AsyncManagedClient {
prepare_regular_mcp_tools_for_model(tools, &self.tool_plugin_provenance)
})
}
fn cached_tool_info_snapshot_while_initializing(&self) -> Option<Vec<ToolInfo>> {
if !self.startup_complete.load(Ordering::Acquire) {
return self.cached_tool_info_snapshot.clone();
}
None
}
}
#[derive(Debug, Clone, thiserror::Error)]
@@ -298,13 +305,39 @@ pub(crate) enum StartupOutcomeError {
// We can't store the original error here because anyhow::Error doesn't implement
// `Clone`.
#[error("MCP startup failed: {error}")]
Failed { error: String },
Failed {
error: String,
is_authentication_required: bool,
},
}
impl StartupOutcomeError {
pub(crate) fn is_authentication_required(&self) -> bool {
match self {
Self::Cancelled => false,
Self::Failed {
is_authentication_required,
..
} => *is_authentication_required,
}
}
}
impl From<anyhow::Error> for StartupOutcomeError {
fn from(error: anyhow::Error) -> Self {
let is_authentication_required = error.chain().any(|source| {
source
.downcast_ref::<AuthError>()
.is_some_and(|auth_error| {
matches!(
auth_error,
AuthError::AuthorizationRequired | AuthError::TokenExpired
)
})
});
Self::Failed {
error: error.to_string(),
is_authentication_required,
}
}
}
@@ -570,6 +603,9 @@ async fn start_server_task(
.is_some();
let list_start = Instant::now();
let fetch_start = Instant::now();
let fetch_ticket = codex_apps_tools_cache_context
.as_ref()
.map(|cache_context| cache_context.begin_fetch(CodexAppsToolsFetchSource::Startup));
let tools = list_tools_for_client_uncached(
&server_name,
is_codex_apps_mcp_server,
@@ -585,21 +621,20 @@ async fn start_server_task(
&[],
);
let server_info = mcp_server_info_from_implementation(initialize_result.server_info);
let codex_apps_tools_cache_context = if is_codex_apps_mcp_server {
write_codex_apps_tools_cache(
codex_apps_tools_cache_context.as_ref(),
&server_info,
&tools,
);
let tools = match (codex_apps_tools_cache_context.as_ref(), fetch_ticket) {
(Some(cache_context), Some(fetch_ticket)) => {
cache_context.publish_if_newest_accepted(fetch_ticket, &server_info, tools)
}
(None, None) => tools,
_ => unreachable!("Codex Apps fetch ticket requires cache context"),
};
if is_codex_apps_mcp_server {
emit_duration(
MCP_TOOLS_LIST_DURATION_METRIC,
list_start.elapsed(),
&[("cache", "miss")],
);
codex_apps_tools_cache_context
} else {
None
};
}
let tools = filter_tools(tools, &tool_filter);
let managed = ManagedClient {
@@ -757,6 +792,17 @@ mod tests {
use pretty_assertions::assert_eq;
use rmcp::model::JsonObject;
use rmcp::model::Meta;
use rmcp::transport::auth::AuthError;
#[test]
fn startup_outcome_error_identifies_authentication_required() {
let error = anyhow::Error::new(AuthError::AuthorizationRequired)
.context("failed to initialize MCP server");
let error = StartupOutcomeError::from(error);
assert!(error.is_authentication_required());
}
#[test]
fn mcp_initialize_advertises_openai_form_only_when_supported() {
+16
View File
@@ -11,6 +11,8 @@ use std::time::Duration;
use codex_exec_server::Environment;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::HttpClient;
use codex_exec_server::ReqwestHttpClient;
use codex_protocol::models::PermissionProfile;
use codex_utils_path_uri::PathUri;
use serde::Deserialize;
@@ -81,6 +83,20 @@ impl McpRuntimeContext {
config.environment_id
))
}
/// Resolves the HTTP capability owned by the server's configured environment.
pub fn resolve_http_client(
&self,
server_name: &str,
config: &codex_config::McpServerConfig,
) -> Result<Arc<dyn HttpClient>, String> {
Ok(self
.resolve_server_environment(server_name, config)?
.map_or_else(
|| Arc::new(ReqwestHttpClient) as Arc<dyn HttpClient>,
|environment| environment.get_http_client(),
))
}
}
pub(crate) fn emit_duration(metric: &str, duration: Duration, tags: &[(&str, &str)]) {
+1
View File
@@ -30,6 +30,7 @@ gethostname = { workspace = true }
indexmap = { workspace = true, features = ["serde"] }
multimap = { workspace = true }
prost = "0.14.3"
regex-lite = { workspace = true }
schemars = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_ignored = { workspace = true }
+181 -16
View File
@@ -19,6 +19,7 @@ use super::requirements_exec_policy::RequirementsExecPolicyToml;
use crate::Constrained;
use crate::ConstraintError;
use crate::ManagedHooksRequirementsToml;
use crate::mcp_requirements::McpServerRequirement;
use crate::mcp_types::AppToolApproval;
use crate::permissions_toml::PermissionProfileToml;
use crate::types::WindowsSandboxModeToml;
@@ -217,18 +218,6 @@ impl ConfigRequirements {
}
}
#[derive(Deserialize, Debug, Clone, PartialEq, Eq)]
#[serde(untagged)]
pub enum McpServerIdentity {
Command { command: String },
Url { url: String },
}
#[derive(Deserialize, Debug, Clone, PartialEq, Eq)]
pub struct McpServerRequirement {
pub identity: McpServerIdentity,
}
#[derive(Deserialize, Debug, Clone, Default, PartialEq, Eq)]
pub struct PluginRequirementsToml {
pub mcp_servers: Option<BTreeMap<String, McpServerRequirement>>,
@@ -1201,6 +1190,25 @@ impl ConfigRequirementsToml {
}
}
fn validate_mcp_server_requirements(
requirements: &BTreeMap<String, McpServerRequirement>,
source: &RequirementSource,
plugin_name: Option<&str>,
) -> Result<(), ConstraintError> {
for (server_name, requirement) in requirements {
requirement
.validate()
.map_err(|reason| ConstraintError::McpServerRequirementParse {
server_name: plugin_name
.map(|plugin_name| format!("{plugin_name}/{server_name}"))
.unwrap_or_else(|| server_name.clone()),
requirement_source: source.clone(),
reason,
})?;
}
Ok(())
}
impl TryFrom<ConfigRequirementsWithSources> for ConfigRequirements {
type Error = ConstraintError;
@@ -1233,6 +1241,25 @@ impl TryFrom<ConfigRequirementsWithSources> for ConfigRequirements {
guardian_policy_config,
} = toml;
if let Some(requirements) = &mcp_servers {
validate_mcp_server_requirements(
&requirements.value,
&requirements.source,
/*plugin_name*/ None,
)?;
}
if let Some(plugin_requirements) = &plugins {
for (plugin_name, plugin) in &plugin_requirements.value {
if let Some(requirements) = &plugin.mcp_servers {
validate_mcp_server_requirements(
requirements,
&plugin_requirements.source,
Some(plugin_name),
)?;
}
}
}
let approval_policy = match allowed_approval_policies {
Some(Sourced {
value: policies,
@@ -1542,6 +1569,9 @@ pub fn sandbox_mode_requirement_for_permission_profile(
mod tests {
use super::*;
use crate::HookEventsToml;
use crate::McpServerCommandMatcher;
use crate::McpServerIdentity;
use crate::McpServerValueMatcher;
use anyhow::Result;
use codex_execpolicy::Decision;
use codex_execpolicy::Evaluation;
@@ -3454,6 +3484,9 @@ command = "python3 /enterprise/hooks/pre.py"
#[test]
fn deserialize_mcp_server_requirements() -> Result<()> {
let toml_str = r#"
[mcp_servers.docs]
description = "ignored legacy field"
[mcp_servers.docs.identity]
command = "codex-mcp"
@@ -3469,7 +3502,7 @@ command = "python3 /enterprise/hooks/pre.py"
BTreeMap::from([
(
"docs".to_string(),
McpServerRequirement {
McpServerRequirement::Identity {
identity: McpServerIdentity::Command {
command: "codex-mcp".to_string(),
},
@@ -3477,7 +3510,7 @@ command = "python3 /enterprise/hooks/pre.py"
),
(
"remote".to_string(),
McpServerRequirement {
McpServerRequirement::Identity {
identity: McpServerIdentity::Url {
url: "https://example.com/mcp".to_string(),
},
@@ -3490,6 +3523,74 @@ command = "python3 /enterprise/hooks/pre.py"
Ok(())
}
#[test]
fn deserialize_mcp_server_matcher_requirements() -> Result<()> {
let toml_str = r#"
[mcp_servers.internal_mcp_proxy.identity]
command = { executable = "company-cli", args = [
{ match = "exact", value = "mcp" },
{ match = "exact", value = "proxy" },
{ match = "exact", value = "--server" },
{ match = "regex", expression = '^https://[A-Za-z0-9-]+\.mcp\.internal\.example\.com(?::443)?(?:/.*)?$' },
] }
"#;
let requirements: ConfigRequirements =
with_unknown_source(from_str(toml_str)?).try_into()?;
assert_eq!(
requirements.mcp_servers,
Some(Sourced::new(
BTreeMap::from([(
"internal_mcp_proxy".to_string(),
McpServerRequirement::Command(McpServerCommandMatcher {
executable: "company-cli".to_string(),
args: vec![
McpServerValueMatcher::Exact {
value: "mcp".to_string(),
},
McpServerValueMatcher::Exact {
value: "proxy".to_string(),
},
McpServerValueMatcher::Exact {
value: "--server".to_string(),
},
McpServerValueMatcher::Regex {
expression: r"^https://[A-Za-z0-9-]+\.mcp\.internal\.example\.com(?::443)?(?:/.*)?$"
.to_string(),
},
],
}),
)]),
RequirementSource::Unknown,
))
);
Ok(())
}
#[test]
fn invalid_mcp_server_requirement_regex_reports_the_server_name_and_source() -> Result<()> {
let toml_str = r#"
[mcp_servers.broken_rule.identity]
url = { match = "regex", expression = "[" }
"#;
let err = ConfigRequirements::try_from(with_unknown_source(from_str(toml_str)?))
.expect_err("invalid matcher regex should fail requirements normalization");
let ConstraintError::McpServerRequirementParse {
server_name,
requirement_source,
reason,
} = err
else {
panic!("unexpected error: {err:?}");
};
assert_eq!(server_name, "broken_rule");
assert_eq!(requirement_source, RequirementSource::Unknown);
assert!(reason.contains("invalid regex `[`"), "{reason}");
Ok(())
}
#[test]
fn deserialize_plugin_mcp_server_requirements() -> Result<()> {
let toml_str = r#"
@@ -3511,7 +3612,7 @@ command = "python3 /enterprise/hooks/pre.py"
PluginRequirementsToml {
mcp_servers: Some(BTreeMap::from([(
"remote".to_string(),
McpServerRequirement {
McpServerRequirement::Identity {
identity: McpServerIdentity::Url {
url: "https://example.com/mcp".to_string(),
},
@@ -3524,7 +3625,7 @@ command = "python3 /enterprise/hooks/pre.py"
PluginRequirementsToml {
mcp_servers: Some(BTreeMap::from([(
"sample".to_string(),
McpServerRequirement {
McpServerRequirement::Identity {
identity: McpServerIdentity::Command {
command: "sample-mcp".to_string(),
},
@@ -3539,6 +3640,70 @@ command = "python3 /enterprise/hooks/pre.py"
Ok(())
}
#[test]
fn deserialize_plugin_mcp_server_matcher_requirement() -> Result<()> {
let toml_str = r#"
[plugins."sample@test".mcp_servers.internal_proxy.identity]
command = { executable = "company-cli", args = [
{ match = "exact", value = "mcp" },
{ match = "regex", expression = '^https://[a-z]+\.example\.com$' },
] }
"#;
let requirements: ConfigRequirements =
with_unknown_source(from_str(toml_str)?).try_into()?;
assert_eq!(
requirements.plugins,
Some(Sourced::new(
BTreeMap::from([(
"sample@test".to_string(),
PluginRequirementsToml {
mcp_servers: Some(BTreeMap::from([(
"internal_proxy".to_string(),
McpServerRequirement::Command(McpServerCommandMatcher {
executable: "company-cli".to_string(),
args: vec![
McpServerValueMatcher::Exact {
value: "mcp".to_string(),
},
McpServerValueMatcher::Regex {
expression: r"^https://[a-z]+\.example\.com$".to_string(),
},
],
}),
)])),
},
)]),
RequirementSource::Unknown,
))
);
Ok(())
}
#[test]
fn invalid_plugin_mcp_server_regex_reports_plugin_and_server_name() -> Result<()> {
let toml_str = r#"
[plugins."sample@test".mcp_servers.broken_rule.identity]
url = { match = "regex", expression = "[" }
"#;
let err = ConfigRequirements::try_from(with_unknown_source(from_str(toml_str)?))
.expect_err("invalid plugin MCP regex should fail requirements normalization");
let ConstraintError::McpServerRequirementParse {
server_name,
requirement_source,
reason,
} = err
else {
panic!("unexpected error: {err:?}");
};
assert_eq!(server_name, "sample@test/broken_rule");
assert_eq!(requirement_source, RequirementSource::Unknown);
assert!(reason.contains("invalid regex `[`"), "{reason}");
Ok(())
}
#[test]
fn deserialize_exec_policy_requirements() -> Result<()> {
let toml_str = r#"
+9
View File
@@ -24,6 +24,15 @@ pub enum ConstraintError {
requirement_source: RequirementSource,
reason: String,
},
#[error(
"invalid requirement for MCP server `{server_name}` (set by {requirement_source}): {reason}"
)]
McpServerRequirementParse {
server_name: String,
requirement_source: RequirementSource,
reason: String,
},
}
impl ConstraintError {
+5 -2
View File
@@ -12,6 +12,7 @@ mod key_aliases;
pub mod loader;
mod marketplace_edit;
mod mcp_edit;
mod mcp_requirements;
mod mcp_types;
mod merge;
mod overrides;
@@ -66,8 +67,6 @@ pub use config_requirements::FilesystemDenyReadPattern;
pub use config_requirements::MarketplaceAllowedSourceKind;
pub use config_requirements::MarketplaceAllowedSourceToml;
pub use config_requirements::MarketplaceRequirementsToml;
pub use config_requirements::McpServerIdentity;
pub use config_requirements::McpServerRequirement;
pub use config_requirements::NetworkConstraints;
pub use config_requirements::NetworkDomainPermissionToml;
pub use config_requirements::NetworkDomainPermissionsToml;
@@ -113,6 +112,10 @@ pub use marketplace_edit::remove_user_marketplace;
pub use marketplace_edit::remove_user_marketplace_config;
pub use mcp_edit::ConfigEditsBuilder;
pub use mcp_edit::load_global_mcp_servers;
pub use mcp_requirements::McpServerCommandMatcher;
pub use mcp_requirements::McpServerIdentity;
pub use mcp_requirements::McpServerRequirement;
pub use mcp_requirements::McpServerValueMatcher;
pub use mcp_types::AppToolApproval;
pub use mcp_types::DEFAULT_MCP_SERVER_ENVIRONMENT_ID;
pub use mcp_types::McpServerAuth;
+163
View File
@@ -0,0 +1,163 @@
use crate::mcp_types::McpServerConfig;
use crate::mcp_types::McpServerTransportConfig;
use regex_lite::Regex;
use serde::Deserialize;
#[derive(Deserialize, Debug, Clone, PartialEq, Eq)]
#[serde(untagged)]
pub enum McpServerIdentity {
Command { command: String },
Url { url: String },
}
/// String matching operations available to managed MCP server matchers.
#[derive(Deserialize, Debug, Clone, PartialEq, Eq)]
#[serde(tag = "match", rename_all = "snake_case", deny_unknown_fields)]
pub enum McpServerValueMatcher {
Exact { value: String },
Prefix { value: String },
Regex { expression: String },
}
impl McpServerValueMatcher {
fn compile_full_regex(expression: &str) -> Result<Regex, String> {
Regex::new(&format!(r"\A(?:{expression})\z")).map_err(|err| {
format!("regex `{expression}` cannot be used for full-value matching: {err}")
})
}
fn validate(&self) -> Result<(), String> {
let Self::Regex { expression } = self else {
return Ok(());
};
Regex::new(expression).map_err(|err| format!("invalid regex `{expression}`: {err}"))?;
Self::compile_full_regex(expression).map(|_| ())
}
fn matches(&self, candidate: &str) -> bool {
match self {
Self::Exact { value } => candidate == value,
Self::Prefix { value } => candidate.starts_with(value),
Self::Regex { expression } => Self::compile_full_regex(expression)
.ok()
.is_some_and(|regex| regex.is_match(candidate)),
}
}
}
#[derive(Deserialize, Debug, Clone, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct McpServerCommandMatcher {
pub executable: String,
pub args: Vec<McpServerValueMatcher>,
}
#[derive(Deserialize, Debug, Clone, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
struct RawMcpServerCommandIdentity {
command: McpServerCommandMatcher,
}
#[derive(Deserialize, Debug, Clone, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
struct RawMcpServerUrlIdentity {
url: McpServerValueMatcher,
}
/// A requirement for one named MCP server.
///
/// The `Identity` variant preserves the released exact-match contract. The
/// command and URL variants are the normalized matcher-based forms accepted
/// under the `identity` key.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum McpServerRequirement {
Identity { identity: McpServerIdentity },
Command(McpServerCommandMatcher),
Url(McpServerValueMatcher),
}
#[derive(Deserialize)]
struct RawMcpServerRequirement {
identity: RawMcpServerIdentity,
}
#[derive(Deserialize)]
#[serde(untagged)]
enum RawMcpServerIdentity {
Exact(McpServerIdentity),
Command(RawMcpServerCommandIdentity),
Url(RawMcpServerUrlIdentity),
}
impl<'de> Deserialize<'de> for McpServerRequirement {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let RawMcpServerRequirement { identity } =
RawMcpServerRequirement::deserialize(deserializer)?;
match identity {
RawMcpServerIdentity::Exact(identity) => Ok(Self::Identity { identity }),
RawMcpServerIdentity::Command(matcher) => Ok(Self::Command(matcher.command)),
RawMcpServerIdentity::Url(matcher) => Ok(Self::Url(matcher.url)),
}
}
}
impl McpServerRequirement {
pub(crate) fn validate(&self) -> Result<(), String> {
match self {
Self::Identity { .. } => Ok(()),
Self::Command(matcher) => {
for (index, arg) in matcher.args.iter().enumerate() {
arg.validate().map_err(|err| {
format!("invalid argument matcher at index {index}: {err}")
})?;
}
Ok(())
}
Self::Url(matcher) => matcher.validate(),
}
}
pub fn matches(&self, server: &McpServerConfig) -> bool {
match (self, &server.transport) {
(
Self::Identity {
identity:
McpServerIdentity::Command {
command: want_command,
},
},
McpServerTransportConfig::Stdio {
command: got_command,
..
},
) => got_command == want_command,
(
Self::Identity {
identity: McpServerIdentity::Url { url: want_url },
},
McpServerTransportConfig::StreamableHttp { url: got_url, .. },
) => got_url == want_url,
(Self::Command(matcher), McpServerTransportConfig::Stdio { command, args, .. }) => {
matcher.executable == *command
&& matcher.args.len() == args.len()
&& matcher
.args
.iter()
.zip(args)
.all(|(matcher, arg)| matcher.matches(arg))
}
(Self::Url(matcher), McpServerTransportConfig::StreamableHttp { url, .. }) => {
matcher.matches(url)
}
_ => false,
}
}
}
#[cfg(test)]
#[path = "mcp_requirements_tests.rs"]
mod tests;
@@ -0,0 +1,218 @@
use super::*;
use crate::mcp_types::McpServerConfig;
use pretty_assertions::assert_eq;
use std::collections::HashMap;
fn stdio_server(command: &str, args: &[&str]) -> McpServerConfig {
McpServerConfig {
auth: Default::default(),
transport: McpServerTransportConfig::Stdio {
command: command.to_string(),
args: args.iter().map(ToString::to_string).collect(),
env: None,
env_vars: Vec::new(),
cwd: None,
},
environment_id: crate::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(),
enabled: true,
required: false,
supports_parallel_tool_calls: false,
disabled_reason: None,
startup_timeout_sec: None,
tool_timeout_sec: None,
default_tools_approval_mode: None,
enabled_tools: None,
disabled_tools: None,
scopes: None,
oauth: None,
oauth_resource: None,
tools: HashMap::new(),
}
}
#[test]
fn command_matcher_matches_exact_positional_arguments() {
let requirement = McpServerRequirement::Command(McpServerCommandMatcher {
executable: "company-cli".to_string(),
args: vec![
McpServerValueMatcher::Exact {
value: "mcp".to_string(),
},
McpServerValueMatcher::Regex {
expression: r"https://[a-z]+\.example\.com".to_string(),
},
],
});
assert!(requirement.matches(&stdio_server(
"company-cli",
&["mcp", "https://pricing.example.com"]
)));
assert!(!requirement.matches(&stdio_server(
"company-cli",
&["https://pricing.example.com", "mcp"]
)));
assert!(!requirement.matches(&stdio_server(
"company-cli",
&["mcp", "https://pricing.example.com", "--verbose"]
)));
assert!(!requirement.matches(&stdio_server(
"/usr/local/bin/company-cli",
&["mcp", "https://pricing.example.com"]
)));
}
#[test]
fn regex_matcher_requires_a_full_value_match() {
let matcher = McpServerValueMatcher::Regex {
expression: "mcp".to_string(),
};
assert!(matcher.matches("mcp"));
assert!(!matcher.matches("mcp-proxy"));
assert!(!matcher.matches("prefix-mcp"));
}
#[test]
fn regex_matcher_allows_a_later_alternative_to_match_the_full_value() {
let matcher = McpServerValueMatcher::Regex {
expression: r"https://api\.example\.com|https://api\.example\.com/mcp".to_string(),
};
assert!(matcher.matches("https://api.example.com/mcp"));
}
#[test]
fn regex_matcher_validation_rejects_expression_that_cannot_be_wrapped() {
let matcher = McpServerValueMatcher::Regex {
expression: "(?x)mcp # trailing comment".to_string(),
};
let err = matcher
.validate()
.expect_err("expression should not be valid for full-value matching");
assert!(
err.contains("cannot be used for full-value matching"),
"{err}"
);
}
#[test]
fn legacy_command_identity_keeps_ignoring_arguments() {
let requirement: McpServerRequirement = toml::from_str(
r#"
[identity]
command = "company-cli"
"#,
)
.expect("legacy command identity");
assert!(requirement.matches(&stdio_server(
"company-cli",
&["any", "arguments", "remain", "allowed"]
)));
assert!(!requirement.matches(&stdio_server("different-cli", &[])));
}
#[test]
fn requirement_deserializes_command_and_url_matcher_shapes() {
let command: McpServerRequirement = toml::from_str(
r#"
[identity]
command = { executable = "company-cli", args = [
{ match = "exact", value = "mcp" },
{ match = "regex", expression = '^https://[a-z]+\.example\.com$' },
] }
"#,
)
.expect("command matcher");
let url: McpServerRequirement = toml::from_str(
r#"
[identity]
url = { match = "prefix", value = "https://mcp.example.com/" }
"#,
)
.expect("URL matcher");
assert_eq!(
command,
McpServerRequirement::Command(McpServerCommandMatcher {
executable: "company-cli".to_string(),
args: vec![
McpServerValueMatcher::Exact {
value: "mcp".to_string(),
},
McpServerValueMatcher::Regex {
expression: r"^https://[a-z]+\.example\.com$".to_string(),
},
],
})
);
assert_eq!(
url,
McpServerRequirement::Url(McpServerValueMatcher::Prefix {
value: "https://mcp.example.com/".to_string(),
})
);
}
#[test]
fn requirement_rejects_matchers_outside_identity() {
for contents in [
r#"
command = "company-cli"
"#,
r#"
command = { executable = "company-cli", args = [] }
"#,
r#"
url = { match = "prefix", value = "https://mcp.example.com/" }
"#,
] {
let err = toml::from_str::<McpServerRequirement>(contents)
.expect_err("MCP server requirements should use the identity key");
assert!(
err.to_string().contains("missing field `identity`"),
"{err}"
);
}
}
#[test]
fn matcher_identity_rejects_unknown_fields() {
for contents in [
r#"
[identity]
unknown = "value"
command = { executable = "company-cli", args = [] }
"#,
r#"
[identity]
command = { executable = "company-cli", args = [], unknown = "value" }
"#,
] {
toml::from_str::<McpServerRequirement>(contents)
.expect_err("matcher identities should reject unknown fields");
}
}
#[test]
fn identity_requirement_keeps_ignoring_unrelated_sibling_fields() {
let requirement: McpServerRequirement = toml::from_str(
r#"
unrelated = "ignored"
[identity]
command = "company-cli"
"#,
)
.expect("legacy identity with unrelated sibling field");
assert_eq!(
requirement,
McpServerRequirement::Identity {
identity: McpServerIdentity::Command {
command: "company-cli".to_string(),
},
}
);
}
+141 -45
View File
@@ -10,6 +10,7 @@ use codex_utils_path_uri::PathUri;
use codex_utils_plugins::DISCOVERABLE_PLUGIN_MANIFEST_PATHS;
use codex_utils_plugins::plugin_namespace_for_root_uri;
use codex_utils_plugins::plugin_namespace_for_skill_uri;
use futures::StreamExt;
use futures::future::join_all;
use crate::model::SkillDependencies;
@@ -22,7 +23,6 @@ use super::ParsedSkillFrontmatter;
use super::SKILLS_FILENAME;
use super::SKILLS_METADATA_DIR;
use super::SKILLS_METADATA_FILENAME;
use super::SkillFileDiscovery;
use super::SkillMetadataFile;
use super::parse_skill_frontmatter_metadata_inner;
use super::resolve_dependencies;
@@ -31,6 +31,25 @@ use super::sanitize_single_line;
use super::validate_len;
const MAX_SKILLS_ENTRIES_PER_ROOT: usize = 20_000;
const MAX_CONCURRENT_SKILL_LOADS: usize = 64;
struct EnvironmentSkillDiscovery {
skills: Vec<DiscoveredEnvironmentSkill>,
plugin_roots: HashSet<PathUri>,
namespace_roots: HashSet<PathUri>,
warnings: Vec<String>,
}
struct DiscoveredEnvironmentSkill {
path: PathUri,
metadata: SkillMetadataDiscovery,
}
enum SkillMetadataDiscovery {
Present(PathUri),
Absent,
Probe(PathUri),
}
/// URI-native metadata for one skill owned by an execution environment.
#[derive(Clone, Debug, PartialEq, Eq)]
@@ -65,28 +84,44 @@ impl EnvironmentSkillMetadata {
async fn parse(
file_system: &dyn ExecutorFileSystem,
path: &PathUri,
skill: &DiscoveredEnvironmentSkill,
plugin_namespace: Option<&str>,
) -> Result<Self, String> {
let contents = file_system
.read_file_text(path, /*sandbox*/ None)
.await
.map_err(|err| format!("failed to read file: {err}"))?;
let (contents, discovered_metadata) = match &skill.metadata {
SkillMetadataDiscovery::Present(metadata_path) => {
let (contents, metadata) = tokio::join!(
read_skill_contents(file_system, &skill.path),
read_skill_metadata(file_system, metadata_path),
);
(contents?, metadata)
}
SkillMetadataDiscovery::Absent | SkillMetadataDiscovery::Probe(_) => (
read_skill_contents(file_system, &skill.path).await?,
(None, None),
),
};
let ParsedSkillFrontmatter {
name: base_name,
description,
short_description,
} = parse_skill_frontmatter_metadata_inner(&contents, || default_skill_name(path))
} = parse_skill_frontmatter_metadata_inner(&contents, || default_skill_name(&skill.path))
.map_err(|err| err.to_string())?;
let name = plugin_namespace
.map(|namespace| format!("{namespace}:{base_name}"))
.unwrap_or(base_name);
validate_len(&name, MAX_QUALIFIED_NAME_LEN, "qualified name")
.map_err(|err| err.to_string())?;
let (dependencies, policy) = load_skill_metadata(file_system, path).await;
let (dependencies, policy) = match &skill.metadata {
SkillMetadataDiscovery::Present(_) | SkillMetadataDiscovery::Absent => {
discovered_metadata
}
SkillMetadataDiscovery::Probe(metadata_path) => {
probe_skill_metadata(file_system, metadata_path).await
}
};
Ok(Self {
path_to_skills_md: path.clone(),
path_to_skills_md: skill.path.clone(),
name,
description,
short_description,
@@ -123,6 +158,7 @@ pub async fn load_environment_skills_from_root(
.await
{
Ok(walk) => {
let inventory_complete = !walk.truncated && walk.errors.is_empty();
let mut warnings = walk
.errors
.into_iter()
@@ -139,10 +175,13 @@ pub async fn load_environment_skills_from_root(
));
}
let mut skill_files = Vec::new();
let mut file_paths = HashSet::new();
let mut directory_paths = HashSet::new();
let mut plugin_roots = HashSet::new();
for entry in walk.entries {
match entry.kind {
WalkEntryKind::Directory => {
directory_paths.insert(entry.path.clone());
if DISCOVERABLE_PLUGIN_MANIFEST_PATHS
.iter()
.any(|path| path.split('/').next() == entry.path.basename().as_deref())
@@ -152,40 +191,53 @@ pub async fn load_environment_skills_from_root(
}
}
WalkEntryKind::File => {
file_paths.insert(entry.path.clone());
if entry.path.basename().as_deref() == Some(SKILLS_FILENAME) {
skill_files.push(entry.path);
}
}
}
}
SkillFileDiscovery {
skill_files,
let skills = skill_files
.into_iter()
.map(|path| DiscoveredEnvironmentSkill {
metadata: discover_skill_metadata(
&path,
&file_paths,
&directory_paths,
inventory_complete,
),
path,
})
.collect();
EnvironmentSkillDiscovery {
skills,
plugin_roots,
namespace_roots: HashSet::from([root.clone()]),
warnings,
}
}
Err(error) if error.kind() == io::ErrorKind::NotFound => SkillFileDiscovery {
skill_files: Vec::new(),
Err(error) if error.kind() == io::ErrorKind::NotFound => EnvironmentSkillDiscovery {
skills: Vec::new(),
plugin_roots: HashSet::new(),
namespace_roots: HashSet::new(),
warnings: Vec::new(),
},
Err(error) => SkillFileDiscovery {
skill_files: Vec::new(),
Err(error) => EnvironmentSkillDiscovery {
skills: Vec::new(),
plugin_roots: HashSet::new(),
namespace_roots: HashSet::new(),
warnings: vec![format!("failed to walk skills root {root}: {error:#}")],
},
};
outcome.warnings.extend(discovery.warnings);
if discovery.skill_files.is_empty() {
if discovery.skills.is_empty() {
return outcome;
}
let mut skill_ancestors = HashSet::new();
for skill_path in &discovery.skill_files {
let mut ancestor = skill_path.parent();
for skill in &discovery.skills {
let mut ancestor = skill.path.parent();
while let Some(path) = ancestor {
skill_ancestors.insert(path.clone());
ancestor = path.parent();
@@ -220,24 +272,33 @@ pub async fn load_environment_skills_from_root(
.filter_map(|(plugin_root, namespace)| namespace.map(|namespace| (plugin_root, namespace)))
.collect::<HashMap<_, _>>();
for path in discovery.skill_files {
let mut ancestor = path.parent();
let plugin_namespace = loop {
let Some(current) = ancestor else {
break None;
// Remote executors can multiplex these independent per-skill reads, so polling a bounded
// number together allows the I/O for each skill and its metadata to happen concurrently.
let skill_results = futures::stream::iter(discovery.skills)
.map(|skill| {
let mut ancestor = skill.path.parent();
let plugin_namespace = loop {
let Some(current) = ancestor else {
break None;
};
if let Some(namespace) = plugin_namespaces.get(&current) {
break Some(namespace.as_str());
}
ancestor = current.parent();
};
if let Some(namespace) = plugin_namespaces.get(&current) {
break Some(namespace.as_str());
let path = skill.path.clone();
async move {
let result =
EnvironmentSkillMetadata::parse(file_system, &skill, plugin_namespace).await;
(path, result)
}
ancestor = current.parent();
};
match EnvironmentSkillMetadata::parse(
file_system,
&path,
/*plugin_namespace*/ plugin_namespace,
)
.await
{
})
.buffered(MAX_CONCURRENT_SKILL_LOADS)
.collect::<Vec<_>>()
.await;
for (path, result) in skill_results {
match result {
Ok(skill) if skill.matches_product_restriction(restriction_product) => {
outcome.skills.push(skill);
}
@@ -257,20 +318,48 @@ pub async fn load_environment_skills_from_root(
outcome
}
async fn load_skill_metadata(
fn discover_skill_metadata(
skill_path: &PathUri,
file_paths: &HashSet<PathUri>,
directory_paths: &HashSet<PathUri>,
inventory_complete: bool,
) -> SkillMetadataDiscovery {
let Some(skill_dir) = skill_path.parent() else {
return SkillMetadataDiscovery::Absent;
};
let Ok(metadata_dir) = skill_dir.join(SKILLS_METADATA_DIR) else {
return SkillMetadataDiscovery::Absent;
};
let Ok(metadata_path) = metadata_dir.join(SKILLS_METADATA_FILENAME) else {
return SkillMetadataDiscovery::Absent;
};
if file_paths.contains(&metadata_path) {
SkillMetadataDiscovery::Present(metadata_path)
} else if inventory_complete && !directory_paths.contains(&metadata_dir) {
SkillMetadataDiscovery::Absent
} else {
// The walk can omit entries after an error or traversal limit. It also omits file
// symlinks, so keep the existing probe when the metadata directory itself was observed.
SkillMetadataDiscovery::Probe(metadata_path)
}
}
async fn read_skill_contents(
file_system: &dyn ExecutorFileSystem,
skill_path: &PathUri,
) -> Result<String, String> {
file_system
.read_file_text(skill_path, /*sandbox*/ None)
.await
.map_err(|err| format!("failed to read file: {err}"))
}
async fn probe_skill_metadata(
file_system: &dyn ExecutorFileSystem,
metadata_path: &PathUri,
) -> (Option<SkillDependencies>, Option<SkillPolicy>) {
let Some(skill_dir) = skill_path.parent() else {
return (None, None);
};
let Ok(metadata_path) =
skill_dir.join(&format!("{SKILLS_METADATA_DIR}/{SKILLS_METADATA_FILENAME}"))
else {
return (None, None);
};
match file_system
.get_metadata(&metadata_path, /*sandbox*/ None)
.get_metadata(metadata_path, /*sandbox*/ None)
.await
{
Ok(metadata) if metadata.is_file => {}
@@ -281,8 +370,15 @@ async fn load_skill_metadata(
return (None, None);
}
}
read_skill_metadata(file_system, metadata_path).await
}
async fn read_skill_metadata(
file_system: &dyn ExecutorFileSystem,
metadata_path: &PathUri,
) -> (Option<SkillDependencies>, Option<SkillPolicy>) {
let contents = match file_system
.read_file_text(&metadata_path, /*sandbox*/ None)
.read_file_text(metadata_path, /*sandbox*/ None)
.await
{
Ok(contents) => contents,
@@ -1,5 +1,7 @@
use std::fs;
use std::sync::Mutex;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use codex_core_skills::loader::EnvironmentSkillMetadata;
use codex_core_skills::loader::load_environment_skills_from_root;
@@ -13,6 +15,8 @@ use codex_exec_server::FileSystemSandboxContext;
use codex_exec_server::LOCAL_FS;
use codex_exec_server::ReadDirectoryEntry;
use codex_exec_server::RemoveOptions;
use codex_exec_server::WalkOptions;
use codex_exec_server::WalkOutcome;
use codex_utils_path_uri::PathUri;
use pretty_assertions::assert_eq;
use tempfile::tempdir;
@@ -20,6 +24,15 @@ use tempfile::tempdir;
struct RecordingFileSystem<'a> {
inner: &'a dyn ExecutorFileSystem,
read_files: Mutex<Vec<PathUri>>,
metadata_files: Mutex<Vec<PathUri>>,
walks: AtomicUsize,
}
#[derive(Debug, PartialEq, Eq)]
struct FileSystemCalls {
walks: usize,
read_files: Vec<PathUri>,
metadata_files: Vec<PathUri>,
}
impl<'a> RecordingFileSystem<'a> {
@@ -27,14 +40,29 @@ impl<'a> RecordingFileSystem<'a> {
Self {
inner,
read_files: Mutex::new(Vec::new()),
metadata_files: Mutex::new(Vec::new()),
walks: AtomicUsize::new(0),
}
}
fn read_files(&self) -> Vec<PathUri> {
self.read_files
fn calls(&self) -> FileSystemCalls {
let mut read_files = self
.read_files
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
.clone();
read_files.sort_by_key(ToString::to_string);
let mut metadata_files = self
.metadata_files
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone();
metadata_files.sort_by_key(ToString::to_string);
FileSystemCalls {
walks: self.walks.load(Ordering::Relaxed),
read_files,
metadata_files,
}
}
}
@@ -90,6 +118,10 @@ impl ExecutorFileSystem for RecordingFileSystem<'_> {
path: &'a PathUri,
sandbox: Option<&'a FileSystemSandboxContext>,
) -> ExecutorFileSystemFuture<'a, FileMetadata> {
self.metadata_files
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push(path.clone());
self.inner.get_metadata(path, sandbox)
}
@@ -101,6 +133,16 @@ impl ExecutorFileSystem for RecordingFileSystem<'_> {
self.inner.read_directory(path, sandbox)
}
fn walk<'a>(
&'a self,
path: &'a PathUri,
options: WalkOptions,
sandbox: Option<&'a FileSystemSandboxContext>,
) -> ExecutorFileSystemFuture<'a, WalkOutcome> {
self.walks.fetch_add(1, Ordering::Relaxed);
self.inner.walk(path, options, sandbox)
}
fn remove<'a>(
&'a self,
path: &'a PathUri,
@@ -200,7 +242,8 @@ async fn loads_nearest_plugin_namespaces_without_reading_unused_sibling_manifest
);
let mut manifest_reads = file_system
.read_files()
.calls()
.read_files
.into_iter()
.filter(|path| path.basename().as_deref() == Some("plugin.json"))
.collect::<Vec<_>>();
@@ -214,3 +257,73 @@ async fn loads_nearest_plugin_namespaces_without_reading_unused_sibling_manifest
expected_manifest_reads.sort_by_key(ToString::to_string);
assert_eq!(manifest_reads, expected_manifest_reads);
}
#[tokio::test]
async fn reuses_walk_inventory_for_missing_skill_metadata() {
const SKILL_COUNT: usize = 66;
let root = tempdir().expect("tempdir");
let manifest_path = root.path().join(".codex-plugin/plugin.json");
fs::create_dir_all(manifest_path.parent().expect("manifest parent")).expect("manifest dir");
fs::write(&manifest_path, r#"{"name":"inventory"}"#).expect("manifest");
let mut skill_paths = Vec::new();
for index in 0..SKILL_COUNT {
let name = format!("skill-{index}");
let skill_path = root.path().join(&name).join("SKILL.md");
fs::create_dir_all(skill_path.parent().expect("skill parent")).expect("skill dir");
fs::write(
&skill_path,
format!("---\nname: {name}\ndescription: {name} skill.\n---\n"),
)
.expect("skill");
skill_paths.push(skill_path);
}
let file_system = RecordingFileSystem::new(LOCAL_FS.as_ref());
let root_uri = PathUri::from_host_native_path(root.path()).expect("root URI");
let outcome = load_environment_skills_from_root(
&file_system,
&root_uri,
/*restriction_product*/ None,
)
.await;
let mut expected_skills = skill_paths
.iter()
.enumerate()
.map(|(index, skill_path)| EnvironmentSkillMetadata {
path_to_skills_md: PathUri::from_host_native_path(skill_path).unwrap(),
name: format!("inventory:skill-{index}"),
description: format!("skill-{index} skill."),
short_description: None,
dependencies: None,
policy: None,
})
.collect::<Vec<_>>();
expected_skills.sort_by(|left, right| {
left.name.cmp(&right.name).then_with(|| {
left.path_to_skills_md
.to_string()
.cmp(&right.path_to_skills_md.to_string())
})
});
assert_eq!(outcome.skills, expected_skills);
assert_eq!(outcome.warnings, Vec::<String>::new());
let mut expected_read_files = skill_paths
.iter()
.map(|path| PathUri::from_host_native_path(path).unwrap())
.collect::<Vec<_>>();
let manifest_uri = PathUri::from_host_native_path(manifest_path).unwrap();
expected_read_files.push(manifest_uri.clone());
expected_read_files.sort_by_key(ToString::to_string);
assert_eq!(
file_system.calls(),
FileSystemCalls {
walks: 1,
read_files: expected_read_files,
metadata_files: vec![manifest_uri],
}
);
}
+23 -1
View File
@@ -810,12 +810,15 @@
"clock_source": {
"$ref": "#/definitions/CurrentTimeSource"
},
"delivery_mode": {
"$ref": "#/definitions/CurrentTimeReminderDeliveryMode"
},
"enabled": {
"type": "boolean"
},
"reminder_interval_seconds": {
"format": "uint64",
"minimum": 1.0,
"minimum": 0.0,
"type": "integer"
},
"sleep_tool": {
@@ -825,6 +828,25 @@
},
"type": "object"
},
"CurrentTimeReminderDeliveryMode": {
"description": "Which inference boundaries may receive current-time reminders.",
"oneOf": [
{
"description": "Allow a reminder before any inference request once the interval is due.",
"enum": [
"any_inference"
],
"type": "string"
},
{
"description": "Allow reminders after user input or tool output; new context windows still force one.",
"enum": [
"after_user_or_tool_output"
],
"type": "string"
}
]
},
"CurrentTimeSource": {
"enum": [
"system",
+14
View File
@@ -1,5 +1,6 @@
use super::residency::is_v2_resident_session_source;
use super::*;
use codex_extension_api::ExtensionDataInit;
const AGENT_NAMES: &str = include_str!("../agent_names.txt");
@@ -433,6 +434,16 @@ impl AgentControl {
))
})?;
let selected_capability_roots = parent_history
.items
.iter()
.find_map(|item| {
let RolloutItem::SessionMeta(meta_line) = item else {
return None;
};
Some(meta_line.meta.selected_capability_roots.clone())
})
.unwrap_or_default();
let mut forked_rollout_items = parent_history.items;
if let SpawnAgentForkMode::LastNTurns(last_n_turns) = fork_mode {
forked_rollout_items =
@@ -504,6 +515,8 @@ impl AgentControl {
{
forked_rollout_items.push(RolloutItem::ResponseItem(subagent_usage_hint_message));
}
let mut thread_extension_init = ExtensionDataInit::new();
thread_extension_init.insert(selected_capability_roots);
state
.fork_thread_with_source(
@@ -517,6 +530,7 @@ impl AgentControl {
inherited_environments,
inherited_exec_policy,
options.environments.clone(),
thread_extension_init,
)
.await
}
+40 -1
View File
@@ -17,6 +17,8 @@ use codex_features::Feature;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_protocol::AgentPath;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_protocol::config_types::ModeKind;
use codex_protocol::models::ContentItem;
use codex_protocol::models::MessagePhase;
@@ -38,6 +40,7 @@ use codex_thread_store::InMemoryThreadStore;
use codex_thread_store::LocalThreadStore;
use codex_thread_store::LocalThreadStoreConfig;
use codex_thread_store::ThreadStore;
use codex_utils_path_uri::PathUri;
use pretty_assertions::assert_eq;
use tempfile::TempDir;
use tokio::time::Duration;
@@ -1446,7 +1449,34 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() {
#[tokio::test]
async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_limit() {
let harness = AgentControlHarness::new().await;
let (parent_thread_id, parent_thread) = harness.start_thread().await;
let selected_capability_roots = vec![SelectedCapabilityRoot {
id: "demo@1".to_string(),
location: CapabilityRootLocation::Environment {
environment_id: "build".to_string(),
path: PathUri::parse("file:///plugins/demo").expect("plugin root URI"),
},
}];
let mut thread_extension_init = ExtensionDataInit::new();
thread_extension_init.insert(selected_capability_roots.clone());
let parent = harness
.manager
.start_thread_with_options(StartThreadOptions {
config: harness.config.clone(),
allow_provider_model_fallback: false,
initial_history: InitialHistory::New,
session_source: None,
thread_source: None,
dynamic_tools: Vec::new(),
metrics_service_name: None,
parent_trace: None,
environments: Vec::new(),
thread_extension_init,
supports_openai_form_elicitation: false,
})
.await
.expect("start parent thread");
let parent_thread_id = parent.thread_id;
let parent_thread = parent.thread;
let startup_turn_context = parent_thread.codex.session.new_default_turn().await;
parent_thread
.codex
@@ -1525,6 +1555,14 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li
!history_contains_text(history.raw_items(), "parent startup developer context"),
"bounded fork should drop parent startup context even when fewer turns exist than requested"
);
assert_eq!(
&child_thread
.codex
.session
.services
.selected_capability_roots,
&selected_capability_roots
);
assert!(
child_thread
.codex
@@ -2213,6 +2251,7 @@ async fn spawn_thread_subagents_persist_parent_originator_across_new_and_truncat
.manager
.start_thread_with_options(StartThreadOptions {
config: harness.config.clone(),
allow_provider_model_fallback: false,
initial_history: InitialHistory::New,
session_source: None,
thread_source: None,
+45 -18
View File
@@ -40,6 +40,7 @@ use crate::guardian::spawn_approval_request_review;
use crate::mcp_tool_call::MCP_TOOL_APPROVAL_ACCEPT;
use crate::mcp_tool_call::MCP_TOOL_APPROVAL_ACCEPT_FOR_SESSION;
use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC;
use crate::mcp_tool_call::McpToolApprovalMetadata;
use crate::mcp_tool_call::build_guardian_mcp_tool_review_request;
use crate::mcp_tool_call::is_mcp_tool_approval_question_id;
use crate::mcp_tool_call::lookup_mcp_tool_metadata;
@@ -60,6 +61,12 @@ use codex_protocol::protocol::MultiAgentVersion;
#[cfg(test)]
use crate::session::completed_session_loop_termination;
#[derive(Clone)]
struct PendingMcpInvocation {
invocation: McpInvocation,
metadata: Option<McpToolApprovalMetadata>,
}
/// Start an interactive sub-Codex thread and return IO channels.
///
/// The returned `events_rx` yields non-approval events emitted by the sub-agent.
@@ -86,6 +93,7 @@ pub(crate) async fn run_codex_thread_interactive(
};
let CodexSpawnOk { codex, .. } = Box::pin(Codex::spawn(CodexSpawnArgs {
config,
allow_provider_model_fallback: false,
user_instructions,
installation_id: parent_session.installation_id.clone(),
auth_manager,
@@ -148,10 +156,10 @@ pub(crate) async fn run_codex_thread_interactive(
let parent_session_clone = Arc::clone(&parent_session);
let parent_ctx_clone = Arc::clone(&parent_ctx);
let codex_for_events = Arc::clone(&codex);
// Cache delegated MCP invocations so guardian can recover the full tool call
// context when the later legacy RequestUserInput approval event only carries
// a call_id plus approval question metadata.
let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::<String, McpInvocation>::new()));
// Cache the child call's MCP metadata at begin time. The later legacy
// RequestUserInput approval event only carries a call_id and question metadata.
let pending_mcp_invocations =
Arc::new(Mutex::new(HashMap::<String, PendingMcpInvocation>::new()));
tokio::spawn(async move {
forward_events(
codex_for_events,
@@ -269,7 +277,7 @@ async fn forward_events(
tx_sub: Sender<Event>,
parent_session: Arc<Session>,
parent_ctx: Arc<TurnContext>,
pending_mcp_invocations: Arc<Mutex<HashMap<String, McpInvocation>>>,
pending_mcp_invocations: Arc<Mutex<HashMap<String, PendingMcpInvocation>>>,
cancel_token: CancellationToken,
) {
let cancelled = cancel_token.cancelled();
@@ -356,10 +364,34 @@ async fn forward_events(
id,
msg: EventMsg::McpToolCallBegin(event),
} => {
// Runtime refreshes are published before a request step is captured, so
// the child runtime at call begin is the one executing this invocation.
// Cache its metadata now; the later approval event has only a call ID.
let metadata = if let Some(turn_context) =
codex.session.turn_context_for_sub_id(&id).await
{
let mcp = codex.session.services.latest_mcp_runtime();
lookup_mcp_tool_metadata(
codex.session.as_ref(),
turn_context.as_ref(),
mcp.manager(),
&event.invocation.server,
&event.invocation.tool,
)
.await
} else {
None
};
pending_mcp_invocations
.lock()
.await
.insert(event.call_id.clone(), event.invocation.clone());
.insert(
event.call_id.clone(),
PendingMcpInvocation {
invocation: event.invocation.clone(),
metadata,
},
);
if !forward_event_or_shutdown(
&codex,
&tx_sub,
@@ -647,7 +679,7 @@ async fn handle_request_user_input(
id: String,
parent_session: &Arc<Session>,
parent_ctx: &Arc<TurnContext>,
pending_mcp_invocations: &Arc<Mutex<HashMap<String, McpInvocation>>>,
pending_mcp_invocations: &Arc<Mutex<HashMap<String, PendingMcpInvocation>>>,
event: RequestUserInputEvent,
cancel_token: &CancellationToken,
) {
@@ -685,12 +717,12 @@ async fn handle_request_user_input(
/// programmatically after running the guardian review.
///
/// The RequestUserInput event only carries `call_id` plus approval question
/// metadata, so this helper joins it back to the cached `McpToolCallBegin`
/// invocation in order to rebuild the full guardian review request.
/// metadata, so this helper joins it back to the child runtime metadata cached at
/// `McpToolCallBegin` in order to rebuild the full guardian review request.
async fn maybe_auto_review_mcp_request_user_input(
parent_session: &Arc<Session>,
parent_ctx: &Arc<TurnContext>,
pending_mcp_invocations: &Arc<Mutex<HashMap<String, McpInvocation>>>,
pending_mcp_invocations: &Arc<Mutex<HashMap<String, PendingMcpInvocation>>>,
event: &RequestUserInputEvent,
cancel_token: &CancellationToken,
) -> Option<RequestUserInputResponse> {
@@ -701,18 +733,13 @@ async fn maybe_auto_review_mcp_request_user_input(
.questions
.iter()
.find(|question| is_mcp_tool_approval_question_id(&question.id))?;
let invocation = pending_mcp_invocations
let pending = pending_mcp_invocations
.lock()
.await
.get(&event.call_id)
.cloned()?;
let metadata = lookup_mcp_tool_metadata(
parent_session.as_ref(),
parent_ctx.as_ref(),
&invocation.server,
&invocation.tool,
)
.await;
let invocation = pending.invocation;
let metadata = pending.metadata;
let approvals_reviewer =
mcp_approvals_reviewer(parent_ctx, &invocation.server, metadata.as_ref());
if !routes_approval_to_guardian_with_reviewer(parent_ctx, approvals_reviewer) {
+14 -8
View File
@@ -411,10 +411,13 @@ async fn delegated_mcp_guardian_abort_returns_synthetic_decline_answer() {
let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([(
"call-1".to_string(),
McpInvocation {
server: "custom_server".to_string(),
tool: "dangerous_tool".to_string(),
arguments: None,
PendingMcpInvocation {
invocation: McpInvocation {
server: "custom_server".to_string(),
tool: "dangerous_tool".to_string(),
arguments: None,
},
metadata: None,
},
)])));
let cancel_token = CancellationToken::new();
@@ -460,10 +463,13 @@ async fn delegated_mcp_user_reviewer_returns_none_without_metadata() {
crate::session::tests::make_session_and_context_with_rx().await;
let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([(
"call-1".to_string(),
McpInvocation {
server: CODEX_APPS_MCP_SERVER_NAME.to_string(),
tool: "dangerous_tool".to_string(),
arguments: None,
PendingMcpInvocation {
invocation: McpInvocation {
server: CODEX_APPS_MCP_SERVER_NAME.to_string(),
tool: "dangerous_tool".to_string(),
arguments: None,
},
metadata: None,
},
)])));
let cancel_token = CancellationToken::new();

Some files were not shown because too many files have changed in this diff Show More