diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs index 39994a464..7fb6863d9 100644 --- a/crates/worker/src/feature/builtin/memory_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -503,9 +503,10 @@ impl MemoryLifecycleTask { }; let manifest = self.extraction_manifest(); + let cancellation_observer = cancellation.clone(); let cancel_observer = move |sender: tokio::sync::mpsc::Sender<()>| { tokio::spawn(async move { - cancellation.cancelled().await; + cancellation_observer.cancelled().await; let _ = sender.send(()).await; }); }; @@ -610,6 +611,21 @@ impl MemoryLifecycleTask { ExtractionDisposition::Completed => {} } + if cancellation.is_cancelled() { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Cancelled, + "memory-extract cancelled before completion side effects", + usage_audit, + extract_audit, + None, + ) + .await; + return Ok(()); + } + if self.target == ExtractionTarget::Subjektiv && let Err(reason) = self.record_subjektiv_session(&capture.session_id) { @@ -634,6 +650,20 @@ impl MemoryLifecycleTask { staging_id: pointer_staging_id, }; let payload = serde_json::to_value(&next_pointer).map_err(hook_internal)?; + if cancellation.is_cancelled() { + audit + .emit( + self.workspace_client.as_ref(), + self.event_tx.as_ref(), + memory::audit::WorkerLifecycleStatus::Cancelled, + "memory-extract cancelled before pointer commit", + usage_audit, + extract_audit, + None, + ) + .await; + return Ok(()); + } if !self .extensions .append_if_current(&capture.location(), self.pointer_domain(), payload) @@ -1010,7 +1040,7 @@ fn model_audit_from_manifest(model: &manifest::ModelManifest) -> memory::audit:: #[cfg(test)] mod tests { use std::pin::Pin; - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent}; @@ -1110,6 +1140,52 @@ mod tests { } } + #[derive(Debug, Default)] + struct BlockingSessionWorkspaceClient { + requests: Mutex>, + session_started: AtomicBool, + release_session: AtomicBool, + } + + impl WorkspaceClient for BlockingSessionWorkspaceClient { + fn workspace_id(&self) -> Option<&str> { + Some("workspace-1") + } + + fn kind(&self) -> &str { + "blocking-subjektiv-session-test" + } + + fn is_available(&self) -> bool { + true + } + + fn execute( + &self, + request: crate::worker::WorkspaceRequest, + ) -> Result { + let is_subjektiv_session = request.path.ends_with("/subjektiv/sessions"); + self.requests.lock().unwrap().push(request); + if !is_subjektiv_session { + return Err(crate::worker::WorkspaceClientError::Unavailable( + "blocking client only accepts subject-session recording".to_string(), + )); + } + self.session_started.store(true, Ordering::Release); + while !self.release_session.load(Ordering::Acquire) { + std::thread::yield_now(); + } + Ok(crate::worker::WorkspaceResponse { + status: 200, + body: serde_json::to_string(&server_api::SubjektivRecordSessionResponse { + subject_id: "subject-1".to_string(), + session_id: "session-1".to_string(), + }) + .unwrap(), + }) + } + } + #[derive(Clone)] struct PendingClient { calls: Arc, @@ -1599,6 +1675,50 @@ permission = "write" ); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn cancellation_during_subject_session_recording_never_commits_pointer() { + let client = ScriptClient::new(vec![finish_empty_events("finish-1"), completed_events()]); + let extension_writes = Arc::new(Mutex::new(Vec::new())); + let (event_tx, mut event_rx) = broadcast::channel(32); + let workspace_client = Arc::new(BlockingSessionWorkspaceClient::default()); + let mut task = test_task( + capture(2, 250), + Box::new(client), + Arc::clone(&extension_writes), + event_tx, + workspace_client.clone(), + ); + task.target = ExtractionTarget::Subjektiv; + task.config.profile.consolidation.request_enabled = false; + let registry = start_background_task(task); + tokio::time::timeout(Duration::from_secs(5), async { + while !workspace_client.session_started.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("subject-session recording should begin"); + + let release_client = workspace_client.clone(); + let release = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(50)).await; + release_client + .release_session + .store(true, Ordering::Release); + }); + let rewrite_guard = registry.begin_session_rewrite().await.unwrap(); + release.await.unwrap(); + drop(rewrite_guard); + registry.shutdown().await.unwrap(); + + assert!(extension_writes.lock().unwrap().is_empty()); + assert_eq!(workspace_client.requests.lock().unwrap().len(), 1); + let events = std::iter::from_fn(|| event_rx.try_recv().ok()).collect::>(); + assert!(events.iter().any(|event| { + matches!(event, Event::MemoryWorker(event) if event.status == "cancelled") + })); + } + #[tokio::test] async fn subjektiv_candidate_remains_staged_when_later_extraction_fails_without_pointer() { let source = capture(2, 250); diff --git a/crates/workspace-server/src/lib.rs b/crates/workspace-server/src/lib.rs index 2210f0866..c23670ced 100644 --- a/crates/workspace-server/src/lib.rs +++ b/crates/workspace-server/src/lib.rs @@ -112,6 +112,8 @@ pub enum Error { }, #[error("unknown local repository `{0}`")] UnknownRepository(String), + #[error("subjektiv subject `{0}` was not found in this Workspace")] + SubjektivSubjectNotFound(String), #[error( "Merge Request reopen requires an authenticated Browser session and explicit confirmation" )] diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 11b0a321e..34e251b37 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -16619,9 +16619,7 @@ async fn scoped_get_subjektiv_subject( let subject = open_subjektiv_store(&api)? .subject(&path.subject_id) .map_err(|error| Error::Store(error.to_string()))? - .ok_or_else(|| { - Error::InvalidInput(format!("unknown subjektiv subject `{}`", path.subject_id)) - })?; + .ok_or_else(|| Error::SubjektivSubjectNotFound(path.subject_id.clone()))?; Ok(Json(subjektiv_subject_response(&api, subject))) } @@ -16648,9 +16646,7 @@ async fn scoped_start_subjektiv_subject_worker( let subject = open_subjektiv_store(&api)? .subject(&path.subject_id) .map_err(|error| Error::Store(error.to_string()))? - .ok_or_else(|| { - Error::InvalidInput(format!("unknown subjektiv subject `{}`", path.subject_id)) - })?; + .ok_or_else(|| Error::SubjektivSubjectNotFound(path.subject_id.clone()))?; if subject.state != crate::subjektiv::SubjectState::Active { return Err(Error::InvalidInput(format!( "subjektiv subject `{}` is retired", @@ -28062,6 +28058,7 @@ fn api_error_status(error: &Error) -> StatusCode { | Error::UnknownRuntime(_) | Error::UnknownWorker { .. } | Error::UnknownRepository(_) + | Error::SubjektivSubjectNotFound(_) | Error::RuntimeBindingNotFound { .. } | Error::WorkspaceIdMismatch => StatusCode::NOT_FOUND, Error::RuntimeOperationFailed { code, .. } if code == "skill_not_found" => { @@ -34049,6 +34046,17 @@ mod tests { let workspace = tempfile::tempdir().unwrap(); init_clean_git_workspace(workspace.path()); let api = test_api(workspace.path()).await; + let missing = scoped_get_subjektiv_subject( + State(api.clone()), + AxumPath(ScopedSubjektivSubjectPath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + subject_id: "missing-subject".to_string(), + }), + ) + .await + .unwrap_err() + .into_response(); + assert_eq!(missing.status(), StatusCode::NOT_FOUND); let subject = open_subjektiv_store(&api) .unwrap() .create_subject(crate::subjektiv::SubjectRole::new("companion").unwrap())