diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index 50bfc4697..89ceb9a89 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -135,9 +135,9 @@ pub use uploaded_file::{ UploadedFileUploadContext, }; pub use worker_metadata::{ - CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata, - WorkerMetadataStore, WorkerPeer, WorkerReclaimedChild, WorkerSpawnedChild, - WorkerSpawnedScopeRule, WorkerStoreError, validate_worker_name, + CombinedStore, FsWorkerStore, SubjektivSessionAttributionState, WorkerActiveSegmentRef, + WorkerAggregateStore, WorkerMetadata, WorkerMetadataStore, WorkerPeer, WorkerReclaimedChild, + WorkerSpawnedChild, WorkerSpawnedScopeRule, WorkerStoreError, validate_worker_name, }; pub use worker_session_store::WorkerSessionStore; diff --git a/crates/session-store/src/worker_metadata.rs b/crates/session-store/src/worker_metadata.rs index 93ab6b62d..4a4e50344 100644 --- a/crates/session-store/src/worker_metadata.rs +++ b/crates/session-store/src/worker_metadata.rs @@ -112,6 +112,37 @@ pub struct WorkerPeer { pub worker_name: String, } +/// Durable completion state for the Host-owned Subjektiv Session attribution +/// lifecycle. The Host writes `OutcomeUnknown` before sending the request, so a +/// process interruption or lost response remains explicitly recoverable by the +/// next Worker create/restore lifecycle instead of leaking into ordinary Runs. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "state", rename_all = "snake_case", deny_unknown_fields)] +pub enum SubjektivSessionAttributionState { + OutcomeUnknown { + session_id: SessionId, + diagnostic: String, + }, + Failed { + session_id: SessionId, + diagnostic: String, + }, + Confirmed { + session_id: SessionId, + subject_id: String, + }, +} + +impl SubjektivSessionAttributionState { + pub fn session_id(&self) -> SessionId { + match self { + Self::OutcomeUnknown { session_id, .. } + | Self::Failed { session_id, .. } + | Self::Confirmed { session_id, .. } => *session_id, + } + } +} + /// Persistent metadata for a Worker name. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct WorkerMetadata { @@ -133,6 +164,10 @@ pub struct WorkerMetadata { pub peers: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub resolved_manifest_snapshot: Option, + /// Host-owned completion evidence for Session-to-Subject attribution. Its + /// absence is also the legacy representation and never grants visibility. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subjektiv_session_attribution: Option, } impl WorkerMetadata { @@ -147,6 +182,7 @@ impl WorkerMetadata { reclaimed_children: Vec::new(), peers: Vec::new(), resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, } } @@ -161,6 +197,18 @@ impl WorkerMetadata { } } +fn replace_active_session(metadata: &mut WorkerMetadata, active: Option) { + if let Some(session_id) = active.as_ref().map(|active| active.session_id) + && metadata + .subjektiv_session_attribution + .as_ref() + .is_some_and(|state| state.session_id() != session_id) + { + metadata.subjektiv_session_attribution = None; + } + metadata.active = active; +} + /// Sync persistence backend for Worker metadata. pub trait WorkerMetadataStore: Send + Sync { /// Create or replace metadata for its `worker_name` key. @@ -236,7 +284,7 @@ pub trait WorkerMetadataStore: Send + Sync { workspace_root: Option, ) -> Result { self.update_by_name(worker_name, |metadata| { - metadata.active = active; + replace_active_session(metadata, active); metadata.resolved_manifest_snapshot = resolved_manifest_snapshot; if let Some(workspace_id) = workspace_id { metadata.workspace_id = Some(workspace_id); @@ -256,7 +304,7 @@ pub trait WorkerMetadataStore: Send + Sync { workspace_id: Option, ) -> Result { self.update_by_name(worker_name, |metadata| { - metadata.active = active; + replace_active_session(metadata, active); metadata.resolved_manifest_snapshot = resolved_manifest_snapshot; if let Some(workspace_id) = workspace_id { metadata.workspace_id = Some(workspace_id); @@ -273,7 +321,7 @@ pub trait WorkerMetadataStore: Send + Sync { workspace_root: Option, ) -> Result { self.update_by_name(worker_name, |metadata| { - metadata.active = active; + replace_active_session(metadata, active); metadata.resolved_manifest_snapshot = resolved_manifest_snapshot; if let Some(workspace_root) = workspace_root { metadata.workspace_root = Some(workspace_root); @@ -281,6 +329,18 @@ pub trait WorkerMetadataStore: Send + Sync { }) } + /// Set the Host-owned Subjektiv Session attribution lifecycle state while + /// preserving the active Session pointer and every unrelated field. + fn set_subjektiv_session_attribution( + &self, + worker_name: &str, + state: SubjektivSessionAttributionState, + ) -> Result { + self.update_by_name(worker_name, |metadata| { + metadata.subjektiv_session_attribution = Some(state); + }) + } + /// Set spawned-child registry state while preserving active pointer and manifest snapshot. fn set_spawned_children( &self, @@ -993,6 +1053,85 @@ mod tests { assert_eq!(restored, metadata); } + #[test] + fn subjektiv_attribution_state_roundtrips_and_active_updates_preserve_it() { + let tmp = tempfile::TempDir::new().unwrap(); + let store = FsWorkerStore::new(tmp.path()).unwrap(); + let session_id = crate::new_session_id(); + store + .write(&WorkerMetadata::new( + "subject-worker", + Some(WorkerActiveSegmentRef::pending_segment(session_id)), + )) + .unwrap(); + store + .set_subjektiv_session_attribution( + "subject-worker", + SubjektivSessionAttributionState::OutcomeUnknown { + session_id, + diagnostic: "request interrupted".to_string(), + }, + ) + .unwrap(); + store + .set_active( + "subject-worker", + Some(WorkerActiveSegmentRef::active_segment( + session_id, + crate::new_segment_id(), + )), + None, + ) + .unwrap(); + + let restored = store.read_by_name("subject-worker").unwrap().unwrap(); + assert!(matches!( + restored.subjektiv_session_attribution, + Some(SubjektivSessionAttributionState::OutcomeUnknown { + session_id: stored_session_id, + ref diagnostic, + }) if stored_session_id == session_id && diagnostic == "request interrupted" + )); + } + + #[test] + fn new_active_session_clears_prior_session_attribution_state() { + let tmp = tempfile::TempDir::new().unwrap(); + let store = FsWorkerStore::new(tmp.path()).unwrap(); + let old_session_id = crate::new_session_id(); + let new_session_id = crate::new_session_id(); + store + .write(&WorkerMetadata::new( + "subject-worker", + Some(WorkerActiveSegmentRef::pending_segment(old_session_id)), + )) + .unwrap(); + store + .set_subjektiv_session_attribution( + "subject-worker", + SubjektivSessionAttributionState::Confirmed { + session_id: old_session_id, + subject_id: "subject-1".to_string(), + }, + ) + .unwrap(); + + store + .set_active( + "subject-worker", + Some(WorkerActiveSegmentRef::pending_segment(new_session_id)), + None, + ) + .unwrap(); + + let restored = store.read_by_name("subject-worker").unwrap().unwrap(); + assert_eq!( + restored.active, + Some(WorkerActiveSegmentRef::pending_segment(new_session_id)) + ); + assert!(restored.subjektiv_session_attribution.is_none()); + } + #[test] fn spawned_scope_rule_defaults_resolved_and_roundtrips_logical_policy() { let legacy: WorkerSpawnedScopeRule = serde_json::from_value(serde_json::json!({ diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 690e52623..dde8cea97 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -1279,6 +1279,9 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { })?; let worker = prepared.worker_mut(); validate_worker_memory_settings(worker.manifest(), &request.request)?; + worker + .finalize_subjektiv_session_attribution(request.request.subjektiv_attached) + .map_err(|error| format!("finalize Worker Session attribution: {error}"))?; let workdir_capabilities = request .workdir_attachments .keys() @@ -1517,6 +1520,9 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { Err(err) => return Err(format!("failed to restore Worker from metadata: {err}")), }; validate_worker_memory_settings(worker.manifest(), &request.request)?; + worker + .finalize_subjektiv_session_attribution(request.request.subjektiv_attached) + .map_err(|error| format!("finalize restored Worker Session attribution: {error}"))?; let flow_transition_enabled = worker.manifest().feature.flow.enabled; let workdir_sessions = restored_workdir_router( &request.workdir_attachments, @@ -5079,6 +5085,7 @@ mod tests { &[ "WorkerBootstrap::new(", ".prepare()", + "finalize_subjektiv_session_attribution(", "worker.bind_workdir_sessions(", "worker.bind_worker_observation_provider(", "install_runtime_flow_transition_feature()", @@ -5089,6 +5096,7 @@ mod tests { restore, &[ "Worker::restore_from_worker_metadata_with_context(", + "finalize_subjektiv_session_attribution(", "worker.bind_workdir_sessions(", "worker.bind_worker_observation_provider(", "install_runtime_flow_transition_feature()", diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 1b83ff393..950a81a55 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1408,7 +1408,6 @@ where && let Some(subjektiv_sessions) = crate::feature::builtin::subjektiv_session::SubjektivSessionFeature::from_resolved_config( &worker.manifest().feature.subjektiv, - worker.committed_session_capture_handle(), worker.workspace_client_handle(), )? { diff --git a/crates/worker/src/discovery.rs b/crates/worker/src/discovery.rs index e27fbb20d..68da48ec7 100644 --- a/crates/worker/src/discovery.rs +++ b/crates/worker/src/discovery.rs @@ -1141,6 +1141,7 @@ mod tests { worker_name: "peer".into(), }], resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }; store.write(&parent).unwrap(); store @@ -1156,6 +1157,7 @@ mod tests { reclaimed_children: Vec::new(), peers: Vec::new(), resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); store @@ -1171,6 +1173,7 @@ mod tests { reclaimed_children: Vec::new(), peers: Vec::new(), resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); store @@ -1183,6 +1186,7 @@ mod tests { reclaimed_children: Vec::new(), peers: Vec::new(), resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); store @@ -1198,6 +1202,7 @@ mod tests { reclaimed_children: Vec::new(), peers: Vec::new(), resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); store @@ -1212,6 +1217,7 @@ mod tests { worker_name: "parent".into(), }], resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); @@ -1420,6 +1426,7 @@ mod tests { worker_name: "target".into(), }], resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); inner.write(&WorkerMetadata::new("target", None)).unwrap(); @@ -1461,6 +1468,7 @@ mod tests { worker_name: "target".into(), }], resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); store @@ -1475,6 +1483,7 @@ mod tests { worker_name: "source".into(), }], resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); let runtime_dir = Arc::new(RuntimeDir::create(&runtime_base, "source").await.unwrap()); @@ -1592,6 +1601,7 @@ mod tests { worker_name: "target".into(), }], resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); store @@ -1606,6 +1616,7 @@ mod tests { worker_name: "source".into(), }], resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); let runtime_dir = Arc::new(RuntimeDir::create(&runtime_base, "source").await.unwrap()); @@ -1719,6 +1730,7 @@ mod tests { reclaimed_children: Vec::new(), peers: Vec::new(), resolved_manifest_snapshot: None, + subjektiv_session_attribution: None, }) .unwrap(); store.write(&WorkerMetadata::new("target", None)).unwrap(); diff --git a/crates/worker/src/feature/builtin/memory_lifecycle.rs b/crates/worker/src/feature/builtin/memory_lifecycle.rs index 3720adb28..5635bfe35 100644 --- a/crates/worker/src/feature/builtin/memory_lifecycle.rs +++ b/crates/worker/src/feature/builtin/memory_lifecycle.rs @@ -644,23 +644,6 @@ impl MemoryLifecycleTask { return Ok(()); } - if self.target == ExtractionTarget::Subjektiv - && let Err(reason) = self.record_subjektiv_session(&capture.session_id) - { - audit - .emit( - self.workspace_client.as_ref(), - self.event_tx.as_ref(), - memory::audit::WorkerLifecycleStatus::Failed, - reason, - usage_audit, - extract_audit, - None, - ) - .await; - return Ok(()); - } - context.generation_fence.ensure_current()?; let next_pointer = memory::ExtractPointerPayload { processed_through_entry: capture.entry_count - 1, @@ -745,35 +728,6 @@ impl FeatureBackgroundTask for MemoryLifecycleTask { } impl MemoryLifecycleTask { - fn record_subjektiv_session(&self, session_id: &str) -> Result<(), String> { - let response = self - .workspace_client - .execute_server_operation( - crate::worker::WorkspaceServerOperation::SubjektivRecordSession( - server_api::SubjektivRecordSessionRequest { - session_id: session_id.to_string(), - }, - ), - ) - .map_err(|error| format!("record subjektiv Session attribution failed: {error}"))?; - if !response.is_success() { - return Err(format!( - "record subjektiv Session attribution returned HTTP {}: {}", - response.status, response.body - )); - } - let output: server_api::SubjektivRecordSessionResponse = - serde_json::from_str(&response.body).map_err(|error| { - format!("decode subjektiv Session attribution response: {error}") - })?; - if output.session_id != session_id { - return Err( - "subjektiv Session attribution response changed session identity".to_string(), - ); - } - Ok(()) - } - async fn request_consolidation(&self, subjektiv: bool) { let audit = WorkerAuditBase::new( memory::audit::AuditWorker::MemoryConsolidation, @@ -1068,7 +1022,7 @@ fn model_audit_from_manifest(model: &manifest::ModelManifest) -> memory::audit:: #[cfg(test)] mod tests { use std::pin::Pin; - use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent}; @@ -1106,7 +1060,6 @@ mod tests { request: crate::worker::WorkspaceRequest, ) -> Result { let is_subjektiv_stage = request.path.ends_with("/subjektiv/staging"); - let is_subjektiv_session = request.path.ends_with("/subjektiv/sessions"); let is_subjektiv_consolidation = request.path.ends_with("/subjektiv/consolidation"); let is_stage_candidate = request .body @@ -1129,16 +1082,6 @@ mod tests { .unwrap(), }); } - if is_subjektiv_session { - return 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(), - }); - } if is_subjektiv_stage { return Ok(crate::worker::WorkspaceResponse { status: 200, @@ -1181,52 +1124,6 @@ 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, @@ -1718,7 +1615,8 @@ permission = "write" .iter() .filter(|request| request.path.ends_with("/subjektiv/sessions")) .count(), - 1 + 0, + "Memory extraction must not own Session attribution", ); assert!(requests.iter().all(|request| { request @@ -1736,7 +1634,7 @@ permission = "write" } #[tokio::test] - async fn empty_subjektiv_extraction_records_session_and_commits_pointer_without_candidate() { + async fn empty_subjektiv_extraction_commits_pointer_without_attribution_or_candidate() { let client = ScriptClient::new(vec![finish_empty_events("finish-1"), completed_events()]); let extension_writes = Arc::new(Mutex::new(Vec::new())); let (event_tx, _) = broadcast::channel(32); @@ -1765,56 +1663,12 @@ permission = "write" .any(|request| request.path.ends_with("/subjektiv/staging")) ); assert!( - requests + !requests .iter() .any(|request| request.path.ends_with("/subjektiv/sessions")) ); } - #[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.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); @@ -2116,7 +1970,6 @@ permission = "write" segment_id: "segment-1".to_string(), session_revision: history_len.try_into().unwrap(), entry_count: history_len, - has_committed_run: true, run_exit: CommittedRunExit::Finished, history: (0..history_len) .map(|index| HistoryEntry { diff --git a/crates/worker/src/feature/builtin/subjektiv_memory.rs b/crates/worker/src/feature/builtin/subjektiv_memory.rs index 4a958fddc..73331e4d3 100644 --- a/crates/worker/src/feature/builtin/subjektiv_memory.rs +++ b/crates/worker/src/feature/builtin/subjektiv_memory.rs @@ -1101,7 +1101,6 @@ mod tests { segment_id: "segment-1".into(), session_revision: 1, entry_count: 0, - has_committed_run: true, run_exit: CommittedRunExit::Finished, history: Vec::new(), usage_history: Vec::new(), @@ -1352,7 +1351,6 @@ mod tests { segment_id: "segment-1".into(), session_revision: 1, entry_count: 3, - has_committed_run: true, run_exit: CommittedRunExit::Interrupted, history: vec![ crate::session_history::history_entry( @@ -1412,7 +1410,6 @@ mod tests { segment_id: "segment-1".into(), session_revision: 2, entry_count: 2, - has_committed_run: true, run_exit: CommittedRunExit::Finished, history: vec![ crate::session_history::history_entry( @@ -1490,7 +1487,6 @@ mod tests { segment_id: "segment-1".into(), session_revision: 2, entry_count: 2, - has_committed_run: true, run_exit: CommittedRunExit::Finished, history: vec![ crate::session_history::history_entry( diff --git a/crates/worker/src/feature/builtin/subjektiv_session.rs b/crates/worker/src/feature/builtin/subjektiv_session.rs index e0fb4b8f1..b7ee46256 100644 --- a/crates/worker/src/feature/builtin/subjektiv_session.rs +++ b/crates/worker/src/feature/builtin/subjektiv_session.rs @@ -10,29 +10,21 @@ use agen::tool::{Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta use async_trait::async_trait; use schemars::JsonSchema; use server_api::{ - SUBJEKTIV_SESSION_MAX_TOOL_CONTENT_BYTES, SubjektivRecordSessionRequest, - SubjektivRecordSessionResponse, SubjektivSessionBackendOperation, + SUBJEKTIV_SESSION_MAX_TOOL_CONTENT_BYTES, SubjektivSessionBackendOperation, SubjektivSessionBackendRequest, SubjektivSessionBackendResponse, SubjektivSessionBackendResult, SubjektivSessionDiagnosticCode, SubjektivSessionErrorResponse, SubjektivSessionListRequest, SubjektivSessionReadRequest, SubjektivSessionSearchRequest, }; -use crate::feature::session::CommittedSessionCaptureHandle; use crate::feature::{ - FeatureDescriptor, FeatureHookPoint, FeatureInstallContext, FeatureInstallError, FeatureModule, - HookDeclaration, ToolContribution, ToolDeclaration, -}; -use crate::hook::{ - Hook, HookError, HookErrorCategory, HookExecutionPolicy, HookFailurePolicy, - HookPreRequestAction, PreLlmRequest, PreRequestContext, RunCommitted, RunCommittedContext, + FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureModule, ToolContribution, + ToolDeclaration, }; use crate::worker::{WorkspaceClient, WorkspaceServerOperation}; const LIST_TOOL: &str = "SubjektivSessionList"; const SEARCH_TOOL: &str = "SubjektivSessionSearch"; const READ_TOOL: &str = "SubjektivSessionRead"; -const COMMIT_HOOK: &str = "record-subjektiv-session-after-commit"; -const RETRY_HOOK: &str = "retry-subjektiv-session-attribution"; const LIST_DESCRIPTION: &str = "List committed retained or archived Sessions attributed by the Host to the connected subject. Filter by exact session_id or storage and continue with the returned opaque cursor; subject, Runtime, Worker, archive, and filesystem authority are never model inputs."; const SEARCH_DESCRIPTION: &str = "Search committed public entries across Host-authorized Sessions attributed to the connected subject, or list one Session's public entries when session_id is supplied and query is omitted. Search is case-insensitive literal matching; system prompts, hidden reasoning, traces, and attachment bodies are never exposed."; @@ -46,13 +38,11 @@ pub(crate) struct SubjektivSessionFeature { #[derive(Clone)] struct SubjektivSessionState { client: Arc, - capture: CommittedSessionCaptureHandle, } impl SubjektivSessionFeature { pub(crate) fn from_resolved_config( config: &manifest::ResolvedSubjektivFeatureConfig, - capture: CommittedSessionCaptureHandle, client: Arc, ) -> std::io::Result> { if !config.execution_enabled() { @@ -68,7 +58,7 @@ impl SubjektivSessionFeature { )); } Ok(Some(Self { - state: SubjektivSessionState { client, capture }, + state: SubjektivSessionState { client }, })) } } @@ -82,14 +72,6 @@ impl FeatureModule for SubjektivSessionFeature { .with_tool(ToolDeclaration::new(LIST_TOOL, LIST_DESCRIPTION)) .with_tool(ToolDeclaration::new(SEARCH_TOOL, SEARCH_DESCRIPTION)) .with_tool(ToolDeclaration::new(READ_TOOL, READ_DESCRIPTION)) - .with_hook(HookDeclaration::new( - COMMIT_HOOK, - FeatureHookPoint::RunCommitted, - )) - .with_hook(HookDeclaration::new( - RETRY_HOOK, - FeatureHookPoint::PreLlmRequest, - )) } fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { @@ -103,19 +85,7 @@ impl FeatureModule for SubjektivSessionFeature { tool_definition(name, description, self.state.clone(), operation), ))?; } - context.hooks().add_run_committed( - COMMIT_HOOK, - HookExecutionPolicy::new(HookFailurePolicy::AttentionRequired, 30_000), - RecordCommittedSessionHook { - state: self.state.clone(), - }, - )?; - context.hooks().add_pre_request( - RETRY_HOOK, - RetryCommittedSessionHook { - state: self.state.clone(), - }, - ) + Ok(()) } } @@ -204,88 +174,6 @@ impl SubjektivSessionState { response.status ))) } - - fn record_current_committed_session(&self) -> Result<(), HookError> { - let capture = self - .capture - .capture() - .map_err(|error| HookError::new(HookErrorCategory::Dependency, error.to_string()))?; - if !capture.has_committed_run { - return Err(HookError::new( - HookErrorCategory::Dependency, - "committed Session capture is not available yet", - )); - } - self.record_session(&capture.session_id) - } - - fn retry_current_committed_session(&self) -> Result<(), HookError> { - let capture = self - .capture - .capture() - .map_err(|error| HookError::new(HookErrorCategory::Dependency, error.to_string()))?; - if !capture.has_committed_run { - return Ok(()); - } - self.record_session(&capture.session_id) - } - - fn record_session(&self, session_id: &str) -> Result<(), HookError> { - let response = self - .client - .execute_server_operation(WorkspaceServerOperation::SubjektivRecordSession( - SubjektivRecordSessionRequest { - session_id: session_id.to_string(), - }, - )) - .map_err(|error| HookError::new(HookErrorCategory::Dependency, error.to_string()))?; - if !response.is_success() { - return Err(HookError::new( - HookErrorCategory::Dependency, - format!( - "record subjektiv Session attribution returned HTTP {}", - response.status - ), - )); - } - let output: SubjektivRecordSessionResponse = - serde_json::from_str(&response.body).map_err(|error| { - HookError::new( - HookErrorCategory::Dependency, - format!("decode subjektiv Session attribution response: {error}"), - ) - })?; - if output.session_id != session_id { - return Err(HookError::new( - HookErrorCategory::Dependency, - "subjektiv Session attribution response changed session identity", - )); - } - Ok(()) - } -} - -struct RecordCommittedSessionHook { - state: SubjektivSessionState, -} - -#[async_trait] -impl Hook for RecordCommittedSessionHook { - async fn call(&self, _input: &RunCommittedContext) -> Result<(), HookError> { - self.state.record_current_committed_session() - } -} - -struct RetryCommittedSessionHook { - state: SubjektivSessionState, -} - -#[async_trait] -impl Hook for RetryCommittedSessionHook { - async fn call(&self, _input: &PreRequestContext) -> Result { - self.state.retry_current_committed_session()?; - Ok(HookPreRequestAction::Continue) - } } fn parse(input: &str, tool: &str) -> Result { @@ -348,12 +236,8 @@ fn response_summary(response: &SubjektivSessionBackendResponse) -> String { #[cfg(test)] mod tests { use super::*; - use std::collections::HashSet; use std::sync::Mutex; - use std::sync::atomic::{AtomicUsize, Ordering}; - use crate::feature::session::{CommittedRunExit, CommittedSessionCapture}; - use crate::hook::{HookHistoryRange, HookInvocationContext, PreRequestInfo, RunCommittedExit}; use crate::worker::{ WorkspaceClientError, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, }; @@ -398,81 +282,6 @@ mod tests { } } - #[derive(Debug, Default)] - struct AttributionClient { - requests: Mutex>, - recorded: Mutex>, - insertions: AtomicUsize, - failures: AtomicUsize, - } - - impl WorkspaceClient for AttributionClient { - fn workspace_id(&self) -> Option<&str> { - Some("workspace-test") - } - - fn kind(&self) -> &str { - "subjektiv-attribution-test" - } - - fn is_available(&self) -> bool { - true - } - - fn execute( - &self, - request: WorkspaceRequest, - ) -> Result { - self.requests.lock().unwrap().push(request.clone()); - if self - .failures - .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| { - remaining.checked_sub(1) - }) - .is_ok() - { - return Ok(WorkspaceResponse { - status: 503, - body: "temporarily unavailable".into(), - }); - } - let input: SubjektivRecordSessionRequest = - serde_json::from_str(request.body.as_deref().unwrap()).unwrap(); - if self - .recorded - .lock() - .unwrap() - .insert(input.session_id.clone()) - { - self.insertions.fetch_add(1, Ordering::SeqCst); - } - Ok(WorkspaceResponse { - status: 200, - body: serde_json::to_string(&SubjektivRecordSessionResponse { - subject_id: "subject-1".into(), - session_id: input.session_id, - }) - .unwrap(), - }) - } - } - - fn capture(has_committed_run: bool) -> CommittedSessionCaptureHandle { - CommittedSessionCaptureHandle::new(move || { - Ok(CommittedSessionCapture { - session_id: "session-1".into(), - segment_id: "segment-1".into(), - session_revision: usize::from(has_committed_run) as u64, - entry_count: usize::from(has_committed_run), - has_committed_run, - run_exit: CommittedRunExit::Finished, - history: Vec::new(), - usage_history: Vec::new(), - extensions: Vec::new(), - }) - }) - } - fn list_response() -> SubjektivSessionBackendResponse { SubjektivSessionBackendResponse::Ok { result: SubjektivSessionBackendResult::List(server_api::SubjektivSessionListResponse { @@ -483,47 +292,16 @@ mod tests { } } - fn state(client: Arc, has_committed_run: bool) -> SubjektivSessionState { - SubjektivSessionState { - client, - capture: capture(has_committed_run), - } + fn state(client: Arc) -> SubjektivSessionState { + SubjektivSessionState { client } } fn feature(client: Arc) -> SubjektivSessionFeature { SubjektivSessionFeature { - state: state(client, true), + state: state(client), } } - fn run_committed_context() -> RunCommittedContext { - RunCommittedContext { - invocation: HookInvocationContext { - workspace_id: Some("workspace-test".into()), - worker_id: "worker-1".into(), - session_id: "session-1".into(), - session_revision: 1, - run_id: Some("run-1".into()), - turn_index: Some(0), - call_id: None, - }, - exit: RunCommittedExit::Finished, - committed_history: HookHistoryRange::default(), - } - } - - fn pre_request_context() -> PreRequestContext { - PreRequestContext::new( - PreRequestInfo { - item_count: 0, - estimated_tokens: None, - turn_index: 0, - tool_calls_this_turn: 0, - }, - None, - ) - } - #[test] fn descriptor_and_definitions_register_exact_session_tools_and_typed_schemas() { let feature = feature(Arc::new(SessionToolClient::new(list_response()))); @@ -540,17 +318,7 @@ mod tests { (READ_TOOL, READ_DESCRIPTION), ] ); - assert_eq!( - descriptor - .hooks - .iter() - .map(|hook| (hook.name.as_str(), hook.point.clone())) - .collect::>(), - vec![ - (COMMIT_HOOK, FeatureHookPoint::RunCommitted), - (RETRY_HOOK, FeatureHookPoint::PreLlmRequest), - ] - ); + assert!(descriptor.hooks.is_empty()); for (name, description, operation, expected_schema) in [ ( @@ -595,7 +363,6 @@ mod tests { let config = manifest::ResolvedSubjektivFeatureConfig::default(); let result = SubjektivSessionFeature::from_resolved_config( &config, - capture(false), Arc::new(SessionToolClient::new(list_response())), ) .unwrap(); @@ -620,7 +387,7 @@ mod tests { ), ] { let error = SubjektivSessionTool { - state: state(client.clone(), true), + state: state(client.clone()), operation, } .execute(input, ToolExecutionContext::new("call", "batch", 0)) @@ -635,7 +402,7 @@ mod tests { async fn tool_routes_one_host_bound_operation_to_session_history() { let client = Arc::new(SessionToolClient::new(list_response())); let tool = SubjektivSessionTool { - state: state(client.clone(), true), + state: state(client.clone()), operation: SessionOperation::List, }; let output = tool @@ -682,7 +449,7 @@ mod tests { }, }; let error_tool = SubjektivSessionTool { - state: state(Arc::new(SessionToolClient::new(typed_error.clone())), true), + state: state(Arc::new(SessionToolClient::new(typed_error.clone()))), operation: SessionOperation::Search, }; let output = error_tool @@ -769,87 +536,4 @@ mod tests { } if content == escaped_content && cursor == "next-page" )); } - - #[tokio::test] - async fn candidate_free_commit_records_session_and_pre_request_waits_for_committed_capture() { - let client = Arc::new(AttributionClient::default()); - let commit_hook = RecordCommittedSessionHook { - state: state(client.clone(), true), - }; - commit_hook.call(&run_committed_context()).await.unwrap(); - let requests = client.requests.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert_eq!(requests[0].path, "/api/w/workspace-test/subjektiv/sessions"); - assert!(!requests[0].path.contains("staging")); - drop(requests); - - let before_commit = Arc::new(AttributionClient::default()); - let premature_commit_hook = RecordCommittedSessionHook { - state: state(before_commit.clone(), false), - }; - assert!( - premature_commit_hook - .call(&run_committed_context()) - .await - .is_err() - ); - let retry_hook = RetryCommittedSessionHook { - state: state(before_commit.clone(), false), - }; - assert_eq!( - retry_hook.call(&pre_request_context()).await.unwrap(), - HookPreRequestAction::Continue - ); - assert!(before_commit.requests.lock().unwrap().is_empty()); - } - - #[tokio::test] - async fn transient_attribution_failure_is_a_hook_failure_then_retries_exactly() { - let client = Arc::new(AttributionClient::default()); - client.failures.store(1, Ordering::SeqCst); - let state = state(client.clone(), true); - let commit_hook = RecordCommittedSessionHook { - state: state.clone(), - }; - let retry_hook = RetryCommittedSessionHook { state }; - - let error = commit_hook - .call(&run_committed_context()) - .await - .unwrap_err(); - assert_eq!(error.category, HookErrorCategory::Dependency); - assert_eq!( - retry_hook.call(&pre_request_context()).await.unwrap(), - HookPreRequestAction::Continue - ); - - let requests = client.requests.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert_eq!(requests[0].body, requests[1].body); - assert_eq!(client.insertions.load(Ordering::SeqCst), 1); - } - - #[tokio::test] - async fn repeated_commit_and_pre_request_use_exact_request_and_backend_dedupes() { - let client = Arc::new(AttributionClient::default()); - let state = state(client.clone(), true); - let commit_hook = RecordCommittedSessionHook { - state: state.clone(), - }; - let retry_hook = RetryCommittedSessionHook { state }; - - commit_hook.call(&run_committed_context()).await.unwrap(); - commit_hook.call(&run_committed_context()).await.unwrap(); - retry_hook.call(&pre_request_context()).await.unwrap(); - - let requests = client.requests.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!( - requests - .iter() - .all(|request| request.body == requests[0].body) - ); - assert_eq!(client.recorded.lock().unwrap().len(), 1); - assert_eq!(client.insertions.load(Ordering::SeqCst), 1); - } } diff --git a/crates/worker/src/feature/session.rs b/crates/worker/src/feature/session.rs index b31a27674..0729f0af5 100644 --- a/crates/worker/src/feature/session.rs +++ b/crates/worker/src/feature/session.rs @@ -25,9 +25,6 @@ pub(crate) struct CommittedSessionCapture { /// Monotonic committed-log revision for the captured Segment. pub(crate) session_revision: u64, pub(crate) entry_count: usize, - /// Whether the durable capture contains a committed run boundary. This is - /// false for a newly materialized Session before its first Run commits. - pub(crate) has_committed_run: bool, pub(crate) run_exit: CommittedRunExit, pub(crate) history: Vec>, pub(crate) usage_history: Vec, diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 45046af93..3f380d7e2 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -17,7 +17,7 @@ use agen::{ use arc_swap::ArcSwap; use session_store::{ LogEntry, PasteArtifactLimits, PromptRenderProvenance, SegmentId, SessionExtension, SessionId, - Store, StoreError, SystemItem, segment_log, + Store, StoreError, SubjektivSessionAttributionState, SystemItem, segment_log, }; use session_store::{ WorkerActiveSegmentRef, WorkerMetadata, WorkerMetadataStore, WorkerReclaimedChild, @@ -3491,12 +3491,6 @@ impl Worker { let entries = store .read_all(location.session_id, location.segment_id) .map_err(|error| FeatureSessionError::Capture(error.to_string()))?; - let has_committed_run = entries.iter().any(|entry| { - matches!( - entry, - LogEntry::RunCompleted { .. } | LogEntry::RunYielded { .. } - ) - }); let run_exit = entries .iter() .rev() @@ -3523,7 +3517,6 @@ impl Worker { segment_id: location.segment_id.to_string(), session_revision: entries.len().try_into().unwrap_or(u64::MAX), entry_count: entries.len(), - has_committed_run, run_exit, history, usage_history: restored.usage_history, @@ -6520,6 +6513,15 @@ impl Worker { } } +fn bounded_attribution_diagnostic(mut diagnostic: String) -> String { + const MAX_CHARS: usize = 1024; + if diagnostic.chars().count() > MAX_CHARS { + diagnostic = diagnostic.chars().take(MAX_CHARS).collect(); + diagnostic.push('…'); + } + diagnostic +} + fn worker_language(cfg: &manifest::EngineManifest) -> &str { let language = cfg.language.trim(); if language.is_empty() { @@ -6583,6 +6585,27 @@ where workspace_context: WorkerWorkspaceContext, filesystem_authority: WorkerFilesystemAuthority, model_client: Option>, + ) -> Result { + Self::from_manifest_with_context_and_model_client_for_session( + manifest, + store, + loader, + workspace_context, + filesystem_authority, + model_client, + None, + ) + .await + } + + async fn from_manifest_with_context_and_model_client_for_session( + manifest: WorkerManifest, + store: St, + loader: PromptCatalogSource, + workspace_context: WorkerWorkspaceContext, + filesystem_authority: WorkerFilesystemAuthority, + model_client: Option>, + existing_session_id: Option, ) -> Result { let common = prepare_worker_common_with_context_and_model_client( &manifest, @@ -6598,8 +6621,9 @@ where // `ensure_segment_head`) so the SegmentStart entry can capture // the rendered system prompt, not the raw template source. The // session_id + segment_id are allocated here so the worker-allocation - // registration can record them from the start. - let session_id = session_store::new_session_id(); + // registration can record them from the start. Pending lifecycle + // recovery reuses its already persisted Session identity. + let session_id = existing_session_id.unwrap_or_else(session_store::new_session_id); let segment_id = session_store::new_segment_id(); // Register this Worker in the machine-wide worker-allocation @@ -6688,6 +6712,168 @@ where Ok(worker) } + /// Complete the Host-owned Session-to-Subject attribution boundary before a + /// runtime-managed Worker is exposed or accepts input. The authoritative + /// attachment bit comes from the Runtime create/restore request; names, + /// Profiles, and current singleton ownership are never used as identity. + /// + /// `OutcomeUnknown` is persisted before network I/O. A lost response or + /// process interruption is therefore recovered only by a later create/restore + /// lifecycle, where the Backend's same-attribution idempotency is safe to use. + pub fn finalize_subjektiv_session_attribution( + &self, + subjektiv_attached: bool, + ) -> Result<(), WorkerError> { + let worker_name = self.manifest.worker.name.as_str(); + let metadata = self.store.read_by_name(worker_name)?.ok_or_else(|| { + WorkerError::SubjektivSessionAttribution { + state: "metadata_missing", + message: "Worker metadata disappeared before Session attribution".to_string(), + } + })?; + let active = + metadata + .active + .as_ref() + .ok_or_else(|| WorkerError::SubjektivSessionAttribution { + state: "session_missing", + message: "Worker metadata has no active Session attribution target".to_string(), + })?; + let session_id = self.session_id(); + if active.session_id != session_id { + return Err(WorkerError::SubjektivSessionAttribution { + state: "session_mismatch", + message: format!( + "Worker metadata Session {} does not match constructed Session {session_id}", + active.session_id + ), + }); + } + + if !subjektiv_attached { + if metadata.subjektiv_session_attribution.is_some() { + return Err(WorkerError::SubjektivSessionAttribution { + state: "attachment_missing", + message: + "persisted Subjektiv Session attribution state requires an authorized attachment" + .to_string(), + }); + } + return Ok(()); + } + + if let Some(state) = metadata.subjektiv_session_attribution.as_ref() { + if state.session_id() != session_id { + return Err(WorkerError::SubjektivSessionAttribution { + state: "session_mismatch", + message: format!( + "persisted Subjektiv attribution targets Session {}, not {session_id}", + state.session_id() + ), + }); + } + if matches!(state, SubjektivSessionAttributionState::Confirmed { .. }) { + return Ok(()); + } + } + + self.persist_subjektiv_session_attribution( + SubjektivSessionAttributionState::OutcomeUnknown { + session_id, + diagnostic: "attribution request started; completion is not yet confirmed" + .to_string(), + }, + )?; + let response = match self.workspace_client().execute_server_operation( + WorkspaceServerOperation::SubjektivRecordSession( + server_api::SubjektivRecordSessionRequest { + session_id: session_id.to_string(), + }, + ), + ) { + Ok(response) => response, + Err(error) => { + let message = bounded_attribution_diagnostic(format!( + "Session attribution transport outcome is unknown: {error}" + )); + self.persist_subjektiv_session_attribution( + SubjektivSessionAttributionState::OutcomeUnknown { + session_id, + diagnostic: message.clone(), + }, + )?; + return Err(WorkerError::SubjektivSessionAttribution { + state: "outcome_unknown", + message, + }); + } + }; + if !response.is_success() { + let message = bounded_attribution_diagnostic(format!( + "Session attribution was rejected with HTTP {}: {}", + response.status, response.body + )); + self.persist_subjektiv_session_attribution(SubjektivSessionAttributionState::Failed { + session_id, + diagnostic: message.clone(), + })?; + return Err(WorkerError::SubjektivSessionAttribution { + state: "failed", + message, + }); + } + let output: server_api::SubjektivRecordSessionResponse = + match serde_json::from_str(&response.body) { + Ok(output) => output, + Err(error) => { + let message = bounded_attribution_diagnostic(format!( + "Session attribution response outcome is unknown: {error}" + )); + self.persist_subjektiv_session_attribution( + SubjektivSessionAttributionState::OutcomeUnknown { + session_id, + diagnostic: message.clone(), + }, + )?; + return Err(WorkerError::SubjektivSessionAttribution { + state: "outcome_unknown", + message, + }); + } + }; + if output.session_id != session_id.to_string() + || output.subject_id.trim().is_empty() + || output.subject_id.chars().any(char::is_control) + { + let message = + "Session attribution response did not confirm the requested immutable identity" + .to_string(); + self.persist_subjektiv_session_attribution( + SubjektivSessionAttributionState::OutcomeUnknown { + session_id, + diagnostic: message.clone(), + }, + )?; + return Err(WorkerError::SubjektivSessionAttribution { + state: "outcome_unknown", + message, + }); + } + self.persist_subjektiv_session_attribution(SubjektivSessionAttributionState::Confirmed { + session_id, + subject_id: output.subject_id, + }) + } + + fn persist_subjektiv_session_attribution( + &self, + state: SubjektivSessionAttributionState, + ) -> Result<(), WorkerError> { + self.store + .set_subjektiv_session_attribution(&self.manifest.worker.name, state)?; + Ok(()) + } + /// Build an in-process Internal Worker without machine-wide allocation or durable Worker metadata. pub(crate) async fn from_internal_manifest_with_context( manifest: WorkerManifest, @@ -7054,13 +7240,14 @@ where })?; let manifest = restore_manifest_from_worker_metadata_snapshot(worker_name, Some(snapshot), fallback)?; - Self::from_manifest_with_context_and_model_client( + Self::from_manifest_with_context_and_model_client_for_session( manifest, store, loader, workspace_context, filesystem_authority, model_client, + Some(active.session_id), ) .await } @@ -8085,6 +8272,12 @@ pub enum WorkerError { #[error("feature lifecycle rejected operation: {0}")] FeatureLifecycle(String), + #[error("Subjektiv Session attribution lifecycle is {state}: {message}")] + SubjektivSessionAttribution { + state: &'static str, + message: String, + }, + #[error("Flow input rejected: {0}")] FlowInput(String), @@ -9487,6 +9680,264 @@ mod build_summary_prompt_tests { } } + #[derive(Clone, Copy, Debug)] + enum AttributionOutcome { + Success, + Rejected, + TransportUnknown, + } + + #[derive(Debug)] + struct AttributionLifecycleClient { + outcomes: Mutex>, + requests: Mutex>, + } + + impl AttributionLifecycleClient { + fn new(outcomes: impl IntoIterator) -> Self { + Self { + outcomes: Mutex::new(outcomes.into_iter().collect()), + requests: Mutex::new(Vec::new()), + } + } + } + + impl WorkspaceClient for AttributionLifecycleClient { + fn workspace_id(&self) -> Option<&str> { + Some("workspace-test") + } + + fn kind(&self) -> &str { + "attribution-lifecycle-test" + } + + fn is_available(&self) -> bool { + true + } + + fn execute( + &self, + request: WorkspaceRequest, + ) -> Result { + self.requests.lock().unwrap().push(request.clone()); + match self + .outcomes + .lock() + .unwrap() + .pop_front() + .unwrap_or(AttributionOutcome::Success) + { + AttributionOutcome::TransportUnknown => Err(WorkspaceClientError::Request( + "connection closed after request send".to_string(), + )), + AttributionOutcome::Rejected => Ok(WorkspaceResponse { + status: 503, + body: "temporarily unavailable".to_string(), + }), + AttributionOutcome::Success => { + let input: server_api::SubjektivRecordSessionRequest = + serde_json::from_str(request.body.as_deref().unwrap()).unwrap(); + Ok(WorkspaceResponse { + status: 200, + body: serde_json::to_string(&server_api::SubjektivRecordSessionResponse { + subject_id: "subject-test".to_string(), + session_id: input.session_id, + }) + .unwrap(), + }) + } + } + } + } + + async fn attribution_lifecycle_worker( + client: Arc, + ) -> ( + Worker< + Box, + session_store::CombinedStore, + >, + session_store::CombinedStore, + tempfile::TempDir, + ) { + let temp = tempfile::tempdir().unwrap(); + let store = session_store::CombinedStore::new( + session_store::FsStore::new(temp.path().join("sessions")).unwrap(), + session_store::FsWorkerStore::new(temp.path().join("workers")).unwrap(), + ); + let mut manifest = minimal_manifest(); + manifest.worker.name = format!("attribution-test-{}", uuid::Uuid::now_v7()); + let mut worker = Worker::new( + manifest, + Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated( + Box::new(NoopClient) as Box + ), + store.clone(), + WorkerWorkspaceContext::with_client( + Some(WorkspaceId::new("workspace-test").unwrap()), + client, + ), + WorkerFilesystemAuthority::None, + Scope::empty(), + ) + .await + .unwrap(); + worker.enable_worker_metadata_write_through().unwrap(); + (worker, store, temp) + } + + #[tokio::test] + async fn session_attribution_finalizes_once_before_runs_and_persists_subject() { + let client = Arc::new(AttributionLifecycleClient::new([ + AttributionOutcome::Success, + ])); + let (worker, store, _temp) = attribution_lifecycle_worker(client.clone()).await; + let session_id = worker.session_id(); + + worker.finalize_subjektiv_session_attribution(true).unwrap(); + worker.finalize_subjektiv_session_attribution(true).unwrap(); + + assert_eq!(client.requests.lock().unwrap().len(), 1); + let metadata = store + .read_by_name(&worker.manifest().worker.name) + .unwrap() + .unwrap(); + assert!(matches!( + metadata.subjektiv_session_attribution, + Some(SubjektivSessionAttributionState::Confirmed { + session_id: stored_session_id, + ref subject_id, + }) if stored_session_id == session_id && subject_id == "subject-test" + )); + } + + #[tokio::test] + async fn unknown_and_failed_attribution_remain_incomplete_until_lifecycle_retry() { + for (first, expected_state) in [ + (AttributionOutcome::TransportUnknown, "outcome_unknown"), + (AttributionOutcome::Rejected, "failed"), + ] { + let client = Arc::new(AttributionLifecycleClient::new([ + first, + AttributionOutcome::Success, + ])); + let (worker, store, _temp) = attribution_lifecycle_worker(client.clone()).await; + let error = worker + .finalize_subjektiv_session_attribution(true) + .unwrap_err(); + assert!(error.to_string().contains(expected_state)); + let incomplete = store + .read_by_name(&worker.manifest().worker.name) + .unwrap() + .unwrap() + .subjektiv_session_attribution + .unwrap(); + assert!(matches!( + (expected_state, incomplete), + ( + "outcome_unknown", + SubjektivSessionAttributionState::OutcomeUnknown { .. } + ) | ("failed", SubjektivSessionAttributionState::Failed { .. }) + )); + + worker.finalize_subjektiv_session_attribution(true).unwrap(); + assert_eq!(client.requests.lock().unwrap().len(), 2); + assert!(matches!( + store + .read_by_name(&worker.manifest().worker.name) + .unwrap() + .unwrap() + .subjektiv_session_attribution, + Some(SubjektivSessionAttributionState::Confirmed { .. }) + )); + } + } + + #[tokio::test] + async fn unattached_worker_never_attempts_session_attribution() { + let client = Arc::new(AttributionLifecycleClient::new([ + AttributionOutcome::Success, + ])); + let (worker, store, _temp) = attribution_lifecycle_worker(client.clone()).await; + + worker + .finalize_subjektiv_session_attribution(false) + .unwrap(); + + assert!(client.requests.lock().unwrap().is_empty()); + assert!( + store + .read_by_name(&worker.manifest().worker.name) + .unwrap() + .unwrap() + .subjektiv_session_attribution + .is_none() + ); + } + + #[tokio::test] + async fn pending_restore_reuses_persisted_session_and_recovers_unknown_attribution() { + let temp = tempfile::tempdir().unwrap(); + let store = session_store::CombinedStore::new( + session_store::FsStore::new(temp.path().join("sessions")).unwrap(), + session_store::FsWorkerStore::new(temp.path().join("workers")).unwrap(), + ); + let mut manifest = minimal_manifest(); + manifest.worker.name = format!("pending-attribution-{}", uuid::Uuid::now_v7()); + let session_id = session_store::new_session_id(); + let snapshot = manifest::write_persisted_worker_manifest_snapshot(&manifest).unwrap(); + store + .set_active_with_workspace_id( + &manifest.worker.name, + Some(WorkerActiveSegmentRef::pending_segment(session_id)), + Some(snapshot), + Some("workspace-test".to_string()), + ) + .unwrap(); + store + .set_subjektiv_session_attribution( + &manifest.worker.name, + SubjektivSessionAttributionState::OutcomeUnknown { + session_id, + diagnostic: "response lost before restart".to_string(), + }, + ) + .unwrap(); + let client = Arc::new(AttributionLifecycleClient::new([ + AttributionOutcome::Success, + ])); + let worker_name = manifest.worker.name.clone(); + let worker = Worker::restore_pending_from_worker_metadata_with_context_and_model_client( + &worker_name, + manifest, + store.clone(), + PromptCatalogSource::builtins_only(), + WorkerWorkspaceContext::with_client( + Some(WorkspaceId::new("workspace-test").unwrap()), + client.clone(), + ), + WorkerFilesystemAuthority::None, + Some(Box::new(NoopClient)), + ) + .await + .unwrap(); + + assert_eq!(worker.session_id(), session_id); + worker.finalize_subjektiv_session_attribution(true).unwrap(); + assert_eq!(client.requests.lock().unwrap().len(), 1); + assert!(matches!( + store + .read_by_name(&worker.manifest().worker.name) + .unwrap() + .unwrap() + .subjektiv_session_attribution, + Some(SubjektivSessionAttributionState::Confirmed { + session_id: stored_session_id, + .. + }) if stored_session_id == session_id + )); + } + fn text_segment(text: &str) -> Segment { Segment::Text { content: text.into(), diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index e0bd7bc16..83b8dad31 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -17283,6 +17283,34 @@ fn subjektiv_session_attribution( .map_err(|error| Error::InvalidInput(error.to_string()).into()) } +fn require_existing_subjektiv_session_attribution( + api: &WorkspaceApi, + workspace_id: &str, + context: &server_api::ServerRequestContext, + session_id: &str, +) -> ApiResult { + let expected = subjektiv_session_attribution(api, workspace_id, context, session_id)?; + let existing = open_subjektiv_store(api)? + .session_attribution(&expected.session_id) + .map_err(subjektiv_store_error)? + .ok_or_else(|| { + Error::RepositoryConflict(format!( + "subjektiv_session_attribution_incomplete: Session {} has not completed its Host attribution lifecycle", + expected.session_id + )) + })?; + if existing.subject_id != expected.subject_id + || existing.runtime_id != expected.runtime_id + || existing.worker_id != expected.worker_id + { + return Err(Error::WorkspacePermissionDenied( + "existing Session attribution does not match the authorized subject Worker".to_string(), + ) + .into()); + } + Ok(existing) +} + fn subjektiv_subject_response( api: &WorkspaceApi, subject: crate::subjektiv::SubjectRecord, @@ -17550,8 +17578,12 @@ async fn scoped_stage_subjektiv_candidate( Json(mut request): Json, ) -> ApiResult> { validate_workspace_scope(&api, &path.workspace_id)?; - let attribution = - subjektiv_session_attribution(&api, &path.workspace_id, &context, &request.session_id)?; + let attribution = require_existing_subjektiv_session_attribution( + &api, + &path.workspace_id, + &context, + &request.session_id, + )?; let subject_id = attribution.subject_id.clone(); let session_id = attribution.session_id.as_str(); for source_ref in &mut request.operation.source_refs { @@ -19066,7 +19098,7 @@ async fn scoped_subjektiv_memory_backend( // Re-derive the same subject and immutable Session attribution for the // write. The request never carries subject, Worker, Runtime, origin, // or arbitrary source JSON from the model. - let attribution = subjektiv_session_attribution( + let attribution = require_existing_subjektiv_session_attribution( &api, &path.workspace_id, &context, @@ -39595,6 +39627,20 @@ mod tests { } ) )); + let Json(recorded_session) = scoped_record_subjektiv_session( + State(api.clone()), + AxumPath(ScopedWorkspacePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + }), + context.clone(), + Json(server_api::SubjektivRecordSessionRequest { + session_id: "session-1".to_string(), + }), + ) + .await + .unwrap(); + assert_eq!(recorded_session.subject_id, subject.id); + let operation = memory::backend::MemoryStageCandidateOperation { source: memory::schema::SourceRef { segment_id: "segment-1".to_string(), diff --git a/crates/workspace-server/src/subjektiv.rs b/crates/workspace-server/src/subjektiv.rs index 37fad4588..1e9564954 100644 --- a/crates/workspace-server/src/subjektiv.rs +++ b/crates/workspace-server/src/subjektiv.rs @@ -1221,9 +1221,9 @@ impl SubjektivStore { }) } - /// Atomically records immutable host attribution for the committed Session - /// and stages its extracted candidate. Exact retries return the first stored - /// records, including their original host timestamps. + /// Atomically verifies immutable Host attribution already completed at the + /// Session lifecycle boundary and stages its extracted candidate. Candidate + /// staging never creates or repairs Session attribution. pub fn stage_candidate_with_attribution( &self, mut record: SubjectStagingRecord, @@ -1243,14 +1243,38 @@ impl SubjektivStore { validate_staging_record(self.workspace_id(), &record)?; validate_staging_admission(&record)?; let candidate_raw = serde_json::to_string(&record)?; - let attribution_raw = serde_json::to_string(&attribution)?; self.database.try_transaction(|transaction| { require_active_subject(transaction, &record.subject_id)?; - let attribution = write_session_attribution(transaction, attribution, attribution_raw)?; + let existing_raw = transaction + .query_row( + "SELECT record_json FROM subject_session_attributions WHERE session_id = ?1", + [&attribution.session_id], + |row| row.get::<_, String>(0), + ) + .optional()? + .ok_or_else(|| { + SubjektivError::InvalidRecord(format!( + "Session {} has not completed Host attribution", + attribution.session_id + )) + })?; + let existing = parse_session_attribution(&existing_raw)?; + if existing.subject_id != attribution.subject_id { + return Err(SubjektivError::SessionSubjectConflict { + session_id: attribution.session_id, + existing_subject_id: existing.subject_id, + requested_subject_id: attribution.subject_id, + }); + } + if !equivalent_session_attribution(&existing, &attribution) { + return Err(SubjektivError::SessionAttributionConflict( + attribution.session_id, + )); + } let candidate = write_staging_candidate(transaction, self.workspace_id(), record, candidate_raw)?; - Ok((candidate, attribution)) + Ok((candidate, existing)) }) } @@ -3657,6 +3681,7 @@ mod tests { "2026-09-28T09:00:00.000Z", ); + store.record_session_attribution(session.clone()).unwrap(); let (first_candidate, stored_session) = store .stage_candidate_with_attribution( candidate(&subject.id, "candidate-1", "workspace-a"), @@ -3871,16 +3896,20 @@ mod tests { }; let (_manager, _workspace, store) = open_store(&root, "workspace-a"); + let lifecycle_attribution = attribution( + &subject_id, + "runtime-1", + "worker-1", + "session-1", + "2026-09-28T10:00:00.000Z", + ); + store + .record_session_attribution(lifecycle_attribution.clone()) + .unwrap(); let (staged, stored_attribution) = store .stage_candidate_with_attribution( candidate(&subject_id, "candidate-after-migration", "workspace-a"), - attribution( - &subject_id, - "runtime-1", - "worker-1", - "session-1", - "2026-09-28T10:00:00.000Z", - ), + lifecycle_attribution, ) .unwrap(); @@ -3977,7 +4006,7 @@ mod tests { } #[test] - fn attributed_staging_is_atomic_and_retry_idempotent() { + fn attributed_staging_requires_lifecycle_record_and_retry_is_idempotent() { let temp = tempfile::tempdir().unwrap(); let (_manager, _workspace, store) = open_store(&temp.path().join("storage"), "workspace-a"); let subject = store.create_subject(role()).unwrap(); @@ -3988,6 +4017,24 @@ mod tests { "session-1", "2026-09-28T10:00:00.000Z", ); + assert!(matches!( + store.stage_candidate_with_attribution( + candidate(&subject.id, "candidate-unattributed", "workspace-a"), + first_attribution.clone(), + ), + Err(SubjektivError::InvalidRecord(message)) + if message.contains("has not completed Host attribution") + )); + assert!(store.session_attribution("session-1").unwrap().is_none()); + assert!( + store + .staging_candidate(&subject.id, "candidate-unattributed") + .unwrap() + .is_none() + ); + store + .record_session_attribution(first_attribution.clone()) + .unwrap(); let (first_candidate, first_attribution) = store .stage_candidate_with_attribution( candidate(&subject.id, "candidate-1", "workspace-a"), @@ -4013,22 +4060,24 @@ mod tests { let mut conflicting_candidate = candidate(&subject.id, "candidate-1", "workspace-a"); conflicting_candidate.claim = "different content".into(); conflicting_candidate.source_refs[0].session_id = Some("session-2".into()); - let result = store.stage_candidate_with_attribution( - conflicting_candidate, - attribution( - &subject.id, - "runtime-1", - "worker-1", - "session-2", - "2026-09-28T12:00:00.000Z", - ), + let conflicting_attribution = attribution( + &subject.id, + "runtime-1", + "worker-1", + "session-2", + "2026-09-28T12:00:00.000Z", ); + store + .record_session_attribution(conflicting_attribution.clone()) + .unwrap(); + let result = + store.stage_candidate_with_attribution(conflicting_candidate, conflicting_attribution); assert!(matches!( result, Err(SubjektivError::CandidateConflict(candidate_id)) if candidate_id == "candidate-1" )); - assert!(store.session_attribution("session-2").unwrap().is_none()); + assert!(store.session_attribution("session-2").unwrap().is_some()); } #[test] @@ -4119,6 +4168,26 @@ mod tests { let (_manager, _workspace, store) = open_store(&temp.path().join("storage"), "workspace-a"); let subject = store.create_subject(role()).unwrap(); let other_subject = store.create_subject(role()).unwrap(); + for lifecycle_attribution in [ + attribution( + &subject.id, + "runtime-1", + "worker-1", + "session-1", + "2026-09-28T10:00:00.000Z", + ), + attribution( + &subject.id, + "runtime-2", + "worker-2", + "session-2", + "2026-09-28T11:00:00.000Z", + ), + ] { + store + .record_session_attribution(lifecycle_attribution) + .unwrap(); + } store .stage_candidate_with_attribution( @@ -4716,17 +4785,18 @@ mod tests { "This attribution must roll back with the stale proposal", )) .unwrap(); + let stale_attribution = attribution( + &subject.id, + "runtime-proposal", + "worker-proposal", + "session-proposal-stale", + "2026-09-28T12:00:00.000Z", + ); + store + .record_session_attribution(stale_attribution.clone()) + .unwrap(); assert!(matches!( - store.stage_candidate_with_attribution( - attributed_stale, - attribution( - &subject.id, - "runtime-proposal", - "worker-proposal", - "session-proposal-stale", - "2026-09-28T12:00:00.000Z", - ), - ), + store.stage_candidate_with_attribution(attributed_stale, stale_attribution,), Err(SubjektivError::RevisionConflict { expected: 1, actual: 2, @@ -4737,7 +4807,7 @@ mod tests { store .session_attribution("session-proposal-stale") .unwrap() - .is_none() + .is_some() ); let reopen = candidate(&subject.id, "candidate-reopen", "workspace-a") .with_revision_proposal(proposal(