worker: split session capture observation features
This commit is contained in:
@@ -1600,9 +1600,12 @@ fn tool_kind(name: &str) -> &'static str {
|
||||
"WebFetch" | "WebSearch" => "web",
|
||||
"SubWorkerSpawn"
|
||||
| "SubWorkerSend"
|
||||
| "SubWorkerReadOutput"
|
||||
| "SubWorkerList"
|
||||
| "SubWorkerStop"
|
||||
| "ListWorkerSessions"
|
||||
| "ViewSessionOverview"
|
||||
| "SearchSessionEntries"
|
||||
| "ReadSessionEntry"
|
||||
| "WorkerList"
|
||||
| "WorkerSpawn"
|
||||
| "WorkerStop"
|
||||
|
||||
@@ -34,8 +34,8 @@ use crate::compact::usage_tracker::UsageTracker;
|
||||
use crate::fs_view::ReadRequirement;
|
||||
#[cfg(test)]
|
||||
use crate::fs_view::slice_lines;
|
||||
use crate::session_reference::{
|
||||
ReadDetail, ReadOptions, ReadSelector, SearchOptions, SessionReferenceView, ToolPart,
|
||||
use crate::session_capture::{
|
||||
ReadDetail, ReadOptions, ReadSelector, SearchOptions, SessionCapture, ToolPart,
|
||||
};
|
||||
|
||||
/// Aggregated output of a compact worker run.
|
||||
@@ -150,7 +150,7 @@ this to verify details before writing the summary.";
|
||||
|
||||
struct SessionLogToolState {
|
||||
items: Arc<Vec<Item>>,
|
||||
view: SessionReferenceView,
|
||||
view: SessionCapture,
|
||||
}
|
||||
|
||||
struct SearchSessionLogTool {
|
||||
@@ -185,6 +185,9 @@ impl Tool for SearchSessionLogTool {
|
||||
tool_name: None,
|
||||
limit: Some(limit),
|
||||
min_entry_index: Some(offset as u64),
|
||||
from: None,
|
||||
through: None,
|
||||
offset: 0,
|
||||
});
|
||||
let blocks = hits
|
||||
.iter()
|
||||
@@ -252,7 +255,7 @@ impl Tool for ReadSessionItemsTool {
|
||||
SessionReadMode::Full => ReadDetail::Full,
|
||||
};
|
||||
let read = if offset >= end {
|
||||
crate::session_reference::ReadResult {
|
||||
crate::session_capture::ReadResult {
|
||||
entries: Vec::new(),
|
||||
truncated: false,
|
||||
}
|
||||
@@ -496,7 +499,7 @@ pub(crate) fn write_summary_tool(ctx: Arc<Mutex<CompactWorkerContext>>) -> ToolD
|
||||
}
|
||||
|
||||
pub(crate) fn search_session_log_tool(items: Arc<Vec<Item>>) -> ToolDefinition {
|
||||
let view = SessionReferenceView::new("compact-target", (*items).clone());
|
||||
let view = SessionCapture::new("compact-target", (*items).clone());
|
||||
let state = Arc::new(SessionLogToolState { items, view });
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(SearchSessionParams);
|
||||
@@ -512,7 +515,7 @@ pub(crate) fn search_session_log_tool(items: Arc<Vec<Item>>) -> ToolDefinition {
|
||||
}
|
||||
|
||||
pub(crate) fn read_session_items_tool(items: Arc<Vec<Item>>) -> ToolDefinition {
|
||||
let view = SessionReferenceView::new("compact-target", (*items).clone());
|
||||
let view = SessionCapture::new("compact-target", (*items).clone());
|
||||
let state = Arc::new(SessionLogToolState { items, view });
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(ReadSessionParams);
|
||||
@@ -808,7 +811,7 @@ mod tests {
|
||||
"very large raw trace body with secret detail",
|
||||
),
|
||||
]);
|
||||
let view = SessionReferenceView::new("test", (*items).clone());
|
||||
let view = SessionCapture::new("test", (*items).clone());
|
||||
let tool: Arc<dyn Tool> = Arc::new(SearchSessionLogTool {
|
||||
state: Arc::new(SessionLogToolState { items, view }),
|
||||
});
|
||||
@@ -828,7 +831,7 @@ mod tests {
|
||||
"read trace",
|
||||
"raw trace detail",
|
||||
)]);
|
||||
let view = SessionReferenceView::new("test", (*items).clone());
|
||||
let view = SessionCapture::new("test", (*items).clone());
|
||||
let tool: Arc<dyn Tool> = Arc::new(ReadSessionItemsTool {
|
||||
state: Arc::new(SessionLogToolState { items, view }),
|
||||
});
|
||||
|
||||
@@ -21,9 +21,7 @@ use crate::shutdown_after_idle::{
|
||||
ShutdownAfterIdleRequest, TicketIntakeReadyShutdownHook, is_ticket_intake_role,
|
||||
take_shutdown_request_after_status,
|
||||
};
|
||||
use crate::spawn::comm_tools::{
|
||||
sub_worker_list_tool, sub_worker_read_output_tool, sub_worker_send_tool, sub_worker_stop_tool,
|
||||
};
|
||||
use crate::spawn::comm_tools::{sub_worker_list_tool, sub_worker_send_tool, sub_worker_stop_tool};
|
||||
use crate::spawn::registry::SpawnedWorkerRegistry;
|
||||
use crate::spawn::tool::sub_worker_spawn_tool;
|
||||
use crate::worker::{SystemItemCommitter, Worker, WorkerError, WorkerRunResult};
|
||||
@@ -59,6 +57,10 @@ impl WorkerHandle {
|
||||
self.event_tx.subscribe()
|
||||
}
|
||||
|
||||
pub fn committed_entries(&self) -> Vec<LogEntry> {
|
||||
self.sink.subscribe_with_snapshot().0
|
||||
}
|
||||
|
||||
pub fn snapshot_event(&self) -> Event {
|
||||
self.snapshot_event_with_entry_subscription().0
|
||||
}
|
||||
@@ -725,6 +727,7 @@ where
|
||||
worker.register_worker_orchestration_instruction();
|
||||
}
|
||||
|
||||
let host_worker_observation_provider = worker.worker_observation_provider();
|
||||
{
|
||||
let workspace_client = worker.workspace_client_handle();
|
||||
let engine = worker.engine_mut();
|
||||
@@ -779,7 +782,11 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
// Worker-orchestration tools (SubWorkerSpawn + the four comm tools) share
|
||||
let mut observation_providers: Vec<
|
||||
Arc<dyn crate::feature::builtin::worker_observation::WorkerObservationProvider>,
|
||||
> = Vec::new();
|
||||
|
||||
// Worker-orchestration tools (SubWorkerSpawn + three control tools) share
|
||||
// the Worker-scoped `SpawnedWorkerRegistry` (also consumed by the main
|
||||
// loop's `WorkerEvent` handler). Expose them only behind the explicit
|
||||
// profile feature and require delegation authority up front so enabling
|
||||
@@ -814,8 +821,26 @@ where
|
||||
));
|
||||
engine.register_tool(sub_worker_list_tool(spawned_registry.clone()));
|
||||
engine.register_tool(sub_worker_send_tool(spawned_registry.clone()));
|
||||
engine.register_tool(sub_worker_read_output_tool(spawned_registry.clone()));
|
||||
engine.register_tool(sub_worker_stop_tool(spawned_registry));
|
||||
engine.register_tool(sub_worker_stop_tool(spawned_registry.clone()));
|
||||
observation_providers.push(Arc::new(
|
||||
crate::feature::builtin::worker_observation::SpawnedSubWorkerObservationProvider::new(
|
||||
spawned_registry,
|
||||
),
|
||||
));
|
||||
}
|
||||
if let Some(provider) = host_worker_observation_provider {
|
||||
observation_providers.push(provider);
|
||||
}
|
||||
if !observation_providers.is_empty() {
|
||||
feature_registry = feature_registry.with_module(
|
||||
crate::feature::builtin::worker_observation::WorkerObservationFeature::new(
|
||||
Arc::new(
|
||||
crate::feature::builtin::worker_observation::CompositeWorkerObservationProvider::new(
|
||||
observation_providers,
|
||||
),
|
||||
),
|
||||
),
|
||||
);
|
||||
}
|
||||
}
|
||||
let _feature_install_report = worker.install_features(feature_registry);
|
||||
|
||||
@@ -7,16 +7,22 @@
|
||||
pub mod manage_workdir;
|
||||
pub mod manage_worker;
|
||||
pub mod memory;
|
||||
pub mod memory_extract;
|
||||
pub mod objective;
|
||||
pub mod session_explore;
|
||||
pub mod task;
|
||||
pub mod ticket;
|
||||
pub mod worker_observation;
|
||||
|
||||
pub(crate) use session_explore::{
|
||||
SessionExploreFeature, SessionExploreState, render_extract_input,
|
||||
};
|
||||
pub(crate) use memory_extract::{MemoryExtractFeature, MemoryExtractState, render_extract_input};
|
||||
pub(crate) use session_explore::{SessionExploreFeature, SessionExploreState};
|
||||
pub use task::{TaskFeature, task_tools_feature};
|
||||
pub use ticket::{
|
||||
TicketFeature, TicketFeatureAccess, ticket_tools_feature, ticket_tools_feature_with_access,
|
||||
ticket_tools_feature_with_backend,
|
||||
};
|
||||
pub use worker_observation::{
|
||||
CompositeWorkerObservationProvider, WorkerObservationError, WorkerObservationFeature,
|
||||
WorkerObservationProvider, WorkerObservationSubject, WorkerObservationSubjectRef,
|
||||
WorkerSessionCapture, WorkspaceClientWorkerObservationProvider,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,458 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use memory::backend::{
|
||||
MemoryBackendOperation, MemoryBackendOperationResult, MemoryStageCandidateOperation,
|
||||
};
|
||||
use memory::extract::{CandidateKind, ExtractedCandidate, StagingEvidence};
|
||||
use memory::schema::{EvidenceKind, SourceEvidenceRef, SourceRef};
|
||||
use schemars::JsonSchema;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::feature::{
|
||||
FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureModule, ToolContribution,
|
||||
ToolDeclaration,
|
||||
};
|
||||
use crate::session_capture::{
|
||||
ReferenceKind, SearchOptions, SessionCapture, SessionEntryEvidence, ToolPart,
|
||||
};
|
||||
use crate::worker::WorkspaceClient;
|
||||
|
||||
use super::memory::WorkspaceMemoryBackendError;
|
||||
|
||||
const STAGE_DESCRIPTION: &str = "Stage one durable Memory candidate using SessionEntryRef values from the co-installed session-explore capture.";
|
||||
const FINISH_DESCRIPTION: &str =
|
||||
"Finish Memory extraction after validating the number of candidates staged during this run.";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct MemoryExtractState {
|
||||
view: Arc<SessionCapture>,
|
||||
workspace_client: Arc<dyn WorkspaceClient>,
|
||||
source: SourceRef,
|
||||
extract_run_id: String,
|
||||
staged: Arc<Mutex<Vec<String>>>,
|
||||
finished: Arc<Mutex<Option<FinishMemoryExtractionParams>>>,
|
||||
}
|
||||
|
||||
impl MemoryExtractState {
|
||||
pub(crate) fn new(
|
||||
view: SessionCapture,
|
||||
workspace_client: Arc<dyn WorkspaceClient>,
|
||||
source: SourceRef,
|
||||
extract_run_id: String,
|
||||
) -> Self {
|
||||
Self {
|
||||
view: Arc::new(view),
|
||||
workspace_client,
|
||||
source,
|
||||
extract_run_id,
|
||||
staged: Arc::new(Mutex::new(Vec::new())),
|
||||
finished: Arc::new(Mutex::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn staged(&self) -> Vec<String> {
|
||||
self.staged
|
||||
.lock()
|
||||
.expect("memory extract staged state poisoned")
|
||||
.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn is_finished(&self) -> bool {
|
||||
self.finished
|
||||
.lock()
|
||||
.expect("memory extract finished state poisoned")
|
||||
.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct MemoryExtractFeature {
|
||||
state: MemoryExtractState,
|
||||
}
|
||||
|
||||
impl MemoryExtractFeature {
|
||||
pub(crate) fn new(state: MemoryExtractState) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
impl FeatureModule for MemoryExtractFeature {
|
||||
fn descriptor(&self) -> FeatureDescriptor {
|
||||
FeatureDescriptor::builtin("memory-extract", "Memory Extract")
|
||||
.with_description(
|
||||
"Memory staging and extraction completion, independent from session exploration.",
|
||||
)
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"StageMemoryCandidate",
|
||||
STAGE_DESCRIPTION,
|
||||
))
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"FinishMemoryExtraction",
|
||||
FINISH_DESCRIPTION,
|
||||
))
|
||||
}
|
||||
|
||||
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
|
||||
context.tools().register(ToolContribution::new(
|
||||
"StageMemoryCandidate",
|
||||
stage_definition(self.state.clone()),
|
||||
))?;
|
||||
context.tools().register(ToolContribution::new(
|
||||
"FinishMemoryExtraction",
|
||||
finish_definition(self.state.clone()),
|
||||
))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn stage_definition(state: MemoryExtractState) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = serde_json::to_value(schemars::schema_for!(StageMemoryCandidateParams))
|
||||
.unwrap_or_else(|_| serde_json::json!({}));
|
||||
let meta = ToolMeta::new("StageMemoryCandidate")
|
||||
.description(STAGE_DESCRIPTION)
|
||||
.input_schema(schema);
|
||||
let tool: Arc<dyn Tool> = Arc::new(StageMemoryCandidateTool {
|
||||
state: state.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
fn finish_definition(state: MemoryExtractState) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = serde_json::to_value(schemars::schema_for!(FinishMemoryExtractionParams))
|
||||
.unwrap_or_else(|_| serde_json::json!({}));
|
||||
let meta = ToolMeta::new("FinishMemoryExtraction")
|
||||
.description(FINISH_DESCRIPTION)
|
||||
.input_schema(schema);
|
||||
let tool: Arc<dyn Tool> = Arc::new(FinishMemoryExtractionTool {
|
||||
state: state.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct StageMemoryCandidateParams {
|
||||
kind: CandidateKind,
|
||||
claim: String,
|
||||
why_useful: String,
|
||||
#[serde(default)]
|
||||
staleness: Option<String>,
|
||||
entry_refs: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct FinishMemoryExtractionParams {
|
||||
staged_count: usize,
|
||||
#[serde(default)]
|
||||
no_candidates_reason: Option<String>,
|
||||
}
|
||||
|
||||
struct StageMemoryCandidateTool {
|
||||
state: MemoryExtractState,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for StageMemoryCandidateTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_context: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: StageMemoryCandidateParams =
|
||||
serde_json::from_str(input_json).map_err(|error| {
|
||||
ToolError::InvalidArgument(format!("invalid StageMemoryCandidate input: {error}"))
|
||||
})?;
|
||||
if params.entry_refs.is_empty() {
|
||||
return Err(ToolError::InvalidArgument(
|
||||
"StageMemoryCandidate requires at least one entry_ref".to_string(),
|
||||
));
|
||||
}
|
||||
let mut evidence = Vec::with_capacity(params.entry_refs.len());
|
||||
let mut source_refs = Vec::with_capacity(params.entry_refs.len());
|
||||
for entry_ref in ¶ms.entry_refs {
|
||||
let projection = self.state.view.evidence_for(entry_ref).ok_or_else(|| {
|
||||
ToolError::InvalidArgument(format!(
|
||||
"unknown SessionEntryRef {entry_ref:?} for this extraction capture"
|
||||
))
|
||||
})?;
|
||||
evidence.push(staging_evidence(&projection));
|
||||
source_refs.push(source_evidence_ref(&projection));
|
||||
}
|
||||
let candidate = ExtractedCandidate {
|
||||
kind: params.kind,
|
||||
claim: params.claim,
|
||||
why_useful: params.why_useful,
|
||||
staleness: params.staleness,
|
||||
evidence_ids: params.entry_refs,
|
||||
};
|
||||
let result = self
|
||||
.state
|
||||
.workspace_client
|
||||
.execute_memory_backend_operation(MemoryBackendOperation::StageCandidate(
|
||||
MemoryStageCandidateOperation {
|
||||
source: self.state.source.clone(),
|
||||
extract_run_id: self.state.extract_run_id.clone(),
|
||||
candidate,
|
||||
evidence,
|
||||
source_refs,
|
||||
},
|
||||
))
|
||||
.await
|
||||
.map_err(map_memory_stage_error)?;
|
||||
let staging_ids = match result {
|
||||
MemoryBackendOperationResult::StagingWritten(output) if output.staging_count == 1 => {
|
||||
output.staging_ids
|
||||
}
|
||||
MemoryBackendOperationResult::StagingWritten(output) => {
|
||||
return Err(ToolError::ExecutionFailed(format!(
|
||||
"StageMemoryCandidate expected one staging record, backend wrote {}",
|
||||
output.staging_count
|
||||
)));
|
||||
}
|
||||
other => {
|
||||
return Err(ToolError::ExecutionFailed(format!(
|
||||
"unexpected Memory backend result for StageMemoryCandidate: {other:?}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
let staging_id = staging_ids.into_iter().next().ok_or_else(|| {
|
||||
ToolError::ExecutionFailed(
|
||||
"StageMemoryCandidate backend did not return a staging id".to_string(),
|
||||
)
|
||||
})?;
|
||||
self.state
|
||||
.staged
|
||||
.lock()
|
||||
.expect("memory extract staged state poisoned")
|
||||
.push(staging_id.clone());
|
||||
Ok(ToolOutput {
|
||||
summary: format!("Staged Memory candidate {staging_id}."),
|
||||
content: Some(format!("staging_id: {staging_id}")),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct FinishMemoryExtractionTool {
|
||||
state: MemoryExtractState,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for FinishMemoryExtractionTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_context: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: FinishMemoryExtractionParams =
|
||||
serde_json::from_str(input_json).map_err(|error| {
|
||||
ToolError::InvalidArgument(format!("invalid FinishMemoryExtraction input: {error}"))
|
||||
})?;
|
||||
let actual = self
|
||||
.state
|
||||
.staged
|
||||
.lock()
|
||||
.expect("memory extract staged state poisoned")
|
||||
.len();
|
||||
if params.staged_count != actual {
|
||||
return Err(ToolError::InvalidArgument(format!(
|
||||
"FinishMemoryExtraction staged_count {} does not match actual staged count {actual}",
|
||||
params.staged_count
|
||||
)));
|
||||
}
|
||||
let reason = params.no_candidates_reason.clone();
|
||||
*self
|
||||
.state
|
||||
.finished
|
||||
.lock()
|
||||
.expect("memory extract finished state poisoned") = Some(params);
|
||||
Ok(ToolOutput {
|
||||
summary: reason
|
||||
.map(|reason| {
|
||||
format!("Finished extraction with {actual} staged candidate(s): {reason}")
|
||||
})
|
||||
.unwrap_or_else(|| {
|
||||
format!("Finished extraction with {actual} staged candidate(s).")
|
||||
}),
|
||||
content: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn map_memory_stage_error(error: WorkspaceMemoryBackendError) -> ToolError {
|
||||
match error {
|
||||
WorkspaceMemoryBackendError::Backend(message) => ToolError::InvalidArgument(message),
|
||||
WorkspaceMemoryBackendError::Http { status, body }
|
||||
if matches!(
|
||||
status,
|
||||
reqwest::StatusCode::BAD_REQUEST | reqwest::StatusCode::UNPROCESSABLE_ENTITY
|
||||
) =>
|
||||
{
|
||||
ToolError::InvalidArgument(body)
|
||||
}
|
||||
error => ToolError::ExecutionFailed(format!("write Memory staging failed: {error}")),
|
||||
}
|
||||
}
|
||||
|
||||
fn evidence_kind(entry: &SessionEntryEvidence) -> EvidenceKind {
|
||||
match (entry.kind, entry.tool_part) {
|
||||
(ReferenceKind::Tool, Some(ToolPart::Input)) => EvidenceKind::new(EvidenceKind::TOOL_CALL),
|
||||
(ReferenceKind::Tool, _) => EvidenceKind::new(EvidenceKind::TOOL_RESULT),
|
||||
_ => EvidenceKind::new(EvidenceKind::MESSAGE),
|
||||
}
|
||||
}
|
||||
|
||||
fn staging_evidence(entry: &SessionEntryEvidence) -> StagingEvidence {
|
||||
StagingEvidence {
|
||||
id: entry.entry_ref.to_string(),
|
||||
kind: evidence_kind(entry),
|
||||
entry_range: Some(entry.entry_range),
|
||||
excerpt: Some(entry.excerpt.clone()),
|
||||
summary: Some(entry.summary.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
fn source_evidence_ref(entry: &SessionEntryEvidence) -> SourceEvidenceRef {
|
||||
SourceEvidenceRef {
|
||||
segment_id: Some(entry.segment_id.clone()),
|
||||
entry_range: Some(entry.entry_range),
|
||||
evidence_id: Some(entry.entry_ref.to_string()),
|
||||
evidence_kind: Some(evidence_kind(entry)),
|
||||
label: Some(entry.label.clone()),
|
||||
summary: Some(entry.summary.clone()),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn render_extract_input(view: &SessionCapture) -> String {
|
||||
let mut output = String::from("# Session overview\n\n");
|
||||
if view.overview().is_empty() {
|
||||
output.push_str("No user/assistant overview entries are available.\n\n");
|
||||
} else {
|
||||
for item in view.overview() {
|
||||
output.push_str(&format!(
|
||||
"- [{} {}] {}\n {}\n intervening_entries: {}\n",
|
||||
item.id,
|
||||
item.kind.as_str(),
|
||||
item.label,
|
||||
truncate_line(&item.text, 500),
|
||||
item.intervening_entries,
|
||||
));
|
||||
}
|
||||
output.push('\n');
|
||||
}
|
||||
output.push_str("# Initial session entry index\n\n");
|
||||
output.push_str("Use ShowOverview, SearchEntries, and ReadEntry to inspect details. Cite only SessionEntryRef values in StageMemoryCandidate.entry_refs.\n\n");
|
||||
let hits = view.search(&SearchOptions {
|
||||
query: String::new(),
|
||||
kind: None,
|
||||
tool_part: None,
|
||||
tool_name: None,
|
||||
limit: Some(50),
|
||||
min_entry_index: None,
|
||||
from: None,
|
||||
through: None,
|
||||
offset: 0,
|
||||
});
|
||||
for hit in hits {
|
||||
output.push_str(&format!(
|
||||
"- [{} {}] {} — {}\n",
|
||||
hit.id,
|
||||
hit.kind.as_str(),
|
||||
hit.label,
|
||||
hit.summary
|
||||
));
|
||||
}
|
||||
output
|
||||
}
|
||||
|
||||
fn truncate_line(text: &str, max_chars: usize) -> String {
|
||||
let normalized = text.replace('\n', " ");
|
||||
if normalized.chars().count() <= max_chars {
|
||||
normalized
|
||||
} else {
|
||||
let mut output = normalized.chars().take(max_chars).collect::<String>();
|
||||
output.push('…');
|
||||
output
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use llm_engine::Item;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn state() -> MemoryExtractState {
|
||||
MemoryExtractState::new(
|
||||
SessionCapture::new("segment-1", vec![Item::user_message("durable decision")]),
|
||||
crate::worker::marker_workspace_client(None, "test-backend"),
|
||||
SourceRef {
|
||||
segment_id: "segment-1".to_string(),
|
||||
range: [0, 0],
|
||||
},
|
||||
"run-1".to_string(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn memory_extract_declares_only_memory_mutation_tools() {
|
||||
let descriptor = MemoryExtractFeature::new(state()).descriptor();
|
||||
assert_eq!(descriptor.id.as_str(), "builtin:memory-extract");
|
||||
assert_eq!(
|
||||
descriptor
|
||||
.tools
|
||||
.iter()
|
||||
.map(|tool| tool.name.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["StageMemoryCandidate", "FinishMemoryExtraction"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_input_uses_session_entry_refs_and_new_tool_names() {
|
||||
let view = SessionCapture::new(
|
||||
"segment-1",
|
||||
vec![
|
||||
Item::user_message("preference"),
|
||||
Item::tool_call("call-1", "Read", "{}"),
|
||||
],
|
||||
);
|
||||
let input = render_extract_input(&view);
|
||||
assert!(input.contains("E00000000"));
|
||||
assert!(input.contains("E00000001"));
|
||||
assert!(input.contains("StageMemoryCandidate.entry_refs"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backend_input_failures_remain_invalid_argument_tool_errors() {
|
||||
let backend = map_memory_stage_error(WorkspaceMemoryBackendError::Backend(
|
||||
"invalid candidate".to_string(),
|
||||
));
|
||||
assert!(matches!(backend, ToolError::InvalidArgument(_)));
|
||||
let http = map_memory_stage_error(WorkspaceMemoryBackendError::Http {
|
||||
status: reqwest::StatusCode::UNPROCESSABLE_ENTITY,
|
||||
body: "invalid candidate".to_string(),
|
||||
});
|
||||
assert!(matches!(http, ToolError::InvalidArgument(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stage_rejects_entry_ref_outside_capture_before_backend_mutation() {
|
||||
let tool = StageMemoryCandidateTool { state: state() };
|
||||
let error = tool
|
||||
.execute(
|
||||
r#"{"kind":"decision","claim":"claim","why_useful":"useful","entry_refs":["E00000009"]}"#,
|
||||
llm_engine::tool::ToolExecutionContext::direct(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(format!("{error:?}").contains("unknown SessionEntryRef"));
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,881 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_engine::Item;
|
||||
use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use session_store::collect_state;
|
||||
|
||||
use crate::feature::{
|
||||
FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureInstructionContribution,
|
||||
FeatureInstructionDeclaration, FeatureInstructionId, FeatureModule, ToolContribution,
|
||||
ToolDeclaration,
|
||||
};
|
||||
use crate::session_capture::{
|
||||
ReadDetail, ReadOptions, ReadSelector, ReferenceKind, SearchOptions, SessionCapture,
|
||||
SessionEntryRef, ToolPart,
|
||||
};
|
||||
use crate::spawn::registry::SpawnedWorkerRegistry;
|
||||
|
||||
const MAX_SUBJECTS: usize = 100;
|
||||
const DEFAULT_PAGE_LIMIT: usize = 20;
|
||||
const MAX_PAGE_LIMIT: usize = 100;
|
||||
const MAX_READ_BYTES: usize = 16 * 1024;
|
||||
const OBSERVATION_INSTRUCTION_ID: &str = "worker-observation.policy";
|
||||
const OBSERVATION_PROMPT_REF: &str = "$yoi/common/worker-observation";
|
||||
#[cfg(test)]
|
||||
const OBSERVATION_PROMPT_SOURCE: &str =
|
||||
include_str!("../../../../../resources/prompts/common/worker-observation.md");
|
||||
|
||||
fn observation_instruction() -> FeatureInstructionDeclaration {
|
||||
FeatureInstructionDeclaration::new(
|
||||
FeatureInstructionId::builtin(OBSERVATION_INSTRUCTION_ID),
|
||||
OBSERVATION_PROMPT_REF,
|
||||
"Worker session observation authority and privacy policy",
|
||||
)
|
||||
.expect("static worker-observation instruction declaration is valid")
|
||||
}
|
||||
|
||||
#[derive(
|
||||
Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, JsonSchema,
|
||||
)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
|
||||
pub enum WorkerObservationSubjectRef {
|
||||
RuntimeWorker {
|
||||
runtime_id: String,
|
||||
worker_id: String,
|
||||
},
|
||||
SubWorker {
|
||||
name: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct WorkerObservationSubject {
|
||||
pub subject: WorkerObservationSubjectRef,
|
||||
pub display_name: String,
|
||||
pub relation: String,
|
||||
pub status: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WorkerSessionCapture {
|
||||
pub segment_id: String,
|
||||
pub items: Vec<Item>,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum WorkerObservationError {
|
||||
#[error("worker session was not found or is not accessible")]
|
||||
NotFound,
|
||||
#[error("worker session observation failed: {0}")]
|
||||
Unavailable(String),
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait WorkerObservationProvider: Send + Sync {
|
||||
/// Returns only subjects already authorized for the current Worker.
|
||||
async fn list_worker_sessions(
|
||||
&self,
|
||||
) -> Result<Vec<WorkerObservationSubject>, WorkerObservationError>;
|
||||
|
||||
/// Reauthorizes and captures the latest committed session for one subject.
|
||||
/// Unauthorized and missing subjects must both return `NotFound`.
|
||||
async fn capture_worker_session(
|
||||
&self,
|
||||
subject: &WorkerObservationSubjectRef,
|
||||
) -> Result<WorkerSessionCapture, WorkerObservationError>;
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct WorkspaceWorkerObservationListResponse {
|
||||
sessions: Vec<WorkerObservationSubject>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct WorkspaceWorkerObservationCaptureResponse {
|
||||
segment_id: String,
|
||||
entries: Vec<serde_json::Value>,
|
||||
}
|
||||
|
||||
pub struct WorkspaceClientWorkerObservationProvider {
|
||||
client: Arc<dyn crate::worker::WorkspaceClient>,
|
||||
}
|
||||
|
||||
impl WorkspaceClientWorkerObservationProvider {
|
||||
pub fn new(client: Arc<dyn crate::worker::WorkspaceClient>) -> Self {
|
||||
Self { client }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkerObservationProvider for WorkspaceClientWorkerObservationProvider {
|
||||
async fn list_worker_sessions(
|
||||
&self,
|
||||
) -> Result<Vec<WorkerObservationSubject>, WorkerObservationError> {
|
||||
let response = self
|
||||
.client
|
||||
.execute(crate::worker::WorkspaceRequest::get(
|
||||
"/worker-observation/sessions",
|
||||
))
|
||||
.map_err(workspace_client_error)?;
|
||||
let body = workspace_response_body(response)?;
|
||||
serde_json::from_str::<WorkspaceWorkerObservationListResponse>(&body)
|
||||
.map(|response| response.sessions)
|
||||
.map_err(|error| WorkerObservationError::Unavailable(error.to_string()))
|
||||
}
|
||||
|
||||
async fn capture_worker_session(
|
||||
&self,
|
||||
subject: &WorkerObservationSubjectRef,
|
||||
) -> Result<WorkerSessionCapture, WorkerObservationError> {
|
||||
let body = serde_json::to_string(subject)
|
||||
.map_err(|error| WorkerObservationError::Unavailable(error.to_string()))?;
|
||||
let response = self
|
||||
.client
|
||||
.execute(crate::worker::WorkspaceRequest::json(
|
||||
crate::worker::WorkspaceRequestMethod::Post,
|
||||
"/worker-observation/session",
|
||||
body,
|
||||
))
|
||||
.map_err(workspace_client_error)?;
|
||||
let body = workspace_response_body(response)?;
|
||||
let response = serde_json::from_str::<WorkspaceWorkerObservationCaptureResponse>(&body)
|
||||
.map_err(|error| WorkerObservationError::Unavailable(error.to_string()))?;
|
||||
let entries = response
|
||||
.entries
|
||||
.into_iter()
|
||||
.map(|entry| {
|
||||
serde_json::from_value(entry)
|
||||
.map_err(|error| WorkerObservationError::Unavailable(error.to_string()))
|
||||
})
|
||||
.collect::<Result<Vec<session_store::LogEntry>, _>>()?;
|
||||
let state = collect_state(&entries);
|
||||
Ok(WorkerSessionCapture {
|
||||
segment_id: response.segment_id,
|
||||
items: state.history,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn workspace_response_body(
|
||||
response: crate::worker::WorkspaceResponse,
|
||||
) -> Result<String, WorkerObservationError> {
|
||||
match response.status {
|
||||
200..=299 => Ok(response.body),
|
||||
403 | 404 => Err(WorkerObservationError::NotFound),
|
||||
status => Err(WorkerObservationError::Unavailable(format!(
|
||||
"Workspace observation request failed with status {status}: {}",
|
||||
response.body
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn workspace_client_error(error: crate::worker::WorkspaceClientError) -> WorkerObservationError {
|
||||
WorkerObservationError::Unavailable(error.to_string())
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct WorkerObservationFeature {
|
||||
provider: Arc<dyn WorkerObservationProvider>,
|
||||
}
|
||||
|
||||
impl WorkerObservationFeature {
|
||||
pub fn new(provider: Arc<dyn WorkerObservationProvider>) -> Self {
|
||||
Self { provider }
|
||||
}
|
||||
}
|
||||
|
||||
impl FeatureModule for WorkerObservationFeature {
|
||||
fn descriptor(&self) -> FeatureDescriptor {
|
||||
FeatureDescriptor::builtin("worker-observation", "Worker Observation")
|
||||
.with_description(
|
||||
"Read-only exploration of explicitly granted active Worker sessions.",
|
||||
)
|
||||
.with_instruction(observation_instruction())
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"ListWorkerSessions",
|
||||
"List bounded summaries of active Worker sessions granted to this Worker.",
|
||||
))
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"ViewSessionOverview",
|
||||
"Show a sparse overview of the latest committed capture for one granted Worker session.",
|
||||
))
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"SearchSessionEntries",
|
||||
"Search or compactly list a bounded range in one granted Worker session.",
|
||||
))
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"ReadSessionEntry",
|
||||
"Read one committed entry from one granted Worker session by SessionEntryRef.",
|
||||
))
|
||||
}
|
||||
|
||||
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
|
||||
context
|
||||
.instructions()
|
||||
.register(FeatureInstructionContribution::new(
|
||||
observation_instruction(),
|
||||
))?;
|
||||
context.tools().register(ToolContribution::new(
|
||||
"ListWorkerSessions",
|
||||
list_definition(self.provider.clone()),
|
||||
))?;
|
||||
context.tools().register(ToolContribution::new(
|
||||
"ViewSessionOverview",
|
||||
overview_definition(self.provider.clone()),
|
||||
))?;
|
||||
context.tools().register(ToolContribution::new(
|
||||
"SearchSessionEntries",
|
||||
search_definition(self.provider.clone()),
|
||||
))?;
|
||||
context.tools().register(ToolContribution::new(
|
||||
"ReadSessionEntry",
|
||||
read_definition(self.provider.clone()),
|
||||
))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CompositeWorkerObservationProvider {
|
||||
providers: Vec<Arc<dyn WorkerObservationProvider>>,
|
||||
}
|
||||
|
||||
impl CompositeWorkerObservationProvider {
|
||||
pub fn new(providers: Vec<Arc<dyn WorkerObservationProvider>>) -> Self {
|
||||
Self { providers }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkerObservationProvider for CompositeWorkerObservationProvider {
|
||||
async fn list_worker_sessions(
|
||||
&self,
|
||||
) -> Result<Vec<WorkerObservationSubject>, WorkerObservationError> {
|
||||
let mut seen = std::collections::HashSet::new();
|
||||
let mut subjects = Vec::new();
|
||||
let mut unavailable = None;
|
||||
for provider in &self.providers {
|
||||
let provider_subjects = match provider.list_worker_sessions().await {
|
||||
Ok(subjects) => subjects,
|
||||
Err(WorkerObservationError::NotFound) => continue,
|
||||
Err(error) => {
|
||||
unavailable.get_or_insert(error);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
for subject in provider_subjects {
|
||||
if seen.insert(subject.subject.clone()) {
|
||||
subjects.push(subject);
|
||||
if subjects.len() == MAX_SUBJECTS {
|
||||
return Ok(subjects);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if subjects.is_empty() {
|
||||
if let Some(error) = unavailable {
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
Ok(subjects)
|
||||
}
|
||||
|
||||
async fn capture_worker_session(
|
||||
&self,
|
||||
subject: &WorkerObservationSubjectRef,
|
||||
) -> Result<WorkerSessionCapture, WorkerObservationError> {
|
||||
for provider in &self.providers {
|
||||
match provider.capture_worker_session(subject).await {
|
||||
Ok(capture) => return Ok(capture),
|
||||
Err(WorkerObservationError::NotFound) => continue,
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
Err(WorkerObservationError::NotFound)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct SpawnedSubWorkerObservationProvider {
|
||||
registry: Arc<SpawnedWorkerRegistry>,
|
||||
}
|
||||
|
||||
impl SpawnedSubWorkerObservationProvider {
|
||||
pub(crate) fn new(registry: Arc<SpawnedWorkerRegistry>) -> Self {
|
||||
Self { registry }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkerObservationProvider for SpawnedSubWorkerObservationProvider {
|
||||
async fn list_worker_sessions(
|
||||
&self,
|
||||
) -> Result<Vec<WorkerObservationSubject>, WorkerObservationError> {
|
||||
let subjects = self
|
||||
.registry
|
||||
.list_internal()
|
||||
.into_iter()
|
||||
.take(MAX_SUBJECTS)
|
||||
.map(|record| WorkerObservationSubject {
|
||||
subject: WorkerObservationSubjectRef::SubWorker {
|
||||
name: record.worker_name.clone(),
|
||||
},
|
||||
display_name: record.worker_name,
|
||||
relation: "subworker".to_string(),
|
||||
status: format!("{:?}", record.session.status()).to_lowercase(),
|
||||
})
|
||||
.collect();
|
||||
Ok(subjects)
|
||||
}
|
||||
|
||||
async fn capture_worker_session(
|
||||
&self,
|
||||
subject: &WorkerObservationSubjectRef,
|
||||
) -> Result<WorkerSessionCapture, WorkerObservationError> {
|
||||
let WorkerObservationSubjectRef::SubWorker { name } = subject else {
|
||||
return Err(WorkerObservationError::NotFound);
|
||||
};
|
||||
let record = self
|
||||
.registry
|
||||
.get_internal(name)
|
||||
.ok_or(WorkerObservationError::NotFound)?;
|
||||
let entries = record.session.entries();
|
||||
let state = collect_state(&entries);
|
||||
Ok(WorkerSessionCapture {
|
||||
segment_id: format!("subworker:{name}"),
|
||||
items: state.history,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn list_definition(provider: Arc<dyn WorkerObservationProvider>) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = serde_json::to_value(schemars::schema_for!(ListWorkerSessionsParams))
|
||||
.unwrap_or_else(|_| serde_json::json!({}));
|
||||
let meta = ToolMeta::new("ListWorkerSessions")
|
||||
.description("List active Worker sessions explicitly granted to this Worker.")
|
||||
.input_schema(schema);
|
||||
let tool: Arc<dyn Tool> = Arc::new(ListWorkerSessionsTool {
|
||||
provider: provider.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
fn overview_definition(provider: Arc<dyn WorkerObservationProvider>) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = serde_json::to_value(schemars::schema_for!(ViewSessionOverviewParams))
|
||||
.unwrap_or_else(|_| serde_json::json!({}));
|
||||
let meta = ToolMeta::new("ViewSessionOverview")
|
||||
.description("Show a sparse bounded index for one granted Worker session.")
|
||||
.input_schema(schema);
|
||||
let tool: Arc<dyn Tool> = Arc::new(ViewSessionOverviewTool {
|
||||
provider: provider.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
fn search_definition(provider: Arc<dyn WorkerObservationProvider>) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = serde_json::to_value(schemars::schema_for!(SearchSessionEntriesParams))
|
||||
.unwrap_or_else(|_| serde_json::json!({}));
|
||||
let meta = ToolMeta::new("SearchSessionEntries")
|
||||
.description("Search or list a bounded range in one granted Worker session.")
|
||||
.input_schema(schema);
|
||||
let tool: Arc<dyn Tool> = Arc::new(SearchSessionEntriesTool {
|
||||
provider: provider.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
fn read_definition(provider: Arc<dyn WorkerObservationProvider>) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = serde_json::to_value(schemars::schema_for!(ReadSessionEntryParams))
|
||||
.unwrap_or_else(|_| serde_json::json!({}));
|
||||
let meta = ToolMeta::new("ReadSessionEntry")
|
||||
.description("Read one entry by SessionEntryRef from one granted Worker session.")
|
||||
.input_schema(schema);
|
||||
let tool: Arc<dyn Tool> = Arc::new(ReadSessionEntryTool {
|
||||
provider: provider.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct ListWorkerSessionsParams {
|
||||
#[serde(default)]
|
||||
limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct ViewSessionOverviewParams {
|
||||
subject: WorkerObservationSubjectRef,
|
||||
#[serde(default)]
|
||||
offset: usize,
|
||||
#[serde(default)]
|
||||
limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct SearchSessionEntriesParams {
|
||||
subject: WorkerObservationSubjectRef,
|
||||
#[serde(default)]
|
||||
query: String,
|
||||
#[serde(default)]
|
||||
kind: Option<String>,
|
||||
#[serde(default)]
|
||||
tool_part: Option<String>,
|
||||
#[serde(default)]
|
||||
tool_name: Option<String>,
|
||||
#[serde(default)]
|
||||
from: Option<String>,
|
||||
#[serde(default)]
|
||||
through: Option<String>,
|
||||
#[serde(default)]
|
||||
offset: usize,
|
||||
#[serde(default)]
|
||||
limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct ReadSessionEntryParams {
|
||||
subject: WorkerObservationSubjectRef,
|
||||
entry_ref: String,
|
||||
#[serde(default = "default_read_mode")]
|
||||
mode: String,
|
||||
}
|
||||
|
||||
fn default_read_mode() -> String {
|
||||
"compact".to_string()
|
||||
}
|
||||
|
||||
struct ListWorkerSessionsTool {
|
||||
provider: Arc<dyn WorkerObservationProvider>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ListWorkerSessionsTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_context: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: ListWorkerSessionsParams = parse_input("ListWorkerSessions", input_json)?;
|
||||
let limit = bounded_limit(params.limit);
|
||||
let mut subjects = self
|
||||
.provider
|
||||
.list_worker_sessions()
|
||||
.await
|
||||
.map_err(tool_error)?;
|
||||
subjects.truncate(limit);
|
||||
let sessions = subjects
|
||||
.iter()
|
||||
.map(|subject| {
|
||||
serde_json::json!({
|
||||
"subject": bounded_subject(&subject.subject),
|
||||
"display_name": truncate_text(&subject.display_name, 200),
|
||||
"relation": truncate_text(&subject.relation, 64),
|
||||
"status": truncate_text(&subject.status, 64),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
json_output(
|
||||
format!("Listed {} Worker session(s).", sessions.len()),
|
||||
serde_json::json!({ "sessions": sessions }),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
struct ViewSessionOverviewTool {
|
||||
provider: Arc<dyn WorkerObservationProvider>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ViewSessionOverviewTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_context: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: ViewSessionOverviewParams = parse_input("ViewSessionOverview", input_json)?;
|
||||
let view = latest_view(&*self.provider, ¶ms.subject).await?;
|
||||
let limit = bounded_limit(params.limit);
|
||||
let entries = view
|
||||
.overview()
|
||||
.iter()
|
||||
.skip(params.offset)
|
||||
.take(limit)
|
||||
.map(|entry| {
|
||||
serde_json::json!({
|
||||
"entry_ref": entry.id,
|
||||
"entry_range": entry.entry_range,
|
||||
"kind": entry.kind.as_str(),
|
||||
"label": entry.label,
|
||||
"text": entry.text,
|
||||
"intervening_entries": entry.intervening_entries,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let has_more = params.offset.saturating_add(entries.len()) < view.overview().len();
|
||||
json_output(
|
||||
format!(
|
||||
"Showing {} Worker session overview entrie(s).",
|
||||
entries.len()
|
||||
),
|
||||
serde_json::json!({
|
||||
"subject": params.subject,
|
||||
"entries": entries,
|
||||
"next_offset": has_more.then_some(params.offset + entries.len()),
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
struct SearchSessionEntriesTool {
|
||||
provider: Arc<dyn WorkerObservationProvider>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SearchSessionEntriesTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_context: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: SearchSessionEntriesParams = parse_input("SearchSessionEntries", input_json)?;
|
||||
let view = latest_view(&*self.provider, ¶ms.subject).await?;
|
||||
let from = params.from.as_deref().map(parse_entry_ref).transpose()?;
|
||||
let through = params.through.as_deref().map(parse_entry_ref).transpose()?;
|
||||
if let (Some(from), Some(through)) = (&from, &through) {
|
||||
if from.source_index() > through.source_index() {
|
||||
return Err(ToolError::InvalidArgument(
|
||||
"SearchSessionEntries from must not be after through".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
let entries = view
|
||||
.search(&SearchOptions {
|
||||
query: params.query,
|
||||
kind: params.kind.as_deref().map(parse_kind).transpose()?,
|
||||
tool_part: params
|
||||
.tool_part
|
||||
.as_deref()
|
||||
.map(parse_tool_part)
|
||||
.transpose()?,
|
||||
tool_name: params.tool_name,
|
||||
limit: Some(bounded_limit(params.limit)),
|
||||
min_entry_index: None,
|
||||
from,
|
||||
through,
|
||||
offset: params.offset,
|
||||
})
|
||||
.into_iter()
|
||||
.map(|entry| {
|
||||
serde_json::json!({
|
||||
"entry_ref": entry.id,
|
||||
"entry_range": entry.entry_range,
|
||||
"kind": entry.kind.as_str(),
|
||||
"tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()),
|
||||
"tool_name": entry.tool_name,
|
||||
"label": entry.label,
|
||||
"text": entry.summary,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
json_output(
|
||||
format!("Found {} Worker session entrie(s).", entries.len()),
|
||||
serde_json::json!({ "subject": params.subject, "entries": entries }),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
struct ReadSessionEntryTool {
|
||||
provider: Arc<dyn WorkerObservationProvider>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ReadSessionEntryTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_context: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: ReadSessionEntryParams = parse_input("ReadSessionEntry", input_json)?;
|
||||
let entry_ref = parse_entry_ref(¶ms.entry_ref)?;
|
||||
let detail = match params.mode.as_str() {
|
||||
"compact" => ReadDetail::Compact,
|
||||
"full" => ReadDetail::Full,
|
||||
other => {
|
||||
return Err(ToolError::InvalidArgument(format!(
|
||||
"invalid mode {other:?}; expected compact or full"
|
||||
)));
|
||||
}
|
||||
};
|
||||
let view = latest_view(&*self.provider, ¶ms.subject).await?;
|
||||
let read = view.read(
|
||||
ReadSelector::Id(entry_ref.as_str()),
|
||||
ReadOptions {
|
||||
include_tools: true,
|
||||
tool_part: ToolPart::Both,
|
||||
detail,
|
||||
max_items: 1,
|
||||
max_bytes: MAX_READ_BYTES,
|
||||
},
|
||||
);
|
||||
let entries = read
|
||||
.entries
|
||||
.into_iter()
|
||||
.map(|entry| {
|
||||
serde_json::json!({
|
||||
"entry_ref": entry.id,
|
||||
"entry_range": entry.entry_range,
|
||||
"kind": entry.kind.as_str(),
|
||||
"tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()),
|
||||
"tool_name": entry.tool_name,
|
||||
"label": entry.label,
|
||||
"text": entry.text,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if entries.is_empty() {
|
||||
return Err(ToolError::ExecutionFailed(
|
||||
"worker session was not found or is not accessible".to_string(),
|
||||
));
|
||||
}
|
||||
json_output(
|
||||
format!("Read {} Worker session entry.", entries.len()),
|
||||
serde_json::json!({
|
||||
"subject": params.subject,
|
||||
"entries": entries,
|
||||
"truncated": read.truncated,
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
async fn latest_view(
|
||||
provider: &dyn WorkerObservationProvider,
|
||||
subject: &WorkerObservationSubjectRef,
|
||||
) -> Result<SessionCapture, ToolError> {
|
||||
let capture = provider
|
||||
.capture_worker_session(subject)
|
||||
.await
|
||||
.map_err(tool_error)?;
|
||||
Ok(SessionCapture::new(capture.segment_id, capture.items))
|
||||
}
|
||||
|
||||
fn parse_input<T: serde::de::DeserializeOwned>(
|
||||
tool_name: &str,
|
||||
input_json: &str,
|
||||
) -> Result<T, ToolError> {
|
||||
serde_json::from_str(input_json)
|
||||
.map_err(|error| ToolError::InvalidArgument(format!("invalid {tool_name} input: {error}")))
|
||||
}
|
||||
|
||||
fn parse_entry_ref(value: &str) -> Result<SessionEntryRef, ToolError> {
|
||||
SessionEntryRef::parse(value)
|
||||
.ok_or_else(|| ToolError::InvalidArgument(format!("invalid SessionEntryRef {value:?}")))
|
||||
}
|
||||
|
||||
fn parse_kind(value: &str) -> Result<ReferenceKind, ToolError> {
|
||||
ReferenceKind::parse(value)
|
||||
.ok_or_else(|| ToolError::InvalidArgument(format!("invalid entry kind {value:?}")))
|
||||
}
|
||||
|
||||
fn parse_tool_part(value: &str) -> Result<ToolPart, ToolError> {
|
||||
ToolPart::parse(value)
|
||||
.ok_or_else(|| ToolError::InvalidArgument(format!("invalid tool_part {value:?}")))
|
||||
}
|
||||
|
||||
fn bounded_subject(subject: &WorkerObservationSubjectRef) -> WorkerObservationSubjectRef {
|
||||
match subject {
|
||||
WorkerObservationSubjectRef::RuntimeWorker {
|
||||
runtime_id,
|
||||
worker_id,
|
||||
} => WorkerObservationSubjectRef::RuntimeWorker {
|
||||
runtime_id: truncate_text(runtime_id, 200),
|
||||
worker_id: truncate_text(worker_id, 200),
|
||||
},
|
||||
WorkerObservationSubjectRef::SubWorker { name } => WorkerObservationSubjectRef::SubWorker {
|
||||
name: truncate_text(name, 200),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate_text(value: &str, max_chars: usize) -> String {
|
||||
if value.chars().count() <= max_chars {
|
||||
value.to_string()
|
||||
} else {
|
||||
let mut truncated = value.chars().take(max_chars).collect::<String>();
|
||||
truncated.push('…');
|
||||
truncated
|
||||
}
|
||||
}
|
||||
|
||||
fn bounded_limit(limit: Option<usize>) -> usize {
|
||||
limit.unwrap_or(DEFAULT_PAGE_LIMIT).clamp(1, MAX_PAGE_LIMIT)
|
||||
}
|
||||
|
||||
fn tool_error(error: WorkerObservationError) -> ToolError {
|
||||
match error {
|
||||
WorkerObservationError::NotFound => ToolError::ExecutionFailed(
|
||||
"worker session was not found or is not accessible".to_string(),
|
||||
),
|
||||
WorkerObservationError::Unavailable(message) => ToolError::ExecutionFailed(message),
|
||||
}
|
||||
}
|
||||
|
||||
fn json_output(summary: String, value: serde_json::Value) -> Result<ToolOutput, ToolError> {
|
||||
let content = serde_json::to_string_pretty(&value)
|
||||
.map_err(|error| ToolError::ExecutionFailed(format!("serialize tool output: {error}")))?;
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: Some(content),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Mutex;
|
||||
|
||||
use llm_engine::Role;
|
||||
|
||||
use crate::feature::{FeatureRegistryBuilder, HookRegistryBuilder};
|
||||
|
||||
use super::*;
|
||||
|
||||
struct FakeProvider {
|
||||
captures: Mutex<Vec<Item>>,
|
||||
}
|
||||
|
||||
fn granted_subject() -> WorkerObservationSubjectRef {
|
||||
WorkerObservationSubjectRef::RuntimeWorker {
|
||||
runtime_id: "runtime-1".to_string(),
|
||||
worker_id: "granted".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkerObservationProvider for FakeProvider {
|
||||
async fn list_worker_sessions(
|
||||
&self,
|
||||
) -> Result<Vec<WorkerObservationSubject>, WorkerObservationError> {
|
||||
Ok(vec![WorkerObservationSubject {
|
||||
subject: granted_subject(),
|
||||
display_name: "Granted".to_string(),
|
||||
relation: "peer".to_string(),
|
||||
status: "idle".to_string(),
|
||||
}])
|
||||
}
|
||||
|
||||
async fn capture_worker_session(
|
||||
&self,
|
||||
subject: &WorkerObservationSubjectRef,
|
||||
) -> Result<WorkerSessionCapture, WorkerObservationError> {
|
||||
if subject != &granted_subject() {
|
||||
return Err(WorkerObservationError::NotFound);
|
||||
}
|
||||
Ok(WorkerSessionCapture {
|
||||
segment_id: "segment".to_string(),
|
||||
items: self.captures.lock().unwrap().clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn message(_id: &str, role: Role, content: &str) -> Item {
|
||||
match role {
|
||||
Role::User => Item::user_message(content),
|
||||
Role::Assistant => Item::assistant_message(content),
|
||||
Role::System => Item::system_message(content),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_source_names_the_worker_observation_contract() {
|
||||
for token in [
|
||||
"ListWorkerSessions",
|
||||
"ViewSessionOverview",
|
||||
"SearchSessionEntries",
|
||||
"ReadSessionEntry",
|
||||
"SessionEntryRef",
|
||||
] {
|
||||
assert!(OBSERVATION_PROMPT_SOURCE.contains(token), "missing {token}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_observation_installs_without_session_explore_or_memory_extract() {
|
||||
let provider = Arc::new(FakeProvider {
|
||||
captures: Mutex::new(Vec::new()),
|
||||
});
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hook_builder = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(WorkerObservationFeature::new(provider))
|
||||
.install_into_pending(&mut pending_tools, &mut hook_builder);
|
||||
assert!(report.reports[0].installed);
|
||||
assert_eq!(
|
||||
report.installed_tool_names(),
|
||||
[
|
||||
"ListWorkerSessions",
|
||||
"ViewSessionOverview",
|
||||
"SearchSessionEntries",
|
||||
"ReadSessionEntry",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_grants_hide_unauthorized_subjects_and_latest_capture_preserves_refs() {
|
||||
let provider = Arc::new(FakeProvider {
|
||||
captures: Mutex::new(vec![message("u1", Role::User, "first")]),
|
||||
});
|
||||
let list = list_definition(provider.clone())().1;
|
||||
let listed = list
|
||||
.execute("{}", llm_engine::tool::ToolExecutionContext::direct())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(listed.content.unwrap().contains("granted"));
|
||||
|
||||
let read = read_definition(provider.clone())().1;
|
||||
let hidden = read
|
||||
.execute(
|
||||
r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"unauthorized"},"entry_ref":"E00000000"}"#,
|
||||
llm_engine::tool::ToolExecutionContext::direct(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(format!("{hidden:?}").contains("not found or is not accessible"));
|
||||
|
||||
provider
|
||||
.captures
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(message("a1", Role::Assistant, "second"));
|
||||
let output = read
|
||||
.execute(
|
||||
r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"granted"},"entry_ref":"E00000000"}"#,
|
||||
llm_engine::tool::ToolExecutionContext::direct(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(output.content.unwrap().contains("first"));
|
||||
|
||||
let output = read
|
||||
.execute(
|
||||
r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"granted"},"entry_ref":"E00000001"}"#,
|
||||
llm_engine::tool::ToolExecutionContext::direct(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(output.content.unwrap().contains("second"));
|
||||
}
|
||||
}
|
||||
@@ -58,8 +58,8 @@ pub fn fire_and_forget(socket: Option<PathBuf>, event: WorkerEvent) {
|
||||
/// Only events classified by `WorkerEvent::should_notify_agent` are injected
|
||||
/// into the parent's LLM context as system messages; control-plane-only events
|
||||
/// keep this renderer for diagnostics/tests. Agent-visible summaries are kept
|
||||
/// deliberately short — the LLM can always call `SubWorkerReadOutput` to fetch more
|
||||
/// detail if the event summary is not enough.
|
||||
/// deliberately short — the LLM can use worker-observation tools to inspect the committed
|
||||
/// session when the event summary is not enough.
|
||||
pub fn render_event(event: &WorkerEvent) -> String {
|
||||
match event {
|
||||
WorkerEvent::TurnEnded { worker_name } => {
|
||||
|
||||
@@ -11,7 +11,7 @@ pub mod model_client;
|
||||
pub mod prompt;
|
||||
pub mod runtime;
|
||||
pub mod segment_log_sink;
|
||||
mod session_reference;
|
||||
mod session_capture;
|
||||
pub mod shared_state;
|
||||
mod shutdown_after_idle;
|
||||
pub mod skill;
|
||||
|
||||
@@ -208,7 +208,6 @@ struct ToolCapabilities {
|
||||
memory_update_document: bool,
|
||||
sub_worker_spawn: bool,
|
||||
sub_worker_send: bool,
|
||||
sub_worker_read_output: bool,
|
||||
sub_worker_stop: bool,
|
||||
sub_worker_list: bool,
|
||||
sub_worker_restore: bool,
|
||||
@@ -224,7 +223,6 @@ impl ToolCapabilities {
|
||||
"MemoryUpdateDocument" => capabilities.memory_update_document = true,
|
||||
"SubWorkerSpawn" => capabilities.sub_worker_spawn = true,
|
||||
"SubWorkerSend" => capabilities.sub_worker_send = true,
|
||||
"SubWorkerReadOutput" => capabilities.sub_worker_read_output = true,
|
||||
"SubWorkerStop" => capabilities.sub_worker_stop = true,
|
||||
"SubWorkerList" => capabilities.sub_worker_list = true,
|
||||
_ => {}
|
||||
@@ -248,7 +246,6 @@ impl ToolCapabilities {
|
||||
fn sub_worker_management(self) -> bool {
|
||||
self.sub_worker_spawn
|
||||
|| self.sub_worker_send
|
||||
|| self.sub_worker_read_output
|
||||
|| self.sub_worker_stop
|
||||
|| self.sub_worker_list
|
||||
|| self.sub_worker_restore
|
||||
|
||||
@@ -1,26 +1,55 @@
|
||||
//! Immutable reference view over a session history slice.
|
||||
//! Workspace- and Memory-independent exploration of an immutable ordered session capture.
|
||||
//!
|
||||
//! This module is shared substrate for internal workers that need to inspect a
|
||||
//! bounded, host-created view of session history without reading the live
|
||||
//! foreground Worker state directly.
|
||||
//! Hosts construct a capture from committed session items. The capture excludes reasoning,
|
||||
//! assigns append-stable `SessionEntryRef` values, and provides sparse overview, bounded
|
||||
//! range/search, read, and generic evidence projections without granting mutation authority.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use llm_engine::{Item, Role};
|
||||
use memory::extract::StagingEvidence;
|
||||
use memory::schema::{EvidenceKind, SourceEvidenceRef};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
const DEFAULT_SEARCH_LIMIT: usize = 20;
|
||||
const MAX_SEARCH_LIMIT: usize = 50;
|
||||
const DEFAULT_READ_MAX_ITEMS: usize = 40;
|
||||
const MAX_READ_MAX_ITEMS: usize = 80;
|
||||
const DEFAULT_READ_MAX_BYTES: usize = 32 * 1024;
|
||||
const OVERVIEW_ANCHOR_STRIDE: usize = 8;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub(crate) struct SessionEntryRef(String);
|
||||
|
||||
impl SessionEntryRef {
|
||||
pub(crate) fn new(source_index: usize) -> Self {
|
||||
Self(format!("E{source_index:08}"))
|
||||
}
|
||||
|
||||
pub(crate) fn parse(value: &str) -> Option<Self> {
|
||||
let reference = Self(value.to_string());
|
||||
reference.source_index()?;
|
||||
Some(reference)
|
||||
}
|
||||
|
||||
pub(crate) fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
|
||||
pub(crate) fn source_index(&self) -> Option<u64> {
|
||||
self.0.strip_prefix('E')?.parse().ok()
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SessionEntryRef {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum ReferenceKind {
|
||||
User,
|
||||
Assistant,
|
||||
System,
|
||||
Tool,
|
||||
}
|
||||
|
||||
@@ -29,7 +58,6 @@ impl ReferenceKind {
|
||||
match self {
|
||||
Self::User => "user",
|
||||
Self::Assistant => "assistant",
|
||||
Self::System => "system",
|
||||
Self::Tool => "tool",
|
||||
}
|
||||
}
|
||||
@@ -38,15 +66,10 @@ impl ReferenceKind {
|
||||
match value {
|
||||
"user" => Some(Self::User),
|
||||
"assistant" | "agent" => Some(Self::Assistant),
|
||||
"system" => Some(Self::System),
|
||||
"tool" => Some(Self::Tool),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn evidence_kind(self) -> EvidenceKind {
|
||||
EvidenceKind::new(EvidenceKind::MESSAGE)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -73,16 +96,17 @@ impl ToolPart {
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct OverviewItem {
|
||||
pub id: String,
|
||||
pub id: SessionEntryRef,
|
||||
pub entry_range: [u64; 2],
|
||||
pub kind: ReferenceKind,
|
||||
pub label: String,
|
||||
pub text: String,
|
||||
pub intervening_entries: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct ReferenceEntry {
|
||||
pub id: String,
|
||||
pub id: SessionEntryRef,
|
||||
pub entry_range: [u64; 2],
|
||||
pub kind: ReferenceKind,
|
||||
pub tool_part: Option<ToolPart>,
|
||||
@@ -92,20 +116,6 @@ pub(crate) struct ReferenceEntry {
|
||||
search_text: String,
|
||||
}
|
||||
|
||||
impl ReferenceEntry {
|
||||
fn evidence_kind(&self) -> EvidenceKind {
|
||||
match (self.kind, self.tool_part) {
|
||||
(ReferenceKind::Tool, Some(ToolPart::Input)) => {
|
||||
EvidenceKind::new(EvidenceKind::TOOL_CALL)
|
||||
}
|
||||
(ReferenceKind::Tool, Some(ToolPart::Output | ToolPart::Both) | None) => {
|
||||
EvidenceKind::new(EvidenceKind::TOOL_RESULT)
|
||||
}
|
||||
_ => self.kind.evidence_kind(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub(crate) struct SearchOptions {
|
||||
pub query: String,
|
||||
@@ -114,11 +124,14 @@ pub(crate) struct SearchOptions {
|
||||
pub tool_name: Option<String>,
|
||||
pub limit: Option<usize>,
|
||||
pub min_entry_index: Option<u64>,
|
||||
pub from: Option<SessionEntryRef>,
|
||||
pub through: Option<SessionEntryRef>,
|
||||
pub offset: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct SearchHit {
|
||||
pub id: String,
|
||||
pub id: SessionEntryRef,
|
||||
pub kind: ReferenceKind,
|
||||
pub tool_part: Option<ToolPart>,
|
||||
pub tool_name: Option<String>,
|
||||
@@ -162,7 +175,7 @@ impl Default for ReadOptions {
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct ReadEntry {
|
||||
pub id: String,
|
||||
pub id: SessionEntryRef,
|
||||
pub kind: ReferenceKind,
|
||||
pub tool_part: Option<ToolPart>,
|
||||
pub tool_name: Option<String>,
|
||||
@@ -178,14 +191,26 @@ pub(crate) struct ReadResult {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct SessionReferenceView {
|
||||
pub(crate) struct SessionEntryEvidence {
|
||||
pub segment_id: String,
|
||||
pub entry_ref: SessionEntryRef,
|
||||
pub entry_range: [u64; 2],
|
||||
pub kind: ReferenceKind,
|
||||
pub tool_part: Option<ToolPart>,
|
||||
pub label: String,
|
||||
pub summary: String,
|
||||
pub excerpt: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct SessionCapture {
|
||||
segment_id: String,
|
||||
items: Arc<Vec<Item>>,
|
||||
overview: Vec<OverviewItem>,
|
||||
index: Vec<ReferenceEntry>,
|
||||
}
|
||||
|
||||
impl SessionReferenceView {
|
||||
impl SessionCapture {
|
||||
pub(crate) fn new(segment_id: impl Into<String>, items: Vec<Item>) -> Self {
|
||||
let segment_id = segment_id.into();
|
||||
let items = Arc::new(items);
|
||||
@@ -199,7 +224,7 @@ impl SessionReferenceView {
|
||||
let kind = match role {
|
||||
Role::User => ReferenceKind::User,
|
||||
Role::Assistant => ReferenceKind::Assistant,
|
||||
Role::System => ReferenceKind::System,
|
||||
Role::System => continue,
|
||||
};
|
||||
let text = content
|
||||
.iter()
|
||||
@@ -208,7 +233,7 @@ impl SessionReferenceView {
|
||||
.join("");
|
||||
let label = format!("{} message", kind.as_str());
|
||||
let summary = truncate_chars(&text, 240);
|
||||
let id = format!("M{idx:04}");
|
||||
let id = SessionEntryRef::new(idx);
|
||||
index.push(ReferenceEntry {
|
||||
id: id.clone(),
|
||||
entry_range,
|
||||
@@ -221,11 +246,12 @@ impl SessionReferenceView {
|
||||
});
|
||||
if matches!(kind, ReferenceKind::User | ReferenceKind::Assistant) {
|
||||
overview.push(OverviewItem {
|
||||
id: format!("O{:04}", overview.len()),
|
||||
id: id.clone(),
|
||||
entry_range,
|
||||
kind,
|
||||
label,
|
||||
text,
|
||||
intervening_entries: 0,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -234,7 +260,7 @@ impl SessionReferenceView {
|
||||
} => {
|
||||
let text = format!("{name}\n{arguments}");
|
||||
index.push(ReferenceEntry {
|
||||
id: format!("T{idx:04}i"),
|
||||
id: SessionEntryRef::new(idx),
|
||||
entry_range,
|
||||
kind: ReferenceKind::Tool,
|
||||
tool_part: Some(ToolPart::Input),
|
||||
@@ -249,7 +275,7 @@ impl SessionReferenceView {
|
||||
} => {
|
||||
let text = format!("{summary}\n{}", content.as_deref().unwrap_or_default());
|
||||
index.push(ReferenceEntry {
|
||||
id: format!("T{idx:04}o"),
|
||||
id: SessionEntryRef::new(idx),
|
||||
entry_range,
|
||||
kind: ReferenceKind::Tool,
|
||||
tool_part: Some(ToolPart::Output),
|
||||
@@ -263,6 +289,30 @@ impl SessionReferenceView {
|
||||
}
|
||||
}
|
||||
|
||||
if overview.len() > 2 {
|
||||
let last = overview.len() - 1;
|
||||
overview = overview
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, entry)| {
|
||||
(index == 0 || index == last || index % OVERVIEW_ANCHOR_STRIDE == 0)
|
||||
.then_some(entry)
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
|
||||
for overview_index in 0..overview.len().saturating_sub(1) {
|
||||
let current_entry = overview[overview_index].entry_range[0];
|
||||
let next_entry = overview[overview_index + 1].entry_range[0];
|
||||
overview[overview_index].intervening_entries = index
|
||||
.iter()
|
||||
.filter(|entry| {
|
||||
let entry_index = entry.entry_range[0];
|
||||
entry_index > current_entry && entry_index < next_entry
|
||||
})
|
||||
.count();
|
||||
}
|
||||
|
||||
Self {
|
||||
segment_id,
|
||||
items,
|
||||
@@ -282,11 +332,21 @@ impl SessionReferenceView {
|
||||
.unwrap_or(DEFAULT_SEARCH_LIMIT)
|
||||
.clamp(1, MAX_SEARCH_LIMIT);
|
||||
let tool_name = options.tool_name.as_deref();
|
||||
let min_entry_index = options.min_entry_index.unwrap_or(0);
|
||||
let min_entry_index = options
|
||||
.from
|
||||
.as_ref()
|
||||
.and_then(SessionEntryRef::source_index)
|
||||
.unwrap_or_else(|| options.min_entry_index.unwrap_or(0));
|
||||
let max_entry_index = options
|
||||
.through
|
||||
.as_ref()
|
||||
.and_then(SessionEntryRef::source_index)
|
||||
.unwrap_or(u64::MAX);
|
||||
let mut skipped = 0usize;
|
||||
let mut hits = Vec::new();
|
||||
|
||||
for entry in &self.index {
|
||||
if entry.entry_range[0] < min_entry_index {
|
||||
if entry.entry_range[0] < min_entry_index || entry.entry_range[0] > max_entry_index {
|
||||
continue;
|
||||
}
|
||||
if let Some(kind) = options.kind {
|
||||
@@ -313,6 +373,10 @@ impl SessionReferenceView {
|
||||
if !query.is_empty() && !entry.search_text.to_lowercase().contains(&query) {
|
||||
continue;
|
||||
}
|
||||
if skipped < options.offset {
|
||||
skipped += 1;
|
||||
continue;
|
||||
}
|
||||
hits.push(SearchHit {
|
||||
id: entry.id.clone(),
|
||||
kind: entry.kind,
|
||||
@@ -338,7 +402,11 @@ impl SessionReferenceView {
|
||||
let mut truncated = false;
|
||||
|
||||
let selected: Vec<&ReferenceEntry> = match selector {
|
||||
ReadSelector::Id(id) => self.index.iter().filter(|entry| entry.id == id).collect(),
|
||||
ReadSelector::Id(id) => self
|
||||
.index
|
||||
.iter()
|
||||
.filter(|entry| entry.id.as_str() == id)
|
||||
.collect(),
|
||||
ReadSelector::EntryRange([start, end]) => self
|
||||
.index
|
||||
.iter()
|
||||
@@ -384,42 +452,32 @@ impl SessionReferenceView {
|
||||
ReadResult { entries, truncated }
|
||||
}
|
||||
|
||||
pub(crate) fn source_ref_for(&self, id: &str) -> Option<SourceEvidenceRef> {
|
||||
let entry = self.index.iter().find(|entry| entry.id == id)?;
|
||||
Some(SourceEvidenceRef {
|
||||
segment_id: Some(self.segment_id.clone()),
|
||||
entry_range: Some(entry.entry_range),
|
||||
evidence_id: Some(entry.id.clone()),
|
||||
evidence_kind: Some(entry.evidence_kind()),
|
||||
label: Some(entry.label.clone()),
|
||||
summary: Some(entry.summary.clone()),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn staging_evidence_for(&self, id: &str) -> Option<StagingEvidence> {
|
||||
let entry = self.index.iter().find(|entry| entry.id == id)?;
|
||||
let read = self.read(
|
||||
ReadSelector::Id(id),
|
||||
ReadOptions {
|
||||
include_tools: true,
|
||||
tool_part: ToolPart::Both,
|
||||
detail: ReadDetail::Compact,
|
||||
max_items: 1,
|
||||
max_bytes: 2 * 1024,
|
||||
},
|
||||
);
|
||||
let excerpt = read
|
||||
pub(crate) fn evidence_for(&self, id: &str) -> Option<SessionEntryEvidence> {
|
||||
let entry = self.index.iter().find(|entry| entry.id.as_str() == id)?;
|
||||
let excerpt = self
|
||||
.read(
|
||||
ReadSelector::Id(id),
|
||||
ReadOptions {
|
||||
include_tools: true,
|
||||
tool_part: ToolPart::Both,
|
||||
detail: ReadDetail::Compact,
|
||||
max_items: 1,
|
||||
max_bytes: 2 * 1024,
|
||||
},
|
||||
)
|
||||
.entries
|
||||
.first()
|
||||
.map(|entry| entry.text.clone())
|
||||
.unwrap_or_else(|| entry.summary.clone());
|
||||
Some(StagingEvidence {
|
||||
id: entry.id.clone(),
|
||||
kind: entry.evidence_kind(),
|
||||
entry_range: Some(entry.entry_range),
|
||||
excerpt: Some(excerpt),
|
||||
summary: Some(entry.summary.clone()),
|
||||
Some(SessionEntryEvidence {
|
||||
segment_id: self.segment_id.clone(),
|
||||
entry_ref: entry.id.clone(),
|
||||
entry_range: entry.entry_range,
|
||||
kind: entry.kind,
|
||||
tool_part: entry.tool_part,
|
||||
label: entry.label.clone(),
|
||||
summary: entry.summary.clone(),
|
||||
excerpt,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -486,7 +544,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn overview_contains_user_and_assistant_only() {
|
||||
let view = SessionReferenceView::new(
|
||||
let view = SessionCapture::new(
|
||||
"segment-1",
|
||||
vec![
|
||||
Item::system_message("sys"),
|
||||
@@ -506,7 +564,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn search_filters_tool_input_and_output() {
|
||||
let view = SessionReferenceView::new(
|
||||
let view = SessionCapture::new(
|
||||
"segment-1",
|
||||
vec![
|
||||
Item::tool_call("c1", "Read", "{\"file\":\"Cargo.toml\"}"),
|
||||
@@ -521,6 +579,9 @@ mod tests {
|
||||
tool_name: Some("Read".into()),
|
||||
limit: None,
|
||||
min_entry_index: None,
|
||||
from: None,
|
||||
through: None,
|
||||
offset: 0,
|
||||
});
|
||||
assert_eq!(input_hits.len(), 1);
|
||||
assert_eq!(input_hits[0].tool_part, Some(ToolPart::Input));
|
||||
@@ -532,6 +593,9 @@ mod tests {
|
||||
tool_name: None,
|
||||
limit: None,
|
||||
min_entry_index: None,
|
||||
from: None,
|
||||
through: None,
|
||||
offset: 0,
|
||||
});
|
||||
assert_eq!(output_hits.len(), 1);
|
||||
assert_eq!(output_hits[0].tool_part, Some(ToolPart::Output));
|
||||
@@ -539,7 +603,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn read_by_entry_range_is_bounded_and_can_skip_tools() {
|
||||
let view = SessionReferenceView::new(
|
||||
let view = SessionCapture::new(
|
||||
"segment-1",
|
||||
vec![
|
||||
Item::user_message("one"),
|
||||
@@ -570,11 +634,117 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn source_ref_uses_entry_range_and_evidence_id() {
|
||||
let view = SessionReferenceView::new("segment-1", vec![Item::user_message("hello")]);
|
||||
let source = view.source_ref_for("M0000").unwrap();
|
||||
assert_eq!(source.segment_id.as_deref(), Some("segment-1"));
|
||||
assert_eq!(source.entry_range, Some([0, 0]));
|
||||
assert_eq!(source.evidence_id.as_deref(), Some("M0000"));
|
||||
fn system_prompt_and_reasoning_are_excluded_from_every_projection() {
|
||||
let view = SessionCapture::new(
|
||||
"segment-1",
|
||||
vec![
|
||||
Item::system_message("raw secret system prompt"),
|
||||
Item::reasoning("private chain of thought"),
|
||||
Item::user_message("visible user entry"),
|
||||
],
|
||||
);
|
||||
let hits = view.search(&SearchOptions {
|
||||
query: String::new(),
|
||||
kind: None,
|
||||
tool_part: None,
|
||||
tool_name: None,
|
||||
limit: None,
|
||||
min_entry_index: None,
|
||||
from: None,
|
||||
through: None,
|
||||
offset: 0,
|
||||
});
|
||||
assert_eq!(hits.len(), 1);
|
||||
assert_eq!(hits[0].id.as_str(), "E00000002");
|
||||
assert!(!hits[0].summary.contains("secret"));
|
||||
assert!(!hits[0].summary.contains("chain of thought"));
|
||||
assert!(
|
||||
view.read(ReadSelector::Id("E00000000"), ReadOptions::default())
|
||||
.entries
|
||||
.is_empty()
|
||||
);
|
||||
assert!(view.evidence_for("E00000000").is_none());
|
||||
assert_eq!(view.overview().len(), 1);
|
||||
assert_eq!(view.overview()[0].id.as_str(), "E00000002");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn overview_is_sparse_and_reports_intervening_non_reasoning_entries() {
|
||||
let items = (0..20)
|
||||
.map(|index| Item::user_message(format!("message-{index}")))
|
||||
.collect::<Vec<_>>();
|
||||
let view = SessionCapture::new("segment-1", items);
|
||||
let refs = view
|
||||
.overview()
|
||||
.iter()
|
||||
.map(|entry| entry.id.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
refs,
|
||||
vec!["E00000000", "E00000008", "E00000016", "E00000019"]
|
||||
);
|
||||
assert_eq!(view.overview()[0].intervening_entries, 7);
|
||||
assert_eq!(view.overview()[1].intervening_entries, 7);
|
||||
assert_eq!(view.overview()[2].intervening_entries, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_preserves_existing_session_entry_refs() {
|
||||
let first = SessionCapture::new(
|
||||
"segment-1",
|
||||
vec![
|
||||
Item::user_message("first"),
|
||||
Item::assistant_message("second"),
|
||||
],
|
||||
);
|
||||
let appended = SessionCapture::new(
|
||||
"segment-1",
|
||||
vec![
|
||||
Item::user_message("first"),
|
||||
Item::assistant_message("second"),
|
||||
Item::user_message("third"),
|
||||
],
|
||||
);
|
||||
let first_refs = first
|
||||
.search(&SearchOptions {
|
||||
query: String::new(),
|
||||
kind: None,
|
||||
tool_part: None,
|
||||
tool_name: None,
|
||||
limit: None,
|
||||
min_entry_index: None,
|
||||
from: None,
|
||||
through: None,
|
||||
offset: 0,
|
||||
})
|
||||
.into_iter()
|
||||
.map(|entry| entry.id)
|
||||
.collect::<Vec<_>>();
|
||||
let appended_refs = appended
|
||||
.search(&SearchOptions {
|
||||
query: String::new(),
|
||||
kind: None,
|
||||
tool_part: None,
|
||||
tool_name: None,
|
||||
limit: None,
|
||||
min_entry_index: None,
|
||||
from: None,
|
||||
through: None,
|
||||
offset: 0,
|
||||
})
|
||||
.into_iter()
|
||||
.take(2)
|
||||
.map(|entry| entry.id)
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(first_refs, appended_refs);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evidence_projection_uses_entry_range_and_session_entry_ref() {
|
||||
let view = SessionCapture::new("segment-1", vec![Item::user_message("hello")]);
|
||||
let source = view.evidence_for("E00000000").unwrap();
|
||||
assert_eq!(source.segment_id, "segment-1");
|
||||
assert_eq!(source.entry_range, [0, 0]);
|
||||
assert_eq!(source.entry_ref.as_str(), "E00000000");
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
//! Parent-facing tools for in-process Internal SubWorker sessions.
|
||||
//!
|
||||
//! All five tools share the same parent-owned `SpawnedWorkerRegistry` of typed session handles.
|
||||
//! All four tools share the same parent-owned `SpawnedWorkerRegistry` of typed session handles.
|
||||
//! There is no Runtime catalog lookup or child socket transport, so a Worker can operate only on
|
||||
//! its direct Internal children. The socket helper at the bottom remains solely for the legacy
|
||||
//! top-level Worker callback protocol and is not part of SubWorker communication.
|
||||
@@ -10,12 +10,10 @@ use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_engine::llm_client::types::{ContentPart, Item, Role};
|
||||
use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use protocol::stream::{JsonLineReader, JsonLineWriter};
|
||||
use protocol::{Event, Method};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use session_store::LogEntry;
|
||||
use tokio::net::UnixStream;
|
||||
|
||||
use crate::spawn::registry::SpawnedWorkerRegistry;
|
||||
@@ -96,7 +94,7 @@ pub fn sub_worker_list_tool(registry: Arc<SpawnedWorkerRegistry>) -> ToolDefinit
|
||||
const SEND_TO_POD_DESCRIPTION: &str = "Send a text message to a previously spawned SubWorker. The SubWorker \
|
||||
processes it as a user turn. Fails if the SubWorker is already executing a \
|
||||
turn — retry after it finishes. Does not wait for the turn to complete; \
|
||||
use `SubWorkerReadOutput` to fetch results afterwards.";
|
||||
use worker-observation tools to inspect its committed session.";
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct SubWorkerSendInput {
|
||||
@@ -146,76 +144,6 @@ pub fn sub_worker_send_tool(registry: Arc<SpawnedWorkerRegistry>) -> ToolDefinit
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SubWorkerReadOutput
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const READ_POD_OUTPUT_DESCRIPTION: &str = "Fetch new assistant text from a SubWorker since the last read. \
|
||||
Uses an internal cursor per-SubWorker so consecutive calls return only \
|
||||
newly-produced output. Returns the SubWorker's current status and the new \
|
||||
text, or reports `stopped` if the SubWorker can no longer be reached.";
|
||||
|
||||
struct SubWorkerReadOutputTool {
|
||||
registry: Arc<SpawnedWorkerRegistry>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SubWorkerReadOutputTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let input: NameInput = serde_json::from_str(input_json).map_err(|e| {
|
||||
ToolError::InvalidArgument(format!("invalid SubWorkerReadOutput input: {e}"))
|
||||
})?;
|
||||
if let Some(record) = self.registry.get_internal(&input.name) {
|
||||
let entries = record.session.entries();
|
||||
let cursor = self.registry.cursor(&input.name).await;
|
||||
let new_entries = if cursor >= entries.len() {
|
||||
&[] as &[LogEntry]
|
||||
} else {
|
||||
&entries[cursor..]
|
||||
};
|
||||
let values = new_entries
|
||||
.iter()
|
||||
.filter_map(|entry| serde_json::to_value(entry).ok())
|
||||
.collect::<Vec<_>>();
|
||||
let new_text = extract_assistant_text(&values);
|
||||
self.registry.set_cursor(&input.name, entries.len()).await;
|
||||
let status = format!("{:?}", record.session.status()).to_lowercase();
|
||||
let summary = if new_text.is_empty() {
|
||||
format!("worker `{}` {status}; no new assistant text", input.name)
|
||||
} else {
|
||||
format!(
|
||||
"worker `{}` {status}: {} new line(s) of assistant text",
|
||||
input.name,
|
||||
new_text.lines().count()
|
||||
)
|
||||
};
|
||||
return Ok(ToolOutput {
|
||||
summary,
|
||||
content: (!new_text.is_empty()).then_some(new_text),
|
||||
});
|
||||
}
|
||||
Err(unknown_worker_err(&input.name))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sub_worker_read_output_tool(registry: Arc<SpawnedWorkerRegistry>) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(NameInput);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("SubWorkerReadOutput")
|
||||
.description(READ_POD_OUTPUT_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(SubWorkerReadOutputTool {
|
||||
registry: registry.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SubWorkerStop
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -323,47 +251,6 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_assistant_text(entries: &[serde_json::Value]) -> String {
|
||||
let mut out = String::new();
|
||||
for value in entries {
|
||||
// The wire payload is the JSON form of `session_store::LogEntry`.
|
||||
// Walk current singular assistant items and the seeded history in
|
||||
// post-compaction `SegmentStart` entries.
|
||||
let Ok(entry) = serde_json::from_value::<LogEntry>(value.clone()) else {
|
||||
continue;
|
||||
};
|
||||
match entry {
|
||||
LogEntry::SegmentStart { history, .. } => {
|
||||
for logged in history {
|
||||
push_assistant_text(&mut out, logged);
|
||||
}
|
||||
}
|
||||
LogEntry::AssistantItem { item, .. } => push_assistant_text(&mut out, item),
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn push_assistant_text(out: &mut String, logged: session_store::LoggedItem) {
|
||||
let item: Item = logged.into();
|
||||
if let Item::Message {
|
||||
role: Role::Assistant,
|
||||
content,
|
||||
..
|
||||
} = item
|
||||
{
|
||||
for part in content {
|
||||
if let ContentPart::Text { text } = part {
|
||||
if !out.is_empty() {
|
||||
out.push_str("\n\n");
|
||||
}
|
||||
out.push_str(&text);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
@@ -1,15 +1,13 @@
|
||||
//! Parent-owned registry of direct Internal SubWorker sessions.
|
||||
//!
|
||||
//! `SubWorkerSpawn` inserts typed `InternalWorkerSessionHandle`s; List/Send/ReadOutput/Stop use
|
||||
//! the same in-memory authority. Internal children are not persisted, restored, discovered as
|
||||
//! `SubWorkerSpawn` inserts typed `InternalWorkerSessionHandle`s; List/Send/Stop and
|
||||
//! worker-observation use the same in-memory authority. Internal children are not persisted, restored, discovered as
|
||||
//! Runtime Workers, or addressed through sockets. Restore consumes any legacy persisted process
|
||||
//! child records only to reclaim their delegated scope and clear obsolete metadata.
|
||||
//!
|
||||
//! `SubWorkerReadOutput` owns a per-child, process-lifetime history cursor so consecutive reads
|
||||
//! yield only new assistant text. Parent registry drop closes all session handles and synchronously
|
||||
//! returns delegated Write deny rules to the parent scope.
|
||||
//! Parent registry drop closes all session handles and synchronously returns delegated Write deny
|
||||
//! rules to the parent scope.
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::collections::HashSet;
|
||||
use std::io;
|
||||
use std::sync::{
|
||||
Arc,
|
||||
@@ -20,7 +18,6 @@ use manifest::{Permission, ScopeRule, SharedScope};
|
||||
use session_store::{
|
||||
WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerStoreError,
|
||||
};
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::internal_worker::InternalWorkerSessionHandle;
|
||||
@@ -95,7 +92,6 @@ impl Drop for InternalSpawnReservation {
|
||||
pub struct SpawnedWorkerRegistry {
|
||||
internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>,
|
||||
internal_names: std::sync::Mutex<HashSet<String>>,
|
||||
cursors: Mutex<HashMap<String, usize>>,
|
||||
parent_scope: Option<SharedScope>,
|
||||
}
|
||||
|
||||
@@ -111,7 +107,6 @@ impl SpawnedWorkerRegistry {
|
||||
Arc::new(Self {
|
||||
internal_records: std::sync::Mutex::new(Vec::new()),
|
||||
internal_names: std::sync::Mutex::new(HashSet::new()),
|
||||
cursors: Mutex::new(HashMap::new()),
|
||||
parent_scope: None,
|
||||
})
|
||||
}
|
||||
@@ -120,7 +115,6 @@ impl SpawnedWorkerRegistry {
|
||||
Arc::new(Self {
|
||||
internal_records: std::sync::Mutex::new(Vec::new()),
|
||||
internal_names: std::sync::Mutex::new(HashSet::new()),
|
||||
cursors: Mutex::new(HashMap::new()),
|
||||
parent_scope: Some(parent_scope),
|
||||
})
|
||||
}
|
||||
@@ -198,7 +192,6 @@ impl SpawnedWorkerRegistry {
|
||||
registry: Arc::new(Self {
|
||||
internal_records: std::sync::Mutex::new(Vec::new()),
|
||||
internal_names: std::sync::Mutex::new(HashSet::new()),
|
||||
cursors: Mutex::new(HashMap::new()),
|
||||
parent_scope,
|
||||
}),
|
||||
reclaimed_unreachable: !persisted_children.is_empty(),
|
||||
@@ -292,20 +285,8 @@ impl SpawnedWorkerRegistry {
|
||||
}
|
||||
removed
|
||||
};
|
||||
self.cursors.lock().await.remove(worker_name);
|
||||
Ok(removed)
|
||||
}
|
||||
|
||||
pub async fn cursor(&self, worker_name: &str) -> usize {
|
||||
*self.cursors.lock().await.get(worker_name).unwrap_or(&0)
|
||||
}
|
||||
|
||||
pub async fn set_cursor(&self, worker_name: &str, value: usize) {
|
||||
self.cursors
|
||||
.lock()
|
||||
.await
|
||||
.insert(worker_name.to_owned(), value);
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for SpawnedWorkerRegistry {
|
||||
|
||||
@@ -409,7 +409,7 @@ impl Tool for SubWorkerSpawnTool {
|
||||
}
|
||||
}
|
||||
parent_notifies.push_notify(
|
||||
format!("SubWorker `{child_name}` turn ended with status {status:?}. Read its output before making completion decisions."),
|
||||
format!("SubWorker `{child_name}` turn ended with status {status:?}. Inspect its committed session with worker-observation tools before making completion decisions."),
|
||||
true,
|
||||
);
|
||||
})),
|
||||
@@ -850,7 +850,9 @@ mod tests {
|
||||
use async_trait::async_trait;
|
||||
use futures::Stream;
|
||||
use llm_engine::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
|
||||
use llm_engine::llm_client::types::ContentPart;
|
||||
use llm_engine::llm_client::{ClientError, LlmClient, Request};
|
||||
use llm_engine::{Item, Role};
|
||||
use manifest::{AuthRef, ModelManifest, SchemeKind, WorkerManifest};
|
||||
use tempfile::TempDir;
|
||||
|
||||
@@ -1028,22 +1030,23 @@ extract_threshold = 4000
|
||||
.contains("reviewer-child")
|
||||
);
|
||||
|
||||
let read = (crate::spawn::comm_tools::sub_worker_read_output_tool(registry.clone()))().1;
|
||||
let first_output = read
|
||||
.execute(r#"{"name":"reviewer-child"}"#, context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
first_output
|
||||
.content
|
||||
.unwrap_or_default()
|
||||
.contains("reviewed")
|
||||
);
|
||||
let second_output = read
|
||||
.execute(r#"{"name":"reviewer-child"}"#, context.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(second_output.content.is_none());
|
||||
let observation =
|
||||
crate::feature::builtin::worker_observation::SpawnedSubWorkerObservationProvider::new(
|
||||
registry.clone(),
|
||||
);
|
||||
let observed_child =
|
||||
crate::feature::builtin::worker_observation::WorkerObservationSubjectRef::SubWorker {
|
||||
name: "reviewer-child".to_string(),
|
||||
};
|
||||
let first_capture = crate::feature::builtin::worker_observation::WorkerObservationProvider::capture_worker_session(
|
||||
&observation,
|
||||
&observed_child,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(first_capture.items.iter().any(|item| {
|
||||
matches!(item, Item::Message { role: Role::Assistant, content, .. } if content.iter().any(|part| matches!(part, ContentPart::Text { text } if text.contains("reviewed"))))
|
||||
}));
|
||||
|
||||
let send = (crate::spawn::comm_tools::sub_worker_send_tool(registry.clone()))().1;
|
||||
send.execute(
|
||||
@@ -1057,6 +1060,13 @@ extract_threshold = 4000
|
||||
crate::internal_worker::InternalWorkerSessionStatus::Idle
|
||||
);
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 2);
|
||||
let latest_capture = crate::feature::builtin::worker_observation::WorkerObservationProvider::capture_worker_session(
|
||||
&observation,
|
||||
&observed_child,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(latest_capture.items.len() > first_capture.items.len());
|
||||
|
||||
fail_requests.store(true, Ordering::SeqCst);
|
||||
send.execute(
|
||||
@@ -1094,9 +1104,9 @@ extract_threshold = 4000
|
||||
.unwrap();
|
||||
assert!(!spawner_scope.snapshot().is_writable(&workspace_root));
|
||||
drop(list);
|
||||
drop(read);
|
||||
drop(send);
|
||||
drop(stop);
|
||||
drop(observation);
|
||||
drop(tool);
|
||||
drop(registry);
|
||||
assert!(spawner_scope.snapshot().is_writable(&workspace_root));
|
||||
|
||||
+42
-12
@@ -32,7 +32,8 @@ use crate::compact::state::CompactState;
|
||||
use crate::compact::usage_tracker::UsageTracker;
|
||||
use crate::feature::builtin::memory::WorkspaceMemoryBackendError;
|
||||
use crate::feature::builtin::{
|
||||
SessionExploreFeature, SessionExploreState, TaskFeature, render_extract_input,
|
||||
MemoryExtractFeature, MemoryExtractState, SessionExploreFeature, SessionExploreState,
|
||||
TaskFeature, WorkerObservationProvider, render_extract_input,
|
||||
};
|
||||
use crate::feature::{
|
||||
FeatureInstructionDeclaration, FeatureInstructionId, FeatureRegistryBuilder,
|
||||
@@ -686,6 +687,9 @@ pub struct Worker<C: LlmClient, St: Store> {
|
||||
/// the narrow snapshot/restore surface Worker needs for compaction and rewind.
|
||||
/// Store/reminder ownership stays inside the Task feature module.
|
||||
task_feature: TaskFeature,
|
||||
/// Host-owned projection of Worker sessions explicitly granted to this Worker.
|
||||
/// The provider reauthorizes every capture and never derives authority from model input.
|
||||
worker_observation_provider: Option<Arc<dyn WorkerObservationProvider>>,
|
||||
/// Parsed system-prompt template awaiting first-turn materialisation.
|
||||
/// `Some` until `ensure_system_prompt_materialized` renders it once,
|
||||
/// then `None` forever — including after compaction.
|
||||
@@ -839,6 +843,7 @@ impl<C: LlmClient + Clone + 'static, St: Store + Clone + 'static> Worker<C, St>
|
||||
usage_history: self.usage_history.clone(),
|
||||
tracker: None,
|
||||
task_feature: self.task_feature.clone(),
|
||||
worker_observation_provider: None,
|
||||
system_prompt_template: None,
|
||||
feature_instructions: self.feature_instructions.clone(),
|
||||
alerter: self.alerter.clone(),
|
||||
@@ -1036,6 +1041,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
|
||||
usage_history: Arc::new(Mutex::new(Vec::<UsageRecord>::new())),
|
||||
tracker: None,
|
||||
task_feature: TaskFeature::new(),
|
||||
worker_observation_provider: None,
|
||||
system_prompt_template: None,
|
||||
feature_instructions: Vec::new(),
|
||||
alerter: None,
|
||||
@@ -1170,6 +1176,19 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
|
||||
self.workspace_context.client_handle()
|
||||
}
|
||||
|
||||
/// Bind the host-owned Worker-session observation projection. The provider
|
||||
/// is responsible for workspace authorization and per-capture revalidation.
|
||||
pub fn bind_worker_observation_provider(
|
||||
&mut self,
|
||||
provider: Option<Arc<dyn WorkerObservationProvider>>,
|
||||
) {
|
||||
self.worker_observation_provider = provider;
|
||||
}
|
||||
|
||||
pub(crate) fn worker_observation_provider(&self) -> Option<Arc<dyn WorkerObservationProvider>> {
|
||||
self.worker_observation_provider.clone()
|
||||
}
|
||||
|
||||
async fn resident_summary_from_workspace_authority(
|
||||
&self,
|
||||
) -> Result<Option<String>, WorkerError> {
|
||||
@@ -3445,15 +3464,21 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
|
||||
segment_id: source_segment_id.to_string(),
|
||||
range: [start_entry as u64, end_entry as u64],
|
||||
};
|
||||
let session_view = crate::session_reference::SessionReferenceView::new(
|
||||
let session_view = crate::session_capture::SessionCapture::new(
|
||||
source_segment_id.to_string(),
|
||||
items_to_extract,
|
||||
);
|
||||
let session_explore_state =
|
||||
SessionExploreState::new(session_view, self.workspace_client_handle(), source);
|
||||
let session_explore_state = SessionExploreState::new(session_view.clone());
|
||||
let memory_extract_state = MemoryExtractState::new(
|
||||
session_view,
|
||||
self.workspace_client_handle(),
|
||||
source,
|
||||
audit.run_id.to_string(),
|
||||
);
|
||||
let input_text = render_extract_input(session_explore_state.view());
|
||||
let features = FeatureRegistryBuilder::new()
|
||||
.with_module(SessionExploreFeature::new(session_explore_state.clone()));
|
||||
.with_module(SessionExploreFeature::new(session_explore_state.clone()))
|
||||
.with_module(MemoryExtractFeature::new(memory_extract_state.clone()));
|
||||
let mut internal_manifest = self.manifest.clone();
|
||||
internal_manifest.model = model.clone();
|
||||
let internal_spec = InternalWorkerSpec {
|
||||
@@ -3469,10 +3494,11 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
|
||||
max_turns: extract_worker_max_turns,
|
||||
features,
|
||||
required_tools: &[
|
||||
"search_evidence",
|
||||
"read_evidence",
|
||||
"stage_candidate",
|
||||
"finish_extraction",
|
||||
"ShowOverview",
|
||||
"SearchEntries",
|
||||
"ReadEntry",
|
||||
"StageMemoryCandidate",
|
||||
"FinishMemoryExtraction",
|
||||
],
|
||||
authority: InternalWorkerAuthority {
|
||||
workspace: self.workspace_context.clone(),
|
||||
@@ -3533,11 +3559,11 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
|
||||
}
|
||||
};
|
||||
|
||||
let staging_results = session_explore_state.staged();
|
||||
if !session_explore_state.is_finished() {
|
||||
let staging_results = memory_extract_state.staged();
|
||||
if !memory_extract_state.is_finished() {
|
||||
tracing::warn!(
|
||||
staged_count = staging_results.len(),
|
||||
"extract worker did not call finish_extraction; advancing pointer with staged output"
|
||||
"extract worker did not call FinishMemoryExtraction; advancing pointer with staged output"
|
||||
);
|
||||
}
|
||||
let staging_id = staging_results.first().cloned().unwrap_or_default();
|
||||
@@ -3925,6 +3951,7 @@ where
|
||||
usage_history: Arc::new(Mutex::new(Vec::new())),
|
||||
tracker: None,
|
||||
task_feature: TaskFeature::new(),
|
||||
worker_observation_provider: None,
|
||||
system_prompt_template: common.system_prompt_template,
|
||||
feature_instructions: common.feature_instructions,
|
||||
alerter: None,
|
||||
@@ -3999,6 +4026,7 @@ where
|
||||
usage_history: Arc::new(Mutex::new(Vec::new())),
|
||||
tracker: None,
|
||||
task_feature: TaskFeature::new(),
|
||||
worker_observation_provider: None,
|
||||
system_prompt_template: common.system_prompt_template,
|
||||
feature_instructions: common.feature_instructions,
|
||||
alerter: None,
|
||||
@@ -4107,6 +4135,7 @@ where
|
||||
usage_history: Arc::new(Mutex::new(Vec::new())),
|
||||
tracker: None,
|
||||
task_feature: TaskFeature::new(),
|
||||
worker_observation_provider: None,
|
||||
system_prompt_template: common.system_prompt_template,
|
||||
feature_instructions: common.feature_instructions,
|
||||
alerter: None,
|
||||
@@ -4399,6 +4428,7 @@ where
|
||||
usage_history: Arc::new(Mutex::new(state.usage_history)),
|
||||
tracker: None,
|
||||
task_feature,
|
||||
worker_observation_provider: None,
|
||||
// Restore replays the saved system_prompt verbatim — no
|
||||
// template re-render on resume.
|
||||
system_prompt_template: None,
|
||||
|
||||
@@ -539,14 +539,14 @@ target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
fn finish_extraction_tool_use_events(call_id: &str) -> Vec<LlmEvent> {
|
||||
fn finish_memory_extraction_tool_use_events(call_id: &str) -> Vec<LlmEvent> {
|
||||
let input = serde_json::json!({
|
||||
"staged_count": 0,
|
||||
"no_candidates_reason": "test run has no durable candidates"
|
||||
})
|
||||
.to_string();
|
||||
vec![
|
||||
LlmEvent::tool_use_start(0, call_id, "finish_extraction"),
|
||||
LlmEvent::tool_use_start(0, call_id, "FinishMemoryExtraction"),
|
||||
LlmEvent::tool_input_delta(0, input),
|
||||
LlmEvent::tool_use_stop(0),
|
||||
LlmEvent::Status(StatusEvent {
|
||||
@@ -559,13 +559,13 @@ fn finish_extraction_tool_use_events(call_id: &str) -> Vec<LlmEvent> {
|
||||
async fn compact_resets_extract_pointer_so_extract_can_fire_again() {
|
||||
// Mock LLM responses, in call order:
|
||||
// [0] first run with usage(1000) so extract threshold (=1) fires.
|
||||
// [1] extract worker invokes finish_extraction with empty output.
|
||||
// [1] extract worker invokes FinishMemoryExtraction with empty output.
|
||||
// [2] extract worker closes after the tool result.
|
||||
// [3] compact worker invokes write_summary.
|
||||
// [4] compact worker closes after the tool result.
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("hi", 1000),
|
||||
finish_extraction_tool_use_events("ec1"),
|
||||
finish_memory_extraction_tool_use_events("ec1"),
|
||||
single_text_events("done"),
|
||||
write_summary_tool_use_events("sc1", "summary"),
|
||||
single_text_events("done"),
|
||||
@@ -701,7 +701,7 @@ permission = "write"
|
||||
async fn extract_large_unprocessed_range_does_not_abort_on_input_occupancy() {
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("recorded", 1000),
|
||||
finish_extraction_tool_use_events("ec-large"),
|
||||
finish_memory_extraction_tool_use_events("ec-large"),
|
||||
single_text_events("done"),
|
||||
]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
|
||||
@@ -722,7 +722,7 @@ async fn extract_large_unprocessed_range_does_not_abort_on_input_occupancy() {
|
||||
async fn spawn_and_wait_drives_extract_to_completion() {
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("hi", 1000),
|
||||
finish_extraction_tool_use_events("ec1"),
|
||||
finish_memory_extraction_tool_use_events("ec1"),
|
||||
single_text_events("done"),
|
||||
]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
|
||||
@@ -750,7 +750,7 @@ async fn detached_extract_does_not_fork_session_log() {
|
||||
// `ensure_head_or_fork` does not spawn a new session.
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("hi", 1000),
|
||||
finish_extraction_tool_use_events("ec1"),
|
||||
finish_memory_extraction_tool_use_events("ec1"),
|
||||
single_text_events("done"),
|
||||
text_events_with_usage("ok", 1000),
|
||||
]);
|
||||
|
||||
Reference in New Issue
Block a user