From d749c50f2d195585d96b66c3a416587acf8a2cc6 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Tue, 19 May 2026 15:33:45 +0100 Subject: [PATCH] Map token pressure at host wall --- epiphany-codex-bridge/src/coordinator.rs | 4 ++-- epiphany-codex-bridge/src/mutation_service.rs | 4 ++-- epiphany-codex-bridge/src/pressure.rs | 15 +++++---------- epiphany-codex-bridge/src/view.rs | 4 ++-- .../epiphany_automation.rs | 7 +++++-- .../epiphany_mutation_routes.rs | 4 +++- .../epiphany_read_routes.rs | 4 +++- .../epiphany_thread_host.rs | 13 +++++++++++++ 8 files changed, 35 insertions(+), 20 deletions(-) diff --git a/epiphany-codex-bridge/src/coordinator.rs b/epiphany-codex-bridge/src/coordinator.rs index 53100c8e..fef6b186 100644 --- a/epiphany-codex-bridge/src/coordinator.rs +++ b/epiphany-codex-bridge/src/coordinator.rs @@ -19,7 +19,6 @@ use codex_app_server_protocol::ThreadEpiphanyRoleStatus; use codex_app_server_protocol::ThreadEpiphanyRolesSource; use codex_app_server_protocol::ThreadEpiphanySceneAction; use codex_app_server_protocol::ThreadEpiphanyViewCoordinator; -use codex_protocol::protocol::TokenUsageInfo; use epiphany_core::EpiphanyCoordinatorAction as CoreEpiphanyCoordinatorAction; use epiphany_core::EpiphanyCoordinatorAutomationAction as CoreEpiphanyCoordinatorAutomationAction; use epiphany_core::EpiphanyCoordinatorCrrcRecommendation; @@ -45,6 +44,7 @@ use epiphany_core::EpiphanyRoleBoardJob; use epiphany_core::EpiphanyRoleBoardJobStatus; use epiphany_core::EpiphanyRoleBoardLane; use epiphany_core::EpiphanyRoleBoardPlanningSummary; +use epiphany_core::EpiphanyTokenUsageSnapshot; use epiphany_core::coordinator_automation_action; use epiphany_core::derive_role_board; use epiphany_core::recommend_coordinator_action; @@ -390,7 +390,7 @@ pub struct EpiphanyCoordinatorAutomationInput<'a> { pub state: &'a EpiphanyThreadState, pub retrieval_override: &'a EpiphanyRetrievalState, pub watcher_snapshot: EpiphanyFreshnessWatcherSnapshot<'a>, - pub token_usage_info: Option<&'a TokenUsageInfo>, + pub token_usage_info: Option<&'a EpiphanyTokenUsageSnapshot>, pub runtime_store_path: &'a Path, pub force_checkpoint_compaction: bool, } diff --git a/epiphany-codex-bridge/src/mutation_service.rs b/epiphany-codex-bridge/src/mutation_service.rs index b43bba68..f5dafc23 100644 --- a/epiphany-codex-bridge/src/mutation_service.rs +++ b/epiphany-codex-bridge/src/mutation_service.rs @@ -14,13 +14,13 @@ use codex_app_server_protocol::ThreadEpiphanyRoleId; use codex_app_server_protocol::ThreadEpiphanyStateUpdatedField; use codex_app_server_protocol::ThreadEpiphanyUpdatePatch; use codex_protocol::error::CodexErr; -use codex_protocol::protocol::TokenUsageInfo as CoreTokenUsageInfo; use epiphany_core::EpiphanyJobInterruptRequest; use epiphany_core::EpiphanyJobInterruptResult; use epiphany_core::EpiphanyJobLaunchRequest; use epiphany_core::EpiphanyJobLaunchResult; use epiphany_core::EpiphanyPromotionInput; use epiphany_core::EpiphanyStateUpdate; +use epiphany_core::EpiphanyTokenUsageSnapshot; use epiphany_core::RuntimeSpineHeartbeatJobOptions; use epiphany_core::RuntimeSpineHeartbeatLaunchPlanOptions; use epiphany_core::apply_epiphany_state_update; @@ -573,7 +573,7 @@ pub async fn launch_thread_epiphany_reorient( state: Option<&EpiphanyThreadState>, retrieval_override: Option<&EpiphanyRetrievalState>, watcher_snapshot: Option>, - token_usage_info: Option<&CoreTokenUsageInfo>, + token_usage_info: Option<&EpiphanyTokenUsageSnapshot>, ) -> Result { let (state_revision, retrieval, graph, watcher) = map_epiphany_freshness(state, retrieval_override, watcher_snapshot); diff --git a/epiphany-codex-bridge/src/pressure.rs b/epiphany-codex-bridge/src/pressure.rs index fcfc4dfd..7c9bd3fb 100644 --- a/epiphany-codex-bridge/src/pressure.rs +++ b/epiphany-codex-bridge/src/pressure.rs @@ -2,26 +2,21 @@ use codex_app_server_protocol::ThreadEpiphanyPressure; use codex_app_server_protocol::ThreadEpiphanyPressureBasis; use codex_app_server_protocol::ThreadEpiphanyPressureLevel; use codex_app_server_protocol::ThreadEpiphanyPressureStatus; -use codex_protocol::protocol::TokenUsageInfo as CoreTokenUsageInfo; use epiphany_core::EpiphanyPressure; use epiphany_core::EpiphanyPressureBasis as CoreEpiphanyPressureBasis; use epiphany_core::EpiphanyPressureLevel as CoreEpiphanyPressureLevel; use epiphany_core::EpiphanyPressureStatus as CoreEpiphanyPressureStatus; -use epiphany_core::EpiphanyTokenUsageSnapshot; +pub use epiphany_core::EpiphanyTokenUsageSnapshot; use epiphany_core::derive_pressure_view; use crate::launch::epiphany_agent_prompt_with_memory; use crate::launch::epiphany_specialist_prompt_config; use crate::launch::pressure_level_label; -pub fn map_epiphany_pressure(info: Option<&CoreTokenUsageInfo>) -> ThreadEpiphanyPressure { - let snapshot = info.map(|info| EpiphanyTokenUsageSnapshot { - total_tokens: info.total_token_usage.total_tokens, - last_turn_tokens: info.last_token_usage.total_tokens, - model_context_window: info.model_context_window, - model_auto_compact_token_limit: info.model_auto_compact_token_limit, - }); - map_core_epiphany_pressure(derive_pressure_view(snapshot.as_ref())) +pub fn map_epiphany_pressure( + snapshot: Option<&EpiphanyTokenUsageSnapshot>, +) -> ThreadEpiphanyPressure { + map_core_epiphany_pressure(derive_pressure_view(snapshot)) } fn map_core_epiphany_pressure(pressure: EpiphanyPressure) -> ThreadEpiphanyPressure { diff --git a/epiphany-codex-bridge/src/view.rs b/epiphany-codex-bridge/src/view.rs index 6873fe1f..650f8a7e 100644 --- a/epiphany-codex-bridge/src/view.rs +++ b/epiphany-codex-bridge/src/view.rs @@ -1,10 +1,10 @@ use std::path::Path; use codex_app_server_protocol::*; -use codex_protocol::protocol::TokenUsageInfo as CoreTokenUsageInfo; use epiphany_core::EpiphanyDistillInput; use epiphany_core::EpiphanyMapProposalInput; use epiphany_core::EpiphanySceneInput; +use epiphany_core::EpiphanyTokenUsageSnapshot; use epiphany_core::derive_scene; use epiphany_core::distill_observation; use epiphany_core::propose_map_update; @@ -152,7 +152,7 @@ pub struct EpiphanyViewResponseInput<'a> { pub state: Option<&'a EpiphanyThreadState>, pub retrieval_override: Option<&'a EpiphanyRetrievalState>, pub watcher_snapshot: Option>, - pub token_usage_info: Option<&'a CoreTokenUsageInfo>, + pub token_usage_info: Option<&'a EpiphanyTokenUsageSnapshot>, pub runtime_store_path: Option<&'a Path>, } diff --git a/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_automation.rs b/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_automation.rs index f41e608c..89d3af2e 100644 --- a/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_automation.rs +++ b/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_automation.rs @@ -25,6 +25,7 @@ use tracing::warn; use super::epiphany_thread_host::EpiphanyCodexThreadHost; use super::epiphany_thread_host::client_visible_live_thread_epiphany_state; +use super::epiphany_thread_host::epiphany_token_usage_snapshot; use crate::outgoing_message::ThreadScopedOutgoingMessageSender; use crate::thread_state::ThreadState; @@ -51,6 +52,7 @@ pub(crate) async fn maybe_run_epiphany_coordinator_automation_for_turn_boundary( .snapshot(&thread_id_text) .await; let token_usage_info = thread.token_usage_info().await; + let token_usage_snapshot = epiphany_token_usage_snapshot(token_usage_info.as_ref()); let runtime_store_path = thread.epiphany_runtime_spine_store_path().await; let verdict = select_epiphany_coordinator_automation(EpiphanyCoordinatorAutomationInput { @@ -58,7 +60,7 @@ pub(crate) async fn maybe_run_epiphany_coordinator_automation_for_turn_boundary( state: &state, retrieval_override: &retrieval_override, watcher_snapshot: epiphany_freshness_watcher_snapshot(&watcher_snapshot), - token_usage_info: token_usage_info.as_ref(), + token_usage_info: token_usage_snapshot.as_ref(), runtime_store_path: runtime_store_path.as_path(), force_checkpoint_compaction, }) @@ -112,7 +114,8 @@ pub(crate) async fn maybe_run_epiphany_pre_compaction_checkpoint_intervention_fo token_usage_info: Option, thread_state: &Arc>, ) { - let pressure = map_epiphany_pressure(token_usage_info.as_ref()); + let token_usage_snapshot = epiphany_token_usage_snapshot(token_usage_info.as_ref()); + let pressure = map_epiphany_pressure(token_usage_snapshot.as_ref()); if !should_run_epiphany_pre_compaction_checkpoint_intervention(&pressure) { return; } diff --git a/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_mutation_routes.rs b/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_mutation_routes.rs index 78dba2a3..6c44b6b9 100644 --- a/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_mutation_routes.rs +++ b/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_mutation_routes.rs @@ -26,6 +26,7 @@ use super::CodexMessageProcessor; use super::ConnectionRequestId; use super::ThreadReadViewError; use super::epiphany_thread_host::EpiphanyCodexThreadHost; +use super::epiphany_thread_host::epiphany_token_usage_snapshot; impl CodexMessageProcessor { async fn load_epiphany_thread( @@ -259,6 +260,7 @@ impl CodexMessageProcessor { .await; let token_usage_info = loaded_thread.token_usage_info().await; let host = EpiphanyCodexThreadHost::new(loaded_thread.as_ref()); + let token_usage_snapshot = epiphany_token_usage_snapshot(token_usage_info.as_ref()); let applied = match launch_thread_epiphany_reorient( &host, &thread_id, @@ -267,7 +269,7 @@ impl CodexMessageProcessor { thread.epiphany_state.as_ref(), Some(&retrieval_override), Some(epiphany_freshness_watcher_snapshot(&watcher_snapshot)), - token_usage_info.as_ref(), + token_usage_snapshot.as_ref(), ) .await { diff --git a/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_read_routes.rs b/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_read_routes.rs index c8268c13..212ca125 100644 --- a/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_read_routes.rs +++ b/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_read_routes.rs @@ -30,6 +30,7 @@ use epiphany_codex_bridge::view::map_epiphany_view_response; use super::CodexMessageProcessor; use super::ConnectionRequestId; use super::ThreadReadViewError; +use super::epiphany_thread_host::epiphany_token_usage_snapshot; use super::latest_token_usage_info_from_rollout_path; impl CodexMessageProcessor { @@ -124,6 +125,7 @@ impl CodexMessageProcessor { } else { None }; + let token_usage_snapshot = epiphany_token_usage_snapshot(token_usage_info.as_ref()); let response = map_epiphany_view_response(EpiphanyViewResponseInput { thread_id: thread_id.clone(), lenses, @@ -133,7 +135,7 @@ impl CodexMessageProcessor { watcher_snapshot: watcher_snapshot .as_ref() .map(epiphany_freshness_watcher_snapshot), - token_usage_info: token_usage_info.as_ref(), + token_usage_info: token_usage_snapshot.as_ref(), runtime_store_path: runtime_store_path.as_deref(), }) .await; diff --git a/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_thread_host.rs b/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_thread_host.rs index d69f7773..5385afe5 100644 --- a/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_thread_host.rs +++ b/vendor/codex/codex-rs/app-server/src/codex_message_processor/epiphany_thread_host.rs @@ -6,7 +6,9 @@ use codex_core::latest_epiphany_state_from_rollout_items; use codex_protocol::error::CodexErr; use codex_protocol::protocol::EpiphanyThreadState; use codex_protocol::protocol::InitialHistory; +use codex_protocol::protocol::TokenUsageInfo; use epiphany_codex_bridge::mutation_service::EpiphanyMutationHost; +use epiphany_codex_bridge::pressure::EpiphanyTokenUsageSnapshot; use epiphany_codex_bridge::state::client_visible_epiphany_state_for_paths; use epiphany_codex_bridge::state::thread_state_mirror_id_from_rollout_path; @@ -91,3 +93,14 @@ pub(super) async fn client_visible_live_thread_epiphany_state( ) -> EpiphanyThreadState { live_thread_epiphany_state(thread).await.unwrap_or(fallback) } + +pub(super) fn epiphany_token_usage_snapshot( + info: Option<&TokenUsageInfo>, +) -> Option { + info.map(|info| EpiphanyTokenUsageSnapshot { + total_tokens: info.total_token_usage.total_tokens, + last_turn_tokens: info.last_token_usage.total_tokens, + model_context_window: info.model_context_window, + model_auto_compact_token_limit: info.model_auto_compact_token_limit, + }) +}