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

fix: finalize subjektiv attribution at session lifecycle

Browse Source
This commit is contained in:
Hare
2026-10-05 12:25:20 +09:00
parent c6d2ad1f4c
commit d2a0b75045
12 changed files with 798 additions and 543 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/session-store/src/lib.rs
+3 -3
View File
@@ -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;
crates/session-store/src/worker_metadata.rs
+142 -3
View File
@@ -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<WorkerPeer>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub resolved_manifest_snapshot: Option<serde_json::Value>,
/// 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<SubjektivSessionAttributionState>,
}
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<WorkerActiveSegmentRef>) {
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<PathBuf>,
) -> Result<WorkerMetadata, WorkerStoreError> {
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<String>,
) -> Result<WorkerMetadata, WorkerStoreError> {
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<PathBuf>,
) -> Result<WorkerMetadata, WorkerStoreError> {
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<WorkerMetadata, WorkerStoreError> {
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!({
crates/worker-runtime/src/worker_backend.rs
+8
View File
@@ -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()",
crates/worker/src/controller.rs
-1
View File
@@ -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(),
)?
{
crates/worker/src/discovery.rs
+12
View File
@@ -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();
crates/worker/src/feature/builtin/memory_lifecycle.rs
+5 -152
View File
@@ -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<crate::worker::WorkspaceResponse, crate::worker::WorkspaceClientError> {
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<Vec<crate::worker::WorkspaceRequest>>,
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<crate::worker::WorkspaceResponse, crate::worker::WorkspaceClientError> {
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<AtomicUsize>,
@@ -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::<Vec<_>>();
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 {
crates/worker/src/feature/builtin/subjektiv_memory.rs
-4
View File
@@ -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(
crates/worker/src/feature/builtin/subjektiv_session.rs
+12 -328
View File
@@ -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<dyn WorkspaceClient>,
capture: CommittedSessionCaptureHandle,
}
impl SubjektivSessionFeature {
pub(crate) fn from_resolved_config(
config: &manifest::ResolvedSubjektivFeatureConfig,
capture: CommittedSessionCaptureHandle,
client: Arc<dyn WorkspaceClient>,
) -> std::io::Result<Option<Self>> {
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<RunCommitted> for RecordCommittedSessionHook {
async fn call(&self, _input: &RunCommittedContext) -> Result<(), HookError> {
self.state.record_current_committed_session()
}
}
struct RetryCommittedSessionHook {
state: SubjektivSessionState,
}
#[async_trait]
impl Hook<PreLlmRequest> for RetryCommittedSessionHook {
async fn call(&self, _input: &PreRequestContext) -> Result<HookPreRequestAction, HookError> {
self.state.retry_current_committed_session()?;
Ok(HookPreRequestAction::Continue)
}
}
fn parse<T: serde::de::DeserializeOwned>(input: &str, tool: &str) -> Result<T, ToolError> {
@@ -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<Vec<WorkspaceRequest>>,
recorded: Mutex<HashSet<String>>,
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<WorkspaceResponse, WorkspaceClientError> {
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<dyn WorkspaceClient>, has_committed_run: bool) -> SubjektivSessionState {
SubjektivSessionState {
client,
capture: capture(has_committed_run),
}
fn state(client: Arc<dyn WorkspaceClient>) -> SubjektivSessionState {
SubjektivSessionState { client }
}
fn feature(client: Arc<dyn WorkspaceClient>) -> 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<_>>(),
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);
}
}
crates/worker/src/feature/session.rs
-3
View File
@@ -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<HistoryEntry<SessionHistoryMetadata>>,
pub(crate) usage_history: Vec<UsageRecord>,
crates/worker/src/worker.rs
+462 -11
View File
@@ -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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
}
}
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<Box<dyn LlmClient>>,
) -> Result<Self, WorkerError> {
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<Box<dyn LlmClient>>,
existing_session_id: Option<SessionId>,
) -> Result<Self, WorkerError> {
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<VecDeque<AttributionOutcome>>,
requests: Mutex<Vec<WorkspaceRequest>>,
}
impl AttributionLifecycleClient {
fn new(outcomes: impl IntoIterator<Item = AttributionOutcome>) -> 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<WorkspaceResponse, WorkspaceClientError> {
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<AttributionLifecycleClient>,
) -> (
Worker<
Box<dyn LlmClient>,
session_store::CombinedStore<session_store::FsStore, session_store::FsWorkerStore>,
>,
session_store::CombinedStore<session_store::FsStore, session_store::FsWorkerStore>,
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<dyn LlmClient>
),
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(),
crates/workspace-server/src/server.rs
+49 -3
View File
@@ -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<crate::subjektiv::SubjectSessionAttribution> {
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<server_api::SubjektivStageCandidateRequest>,
) -> ApiResult<Json<server_api::SubjektivStageCandidateResponse>> {
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(),
crates/workspace-server/src/subjektiv.rs
+105 -35
View File
@@ -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(