worker: split session capture observation features
This commit is contained in:
@@ -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"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user