Implemented thread-level atomic elicitation counter for stopwatch pausing

While trying to build out CLI-Tools for the agent to use under skills we
have found that those tools sometimes need to invoke a user elicitation.
These elicitations are handled out of band of the codex app-server but
need to indicate to the exec manager that the command running is
not going to progress on the usual timeout horizon.

Model calls universal exec:
`$ download-credit-card-history --start-date 2026-01-19 --end-date 2026-02-19 > credit_history.jsonl`

download-cred-card-history might hit a hosted/preauthenticated service to fetch data.  That service might decide that the request requires an end user approval the access to the personal data.  It should be able to signal to the running thread that the command in question is blocked on user elicitation.  In that case we want the exec to continue, but the timeout to not expire on the tool call, essentially freezing time until the user approves or rejects the command at which point the tool would signal the app-server to decrement the outstanding elicitation count.  Now timeouts would proceed as normal.

- New v2 RPC methods:
    - thread/increment_elicitation
    - thread/decrement_elicitation
- Protocol updates in:
    - codex-rs/app-server-protocol/src/protocol/common.rs
    - codex-rs/app-server-protocol/src/protocol/v2.rs
- App-server handlers wired in:
    - codex-rs/app-server/src/codex_message_processor.rs

- Counter starts at 0 per thread.
- increment atomically increases the counter.
- decrement atomically decreases the counter; decrement at 0 returns invalid request.
- Transition rules:
    - 0 -> 1: broadcast pause state, pausing all active stopwatches immediately.
    - >0 -> >0: remain paused.
    - 1 -> 0: broadcast unpause state, resuming stopwatches.
- Core thread/session logic:
    - codex-rs/core/src/codex_thread.rs
    - codex-rs/core/src/codex.rs
    - codex-rs/core/src/mcp_connection_manager.rs

- Added centralized stopwatch tracking/controller:
    - codex-rs/exec-server/src/posix/stopwatch_controller.rs
- Hooked pause/unpause broadcast handling + stopwatch registration:
    - codex-rs/exec-server/src/posix/mcp.rs
    - codex-rs/exec-server/src/posix/stopwatch.rs
    - codex-rs/exec-server/src/posix.rs
This commit is contained in:
Channing Conger
2026-02-19 20:24:53 -08:00
parent 366ecaf17a
commit 54d4da7f29
24 changed files with 736 additions and 25 deletions
@@ -2588,6 +2588,17 @@
],
"type": "object"
},
"ThreadDecrementElicitationParams": {
"properties": {
"threadId": {
"type": "string"
}
},
"required": [
"threadId"
],
"type": "object"
},
"ThreadForkParams": {
"description": "There are two ways to fork a thread: 1. By thread_id: load the thread from disk by thread_id and fork it into a new thread. 2. By path: load the thread from disk by path and fork it into a new thread.\n\nIf using path, the thread_id param will be ignored.\n\nPrefer using thread_id whenever possible.",
"properties": {
@@ -2661,6 +2672,17 @@
"ThreadId": {
"type": "string"
},
"ThreadIncrementElicitationParams": {
"properties": {
"threadId": {
"type": "string"
}
},
"required": [
"threadId"
],
"type": "object"
},
"ThreadListParams": {
"properties": {
"archived": {
@@ -3523,6 +3545,54 @@
"title": "Thread/archiveRequest",
"type": "object"
},
{
"properties": {
"id": {
"$ref": "#/definitions/RequestId"
},
"method": {
"enum": [
"thread/increment_elicitation"
],
"title": "Thread/incrementElicitationRequestMethod",
"type": "string"
},
"params": {
"$ref": "#/definitions/ThreadIncrementElicitationParams"
}
},
"required": [
"id",
"method",
"params"
],
"title": "Thread/incrementElicitationRequest",
"type": "object"
},
{
"properties": {
"id": {
"$ref": "#/definitions/RequestId"
},
"method": {
"enum": [
"thread/decrement_elicitation"
],
"title": "Thread/decrementElicitationRequestMethod",
"type": "string"
},
"params": {
"$ref": "#/definitions/ThreadDecrementElicitationParams"
}
},
"required": [
"id",
"method",
"params"
],
"title": "Thread/decrementElicitationRequest",
"type": "object"
},
{
"properties": {
"id": {
@@ -570,6 +570,54 @@
"title": "Thread/archiveRequest",
"type": "object"
},
{
"properties": {
"id": {
"$ref": "#/definitions/RequestId"
},
"method": {
"enum": [
"thread/increment_elicitation"
],
"title": "Thread/incrementElicitationRequestMethod",
"type": "string"
},
"params": {
"$ref": "#/definitions/v2/ThreadIncrementElicitationParams"
}
},
"required": [
"id",
"method",
"params"
],
"title": "Thread/incrementElicitationRequest",
"type": "object"
},
{
"properties": {
"id": {
"$ref": "#/definitions/RequestId"
},
"method": {
"enum": [
"thread/decrement_elicitation"
],
"title": "Thread/decrementElicitationRequestMethod",
"type": "string"
},
"params": {
"$ref": "#/definitions/v2/ThreadDecrementElicitationParams"
}
},
"required": [
"id",
"method",
"params"
],
"title": "Thread/decrementElicitationRequest",
"type": "object"
},
{
"properties": {
"id": {
@@ -15359,6 +15407,38 @@
"title": "ThreadCompactStartResponse",
"type": "object"
},
"ThreadDecrementElicitationParams": {
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"threadId": {
"type": "string"
}
},
"required": [
"threadId"
],
"title": "ThreadDecrementElicitationParams",
"type": "object"
},
"ThreadDecrementElicitationResponse": {
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"count": {
"format": "uint64",
"minimum": 0.0,
"type": "integer"
},
"paused": {
"type": "boolean"
}
},
"required": [
"count",
"paused"
],
"title": "ThreadDecrementElicitationResponse",
"type": "object"
},
"ThreadForkParams": {
"$schema": "http://json-schema.org/draft-07/schema#",
"description": "There are two ways to fork a thread: 1. By thread_id: load the thread from disk by thread_id and fork it into a new thread. 2. By path: load the thread from disk by path and fork it into a new thread.\n\nIf using path, the thread_id param will be ignored.\n\nPrefer using thread_id whenever possible.",
@@ -15477,6 +15557,38 @@
"ThreadId": {
"type": "string"
},
"ThreadIncrementElicitationParams": {
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"threadId": {
"type": "string"
}
},
"required": [
"threadId"
],
"title": "ThreadIncrementElicitationParams",
"type": "object"
},
"ThreadIncrementElicitationResponse": {
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"count": {
"format": "uint64",
"minimum": 0.0,
"type": "integer"
},
"paused": {
"type": "boolean"
}
},
"required": [
"count",
"paused"
],
"title": "ThreadIncrementElicitationResponse",
"type": "object"
},
"ThreadItem": {
"oneOf": [
{
@@ -0,0 +1,13 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"threadId": {
"type": "string"
}
},
"required": [
"threadId"
],
"title": "ThreadDecrementElicitationParams",
"type": "object"
}
@@ -0,0 +1,19 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"count": {
"format": "uint64",
"minimum": 0.0,
"type": "integer"
},
"paused": {
"type": "boolean"
}
},
"required": [
"count",
"paused"
],
"title": "ThreadDecrementElicitationResponse",
"type": "object"
}
@@ -0,0 +1,13 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"threadId": {
"type": "string"
}
},
"required": [
"threadId"
],
"title": "ThreadIncrementElicitationParams",
"type": "object"
}
@@ -0,0 +1,19 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"count": {
"format": "uint64",
"minimum": 0.0,
"type": "integer"
},
"paused": {
"type": "boolean"
}
},
"required": [
"count",
"paused"
],
"title": "ThreadIncrementElicitationResponse",
"type": "object"
}
File diff suppressed because one or more lines are too long
@@ -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 ThreadDecrementElicitationParams = { threadId: string, };
@@ -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 ThreadDecrementElicitationResponse = { count: bigint, paused: boolean, };
@@ -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 ThreadIncrementElicitationParams = { threadId: string, };
@@ -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 ThreadIncrementElicitationResponse = { count: bigint, paused: boolean, };
@@ -152,8 +152,12 @@ export type { ThreadArchiveResponse } from "./ThreadArchiveResponse";
export type { ThreadArchivedNotification } from "./ThreadArchivedNotification";
export type { ThreadCompactStartParams } from "./ThreadCompactStartParams";
export type { ThreadCompactStartResponse } from "./ThreadCompactStartResponse";
export type { ThreadDecrementElicitationParams } from "./ThreadDecrementElicitationParams";
export type { ThreadDecrementElicitationResponse } from "./ThreadDecrementElicitationResponse";
export type { ThreadForkParams } from "./ThreadForkParams";
export type { ThreadForkResponse } from "./ThreadForkResponse";
export type { ThreadIncrementElicitationParams } from "./ThreadIncrementElicitationParams";
export type { ThreadIncrementElicitationResponse } from "./ThreadIncrementElicitationResponse";
export type { ThreadItem } from "./ThreadItem";
export type { ThreadListParams } from "./ThreadListParams";
export type { ThreadListResponse } from "./ThreadListResponse";
@@ -202,6 +202,14 @@ client_request_definitions! {
params: v2::ThreadArchiveParams,
response: v2::ThreadArchiveResponse,
},
ThreadIncrementElicitation => "thread/increment_elicitation" {
params: v2::ThreadIncrementElicitationParams,
response: v2::ThreadIncrementElicitationResponse,
},
ThreadDecrementElicitation => "thread/decrement_elicitation" {
params: v2::ThreadDecrementElicitationParams,
response: v2::ThreadDecrementElicitationResponse,
},
ThreadSetName => "thread/name/set" {
params: v2::ThreadSetNameParams,
response: v2::ThreadSetNameResponse,
@@ -1692,6 +1692,36 @@ pub struct ThreadArchiveParams {
#[ts(export_to = "v2/")]
pub struct ThreadArchiveResponse {}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
pub struct ThreadIncrementElicitationParams {
pub thread_id: String,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
pub struct ThreadIncrementElicitationResponse {
pub count: u64,
pub paused: bool,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
pub struct ThreadDecrementElicitationParams {
pub thread_id: String,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
pub struct ThreadDecrementElicitationResponse {
pub count: u64,
pub paused: bool,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
+14
View File
@@ -122,6 +122,8 @@ Example with notification opt-out:
- `thread/read` — read a stored thread by id without resuming it; optionally include turns via `includeTurns`. The returned `thread` includes `status` (`ThreadStatus`), defaulting to `notLoaded` when the thread is not currently loaded.
- `thread/status/changed` — notification emitted when a loaded thread’s status changes (`threadId` + new `status`).
- `thread/archive` — move a thread’s rollout file into the archived directory; returns `{}` on success and emits `thread/archived`.
- `thread/increment_elicitation` — atomically increment a thread-scoped out-of-band elicitation counter and pause shell-tool timeout stopwatches while the counter is above zero; returns `{ count, paused }`.
- `thread/decrement_elicitation` — atomically decrement that counter and resume shell-tool timeout stopwatches when it reaches zero; returns `{ count, paused }`.
- `thread/name/set` — set or update a thread’s user-facing name; returns `{}` on success. Thread names are not required to be unique; name lookups resolve to the most recently updated thread.
- `thread/unarchive` — move an archived rollout file back into the sessions directory; returns the restored `thread` on success and emits `thread/unarchived`.
- `thread/compact/start` — trigger conversation history compaction for a thread; returns `{}` immediately while progress streams through standard turn/item notifications.
@@ -309,6 +311,18 @@ Use `thread/unarchive` to move an archived rollout back into the sessions direct
{ "method": "thread/unarchived", "params": { "threadId": "thr_b" } }
```
### Example: Pause command timeout stopwatches for out-of-band approval
Use `thread/increment_elicitation` and `thread/decrement_elicitation` to bracket external approval flows (for example, a downstream CLI waiting on user approval). While `count > 0`, active shell-tool timeout stopwatches are paused.
```json
{ "method": "thread/increment_elicitation", "id": 24, "params": { "threadId": "thr_b" } }
{ "id": 24, "result": { "count": 1, "paused": true } }
{ "method": "thread/decrement_elicitation", "id": 25, "params": { "threadId": "thr_b" } }
{ "id": 25, "result": { "count": 0, "paused": false } }
```
### Example: Trigger thread compaction
Use `thread/compact/start` to trigger manual history compaction for a thread. The request returns immediately with `{}`.
@@ -128,8 +128,12 @@ use codex_app_server_protocol::ThreadBackgroundTerminalsCleanParams;
use codex_app_server_protocol::ThreadBackgroundTerminalsCleanResponse;
use codex_app_server_protocol::ThreadCompactStartParams;
use codex_app_server_protocol::ThreadCompactStartResponse;
use codex_app_server_protocol::ThreadDecrementElicitationParams;
use codex_app_server_protocol::ThreadDecrementElicitationResponse;
use codex_app_server_protocol::ThreadForkParams;
use codex_app_server_protocol::ThreadForkResponse;
use codex_app_server_protocol::ThreadIncrementElicitationParams;
use codex_app_server_protocol::ThreadIncrementElicitationResponse;
use codex_app_server_protocol::ThreadItem;
use codex_app_server_protocol::ThreadListParams;
use codex_app_server_protocol::ThreadListResponse;
@@ -544,6 +548,14 @@ impl CodexMessageProcessor {
self.thread_archive(to_connection_request_id(request_id), params)
.await;
}
ClientRequest::ThreadIncrementElicitation { request_id, params } => {
self.thread_increment_elicitation(to_connection_request_id(request_id), params)
.await;
}
ClientRequest::ThreadDecrementElicitation { request_id, params } => {
self.thread_decrement_elicitation(to_connection_request_id(request_id), params)
.await;
}
ClientRequest::ThreadSetName { request_id, params } => {
self.thread_set_name(to_connection_request_id(request_id), params)
.await;
@@ -2169,6 +2181,79 @@ impl CodexMessageProcessor {
}
}
async fn thread_increment_elicitation(
&self,
request_id: ConnectionRequestId,
params: ThreadIncrementElicitationParams,
) {
let (_, thread) = match self.load_thread(&params.thread_id).await {
Ok(value) => value,
Err(error) => {
self.outgoing.send_error(request_id, error).await;
return;
}
};
match thread.increment_out_of_band_elicitation_count().await {
Ok(count) => {
self.outgoing
.send_response(
request_id,
ThreadIncrementElicitationResponse {
count,
paused: count > 0,
},
)
.await;
}
Err(err) => {
self.send_internal_error(
request_id,
format!("failed to increment out-of-band elicitation counter: {err}"),
)
.await;
}
}
}
async fn thread_decrement_elicitation(
&self,
request_id: ConnectionRequestId,
params: ThreadDecrementElicitationParams,
) {
let (_, thread) = match self.load_thread(&params.thread_id).await {
Ok(value) => value,
Err(error) => {
self.outgoing.send_error(request_id, error).await;
return;
}
};
match thread.decrement_out_of_band_elicitation_count().await {
Ok(count) => {
self.outgoing
.send_response(
request_id,
ThreadDecrementElicitationResponse {
count,
paused: count > 0,
},
)
.await;
}
Err(CodexErr::InvalidRequest(message)) => {
self.send_invalid_request_error(request_id, message).await;
}
Err(err) => {
self.send_internal_error(
request_id,
format!("failed to decrement out-of-band elicitation counter: {err}"),
)
.await;
}
}
}
async fn thread_set_name(&self, request_id: ConnectionRequestId, params: ThreadSetNameParams) {
let ThreadSetNameParams { thread_id, name } = params;
let Some(name) = codex_core::util::normalize_thread_name(&name) else {
+14
View File
@@ -2399,6 +2399,20 @@ impl Session {
.await
}
pub(crate) async fn notify_out_of_band_elicitation_pause_state(
&self,
paused: bool,
) -> anyhow::Result<()> {
self.services
.mcp_connection_manager
.read()
.await
.notify_out_of_band_elicitation_state_change(&crate::OutOfBandElicitationState {
paused,
})
.await
}
/// Records input items: always append to conversation history and
/// persist these response items to rollout.
pub(crate) async fn record_conversation_items(
+53
View File
@@ -1,6 +1,7 @@
use crate::agent::AgentStatus;
use crate::codex::Codex;
use crate::codex::SteerInputError;
use crate::error::CodexErr;
use crate::error::Result as CodexResult;
use crate::features::Feature;
use crate::file_watcher::WatchRegistration;
@@ -18,6 +19,7 @@ use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::TokenUsage;
use codex_protocol::user_input::UserInput;
use std::path::PathBuf;
use tokio::sync::Mutex;
use tokio::sync::watch;
use crate::state_db::StateDbHandle;
@@ -37,6 +39,7 @@ pub struct ThreadConfigSnapshot {
pub struct CodexThread {
pub(crate) codex: Codex,
rollout_path: Option<PathBuf>,
out_of_band_elicitation_count: Mutex<u64>,
_watch_registration: WatchRegistration,
}
@@ -51,6 +54,7 @@ impl CodexThread {
Self {
codex,
rollout_path,
out_of_band_elicitation_count: Mutex::new(0),
_watch_registration: watch_registration,
}
}
@@ -130,4 +134,53 @@ impl CodexThread {
pub fn enabled(&self, feature: Feature) -> bool {
self.codex.enabled(feature)
}
pub async fn increment_out_of_band_elicitation_count(&self) -> CodexResult<u64> {
let mut guard = self.out_of_band_elicitation_count.lock().await;
let was_zero = *guard == 0;
*guard = guard.checked_add(1).ok_or_else(|| {
CodexErr::Fatal("out-of-band elicitation count overflowed".to_string())
})?;
if was_zero
&& let Err(err) = self
.codex
.session
.notify_out_of_band_elicitation_pause_state(true)
.await
{
*guard -= 1;
return Err(CodexErr::Fatal(format!(
"failed to pause out-of-band elicitation state: {err:#}"
)));
}
Ok(*guard)
}
pub async fn decrement_out_of_band_elicitation_count(&self) -> CodexResult<u64> {
let mut guard = self.out_of_band_elicitation_count.lock().await;
if *guard == 0 {
return Err(CodexErr::InvalidRequest(
"out-of-band elicitation count is already zero".to_string(),
));
}
*guard -= 1;
let now_zero = *guard == 0;
if now_zero
&& let Err(err) = self
.codex
.session
.notify_out_of_band_elicitation_pause_state(false)
.await
{
*guard += 1;
return Err(CodexErr::Fatal(format!(
"failed to resume out-of-band elicitation state: {err:#}"
)));
}
Ok(*guard)
}
}
+3
View File
@@ -44,8 +44,11 @@ mod mcp_connection_manager;
pub mod models_manager;
mod network_policy_decision;
pub mod network_proxy_loader;
pub use mcp_connection_manager::MCP_OUT_OF_BAND_ELICITATION_STATE_CAPABILITY;
pub use mcp_connection_manager::MCP_OUT_OF_BAND_ELICITATION_STATE_METHOD;
pub use mcp_connection_manager::MCP_SANDBOX_STATE_CAPABILITY;
pub use mcp_connection_manager::MCP_SANDBOX_STATE_METHOD;
pub use mcp_connection_manager::OutOfBandElicitationState;
pub use mcp_connection_manager::SandboxState;
mod mcp_tool_call;
mod memories;
@@ -289,6 +289,7 @@ struct ManagedClient {
tool_filter: ToolFilter,
tool_timeout: Option<Duration>,
server_supports_sandbox_state_capability: bool,
server_supports_out_of_band_elicitation_state_capability: bool,
}
impl ManagedClient {
@@ -307,6 +308,24 @@ impl ManagedClient {
.await?;
Ok(())
}
async fn notify_out_of_band_elicitation_state_change(
&self,
state: &OutOfBandElicitationState,
) -> Result<()> {
if !self.server_supports_out_of_band_elicitation_state_capability {
return Ok(());
}
let _response = self
.client
.send_custom_request(
MCP_OUT_OF_BAND_ELICITATION_STATE_METHOD,
Some(serde_json::to_value(state)?),
)
.await?;
Ok(())
}
}
#[derive(Clone)]
@@ -360,13 +379,27 @@ impl AsyncManagedClient {
let managed = self.client().await?;
managed.notify_sandbox_state_change(sandbox_state).await
}
async fn notify_out_of_band_elicitation_state_change(
&self,
state: &OutOfBandElicitationState,
) -> Result<()> {
let managed = self.client().await?;
managed
.notify_out_of_band_elicitation_state_change(state)
.await
}
}
pub const MCP_SANDBOX_STATE_CAPABILITY: &str = "codex/sandbox-state";
pub const MCP_OUT_OF_BAND_ELICITATION_STATE_CAPABILITY: &str =
"codex/out-of-band-elicitation-state";
/// Custom MCP request to push sandbox state updates.
/// When used, the `params` field of the notification is [`SandboxState`].
pub const MCP_SANDBOX_STATE_METHOD: &str = "codex/sandbox-state/update";
pub const MCP_OUT_OF_BAND_ELICITATION_STATE_METHOD: &str =
"codex/out-of-band-elicitation-state/update";
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -378,6 +411,12 @@ pub struct SandboxState {
pub use_linux_sandbox_bwrap: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OutOfBandElicitationState {
pub paused: bool,
}
/// A thin wrapper around a set of running [`RmcpClient`] instances.
pub(crate) struct McpConnectionManager {
clients: HashMap<String, AsyncManagedClient>,
@@ -885,6 +924,41 @@ impl McpConnectionManager {
Ok(())
}
pub async fn notify_out_of_band_elicitation_state_change(
&self,
state: &OutOfBandElicitationState,
) -> Result<()> {
let mut join_set = JoinSet::new();
for async_managed_client in self.clients.values() {
let state = state.clone();
let async_managed_client = async_managed_client.clone();
join_set.spawn(async move {
async_managed_client
.notify_out_of_band_elicitation_state_change(&state)
.await
});
}
while let Some(join_res) = join_set.join_next().await {
match join_res {
Ok(Ok(())) => {}
Ok(Err(err)) => {
warn!(
"Failed to notify out-of-band elicitation state change to MCP server: {err:#}",
);
}
Err(err) => {
warn!(
"Task panic when notifying out-of-band elicitation state change to MCP server: {err:#}",
);
}
}
}
Ok(())
}
}
async fn emit_update(
@@ -1123,6 +1197,12 @@ async fn start_server_task(
.as_ref()
.and_then(|exp| exp.get(MCP_SANDBOX_STATE_CAPABILITY))
.is_some();
let server_supports_out_of_band_elicitation_state_capability = initialize_result
.capabilities
.experimental
.as_ref()
.and_then(|exp| exp.get(MCP_OUT_OF_BAND_ELICITATION_STATE_CAPABILITY))
.is_some();
let managed = ManagedClient {
client: Arc::clone(&client),
@@ -1130,6 +1210,7 @@ async fn start_server_task(
tool_timeout: Some(tool_timeout),
tool_filter,
server_supports_sandbox_state_capability,
server_supports_out_of_band_elicitation_state_capability,
};
Ok(managed)
+1
View File
@@ -82,6 +82,7 @@ mod mcp;
mod mcp_escalation_policy;
mod socket;
mod stopwatch;
mod stopwatch_controller;
pub use mcp::ExecResult;
+60 -22
View File
@@ -4,8 +4,11 @@ use std::time::Duration;
use anyhow::Context as _;
use anyhow::Result;
use codex_core::MCP_OUT_OF_BAND_ELICITATION_STATE_CAPABILITY;
use codex_core::MCP_OUT_OF_BAND_ELICITATION_STATE_METHOD;
use codex_core::MCP_SANDBOX_STATE_CAPABILITY;
use codex_core::MCP_SANDBOX_STATE_METHOD;
use codex_core::OutOfBandElicitationState;
use codex_core::SandboxState;
use codex_core::protocol::SandboxPolicy;
use codex_execpolicy::Policy;
@@ -32,11 +35,13 @@ use crate::posix::escalate_server::EscalateServer;
use crate::posix::escalate_server::{self};
use crate::posix::mcp_escalation_policy::McpEscalationPolicy;
use crate::posix::stopwatch::Stopwatch;
use crate::posix::stopwatch_controller::StopwatchController;
/// Path to our patched bash.
const CODEX_BASH_PATH_ENV_VAR: &str = "CODEX_BASH_PATH";
const SANDBOX_STATE_CAPABILITY_VERSION: &str = "1.0.0";
const OUT_OF_BAND_ELICITATION_STATE_CAPABILITY_VERSION: &str = "1.0.0";
pub(crate) fn get_bash_path() -> Result<PathBuf> {
std::env::var(CODEX_BASH_PATH_ENV_VAR)
@@ -83,6 +88,7 @@ pub struct ExecTool {
policy: Arc<RwLock<Policy>>,
preserve_program_paths: bool,
sandbox_state: Arc<RwLock<Option<SandboxState>>>,
stopwatch_controller: StopwatchController,
}
#[tool_router]
@@ -100,6 +106,7 @@ impl ExecTool {
policy,
preserve_program_paths,
sandbox_state: Arc::new(RwLock::new(None)),
stopwatch_controller: StopwatchController::default(),
}
}
@@ -116,6 +123,7 @@ impl ExecTool {
.unwrap_or(codex_core::exec::DEFAULT_EXEC_COMMAND_TIMEOUT_MS),
);
let stopwatch = Stopwatch::new(effective_timeout);
let stopwatch_id = self.stopwatch_controller.register(stopwatch.clone()).await;
let cancel_token = stopwatch.cancellation_token();
let sandbox_state =
self.sandbox_state
@@ -141,8 +149,9 @@ impl ExecTool {
let result = escalate_server
.exec(params, cancel_token, &sandbox_state)
.await
.map_err(|e| McpError::internal_error(e.to_string(), None))?;
.await;
self.stopwatch_controller.unregister(stopwatch_id).await;
let result = result.map_err(|e| McpError::internal_error(e.to_string(), None))?;
Ok(CallToolResult::success(vec![Content::json(
ExecResult::from(result),
)?]))
@@ -169,6 +178,15 @@ impl ServerHandler for ExecTool {
MCP_SANDBOX_STATE_CAPABILITY.to_string(),
sandbox_state_capability,
);
let mut out_of_band_elicitation_state_capability = JsonObject::new();
out_of_band_elicitation_state_capability.insert(
"version".to_string(),
serde_json::Value::String(OUT_OF_BAND_ELICITATION_STATE_CAPABILITY_VERSION.to_string()),
);
experimental_capabilities.insert(
MCP_OUT_OF_BAND_ELICITATION_STATE_CAPABILITY.to_string(),
out_of_band_elicitation_state_capability,
);
ServerInfo {
protocol_version: ProtocolVersion::V_2025_06_18,
capabilities: ServerCapabilities::builder()
@@ -197,27 +215,47 @@ impl ServerHandler for ExecTool {
_context: rmcp::service::RequestContext<rmcp::RoleServer>,
) -> Result<CustomResult, McpError> {
let CustomRequest { method, params, .. } = request;
if method != MCP_SANDBOX_STATE_METHOD {
return Err(McpError::method_not_found::<CodexSandboxStateUpdateMethod>());
match method.as_ref() {
MCP_SANDBOX_STATE_METHOD => {
let Some(params) = params else {
return Err(McpError::invalid_params(
"missing params for sandbox state request".to_string(),
None,
));
};
let Ok(sandbox_state) = serde_json::from_value::<SandboxState>(params.clone())
else {
return Err(McpError::invalid_params(
"failed to deserialize sandbox state".to_string(),
Some(params),
));
};
*self.sandbox_state.write().await = Some(sandbox_state);
Ok(CustomResult::new(json!({})))
}
MCP_OUT_OF_BAND_ELICITATION_STATE_METHOD => {
let Some(params) = params else {
return Err(McpError::invalid_params(
"missing params for out-of-band elicitation state request".to_string(),
None,
));
};
let Ok(state) = serde_json::from_value::<OutOfBandElicitationState>(params.clone())
else {
return Err(McpError::invalid_params(
"failed to deserialize out-of-band elicitation state".to_string(),
Some(params),
));
};
self.stopwatch_controller.set_paused(state.paused).await;
Ok(CustomResult::new(json!({})))
}
_ => Err(McpError::method_not_found::<CodexSandboxStateUpdateMethod>()),
}
let Some(params) = params else {
return Err(McpError::invalid_params(
"missing params for sandbox state request".to_string(),
None,
));
};
let Ok(sandbox_state) = serde_json::from_value::<SandboxState>(params.clone()) else {
return Err(McpError::invalid_params(
"failed to deserialize sandbox state".to_string(),
Some(params),
));
};
*self.sandbox_state.write().await = Some(sandbox_state);
Ok(CustomResult::new(json!({})))
}
}
+2 -2
View File
@@ -90,7 +90,7 @@ impl Stopwatch {
result
}
async fn pause(&self) {
pub(crate) async fn pause(&self) {
let mut guard = self.inner.lock().await;
guard.active_pauses += 1;
if guard.active_pauses == 1
@@ -101,7 +101,7 @@ impl Stopwatch {
}
}
async fn resume(&self) {
pub(crate) async fn resume(&self) {
let mut guard = self.inner.lock().await;
if guard.active_pauses == 0 {
return;
@@ -0,0 +1,112 @@
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use crate::posix::stopwatch::Stopwatch;
#[derive(Clone, Debug, Default)]
pub(crate) struct StopwatchController {
state: Arc<Mutex<StopwatchControllerState>>,
operation_lock: Arc<Mutex<()>>,
}
#[derive(Debug, Default)]
struct StopwatchControllerState {
paused: bool,
next_stopwatch_id: u64,
stopwatches: HashMap<u64, Stopwatch>,
}
impl StopwatchController {
pub(crate) async fn register(&self, stopwatch: Stopwatch) -> u64 {
let _operation_guard = self.operation_lock.lock().await;
let (stopwatch_id, paused) = {
let mut guard = self.state.lock().await;
let stopwatch_id = guard.next_stopwatch_id;
guard.next_stopwatch_id += 1;
guard.stopwatches.insert(stopwatch_id, stopwatch.clone());
(stopwatch_id, guard.paused)
};
if paused {
stopwatch.pause().await;
}
stopwatch_id
}
pub(crate) async fn unregister(&self, stopwatch_id: u64) {
let _operation_guard = self.operation_lock.lock().await;
self.state.lock().await.stopwatches.remove(&stopwatch_id);
}
pub(crate) async fn set_paused(&self, paused: bool) {
let _operation_guard = self.operation_lock.lock().await;
let stopwatches = {
let mut guard = self.state.lock().await;
if guard.paused == paused {
return;
}
guard.paused = paused;
guard.stopwatches.values().cloned().collect::<Vec<_>>()
};
for stopwatch in stopwatches {
if paused {
stopwatch.pause().await;
} else {
stopwatch.resume().await;
}
}
}
}
#[cfg(test)]
mod tests {
use super::StopwatchController;
use crate::posix::stopwatch::Stopwatch;
use tokio::time::Duration;
use tokio::time::timeout;
#[tokio::test]
async fn pausing_controller_pauses_registered_stopwatch() {
let controller = StopwatchController::default();
let stopwatch = Stopwatch::new(Duration::from_millis(50));
let token = stopwatch.cancellation_token();
let stopwatch_id = controller.register(stopwatch).await;
controller.set_paused(true).await;
assert!(
timeout(Duration::from_millis(30), token.cancelled())
.await
.is_err()
);
controller.set_paused(false).await;
controller.unregister(stopwatch_id).await;
token.cancelled().await;
}
#[tokio::test]
async fn registering_while_paused_starts_paused() {
let controller = StopwatchController::default();
controller.set_paused(true).await;
let stopwatch = Stopwatch::new(Duration::from_millis(50));
let token = stopwatch.cancellation_token();
let stopwatch_id = controller.register(stopwatch).await;
assert!(
timeout(Duration::from_millis(30), token.cancelled())
.await
.is_err()
);
controller.set_paused(false).await;
controller.unregister(stopwatch_id).await;
token.cancelled().await;
}
}