Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: async workspace memory backend

Browse Source
This commit is contained in:
Keisuke Hirata 2026-07-22 19:44:52 +09:00
parent aa542b38d9
commit fe2a39c12d
No known key found for this signature in database
3 changed files with 167 additions and 133 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
crates/worker/src/feature/builtin/memory.rs
View File

@ -34,15 +34,15 @@ impl WorkspaceHttpMemoryBackend {
} }
} }
pub fn execute_operation( pub async fn execute_operation(
&self, &self,
operation: MemoryBackendOperation, operation: MemoryBackendOperation,
) -> Result<MemoryBackendOperationResult, WorkspaceMemoryBackendError> { ) -> Result<MemoryBackendOperationResult, WorkspaceMemoryBackendError> {
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<ToolOutput, ToolError> { async fn execute(&self, operation: MemoryBackendOperation) -> Result<ToolOutput, ToolError> {
match self.execute_operation(operation) { match self.execute_operation(operation).await {
Ok(MemoryBackendOperationResult::ToolOutput(output)) => Ok(tool_output(output)), Ok(MemoryBackendOperationResult::ToolOutput(output)) => Ok(tool_output(output)),
Ok(result) => Err(ToolError::ExecutionFailed(format!( Ok(result) => Err(ToolError::ExecutionFailed(format!(
"unexpected memory backend result for model-visible tool: {result:?}" "unexpected memory backend result for model-visible tool: {result:?}"
@ -70,7 +70,7 @@ pub enum WorkspaceMemoryBackendError {
} }
impl WorkspaceClient { impl WorkspaceClient {
pub fn execute_memory_backend_operation( pub async fn execute_memory_backend_operation(
&self, &self,
operation: MemoryBackendOperation, operation: MemoryBackendOperation,
) -> Result<MemoryBackendOperationResult, WorkspaceMemoryBackendError> { ) -> Result<MemoryBackendOperationResult, WorkspaceMemoryBackendError> {
@ -78,7 +78,7 @@ impl WorkspaceClient {
WorkspaceClient::Http { WorkspaceClient::Http {
workspace_id, workspace_id,
base_url, 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 { WorkspaceClient::Available { kind } => Err(WorkspaceMemoryBackendError::Unavailable {
reason: format!( reason: format!(
"workspace client kind `{kind}` does not expose the Backend Workspace API" "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, workspace_id: &str,
base_url: &str, base_url: &str,
operation: MemoryBackendOperation, operation: MemoryBackendOperation,
@ -103,12 +103,13 @@ fn execute_http_memory_backend(
base_url.trim_end_matches('/'), base_url.trim_end_matches('/'),
workspace_id workspace_id
); );
let response = reqwest::blocking::Client::new() let response = reqwest::Client::new()
.post(url) .post(url)
.json(&operation) .json(&operation)
.send()?; .send()
.await?;
let status = response.status(); let status = response.status();
let body = response.text()?; let body = response.text().await?;
if !status.is_success() { if !status.is_success() {
return Err(WorkspaceMemoryBackendError::Http { status, body }); return Err(WorkspaceMemoryBackendError::Http { status, body });
} }
@ -220,7 +221,7 @@ impl Tool for WorkspaceHttpMemoryTool {
_ctx: ToolExecutionContext, _ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> { ) -> Result<ToolOutput, ToolError> {
let operation = (self.build)(input_json)?; let operation = (self.build)(input_json)?;
self.backend.execute(operation) self.backend.execute(operation).await
} }
} }

5
crates/worker/src/feature/builtin/session_explore.rs
View File

@ -433,6 +433,7 @@ impl Tool for StageCandidateTool {
source_refs, source_refs,
}, },
)) ))
.await
.map_err(|e| ToolError::ExecutionFailed(format!("write staging failed: {e}")))?; .map_err(|e| ToolError::ExecutionFailed(format!("write staging failed: {e}")))?;
let ids = match result { let ids = match result {
MemoryBackendOperationResult::StagingWritten(output) if output.staging_count == 1 => { MemoryBackendOperationResult::StagingWritten(output) if output.staging_count == 1 => {
@ -713,7 +714,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn stage_candidate_writes_staging_record_with_source_evidence() { async fn stage_candidate_writes_staging_record_with_source_evidence() {
let (client, request_rx) = stub_memory_backend_response( 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( let state = SessionExploreState::new(
SessionReferenceView::new("segment-1", vec![Item::user_message("durable decision")]), SessionReferenceView::new("segment-1", vec![Item::user_message("durable decision")]),
@ -745,7 +746,7 @@ mod tests {
vec!["00000000-0000-7000-8000-000000000001".to_string()] vec!["00000000-0000-7000-8000-000000000001".to_string()]
); );
let request = request_rx.recv().unwrap(); let request = request_rx.recv().unwrap();
assert!(request.contains("\"StageCandidate\"")); assert!(request.contains("\"operation\":\"stage_candidate\""));
assert!(request.contains("\"kind\":\"decision\"")); assert!(request.contains("\"kind\":\"decision\""));
assert!(request.contains("\"id\":\"M0000\"")); assert!(request.contains("\"id\":\"M0000\""));
assert!(request.contains("\"evidence_id\":\"M0000\"")); assert!(request.contains("\"evidence_id\":\"M0000\""));

272
crates/worker/src/worker.rs
View File

@ -897,12 +897,17 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
self.workspace_context.client() self.workspace_context.client()
} }
fn resident_summary_from_workspace_authority(&self) -> Result<Option<String>, WorkerError> { async fn resident_summary_from_workspace_authority(
let result = self.workspace_client().execute_memory_backend_operation( &self,
memory::backend::MemoryBackendOperation::ResidentSummary( ) -> Result<Option<String>, WorkerError> {
memory::backend::MemoryResidentSummaryOperation::default(), let result = self
), .workspace_client()
)?; .execute_memory_backend_operation(
memory::backend::MemoryBackendOperation::ResidentSummary(
memory::backend::MemoryResidentSummaryOperation::default(),
),
)
.await?;
match result { match result {
memory::backend::MemoryBackendOperationResult::ToolOutput(output) => Ok(output.content), memory::backend::MemoryBackendOperationResult::ToolOutput(output) => Ok(output.content),
other => Err(WorkerError::FeatureInstall(format!( other => Err(WorkerError::FeatureInstall(format!(
@ -1465,7 +1470,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// Subsequent invocations are no-ops: the template field is /// Subsequent invocations are no-ops: the template field is
/// consumed with `Option::take()`, so the materialised value /// consumed with `Option::take()`, so the materialised value
/// persists across all later turns and compaction. /// 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 { let Some(template) = self.system_prompt_template.take() else {
return Ok(()); return Ok(());
}; };
@ -1498,7 +1503,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.as_ref() .as_ref()
.is_some_and(|m| m.inject_summary.unwrap_or(true)); .is_some_and(|m| m.inject_summary.unwrap_or(true));
let resident_summary: Option<String> = if inject_summary { let resident_summary: Option<String> = if inject_summary {
match self.resident_summary_from_workspace_authority() { match self.resident_summary_from_workspace_authority().await {
Ok(summary) => summary, Ok(summary) => summary,
Err(error) => { Err(error) => {
tracing::debug!(%error, "resident memory summary unavailable"); tracing::debug!(%error, "resident memory summary unavailable");
@ -1579,7 +1584,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// first so extract sees a stable history range). /// first so extract sees a stable history range).
async fn prepare_for_run(&mut self) -> Result<(), WorkerError> { async fn prepare_for_run(&mut self) -> Result<(), WorkerError> {
self.ensure_interceptor_installed(); self.ensure_interceptor_installed();
self.ensure_system_prompt_materialized()?; self.ensure_system_prompt_materialized().await?;
self.cleanup_finished_memory_task(); self.cleanup_finished_memory_task();
self.ensure_segment_head()?; self.ensure_segment_head()?;
if self.should_pre_run_compact() { if self.should_pre_run_compact() {
@ -2897,7 +2902,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
None, None,
None, None,
None, None,
); )
.await;
return Ok(()); return Ok(());
}; };
@ -2926,7 +2932,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
None, None,
None, None,
None, None,
); )
.await;
return Ok(()); return Ok(());
} }
let result = self.run_extract_once(&memory_cfg, threshold).await; let result = self.run_extract_once(&memory_cfg, threshold).await;
@ -2997,7 +3004,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
None, None,
None, None,
None, None,
); ).await;
return Ok(ExtractDecision::Skipped); return Ok(ExtractDecision::Skipped);
} }
@ -3008,18 +3015,23 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.history() .history()
.len(); .len();
if current_history_len <= processed_history_len { if current_history_len <= processed_history_len {
audit.emit( audit
self.workspace_client(), .emit(
event_tx, self.workspace_client(),
memory::audit::WorkerLifecycleStatus::Skipped, event_tx,
"no_new_history_items", memory::audit::WorkerLifecycleStatus::Skipped,
None, "no_new_history_items",
Some(memory::audit::ExtractAudit { None,
history_range: Some([processed_history_len as u64, current_history_len as u64]), Some(memory::audit::ExtractAudit {
..Default::default() history_range: Some([
}), processed_history_len as u64,
None, current_history_len as u64,
); ]),
..Default::default()
}),
None,
)
.await;
return Ok(ExtractDecision::Skipped); return Ok(ExtractDecision::Skipped);
} }
@ -3031,15 +3043,17 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.read_all(self.session_id(), self.segment_id())? .read_all(self.session_id(), self.segment_id())?
.len(); .len();
if entries_now == 0 { if entries_now == 0 {
audit.emit( audit
self.workspace_client(), .emit(
event_tx, self.workspace_client(),
memory::audit::WorkerLifecycleStatus::Skipped, event_tx,
"empty_segment_log", memory::audit::WorkerLifecycleStatus::Skipped,
None, "empty_segment_log",
None, None,
None, None,
); None,
)
.await;
return Ok(ExtractDecision::Skipped); return Ok(ExtractDecision::Skipped);
} }
let end_entry = entries_now - 1; let end_entry = entries_now - 1;
@ -3048,21 +3062,26 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.map(|p| p.processed_through_entry + 1) .map(|p| p.processed_through_entry + 1)
.unwrap_or(0); .unwrap_or(0);
if start_entry > end_entry { if start_entry > end_entry {
audit.emit( audit
self.workspace_client(), .emit(
event_tx, self.workspace_client(),
memory::audit::WorkerLifecycleStatus::Skipped, event_tx,
"no_new_segment_entries", memory::audit::WorkerLifecycleStatus::Skipped,
None, "no_new_segment_entries",
Some(memory::audit::ExtractAudit { None,
session_id: Some(self.session_id().to_string()), Some(memory::audit::ExtractAudit {
segment_id: Some(self.segment_id().to_string()), session_id: Some(self.session_id().to_string()),
entry_range: Some([start_entry as u64, end_entry as u64]), segment_id: Some(self.segment_id().to_string()),
history_range: Some([processed_history_len as u64, current_history_len as u64]), entry_range: Some([start_entry as u64, end_entry as u64]),
..Default::default() history_range: Some([
}), processed_history_len as u64,
None, current_history_len as u64,
); ]),
..Default::default()
}),
None,
)
.await;
return Ok(ExtractDecision::Skipped); return Ok(ExtractDecision::Skipped);
} }
@ -3073,15 +3092,19 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
history_range: Some([processed_history_len as u64, current_history_len as u64]), history_range: Some([processed_history_len as u64, current_history_len as u64]),
..Default::default() ..Default::default()
}; };
audit.emit( audit
self.workspace_client(), .emit(
event_tx, self.workspace_client(),
memory::audit::WorkerLifecycleStatus::Started, event_tx,
format!("token_threshold_reached tokens_since={tokens_since} threshold={threshold}"), memory::audit::WorkerLifecycleStatus::Started,
None, format!(
Some(extract_audit_base.clone()), "token_threshold_reached tokens_since={tokens_since} threshold={threshold}"
None, ),
); None,
Some(extract_audit_base.clone()),
None,
)
.await;
let items_to_extract = self.engine.as_ref().expect("worker present").history() let items_to_extract = self.engine.as_ref().expect("worker present").history()
[processed_history_len..current_history_len] [processed_history_len..current_history_len]
@ -3094,15 +3117,17 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
let client = match self.build_extractor_client(memory_cfg) { let client = match self.build_extractor_client(memory_cfg) {
Ok(client) => client, Ok(client) => client,
Err(err) => { Err(err) => {
audit.emit( audit
self.workspace_client(), .emit(
event_tx, self.workspace_client(),
memory::audit::WorkerLifecycleStatus::Failed, event_tx,
format!("client_build_failed: {err}"), memory::audit::WorkerLifecycleStatus::Failed,
None, format!("client_build_failed: {err}"),
Some(extract_audit_base), None,
None, Some(extract_audit_base),
); None,
)
.await;
return Err(err); return Err(err);
} }
}; };
@ -3110,15 +3135,17 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
let extract_system_prompt = match self.prompts.memory_extract_system(memory_language) { let extract_system_prompt = match self.prompts.memory_extract_system(memory_language) {
Ok(prompt) => prompt, Ok(prompt) => prompt,
Err(err) => { Err(err) => {
audit.emit( audit
self.workspace_client(), .emit(
event_tx, self.workspace_client(),
memory::audit::WorkerLifecycleStatus::Failed, event_tx,
format!("prompt_render_failed: {err}"), memory::audit::WorkerLifecycleStatus::Failed,
None, format!("prompt_render_failed: {err}"),
Some(extract_audit_base), None,
None, Some(extract_audit_base),
); None,
)
.await;
return Err(WorkerError::PromptCatalog(err)); return Err(WorkerError::PromptCatalog(err));
} }
}; };
@ -3151,15 +3178,17 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.iter() .iter()
.any(|installed| installed == name) .any(|installed| installed == name)
}) { }) {
audit.emit( audit
self.workspace_client(), .emit(
event_tx, self.workspace_client(),
memory::audit::WorkerLifecycleStatus::Failed, event_tx,
"session_explore_feature_install_failed", memory::audit::WorkerLifecycleStatus::Failed,
None, "session_explore_feature_install_failed",
Some(extract_audit_base), None,
None, Some(extract_audit_base),
); None,
)
.await;
return Err(WorkerError::FeatureInstall( return Err(WorkerError::FeatureInstall(
"session-explore feature install failed".to_string(), "session-explore feature install failed".to_string(),
)); ));
@ -3178,15 +3207,17 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
Ok(result) => result.usage.as_ref().map(usage_audit_from_event), Ok(result) => result.usage.as_ref().map(usage_audit_from_event),
Err(err) => { Err(err) => {
let usage = err.usage.as_ref().map(usage_audit_from_event); let usage = err.usage.as_ref().map(usage_audit_from_event);
audit.emit( audit
self.workspace_client(), .emit(
event_tx, self.workspace_client(),
lifecycle_status_for_worker_error(&err.source), event_tx,
format!("worker_failed: {}", err.source), lifecycle_status_for_worker_error(&err.source),
usage, format!("worker_failed: {}", err.source),
Some(extract_audit_base), usage,
None, Some(extract_audit_base),
); None,
)
.await;
return Err(WorkerError::Engine(err.source)); return Err(WorkerError::Engine(err.source));
} }
}; };
@ -3228,15 +3259,17 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
} else { } else {
"completed_staging_written" "completed_staging_written"
}; };
audit.emit( audit
self.workspace_client(), .emit(
event_tx, self.workspace_client(),
memory::audit::WorkerLifecycleStatus::Completed, event_tx,
reason, memory::audit::WorkerLifecycleStatus::Completed,
usage, reason,
Some(extract_audit), usage,
None, Some(extract_audit),
); None,
)
.await;
Ok(ExtractDecision::Completed) Ok(ExtractDecision::Completed)
} }
@ -3274,7 +3307,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
None, None,
None, None,
None, None,
); )
.await;
if reason == "consolidation_backend_operation_unavailable" { if reason == "consolidation_backend_operation_unavailable" {
tracing::debug!( tracing::debug!(
"workspace memory consolidation skipped: backend operation is unavailable" "workspace memory consolidation skipped: backend operation is unavailable"
@ -3357,7 +3391,7 @@ impl WorkerAuditBase {
} }
} }
fn emit( async fn emit(
&self, &self,
workspace_client: &WorkspaceClient, workspace_client: &WorkspaceClient,
event_tx: Option<&broadcast::Sender<Event>>, event_tx: Option<&broadcast::Sender<Event>>,
@ -3379,15 +3413,15 @@ impl WorkerAuditBase {
extract, extract,
consolidation, consolidation,
}; };
let _ = workspace_client.execute_memory_backend_operation( let _ = workspace_client
memory::backend::MemoryBackendOperation::AppendAudit( .execute_memory_backend_operation(memory::backend::MemoryBackendOperation::AppendAudit(
memory::backend::MemoryAppendAuditOperation { memory::backend::MemoryAppendAuditOperation {
event: memory::audit::AuditEvent::new( event: memory::audit::AuditEvent::new(
memory::audit::AuditPayload::WorkerLifecycle(payload), memory::audit::AuditPayload::WorkerLifecycle(payload),
), ),
}, },
), ))
); .await;
if should_emit_memory_worker_event(self.worker, status, &reason) { if should_emit_memory_worker_event(self.worker, status, &reason) {
emit_memory_worker_event( emit_memory_worker_event(
event_tx, event_tx,
@ -5552,7 +5586,7 @@ mod build_summary_prompt_tests {
) )
.unwrap(); .unwrap();
worker.set_system_prompt_template(template); 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() worker.engine().get_system_prompt().unwrap().to_string()
} }
@ -5583,17 +5617,15 @@ mod build_summary_prompt_tests {
let mut buffer = [0_u8; 1024]; let mut buffer = [0_u8; 1024];
let _ = stream.read(&mut buffer).unwrap(); let _ = stream.read(&mut buffer).unwrap();
let body = serde_json::json!({ let body = serde_json::json!({
"Ok": { "status": "ok",
"result": { "result": {
"ToolOutput": { "kind": "tool_output",
"summary": if content.is_some() { "summary": if content.is_some() {
"resident memory summary collected" "resident memory summary collected"
} else { } else {
"resident memory summary unavailable" "resident memory summary unavailable"
}, },
"content": content, "content": content,
}
}
} }
}) })
.to_string(); .to_string();