From fe2a39c12dda42920ab7c29f5a4ee7d8c665cef4 Mon Sep 17 00:00:00 2001 From: Hare Date: Wed, 22 Jul 2026 19:44:52 +0900 Subject: [PATCH] fix: async workspace memory backend --- crates/worker/src/feature/builtin/memory.rs | 23 +- .../src/feature/builtin/session_explore.rs | 5 +- crates/worker/src/worker.rs | 272 ++++++++++-------- 3 files changed, 167 insertions(+), 133 deletions(-) diff --git a/crates/worker/src/feature/builtin/memory.rs b/crates/worker/src/feature/builtin/memory.rs index b7542e07..9014e3bc 100644 --- a/crates/worker/src/feature/builtin/memory.rs +++ b/crates/worker/src/feature/builtin/memory.rs @@ -34,15 +34,15 @@ impl WorkspaceHttpMemoryBackend { } } - pub fn execute_operation( + pub async fn execute_operation( &self, operation: MemoryBackendOperation, ) -> Result { - execute_http_memory_backend(&self.workspace_id, &self.base_url, operation) + execute_http_memory_backend(&self.workspace_id, &self.base_url, operation).await } - fn execute(&self, operation: MemoryBackendOperation) -> Result { - match self.execute_operation(operation) { + async fn execute(&self, operation: MemoryBackendOperation) -> Result { + match self.execute_operation(operation).await { Ok(MemoryBackendOperationResult::ToolOutput(output)) => Ok(tool_output(output)), Ok(result) => Err(ToolError::ExecutionFailed(format!( "unexpected memory backend result for model-visible tool: {result:?}" @@ -70,7 +70,7 @@ pub enum WorkspaceMemoryBackendError { } impl WorkspaceClient { - pub fn execute_memory_backend_operation( + pub async fn execute_memory_backend_operation( &self, operation: MemoryBackendOperation, ) -> Result { @@ -78,7 +78,7 @@ impl WorkspaceClient { WorkspaceClient::Http { workspace_id, base_url, - } => execute_http_memory_backend(workspace_id, base_url, operation), + } => execute_http_memory_backend(workspace_id, base_url, operation).await, WorkspaceClient::Available { kind } => Err(WorkspaceMemoryBackendError::Unavailable { reason: format!( "workspace client kind `{kind}` does not expose the Backend Workspace API" @@ -93,7 +93,7 @@ impl WorkspaceClient { } } -fn execute_http_memory_backend( +async fn execute_http_memory_backend( workspace_id: &str, base_url: &str, operation: MemoryBackendOperation, @@ -103,12 +103,13 @@ fn execute_http_memory_backend( base_url.trim_end_matches('/'), workspace_id ); - let response = reqwest::blocking::Client::new() + let response = reqwest::Client::new() .post(url) .json(&operation) - .send()?; + .send() + .await?; let status = response.status(); - let body = response.text()?; + let body = response.text().await?; if !status.is_success() { return Err(WorkspaceMemoryBackendError::Http { status, body }); } @@ -220,7 +221,7 @@ impl Tool for WorkspaceHttpMemoryTool { _ctx: ToolExecutionContext, ) -> Result { let operation = (self.build)(input_json)?; - self.backend.execute(operation) + self.backend.execute(operation).await } } diff --git a/crates/worker/src/feature/builtin/session_explore.rs b/crates/worker/src/feature/builtin/session_explore.rs index 23cfaa94..5585a332 100644 --- a/crates/worker/src/feature/builtin/session_explore.rs +++ b/crates/worker/src/feature/builtin/session_explore.rs @@ -433,6 +433,7 @@ impl Tool for StageCandidateTool { source_refs, }, )) + .await .map_err(|e| ToolError::ExecutionFailed(format!("write staging failed: {e}")))?; let ids = match result { MemoryBackendOperationResult::StagingWritten(output) if output.staging_count == 1 => { @@ -713,7 +714,7 @@ mod tests { #[tokio::test] async fn stage_candidate_writes_staging_record_with_source_evidence() { let (client, request_rx) = stub_memory_backend_response( - r#"{"Ok":{"result":{"StagingWritten":{"staging_count":1,"staging_ids":["00000000-0000-7000-8000-000000000001"]}}}}"#, + r#"{"status":"ok","result":{"kind":"staging_written","staging_count":1,"staging_ids":["00000000-0000-7000-8000-000000000001"]}}"#, ); let state = SessionExploreState::new( SessionReferenceView::new("segment-1", vec![Item::user_message("durable decision")]), @@ -745,7 +746,7 @@ mod tests { vec!["00000000-0000-7000-8000-000000000001".to_string()] ); let request = request_rx.recv().unwrap(); - assert!(request.contains("\"StageCandidate\"")); + assert!(request.contains("\"operation\":\"stage_candidate\"")); assert!(request.contains("\"kind\":\"decision\"")); assert!(request.contains("\"id\":\"M0000\"")); assert!(request.contains("\"evidence_id\":\"M0000\"")); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 95698ae7..a5a43b91 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -897,12 +897,17 @@ impl Worker { self.workspace_context.client() } - fn resident_summary_from_workspace_authority(&self) -> Result, WorkerError> { - let result = self.workspace_client().execute_memory_backend_operation( - memory::backend::MemoryBackendOperation::ResidentSummary( - memory::backend::MemoryResidentSummaryOperation::default(), - ), - )?; + async fn resident_summary_from_workspace_authority( + &self, + ) -> Result, WorkerError> { + let result = self + .workspace_client() + .execute_memory_backend_operation( + memory::backend::MemoryBackendOperation::ResidentSummary( + memory::backend::MemoryResidentSummaryOperation::default(), + ), + ) + .await?; match result { memory::backend::MemoryBackendOperationResult::ToolOutput(output) => Ok(output.content), other => Err(WorkerError::FeatureInstall(format!( @@ -1465,7 +1470,7 @@ impl Worker { /// Subsequent invocations are no-ops: the template field is /// consumed with `Option::take()`, so the materialised value /// persists across all later turns and compaction. - fn ensure_system_prompt_materialized(&mut self) -> Result<(), WorkerError> { + async fn ensure_system_prompt_materialized(&mut self) -> Result<(), WorkerError> { let Some(template) = self.system_prompt_template.take() else { return Ok(()); }; @@ -1498,7 +1503,7 @@ impl Worker { .as_ref() .is_some_and(|m| m.inject_summary.unwrap_or(true)); let resident_summary: Option = if inject_summary { - match self.resident_summary_from_workspace_authority() { + match self.resident_summary_from_workspace_authority().await { Ok(summary) => summary, Err(error) => { tracing::debug!(%error, "resident memory summary unavailable"); @@ -1579,7 +1584,7 @@ impl Worker { /// first so extract sees a stable history range). async fn prepare_for_run(&mut self) -> Result<(), WorkerError> { self.ensure_interceptor_installed(); - self.ensure_system_prompt_materialized()?; + self.ensure_system_prompt_materialized().await?; self.cleanup_finished_memory_task(); self.ensure_segment_head()?; if self.should_pre_run_compact() { @@ -2897,7 +2902,8 @@ impl Worker { None, None, None, - ); + ) + .await; return Ok(()); }; @@ -2926,7 +2932,8 @@ impl Worker { None, None, None, - ); + ) + .await; return Ok(()); } let result = self.run_extract_once(&memory_cfg, threshold).await; @@ -2997,7 +3004,7 @@ impl Worker { None, None, None, - ); + ).await; return Ok(ExtractDecision::Skipped); } @@ -3008,18 +3015,23 @@ impl Worker { .history() .len(); if current_history_len <= processed_history_len { - audit.emit( - self.workspace_client(), - event_tx, - memory::audit::WorkerLifecycleStatus::Skipped, - "no_new_history_items", - None, - Some(memory::audit::ExtractAudit { - history_range: Some([processed_history_len as u64, current_history_len as u64]), - ..Default::default() - }), - None, - ); + audit + .emit( + self.workspace_client(), + event_tx, + memory::audit::WorkerLifecycleStatus::Skipped, + "no_new_history_items", + None, + Some(memory::audit::ExtractAudit { + history_range: Some([ + processed_history_len as u64, + current_history_len as u64, + ]), + ..Default::default() + }), + None, + ) + .await; return Ok(ExtractDecision::Skipped); } @@ -3031,15 +3043,17 @@ impl Worker { .read_all(self.session_id(), self.segment_id())? .len(); if entries_now == 0 { - audit.emit( - self.workspace_client(), - event_tx, - memory::audit::WorkerLifecycleStatus::Skipped, - "empty_segment_log", - None, - None, - None, - ); + audit + .emit( + self.workspace_client(), + event_tx, + memory::audit::WorkerLifecycleStatus::Skipped, + "empty_segment_log", + None, + None, + None, + ) + .await; return Ok(ExtractDecision::Skipped); } let end_entry = entries_now - 1; @@ -3048,21 +3062,26 @@ impl Worker { .map(|p| p.processed_through_entry + 1) .unwrap_or(0); if start_entry > end_entry { - audit.emit( - self.workspace_client(), - event_tx, - memory::audit::WorkerLifecycleStatus::Skipped, - "no_new_segment_entries", - None, - Some(memory::audit::ExtractAudit { - session_id: Some(self.session_id().to_string()), - segment_id: Some(self.segment_id().to_string()), - entry_range: Some([start_entry as u64, end_entry as u64]), - history_range: Some([processed_history_len as u64, current_history_len as u64]), - ..Default::default() - }), - None, - ); + audit + .emit( + self.workspace_client(), + event_tx, + memory::audit::WorkerLifecycleStatus::Skipped, + "no_new_segment_entries", + None, + Some(memory::audit::ExtractAudit { + session_id: Some(self.session_id().to_string()), + segment_id: Some(self.segment_id().to_string()), + entry_range: Some([start_entry as u64, end_entry as u64]), + history_range: Some([ + processed_history_len as u64, + current_history_len as u64, + ]), + ..Default::default() + }), + None, + ) + .await; return Ok(ExtractDecision::Skipped); } @@ -3073,15 +3092,19 @@ impl Worker { history_range: Some([processed_history_len as u64, current_history_len as u64]), ..Default::default() }; - audit.emit( - self.workspace_client(), - event_tx, - memory::audit::WorkerLifecycleStatus::Started, - format!("token_threshold_reached tokens_since={tokens_since} threshold={threshold}"), - None, - Some(extract_audit_base.clone()), - None, - ); + audit + .emit( + self.workspace_client(), + event_tx, + memory::audit::WorkerLifecycleStatus::Started, + format!( + "token_threshold_reached tokens_since={tokens_since} threshold={threshold}" + ), + None, + Some(extract_audit_base.clone()), + None, + ) + .await; let items_to_extract = self.engine.as_ref().expect("worker present").history() [processed_history_len..current_history_len] @@ -3094,15 +3117,17 @@ impl Worker { let client = match self.build_extractor_client(memory_cfg) { Ok(client) => client, Err(err) => { - audit.emit( - self.workspace_client(), - event_tx, - memory::audit::WorkerLifecycleStatus::Failed, - format!("client_build_failed: {err}"), - None, - Some(extract_audit_base), - None, - ); + audit + .emit( + self.workspace_client(), + event_tx, + memory::audit::WorkerLifecycleStatus::Failed, + format!("client_build_failed: {err}"), + None, + Some(extract_audit_base), + None, + ) + .await; return Err(err); } }; @@ -3110,15 +3135,17 @@ impl Worker { let extract_system_prompt = match self.prompts.memory_extract_system(memory_language) { Ok(prompt) => prompt, Err(err) => { - audit.emit( - self.workspace_client(), - event_tx, - memory::audit::WorkerLifecycleStatus::Failed, - format!("prompt_render_failed: {err}"), - None, - Some(extract_audit_base), - None, - ); + audit + .emit( + self.workspace_client(), + event_tx, + memory::audit::WorkerLifecycleStatus::Failed, + format!("prompt_render_failed: {err}"), + None, + Some(extract_audit_base), + None, + ) + .await; return Err(WorkerError::PromptCatalog(err)); } }; @@ -3151,15 +3178,17 @@ impl Worker { .iter() .any(|installed| installed == name) }) { - audit.emit( - self.workspace_client(), - event_tx, - memory::audit::WorkerLifecycleStatus::Failed, - "session_explore_feature_install_failed", - None, - Some(extract_audit_base), - None, - ); + audit + .emit( + self.workspace_client(), + event_tx, + memory::audit::WorkerLifecycleStatus::Failed, + "session_explore_feature_install_failed", + None, + Some(extract_audit_base), + None, + ) + .await; return Err(WorkerError::FeatureInstall( "session-explore feature install failed".to_string(), )); @@ -3178,15 +3207,17 @@ impl Worker { Ok(result) => result.usage.as_ref().map(usage_audit_from_event), Err(err) => { let usage = err.usage.as_ref().map(usage_audit_from_event); - audit.emit( - self.workspace_client(), - event_tx, - lifecycle_status_for_worker_error(&err.source), - format!("worker_failed: {}", err.source), - usage, - Some(extract_audit_base), - None, - ); + audit + .emit( + self.workspace_client(), + event_tx, + lifecycle_status_for_worker_error(&err.source), + format!("worker_failed: {}", err.source), + usage, + Some(extract_audit_base), + None, + ) + .await; return Err(WorkerError::Engine(err.source)); } }; @@ -3228,15 +3259,17 @@ impl Worker { } else { "completed_staging_written" }; - audit.emit( - self.workspace_client(), - event_tx, - memory::audit::WorkerLifecycleStatus::Completed, - reason, - usage, - Some(extract_audit), - None, - ); + audit + .emit( + self.workspace_client(), + event_tx, + memory::audit::WorkerLifecycleStatus::Completed, + reason, + usage, + Some(extract_audit), + None, + ) + .await; Ok(ExtractDecision::Completed) } @@ -3274,7 +3307,8 @@ impl Worker { None, None, None, - ); + ) + .await; if reason == "consolidation_backend_operation_unavailable" { tracing::debug!( "workspace memory consolidation skipped: backend operation is unavailable" @@ -3357,7 +3391,7 @@ impl WorkerAuditBase { } } - fn emit( + async fn emit( &self, workspace_client: &WorkspaceClient, event_tx: Option<&broadcast::Sender>, @@ -3379,15 +3413,15 @@ impl WorkerAuditBase { extract, consolidation, }; - let _ = workspace_client.execute_memory_backend_operation( - memory::backend::MemoryBackendOperation::AppendAudit( + let _ = workspace_client + .execute_memory_backend_operation(memory::backend::MemoryBackendOperation::AppendAudit( memory::backend::MemoryAppendAuditOperation { event: memory::audit::AuditEvent::new( memory::audit::AuditPayload::WorkerLifecycle(payload), ), }, - ), - ); + )) + .await; if should_emit_memory_worker_event(self.worker, status, &reason) { emit_memory_worker_event( event_tx, @@ -5552,7 +5586,7 @@ mod build_summary_prompt_tests { ) .unwrap(); worker.set_system_prompt_template(template); - worker.ensure_system_prompt_materialized().unwrap(); + worker.ensure_system_prompt_materialized().await.unwrap(); worker.engine().get_system_prompt().unwrap().to_string() } @@ -5583,17 +5617,15 @@ mod build_summary_prompt_tests { let mut buffer = [0_u8; 1024]; let _ = stream.read(&mut buffer).unwrap(); let body = serde_json::json!({ - "Ok": { - "result": { - "ToolOutput": { - "summary": if content.is_some() { - "resident memory summary collected" - } else { - "resident memory summary unavailable" - }, - "content": content, - } - } + "status": "ok", + "result": { + "kind": "tool_output", + "summary": if content.is_some() { + "resident memory summary collected" + } else { + "resident memory summary unavailable" + }, + "content": content, } }) .to_string();