diff --git a/crates/session-store/src/lib.rs b/crates/session-store/src/lib.rs index 17035316..93859748 100644 --- a/crates/session-store/src/lib.rs +++ b/crates/session-store/src/lib.rs @@ -53,7 +53,9 @@ pub use segment::{ }; pub use segment_log::{LogEntry, RestoredState, SegmentOrigin, SessionExtension, collect_state}; pub use store::{Store, StoreError}; -pub use system_item::{SystemItem, SystemReminder, SystemReminderSource, render_worker_event}; +pub use system_item::{ + PromptRenderProvenance, SystemItem, SystemReminder, SystemReminderSource, render_worker_event, +}; pub use worker_metadata::{ CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata, WorkerMetadataStore, WorkerPeer, WorkerReclaimedChild, WorkerSpawnedChild, diff --git a/crates/session-store/src/system_item.rs b/crates/session-store/src/system_item.rs index e4fe4da5..ed662a62 100644 --- a/crates/session-store/src/system_item.rs +++ b/crates/session-store/src/system_item.rs @@ -82,6 +82,7 @@ impl SystemReminder { SystemReminderSource::TaskInactivity => SystemItem::TaskReminder { source: self.source, body: self.rendered_body(), + prompt_provenance: None, }, } } @@ -102,6 +103,16 @@ fn render_system_reminder(body: &str) -> String { format!("{SYSTEM_REMINDER_OPEN}\n{body}\n{SYSTEM_REMINDER_CLOSE}") } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct PromptRenderProvenance { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub workspace_id: Option, + pub config_revision: u64, + pub source_digest: String, + pub projection_digest: String, + pub logical_name: String, +} + /// One agent-injected system item, tagged by origin. /// /// Each variant carries the kind-specific raw data clients use for @@ -124,13 +135,23 @@ pub enum SystemItem { /// `Method::Notify`. `message` is the raw caller-supplied text; /// `body` is the wrapped LLM-context form (Worker renders it via /// `notify_wrapper` at commit time). - Notification { message: String, body: String }, + Notification { + message: String, + body: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + prompt_provenance: Option, + }, /// Lifecycle event reported by a child Worker via `Method::WorkerEvent`. /// `event` is the typed payload (so the TUI can render per-child /// banners without re-parsing); `body` is the wrapped LLM-context /// form (same `notify_wrapper` path as `Notification`). - WorkerEvent { event: WorkerEvent, body: String }, + WorkerEvent { + event: WorkerEvent, + body: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + prompt_provenance: Option, + }, /// `@` file reference resolution. `body` is the rendered /// LLM-context text (`[File: ]\n…` for regular files, @@ -162,12 +183,18 @@ pub enum SystemItem { #[serde(default = "default_task_reminder_source")] source: SystemReminderSource, body: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + prompt_provenance: Option, }, /// Synthetic note inserted after an interrupted turn before the next /// user input. `body` is the exact LLM-context text explaining that the /// previous turn was cut short. - Interrupt { body: String }, + Interrupt { + body: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + prompt_provenance: Option, + }, } impl SystemItem { @@ -184,7 +211,7 @@ impl SystemItem { format!("Ignored legacy procedure item: /{slug}") } SystemItem::TaskReminder { body, .. } => body.clone(), - SystemItem::Interrupt { body } => body.clone(), + SystemItem::Interrupt { body, .. } => body.clone(), } } @@ -237,11 +264,36 @@ pub fn render_worker_event(event: &WorkerEvent) -> String { mod tests { use super::*; + #[test] + fn legacy_prompt_rendered_items_default_missing_provenance() { + let notification: SystemItem = serde_json::from_str( + r#"{"kind":"notification","message":"legacy","body":"legacy body"}"#, + ) + .unwrap(); + let interrupt: SystemItem = + serde_json::from_str(r#"{"kind":"interrupt","body":"legacy interrupt"}"#).unwrap(); + assert!(matches!( + notification, + SystemItem::Notification { + prompt_provenance: None, + .. + } + )); + assert!(matches!( + interrupt, + SystemItem::Interrupt { + prompt_provenance: None, + .. + } + )); + } + #[test] fn notification_history_text_returns_stored_body() { let item = SystemItem::Notification { message: "child done".into(), body: "[Notification]\nchild done\n\n(non-blocking hint…)".into(), + prompt_provenance: None, }; assert_eq!( item.history_text(), @@ -256,6 +308,7 @@ mod tests { worker_name: "child".into(), }, body: "[Notification]\npod `child` finished a turn\n\n(non-blocking hint…)".into(), + prompt_provenance: None, }; assert!(item.history_text().starts_with("[Notification]\n")); assert!(item.history_text().contains("`child`")); @@ -292,7 +345,7 @@ mod tests { fn system_reminder_source_is_retained_in_system_item() { let item = SystemReminder::task_inactivity("remember tasks").into_system_item(); match item { - SystemItem::TaskReminder { source, body } => { + SystemItem::TaskReminder { source, body, .. } => { assert_eq!(source, SystemReminderSource::TaskInactivity); assert_eq!( body, @@ -352,6 +405,7 @@ mod tests { worker_name: "child".into(), }, body: "[Notification] worker `child` finished a turn".into(), + prompt_provenance: None, }; let json = serde_json::to_string(&item).unwrap(); let parsed: SystemItem = serde_json::from_str(&json).unwrap(); @@ -359,6 +413,7 @@ mod tests { SystemItem::WorkerEvent { event: WorkerEvent::TurnEnded { worker_name }, body, + .. } => { assert_eq!(worker_name, "child"); assert!(body.contains("`child`")); diff --git a/crates/tui/src/app.rs b/crates/tui/src/app.rs index a765ea71..5f510610 100644 --- a/crates/tui/src/app.rs +++ b/crates/tui/src/app.rs @@ -2165,7 +2165,7 @@ impl App { session_store::SystemItem::FileAttachment { body, .. } | session_store::SystemItem::SkillActivation { body, .. } | session_store::SystemItem::TaskReminder { body, .. } - | session_store::SystemItem::Interrupt { body } => { + | session_store::SystemItem::Interrupt { body, .. } => { self.task_store.apply_system_message_text(&body); self.blocks.push(Block::SystemMessage { text: body }); } diff --git a/crates/worker/src/feature/builtin/task/mod.rs b/crates/worker/src/feature/builtin/task/mod.rs index c0317983..4cce2231 100644 --- a/crates/worker/src/feature/builtin/task/mod.rs +++ b/crates/worker/src/feature/builtin/task/mod.rs @@ -334,7 +334,7 @@ mod tests { } let queued = pending.lock().expect("pending queue poisoned"); - let SystemItem::TaskReminder { source, body } = &queued[0] else { + let SystemItem::TaskReminder { source, body, .. } = &queued[0] else { panic!("unexpected system item: {:?}", queued[0]); }; assert_eq!(*source, SystemReminderSource::TaskInactivity); diff --git a/crates/worker/src/ipc/interceptor.rs b/crates/worker/src/ipc/interceptor.rs index 55266d31..8be84982 100644 --- a/crates/worker/src/ipc/interceptor.rs +++ b/crates/worker/src/ipc/interceptor.rs @@ -21,7 +21,6 @@ use llm_engine::interceptor::{ }; use llm_engine::tool::ToolOutput; use tracing::info; -use tracing::warn; use crate::compact::state::CompactState; use crate::compact::usage_tracker::UsageTracker; @@ -32,7 +31,7 @@ use crate::hook::{ HookRegistry, HookTurnEndAction, PreRequestContext, PreRequestInfo, PromptSubmitInfo, SystemItemAppendHandle, ToolCallSummary, ToolResultSummary, TurnEndInfo, }; -use crate::ipc::notify_buffer::{NotifyBuffer, build_system_item}; +use crate::ipc::notify_buffer::{NotifyBuffer, build_system_item_with_provenance}; use crate::prompt::catalog::PromptCatalog; use crate::worker::SystemItemCommitter; use llm_engine::token_counter::total_tokens; @@ -66,6 +65,8 @@ pub(crate) struct WorkerInterceptor { /// Prompt catalog used to render pending notification entries into the /// same system-message text that will be persisted in history. prompts: Arc>, + /// Workspace scope associated with Prompt projection provenance. + prompt_workspace_id: Option, /// Type-erased commit handle. The interceptor uses it to commit /// `LogEntry::SystemItem` entries directly (sync) before /// returning the corresponding `Item::system_message`s up to the @@ -96,6 +97,7 @@ impl WorkerInterceptor { pending_notifies, pending_attachments, prompts, + prompt_workspace_id: None, log_writer, next_turn_index: AtomicUsize::new(0), tool_calls_this_turn: AtomicUsize::new(0), @@ -107,6 +109,11 @@ impl WorkerInterceptor { self } + pub(crate) fn with_prompt_workspace_id(mut self, workspace_id: Option) -> Self { + self.prompt_workspace_id = workspace_id; + self + } + /// Commit each `SystemItem` as its own `LogEntry::SystemItem` /// entry through the attached writer (no-op when no writer is /// wired). Sync — writes complete before the matching @@ -163,6 +170,32 @@ impl WorkerInterceptor { } false } + fn attach_prompt_provenance(&self, items: &mut [SystemItem]) { + let prompts = self.prompts.load(); + let projection = prompts.projection(); + let provenance = |logical_name: &str| session_store::PromptRenderProvenance { + workspace_id: self.prompt_workspace_id.clone(), + config_revision: projection.config_revision, + source_digest: projection.source_digest.clone(), + projection_digest: projection.catalog_digest.clone(), + logical_name: logical_name.to_string(), + }; + for item in items { + match item { + SystemItem::TaskReminder { + prompt_provenance, .. + } if prompt_provenance.is_none() => { + *prompt_provenance = Some(provenance("internal.task_reminder")); + } + SystemItem::Interrupt { + prompt_provenance, .. + } if prompt_provenance.is_none() => { + *prompt_provenance = Some(provenance("internal.interrupt_system_note")); + } + _ => {} + } + } + } } #[async_trait] @@ -181,7 +214,7 @@ impl Interceptor for WorkerInterceptor { return action.into(); } } - let extras: Vec = std::mem::take( + let mut extras: Vec = std::mem::take( &mut *self .pending_attachments .lock() @@ -195,6 +228,7 @@ impl Interceptor for WorkerInterceptor { // commits land BEFORE the worker pushes its // `Item::system_message`s, so on-disk order matches // worker-history order. + self.attach_prompt_provenance(&mut extras); let items: Vec = extras.iter().map(SystemItem::to_history_item).collect(); match self.commit_system_items(&extras) { Ok(()) => PromptAction::ContinueWith(items), @@ -210,31 +244,22 @@ impl Interceptor for WorkerInterceptor { } let prompts = self.prompts.load_full(); + let projection = prompts.projection(); + let provenance = session_store::PromptRenderProvenance { + workspace_id: self.prompt_workspace_id.clone(), + config_revision: projection.config_revision, + source_digest: projection.source_digest.clone(), + projection_digest: projection.catalog_digest.clone(), + logical_name: "internal.notify_wrapper".to_string(), + }; let mut system_items: Vec = Vec::with_capacity(drained.len()); let mut items: Vec = Vec::with_capacity(drained.len()); for entry in drained { - match build_system_item(&entry, &prompts) { - Ok(system_item) => { - items.push(system_item.to_history_item()); - system_items.push(system_item); - } - Err(e) => { - // A render failure here would starve the LLM of - // the notify text. Fall back to a raw item so the - // trigger still lands in history; the entry will - // simply be skipped from the SystemItem batch. - warn!(error = %e, "failed to render notify_wrapper; using raw message"); - let fallback = match &entry { - super::notify_buffer::PendingNotify::Notify { message, .. } => { - message.clone() - } - super::notify_buffer::PendingNotify::WorkerEvent { event } => { - session_store::render_worker_event(event) - } - }; - items.push(Item::system_message(fallback)); - } - } + let system_item = + build_system_item_with_provenance(&entry, &prompts, Some(provenance.clone())) + .map_err(|error| format!("failed to render notify_wrapper: {error}"))?; + items.push(system_item.to_history_item()); + system_items.push(system_item); } self.commit_system_items(&system_items) .map_err(|error| format!("session persistence failed: {error}"))?; @@ -265,11 +290,12 @@ impl Interceptor for WorkerInterceptor { } } - let system_items: Vec = std::mem::take( + let mut system_items: Vec = std::mem::take( &mut *pending_hook_system_items .lock() .expect("pending hook system-item queue poisoned"), ); + self.attach_prompt_provenance(&mut system_items); let appended_items: Vec = system_items .iter() .map(SystemItem::to_history_item) @@ -1033,16 +1059,26 @@ mod tests { .lock() .expect("committed system-item list poisoned"); assert_eq!(committed.len(), 1); - let SystemItem::TaskReminder { body, .. } = &committed[0] else { - panic!("expected task reminder, got {:?}", committed[0]); + let SystemItem::TaskReminder { + body, + prompt_provenance: Some(provenance), + .. + } = &committed[0] + else { + panic!( + "expected task reminder with Prompt provenance, got {:?}", + committed[0] + ); }; assert!(body.contains("track active work")); + assert_eq!(provenance.logical_name, "internal.task_reminder"); } #[tokio::test] async fn pending_notifications_use_the_latest_prompt_projection() { let prompts = test_prompts(); let buffer = NotifyBuffer::new(); + let committed = Arc::new(Mutex::new(Vec::new())); let interceptor = WorkerInterceptor::new( Arc::new(HookRegistryBuilder::new().build()), None, @@ -1050,8 +1086,11 @@ mod tests { buffer.clone(), Arc::new(Mutex::new(Vec::new())), prompts.clone(), - None, - ); + Some(Arc::new(RecordingSystemItemCommitter { + committed: committed.clone(), + })), + ) + .with_prompt_workspace_id(Some("workspace-a".to_string())); let current = prompts.load_full(); let projection = current.projection(); @@ -1076,6 +1115,18 @@ mod tests { let appends = interceptor.pending_history_appends().await.unwrap(); assert_eq!(appends.len(), 1); assert!(format!("{:?}", appends[0]).contains("CURRENT-PROJECTION updated")); + let committed = committed.lock().unwrap(); + let SystemItem::Notification { + prompt_provenance: Some(provenance), + .. + } = &committed[0] + else { + panic!("notification Prompt provenance was not committed"); + }; + assert_eq!(provenance.workspace_id.as_deref(), Some("workspace-a")); + assert_eq!(provenance.config_revision, 2); + assert_eq!(provenance.source_digest, "source-2"); + assert_eq!(provenance.logical_name, "internal.notify_wrapper"); } #[tokio::test] diff --git a/crates/worker/src/ipc/notify_buffer.rs b/crates/worker/src/ipc/notify_buffer.rs index 64c0483c..b07df6cd 100644 --- a/crates/worker/src/ipc/notify_buffer.rs +++ b/crates/worker/src/ipc/notify_buffer.rs @@ -111,9 +111,18 @@ impl NotifyBuffer { /// Render one pending entry into a typed `SystemItem`. The /// `notify_wrapper` prompt produces the LLM-context body for both /// `Notify` (raw message) and `WorkerEvent` (rendered event line). +#[cfg(test)] pub(crate) fn build_system_item( entry: &PendingNotify, prompts: &PromptCatalog, +) -> Result { + build_system_item_with_provenance(entry, prompts, None) +} + +pub(crate) fn build_system_item_with_provenance( + entry: &PendingNotify, + prompts: &PromptCatalog, + prompt_provenance: Option, ) -> Result { match entry { PendingNotify::Notify { message, .. } => { @@ -121,6 +130,7 @@ pub(crate) fn build_system_item( Ok(SystemItem::Notification { message: message.clone(), body, + prompt_provenance, }) } PendingNotify::WorkerEvent { event } => { @@ -129,6 +139,7 @@ pub(crate) fn build_system_item( Ok(SystemItem::WorkerEvent { event: event.clone(), body, + prompt_provenance, }) } } @@ -178,7 +189,7 @@ mod tests { let catalog = PromptCatalog::builtins_only().unwrap(); let item = build_system_item(&entry, &catalog).unwrap(); match item { - SystemItem::Notification { message, body } => { + SystemItem::Notification { message, body, .. } => { assert_eq!(message, "hello"); assert!(body.contains("[Notification]")); assert!(body.contains("hello")); @@ -198,7 +209,7 @@ mod tests { let catalog = PromptCatalog::builtins_only().unwrap(); let item = build_system_item(&entry, &catalog).unwrap(); match item { - SystemItem::WorkerEvent { event, body } => { + SystemItem::WorkerEvent { event, body, .. } => { assert!( matches!(event, WorkerEvent::TurnEnded { ref worker_name } if worker_name == "child") ); diff --git a/crates/worker/src/segment_log_sink.rs b/crates/worker/src/segment_log_sink.rs index 0951d80c..45c6b2b7 100644 --- a/crates/worker/src/segment_log_sink.rs +++ b/crates/worker/src/segment_log_sink.rs @@ -281,6 +281,7 @@ mod tests { item: session_store::SystemItem::Notification { message: text.to_owned(), body: format!("[Notification] {text}"), + prompt_provenance: None, }, } } diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 374a496f..dd4ec9f2 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -13,8 +13,8 @@ use llm_engine::llm_client::types::Role; use llm_engine::state::Mutable; use llm_engine::{Engine, EngineError, EngineResult, ToolOutputLimits, UsageRecord}; use session_store::{ - LogEntry, SegmentId, SessionExtension, SessionId, Store, StoreError, SystemItem, segment_log, - to_logged, + LogEntry, PromptRenderProvenance, SegmentId, SessionExtension, SessionId, Store, StoreError, + SystemItem, segment_log, to_logged, }; use session_store::{ WorkerActiveSegmentRef, WorkerMetadata, WorkerMetadataStore, WorkerReclaimedChild, @@ -1285,6 +1285,21 @@ impl Worker { Arc::clone(&self.prompts) } + fn prompt_render_provenance(&self, logical_name: &str) -> PromptRenderProvenance { + let prompts = self.prompts.load(); + let projection = prompts.projection(); + PromptRenderProvenance { + workspace_id: self + .workspace_context + .workspace_id() + .map(|workspace_id| workspace_id.as_str().to_string()), + config_revision: projection.config_revision, + source_digest: projection.source_digest.clone(), + projection_digest: projection.catalog_digest.clone(), + logical_name: logical_name.to_string(), + } + } + fn refresh_prompt_projection_for_future_operations(&self) -> Result<(), WorkerError> { // The launch catalog remains authoritative until the initial system // Prompt has been rendered and committed. Later operation boundaries @@ -2037,7 +2052,12 @@ impl Worker { self.prompts.clone(), self.log_writer.clone(), ) - .with_usage_tracker(self.usage_tracker.clone()); + .with_usage_tracker(self.usage_tracker.clone()) + .with_prompt_workspace_id( + self.workspace_context + .workspace_id() + .map(|workspace_id| workspace_id.as_str().to_string()), + ); self.engine_mut().set_interceptor(interceptor); self.interceptor_installed = true; } @@ -2527,10 +2547,13 @@ impl Worker { if !closures.is_empty() { self.engine_mut().append_history(closures)?; } + let interrupt_prompt_provenance = + self.prompt_render_provenance("internal.interrupt_system_note"); self.commit_entry(LogEntry::SystemItem { ts: segment_log::now_millis(), item: SystemItem::Interrupt { body: system_note.clone(), + prompt_provenance: Some(interrupt_prompt_provenance), }, })?; self.engine_mut() @@ -6879,7 +6902,7 @@ mod build_summary_prompt_tests { matches!( entry, LogEntry::SystemItem { - item: SystemItem::Interrupt { body }, + item: SystemItem::Interrupt { body, .. }, .. } if body == &interrupt_note )