From d64b9e4bf65f62f511451a0c217ddeac233ad3df Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Mon, 6 Jul 2026 16:36:53 +0800 Subject: [PATCH 1/4] fix: fix cargo doc warnnings --- crates/core/src/history/compaction.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/core/src/history/compaction.rs b/crates/core/src/history/compaction.rs index 2651482a..d880ad63 100644 --- a/crates/core/src/history/compaction.rs +++ b/crates/core/src/history/compaction.rs @@ -4,14 +4,14 @@ //! Two compaction modes (`CompactionKind`) choose different preserve strategies: //! //! * **Auto** — token-budget threshold in the query loop (`query.rs`). -//! Preserves a tail window of roughly [`COMPACT_USER_MESSAGE_MAX_TOKENS`] -//! estimated tokens via [`split_by_user_message_budget`], regardless of user +//! Preserves a tail window of roughly `COMPACT_USER_MESSAGE_MAX_TOKENS` +//! estimated tokens via `split_by_user_message_budget`, regardless of user //! message boundaries. Example: `[user1, asst1, user2, asst2, user3]` may //! become `[summary, asst2, user3]` when `asst2` and `user3` fit the tail //! budget but `user2` does not. //! * **Proactive** — `/compact` or provider `context_too_long` retry. //! Preserves from the latest user message onward via -//! [`preserve_suffix_from_latest_user_message`]. Example: the same history +//! `preserve_suffix_from_latest_user_message`. Example: the same history //! becomes `[summary, user3]` only. //! //! The compaction flow: @@ -129,12 +129,12 @@ pub enum CompactionKind { /// Automatic compaction when context pressure is high. /// /// Skips when [`should_compact`] says the session is already within budget. - /// Preserve strategy: [`split_by_user_message_budget`] over the tail - /// [`COMPACT_USER_MESSAGE_MAX_TOKENS`] window (items, not user turns). + /// Preserve strategy: `split_by_user_message_budget` over the tail + /// `COMPACT_USER_MESSAGE_MAX_TOKENS` window (items, not user turns). Auto, /// Forced compaction that always runs. /// - /// Preserve strategy: [`preserve_suffix_from_latest_user_message`] — from + /// Preserve strategy: `preserve_suffix_from_latest_user_message` — from /// the last `Role::User` item through the end of history. Proactive, } From 1188fdbd974af5fd413169c2b69ec22d5e446dd7 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Mon, 6 Jul 2026 21:55:26 +0800 Subject: [PATCH 2/4] feat: add --yolo and --dangerously-skip-permissions for full-access TUI startup --- .../docs/reference/command-line-options.mdx | 2 + crates/cli/src/agent_command.rs | 70 +++++++++++++++++-- crates/cli/src/main.rs | 44 ++++++++++++ crates/server/tests/subagent_lifecycle.rs | 45 ++++++++++-- 4 files changed, 152 insertions(+), 9 deletions(-) diff --git a/apps/web/content/docs/reference/command-line-options.mdx b/apps/web/content/docs/reference/command-line-options.mdx index 5cf5d900..4bd10e55 100644 --- a/apps/web/content/docs/reference/command-line-options.mdx +++ b/apps/web/content/docs/reference/command-line-options.mdx @@ -12,6 +12,7 @@ Global flags can be used with interactive mode and subcommands. | --- | --- | --- | | `--model` | string | Override the model used for this process. | | `--log-level` | `trace`, `debug`, `info`, `warn`, `error` | Override process logging level. | +| `--dangerously-skip-permissions`, `--yolo` | none | Start the interactive TUI with full-access permissions, skipping approval prompts. | | `--help` | none | Print command help. | | `--version` | none | Print Devo version. | @@ -19,6 +20,7 @@ Example: ```bash devo --model deepseek-v4-pro --log-level debug +devo --yolo ``` ## Commands diff --git a/crates/cli/src/agent_command.rs b/crates/cli/src/agent_command.rs index 1850bbaa..5522bd22 100644 --- a/crates/cli/src/agent_command.rs +++ b/crates/cli/src/agent_command.rs @@ -23,11 +23,13 @@ use devo_util_paths::find_devo_home; /// when a provider config already exists. `exit_after_onboarding` exits after a /// successful onboarding save instead of continuing into the interactive TUI. /// `log_level` is forwarded to the background server process. +/// `dangerously_skip_permissions` starts the session with full-access permissions. pub(crate) async fn run_agent( force_onboarding: bool, exit_after_onboarding: bool, log_level: Option<&str>, initial_session_id: Option, + dangerously_skip_permissions: bool, ) -> Result { let cwd = std::env::current_dir()?; let config_home = find_devo_home().context("could not determine devo home directory")?; @@ -45,11 +47,8 @@ pub(crate) async fn run_agent( .collect(); let app_config = FileSystemAppConfigLoader::new(config_home.clone()).load(Some(&cwd))?; let project_key = project_config_key(&cwd); - let permission_preset = app_config - .projects - .get(&project_key) - .and_then(|config| config.permission_preset) - .unwrap_or(PermissionPreset::Default); + let permission_preset = + initial_permission_preset(&app_config, &project_key, dangerously_skip_permissions); let (onboarding_mode, resolved) = resolve_initial_provider_settings( force_onboarding, &app_config, @@ -116,6 +115,24 @@ pub(crate) async fn run_agent( Ok(exit) } +fn initial_permission_preset( + app_config: &AppConfig, + project_key: &str, + dangerously_skip_permissions: bool, +) -> PermissionPreset { + let mut permission_preset = app_config + .projects + .get(project_key) + .and_then(|config| config.permission_preset) + .unwrap_or(PermissionPreset::Default); + + if dangerously_skip_permissions { + permission_preset = PermissionPreset::FullAccess; + } + + permission_preset +} + /// Resolves the initial provider settings and whether onboarding should be shown. /// /// `force_onboarding` requests onboarding regardless of stored configuration. @@ -243,6 +260,7 @@ mod tests { use pretty_assertions::assert_eq; + use super::initial_permission_preset; use super::resolve_initial_provider_settings; use super::saved_model_entries; use devo_core::AppConfig; @@ -251,10 +269,12 @@ mod tests { use devo_core::Model; use devo_core::ModelBindingConfig; use devo_core::PresetModelCatalog; + use devo_core::ProjectConfig; use devo_core::ProviderConfigSection; use devo_core::ProviderDefaultsConfig; use devo_core::ProviderVendorConfig; use devo_core::ResolvedProviderSettings; + use devo_protocol::PermissionPreset; use devo_protocol::ProviderWireApi; use devo_tui::SavedModelEntry; @@ -283,6 +303,46 @@ mod tests { } } + #[test] + fn initial_permission_preset_uses_project_config_when_flag_is_false() { + let mut app_config = AppConfig::default(); + app_config.projects.insert( + "project-key".to_string(), + ProjectConfig { + permission_preset: Some(PermissionPreset::ReadOnly), + }, + ); + + assert_eq!( + initial_permission_preset( + &app_config, + "project-key", + /*dangerously_skip_permissions*/ false, + ), + PermissionPreset::ReadOnly, + ); + } + + #[test] + fn initial_permission_preset_overrides_to_full_access_when_flag_is_true() { + let mut app_config = AppConfig::default(); + app_config.projects.insert( + "project-key".to_string(), + ProjectConfig { + permission_preset: Some(PermissionPreset::ReadOnly), + }, + ); + + assert_eq!( + initial_permission_preset( + &app_config, + "project-key", + /*dangerously_skip_permissions*/ true, + ), + PermissionPreset::FullAccess, + ); + } + #[test] fn resolve_initial_provider_settings_uses_catalog_fallback_during_onboarding() { let actual = resolve_initial_provider_settings( diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 2550e188..39ca5b4b 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -59,6 +59,14 @@ struct Cli { .try_map(|level| level.parse::()) )] log_level: Option, + + /// Start with full-access permissions, skipping approval prompts. + #[arg( + long = "dangerously-skip-permissions", + visible_alias = "yolo", + global = true + )] + dangerously_skip_permissions: bool, } fn main() -> Result<()> { @@ -186,6 +194,7 @@ async fn run_cli() -> Result<()> { /*exit_after_onboarding*/ true, log_level.as_deref(), None, + cli.dangerously_skip_permissions, ) .await?; for line in onboarding_exit_messages(&exit, /*color_enabled*/ true) { @@ -211,6 +220,7 @@ async fn run_cli() -> Result<()> { /*exit_after_onboarding*/ false, log_level.as_deref(), Some(*session_id), + cli.dangerously_skip_permissions, ) .await?; for line in exit_messages(&exit, /*color_enabled*/ true) { @@ -240,6 +250,7 @@ async fn run_cli() -> Result<()> { /*exit_after_onboarding*/ false, log_level.as_deref(), None, + cli.dangerously_skip_permissions, ) .await?; let exit_lines = exit_messages(&exit, /*color_enabled*/ true); @@ -420,6 +431,33 @@ mod tests { } } + #[test] + fn cli_parses_dangerously_skip_permissions_flag() { + let cli = Cli::try_parse_from(["devo", "--dangerously-skip-permissions"]) + .expect("parse dangerously-skip-permissions"); + + assert!(cli.command.is_none()); + assert!(cli.dangerously_skip_permissions); + } + + #[test] + fn cli_parses_yolo_alias_for_dangerously_skip_permissions() { + let cli = Cli::try_parse_from(["devo", "--yolo"]).expect("parse yolo"); + + assert!(cli.command.is_none()); + assert!(cli.dangerously_skip_permissions); + } + + #[test] + fn cli_parses_yolo_alias_on_resume_subcommand() { + let session_id = SessionId::new(); + let cli = Cli::try_parse_from(["devo", "resume", &session_id.to_string(), "--yolo"]) + .expect("parse resume with yolo"); + + assert!(matches!(cli.command, Some(Command::Resume { .. }))); + assert!(cli.dangerously_skip_permissions); + } + #[test] fn cli_rejects_unsupported_log_levels() { let err = Cli::try_parse_from(["devo", "--log-level", "off"]).expect_err("reject off"); @@ -440,6 +478,7 @@ mod tests { command: None, model: None, log_level: Some(level), + dangerously_skip_permissions: false, }; assert_eq!( @@ -462,11 +501,13 @@ mod tests { command: None, model: None, log_level: None, + dangerously_skip_permissions: false, }, Cli { command: Some(Command::Onboard), model: None, log_level: None, + dangerously_skip_permissions: false, }, Cli { command: Some(Command::Prompt { @@ -475,6 +516,7 @@ mod tests { }), model: None, log_level: None, + dangerously_skip_permissions: false, }, ] { assert_eq!( @@ -493,6 +535,7 @@ mod tests { command: Some(Command::Doctor), model: None, log_level: None, + dangerously_skip_permissions: false, }; let server = Cli { command: Some(Command::Server { @@ -502,6 +545,7 @@ mod tests { }), model: None, log_level: None, + dangerously_skip_permissions: false, }; assert_eq!( diff --git a/crates/server/tests/subagent_lifecycle.rs b/crates/server/tests/subagent_lifecycle.rs index 6c5b93da..b65afd39 100644 --- a/crates/server/tests/subagent_lifecycle.rs +++ b/crates/server/tests/subagent_lifecycle.rs @@ -953,8 +953,22 @@ async fn fork_all_inherits_stable_parent_context() -> Result<()> { let requests = provider.requests(); assert_eq!(requests.len(), 4); - let completed_child_texts = message_texts(&requests[1]); - assert_subagent_request_hides_agent_tools(&requests[1]); + // Parent turn 2 and the active-fork child start concurrently, so provider + // request order is not stable; locate each turn by its task text. + let completed_child_request = find_unique_request_with_text(&requests, "use inherited context"); + let active_child_request = find_unique_request_with_text(&requests, "fork while parent active"); + let parent_active_turn_request = + find_unique_request_with_text(&requests, "active parent text should not be inherited yet"); + assert!( + parent_active_turn_request + .tools + .as_ref() + .is_some_and(|tools| { tools.iter().any(|tool| tool.name == "spawn_agent") }), + "parent active turn should expose agent coordination tools" + ); + + let completed_child_texts = message_texts(completed_child_request); + assert_subagent_request_hides_agent_tools(completed_child_request); assert!( completed_child_texts .iter() @@ -980,8 +994,8 @@ async fn fork_all_inherits_stable_parent_context() -> Result<()> { ); assert_subagent_reminder_before_task(&completed_child_texts, "use inherited context"); - let active_child_texts = message_texts(&requests[3]); - assert_subagent_request_hides_agent_tools(&requests[3]); + let active_child_texts = message_texts(active_child_request); + assert_subagent_request_hides_agent_tools(active_child_request); assert!( active_child_texts .iter() @@ -1078,6 +1092,29 @@ async fn fork_none_omits_parent_context_and_places_reminder_before_task() -> Res Ok(()) } +fn request_contains_text(request: &ModelRequest, needle: &str) -> bool { + message_texts(request) + .iter() + .any(|text| text.contains(needle)) +} + +fn find_unique_request_with_text<'a>( + requests: &'a [ModelRequest], + needle: &str, +) -> &'a ModelRequest { + let matches: Vec<_> = requests + .iter() + .filter(|request| request_contains_text(request, needle)) + .collect(); + assert_eq!( + matches.len(), + 1, + "expected exactly one request containing {needle:?}, got {} matching request(s)", + matches.len(), + ); + matches[0] +} + fn assert_generated_name(name: &str) { let Some((adjective, noun)) = name.split_once('-') else { panic!("generated name should be adjective-noun: {name}"); From ad73d6f4b7d91d362a9a8d51cf1aa1ce1a9353a1 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Mon, 6 Jul 2026 23:58:27 +0800 Subject: [PATCH 3/4] fix: lazy-load sessions --- crates/core/src/history/compaction.rs | 2 +- crates/server/AGENTS.md | 9 + crates/server/src/bootstrap.rs | 6 +- crates/server/src/db.rs | 466 ++++++++++++------ crates/server/src/persistence.rs | 359 +++++++++++++- crates/server/src/runtime.rs | 9 + crates/server/src/runtime/agents.rs | 6 +- .../runtime/handlers/acp/session_support.rs | 25 +- .../runtime/handlers/acp_config_options.rs | 2 +- crates/server/src/runtime/handlers/session.rs | 106 +++- crates/server/src/runtime/lifecycle.rs | 326 +++++++----- .../src/runtime/session_actor/registry.rs | 16 +- crates/server/src/runtime/session_cache.rs | 297 +++++++++++ crates/server/tests/acp_session_lifecycle.rs | 102 +++- crates/server/tests/persistence_resume.rs | 402 +++++++++++++++ 15 files changed, 1820 insertions(+), 313 deletions(-) create mode 100644 crates/server/src/runtime/session_cache.rs diff --git a/crates/core/src/history/compaction.rs b/crates/core/src/history/compaction.rs index d880ad63..6e934dc4 100644 --- a/crates/core/src/history/compaction.rs +++ b/crates/core/src/history/compaction.rs @@ -434,7 +434,7 @@ mod tests { #[test] fn compact_builds_merged_messages() { - let items = vec![ + let items = [ ResponseItem::Message(Message::user("hello")), // Split assistant turn — text + two tool calls ResponseItem::Message(Message::assistant_text("ok")), diff --git a/crates/server/AGENTS.md b/crates/server/AGENTS.md index b8960de7..0e77f882 100644 --- a/crates/server/AGENTS.md +++ b/crates/server/AGENTS.md @@ -39,3 +39,12 @@ The server runtime uses **one session actor per session**. Durable session state - **Runtime concurrency changes need integration coverage** in `crates/server/tests/`: interrupt mid-stream, queued follow-ups, goal lifecycle interrupts, persistence/resume, research. - **Prefer waiting on observable protocol outcomes** (notifications, terminal status) over sleeping or polling internal maps. - Follow existing test conventions: `pretty_assertions::assert_eq`, compare whole objects where possible, platform-aware paths when touching filesystem behavior. + +### Session persistence layers + +- **Rollout JSONL** under `~/.devo/sessions/` is the canonical conversation history. +- **SQLite** (`devo.db` `sessions` table) stores a lightweight index (`rollout_path`, `parent_session_id`, title, cwd, timestamps) used by `session/list` and resume decisions. +- **In-memory session actors** are loaded on demand via `get_or_load_parent_session`; root sessions are LRU-evicted (capacity 16) when unpinned. +- **`session/list`** returns durable user-visible sessions only (non-ephemeral, no `agent_path`; includes forks with `parent_session_id`); subagent rows are indexed but hidden from list. +- **`session/resume`** loads parent sessions lazily from rollout files. Subagent session ids cannot be resumed directly; missing rollout files fail with an explicit restore error. +- **Startup** runs `index_rollout_metadata` in the background instead of replaying every rollout into memory. diff --git a/crates/server/src/bootstrap.rs b/crates/server/src/bootstrap.rs index 79f590cc..41c2365f 100644 --- a/crates/server/src/bootstrap.rs +++ b/crates/server/src/bootstrap.rs @@ -252,9 +252,9 @@ pub async fn run_server_process( serde_json::Map::from_iter([("trigger".to_string(), serde_json::json!("init"))]), ) .await; - tracing::info!("starting persisted session restore"); - runtime.load_persisted_sessions().await?; - tracing::info!("persisted session restore completed"); + if runtime.backfill_session_index_if_required()? { + tracing::info!("rollout metadata index backfill completed"); + } let shutdown_signal = tokio_util::sync::CancellationToken::new(); let internal_proxy_control = InternalProxyControl::new(shutdown_signal.clone()); diff --git a/crates/server/src/db.rs b/crates/server/src/db.rs index 4ab5a56c..43e900ee 100644 --- a/crates/server/src/db.rs +++ b/crates/server/src/db.rs @@ -31,6 +31,22 @@ impl QueueType { } } +/// SQLite index row used for lazy session list and resume. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionIndexRecord { + pub metadata: SessionMetadata, + pub rollout_path: Option, +} + +/// Source of a session metadata upsert. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SessionUpsertSource { + /// Live runtime writes should preserve existing non-null fields when the update omits them. + RuntimeLive, + /// Rollout index writes rebuild SQLite from canonical rollout metadata. + RolloutIndex, +} + /// Session-level token statistics. #[derive(Debug, Clone)] pub struct SessionStats { @@ -211,13 +227,47 @@ impl Database { ) .context("failed to backfill total_tokens column")?; } + if !sessions_has_column(&conn, "rollout_path")? { + conn.execute("ALTER TABLE sessions ADD COLUMN rollout_path TEXT", []) + .context("failed to add rollout_path column")?; + } + if !sessions_has_column(&conn, "parent_session_id")? { + conn.execute("ALTER TABLE sessions ADD COLUMN parent_session_id TEXT", []) + .context("failed to add parent_session_id column")?; + } + if !sessions_has_column(&conn, "agent_path")? { + conn.execute("ALTER TABLE sessions ADD COLUMN agent_path TEXT", []) + .context("failed to add agent_path column")?; + } Ok(()) } // === Session CRUD === - /// Inserts or updates a session's metadata. - pub fn upsert_session(&self, meta: &SessionMetadata) -> Result<()> { + /// Inserts or updates a session's metadata and optional rollout index fields. + pub fn upsert_session( + &self, + meta: &SessionMetadata, + rollout_path: Option<&std::path::Path>, + ) -> Result<()> { + self.upsert_session_with_source(meta, rollout_path, SessionUpsertSource::RuntimeLive) + } + + /// Inserts or updates session metadata using rollout-index semantics. + pub fn upsert_rollout_index_session( + &self, + meta: &SessionMetadata, + rollout_path: Option<&std::path::Path>, + ) -> Result<()> { + self.upsert_session_with_source(meta, rollout_path, SessionUpsertSource::RolloutIndex) + } + + fn upsert_session_with_source( + &self, + meta: &SessionMetadata, + rollout_path: Option<&std::path::Path>, + source: SessionUpsertSource, + ) -> Result<()> { let conn = self.conn.lock().expect("database mutex poisoned"); let additional_directories = serde_json::to_string(&meta.additional_directories) .context("failed to serialize session additional directories")?; @@ -226,18 +276,51 @@ impl Database { SessionTitleState::Provisional => "provisional", SessionTitleState::Final(_) => "final", }; - conn.execute( - "INSERT INTO sessions (id, title, title_state, model, thinking, cwd, additional_directories, ephemeral, created_at, updated_at, last_activity_at, schema_version) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, 2) - ON CONFLICT(id) DO UPDATE SET - title = excluded.title, - title_state = excluded.title_state, - model = excluded.model, - thinking = excluded.thinking, + let rollout_path_str = rollout_path.map(|path| path.to_string_lossy().into_owned()); + let parent_session_id = meta.parent_session_id.map(|id| id.to_string()); + let agent_path = meta.agent_path.clone(); + let update_clause = match source { + SessionUpsertSource::RuntimeLive => { + "title = COALESCE(excluded.title, sessions.title), + title_state = CASE + WHEN excluded.title IS NOT NULL THEN excluded.title_state + ELSE sessions.title_state + END, + model = COALESCE(excluded.model, sessions.model), + thinking = COALESCE(excluded.thinking, sessions.thinking), cwd = excluded.cwd, additional_directories = excluded.additional_directories, updated_at = excluded.updated_at, - last_activity_at = excluded.last_activity_at", + last_activity_at = excluded.last_activity_at, + parent_session_id = COALESCE(excluded.parent_session_id, sessions.parent_session_id), + agent_path = COALESCE(excluded.agent_path, sessions.agent_path), + rollout_path = COALESCE(excluded.rollout_path, sessions.rollout_path)" + } + SessionUpsertSource::RolloutIndex => { + "title = COALESCE(excluded.title, sessions.title), + title_state = CASE + WHEN excluded.title IS NOT NULL THEN excluded.title_state + ELSE sessions.title_state + END, + model = COALESCE(excluded.model, sessions.model), + thinking = COALESCE(excluded.thinking, sessions.thinking), + cwd = excluded.cwd, + additional_directories = excluded.additional_directories, + created_at = excluded.created_at, + updated_at = excluded.updated_at, + last_activity_at = excluded.last_activity_at, + parent_session_id = COALESCE(excluded.parent_session_id, sessions.parent_session_id), + agent_path = COALESCE(excluded.agent_path, sessions.agent_path), + rollout_path = COALESCE(excluded.rollout_path, sessions.rollout_path)" + } + }; + let sql = format!( + "INSERT INTO sessions (id, title, title_state, model, thinking, cwd, additional_directories, ephemeral, created_at, updated_at, last_activity_at, rollout_path, parent_session_id, agent_path, schema_version) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, 3) + ON CONFLICT(id) DO UPDATE SET {update_clause}" + ); + conn.execute( + &sql, params![ meta.session_id.to_string(), meta.title, @@ -250,81 +333,40 @@ impl Database { meta.created_at.timestamp(), meta.updated_at.timestamp(), meta.last_activity_at.timestamp(), + rollout_path_str, + parent_session_id, + agent_path, ], ) .context("failed to upsert session")?; Ok(()) } + /// Returns true when durable sessions need rollout metadata backfilled into SQLite. + pub fn session_index_backfill_required(&self) -> Result { + let conn = self.conn.lock().expect("database mutex poisoned"); + let count: i64 = conn + .query_row( + "SELECT COUNT(*) FROM sessions + WHERE ephemeral = 0 AND (rollout_path IS NULL OR rollout_path = '')", + [], + |row| row.get(0), + ) + .context("failed to check session index backfill requirement")?; + Ok(count > 0) + } + /// Retrieves a session's metadata by ID. pub fn get_session(&self, id: &SessionId) -> Result> { let conn = self.conn.lock().expect("database mutex poisoned"); let mut stmt = conn .prepare( - "SELECT id, title, title_state, model, thinking, cwd, additional_directories, ephemeral, created_at, updated_at, last_activity_at + "SELECT id, title, title_state, model, thinking, cwd, additional_directories, ephemeral, created_at, updated_at, last_activity_at, parent_session_id, agent_path FROM sessions WHERE id = ?1", ) .context("failed to prepare get_session statement")?; let result = stmt.query_row(params![id.to_string()], |row| { - let id_str: String = row.get(0)?; - let title: Option = row.get(1)?; - let title_state_str: String = row.get(2)?; - let model: Option = row.get(3)?; - let thinking: Option = row.get(4)?; - let cwd_str: String = row.get(5)?; - let additional_directories_str: String = row.get(6)?; - let ephemeral: i32 = row.get(7)?; - let created_at: i64 = row.get(8)?; - let updated_at: i64 = row.get(9)?; - let last_activity_at: i64 = row.get(10)?; - - let title_state = match title_state_str.as_str() { - "provisional" => SessionTitleState::Provisional, - "final" => { - SessionTitleState::Final(devo_protocol::SessionTitleFinalSource::ModelGenerated) - } - _ => SessionTitleState::Unset, - }; - - Ok(SessionMetadata { - session_id: parse_session_id_column(id_str, 0)?, - cwd: PathBuf::from(&cwd_str), - additional_directories: parse_additional_directories_column( - additional_directories_str, - 6, - )?, - created_at: Utc - .timestamp_opt(created_at, 0) - .single() - .unwrap_or_else(Utc::now), - updated_at: Utc - .timestamp_opt(updated_at, 0) - .single() - .unwrap_or_else(Utc::now), - last_activity_at: Utc - .timestamp_opt(last_activity_at, 0) - .single() - .unwrap_or_else(Utc::now), - title, - title_state, - parent_session_id: None, - agent_path: None, - agent_nickname: None, - agent_role: None, - ephemeral: ephemeral != 0, - model, - model_binding_id: None, - reasoning_effort_selection: thinking, - reasoning_effort: None, - total_input_tokens: 0, - total_output_tokens: 0, - total_tokens: 0, - total_cache_creation_tokens: 0, - total_cache_read_tokens: 0, - prompt_token_estimate: 0, - last_query_total_tokens: 0, - status: SessionRuntimeStatus::Idle, - }) + parse_session_metadata_row(row, false) }); match result { @@ -334,77 +376,67 @@ impl Database { } } + /// Returns resume/list index fields for a session id. + pub fn get_session_index(&self, id: &SessionId) -> Result> { + let conn = self.conn.lock().expect("database mutex poisoned"); + let mut stmt = conn + .prepare( + "SELECT id, title, title_state, model, thinking, cwd, additional_directories, ephemeral, created_at, updated_at, last_activity_at, parent_session_id, agent_path, rollout_path + FROM sessions WHERE id = ?1", + ) + .context("failed to prepare get_session_index statement")?; + let result = stmt.query_row(params![id.to_string()], |row| { + let metadata = parse_session_metadata_row(row, true)?; + let rollout_path = row + .get::<_, Option>(13)? + .map(PathBuf::from) + .filter(|path| !path.as_os_str().is_empty()); + Ok(SessionIndexRecord { + metadata, + rollout_path, + }) + }); + + match result { + Ok(index) => Ok(Some(index)), + Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None), + Err(e) => Err(e.into()), + } + } + + /// Lists durable user-visible sessions (roots and forks, excluding subagents). + pub fn list_root_sessions(&self) -> Result> { + let conn = self.conn.lock().expect("database mutex poisoned"); + let mut stmt = conn + .prepare( + "SELECT id, title, title_state, model, thinking, cwd, additional_directories, ephemeral, created_at, updated_at, last_activity_at, parent_session_id, agent_path + FROM sessions + WHERE ephemeral = 0 AND agent_path IS NULL + ORDER BY last_activity_at DESC, updated_at DESC", + ) + .context("failed to prepare list_root_sessions statement")?; + let rows = stmt + .query_map([], |row| parse_session_metadata_row(row, false)) + .context("failed to query root sessions")?; + + let mut sessions = Vec::new(); + for row in rows { + sessions.push(row?); + } + Ok(sessions) + } + /// Lists all sessions ordered by most recently updated. pub fn list_sessions(&self) -> Result> { let conn = self.conn.lock().expect("database mutex poisoned"); let mut stmt = conn .prepare( - "SELECT id, title, title_state, model, thinking, cwd, additional_directories, ephemeral, created_at, updated_at, last_activity_at + "SELECT id, title, title_state, model, thinking, cwd, additional_directories, ephemeral, created_at, updated_at, last_activity_at, parent_session_id, agent_path FROM sessions ORDER BY last_activity_at DESC, updated_at DESC", ) .context("failed to prepare list_sessions statement")?; let rows = stmt - .query_map([], |row| { - let id_str: String = row.get(0)?; - let title: Option = row.get(1)?; - let title_state_str: String = row.get(2)?; - let model: Option = row.get(3)?; - let thinking: Option = row.get(4)?; - let cwd_str: String = row.get(5)?; - let additional_directories_str: String = row.get(6)?; - let ephemeral: i32 = row.get(7)?; - let created_at: i64 = row.get(8)?; - let updated_at: i64 = row.get(9)?; - let last_activity_at: i64 = row.get(10)?; - - let title_state = match title_state_str.as_str() { - "provisional" => SessionTitleState::Provisional, - "final" => SessionTitleState::Final( - devo_protocol::SessionTitleFinalSource::ModelGenerated, - ), - _ => SessionTitleState::Unset, - }; - - Ok(SessionMetadata { - session_id: parse_session_id_column(id_str, 0)?, - cwd: PathBuf::from(&cwd_str), - additional_directories: parse_additional_directories_column( - additional_directories_str, - 6, - )?, - created_at: Utc - .timestamp_opt(created_at, 0) - .single() - .unwrap_or_else(Utc::now), - updated_at: Utc - .timestamp_opt(updated_at, 0) - .single() - .unwrap_or_else(Utc::now), - last_activity_at: Utc - .timestamp_opt(last_activity_at, 0) - .single() - .unwrap_or_else(Utc::now), - title, - title_state, - parent_session_id: None, - agent_path: None, - agent_nickname: None, - agent_role: None, - ephemeral: ephemeral != 0, - model, - model_binding_id: None, - reasoning_effort_selection: thinking, - reasoning_effort: None, - total_input_tokens: 0, - total_output_tokens: 0, - total_tokens: 0, - total_cache_creation_tokens: 0, - total_cache_read_tokens: 0, - prompt_token_estimate: 0, - last_query_total_tokens: 0, - status: SessionRuntimeStatus::Idle, - }) - }) + .query_map([], |row| parse_session_metadata_row(row, false)) .context("failed to query sessions")?; let mut sessions = Vec::new(); @@ -696,6 +728,88 @@ impl Database { } } +fn sessions_has_column(conn: &Connection, column: &str) -> Result { + let mut stmt = conn + .prepare("PRAGMA table_info(sessions)") + .context("failed to inspect sessions schema")?; + let columns = stmt + .query_map([], |row| row.get::<_, String>(1)) + .context("failed to read sessions schema")?; + for column_name in columns { + if column_name? == column { + return Ok(true); + } + } + Ok(false) +} + +fn parse_session_metadata_row( + row: &rusqlite::Row<'_>, + _include_parent: bool, +) -> rusqlite::Result { + let id_str: String = row.get(0)?; + let title: Option = row.get(1)?; + let title_state_str: String = row.get(2)?; + let model: Option = row.get(3)?; + let thinking: Option = row.get(4)?; + let cwd_str: String = row.get(5)?; + let additional_directories_str: String = row.get(6)?; + let ephemeral: i32 = row.get(7)?; + let created_at: i64 = row.get(8)?; + let updated_at: i64 = row.get(9)?; + let last_activity_at: i64 = row.get(10)?; + let parent_session_id = row + .get::<_, Option>(11)? + .map(|value| parse_session_id_column(value, 11)) + .transpose()?; + let agent_path = row + .get::<_, Option>(12)? + .filter(|path| !path.is_empty()); + + let title_state = match title_state_str.as_str() { + "provisional" => SessionTitleState::Provisional, + "final" => SessionTitleState::Final(devo_protocol::SessionTitleFinalSource::ModelGenerated), + _ => SessionTitleState::Unset, + }; + + Ok(SessionMetadata { + session_id: parse_session_id_column(id_str, 0)?, + cwd: PathBuf::from(&cwd_str), + additional_directories: parse_additional_directories_column(additional_directories_str, 6)?, + created_at: Utc + .timestamp_opt(created_at, 0) + .single() + .unwrap_or_else(Utc::now), + updated_at: Utc + .timestamp_opt(updated_at, 0) + .single() + .unwrap_or_else(Utc::now), + last_activity_at: Utc + .timestamp_opt(last_activity_at, 0) + .single() + .unwrap_or_else(Utc::now), + title, + title_state, + parent_session_id, + agent_path, + agent_nickname: None, + agent_role: None, + ephemeral: ephemeral != 0, + model, + model_binding_id: None, + reasoning_effort_selection: thinking, + reasoning_effort: None, + total_input_tokens: 0, + total_output_tokens: 0, + total_tokens: 0, + total_cache_creation_tokens: 0, + total_cache_read_tokens: 0, + prompt_token_estimate: 0, + last_query_total_tokens: 0, + status: SessionRuntimeStatus::Idle, + }) +} + fn parse_session_id_column(id: String, column: usize) -> rusqlite::Result { SessionId::from_str(&id).map_err(|error| { rusqlite::Error::FromSqlConversionFailure( @@ -802,7 +916,7 @@ mod tests { let (db, _dir) = test_db(); let mut meta = sample_session("session-1"); meta.additional_directories = vec![PathBuf::from("/tmp/shared")]; - db.upsert_session(&meta).expect("upsert"); + db.upsert_session(&meta, None).expect("upsert"); let retrieved = db.get_session(&meta.session_id).expect("get"); assert!(retrieved.is_some()); @@ -825,8 +939,8 @@ mod tests { meta1.last_activity_at = baseline; meta2.updated_at = baseline + chrono::Duration::seconds(10); meta2.last_activity_at = baseline - chrono::Duration::seconds(10); - db.upsert_session(&meta1).expect("upsert"); - db.upsert_session(&meta2).expect("upsert"); + db.upsert_session(&meta1, None).expect("upsert"); + db.upsert_session(&meta2, None).expect("upsert"); let sessions = db.list_sessions().expect("list"); assert_eq!(sessions.len(), 2); @@ -869,7 +983,7 @@ mod tests { fn delete_session_cascades() { let (db, _dir) = test_db(); let meta = sample_session("session-1"); - db.upsert_session(&meta).expect("upsert"); + db.upsert_session(&meta, None).expect("upsert"); db.delete_session(&meta.session_id).expect("delete"); let retrieved = db.get_session(&meta.session_id).expect("get"); @@ -880,7 +994,7 @@ mod tests { fn update_and_get_stats() { let (db, _dir) = test_db(); let meta = sample_session("session-1"); - db.upsert_session(&meta).expect("upsert"); + db.upsert_session(&meta, None).expect("upsert"); let stats = SessionStats { total_input_tokens: 1000, @@ -906,7 +1020,7 @@ mod tests { fn push_and_drain_pending() { let (db, _dir) = test_db(); let meta = sample_session("session-1"); - db.upsert_session(&meta).expect("upsert"); + db.upsert_session(&meta, None).expect("upsert"); let item1 = PendingInputItem::new( PendingInputKind::UserText { @@ -952,7 +1066,7 @@ mod tests { fn queue_types_are_isolated() { let (db, _dir) = test_db(); let meta = sample_session("session-1"); - db.upsert_session(&meta).expect("upsert"); + db.upsert_session(&meta, None).expect("upsert"); let turn_item = PendingInputItem::new( PendingInputKind::UserText { @@ -1000,7 +1114,7 @@ mod tests { fn remove_pending_by_id_only_removes_matching_item() { let (db, _dir) = test_db(); let meta = sample_session("session-1"); - db.upsert_session(&meta).expect("upsert"); + db.upsert_session(&meta, None).expect("upsert"); let first = PendingInputItem::new( PendingInputKind::UserText { @@ -1042,11 +1156,79 @@ mod tests { fn drain_pending_empty_returns_empty() { let (db, _dir) = test_db(); let meta = sample_session("session-1"); - db.upsert_session(&meta).expect("upsert"); + db.upsert_session(&meta, None).expect("upsert"); let drained = db .drain_pending(&meta.session_id, QueueType::Turn) .expect("drain"); assert!(drained.is_empty()); } + + #[test] + fn session_index_backfill_required_detects_missing_rollout_path() { + let (db, _dir) = test_db(); + let meta = sample_session("session-1"); + + assert!( + !db.session_index_backfill_required() + .expect("check empty db") + ); + db.upsert_session(&meta, None) + .expect("upsert without rollout path"); + + assert!( + db.session_index_backfill_required() + .expect("check missing rollout path") + ); + db.upsert_rollout_index_session(&meta, Some("/tmp/session.jsonl".as_ref())) + .expect("upsert rollout path"); + assert!( + !db.session_index_backfill_required() + .expect("check populated rollout path") + ); + } + + #[test] + fn list_root_sessions_excludes_subagents_and_ephemeral() { + let (db, _dir) = test_db(); + let root_id = SessionId::new(); + let subagent_id = SessionId::new(); + let ephemeral_id = SessionId::new(); + let rollout_path = PathBuf::from("/tmp/root.jsonl"); + + let mut root = sample_session(&root_id.to_string()); + root.session_id = root_id; + db.upsert_session(&root, Some(rollout_path.as_path())) + .expect("upsert root"); + + let mut subagent = sample_session(&subagent_id.to_string()); + subagent.session_id = subagent_id; + subagent.parent_session_id = Some(root_id); + subagent.agent_path = Some("root/review".into()); + db.upsert_session(&subagent, Some("/tmp/subagent.jsonl".as_ref())) + .expect("upsert subagent"); + + let mut ephemeral = sample_session(&ephemeral_id.to_string()); + ephemeral.session_id = ephemeral_id; + ephemeral.ephemeral = true; + db.upsert_session(&ephemeral, None) + .expect("upsert ephemeral"); + + let roots = db.list_root_sessions().expect("list root sessions"); + assert_eq!(roots.len(), 1); + assert_eq!(roots[0].session_id, root_id); + + let index = db + .get_session_index(&root_id) + .expect("get index") + .expect("root index"); + assert_eq!(index.rollout_path, Some(rollout_path)); + assert_eq!(index.metadata.parent_session_id, None); + + let subagent_index = db + .get_session_index(&subagent_id) + .expect("get subagent index") + .expect("subagent index"); + assert_eq!(subagent_index.metadata.parent_session_id, Some(root_id)); + } } diff --git a/crates/server/src/persistence.rs b/crates/server/src/persistence.rs index 45e9bc65..e1b2ceb6 100644 --- a/crates/server/src/persistence.rs +++ b/crates/server/src/persistence.rs @@ -360,6 +360,93 @@ impl RolloutStore { Ok(sessions) } + /// Indexes rollout SessionMeta headers into SQLite without replaying turns. + pub(crate) fn index_rollout_metadata(&self, db: &crate::db::Database) -> Result<()> { + let mut canonical = + HashMap::, PathBuf, SessionRecord)>::new(); + for rollout_path in self.rollout_paths()? { + match read_rollout_index_fields(&rollout_path) { + Ok((record, last_activity_at)) => match canonical.get(&record.id) { + Some((existing_activity, existing_path, _)) => { + if last_activity_at > *existing_activity { + tracing::warn!( + session_id = %record.id, + kept_rollout_path = %rollout_path.display(), + replaced_rollout_path = %existing_path.display(), + "duplicate rollout for session id; keeping newest last_activity_at" + ); + canonical.insert(record.id, (last_activity_at, rollout_path, record)); + } else { + tracing::warn!( + session_id = %record.id, + kept_rollout_path = %existing_path.display(), + ignored_rollout_path = %rollout_path.display(), + "duplicate rollout for session id; keeping newest last_activity_at" + ); + } + } + None => { + canonical.insert(record.id, (last_activity_at, rollout_path, record)); + } + }, + Err(error) => { + tracing::warn!( + rollout_path = %rollout_path.display(), + error = %error, + "failed to index rollout metadata; skipping file" + ); + } + } + } + + for (session_id, (last_activity_at, rollout_path, record)) in canonical { + let metadata = session_metadata_from_record(&record, last_activity_at); + if let Err(error) = + db.upsert_rollout_index_session(&metadata, Some(rollout_path.as_path())) + { + tracing::warn!( + session_id = %session_id, + error = %error, + "failed to upsert indexed session metadata" + ); + } + } + Ok(()) + } + + /// Finds a rollout file by session id suffix when the SQLite index is stale. + pub(crate) fn find_rollout_by_session_id( + &self, + session_id: &SessionId, + ) -> Result> { + let suffix = format!("-{session_id}.jsonl"); + for rollout_path in self.rollout_paths()? { + let Some(file_name) = rollout_path.file_name().and_then(|name| name.to_str()) else { + continue; + }; + if file_name.ends_with(&suffix) { + return Ok(Some(rollout_path)); + } + } + Ok(None) + } + + /// Resolves a session cwd from the SQLite index, falling back to rollout SessionMeta. + pub(crate) fn resolve_indexed_session_cwd( + &self, + db: &crate::db::Database, + session_id: &SessionId, + ) -> Result> { + if let Some(index) = db.get_session_index(session_id)? { + return Ok(Some(index.metadata.cwd)); + } + if let Some(rollout_path) = self.find_rollout_by_session_id(session_id)? { + let (record, _) = read_rollout_index_fields(&rollout_path)?; + return Ok(Some(record.cwd)); + } + Ok(None) + } + /// Deletes canonical rollout files for a session. pub(crate) fn delete_session_rollouts(&self, session_id: &SessionId) -> Result { let suffix = format!("-{session_id}.jsonl"); @@ -394,7 +481,7 @@ impl RolloutStore { Ok(deleted) } - async fn load_session_from_rollout( + pub(crate) async fn load_session_from_rollout( &self, rollout_path: &Path, deps: &ServerRuntimeDependencies, @@ -433,7 +520,7 @@ impl RolloutStore { .with_context(|| format!("replay rollout {}", rollout_path.display())) } - fn rollout_paths(&self) -> Result> { + pub(crate) fn rollout_paths(&self) -> Result> { let mut files = Vec::new(); let root = self.data_root.join("sessions"); if !root.exists() { @@ -1493,6 +1580,83 @@ fn apply_prompt_turn_item( } } +fn read_rollout_index_fields(path: &Path) -> Result<(SessionRecord, chrono::DateTime)> { + let file = File::open(path).with_context(|| format!("open rollout file {}", path.display()))?; + let reader = BufReader::new(file); + let mut session: Option = None; + let mut last_activity_at: Option> = None; + + for line in reader.lines() { + let line = line.with_context(|| format!("read line from {}", path.display()))?; + if line.trim().is_empty() { + continue; + } + let parsed = match serde_json::from_str::(&line) { + Ok(parsed) => parsed, + Err(_) => continue, + }; + match parsed { + RolloutLine::SessionMeta(meta_line) => { + let mut record = meta_line.session; + if record.last_activity_at.is_none() { + record.last_activity_at = Some(record.created_at); + } + last_activity_at = record.last_activity_at; + session = Some(record); + } + RolloutLine::SessionTitleUpdated(line) => { + if let Some(record) = session.as_mut() { + record.title = Some(line.title); + record.title_state = line.title_state; + record.updated_at = line.timestamp; + last_activity_at = Some(line.timestamp); + } + } + _ => {} + } + } + + let session = session + .with_context(|| format!("missing SessionMeta line in rollout {}", path.display()))?; + let last_activity_at = last_activity_at + .or(session.last_activity_at) + .unwrap_or(session.created_at); + Ok((session, last_activity_at)) +} + +fn session_metadata_from_record( + record: &SessionRecord, + last_activity_at: chrono::DateTime, +) -> SessionMetadata { + SessionMetadata { + session_id: record.id, + cwd: record.cwd.clone(), + additional_directories: record.additional_directories.clone(), + created_at: record.created_at, + updated_at: record.updated_at, + last_activity_at, + title: record.title.clone(), + title_state: record.title_state.clone(), + parent_session_id: record.parent_session_id, + agent_path: record.agent_path.clone(), + agent_nickname: record.agent_nickname.clone(), + agent_role: record.agent_role.clone(), + ephemeral: false, + model: record.model.clone(), + model_binding_id: record.model_binding_id.clone(), + reasoning_effort_selection: record.reasoning_effort_selection.clone(), + reasoning_effort: None, + total_input_tokens: 0, + total_output_tokens: 0, + total_tokens: 0, + total_cache_creation_tokens: 0, + total_cache_read_tokens: 0, + prompt_token_estimate: 0, + last_query_total_tokens: 0, + status: SessionRuntimeStatus::Idle, + } +} + fn collect_rollout_files(root: &Path, files: &mut Vec) -> Result<()> { for entry in std::fs::read_dir(root).with_context(|| format!("read dir {}", root.display()))? { let entry = entry.with_context(|| format!("read entry in {}", root.display()))?; @@ -1890,6 +2054,197 @@ mod tests { ); } + #[test] + fn index_rollout_metadata_reads_session_meta_only() { + use chrono::Utc; + use pretty_assertions::assert_eq; + use std::io::Write; + use tempfile::TempDir; + + let dir = TempDir::new().expect("temp dir"); + let data_root = dir.path().to_path_buf(); + let session_id = SessionId::new(); + let now = Utc::now(); + let rollout_store = super::RolloutStore::new(data_root.clone()); + let record = rollout_store.create_session_record( + session_id, + now, + data_root.clone(), + Vec::new(), + Some("Indexed session".into()), + Some("test-model".into()), + None, + None, + "test-provider".into(), + None, + ); + rollout_store + .append_session_meta(&record) + .expect("append session meta"); + std::fs::OpenOptions::new() + .append(true) + .open(&record.rollout_path) + .expect("open rollout") + .write_all(b"{\"type\":\"turn\",\"payload\":{}}\n") + .expect("append non-meta line"); + + let db = crate::db::Database::open(data_root.join("index.db")).expect("open db"); + rollout_store + .index_rollout_metadata(&db) + .expect("index rollout metadata"); + + let index = db + .get_session_index(&session_id) + .expect("get index") + .expect("indexed session"); + assert_eq!(index.metadata.title.as_deref(), Some("Indexed session")); + assert_eq!(index.rollout_path, Some(record.rollout_path)); + assert_eq!(index.metadata.parent_session_id, None); + + let roots = db.list_root_sessions().expect("list roots"); + assert_eq!(roots.len(), 1); + assert_eq!(roots[0].session_id, session_id); + } + + #[test] + fn index_rollout_metadata_overwrites_stale_sqlite_title_when_rollout_has_title() { + use chrono::Utc; + use devo_protocol::SessionMetadata; + use devo_protocol::SessionRuntimeStatus; + use devo_protocol::SessionTitleState; + use pretty_assertions::assert_eq; + use tempfile::TempDir; + + let dir = TempDir::new().expect("temp dir"); + let data_root = dir.path().to_path_buf(); + let session_id = SessionId::new(); + let now = Utc::now(); + let rollout_store = super::RolloutStore::new(data_root.clone()); + let record = rollout_store.create_session_record( + session_id, + now, + data_root.clone(), + Vec::new(), + None, + Some("test-model".into()), + None, + None, + "test-provider".into(), + None, + ); + rollout_store + .append_session_meta(&record) + .expect("append session meta"); + rollout_store + .append_title_update( + &record, + "Canonical rollout title".into(), + SessionTitleState::Final(devo_core::SessionTitleFinalSource::ModelGenerated), + None, + ) + .expect("append title update"); + + let db = crate::db::Database::open(data_root.join("index.db")).expect("open db"); + db.upsert_session( + &SessionMetadata { + session_id, + cwd: data_root.clone(), + additional_directories: Vec::new(), + created_at: now, + updated_at: now, + last_activity_at: now, + title: Some("Existing title".into()), + title_state: SessionTitleState::Final( + devo_core::SessionTitleFinalSource::ModelGenerated, + ), + parent_session_id: None, + agent_path: None, + agent_nickname: None, + agent_role: None, + ephemeral: false, + model: Some("test-model".into()), + model_binding_id: None, + reasoning_effort_selection: None, + reasoning_effort: None, + total_input_tokens: 0, + total_output_tokens: 0, + total_tokens: 0, + total_cache_creation_tokens: 0, + total_cache_read_tokens: 0, + prompt_token_estimate: 0, + last_query_total_tokens: 0, + status: SessionRuntimeStatus::Idle, + }, + None, + ) + .expect("seed sqlite title"); + + rollout_store + .index_rollout_metadata(&db) + .expect("index rollout metadata"); + + let index = db + .get_session_index(&session_id) + .expect("get index") + .expect("indexed session"); + assert_eq!( + index.metadata.title.as_deref(), + Some("Canonical rollout title") + ); + } + + #[test] + fn index_rollout_metadata_reads_session_title_updates() { + use chrono::Utc; + use devo_core::SessionTitleFinalSource; + use devo_core::SessionTitleState; + use pretty_assertions::assert_eq; + use tempfile::TempDir; + + let dir = TempDir::new().expect("temp dir"); + let data_root = dir.path().to_path_buf(); + let session_id = SessionId::new(); + let now = Utc::now(); + let rollout_store = super::RolloutStore::new(data_root.clone()); + let record = rollout_store.create_session_record( + session_id, + now, + data_root.clone(), + Vec::new(), + None, + Some("test-model".into()), + None, + None, + "test-provider".into(), + None, + ); + rollout_store + .append_session_meta(&record) + .expect("append session meta"); + rollout_store + .append_title_update( + &record, + "Updated from rollout".into(), + SessionTitleState::Final(SessionTitleFinalSource::ModelGenerated), + None, + ) + .expect("append title update"); + + let db = crate::db::Database::open(data_root.join("index.db")).expect("open db"); + rollout_store + .index_rollout_metadata(&db) + .expect("index rollout metadata"); + + let index = db + .get_session_index(&session_id) + .expect("get index") + .expect("indexed session"); + assert_eq!( + index.metadata.title.as_deref(), + Some("Updated from rollout") + ); + } + #[test] fn replay_backfills_tool_result_name_from_prior_tool_call() { let mut messages = Vec::new(); diff --git a/crates/server/src/runtime.rs b/crates/server/src/runtime.rs index 2fc1ec7c..59c8cb7d 100644 --- a/crates/server/src/runtime.rs +++ b/crates/server/src/runtime.rs @@ -179,6 +179,7 @@ mod research_streaming; mod research_tool_runtime; mod research_tools; mod session_actor; +mod session_cache; mod session_interactive; mod skills; mod subagent_usage; @@ -248,6 +249,10 @@ pub struct ServerRuntime { active_workspace_baselines: Mutex>, /// Weak back-reference used when session actors need the owning runtime `Arc`. self_weak: std::sync::Weak, + /// LRU order for loaded root session actors. + session_lru: Mutex, + /// Per-session gate that serializes lazy parent session hydration. + parent_session_load_gate: Arc, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -352,6 +357,10 @@ impl ServerRuntime { command_exec_manager: command_exec::CommandExecManager::new(), active_workspace_baselines: Mutex::new(HashMap::new()), self_weak: self_weak.clone(), + session_lru: Mutex::new(session_cache::ParentSessionLru::new( + session_cache::PARENT_SESSION_LRU_CAPACITY, + )), + parent_session_load_gate: Arc::new(session_cache::SessionLoadGate::default()), }) } } diff --git a/crates/server/src/runtime/agents.rs b/crates/server/src/runtime/agents.rs index 370d153e..913078ae 100644 --- a/crates/server/src/runtime/agents.rs +++ b/crates/server/src/runtime/agents.rs @@ -128,6 +128,7 @@ impl ServerRuntime { Some(record) }; + let rollout_path_for_db = record.as_ref().map(|entry| entry.rollout_path.clone()); let mut core_session = runtime_context.new_session_state( child_session_id, parent_summary.cwd.clone(), @@ -281,7 +282,10 @@ impl ServerRuntime { ) .await; if !summary.ephemeral - && let Err(error) = self.deps.db.upsert_session(&summary) + && let Err(error) = self.deps.db.upsert_session( + &summary, + rollout_path_for_db.as_deref().map(std::path::Path::new), + ) { tracing::warn!( session_id = %child_session_id, diff --git a/crates/server/src/runtime/handlers/acp/session_support.rs b/crates/server/src/runtime/handlers/acp/session_support.rs index 15eb7065..1eba3fc1 100644 --- a/crates/server/src/runtime/handlers/acp/session_support.rs +++ b/crates/server/src/runtime/handlers/acp/session_support.rs @@ -9,22 +9,15 @@ impl ServerRuntime { session_id: SessionId, cwd: &Path, ) -> Result<(), (AcpErrorCode, String)> { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let stored_cwd = self + .resolve_indexed_session_cwd(session_id) + .map_err(|error| (AcpErrorCode::ServerError, error.to_string()))?; + let Some(stored_cwd) = stored_cwd else { return Err(( AcpErrorCode::ServerError, "session does not exist".to_string(), )); }; - let stored_cwd = session_arc - .summary() - .await - .ok_or_else(|| { - ( - AcpErrorCode::ServerError, - "session actor unavailable".to_string(), - ) - })? - .cwd; if stored_cwd != cwd { return Err(( AcpErrorCode::InvalidParams, @@ -34,6 +27,14 @@ impl ServerRuntime { Ok(()) } + fn resolve_indexed_session_cwd( + &self, + session_id: SessionId, + ) -> anyhow::Result> { + self.rollout_store + .resolve_indexed_session_cwd(&self.deps.db, &session_id) + } + pub(super) async fn acp_session_tool_registry( &self, method: &str, @@ -111,7 +112,7 @@ impl ServerRuntime { } if !updated_summary.ephemeral - && let Err(error) = self.deps.db.upsert_session(&updated_summary) + && let Err(error) = self.deps.db.upsert_session(&updated_summary, None) { tracing::warn!( session_id = %session_id, diff --git a/crates/server/src/runtime/handlers/acp_config_options.rs b/crates/server/src/runtime/handlers/acp_config_options.rs index ea7eb3ce..d14347f2 100644 --- a/crates/server/src/runtime/handlers/acp_config_options.rs +++ b/crates/server/src/runtime/handlers/acp_config_options.rs @@ -250,7 +250,7 @@ impl ServerRuntime { if let Some(updated_session) = updated_session && !updated_session.ephemeral - && let Err(error) = self.deps.db.upsert_session(&updated_session) + && let Err(error) = self.deps.db.upsert_session(&updated_session, None) { tracing::warn!( session_id = %params.session_id, diff --git a/crates/server/src/runtime/handlers/session.rs b/crates/server/src/runtime/handlers/session.rs index 00dedd9f..f67826f7 100644 --- a/crates/server/src/runtime/handlers/session.rs +++ b/crates/server/src/runtime/handlers/session.rs @@ -106,6 +106,7 @@ impl ServerRuntime { let config = core_session.config.clone(); let pending_turn_queue = Arc::clone(&core_session.pending_turn_queue); let btw_input_queue = Arc::clone(&core_session.btw_input_queue); + let rollout_path_for_db = record.as_ref().map(|entry| entry.rollout_path.clone()); let actor_state = SessionActorState { runtime_context, record, @@ -134,10 +135,16 @@ impl ServerRuntime { self.insert_session_actor(actor_state).await; self.subscribe_connection_to_session(connection_id, session_id, None) .await; + self.runtime_arc() + .after_root_session_insert(session_id) + .await; // Persist session metadata to SQLite (skip for ephemeral sessions) if !summary.ephemeral - && let Err(err) = self.deps.db.upsert_session(&summary) + && let Err(err) = self + .deps + .db + .upsert_session(&summary, rollout_path_for_db.as_deref()) { tracing::warn!( session_id = %session_id, @@ -177,7 +184,35 @@ impl ServerRuntime { } pub(crate) async fn list_session_summaries(&self) -> Vec { - self.list_session_summaries_from_actors().await + let mut sessions_by_id = match self.deps.db.list_root_sessions() { + Ok(sessions) => sessions + .into_iter() + .map(|session| (session.session_id, session)) + .collect::>(), + Err(error) => { + tracing::warn!(error = %error, "failed to list root sessions from database"); + std::collections::HashMap::new() + } + }; + + for handle in self.list_session_handles().await { + let Some(runtime_summary) = handle.summary().await else { + continue; + }; + if runtime_summary.ephemeral || runtime_summary.agent_path.is_some() { + continue; + } + sessions_by_id.insert(runtime_summary.session_id, runtime_summary); + } + + let mut sessions = sessions_by_id.into_values().collect::>(); + sessions.sort_by(|left, right| { + right + .last_activity_at + .cmp(&left.last_activity_at) + .then_with(|| right.updated_at.cmp(&left.updated_at)) + }); + sessions } pub(crate) async fn handle_session_metadata_update( @@ -229,7 +264,7 @@ impl ServerRuntime { // Persist updated session metadata to SQLite if !updated_session.ephemeral - && let Err(err) = self.deps.db.upsert_session(&updated_session) + && let Err(err) = self.deps.db.upsert_session(&updated_session, None) { tracing::warn!( session_id = %params.session_id, @@ -366,7 +401,7 @@ impl ServerRuntime { // Persist updated session metadata to SQLite if !summary.ephemeral - && let Err(err) = self.deps.db.upsert_session(&summary) + && let Err(err) = self.deps.db.upsert_session(&summary, None) { tracing::warn!( session_id = %params.session_id, @@ -419,12 +454,44 @@ impl ServerRuntime { params: SessionResumeParams, tool_registry_update: RuntimeSessionToolRegistryUpdate, ) -> serde_json::Value { - let Some(session_handle) = self.session(params.session_id).await else { - return self.error_response( - request_id, - ProtocolErrorCode::SessionNotFound, - "session does not exist", - ); + let session_handle = match self + .runtime_arc() + .get_or_load_parent_session(params.session_id) + .await + { + Ok(handle) => handle, + Err(crate::runtime::session_cache::LoadSessionError::SessionNotFound) => { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + } + Err(crate::runtime::session_cache::LoadSessionError::SubagentNotResumable { + parent_session_id, + }) => { + return self.error_response( + request_id, + ProtocolErrorCode::InvalidParams, + format!( + "subagent sessions cannot be resumed directly; resume the parent session {parent_session_id} instead" + ), + ); + } + Err(crate::runtime::session_cache::LoadSessionError::RolloutMissing) => { + return self.error_response( + request_id, + ProtocolErrorCode::InternalError, + "session metadata exists but rollout file is missing; session cannot be restored", + ); + } + Err(crate::runtime::session_cache::LoadSessionError::RestoreFailed(message)) => { + return self.error_response( + request_id, + ProtocolErrorCode::InternalError, + format!("failed to restore session: {message}"), + ); + } }; match tool_registry_update { RuntimeSessionToolRegistryUpdate::KeepCurrent => {} @@ -563,10 +630,29 @@ impl ServerRuntime { forked_runtime.record = Some(record); } let summary = forked_runtime.summary.clone(); + let rollout_path_for_db = forked_runtime + .record + .as_ref() + .map(|entry| entry.rollout_path.clone()); self.insert_session_actor(SessionActorState::from_runtime_session(forked_runtime)) .await; self.subscribe_connection_to_session(connection_id, forked_id, None) .await; + self.runtime_arc() + .after_root_session_insert(forked_id) + .await; + if !summary.ephemeral + && let Err(err) = self + .deps + .db + .upsert_session(&summary, rollout_path_for_db.as_deref()) + { + tracing::warn!( + session_id = %forked_id, + error = %err, + "failed to persist forked session metadata to database" + ); + } tracing::info!( connection_id, source_session_id = %params.session_id, diff --git a/crates/server/src/runtime/lifecycle.rs b/crates/server/src/runtime/lifecycle.rs index 2a65e6d4..de3f10fa 100644 --- a/crates/server/src/runtime/lifecycle.rs +++ b/crates/server/src/runtime/lifecycle.rs @@ -1,177 +1,233 @@ use super::*; +use std::path::Path; + +use crate::execution::RuntimeSession; +use crate::runtime::session_actor::SessionActorState; impl ServerRuntime { - /// Loads durable sessions from rollout files and installs them into the runtime map. - /// Also restores token stats and pending queues from SQLite. - pub async fn load_persisted_sessions(self: &Arc) -> anyhow::Result<()> { - let mut sessions = self.rollout_store.load_sessions(&self.deps).await?; - tracing::info!(session_count = sessions.len(), "loaded persisted sessions"); - let mut restored_goal_stores = std::collections::HashMap::new(); + /// Replays one rollout file and applies SQLite side effects for a single session. + pub(crate) async fn hydrate_runtime_session( + self: &Arc, + session_id: SessionId, + rollout_path: &Path, + ) -> anyhow::Result { + let mut runtime_session = self + .rollout_store + .load_session_from_rollout(rollout_path, &self.deps) + .await?; + if runtime_session.summary.session_id != session_id { + anyhow::bail!( + "rollout session id mismatch: expected {session_id}, got {}", + runtime_session.summary.session_id + ); + } + self.apply_persisted_session_side_effects(session_id, &mut runtime_session) + .await?; + Ok(runtime_session) + } - for (session_id, runtime_session) in sessions.iter_mut() { - if !runtime_session.summary.ephemeral - && let Err(err) = self.deps.db.upsert_session(&runtime_session.summary) - { - tracing::warn!( + async fn apply_persisted_session_side_effects( + self: &Arc, + session_id: SessionId, + runtime_session: &mut RuntimeSession, + ) -> anyhow::Result<()> { + if !runtime_session.summary.ephemeral + && let Err(err) = self.deps.db.upsert_session( + &runtime_session.summary, + runtime_session + .record + .as_ref() + .map(|record| record.rollout_path.as_path()), + ) + { + tracing::warn!( + session_id = %session_id, + error = %err, + "failed to seed restored session metadata to database" + ); + } + + match self.deps.db.get_stats(&session_id) { + Ok(Some(stats)) => { + runtime_session.summary.total_input_tokens = stats.total_input_tokens; + runtime_session.summary.total_output_tokens = stats.total_output_tokens; + runtime_session.summary.total_tokens = stats.total_tokens; + runtime_session.summary.total_cache_creation_tokens = + stats.total_cache_creation_tokens; + runtime_session.summary.total_cache_read_tokens = stats.total_cache_read_tokens; + runtime_session.summary.prompt_token_estimate = stats.prompt_token_estimate; + if let Ok(mut core) = runtime_session.core_session.try_lock() { + core.total_input_tokens = stats.total_input_tokens; + core.total_output_tokens = stats.total_output_tokens; + core.total_tokens = stats.total_tokens; + core.total_cache_creation_tokens = stats.total_cache_creation_tokens; + core.total_cache_read_tokens = stats.total_cache_read_tokens; + core.last_input_tokens = stats.last_input_tokens; + core.last_turn_tokens = core.last_turn_tokens.max(stats.last_input_tokens); + core.prompt_token_estimate = stats.prompt_token_estimate; + } + tracing::debug!( session_id = %session_id, - error = %err, - "failed to seed restored session metadata to database" + "restored token stats from database" ); - continue; } - - match self.deps.db.get_stats(session_id) { - Ok(Some(stats)) => { - runtime_session.summary.total_input_tokens = stats.total_input_tokens; - runtime_session.summary.total_output_tokens = stats.total_output_tokens; - runtime_session.summary.total_tokens = stats.total_tokens; - runtime_session.summary.total_cache_creation_tokens = - stats.total_cache_creation_tokens; - runtime_session.summary.total_cache_read_tokens = stats.total_cache_read_tokens; - runtime_session.summary.prompt_token_estimate = stats.prompt_token_estimate; - if let Ok(mut core) = runtime_session.core_session.try_lock() { - core.total_input_tokens = stats.total_input_tokens; - core.total_output_tokens = stats.total_output_tokens; - core.total_tokens = stats.total_tokens; - core.total_cache_creation_tokens = stats.total_cache_creation_tokens; - core.total_cache_read_tokens = stats.total_cache_read_tokens; - core.last_input_tokens = stats.last_input_tokens; - core.last_turn_tokens = core.last_turn_tokens.max(stats.last_input_tokens); - core.prompt_token_estimate = stats.prompt_token_estimate; - } - tracing::debug!( - session_id = %session_id, - "restored token stats from database" - ); - } - Ok(None) => { - let stats = crate::db::SessionStats { - total_input_tokens: runtime_session.summary.total_input_tokens, - total_output_tokens: runtime_session.summary.total_output_tokens, - total_tokens: runtime_session.summary.total_tokens, - total_cache_creation_tokens: runtime_session - .summary - .total_cache_creation_tokens, - total_cache_read_tokens: runtime_session.summary.total_cache_read_tokens, - last_input_tokens: 0, - turn_count: 0, - prompt_token_estimate: runtime_session.summary.prompt_token_estimate, - }; - if let Err(err) = self.deps.db.update_stats(session_id, &stats) { - tracing::warn!( - session_id = %session_id, - error = %err, - "failed to persist initial token stats to database" - ); - } - } - Err(err) => { + Ok(None) => { + let stats = crate::db::SessionStats { + total_input_tokens: runtime_session.summary.total_input_tokens, + total_output_tokens: runtime_session.summary.total_output_tokens, + total_tokens: runtime_session.summary.total_tokens, + total_cache_creation_tokens: runtime_session + .summary + .total_cache_creation_tokens, + total_cache_read_tokens: runtime_session.summary.total_cache_read_tokens, + last_input_tokens: 0, + turn_count: 0, + prompt_token_estimate: runtime_session.summary.prompt_token_estimate, + }; + if let Err(err) = self.deps.db.update_stats(&session_id, &stats) { tracing::warn!( session_id = %session_id, error = %err, - "failed to load token stats from database" + "failed to persist initial token stats to database" ); } } + Err(err) => { + tracing::warn!( + session_id = %session_id, + error = %err, + "failed to load token stats from database" + ); + } + } - match self - .deps - .db - .drain_pending(session_id, crate::db::QueueType::Turn) - { - Ok(items) => { - if !items.is_empty() { - let core_session = runtime_session.core_session.lock().await; - let mut queue = core_session - .pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned"); - queue.extend(items); - tracing::debug!( - session_id = %session_id, - pending_count = queue.len(), - "restored pending turn queue from database" - ); - } - } - Err(err) => { - tracing::warn!( + match self + .deps + .db + .drain_pending(&session_id, crate::db::QueueType::Turn) + { + Ok(items) => { + if !items.is_empty() { + let core_session = runtime_session.core_session.lock().await; + let mut queue = core_session + .pending_turn_queue + .lock() + .expect("pending turn queue mutex should not be poisoned"); + queue.extend(items); + tracing::debug!( session_id = %session_id, - error = %err, - "failed to load pending turn queue from database" + pending_count = queue.len(), + "restored pending turn queue from database" ); } } - - if let Err(err) = self - .deps - .db - .clear_pending(session_id, crate::db::QueueType::Btw) - { + Err(err) => { tracing::warn!( session_id = %session_id, error = %err, - "failed to clear stale btw inputs from database" + "failed to load pending turn queue from database" ); } + } - match self.goal_durable_store.replay_goal_store(*session_id).await { - Ok(Some(mut goal_store)) => { - if let Some(goal) = goal_store.get() - && goal.status == crate::goal::GoalStatus::Active - { - let previous_status = goal.status; - match goal_store.set_status(devo_protocol::ThreadGoalStatus::Paused) { - Ok(paused_goal) => { - if let Err(error) = self - .goal_durable_store - .append_status_changed( - &paused_goal, - previous_status, - Some( - "Goal paused because the session was restored without explicit resume." - .to_string(), - ), - ) - .await - { - tracing::warn!( - session_id = %session_id, - error = %error, - "failed to persist restored goal pause record" - ); - } - } - Err(error) => { + if let Err(err) = self + .deps + .db + .clear_pending(&session_id, crate::db::QueueType::Btw) + { + tracing::warn!( + session_id = %session_id, + error = %err, + "failed to clear stale btw inputs from database" + ); + } + + match self.goal_durable_store.replay_goal_store(session_id).await { + Ok(Some(mut goal_store)) => { + if let Some(goal) = goal_store.get() + && goal.status == crate::goal::GoalStatus::Active + { + let previous_status = goal.status; + match goal_store.set_status(devo_protocol::ThreadGoalStatus::Paused) { + Ok(paused_goal) => { + if let Err(error) = self + .goal_durable_store + .append_status_changed( + &paused_goal, + previous_status, + Some( + "Goal paused because the session was restored without explicit resume." + .to_string(), + ), + ) + .await + { tracing::warn!( session_id = %session_id, error = %error, - "failed to pause restored active goal" + "failed to persist restored goal pause record" ); } } + Err(error) => { + tracing::warn!( + session_id = %session_id, + error = %error, + "failed to pause restored active goal" + ); + } } - restored_goal_stores.insert(*session_id, goal_store); - } - Ok(None) => {} - Err(error) => { - tracing::warn!( - session_id = %session_id, - error = %error, - "failed to replay durable goal records" - ); } + self.goal_stores.lock().await.insert(session_id, goal_store); + } + Ok(None) => {} + Err(error) => { + tracing::warn!( + session_id = %session_id, + error = %error, + "failed to replay durable goal records" + ); } } - for (_, runtime_session) in sessions { - self.insert_session_actor(SessionActorState::from_runtime_session(runtime_session)) - .await; + Ok(()) + } + + /// Loads durable sessions from rollout files and installs them into the runtime map. + /// Used by integration tests and bulk-restore tooling; production startup indexes only. + pub async fn load_persisted_sessions(self: &Arc) -> anyhow::Result<()> { + let sessions = self.rollout_store.load_sessions(&self.deps).await?; + tracing::info!(session_count = sessions.len(), "loaded persisted sessions"); + + for (session_id, mut runtime_session) in sessions { + self.apply_persisted_session_side_effects(session_id, &mut runtime_session) + .await?; + if runtime_session.summary.parent_session_id.is_none() { + self.insert_root_session_actor(runtime_session) + .await + .map_err(|error| anyhow::anyhow!("{error}"))?; + } else { + self.insert_session_actor(SessionActorState::from_runtime_session(runtime_session)) + .await; + } } - self.goal_stores.lock().await.extend(restored_goal_stores); Ok(()) } + /// Rebuilds the SQLite session index from rollout SessionMeta headers. + pub fn refresh_session_index(&self) -> anyhow::Result<()> { + self.rollout_store.index_rollout_metadata(&self.deps.db) + } + + /// Backfills rollout metadata into SQLite only when legacy rows lack index fields. + pub fn backfill_session_index_if_required(&self) -> anyhow::Result { + if !self.deps.db.session_index_backfill_required()? { + return Ok(false); + } + self.rollout_store.index_rollout_metadata(&self.deps.db)?; + Ok(true) + } + /// Completes deferred (in-progress) items for all active turns and /// persists interrupted turn records. Called on graceful shutdown. pub async fn shutdown(self: &Arc) { diff --git a/crates/server/src/runtime/session_actor/registry.rs b/crates/server/src/runtime/session_actor/registry.rs index 8286a8d2..597964b0 100644 --- a/crates/server/src/runtime/session_actor/registry.rs +++ b/crates/server/src/runtime/session_actor/registry.rs @@ -23,11 +23,18 @@ impl ServerRuntime { ) -> super::SessionHandle { let runtime = self.runtime_arc(); let session_id = state.session_id(); + if let Some(existing) = self.sessions.lock().await.get(&session_id).cloned() { + return existing; + } + let handle = super::SessionHandle::spawn(session_id, state, runtime); - self.sessions - .lock() - .await - .insert(session_id, handle.clone()); + let mut sessions = self.sessions.lock().await; + if let Some(existing) = sessions.get(&session_id).cloned() { + drop(sessions); + handle.shutdown().await; + return existing; + } + sessions.insert(session_id, handle.clone()); handle } @@ -45,6 +52,7 @@ impl ServerRuntime { self.sessions.lock().await.values().cloned().collect() } + #[allow(dead_code)] pub(crate) async fn list_session_summaries_from_actors( &self, ) -> Vec { diff --git a/crates/server/src/runtime/session_cache.rs b/crates/server/src/runtime/session_cache.rs new file mode 100644 index 00000000..33725c72 --- /dev/null +++ b/crates/server/src/runtime/session_cache.rs @@ -0,0 +1,297 @@ +use std::collections::{HashMap, VecDeque}; +use std::sync::Arc; + +use tokio::sync::Mutex; + +use devo_protocol::SessionId; + +use crate::execution::RuntimeSession; +use crate::runtime::ServerRuntime; +use crate::runtime::session_actor::SessionActorState; +use crate::runtime::session_actor::SessionHandle; + +pub(crate) const PARENT_SESSION_LRU_CAPACITY: usize = 16; + +#[derive(Debug, Default)] +pub(crate) struct SessionLoadGate { + locks: Mutex>>>, +} + +pub(crate) struct SessionLoadPermit { + session_id: SessionId, + lock: Arc>, + gate: Arc, + guard: Option>, +} + +impl SessionLoadGate { + pub(crate) async fn acquire(self: &Arc, session_id: SessionId) -> SessionLoadPermit { + let lock = { + let mut locks = self.locks.lock().await; + locks + .entry(session_id) + .or_insert_with(|| Arc::new(Mutex::new(()))) + .clone() + }; + let guard = lock.clone().lock_owned().await; + SessionLoadPermit { + session_id, + lock, + gate: Arc::clone(self), + guard: Some(guard), + } + } +} + +impl Drop for SessionLoadPermit { + fn drop(&mut self) { + drop(self.guard.take()); + if Arc::strong_count(&self.lock) != 2 { + return; + } + let session_id = self.session_id; + let lock = Arc::clone(&self.lock); + let gate = Arc::clone(&self.gate); + tokio::spawn(async move { + let mut locks = gate.locks.lock().await; + if locks.get(&session_id).is_some_and(|existing| { + Arc::ptr_eq(existing, &lock) && Arc::strong_count(existing) == 2 + }) { + locks.remove(&session_id); + } + }); + } +} + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub(crate) enum LoadSessionError { + #[error("session not found")] + SessionNotFound, + #[error( + "subagent sessions cannot be resumed directly; resume the parent session {parent_session_id} instead" + )] + SubagentNotResumable { parent_session_id: SessionId }, + #[error("session metadata exists but rollout file is missing; session cannot be restored")] + RolloutMissing, + #[error("failed to restore session: {0}")] + RestoreFailed(String), +} + +#[derive(Debug, Clone, Default)] +pub(crate) struct ParentSessionLru { + order: VecDeque, + capacity: usize, +} + +impl ParentSessionLru { + pub(crate) fn new(capacity: usize) -> Self { + Self { + order: VecDeque::new(), + capacity, + } + } + + pub(crate) fn touch(&mut self, session_id: SessionId) { + self.order.retain(|id| *id != session_id); + self.order.push_front(session_id); + } + + pub(crate) fn remove(&mut self, session_id: SessionId) { + self.order.retain(|id| *id != session_id); + } + + pub(crate) fn len(&self) -> usize { + self.order.len() + } +} + +impl ServerRuntime { + pub(crate) async fn get_or_load_parent_session( + self: &Arc, + session_id: SessionId, + ) -> Result { + if let Some(handle) = self.session(session_id).await { + self.touch_parent_session_lru(session_id).await; + return Ok(handle); + } + + let _load_permit = self.parent_session_load_gate.acquire(session_id).await; + if let Some(handle) = self.session(session_id).await { + self.touch_parent_session_lru(session_id).await; + return Ok(handle); + } + + let index = self + .deps + .db + .get_session_index(&session_id) + .map_err(|error| LoadSessionError::RestoreFailed(error.to_string()))? + .ok_or(LoadSessionError::SessionNotFound)?; + + if index.metadata.agent_path.is_some() { + let parent_session_id = index + .metadata + .parent_session_id + .ok_or(LoadSessionError::SessionNotFound)?; + return Err(LoadSessionError::SubagentNotResumable { parent_session_id }); + } + + let stored_rollout_path = index.rollout_path.clone(); + let rollout_path = match stored_rollout_path { + Some(ref path) if path.exists() => path.clone(), + Some(_) | None => self + .rollout_store + .find_rollout_by_session_id(&session_id) + .map_err(|error| LoadSessionError::RestoreFailed(error.to_string()))? + .filter(|path| path.exists()) + .ok_or(LoadSessionError::RolloutMissing)?, + }; + + if stored_rollout_path.as_ref() != Some(&rollout_path) + && let Err(error) = self + .deps + .db + .upsert_rollout_index_session(&index.metadata, Some(rollout_path.as_path())) + { + tracing::warn!( + session_id = %session_id, + error = %error, + "failed to backfill rollout_path after suffix lookup" + ); + } + + let runtime_session = self + .hydrate_runtime_session(session_id, &rollout_path) + .await + .map_err(|error| LoadSessionError::RestoreFailed(error.to_string()))?; + self.insert_root_session_actor(runtime_session).await + } + + pub(crate) async fn insert_root_session_actor( + self: &Arc, + runtime_session: RuntimeSession, + ) -> Result { + let session_id = runtime_session.summary.session_id; + let handle = self + .insert_session_actor(SessionActorState::from_runtime_session(runtime_session)) + .await; + self.touch_parent_session_lru(session_id).await; + self.evict_parent_sessions_if_needed(Some(session_id)).await; + Ok(handle) + } + + pub(crate) async fn after_root_session_insert(self: &Arc, session_id: SessionId) { + self.touch_parent_session_lru(session_id).await; + self.evict_parent_sessions_if_needed(Some(session_id)).await; + } + + pub(crate) async fn touch_parent_session_lru(&self, session_id: SessionId) { + let mut lru = self.session_lru.lock().await; + lru.touch(session_id); + } + + async fn evict_parent_sessions_if_needed(self: &Arc, exclude: Option) { + let candidates: Vec = { + let lru = self.session_lru.lock().await; + if lru.len() <= lru.capacity { + return; + } + lru.order.iter().rev().copied().collect() + }; + for session_id in candidates { + if exclude == Some(session_id) { + continue; + } + if self.is_parent_session_pinned(session_id).await { + continue; + } + self.evict_parent_session_cascade(session_id).await; + return; + } + } + + pub(crate) async fn is_parent_session_pinned(&self, session_id: SessionId) -> bool { + if self.active_turns.has_session(session_id).await { + return true; + } + let connections = self.connections.lock().await; + for connection in connections.values() { + if connection + .subscriptions + .iter() + .any(|subscription| subscription.session_id == Some(session_id)) + { + return true; + } + } + false + } + + pub(crate) async fn evict_parent_session_cascade(&self, parent_session_id: SessionId) { + let child_session_ids = { + let registries = self.agent_registries.lock().await; + registries + .get(&parent_session_id) + .map(|registry| registry.children_of(parent_session_id)) + .unwrap_or_default() + }; + + for child_session_id in child_session_ids { + if let Some(handle) = self.remove_session_actor(child_session_id).await { + handle.shutdown().await; + } + self.goal_stores.lock().await.remove(&child_session_id); + } + + if let Some(handle) = self.remove_session_actor(parent_session_id).await { + handle.shutdown().await; + } + self.goal_stores.lock().await.remove(&parent_session_id); + self.session_lru.lock().await.remove(parent_session_id); + self.agent_registries + .lock() + .await + .remove(&parent_session_id); + self.agent_mailboxes.lock().await.remove(&parent_session_id); + self.agent_output_buffers + .lock() + .await + .remove(&parent_session_id); + self.agent_wait_cursors + .lock() + .await + .remove(&parent_session_id); + self.research_child_agents + .lock() + .await + .remove(&parent_session_id); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parent_session_lru_tracks_recency_for_eviction() { + let mut lru = ParentSessionLru::new(2); + let first = SessionId::new(); + let second = SessionId::new(); + let third = SessionId::new(); + lru.touch(first); + lru.touch(second); + lru.touch(third); + assert_eq!(lru.len(), 3); + lru.touch(second); + assert_eq!(lru.len(), 3); + } + + #[test] + fn parent_session_lru_capacity_allows_exact_capacity() { + let mut lru = ParentSessionLru::new(PARENT_SESSION_LRU_CAPACITY); + for _ in 0..PARENT_SESSION_LRU_CAPACITY { + lru.touch(SessionId::new()); + } + assert_eq!(lru.len(), PARENT_SESSION_LRU_CAPACITY); + } +} diff --git a/crates/server/tests/acp_session_lifecycle.rs b/crates/server/tests/acp_session_lifecycle.rs index 9aef8537..78769142 100644 --- a/crates/server/tests/acp_session_lifecycle.rs +++ b/crates/server/tests/acp_session_lifecycle.rs @@ -1072,7 +1072,11 @@ async fn create_acp_session( ) .await .context("session/new response")?; - let response: AcpSuccessResponse = serde_json::from_value(response)?; + if response.get("result").is_none() { + anyhow::bail!("session/new error response: {response}"); + } + let response: AcpSuccessResponse = + serde_json::from_value(response).context("decode session/new success response")?; Ok(response.result.session_id) } @@ -1097,7 +1101,15 @@ async fn list_acp_sessions( ) .await .context("session/list response")?; - Ok(serde_json::from_value::>(response)?.result) + if response.get("result").is_none() { + panic!("session/list error response: {response}"); + } + let response_for_error = response.clone(); + serde_json::from_value::>(response) + .map(|success| success.result) + .map_err(|error| { + anyhow::anyhow!("session/list decode failed: {error}; response={response_for_error}") + }) } async fn wait_for_response( @@ -1284,6 +1296,92 @@ async fn assert_auth_required( Ok(()) } +#[tokio::test] +async fn acp_session_list_includes_live_actor_missing_from_database() -> Result<()> { + let data_root = TempDir::new()?; + let runtime = build_runtime(data_root.path())?; + let (connection_id, _notifications_rx) = initialize_acp_connection(&runtime).await?; + let cwd = data_root.path().join("repo"); + std::fs::create_dir_all(&cwd)?; + let session_id = create_acp_session(&runtime, connection_id, &cwd, 10).await?; + + let db = devo_server::db::Database::open(data_root.path().join("test_persistence.db"))?; + db.delete_session(&session_id)?; + + let result = list_acp_sessions(&runtime, connection_id, 11, None, None).await?; + let listed = result + .sessions + .iter() + .filter_map(|session| session.meta.as_ref()) + .filter_map(|meta| meta.get(devo_server::DEVO_SESSION_META)) + .map(|value| serde_json::from_value::(value.clone())) + .collect::, _>>()?; + assert_eq!( + listed + .iter() + .filter(|session| session.session_id == session_id) + .count(), + 1 + ); + Ok(()) +} + +#[tokio::test] +async fn acp_session_resume_lazy_loads_without_memory_actor() -> Result<()> { + let data_root = TempDir::new()?; + let runtime = build_runtime(data_root.path())?; + let (connection_id, mut notifications_rx) = initialize_acp_connection(&runtime).await?; + let cwd = data_root.path().join("repo"); + std::fs::create_dir_all(&cwd)?; + let session_id = create_acp_session(&runtime, connection_id, &cwd, 10).await?; + + let prompt_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 11, + "method": "session/prompt", + "params": { + "sessionId": session_id, + "prompt": [ + { + "type": "text", + "text": "persist for lazy ACP resume" + } + ] + } + }), + ) + .await; + assert_eq!(prompt_response, None); + let _: AcpSuccessResponse = + wait_for_response(&mut notifications_rx, 11).await?; + + let restored_runtime = build_runtime(data_root.path())?; + restored_runtime.refresh_session_index()?; + let (restored_connection_id, mut restored_notifications_rx) = + initialize_acp_connection(&restored_runtime).await?; + let resume_response = restored_runtime + .handle_incoming( + restored_connection_id, + serde_json::json!({ + "id": 12, + "method": "session/resume", + "params": { + "sessionId": session_id, + "cwd": path_value(&cwd), + "mcpServers": [] + } + }), + ) + .await + .context("session/resume response")?; + assert!(resume_response["result"].is_object()); + let _: AcpSuccessResponse = serde_json::from_value(resume_response)?; + assert_no_replayed_history(&mut restored_notifications_rx).await?; + Ok(()) +} + async fn assert_legacy_session_method_removed( runtime: &Arc, connection_id: u64, diff --git a/crates/server/tests/persistence_resume.rs b/crates/server/tests/persistence_resume.rs index d7a96f0b..a69e8606 100644 --- a/crates/server/tests/persistence_resume.rs +++ b/crates/server/tests/persistence_resume.rs @@ -1706,3 +1706,405 @@ async fn interrupt_mid_stream_does_not_duplicate_last_item_on_resume() -> Result Ok(()) } + +#[tokio::test] +async fn lazy_resume_loads_parent_session_from_rollout_on_map_miss() -> Result<()> { + let data_root = TempDir::new()?; + let runtime = build_runtime(data_root.path())?; + let (connection_id, mut notifications_rx) = initialize_connection(&runtime).await?; + + let start_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 1, + "method": "session/start", + "params": { + "cwd": data_root.path(), + "ephemeral": false, + "title": "Lazy resume session", + "model": "test-model" + } + }), + ) + .await + .context("session/start response")?; + let session_id = serde_json::from_value::< + devo_server::SuccessResponse, + >(start_response)? + .result + .session + .session_id; + + let turn_start_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 2, + "method": "_devo/turn/start", + "params": { + "session_id": session_id, + "input": [{ "type": "text", "text": "persist for lazy resume" }], + "model": null, + "sandbox": null, + "approval_policy": null, + "cwd": null + } + }), + ) + .await + .context("turn/start response")?; + let _: devo_server::SuccessResponse = + serde_json::from_value(turn_start_response)?; + wait_for_turn_completed(&mut notifications_rx).await?; + + let rebuilt_runtime = build_runtime(data_root.path())?; + rebuilt_runtime.refresh_session_index()?; + + let (rebuilt_connection_id, _rebuilt_notifications_rx) = + initialize_connection(&rebuilt_runtime).await?; + let resume_response = rebuilt_runtime + .handle_incoming( + rebuilt_connection_id, + serde_json::json!({ + "id": 3, + "method": "_devo/session/resume", + "params": { "session_id": session_id } + }), + ) + .await + .context("session/resume response")?; + let resume_result = serde_json::from_value::< + devo_server::SuccessResponse, + >(resume_response)? + .result; + assert_eq!(resume_result.session.session_id, session_id); + Ok(()) +} + +#[tokio::test] +async fn lazy_resume_after_compat_backfill_without_refresh_session_index() -> Result<()> { + let data_root = TempDir::new()?; + let runtime = build_runtime(data_root.path())?; + let (connection_id, mut notifications_rx) = initialize_connection(&runtime).await?; + + let start_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 1, + "method": "session/start", + "params": { + "cwd": data_root.path(), + "ephemeral": false, + "title": "Lazy resume backfill", + "model": "test-model" + } + }), + ) + .await + .context("session/start response")?; + let session_id = serde_json::from_value::< + devo_server::SuccessResponse, + >(start_response)? + .result + .session + .session_id; + + let turn_start_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 2, + "method": "_devo/turn/start", + "params": { + "session_id": session_id, + "input": [{ "type": "text", "text": "persist for compat backfill" }], + "model": null, + "sandbox": null, + "approval_policy": null, + "cwd": null + } + }), + ) + .await + .context("turn/start response")?; + let _: devo_server::SuccessResponse = + serde_json::from_value(turn_start_response)?; + wait_for_turn_completed(&mut notifications_rx).await?; + + let db_path = data_root.path().join("test_persistence.db"); + { + let conn = rusqlite::Connection::open(&db_path)?; + conn.execute( + "UPDATE sessions SET rollout_path = NULL WHERE id = ?1", + rusqlite::params![session_id.to_string()], + )?; + } + + let rebuilt_runtime = build_runtime(data_root.path())?; + assert!(rebuilt_runtime.backfill_session_index_if_required()?); + + let (rebuilt_connection_id, _) = initialize_connection(&rebuilt_runtime).await?; + let resume_response = rebuilt_runtime + .handle_incoming( + rebuilt_connection_id, + serde_json::json!({ + "id": 3, + "method": "_devo/session/resume", + "params": { "session_id": session_id } + }), + ) + .await + .context("session/resume response")?; + let resume_result = serde_json::from_value::< + devo_server::SuccessResponse, + >(resume_response)? + .result; + assert_eq!(resume_result.session.session_id, session_id); + Ok(()) +} + +#[tokio::test] +async fn concurrent_lazy_resume_single_actor() -> Result<()> { + let data_root = TempDir::new()?; + let runtime = build_runtime(data_root.path())?; + let (connection_id, mut notifications_rx) = initialize_connection(&runtime).await?; + + let start_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 1, + "method": "session/start", + "params": { + "cwd": data_root.path(), + "ephemeral": false, + "title": "Concurrent lazy resume", + "model": "test-model" + } + }), + ) + .await + .context("session/start response")?; + let session_id = serde_json::from_value::< + devo_server::SuccessResponse, + >(start_response)? + .result + .session + .session_id; + + let turn_start_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 2, + "method": "_devo/turn/start", + "params": { + "session_id": session_id, + "input": [{ "type": "text", "text": "persist for concurrent resume" }], + "model": null, + "sandbox": null, + "approval_policy": null, + "cwd": null + } + }), + ) + .await + .context("turn/start response")?; + let _: devo_server::SuccessResponse = + serde_json::from_value(turn_start_response)?; + wait_for_turn_completed(&mut notifications_rx).await?; + + let rebuilt_runtime = build_runtime(data_root.path())?; + rebuilt_runtime.refresh_session_index()?; + let (connection_a, _) = initialize_connection(&rebuilt_runtime).await?; + let (connection_b, _) = initialize_connection(&rebuilt_runtime).await?; + + let request_a = rebuilt_runtime.handle_incoming( + connection_a, + serde_json::json!({ + "id": 3, + "method": "_devo/session/resume", + "params": { "session_id": session_id } + }), + ); + let request_b = rebuilt_runtime.handle_incoming( + connection_b, + serde_json::json!({ + "id": 4, + "method": "_devo/session/resume", + "params": { "session_id": session_id } + }), + ); + let (response_a, response_b) = tokio::join!(request_a, request_b); + let response_a = response_a.context("first session/resume response")?; + let response_b = response_b.context("second session/resume response")?; + let result_a = serde_json::from_value::< + devo_server::SuccessResponse, + >(response_a)? + .result; + let result_b = serde_json::from_value::< + devo_server::SuccessResponse, + >(response_b)? + .result; + assert_eq!(result_a.session.session_id, session_id); + assert_eq!(result_b.session.session_id, session_id); + + Ok(()) +} + +#[tokio::test] +async fn lazy_resume_rejects_subagent_session_ids() -> Result<()> { + let data_root = TempDir::new()?; + let parent_id = SessionId::new(); + let child_id = SessionId::new(); + let now = chrono::Utc::now(); + let db = devo_server::db::Database::open(data_root.path().join("test_persistence.db"))?; + let mut parent = sample_indexed_session(parent_id, data_root.path(), now, None); + parent.title = Some("Parent".into()); + db.upsert_session(&parent, Some("/tmp/parent.jsonl".as_ref()))?; + let mut child = sample_indexed_session(child_id, data_root.path(), now, Some(parent_id)); + child.title = Some("Child".into()); + child.agent_path = Some("root/subagent".into()); + db.upsert_session(&child, Some("/tmp/child.jsonl".as_ref()))?; + let runtime = build_runtime(data_root.path())?; + + let (connection_id, _) = initialize_connection(&runtime).await?; + let resume_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 1, + "method": "_devo/session/resume", + "params": { "session_id": child_id } + }), + ) + .await + .context("session/resume response")?; + let error = serde_json::from_value::(resume_response)?; + assert_eq!( + error.error.code, + devo_server::ProtocolErrorCode::InvalidParams + ); + assert!( + error + .error + .message + .contains("subagent sessions cannot be resumed directly") + ); + Ok(()) +} + +#[tokio::test] +async fn lazy_resume_fails_when_rollout_file_is_missing() -> Result<()> { + let data_root = TempDir::new()?; + let runtime = build_runtime(data_root.path())?; + let (connection_id, mut notifications_rx) = initialize_connection(&runtime).await?; + + let start_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 1, + "method": "session/start", + "params": { + "cwd": data_root.path(), + "ephemeral": false, + "title": "Missing rollout", + "model": "test-model" + } + }), + ) + .await + .context("session/start response")?; + let session_id = serde_json::from_value::< + devo_server::SuccessResponse, + >(start_response)? + .result + .session + .session_id; + + let turn_start_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 2, + "method": "_devo/turn/start", + "params": { + "session_id": session_id, + "input": [{ "type": "text", "text": "create rollout then delete it" }], + "model": null, + "sandbox": null, + "approval_policy": null, + "cwd": null + } + }), + ) + .await + .context("turn/start response")?; + let _: devo_server::SuccessResponse = + serde_json::from_value(turn_start_response)?; + wait_for_turn_completed(&mut notifications_rx).await?; + + let db = devo_server::db::Database::open(data_root.path().join("test_persistence.db"))?; + let index = db.get_session_index(&session_id)?.expect("indexed session"); + let rollout_path = index.rollout_path.expect("rollout path"); + std::fs::remove_file(&rollout_path)?; + + let rebuilt_runtime = build_runtime(data_root.path())?; + let (rebuilt_connection_id, _) = initialize_connection(&rebuilt_runtime).await?; + let resume_response = rebuilt_runtime + .handle_incoming( + rebuilt_connection_id, + serde_json::json!({ + "id": 3, + "method": "_devo/session/resume", + "params": { "session_id": session_id } + }), + ) + .await + .context("session/resume response")?; + let error = serde_json::from_value::(resume_response)?; + assert_eq!( + error.error.code, + devo_server::ProtocolErrorCode::InternalError + ); + assert!(error.error.message.contains("rollout file is missing")); + Ok(()) +} + +fn sample_indexed_session( + session_id: SessionId, + cwd: &std::path::Path, + now: chrono::DateTime, + parent_session_id: Option, +) -> devo_server::SessionMetadata { + devo_server::SessionMetadata { + session_id, + cwd: cwd.to_path_buf(), + additional_directories: Vec::new(), + created_at: now, + updated_at: now, + last_activity_at: now, + title: None, + title_state: devo_core::SessionTitleState::Unset, + parent_session_id, + agent_path: None, + agent_nickname: None, + agent_role: None, + ephemeral: false, + model: Some("test-model".into()), + model_binding_id: None, + reasoning_effort_selection: None, + reasoning_effort: None, + total_input_tokens: 0, + total_output_tokens: 0, + total_tokens: 0, + total_cache_creation_tokens: 0, + total_cache_read_tokens: 0, + prompt_token_estimate: 0, + last_query_total_tokens: 0, + status: devo_protocol::SessionRuntimeStatus::Idle, + } +} From a3667664ae63166cd9cf987b4190fbcf6c3a9673 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Tue, 7 Jul 2026 00:10:41 +0800 Subject: [PATCH 4/4] fix: Harden lazy session resume state --- crates/server/src/runtime/research.rs | 19 +- .../server/src/runtime/session_actor/turn.rs | 8 +- crates/server/src/runtime/subagent_usage.rs | 116 +++++++++- crates/server/tests/persistence_resume.rs | 198 ++++++++++++++++++ 4 files changed, 326 insertions(+), 15 deletions(-) diff --git a/crates/server/src/runtime/research.rs b/crates/server/src/runtime/research.rs index cc5bd45b..153482f3 100644 --- a/crates/server/src/runtime/research.rs +++ b/crates/server/src/runtime/research.rs @@ -378,18 +378,25 @@ impl ServerRuntime { self.capture_turn_workspace_baseline(session_id, turn.turn_id, PathBuf::from(cwd.clone())) .await; let usage_ledger = self.research_usage_ledger(session_id).await; - self.begin_parent_usage_turn( - session_id, - turn.turn_id, - Some(turn_config.model.context_window as u64), - ) - .await; // Research runs outside the session actor (unlike execute_turn_in_actor). // Register the same TurnInlineState / active_stream path so item persist, // usage, and hooks never flood the actor mailbox (capacity 64). Without // this, every token/tool item does blocking send().await and can stall // client inbound handlers that also need the mailbox. self.begin_research_turn_stream(session_id, &turn).await; + let usage_context_window = Some(turn_config.model.context_window as u64); + if let Some(summary) = self.session_summary_snapshot(session_id).await { + self.begin_parent_usage_turn_with_base( + session_id, + turn.turn_id, + crate::runtime::subagent_usage::UsageTotals::from_session_summary(&summary), + usage_context_window, + ) + .await; + } else { + self.begin_parent_usage_turn(session_id, turn.turn_id, usage_context_window) + .await; + } let result = self .run_research_pipeline( session_id, diff --git a/crates/server/src/runtime/session_actor/turn.rs b/crates/server/src/runtime/session_actor/turn.rs index d0003e8e..686d292c 100644 --- a/crates/server/src/runtime/session_actor/turn.rs +++ b/crates/server/src/runtime/session_actor/turn.rs @@ -4,6 +4,7 @@ use tokio::sync::mpsc; use crate::runtime::ServerRuntime; use crate::runtime::session_actor::state::SessionActorState; +use crate::runtime::subagent_usage::UsageTotals; use crate::runtime::turn_exec::{ ExecuteTurnRequest, FinalizeTurnParams, QUERY_EVENT_CHANNEL_CAPACITY, TurnModelQueryParams, spawn_turn_event_stream, @@ -57,7 +58,12 @@ pub(super) async fn execute_turn_in_actor( // updates on the wrong session. if usage_parent_session_id.is_none() { runtime - .begin_parent_usage_turn(session_id, turn.turn_id, usage_context_window) + .begin_parent_usage_turn_with_base( + session_id, + turn.turn_id, + UsageTotals::from_session_summary(&state.summary), + usage_context_window, + ) .await; } diff --git a/crates/server/src/runtime/subagent_usage.rs b/crates/server/src/runtime/subagent_usage.rs index c6f29d1b..8ec0db51 100644 --- a/crates/server/src/runtime/subagent_usage.rs +++ b/crates/server/src/runtime/subagent_usage.rs @@ -172,20 +172,35 @@ impl SubagentUsageState { base_session_totals: UsageTotals, context_window: Option, ) { + let key = ParentTurnKey { + session_id, + turn_id, + }; + if let Some(turn) = self.parent_turns.get_mut(&key) { + if turn.base_session_totals == UsageTotals::default() + && turn.parent_turn_usage == UsageTotals::default() + && turn.inflight_usage == UsageTotals::default() + && base_session_totals != UsageTotals::default() + { + turn.base_session_totals = base_session_totals; + } + if context_window.is_some() { + turn.context_window = context_window; + } + return; + } let base_child_totals = self.child_totals_for_parent_session(session_id); - self.parent_turns - .entry(ParentTurnKey { - session_id, - turn_id, - }) - .or_insert(ParentTurnUsage { + self.parent_turns.insert( + key, + ParentTurnUsage { base_session_totals, base_child_totals, parent_turn_usage: UsageTotals::default(), inflight_usage: UsageTotals::default(), latest_query_usage: UsageTotals::default(), context_window, - }); + }, + ); } pub(super) fn register_child_owner( @@ -409,7 +424,30 @@ impl ServerRuntime { .session_summary_snapshot(session_id) .await .map(|summary| UsageTotals::from_session_summary(&summary)) - .unwrap_or_default(); + .unwrap_or_else(|| { + tracing::warn!( + session_id = %session_id, + turn_id = %turn_id, + "starting usage ledger with zero base because session summary snapshot was unavailable" + ); + UsageTotals::default() + }); + self.begin_parent_usage_turn_with_base( + session_id, + turn_id, + base_session_totals, + context_window, + ) + .await; + } + + pub(super) async fn begin_parent_usage_turn_with_base( + &self, + session_id: SessionId, + turn_id: TurnId, + base_session_totals: UsageTotals, + context_window: Option, + ) { self.subagent_usage.lock().await.begin_parent_turn( session_id, turn_id, @@ -658,6 +696,68 @@ mod tests { } } + #[test] + fn parent_turn_zero_base_can_be_corrected_before_usage_records() { + let mut state = SubagentUsageState::default(); + let parent_session_id = SessionId::new(); + let parent_turn_id = TurnId::new(); + + state.begin_parent_turn( + parent_session_id, + parent_turn_id, + UsageTotals::default(), + None, + ); + state.begin_parent_turn(parent_session_id, parent_turn_id, totals(100, 10), None); + + let snapshot = state + .record_parent_turn_usage( + parent_session_id, + parent_turn_id, + totals(5, 1), + None, + UsageUpdateKind::InFlight, + ) + .expect("parent snapshot"); + assert_eq!(snapshot.session_totals, totals(105, 11)); + } + + #[test] + fn parent_turn_base_is_not_changed_after_usage_records() { + let mut state = SubagentUsageState::default(); + let parent_session_id = SessionId::new(); + let parent_turn_id = TurnId::new(); + + state.begin_parent_turn( + parent_session_id, + parent_turn_id, + UsageTotals::default(), + None, + ); + let first_snapshot = state + .record_parent_turn_usage( + parent_session_id, + parent_turn_id, + totals(5, 1), + None, + UsageUpdateKind::InFlight, + ) + .expect("first parent snapshot"); + assert_eq!(first_snapshot.session_totals, totals(5, 1)); + + state.begin_parent_turn(parent_session_id, parent_turn_id, totals(100, 10), None); + let second_snapshot = state + .record_parent_turn_usage( + parent_session_id, + parent_turn_id, + totals(6, 2), + None, + UsageUpdateKind::InFlight, + ) + .expect("second parent snapshot"); + assert_eq!(second_snapshot.session_totals, totals(6, 2)); + } + #[test] fn child_usage_replaces_latest_inflight_without_double_counting() { let mut state = SubagentUsageState::default(); diff --git a/crates/server/tests/persistence_resume.rs b/crates/server/tests/persistence_resume.rs index a69e8606..ef504fe3 100644 --- a/crates/server/tests/persistence_resume.rs +++ b/crates/server/tests/persistence_resume.rs @@ -28,6 +28,7 @@ use devo_core::SkillsConfig; use devo_core::TurnLine; use devo_core::TurnRecord; use devo_core::tools::ToolRegistry; +use devo_protocol::DEVO_TURN_USAGE_META; use devo_protocol::Model; use devo_protocol::ModelRequest; use devo_protocol::ModelResponse; @@ -41,6 +42,7 @@ use devo_protocol::SessionId; use devo_protocol::StopReason; use devo_protocol::StreamEvent; use devo_protocol::TurnStatus; +use devo_protocol::TurnUsageUpdatedPayload; use devo_protocol::Usage; use devo_provider::ModelProviderSDK; use devo_provider::SingleProviderRouter; @@ -50,6 +52,33 @@ use devo_server::ServerRuntimeDependencies; struct SingleReplyProvider; +struct UsageReplyProvider { + input_tokens: usize, + output_tokens: usize, + cache_read_input_tokens: Option, +} + +impl UsageReplyProvider { + fn new(input_tokens: usize, output_tokens: usize) -> Self { + Self { + input_tokens, + output_tokens, + cache_read_input_tokens: Some(input_tokens / 2), + } + } + + fn usage(&self) -> Usage { + Usage { + input_tokens: self.input_tokens, + output_tokens: self.output_tokens, + cache_creation_input_tokens: None, + cache_read_input_tokens: self.cache_read_input_tokens, + reasoning_output_tokens: None, + total_tokens: Some(self.input_tokens + self.output_tokens), + } + } +} + #[derive(Default)] struct CapturingProvider { requests: Mutex>, @@ -93,6 +122,46 @@ impl ModelProviderSDK for SingleReplyProvider { } } +#[async_trait] +impl ModelProviderSDK for UsageReplyProvider { + async fn completion(&self, _request: ModelRequest) -> Result { + Ok(ModelResponse { + id: "title-usage".into(), + content: vec![ResponseContent::Text("Generated usage title".to_string())], + stop_reason: Some(StopReason::EndTurn), + usage: self.usage(), + metadata: ResponseMetadata::default(), + }) + } + + async fn completion_stream( + &self, + _request: ModelRequest, + ) -> Result> + Send>>> { + let usage = self.usage(); + Ok(Box::pin(stream::iter(vec![ + Ok(StreamEvent::UsageDelta(usage.clone())), + Ok(StreamEvent::TextDelta { + index: 0, + text: "Usage reply".into(), + }), + Ok(StreamEvent::MessageDone { + response: ModelResponse { + id: "resp-usage".into(), + content: vec![ResponseContent::Text("Usage reply".into())], + stop_reason: Some(StopReason::EndTurn), + usage, + metadata: ResponseMetadata::default(), + }, + }), + ]))) + } + + fn name(&self) -> &str { + "usage-reply-test-provider" + } +} + #[async_trait] impl ModelProviderSDK for CapturingProvider { async fn completion(&self, request: ModelRequest) -> Result { @@ -1491,6 +1560,32 @@ async fn wait_for_turn_completed( Ok(()) } +async fn wait_for_turn_usage_updated( + notifications_rx: &mut mpsc::Receiver, +) -> Result { + timeout(Duration::from_secs(5), async { + while let Some(value) = notifications_rx.recv().await { + let normalized = legacy_event_from_acp_notification(value.clone()); + if normalized.get("method") == Some(&serde_json::json!("turn/usage/updated")) { + return serde_json::from_value(normalized["params"].clone()) + .context("decode legacy turn/usage/updated payload"); + } + if value.get("method") == Some(&serde_json::json!("session/update")) { + let Some(meta) = value["params"]["update"]["_meta"].as_object() else { + continue; + }; + if let Some(payload) = meta.get(DEVO_TURN_USAGE_META) { + return serde_json::from_value(payload.clone()) + .context("decode ACP turn usage meta payload"); + } + } + } + anyhow::bail!("notification channel closed before turn/usage/updated") + }) + .await + .context("timed out waiting for turn/usage/updated")? +} + async fn wait_for_title_update( notifications_rx: &mut mpsc::Receiver, expected_title: &str, @@ -1707,6 +1802,109 @@ async fn interrupt_mid_stream_does_not_duplicate_last_item_on_resume() -> Result Ok(()) } +#[tokio::test] +async fn first_usage_update_after_resume_preserves_historical_session_totals() -> Result<()> { + let data_root = TempDir::new()?; + let provider = Arc::new(UsageReplyProvider::new(100, 25)); + let runtime = build_runtime_with_provider(data_root.path(), provider.clone())?; + let (connection_id, mut notifications_rx) = initialize_connection(&runtime).await?; + + let start_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 1, + "method": "session/start", + "params": { + "cwd": data_root.path(), + "ephemeral": false, + "title": "Usage resume base", + "model": "test-model" + } + }), + ) + .await + .context("session/start response")?; + let session_id = serde_json::from_value::< + devo_server::SuccessResponse, + >(start_response)? + .result + .session + .session_id; + + let first_turn_response = runtime + .handle_incoming( + connection_id, + serde_json::json!({ + "id": 2, + "method": "_devo/turn/start", + "params": { + "session_id": session_id, + "input": [{ "type": "text", "text": "persist usage totals" }], + "model": null, + "sandbox": null, + "approval_policy": null, + "cwd": null + } + }), + ) + .await + .context("first turn/start response")?; + let _: devo_server::SuccessResponse = + serde_json::from_value(first_turn_response)?; + let first_usage = wait_for_turn_usage_updated(&mut notifications_rx).await?; + assert_eq!(first_usage.total_input_tokens, 100); + wait_for_turn_completed(&mut notifications_rx).await?; + + let rebuilt_runtime = build_runtime_with_provider(data_root.path(), provider)?; + rebuilt_runtime.refresh_session_index()?; + let (rebuilt_connection_id, mut rebuilt_notifications_rx) = + initialize_connection(&rebuilt_runtime).await?; + let resume_response = rebuilt_runtime + .handle_incoming( + rebuilt_connection_id, + serde_json::json!({ + "id": 3, + "method": "_devo/session/resume", + "params": { "session_id": session_id } + }), + ) + .await + .context("session/resume response")?; + let resumed = serde_json::from_value::< + devo_server::SuccessResponse, + >(resume_response)? + .result + .session; + assert_eq!(resumed.total_input_tokens, 100); + + let second_turn_response = rebuilt_runtime + .handle_incoming( + rebuilt_connection_id, + serde_json::json!({ + "id": 4, + "method": "_devo/turn/start", + "params": { + "session_id": session_id, + "input": [{ "type": "text", "text": "post resume usage" }], + "model": null, + "sandbox": null, + "approval_policy": null, + "cwd": null + } + }), + ) + .await + .context("second turn/start response")?; + let _: devo_server::SuccessResponse = + serde_json::from_value(second_turn_response)?; + let post_resume_usage = wait_for_turn_usage_updated(&mut rebuilt_notifications_rx).await?; + assert_eq!(post_resume_usage.total_input_tokens, 200); + assert_eq!(post_resume_usage.total_output_tokens, 50); + assert_eq!(post_resume_usage.total_cache_read_tokens, 100); + Ok(()) +} + #[tokio::test] async fn lazy_resume_loads_parent_session_from_rollout_on_map_miss() -> Result<()> { let data_root = TempDir::new()?;