From af2d874883b0c39e9da02d457c8aa93f76827f58 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 22:10:28 +0800 Subject: [PATCH] fix: harden handoff status transitions --- crates/dbx-core/src/storage.rs | 94 +++++++++++++++++++------ src-tauri/src/commands/agent_runtime.rs | 27 +++++-- src/lib/agentHandoff.ts | 9 ++- src/stores/agentRuntimeStore.ts | 4 +- tests/agentHandoff.test.ts | 17 +++++ 5 files changed, 120 insertions(+), 31 deletions(-) diff --git a/crates/dbx-core/src/storage.rs b/crates/dbx-core/src/storage.rs index 3950072916..dddf77d36b 100644 --- a/crates/dbx-core/src/storage.rs +++ b/crates/dbx-core/src/storage.rs @@ -144,11 +144,10 @@ async fn ensure_history_columns(pool: &SqlitePool) -> Result<(), String> { } async fn ensure_handoffs_sequence(pool: &SqlitePool) -> Result<(), String> { - let (seq_columns,): (i64,) = - sqlx::query_as("SELECT COUNT(*) FROM pragma_table_info('handoffs') WHERE name = 'seq'") - .fetch_one(pool) - .await - .map_err(|e| e.to_string())?; + let (seq_columns,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM pragma_table_info('handoffs') WHERE name = 'seq'") + .fetch_one(pool) + .await + .map_err(|e| e.to_string())?; if seq_columns > 0 { return Ok(()); @@ -322,29 +321,26 @@ impl Storage { } pub async fn update_handoff_status(&self, id: &str, status: HandoffStatus) -> Result { - let row: Option<(String,)> = sqlx::query_as("SELECT payload_json FROM handoffs WHERE id = ?") - .bind(id) - .fetch_optional(&self.db) - .await - .map_err(|e| e.to_string())?; - let Some((payload_json,)) = row else { + let from_statuses = allowed_handoff_status_transitions(&status); + if from_statuses.is_empty() { return Ok(false); - }; + } - let mut item: HandoffItem = serde_json::from_str(&payload_json).map_err(|e| e.to_string())?; - item.status = status; - let updated_json = serde_json::to_string(&item).map_err(|e| e.to_string())?; - let status_value = handoff_status_value(&item.status)?; + let status_value = handoff_status_value(&status)?; - sqlx::query("UPDATE handoffs SET payload_json = ?, status = ? WHERE id = ?") - .bind(updated_json) + let result = sqlx::query( + "UPDATE handoffs SET payload_json = json_set(payload_json, '$.status', ?), status = ? \ + WHERE id = ? AND status IN (SELECT value FROM json_each(?))", + ) + .bind(status_value.clone()) .bind(status_value) .bind(id) + .bind(serde_json::to_string(&from_statuses).map_err(|e| e.to_string())?) .execute(&self.db) .await .map_err(|e| e.to_string())?; - Ok(true) + Ok(result.rows_affected() > 0) } pub async fn load_pending_handoffs(&self) -> Result, String> { @@ -368,6 +364,17 @@ fn handoff_status_value(status: &HandoffStatus) -> Result { .ok_or_else(|| "Failed to serialize handoff status".to_string()) } +fn allowed_handoff_status_transitions(status: &HandoffStatus) -> &'static [&'static str] { + match status { + HandoffStatus::Queued => &[], + HandoffStatus::Shown => &["queued", "shown"], + HandoffStatus::Approved => &["queued", "shown", "approved"], + HandoffStatus::Rejected => &["queued", "shown"], + HandoffStatus::Executed => &["approved", "executed"], + HandoffStatus::Failed => &["approved", "executed", "failed"], + } +} + // --------------------------------------------------------------------------- // AI Config // --------------------------------------------------------------------------- @@ -1088,8 +1095,10 @@ mod handoff_tests { async fn handoffs_table_has_stable_autoincrement_sequence() { let storage = open_temp_storage().await; - let columns: Vec<(String,)> = - sqlx::query_as("SELECT name FROM pragma_table_info('handoffs')").fetch_all(&storage.db).await.unwrap(); + let columns: Vec<(String,)> = sqlx::query_as("SELECT name FROM pragma_table_info('handoffs')") + .fetch_all(&storage.db) + .await + .unwrap(); assert!(columns.iter().any(|(name,)| name == "seq")); } @@ -1125,12 +1134,15 @@ mod handoff_tests { async fn update_handoff_status_updates_payload_and_pending_visibility() { let storage = open_temp_storage().await; let item = queued_handoff("review me"); + let queued_reject = queued_handoff("reject from queued"); let missing = uuid::Uuid::new_v4().to_string(); storage.save_handoff(&item).await.unwrap(); + storage.save_handoff(&queued_reject).await.unwrap(); assert!(storage.update_handoff_status(&item.id, HandoffStatus::Shown).await.unwrap()); assert!(!storage.update_handoff_status(&missing, HandoffStatus::Rejected).await.unwrap()); + assert!(storage.update_handoff_status(&queued_reject.id, HandoffStatus::Rejected).await.unwrap()); let shown = storage.load_pending_handoffs().await.unwrap(); assert_eq!(shown.len(), 1); @@ -1139,6 +1151,46 @@ mod handoff_tests { assert!(storage.update_handoff_status(&item.id, HandoffStatus::Rejected).await.unwrap()); assert!(storage.load_pending_handoffs().await.unwrap().is_empty()); } + + #[tokio::test] + async fn update_handoff_status_does_not_let_shown_overwrite_rejected() { + let storage = open_temp_storage().await; + let item = queued_handoff("reject wins"); + + storage.save_handoff(&item).await.unwrap(); + + assert!(storage.update_handoff_status(&item.id, HandoffStatus::Rejected).await.unwrap()); + assert!(!storage.update_handoff_status(&item.id, HandoffStatus::Shown).await.unwrap()); + + assert!(storage.load_pending_handoffs().await.unwrap().is_empty()); + } + + #[tokio::test] + async fn update_handoff_status_only_marks_shown_from_queued_or_shown() { + let storage = open_temp_storage().await; + let queued = queued_handoff("queued"); + let mut shown = queued_handoff("shown"); + shown.status = HandoffStatus::Shown; + let mut approved = queued_handoff("approved"); + approved.status = HandoffStatus::Approved; + let mut executed = queued_handoff("executed"); + executed.status = HandoffStatus::Executed; + let mut failed = queued_handoff("failed"); + failed.status = HandoffStatus::Failed; + + for item in [&queued, &shown, &approved, &executed, &failed] { + storage.save_handoff(item).await.unwrap(); + } + + assert!(storage.update_handoff_status(&queued.id, HandoffStatus::Shown).await.unwrap()); + assert!(storage.update_handoff_status(&shown.id, HandoffStatus::Shown).await.unwrap()); + assert!(!storage.update_handoff_status(&approved.id, HandoffStatus::Shown).await.unwrap()); + assert!(!storage.update_handoff_status(&executed.id, HandoffStatus::Shown).await.unwrap()); + assert!(!storage.update_handoff_status(&failed.id, HandoffStatus::Shown).await.unwrap()); + + let loaded = storage.load_pending_handoffs().await.unwrap(); + assert_eq!(loaded.iter().map(|item| item.id.as_str()).collect::>(), vec![queued.id, shown.id]); + } } // --------------------------------------------------------------------------- diff --git a/src-tauri/src/commands/agent_runtime.rs b/src-tauri/src/commands/agent_runtime.rs index d88a2ef287..74be82af6f 100644 --- a/src-tauri/src/commands/agent_runtime.rs +++ b/src-tauri/src/commands/agent_runtime.rs @@ -325,7 +325,10 @@ async fn route_request(first_line: &str, body: &str, state: &AgentRuntimeState) if let Some(limit) = query_limit(first_line) { truncate_result_rows(&mut body, limit); } - return RuntimeResponse { status: "200 OK", body }; + return RuntimeResponse { + status: "200 OK", + body, + }; } if first_line.starts_with("POST /handoff ") { @@ -347,12 +350,7 @@ async fn route_request(first_line: &str, body: &str, state: &AgentRuntimeState) RuntimeResponse { status: "404 Not Found", body: serde_json::json!({"error": "not found"}) } } -async fn update_handoff_status( - app_state: &AppState, - runtime: &AgentRuntimeState, - id: &str, - status: HandoffStatus, -) -> Result { +async fn update_handoff_status(app_state: &AppState, runtime: &AgentRuntimeState, id: &str, status: HandoffStatus) -> Result { let stored = app_state.storage.update_handoff_status(id, status.clone()).await?; let runtime_updated = update_runtime_handoff_status(runtime, id, status).await; Ok(stored || runtime_updated) @@ -361,12 +359,26 @@ async fn update_handoff_status( async fn update_runtime_handoff_status(state: &AgentRuntimeState, id: &str, status: HandoffStatus) -> bool { let mut handoffs = state.handoffs.write().await; if let Some(item) = handoffs.iter_mut().find(|item| item.id == id) { + if !can_update_runtime_handoff_status(&item.status, &status) { + return false; + } item.status = status; return true; } false } +fn can_update_runtime_handoff_status(current: &HandoffStatus, next: &HandoffStatus) -> bool { + matches!( + (current, next), + (HandoffStatus::Queued | HandoffStatus::Shown, HandoffStatus::Shown) + | (HandoffStatus::Queued | HandoffStatus::Shown, HandoffStatus::Rejected) + | (HandoffStatus::Queued | HandoffStatus::Shown | HandoffStatus::Approved, HandoffStatus::Approved) + | (HandoffStatus::Approved | HandoffStatus::Executed, HandoffStatus::Executed) + | (HandoffStatus::Approved | HandoffStatus::Executed | HandoffStatus::Failed, HandoffStatus::Failed) + ) +} + async fn pending_runtime_handoffs(state: &AgentRuntimeState) -> Vec { state .handoffs @@ -632,6 +644,7 @@ mod tests { assert!(update_runtime_handoff_status(&state, &id, dbx_core::handoff::HandoffStatus::Rejected).await); assert!(pending_runtime_handoffs(&state).await.is_empty()); + assert!(!update_runtime_handoff_status(&state, &id, dbx_core::handoff::HandoffStatus::Shown).await); } #[test] diff --git a/src/lib/agentHandoff.ts b/src/lib/agentHandoff.ts index e50d2ee634..bedefe5edc 100644 --- a/src/lib/agentHandoff.ts +++ b/src/lib/agentHandoff.ts @@ -22,8 +22,13 @@ export function isPendingHandoff(item: AgentHandoffItem): boolean { return item.status === "queued" || item.status === "shown"; } -export function mergeLoadedHandoffs(items: AgentHandoffItem[]): AgentHandoffItem[] { - return items.filter(isPendingHandoff).sort((a, b) => Date.parse(a.createdAt) - Date.parse(b.createdAt)); +export function mergeLoadedHandoffs( + items: AgentHandoffItem[], + ignoredIds: ReadonlySet = new Set(), +): AgentHandoffItem[] { + return items + .filter((item) => isPendingHandoff(item) && !ignoredIds.has(item.id)) + .sort((a, b) => Date.parse(a.createdAt) - Date.parse(b.createdAt)); } export function updateHandoffStatus( diff --git a/src/stores/agentRuntimeStore.ts b/src/stores/agentRuntimeStore.ts index e687c01c9f..8843dfa2d6 100644 --- a/src/stores/agentRuntimeStore.ts +++ b/src/stores/agentRuntimeStore.ts @@ -17,6 +17,7 @@ export const useAgentRuntimeStore = defineStore("agentRuntime", () => { const handoffs = ref([]); const activeHandoffId = ref(null); const handoffDialogOpen = ref(false); + const locallyClosedHandoffIds = new Set(); let timer: ReturnType | null = null; const activeHandoff = computed(() => handoffs.value.find((item) => item.id === activeHandoffId.value) ?? null); @@ -59,7 +60,7 @@ export const useAgentRuntimeStore = defineStore("agentRuntime", () => { } async function loadHandoffs() { - const loaded = mergeLoadedHandoffs(await api.agentRuntimeLoadHandoffs()); + const loaded = mergeLoadedHandoffs(await api.agentRuntimeLoadHandoffs(), locallyClosedHandoffIds); handoffs.value = loaded; const state = deriveHandoffDialogState(loaded, activeHandoffId.value); activeHandoffId.value = state.active?.id ?? null; @@ -76,6 +77,7 @@ export const useAgentRuntimeStore = defineStore("agentRuntime", () => { async function rejectHandoff(id: string) { await api.agentRuntimeRejectHandoff(id); + locallyClosedHandoffIds.add(id); handoffs.value = updateHandoffStatus(handoffs.value, id, "rejected"); const state = deriveHandoffDialogState(handoffs.value, activeHandoffId.value === id ? null : activeHandoffId.value); activeHandoffId.value = state.active?.id ?? null; diff --git a/tests/agentHandoff.test.ts b/tests/agentHandoff.test.ts index aa677d0fbd..ebc901eac1 100644 --- a/tests/agentHandoff.test.ts +++ b/tests/agentHandoff.test.ts @@ -55,3 +55,20 @@ test("updateHandoffStatus marks shown and removes rejected handoffs from pending const rejected = updateHandoffStatus(shown, "1", "rejected"); assert.deepEqual(rejected, []); }); + +test("updateHandoffStatus does not let shown overwrite rejected handoffs", () => { + const rejected = updateHandoffStatus([handoff("1", "shown")], "1", "rejected"); + + const staleShown = updateHandoffStatus(rejected, "1", "shown"); + + assert.deepEqual(staleShown, []); +}); + +test("mergeLoadedHandoffs can ignore locally closed handoffs from stale loads", () => { + const merged = mergeLoadedHandoffs([handoff("1", "shown"), handoff("2", "queued")], new Set(["1"])); + + assert.deepEqual( + merged.map((item) => item.id), + ["2"], + ); +});