diff --git a/prompts/en/fragments/system/memory_persistence.md.j2 b/prompts/en/fragments/system/memory_persistence.md.j2 index 83fce0bed..4538816ea 100644 --- a/prompts/en/fragments/system/memory_persistence.md.j2 +++ b/prompts/en/fragments/system/memory_persistence.md.j2 @@ -1 +1 @@ -Review the recent conversation and persist any important information as memories. Start by recalling existing memories related to the topics discussed, then save new or updated memories with appropriate associations. \ No newline at end of file +Review the recent conversation and persist important information as memories. Required order: use memory_recall first, then memory_save when appropriate, then finish by calling memory_persistence_complete. Never invent memory IDs. diff --git a/prompts/en/fragments/system/memory_persistence_contract_retry.md.j2 b/prompts/en/fragments/system/memory_persistence_contract_retry.md.j2 new file mode 100644 index 000000000..3da91e069 --- /dev/null +++ b/prompts/en/fragments/system/memory_persistence_contract_retry.md.j2 @@ -0,0 +1 @@ +You must finish this memory-persistence run by calling memory_persistence_complete. First recall relevant memories, then save real memories if needed, then call memory_persistence_complete with either outcome="saved" and exact saved_memory_ids from successful memory_save calls in this run, or outcome="no_memories" with a short reason and no saved IDs. Do not invent memory IDs. diff --git a/prompts/en/memory_persistence.md.j2 b/prompts/en/memory_persistence.md.j2 index a35c11852..42614e418 100644 --- a/prompts/en/memory_persistence.md.j2 +++ b/prompts/en/memory_persistence.md.j2 @@ -27,10 +27,15 @@ This is an automatic process triggered periodically during conversation. You are - Use `related_to` for topical connections - Use `part_of` when a detail belongs to a larger concept already in memory +4. **Finish with the terminal tool.** You must call `memory_persistence_complete` before finishing: + - Use `outcome: "saved"` with `saved_memory_ids` that exactly match IDs returned by successful `memory_save` calls in this run. + - Use `outcome: "no_memories"` when nothing is worth saving, with a short `reason` and no saved IDs. + ## Rules 1. Do NOT save trivial information — greetings, small talk, tool call details, intermediate reasoning. 2. Do NOT duplicate memories. If you recall an existing memory that covers the same fact, either skip it or save an update with an `updates` association. 3. Be concise in memory content. Memories are structured data, not conversation transcripts. 4. Focus on the most recent portion of the conversation — older content has likely already been captured by previous persistence runs. -5. Return a brief summary of what you saved (for logging only — the channel will not see this). +5. Do not invent memory IDs. Every ID in `saved_memory_ids` must come from a real successful `memory_save` call in this run. +6. Do not return plain text as the terminal result. End the run by calling `memory_persistence_complete`. diff --git a/prompts/en/tools/memory_persistence_complete_description.md.j2 b/prompts/en/tools/memory_persistence_complete_description.md.j2 new file mode 100644 index 000000000..e6a121db1 --- /dev/null +++ b/prompts/en/tools/memory_persistence_complete_description.md.j2 @@ -0,0 +1 @@ +Finish a memory persistence run with a terminal outcome. Use outcome="saved" only with exact memory IDs returned by successful memory_save calls in this run, or outcome="no_memories" with a short reason and no IDs. diff --git a/src/agent/branch.rs b/src/agent/branch.rs index 001eab7ea..205221579 100644 --- a/src/agent/branch.rs +++ b/src/agent/branch.rs @@ -5,14 +5,18 @@ use crate::error::Result; use crate::hooks::SpacebotHook; use crate::llm::SpacebotModel; use crate::llm::routing::is_context_overflow_error; +use crate::tools::MemoryPersistenceContractState; use crate::{AgentDeps, BranchId, ChannelId, ProcessEvent, ProcessId, ProcessType}; use rig::agent::AgentBuilder; use rig::completion::CompletionModel; use rig::tool::server::ToolServerHandle; +use std::sync::Arc; use uuid::Uuid; /// Max consecutive context overflow recoveries before giving up. const MAX_OVERFLOW_RETRIES: usize = 2; +/// Max retries when a memory persistence branch misses terminal completion contract. +const MAX_MEMORY_CONTRACT_RETRIES: usize = 2; /// A branch is a fork of a channel's context for thinking. pub struct Branch { @@ -29,6 +33,14 @@ pub struct Branch { pub tool_server: ToolServerHandle, /// Maximum LLM turns before the branch is forced to conclude. pub max_turns: usize, + /// Optional completion contract state used only by silent memory-persistence branches. + pub memory_persistence_contract: Option>, +} + +#[derive(Debug, Clone)] +pub struct BranchExecutionConfig { + pub max_turns: usize, + pub memory_persistence_contract: Option>, } impl Branch { @@ -40,17 +52,20 @@ impl Branch { system_prompt: impl Into, history: Vec, tool_server: ToolServerHandle, - max_turns: usize, + execution_config: BranchExecutionConfig, ) -> Self { let id = Uuid::new_v4(); let process_id = ProcessId::Branch(id); - let hook = SpacebotHook::new( + let mut hook = SpacebotHook::new( deps.agent_id.clone(), process_id, ProcessType::Branch, Some(channel_id.clone()), deps.event_tx.clone(), ); + if let Some(contract_state) = &execution_config.memory_persistence_contract { + hook = hook.with_memory_persistence_contract(contract_state.clone()); + } Self { id, @@ -61,7 +76,8 @@ impl Branch { system_prompt: system_prompt.into(), history, tool_server, - max_turns, + max_turns: execution_config.max_turns, + memory_persistence_contract: execution_config.memory_persistence_contract, } } @@ -102,8 +118,13 @@ impl Branch { let mut current_prompt = prompt; let mut overflow_retries = 0; + let mut memory_contract_retries = 0; + let enforce_memory_contract = self.memory_persistence_contract.is_some(); let conclusion = loop { + if enforce_memory_contract { + self.hook.set_completion_contract_request_active(true); + } match self .hook .prompt_once(&agent, &mut self.history, ¤t_prompt) @@ -111,6 +132,15 @@ impl Branch { { Ok(response) => break response, Err(rig::completion::PromptError::MaxTurnsError { .. }) => { + self.hook.set_completion_contract_request_active(false); + if enforce_memory_contract { + tracing::warn!( + branch_id = %self.id, + "memory persistence branch exceeded turn limit without completing contract" + ); + break "Memory persistence branch exceeded turn limit without completing the memory persistence contract." + .to_string(); + } let partial = crate::agent::extract_last_assistant_text(&self.history) .unwrap_or_else(|| { "Branch exhausted its turns without a final conclusion.".into() @@ -118,11 +148,51 @@ impl Branch { tracing::warn!(branch_id = %self.id, "branch hit max turns, returning partial result"); break partial; } + Err(rig::completion::PromptError::PromptCancelled { reason, .. }) + if enforce_memory_contract + && SpacebotHook::is_memory_persistence_contract_reason(&reason) => + { + self.hook.set_completion_contract_request_active(false); + if matches!( + self.history.last(), + Some(rig::message::Message::Assistant { .. }) + ) { + self.history.pop(); + } + memory_contract_retries += 1; + if memory_contract_retries > MAX_MEMORY_CONTRACT_RETRIES { + tracing::warn!( + branch_id = %self.id, + retries = MAX_MEMORY_CONTRACT_RETRIES, + "memory persistence completion contract retries exhausted" + ); + break "Memory persistence branch failed to produce a terminal completion outcome." + .to_string(); + } + + tracing::warn!( + branch_id = %self.id, + attempt = memory_contract_retries, + "memory persistence branch missing terminal completion outcome, retrying" + ); + let prompt_engine = self.deps.runtime_config.prompts.load(); + current_prompt = prompt_engine + .render_system_memory_persistence_contract_retry() + .unwrap_or_else(|_| { + SpacebotHook::MEMORY_PERSISTENCE_CONTRACT_PROMPT.to_string() + }); + } Err(rig::completion::PromptError::PromptCancelled { reason, .. }) => { + if enforce_memory_contract { + self.hook.set_completion_contract_request_active(false); + } tracing::info!(branch_id = %self.id, %reason, "branch cancelled"); break format!("Branch was cancelled: {reason}"); } Err(error) if is_context_overflow_error(&error.to_string()) => { + if enforce_memory_contract { + self.hook.set_completion_contract_request_active(false); + } overflow_retries += 1; if overflow_retries > MAX_OVERFLOW_RETRIES { tracing::error!( @@ -146,12 +216,19 @@ impl Branch { "Continue where you left off. Older context has been compacted.".into(); } Err(error) => { + if enforce_memory_contract { + self.hook.set_completion_contract_request_active(false); + } tracing::error!(branch_id = %self.id, %error, "branch LLM call failed"); return Err(crate::error::AgentError::Other(error.into()).into()); } } }; + if enforce_memory_contract { + self.hook.set_completion_contract_request_active(false); + } + // Scrub tool secret values from the conclusion before sending to the // channel. Branches can spawn workers whose output may contain secrets. // Layer 1: exact-match redaction of known secrets from the store. diff --git a/src/agent/channel_dispatch.rs b/src/agent/channel_dispatch.rs index fa3cc29a2..7fa51f6b2 100644 --- a/src/agent/channel_dispatch.rs +++ b/src/agent/channel_dispatch.rs @@ -4,11 +4,12 @@ //! background processes: `spawn_branch_from_state`, `spawn_worker_from_state`, //! and `spawn_opencode_worker_from_state`. -use crate::agent::branch::Branch; +use crate::agent::branch::{Branch, BranchExecutionConfig}; use crate::agent::channel::ChannelState; use crate::agent::channel_prompt::TemporalContext; use crate::agent::worker::Worker; use crate::error::{AgentError, Error as SpacebotError}; +use crate::tools::{BranchToolProfile, MemoryPersistenceContractState}; use crate::{AgentDeps, BranchId, ChannelId, ProcessEvent, WorkerId}; use futures::FutureExt as _; use std::sync::Arc; @@ -112,6 +113,11 @@ fn build_worker_status_text( Some(system_info.render_for_worker(¤t_time_line)) } +#[derive(Debug, Clone)] +struct BranchSpawnOptions { + profile: BranchToolProfile, +} + /// Spawn a branch from a ChannelState. Used by the BranchTool. pub async fn spawn_branch_from_state( state: &ChannelState, @@ -134,6 +140,9 @@ pub async fn spawn_branch_from_state( &system_prompt, &description, "branch", + BranchSpawnOptions { + profile: BranchToolProfile::Default, + }, ) .await } @@ -147,6 +156,8 @@ pub(crate) async fn spawn_memory_persistence_branch( state: &ChannelState, deps: &AgentDeps, ) -> std::result::Result { + let contract_state = Arc::new(MemoryPersistenceContractState::default()); + let prompt_engine = deps.runtime_config.prompts.load(); let system_prompt = prompt_engine .render_static("memory_persistence") @@ -162,6 +173,9 @@ pub(crate) async fn spawn_memory_persistence_branch( &system_prompt, "persisting memories...", "memory_persistence_branch", + BranchSpawnOptions { + profile: BranchToolProfile::MemoryPersistence { contract_state }, + }, ) .await } @@ -215,7 +229,14 @@ async fn spawn_branch( system_prompt: &str, status_label: &str, dispatch_type: &'static str, + branch_options: BranchSpawnOptions, ) -> std::result::Result { + let BranchSpawnOptions { profile } = branch_options; + let memory_persistence_contract = match &profile { + BranchToolProfile::MemoryPersistence { contract_state } => Some(contract_state.clone()), + BranchToolProfile::Default => None, + }; + let max_branches = **state.deps.runtime_config.max_concurrent_branches.load(); { let branches = state.active_branches.read().await; @@ -243,6 +264,7 @@ async fn spawn_branch( state.conversation_logger.clone(), state.channel_store.clone(), crate::conversation::ProcessRunLogger::new(state.deps.sqlite_pool.clone()), + profile, ); let branch_max_turns = **state.deps.runtime_config.branch_max_turns.load(); @@ -253,7 +275,10 @@ async fn spawn_branch( system_prompt, history, tool_server, - branch_max_turns, + BranchExecutionConfig { + max_turns: branch_max_turns, + memory_persistence_contract, + }, ); let branch_id = branch.id; diff --git a/src/agent/ingestion.rs b/src/agent/ingestion.rs index 568c39acc..24d845ca5 100644 --- a/src/agent/ingestion.rs +++ b/src/agent/ingestion.rs @@ -14,6 +14,7 @@ use crate::ProcessType; use crate::config::IngestionConfig; use crate::hooks::SpacebotHook; use crate::llm::SpacebotModel; +use crate::tools::MemoryPersistenceContractState; use anyhow::Context as _; use rig::agent::AgentBuilder; @@ -24,6 +25,7 @@ use sqlx::SqlitePool; use std::collections::HashSet; use std::path::{Path, PathBuf}; +use std::sync::Arc; use std::time::Duration; use uuid::Uuid; @@ -484,6 +486,7 @@ async fn process_chunk( let conversation_logger = crate::conversation::history::ConversationLogger::new(deps.sqlite_pool.clone()); let channel_store = crate::conversation::ChannelStore::new(deps.sqlite_pool.clone()); + let contract_state = Arc::new(MemoryPersistenceContractState::default()); let tool_server: ToolServerHandle = crate::tools::create_branch_tool_server( None, deps.agent_id.clone(), @@ -494,6 +497,9 @@ async fn process_chunk( conversation_logger, channel_store, crate::conversation::ProcessRunLogger::new(deps.sqlite_pool.clone()), + crate::tools::BranchToolProfile::MemoryPersistence { + contract_state: contract_state.clone(), + }, ); let agent = AgentBuilder::new(model) @@ -517,6 +523,14 @@ async fn process_chunk( let result = hook.prompt_once(&agent, &mut history, &user_prompt).await; classify_chunk_prompt_result(result, filename, chunk_number, total_chunks)?; + if !contract_state.has_terminal_outcome() { + tracing::warn!( + file = %filename, + chunk = %format!("{chunk_number}/{total_chunks}"), + "ingestion chunk completed without memory_persistence_complete signal" + ); + } + Ok(()) } diff --git a/src/agent/worker.rs b/src/agent/worker.rs index a2b887e92..bfaf89170 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -3,7 +3,7 @@ use crate::agent::compactor::estimate_history_tokens; use crate::config::BrowserConfig; use crate::error::Result; -use crate::hooks::{SpacebotHook, ToolNudgePolicy}; +use crate::hooks::SpacebotHook; use crate::llm::SpacebotModel; use crate::llm::routing::{is_context_overflow_error, is_retriable_error}; use crate::{AgentDeps, ChannelId, ProcessId, ProcessType, WorkerId}; @@ -546,38 +546,28 @@ impl Worker { let mut follow_up_prompt = follow_up.clone(); let mut follow_up_overflow_retries = 0; let mut follow_up_transient_retries = 0u32; - let follow_up_hook = self - .hook - .clone() - .with_tool_nudge_policy(ToolNudgePolicy::Disabled); let follow_up_result: std::result::Result = loop { - match follow_up_hook - .prompt_once(&agent, &mut history, &follow_up_prompt) + match self + .hook + .prompt_with_tool_nudge_retry(&agent, &mut history, &follow_up_prompt) .await { Ok(response) => break Ok(response), Err(rig::completion::PromptError::PromptCancelled { ref reason, .. - }) if SpacebotHook::is_context_injection_reason(reason) => { - // Context injection during a follow-up: drain - // buffered messages, append to history, and - // re-prompt — same as the main task loop. - let injected = follow_up_hook.take_injected_messages(); - for message in &injected { - tracing::info!( - worker_id = %self.id, - "injecting context into worker follow-up history" - ); - history.push(rig::message::Message::user(format!( - "[Context update from the user]: {message}" - ))); - } - follow_up_prompt = "New context has been provided above. \ - Incorporate this information and continue working \ - on your task. Do not repeat completed work." - .to_string(); - continue; + }) if SpacebotHook::is_tool_nudge_reason(reason) => { + let failure_reason = format!( + "follow-up ended without terminal outcome after {} nudge retries", + SpacebotHook::TOOL_NUDGE_MAX_RETRIES + ); + self.write_failure_log(&history, &failure_reason); + tracing::warn!( + worker_id = %self.id, + %reason, + "follow-up completion contract retries exhausted" + ); + break Err(failure_reason); } Err(error) if is_context_overflow_error(&error.to_string()) => { follow_up_overflow_retries += 1; diff --git a/src/hooks/spacebot.rs b/src/hooks/spacebot.rs index 9c4c323bb..c058c6117 100644 --- a/src/hooks/spacebot.rs +++ b/src/hooks/spacebot.rs @@ -1,9 +1,11 @@ //! SpacebotHook: Prompt hook for channels, branches, and workers. use crate::hooks::loop_guard::{LoopGuard, LoopGuardConfig, LoopGuardVerdict}; +use crate::tools::{MemoryPersistenceContractState, MemoryPersistenceTerminalOutcome}; use crate::{AgentId, ChannelId, ProcessEvent, ProcessId, ProcessType}; use rig::agent::{HookAction, PromptHook, ToolCallHookAction}; use rig::completion::{CompletionModel, CompletionResponse, Message, Prompt, PromptError}; +use std::sync::Arc; use tokio::sync::broadcast; /// Controls whether hook-driven tool nudge retries are enabled. @@ -38,6 +40,7 @@ pub struct SpacebotHook { tool_nudge_policy: ToolNudgePolicy, completion_calls: std::sync::Arc, nudge_request_active: std::sync::Arc, + completion_contract_request_active: std::sync::Arc, /// Set to `true` when the worker calls `set_status` with `kind: "outcome"`. /// Once signaled, the nudge system allows text-only responses to pass /// through as legitimate completions. @@ -57,6 +60,7 @@ pub struct SpacebotHook { /// `prompt_with_tool_nudge_retry` loop reads and clears this buffer to /// append the messages to history before re-prompting. injected_messages: std::sync::Arc>>, + memory_persistence_contract: Option>, } impl SpacebotHook { @@ -68,8 +72,19 @@ impl SpacebotHook { pub const TOOL_NUDGE_REASON: &str = "spacebot_tool_nudge_retry"; /// PromptCancelled reason used when injected context is pending. pub const CONTEXT_INJECTION_REASON: &str = "spacebot_context_injection"; + /// PromptCancelled reason used for memory-persistence contract retries. + pub const MEMORY_PERSISTENCE_CONTRACT_REASON: &str = + "spacebot_memory_persistence_contract_retry"; /// Maximum nudge retries per prompt request. pub const TOOL_NUDGE_MAX_RETRIES: usize = 2; + /// Maximum completion-contract retries per prompt request. + pub const MEMORY_PERSISTENCE_CONTRACT_MAX_RETRIES: usize = 2; + /// Prompt used to nudge memory-persistence branches toward a terminal tool outcome. + pub const MEMORY_PERSISTENCE_CONTRACT_PROMPT: &str = "You must finish this memory-persistence run by calling memory_persistence_complete. \ + First recall relevant memories, then save real memories if needed, then call \ + memory_persistence_complete with either outcome=\"saved\" and exact saved_memory_ids \ + from successful memory_save calls in this run, or outcome=\"no_memories\" with a short \ + reason and no saved IDs. Do not invent memory IDs."; /// Create a new hook. pub fn new( @@ -89,6 +104,9 @@ impl SpacebotHook { tool_nudge_policy: ToolNudgePolicy::for_process(process_type), completion_calls: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)), nudge_request_active: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)), + completion_contract_request_active: std::sync::Arc::new( + std::sync::atomic::AtomicBool::new(false), + ), outcome_signaled: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)), nudge_attempts: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)), loop_guard: std::sync::Arc::new(std::sync::Mutex::new(LoopGuard::new( @@ -96,6 +114,7 @@ impl SpacebotHook { ))), inject_rx: None, injected_messages: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())), + memory_persistence_contract: None, } } @@ -105,6 +124,14 @@ impl SpacebotHook { self } + pub fn with_memory_persistence_contract( + mut self, + contract_state: Arc, + ) -> Self { + self.memory_persistence_contract = Some(contract_state); + self + } + /// Attach a context injection receiver to this hook. /// /// When set, `on_completion_call` will drain pending messages from the @@ -144,6 +171,11 @@ impl SpacebotHook { .store(active, std::sync::atomic::Ordering::Relaxed); } + pub fn set_completion_contract_request_active(&self, active: bool) { + self.completion_contract_request_active + .store(active, std::sync::atomic::Ordering::Relaxed); + } + /// Return true if a PromptCancelled reason indicates a tool nudge retry. pub fn is_tool_nudge_reason(reason: &str) -> bool { reason == Self::TOOL_NUDGE_REASON @@ -154,6 +186,10 @@ impl SpacebotHook { reason == Self::CONTEXT_INJECTION_REASON } + pub fn is_memory_persistence_contract_reason(reason: &str) -> bool { + reason == Self::MEMORY_PERSISTENCE_CONTRACT_REASON + } + /// Drain and return all buffered injected messages. pub fn take_injected_messages(&self) -> Vec { self.injected_messages @@ -177,6 +213,7 @@ impl SpacebotHook { { self.reset_tool_nudge_state(); self.set_tool_nudge_request_active(true); + self.set_completion_contract_request_active(false); let mut current_prompt = std::borrow::Cow::Borrowed(prompt); let mut using_tool_nudge_prompt = false; @@ -239,6 +276,7 @@ impl SpacebotHook { if attempts >= Self::TOOL_NUDGE_MAX_RETRIES { // Retries exhausted — propagate the cancellation. self.set_tool_nudge_request_active(false); + self.set_completion_contract_request_active(false); return result; } Self::prune_tool_nudge_retry_history( @@ -265,6 +303,7 @@ impl SpacebotHook { ); } self.set_tool_nudge_request_active(false); + self.set_completion_contract_request_active(false); return result; } } @@ -538,6 +577,74 @@ impl SpacebotHook { ) }) } + + fn should_reject_memory_persistence_completion( + &self, + response: &CompletionResponse, + ) -> bool + where + M: CompletionModel, + { + let Some(contract_state) = &self.memory_persistence_contract else { + return false; + }; + if !self + .completion_contract_request_active + .load(std::sync::atomic::Ordering::Relaxed) + { + return false; + } + if contract_state.has_terminal_outcome() { + return false; + } + + !response + .choice + .iter() + .any(|content| matches!(content, rig::message::AssistantContent::ToolCall(_))) + } + + fn parse_memory_persistence_terminal_outcome( + result: &str, + ) -> Option { + let parsed = serde_json::from_str::(result).ok()?; + if parsed.get("success").and_then(|value| value.as_bool()) != Some(true) { + return None; + } + + match parsed.get("outcome").and_then(|value| value.as_str()) { + Some("saved") => { + let saved_memory_ids = parsed + .get("saved_memory_ids") + .and_then(|value| value.as_array()) + .map(|values| { + values + .iter() + .filter_map(|value| value.as_str()) + .map(str::trim) + .filter(|memory_id| !memory_id.is_empty()) + .map(ToOwned::to_owned) + .collect::>() + }) + .unwrap_or_default(); + if saved_memory_ids.is_empty() { + return None; + } + Some(MemoryPersistenceTerminalOutcome::Saved { saved_memory_ids }) + } + Some("no_memories") => { + let reason = parsed + .get("reason") + .and_then(|value| value.as_str()) + .map(str::trim) + .filter(|reason| !reason.is_empty())?; + Some(MemoryPersistenceTerminalOutcome::NoMemories { + reason: reason.to_string(), + }) + } + _ => None, + } + } } // Timer map for tool call duration measurement. Entries are inserted in @@ -611,6 +718,12 @@ where }; } + if self.should_reject_memory_persistence_completion::(response) { + return HookAction::Terminate { + reason: Self::MEMORY_PERSISTENCE_CONTRACT_REASON.into(), + }; + } + // Emit text content from worker completion responses so the live // transcript can show the model's reasoning between tool calls. if self.process_type == ProcessType::Worker { @@ -803,11 +916,20 @@ where guard.record_outcome(tool_name, _args, result); } + let is_tool_error = result.starts_with("Toolset error:"); + + if !is_tool_error + && tool_name == "memory_persistence_complete" + && let Some(contract_state) = &self.memory_persistence_contract + && let Some(outcome) = Self::parse_memory_persistence_terminal_outcome(result) + { + contract_state.set_terminal_outcome(outcome); + } + // A successful tool call proves the worker is still productive. // Reset the consecutive nudge counter so a brief narration blip // after many tool calls doesn't exhaust the retry budget. // Tool errors (from Rig's error path) don't count as productive. - let is_tool_error = result.starts_with("Toolset error:"); if self.tool_nudge_policy.is_enabled() && !is_tool_error { self.nudge_attempts .store(0, std::sync::atomic::Ordering::Relaxed); @@ -856,11 +978,13 @@ mod tests { use crate::ProcessEvent; use crate::llm::SpacebotModel; use crate::llm::model::RawResponse; + use crate::tools::MemoryPersistenceContractState; use crate::{ProcessId, ProcessType}; use rig::OneOrMany; use rig::agent::{HookAction, PromptHook}; use rig::completion::{CompletionResponse, Message, Usage}; use rig::message::AssistantContent; + use std::sync::Arc; fn make_hook() -> SpacebotHook { let (event_tx, _event_rx) = tokio::sync::broadcast::channel(8); @@ -873,6 +997,20 @@ mod tests { ) } + fn make_memory_persistence_hook() -> (SpacebotHook, Arc) { + let (event_tx, _event_rx) = tokio::sync::broadcast::channel(8); + let contract_state = Arc::new(MemoryPersistenceContractState::default()); + let hook = SpacebotHook::new( + std::sync::Arc::::from("agent"), + ProcessId::Branch(uuid::Uuid::new_v4()), + ProcessType::Branch, + None, + event_tx, + ) + .with_memory_persistence_contract(contract_state.clone()); + (hook, contract_state) + } + fn prompt_message() -> Message { Message::from("test prompt") } @@ -1621,4 +1759,156 @@ mod tests { "Nudge should still work when inject_rx is attached but empty" ); } + + #[tokio::test] + async fn memory_persistence_plain_text_completion_is_rejected_without_terminal_tool() { + let (hook, _contract_state) = make_memory_persistence_hook(); + let prompt = prompt_message(); + hook.set_completion_contract_request_active(true); + + let _ = + >::on_completion_call(&hook, &prompt, &[]) + .await; + let action = >::on_completion_response( + &hook, + &prompt, + &text_response("Saved the memories."), + ) + .await; + + assert!(matches!( + action, + HookAction::Terminate { ref reason } + if reason == SpacebotHook::MEMORY_PERSISTENCE_CONTRACT_REASON + )); + } + + #[tokio::test] + async fn memory_persistence_fabricated_saved_ids_are_rejected() { + let (hook, contract_state) = make_memory_persistence_hook(); + let prompt = prompt_message(); + hook.set_completion_contract_request_active(true); + + let _ = >::on_tool_result( + &hook, + "memory_save", + None, + "internal_1", + "{}", + "{\"success\":true,\"memory_id\":\"mem_real_1\"}", + ) + .await; + + let _ = >::on_tool_result( + &hook, + "memory_persistence_complete", + None, + "internal_2", + "{\"outcome\":\"saved\",\"saved_memory_ids\":[\"mem_fake\"]}", + "Toolset error: memory_persistence_complete failed: saved_memory_ids mismatch", + ) + .await; + + assert!( + !contract_state.has_terminal_outcome(), + "terminal outcome must not be recorded for fabricated IDs" + ); + + let _ = + >::on_completion_call(&hook, &prompt, &[]) + .await; + let action = >::on_completion_response( + &hook, + &prompt, + &text_response("Done."), + ) + .await; + + assert!(matches!( + action, + HookAction::Terminate { ref reason } + if reason == SpacebotHook::MEMORY_PERSISTENCE_CONTRACT_REASON + )); + } + + #[tokio::test] + async fn memory_persistence_saved_outcome_accepts_real_memory_save_ids() { + let (hook, contract_state) = make_memory_persistence_hook(); + let prompt = prompt_message(); + hook.set_completion_contract_request_active(true); + + let _ = >::on_tool_result( + &hook, + "memory_save", + None, + "internal_1", + "{}", + "{\"success\":true,\"memory_id\":\"mem_real_1\"}", + ) + .await; + let _ = >::on_tool_result( + &hook, + "memory_save", + None, + "internal_2", + "{}", + "{\"success\":true,\"memory_id\":\"mem_real_2\"}", + ) + .await; + + let _ = >::on_tool_result( + &hook, + "memory_persistence_complete", + None, + "internal_3", + "{}", + "{\"success\":true,\"outcome\":\"saved\",\"saved_memory_ids\":[\"mem_real_1\",\"mem_real_2\"]}", + ) + .await; + + assert!(contract_state.has_terminal_outcome()); + + let _ = + >::on_completion_call(&hook, &prompt, &[]) + .await; + let action = >::on_completion_response( + &hook, + &prompt, + &text_response("Persisted memories."), + ) + .await; + + assert!(matches!(action, HookAction::Continue)); + } + + #[tokio::test] + async fn memory_persistence_no_memories_outcome_is_accepted_without_saves() { + let (hook, contract_state) = make_memory_persistence_hook(); + let prompt = prompt_message(); + hook.set_completion_contract_request_active(true); + + let _ = >::on_tool_result( + &hook, + "memory_persistence_complete", + None, + "internal_1", + "{}", + "{\"success\":true,\"outcome\":\"no_memories\",\"saved_memory_ids\":[],\"reason\":\"No durable facts found\"}", + ) + .await; + + assert!(contract_state.has_terminal_outcome()); + + let _ = + >::on_completion_call(&hook, &prompt, &[]) + .await; + let action = >::on_completion_response( + &hook, + &prompt, + &text_response("No memories persisted."), + ) + .await; + + assert!(matches!(action, HookAction::Continue)); + } } diff --git a/src/llm/model.rs b/src/llm/model.rs index c72d4eccb..9dadcd93f 100644 --- a/src/llm/model.rs +++ b/src/llm/model.rs @@ -1329,9 +1329,14 @@ fn convert_messages_to_openai(messages: &OneOrMany) -> Vec { + let tool_call_id = tr + .call_id + .as_deref() + .filter(|call_id| !call_id.is_empty()) + .unwrap_or(&tr.id); tool_results.push(serde_json::json!({ "role": "tool", - "tool_call_id": tr.id, + "tool_call_id": tool_call_id, "content": tool_result_content_to_string(&tr.content), })); } @@ -1367,11 +1372,18 @@ fn convert_messages_to_openai(messages: &OneOrMany) -> Vec { - // OpenAI expects arguments as a JSON string + // OpenAI expects arguments as a JSON string. + // Prefer call_id (set when replaying Responses-API tool calls + // through chat-completions) to keep assistant and tool IDs aligned. + let preferred_id = tc + .call_id + .as_deref() + .filter(|c| !c.is_empty()) + .unwrap_or(&tc.id); let args_string = serde_json::to_string(&tc.function.arguments) .unwrap_or_else(|_| "{}".to_string()); tool_calls.push(serde_json::json!({ - "id": tc.id, + "id": preferred_id, "type": "function", "function": { "name": tc.function.name, @@ -1420,9 +1432,14 @@ fn convert_messages_to_openai_responses(messages: &OneOrMany) -> Vec { + let call_id = tool_result + .call_id + .as_deref() + .filter(|call_id| !call_id.is_empty()) + .unwrap_or(&tool_result.id); result.push(serde_json::json!({ "type": "function_call_output", - "call_id": tool_result.id, + "call_id": call_id, "output": tool_result_content_to_string(&tool_result.content), })); } @@ -1451,11 +1468,16 @@ fn convert_messages_to_openai_responses(messages: &OneOrMany) -> Vec { let arguments = serde_json::to_string(&tool_call.function.arguments) .unwrap_or_else(|_| "{}".to_string()); + let call_id = tool_call + .call_id + .as_deref() + .filter(|call_id| !call_id.is_empty()) + .unwrap_or(&tool_call.id); result.push(serde_json::json!({ "type": "function_call", "name": tool_call.function.name, "arguments": arguments, - "call_id": tool_call.id, + "call_id": call_id, })); } _ => {} @@ -2570,6 +2592,61 @@ fn parse_openai_tool_call(tool_call: &serde_json::Value, fallback_id: String) -> Some(make_tool_call(id, name.to_string(), arguments)) } +fn extract_text_content_from_responses_output_item( + value: &serde_json::Value, + text_parts: &mut Vec, +) { + match value { + serde_json::Value::Array(items) => { + for item in items { + extract_text_content_from_responses_output_item(item, text_parts); + } + } + serde_json::Value::Object(map) => { + if matches!( + map.get("type").and_then(serde_json::Value::as_str), + Some("function_call") | Some("function_call_output") + ) { + return; + } + + if let Some(text) = map.get("text").and_then(serde_json::Value::as_str) + && !text.trim().is_empty() + { + text_parts.push(text.to_string()); + } + if let Some(summary) = map.get("summary") { + collect_openai_text_content(summary, text_parts); + } + if let Some(refusal) = map.get("refusal") { + collect_openai_text_content(refusal, text_parts); + } + if let Some(content) = map.get("content") { + extract_text_content_from_responses_output_item(content, text_parts); + } + } + _ => {} + } +} + +fn make_openai_responses_tool_call( + id: String, + call_id: Option, + name: String, + arguments: serde_json::Value, +) -> ToolCall { + ToolCall { + id, + call_id, + function: ToolFunction { + name: name.trim().to_string(), + arguments, + }, + signature: None, + additional_params: None, + } +} + fn parse_openai_responses_response( body: serde_json::Value, provider_label: &str, @@ -2579,19 +2656,34 @@ fn parse_openai_responses_response( .ok_or_else(|| CompletionError::ResponseError("missing output array".into()))?; let mut assistant_content = Vec::new(); + let mut fallback_text_parts = Vec::new(); for (index, output_item) in output_items.iter().enumerate() { match output_item["type"].as_str() { Some("message") => { if let Some(content_items) = output_item["content"].as_array() { + let mut message_output_text = Vec::new(); + let mut message_fallback_text = Vec::new(); + for content_item in content_items { if content_item["type"].as_str() == Some("output_text") && let Some(text) = content_item["text"].as_str() && !text.is_empty() { - assistant_content.push(AssistantContent::Text(Text { - text: text.to_string(), - })); + message_output_text.push(text.to_string()); + } + + extract_text_content_from_responses_output_item( + content_item, + &mut message_fallback_text, + ); + } + + if message_output_text.is_empty() { + fallback_text_parts.extend(message_fallback_text); + } else { + for text in message_output_text { + assistant_content.push(AssistantContent::Text(Text { text })); } } } @@ -2599,29 +2691,53 @@ fn parse_openai_responses_response( Some("function_call") => { let call_id = output_item["call_id"] .as_str() - .or_else(|| output_item["id"].as_str()) + .filter(|id| !id.is_empty()) + .map(ToOwned::to_owned); + let id = output_item["id"] + .as_str() .filter(|id| !id.is_empty()) .map(ToOwned::to_owned) + .or_else(|| call_id.clone()) .unwrap_or_else(|| format!("function_call_{index}")); let name = output_item["name"].as_str().unwrap_or("").to_string(); let arguments = parse_openai_tool_arguments(&output_item["arguments"]); - assistant_content.push(AssistantContent::ToolCall(make_tool_call( - call_id, name, arguments, - ))); + assistant_content.push(AssistantContent::ToolCall( + make_openai_responses_tool_call(id, call_id, name, arguments), + )); + } + _ => { + extract_text_content_from_responses_output_item( + output_item, + &mut fallback_text_parts, + ); } - _ => {} + } + } + + let has_text = assistant_content + .iter() + .any(|content| matches!(content, AssistantContent::Text(_))); + if !has_text { + for text in fallback_text_parts { + assistant_content.push(AssistantContent::Text(Text { text })); } } let choice = OneOrMany::many(assistant_content).map_err(|_| { + let output_types = output_items + .iter() + .map(|item| item["type"].as_str().unwrap_or("")) + .collect::>() + .join(", "); tracing::warn!( provider = %provider_label, output_items = output_items.len(), + output_types = %output_types, "empty response from responses API" ); CompletionError::ResponseError(format!( - "empty response from {provider_label} Responses API" + "empty or unsupported response from {provider_label} Responses API; expected text-bearing message content (output_text/text/summary/refusal/content) or function_call output items; received output types: {output_types}" )) })?; @@ -3010,6 +3126,146 @@ mod tests { assert!(error.to_string().contains("finish_reason: stop")); } + #[test] + fn convert_messages_to_openai_tool_result_prefers_call_id_over_id() { + let messages = OneOrMany::one(Message::User { + content: OneOrMany::one(UserContent::ToolResult(rig::message::ToolResult { + id: "legacy-id".to_string(), + call_id: Some("stable-call-id".to_string()), + content: OneOrMany::one(rig::message::ToolResultContent::text("ok")), + })), + }); + + let converted = convert_messages_to_openai(&messages); + assert_eq!(converted.len(), 1); + assert_eq!(converted[0]["role"], "tool"); + assert_eq!(converted[0]["tool_call_id"], "stable-call-id"); + } + + #[test] + fn convert_messages_to_openai_responses_function_call_output_prefers_call_id_over_id() { + let messages = OneOrMany::one(Message::User { + content: OneOrMany::one(UserContent::ToolResult(rig::message::ToolResult { + id: "legacy-id".to_string(), + call_id: Some("stable-call-id".to_string()), + content: OneOrMany::one(rig::message::ToolResultContent::text("ok")), + })), + }); + + let converted = convert_messages_to_openai_responses(&messages); + assert_eq!(converted.len(), 1); + assert_eq!(converted[0]["type"], "function_call_output"); + assert_eq!(converted[0]["call_id"], "stable-call-id"); + } + + #[test] + fn convert_messages_to_openai_responses_function_call_prefers_call_id_over_id() { + let messages = OneOrMany::one(Message::Assistant { + content: OneOrMany::one(AssistantContent::ToolCall(ToolCall { + id: "legacy-id".to_string(), + call_id: Some("stable-call-id".to_string()), + function: ToolFunction { + name: "reply".to_string(), + arguments: serde_json::json!({"content": "ok"}), + }, + signature: None, + additional_params: None, + })), + id: None, + }); + + let converted = convert_messages_to_openai_responses(&messages); + assert_eq!(converted.len(), 1); + assert_eq!(converted[0]["type"], "function_call"); + assert_eq!(converted[0]["call_id"], "stable-call-id"); + } + + #[test] + fn parse_openai_responses_response_parses_fallback_text_without_output_text() { + let body = serde_json::json!({ + "output": [{ + "type": "message", + "content": [{ + "type": "reasoning", + "summary": [ + {"text": "step 1"}, + {"text": "step 2"} + ] + }] + }], + "usage": { + "input_tokens": 3, + "output_tokens": 2, + "input_tokens_details": {"cached_tokens": 0} + } + }); + + let response = + parse_openai_responses_response(body, "OpenAI").expect("fallback text should parse"); + let texts: Vec<_> = response + .choice + .iter() + .filter_map(|content| match content { + AssistantContent::Text(text) => Some(text.text.clone()), + _ => None, + }) + .collect(); + + assert_eq!(texts, vec!["step 1".to_string(), "step 2".to_string()]); + } + + #[test] + fn parse_openai_responses_response_preserves_function_call_call_id_from_completed_response() { + let body = serde_json::json!({ + "output": [{ + "type": "function_call", + "id": "legacy-id", + "call_id": "stable-call-id", + "name": "reply", + "arguments": "{\"content\":\"ok\"}" + }], + "usage": { + "input_tokens": 3, + "output_tokens": 2, + "input_tokens_details": {"cached_tokens": 0} + } + }); + + let response = + parse_openai_responses_response(body, "OpenAI").expect("function call should parse"); + match response.choice.first_ref() { + AssistantContent::ToolCall(tool_call) => { + assert_eq!(tool_call.id, "legacy-id"); + assert_eq!(tool_call.call_id.as_deref(), Some("stable-call-id")); + assert_eq!(tool_call.function.name, "reply"); + } + _ => panic!("expected tool call"), + } + } + + #[test] + fn parse_openai_responses_response_unsupported_empty_error_is_actionable_and_provider_specific() + { + let body = serde_json::json!({ + "output": [{ + "type": "unknown_shape", + "foo": "bar" + }], + "usage": { + "input_tokens": 1, + "output_tokens": 0, + "input_tokens_details": {"cached_tokens": 0} + } + }); + + let error = + parse_openai_responses_response(body, "OpenAI").expect_err("should be unsupported"); + let error_text = error.to_string(); + assert!(error_text.contains("OpenAI Responses API")); + assert!(error_text.contains("output_text/text/summary/refusal/content")); + assert!(error_text.contains("unknown_shape")); + } + #[test] fn parse_openai_chat_sse_response_reconstructs_tool_calls() { let sse = concat!( diff --git a/src/prompts/engine.rs b/src/prompts/engine.rs index 16969accb..e706b59d1 100644 --- a/src/prompts/engine.rs +++ b/src/prompts/engine.rs @@ -363,6 +363,11 @@ impl PromptEngine { self.render_static("fragments/system/memory_persistence") } + /// Retry nudge sent to a memory-persistence branch that missed its terminal completion call. + pub fn render_system_memory_persistence_contract_retry(&self) -> Result { + self.render_static("fragments/system/memory_persistence_contract_retry") + } + /// Render the profile synthesis prompt with identity and bulletin context. pub fn render_system_profile_synthesis( &self, diff --git a/src/prompts/text.rs b/src/prompts/text.rs index 717e4c4a4..8ec497891 100644 --- a/src/prompts/text.rs +++ b/src/prompts/text.rs @@ -168,6 +168,9 @@ fn lookup(lang: &str, key: &str) -> &'static str { ("en", "tools/memory_save") => { include_str!("../../prompts/en/tools/memory_save_description.md.j2") } + ("en", "tools/memory_persistence_complete") => { + include_str!("../../prompts/en/tools/memory_persistence_complete_description.md.j2") + } ("en", "tools/memory_recall") => { include_str!("../../prompts/en/tools/memory_recall_description.md.j2") } diff --git a/src/tools.rs b/src/tools.rs index a688d46f8..eeaecd3ac 100644 --- a/src/tools.rs +++ b/src/tools.rs @@ -40,6 +40,7 @@ pub mod file; pub mod install_skill; pub mod mcp; pub mod memory_delete; +pub mod memory_persistence_complete; pub mod memory_recall; pub mod memory_save; pub mod project_manage; @@ -99,6 +100,11 @@ pub use mcp::{McpToolAdapter, McpToolError, McpToolOutput}; pub use memory_delete::{ MemoryDeleteArgs, MemoryDeleteError, MemoryDeleteOutput, MemoryDeleteTool, }; +pub use memory_persistence_complete::{ + MemoryPersistenceCompleteArgs, MemoryPersistenceCompleteError, MemoryPersistenceCompleteOutput, + MemoryPersistenceCompleteTool, MemoryPersistenceContractState, + MemoryPersistenceTerminalOutcome, +}; pub use memory_recall::{ MemoryOutput, MemoryRecallArgs, MemoryRecallError, MemoryRecallOutput, MemoryRecallTool, }; @@ -176,6 +182,14 @@ use std::path::PathBuf; use std::sync::Arc; use tokio::sync::{broadcast, mpsc}; +#[derive(Debug, Clone)] +pub enum BranchToolProfile { + Default, + MemoryPersistence { + contract_state: Arc, + }, +} + /// Deserialize a `u64` that may arrive as either a JSON number or a JSON string. /// /// LLMs sometimes send `"timeout_seconds": "400"` instead of `"timeout_seconds": 400`. @@ -494,13 +508,19 @@ pub fn create_branch_tool_server( conversation_logger: crate::conversation::history::ConversationLogger, channel_store: crate::conversation::ChannelStore, run_logger: crate::conversation::history::ProcessRunLogger, + profile: BranchToolProfile, ) -> ToolServerHandle { + let mut memory_save = memory_save_with_events( + memory_search.clone(), + agent_id.clone(), + memory_event_tx.clone(), + ); + if let BranchToolProfile::MemoryPersistence { contract_state } = &profile { + memory_save = memory_save.with_contract_state(contract_state.clone()); + } + let mut server = ToolServer::new() - .tool(memory_save_with_events( - memory_search.clone(), - agent_id.clone(), - memory_event_tx.clone(), - )) + .tool(memory_save) .tool(MemoryRecallTool::new(memory_search.clone())) .tool(MemoryDeleteTool::new(memory_search)) .tool(ChannelRecallTool::new(conversation_logger, channel_store)) @@ -515,6 +535,10 @@ pub fn create_branch_tool_server( .tool(TaskListTool::new(task_store.clone(), agent_id.to_string())) .tool(TaskUpdateTool::for_branch(task_store, agent_id.clone())); + if let BranchToolProfile::MemoryPersistence { contract_state } = profile { + server = server.tool(MemoryPersistenceCompleteTool::new(contract_state)); + } + if let Some(state) = state { server = server.tool(SpawnWorkerTool::new(state)); } diff --git a/src/tools/memory_persistence_complete.rs b/src/tools/memory_persistence_complete.rs new file mode 100644 index 000000000..df9e2dd69 --- /dev/null +++ b/src/tools/memory_persistence_complete.rs @@ -0,0 +1,264 @@ +//! Terminal completion tool for memory persistence branches. + +use rig::completion::ToolDefinition; +use rig::tool::Tool; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; +use std::sync::{Arc, Mutex}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum MemoryPersistenceTerminalOutcome { + Saved { saved_memory_ids: Vec }, + NoMemories { reason: String }, +} + +#[derive(Debug, Default)] +pub struct MemoryPersistenceContractState { + saved_memory_ids: Mutex>, + terminal_outcome: Mutex>, +} + +impl MemoryPersistenceContractState { + pub fn record_saved_memory_id(&self, memory_id: impl Into) { + if let Ok(mut saved_memory_ids) = self.saved_memory_ids.lock() { + saved_memory_ids.insert(memory_id.into()); + } + } + + pub fn saved_memory_ids(&self) -> Vec { + self.saved_memory_ids + .lock() + .map(|saved_memory_ids| saved_memory_ids.iter().cloned().collect::>()) + .unwrap_or_default() + } + + pub fn set_terminal_outcome(&self, outcome: MemoryPersistenceTerminalOutcome) { + if let Ok(mut terminal_outcome) = self.terminal_outcome.lock() { + *terminal_outcome = Some(outcome); + } + } + + pub fn terminal_outcome(&self) -> Option { + self.terminal_outcome + .lock() + .ok() + .and_then(|terminal_outcome| terminal_outcome.clone()) + } + + pub fn has_terminal_outcome(&self) -> bool { + self.terminal_outcome + .lock() + .ok() + .map(|guard| guard.is_some()) + .unwrap_or(false) + } +} + +#[derive(Debug, Clone)] +pub struct MemoryPersistenceCompleteTool { + state: Arc, +} + +impl MemoryPersistenceCompleteTool { + pub fn new(state: Arc) -> Self { + Self { state } + } +} + +#[derive(Debug, thiserror::Error)] +#[error("memory_persistence_complete failed: {0}")] +pub struct MemoryPersistenceCompleteError(String); + +#[derive(Debug, Deserialize, JsonSchema)] +pub struct MemoryPersistenceCompleteArgs { + /// Terminal branch outcome. Use "saved" when at least one memory was + /// saved in this run, otherwise use "no_memories". + pub outcome: String, + /// Required when outcome is "saved". Must exactly match the memory IDs + /// returned by successful memory_save tool calls in this run. + #[serde(default)] + pub saved_memory_ids: Vec, + /// Required when outcome is "no_memories". Explain briefly why nothing + /// was worth saving. + #[serde(default)] + pub reason: Option, +} + +#[derive(Debug, Serialize)] +pub struct MemoryPersistenceCompleteOutput { + pub success: bool, + pub outcome: String, + pub saved_memory_ids: Vec, + pub reason: Option, +} + +impl Tool for MemoryPersistenceCompleteTool { + const NAME: &'static str = "memory_persistence_complete"; + + type Error = MemoryPersistenceCompleteError; + type Args = MemoryPersistenceCompleteArgs; + type Output = MemoryPersistenceCompleteOutput; + + async fn definition(&self, _prompt: String) -> ToolDefinition { + ToolDefinition { + name: Self::NAME.to_string(), + description: crate::prompts::text::get("tools/memory_persistence_complete").to_string(), + parameters: serde_json::json!({ + "type": "object", + "properties": { + "outcome": { + "type": "string", + "enum": ["saved", "no_memories"], + "description": "Terminal memory persistence outcome for this run" + }, + "saved_memory_ids": { + "type": "array", + "items": { "type": "string" }, + "description": "Required for outcome=saved. Must exactly match successful memory_save IDs from this run" + }, + "reason": { + "type": "string", + "description": "Required for outcome=no_memories. Brief reason why no memories were saved" + } + }, + "required": ["outcome"] + }), + } + } + + async fn call(&self, args: Self::Args) -> Result { + let outcome = args.outcome.trim(); + let recorded_ids = self.state.saved_memory_ids(); + + match outcome { + "saved" => { + if args.saved_memory_ids.is_empty() { + return Err(MemoryPersistenceCompleteError( + "outcome 'saved' requires non-empty saved_memory_ids".to_string(), + )); + } + if recorded_ids.is_empty() { + return Err(MemoryPersistenceCompleteError( + "outcome 'saved' is invalid because no successful memory_save calls were recorded in this run" + .to_string(), + )); + } + + let provided_ids = args.saved_memory_ids; + let provided_set = provided_ids.iter().cloned().collect::>(); + if provided_set.len() != provided_ids.len() { + return Err(MemoryPersistenceCompleteError( + "saved_memory_ids must not contain duplicates".to_string(), + )); + } + + let recorded_set = recorded_ids.iter().cloned().collect::>(); + if provided_set != recorded_set { + return Err(MemoryPersistenceCompleteError(format!( + "saved_memory_ids mismatch: expected {:?}, got {:?}", + recorded_ids, provided_ids + ))); + } + + Ok(MemoryPersistenceCompleteOutput { + success: true, + outcome: "saved".to_string(), + saved_memory_ids: provided_ids, + reason: None, + }) + } + "no_memories" => { + if !args.saved_memory_ids.is_empty() { + return Err(MemoryPersistenceCompleteError( + "outcome 'no_memories' must not include saved_memory_ids".to_string(), + )); + } + if !recorded_ids.is_empty() { + return Err(MemoryPersistenceCompleteError(format!( + "outcome 'no_memories' is invalid because successful memory_save calls were recorded: {:?}", + recorded_ids + ))); + } + + let reason = args.reason.unwrap_or_default(); + if reason.trim().len() < 3 { + return Err(MemoryPersistenceCompleteError( + "outcome 'no_memories' requires a short reason".to_string(), + )); + } + + Ok(MemoryPersistenceCompleteOutput { + success: true, + outcome: "no_memories".to_string(), + saved_memory_ids: Vec::new(), + reason: Some(reason.trim().to_string()), + }) + } + _ => Err(MemoryPersistenceCompleteError(format!( + "invalid outcome '{outcome}'; expected 'saved' or 'no_memories'" + ))), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn saved_outcome_rejects_fabricated_ids() { + let state = Arc::new(MemoryPersistenceContractState::default()); + state.record_saved_memory_id("mem_real_1"); + let tool = MemoryPersistenceCompleteTool::new(state); + + let error = tool + .call(MemoryPersistenceCompleteArgs { + outcome: "saved".to_string(), + saved_memory_ids: vec!["mem_fake".to_string()], + reason: None, + }) + .await + .expect_err("fabricated ids should fail"); + + assert!(error.to_string().contains("saved_memory_ids mismatch")); + } + + #[tokio::test] + async fn saved_outcome_accepts_exact_recorded_ids() { + let state = Arc::new(MemoryPersistenceContractState::default()); + state.record_saved_memory_id("mem_1"); + state.record_saved_memory_id("mem_2"); + let tool = MemoryPersistenceCompleteTool::new(state); + + let output = tool + .call(MemoryPersistenceCompleteArgs { + outcome: "saved".to_string(), + saved_memory_ids: vec!["mem_2".to_string(), "mem_1".to_string()], + reason: None, + }) + .await + .expect("exact ids should pass"); + + assert!(output.success); + assert_eq!(output.outcome, "saved"); + } + + #[tokio::test] + async fn no_memories_outcome_accepts_short_reason_without_saves() { + let state = Arc::new(MemoryPersistenceContractState::default()); + let tool = MemoryPersistenceCompleteTool::new(state); + + let output = tool + .call(MemoryPersistenceCompleteArgs { + outcome: "no_memories".to_string(), + saved_memory_ids: Vec::new(), + reason: Some("No durable facts in recent turns".to_string()), + }) + .await + .expect("no_memories should pass with reason"); + + assert!(output.success); + assert_eq!(output.outcome, "no_memories"); + } +} diff --git a/src/tools/memory_save.rs b/src/tools/memory_save.rs index 9188db222..fbefbeaf6 100644 --- a/src/tools/memory_save.rs +++ b/src/tools/memory_save.rs @@ -19,6 +19,7 @@ const MAX_MEMORY_CONTENT_BYTES: usize = 50_000; pub struct MemorySaveTool { memory_search: Arc, event_context: Option, + contract_state: Option>, } #[derive(Debug, Clone)] @@ -33,6 +34,7 @@ impl MemorySaveTool { Self { memory_search, event_context: None, + contract_state: None, } } @@ -48,6 +50,14 @@ impl MemorySaveTool { }); self } + + pub fn with_contract_state( + mut self, + contract_state: Arc, + ) -> Self { + self.contract_state = Some(contract_state); + self + } } /// Error type for memory save tool. @@ -289,19 +299,77 @@ impl Tool for MemorySaveTool { } } - // Generate and store embedding (async to avoid blocking the tokio runtime) - let embedding = self + // Generate and store embedding. On failure, compensate by deleting the + // SQLite row (and any associations already written) so there is no orphan. + let embedding = match self .memory_search .embedding_model_arc() .embed_one(&args.content) .await - .map_err(|e| MemorySaveError(format!("Failed to generate embedding: {e}")))?; + { + Ok(emb) => emb, + Err(embed_err) => { + if let Err(assoc_err) = self + .memory_search + .store() + .delete_associations_for_memory(&memory.id) + .await + { + tracing::error!( + memory_id = %memory.id, + error = %assoc_err, + "compensating association delete failed after embedding generation error" + ); + } + if let Err(del_err) = self.memory_search.store().delete(&memory.id).await { + tracing::error!( + memory_id = %memory.id, + %del_err, + "compensating delete failed after embedding generation error" + ); + } + return Err(MemorySaveError(format!( + "Failed to generate embedding: {embed_err}" + ))); + } + }; - self.memory_search + match self + .memory_search .embedding_table() .store(&memory.id, &args.content, &embedding) .await - .map_err(|e| MemorySaveError(format!("Failed to store embedding: {e}")))?; + { + Ok(()) => { + if let Some(contract_state) = &self.contract_state { + contract_state.record_saved_memory_id(memory.id.clone()); + } + } + Err(embed_err) => { + if let Err(assoc_err) = self + .memory_search + .store() + .delete_associations_for_memory(&memory.id) + .await + { + tracing::error!( + memory_id = %memory.id, + error = %assoc_err, + "compensating association delete failed after embedding store error" + ); + } + if let Err(del_err) = self.memory_search.store().delete(&memory.id).await { + tracing::error!( + memory_id = %memory.id, + %del_err, + "compensating delete failed after embedding store error" + ); + } + return Err(MemorySaveError(format!( + "Failed to store embedding: {embed_err}" + ))); + } + } // Ensure the FTS index exists so full_text_search queries work. // Safe to call repeatedly — no-ops if the index already exists. diff --git a/tests/context_dump.rs b/tests/context_dump.rs index c7b42b1c5..a7595719d 100644 --- a/tests/context_dump.rs +++ b/tests/context_dump.rs @@ -323,6 +323,7 @@ async fn dump_branch_context() { conversation_logger, channel_store, run_logger, + spacebot::tools::BranchToolProfile::Default, ); let tool_defs = branch_tool_server @@ -525,6 +526,7 @@ async fn dump_all_contexts() { conversation_logger, channel_store, run_logger, + spacebot::tools::BranchToolProfile::Default, ); let branch_tool_defs = branch_tool_server.get_tool_defs(None).await.unwrap(); let branch_tools_text = format_tool_defs(&branch_tool_defs);