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,
operation: MemoryBackendOperation,
) -> 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> {
match self.execute_operation(operation) {
async fn execute(&self, operation: MemoryBackendOperation) -> Result<ToolOutput, ToolError> {
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<MemoryBackendOperationResult, WorkspaceMemoryBackendError> {
@ -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<ToolOutput, ToolError> {
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,
},
))
.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\""));

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

@ -897,12 +897,17 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
self.workspace_context.client()
}
fn resident_summary_from_workspace_authority(&self) -> Result<Option<String>, WorkerError> {
let result = self.workspace_client().execute_memory_backend_operation(
async fn resident_summary_from_workspace_authority(
&self,
) -> Result<Option<String>, 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<C: LlmClient, St: Store> Worker<C, St> {
/// 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<C: LlmClient, St: Store> Worker<C, St> {
.as_ref()
.is_some_and(|m| m.inject_summary.unwrap_or(true));
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,
Err(error) => {
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).
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<C: LlmClient, St: Store> Worker<C, St> {
None,
None,
None,
);
)
.await;
return Ok(());
};
@ -2926,7 +2932,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
None,
None,
None,
);
)
.await;
return Ok(());
}
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,
);
).await;
return Ok(ExtractDecision::Skipped);
}
@ -3008,18 +3015,23 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.history()
.len();
if current_history_len <= processed_history_len {
audit.emit(
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]),
history_range: Some([
processed_history_len as u64,
current_history_len as u64,
]),
..Default::default()
}),
None,
);
)
.await;
return Ok(ExtractDecision::Skipped);
}
@ -3031,7 +3043,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.read_all(self.session_id(), self.segment_id())?
.len();
if entries_now == 0 {
audit.emit(
audit
.emit(
self.workspace_client(),
event_tx,
memory::audit::WorkerLifecycleStatus::Skipped,
@ -3039,7 +3052,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
None,
None,
None,
);
)
.await;
return Ok(ExtractDecision::Skipped);
}
let end_entry = entries_now - 1;
@ -3048,7 +3062,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.map(|p| p.processed_through_entry + 1)
.unwrap_or(0);
if start_entry > end_entry {
audit.emit(
audit
.emit(
self.workspace_client(),
event_tx,
memory::audit::WorkerLifecycleStatus::Skipped,
@ -3058,11 +3073,15 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
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]),
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<C: LlmClient, St: Store> Worker<C, St> {
history_range: Some([processed_history_len as u64, current_history_len as u64]),
..Default::default()
};
audit.emit(
audit
.emit(
self.workspace_client(),
event_tx,
memory::audit::WorkerLifecycleStatus::Started,
format!("token_threshold_reached tokens_since={tokens_since} threshold={threshold}"),
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,7 +3117,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
let client = match self.build_extractor_client(memory_cfg) {
Ok(client) => client,
Err(err) => {
audit.emit(
audit
.emit(
self.workspace_client(),
event_tx,
memory::audit::WorkerLifecycleStatus::Failed,
@ -3102,7 +3126,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
None,
Some(extract_audit_base),
None,
);
)
.await;
return Err(err);
}
};
@ -3110,7 +3135,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
let extract_system_prompt = match self.prompts.memory_extract_system(memory_language) {
Ok(prompt) => prompt,
Err(err) => {
audit.emit(
audit
.emit(
self.workspace_client(),
event_tx,
memory::audit::WorkerLifecycleStatus::Failed,
@ -3118,7 +3144,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
None,
Some(extract_audit_base),
None,
);
)
.await;
return Err(WorkerError::PromptCatalog(err));
}
};
@ -3151,7 +3178,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.iter()
.any(|installed| installed == name)
}) {
audit.emit(
audit
.emit(
self.workspace_client(),
event_tx,
memory::audit::WorkerLifecycleStatus::Failed,
@ -3159,7 +3187,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
None,
Some(extract_audit_base),
None,
);
)
.await;
return Err(WorkerError::FeatureInstall(
"session-explore feature install failed".to_string(),
));
@ -3178,7 +3207,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
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(
audit
.emit(
self.workspace_client(),
event_tx,
lifecycle_status_for_worker_error(&err.source),
@ -3186,7 +3216,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
usage,
Some(extract_audit_base),
None,
);
)
.await;
return Err(WorkerError::Engine(err.source));
}
};
@ -3228,7 +3259,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
} else {
"completed_staging_written"
};
audit.emit(
audit
.emit(
self.workspace_client(),
event_tx,
memory::audit::WorkerLifecycleStatus::Completed,
@ -3236,7 +3268,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
usage,
Some(extract_audit),
None,
);
)
.await;
Ok(ExtractDecision::Completed)
}
@ -3274,7 +3307,8 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
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<Event>>,
@ -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,9 +5617,9 @@ mod build_summary_prompt_tests {
let mut buffer = [0_u8; 1024];
let _ = stream.read(&mut buffer).unwrap();
let body = serde_json::json!({
"Ok": {
"status": "ok",
"result": {
"ToolOutput": {
"kind": "tool_output",
"summary": if content.is_some() {
"resident memory summary collected"
} else {
@ -5593,8 +5627,6 @@ mod build_summary_prompt_tests {
},
"content": content,
}
}
}
})
.to_string();
let response = format!(