Files
codex/codex-rs/app-server/src/message_processor_tracing_tests.rs
Ruslan Nigmatullin 4d201e340e state: pass state db handles through consumers (#20561)
## Why

SQLite state was still being opened from consumer paths, including lazy
`OnceCell`-backed thread-store call sites. That let one process
construct multiple state DB connections for the same Codex home, which
makes SQLite lock contention and `database is locked` failures much
easier to hit.

State DB lifetime should be chosen by main-like entrypoints and tests,
then passed through explicitly. Consumers should use the supplied
`Option<StateDbHandle>` or `StateDbHandle` and keep their existing
filesystem fallback or error behavior when no handle is available.

The startup path also needs to keep the rollout crate in charge of
SQLite state initialization. Opening `codex_state::StateRuntime`
directly bypasses rollout metadata backfill, so entrypoints should
initialize through `codex_rollout::state_db` and receive a handle only
after required rollout backfills have completed.

## What Changed

- Initialize the state DB in main-like entrypoints for CLI, TUI,
app-server, exec, MCP server, and the thread-manager sample.
- Pass `Option<StateDbHandle>` through `ThreadManager`,
`LocalThreadStore`, app-server processors, TUI app wiring, rollout
listing/recording, personality migration, shell snapshot cleanup,
session-name lookup, and memory/device-key consumers.
- Remove the lazy local state DB wrapper from the thread store so
non-test consumers use only the supplied handle or their existing
fallback path.
- Make `codex_rollout::state_db::init` the local state startup path: it
opens/migrates SQLite, runs rollout metadata backfill when needed, waits
for concurrent backfill workers up to a bounded timeout, verifies
completion, and then returns the initialized handle.
- Keep optional/non-owning SQLite helpers, such as remote TUI local
reads, as open-only paths that do not run startup backfill.
- Switch app-server startup from direct
`codex_state::StateRuntime::init` to the rollout state initializer so
app-server cannot skip rollout backfill.
- Collapse split rollout lookup/list APIs so callers use the normal
methods with an optional state handle instead of `_with_state_db`
variants.
- Restore `getConversationSummary(ThreadId)` to delegate through
`ThreadStore::read_thread` instead of a LocalThreadStore-specific
rollout path special case.
- Keep DB-backed rollout path lookup keyed on the DB row and file
existence, without imposing the filesystem filename convention on
existing DB rows.
- Verify readable DB-backed rollout paths against `session_meta.id`
before returning them, so a stale SQLite row that points at another
thread's JSONL falls back to filesystem search and read-repairs the DB
row.
- Keep `debug prompt-input` filesystem-only so a one-off debug command
does not initialize or backfill SQLite state just to print prompt input.
- Keep goal-session test Codex homes alive only in the goal-specific
helper, rather than leaking tempdirs from the shared session test
helper.
- Update tests and call sites to pass explicit state handles where DB
behavior is expected and explicit `None` where filesystem-only behavior
is intended.

## Validation

- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo check -p
codex-rollout -p codex-thread-store -p codex-app-server -p codex-core -p
codex-tui -p codex-exec -p codex-cli --tests`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-rollout state_db_`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-rollout find_thread_path`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-rollout find_thread_path -- --nocapture`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-rollout try_init_ -- --nocapture`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-rollout`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo clippy -p
codex-rollout --lib -- -D warnings`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-thread-store
read_thread_falls_back_when_sqlite_path_points_to_another_thread --
--nocapture`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-thread-store`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
shell_snapshot`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
--test all personality_migration`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
--test all rollout_list_find`
- `RUST_MIN_STACK=8388608 CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
--test all rollout_list_find::find_prefers_sqlite_path_by_id --
--nocapture`
- `RUST_MIN_STACK=8388608 CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
--test all rollout_list_find -- --nocapture`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
interrupt_accounts_active_goal_before_pausing`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-app-server get_auth_status -- --test-threads=1`
- `CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-app-server --lib`
- `CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo check -p codex-rollout
-p codex-app-server --tests`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db just fix -p codex-rollout
-p codex-thread-store -p codex-core -p codex-app-server -p codex-tui -p
codex-exec -p codex-cli`
- `CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db just fix -p codex-rollout -p
codex-app-server`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db just fix -p
codex-rollout`
- `CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db just fix -p codex-core`
- `just argument-comment-lint -p codex-core`
- `just argument-comment-lint -p codex-rollout`

Focused coverage added in `codex-rollout`:

- `recorder::tests::state_db_init_backfills_before_returning` verifies
the rollout metadata row exists before startup init returns.
- `state_db::tests::try_init_waits_for_concurrent_startup_backfill`
verifies startup waits for another worker to finish backfill instead of
disabling the handle for the process.
-
`state_db::tests::try_init_times_out_waiting_for_stuck_startup_backfill`
verifies startup does not hang indefinitely on a stuck backfill lease.
-
`tests::find_thread_path_accepts_existing_state_db_path_without_canonical_filename`
verifies DB-backed lookup accepts valid existing rollout paths even when
the filename does not include the thread UUID.
-
`tests::find_thread_path_falls_back_when_db_path_points_to_another_thread`
verifies DB-backed lookup ignores a stale row whose existing path
belongs to another thread and read-repairs the row after filesystem
fallback.

Focused coverage updated in `codex-core`:

- `rollout_list_find::find_prefers_sqlite_path_by_id` now uses a
DB-preferred rollout file with matching `session_meta.id`, so it still
verifies that valid SQLite paths win without depending on stale/empty
rollout contents.

`cargo test -p codex-app-server thread_list_respects_search_term_filter
-- --test-threads=1 --nocapture` was attempted locally but timed out
waiting for the app-server test harness `initialize` response before
reaching the changed thread-list code path.

`bazel test //codex-rs/thread-store:thread-store-unit-tests
--test_output=errors` was attempted locally after the thread-store fix,
but this container failed before target analysis while fetching `v8+`
through BuildBuddy/direct GitHub. The equivalent local crate coverage,
including `cargo test -p codex-thread-store`, passes.

A plain local `cargo check -p codex-rollout -p codex-app-server --tests`
also requires system `libcap.pc` for `codex-linux-sandbox`; the
follow-up app-server check above used `CODEX_SKIP_VENDORED_BWRAP=1` in
this container.
2026-05-04 11:46:03 -07:00

810 lines
27 KiB
Rust

use super::ConnectionSessionState;
use super::MessageProcessor;
use super::MessageProcessorArgs;
use crate::analytics_utils::analytics_events_client_from_config;
use crate::config_manager::ConfigManager;
use crate::outgoing_message::ConnectionId;
use crate::outgoing_message::OutgoingMessageSender;
use crate::transport::AppServerTransport;
use crate::transport::ConnectionOrigin;
use anyhow::Result;
use app_test_support::create_mock_responses_server_repeating_assistant;
use app_test_support::write_mock_responses_config_toml;
use codex_analytics::AppServerRpcTransport;
use codex_app_server_protocol::ClientInfo;
use codex_app_server_protocol::ClientRequest;
use codex_app_server_protocol::DeviceKeySignParams;
use codex_app_server_protocol::DeviceKeySignPayload;
use codex_app_server_protocol::InitializeCapabilities;
use codex_app_server_protocol::InitializeParams;
use codex_app_server_protocol::InitializeResponse;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::JSONRPCRequest;
use codex_app_server_protocol::RemoteControlClientConnectionAudience;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse;
use codex_app_server_protocol::TurnStartParams;
use codex_app_server_protocol::TurnStartResponse;
use codex_app_server_protocol::UserInput;
use codex_arg0::Arg0DispatchPaths;
use codex_config::CloudRequirementsLoader;
use codex_config::LoaderOverrides;
use codex_core::config::Config;
use codex_core::config::ConfigBuilder;
use codex_exec_server::EnvironmentManager;
use codex_feedback::CodexFeedback;
use codex_login::AuthManager;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::W3cTraceContext;
use opentelemetry::global;
use opentelemetry::trace::SpanId;
use opentelemetry::trace::SpanKind;
use opentelemetry::trace::TraceId;
use opentelemetry::trace::TracerProvider as _;
use opentelemetry_sdk::propagation::TraceContextPropagator;
use opentelemetry_sdk::trace::InMemorySpanExporter;
use opentelemetry_sdk::trace::SdkTracerProvider;
use opentelemetry_sdk::trace::SpanData;
use pretty_assertions::assert_eq;
use serial_test::serial;
use std::collections::BTreeMap;
use std::future::Future;
use std::path::Path;
use std::sync::Arc;
use std::sync::OnceLock;
use tempfile::TempDir;
use tokio::sync::mpsc;
use tracing_subscriber::layer::SubscriberExt;
use wiremock::MockServer;
const TEST_CONNECTION_ID: ConnectionId = ConnectionId(7);
struct TestTracing {
exporter: InMemorySpanExporter,
provider: SdkTracerProvider,
}
struct RemoteTrace {
trace_id: TraceId,
parent_span_id: SpanId,
context: W3cTraceContext,
}
impl RemoteTrace {
fn new(trace_id: &str, parent_span_id: &str) -> Self {
let trace_id = TraceId::from_hex(trace_id).expect("trace id");
let parent_span_id = SpanId::from_hex(parent_span_id).expect("parent span id");
let context = W3cTraceContext {
traceparent: Some(format!("00-{trace_id}-{parent_span_id}-01")),
tracestate: Some("vendor=value".to_string()),
};
Self {
trace_id,
parent_span_id,
context,
}
}
}
fn init_test_tracing() -> &'static TestTracing {
static TEST_TRACING: OnceLock<TestTracing> = OnceLock::new();
TEST_TRACING.get_or_init(|| {
let exporter = InMemorySpanExporter::default();
let provider = SdkTracerProvider::builder()
.with_simple_exporter(exporter.clone())
.build();
let tracer = provider.tracer("codex-app-server-message-processor-tests");
global::set_text_map_propagator(TraceContextPropagator::new());
let subscriber =
tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer));
tracing::subscriber::set_global_default(subscriber)
.expect("global tracing subscriber should only be installed once");
TestTracing { exporter, provider }
})
}
fn request_from_client_request(request: ClientRequest) -> JSONRPCRequest {
serde_json::from_value(serde_json::to_value(request).expect("serialize client request"))
.expect("client request should convert to JSON-RPC")
}
struct TracingHarness {
_server: MockServer,
_codex_home: TempDir,
processor: Arc<MessageProcessor>,
outgoing_rx: mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
session: Arc<ConnectionSessionState>,
tracing: &'static TestTracing,
}
impl TracingHarness {
async fn new() -> Result<Self> {
Self::new_with_origin(ConnectionOrigin::WebSocket).await
}
async fn new_with_origin(origin: ConnectionOrigin) -> Result<Self> {
let server = create_mock_responses_server_repeating_assistant("Done").await;
let codex_home = TempDir::new()?;
let config = Arc::new(build_test_config(codex_home.path(), &server.uri()).await?);
let (processor, outgoing_rx) = build_test_processor(config).await;
let tracing = init_test_tracing();
tracing.exporter.reset();
tracing::callsite::rebuild_interest_cache();
let mut harness = Self {
_server: server,
_codex_home: codex_home,
processor,
outgoing_rx,
session: Arc::new(ConnectionSessionState::new(origin)),
tracing,
};
let _: InitializeResponse = harness
.request(
ClientRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_info: ClientInfo {
name: "codex-app-server-tests".to_string(),
title: None,
version: "0.1.0".to_string(),
},
capabilities: Some(InitializeCapabilities {
experimental_api: true,
..Default::default()
}),
},
},
/*trace*/ None,
)
.await;
assert!(harness.session.initialized());
Ok(harness)
}
fn reset_tracing(&self) {
self.tracing.exporter.reset();
}
async fn shutdown(self) {
self.processor.shutdown_threads().await;
self.processor.drain_background_tasks().await;
}
async fn request<T>(&mut self, request: ClientRequest, trace: Option<W3cTraceContext>) -> T
where
T: serde::de::DeserializeOwned,
{
let request_id = match request.id() {
RequestId::Integer(request_id) => *request_id,
request_id => panic!("expected integer request id in test harness, got {request_id:?}"),
};
let mut request = request_from_client_request(request);
request.trace = trace;
self.processor
.process_request(
TEST_CONNECTION_ID,
request,
&AppServerTransport::Stdio,
Arc::clone(&self.session),
)
.await;
read_response(&mut self.outgoing_rx, request_id).await
}
async fn request_error(
&mut self,
request: ClientRequest,
trace: Option<W3cTraceContext>,
) -> JSONRPCErrorError {
let request_id = match request.id() {
RequestId::Integer(request_id) => *request_id,
request_id => panic!("expected integer request id in test harness, got {request_id:?}"),
};
let mut request = request_from_client_request(request);
request.trace = trace;
self.processor
.process_request(
TEST_CONNECTION_ID,
request,
&AppServerTransport::Stdio,
Arc::clone(&self.session),
)
.await;
read_error(&mut self.outgoing_rx, request_id).await
}
async fn start_thread(
&mut self,
request_id: i64,
trace: Option<W3cTraceContext>,
) -> ThreadStartResponse {
let response = self
.request(
ClientRequest::ThreadStart {
request_id: RequestId::Integer(request_id),
params: ThreadStartParams {
ephemeral: Some(true),
..ThreadStartParams::default()
},
},
trace,
)
.await;
read_thread_started_notification(&mut self.outgoing_rx).await;
response
}
}
async fn build_test_config(codex_home: &Path, server_uri: &str) -> Result<Config> {
write_mock_responses_config_toml(
codex_home,
server_uri,
&BTreeMap::new(),
/*auto_compact_limit*/ 8_192,
Some(false),
"mock_provider",
"compact",
)?;
Ok(ConfigBuilder::default()
.codex_home(codex_home.to_path_buf())
.build()
.await?)
}
async fn build_test_processor(
config: Arc<Config>,
) -> (
Arc<MessageProcessor>,
mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
) {
let (outgoing_tx, outgoing_rx) = mpsc::channel(16);
let auth_manager =
AuthManager::shared_from_config(config.as_ref(), /*enable_codex_api_key_env*/ false).await;
let config_manager = ConfigManager::new(
config.codex_home.to_path_buf(),
Vec::new(),
LoaderOverrides::default(),
CloudRequirementsLoader::default(),
Arg0DispatchPaths::default(),
Arc::new(codex_config::NoopThreadConfigLoader),
);
let analytics_events_client =
analytics_events_client_from_config(Arc::clone(&auth_manager), config.as_ref());
let outgoing = Arc::new(OutgoingMessageSender::new(
outgoing_tx,
analytics_events_client.clone(),
));
let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs {
outgoing,
analytics_events_client,
arg0_paths: Arg0DispatchPaths::default(),
config,
config_manager,
environment_manager: Arc::new(EnvironmentManager::default_for_tests()),
feedback: CodexFeedback::new(),
log_db: None,
state_db: None,
config_warnings: Vec::new(),
session_source: SessionSource::VSCode,
auth_manager,
rpc_transport: AppServerRpcTransport::Stdio,
remote_control_handle: None,
plugin_startup_tasks: crate::PluginStartupTasks::Start,
}));
(processor, outgoing_rx)
}
fn run_current_thread_test_with_stack<F>(name: &str, future: F) -> Result<()>
where
F: Future<Output = Result<()>> + Send + 'static,
{
const TEST_STACK_SIZE_BYTES: usize = 4 * 1024 * 1024;
let handle = std::thread::Builder::new()
.name(name.to_string())
.stack_size(TEST_STACK_SIZE_BYTES)
.spawn(move || -> Result<()> {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()?;
runtime.block_on(Box::pin(future))
})?;
match handle.join() {
Ok(result) => result,
Err(_) => Err(anyhow::anyhow!("{name} thread panicked")),
}
}
fn span_attr<'a>(span: &'a SpanData, key: &str) -> Option<&'a str> {
span.attributes
.iter()
.find(|kv| kv.key.as_str() == key)
.and_then(|kv| match &kv.value {
opentelemetry::Value::String(value) => Some(value.as_str()),
_ => None,
})
}
fn find_rpc_span_with_trace<'a>(
spans: &'a [SpanData],
kind: SpanKind,
method: &str,
trace_id: TraceId,
) -> &'a SpanData {
spans
.iter()
.find(|span| {
span.span_kind == kind
&& span_attr(span, "rpc.system") == Some("jsonrpc")
&& span_attr(span, "rpc.method") == Some(method)
&& span.span_context.trace_id() == trace_id
})
.unwrap_or_else(|| {
panic!(
"missing {kind:?} span for rpc.method={method} trace={trace_id}; exported spans:\n{}",
format_spans(spans)
)
})
}
fn find_span_with_trace<'a, F>(
spans: &'a [SpanData],
trace_id: TraceId,
description: &str,
predicate: F,
) -> &'a SpanData
where
F: Fn(&SpanData) -> bool,
{
spans
.iter()
.find(|span| span.span_context.trace_id() == trace_id && predicate(span))
.unwrap_or_else(|| {
panic!(
"missing span matching {description} for trace={trace_id}; exported spans:\n{}",
format_spans(spans)
)
})
}
fn format_spans(spans: &[SpanData]) -> String {
spans
.iter()
.map(|span| {
let rpc_method = span_attr(span, "rpc.method").unwrap_or("-");
format!(
"name={} span_id={} kind={:?} parent={} trace={} rpc.method={}",
span.name,
span.span_context.span_id(),
span.span_kind,
span.parent_span_id,
span.span_context.trace_id(),
rpc_method
)
})
.collect::<Vec<_>>()
.join("\n")
}
fn span_depth_from_ancestor(
spans: &[SpanData],
child: &SpanData,
ancestor: &SpanData,
) -> Option<usize> {
let ancestor_span_id = ancestor.span_context.span_id();
let mut parent_span_id = child.parent_span_id;
let mut depth = 1;
while parent_span_id != SpanId::INVALID {
if parent_span_id == ancestor_span_id {
return Some(depth);
}
let Some(parent_span) = spans
.iter()
.find(|span| span.span_context.span_id() == parent_span_id)
else {
break;
};
parent_span_id = parent_span.parent_span_id;
depth += 1;
}
None
}
fn assert_span_descends_from(spans: &[SpanData], child: &SpanData, ancestor: &SpanData) {
if span_depth_from_ancestor(spans, child, ancestor).is_some() {
return;
}
panic!(
"span {} does not descend from {}; exported spans:\n{}",
child.name,
ancestor.name,
format_spans(spans)
);
}
fn assert_has_internal_descendant_at_min_depth(
spans: &[SpanData],
ancestor: &SpanData,
min_depth: usize,
) {
if spans.iter().any(|span| {
span.span_kind == SpanKind::Internal
&& span.span_context.trace_id() == ancestor.span_context.trace_id()
&& span_depth_from_ancestor(spans, span, ancestor)
.is_some_and(|depth| depth >= min_depth)
}) {
return;
}
panic!(
"missing internal descendant at depth >= {min_depth} below {}; exported spans:\n{}",
ancestor.name,
format_spans(spans)
);
}
async fn read_response<T: serde::de::DeserializeOwned>(
outgoing_rx: &mut mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
request_id: i64,
) -> T {
loop {
let envelope = tokio::time::timeout(std::time::Duration::from_secs(5), outgoing_rx.recv())
.await
.expect("timed out waiting for response")
.expect("outgoing channel closed");
let crate::outgoing_message::OutgoingEnvelope::ToConnection {
connection_id,
message,
..
} = envelope
else {
continue;
};
if connection_id != TEST_CONNECTION_ID {
continue;
}
let crate::outgoing_message::OutgoingMessage::Response(response) = message else {
continue;
};
if response.id != RequestId::Integer(request_id) {
continue;
}
return serde_json::from_value(response.result)
.expect("response payload should deserialize");
}
}
async fn read_error(
outgoing_rx: &mut mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
request_id: i64,
) -> JSONRPCErrorError {
loop {
let envelope = tokio::time::timeout(std::time::Duration::from_secs(5), outgoing_rx.recv())
.await
.expect("timed out waiting for error")
.expect("outgoing channel closed");
let crate::outgoing_message::OutgoingEnvelope::ToConnection {
connection_id,
message,
..
} = envelope
else {
continue;
};
if connection_id != TEST_CONNECTION_ID {
continue;
}
let crate::outgoing_message::OutgoingMessage::Error(error) = message else {
continue;
};
if error.id != RequestId::Integer(request_id) {
continue;
}
return error.error;
}
}
async fn read_thread_started_notification(
outgoing_rx: &mut mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
) {
loop {
let envelope = tokio::time::timeout(std::time::Duration::from_secs(5), outgoing_rx.recv())
.await
.expect("timed out waiting for thread/started notification")
.expect("outgoing channel closed");
match envelope {
crate::outgoing_message::OutgoingEnvelope::ToConnection {
connection_id,
message,
..
} => {
if connection_id != TEST_CONNECTION_ID {
continue;
}
let crate::outgoing_message::OutgoingMessage::AppServerNotification(notification) =
message
else {
continue;
};
if matches!(
notification,
codex_app_server_protocol::ServerNotification::ThreadStarted(_)
) {
return;
}
}
crate::outgoing_message::OutgoingEnvelope::Broadcast { message } => {
let crate::outgoing_message::OutgoingMessage::AppServerNotification(notification) =
message
else {
continue;
};
if matches!(
notification,
codex_app_server_protocol::ServerNotification::ThreadStarted(_)
) {
return;
}
}
}
}
}
async fn wait_for_exported_spans<F>(tracing: &TestTracing, predicate: F) -> Vec<SpanData>
where
F: Fn(&[SpanData]) -> bool,
{
let mut last_spans = Vec::new();
for _ in 0..200 {
tokio::task::yield_now().await;
tracing
.provider
.force_flush()
.expect("force flush should succeed");
let spans = tracing.exporter.get_finished_spans().expect("span export");
last_spans = spans.clone();
if predicate(&spans) {
return spans;
}
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
}
panic!(
"timed out waiting for expected exported spans:\n{}",
format_spans(&last_spans)
);
}
async fn wait_for_new_exported_spans<F>(
tracing: &TestTracing,
baseline_len: usize,
predicate: F,
) -> Vec<SpanData>
where
F: Fn(&[SpanData]) -> bool,
{
let spans = wait_for_exported_spans(tracing, |spans| {
spans.len() > baseline_len && predicate(&spans[baseline_len..])
})
.await;
spans.into_iter().skip(baseline_len).collect()
}
#[test]
#[serial(app_server_tracing)]
fn thread_start_jsonrpc_span_exports_server_span_and_parents_children() -> Result<()> {
run_current_thread_test_with_stack(
"thread_start_jsonrpc_span_exports_server_span_and_parents_children",
async {
let mut harness = TracingHarness::new().await?;
let RemoteTrace {
trace_id: remote_trace_id,
parent_span_id: remote_parent_span_id,
context: remote_trace,
..
} = RemoteTrace::new("00000000000000000000000000000011", "0000000000000022");
let _: ThreadStartResponse = harness
.start_thread(/*request_id*/ 20_002, /*trace*/ None)
.await;
let untraced_spans = wait_for_exported_spans(harness.tracing, |spans| {
spans.iter().any(|span| {
span.span_kind == SpanKind::Server
&& span_attr(span, "rpc.method") == Some("thread/start")
})
})
.await;
let untraced_server_span = find_rpc_span_with_trace(
&untraced_spans,
SpanKind::Server,
"thread/start",
untraced_spans
.iter()
.rev()
.find(|span| {
span.span_kind == SpanKind::Server
&& span_attr(span, "rpc.system") == Some("jsonrpc")
&& span_attr(span, "rpc.method") == Some("thread/start")
})
.unwrap_or_else(|| {
panic!(
"missing latest thread/start server span; exported spans:\n{}",
format_spans(&untraced_spans)
)
})
.span_context
.trace_id(),
);
assert_has_internal_descendant_at_min_depth(
&untraced_spans,
untraced_server_span,
/*min_depth*/ 1,
);
let baseline_len = untraced_spans.len();
let _: ThreadStartResponse = harness
.start_thread(/*request_id*/ 20_003, Some(remote_trace))
.await;
let spans = wait_for_new_exported_spans(harness.tracing, baseline_len, |spans| {
spans.iter().any(|span| {
span.span_kind == SpanKind::Server
&& span_attr(span, "rpc.method") == Some("thread/start")
&& span.span_context.trace_id() == remote_trace_id
}) && spans.iter().any(|span| {
span.name.as_ref() == "app_server.thread_start.notify_started"
&& span.span_context.trace_id() == remote_trace_id
})
})
.await;
let server_request_span =
find_rpc_span_with_trace(&spans, SpanKind::Server, "thread/start", remote_trace_id);
assert_eq!(server_request_span.name.as_ref(), "thread/start");
assert_eq!(server_request_span.parent_span_id, remote_parent_span_id);
assert!(server_request_span.parent_span_is_remote);
assert_eq!(server_request_span.span_context.trace_id(), remote_trace_id);
assert_ne!(server_request_span.span_context.span_id(), SpanId::INVALID);
assert_has_internal_descendant_at_min_depth(
&spans,
server_request_span,
/*min_depth*/ 1,
);
assert_has_internal_descendant_at_min_depth(
&spans,
server_request_span,
/*min_depth*/ 2,
);
harness.shutdown().await;
Ok(())
},
)
}
#[tokio::test(flavor = "current_thread")]
#[serial(app_server_tracing)]
async fn remote_control_origin_rejects_device_key_requests() -> Result<()> {
let mut harness = TracingHarness::new_with_origin(ConnectionOrigin::RemoteControl).await?;
let error = harness
.request_error(
ClientRequest::DeviceKeySign {
request_id: RequestId::Integer(20_004),
params: DeviceKeySignParams {
key_id: "dk_123".to_string(),
payload: DeviceKeySignPayload::RemoteControlClientConnection {
nonce: "nonce-123".to_string(),
audience:
RemoteControlClientConnectionAudience::RemoteControlClientWebsocket,
session_id: "wssess_123".to_string(),
target_origin: "https://chatgpt.com".to_string(),
target_path: "/api/codex/remote/control/client".to_string(),
account_user_id: "acct_123".to_string(),
client_id: "cli_123".to_string(),
token_expires_at: 4_102_444_800,
token_sha256_base64url: "47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU"
.to_string(),
scopes: vec!["remote_control_controller_websocket".to_string()],
},
},
},
/*trace*/ None,
)
.await;
assert_eq!(error.code, crate::error_code::INVALID_REQUEST_ERROR_CODE);
assert_eq!(
error.message,
"device/key/sign is not available over remote transports"
);
harness.shutdown().await;
Ok(())
}
#[tokio::test(flavor = "current_thread")]
#[serial(app_server_tracing)]
async fn turn_start_jsonrpc_span_parents_core_turn_spans() -> Result<()> {
let mut harness = TracingHarness::new().await?;
let thread_start_response = harness.start_thread(/*request_id*/ 2, /*trace*/ None).await;
let thread_id = thread_start_response.thread.id.clone();
harness.reset_tracing();
let RemoteTrace {
trace_id: remote_trace_id,
parent_span_id: remote_parent_span_id,
context: remote_trace,
} = RemoteTrace::new("00000000000000000000000000000077", "0000000000000088");
let turn_start_response: TurnStartResponse = harness
.request(
ClientRequest::TurnStart {
request_id: RequestId::Integer(3),
params: TurnStartParams {
environments: None,
thread_id,
input: vec![UserInput::Text {
text: "hello".to_string(),
text_elements: Vec::new(),
}],
responsesapi_client_metadata: None,
cwd: None,
approval_policy: None,
sandbox_policy: None,
permissions: None,
approvals_reviewer: None,
model: None,
service_tier: None,
effort: None,
summary: None,
personality: None,
output_schema: None,
collaboration_mode: None,
},
},
Some(remote_trace),
)
.await;
let spans = wait_for_exported_spans(harness.tracing, |spans| {
spans.iter().any(|span| {
span.span_kind == SpanKind::Server
&& span_attr(span, "rpc.method") == Some("turn/start")
&& span.span_context.trace_id() == remote_trace_id
}) && spans.iter().any(|span| {
span_attr(span, "codex.op") == Some("user_input")
&& span.span_context.trace_id() == remote_trace_id
})
})
.await;
let server_request_span =
find_rpc_span_with_trace(&spans, SpanKind::Server, "turn/start", remote_trace_id);
let core_turn_span =
find_span_with_trace(&spans, remote_trace_id, "codex.op=user_input", |span| {
span_attr(span, "codex.op") == Some("user_input")
});
assert_eq!(server_request_span.parent_span_id, remote_parent_span_id);
assert!(server_request_span.parent_span_is_remote);
assert_eq!(server_request_span.span_context.trace_id(), remote_trace_id);
assert_eq!(
span_attr(server_request_span, "turn.id"),
Some(turn_start_response.turn.id.as_str())
);
assert_span_descends_from(&spans, core_turn_span, server_request_span);
harness.shutdown().await;
Ok(())
}