Add cancellation for experimental thread samples (#55441)

## What changed

Add the experimental `thread/sample/cancel` request with a `requestId` identifying a `thread/sample` request on the same connection. Cancellation returns an empty response, stops the matching sample with `sample canceled`, and cleans up its ephemeral fork. Repeated cancellation and unknown request IDs are harmless no-ops; other request types are unaffected.

## Testing

Add coverage for connection and request ID isolation, request type filtering, and repeated cancellation. Extend the sample integration test to verify cancellation keeps the parent responsive and reclaims the child thread.

GitOrigin-RevId: 92d12e5c0190e1d3b06cb482543615a3ec77bdc8

codex-rs/app-server-protocol/src/protocol/common.rs

@@ -574,6 +574,12 @@ client_request_definitions! {
         serialization: None,
         response: v2::ThreadSampleResponse,
     },
+    #[experimental("thread/sample/cancel")]
+    ThreadSampleCancel => "thread/sample/cancel" {
+        params: v2::ThreadSampleCancelParams,
+        serialization: None,
+        response: v2::ThreadSampleCancelResponse,
+    },
     ThreadArchive => "thread/archive" {
         params: v2::ThreadArchiveParams,
         serialization: thread_id(params.thread_id),

codex-rs/app-server-protocol/src/protocol/v2/thread_sample.rs

@@ -1,6 +1,7 @@
 //! One model response on an ephemeral fork that inherits the source's configuration and context.
 
 use crate::JsonSchema;
+use crate::RequestId;
 use crate::TS;
 use codex_protocol::openai_models::ReasoningEffort;
 use serde::Deserialize;
@@ -33,3 +34,15 @@ pub struct ThreadSampleResponse {
     pub cached_input_tokens: Option<i64>,
     pub output_tokens: Option<i64>,
 }
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+#[ts(export_to = "v2/")]
+pub struct ThreadSampleCancelParams {
+    pub request_id: RequestId,
+}
+
+#[derive(Serialize, Deserialize, Debug, Default, Clone, PartialEq, Eq, JsonSchema, TS)]
+#[serde(rename_all = "camelCase")]
+#[ts(export_to = "v2/")]
+pub struct ThreadSampleCancelResponse {}

codex-rs/app-server/src/message_processor.rs

@@ -72,6 +72,7 @@ use codex_app_server_protocol::JSONRPCErrorError;
 use codex_app_server_protocol::JSONRPCNotification;
 use codex_app_server_protocol::JSONRPCRequest;
 use codex_app_server_protocol::JSONRPCResponse;
+use codex_app_server_protocol::ThreadSampleCancelResponse;
 use codex_app_server_protocol::UserVerificationCancelResponse;
 use codex_app_server_protocol::experimental_required_message;
 use codex_arg0::Arg0DispatchPaths;
@@ -1344,6 +1345,15 @@ impl MessageProcessor {
                 .thread_sample(params, Arc::clone(&session.rpc_gate), request_context)
                 .await
                 .map(|response| Some(response.into())),
+            ClientRequest::ThreadSampleCancel { params, .. } => {
+                self.outgoing
+                    .cancel_thread_sample_request(&ConnectionRequestId {
+                        connection_id,
+                        request_id: params.request_id,
+                    })
+                    .await;
+                Ok(Some(ThreadSampleCancelResponse {}.into()))
+            }
             ClientRequest::ThreadArchive { params, .. } => {
                 self.thread_processor
                     .thread_archive(request_id.clone(), params)

codex-rs/app-server/src/outgoing_message.rs

@@ -75,6 +75,7 @@ pub(crate) struct RequestContext {
 enum RequestCancellationScope {
     Unavailable,
     UserVerification,
+    ThreadSample,
 }
 
 impl RequestContext {
@@ -92,6 +93,7 @@ impl RequestContext {
                 | "userVerification/enroll"
                 | "userVerification/delete"
                 | "userVerification/verify" => RequestCancellationScope::UserVerification,
+                "thread/sample" => RequestCancellationScope::ThreadSample,
                 _ => RequestCancellationScope::Unavailable,
             },
             span,
@@ -278,6 +280,15 @@ impl OutgoingMessageSender {
         }
     }
 
+    pub(crate) async fn cancel_thread_sample_request(&self, request_id: &ConnectionRequestId) {
+        let contexts = self.request_contexts.lock().await;
+        if let Some(context) = contexts.get(request_id)
+            && context.cancellation_scope == RequestCancellationScope::ThreadSample
+        {
+            context.cancellation.cancel();
+        }
+    }
+
     pub(crate) async fn connection_closed(&self, connection_id: ConnectionId) {
         self.disconnect_user_verification_connection(connection_id)
             .await;
@@ -908,6 +919,10 @@ fn timestamped_server_notification(notification: ServerNotification) -> Outgoing
     })
 }
 
+#[cfg(test)]
+#[path = "thread_sample_cancel_context_tests.rs"]
+mod thread_sample_cancel_context_tests;
+
 #[cfg(test)]
 #[path = "user_verification_ownership_tests.rs"]
 mod user_verification_ownership_tests;

codex-rs/app-server/src/request_processors/thread_processor.rs

@@ -637,6 +637,7 @@ impl ThreadRequestProcessor {
                 samples: Arc::clone(&self.sample_responses),
                 params,
                 trace: context.request_trace(),
+                cancellation: context.cancellation,
             },
             &self.background_tasks,
             rpc_gate,

codex-rs/app-server/src/thread_sample.rs

@@ -49,6 +49,7 @@ use std::sync::Mutex;
 use std::time::Duration;
 use tokio::sync::oneshot;
 use tokio::time::Instant;
+use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 
 const SAMPLE_TIMEOUT: Duration = Duration::from_secs(30);
@@ -68,6 +69,7 @@ pub(crate) struct SampleTask {
     pub(crate) samples: Arc<SampleResponseRegistry>,
     pub(crate) params: ThreadSampleParams,
     pub(crate) trace: Option<W3cTraceContext>,
+    pub(crate) cancellation: CancellationToken,
 }
 
 pub(crate) async fn generate_sample(
@@ -109,6 +111,7 @@ impl SampleTask {
             samples,
             params,
             trace,
+            cancellation,
         } = self;
         let id = parent.startup_metadata().thread_id;
         if parent.agent_status().await == AgentStatus::Running {
@@ -195,6 +198,9 @@ impl SampleTask {
         options.session_source = Some(SessionSource::Custom("thread_sample".to_string()));
         options.thread_source = Some(ThreadSource::Feature("thread_sample".to_string()));
         options.parent_trace = trace.clone();
+        if cancellation.is_cancelled() {
+            return Err(internal_error("sample canceled"));
+        }
         if Instant::now() >= deadline {
             return Err(internal_error("sample timed out"));
         }
@@ -280,6 +286,7 @@ impl SampleTask {
         };
         let result = tokio::select! {
             biased;
+            _ = cancellation.cancelled() => Err(internal_error("sample canceled")),
             _ = tokio::time::sleep_until(deadline) => Err(internal_error("sample timed out")),
             result = sample => result,
         };

codex-rs/app-server/src/thread_sample_cancel_context_tests.rs

@@ -0,0 +1,56 @@
+//! Checks that sample cancellation only affects matching requests on the same connection.
+
+use super::*;
+use pretty_assertions::assert_eq;
+
+#[tokio::test]
+async fn thread_sample_cancel_only_signals_matching_request_contexts() {
+    let (tx, _rx) = mpsc::channel(/*buffer*/ 1);
+    let outgoing = OutgoingMessageSender::new(tx, AnalyticsEventsClient::disabled());
+    for method in [
+        "thread/sample",
+        "thread/sample/cancel",
+        "userVerification/status",
+        "userVerification/enroll",
+        "userVerification/delete",
+        "userVerification/verify",
+        "userVerification/cancel",
+        "turn/start",
+    ] {
+        let target = ConnectionRequestId {
+            connection_id: ConnectionId(1),
+            request_id: RequestId::Integer(7),
+        };
+        let context = RequestContext::new(
+            target.clone(),
+            method,
+            Span::none(),
+            /*parent_trace*/ None,
+        );
+        let cancellation = context.cancellation.clone();
+        outgoing.register_request_context(context).await;
+        for wrong_target in [
+            ConnectionRequestId {
+                connection_id: ConnectionId(2),
+                request_id: RequestId::Integer(7),
+            },
+            ConnectionRequestId {
+                connection_id: ConnectionId(1),
+                request_id: RequestId::String("7".into()),
+            },
+            ConnectionRequestId {
+                connection_id: ConnectionId(1),
+                request_id: RequestId::Integer(8),
+            },
+        ] {
+            outgoing.cancel_thread_sample_request(&wrong_target).await;
+            assert!(!cancellation.is_cancelled());
+        }
+        outgoing.cancel_thread_sample_request(&target).await;
+        outgoing.cancel_thread_sample_request(&target).await;
+        assert_eq!(cancellation.is_cancelled(), method == "thread/sample");
+        assert_eq!(outgoing.request_context_count().await, 1);
+        outgoing.connection_closed(ConnectionId(1)).await;
+        assert_eq!(outgoing.request_context_count().await, 0);
+    }
+}

codex-rs/app-server/tests/suite/v2/thread_sample.rs

@@ -12,6 +12,7 @@ use codex_app_server_protocol::ThreadForkParams;
 use codex_app_server_protocol::ThreadForkResponse;
 use codex_app_server_protocol::ThreadHistoryMode;
 use codex_app_server_protocol::ThreadLoadedListResponse;
+use codex_app_server_protocol::ThreadSampleCancelResponse;
 use codex_app_server_protocol::ThreadSampleParams;
 use codex_app_server_protocol::ThreadSampleResponse;
 use codex_app_server_protocol::ThreadStartParams;
@@ -732,8 +733,15 @@ async fn sample_does_not_hide_stream_errors_after_tool_calls() -> Result<()> {
     Ok(())
 }
 
+enum SampleStop {
+    Timeout,
+    Cancel,
+}
+
+#[test_case(SampleStop::Timeout; "timeout")]
+#[test_case(SampleStop::Cancel; "cancel")]
 #[tokio::test]
-async fn sample_wait_keeps_parent_responsive_and_timeout_reclaims_child() -> Result<()> {
+async fn sample_wait_keeps_parent_responsive_and_reclaims_child(stop: SampleStop) -> Result<()> {
     let server = responses::start_mock_server().await;
     responses::mount_sse_once(
         &server,
@@ -781,16 +789,40 @@ async fn sample_wait_keeps_parent_responsive_and_timeout_reclaims_child() -> Res
         }),
     )
     .await??;
+    let (expected_error, wait) = match stop {
+        SampleStop::Timeout => ("sample timed out", WAIT),
+        SampleStop::Cancel => {
+            let cancel_id = client
+                .send_raw_request(
+                    "thread/sample/cancel",
+                    Some(json!({"requestId": sample_id})),
+                )
+                .await?;
+            let response: ThreadSampleCancelResponse =
+                timeout(Duration::from_secs(10), client.read_response(cancel_id)).await??;
+            assert_eq!(response, ThreadSampleCancelResponse {});
+            ("sample canceled", Duration::from_secs(10))
+        }
+    };
     let error = timeout(
-        WAIT,
+        wait,
         client.read_stream_until_error_message(RequestId::Integer(sample_id)),
     )
     .await??;
-    assert!(
-        error.error.message.contains("timed out"),
-        "{}",
-        error.error.message
-    );
+    assert_eq!(error.error.message, expected_error);
+    if matches!(stop, SampleStop::Cancel) {
+        for request_id in [json!(sample_id), json!("unknown")] {
+            let cancel_id = client
+                .send_raw_request(
+                    "thread/sample/cancel",
+                    Some(json!({"requestId": request_id})),
+                )
+                .await?;
+            let response: ThreadSampleCancelResponse =
+                timeout(Duration::from_secs(10), client.read_response(cancel_id)).await??;
+            assert_eq!(response, ThreadSampleCancelResponse {});
+        }
+    }
     assert_eq!(mock.requests().len(), 1);
     timeout(WAIT, async {
         loop {