From ea1a17b917f8feb57735aabc1f8a46dc8886f555 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9D=92=E9=9B=B2?= <137844255@qq.com> Date: Thu, 6 Aug 2026 14:08:01 +0800 Subject: [PATCH] fix: restore context for resumed sessions --- crates/orca-runtime/src/runtime_host.rs | 87 ++++++++++++++++++- crates/orca-runtime/src/thread_store.rs | 4 +- .../orca-runtime/src/thread_store/writer.rs | 34 ++++++++ crates/orca-tui/src/app.rs | 50 ++++++++++- crates/orca-tui/src/surface_projection.rs | 3 +- crates/orca-tui/src/types.rs | 31 ++++++- crates/orca-tui/src/ui.rs | 37 +++++++- 7 files changed, 236 insertions(+), 10 deletions(-) diff --git a/crates/orca-runtime/src/runtime_host.rs b/crates/orca-runtime/src/runtime_host.rs index 3e131c36..6b050ff2 100644 --- a/crates/orca-runtime/src/runtime_host.rs +++ b/crates/orca-runtime/src/runtime_host.rs @@ -7575,6 +7575,12 @@ fn bootstrap_recorded_surface( incarnation_bytes[8] = 0x80 | (incarnation_bytes[8] & 0x3f); let incarnation = surface::SurfaceIncarnation::try_from_bytes(incarnation_bytes) .expect("normalized thread UUID is v7"); + let restored_context_tokens = + crate::thread_store::read_latest_context_tokens(&path).map_err(|error| { + RuntimeHostError::ThreadStartFailed { + message: format!("failed to restore latest provider context: {error}"), + } + })?; let owner_lease = surface_owner.owner_lease; let current_owner_epoch = surface::ThreadOwnerEpoch::new(owner_lease.owner_epoch()); @@ -7586,6 +7592,7 @@ fn bootstrap_recorded_surface( surface::ThreadPersistence::RecordedCatalogued, config, title, + restored_context_tokens, )?; let ledger = surface::JsonlSurfaceCommitLedger::new(path, snapshot.cursor.clone()); let mut coordinator = surface::RuntimeCommitCoordinator::recover_with_owned_lease( @@ -7755,6 +7762,7 @@ fn bootstrap_ephemeral_surface( surface::ThreadPersistence::EphemeralNonCataloguedOneShot { close_after }, config, title, + None, )?; let ledger = surface::RuntimeSurfaceCommitLedger::Ephemeral( surface::InMemorySurfaceCommitLedger::new(snapshot.cursor.clone()), @@ -7888,6 +7896,7 @@ fn initial_surface_snapshot( persistence: surface::ThreadPersistence, config: &RunConfig, title: &str, + restored_context_tokens: Option, ) -> Result { let cwd = config.cwd.clone().unwrap_or_else(|| { std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("/")) @@ -7999,6 +8008,11 @@ fn initial_surface_snapshot( let permission_rules_digest = surface_sha256( &serde_json::to_vec(&permission_rules).expect("surface permission rules are serializable"), ); + let context_limit_tokens = config + .model_runtime + .context_window + .unwrap_or_else(|| orca_core::model::max_context_tokens(config.model.as_deref())) + .max(1) as u64; let settings = surface::SurfaceRuntimeSettings { model: surface::NonEmptyText::try_new( config @@ -8083,8 +8097,8 @@ fn initial_surface_snapshot( }, context: surface::SurfaceContextSnapshot { revision: surface::ContextRevision::try_new(1).expect("one is a valid revision"), - used_tokens: 0, - limit_tokens: 128_000, + used_tokens: restored_context_tokens.unwrap_or(0), + limit_tokens: context_limit_tokens, compaction: surface::CompactionState::Idle, fragments: Vec::new(), provider_replay: surface::ProviderReplayHealth::None, @@ -50281,6 +50295,75 @@ mod tests { } } + #[test] + fn resumed_legacy_usage_restores_latest_provider_context() { + let _env = crate::history::lock_test_env(); + let home = tempfile::tempdir().unwrap(); + let _home = OrcaHomeRestore::set(home.path()); + let cwd = tempfile::tempdir().unwrap(); + + let host = RuntimeHost::start_with_executor(Arc::new(PanicExecutor)) + .expect("start legacy context fixture host"); + let thread = host + .start_thread( + surface_test_config(cwd.path().to_path_buf(), HistoryMode::Record), + "legacy context fixture", + ) + .expect("start legacy context fixture thread"); + let thread_id = thread.thread_id().to_string(); + let transcript_path = SessionStore::new() + .load_session(&thread_id) + .expect("load legacy context fixture") + .path; + thread + .shutdown() + .expect("shutdown legacy context fixture thread"); + host.shutdown() + .expect("shutdown legacy context fixture host"); + + let mut writer = crate::thread_store::SessionWriter::append_to_existing(transcript_path) + .expect("open legacy context fixture transcript"); + writer + .append_usage(UsageTotals { + input_tokens: 929_128, + output_tokens: 10_260, + cache_tokens: 893_696, + estimated_cost_usd: 0.063661744, + }) + .expect("append penultimate usage snapshot"); + writer + .append_usage(UsageTotals { + input_tokens: 970_611, + output_tokens: 10_627, + cache_tokens: 935_040, + estimated_cost_usd: 0.065860634, + }) + .expect("append latest usage snapshot"); + drop(writer); + + let resumed_host = RuntimeHost::start_with_executor(Arc::new(PanicExecutor)) + .expect("start resumed legacy context host"); + let resumed = resumed_host + .start_thread( + surface_test_config(cwd.path().to_path_buf(), HistoryMode::Resume(thread_id)), + "resume legacy context fixture", + ) + .expect("resume legacy context fixture"); + let snapshot = fresh_surface_attachment(&resumed.surface()) + .baseline + .snapshot; + + assert_eq!(snapshot.context.used_tokens, 41_483); + assert_eq!(snapshot.context.limit_tokens, 1_000_000); + + resumed + .shutdown() + .expect("shutdown resumed legacy context thread"); + resumed_host + .shutdown() + .expect("shutdown resumed legacy context host"); + } + #[test] fn background_continuation_reuses_pending_response_turn_identity() { let registry = TaskRegistry::new("background-continuation-identity".to_string()); diff --git a/crates/orca-runtime/src/thread_store.rs b/crates/orca-runtime/src/thread_store.rs index 13894c04..3ee46cce 100644 --- a/crates/orca-runtime/src/thread_store.rs +++ b/crates/orca-runtime/src/thread_store.rs @@ -31,7 +31,9 @@ pub use types::{ TurnItemsView, }; pub use writer::SessionWriter; -pub(crate) use writer::{read_manual_compaction_snapshot, redact_sensitive_text}; +pub(crate) use writer::{ + read_latest_context_tokens, read_manual_compaction_snapshot, redact_sensitive_text, +}; pub(crate) fn resume_conversation( transcript: &SessionTranscript, diff --git a/crates/orca-runtime/src/thread_store/writer.rs b/crates/orca-runtime/src/thread_store/writer.rs index 771e731f..ffeb8dac 100644 --- a/crates/orca-runtime/src/thread_store/writer.rs +++ b/crates/orca-runtime/src/thread_store/writer.rs @@ -431,6 +431,27 @@ pub(crate) fn read_transcript(path: &Path) -> io::Result { }) } +pub(crate) fn read_latest_context_tokens(path: &Path) -> io::Result> { + let mut previous_input_tokens = 0; + let mut latest_context_tokens = None; + for record in read_records(path)? { + match record { + SessionRecord::Usage(usage) => { + latest_context_tokens = Some( + usage + .input_tokens + .checked_sub(previous_input_tokens) + .unwrap_or(usage.input_tokens), + ); + previous_input_tokens = usage.input_tokens; + } + SessionRecord::UsageBaseline(_) => previous_input_tokens = 0, + _ => {} + } + } + Ok(latest_context_tokens) +} + fn redact_session_record(record: &SessionRecord) -> SessionRecord { let mut redacted = record.clone(); match &mut redacted { @@ -1304,6 +1325,19 @@ mod tests { assert_eq!(values, vec![8_000]); } + #[test] + fn latest_context_tokens_are_the_delta_between_usage_snapshots() { + let (_directory, path, mut writer) = new_transcript(); + writer + .append_usage(usage(929_128, 10_260, 893_696, 0.06)) + .expect("write penultimate usage snapshot"); + writer + .append_usage(usage(970_611, 10_627, 935_040, 0.06)) + .expect("write latest usage snapshot"); + + assert_eq!(read_latest_context_tokens(&path).unwrap(), Some(41_483)); + } + #[test] fn thread_store_append_probe_child() { let Some(path) = std::env::var_os(APPEND_PROBE_PATH).map(PathBuf::from) else { diff --git a/crates/orca-tui/src/app.rs b/crates/orca-tui/src/app.rs index b8de71e3..646dd99c 100644 --- a/crates/orca-tui/src/app.rs +++ b/crates/orca-tui/src/app.rs @@ -58,7 +58,8 @@ use crate::surface_actions::{TuiHostActions, TuiSurfaceActions}; use crate::terminal_presentation::{TerminalPresentation, TerminalPresentationProfile}; use crate::theme::Theme; use crate::types::{ - AppState, AppStatus, AttachedTuiEvent, ChatMessage, SessionAttachmentId, TuiEvent, UserAction, + AppState, AppStatus, AttachedTuiEvent, ChatMessage, SessionAttachmentId, + SurfaceProjectionState, TuiEvent, UserAction, }; use crate::ui; use crate::vim::{PendingInsertEscapeFlow, VimState}; @@ -5350,6 +5351,48 @@ done }); } + #[test] + fn resumed_legacy_usage_projects_context_before_next_turn() { + with_orca_home(|home| { + let mut writer = + history::SessionWriter::start(home, "mock", Some("auto".to_string()), "context") + .expect("create legacy context transcript"); + writer + .append_usage(orca_core::cost_types::UsageTotals { + input_tokens: 929_128, + output_tokens: 10_260, + cache_tokens: 893_696, + estimated_cost_usd: 0.063661744, + }) + .expect("append penultimate usage snapshot"); + writer + .append_usage(orca_core::cost_types::UsageTotals { + input_tokens: 970_611, + output_tokens: 10_627, + cache_tokens: 935_040, + estimated_cost_usd: 0.065860634, + }) + .expect("append latest usage snapshot"); + drop(writer); + let session_id = history::load_session("latest") + .expect("load legacy context transcript") + .meta + .session_id; + + let mut harness = + HostedTuiHarness::start(test_config(HistoryMode::Resume(session_id)), None); + let event = + harness.recv_until(|event| matches!(event, TuiEvent::SurfaceProjectionSynced(_))); + let TuiEvent::SurfaceProjectionSynced(projection) = event else { + unreachable!("predicate accepted only a surface projection") + }; + assert_eq!(projection.context_used_tokens, 41_483); + assert_eq!(projection.context_limit_tokens, 1_000_000); + + harness.shutdown(); + }); + } + #[test] fn goal_auto_continuation_pauses_after_three_no_progress_turns() { with_orca_home(|_home| { @@ -9170,6 +9213,11 @@ fn emit_typed_history_snapshot( label: label.to_string(), }) .map_err(|error| error.to_string())?; + event_tx + .send(TuiEvent::SurfaceProjectionSynced(Box::new( + SurfaceProjectionState::from_surface_snapshot(&snapshot), + ))) + .map_err(|error| error.to_string())?; let tasks = crate::surface_projection::workflow_task_summaries(&snapshot); if !tasks.is_empty() { event_tx diff --git a/crates/orca-tui/src/surface_projection.rs b/crates/orca-tui/src/surface_projection.rs index 62ccc2d0..545e2b9f 100644 --- a/crates/orca-tui/src/surface_projection.rs +++ b/crates/orca-tui/src/surface_projection.rs @@ -165,12 +165,13 @@ pub(crate) enum SurfaceProjectionError { pub(crate) type TuiStreamDeliveryWatermark = BTreeMap; impl SurfaceProjectionState { - fn from_surface_snapshot(snapshot: &orca_runtime::surface::SurfaceSnapshot) -> Self { + pub(crate) fn from_surface_snapshot(snapshot: &orca_runtime::surface::SurfaceSnapshot) -> Self { Self { session_id: surface_thread_id_text(&snapshot.thread.thread_id), title: snapshot.thread.title.as_str().to_string(), usage_revision: snapshot.usage.revision.get(), usage: core_usage_totals(&snapshot.usage.thread_total), + context_revision: snapshot.context.revision.get(), context_used_tokens: usize::try_from(snapshot.context.used_tokens) .unwrap_or(usize::MAX), context_limit_tokens: usize::try_from(snapshot.context.limit_tokens) diff --git a/crates/orca-tui/src/types.rs b/crates/orca-tui/src/types.rs index 528086f1..b79f6e9a 100644 --- a/crates/orca-tui/src/types.rs +++ b/crates/orca-tui/src/types.rs @@ -326,6 +326,7 @@ pub struct SurfaceProjectionState { pub(crate) title: String, pub(crate) usage_revision: u64, pub(crate) usage: UsageTotals, + pub(crate) context_revision: u64, pub(crate) context_used_tokens: usize, pub(crate) context_limit_tokens: usize, pub(crate) workflow_tasks: Vec, @@ -928,6 +929,8 @@ pub struct AppState { pub session_picker_error: Option, pub usage: UsageTotals, usage_revision: Option, + context_revision: Option, + context_observed: bool, pub context_used_tokens: usize, pub context_limit_tokens: usize, pub slash_menu: Option, @@ -1090,6 +1093,8 @@ impl AppState { session_picker_error: None, usage: UsageTotals::default(), usage_revision: None, + context_revision: None, + context_observed: false, context_used_tokens: 0, context_limit_tokens: 0, slash_menu: None, @@ -1735,8 +1740,19 @@ impl AppState { self.current_session_title = Some(projection.title.clone()); self.usage = projection.usage.clone(); self.usage_revision = Some(projection.usage_revision); - self.context_used_tokens = projection.context_used_tokens; - self.context_limit_tokens = projection.context_limit_tokens; + let revision_advanced = self + .context_revision + .is_none_or(|revision| projection.context_revision > revision); + let should_apply_context = + revision_advanced && (!self.context_observed || self.context_revision.is_some()); + if revision_advanced { + self.context_revision = Some(projection.context_revision); + } + if should_apply_context { + self.context_used_tokens = projection.context_used_tokens; + self.context_limit_tokens = projection.context_limit_tokens; + self.context_observed = false; + } self.current_goal = projection.current_goal.clone(); self.active_surface_operation_id = projection.foreground_operation_id.clone(); self.apply_workflow_tasks_update(projection.workflow_tasks.clone()); @@ -1754,8 +1770,11 @@ impl AppState { ); debug_assert_eq!(self.usage, projection.usage); debug_assert_eq!(self.usage_revision, Some(projection.usage_revision)); - debug_assert_eq!(self.context_used_tokens, projection.context_used_tokens); - debug_assert_eq!(self.context_limit_tokens, projection.context_limit_tokens); + if !self.context_observed { + debug_assert_eq!(self.context_revision, Some(projection.context_revision)); + debug_assert_eq!(self.context_used_tokens, projection.context_used_tokens); + debug_assert_eq!(self.context_limit_tokens, projection.context_limit_tokens); + } debug_assert_eq!( self.workflow_panel.tasks, sort_workflow_tasks_for_panel(projection.workflow_tasks.clone()) @@ -1840,6 +1859,8 @@ impl AppState { self.usage_revision = None; self.context_used_tokens = 0; self.context_limit_tokens = 0; + self.context_revision = None; + self.context_observed = false; self.approval_dialog = None; self.pending_input = None; self.approval_allowlist.clear(); @@ -3134,6 +3155,7 @@ impl AppState { } => { self.context_used_tokens = used_tokens; self.context_limit_tokens = limit_tokens; + self.context_observed = true; } TuiEvent::CompactionStarted => { self.set_status(AppStatus::Compacting); @@ -5670,6 +5692,7 @@ mod tests { cache_tokens: 7, estimated_cost_usd: 0.007, }, + context_revision: 1, context_used_tokens: 700, context_limit_tokens: 1_000, workflow_tasks: vec![workflow_task_summary("task-1", "Canonical task")], diff --git a/crates/orca-tui/src/ui.rs b/crates/orca-tui/src/ui.rs index 66c5f70c..f0d679af 100644 --- a/crates/orca-tui/src/ui.rs +++ b/crates/orca-tui/src/ui.rs @@ -4408,7 +4408,10 @@ fn truncate_lines(text: &str, max_lines: usize) -> String { #[cfg(test)] mod tests { use super::*; - use crate::types::{SlashMenu, SlashMenuItem, TuiEvent, TuiInteractionKey, TuiInteractionKind}; + use crate::types::{ + SlashMenu, SlashMenuItem, SurfaceProjectionState, TuiEvent, TuiInteractionKey, + TuiInteractionKind, + }; use chrono::Utc; use crossbeam_channel as mpsc; use orca_core::config::{AdditionalWorkingDirectory, ThemeName}; @@ -6401,6 +6404,38 @@ mod tests { assert!(!narrow.contains("blade-deepseek")); } + #[test] + fn provider_context_survives_same_revision_surface_sync_in_footer() { + let mut state = test_state(); + let snapshot = SurfaceProjectionState { + session_id: "context-session".to_string(), + title: "Context session".to_string(), + usage_revision: 1, + usage: orca_core::cost_types::UsageTotals::default(), + context_revision: 1, + context_used_tokens: 0, + context_limit_tokens: 128_000, + workflow_tasks: Vec::new(), + current_goal: None, + foreground_operation_id: None, + }; + state.update(TuiEvent::SurfaceProjectionSynced(Box::new( + snapshot.clone(), + ))); + state.update(TuiEvent::ContextUpdated { + used_tokens: 41_483, + limit_tokens: 1_000_000, + }); + state.update(TuiEvent::SurfaceProjectionSynced(Box::new(snapshot))); + + let theme = Theme::named(orca_core::config::ThemeName::Dark); + assert!( + status_line(&state, &theme, 100) + .to_string() + .contains("context 4%") + ); + } + #[test] fn status_line_reserves_known_context_before_truncating_a_long_model() { let mut state = test_state();