mirror of
https://github.com/openai/codex.git
synced 2026-04-24 14:45:27 +00:00
ec8d4bfc77
## Problem When a user resumed or forked a session, the TUI could render the restored thread history immediately, but it did not receive token usage until a later model turn emitted a fresh usage event. That left the context/status UI blank or stale during the exact window where the user expects resumed state to look complete. Core already reconstructed token usage from the rollout; the missing behavior was app-server lifecycle replay to the client that just attached. ## Mental model Token usage has two representations. The rollout is the durable source of historical `TokenCount` events, and the core session cache is the in-memory snapshot reconstructed from that rollout on resume or fork. App-server v2 clients do not read core state directly; they learn about usage through `thread/tokenUsage/updated`. The fix keeps those roles separate: core exposes the restored `TokenUsageInfo`, and app-server sends one targeted notification after a successful `thread/resume` or `thread/fork` response when that restored snapshot exists. This notification is not a new model event. It is a replay of already-persisted state for the client that just attached. That distinction matters because using the normal core event path here would risk duplicating `TokenCount` entries in the rollout and making future resumes count historical usage twice. ## Non-goals This change does not add a new protocol method or payload shape. It reuses the existing v2 `thread/tokenUsage/updated` notification and the TUI’s existing handler for that notification. This change does not alter how token usage is computed, accumulated, compacted, or written during turns. It only exposes the token usage that resume and fork reconstruction already restored. This change does not broadcast historical usage replay to every subscribed client. The replay is intentionally scoped to the connection that requested resume or fork so already-attached clients are not surprised by an old usage update while they may be rendering live activity. ## Tradeoffs Sending the usage notification after the JSON-RPC response preserves a clear lifecycle order: the client first receives the thread object, then receives restored usage for that thread. The tradeoff is that usage is still a notification rather than part of the `thread/resume` or `thread/fork` response. That keeps the protocol shape stable and avoids duplicating usage fields across response types, but clients must continue listening for notifications after receiving the response. The helper selects the latest non-in-progress turn id for the replayed usage notification. This is conservative because restored usage belongs to completed persisted accounting, not to newly attached in-flight work. The fallback to the last turn preserves a stable wire payload for unusual histories, but histories with no meaningful completed turn still have a weak attribution story. ## Architecture Core already seeds `Session` token state from the last persisted rollout `TokenCount` during `InitialHistory::Resumed` and `InitialHistory::Forked`. The new core accessor exposes the complete `TokenUsageInfo` through `CodexThread` without giving app-server direct session mutation authority. App-server calls that accessor from three lifecycle paths: cold `thread/resume`, running-thread resume/rejoin, and `thread/fork`. In each path, the server sends the normal response first, then calls a shared helper that converts core usage into `ThreadTokenUsageUpdatedNotification` and sends it only to the requesting connection. The tests build fake rollouts with a user turn plus a persisted token usage event. They then exercise `thread/resume` and `thread/fork` without starting another model turn, proving that restored usage arrives before any next-turn token event could be produced. ## Observability The primary debug path is the app-server JSON-RPC stream. After `thread/resume` or `thread/fork`, a client should see the response followed by `thread/tokenUsage/updated` when the source rollout includes token usage. If the notification is absent, check whether the rollout contains an `event_msg` payload of type `token_count`, whether core reconstruction seeded `Session::token_usage_info`, and whether the connection stayed attached long enough to receive the targeted notification. The notification is sent through the existing `OutgoingMessageSender::send_server_notification_to_connections` path, so existing app-server tracing around server notifications still applies. Because this is a replay, not a model turn event, debugging should start at the resume/fork handlers rather than the turn event translation in `bespoke_event_handling`. ## Tests The focused regression coverage is `cargo test -p codex-app-server emits_restored_token_usage`, which covers both resume and fork. The core reconstruction guard is `cargo test -p codex-core record_initial_history_seeds_token_info_from_rollout`. Formatting and lint/fix passes were run with `just fmt`, `just fix -p codex-core`, and `just fix -p codex-app-server`. Full crate test runs surfaced pre-existing unrelated failures in command execution and plugin marketplace tests; the new token usage tests passed in focused runs and within the app-server suite before the unrelated command execution failure.
52 lines
2.2 KiB
Rust
52 lines
2.2 KiB
Rust
mod analytics_server;
|
|
mod auth_fixtures;
|
|
mod config;
|
|
mod mcp_process;
|
|
mod mock_model_server;
|
|
mod models_cache;
|
|
mod responses;
|
|
mod rollout;
|
|
|
|
pub use analytics_server::start_analytics_events_server;
|
|
pub use auth_fixtures::ChatGptAuthFixture;
|
|
pub use auth_fixtures::ChatGptIdTokenClaims;
|
|
pub use auth_fixtures::encode_id_token;
|
|
pub use auth_fixtures::write_chatgpt_auth;
|
|
use codex_app_server_protocol::JSONRPCResponse;
|
|
pub use config::write_mock_responses_config_toml;
|
|
pub use config::write_mock_responses_config_toml_with_chatgpt_base_url;
|
|
pub use core_test_support::PathBufExt;
|
|
pub use core_test_support::format_with_current_shell;
|
|
pub use core_test_support::format_with_current_shell_display;
|
|
pub use core_test_support::format_with_current_shell_display_non_login;
|
|
pub use core_test_support::format_with_current_shell_non_login;
|
|
pub use core_test_support::test_absolute_path;
|
|
pub use core_test_support::test_path_buf_with_windows;
|
|
pub use core_test_support::test_tmp_path;
|
|
pub use core_test_support::test_tmp_path_buf;
|
|
pub use mcp_process::DEFAULT_CLIENT_NAME;
|
|
pub use mcp_process::McpProcess;
|
|
pub use mock_model_server::create_mock_responses_server_repeating_assistant;
|
|
pub use mock_model_server::create_mock_responses_server_sequence;
|
|
pub use mock_model_server::create_mock_responses_server_sequence_unchecked;
|
|
pub use models_cache::write_models_cache;
|
|
pub use models_cache::write_models_cache_with_models;
|
|
pub use responses::create_apply_patch_sse_response;
|
|
pub use responses::create_exec_command_sse_response;
|
|
pub use responses::create_final_assistant_message_sse_response;
|
|
pub use responses::create_request_permissions_sse_response;
|
|
pub use responses::create_request_user_input_sse_response;
|
|
pub use responses::create_shell_command_sse_response;
|
|
pub use rollout::create_fake_rollout;
|
|
pub use rollout::create_fake_rollout_with_source;
|
|
pub use rollout::create_fake_rollout_with_text_elements;
|
|
pub use rollout::create_fake_rollout_with_token_usage;
|
|
pub use rollout::rollout_path;
|
|
use serde::de::DeserializeOwned;
|
|
|
|
pub fn to_response<T: DeserializeOwned>(response: JSONRPCResponse) -> anyhow::Result<T> {
|
|
let value = serde_json::to_value(response.result)?;
|
|
let codex_response = serde_json::from_value(value)?;
|
|
Ok(codex_response)
|
|
}
|