codex/codex-rs/core/tests/suite/sqlite_state.rs

use anyhow::Result;
use codex_config::types::McpServerConfig;
use codex_config::types::McpServerTransportConfig;
use codex_core::config::Config;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_protocol::ThreadId;
use codex_protocol::config_types::WebSearchMode;
use codex_protocol::dynamic_tools::DynamicToolFunctionSpec;
use codex_protocol::dynamic_tools::DynamicToolNamespaceSpec;
use codex_protocol::dynamic_tools::DynamicToolNamespaceTool;
use codex_protocol::dynamic_tools::DynamicToolSpec;
use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::RolloutLine;
use codex_protocol::protocol::SessionMeta;
use codex_protocol::protocol::SessionMetaLine;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::UserMessageEvent;
use codex_protocol::user_input::UserInput;
use codex_web_search_extension::install as install_web_search_extension;
use core_test_support::responses;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_function_call;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::ev_web_search_call_done;
use core_test_support::responses::mount_sse_once;
use core_test_support::responses::mount_sse_sequence;
use core_test_support::responses::start_mock_server;
use core_test_support::skip_if_no_network;
use core_test_support::stdio_server_bin;
use core_test_support::test_codex::local_selections;
use core_test_support::test_codex::test_codex;
use core_test_support::test_codex::turn_permission_fields;
use core_test_support::wait_for_event;
use core_test_support::wait_for_event_match;
use core_test_support::wait_for_mcp_server;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::HashMap;
use std::fs;
use std::sync::Arc;
use tokio::time::Duration;
use tracing_subscriber::prelude::*;
use uuid::Uuid;
use wiremock::Mock;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn new_thread_is_recorded_in_state_db() -> Result<()> {
    let server = start_mock_server().await;
    let mut builder = test_codex().with_config(|config| {
        config
            .features
            .enable(Feature::Sqlite)
            .expect("test config should allow feature update");
    });
    let test = builder.build(&server).await?;

    let thread_id = test.session_configured.thread_id;
    let rollout_path = test.codex.rollout_path().expect("rollout path");
    let db_path = codex_state::state_db_path(test.config.sqlite_home.as_path());

    for _ in 0..100 {
        if tokio::fs::try_exists(&db_path).await.unwrap_or(false) {
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    let db = test.codex.state_db().expect("state db enabled");
    assert!(
        !rollout_path.exists(),
        "fresh thread rollout should not be materialized before first user message"
    );

    let initial_metadata = db.get_thread(thread_id).await?;
    assert!(
        initial_metadata.is_none(),
        "fresh thread should not be recorded in state db before first user message"
    );

    test.submit_turn("materialize rollout").await?;

    let mut metadata = None;
    for _ in 0..100 {
        metadata = db.get_thread(thread_id).await?;
        if metadata.is_some() {
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    let metadata = metadata.expect("thread should exist in state db");
    assert_eq!(metadata.id, thread_id);
    assert_eq!(metadata.rollout_path, rollout_path.to_path_buf());
    assert!(
        rollout_path.exists(),
        "rollout should be materialized after first user message"
    );

    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn resume_restores_dynamic_tools_from_rollout_with_sqlite_enabled() -> Result<()> {
    let server = start_mock_server().await;
    let mock = mount_sse_sequence(
        &server,
        vec![
            responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
            responses::sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
        ],
    )
    .await;

    let namespace = "resume_tools";
    let namespace_description = "Tools available after resume.";
    let tool_name = "resume_lookup";
    let tool_description = "Look up a value after resume.";
    let input_schema = json!({
        "type": "object",
        "properties": { "query": { "type": "string" } },
        "required": ["query"],
        "additionalProperties": false,
    });
    let dynamic_tool = DynamicToolSpec::Namespace(DynamicToolNamespaceSpec {
        name: namespace.to_string(),
        description: namespace_description.to_string(),
        tools: vec![DynamicToolNamespaceTool::Function(
            DynamicToolFunctionSpec {
                name: tool_name.to_string(),
                description: tool_description.to_string(),
                input_schema: input_schema.clone(),
                defer_loading: false,
            },
        )],
    });
    let mut builder = test_codex().with_config(|config| {
        config
            .features
            .enable(Feature::Sqlite)
            .expect("test config should allow feature update");
    });
    let base_test = builder.build(&server).await?;
    let started = base_test
        .thread_manager
        .start_thread_with_tools(base_test.config.clone(), vec![dynamic_tool])
        .await?;
    let rollout_path = started
        .session_configured
        .rollout_path
        .clone()
        .expect("rollout path");

    started
        .thread
        .submit(Op::UserInput {
            items: vec![UserInput::Text {
                text: "persist this thread".to_string(),
                text_elements: Vec::new(),
            }],
            final_output_json_schema: None,
            responsesapi_client_metadata: None,
            additional_context: Default::default(),
            thread_settings: Default::default(),
        })
        .await?;
    wait_for_event(&started.thread, |event| {
        matches!(event, EventMsg::TurnComplete(_))
    })
    .await;

    let mut resume_builder = test_codex().with_config(|config| {
        config
            .features
            .enable(Feature::Sqlite)
            .expect("test config should allow feature update");
    });
    let resumed = resume_builder
        .resume(&server, base_test.home.clone(), rollout_path)
        .await?;
    resumed.submit_turn("use the restored tool").await?;

    let requests = mock.requests();
    assert_eq!(requests.len(), 2);
    let resumed_body = requests[1].body_json();
    let tools = resumed_body
        .get("tools")
        .and_then(serde_json::Value::as_array)
        .expect("resumed request tools");
    let restored_namespace = tools
        .iter()
        .find(|tool| tool.get("name") == Some(&json!(namespace)))
        .expect("dynamic tool namespace should be restored from rollout metadata");
    assert_eq!(
        restored_namespace,
        &json!({
            "type": "namespace",
            "name": namespace,
            "description": namespace_description,
            "tools": [{
                "type": "function",
                "name": tool_name,
                "description": tool_description,
                "strict": false,
                "parameters": input_schema,
            }],
        })
    );

    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn resume_restores_legacy_dynamic_tools_from_rollout_with_sqlite_enabled() -> Result<()> {
    let server = start_mock_server().await;
    let mock = mount_sse_sequence(
        &server,
        vec![
            responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
            responses::sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
        ],
    )
    .await;

    let namespace = "resume_tools";
    let tool_name = "resume_lookup";
    let tool_description = "Look up a value after resume.";
    let input_schema = json!({
        "type": "object",
        "properties": { "query": { "type": "string" } },
        "required": ["query"],
        "additionalProperties": false,
    });
    let mut builder = test_codex().with_config(|config| {
        config
            .features
            .enable(Feature::Sqlite)
            .expect("test config should allow feature update");
    });
    let base_test = builder.build(&server).await?;
    let started = base_test
        .thread_manager
        .start_thread_with_tools(base_test.config.clone(), Vec::new())
        .await?;
    let rollout_path = started
        .session_configured
        .rollout_path
        .clone()
        .expect("rollout path");

    started
        .thread
        .submit(Op::UserInput {
            items: vec![UserInput::Text {
                text: "persist this thread".to_string(),
                text_elements: Vec::new(),
            }],
            final_output_json_schema: None,
            responsesapi_client_metadata: None,
            additional_context: Default::default(),
            thread_settings: Default::default(),
        })
        .await?;
    wait_for_event(&started.thread, |event| {
        matches!(event, EventMsg::TurnComplete(_))
    })
    .await;
    started.thread.submit(Op::Shutdown).await?;
    wait_for_event(&started.thread, |event| {
        matches!(event, EventMsg::ShutdownComplete)
    })
    .await;

    let mut rollout_lines = fs::read_to_string(&rollout_path)?
        .lines()
        .map(serde_json::from_str::<serde_json::Value>)
        .collect::<serde_json::Result<Vec<_>>>()?;
    rollout_lines.first_mut().expect("session metadata line")["payload"]["dynamic_tools"] = json!([{
        "namespace": namespace,
        "name": tool_name,
        "description": tool_description,
        "inputSchema": input_schema,
        "exposeToContext": true,
    }]);
    let rollout = rollout_lines
        .iter()
        .map(serde_json::to_string)
        .collect::<serde_json::Result<Vec<_>>>()?
        .join("\n");
    fs::write(&rollout_path, format!("{rollout}\n"))?;

    let mut resume_builder = test_codex().with_config(|config| {
        config
            .features
            .enable(Feature::Sqlite)
            .expect("test config should allow feature update");
    });
    let resumed = resume_builder
        .resume(&server, base_test.home.clone(), rollout_path)
        .await?;
    resumed.submit_turn("use the restored tool").await?;

    let requests = mock.requests();
    assert_eq!(requests.len(), 2);
    let resumed_body = requests[1].body_json();
    let tools = resumed_body
        .get("tools")
        .and_then(serde_json::Value::as_array)
        .expect("resumed request tools");
    let restored_namespace = tools
        .iter()
        .find(|tool| tool.get("name") == Some(&json!(namespace)))
        .expect("dynamic tool namespace should be restored from rollout metadata");
    assert_eq!(
        restored_namespace,
        &json!({
            "type": "namespace",
            "name": namespace,
            "description": "Tools in the resume_tools namespace.",
            "tools": [{
                "type": "function",
                "name": tool_name,
                "description": tool_description,
                "strict": false,
                "parameters": input_schema,
            }],
        })
    );

    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn backfill_scans_existing_rollouts() -> Result<()> {
    let server = start_mock_server().await;

    let uuid = Uuid::now_v7();
    let thread_id = ThreadId::from_string(&uuid.to_string())?;
    let rollout_rel_path = format!("sessions/2026/01/27/rollout-2026-01-27T12-00-00-{uuid}.jsonl");
    let rollout_rel_path_for_hook = rollout_rel_path.clone();

    let mut builder = test_codex()
        .with_pre_build_hook(move |codex_home| {
            let rollout_path = codex_home.join(&rollout_rel_path_for_hook);
            let parent = rollout_path
                .parent()
                .expect("rollout path should have parent");
            fs::create_dir_all(parent).expect("should create rollout directory");
            let session_meta_line = SessionMetaLine {
                meta: SessionMeta {
                    session_id: thread_id.into(),
                    id: thread_id,
                    forked_from_id: None,
                    parent_thread_id: None,
                    timestamp: "2026-01-27T12:00:00Z".to_string(),
                    cwd: codex_home.to_path_buf(),
                    originator: "test".to_string(),
                    cli_version: "test".to_string(),
                    source: SessionSource::default(),
                    thread_source: None,
                    agent_path: None,
                    agent_nickname: None,
                    agent_role: None,
                    model_provider: None,
                    base_instructions: None,
                    dynamic_tools: None,
                    selected_capability_roots: Vec::new(),
                    memory_mode: None,
                    history_mode: Default::default(),
                    multi_agent_version: None,
                    context_window: None,
                },
                git: None,
            };

            let lines = [
                RolloutLine {
                    timestamp: "2026-01-27T12:00:00Z".to_string(),
                    item: RolloutItem::SessionMeta(session_meta_line),
                },
                RolloutLine {
                    timestamp: "2026-01-27T12:00:01Z".to_string(),
                    item: RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
                        client_id: None,
                        message: "hello from backfill".to_string(),
                        images: None,
                        local_images: Vec::new(),
                        text_elements: Vec::new(),
                        ..Default::default()
                    })),
                },
            ];

            let jsonl = lines
                .iter()
                .map(|line| serde_json::to_string(line).expect("rollout line should serialize"))
                .collect::<Vec<_>>()
                .join("\n");
            fs::write(&rollout_path, format!("{jsonl}\n")).expect("should write rollout file");
        })
        .with_config(|config| {
            config
                .features
                .enable(Feature::Sqlite)
                .expect("test config should allow feature update");
        });

    let test = builder.build(&server).await?;

    let db_path = codex_state::state_db_path(test.config.sqlite_home.as_path());
    let rollout_path = test.config.codex_home.join(&rollout_rel_path);
    let default_provider = test.config.model_provider_id.clone();

    for _ in 0..20 {
        if tokio::fs::try_exists(&db_path).await.unwrap_or(false) {
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    let db = test.codex.state_db().expect("state db enabled");

    let mut metadata = None;
    for _ in 0..40 {
        metadata = db.get_thread(thread_id).await?;
        if metadata.is_some() {
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    let metadata = metadata.expect("backfilled thread should exist in state db");
    assert_eq!(metadata.id, thread_id);
    assert_eq!(metadata.rollout_path, rollout_path.to_path_buf());
    assert_eq!(metadata.model_provider, default_provider);
    assert!(metadata.first_user_message.is_some());

    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn user_messages_persist_in_state_db() -> Result<()> {
    let server = start_mock_server().await;
    mount_sse_sequence(
        &server,
        vec![
            responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
            responses::sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
        ],
    )
    .await;

    let mut builder = test_codex().with_config(|config| {
        config
            .features
            .enable(Feature::Sqlite)
            .expect("test config should allow feature update");
    });
    let test = builder.build(&server).await?;

    let db_path = codex_state::state_db_path(test.config.sqlite_home.as_path());
    for _ in 0..100 {
        if tokio::fs::try_exists(&db_path).await.unwrap_or(false) {
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    test.submit_turn("hello from sqlite").await?;
    test.submit_turn("another message").await?;

    let db = test.codex.state_db().expect("state db enabled");
    let thread_id = test.session_configured.thread_id;

    let mut metadata = None;
    for _ in 0..100 {
        metadata = db.get_thread(thread_id).await?;
        if metadata
            .as_ref()
            .map(|entry| entry.first_user_message.is_some())
            .unwrap_or(false)
        {
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    let metadata = metadata.expect("thread should exist in state db");
    assert!(metadata.first_user_message.is_some());

    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn web_search_marks_thread_memory_mode_polluted_when_configured() -> Result<()> {
    let server = start_mock_server().await;
    mount_sse_sequence(
        &server,
        vec![responses::sse(vec![
            ev_response_created("resp-1"),
            ev_web_search_call_done("ws-1", "completed", "weather seattle"),
            ev_completed("resp-1"),
        ])],
    )
    .await;

    let mut builder = test_codex().with_config(|config| {
        config
            .features
            .enable(Feature::Sqlite)
            .expect("test config should allow feature update");
        config.memories.disable_on_external_context = true;
    });
    let test = builder.build(&server).await?;
    let db = test.codex.state_db().expect("state db enabled");
    let thread_id = test.session_configured.thread_id;

    test.submit_turn("search the web").await?;

    let mut memory_mode = None;
    for _ in 0..100 {
        memory_mode = db.get_thread_memory_mode(thread_id).await?;
        if memory_mode.as_deref() == Some("polluted") {
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    assert_eq!(memory_mode.as_deref(), Some("polluted"));
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn standalone_web_search_marks_thread_memory_mode_polluted_when_configured() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let server = start_mock_server().await;
    Mock::given(method("POST"))
        .and(path("/v1/alpha/search"))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
            "output": "Search result",
        })))
        .expect(1)
        .mount(&server)
        .await;
    mount_sse_sequence(
        &server,
        vec![
            responses::sse(vec![
                ev_response_created("resp-1"),
                responses::ev_function_call_with_namespace(
                    "web-run-1",
                    "web",
                    "run",
                    &json!({
                        "search_query": [{"q": "standalone web search"}],
                    })
                    .to_string(),
                ),
                ev_completed("resp-1"),
            ]),
            responses::sse(vec![
                responses::ev_assistant_message("msg-1", "done"),
                ev_completed("resp-2"),
            ]),
        ],
    )
    .await;

    let auth = CodexAuth::from_api_key("dummy");
    let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone());
    let mut extension_builder = ExtensionRegistryBuilder::<Config>::new();
    install_web_search_extension(&mut extension_builder, auth_manager);
    let mut builder = test_codex()
        .with_auth(auth)
        .with_extensions(Arc::new(extension_builder.build()))
        .with_config(|config| {
            config
                .features
                .enable(Feature::Sqlite)
                .expect("test config should allow feature update");
            config
                .features
                .enable(Feature::StandaloneWebSearch)
                .expect("standalone web search should be enabled");
            config.memories.disable_on_external_context = true;
            config
                .web_search_mode
                .set(WebSearchMode::Live)
                .expect("web search mode should be accepted");
        });
    let test = builder.build(&server).await?;
    let db = test.codex.state_db().expect("state db enabled");
    let thread_id = test.session_configured.thread_id;

    test.submit_turn("search the web").await?;

    let mut memory_mode = None;
    for _ in 0..100 {
        memory_mode = db.get_thread_memory_mode(thread_id).await?;
        if memory_mode.as_deref() == Some("polluted") {
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    assert_eq!(memory_mode.as_deref(), Some("polluted"));
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn mcp_call_marks_thread_memory_mode_polluted_when_configured() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let server = start_mock_server().await;
    let call_id = "call-123";
    let server_name = "rmcp";
    let namespace = format!("mcp__{server_name}");
    mount_sse_once(
        &server,
        responses::sse(vec![
            ev_response_created("resp-1"),
            responses::ev_function_call_with_namespace(
                call_id,
                &namespace,
                "echo",
                "{\"message\":\"ping\"}",
            ),
            ev_completed("resp-1"),
        ]),
    )
    .await;
    mount_sse_once(
        &server,
        responses::sse(vec![
            responses::ev_assistant_message("msg-1", "rmcp echo tool completed."),
            ev_completed("resp-2"),
        ]),
    )
    .await;

    let rmcp_test_server_bin = stdio_server_bin()?;
    let mut builder = test_codex().with_config(move |config| {
        config
            .features
            .enable(Feature::Sqlite)
            .expect("test config should allow feature update");
        config.memories.disable_on_external_context = true;

        let mut servers = config.mcp_servers.get().clone();
        servers.insert(
            server_name.to_string(),
            McpServerConfig {
                auth: Default::default(),
                transport: McpServerTransportConfig::Stdio {
                    command: rmcp_test_server_bin,
                    args: Vec::new(),
                    env: Some(HashMap::from([(
                        "MCP_TEST_VALUE".to_string(),
                        "propagated-env".to_string(),
                    )])),
                    env_vars: Vec::new(),
                    cwd: None,
                },
                environment_id: "local".to_string(),
                enabled: true,
                required: false,
                supports_parallel_tool_calls: false,
                disabled_reason: None,
                startup_timeout_sec: Some(Duration::from_secs(10)),
                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(),
            },
        );
        config
            .mcp_servers
            .set(servers)
            .expect("test mcp servers should accept any configuration");
    });
    let test = builder.build(&server).await?;
    wait_for_mcp_server(&test.codex, server_name).await?;
    let db = test.codex.state_db().expect("state db enabled");
    let thread_id = test.session_configured.thread_id;
    let cwd = test.config.cwd.clone();
    let (sandbox_policy, permission_profile) =
        turn_permission_fields(PermissionProfile::read_only(), cwd.as_path());

    test.codex
        .submit(Op::UserInput {
            items: vec![UserInput::Text {
                text: "call the rmcp echo tool".to_string(),
                text_elements: Vec::new(),
            }],
            final_output_json_schema: None,
            responsesapi_client_metadata: None,
            additional_context: Default::default(),
            thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {
                environments: Some(local_selections(cwd)),
                approval_policy: Some(AskForApproval::Never),
                sandbox_policy: Some(sandbox_policy),
                permission_profile,
                collaboration_mode: Some(codex_protocol::config_types::CollaborationMode {
                    mode: codex_protocol::config_types::ModeKind::Default,
                    settings: codex_protocol::config_types::Settings {
                        model: test.session_configured.model.clone(),
                        reasoning_effort: None,
                        developer_instructions: None,
                    },
                }),
                ..Default::default()
            },
        })
        .await?;
    wait_for_event(&test.codex, |event| {
        matches!(event, EventMsg::McpToolCallEnd(_))
    })
    .await;
    wait_for_event_match(&test.codex, |event| match event {
        EventMsg::Error(err) => Some(Err(anyhow::anyhow!(err.message.clone()))),
        EventMsg::TurnComplete(_) => Some(Ok(())),
        _ => None,
    })
    .await?;

    let mut memory_mode = None;
    for _ in 0..100 {
        memory_mode = db.get_thread_memory_mode(thread_id).await?;
        if memory_mode.as_deref() == Some("polluted") {
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    assert_eq!(memory_mode.as_deref(), Some("polluted"));
    Ok(())
}

#[tokio::test(flavor = "current_thread")]
async fn tool_call_logs_include_thread_id() -> Result<()> {
    let server = start_mock_server().await;
    let call_id = "call-1";
    let args = json!({
        "command": "echo hello",
        "timeout_ms": 1_000,
        "login": false,
    });
    let args_json = serde_json::to_string(&args)?;
    mount_sse_sequence(
        &server,
        vec![
            responses::sse(vec![
                ev_response_created("resp-1"),
                ev_function_call(call_id, "shell_command", &args_json),
                ev_completed("resp-1"),
            ]),
            responses::sse(vec![ev_completed("resp-2")]),
        ],
    )
    .await;

    let mut builder = test_codex().with_config(|config| {
        config
            .features
            .enable(Feature::Sqlite)
            .expect("test config should allow feature update");
    });
    let test = builder.build(&server).await?;
    let db = test.codex.state_db().expect("state db enabled");
    let expected_thread_id = test.session_configured.thread_id.to_string();

    test.submit_turn("run a shell command").await?;

    let log_db_layer = codex_state::log_db::start(db.clone());
    let subscriber = tracing_subscriber::registry().with(log_db_layer.clone());
    let dispatch = tracing::Dispatch::new(subscriber);
    tracing::dispatcher::with_default(&dispatch, || {
        let span = tracing::info_span!("test_log_span", thread_id = %expected_thread_id);
        let _entered = span.enter();
        tracing::info!("ToolCall: shell_command {{\"command\":\"echo hello\"}}");
    });
    log_db_layer.flush().await;

    let mut found = None;
    for _ in 0..80 {
        let query = codex_state::LogQuery {
            descending: true,
            limit: Some(20),
            ..Default::default()
        };
        let rows = db.query_logs(&query).await?;
        if let Some(row) = rows.into_iter().find(|row| {
            row.message
                .as_deref()
                .is_some_and(|m| m.contains("ToolCall:"))
        }) {
            let thread_id = row.thread_id;
            let message = row.message;
            found = Some((thread_id, message));
            break;
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }

    let (thread_id, message) = found.expect("expected ToolCall log row");
    assert_eq!(thread_id, Some(expected_thread_id));
    assert!(
        message
            .as_deref()
            .is_some_and(|text| text.contains("ToolCall:")),
        "expected ToolCall message, got {message:?}"
    );

    Ok(())
}