refactor: rename pod crate to worker
This commit is contained in:
@@ -0,0 +1,736 @@
|
||||
//! Durable active workflow invocation state.
|
||||
//!
|
||||
//! Workflow bodies are resolved at invocation time and snapshotted here. The
|
||||
//! snapshot, not whatever resource version is installed later, is the procedural
|
||||
//! authority that survives compaction for the currently governed task.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_engine::Item;
|
||||
use llm_engine::tool::{
|
||||
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use session_store::{LogEntry, SystemItem, segment_log};
|
||||
|
||||
pub const DOMAIN: &str = "worker.active_workflows";
|
||||
pub const REHYDRATION_MESSAGE_PREFIX: &str = "[Active workflow snapshot]";
|
||||
pub const INACTIVE_MESSAGE_PREFIX: &str = "[Active workflow state]";
|
||||
const SCHEMA_VERSION: u32 = 1;
|
||||
|
||||
pub type LogEntryCommitter = Arc<dyn Fn(LogEntry) + Send + Sync>;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ActiveWorkflowSnapshot {
|
||||
pub schema_version: u32,
|
||||
pub workflows: Vec<ActiveWorkflowRecord>,
|
||||
}
|
||||
|
||||
impl Default for ActiveWorkflowSnapshot {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
schema_version: SCHEMA_VERSION,
|
||||
workflows: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ActiveWorkflowRecord {
|
||||
pub slug: String,
|
||||
pub status: ActiveWorkflowStatus,
|
||||
pub invocation: WorkflowInvocationInfo,
|
||||
pub task_scope: String,
|
||||
pub body_snapshot_policy: WorkflowBodySnapshotPolicy,
|
||||
pub guidance_snapshot: String,
|
||||
pub obligations: Vec<String>,
|
||||
pub checkpoints: Vec<WorkflowCheckpoint>,
|
||||
pub updated_at_ms: u64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub completion: Option<WorkflowCompletionInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ActiveWorkflowStatus {
|
||||
Active,
|
||||
Completed,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ActiveWorkflowStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
let s = match self {
|
||||
Self::Active => "active",
|
||||
Self::Completed => "completed",
|
||||
Self::Cancelled => "cancelled",
|
||||
};
|
||||
f.write_str(s)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct WorkflowInvocationInfo {
|
||||
pub source: WorkflowInvocationSource,
|
||||
pub invoked_at_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum WorkflowInvocationSource {
|
||||
UserWorkflowInvokeSegment,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum WorkflowBodySnapshotPolicy {
|
||||
SnapshottedAtInvocation,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct WorkflowCheckpoint {
|
||||
pub label: String,
|
||||
pub status: WorkflowCheckpointStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum WorkflowCheckpointStatus {
|
||||
Open,
|
||||
Done,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct WorkflowCompletionInfo {
|
||||
pub completed_at_ms: u64,
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ActiveWorkflowStore {
|
||||
inner: Arc<Mutex<ActiveWorkflowSnapshot>>,
|
||||
}
|
||||
|
||||
impl ActiveWorkflowStore {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn snapshot(&self) -> ActiveWorkflowSnapshot {
|
||||
self.inner.lock().unwrap_or_else(|e| e.into_inner()).clone()
|
||||
}
|
||||
|
||||
pub fn replace_with(&self, snapshot: ActiveWorkflowSnapshot) {
|
||||
*self.inner.lock().unwrap_or_else(|e| e.into_inner()) = snapshot;
|
||||
}
|
||||
|
||||
pub fn active_records(&self) -> Vec<ActiveWorkflowRecord> {
|
||||
self.snapshot()
|
||||
.workflows
|
||||
.into_iter()
|
||||
.filter(|record| record.status == ActiveWorkflowStatus::Active)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn activate_from_system_items(
|
||||
&self,
|
||||
items: &[SystemItem],
|
||||
task_scope: String,
|
||||
invoked_at_ms: u64,
|
||||
) -> bool {
|
||||
let mut grouped: BTreeMap<String, Vec<String>> = BTreeMap::new();
|
||||
for item in items {
|
||||
if let SystemItem::Workflow { slug, body } = item {
|
||||
grouped.entry(slug.clone()).or_default().push(body.clone());
|
||||
}
|
||||
}
|
||||
if grouped.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let mut snapshot = self.snapshot();
|
||||
snapshot.schema_version = SCHEMA_VERSION;
|
||||
for (slug, bodies) in grouped {
|
||||
let guidance_snapshot = bodies.join("\n\n---\n\n");
|
||||
let obligations = extract_obligations(&guidance_snapshot);
|
||||
let checkpoints = obligations
|
||||
.iter()
|
||||
.take(32)
|
||||
.map(|label| WorkflowCheckpoint {
|
||||
label: label.clone(),
|
||||
status: WorkflowCheckpointStatus::Open,
|
||||
})
|
||||
.collect();
|
||||
let record = ActiveWorkflowRecord {
|
||||
slug: slug.clone(),
|
||||
status: ActiveWorkflowStatus::Active,
|
||||
invocation: WorkflowInvocationInfo {
|
||||
source: WorkflowInvocationSource::UserWorkflowInvokeSegment,
|
||||
invoked_at_ms,
|
||||
},
|
||||
task_scope: truncate_chars(&task_scope, 2_000),
|
||||
body_snapshot_policy: WorkflowBodySnapshotPolicy::SnapshottedAtInvocation,
|
||||
guidance_snapshot,
|
||||
obligations,
|
||||
checkpoints,
|
||||
updated_at_ms: invoked_at_ms,
|
||||
completion: None,
|
||||
};
|
||||
upsert_record(&mut snapshot.workflows, record);
|
||||
}
|
||||
self.replace_with(snapshot);
|
||||
true
|
||||
}
|
||||
|
||||
pub fn set_status(
|
||||
&self,
|
||||
slug: &str,
|
||||
status: ActiveWorkflowStatus,
|
||||
reason: String,
|
||||
now_ms: u64,
|
||||
) -> Result<ActiveWorkflowRecord, String> {
|
||||
let mut snapshot = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let record = snapshot
|
||||
.workflows
|
||||
.iter_mut()
|
||||
.find(|record| record.slug == slug)
|
||||
.ok_or_else(|| format!("active workflow `{slug}` not found"))?;
|
||||
record.status = status;
|
||||
record.updated_at_ms = now_ms;
|
||||
record.completion = Some(WorkflowCompletionInfo {
|
||||
completed_at_ms: now_ms,
|
||||
reason,
|
||||
});
|
||||
for checkpoint in &mut record.checkpoints {
|
||||
checkpoint.status = match status {
|
||||
ActiveWorkflowStatus::Active => WorkflowCheckpointStatus::Open,
|
||||
ActiveWorkflowStatus::Completed => WorkflowCheckpointStatus::Done,
|
||||
ActiveWorkflowStatus::Cancelled => WorkflowCheckpointStatus::Cancelled,
|
||||
};
|
||||
}
|
||||
Ok(record.clone())
|
||||
}
|
||||
|
||||
pub fn snapshot_text(&self) -> Option<String> {
|
||||
let active = self.active_records();
|
||||
(!active.is_empty()).then(|| render_snapshot_text(&active))
|
||||
}
|
||||
|
||||
pub fn rehydration_message(&self) -> Option<String> {
|
||||
let active = self.active_records();
|
||||
(!active.is_empty()).then(|| render_rehydration_message(&active))
|
||||
}
|
||||
|
||||
pub fn sanitize_context(&self, context: &mut Vec<Item>) -> usize {
|
||||
let removed = strip_rehydration_messages(context);
|
||||
if let Some(message) = self.rehydration_message() {
|
||||
context.push(Item::system_message(message));
|
||||
} else if removed > 0 || context.iter().any(has_active_workflow_hint) {
|
||||
context.push(Item::system_message(inactive_workflow_message()));
|
||||
}
|
||||
removed
|
||||
}
|
||||
|
||||
pub fn extension_entry(&self) -> LogEntry {
|
||||
LogEntry::Extension {
|
||||
ts: segment_log::now_millis(),
|
||||
domain: DOMAIN.into(),
|
||||
payload: serde_json::to_value(self.snapshot())
|
||||
.expect("ActiveWorkflowSnapshot is always JSON-serializable"),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn restore_from_history_and_extensions(
|
||||
&self,
|
||||
_history: &[Item],
|
||||
extensions: &[(String, serde_json::Value)],
|
||||
) {
|
||||
let (snapshot, diagnostics) = fold_extensions(extensions);
|
||||
for diagnostic in diagnostics {
|
||||
tracing::warn!(diagnostic, "failed to restore active workflow state");
|
||||
}
|
||||
self.replace_with(snapshot);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn fold_extensions(
|
||||
extensions: &[(String, serde_json::Value)],
|
||||
) -> (ActiveWorkflowSnapshot, Vec<String>) {
|
||||
let mut latest = None;
|
||||
let mut diagnostics = Vec::new();
|
||||
for (domain, payload) in extensions {
|
||||
if domain != DOMAIN {
|
||||
continue;
|
||||
}
|
||||
match serde_json::from_value::<ActiveWorkflowSnapshot>(payload.clone()) {
|
||||
Ok(snapshot) if snapshot.schema_version == SCHEMA_VERSION => latest = Some(snapshot),
|
||||
Ok(snapshot) => {
|
||||
latest = None;
|
||||
diagnostics.push(format!(
|
||||
"unsupported active workflow schema_version {}",
|
||||
snapshot.schema_version
|
||||
));
|
||||
}
|
||||
Err(err) => {
|
||||
latest = None;
|
||||
diagnostics.push(format!("corrupt active workflow payload: {err}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
(latest.unwrap_or_default(), diagnostics)
|
||||
}
|
||||
|
||||
pub fn strip_rehydration_messages(items: &mut Vec<Item>) -> usize {
|
||||
let before = items.len();
|
||||
items.retain(|item| !is_rehydration_message(item));
|
||||
before - items.len()
|
||||
}
|
||||
|
||||
pub fn is_rehydration_message(item: &Item) -> bool {
|
||||
item_system_text(item)
|
||||
.map(|text| text.trim_start().starts_with(REHYDRATION_MESSAGE_PREFIX))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn has_active_workflow_hint(item: &Item) -> bool {
|
||||
item_system_text(item)
|
||||
.map(|text| {
|
||||
text.contains("Active Workflow Invocation State")
|
||||
|| text.contains("ActiveWorkflowStore:")
|
||||
|| text.contains(REHYDRATION_MESSAGE_PREFIX)
|
||||
})
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn item_system_text(item: &Item) -> Option<String> {
|
||||
match item {
|
||||
Item::Message { role, content, .. } if *role == llm_engine::Role::System => Some(
|
||||
content
|
||||
.iter()
|
||||
.map(|part| part.as_text())
|
||||
.collect::<String>(),
|
||||
),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn inactive_workflow_message() -> String {
|
||||
format!(
|
||||
"{INACTIVE_MESSAGE_PREFIX}\n\n\
|
||||
No currently valid active workflow invocation state is active. Ignore older compacted \
|
||||
history or summaries that appear to describe active workflow obligations; only validated \
|
||||
typed `{DOMAIN}` records with status `active` establish active workflow guidance."
|
||||
)
|
||||
}
|
||||
|
||||
pub fn active_workflow_tools(
|
||||
store: ActiveWorkflowStore,
|
||||
committer: Option<LogEntryCommitter>,
|
||||
) -> Vec<ToolDefinition> {
|
||||
vec![
|
||||
list_tool(store.clone()),
|
||||
status_tool(
|
||||
store.clone(),
|
||||
ActiveWorkflowStatus::Completed,
|
||||
committer.clone(),
|
||||
),
|
||||
status_tool(store, ActiveWorkflowStatus::Cancelled, committer),
|
||||
]
|
||||
}
|
||||
|
||||
fn list_tool(store: ActiveWorkflowStore) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
(
|
||||
ToolMeta::new("ActiveWorkflowList")
|
||||
.description("List durable active workflow invocations and their status")
|
||||
.input_schema(
|
||||
json!({"type":"object","properties":{},"additionalProperties":false}),
|
||||
),
|
||||
Arc::new(ActiveWorkflowListTool {
|
||||
store: store.clone(),
|
||||
}) as Arc<dyn Tool>,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn status_tool(
|
||||
store: ActiveWorkflowStore,
|
||||
status: ActiveWorkflowStatus,
|
||||
committer: Option<LogEntryCommitter>,
|
||||
) -> ToolDefinition {
|
||||
let name = match status {
|
||||
ActiveWorkflowStatus::Completed => "ActiveWorkflowComplete",
|
||||
ActiveWorkflowStatus::Cancelled => "ActiveWorkflowCancel",
|
||||
ActiveWorkflowStatus::Active => unreachable!("active status tool is not exposed"),
|
||||
};
|
||||
let description = match status {
|
||||
ActiveWorkflowStatus::Completed => {
|
||||
"Mark an active workflow as completed when its governed task is finished"
|
||||
}
|
||||
ActiveWorkflowStatus::Cancelled => {
|
||||
"Cancel an active workflow when the governed task is explicitly abandoned"
|
||||
}
|
||||
ActiveWorkflowStatus::Active => unreachable!("active status tool is not exposed"),
|
||||
};
|
||||
let store_for_tool = store.clone();
|
||||
let committer_for_tool = committer.clone();
|
||||
Arc::new(move || {
|
||||
(
|
||||
ToolMeta::new(name)
|
||||
.description(description)
|
||||
.input_schema(json!({
|
||||
"type":"object",
|
||||
"properties":{
|
||||
"slug":{"type":"string","description":"Workflow slug to update"},
|
||||
"reason":{"type":"string","description":"Brief completion/cancellation reason"}
|
||||
},
|
||||
"required":["slug"],
|
||||
"additionalProperties":false
|
||||
})),
|
||||
Arc::new(ActiveWorkflowStatusTool {
|
||||
store: store_for_tool.clone(),
|
||||
status,
|
||||
committer: committer_for_tool.clone(),
|
||||
}) as Arc<dyn Tool>,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
struct ActiveWorkflowListTool {
|
||||
store: ActiveWorkflowStore,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ActiveWorkflowListTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
_input_json: &str,
|
||||
_ctx: ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let snapshot = self.store.snapshot();
|
||||
let content = serde_json::to_string_pretty(&snapshot)
|
||||
.map_err(|err| ToolError::Internal(err.to_string()))?;
|
||||
let active = snapshot
|
||||
.workflows
|
||||
.iter()
|
||||
.filter(|record| record.status == ActiveWorkflowStatus::Active)
|
||||
.count();
|
||||
Ok(ToolOutput {
|
||||
summary: format!(
|
||||
"ActiveWorkflowStore: {} workflow(s), {active} active",
|
||||
snapshot.workflows.len()
|
||||
),
|
||||
content: Some(content),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct ActiveWorkflowStatusTool {
|
||||
store: ActiveWorkflowStore,
|
||||
status: ActiveWorkflowStatus,
|
||||
committer: Option<LogEntryCommitter>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ActiveWorkflowStatusTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: WorkflowStatusParams = serde_json::from_str(input_json)
|
||||
.map_err(|err| ToolError::InvalidArgument(err.to_string()))?;
|
||||
let reason = params.reason.unwrap_or_else(|| self.status.to_string());
|
||||
let record = self
|
||||
.store
|
||||
.set_status(¶ms.slug, self.status, reason, segment_log::now_millis())
|
||||
.map_err(ToolError::InvalidArgument)?;
|
||||
if let Some(committer) = &self.committer {
|
||||
committer(self.store.extension_entry());
|
||||
}
|
||||
let content = serde_json::to_string_pretty(&record)
|
||||
.map_err(|err| ToolError::Internal(err.to_string()))?;
|
||||
Ok(ToolOutput {
|
||||
summary: format!("workflow {} marked {}", record.slug, record.status),
|
||||
content: Some(content),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct WorkflowStatusParams {
|
||||
slug: String,
|
||||
#[serde(default)]
|
||||
reason: Option<String>,
|
||||
}
|
||||
|
||||
fn upsert_record(records: &mut Vec<ActiveWorkflowRecord>, record: ActiveWorkflowRecord) {
|
||||
if let Some(existing) = records
|
||||
.iter_mut()
|
||||
.find(|existing| existing.slug == record.slug)
|
||||
{
|
||||
*existing = record;
|
||||
} else {
|
||||
records.push(record);
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_obligations(body: &str) -> Vec<String> {
|
||||
let mut obligations = Vec::new();
|
||||
for line in body.lines() {
|
||||
let trimmed = line.trim();
|
||||
let candidate = trimmed
|
||||
.strip_prefix("- ")
|
||||
.or_else(|| trimmed.strip_prefix("* "))
|
||||
.or_else(|| trimmed.strip_prefix("• "))
|
||||
.unwrap_or(trimmed);
|
||||
let lower = candidate.to_ascii_lowercase();
|
||||
let looks_obligating = lower.contains("must")
|
||||
|| lower.contains("require")
|
||||
|| lower.contains("obligation")
|
||||
|| lower.contains("review")
|
||||
|| lower.contains("merge")
|
||||
|| lower.contains("close")
|
||||
|| lower.contains("report")
|
||||
|| lower.contains("handoff");
|
||||
if looks_obligating && !candidate.is_empty() {
|
||||
obligations.push(truncate_chars(candidate, 240));
|
||||
}
|
||||
if obligations.len() >= 32 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if obligations.is_empty() {
|
||||
obligations
|
||||
.push("Follow the snapshotted workflow body until completion or cancellation".into());
|
||||
}
|
||||
obligations
|
||||
}
|
||||
|
||||
fn render_snapshot_text(records: &[ActiveWorkflowRecord]) -> String {
|
||||
let json = serde_json::to_string_pretty(&ActiveWorkflowSnapshot {
|
||||
schema_version: SCHEMA_VERSION,
|
||||
workflows: records.to_vec(),
|
||||
})
|
||||
.unwrap_or_else(|_| String::from("{\"schema_version\":1,\"workflows\":[]}"));
|
||||
format!(
|
||||
"ActiveWorkflowStore: {} active workflow(s)\n\n```json\n{}\n```",
|
||||
records.len(),
|
||||
json
|
||||
)
|
||||
}
|
||||
|
||||
fn render_rehydration_message(records: &[ActiveWorkflowRecord]) -> String {
|
||||
let mut out = format!(
|
||||
"{REHYDRATION_MESSAGE_PREFIX}\n\n\
|
||||
The following workflow invocation state is durable state carried across compaction. \
|
||||
Continue to follow each active workflow's snapshotted guidance until the governed task \
|
||||
is completed with ActiveWorkflowComplete or explicitly cancelled with ActiveWorkflowCancel. \
|
||||
Missing or obsolete workflow resources must not replace these invocation snapshots.\n"
|
||||
);
|
||||
for record in records {
|
||||
out.push_str(&format!(
|
||||
"\n## /{} ({})\n- invoked_at_ms: {}\n- invocation_source: {:?}\n- body_snapshot_policy: {:?}\n- task_scope: {}\n\n### Current obligations/checkpoints\n",
|
||||
record.slug,
|
||||
record.status,
|
||||
record.invocation.invoked_at_ms,
|
||||
record.invocation.source,
|
||||
record.body_snapshot_policy,
|
||||
record.task_scope.replace('\n', " "),
|
||||
));
|
||||
for checkpoint in &record.checkpoints {
|
||||
out.push_str(&format!(
|
||||
"- [{}] {}\n",
|
||||
checkpoint.status_label(),
|
||||
checkpoint.label
|
||||
));
|
||||
}
|
||||
out.push_str("\n### Snapshotted workflow guidance\n");
|
||||
out.push_str(record.guidance_snapshot.trim_end());
|
||||
out.push_str("\n");
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
impl WorkflowCheckpoint {
|
||||
fn status_label(&self) -> &'static str {
|
||||
match self.status {
|
||||
WorkflowCheckpointStatus::Open => "open",
|
||||
WorkflowCheckpointStatus::Done => "done",
|
||||
WorkflowCheckpointStatus::Cancelled => "cancelled",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate_chars(text: &str, max_chars: usize) -> String {
|
||||
let mut out = String::new();
|
||||
for (idx, ch) in text.chars().enumerate() {
|
||||
if idx >= max_chars {
|
||||
out.push('…');
|
||||
return out;
|
||||
}
|
||||
out.push(ch);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn store_with_active_workflow() -> ActiveWorkflowStore {
|
||||
let store = ActiveWorkflowStore::new();
|
||||
assert!(store.activate_from_system_items(
|
||||
&[SystemItem::Workflow {
|
||||
slug: "multi-agent-workflow".into(),
|
||||
body: "# Multi-agent workflow\n- Delegate implementation to coder.\n- Require external review before merge.\n- Close the Ticket after merge and report evidence.\n".into(),
|
||||
}],
|
||||
"/multi-agent-workflow implement ticket".into(),
|
||||
42,
|
||||
));
|
||||
store
|
||||
}
|
||||
|
||||
fn active_extension(store: &ActiveWorkflowStore) -> (String, serde_json::Value) {
|
||||
(
|
||||
DOMAIN.to_string(),
|
||||
serde_json::to_value(store.snapshot()).expect("snapshot json"),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn active_workflow_guidance_carries_merge_close_obligations() {
|
||||
let store = store_with_active_workflow();
|
||||
let msg = store.rehydration_message().unwrap();
|
||||
|
||||
assert!(msg.contains("multi-agent-workflow"));
|
||||
assert!(msg.contains("external review before merge"));
|
||||
assert!(msg.contains("Close the Ticket after merge"));
|
||||
assert!(msg.contains("Snapshotted workflow guidance"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compacted_rehydration_message_is_removed_when_typed_state_missing_or_invalid() {
|
||||
for extensions in [
|
||||
Vec::new(),
|
||||
vec![(DOMAIN.to_string(), json!({"schema_version":"bad"}))],
|
||||
vec![(
|
||||
DOMAIN.to_string(),
|
||||
json!({"schema_version":999,"workflows":[]}),
|
||||
)],
|
||||
] {
|
||||
let original = store_with_active_workflow();
|
||||
let stale_message = original.rehydration_message().unwrap();
|
||||
let mut context = vec![
|
||||
Item::system_message(stale_message),
|
||||
Item::user_message("continue"),
|
||||
];
|
||||
let restored = ActiveWorkflowStore::new();
|
||||
|
||||
restored.restore_from_history_and_extensions(&context, &extensions);
|
||||
let removed = restored.sanitize_context(&mut context);
|
||||
|
||||
assert_eq!(removed, 1);
|
||||
assert!(restored.active_records().is_empty());
|
||||
assert!(!context.iter().any(is_rehydration_message));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completion_or_cancellation_suppresses_old_compacted_guidance() {
|
||||
for status in [
|
||||
ActiveWorkflowStatus::Completed,
|
||||
ActiveWorkflowStatus::Cancelled,
|
||||
] {
|
||||
let store = store_with_active_workflow();
|
||||
let stale_message = store.rehydration_message().unwrap();
|
||||
let mut context = vec![
|
||||
Item::system_message(stale_message),
|
||||
Item::user_message("continue"),
|
||||
];
|
||||
|
||||
store
|
||||
.set_status("multi-agent-workflow", status, status.to_string(), 84)
|
||||
.expect("workflow exists");
|
||||
let removed = store.sanitize_context(&mut context);
|
||||
|
||||
assert_eq!(removed, 1);
|
||||
assert!(!context.iter().any(is_rehydration_message));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unmatched_status_tool_calls_do_not_mutate_restored_state() {
|
||||
let store = store_with_active_workflow();
|
||||
let extensions = vec![active_extension(&store)];
|
||||
let history = vec![
|
||||
Item::tool_call(
|
||||
"call-1",
|
||||
"ActiveWorkflowCancel",
|
||||
json!({"slug":"multi-agent-workflow","reason":"not durable"}).to_string(),
|
||||
),
|
||||
Item::tool_result_error("call-1", "error: failed"),
|
||||
];
|
||||
let restored = ActiveWorkflowStore::new();
|
||||
|
||||
restored.restore_from_history_and_extensions(&history, &extensions);
|
||||
|
||||
assert_eq!(restored.active_records().len(), 1);
|
||||
assert_eq!(
|
||||
restored.snapshot().workflows[0].status,
|
||||
ActiveWorkflowStatus::Active
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn status_tool_persists_typed_extension_on_success() {
|
||||
let store = store_with_active_workflow();
|
||||
let committed = Arc::new(Mutex::new(Vec::<LogEntry>::new()));
|
||||
let committed_for_tool = committed.clone();
|
||||
let tools = active_workflow_tools(
|
||||
store.clone(),
|
||||
Some(Arc::new(move |entry| {
|
||||
committed_for_tool
|
||||
.lock()
|
||||
.expect("committed entries mutex poisoned")
|
||||
.push(entry);
|
||||
})),
|
||||
);
|
||||
let (_, tool) = tools[1]();
|
||||
|
||||
tool.execute(
|
||||
&json!({"slug":"multi-agent-workflow","reason":"review complete"}).to_string(),
|
||||
ToolExecutionContext::default(),
|
||||
)
|
||||
.await
|
||||
.expect("status tool succeeds");
|
||||
|
||||
let committed = committed.lock().expect("committed entries mutex poisoned");
|
||||
let LogEntry::Extension {
|
||||
domain, payload, ..
|
||||
} = committed.last().expect("extension committed")
|
||||
else {
|
||||
panic!("expected typed active workflow extension");
|
||||
};
|
||||
assert_eq!(domain, DOMAIN);
|
||||
let snapshot: ActiveWorkflowSnapshot = serde_json::from_value(payload.clone()).unwrap();
|
||||
assert_eq!(
|
||||
snapshot.workflows[0].status,
|
||||
ActiveWorkflowStatus::Completed
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn corrupt_extension_fails_closed_with_diagnostic() {
|
||||
let entries = vec![(DOMAIN.to_string(), json!({"schema_version":"bad"}))];
|
||||
|
||||
let (snapshot, diagnostics) = fold_extensions(&entries);
|
||||
|
||||
assert!(snapshot.workflows.is_empty());
|
||||
assert_eq!(diagnostics.len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
//! Sync buffer for `session_metrics::Metric` values queued from inside
|
||||
//! Engine callbacks (which run synchronously and cannot themselves
|
||||
//! perform `async` store writes).
|
||||
//!
|
||||
//! Worker drains this buffer in `persist_turn` and writes each metric via
|
||||
//! `session_metrics::record_metric`, alongside the regular `LlmUsage`
|
||||
//! entries.
|
||||
|
||||
use std::sync::Mutex;
|
||||
|
||||
use session_metrics::Metric;
|
||||
|
||||
pub(crate) struct MetricsTracker {
|
||||
pending: Mutex<Vec<Metric>>,
|
||||
}
|
||||
|
||||
impl MetricsTracker {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
pending: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Queue a metric for the next `persist_turn` flush.
|
||||
pub(crate) fn push(&self, metric: Metric) {
|
||||
self.pending.lock().unwrap().push(metric);
|
||||
}
|
||||
|
||||
/// Drain all queued metrics. Called by Worker after a run completes.
|
||||
pub(crate) fn drain(&self) -> Vec<Metric> {
|
||||
std::mem::take(&mut *self.pending.lock().unwrap())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn push_then_drain_returns_in_order_and_clears() {
|
||||
let t = MetricsTracker::new();
|
||||
t.push(Metric::now("a"));
|
||||
t.push(Metric::now("b"));
|
||||
let drained = t.drain();
|
||||
assert_eq!(drained.len(), 2);
|
||||
assert_eq!(drained[0].name, "a");
|
||||
assert_eq!(drained[1].name, "b");
|
||||
assert!(t.drain().is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
pub(crate) mod metrics_tracker;
|
||||
pub(crate) mod prune;
|
||||
pub(crate) mod state;
|
||||
pub(crate) mod token_counter;
|
||||
pub(crate) mod usage_tracker;
|
||||
pub(crate) mod worker;
|
||||
@@ -0,0 +1,124 @@
|
||||
//! Prune integration — wires the Engine's prune projection to the Worker's
|
||||
//! usage-history-backed token accounting.
|
||||
//!
|
||||
//! Engine 自身がコンテキスト射影を行う（`worker.rs` の `request_context` 構築
|
||||
//! 直後）。Engine は usage 履歴を知らないので、`min_savings` 判定に使う savings
|
||||
//! の見積もりはコールバックで外部から注入する。このモジュールはそのコールバック
|
||||
//! を組み立てて Engine に差し込むための `impl Worker` を提供する。
|
||||
//!
|
||||
//! 同じ経路で `PruneObserver` も install し、評価のたびに `prune.fire` /
|
||||
//! `prune.skip` metric を `MetricsTracker` に積む。`Fired` 時は uuid を
|
||||
//! `UsageTracker` にも stash しておき、後続の `LlmUsage` と組で
|
||||
//! `prune.post_request` を吐けるようにする。
|
||||
|
||||
use llm_engine::Item;
|
||||
use llm_engine::llm_client::client::LlmClient;
|
||||
use llm_engine::prune::{
|
||||
PruneConfig, PruneDecision, PruneObserver, SavingsEstimator, TokenEstimator,
|
||||
};
|
||||
use session_metrics::Metric;
|
||||
use session_store::Store;
|
||||
|
||||
use crate::Worker;
|
||||
use crate::compact::token_counter::{
|
||||
EstimateSource, savings_for_prune_impl, token_estimates_for_prune_impl,
|
||||
};
|
||||
|
||||
impl<C: LlmClient, St: Store> Worker<C, St> {
|
||||
/// Enable prune projection on the underlying Engine.
|
||||
///
|
||||
/// Registers the config and token/savings-estimator closures on the Engine.
|
||||
/// The estimators combine persisted [`Worker::usage_history_handle`] records
|
||||
/// with in-flight `UsageTracker` records so multi-request tool loops can
|
||||
/// prune before the surrounding Worker run finishes.
|
||||
///
|
||||
/// Measurement-less estimates (before the first LLM call, or immediately
|
||||
/// after a compact) return `0` from the estimator, which naturally
|
||||
/// prevents the prune projection from firing until usage data exists.
|
||||
///
|
||||
/// Also installs a [`PruneObserver`] that pushes `prune.fire` /
|
||||
/// `prune.skip` metrics into the shared [`MetricsTracker`]. On `Fired`
|
||||
/// the observer additionally stashes a fresh correlation_id in
|
||||
/// [`UsageTracker`] so the next `LlmUsage` can be paired with a
|
||||
/// `prune.post_request` metric carrying the same id.
|
||||
pub fn attach_prune(&mut self, config: PruneConfig) {
|
||||
let usage_history_for_tokens = self.usage_history_handle();
|
||||
let usage_tracker_for_tokens = self.usage_tracker_handle();
|
||||
let token_estimator: TokenEstimator = Box::new(move |history: &[Item]| {
|
||||
let mut snapshot = usage_history_for_tokens
|
||||
.lock()
|
||||
.expect("usage_history poisoned")
|
||||
.clone();
|
||||
snapshot.extend(usage_tracker_for_tokens.records());
|
||||
token_estimates_for_prune_impl(history, &snapshot)
|
||||
});
|
||||
|
||||
let usage_history_for_savings = self.usage_history_handle();
|
||||
let usage_tracker_for_savings = self.usage_tracker_handle();
|
||||
let estimator: SavingsEstimator = Box::new(move |history: &[Item], indices| {
|
||||
let mut snapshot = usage_history_for_savings
|
||||
.lock()
|
||||
.expect("usage_history poisoned")
|
||||
.clone();
|
||||
snapshot.extend(usage_tracker_for_savings.records());
|
||||
let est = savings_for_prune_impl(history, &snapshot, indices);
|
||||
match est.source {
|
||||
EstimateSource::NoData => 0,
|
||||
_ => est.tokens,
|
||||
}
|
||||
});
|
||||
|
||||
let metrics = self.metrics_tracker_handle();
|
||||
let usage_tracker = self.usage_tracker_handle();
|
||||
let observer: PruneObserver = Box::new(move |eval| match &eval.decision {
|
||||
PruneDecision::Fired { .. } => {
|
||||
let correlation_id = uuid::Uuid::now_v7().to_string();
|
||||
let mut metric = Metric::now("prune.fire")
|
||||
.with_value(eval.estimated_savings as f64)
|
||||
.with_correlation_id(&correlation_id)
|
||||
.with_dimension("candidate_count", eval.candidate_count.to_string());
|
||||
if let Some(protected_start) = eval.protected_start_index {
|
||||
metric =
|
||||
metric.with_dimension("protected_start_index", protected_start.to_string());
|
||||
}
|
||||
metrics.push(metric);
|
||||
usage_tracker.note_correlation_id(correlation_id);
|
||||
}
|
||||
PruneDecision::SkippedNoCandidates => {
|
||||
metrics.push(Metric::now("prune.skip").with_dimension("reason", "no_candidates"));
|
||||
}
|
||||
PruneDecision::SkippedBelowMinSavings => {
|
||||
let mut metric = Metric::now("prune.skip")
|
||||
.with_dimension("reason", "below_min_savings")
|
||||
.with_dimension("candidate_count", eval.candidate_count.to_string())
|
||||
.with_value(eval.estimated_savings as f64);
|
||||
if let Some(protected_start) = eval.protected_start_index {
|
||||
metric =
|
||||
metric.with_dimension("protected_start_index", protected_start.to_string());
|
||||
}
|
||||
metrics.push(metric);
|
||||
}
|
||||
});
|
||||
|
||||
let worker = self.engine_mut();
|
||||
worker.set_prune_config(Some(config));
|
||||
worker.set_token_estimator(Some(token_estimator));
|
||||
worker.set_savings_estimator(Some(estimator));
|
||||
worker.set_prune_observer(Some(observer));
|
||||
}
|
||||
|
||||
/// If the manifest has a `[compaction]` section, build a `PruneConfig`
|
||||
/// from its `prune_*` fields and call [`attach_prune`](Self::attach_prune).
|
||||
/// Otherwise no-op. Called from all Worker constructors so prune is
|
||||
/// active whenever the manifest asks for it.
|
||||
pub(crate) fn apply_prune_from_manifest(&mut self) {
|
||||
let Some(compaction) = self.manifest().compaction.as_ref() else {
|
||||
return;
|
||||
};
|
||||
let config = PruneConfig {
|
||||
protected_tokens: compaction.prune_protected_tokens,
|
||||
min_savings: compaction.prune_min_savings,
|
||||
};
|
||||
self.attach_prune(config);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
//! Shared state for compaction decisions.
|
||||
//!
|
||||
//! Holds the two configured thresholds and circuit-breaker / thrash-detection
|
||||
//! flags shared between:
|
||||
//! - `WorkerInterceptor` (reads `request_threshold` — the *safety net* for
|
||||
//! between-requests yielding)
|
||||
//! - `Worker::try_pre_run_compact` (reads `post_run_threshold` — the
|
||||
//! *proactive* check before the next turn starts)
|
||||
//! - `Worker::run()` / `resume()` (circuit breaker, thrash detection)
|
||||
//!
|
||||
//! Current occupancy (input-token count) is **not** stored here. The single
|
||||
//! source of truth is `session_store::UsageRecord` (persisted per LLM call)
|
||||
//! projected through `Worker::total_tokens()`. Callers pass the current
|
||||
//! occupancy to `exceeds_*` at check time.
|
||||
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
|
||||
const MAX_COMPACT_FAILURES: usize = 3;
|
||||
|
||||
/// Shared mutable state for compaction decisions.
|
||||
pub(crate) struct CompactState {
|
||||
/// Between-turns threshold (proactive). Checked before the next turn
|
||||
/// starts. `None` disables the pre-run check.
|
||||
post_run_threshold: Option<u64>,
|
||||
/// Between-requests threshold (safety net). Checked inside a turn
|
||||
/// before each LLM request. `None` disables the request check.
|
||||
request_threshold: Option<u64>,
|
||||
/// Token budget retained verbatim at the tail after compaction.
|
||||
retained_tokens: u64,
|
||||
/// Consecutive compact failures. At `MAX_COMPACT_FAILURES`, compaction is disabled.
|
||||
consecutive_failures: AtomicUsize,
|
||||
/// `true` immediately after a successful compact, cleared on next normal completion.
|
||||
just_compacted: AtomicBool,
|
||||
/// `true` when circuit breaker has tripped.
|
||||
disabled: AtomicBool,
|
||||
}
|
||||
|
||||
impl CompactState {
|
||||
pub(crate) fn new(
|
||||
post_run_threshold: Option<u64>,
|
||||
request_threshold: Option<u64>,
|
||||
retained_tokens: u64,
|
||||
) -> Self {
|
||||
Self {
|
||||
post_run_threshold,
|
||||
request_threshold,
|
||||
retained_tokens,
|
||||
consecutive_failures: AtomicUsize::new(0),
|
||||
just_compacted: AtomicBool::new(false),
|
||||
disabled: AtomicBool::new(false),
|
||||
}
|
||||
}
|
||||
|
||||
/// Configured between-requests threshold (if any).
|
||||
pub(crate) fn request_threshold(&self) -> Option<u64> {
|
||||
self.request_threshold
|
||||
}
|
||||
|
||||
/// Token budget retained verbatim at the tail after compaction.
|
||||
pub(crate) fn retained_tokens(&self) -> u64 {
|
||||
self.retained_tokens
|
||||
}
|
||||
|
||||
/// Whether compaction has been disabled by the circuit breaker.
|
||||
pub(crate) fn is_disabled(&self) -> bool {
|
||||
self.disabled.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Whether `current_tokens` exceeds the between-requests threshold.
|
||||
/// Returns `false` when `request_threshold` is unset.
|
||||
pub(crate) fn exceeds_request(&self, current_tokens: u64) -> bool {
|
||||
self.request_threshold
|
||||
.map(|t| current_tokens > t)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Whether `current_tokens` exceeds the post-run threshold.
|
||||
/// Returns `false` when `post_run_threshold` is unset.
|
||||
pub(crate) fn exceeds_post_run(&self, current_tokens: u64) -> bool {
|
||||
self.post_run_threshold
|
||||
.map(|t| current_tokens > t)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// Whether a compact just completed (for thrash detection).
|
||||
pub(crate) fn just_compacted(&self) -> bool {
|
||||
self.just_compacted.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Set or clear the just_compacted flag.
|
||||
pub(crate) fn set_just_compacted(&self, val: bool) {
|
||||
self.just_compacted.store(val, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Record a successful compaction: reset failure counter, set just_compacted.
|
||||
pub(crate) fn record_compact_success(&self) {
|
||||
self.consecutive_failures.store(0, Ordering::Relaxed);
|
||||
self.just_compacted.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Record a compaction failure. Disables compaction after MAX_COMPACT_FAILURES.
|
||||
pub(crate) fn record_compact_failure(&self) {
|
||||
let prev = self.consecutive_failures.fetch_add(1, Ordering::Relaxed);
|
||||
if prev + 1 >= MAX_COMPACT_FAILURES {
|
||||
self.disabled.store(true, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn both_thresholds_configured() {
|
||||
let state = CompactState::new(Some(80_000), Some(90_000), 8_000);
|
||||
assert_eq!(state.request_threshold(), Some(90_000));
|
||||
assert_eq!(state.retained_tokens(), 8_000);
|
||||
|
||||
assert!(!state.exceeds_request(70_000));
|
||||
assert!(!state.exceeds_post_run(70_000));
|
||||
|
||||
assert!(!state.exceeds_request(85_000));
|
||||
assert!(state.exceeds_post_run(85_000));
|
||||
|
||||
assert!(state.exceeds_request(95_000));
|
||||
assert!(state.exceeds_post_run(95_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn post_run_only() {
|
||||
let state = CompactState::new(Some(80_000), None, 8_000);
|
||||
// request check always false when threshold is None.
|
||||
assert!(!state.exceeds_request(1_000_000));
|
||||
assert!(state.exceeds_post_run(85_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_only() {
|
||||
let state = CompactState::new(None, Some(90_000), 8_000);
|
||||
assert!(!state.exceeds_post_run(1_000_000));
|
||||
assert!(state.exceeds_request(95_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn both_none_disables_all_checks() {
|
||||
let state = CompactState::new(None, None, 8_000);
|
||||
assert!(!state.exceeds_request(1_000_000));
|
||||
assert!(!state.exceeds_post_run(1_000_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn circuit_breaker_trips_after_max_failures() {
|
||||
let state = CompactState::new(Some(80_000), Some(90_000), 8_000);
|
||||
assert!(!state.is_disabled());
|
||||
|
||||
state.record_compact_failure();
|
||||
assert!(!state.is_disabled());
|
||||
state.record_compact_failure();
|
||||
assert!(!state.is_disabled());
|
||||
state.record_compact_failure();
|
||||
assert!(state.is_disabled());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn success_resets_failure_count() {
|
||||
let state = CompactState::new(Some(80_000), Some(90_000), 8_000);
|
||||
state.record_compact_failure();
|
||||
state.record_compact_failure();
|
||||
assert!(!state.is_disabled());
|
||||
|
||||
state.record_compact_success();
|
||||
assert!(state.just_compacted());
|
||||
|
||||
state.record_compact_failure();
|
||||
state.record_compact_failure();
|
||||
assert!(!state.is_disabled());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn just_compacted_lifecycle() {
|
||||
let state = CompactState::new(Some(80_000), Some(90_000), 8_000);
|
||||
assert!(!state.just_compacted());
|
||||
|
||||
state.record_compact_success();
|
||||
assert!(state.just_compacted());
|
||||
|
||||
state.set_just_compacted(false);
|
||||
assert!(!state.just_compacted());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,557 @@
|
||||
//! Compact / prune 専用のトークン会計補助。
|
||||
//!
|
||||
//! 汎用部分（`prefix_bytes`, `tokens_at`, `total_tokens`, `total_tokens_at`）は
|
||||
//! [`llm_engine::token_counter`] にあり、`UsageRecord` の列と現在の history から
|
||||
//! pure に推定する。本モジュールは compact / prune 固有のロジック
|
||||
//! （`split_for_retained`, `savings_for_prune`）と、Worker 上の公開 API に
|
||||
//! 限定する。
|
||||
//!
|
||||
//! # 方針
|
||||
//!
|
||||
//! - ローカルトークナイザは持たない。実測値があればそれを採用し、
|
||||
//! measurement 間はバイト数で按分、最新 measurement より先は byte/4 で外挿する
|
||||
//! - Compact の retained split では、request-time pruning / projection 後の
|
||||
//! `UsageRecord` を persisted history prefix の単調系列として扱わない。
|
||||
//! 現在の prompt occupancy 推定を raw serialized bytes に配分し、末尾の
|
||||
//! persisted tail サイズで cut を決める。
|
||||
//! - 推定の出どころは [`EstimateSource`] で呼び出し側に明示する。
|
||||
//! 課金判断には使えないが、compact / prune の閾値判定には十分な精度
|
||||
|
||||
use llm_engine::llm_client::client::LlmClient;
|
||||
use llm_engine::token_counter::{item_bytes, prefix_bytes, tokens_at};
|
||||
use llm_engine::{Item, UsageRecord};
|
||||
use session_store::Store;
|
||||
|
||||
pub use llm_engine::token_counter::{EstimateSource, TokenEstimate};
|
||||
|
||||
use crate::Worker;
|
||||
|
||||
/// history を分割する位置。
|
||||
///
|
||||
/// `items[..index]` が捨てる／要約される側、`items[index..]` が残る側。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct SplitPoint {
|
||||
pub index: usize,
|
||||
pub source: EstimateSource,
|
||||
}
|
||||
|
||||
fn split_for_retained_impl(history: &[Item], records: &[UsageRecord], retained: u64) -> SplitPoint {
|
||||
let prefix = prefix_bytes(history);
|
||||
let current = tokens_at(history, records, history.len(), &prefix);
|
||||
if current.tokens <= retained {
|
||||
return SplitPoint {
|
||||
index: 0,
|
||||
source: current.source,
|
||||
};
|
||||
}
|
||||
|
||||
let cut_index = split_index_by_retained_bytes(&prefix, current.tokens, retained);
|
||||
SplitPoint {
|
||||
index: balance_to_pair_boundary(history, cut_index),
|
||||
source: current.source,
|
||||
}
|
||||
}
|
||||
|
||||
fn split_index_by_retained_bytes(prefix: &[u64], total_tokens: u64, retained_tokens: u64) -> usize {
|
||||
debug_assert!(!prefix.is_empty());
|
||||
|
||||
let len = prefix.len() - 1;
|
||||
if len == 0 {
|
||||
return 0;
|
||||
}
|
||||
if retained_tokens == 0 {
|
||||
return len;
|
||||
}
|
||||
|
||||
let total_bytes = *prefix.last().unwrap_or(&0);
|
||||
if total_bytes == 0 || total_tokens == 0 {
|
||||
return 0;
|
||||
}
|
||||
|
||||
let raw_fallback_tokens = ceil_div_u128(total_bytes as u128, 4) as u64;
|
||||
let rate_tokens = total_tokens.max(raw_fallback_tokens);
|
||||
let target_retained_bytes = ceil_div_u128(
|
||||
retained_tokens as u128 * total_bytes as u128,
|
||||
rate_tokens as u128,
|
||||
)
|
||||
.min(total_bytes as u128) as u64;
|
||||
|
||||
// Drop as many complete Items as possible while keeping the raw persisted
|
||||
// suffix at or above the retained budget. This is monotonic in serialized
|
||||
// history size and intentionally does not inspect per-history_len
|
||||
// UsageRecords: request-time usage can move up and down after pruning /
|
||||
// projection, so it is not a valid prefix series for retained split. The
|
||||
// byte/4 fallback is kept as a lower bound for raw persisted size so a
|
||||
// heavily-pruned request measurement cannot justify retaining megabytes of
|
||||
// history.
|
||||
let mut cut = 0;
|
||||
for (idx, bytes_before) in prefix.iter().enumerate().take(len + 1) {
|
||||
let suffix_bytes = total_bytes.saturating_sub(*bytes_before);
|
||||
if suffix_bytes >= target_retained_bytes {
|
||||
cut = idx;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
cut
|
||||
}
|
||||
|
||||
fn ceil_div_u128(n: u128, d: u128) -> u128 {
|
||||
debug_assert!(d > 0);
|
||||
if n == 0 { 0 } else { ((n - 1) / d) + 1 }
|
||||
}
|
||||
|
||||
/// `history[cut..]` が `ToolCall` / `ToolResult` のペア境界を尊重するよう
|
||||
/// `cut` を後退させる。
|
||||
///
|
||||
/// LLM API は「`ToolResult` を送るならその `ToolCall` も同じ request に
|
||||
/// 含まれていなければならない」というバリデーションを持つ。トークン数
|
||||
/// だけで切った `cut` は並列 tool 呼び出しの途中に落ちうるので、retained
|
||||
/// 側の先頭に対応 `ToolCall` を持たない `ToolResult`（orphan）が残ると
|
||||
/// 次セッション初回 request が API バリデーションで弾かれる。
|
||||
///
|
||||
/// 対策は「retained に入る `ToolResult` について、対応 `ToolCall` も
|
||||
/// retained に含まれる位置まで `cut` を引き下げる」こと。retained_tokens
|
||||
/// 予算は超えうるが、ここでは直接 LLM に投げる訳ではなく次の
|
||||
/// `pre_llm_request` で再評価されるだけなので safe。
|
||||
///
|
||||
/// アルゴリズム: history を末尾から走査し、retained 範囲内の `ToolResult`
|
||||
/// に出会うたびに対応 `ToolCall` の位置で `cut` を min 更新する。`cut` が
|
||||
/// 下がると以前は要約側だった位置が retained に入るので、後続走査で連鎖的
|
||||
/// に正しい位置まで引き下がる。`ToolCall` の `call_id` はユニークなので
|
||||
/// 事前にマップ化して O(n) で済ます。
|
||||
fn balance_to_pair_boundary(history: &[Item], cut: usize) -> usize {
|
||||
let mut idx = cut.min(history.len());
|
||||
if idx == 0 {
|
||||
return 0;
|
||||
}
|
||||
|
||||
let call_positions: std::collections::HashMap<&str, usize> = history
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(i, item)| match item {
|
||||
Item::ToolCall { call_id, .. } => Some((call_id.as_str(), i)),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let mut k = history.len();
|
||||
while k > 0 {
|
||||
k -= 1;
|
||||
if k >= idx {
|
||||
if let Item::ToolResult { call_id, .. } = &history[k] {
|
||||
if let Some(&call_pos) = call_positions.get(call_id.as_str()) {
|
||||
if call_pos < idx {
|
||||
idx = call_pos;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
idx
|
||||
}
|
||||
|
||||
/// 1 つの ToolResult 項目について、`content` を `None` に射影したとき
|
||||
/// 減少するシリアライズ後バイト数。ToolResult 以外や既に content=None
|
||||
/// の item は 0 を返す。
|
||||
fn tool_result_content_bytes(item: &Item) -> u64 {
|
||||
if !matches!(
|
||||
item,
|
||||
Item::ToolResult {
|
||||
content: Some(_),
|
||||
..
|
||||
}
|
||||
) {
|
||||
return 0;
|
||||
}
|
||||
let mut cleared = item.clone();
|
||||
if let Item::ToolResult { content, .. } = &mut cleared {
|
||||
*content = None;
|
||||
}
|
||||
item_bytes(item).saturating_sub(item_bytes(&cleared))
|
||||
}
|
||||
|
||||
/// Prefix-boundary token estimates used by Prune to find its protected suffix.
|
||||
///
|
||||
/// Returns `history.len() + 1` entries where entry `i` estimates
|
||||
/// `history[..i]`. This shares the same [`tokens_at`] accounting as compact's
|
||||
/// retained-tail split and prune's savings estimate.
|
||||
pub(crate) fn token_estimates_for_prune_impl(
|
||||
history: &[Item],
|
||||
records: &[UsageRecord],
|
||||
) -> Vec<TokenEstimate> {
|
||||
let prefix = prefix_bytes(history);
|
||||
(0..=history.len())
|
||||
.map(|idx| tokens_at(history, records, idx, &prefix))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Prune 射影（`ToolResult.content = None`）で節約されるトークン数の推定。
|
||||
///
|
||||
/// `indices` は [`llm_engine::prune::prunable_indices`] が返す候補列を
|
||||
/// 想定する。各候補の content バイト差分を合算し、usage 履歴由来の
|
||||
/// tokens/byte レートでトークン数に換算する。範囲を「丸ごと drop」する
|
||||
/// のではなく、item 自体（summary 等）は残したままの値を返す点が
|
||||
/// `tokens_at` ベースの計算と異なる。
|
||||
pub(crate) fn savings_for_prune_impl(
|
||||
history: &[Item],
|
||||
records: &[UsageRecord],
|
||||
indices: &[usize],
|
||||
) -> TokenEstimate {
|
||||
let removed_bytes: u64 = indices
|
||||
.iter()
|
||||
.filter_map(|&i| history.get(i))
|
||||
.map(tool_result_content_bytes)
|
||||
.sum();
|
||||
|
||||
if removed_bytes == 0 {
|
||||
return TokenEstimate {
|
||||
tokens: 0,
|
||||
source: EstimateSource::Measured,
|
||||
};
|
||||
}
|
||||
|
||||
if records.is_empty() {
|
||||
return TokenEstimate {
|
||||
tokens: removed_bytes / 4,
|
||||
source: EstimateSource::NoData,
|
||||
};
|
||||
}
|
||||
|
||||
// 最新の measurement を使って tokens/byte を求め、バイト差分を換算する。
|
||||
// 実測値そのものではなく比率しか使わないので、history_len と
|
||||
// record.history_len が一致しなくても rate は正しい。
|
||||
let prefix = prefix_bytes(history);
|
||||
let last = records.last().expect("records non-empty");
|
||||
let ref_bytes = prefix[last.history_len.min(history.len())];
|
||||
if ref_bytes == 0 || last.input_total_tokens == 0 {
|
||||
return TokenEstimate {
|
||||
tokens: 0,
|
||||
source: EstimateSource::Extrapolated,
|
||||
};
|
||||
}
|
||||
let tokens =
|
||||
(removed_bytes as u128 * last.input_total_tokens as u128 / ref_bytes as u128) as u64;
|
||||
let source = if last.history_len == history.len() {
|
||||
EstimateSource::Measured
|
||||
} else {
|
||||
EstimateSource::Extrapolated
|
||||
};
|
||||
TokenEstimate { tokens, source }
|
||||
}
|
||||
|
||||
// ── Worker に生やす公開 API ───────────────────────────────────────────────
|
||||
|
||||
impl<C: LlmClient, St: Store> Worker<C, St> {
|
||||
/// 現在の history 全体の推定トークン数。
|
||||
///
|
||||
/// 最後の measurement と、その後に追加された未測定分の byte/4 外挿。
|
||||
pub fn total_tokens(&self) -> TokenEstimate {
|
||||
let usage = self.usage_history();
|
||||
llm_engine::token_counter::total_tokens(self.history(), &usage)
|
||||
}
|
||||
|
||||
/// 任意の history index 時点でのプロンプト全長推定。
|
||||
///
|
||||
/// `total_tokens()` と同じ accounting を任意位置で評価する版。
|
||||
/// memory extract trigger が
|
||||
/// `total_tokens_at(now) - total_tokens_at(pointer)` で
|
||||
/// pointer 以降に増えたプロンプト長を測るのに使う。
|
||||
pub fn total_tokens_at(&self, history_len: usize) -> TokenEstimate {
|
||||
let usage = self.usage_history();
|
||||
llm_engine::token_counter::total_tokens_at(self.history(), &usage, history_len)
|
||||
}
|
||||
|
||||
/// 末尾から `retained` トークン以上を残すための分割位置。
|
||||
///
|
||||
/// `history[..cut.index]` が要約／破棄される側、`history[cut.index..]` が残る側。
|
||||
pub fn split_for_retained(&self, retained: u64) -> SplitPoint {
|
||||
let usage = self.usage_history();
|
||||
split_for_retained_impl(self.history(), &usage, retained)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn msg(text: &str) -> Item {
|
||||
Item::user_message(text)
|
||||
}
|
||||
|
||||
fn record(history_len: usize, tokens: u64) -> UsageRecord {
|
||||
UsageRecord {
|
||||
history_len,
|
||||
input_total_tokens: tokens,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_returns_zero_when_current_below_retained() {
|
||||
let history = vec![msg("a"), msg("b")];
|
||||
let records = vec![record(2, 50)];
|
||||
let cut = split_for_retained_impl(&history, &records, 1000);
|
||||
assert_eq!(cut.index, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_uses_current_occupancy_as_raw_byte_rate() {
|
||||
// Compact retained split does not treat the intermediate record at
|
||||
// len=2 as a raw prefix boundary. It uses the current occupancy
|
||||
// estimate (len=4 → 300) as a serialized-byte rate and keeps the
|
||||
// smallest item-granular suffix whose raw size covers retained=200.
|
||||
let history = vec![msg("a"), msg("b"), msg("c"), msg("d")];
|
||||
let records = vec![record(2, 100), record(4, 300)];
|
||||
let cut = split_for_retained_impl(&history, &records, 200);
|
||||
assert_eq!(cut.index, 1);
|
||||
assert_eq!(cut.source, EstimateSource::Measured);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_does_not_use_non_current_measurements_as_cut_boundaries() {
|
||||
let history = vec![msg("aaaaaa"), msg("bbbbbb"), msg("cccccc"), msg("dddddd")];
|
||||
let records = vec![record(1, 50), record(4, 400)];
|
||||
let cut = split_for_retained_impl(&history, &records, 250);
|
||||
assert_eq!(cut.index, 1);
|
||||
assert_eq!(cut.source, EstimateSource::Measured);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_ignores_non_monotonic_usage_spike_for_retained_tail() {
|
||||
let history: Vec<Item> = (0..20)
|
||||
.map(|idx| msg(&format!("message-{idx}-{}", "x".repeat(100))))
|
||||
.collect();
|
||||
let records = vec![
|
||||
record(2, 900), // request-time spike after pruning/projection
|
||||
record(20, 1000),
|
||||
];
|
||||
let cut = split_for_retained_impl(&history, &records, 100);
|
||||
|
||||
// The old prefix-crossing logic picked index 2 because 900 >=
|
||||
// 1000-100, retaining almost the whole persisted history. The compact
|
||||
// split must instead use raw suffix size and keep only the tail needed
|
||||
// for the retained budget.
|
||||
assert!(cut.index > 10, "cut.index = {}", cut.index);
|
||||
assert_eq!(cut.source, EstimateSource::Measured);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_all_when_retained_zero() {
|
||||
let history = vec![msg("a"), msg("b")];
|
||||
let records = vec![record(2, 100)];
|
||||
let cut = split_for_retained_impl(&history, &records, 0);
|
||||
assert_eq!(cut.index, 2);
|
||||
}
|
||||
|
||||
fn tool_result_with(summary: &str, content: Option<&str>) -> Item {
|
||||
match content {
|
||||
Some(c) => Item::tool_result_with_content("call", summary, c),
|
||||
None => Item::tool_result("call", summary),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_estimates_for_prune_returns_every_prefix_boundary() {
|
||||
let history = vec![msg("a"), msg("b"), msg("c")];
|
||||
let estimates = token_estimates_for_prune_impl(&history, &[record(3, 300)]);
|
||||
assert_eq!(estimates.len(), history.len() + 1);
|
||||
assert_eq!(estimates[0].tokens, 0);
|
||||
assert_eq!(estimates[3].tokens, 300);
|
||||
assert_eq!(estimates[3].source, EstimateSource::Measured);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_estimates_for_prune_propagates_no_data() {
|
||||
let history = vec![msg("a"), msg("b")];
|
||||
let estimates = token_estimates_for_prune_impl(&history, &[]);
|
||||
assert_eq!(estimates.len(), history.len() + 1);
|
||||
assert_eq!(estimates[0].source, EstimateSource::Measured);
|
||||
assert_eq!(estimates[1].source, EstimateSource::NoData);
|
||||
assert_eq!(estimates[2].source, EstimateSource::NoData);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn savings_for_prune_skips_non_toolresult_indices() {
|
||||
let history = vec![msg("a"), msg("b"), msg("c")];
|
||||
// indices point at plain messages, not ToolResult → 0 savings.
|
||||
let est = savings_for_prune_impl(&history, &[record(3, 300)], &[0, 1, 2]);
|
||||
assert_eq!(est.tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn savings_for_prune_skips_content_none_items() {
|
||||
let history = vec![
|
||||
msg("user"),
|
||||
tool_result_with("s1", None),
|
||||
tool_result_with("s2", None),
|
||||
];
|
||||
let est = savings_for_prune_impl(&history, &[record(3, 300)], &[1, 2]);
|
||||
assert_eq!(est.tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn savings_for_prune_counts_only_content_delta() {
|
||||
// 1 item with big content vs the same structure without content.
|
||||
let big = "x".repeat(400);
|
||||
let history = vec![
|
||||
msg("user"),
|
||||
tool_result_with("summary", Some(&big)),
|
||||
msg("tail"),
|
||||
];
|
||||
// 1 record at end so rate = tokens / total_bytes
|
||||
let total_bytes: u64 = history.iter().map(item_bytes).sum();
|
||||
let records = vec![record(history.len(), total_bytes)]; // rate = 1 tok/byte
|
||||
let est = savings_for_prune_impl(&history, &records, &[1]);
|
||||
// saved bytes ≈ size of the big content payload; with rate=1 it
|
||||
// should be close to 400 and far from the full item bytes.
|
||||
let full_item_bytes = item_bytes(&history[1]);
|
||||
assert!(est.tokens > 0);
|
||||
assert!(est.tokens < full_item_bytes);
|
||||
assert!(est.tokens >= 400);
|
||||
assert_eq!(est.source, EstimateSource::Measured);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn savings_for_prune_no_records_falls_back_to_bytes() {
|
||||
let history = vec![msg("u"), tool_result_with("s", Some("hello world"))];
|
||||
let est = savings_for_prune_impl(&history, &[], &[1]);
|
||||
assert_eq!(est.source, EstimateSource::NoData);
|
||||
assert!(est.tokens > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn savings_for_prune_extrapolated_when_history_grew_past_measurement() {
|
||||
let big = "x".repeat(200);
|
||||
let history = vec![
|
||||
msg("u1"),
|
||||
tool_result_with("s", Some(&big)),
|
||||
msg("u2"), // added after the last measurement
|
||||
];
|
||||
let records = vec![record(2, 100)];
|
||||
let est = savings_for_prune_impl(&history, &records, &[1]);
|
||||
assert_eq!(est.source, EstimateSource::Extrapolated);
|
||||
assert!(est.tokens > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn savings_for_prune_empty_indices_is_zero() {
|
||||
let history = vec![msg("a")];
|
||||
let est = savings_for_prune_impl(&history, &[record(1, 100)], &[]);
|
||||
assert_eq!(est.tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn savings_for_prune_ignores_out_of_range_indices() {
|
||||
let history = vec![msg("a")];
|
||||
let est = savings_for_prune_impl(&history, &[record(1, 100)], &[99]);
|
||||
assert_eq!(est.tokens, 0);
|
||||
}
|
||||
|
||||
fn tc(call_id: &str) -> Item {
|
||||
Item::tool_call(call_id, "Read", "{}")
|
||||
}
|
||||
|
||||
fn tr(call_id: &str) -> Item {
|
||||
Item::tool_result(call_id, "summary")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn balance_noop_on_clean_message_boundary() {
|
||||
let history = vec![msg("a"), msg("b"), msg("c")];
|
||||
assert_eq!(balance_to_pair_boundary(&history, 2), 2);
|
||||
assert_eq!(balance_to_pair_boundary(&history, 0), 0);
|
||||
assert_eq!(balance_to_pair_boundary(&history, 3), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn balance_retreats_from_inside_parallel_tool_results() {
|
||||
// [Msg, TC_a, TC_b, TC_c, TR_a, TR_b, TR_c]
|
||||
// cut=5 → retained=[TR_b, TR_c]。TR_c の TC は idx=3、TR_b は idx=2 →
|
||||
// idx=2 まで後退。だが retained に TR_a (idx=4) が新たに入り、その TC_a
|
||||
// は idx=1 でまだ外 → 連鎖後退で最終的に idx=1。retained は
|
||||
// [TC_a, TC_b, TC_c, TR_a, TR_b, TR_c]。
|
||||
let history = vec![
|
||||
msg("u"),
|
||||
tc("a"),
|
||||
tc("b"),
|
||||
tc("c"),
|
||||
tr("a"),
|
||||
tr("b"),
|
||||
tr("c"),
|
||||
];
|
||||
assert_eq!(balance_to_pair_boundary(&history, 5), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn balance_retreats_between_call_and_result() {
|
||||
// [TC_a, TR_a, TC_b, TR_b]。cut=3 → retained=[TR_b] orphan。
|
||||
// TC_b は idx=2 → cut=2。retained=[TC_b, TR_b]。
|
||||
let history = vec![tc("a"), tr("a"), tc("b"), tr("b")];
|
||||
assert_eq!(balance_to_pair_boundary(&history, 3), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn balance_cascades_through_nested_pairs() {
|
||||
// [TC_a, TC_b, TR_b, TR_a, TC_c, TR_c]。cut=3 → retained=[TR_a, TC_c, TR_c]。
|
||||
// TR_a の TC は idx=0 → cut=0。retained=full。
|
||||
let history = vec![tc("a"), tc("b"), tr("b"), tr("a"), tc("c"), tr("c")];
|
||||
assert_eq!(balance_to_pair_boundary(&history, 3), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn balance_noop_when_cut_at_pair_boundary() {
|
||||
// [TC_a, TR_a, Msg, TC_b, TR_b]。cut=2 → retained=[Msg, TC_b, TR_b] balanced。
|
||||
let history = vec![tc("a"), tr("a"), msg("u"), tc("b"), tr("b")];
|
||||
assert_eq!(balance_to_pair_boundary(&history, 2), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn balance_handles_orphan_result_without_matching_call() {
|
||||
// ToolCall がそもそも存在しない ToolResult は触らない（壊れた history は
|
||||
// ここでは直しようがない）。cut=1 → そのまま 1 を返す。
|
||||
let history = vec![msg("u"), tr("zombie")];
|
||||
assert_eq!(balance_to_pair_boundary(&history, 1), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn balance_keeps_cut_when_call_is_inside_retained() {
|
||||
// [Msg, TC_a, TR_a]。cut=1 → retained=[TC_a, TR_a]。TR_a の call_pos=1 >= idx=1。OK。
|
||||
let history = vec![msg("u"), tc("a"), tr("a")];
|
||||
assert_eq!(balance_to_pair_boundary(&history, 1), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_for_retained_aligns_to_pair_boundary() {
|
||||
// 並列 TC*3 / TR*3 ターン後に Msg を 1 件足し、retained=Msg のサイズ相当に
|
||||
// 設定。トークン的には cut=末尾近くだが、orphan を避けるため TC 群の手前
|
||||
// まで後退するはず。
|
||||
let history = vec![
|
||||
msg("user"),
|
||||
tc("a"),
|
||||
tc("b"),
|
||||
tc("c"),
|
||||
tr("a"),
|
||||
tr("b"),
|
||||
tr("c"),
|
||||
msg("tail"),
|
||||
];
|
||||
let total_bytes: u64 = history.iter().map(item_bytes).sum();
|
||||
let records = vec![record(history.len(), total_bytes)]; // rate = 1 tok/byte
|
||||
// tail の item_bytes 相当のみ retain したい。
|
||||
let tail_tokens = item_bytes(&history[7]);
|
||||
let cut = split_for_retained_impl(&history, &records, tail_tokens);
|
||||
// token 単独だと cut は 7（tail のみ retained）になるが、retained 先頭が
|
||||
// Msg なら balance しなくて OK。balance helper の no-op を確認する意味も込めて
|
||||
// index == 7 を期待する。
|
||||
assert_eq!(cut.index, 7);
|
||||
|
||||
// 逆に retained をやや増やしてトークン的に cut=6（TR_c のみ retained）に
|
||||
// させると、TR_c は orphan なので balance が 1 まで後退するはず。
|
||||
let big_retain = tail_tokens + item_bytes(&history[6]);
|
||||
let cut = split_for_retained_impl(&history, &records, big_retain);
|
||||
assert_eq!(cut.index, 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
//! Tracks per-LLM-request Usage measurements within a Worker run.
|
||||
//!
|
||||
//! Bridge between two sync touchpoints in the Engine lifecycle:
|
||||
//!
|
||||
//! - **`pre_llm_request` hook** (async, but synchronously accessed via the
|
||||
//! tracker): captures `history.len()` at the moment a request goes out.
|
||||
//! - **`on_usage` callback** (sync closure): receives the aggregated final
|
||||
//! `UsageEvent` for that request after the stream completes.
|
||||
//!
|
||||
//! Pairing the two yields one `UsageRecord` per LLM call. Worker drains them
|
||||
//! in `persist_turn` and writes them as `LogEntry::LlmUsage` entries.
|
||||
//!
|
||||
//! Multiple LLM calls per Worker run (tool loop) are supported: each call
|
||||
//! produces its own `(history_len, UsageEvent)` pair, and the records are
|
||||
//! buffered in chronological order.
|
||||
|
||||
use std::sync::Mutex;
|
||||
|
||||
use llm_engine::UsageRecord;
|
||||
use llm_engine::timeline::event::UsageEvent;
|
||||
|
||||
/// One drained measurement: the underlying `UsageRecord` plus an optional
|
||||
/// `correlation_id` stamped by the prune projection (or any other future
|
||||
/// upstream observer) so that downstream metrics emitted alongside this
|
||||
/// record can be joined to it after the fact.
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct RecordedUsage {
|
||||
pub(crate) record: UsageRecord,
|
||||
pub(crate) correlation_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Shared between the pre-request hook, the `on_usage` callback, and Worker.
|
||||
pub(crate) struct UsageTracker {
|
||||
/// `history.len()` captured at the most recent `pre_llm_request`.
|
||||
/// Cleared when paired with an incoming `on_usage` event.
|
||||
pending_history_len: Mutex<Option<usize>>,
|
||||
/// Optional `correlation_id` set by an upstream observer (currently
|
||||
/// the prune projection on `Fired`). Paired into the next
|
||||
/// `RecordedUsage` and cleared. Skips that don't fire leave this
|
||||
/// `None`, so the resulting record carries no correlation.
|
||||
pending_correlation_id: Mutex<Option<String>>,
|
||||
/// Records accumulated during the current run; drained by Worker.
|
||||
pending_records: Mutex<Vec<RecordedUsage>>,
|
||||
}
|
||||
|
||||
impl UsageTracker {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
pending_history_len: Mutex::new(None),
|
||||
pending_correlation_id: Mutex::new(None),
|
||||
pending_records: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Called from a `pre_llm_request` hook with the current history length.
|
||||
pub(crate) fn note_request(&self, history_len: usize) {
|
||||
*self.pending_history_len.lock().unwrap() = Some(history_len);
|
||||
}
|
||||
|
||||
/// Stash a `correlation_id` to be paired into the next `RecordedUsage`.
|
||||
/// Currently invoked by the prune observer on `Fired` so that the
|
||||
/// `prune.fire` metric and the `prune.post_request` metric (emitted
|
||||
/// alongside the resulting `LlmUsage`) carry the same join key.
|
||||
///
|
||||
/// Overwrites any previous unconsumed value — by construction the
|
||||
/// observer fires at most once per outgoing LLM request, immediately
|
||||
/// before the pre-request hook captures `history_len`.
|
||||
pub(crate) fn note_correlation_id(&self, id: String) {
|
||||
*self.pending_correlation_id.lock().unwrap() = Some(id);
|
||||
}
|
||||
|
||||
/// Called from the `on_usage` callback with the aggregated final
|
||||
/// UsageEvent. If a `history_len` was previously stashed via
|
||||
/// `note_request`, builds a `RecordedUsage` and pushes it onto the
|
||||
/// buffer. If not (e.g. test code that fires Usage outside a request),
|
||||
/// drops the event.
|
||||
pub(crate) fn record_usage(&self, event: &UsageEvent) {
|
||||
let history_len = match self.pending_history_len.lock().unwrap().take() {
|
||||
Some(n) => n,
|
||||
None => return,
|
||||
};
|
||||
let correlation_id = self.pending_correlation_id.lock().unwrap().take();
|
||||
// UsageEvent.input_tokens は scheme 層で「占有量（プロンプト全長）」に
|
||||
// 正規化済みである前提（Anthropic は cache_read + cache_creation を
|
||||
// 加算して emit する）。
|
||||
let input_total = event.input_tokens.unwrap_or(0);
|
||||
let cache_read = event.cache_read_input_tokens.unwrap_or(0);
|
||||
let cache_write = event.cache_creation_input_tokens.unwrap_or(0);
|
||||
let output = event.output_tokens.unwrap_or(0);
|
||||
self.pending_records.lock().unwrap().push(RecordedUsage {
|
||||
record: UsageRecord {
|
||||
history_len,
|
||||
input_total_tokens: input_total,
|
||||
cache_read_tokens: cache_read,
|
||||
cache_write_tokens: cache_write,
|
||||
output_tokens: output,
|
||||
},
|
||||
correlation_id,
|
||||
});
|
||||
}
|
||||
|
||||
/// Return a clone of the accumulated `UsageRecord`s without clearing them.
|
||||
/// Used by request-time circuit breakers that need the same occupancy
|
||||
/// projection as Worker persistence while the run is still active.
|
||||
pub(crate) fn records(&self) -> Vec<UsageRecord> {
|
||||
self.pending_records
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|r| r.record.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Drain accumulated records. Called by Worker after a run completes,
|
||||
/// before persisting the turn.
|
||||
pub(crate) fn drain(&self) -> Vec<RecordedUsage> {
|
||||
std::mem::take(&mut *self.pending_records.lock().unwrap())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_event(input: u64, cache_read: u64, cache_write: u64, output: u64) -> UsageEvent {
|
||||
UsageEvent {
|
||||
input_tokens: Some(input),
|
||||
output_tokens: Some(output),
|
||||
total_tokens: Some(input + output),
|
||||
cache_read_input_tokens: Some(cache_read),
|
||||
cache_creation_input_tokens: Some(cache_write),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pairs_history_len_with_usage_event() {
|
||||
let tracker = UsageTracker::new();
|
||||
tracker.note_request(5);
|
||||
tracker.record_usage(&make_event(1000, 800, 100, 42));
|
||||
|
||||
let records = tracker.drain();
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].record.history_len, 5);
|
||||
assert_eq!(records[0].record.input_total_tokens, 1000);
|
||||
assert_eq!(records[0].record.cache_read_tokens, 800);
|
||||
assert_eq!(records[0].record.cache_write_tokens, 100);
|
||||
assert_eq!(records[0].record.output_tokens, 42);
|
||||
assert!(records[0].correlation_id.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn records_clones_without_clearing() {
|
||||
let tracker = UsageTracker::new();
|
||||
tracker.note_request(1);
|
||||
tracker.record_usage(&make_event(10, 0, 0, 5));
|
||||
|
||||
let records = tracker.records();
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].history_len, 1);
|
||||
assert_eq!(records[0].input_total_tokens, 10);
|
||||
assert_eq!(tracker.records().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drain_clears_buffer() {
|
||||
let tracker = UsageTracker::new();
|
||||
tracker.note_request(1);
|
||||
tracker.record_usage(&make_event(10, 0, 0, 5));
|
||||
assert_eq!(tracker.drain().len(), 1);
|
||||
assert_eq!(tracker.drain().len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_without_pending_history_len_is_dropped() {
|
||||
let tracker = UsageTracker::new();
|
||||
tracker.record_usage(&make_event(10, 0, 0, 5));
|
||||
assert_eq!(tracker.drain().len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_requests_in_one_run() {
|
||||
let tracker = UsageTracker::new();
|
||||
tracker.note_request(5);
|
||||
tracker.record_usage(&make_event(100, 0, 0, 20));
|
||||
tracker.note_request(10);
|
||||
tracker.record_usage(&make_event(200, 50, 0, 30));
|
||||
|
||||
let records = tracker.drain();
|
||||
assert_eq!(records.len(), 2);
|
||||
assert_eq!(records[0].record.history_len, 5);
|
||||
assert_eq!(records[1].record.history_len, 10);
|
||||
assert_eq!(records[1].record.cache_read_tokens, 50);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn correlation_id_pairs_with_next_record_only() {
|
||||
let tracker = UsageTracker::new();
|
||||
// Stash an ID, then run a request → the ID should land on this record.
|
||||
tracker.note_correlation_id("abc".into());
|
||||
tracker.note_request(5);
|
||||
tracker.record_usage(&make_event(100, 0, 0, 20));
|
||||
// Next request without a fresh stash → no correlation_id.
|
||||
tracker.note_request(10);
|
||||
tracker.record_usage(&make_event(200, 50, 0, 30));
|
||||
|
||||
let records = tracker.drain();
|
||||
assert_eq!(records.len(), 2);
|
||||
assert_eq!(records[0].correlation_id.as_deref(), Some("abc"));
|
||||
assert!(records[1].correlation_id.is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,878 @@
|
||||
//! Compact worker state and the four tools that drive it.
|
||||
//!
|
||||
//! The compact worker is a disposable `Engine` instance spun up by
|
||||
//! [`Worker::compact`]. It receives the history to summarise plus a list of
|
||||
//! default reference files (from the session-lifetime `Tracker`) and runs
|
||||
//! a tool-driven LLM loop. The tools here let it:
|
||||
//!
|
||||
//! - `read_file` — inspect referenced files (reuses `tools::read_tool`)
|
||||
//! - `mark_read_required(path, offset?, limit?)` — nominate a file whose
|
||||
//! contents should be injected into the compacted context as an
|
||||
//! auto-read system message
|
||||
//! - `add_reference(path)` — nominate a file the next session should
|
||||
//! know about by name only (contents not included)
|
||||
//! - `write_summary(text)` — deliver (or overwrite) the structured summary
|
||||
//!
|
||||
//! Everything the worker decides ends up in [`CompactWorkerContext`],
|
||||
//! which `Worker::compact` drains after the loop and turns into the
|
||||
//! compacted session's opening system messages.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_engine::Item;
|
||||
use llm_engine::interceptor::{Interceptor, PreRequestAction, PreToolAction, ToolCallInfo};
|
||||
use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput, ToolResult};
|
||||
use serde::Deserialize;
|
||||
use tools::ScopedFs;
|
||||
|
||||
use crate::compact::usage_tracker::UsageTracker;
|
||||
use crate::fs_view::{ReadRequirement, slice_lines};
|
||||
|
||||
/// Aggregated output of a compact worker run.
|
||||
#[derive(Debug, Default, Clone)]
|
||||
pub(crate) struct CompactWorkerContext {
|
||||
pub read_required: Vec<ReadRequirement>,
|
||||
pub references: Vec<PathBuf>,
|
||||
pub summary: Option<String>,
|
||||
/// Tokens already consumed by `mark_read_required` calls.
|
||||
pub auto_read_consumed: u64,
|
||||
/// Aggregate cap. `0` treats the budget as disabled.
|
||||
pub auto_read_budget: u64,
|
||||
}
|
||||
|
||||
impl CompactWorkerContext {
|
||||
pub(crate) fn with_budget(auto_read_budget: u64) -> Self {
|
||||
Self {
|
||||
auto_read_budget,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn remaining_budget(&self) -> u64 {
|
||||
self.auto_read_budget
|
||||
.saturating_sub(self.auto_read_consumed)
|
||||
}
|
||||
}
|
||||
|
||||
/// Input to `mark_read_required`.
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct MarkParams {
|
||||
/// Absolute path to the file.
|
||||
pub file_path: PathBuf,
|
||||
/// 0-based line offset.
|
||||
#[serde(default)]
|
||||
pub offset: Option<usize>,
|
||||
/// Maximum number of lines to inject.
|
||||
#[serde(default)]
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
/// Input to `add_reference`.
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct ReferenceParams {
|
||||
/// Absolute path to the file.
|
||||
pub file_path: PathBuf,
|
||||
}
|
||||
|
||||
/// Input to `write_summary`.
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct SummaryParams {
|
||||
/// Full structured summary text (overwrites any previous call).
|
||||
pub text: String,
|
||||
}
|
||||
|
||||
/// Input to `search_session_log`.
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct SearchSessionParams {
|
||||
/// Case-insensitive substring to search in compact-target history.
|
||||
pub query: String,
|
||||
/// 0-based item offset to start searching from.
|
||||
#[serde(default)]
|
||||
pub offset: Option<usize>,
|
||||
/// Maximum number of hits to return.
|
||||
#[serde(default)]
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
/// Input to `read_session_items`.
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct ReadSessionParams {
|
||||
/// 0-based compact-target history item offset.
|
||||
pub offset: usize,
|
||||
/// Maximum number of items to return.
|
||||
pub limit: usize,
|
||||
/// `compact` omits tool arguments/full results; `full` includes message text and tool result content.
|
||||
#[serde(default = "default_session_read_mode")]
|
||||
pub mode: String,
|
||||
}
|
||||
|
||||
fn default_session_read_mode() -> String {
|
||||
"compact".to_string()
|
||||
}
|
||||
|
||||
const SESSION_TOOL_MAX_OUTPUT_TOKENS: u64 = 12_000;
|
||||
const SESSION_SEARCH_MAX_RESULTS: usize = 50;
|
||||
const SESSION_READ_MAX_ITEMS: usize = 80;
|
||||
|
||||
const MARK_DESCRIPTION: &str = "Inject a file's contents into the compacted context so the \
|
||||
next session starts with it already read. Use this for files the next task needs in full. \
|
||||
Optionally specify `offset` (0-based line) and `limit` (line count) to inject only a slice. \
|
||||
Counts against `auto_read_budget`; overflow returns an error and the mark is not recorded. \
|
||||
Paths must be absolute.";
|
||||
|
||||
const REFERENCE_DESCRIPTION: &str = "Record a file path as a named reference in the compacted \
|
||||
context without injecting its contents. Use for files that are contextually relevant but \
|
||||
whose current content the next session can fetch on demand.";
|
||||
|
||||
const SUMMARY_DESCRIPTION: &str = "Provide the final structured summary text. Subsequent calls \
|
||||
replace the previous content; only the last call is used. Must be called before the compact run \
|
||||
ends or compaction fails.";
|
||||
|
||||
const SEARCH_SESSION_DESCRIPTION: &str = "Search the compact-target session history by \
|
||||
case-insensitive substring. Returns item indexes and compact snippets. Use this when the initial \
|
||||
overview is not enough to identify which part of the session matters. Results are bounded; narrow \
|
||||
the query if important details are omitted.";
|
||||
|
||||
const READ_SESSION_DESCRIPTION: &str = "Read a bounded range of compact-target session history \
|
||||
items by 0-based index. mode='compact' omits tool arguments, full tool results, and reasoning \
|
||||
bodies; mode='full' includes message text and tool result content but still remains bounded. Use \
|
||||
this to verify details before writing the summary.";
|
||||
|
||||
struct SessionLogToolState {
|
||||
items: Arc<Vec<Item>>,
|
||||
}
|
||||
|
||||
struct SearchSessionLogTool {
|
||||
state: Arc<SessionLogToolState>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SearchSessionLogTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: SearchSessionParams = serde_json::from_str(input_json).map_err(|e| {
|
||||
ToolError::InvalidArgument(format!("invalid search_session_log input: {e}"))
|
||||
})?;
|
||||
let query = params.query.trim().to_lowercase();
|
||||
if query.is_empty() {
|
||||
return Err(ToolError::InvalidArgument(
|
||||
"search_session_log query must not be empty".to_string(),
|
||||
));
|
||||
}
|
||||
let offset = params.offset.unwrap_or(0).min(self.state.items.len());
|
||||
let limit = params
|
||||
.limit
|
||||
.unwrap_or(20)
|
||||
.clamp(1, SESSION_SEARCH_MAX_RESULTS);
|
||||
let mut hits = Vec::new();
|
||||
for (idx, item) in self.state.items.iter().enumerate().skip(offset) {
|
||||
let haystack = session_item_search_text(item).to_lowercase();
|
||||
if haystack.contains(&query) {
|
||||
hits.push(format_session_item(
|
||||
idx,
|
||||
item,
|
||||
SessionReadMode::Compact,
|
||||
600,
|
||||
));
|
||||
if hits.len() >= limit {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut content = hits.join("\n\n");
|
||||
let truncated = truncate_to_token_budget(&mut content, SESSION_TOOL_MAX_OUTPUT_TOKENS);
|
||||
let summary = if hits.is_empty() {
|
||||
format!("No session log hits for {query:?} from item offset {offset}.")
|
||||
} else if truncated {
|
||||
format!(
|
||||
"Found {} session log hit(s) for {query:?}; output truncated. Narrow the query.",
|
||||
hits.len()
|
||||
)
|
||||
} else {
|
||||
format!("Found {} session log hit(s) for {query:?}.", hits.len())
|
||||
};
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: (!content.is_empty()).then_some(content),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct ReadSessionItemsTool {
|
||||
state: Arc<SessionLogToolState>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ReadSessionItemsTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: ReadSessionParams = serde_json::from_str(input_json).map_err(|e| {
|
||||
ToolError::InvalidArgument(format!("invalid read_session_items input: {e}"))
|
||||
})?;
|
||||
let mode = SessionReadMode::parse(¶ms.mode)?;
|
||||
let offset = params.offset.min(self.state.items.len());
|
||||
let limit = params.limit.clamp(1, SESSION_READ_MAX_ITEMS);
|
||||
let end = offset.saturating_add(limit).min(self.state.items.len());
|
||||
let mut blocks = Vec::new();
|
||||
for idx in offset..end {
|
||||
blocks.push(format_session_item(
|
||||
idx,
|
||||
&self.state.items[idx],
|
||||
mode,
|
||||
4_000,
|
||||
));
|
||||
}
|
||||
let mut content = blocks.join("\n\n");
|
||||
let truncated = truncate_to_token_budget(&mut content, SESSION_TOOL_MAX_OUTPUT_TOKENS);
|
||||
let summary = if truncated {
|
||||
format!(
|
||||
"Read session items {offset}..{end} in {mode:?} mode; output truncated. Narrow the range."
|
||||
)
|
||||
} else {
|
||||
format!("Read session items {offset}..{end} in {mode:?} mode.")
|
||||
};
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: (!content.is_empty()).then_some(content),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum SessionReadMode {
|
||||
Compact,
|
||||
Full,
|
||||
}
|
||||
|
||||
impl SessionReadMode {
|
||||
fn parse(value: &str) -> Result<Self, ToolError> {
|
||||
match value {
|
||||
"compact" => Ok(Self::Compact),
|
||||
"full" => Ok(Self::Full),
|
||||
other => Err(ToolError::InvalidArgument(format!(
|
||||
"invalid read_session_items mode {other:?}; expected 'compact' or 'full'"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn session_item_search_text(item: &Item) -> String {
|
||||
match item {
|
||||
Item::Message { role, content, .. } => format!(
|
||||
"{:?} {}",
|
||||
role,
|
||||
content
|
||||
.iter()
|
||||
.map(|p| p.as_text())
|
||||
.collect::<Vec<_>>()
|
||||
.join("")
|
||||
),
|
||||
Item::ToolCall {
|
||||
name, arguments, ..
|
||||
} => format!("tool_call {name} {arguments}"),
|
||||
Item::ToolResult {
|
||||
summary, content, ..
|
||||
} => format!(
|
||||
"tool_result {summary} {}",
|
||||
content.as_deref().unwrap_or_default()
|
||||
),
|
||||
Item::Reasoning { text, summary, .. } => format!("reasoning {text} {}", summary.join(" ")),
|
||||
}
|
||||
}
|
||||
|
||||
fn format_session_item(idx: usize, item: &Item, mode: SessionReadMode, max_chars: usize) -> String {
|
||||
match item {
|
||||
Item::Message { role, content, .. } => {
|
||||
let text = content
|
||||
.iter()
|
||||
.map(|p| p.as_text())
|
||||
.collect::<Vec<_>>()
|
||||
.join("");
|
||||
format!(
|
||||
"[{idx} Message {:?}] {}",
|
||||
role,
|
||||
truncate_chars(&text, max_chars)
|
||||
)
|
||||
}
|
||||
Item::ToolCall {
|
||||
name, arguments, ..
|
||||
} => match mode {
|
||||
SessionReadMode::Compact => format!("[{idx} ToolCall] {name} (arguments omitted)"),
|
||||
SessionReadMode::Full => format!(
|
||||
"[{idx} ToolCall] {name}\narguments: {}",
|
||||
truncate_chars(arguments, max_chars)
|
||||
),
|
||||
},
|
||||
Item::ToolResult {
|
||||
summary,
|
||||
content,
|
||||
is_error,
|
||||
..
|
||||
} => match mode {
|
||||
SessionReadMode::Compact => format!(
|
||||
"[{idx} ToolResult{}] {} (content omitted)",
|
||||
if *is_error { " error" } else { "" },
|
||||
truncate_chars(summary, 800)
|
||||
),
|
||||
SessionReadMode::Full => format!(
|
||||
"[{idx} ToolResult{}] {}\ncontent: {}",
|
||||
if *is_error { " error" } else { "" },
|
||||
truncate_chars(summary, 800),
|
||||
truncate_chars(content.as_deref().unwrap_or(""), max_chars)
|
||||
),
|
||||
},
|
||||
Item::Reasoning { summary, .. } => match mode {
|
||||
SessionReadMode::Compact => format!(
|
||||
"[{idx} Reasoning] {} (body omitted)",
|
||||
truncate_chars(&summary.join(" "), 800)
|
||||
),
|
||||
SessionReadMode::Full => format!(
|
||||
"[{idx} Reasoning] {} (body omitted)",
|
||||
truncate_chars(&summary.join(" "), 800)
|
||||
),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate_chars(text: &str, max_chars: usize) -> String {
|
||||
if text.chars().count() <= max_chars {
|
||||
return text.to_string();
|
||||
}
|
||||
let mut out = text.chars().take(max_chars).collect::<String>();
|
||||
out.push_str("… [truncated]");
|
||||
out
|
||||
}
|
||||
|
||||
fn truncate_to_token_budget(text: &mut String, max_tokens: u64) -> bool {
|
||||
let max_bytes = max_tokens.saturating_mul(4) as usize;
|
||||
if text.len() <= max_bytes {
|
||||
return false;
|
||||
}
|
||||
let mut cut = 0;
|
||||
for (idx, _) in text.char_indices() {
|
||||
if idx > max_bytes {
|
||||
break;
|
||||
}
|
||||
cut = idx;
|
||||
}
|
||||
text.truncate(cut);
|
||||
text.push_str("\n… [session tool output truncated]");
|
||||
true
|
||||
}
|
||||
|
||||
struct MarkReadRequiredTool {
|
||||
fs: ScopedFs,
|
||||
ctx: Arc<Mutex<CompactWorkerContext>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for MarkReadRequiredTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: MarkParams = serde_json::from_str(input_json).map_err(|e| {
|
||||
ToolError::InvalidArgument(format!("invalid mark_read_required input: {e}"))
|
||||
})?;
|
||||
|
||||
// Read the file through the shared ScopedFs so scope and I/O
|
||||
// errors surface the same way the regular `read_file` tool does.
|
||||
let bytes = self
|
||||
.fs
|
||||
.read_bytes(¶ms.file_path)
|
||||
.map_err(|e| ToolError::ExecutionFailed(format!("read failed: {e}")))?;
|
||||
let text = String::from_utf8_lossy(&bytes);
|
||||
let slice = slice_lines(&text, params.offset.unwrap_or(0), params.limit);
|
||||
let estimated_tokens = estimate_tokens(slice.len());
|
||||
|
||||
let mut guard = self.ctx.lock().expect("compact worker context poisoned");
|
||||
let budget = guard.auto_read_budget;
|
||||
let would_consume = guard.auto_read_consumed.saturating_add(estimated_tokens);
|
||||
if budget > 0 && would_consume > budget {
|
||||
return Err(ToolError::ExecutionFailed(format!(
|
||||
"auto-read budget exhausted ({budget} tokens). Remove an existing mark or use \
|
||||
add_reference instead."
|
||||
)));
|
||||
}
|
||||
guard.read_required.push(ReadRequirement {
|
||||
path: params.file_path.clone(),
|
||||
offset: params.offset,
|
||||
limit: params.limit,
|
||||
});
|
||||
guard.auto_read_consumed = would_consume;
|
||||
let remaining = guard.remaining_budget();
|
||||
drop(guard);
|
||||
|
||||
let mut summary = format!(
|
||||
"Marked {} for auto-read (≈{estimated_tokens} tokens). \
|
||||
Budget: {remaining}/{budget} tokens remaining.",
|
||||
params.file_path.display()
|
||||
);
|
||||
if budget > 0 && remaining * 2 <= budget {
|
||||
summary.push_str(
|
||||
"\nNote: auto-read budget is at least half consumed. \
|
||||
Consider calling write_summary and finishing up soon.",
|
||||
);
|
||||
}
|
||||
Ok(ToolOutput {
|
||||
summary,
|
||||
content: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct AddReferenceTool {
|
||||
ctx: Arc<Mutex<CompactWorkerContext>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for AddReferenceTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: ReferenceParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid add_reference input: {e}")))?;
|
||||
let mut guard = self.ctx.lock().expect("compact worker context poisoned");
|
||||
if !guard
|
||||
.references
|
||||
.iter()
|
||||
.any(|p| p.as_path() == params.file_path.as_path())
|
||||
{
|
||||
guard.references.push(params.file_path.clone());
|
||||
}
|
||||
Ok(ToolOutput {
|
||||
summary: format!("Added reference {}", params.file_path.display()),
|
||||
content: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct WriteSummaryTool {
|
||||
ctx: Arc<Mutex<CompactWorkerContext>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for WriteSummaryTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: SummaryParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid write_summary input: {e}")))?;
|
||||
let mut guard = self.ctx.lock().expect("compact worker context poisoned");
|
||||
let overwritten = guard.summary.is_some();
|
||||
guard.summary = Some(params.text);
|
||||
drop(guard);
|
||||
let note = if overwritten {
|
||||
"Summary replaced."
|
||||
} else {
|
||||
"Summary recorded."
|
||||
};
|
||||
Ok(ToolOutput {
|
||||
summary: note.to_string(),
|
||||
content: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn mark_read_required_tool(
|
||||
fs: ScopedFs,
|
||||
ctx: Arc<Mutex<CompactWorkerContext>>,
|
||||
) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(MarkParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("mark_read_required")
|
||||
.description(MARK_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(MarkReadRequiredTool {
|
||||
fs: fs.clone(),
|
||||
ctx: ctx.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn add_reference_tool(ctx: Arc<Mutex<CompactWorkerContext>>) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(ReferenceParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("add_reference")
|
||||
.description(REFERENCE_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(AddReferenceTool { ctx: ctx.clone() });
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn write_summary_tool(ctx: Arc<Mutex<CompactWorkerContext>>) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(SummaryParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("write_summary")
|
||||
.description(SUMMARY_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(WriteSummaryTool { ctx: ctx.clone() });
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn search_session_log_tool(items: Arc<Vec<Item>>) -> ToolDefinition {
|
||||
let state = Arc::new(SessionLogToolState { items });
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(SearchSessionParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("search_session_log")
|
||||
.description(SEARCH_SESSION_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(SearchSessionLogTool {
|
||||
state: state.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn read_session_items_tool(items: Arc<Vec<Item>>) -> ToolDefinition {
|
||||
let state = Arc::new(SessionLogToolState { items });
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(ReadSessionParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("read_session_items")
|
||||
.description(READ_SESSION_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(ReadSessionItemsTool {
|
||||
state: state.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
/// Interceptor that monitors compact-worker context occupancy.
|
||||
///
|
||||
/// `max_input_tokens` remains the hard circuit breaker. Before that point,
|
||||
/// the interceptor can persist a system warning into worker history telling
|
||||
/// the model to stop broad exploration and call `write_summary`, and can block
|
||||
/// additional exploratory tool calls once the final reserve is reached.
|
||||
pub(crate) struct CompactWorkerInterceptor {
|
||||
pub usage_tracker: Arc<UsageTracker>,
|
||||
pub max_input_tokens: u64,
|
||||
pub finish_warning_remaining_tokens: u64,
|
||||
pub final_reserve_tokens: u64,
|
||||
pub on_warning: Option<Arc<dyn Fn(String) + Send + Sync>>,
|
||||
warning_sent: AtomicBool,
|
||||
last_remaining_tokens: AtomicU64,
|
||||
}
|
||||
|
||||
impl CompactWorkerInterceptor {
|
||||
pub(crate) fn new(
|
||||
usage_tracker: Arc<UsageTracker>,
|
||||
max_input_tokens: u64,
|
||||
finish_warning_remaining_tokens: u64,
|
||||
final_reserve_tokens: u64,
|
||||
on_warning: Option<Arc<dyn Fn(String) + Send + Sync>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
usage_tracker,
|
||||
max_input_tokens,
|
||||
finish_warning_remaining_tokens,
|
||||
final_reserve_tokens,
|
||||
on_warning,
|
||||
warning_sent: AtomicBool::new(false),
|
||||
last_remaining_tokens: AtomicU64::new(max_input_tokens),
|
||||
}
|
||||
}
|
||||
|
||||
fn maybe_emit_warning(&self, remaining: u64) -> Option<Item> {
|
||||
let warning_threshold = self.finish_warning_remaining_tokens;
|
||||
let reserve_threshold = self.final_reserve_tokens;
|
||||
let should_warn = (warning_threshold > 0 && remaining <= warning_threshold)
|
||||
|| (reserve_threshold > 0 && remaining <= reserve_threshold);
|
||||
if !should_warn || self.warning_sent.swap(true, Ordering::AcqRel) {
|
||||
return None;
|
||||
}
|
||||
|
||||
let message = format!(
|
||||
"compact worker context budget is low ({remaining}/{} tokens remaining). \
|
||||
Stop broad exploration now, read only if absolutely necessary, then call \
|
||||
`write_summary` with the final structured summary.",
|
||||
self.max_input_tokens
|
||||
);
|
||||
if let Some(cb) = self.on_warning.as_ref() {
|
||||
cb(message.clone());
|
||||
}
|
||||
Some(Item::system_message(format!(
|
||||
"[Compact worker budget warning]\n\n{message}"
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Interceptor for CompactWorkerInterceptor {
|
||||
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction {
|
||||
let records = self.usage_tracker.records();
|
||||
let estimate = llm_engine::token_counter::total_tokens(context, &records);
|
||||
if estimate.tokens > self.max_input_tokens {
|
||||
return PreRequestAction::Cancel(format!(
|
||||
"compact worker input occupancy exceeded {} tokens",
|
||||
self.max_input_tokens
|
||||
));
|
||||
}
|
||||
|
||||
let remaining = self.max_input_tokens.saturating_sub(estimate.tokens);
|
||||
self.last_remaining_tokens
|
||||
.store(remaining, Ordering::Release);
|
||||
if let Some(item) = self.maybe_emit_warning(remaining) {
|
||||
self.usage_tracker.note_request(context.len() + 1);
|
||||
return PreRequestAction::ContinueWith(vec![item]);
|
||||
}
|
||||
|
||||
self.usage_tracker.note_request(context.len());
|
||||
PreRequestAction::Continue
|
||||
}
|
||||
|
||||
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction {
|
||||
if self.final_reserve_tokens == 0 || info.call.name == "write_summary" {
|
||||
return PreToolAction::Continue;
|
||||
}
|
||||
let remaining = self.last_remaining_tokens.load(Ordering::Acquire);
|
||||
if remaining > self.final_reserve_tokens {
|
||||
return PreToolAction::Continue;
|
||||
}
|
||||
PreToolAction::SyntheticResult(ToolResult::error(
|
||||
info.call.id.clone(),
|
||||
"compact worker final reserve reached; do not perform more exploratory tool reads. Call `write_summary` now.",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Crude bytes→tokens estimate; good enough for budget accounting.
|
||||
fn estimate_tokens(bytes: usize) -> u64 {
|
||||
(bytes as u64).div_ceil(4)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use manifest::Scope;
|
||||
|
||||
fn make_fs(tmp: &std::path::Path) -> ScopedFs {
|
||||
let scope = Scope::writable(tmp.to_path_buf()).unwrap();
|
||||
ScopedFs::new(scope, tmp.to_path_buf())
|
||||
}
|
||||
|
||||
fn make_usage(input: u64) -> llm_engine::timeline::event::UsageEvent {
|
||||
llm_engine::timeline::event::UsageEvent {
|
||||
input_tokens: Some(input),
|
||||
output_tokens: Some(0),
|
||||
total_tokens: Some(input),
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_input_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compact_worker_interceptor_uses_occupancy_not_cumulative_usage() {
|
||||
let tracker = Arc::new(UsageTracker::new());
|
||||
let interceptor = CompactWorkerInterceptor::new(tracker.clone(), 150, 0, 0, None);
|
||||
let mut context = vec![Item::user_message("hello")];
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
PreRequestAction::Continue
|
||||
));
|
||||
tracker.record_usage(&make_usage(100));
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
PreRequestAction::Continue
|
||||
));
|
||||
tracker.record_usage(&make_usage(100));
|
||||
|
||||
// Two 100-token requests would exceed a cumulative 150-token cap, but
|
||||
// current occupancy is still the latest 100-token measurement.
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
PreRequestAction::Continue
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compact_worker_interceptor_warns_before_hard_cap() {
|
||||
let tracker = Arc::new(UsageTracker::new());
|
||||
let warnings = Arc::new(Mutex::new(Vec::new()));
|
||||
let captured = warnings.clone();
|
||||
let interceptor = CompactWorkerInterceptor::new(
|
||||
tracker.clone(),
|
||||
150,
|
||||
60,
|
||||
20,
|
||||
Some(Arc::new(move |message| {
|
||||
captured.lock().unwrap().push(message);
|
||||
})),
|
||||
);
|
||||
let mut context = vec![Item::user_message("hello")];
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
PreRequestAction::Continue
|
||||
));
|
||||
tracker.record_usage(&make_usage(100));
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
PreRequestAction::ContinueWith(items)
|
||||
if items.len() == 1 && items[0].as_text().unwrap_or_default().contains("write_summary")
|
||||
));
|
||||
assert_eq!(warnings.lock().unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compact_worker_interceptor_cancels_when_occupancy_exceeds_cap() {
|
||||
let tracker = Arc::new(UsageTracker::new());
|
||||
let interceptor = CompactWorkerInterceptor::new(tracker.clone(), 99, 0, 0, None);
|
||||
let mut context = vec![Item::user_message("hello")];
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
PreRequestAction::Continue
|
||||
));
|
||||
tracker.record_usage(&make_usage(100));
|
||||
|
||||
assert!(matches!(
|
||||
interceptor.pre_llm_request(&mut context).await,
|
||||
PreRequestAction::Cancel(message) if message.contains("occupancy")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mark_read_required_records_and_deducts_budget() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let path = tmp.path().join("hello.txt");
|
||||
std::fs::write(&path, "hello world\n").unwrap();
|
||||
|
||||
let ctx = Arc::new(Mutex::new(CompactWorkerContext::with_budget(1_000)));
|
||||
let tool: Arc<dyn Tool> = Arc::new(MarkReadRequiredTool {
|
||||
fs: make_fs(tmp.path()),
|
||||
ctx: ctx.clone(),
|
||||
});
|
||||
let input = serde_json::json!({ "file_path": path.to_str().unwrap() }).to_string();
|
||||
let out = tool.execute(&input, Default::default()).await.unwrap();
|
||||
|
||||
assert!(out.summary.starts_with("Marked"));
|
||||
let guard = ctx.lock().unwrap();
|
||||
assert_eq!(guard.read_required.len(), 1);
|
||||
assert!(guard.auto_read_consumed > 0);
|
||||
assert!(guard.auto_read_consumed <= 1_000);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mark_read_required_rejects_over_budget() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let path = tmp.path().join("big.txt");
|
||||
std::fs::write(&path, "x".repeat(4_096)).unwrap(); // ≈1024 tokens
|
||||
|
||||
let ctx = Arc::new(Mutex::new(CompactWorkerContext::with_budget(100)));
|
||||
let tool: Arc<dyn Tool> = Arc::new(MarkReadRequiredTool {
|
||||
fs: make_fs(tmp.path()),
|
||||
ctx: ctx.clone(),
|
||||
});
|
||||
let input = serde_json::json!({ "file_path": path.to_str().unwrap() }).to_string();
|
||||
let res = tool.execute(&input, Default::default()).await;
|
||||
|
||||
assert!(matches!(res, Err(ToolError::ExecutionFailed(_))));
|
||||
let guard = ctx.lock().unwrap();
|
||||
assert!(guard.read_required.is_empty());
|
||||
assert_eq!(guard.auto_read_consumed, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_summary_overwrites_previous_call() {
|
||||
let ctx = Arc::new(Mutex::new(CompactWorkerContext::with_budget(0)));
|
||||
let tool: Arc<dyn Tool> = Arc::new(WriteSummaryTool { ctx: ctx.clone() });
|
||||
|
||||
let first = serde_json::json!({ "text": "first" }).to_string();
|
||||
let out1 = tool.execute(&first, Default::default()).await.unwrap();
|
||||
assert!(out1.summary.contains("recorded"));
|
||||
|
||||
let second = serde_json::json!({ "text": "second" }).to_string();
|
||||
let out2 = tool.execute(&second, Default::default()).await.unwrap();
|
||||
assert!(out2.summary.contains("replaced"));
|
||||
|
||||
assert_eq!(ctx.lock().unwrap().summary.as_deref(), Some("second"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn add_reference_deduplicates() {
|
||||
let ctx = Arc::new(Mutex::new(CompactWorkerContext::with_budget(0)));
|
||||
let tool: Arc<dyn Tool> = Arc::new(AddReferenceTool { ctx: ctx.clone() });
|
||||
|
||||
let p = "/abs/path.rs";
|
||||
let input = serde_json::json!({ "file_path": p }).to_string();
|
||||
tool.execute(&input, Default::default()).await.unwrap();
|
||||
tool.execute(&input, Default::default()).await.unwrap();
|
||||
|
||||
let guard = ctx.lock().unwrap();
|
||||
assert_eq!(guard.references.len(), 1);
|
||||
assert_eq!(guard.references[0], PathBuf::from(p));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn search_session_log_returns_bounded_hits_without_full_tool_content() {
|
||||
let items = Arc::new(vec![
|
||||
Item::user_message("investigate compact failure"),
|
||||
Item::tool_result_with_content(
|
||||
"call-1",
|
||||
"read trace with compact failure",
|
||||
"very large raw trace body with secret detail",
|
||||
),
|
||||
]);
|
||||
let tool: Arc<dyn Tool> = Arc::new(SearchSessionLogTool {
|
||||
state: Arc::new(SessionLogToolState { items }),
|
||||
});
|
||||
let input = serde_json::json!({ "query": "compact", "limit": 10 }).to_string();
|
||||
let out = tool.execute(&input, Default::default()).await.unwrap();
|
||||
let content = out.content.unwrap();
|
||||
|
||||
assert!(content.contains("investigate compact failure"));
|
||||
assert!(content.contains("read trace with compact failure"));
|
||||
assert!(!content.contains("secret detail"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_session_items_full_mode_can_read_tool_result_content() {
|
||||
let items = Arc::new(vec![Item::tool_result_with_content(
|
||||
"call-1",
|
||||
"read trace",
|
||||
"raw trace detail",
|
||||
)]);
|
||||
let tool: Arc<dyn Tool> = Arc::new(ReadSessionItemsTool {
|
||||
state: Arc::new(SessionLogToolState { items }),
|
||||
});
|
||||
let input = serde_json::json!({ "offset": 0, "limit": 1, "mode": "full" }).to_string();
|
||||
let out = tool.execute(&input, Default::default()).await.unwrap();
|
||||
let content = out.content.unwrap();
|
||||
|
||||
assert!(content.contains("raw trace detail"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slice_lines_handles_offset_and_limit() {
|
||||
let text = "a\nb\nc\nd";
|
||||
assert_eq!(slice_lines(text, 0, None), "a\nb\nc\nd");
|
||||
assert_eq!(slice_lines(text, 1, Some(2)), "b\nc");
|
||||
assert_eq!(slice_lines(text, 10, None), "");
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,14 @@
|
||||
//! Built-in internal feature modules.
|
||||
//!
|
||||
//! These modules are compiled into the Worker host and contribute through the
|
||||
//! same descriptor-approved registry path used by feature modules. They are not
|
||||
//! an external plugin-loading surface.
|
||||
|
||||
pub mod task;
|
||||
pub mod ticket;
|
||||
|
||||
pub use task::{TaskFeature, task_tools_feature};
|
||||
pub use ticket::{
|
||||
TicketFeature, TicketFeatureAccess, ticket_tools_feature, ticket_tools_feature_with_access,
|
||||
ticket_tools_feature_with_options,
|
||||
};
|
||||
@@ -0,0 +1,573 @@
|
||||
//! Task tools built-in feature module.
|
||||
//!
|
||||
//! The built-in Task feature owns the session-lifetime [`TaskStore`] shared by
|
||||
//! the Task tools and reminder hooks. Worker hosts install this module through the
|
||||
//! feature contribution boundary and use its narrow snapshot surface for
|
||||
//! restore/rewind/compaction compatibility.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_engine::Item;
|
||||
|
||||
mod store;
|
||||
mod tool_impl;
|
||||
|
||||
pub(crate) use self::tool_impl::task_tools;
|
||||
use store::snapshot_overview;
|
||||
pub(crate) use store::{TaskEntry, TaskStatus, TaskStore};
|
||||
|
||||
use crate::feature::{
|
||||
FeatureDescriptor, FeatureHookPoint, FeatureInstallContext, FeatureInstallError, FeatureModule,
|
||||
HookDeclaration, ToolContribution, ToolDeclaration,
|
||||
};
|
||||
use crate::hook::{
|
||||
Hook, HookPreRequestAction, HookPreToolAction, PreLlmRequest, PreRequestContext, PreToolCall,
|
||||
ToolCallSummary,
|
||||
};
|
||||
|
||||
const TASK_REMINDER_REQUEST_THRESHOLD: usize = 24;
|
||||
const TASK_REMINDER_COOLDOWN_REQUESTS: usize = 24;
|
||||
const TASK_MANAGEMENT_TOOL_NAMES: [&str; 2] = ["TaskCreate", "TaskUpdate"];
|
||||
|
||||
/// Construct the built-in Task feature module with a fresh session store.
|
||||
///
|
||||
/// The returned module contributes `TaskCreate`, `TaskUpdate`, `TaskGet`, and
|
||||
/// `TaskList` through descriptor-approved tool registration, plus built-in hooks
|
||||
/// that maintain Task-reminder state. Normal ToolRegistry and PreToolCall
|
||||
/// permission policy still applies at call time.
|
||||
pub fn task_tools_feature() -> TaskFeature {
|
||||
TaskFeature::new()
|
||||
}
|
||||
|
||||
/// Built-in Task feature state and contribution module.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct TaskFeature {
|
||||
state: Arc<TaskFeatureState>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct TaskFeatureState {
|
||||
task_store: TaskStore,
|
||||
reminder_state: TaskReminderState,
|
||||
}
|
||||
|
||||
impl TaskFeature {
|
||||
pub fn new() -> Self {
|
||||
Self::from_store(TaskStore::new())
|
||||
}
|
||||
|
||||
pub fn from_history(history: &[Item]) -> Self {
|
||||
Self::from_store(TaskStore::from_history(history))
|
||||
}
|
||||
|
||||
fn from_store(task_store: TaskStore) -> Self {
|
||||
Self {
|
||||
state: Arc::new(TaskFeatureState {
|
||||
task_store,
|
||||
reminder_state: TaskReminderState::new(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Restore the feature-owned store by replaying durable history into the
|
||||
/// existing shared store handle. Existing Task tool instances and hooks keep
|
||||
/// pointing at the same feature-owned store after rewind.
|
||||
pub fn restore_from_history(&self, history: &[Item]) {
|
||||
let restored = TaskStore::from_history(history);
|
||||
self.state.task_store.replace_with(restored.list());
|
||||
}
|
||||
|
||||
/// Feature-owned snapshot text used by compaction to preserve Task state.
|
||||
pub fn snapshot_text(&self) -> String {
|
||||
self.state.task_store.snapshot_text()
|
||||
}
|
||||
|
||||
/// Feature-owned compact summary used for the synthetic TaskList result.
|
||||
pub fn snapshot_overview(&self) -> String {
|
||||
snapshot_overview(&self.state.task_store.list())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn task_store(&self) -> TaskStore {
|
||||
self.state.task_store.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TaskFeature {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl FeatureModule for TaskFeature {
|
||||
fn descriptor(&self) -> FeatureDescriptor {
|
||||
FeatureDescriptor::builtin("task-tools", "Task tools")
|
||||
.with_description("Session-lifetime task tracking builtin tools")
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"TaskCreate",
|
||||
"Create a session-lifetime user-visible task",
|
||||
))
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"TaskUpdate",
|
||||
"Update a session-lifetime user-visible task",
|
||||
))
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"TaskGet",
|
||||
"Get one session-lifetime user-visible task",
|
||||
))
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"TaskList",
|
||||
"List session-lifetime user-visible tasks",
|
||||
))
|
||||
.with_hook(HookDeclaration::new(
|
||||
"task-reminder-pre-request",
|
||||
FeatureHookPoint::PreRequest,
|
||||
))
|
||||
.with_hook(HookDeclaration::new(
|
||||
"task-reminder-tool-usage",
|
||||
FeatureHookPoint::PreToolCall,
|
||||
))
|
||||
}
|
||||
|
||||
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
|
||||
let names = ["TaskCreate", "TaskList", "TaskGet", "TaskUpdate"];
|
||||
for (name, definition) in names
|
||||
.into_iter()
|
||||
.zip(task_tools(self.state.task_store.clone()))
|
||||
{
|
||||
context
|
||||
.tools()
|
||||
.register(ToolContribution::new(name, definition))?;
|
||||
}
|
||||
|
||||
context.hooks().add_pre_request(
|
||||
"task-reminder-pre-request",
|
||||
TaskReminderPreRequestHook {
|
||||
state: Arc::clone(&self.state),
|
||||
},
|
||||
)?;
|
||||
context.hooks().add_pre_tool_call(
|
||||
"task-reminder-tool-usage",
|
||||
TaskReminderToolUsageHook {
|
||||
state: Arc::clone(&self.state),
|
||||
},
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct TaskReminderState {
|
||||
requests_since_last_task_management: AtomicUsize,
|
||||
requests_since_last_reminder: AtomicUsize,
|
||||
}
|
||||
|
||||
impl Default for TaskReminderState {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
requests_since_last_task_management: AtomicUsize::new(0),
|
||||
requests_since_last_reminder: AtomicUsize::new(TASK_REMINDER_COOLDOWN_REQUESTS),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TaskReminderState {
|
||||
fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
fn note_request(&self) -> (usize, usize) {
|
||||
let since_task_management = self
|
||||
.requests_since_last_task_management
|
||||
.fetch_add(1, Ordering::Relaxed)
|
||||
.saturating_add(1);
|
||||
let since_reminder = self
|
||||
.requests_since_last_reminder
|
||||
.fetch_add(1, Ordering::Relaxed)
|
||||
.saturating_add(1);
|
||||
(since_task_management, since_reminder)
|
||||
}
|
||||
|
||||
fn note_task_management(&self) {
|
||||
self.requests_since_last_task_management
|
||||
.store(0, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
fn note_reminder(&self) {
|
||||
self.requests_since_last_reminder
|
||||
.store(0, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
struct TaskReminderPreRequestHook {
|
||||
state: Arc<TaskFeatureState>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreLlmRequest> for TaskReminderPreRequestHook {
|
||||
async fn call(&self, input: &PreRequestContext) -> HookPreRequestAction {
|
||||
let active_tasks: Vec<TaskEntry> = self
|
||||
.state
|
||||
.task_store
|
||||
.list()
|
||||
.into_iter()
|
||||
.filter(|task| matches!(task.status, TaskStatus::Pending | TaskStatus::Inprogress))
|
||||
.collect();
|
||||
if active_tasks.is_empty() {
|
||||
return HookPreRequestAction::Continue;
|
||||
}
|
||||
|
||||
let (since_task_management, since_reminder) = self.state.reminder_state.note_request();
|
||||
if since_task_management < TASK_REMINDER_REQUEST_THRESHOLD
|
||||
|| since_reminder < TASK_REMINDER_COOLDOWN_REQUESTS
|
||||
{
|
||||
return HookPreRequestAction::Continue;
|
||||
}
|
||||
|
||||
if let Some(system_items) = input.system_items() {
|
||||
self.state.reminder_state.note_reminder();
|
||||
system_items.append_task_reminder(render_task_reminder_body(&active_tasks));
|
||||
}
|
||||
HookPreRequestAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
struct TaskReminderToolUsageHook {
|
||||
state: Arc<TaskFeatureState>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreToolCall> for TaskReminderToolUsageHook {
|
||||
async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction {
|
||||
if is_task_management_tool(&input.tool_name) {
|
||||
self.state.reminder_state.note_task_management();
|
||||
}
|
||||
HookPreToolAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
fn is_task_management_tool(name: &str) -> bool {
|
||||
TASK_MANAGEMENT_TOOL_NAMES.contains(&name)
|
||||
}
|
||||
|
||||
fn render_task_reminder_body(active_tasks: &[TaskEntry]) -> String {
|
||||
let mut body = String::from(
|
||||
"Active session tasks are still open. If progress changed, call TaskUpdate.\n",
|
||||
);
|
||||
for task in active_tasks {
|
||||
body.push_str(&format!(
|
||||
"- taskid {} ({}) {}\n",
|
||||
task.taskid, task.status, task.subject
|
||||
));
|
||||
}
|
||||
body.trim_end_matches('\n').to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use session_store::{SystemItem, SystemReminderSource};
|
||||
|
||||
use super::*;
|
||||
use crate::hook::{PreRequestInfo, SystemItemAppendHandle};
|
||||
|
||||
fn pre_request_context(pending: Arc<Mutex<Vec<SystemItem>>>) -> PreRequestContext {
|
||||
PreRequestContext::new(
|
||||
PreRequestInfo {
|
||||
item_count: 1,
|
||||
estimated_tokens: None,
|
||||
turn_index: 0,
|
||||
tool_calls_this_turn: 0,
|
||||
},
|
||||
Some(SystemItemAppendHandle::new(pending)),
|
||||
)
|
||||
}
|
||||
|
||||
fn tool_summary(name: &str) -> ToolCallSummary {
|
||||
ToolCallSummary {
|
||||
call_id: "call-id".into(),
|
||||
tool_name: name.into(),
|
||||
arguments: serde_json::json!({}),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn task_reminder_hook_appends_after_inactive_request_threshold() {
|
||||
let feature = TaskFeature::new();
|
||||
feature
|
||||
.task_store()
|
||||
.create("keep going".into(), "long task description".into());
|
||||
let hook = TaskReminderPreRequestHook {
|
||||
state: Arc::clone(&feature.state),
|
||||
};
|
||||
let pending = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD - 1 {
|
||||
let _ = hook.call(&pre_request_context(Arc::clone(&pending))).await;
|
||||
assert!(pending.lock().expect("pending queue poisoned").is_empty());
|
||||
}
|
||||
let _ = hook.call(&pre_request_context(Arc::clone(&pending))).await;
|
||||
|
||||
let queued = pending.lock().expect("pending queue poisoned");
|
||||
assert_eq!(queued.len(), 1);
|
||||
let SystemItem::TaskReminder { body, .. } = &queued[0] else {
|
||||
panic!("unexpected system item: {:?}", queued[0]);
|
||||
};
|
||||
assert_eq!(body.matches("<system-reminder>").count(), 1);
|
||||
assert_eq!(body.matches("</system-reminder>").count(), 1);
|
||||
assert!(body.contains("taskid 1"));
|
||||
assert!(body.contains("pending"));
|
||||
assert!(body.contains("keep going"));
|
||||
assert!(!body.contains("long task description"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn task_reminder_hook_retains_source() {
|
||||
let feature = TaskFeature::new();
|
||||
feature.task_store().create("typed".into(), String::new());
|
||||
let hook = TaskReminderPreRequestHook {
|
||||
state: Arc::clone(&feature.state),
|
||||
};
|
||||
let pending = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD {
|
||||
let _ = hook.call(&pre_request_context(Arc::clone(&pending))).await;
|
||||
}
|
||||
|
||||
let queued = pending.lock().expect("pending queue poisoned");
|
||||
let SystemItem::TaskReminder { source, body } = &queued[0] else {
|
||||
panic!("unexpected system item: {:?}", queued[0]);
|
||||
};
|
||||
assert_eq!(*source, SystemReminderSource::TaskInactivity);
|
||||
assert_eq!(body.matches("<system-reminder>").count(), 1);
|
||||
assert_eq!(body.matches("</system-reminder>").count(), 1);
|
||||
assert!(body.contains("typed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_task_reminder_body_is_unwrapped_for_system_reminder_helper() {
|
||||
let feature = TaskFeature::new();
|
||||
let task = feature.task_store().create("body".into(), String::new());
|
||||
let body = render_task_reminder_body(&[task]);
|
||||
|
||||
assert!(!body.contains("<system-reminder>"));
|
||||
assert!(!body.contains("</system-reminder>"));
|
||||
assert!(body.contains("TaskUpdate"));
|
||||
assert!(body.contains("taskid 1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_reminder_state_starts_with_initial_cooldown_elapsed() {
|
||||
let state = TaskReminderState::new();
|
||||
|
||||
assert_eq!(
|
||||
state.requests_since_last_reminder.load(Ordering::Relaxed),
|
||||
TASK_REMINDER_COOLDOWN_REQUESTS
|
||||
);
|
||||
assert_eq!(
|
||||
state
|
||||
.requests_since_last_task_management
|
||||
.load(Ordering::Relaxed),
|
||||
0
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn task_management_tool_call_resets_reminder_inactivity_counter() {
|
||||
let feature = TaskFeature::new();
|
||||
feature
|
||||
.task_store()
|
||||
.create("track me".into(), String::new());
|
||||
let pre_request = TaskReminderPreRequestHook {
|
||||
state: Arc::clone(&feature.state),
|
||||
};
|
||||
let pre_tool = TaskReminderToolUsageHook {
|
||||
state: Arc::clone(&feature.state),
|
||||
};
|
||||
let pending = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD - 1 {
|
||||
let _ = pre_request
|
||||
.call(&pre_request_context(Arc::clone(&pending)))
|
||||
.await;
|
||||
assert!(pending.lock().expect("pending queue poisoned").is_empty());
|
||||
}
|
||||
let _ = pre_tool.call(&tool_summary("TaskUpdate")).await;
|
||||
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD - 1 {
|
||||
let _ = pre_request
|
||||
.call(&pre_request_context(Arc::clone(&pending)))
|
||||
.await;
|
||||
assert!(pending.lock().expect("pending queue poisoned").is_empty());
|
||||
}
|
||||
let _ = pre_request
|
||||
.call(&pre_request_context(Arc::clone(&pending)))
|
||||
.await;
|
||||
assert_eq!(pending.lock().expect("pending queue poisoned").len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn task_reminder_respects_cooldown_after_reminder() {
|
||||
let feature = TaskFeature::new();
|
||||
feature
|
||||
.task_store()
|
||||
.create("cooldown".into(), String::new());
|
||||
let hook = TaskReminderPreRequestHook {
|
||||
state: Arc::clone(&feature.state),
|
||||
};
|
||||
let pending = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD {
|
||||
let _ = hook.call(&pre_request_context(Arc::clone(&pending))).await;
|
||||
}
|
||||
pending.lock().expect("pending queue poisoned").clear();
|
||||
for _ in 0..TASK_REMINDER_COOLDOWN_REQUESTS - 1 {
|
||||
let _ = hook.call(&pre_request_context(Arc::clone(&pending))).await;
|
||||
assert!(pending.lock().expect("pending queue poisoned").is_empty());
|
||||
}
|
||||
let _ = hook.call(&pre_request_context(Arc::clone(&pending))).await;
|
||||
assert_eq!(pending.lock().expect("pending queue poisoned").len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn task_reminder_is_silent_when_no_active_tasks_exist() {
|
||||
let feature = TaskFeature::new();
|
||||
let done = feature
|
||||
.task_store()
|
||||
.create("done".into(), String::new())
|
||||
.taskid;
|
||||
feature
|
||||
.task_store()
|
||||
.update(done, Some(TaskStatus::Completed), None, None)
|
||||
.expect("complete task");
|
||||
let hook = TaskReminderPreRequestHook {
|
||||
state: Arc::clone(&feature.state),
|
||||
};
|
||||
let pending = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD * 2 {
|
||||
let _ = hook.call(&pre_request_context(Arc::clone(&pending))).await;
|
||||
assert!(pending.lock().expect("pending queue poisoned").is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inactive_requests_without_active_tasks_do_not_prime_task_reminder() {
|
||||
let feature = TaskFeature::new();
|
||||
let hook = TaskReminderPreRequestHook {
|
||||
state: Arc::clone(&feature.state),
|
||||
};
|
||||
let pending = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD * 2 {
|
||||
let _ = hook.call(&pre_request_context(Arc::clone(&pending))).await;
|
||||
assert!(pending.lock().expect("pending queue poisoned").is_empty());
|
||||
}
|
||||
|
||||
feature
|
||||
.task_store()
|
||||
.create("new active".into(), String::new());
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD - 1 {
|
||||
let _ = hook.call(&pre_request_context(Arc::clone(&pending))).await;
|
||||
assert!(pending.lock().expect("pending queue poisoned").is_empty());
|
||||
}
|
||||
let _ = hook.call(&pre_request_context(Arc::clone(&pending))).await;
|
||||
assert_eq!(pending.lock().expect("pending queue poisoned").len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn task_create_reset_does_not_block_first_reminder_cooldown() {
|
||||
let feature = TaskFeature::new();
|
||||
let pre_request = TaskReminderPreRequestHook {
|
||||
state: Arc::clone(&feature.state),
|
||||
};
|
||||
let pre_tool = TaskReminderToolUsageHook {
|
||||
state: Arc::clone(&feature.state),
|
||||
};
|
||||
let pending = Arc::new(Mutex::new(Vec::new()));
|
||||
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD * 2 {
|
||||
let _ = pre_request
|
||||
.call(&pre_request_context(Arc::clone(&pending)))
|
||||
.await;
|
||||
assert!(pending.lock().expect("pending queue poisoned").is_empty());
|
||||
}
|
||||
|
||||
let _ = pre_tool.call(&tool_summary("TaskCreate")).await;
|
||||
feature
|
||||
.task_store()
|
||||
.create("created after idle".into(), String::new());
|
||||
assert_eq!(
|
||||
feature
|
||||
.state
|
||||
.reminder_state
|
||||
.requests_since_last_reminder
|
||||
.load(Ordering::Relaxed),
|
||||
TASK_REMINDER_COOLDOWN_REQUESTS,
|
||||
"TaskCreate reset must not clear the initial reminder cooldown"
|
||||
);
|
||||
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD - 1 {
|
||||
let _ = pre_request
|
||||
.call(&pre_request_context(Arc::clone(&pending)))
|
||||
.await;
|
||||
assert!(pending.lock().expect("pending queue poisoned").is_empty());
|
||||
}
|
||||
let _ = pre_request
|
||||
.call(&pre_request_context(Arc::clone(&pending)))
|
||||
.await;
|
||||
assert_eq!(pending.lock().expect("pending queue poisoned").len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_system_item_handle_does_not_mark_reminder_sent() {
|
||||
let feature = TaskFeature::new();
|
||||
feature.task_store().create("handle".into(), String::new());
|
||||
let hook = TaskReminderPreRequestHook {
|
||||
state: Arc::clone(&feature.state),
|
||||
};
|
||||
let no_handle = PreRequestContext::new(
|
||||
PreRequestInfo {
|
||||
item_count: 1,
|
||||
estimated_tokens: None,
|
||||
turn_index: 0,
|
||||
tool_calls_this_turn: 0,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
for _ in 0..TASK_REMINDER_REQUEST_THRESHOLD {
|
||||
let _ = hook.call(&no_handle).await;
|
||||
}
|
||||
assert_eq!(
|
||||
feature
|
||||
.state
|
||||
.reminder_state
|
||||
.requests_since_last_reminder
|
||||
.load(Ordering::Relaxed),
|
||||
TASK_REMINDER_COOLDOWN_REQUESTS + TASK_REMINDER_REQUEST_THRESHOLD,
|
||||
"without a handle the hook must not record a reminder as emitted"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restore_from_history_keeps_existing_store_handle_for_installed_tools() {
|
||||
let feature = TaskFeature::new();
|
||||
let handle = feature.task_store();
|
||||
handle.create("old".into(), String::new());
|
||||
let history = vec![Item::tool_call(
|
||||
"c1",
|
||||
"TaskCreate",
|
||||
r#"{"subject":"restored","description":"from history"}"#,
|
||||
)];
|
||||
|
||||
feature.restore_from_history(&history);
|
||||
|
||||
let tasks = handle.list();
|
||||
assert_eq!(tasks.len(), 1);
|
||||
assert_eq!(tasks[0].subject, "restored");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,458 @@
|
||||
//! Task domain state and snapshot/replay support.
|
||||
//!
|
||||
//! The store survives compaction and Worker restart by replaying TaskCreate /
|
||||
//! TaskUpdate tool-call arguments and compacted TaskStore snapshots from
|
||||
//! persisted history.
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use llm_engine::Item;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum TaskStatus {
|
||||
Pending,
|
||||
Inprogress,
|
||||
Completed,
|
||||
Deleted,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TaskStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
let s = match self {
|
||||
Self::Pending => "pending",
|
||||
Self::Inprogress => "inprogress",
|
||||
Self::Completed => "completed",
|
||||
Self::Deleted => "deleted",
|
||||
};
|
||||
f.write_str(s)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
pub struct TaskEntry {
|
||||
pub taskid: u64,
|
||||
pub status: TaskStatus,
|
||||
pub subject: String,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct Inner {
|
||||
next_taskid: u64,
|
||||
tasks: Vec<TaskEntry>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TaskStore {
|
||||
inner: Arc<Mutex<Inner>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
pub struct TaskSnapshot {
|
||||
pub tasks: Vec<TaskEntry>,
|
||||
}
|
||||
|
||||
impl TaskStore {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(Inner {
|
||||
next_taskid: 1,
|
||||
tasks: Vec::new(),
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn create(&self, subject: String, description: String) -> TaskEntry {
|
||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let task = TaskEntry {
|
||||
taskid: inner.next_taskid,
|
||||
status: TaskStatus::Pending,
|
||||
subject,
|
||||
description,
|
||||
};
|
||||
inner.next_taskid = inner.next_taskid.saturating_add(1);
|
||||
inner.tasks.push(task.clone());
|
||||
task
|
||||
}
|
||||
|
||||
pub fn list(&self) -> Vec<TaskEntry> {
|
||||
self.inner
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.tasks
|
||||
.clone()
|
||||
}
|
||||
|
||||
pub fn get(&self, taskid: u64) -> Option<TaskEntry> {
|
||||
self.inner
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.tasks
|
||||
.iter()
|
||||
.find(|t| t.taskid == taskid)
|
||||
.cloned()
|
||||
}
|
||||
|
||||
pub fn update(
|
||||
&self,
|
||||
taskid: u64,
|
||||
status: Option<TaskStatus>,
|
||||
subject: Option<String>,
|
||||
description: Option<String>,
|
||||
) -> Result<TaskEntry, TaskStoreError> {
|
||||
if status.is_none() && subject.is_none() && description.is_none() {
|
||||
return Err(TaskStoreError::NoFields);
|
||||
}
|
||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let task = inner
|
||||
.tasks
|
||||
.iter_mut()
|
||||
.find(|t| t.taskid == taskid)
|
||||
.ok_or(TaskStoreError::Missing(taskid))?;
|
||||
if let Some(status) = status {
|
||||
task.status = status;
|
||||
}
|
||||
if let Some(subject) = subject {
|
||||
task.subject = subject;
|
||||
}
|
||||
if let Some(description) = description {
|
||||
task.description = description;
|
||||
}
|
||||
Ok(task.clone())
|
||||
}
|
||||
|
||||
pub fn snapshot(&self) -> TaskSnapshot {
|
||||
TaskSnapshot { tasks: self.list() }
|
||||
}
|
||||
|
||||
pub fn replay_history(&self, history: &[Item]) {
|
||||
for item in history {
|
||||
match item {
|
||||
Item::Message { content, .. } => {
|
||||
for part in content {
|
||||
let text = part.as_text();
|
||||
if let Some(snapshot) = parse_compact_snapshot_text(text) {
|
||||
self.replace_with(snapshot);
|
||||
}
|
||||
}
|
||||
}
|
||||
Item::ToolCall {
|
||||
name, arguments, ..
|
||||
} => match name.as_str() {
|
||||
"TaskCreate" => {
|
||||
if let Ok(params) =
|
||||
serde_json::from_str::<ReplayTaskCreateParams>(arguments)
|
||||
{
|
||||
let _ = self.create(params.subject, params.description);
|
||||
}
|
||||
}
|
||||
"TaskUpdate" => {
|
||||
if let Ok(params) =
|
||||
serde_json::from_str::<ReplayTaskUpdateParams>(arguments)
|
||||
{
|
||||
let _ = self.update(
|
||||
params.taskid,
|
||||
params.status,
|
||||
params.subject,
|
||||
params.description,
|
||||
);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
},
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn replace_with(&self, tasks: Vec<TaskEntry>) {
|
||||
let next_taskid = tasks
|
||||
.iter()
|
||||
.map(|t| t.taskid)
|
||||
.max()
|
||||
.unwrap_or(0)
|
||||
.saturating_add(1)
|
||||
.max(1);
|
||||
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
||||
inner.tasks = tasks;
|
||||
inner.next_taskid = next_taskid;
|
||||
}
|
||||
|
||||
pub fn from_history(history: &[Item]) -> Self {
|
||||
let store = Self::new();
|
||||
store.replay_history(history);
|
||||
store
|
||||
}
|
||||
|
||||
pub fn snapshot_text(&self) -> String {
|
||||
let snapshot = self.snapshot();
|
||||
render_snapshot(&snapshot.tasks)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum TaskStoreError {
|
||||
Missing(u64),
|
||||
NoFields,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TaskStoreError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Missing(id) => write!(f, "taskid {id} not found"),
|
||||
Self::NoFields => {
|
||||
f.write_str("at least one of status, subject, description is required")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for TaskStoreError {}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ReplayTaskCreateParams {
|
||||
subject: String,
|
||||
description: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ReplayTaskUpdateParams {
|
||||
taskid: u64,
|
||||
#[serde(default)]
|
||||
status: Option<TaskStatus>,
|
||||
#[serde(default)]
|
||||
subject: Option<String>,
|
||||
#[serde(default)]
|
||||
description: Option<String>,
|
||||
}
|
||||
|
||||
pub fn snapshot_overview(tasks: &[TaskEntry]) -> String {
|
||||
let pending = tasks
|
||||
.iter()
|
||||
.filter(|t| t.status == TaskStatus::Pending)
|
||||
.count();
|
||||
let inprogress = tasks
|
||||
.iter()
|
||||
.filter(|t| t.status == TaskStatus::Inprogress)
|
||||
.count();
|
||||
let completed = tasks
|
||||
.iter()
|
||||
.filter(|t| t.status == TaskStatus::Completed)
|
||||
.count();
|
||||
let deleted = tasks
|
||||
.iter()
|
||||
.filter(|t| t.status == TaskStatus::Deleted)
|
||||
.count();
|
||||
format!(
|
||||
"TaskStore: {} task(s) (pending: {pending}, inprogress: {inprogress}, completed: {completed}, deleted: {deleted})",
|
||||
tasks.len()
|
||||
)
|
||||
}
|
||||
|
||||
pub fn render_snapshot(tasks: &[TaskEntry]) -> String {
|
||||
let snapshot = TaskSnapshot {
|
||||
tasks: tasks.to_vec(),
|
||||
};
|
||||
let json =
|
||||
serde_json::to_string_pretty(&snapshot).unwrap_or_else(|_| String::from("{\"tasks\":[]}"));
|
||||
format!("{}\n\n```json\n{}\n```\n", snapshot_overview(tasks), json)
|
||||
}
|
||||
|
||||
pub(super) fn parse_compact_snapshot_text(text: &str) -> Option<Vec<TaskEntry>> {
|
||||
if !text.starts_with("[Session TaskStore snapshot]") {
|
||||
return None;
|
||||
}
|
||||
let start_marker = "```json\n";
|
||||
let end_marker = "\n```";
|
||||
let start = text.find(start_marker)? + start_marker.len();
|
||||
let rest = &text[start..];
|
||||
let end = rest.find(end_marker)?;
|
||||
let snapshot: TaskSnapshot = serde_json::from_str(&rest[..end]).ok()?;
|
||||
Some(snapshot.tasks)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn replay_history_reconstructs_store_and_ignores_malformed_calls() {
|
||||
let history = vec![
|
||||
Item::tool_call("c1", "TaskCreate", r#"{"subject":"a","description":"A"}"#),
|
||||
Item::tool_call("bad", "TaskCreate", r#"{"subject":1}"#),
|
||||
Item::tool_call("c2", "TaskCreate", r#"{"subject":"b","description":"B"}"#),
|
||||
Item::tool_call("u1", "TaskUpdate", r#"{"taskid":2,"status":"completed"}"#),
|
||||
Item::tool_call("bad2", "TaskUpdate", r#"{"taskid":99,"status":"deleted"}"#),
|
||||
];
|
||||
let store = TaskStore::from_history(&history);
|
||||
let tasks = store.list();
|
||||
assert_eq!(tasks.len(), 2);
|
||||
assert_eq!(tasks[0].taskid, 1);
|
||||
assert_eq!(tasks[0].status, TaskStatus::Pending);
|
||||
assert_eq!(tasks[1].taskid, 2);
|
||||
assert_eq!(tasks[1].status, TaskStatus::Completed);
|
||||
}
|
||||
|
||||
/// Wrap snapshot text the way `Worker::try_pre_run_compact` does, so tests
|
||||
/// exercise the exact format that goes through the session log.
|
||||
fn wrap_snapshot_system_message(snapshot: &str) -> String {
|
||||
format!(
|
||||
"[Session TaskStore snapshot]\n\n{snapshot}\n\n\
|
||||
This is the complete session task list preserved across compaction. \
|
||||
The following TaskList tool result presents the same state through the tool lane."
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replay_history_uses_compact_snapshot_and_continues_updates() {
|
||||
let pre = TaskStore::new();
|
||||
pre.create("kept".into(), "from compact".into());
|
||||
pre.update(1, Some(TaskStatus::Inprogress), None, None)
|
||||
.unwrap();
|
||||
let history = vec![
|
||||
Item::system_message(wrap_snapshot_system_message(&pre.snapshot_text())),
|
||||
Item::tool_call("u1", "TaskUpdate", r#"{"taskid":1,"status":"completed"}"#),
|
||||
Item::tool_call(
|
||||
"c2",
|
||||
"TaskCreate",
|
||||
r#"{"subject":"new","description":"after compact"}"#,
|
||||
),
|
||||
];
|
||||
let store = TaskStore::from_history(&history);
|
||||
let tasks = store.list();
|
||||
assert_eq!(tasks.len(), 2);
|
||||
assert_eq!(tasks[0].taskid, 1);
|
||||
assert_eq!(tasks[0].status, TaskStatus::Completed);
|
||||
assert_eq!(tasks[1].taskid, 2);
|
||||
assert_eq!(tasks[1].subject, "new");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_snapshot_supersedes_pre_compact_taskcreates_in_retained() {
|
||||
// Mirrors the post-compact layout: pre-compact `TaskCreate` calls are
|
||||
// preserved verbatim in retained_items, and the snapshot trails them.
|
||||
// The trailing snapshot must reset the store to the captured state so
|
||||
// pre-compact `TaskCreate`s do not surface as duplicates.
|
||||
let pre = TaskStore::new();
|
||||
pre.create("A".into(), "A-desc".into());
|
||||
pre.update(1, Some(TaskStatus::Completed), None, None)
|
||||
.unwrap();
|
||||
pre.create("B".into(), "B-desc".into());
|
||||
pre.update(2, Some(TaskStatus::Inprogress), None, None)
|
||||
.unwrap();
|
||||
let history = vec![
|
||||
Item::tool_call(
|
||||
"c1",
|
||||
"TaskCreate",
|
||||
r#"{"subject":"A","description":"A-desc"}"#,
|
||||
),
|
||||
Item::tool_call("u1", "TaskUpdate", r#"{"taskid":1,"status":"completed"}"#),
|
||||
Item::tool_call(
|
||||
"c2",
|
||||
"TaskCreate",
|
||||
r#"{"subject":"B","description":"B-desc"}"#,
|
||||
),
|
||||
Item::tool_call("u2", "TaskUpdate", r#"{"taskid":2,"status":"inprogress"}"#),
|
||||
Item::system_message(wrap_snapshot_system_message(&pre.snapshot_text())),
|
||||
Item::tool_call("compact-tasklist", "TaskList", "{}"),
|
||||
Item::tool_call(
|
||||
"c3",
|
||||
"TaskCreate",
|
||||
r#"{"subject":"C","description":"after compact"}"#,
|
||||
),
|
||||
];
|
||||
let store = TaskStore::from_history(&history);
|
||||
let tasks = store.list();
|
||||
assert_eq!(tasks.len(), 3);
|
||||
assert_eq!(tasks[0].taskid, 1);
|
||||
assert_eq!(tasks[0].subject, "A");
|
||||
assert_eq!(tasks[0].status, TaskStatus::Completed);
|
||||
assert_eq!(tasks[1].taskid, 2);
|
||||
assert_eq!(tasks[1].subject, "B");
|
||||
assert_eq!(tasks[1].status, TaskStatus::Inprogress);
|
||||
assert_eq!(tasks[2].taskid, 3);
|
||||
assert_eq!(tasks[2].subject, "C");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_round_trips_multiline_subject_and_description() {
|
||||
// Subject / description with embedded newlines and shape-breaking
|
||||
// characters must survive snapshot serialization unchanged.
|
||||
let pre = TaskStore::new();
|
||||
pre.create(
|
||||
"subject with\nembedded newline\n- bullet".into(),
|
||||
"desc:\n status: not-actually-a-field\n ```code fence```".into(),
|
||||
);
|
||||
pre.update(1, Some(TaskStatus::Inprogress), None, None)
|
||||
.unwrap();
|
||||
|
||||
let history = vec![Item::system_message(wrap_snapshot_system_message(
|
||||
&pre.snapshot_text(),
|
||||
))];
|
||||
let store = TaskStore::from_history(&history);
|
||||
let tasks = store.list();
|
||||
assert_eq!(tasks.len(), 1);
|
||||
assert_eq!(tasks[0].subject, "subject with\nembedded newline\n- bullet");
|
||||
assert_eq!(
|
||||
tasks[0].description,
|
||||
"desc:\n status: not-actually-a-field\n ```code fence```"
|
||||
);
|
||||
assert_eq!(tasks[0].status, TaskStatus::Inprogress);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn synthetic_compact_tasklist_pair_is_well_formed() {
|
||||
// Mirrors `Worker::try_pre_run_compact`'s synthetic insertion:
|
||||
// a system snapshot message followed by a TaskList tool_call/tool_result
|
||||
// pair sharing the `compact-tasklist` id. Verify the structural
|
||||
// contract every provider request builder relies on (matched call_id,
|
||||
// tool name, content recoverable to the same TaskStore state).
|
||||
let pre = TaskStore::new();
|
||||
pre.create("plan".into(), "do A then B".into());
|
||||
|
||||
let snapshot_text = pre.snapshot_text();
|
||||
let system = Item::system_message(wrap_snapshot_system_message(&snapshot_text));
|
||||
let call = Item::tool_call("compact-tasklist", "TaskList", "{}");
|
||||
let result = Item::tool_result_with_content(
|
||||
"compact-tasklist",
|
||||
snapshot_overview(&pre.list()),
|
||||
snapshot_text.clone(),
|
||||
);
|
||||
|
||||
// The system message embeds a parseable snapshot.
|
||||
let extracted = system
|
||||
.as_text()
|
||||
.and_then(parse_compact_snapshot_text)
|
||||
.expect("system message should parse as snapshot");
|
||||
assert_eq!(extracted, pre.list());
|
||||
|
||||
// The synthetic call/result pair shares one call_id and carries the
|
||||
// expected tool name + detailed content.
|
||||
match (&call, &result) {
|
||||
(
|
||||
Item::ToolCall {
|
||||
call_id: c_id,
|
||||
name,
|
||||
..
|
||||
},
|
||||
Item::ToolResult {
|
||||
call_id: r_id,
|
||||
content,
|
||||
..
|
||||
},
|
||||
) => {
|
||||
assert_eq!(c_id.as_str(), r_id.as_str());
|
||||
assert_eq!(c_id.as_str(), "compact-tasklist");
|
||||
assert_eq!(name, "TaskList");
|
||||
assert_eq!(content.as_deref(), Some(snapshot_text.as_str()));
|
||||
}
|
||||
other => panic!("unexpected synthetic pair shape: {other:?}"),
|
||||
}
|
||||
|
||||
// Replaying the full triple reconstructs the same TaskStore.
|
||||
let store = TaskStore::from_history(&[system, call, result]);
|
||||
assert_eq!(store.list(), pre.list());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,313 @@
|
||||
//! Task built-in tool implementations.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::store::{TaskEntry, TaskStatus, TaskStore, render_snapshot, snapshot_overview};
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct TaskCreateParams {
|
||||
/// One-line task subject.
|
||||
subject: String,
|
||||
/// Detailed task description.
|
||||
description: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct TaskListParams {}
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct TaskGetParams {
|
||||
taskid: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct TaskUpdateParams {
|
||||
taskid: u64,
|
||||
#[serde(default)]
|
||||
status: Option<TaskStatus>,
|
||||
#[serde(default)]
|
||||
subject: Option<String>,
|
||||
#[serde(default)]
|
||||
description: Option<String>,
|
||||
}
|
||||
|
||||
struct TaskCreateTool {
|
||||
store: TaskStore,
|
||||
}
|
||||
|
||||
struct TaskListTool {
|
||||
store: TaskStore,
|
||||
}
|
||||
|
||||
struct TaskGetTool {
|
||||
store: TaskStore,
|
||||
}
|
||||
|
||||
struct TaskUpdateTool {
|
||||
store: TaskStore,
|
||||
}
|
||||
|
||||
const CREATE_DESCRIPTION: &str = "Create a session-lifetime task only when user-visible \
|
||||
progress tracking is genuinely useful: multiple active tasks must be remembered, or the work \
|
||||
will involve long edits, long-running commands, extended investigation, or interruption-prone \
|
||||
coordination. Do not create a task just because a request has several steps, and do not create \
|
||||
one for short questions, quick checks, single reviews, or one-off commands. Prefer updating an \
|
||||
existing active task over creating a duplicate. Input only `subject` and `description`; `taskid` \
|
||||
is assigned automatically and initial `status` is `pending`.";
|
||||
const LIST_DESCRIPTION: &str = "List every session-lifetime task, including completed and \
|
||||
deleted entries. Tasks are user-visible real-time status for short-term current-work tracking. \
|
||||
Takes an empty object as input.";
|
||||
const GET_DESCRIPTION: &str = "Get one session-lifetime task by `taskid`. Tasks are \
|
||||
user-visible real-time status for short-term current-work tracking. Returns an error if the task \
|
||||
does not exist.";
|
||||
const UPDATE_DESCRIPTION: &str = "Update an existing session-lifetime task when meaningful \
|
||||
progress changes between substantial steps. Tasks are user-visible real-time status, so avoid \
|
||||
churn for trivial substeps. Keep status current with `pending`, `inprogress`, `completed`, or \
|
||||
`deleted`. Provide `taskid` and at least one of `status`, `subject`, or `description`; deletion is \
|
||||
logical (`status = deleted`). If an unexpected problem blocks progress, do not force the next \
|
||||
step: leave the task as-is, summarize the problem to the user, and end the turn.";
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskCreateTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: TaskCreateParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid TaskCreate input: {e}")))?;
|
||||
let created = self.store.create(params.subject, params.description);
|
||||
let tasks = self.store.list();
|
||||
Ok(task_output(
|
||||
format!(
|
||||
"Created task {} ({})\n{}",
|
||||
created.taskid,
|
||||
created.status,
|
||||
snapshot_overview(&tasks)
|
||||
),
|
||||
&created,
|
||||
&tasks,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskListTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let _: TaskListParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid TaskList input: {e}")))?;
|
||||
let tasks = self.store.list();
|
||||
Ok(ToolOutput {
|
||||
summary: snapshot_overview(&tasks),
|
||||
content: Some(render_snapshot(&tasks)),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskGetTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: TaskGetParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid TaskGet input: {e}")))?;
|
||||
let task = self.store.get(params.taskid).ok_or_else(|| {
|
||||
ToolError::ExecutionFailed(format!("taskid {} not found", params.taskid))
|
||||
})?;
|
||||
let content = serde_json::to_string_pretty(&task).unwrap_or_else(|_| format!("{task:?}"));
|
||||
Ok(ToolOutput {
|
||||
summary: format!("Task {} ({}) {}", task.taskid, task.status, task.subject),
|
||||
content: Some(content),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskUpdateTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: TaskUpdateParams = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid TaskUpdate input: {e}")))?;
|
||||
let updated = self
|
||||
.store
|
||||
.update(
|
||||
params.taskid,
|
||||
params.status,
|
||||
params.subject,
|
||||
params.description,
|
||||
)
|
||||
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
|
||||
let tasks = self.store.list();
|
||||
Ok(task_output(
|
||||
format!(
|
||||
"Updated task {} ({})\n{}",
|
||||
updated.taskid,
|
||||
updated.status,
|
||||
snapshot_overview(&tasks)
|
||||
),
|
||||
&updated,
|
||||
&tasks,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn task_output(summary: String, task: &TaskEntry, tasks: &[TaskEntry]) -> ToolOutput {
|
||||
let content = serde_json::json!({
|
||||
"task": task,
|
||||
"snapshot": { "tasks": tasks },
|
||||
});
|
||||
ToolOutput {
|
||||
summary,
|
||||
content: Some(serde_json::to_string_pretty(&content).unwrap_or_default()),
|
||||
}
|
||||
}
|
||||
fn task_create_tool(store: TaskStore) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(TaskCreateParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("TaskCreate")
|
||||
.description(CREATE_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(TaskCreateTool {
|
||||
store: store.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
fn task_list_tool(store: TaskStore) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(TaskListParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("TaskList")
|
||||
.description(LIST_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(TaskListTool {
|
||||
store: store.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
fn task_get_tool(store: TaskStore) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(TaskGetParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("TaskGet")
|
||||
.description(GET_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(TaskGetTool {
|
||||
store: store.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
fn task_update_tool(store: TaskStore) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(TaskUpdateParams);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("TaskUpdate")
|
||||
.description(UPDATE_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(TaskUpdateTool {
|
||||
store: store.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn task_tools(store: TaskStore) -> Vec<ToolDefinition> {
|
||||
vec![
|
||||
task_create_tool(store.clone()),
|
||||
task_list_tool(store.clone()),
|
||||
task_get_tool(store.clone()),
|
||||
task_update_tool(store),
|
||||
]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn tool(def: ToolDefinition) -> Arc<dyn Tool> {
|
||||
let (_, tool) = def();
|
||||
tool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn task_tools_create_list_get_update() {
|
||||
let store = TaskStore::new();
|
||||
let create = tool(task_create_tool(store.clone()));
|
||||
let list = tool(task_list_tool(store.clone()));
|
||||
let get = tool(task_get_tool(store.clone()));
|
||||
let update = tool(task_update_tool(store.clone()));
|
||||
|
||||
let out = create
|
||||
.execute(
|
||||
r#"{"subject":"implement","description":"write code"}"#,
|
||||
Default::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(out.summary.contains("Created task 1"));
|
||||
assert_eq!(store.get(1).unwrap().status, TaskStatus::Pending);
|
||||
|
||||
let out = update
|
||||
.execute(
|
||||
r#"{"taskid":1,"status":"inprogress","subject":"implement tasks"}"#,
|
||||
Default::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(out.summary.contains("Updated task 1"));
|
||||
let task = store.get(1).unwrap();
|
||||
assert_eq!(task.status, TaskStatus::Inprogress);
|
||||
assert_eq!(task.subject, "implement tasks");
|
||||
|
||||
let out = get
|
||||
.execute(r#"{"taskid":1}"#, Default::default())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(out.summary.contains("Task 1 (inprogress)"));
|
||||
assert!(out.content.unwrap().contains("implement tasks"));
|
||||
|
||||
let out = list.execute("{}", Default::default()).await.unwrap();
|
||||
assert!(out.summary.contains("1 task(s)"));
|
||||
let content = out.content.unwrap();
|
||||
assert!(content.contains("\"taskid\": 1"));
|
||||
assert!(content.contains("```json"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn task_update_validates_existing_and_at_least_one_field() {
|
||||
let store = TaskStore::new();
|
||||
store.create("s".into(), "d".into());
|
||||
let update = tool(task_update_tool(store));
|
||||
|
||||
let err = update
|
||||
.execute(r#"{"taskid":1}"#, Default::default())
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.to_string().contains("at least one"));
|
||||
|
||||
let err = update
|
||||
.execute(r#"{"taskid":99,"status":"deleted"}"#, Default::default())
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.to_string().contains("taskid 99 not found"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,614 @@
|
||||
//! Built-in Ticket feature adapter.
|
||||
//!
|
||||
//! The ticket crate owns Ticket domain logic and Tool implementations. This
|
||||
//! module only resolves the local backend root, declares the built-in feature,
|
||||
//! and contributes those tools through the normal feature registry path.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use ticket::{
|
||||
LocalTicketBackend,
|
||||
config::{DEFAULT_TICKET_BACKEND_RELATIVE_PATH, TicketConfig},
|
||||
tool::{
|
||||
TICKET_BASE_READ_ONLY_TOOL_NAMES, TICKET_BASE_TOOL_NAMES,
|
||||
TICKET_ORCHESTRATION_READ_ONLY_TOOL_NAMES, TICKET_ORCHESTRATION_TOOL_NAMES,
|
||||
TICKET_READ_ONLY_TOOL_NAMES, TICKET_TOOL_NAMES, ticket_tool_description, ticket_tools,
|
||||
},
|
||||
};
|
||||
|
||||
use crate::feature::{
|
||||
FeatureDescriptor, FeatureDiagnostic, FeatureInstallContext, FeatureInstallError,
|
||||
FeatureModule, ToolContribution, ToolDeclaration,
|
||||
};
|
||||
|
||||
const FEATURE_ID: &str = "ticket";
|
||||
const FEATURE_NAME: &str = "Ticket tools";
|
||||
const FEATURE_DESCRIPTION: &str = "Typed local Ticket work-item operations over a bounded backend root. \
|
||||
The tools operate through the ticket crate backend and do not grant generic filesystem write scope.";
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum TicketFeatureAccess {
|
||||
/// Status/diagnostic access for views such as Companion that must not mutate Tickets.
|
||||
ReadOnly,
|
||||
/// Full Ticket lifecycle access, including the read-only tools and all mutating Ticket tools.
|
||||
Lifecycle,
|
||||
}
|
||||
|
||||
impl TicketFeatureAccess {
|
||||
pub fn base_tool_names(self) -> &'static [&'static str] {
|
||||
match self {
|
||||
Self::ReadOnly => &TICKET_BASE_READ_ONLY_TOOL_NAMES,
|
||||
Self::Lifecycle => &TICKET_BASE_TOOL_NAMES,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn orchestration_tool_names(self) -> &'static [&'static str] {
|
||||
match self {
|
||||
Self::ReadOnly => &TICKET_ORCHESTRATION_READ_ONLY_TOOL_NAMES,
|
||||
Self::Lifecycle => &TICKET_ORCHESTRATION_TOOL_NAMES,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct TicketFeature {
|
||||
backend_root: PathBuf,
|
||||
record_language: Option<String>,
|
||||
config_error: Option<String>,
|
||||
access: TicketFeatureAccess,
|
||||
include_base_tools: bool,
|
||||
include_orchestration_tools: bool,
|
||||
}
|
||||
|
||||
impl TicketFeature {
|
||||
pub fn new(backend_root: impl Into<PathBuf>) -> Self {
|
||||
Self::new_with_access(backend_root, TicketFeatureAccess::Lifecycle)
|
||||
}
|
||||
|
||||
pub fn new_with_access(backend_root: impl Into<PathBuf>, access: TicketFeatureAccess) -> Self {
|
||||
Self::new_with_options(backend_root, Some(access), true)
|
||||
}
|
||||
|
||||
pub fn new_with_options(
|
||||
backend_root: impl Into<PathBuf>,
|
||||
access: Option<TicketFeatureAccess>,
|
||||
include_orchestration_tools: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
backend_root: backend_root.into(),
|
||||
record_language: None,
|
||||
config_error: None,
|
||||
access: access.unwrap_or(TicketFeatureAccess::Lifecycle),
|
||||
include_base_tools: access.is_some(),
|
||||
include_orchestration_tools,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn for_workspace(workspace: impl AsRef<Path>) -> Self {
|
||||
Self::for_workspace_with_access(workspace, TicketFeatureAccess::Lifecycle)
|
||||
}
|
||||
|
||||
pub fn for_workspace_with_access(
|
||||
workspace: impl AsRef<Path>,
|
||||
access: TicketFeatureAccess,
|
||||
) -> Self {
|
||||
Self::for_workspace_with_options(workspace, Some(access), true)
|
||||
}
|
||||
|
||||
pub fn for_workspace_with_options(
|
||||
workspace: impl AsRef<Path>,
|
||||
access: Option<TicketFeatureAccess>,
|
||||
include_orchestration_tools: bool,
|
||||
) -> Self {
|
||||
let workspace = workspace.as_ref();
|
||||
match TicketConfig::load_workspace(workspace) {
|
||||
Ok(config) => {
|
||||
let backend_root = config.backend_root().to_path_buf();
|
||||
let record_language = config.ticket_record_language().map(str::to_string);
|
||||
let mut feature =
|
||||
Self::new_with_options(backend_root, access, include_orchestration_tools);
|
||||
feature.record_language = record_language;
|
||||
feature
|
||||
}
|
||||
Err(error) => {
|
||||
let access_value = access.unwrap_or(TicketFeatureAccess::Lifecycle);
|
||||
Self {
|
||||
backend_root: workspace.join(DEFAULT_TICKET_BACKEND_RELATIVE_PATH),
|
||||
record_language: None,
|
||||
config_error: Some(error.to_string()),
|
||||
access: access_value,
|
||||
include_base_tools: access.is_some(),
|
||||
include_orchestration_tools,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn backend_root(&self) -> &Path {
|
||||
&self.backend_root
|
||||
}
|
||||
|
||||
pub fn access(&self) -> TicketFeatureAccess {
|
||||
self.access
|
||||
}
|
||||
|
||||
fn enabled_tool_names(&self) -> Vec<&'static str> {
|
||||
if self.include_base_tools && self.include_orchestration_tools {
|
||||
return match self.access {
|
||||
TicketFeatureAccess::ReadOnly => TICKET_READ_ONLY_TOOL_NAMES.to_vec(),
|
||||
TicketFeatureAccess::Lifecycle => TICKET_TOOL_NAMES.to_vec(),
|
||||
};
|
||||
}
|
||||
let mut names = Vec::new();
|
||||
if self.include_base_tools {
|
||||
names.extend_from_slice(self.access.base_tool_names());
|
||||
}
|
||||
if self.include_orchestration_tools {
|
||||
names.extend_from_slice(self.access.orchestration_tool_names());
|
||||
}
|
||||
names
|
||||
}
|
||||
|
||||
fn usable_backend_root(&self) -> Result<PathBuf, String> {
|
||||
let root = self
|
||||
.backend_root
|
||||
.canonicalize()
|
||||
.map_err(|error| format!("ticket backend root is not usable: {error}"))?;
|
||||
if !root.is_dir() {
|
||||
return Err("ticket backend root is not a directory".to_string());
|
||||
}
|
||||
Ok(root)
|
||||
}
|
||||
}
|
||||
|
||||
impl FeatureModule for TicketFeature {
|
||||
fn descriptor(&self) -> FeatureDescriptor {
|
||||
let mut descriptor = FeatureDescriptor::builtin(FEATURE_ID, FEATURE_NAME)
|
||||
.with_description(FEATURE_DESCRIPTION);
|
||||
let enabled_tool_names = self.enabled_tool_names();
|
||||
for name in &enabled_tool_names {
|
||||
descriptor = descriptor.with_tool(ToolDeclaration::new(
|
||||
*name,
|
||||
ticket_tool_description(name, self.record_language.as_deref()),
|
||||
));
|
||||
}
|
||||
descriptor
|
||||
}
|
||||
|
||||
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
|
||||
if let Some(error) = &self.config_error {
|
||||
context
|
||||
.diagnostics()
|
||||
.push(FeatureDiagnostic::warning(format!(
|
||||
"Ticket tools not registered: {error}"
|
||||
)));
|
||||
return Ok(());
|
||||
}
|
||||
let usable_root = match self.usable_backend_root() {
|
||||
Ok(root) => root,
|
||||
Err(reason) => {
|
||||
context
|
||||
.diagnostics()
|
||||
.push(FeatureDiagnostic::warning(format!(
|
||||
"Ticket tools not registered: {reason}; root={} ",
|
||||
self.backend_root.display()
|
||||
)));
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
let backend = LocalTicketBackend::new(usable_root)
|
||||
.with_record_language(self.record_language.as_deref());
|
||||
let allowed_tool_names = self.enabled_tool_names();
|
||||
let mut tools = context.tools();
|
||||
for definition in ticket_tools(backend) {
|
||||
let (meta, _) = definition();
|
||||
let name = meta.name.clone();
|
||||
if !allowed_tool_names
|
||||
.iter()
|
||||
.any(|allowed| *allowed == name.as_str())
|
||||
{
|
||||
continue;
|
||||
}
|
||||
tools.register(ToolContribution::new(name, definition))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn ticket_tools_feature(workspace: impl AsRef<Path>) -> TicketFeature {
|
||||
TicketFeature::for_workspace(workspace)
|
||||
}
|
||||
|
||||
pub fn ticket_tools_feature_with_access(
|
||||
workspace: impl AsRef<Path>,
|
||||
access: TicketFeatureAccess,
|
||||
) -> TicketFeature {
|
||||
TicketFeature::for_workspace_with_access(workspace, access)
|
||||
}
|
||||
|
||||
pub fn ticket_tools_feature_with_options(
|
||||
workspace: impl AsRef<Path>,
|
||||
access: Option<TicketFeatureAccess>,
|
||||
include_orchestration_tools: bool,
|
||||
) -> TicketFeature {
|
||||
TicketFeature::for_workspace_with_options(workspace, access, include_orchestration_tools)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::feature::{FeatureRegistryBuilder, FeatureRuntimeKind};
|
||||
use crate::hook::HookRegistryBuilder;
|
||||
use tempfile::TempDir;
|
||||
use ticket::tool::{
|
||||
TICKET_BASE_TOOL_NAMES, TICKET_ORCHESTRATION_TOOL_NAMES, TICKET_READ_ONLY_TOOL_NAMES,
|
||||
TICKET_TOOL_NAMES,
|
||||
};
|
||||
|
||||
fn make_ticket_root(root: &Path) {
|
||||
std::fs::create_dir_all(root).unwrap();
|
||||
}
|
||||
|
||||
fn write_ticket_config(workspace: &Path, content: &str) {
|
||||
let yoi_dir = workspace.join(".yoi");
|
||||
std::fs::create_dir_all(&yoi_dir).unwrap();
|
||||
std::fs::write(yoi_dir.join("ticket.config.toml"), content).unwrap();
|
||||
}
|
||||
|
||||
fn pending_tool_description(
|
||||
pending_tools: &[llm_engine::tool::ToolDefinition],
|
||||
name: &str,
|
||||
) -> String {
|
||||
pending_tools
|
||||
.iter()
|
||||
.find_map(|definition| {
|
||||
let (meta, _) = definition();
|
||||
(meta.name == name).then_some(meta.description)
|
||||
})
|
||||
.expect("tool exists")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn descriptor_declares_ticket_tools() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let feature = ticket_tools_feature(temp.path());
|
||||
let descriptor = feature.descriptor();
|
||||
assert_eq!(descriptor.id.to_string(), "builtin:ticket");
|
||||
assert_eq!(descriptor.runtime, FeatureRuntimeKind::Builtin);
|
||||
assert_eq!(descriptor.tools.len(), TICKET_TOOL_NAMES.len());
|
||||
assert_eq!(
|
||||
descriptor
|
||||
.tools
|
||||
.iter()
|
||||
.map(|tool| tool.name.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
TICKET_TOOL_NAMES
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_only_descriptor_declares_only_state_tools() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let feature = ticket_tools_feature_with_access(temp.path(), TicketFeatureAccess::ReadOnly);
|
||||
let descriptor = feature.descriptor();
|
||||
assert_eq!(feature.access(), TicketFeatureAccess::ReadOnly);
|
||||
assert_eq!(descriptor.tools.len(), TICKET_READ_ONLY_TOOL_NAMES.len());
|
||||
assert_eq!(
|
||||
descriptor
|
||||
.tools
|
||||
.iter()
|
||||
.map(|tool| tool.name.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
TICKET_READ_ONLY_TOOL_NAMES
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn descriptor_can_expose_base_ticket_without_orchestration_tools() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let feature = ticket_tools_feature_with_options(
|
||||
temp.path(),
|
||||
Some(TicketFeatureAccess::Lifecycle),
|
||||
false,
|
||||
);
|
||||
let descriptor = feature.descriptor();
|
||||
assert_eq!(
|
||||
descriptor
|
||||
.tools
|
||||
.iter()
|
||||
.map(|tool| tool.name.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
TICKET_BASE_TOOL_NAMES
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn descriptor_can_expose_orchestration_only_tools() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let feature = ticket_tools_feature_with_options(temp.path(), None, true);
|
||||
let descriptor = feature.descriptor();
|
||||
assert_eq!(
|
||||
descriptor
|
||||
.tools
|
||||
.iter()
|
||||
.map(|tool| tool.name.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
TICKET_ORCHESTRATION_TOOL_NAMES
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_only_installation_does_not_expose_mutating_tools() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
make_ticket_root(&temp.path().join(DEFAULT_TICKET_BACKEND_RELATIVE_PATH));
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hooks = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(ticket_tools_feature_with_access(
|
||||
temp.path(),
|
||||
TicketFeatureAccess::ReadOnly,
|
||||
))
|
||||
.install_into_pending(&mut pending_tools, &mut hooks);
|
||||
|
||||
assert_eq!(pending_tools.len(), TICKET_READ_ONLY_TOOL_NAMES.len());
|
||||
assert_eq!(
|
||||
report.reports[0].installed_tools,
|
||||
TICKET_READ_ONLY_TOOL_NAMES
|
||||
);
|
||||
let pending_names = pending_tools
|
||||
.iter()
|
||||
.map(|definition| definition().0.name)
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(pending_names, TICKET_READ_ONLY_TOOL_NAMES);
|
||||
for name in ticket::tool::TICKET_MUTATING_TOOL_NAMES {
|
||||
assert!(
|
||||
!report.reports[0]
|
||||
.installed_tools
|
||||
.iter()
|
||||
.any(|tool| tool == name)
|
||||
);
|
||||
assert!(!pending_names.iter().any(|tool| tool == name));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_only_companion_style_context_exposes_ticket_language_guidance() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
write_ticket_config(
|
||||
temp.path(),
|
||||
r#"
|
||||
[ticket]
|
||||
language = "Japanese"
|
||||
"#,
|
||||
);
|
||||
make_ticket_root(&temp.path().join(DEFAULT_TICKET_BACKEND_RELATIVE_PATH));
|
||||
let feature = ticket_tools_feature_with_access(temp.path(), TicketFeatureAccess::ReadOnly);
|
||||
let descriptor = feature.descriptor();
|
||||
let descriptor_description = descriptor
|
||||
.tools
|
||||
.iter()
|
||||
.find(|tool| tool.name == "TicketShow")
|
||||
.expect("TicketShow declared")
|
||||
.description
|
||||
.clone();
|
||||
assert!(descriptor_description.contains("Ticket record language: Japanese"));
|
||||
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hooks = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(feature)
|
||||
.install_into_pending(&mut pending_tools, &mut hooks);
|
||||
|
||||
assert_eq!(pending_tools.len(), TICKET_READ_ONLY_TOOL_NAMES.len());
|
||||
assert_eq!(
|
||||
report.reports[0].installed_tools,
|
||||
TICKET_READ_ONLY_TOOL_NAMES
|
||||
);
|
||||
let description = pending_tool_description(&pending_tools, "TicketShow");
|
||||
assert!(description.contains("Ticket record language: Japanese"));
|
||||
assert!(description.contains("distinct from worker.language"));
|
||||
assert!(description.contains("Preserve protocol literals"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifecycle_installation_exposes_lifecycle_tools() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
make_ticket_root(&temp.path().join(DEFAULT_TICKET_BACKEND_RELATIVE_PATH));
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hooks = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(ticket_tools_feature_with_access(
|
||||
temp.path(),
|
||||
TicketFeatureAccess::Lifecycle,
|
||||
))
|
||||
.install_into_pending(&mut pending_tools, &mut hooks);
|
||||
|
||||
assert_eq!(pending_tools.len(), TICKET_TOOL_NAMES.len());
|
||||
assert_eq!(report.reports[0].installed_tools, TICKET_TOOL_NAMES);
|
||||
for name in ticket::tool::TICKET_MUTATING_TOOL_NAMES {
|
||||
assert!(
|
||||
report.reports[0]
|
||||
.installed_tools
|
||||
.iter()
|
||||
.any(|tool| tool == name)
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
report.reports[0]
|
||||
.installed_tools
|
||||
.iter()
|
||||
.any(|tool| tool == "TicketIntakeReady")
|
||||
);
|
||||
assert!(
|
||||
report.reports[0]
|
||||
.installed_tools
|
||||
.iter()
|
||||
.any(|tool| tool == "TicketWorkflowState")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifecycle_ticket_role_style_context_exposes_ticket_language_guidance() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
write_ticket_config(
|
||||
temp.path(),
|
||||
r#"
|
||||
[ticket]
|
||||
language = "Japanese"
|
||||
"#,
|
||||
);
|
||||
make_ticket_root(&temp.path().join(DEFAULT_TICKET_BACKEND_RELATIVE_PATH));
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hooks = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(ticket_tools_feature_with_access(
|
||||
temp.path(),
|
||||
TicketFeatureAccess::Lifecycle,
|
||||
))
|
||||
.install_into_pending(&mut pending_tools, &mut hooks);
|
||||
|
||||
assert_eq!(pending_tools.len(), TICKET_TOOL_NAMES.len());
|
||||
assert_eq!(report.reports[0].installed_tools, TICKET_TOOL_NAMES);
|
||||
let description = pending_tool_description(&pending_tools, "TicketComment");
|
||||
assert!(description.contains("Ticket record language: Japanese"));
|
||||
assert!(description.contains("durable Ticket record and Ticket tool body text"));
|
||||
assert!(description.contains("distinct from worker.language"));
|
||||
assert!(description.contains("memory.language"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn installs_ticket_tools_when_default_root_is_usable() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
make_ticket_root(&temp.path().join(DEFAULT_TICKET_BACKEND_RELATIVE_PATH));
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hooks = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(ticket_tools_feature(temp.path()))
|
||||
.install_into_pending(&mut pending_tools, &mut hooks);
|
||||
|
||||
assert_eq!(pending_tools.len(), TICKET_TOOL_NAMES.len());
|
||||
assert_eq!(report.reports.len(), 1);
|
||||
assert!(report.reports[0].installed);
|
||||
assert_eq!(report.reports[0].installed_tools, TICKET_TOOL_NAMES);
|
||||
assert!(report.reports[0].skipped.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn installs_ticket_tools_with_configured_backend_root() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
write_ticket_config(
|
||||
temp.path(),
|
||||
r#"
|
||||
[backend]
|
||||
provider = "builtin:yoi_local"
|
||||
root = "tickets"
|
||||
|
||||
[roles.coder]
|
||||
profile = "project:coder"
|
||||
"#,
|
||||
);
|
||||
make_ticket_root(&temp.path().join("tickets"));
|
||||
|
||||
let feature = ticket_tools_feature(temp.path());
|
||||
assert_eq!(feature.backend_root(), temp.path().join("tickets"));
|
||||
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hooks = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(feature)
|
||||
.install_into_pending(&mut pending_tools, &mut hooks);
|
||||
|
||||
assert_eq!(pending_tools.len(), TICKET_TOOL_NAMES.len());
|
||||
assert!(report.reports[0].diagnostics.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_ticket_config_fails_closed() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
make_ticket_root(&temp.path().join(DEFAULT_TICKET_BACKEND_RELATIVE_PATH));
|
||||
write_ticket_config(
|
||||
temp.path(),
|
||||
r#"
|
||||
[roles.operator]
|
||||
profile = "inherit"
|
||||
"#,
|
||||
);
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hooks = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(ticket_tools_feature(temp.path()))
|
||||
.install_into_pending(&mut pending_tools, &mut hooks);
|
||||
|
||||
assert!(pending_tools.is_empty());
|
||||
assert!(report.reports[0].installed_tools.is_empty());
|
||||
assert_eq!(report.reports[0].diagnostics.len(), 1);
|
||||
let message = &report.reports[0].diagnostics[0].message;
|
||||
assert!(message.contains("Ticket tools not registered"));
|
||||
assert!(message.contains("unsupported Ticket role `operator`"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unsupported_ticket_backend_provider_fails_closed() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
make_ticket_root(&temp.path().join(DEFAULT_TICKET_BACKEND_RELATIVE_PATH));
|
||||
write_ticket_config(
|
||||
temp.path(),
|
||||
r#"
|
||||
[backend]
|
||||
provider = "github"
|
||||
"#,
|
||||
);
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hooks = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(ticket_tools_feature(temp.path()))
|
||||
.install_into_pending(&mut pending_tools, &mut hooks);
|
||||
|
||||
assert!(pending_tools.is_empty());
|
||||
assert!(report.reports[0].installed_tools.is_empty());
|
||||
assert_eq!(report.reports[0].diagnostics.len(), 1);
|
||||
let message = &report.reports[0].diagnostics[0].message;
|
||||
assert!(message.contains("Ticket tools not registered"));
|
||||
assert!(message.contains("unsupported Ticket backend provider `github`"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn does_not_register_ticket_tools_when_root_is_missing() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hooks = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(ticket_tools_feature(temp.path()))
|
||||
.install_into_pending(&mut pending_tools, &mut hooks);
|
||||
|
||||
assert!(pending_tools.is_empty());
|
||||
assert_eq!(report.reports.len(), 1);
|
||||
assert!(report.reports[0].installed);
|
||||
assert!(report.reports[0].installed_tools.is_empty());
|
||||
assert_eq!(report.reports[0].diagnostics.len(), 1);
|
||||
assert!(
|
||||
report.reports[0].diagnostics[0]
|
||||
.message
|
||||
.contains("Ticket tools not registered")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registers_ticket_tools_for_flat_backend_root() {
|
||||
let temp = TempDir::new().unwrap();
|
||||
let root = temp.path().join(DEFAULT_TICKET_BACKEND_RELATIVE_PATH);
|
||||
std::fs::create_dir_all(&root).unwrap();
|
||||
let mut pending_tools = Vec::new();
|
||||
let mut hooks = HookRegistryBuilder::default();
|
||||
let report = FeatureRegistryBuilder::new()
|
||||
.with_module(ticket_tools_feature(temp.path()))
|
||||
.install_into_pending(&mut pending_tools, &mut hooks);
|
||||
|
||||
assert_eq!(pending_tools.len(), TICKET_TOOL_NAMES.len());
|
||||
assert_eq!(report.reports[0].installed_tools, TICKET_TOOL_NAMES);
|
||||
assert!(report.reports[0].diagnostics.is_empty());
|
||||
assert!(!root.join("open").exists());
|
||||
assert!(!root.join("pending").exists());
|
||||
assert!(!root.join("closed").exists());
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,686 @@
|
||||
//! Worker 視点のファイルシステム操作。
|
||||
//!
|
||||
//! `ScopedFs` の上に「Worker が読み取りたい / 列挙したい」操作を集約する軽い wrapper。
|
||||
//!
|
||||
//! - `ReadRequirement` と `render_auto_read` — compact worker が `mark_read_required`
|
||||
//! で nominate したファイルを再読し、`[Auto-read file: ...]` system message に
|
||||
//! 変換する経路。`Worker::compact` から呼ばれる。
|
||||
//! - `slice_lines` — 行 offset / limit でテキストを切り出す純粋ヘルパ。
|
||||
//! compact tool 側の `mark_read_required` でも使用。
|
||||
//! - `list_file_completions` — TUI 補完用、prefix マッチでファイル候補を列挙する経路。
|
||||
//! IPC `Method::ListCompletions` 経由で呼ばれる前提（Phase 2 で接続）。
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use llm_engine::Item;
|
||||
use manifest::Scope;
|
||||
use tools::scoped_fs::first_symlink;
|
||||
use tools::{ScopedFs, ToolsError};
|
||||
use tracing::warn;
|
||||
|
||||
/// 補完候補1件の最大数。`list_file_completions` がこの値を超えたら打ち切り。
|
||||
const COMPLETION_LIMIT: usize = 100;
|
||||
/// submit-time directory FileRef の shallow listing で返す最大 entry 数。
|
||||
/// TUI completion と同じ浅い一覧という意味論に揃えるため、同じ上限を使う。
|
||||
const DIR_FILE_REF_ENTRY_LIMIT: usize = COMPLETION_LIMIT;
|
||||
|
||||
/// Compact worker が `mark_read_required` で nominate した「次セッション開始時に
|
||||
/// 自動で再読すべきファイル」のエントリ。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ReadRequirement {
|
||||
pub path: PathBuf,
|
||||
/// 0-based line offset. `None` means from the start of the file.
|
||||
pub offset: Option<usize>,
|
||||
/// Maximum number of lines. `None` means to the end of the file.
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
/// Worker から見えるファイルシステム操作の入口。Clone は cheap（`ScopedFs` 内 `Arc`）。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WorkerFsView {
|
||||
fs: ScopedFs,
|
||||
}
|
||||
|
||||
/// `list_file_completions` が返す候補1件。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct FileCandidate {
|
||||
/// 入力 prefix と整合する形のパス（prefix が absolute なら absolute、
|
||||
/// relative なら cwd 相対）。
|
||||
pub path: String,
|
||||
pub is_dir: bool,
|
||||
}
|
||||
|
||||
/// `resolve_file_ref` の失敗理由。Worker 側で Alert に振り分けるために
|
||||
/// ScopedFs / 内部判定の両方を区別できるよう保持する。
|
||||
#[derive(Debug)]
|
||||
pub enum ResolveError {
|
||||
/// Path resolution / scope check failed via `ScopedFs`.
|
||||
Fs(ToolsError),
|
||||
/// File contents are not valid UTF-8 (binary / non-text).
|
||||
Binary { path: PathBuf },
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ResolveError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
ResolveError::Fs(e) => write!(f, "{e}"),
|
||||
ResolveError::Binary { path } => {
|
||||
write!(f, "file is not valid UTF-8 text: {}", path.display())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ResolveError {}
|
||||
|
||||
impl WorkerFsView {
|
||||
pub fn new(fs: ScopedFs) -> Self {
|
||||
Self { fs }
|
||||
}
|
||||
|
||||
pub fn fs(&self) -> &ScopedFs {
|
||||
&self.fs
|
||||
}
|
||||
|
||||
/// `requirements` の各エントリを `ScopedFs` 経由で再読し、
|
||||
/// `[Auto-read file: <path>:<range>]\n<body>` 形式の system message に変換する。
|
||||
/// 読み取り失敗（NotFound / OutOfScope 等）は warn で記録してスキップする
|
||||
/// — compact 全体を落とさないため。
|
||||
pub fn render_auto_read(&self, requirements: &[ReadRequirement]) -> Vec<Item> {
|
||||
let mut out = Vec::with_capacity(requirements.len());
|
||||
for req in requirements {
|
||||
match self.fs.read_bytes(&req.path) {
|
||||
Ok(bytes) => {
|
||||
let text = String::from_utf8_lossy(&bytes).into_owned();
|
||||
let body = slice_lines(&text, req.offset.unwrap_or(0), req.limit);
|
||||
let range = format_range(req.offset, req.limit);
|
||||
out.push(Item::system_message(format!(
|
||||
"[Auto-read file: {}{range}]\n{body}",
|
||||
req.path.display()
|
||||
)));
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
path = %req.path.display(),
|
||||
error = %e,
|
||||
"auto-read target could not be read; skipping",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// `path` を ScopedFs 経由で解決し、submit 時の `Segment::FileRef`
|
||||
/// attachment 用 system message を返す。
|
||||
///
|
||||
/// - `path` は relative なら cwd 相対、absolute なら absolute として解釈
|
||||
/// - 通常ディレクトリは浅い entry listing として `[Dir: <path>]\n<body>` に展開する
|
||||
/// - ディレクトリ listing は hidden / gitignore を特別扱いせず、scope 上 readable な
|
||||
/// 直下 entry だけを最大 `DIR_FILE_REF_ENTRY_LIMIT` 件返す
|
||||
/// - ファイル本文またはディレクトリ listing 本文が `max_bytes` を超える場合は切り詰める
|
||||
/// - 非 UTF-8 (バイナリ) は `ResolveError::Binary` で拒否
|
||||
/// - スコープ外 / NotFound / symlink directory 等は `ResolveError::Fs` で返す
|
||||
pub fn resolve_file_ref(&self, path: &str, max_bytes: usize) -> Result<Item, ResolveError> {
|
||||
let p = Path::new(path);
|
||||
let abs = if p.is_absolute() {
|
||||
p.to_path_buf()
|
||||
} else {
|
||||
self.fs.cwd().join(p)
|
||||
};
|
||||
|
||||
// 通常ディレクトリだけを FileRef listing として扱う。symlink を含むパスは
|
||||
// `ScopedFs::read_bytes` に委ね、既存の symlink 診断
|
||||
// (`SymlinkTargetIsDirectory` / `SymlinkOutOfScope` 等) を保つ。
|
||||
if first_symlink(&abs).is_none() {
|
||||
let scope = self.fs.scope();
|
||||
if !scope.is_readable(&abs) {
|
||||
return Err(ResolveError::Fs(ToolsError::OutOfScope(abs)));
|
||||
}
|
||||
let meta = metadata_for_file_ref(&abs).map_err(ResolveError::Fs)?;
|
||||
if meta.is_dir() {
|
||||
return render_dir_file_ref(path, &abs, max_bytes, scope.as_ref());
|
||||
}
|
||||
}
|
||||
|
||||
let bytes = self.fs.read_bytes(&abs).map_err(ResolveError::Fs)?;
|
||||
let total = bytes.len();
|
||||
let (body_bytes, truncated) = if total > max_bytes {
|
||||
(&bytes[..max_bytes], true)
|
||||
} else {
|
||||
(bytes.as_slice(), false)
|
||||
};
|
||||
let body = std::str::from_utf8(body_bytes)
|
||||
.map_err(|_| ResolveError::Binary { path: abs.clone() })?;
|
||||
let mut text = format!("[File: {path}]\n{body}");
|
||||
if truncated {
|
||||
text.push_str(&format!(
|
||||
"\n[...truncated, {total} bytes total — use read_file for the rest]"
|
||||
));
|
||||
}
|
||||
Ok(Item::system_message(text))
|
||||
}
|
||||
|
||||
/// `prefix` にマッチするファイル / ディレクトリを scope 内で浅く列挙する。
|
||||
///
|
||||
/// - `prefix` が空 or `cwd` 相対のときは cwd 直下を見る
|
||||
/// - `prefix` が末尾 `/` のときはそのディレクトリ直下を全列挙
|
||||
/// - 末尾が名前部分のときは、その名前を starts_with でフィルタ
|
||||
/// - scope 上 readable なエントリのみ返す
|
||||
/// - ディレクトリ → ファイル の順、各グループ内は名前昇順
|
||||
/// - 上限 `COMPLETION_LIMIT` 件で打ち切り（深い列挙はしない）
|
||||
pub fn list_file_completions(&self, prefix: &str) -> Vec<FileCandidate> {
|
||||
let cwd = self.fs.cwd();
|
||||
let scope = self.fs.scope();
|
||||
let (dir, name_prefix, is_absolute) = split_prefix(prefix, cwd);
|
||||
|
||||
let read_dir = match std::fs::read_dir(&dir) {
|
||||
Ok(rd) => rd,
|
||||
Err(_) => return Vec::new(),
|
||||
};
|
||||
|
||||
let mut out = Vec::new();
|
||||
for entry in read_dir.flatten() {
|
||||
let file_name = entry.file_name();
|
||||
let name = file_name.to_string_lossy();
|
||||
if !name.starts_with(&name_prefix) {
|
||||
continue;
|
||||
}
|
||||
let path = entry.path();
|
||||
if !scope.is_readable(&path) {
|
||||
continue;
|
||||
}
|
||||
let is_dir = entry.file_type().map(|t| t.is_dir()).unwrap_or(false);
|
||||
let display = if is_absolute {
|
||||
path.display().to_string()
|
||||
} else {
|
||||
path.strip_prefix(cwd)
|
||||
.map(|p| p.display().to_string())
|
||||
.unwrap_or_else(|_| path.display().to_string())
|
||||
};
|
||||
out.push(FileCandidate {
|
||||
path: display,
|
||||
is_dir,
|
||||
});
|
||||
}
|
||||
|
||||
out.sort_by(|a, b| match (a.is_dir, b.is_dir) {
|
||||
(true, false) => std::cmp::Ordering::Less,
|
||||
(false, true) => std::cmp::Ordering::Greater,
|
||||
_ => a.path.cmp(&b.path),
|
||||
});
|
||||
out.truncate(COMPLETION_LIMIT);
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
/// `text` の `offset` 行目から `limit` 行（None なら末尾まで）を、元の改行で繋いで返す。
|
||||
pub fn slice_lines(text: &str, offset: usize, limit: Option<usize>) -> String {
|
||||
let lines: Vec<&str> = text.lines().collect();
|
||||
let start = offset.min(lines.len());
|
||||
let end = limit
|
||||
.map(|n| start.saturating_add(n).min(lines.len()))
|
||||
.unwrap_or(lines.len());
|
||||
lines[start..end].join("\n")
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct DirListingEntry {
|
||||
display: String,
|
||||
kind_rank: u8,
|
||||
}
|
||||
|
||||
fn metadata_for_file_ref(path: &Path) -> Result<std::fs::Metadata, ToolsError> {
|
||||
std::fs::metadata(path).map_err(|e| match e.kind() {
|
||||
std::io::ErrorKind::NotFound => ToolsError::NotFound(path.to_path_buf()),
|
||||
_ => ToolsError::io(path, e),
|
||||
})
|
||||
}
|
||||
|
||||
fn render_dir_file_ref(
|
||||
original_path: &str,
|
||||
abs: &Path,
|
||||
max_bytes: usize,
|
||||
scope: &Scope,
|
||||
) -> Result<Item, ResolveError> {
|
||||
let read_dir = std::fs::read_dir(abs).map_err(|e| ResolveError::Fs(ToolsError::io(abs, e)))?;
|
||||
let mut entries = Vec::new();
|
||||
|
||||
for entry in read_dir {
|
||||
let entry = entry.map_err(|e| ResolveError::Fs(ToolsError::io(abs, e)))?;
|
||||
let path = entry.path();
|
||||
if !scope.is_readable(&path) {
|
||||
continue;
|
||||
}
|
||||
let file_type = match entry.file_type() {
|
||||
Ok(ft) => ft,
|
||||
Err(e) => return Err(ResolveError::Fs(ToolsError::io(&path, e))),
|
||||
};
|
||||
let mut display = entry.file_name().to_string_lossy().into_owned();
|
||||
let kind_rank = if file_type.is_dir() {
|
||||
display.push('/');
|
||||
0
|
||||
} else if file_type.is_symlink() {
|
||||
display.push('@');
|
||||
1
|
||||
} else {
|
||||
2
|
||||
};
|
||||
entries.push(DirListingEntry { display, kind_rank });
|
||||
}
|
||||
|
||||
entries.sort_by(|a, b| {
|
||||
a.kind_rank
|
||||
.cmp(&b.kind_rank)
|
||||
.then_with(|| a.display.cmp(&b.display))
|
||||
});
|
||||
|
||||
let total_entries = entries.len();
|
||||
let entry_truncated = total_entries > DIR_FILE_REF_ENTRY_LIMIT;
|
||||
let body = if total_entries == 0 {
|
||||
"(empty directory)".to_string()
|
||||
} else {
|
||||
entries
|
||||
.iter()
|
||||
.take(DIR_FILE_REF_ENTRY_LIMIT)
|
||||
.map(|e| e.display.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
};
|
||||
let body_total_bytes = body.len();
|
||||
let (body, byte_truncated) = truncate_utf8_bytes(&body, max_bytes);
|
||||
|
||||
let mut text = format!("[Dir: {original_path}]\n{body}");
|
||||
if entry_truncated || byte_truncated {
|
||||
text.push('\n');
|
||||
text.push_str(&dir_listing_truncation_hint(
|
||||
entry_truncated,
|
||||
byte_truncated,
|
||||
total_entries,
|
||||
body_total_bytes,
|
||||
));
|
||||
}
|
||||
Ok(Item::system_message(text))
|
||||
}
|
||||
|
||||
fn truncate_utf8_bytes(s: &str, max_bytes: usize) -> (&str, bool) {
|
||||
if s.len() <= max_bytes {
|
||||
return (s, false);
|
||||
}
|
||||
let mut end = max_bytes;
|
||||
while !s.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
(&s[..end], true)
|
||||
}
|
||||
|
||||
fn dir_listing_truncation_hint(
|
||||
entry_truncated: bool,
|
||||
byte_truncated: bool,
|
||||
total_entries: usize,
|
||||
body_total_bytes: usize,
|
||||
) -> String {
|
||||
match (entry_truncated, byte_truncated) {
|
||||
(true, true) => format!(
|
||||
"[...truncated, {total_entries} readable entries total; first {DIR_FILE_REF_ENTRY_LIMIT} entries were {body_total_bytes} bytes before byte cap — use Glob for more]"
|
||||
),
|
||||
(true, false) => {
|
||||
format!("[...truncated, {total_entries} readable entries total — use Glob for more]")
|
||||
}
|
||||
(false, true) => {
|
||||
format!("[...truncated, {body_total_bytes} bytes total — use Glob or Read for more]")
|
||||
}
|
||||
(false, false) => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn format_range(offset: Option<usize>, limit: Option<usize>) -> String {
|
||||
match (offset, limit) {
|
||||
(None, None) => String::new(),
|
||||
(Some(off), None) => format!(":{}-", off + 1),
|
||||
(None, Some(lim)) => format!(":1-{lim}"),
|
||||
(Some(off), Some(lim)) => format!(":{}-{}", off + 1, off.saturating_add(lim)),
|
||||
}
|
||||
}
|
||||
|
||||
fn split_prefix(prefix: &str, cwd: &Path) -> (PathBuf, String, bool) {
|
||||
let is_absolute = Path::new(prefix).is_absolute();
|
||||
let p = Path::new(prefix);
|
||||
let (parent, name) = if prefix.is_empty() || prefix.ends_with('/') {
|
||||
(p.to_path_buf(), String::new())
|
||||
} else {
|
||||
let parent = p.parent().map(|p| p.to_path_buf()).unwrap_or_default();
|
||||
let name = p
|
||||
.file_name()
|
||||
.map(|n| n.to_string_lossy().into_owned())
|
||||
.unwrap_or_default();
|
||||
(parent, name)
|
||||
};
|
||||
let dir = if is_absolute {
|
||||
parent
|
||||
} else if parent.as_os_str().is_empty() {
|
||||
cwd.to_path_buf()
|
||||
} else {
|
||||
cwd.join(parent)
|
||||
};
|
||||
(dir, name, is_absolute)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use llm_engine::ContentPart;
|
||||
use manifest::{Permission, Scope, ScopeConfig, ScopeRule};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn fs_for(dir: &TempDir) -> ScopedFs {
|
||||
ScopedFs::new(
|
||||
Scope::writable(dir.path()).unwrap(),
|
||||
dir.path().to_path_buf(),
|
||||
)
|
||||
}
|
||||
|
||||
fn touch(path: &Path, content: &str) {
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent).unwrap();
|
||||
}
|
||||
std::fs::write(path, content).unwrap();
|
||||
}
|
||||
|
||||
fn system_text(item: &Item) -> &str {
|
||||
let Item::Message { content, .. } = item else {
|
||||
panic!("expected message item");
|
||||
};
|
||||
let Some(ContentPart::Text { text }) = content.first() else {
|
||||
panic!("expected text content");
|
||||
};
|
||||
text
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slice_lines_handles_offset_and_limit() {
|
||||
let text = "a\nb\nc\nd";
|
||||
assert_eq!(slice_lines(text, 0, None), "a\nb\nc\nd");
|
||||
assert_eq!(slice_lines(text, 1, Some(2)), "b\nc");
|
||||
assert_eq!(slice_lines(text, 10, None), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_auto_read_emits_system_messages_with_range_label() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let file = dir.path().join("hello.txt");
|
||||
std::fs::write(&file, "alpha\nbeta\ngamma\n").unwrap();
|
||||
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
let items = view.render_auto_read(&[ReadRequirement {
|
||||
path: file.clone(),
|
||||
offset: Some(1),
|
||||
limit: Some(1),
|
||||
}]);
|
||||
|
||||
assert_eq!(items.len(), 1);
|
||||
let rendered = format!("{:?}", items[0]);
|
||||
assert!(rendered.contains("Auto-read file"));
|
||||
assert!(rendered.contains(":2-2"));
|
||||
assert!(rendered.contains("beta"));
|
||||
assert!(!rendered.contains("alpha"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_file_ref_emits_system_message_with_path_header() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
std::fs::write(dir.path().join("hello.txt"), "hello world").unwrap();
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let item = view.resolve_file_ref("hello.txt", 1024).unwrap();
|
||||
let text = format!("{item:?}");
|
||||
assert!(text.contains("[File: hello.txt]"));
|
||||
assert!(text.contains("hello world"));
|
||||
assert!(!text.contains("truncated"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_file_ref_truncates_with_hint_when_over_cap() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let body = "x".repeat(2048);
|
||||
std::fs::write(dir.path().join("big.txt"), &body).unwrap();
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let item = view.resolve_file_ref("big.txt", 256).unwrap();
|
||||
let text = format!("{item:?}");
|
||||
assert!(text.contains("[File: big.txt]"));
|
||||
assert!(text.contains("truncated"));
|
||||
assert!(text.contains("2048 bytes total"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_file_ref_lists_directory_shallow_entries() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
std::fs::create_dir_all(dir.path().join("docs/sub")).unwrap();
|
||||
touch(&dir.path().join("docs/.hidden"), "hidden");
|
||||
touch(&dir.path().join("docs/.gitignore"), "ignored.txt\n");
|
||||
touch(
|
||||
&dir.path().join("docs/ignored.txt"),
|
||||
"not ignored for FileRef",
|
||||
);
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let item = view.resolve_file_ref("docs", 4096).unwrap();
|
||||
let text = system_text(&item);
|
||||
assert!(text.starts_with("[Dir: docs]\n"));
|
||||
assert!(text.contains("sub/"));
|
||||
assert!(text.contains(".hidden"));
|
||||
assert!(text.contains(".gitignore"));
|
||||
assert!(text.contains("ignored.txt"));
|
||||
|
||||
let sub_pos = text.find("sub/").unwrap();
|
||||
let hidden_pos = text.find(".hidden").unwrap();
|
||||
assert!(
|
||||
sub_pos < hidden_pos,
|
||||
"directories should sort before files:\n{text}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_file_ref_directory_listing_filters_unreadable_entries() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let docs = dir.path().join("docs");
|
||||
let secret = docs.join("secret");
|
||||
std::fs::create_dir_all(&secret).unwrap();
|
||||
touch(&docs.join("visible.txt"), "ok");
|
||||
touch(&secret.join("hidden.txt"), "nope");
|
||||
|
||||
let cfg = ScopeConfig {
|
||||
allow: vec![ScopeRule {
|
||||
target: dir.path().to_path_buf(),
|
||||
permission: Permission::Write,
|
||||
recursive: true,
|
||||
}],
|
||||
deny: vec![ScopeRule {
|
||||
target: secret.clone(),
|
||||
permission: Permission::Read,
|
||||
recursive: true,
|
||||
}],
|
||||
};
|
||||
let scope = Scope::from_config(&cfg).unwrap();
|
||||
let fs = ScopedFs::new(scope, dir.path().to_path_buf());
|
||||
let view = WorkerFsView::new(fs);
|
||||
|
||||
let item = view.resolve_file_ref("docs", 4096).unwrap();
|
||||
let text = system_text(&item);
|
||||
assert!(text.contains("visible.txt"));
|
||||
assert!(!text.contains("secret"));
|
||||
assert!(!text.contains("hidden.txt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_file_ref_directory_listing_uses_upload_byte_cap() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
std::fs::create_dir(dir.path().join("docs")).unwrap();
|
||||
touch(&dir.path().join("docs/very-long-file-name.txt"), "");
|
||||
touch(&dir.path().join("docs/another-long-file-name.txt"), "");
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let item = view.resolve_file_ref("docs", 10).unwrap();
|
||||
let text = system_text(&item);
|
||||
assert!(text.starts_with("[Dir: docs]\n"));
|
||||
assert!(text.contains("truncated"));
|
||||
assert!(text.contains("bytes total"));
|
||||
assert!(text.contains("use Glob or Read for more"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_file_ref_directory_listing_uses_completion_entry_limit() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
std::fs::create_dir(dir.path().join("docs")).unwrap();
|
||||
for i in 0..(DIR_FILE_REF_ENTRY_LIMIT + 5) {
|
||||
touch(&dir.path().join(format!("docs/file-{i:03}.txt")), "");
|
||||
}
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let item = view.resolve_file_ref("docs", 4096).unwrap();
|
||||
let text = system_text(&item);
|
||||
assert!(text.contains("105 readable entries total"));
|
||||
assert!(text.contains("file-099.txt"));
|
||||
assert!(!text.contains("file-100.txt"));
|
||||
assert!(text.contains("use Glob for more"));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn resolve_file_ref_directory_listing_marks_readable_symlink_entries() {
|
||||
use std::os::unix::fs::symlink;
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
std::fs::create_dir(dir.path().join("docs")).unwrap();
|
||||
touch(&dir.path().join("docs/target.txt"), "target");
|
||||
symlink("target.txt", dir.path().join("docs/link.txt")).unwrap();
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let item = view.resolve_file_ref("docs", 4096).unwrap();
|
||||
let text = system_text(&item);
|
||||
assert!(text.contains("link.txt@"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_file_ref_rejects_binary_with_binary_error() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
std::fs::write(dir.path().join("blob.bin"), [0xff, 0xfe, 0x00, 0x80]).unwrap();
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let err = view.resolve_file_ref("blob.bin", 1024).unwrap_err();
|
||||
assert!(matches!(err, ResolveError::Binary { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_file_ref_returns_fs_error_for_out_of_scope() {
|
||||
let outer = TempDir::new().unwrap();
|
||||
let inner = outer.path().join("scoped");
|
||||
std::fs::create_dir(&inner).unwrap();
|
||||
std::fs::write(outer.path().join("secret.txt"), "nope").unwrap();
|
||||
let scope = Scope::writable(&inner).unwrap();
|
||||
let fs = ScopedFs::new(scope, inner.clone());
|
||||
let view = WorkerFsView::new(fs);
|
||||
|
||||
// Absolute path outside of scope.
|
||||
let outside = outer.path().join("secret.txt");
|
||||
let err = view
|
||||
.resolve_file_ref(outside.to_str().unwrap(), 1024)
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, ResolveError::Fs(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_auto_read_skips_unreadable_targets() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
let items = view.render_auto_read(&[ReadRequirement {
|
||||
path: dir.path().join("missing.txt"),
|
||||
offset: None,
|
||||
limit: None,
|
||||
}]);
|
||||
assert!(items.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_file_completions_lists_pwd_when_prefix_empty() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
touch(&dir.path().join("alpha.rs"), "");
|
||||
touch(&dir.path().join("beta.rs"), "");
|
||||
std::fs::create_dir(dir.path().join("subdir")).unwrap();
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let cands = view.list_file_completions("");
|
||||
// ディレクトリ first
|
||||
let names: Vec<&str> = cands.iter().map(|c| c.path.as_str()).collect();
|
||||
assert_eq!(names, vec!["subdir", "alpha.rs", "beta.rs"]);
|
||||
assert!(cands[0].is_dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_file_completions_filters_by_name_prefix() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
touch(&dir.path().join("alpha.rs"), "");
|
||||
touch(&dir.path().join("beta.rs"), "");
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let cands = view.list_file_completions("al");
|
||||
assert_eq!(cands.len(), 1);
|
||||
assert_eq!(cands[0].path, "alpha.rs");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_file_completions_descends_into_subdir_with_trailing_slash() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
touch(&dir.path().join("sub/x.rs"), "");
|
||||
touch(&dir.path().join("sub/y.rs"), "");
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let cands = view.list_file_completions("sub/");
|
||||
let names: Vec<&str> = cands.iter().map(|c| c.path.as_str()).collect();
|
||||
assert_eq!(names, vec!["sub/x.rs", "sub/y.rs"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_file_completions_filters_out_non_readable_under_scope() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let secret = dir.path().join("secret");
|
||||
std::fs::create_dir(&secret).unwrap();
|
||||
touch(&dir.path().join("visible.rs"), "");
|
||||
touch(&secret.join("hidden.rs"), "");
|
||||
|
||||
let cfg = ScopeConfig {
|
||||
allow: vec![ScopeRule {
|
||||
target: dir.path().to_path_buf(),
|
||||
permission: Permission::Write,
|
||||
recursive: true,
|
||||
}],
|
||||
deny: vec![ScopeRule {
|
||||
target: secret.clone(),
|
||||
permission: Permission::Read,
|
||||
recursive: true,
|
||||
}],
|
||||
};
|
||||
let scope = Scope::from_config(&cfg).unwrap();
|
||||
let fs = ScopedFs::new(scope, dir.path().to_path_buf());
|
||||
let view = WorkerFsView::new(fs);
|
||||
|
||||
let cands = view.list_file_completions("");
|
||||
let names: Vec<&str> = cands.iter().map(|c| c.path.as_str()).collect();
|
||||
assert!(names.contains(&"visible.rs"));
|
||||
assert!(!names.contains(&"secret"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_file_completions_supports_absolute_prefix() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
touch(&dir.path().join("a.rs"), "");
|
||||
let view = WorkerFsView::new(fs_for(&dir));
|
||||
|
||||
let prefix = format!("{}/", dir.path().display());
|
||||
let cands = view.list_file_completions(&prefix);
|
||||
assert_eq!(cands.len(), 1);
|
||||
assert!(cands[0].path.starts_with('/'));
|
||||
assert!(cands[0].path.ends_with("a.rs"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,500 @@
|
||||
//! Worker-layer hook infrastructure
|
||||
//!
|
||||
//! Hooks are the **public** orchestration extension point. They receive
|
||||
//! event-specific context values about each event in the Engine execution loop
|
||||
//! and return a safe public control-flow action. Contexts may carry narrow
|
||||
//! host-created handles for approved side effects; hook return values remain
|
||||
//! flow-control decisions only.
|
||||
//!
|
||||
//! Hooks intentionally cannot mutate the Engine's context, history, tool
|
||||
//! call, or tool result. Internal mechanisms that need such access (e.g.
|
||||
//! compaction, notification injection, output truncation) implement
|
||||
//! `llm_engine::Interceptor` directly inside Worker, never via this trait.
|
||||
//!
|
||||
//! This separation lets Hooks be exposed safely to user-facing
|
||||
//! extension surfaces (scripting, plugins) in the future without
|
||||
//! exposing the underlying mutable state.
|
||||
|
||||
use std::ops::Deref;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_engine::interceptor::{
|
||||
PostToolAction, PreRequestAction, PreToolAction, PromptAction, TurnEndAction,
|
||||
};
|
||||
use llm_engine::tool::{ToolOutput, ToolResult};
|
||||
use serde_json::Value;
|
||||
use session_store::{SystemItem, SystemReminder};
|
||||
|
||||
/// Hook-facing prompt-submit action.
|
||||
///
|
||||
/// A strict subset of [`PromptAction`]: Hooks may continue or cancel
|
||||
/// the submit, but cannot inject items into history. The
|
||||
/// `ContinueWith(Vec<Item>)` variant is reserved for the internal
|
||||
/// `Interceptor` so that Hook (the public extension surface) stays
|
||||
/// read-only by construction (see module-level doc).
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum HookPromptAction {
|
||||
/// Proceed normally.
|
||||
Continue,
|
||||
/// Cancel this submitted prompt with a reason.
|
||||
Cancel(String),
|
||||
}
|
||||
|
||||
impl From<HookPromptAction> for PromptAction {
|
||||
fn from(action: HookPromptAction) -> Self {
|
||||
match action {
|
||||
HookPromptAction::Continue => PromptAction::Continue,
|
||||
HookPromptAction::Cancel(reason) => PromptAction::Cancel(reason),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Hook-facing pre-LLM-request action.
|
||||
///
|
||||
/// Public hooks may observe the request boundary, cancel the run, or yield
|
||||
/// control back to the caller. They cannot return
|
||||
/// `PreRequestAction::ContinueWith(Vec<Item>)`; model-visible request/history
|
||||
/// additions must use durable host-owned paths such as notifications or
|
||||
/// system-item commits.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum HookPreRequestAction {
|
||||
/// Proceed normally.
|
||||
Continue,
|
||||
/// Cancel the run with a reason.
|
||||
Cancel(String),
|
||||
/// Yield control to the caller for host-owned processing/resume.
|
||||
Yield,
|
||||
}
|
||||
|
||||
impl From<HookPreRequestAction> for PreRequestAction {
|
||||
fn from(action: HookPreRequestAction) -> Self {
|
||||
match action {
|
||||
HookPreRequestAction::Continue => PreRequestAction::Continue,
|
||||
HookPreRequestAction::Cancel(reason) => PreRequestAction::Cancel(reason),
|
||||
HookPreRequestAction::Yield => PreRequestAction::Yield,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Hook-facing pre-tool-call action.
|
||||
///
|
||||
/// Hooks may continue, pause/abort the call, or deny it with an error
|
||||
/// string that Worker converts into a synthetic tool result for the current
|
||||
/// tool call. Hooks cannot express the internal no-result skip path, mutate
|
||||
/// the tool call arguments, or construct arbitrary `ToolResult` values.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum HookPreToolAction {
|
||||
/// Proceed with tool execution.
|
||||
Continue,
|
||||
/// Deny this tool call and commit a synthetic error result.
|
||||
Deny(String),
|
||||
/// Abort the entire run.
|
||||
Abort(String),
|
||||
/// Pause execution.
|
||||
Pause,
|
||||
}
|
||||
|
||||
impl HookPreToolAction {
|
||||
pub(crate) fn into_worker_action(self, call_id: String) -> PreToolAction {
|
||||
match self {
|
||||
HookPreToolAction::Continue => PreToolAction::Continue,
|
||||
HookPreToolAction::Deny(reason) => {
|
||||
PreToolAction::SyntheticResult(ToolResult::error(call_id, reason))
|
||||
}
|
||||
HookPreToolAction::Abort(reason) => PreToolAction::Abort(reason),
|
||||
HookPreToolAction::Pause => PreToolAction::Pause,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Hook-facing post-tool-call action.
|
||||
///
|
||||
/// Post-tool hooks are observational except that they may abort the run. They
|
||||
/// cannot rewrite the tool output; adding an explicit bounded transform would
|
||||
/// require a separate safe public type.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum HookPostToolAction {
|
||||
/// Proceed normally.
|
||||
Continue,
|
||||
/// Abort the entire run.
|
||||
Abort(String),
|
||||
}
|
||||
|
||||
impl From<HookPostToolAction> for PostToolAction {
|
||||
fn from(action: HookPostToolAction) -> Self {
|
||||
match action {
|
||||
HookPostToolAction::Continue => PostToolAction::Continue,
|
||||
HookPostToolAction::Abort(reason) => PostToolAction::Abort(reason),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Hook-facing turn-end action.
|
||||
///
|
||||
/// Turn-end hooks may observe a completed turn and optionally pause further
|
||||
/// execution. They cannot return
|
||||
/// `TurnEndAction::ContinueWithMessages(Vec<Item>)`; public hooks must not
|
||||
/// append arbitrary model-visible messages at turn boundaries.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum HookTurnEndAction {
|
||||
/// Finish the turn normally.
|
||||
Finish,
|
||||
/// Pause execution.
|
||||
Pause,
|
||||
}
|
||||
|
||||
impl From<HookTurnEndAction> for TurnEndAction {
|
||||
fn from(action: HookTurnEndAction) -> Self {
|
||||
match action {
|
||||
HookTurnEndAction::Finish => TurnEndAction::Finish,
|
||||
HookTurnEndAction::Pause => TurnEndAction::Pause,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Hook context handles
|
||||
// =============================================================================
|
||||
|
||||
/// Host-created handle for appending approved durable [`SystemItem`] requests.
|
||||
///
|
||||
/// Hook code can use this handle only when the Worker host includes it in an
|
||||
/// event-specific context. The handle queues typed requests; the host drains the
|
||||
/// queue, commits each entry through `LogEntry::SystemItem`, and only then makes
|
||||
/// the matching system message visible to the model. It deliberately exposes no
|
||||
/// raw `llm_engine::Item`, history writer, event sender, `Worker`, `Engine`, or
|
||||
/// notification buffer.
|
||||
pub struct SystemItemAppendHandle {
|
||||
pending: Arc<Mutex<Vec<SystemItem>>>,
|
||||
}
|
||||
|
||||
impl SystemItemAppendHandle {
|
||||
pub(crate) fn new(pending: Arc<Mutex<Vec<SystemItem>>>) -> Self {
|
||||
Self { pending }
|
||||
}
|
||||
|
||||
/// Queue a task-inactivity reminder for durable model-visible append.
|
||||
///
|
||||
/// The body should be the unwrapped reminder text; the host-side
|
||||
/// `SystemReminder` renderer wraps it exactly once in `<system-reminder>`
|
||||
/// tags before commit.
|
||||
pub fn append_task_reminder(&self, body: impl Into<String>) {
|
||||
let item = SystemReminder::task_inactivity(body).into_system_item();
|
||||
self.pending
|
||||
.lock()
|
||||
.expect("system-item append queue poisoned")
|
||||
.push(item);
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Hook input summary/context types (read-only)
|
||||
// =============================================================================
|
||||
|
||||
/// Information passed to `OnPromptSubmit` hooks.
|
||||
pub struct PromptSubmitInfo {
|
||||
/// Concatenated text content of the user's input message.
|
||||
pub input_text: String,
|
||||
/// 0-based turn index this prompt opens.
|
||||
pub turn_index: usize,
|
||||
}
|
||||
|
||||
/// Summary information included in `PreLlmRequest` contexts.
|
||||
pub struct PreRequestInfo {
|
||||
/// Number of items currently in the Engine context.
|
||||
pub item_count: usize,
|
||||
/// Most recently observed `input_tokens` from the LLM provider.
|
||||
/// `None` when the Worker has no compaction state attached, or when
|
||||
/// no LLM call has completed yet.
|
||||
pub estimated_tokens: Option<u64>,
|
||||
/// Current turn index (0-based).
|
||||
pub turn_index: usize,
|
||||
/// Tool calls already executed in this turn.
|
||||
pub tool_calls_this_turn: usize,
|
||||
}
|
||||
|
||||
/// Context passed to `PreLlmRequest` hooks.
|
||||
///
|
||||
/// The summary remains read-only. When the host grants durable system-item
|
||||
/// append authority for this request, `system_items()` exposes a typed append
|
||||
/// handle; otherwise it returns `None` and hooks cannot produce model-visible
|
||||
/// additions.
|
||||
pub struct PreRequestContext {
|
||||
info: PreRequestInfo,
|
||||
system_items: Option<SystemItemAppendHandle>,
|
||||
}
|
||||
|
||||
impl PreRequestContext {
|
||||
pub(crate) fn new(info: PreRequestInfo, system_items: Option<SystemItemAppendHandle>) -> Self {
|
||||
Self { info, system_items }
|
||||
}
|
||||
|
||||
/// Read-only request summary.
|
||||
pub fn info(&self) -> &PreRequestInfo {
|
||||
&self.info
|
||||
}
|
||||
|
||||
/// Host-provided durable system-item append handle, when available.
|
||||
pub fn system_items(&self) -> Option<&SystemItemAppendHandle> {
|
||||
self.system_items.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl Deref for PreRequestContext {
|
||||
type Target = PreRequestInfo;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.info
|
||||
}
|
||||
}
|
||||
|
||||
/// Information passed to `PreToolCall` hooks.
|
||||
pub struct ToolCallSummary {
|
||||
/// Provider-assigned tool call id.
|
||||
pub call_id: String,
|
||||
/// Registered tool name.
|
||||
pub tool_name: String,
|
||||
/// Tool arguments as a JSON value (cloned).
|
||||
///
|
||||
/// LLM-generated arguments are bounded by max_tokens, so cloning
|
||||
/// is cheap relative to tool execution. Structural access is
|
||||
/// required for permission decisions (e.g. inspecting a `path`
|
||||
/// field), which a stringified preview would not support.
|
||||
pub arguments: Value,
|
||||
}
|
||||
|
||||
/// Information passed to `PostToolCall` hooks.
|
||||
pub struct ToolResultSummary {
|
||||
/// Provider-assigned tool call id this result corresponds to.
|
||||
pub call_id: String,
|
||||
/// Registered tool name.
|
||||
pub tool_name: String,
|
||||
/// Whether the tool reported an error.
|
||||
pub is_error: bool,
|
||||
/// Tool output (`summary` always present, `content` may be `None`).
|
||||
pub output: ToolOutput,
|
||||
}
|
||||
|
||||
/// Information passed to `OnTurnEnd` hooks.
|
||||
pub struct TurnEndInfo {
|
||||
/// Turn that just ended (0-based).
|
||||
pub turn_index: usize,
|
||||
/// Tool calls executed in this turn.
|
||||
pub tool_calls_count: usize,
|
||||
/// Preview of the assistant's final text in this turn.
|
||||
/// Truncated at a UTF-8 boundary; empty when no assistant text exists.
|
||||
pub final_text_preview: String,
|
||||
}
|
||||
|
||||
/// Information passed to `OnAbort` hooks.
|
||||
pub struct AbortInfo {
|
||||
/// Reason supplied by the aborter.
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Hook Event Kinds
|
||||
// =============================================================================
|
||||
|
||||
/// Marker trait for hook event kinds.
|
||||
///
|
||||
/// Each event kind specifies its read-only input and the safe public
|
||||
/// control-flow action returned by hooks.
|
||||
pub trait HookEventKind: Send + Sync + 'static {
|
||||
/// Read-only input passed to the hook.
|
||||
type Input: Send + Sync;
|
||||
/// Control-flow action returned by the hook.
|
||||
type Output;
|
||||
}
|
||||
|
||||
/// After receiving user input, before adding to history; may continue or cancel.
|
||||
pub struct OnPromptSubmit;
|
||||
/// Before each LLM request; may continue, cancel, or yield.
|
||||
pub struct PreLlmRequest;
|
||||
/// Before each tool is executed; may continue, deny with a synthetic result,
|
||||
/// abort, or pause.
|
||||
pub struct PreToolCall;
|
||||
/// After each tool completes; observational except it may abort the run.
|
||||
pub struct PostToolCall;
|
||||
/// When a turn ends with no tool calls; observational except it may pause.
|
||||
pub struct OnTurnEnd;
|
||||
/// When execution is interrupted; observational only.
|
||||
pub struct OnAbort;
|
||||
|
||||
impl HookEventKind for OnPromptSubmit {
|
||||
type Input = PromptSubmitInfo;
|
||||
type Output = HookPromptAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for PreLlmRequest {
|
||||
type Input = PreRequestContext;
|
||||
type Output = HookPreRequestAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for PreToolCall {
|
||||
type Input = ToolCallSummary;
|
||||
type Output = HookPreToolAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for PostToolCall {
|
||||
type Input = ToolResultSummary;
|
||||
type Output = HookPostToolAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for OnTurnEnd {
|
||||
type Input = TurnEndInfo;
|
||||
type Output = HookTurnEndAction;
|
||||
}
|
||||
|
||||
impl HookEventKind for OnAbort {
|
||||
type Input = AbortInfo;
|
||||
type Output = ();
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Hook Trait
|
||||
// =============================================================================
|
||||
|
||||
/// Async hook for a specific event kind.
|
||||
///
|
||||
/// Hooks receive a shared reference to the event's read-only input
|
||||
/// and return a safe public control-flow action. Multiple hooks can be
|
||||
/// registered per event; they are evaluated in registration order and
|
||||
/// short-circuit on the first non-continue action.
|
||||
#[async_trait]
|
||||
pub trait Hook<E: HookEventKind>: Send + Sync {
|
||||
async fn call(&self, input: &E::Input) -> E::Output;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Hook Registry
|
||||
// =============================================================================
|
||||
|
||||
/// Builder for constructing a frozen `HookRegistry`.
|
||||
///
|
||||
/// Hooks are added during setup, then `build()` produces an immutable
|
||||
/// registry that can be shared via `Arc`.
|
||||
#[derive(Default)]
|
||||
pub struct HookRegistryBuilder {
|
||||
on_prompt_submit: Vec<Box<dyn Hook<OnPromptSubmit>>>,
|
||||
pre_llm_request: Vec<Box<dyn Hook<PreLlmRequest>>>,
|
||||
pre_tool_call: Vec<Box<dyn Hook<PreToolCall>>>,
|
||||
post_tool_call: Vec<Box<dyn Hook<PostToolCall>>>,
|
||||
on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
|
||||
on_abort: Vec<Box<dyn Hook<OnAbort>>>,
|
||||
}
|
||||
|
||||
impl HookRegistryBuilder {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn add_on_prompt_submit(&mut self, hook: impl Hook<OnPromptSubmit> + 'static) {
|
||||
self.on_prompt_submit.push(Box::new(hook));
|
||||
}
|
||||
|
||||
pub fn add_pre_llm_request(&mut self, hook: impl Hook<PreLlmRequest> + 'static) {
|
||||
self.pre_llm_request.push(Box::new(hook));
|
||||
}
|
||||
|
||||
pub fn add_pre_tool_call(&mut self, hook: impl Hook<PreToolCall> + 'static) {
|
||||
self.pre_tool_call.push(Box::new(hook));
|
||||
}
|
||||
|
||||
pub fn add_post_tool_call(&mut self, hook: impl Hook<PostToolCall> + 'static) {
|
||||
self.post_tool_call.push(Box::new(hook));
|
||||
}
|
||||
|
||||
pub fn add_on_turn_end(&mut self, hook: impl Hook<OnTurnEnd> + 'static) {
|
||||
self.on_turn_end.push(Box::new(hook));
|
||||
}
|
||||
|
||||
pub fn add_on_abort(&mut self, hook: impl Hook<OnAbort> + 'static) {
|
||||
self.on_abort.push(Box::new(hook));
|
||||
}
|
||||
|
||||
/// Freeze the builder into an immutable registry.
|
||||
pub fn build(self) -> HookRegistry {
|
||||
HookRegistry {
|
||||
on_prompt_submit: self.on_prompt_submit,
|
||||
pre_llm_request: self.pre_llm_request,
|
||||
pre_tool_call: self.pre_tool_call,
|
||||
post_tool_call: self.post_tool_call,
|
||||
on_turn_end: self.on_turn_end,
|
||||
on_abort: self.on_abort,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Frozen registry of hooks. Constructed via [`HookRegistryBuilder::build()`].
|
||||
pub struct HookRegistry {
|
||||
pub(crate) on_prompt_submit: Vec<Box<dyn Hook<OnPromptSubmit>>>,
|
||||
pub(crate) pre_llm_request: Vec<Box<dyn Hook<PreLlmRequest>>>,
|
||||
pub(crate) pre_tool_call: Vec<Box<dyn Hook<PreToolCall>>>,
|
||||
pub(crate) post_tool_call: Vec<Box<dyn Hook<PostToolCall>>>,
|
||||
pub(crate) on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
|
||||
pub(crate) on_abort: Vec<Box<dyn Hook<OnAbort>>>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn system_item_append_handle_queues_only_approved_task_reminder_items() {
|
||||
let pending = Arc::new(Mutex::new(Vec::new()));
|
||||
let handle = SystemItemAppendHandle::new(Arc::clone(&pending));
|
||||
|
||||
handle.append_task_reminder("remember tasks");
|
||||
|
||||
let queued = pending.lock().expect("pending queue poisoned");
|
||||
assert_eq!(queued.len(), 1);
|
||||
match &queued[0] {
|
||||
SystemItem::TaskReminder { body, .. } => {
|
||||
assert_eq!(body.matches("<system-reminder>").count(), 1);
|
||||
assert!(body.contains("remember tasks"));
|
||||
}
|
||||
other => panic!("unexpected system item: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pre_request_context_exposes_handle_only_when_host_supplies_one() {
|
||||
let info = PreRequestInfo {
|
||||
item_count: 3,
|
||||
estimated_tokens: Some(42),
|
||||
turn_index: 1,
|
||||
tool_calls_this_turn: 2,
|
||||
};
|
||||
let context = PreRequestContext::new(info, None);
|
||||
|
||||
assert_eq!(context.item_count, 3);
|
||||
assert_eq!(context.info().estimated_tokens, Some(42));
|
||||
assert!(context.system_items().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn public_pre_tool_hook_actions_cannot_emit_internal_no_result_skip() {
|
||||
let continue_action = HookPreToolAction::Continue.into_worker_action("call_1".into());
|
||||
assert!(matches!(continue_action, PreToolAction::Continue));
|
||||
|
||||
let deny_action =
|
||||
HookPreToolAction::Deny("blocked".into()).into_worker_action("call_2".into());
|
||||
match deny_action {
|
||||
PreToolAction::SyntheticResult(result) => {
|
||||
assert_eq!(result.tool_use_id, "call_2");
|
||||
assert_eq!(result.summary, "blocked");
|
||||
assert!(result.is_error);
|
||||
}
|
||||
other => panic!("public deny must produce synthetic result, got {other:?}"),
|
||||
}
|
||||
|
||||
let abort_action =
|
||||
HookPreToolAction::Abort("stop".into()).into_worker_action("call_3".into());
|
||||
assert!(matches!(abort_action, PreToolAction::Abort(reason) if reason == "stop"));
|
||||
|
||||
let pause_action = HookPreToolAction::Pause.into_worker_action("call_4".into());
|
||||
assert!(matches!(pause_action, PreToolAction::Pause));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,599 @@
|
||||
use std::sync::{Arc, Mutex, MutexGuard};
|
||||
|
||||
use protocol::{Event, InFlightBlock, InFlightSnapshot, InFlightToolCallState};
|
||||
use session_store::{LoggedContentPart, LoggedItem};
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct InFlightBlockId(u64);
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct InFlightEvents {
|
||||
inner: Arc<Mutex<InFlightInner>>,
|
||||
event_tx: broadcast::Sender<Event>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct InFlightInner {
|
||||
next_block_id: u64,
|
||||
blocks: Vec<TrackedBlock>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
enum TrackedBlock {
|
||||
Text {
|
||||
block_id: InFlightBlockId,
|
||||
text: String,
|
||||
finished: bool,
|
||||
},
|
||||
Thinking {
|
||||
block_id: InFlightBlockId,
|
||||
text: String,
|
||||
finished: bool,
|
||||
},
|
||||
ToolCall {
|
||||
block_id: InFlightBlockId,
|
||||
id: String,
|
||||
name: String,
|
||||
args: String,
|
||||
state: InFlightToolCallState,
|
||||
},
|
||||
}
|
||||
|
||||
impl InFlightEvents {
|
||||
pub(crate) fn new(event_tx: broadcast::Sender<Event>) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(InFlightInner {
|
||||
next_block_id: 1,
|
||||
blocks: Vec::new(),
|
||||
})),
|
||||
event_tx,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn snapshot_guard(&self) -> MutexGuard<'_, InFlightInner> {
|
||||
self.inner.lock().expect("in-flight event mutex poisoned")
|
||||
}
|
||||
|
||||
pub(crate) fn start_text_block(&self) -> InFlightBlockId {
|
||||
let mut inner = self.lock();
|
||||
let block_id = inner.next_id();
|
||||
inner.blocks.push(TrackedBlock::Text {
|
||||
block_id,
|
||||
text: String::new(),
|
||||
finished: false,
|
||||
});
|
||||
block_id
|
||||
}
|
||||
|
||||
pub(crate) fn text_delta(&self, block_id: InFlightBlockId, text: String) {
|
||||
let mut inner = self.lock();
|
||||
if let Some(TrackedBlock::Text {
|
||||
text: current,
|
||||
finished,
|
||||
..
|
||||
}) = inner.find_block_mut(block_id)
|
||||
{
|
||||
current.push_str(&text);
|
||||
*finished = false;
|
||||
}
|
||||
let _ = self.event_tx.send(Event::TextDelta { text });
|
||||
}
|
||||
|
||||
pub(crate) fn text_done(&self, block_id: InFlightBlockId, text: String) {
|
||||
let mut inner = self.lock();
|
||||
if let Some(TrackedBlock::Text {
|
||||
text: current,
|
||||
finished,
|
||||
..
|
||||
}) = inner.find_block_mut(block_id)
|
||||
{
|
||||
if current.is_empty() {
|
||||
*current = text.clone();
|
||||
}
|
||||
*finished = true;
|
||||
}
|
||||
let _ = self.event_tx.send(Event::TextDone { text });
|
||||
}
|
||||
|
||||
pub(crate) fn thinking_start(&self) -> InFlightBlockId {
|
||||
let mut inner = self.lock();
|
||||
let block_id = inner.next_id();
|
||||
inner.blocks.push(TrackedBlock::Thinking {
|
||||
block_id,
|
||||
text: String::new(),
|
||||
finished: false,
|
||||
});
|
||||
let _ = self.event_tx.send(Event::ThinkingStart);
|
||||
block_id
|
||||
}
|
||||
|
||||
pub(crate) fn thinking_delta(&self, block_id: InFlightBlockId, text: String) {
|
||||
let mut inner = self.lock();
|
||||
if let Some(TrackedBlock::Thinking {
|
||||
text: current,
|
||||
finished,
|
||||
..
|
||||
}) = inner.find_block_mut(block_id)
|
||||
{
|
||||
current.push_str(&text);
|
||||
*finished = false;
|
||||
}
|
||||
let _ = self.event_tx.send(Event::ThinkingDelta { text });
|
||||
}
|
||||
|
||||
pub(crate) fn thinking_done(&self, block_id: InFlightBlockId, text: String) {
|
||||
let mut inner = self.lock();
|
||||
if let Some(TrackedBlock::Thinking {
|
||||
text: current,
|
||||
finished,
|
||||
..
|
||||
}) = inner.find_block_mut(block_id)
|
||||
{
|
||||
if current.is_empty() {
|
||||
*current = text.clone();
|
||||
}
|
||||
*finished = true;
|
||||
}
|
||||
let _ = self.event_tx.send(Event::ThinkingDone { text });
|
||||
}
|
||||
|
||||
pub(crate) fn tool_call_start(&self, id: String, name: String) -> InFlightBlockId {
|
||||
let mut inner = self.lock();
|
||||
let block_id = inner.next_id();
|
||||
inner.blocks.push(TrackedBlock::ToolCall {
|
||||
block_id,
|
||||
id: id.clone(),
|
||||
name: name.clone(),
|
||||
args: String::new(),
|
||||
state: InFlightToolCallState::Pending,
|
||||
});
|
||||
let _ = self.event_tx.send(Event::ToolCallStart { id, name });
|
||||
block_id
|
||||
}
|
||||
|
||||
pub(crate) fn tool_call_args_delta(
|
||||
&self,
|
||||
block_id: InFlightBlockId,
|
||||
id: String,
|
||||
delta: String,
|
||||
) {
|
||||
let mut inner = self.lock();
|
||||
if let Some(TrackedBlock::ToolCall { args, state, .. }) = inner.find_block_mut(block_id) {
|
||||
args.push_str(&delta);
|
||||
*state = InFlightToolCallState::StreamingArgs;
|
||||
}
|
||||
let _ = self
|
||||
.event_tx
|
||||
.send(Event::ToolCallArgsDelta { id, json: delta });
|
||||
}
|
||||
|
||||
pub(crate) fn tool_call_done(&self, block_id: InFlightBlockId, id: String, args: String) {
|
||||
let mut inner = self.lock();
|
||||
let mut name = String::new();
|
||||
if let Some(TrackedBlock::ToolCall {
|
||||
name: current_name,
|
||||
args: current,
|
||||
state,
|
||||
..
|
||||
}) = inner.find_block_mut(block_id)
|
||||
{
|
||||
name = current_name.clone();
|
||||
if current.is_empty() {
|
||||
*current = args.clone();
|
||||
}
|
||||
*state = InFlightToolCallState::Done;
|
||||
}
|
||||
let _ = self.event_tx.send(Event::ToolCallDone {
|
||||
id,
|
||||
name,
|
||||
arguments: args,
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) fn clear_for_committed_item_then<R>(
|
||||
&self,
|
||||
item: &LoggedItem,
|
||||
f: impl FnOnce() -> R,
|
||||
) -> R {
|
||||
let mut inner = self.lock();
|
||||
inner.clear_for_committed_item(item);
|
||||
f()
|
||||
}
|
||||
|
||||
fn lock(&self) -> MutexGuard<'_, InFlightInner> {
|
||||
self.inner.lock().expect("in-flight event mutex poisoned")
|
||||
}
|
||||
}
|
||||
|
||||
impl InFlightInner {
|
||||
fn next_id(&mut self) -> InFlightBlockId {
|
||||
let id = InFlightBlockId(self.next_block_id);
|
||||
self.next_block_id = self.next_block_id.saturating_add(1);
|
||||
id
|
||||
}
|
||||
|
||||
fn find_block_mut(&mut self, block_id: InFlightBlockId) -> Option<&mut TrackedBlock> {
|
||||
self.blocks
|
||||
.iter_mut()
|
||||
.find(|block| block.block_id() == block_id)
|
||||
}
|
||||
|
||||
fn clear_for_committed_item(&mut self, item: &LoggedItem) {
|
||||
match item {
|
||||
LoggedItem::Message { role, content }
|
||||
if matches!(role, session_store::LoggedRole::Assistant) =>
|
||||
{
|
||||
let text = content
|
||||
.iter()
|
||||
.filter_map(|part| match part {
|
||||
LoggedContentPart::Text { text } => Some(text.as_str()),
|
||||
LoggedContentPart::Refusal { refusal } => Some(refusal.as_str()),
|
||||
})
|
||||
.collect::<String>();
|
||||
if !text.is_empty() {
|
||||
self.remove_first_text_matching(&text);
|
||||
}
|
||||
}
|
||||
LoggedItem::Reasoning {
|
||||
text,
|
||||
summary,
|
||||
encrypted_content,
|
||||
..
|
||||
} => {
|
||||
let mut removed = false;
|
||||
if !text.is_empty() {
|
||||
removed |= self.remove_first_thinking_matching(text);
|
||||
}
|
||||
for summary_text in summary {
|
||||
if !summary_text.is_empty() {
|
||||
removed |= self.remove_first_thinking_matching(summary_text);
|
||||
}
|
||||
}
|
||||
if !removed && encrypted_content.is_some() {
|
||||
self.remove_first_empty_finished_thinking();
|
||||
}
|
||||
}
|
||||
LoggedItem::ToolCall { call_id, .. } => {
|
||||
self.remove_tool_call(call_id);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn snapshot(&self) -> InFlightSnapshot {
|
||||
InFlightSnapshot {
|
||||
blocks: self
|
||||
.blocks
|
||||
.iter()
|
||||
.filter_map(TrackedBlock::to_snapshot_block)
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_first_text_matching(&mut self, committed: &str) -> bool {
|
||||
if let Some(index) = self.blocks.iter().position(|block| match block {
|
||||
TrackedBlock::Text { text, .. } => text == committed,
|
||||
_ => false,
|
||||
}) {
|
||||
self.blocks.remove(index);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_first_thinking_matching(&mut self, committed: &str) -> bool {
|
||||
if let Some(index) = self.blocks.iter().position(|block| match block {
|
||||
TrackedBlock::Thinking { text, .. } => text == committed,
|
||||
_ => false,
|
||||
}) {
|
||||
self.blocks.remove(index);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_first_empty_finished_thinking(&mut self) -> bool {
|
||||
if let Some(index) = self.blocks.iter().position(|block| match block {
|
||||
TrackedBlock::Thinking { text, finished, .. } => text.is_empty() && *finished,
|
||||
_ => false,
|
||||
}) {
|
||||
self.blocks.remove(index);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_tool_call(&mut self, call_id: &str) {
|
||||
if let Some(index) = self.blocks.iter().position(|block| match block {
|
||||
TrackedBlock::ToolCall { id, .. } => id == call_id,
|
||||
_ => false,
|
||||
}) {
|
||||
self.blocks.remove(index);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TrackedBlock {
|
||||
fn block_id(&self) -> InFlightBlockId {
|
||||
match self {
|
||||
TrackedBlock::Text { block_id, .. }
|
||||
| TrackedBlock::Thinking { block_id, .. }
|
||||
| TrackedBlock::ToolCall { block_id, .. } => *block_id,
|
||||
}
|
||||
}
|
||||
|
||||
fn to_snapshot_block(&self) -> Option<InFlightBlock> {
|
||||
match self {
|
||||
TrackedBlock::Text { text, finished, .. } => {
|
||||
if text.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(InFlightBlock::Text {
|
||||
text: text.clone(),
|
||||
finished: *finished,
|
||||
})
|
||||
}
|
||||
}
|
||||
TrackedBlock::Thinking { text, finished, .. } => {
|
||||
if text.is_empty() && *finished {
|
||||
None
|
||||
} else {
|
||||
Some(InFlightBlock::Thinking {
|
||||
text: text.clone(),
|
||||
finished: *finished,
|
||||
})
|
||||
}
|
||||
}
|
||||
TrackedBlock::ToolCall {
|
||||
id,
|
||||
name,
|
||||
args,
|
||||
state,
|
||||
..
|
||||
} => Some(InFlightBlock::ToolCall {
|
||||
id: id.clone(),
|
||||
name: name.clone(),
|
||||
args: args.clone(),
|
||||
state: *state,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn snapshot_from_guard(guard: &MutexGuard<'_, InFlightInner>) -> InFlightSnapshot {
|
||||
guard.snapshot()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn snapshot_boundary_does_not_duplicate_or_gap_delta_sent_after_subscribe() {
|
||||
let (event_tx, _) = broadcast::channel(16);
|
||||
let in_flight = InFlightEvents::new(event_tx.clone());
|
||||
let block_id = in_flight.start_text_block();
|
||||
in_flight.text_delta(block_id, "hel".into());
|
||||
|
||||
let guard = in_flight.snapshot_guard();
|
||||
let mut rx = event_tx.subscribe();
|
||||
let snapshot = snapshot_from_guard(&guard);
|
||||
drop(guard);
|
||||
|
||||
in_flight.text_delta(block_id, "lo".into());
|
||||
|
||||
assert_eq!(
|
||||
snapshot.blocks,
|
||||
vec![InFlightBlock::Text {
|
||||
text: "hel".into(),
|
||||
finished: false,
|
||||
}]
|
||||
);
|
||||
assert!(matches!(
|
||||
rx.try_recv().unwrap(),
|
||||
Event::TextDelta { text } if text == "lo"
|
||||
));
|
||||
assert!(rx.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_log_and_in_flight_snapshot_prevents_mirror_only_assistant_gap() {
|
||||
use std::sync::mpsc;
|
||||
use std::thread;
|
||||
|
||||
use crate::segment_log_sink::SegmentLogSink;
|
||||
use session_store::{LogEntry, LoggedRole};
|
||||
|
||||
let (event_tx, _) = broadcast::channel(16);
|
||||
let sink = SegmentLogSink::new();
|
||||
let in_flight = InFlightEvents::new(event_tx);
|
||||
let block_id = in_flight.start_text_block();
|
||||
in_flight.text_delta(block_id, "done".into());
|
||||
in_flight.text_done(block_id, "done".into());
|
||||
|
||||
let assistant_item = LoggedItem::Message {
|
||||
role: LoggedRole::Assistant,
|
||||
content: vec![LoggedContentPart::Text {
|
||||
text: "done".into(),
|
||||
}],
|
||||
};
|
||||
let assistant_entry = LogEntry::AssistantItem {
|
||||
ts: 1,
|
||||
item: assistant_item.clone(),
|
||||
};
|
||||
|
||||
let in_flight_guard = in_flight.snapshot_guard();
|
||||
let in_flight_for_commit = in_flight.clone();
|
||||
let sink_for_commit = sink.clone();
|
||||
let (committed_tx, committed_rx) = mpsc::channel();
|
||||
let commit_thread = thread::spawn(move || {
|
||||
// This mirrors Worker::append_entry ordering: clear in-flight first,
|
||||
// then publish the finalized AssistantItem. AssistantItem entries
|
||||
// are mirror-only and are not delivered as live entry events.
|
||||
in_flight_for_commit.clear_for_committed_item_then(&assistant_item, || {
|
||||
sink_for_commit.publish(assistant_entry);
|
||||
});
|
||||
committed_tx.send(()).unwrap();
|
||||
});
|
||||
|
||||
let (entries_snapshot, mut entry_rx) = sink.subscribe_with_snapshot();
|
||||
let in_flight_snapshot = snapshot_from_guard(&in_flight_guard);
|
||||
drop(in_flight_guard);
|
||||
|
||||
committed_rx.recv().unwrap();
|
||||
commit_thread.join().unwrap();
|
||||
|
||||
assert!(entries_snapshot.is_empty());
|
||||
assert!(matches!(
|
||||
in_flight_snapshot.blocks.as_slice(),
|
||||
[InFlightBlock::Text { text, finished: true }] if text == "done"
|
||||
));
|
||||
assert!(entry_rx.try_recv().is_err());
|
||||
let post_commit_guard = in_flight.snapshot_guard();
|
||||
assert!(snapshot_from_guard(&post_commit_guard).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn committed_assistant_snapshot_does_not_duplicate_in_flight_block() {
|
||||
use crate::segment_log_sink::SegmentLogSink;
|
||||
use session_store::{LogEntry, LoggedRole};
|
||||
|
||||
let (event_tx, _) = broadcast::channel(16);
|
||||
let sink = SegmentLogSink::new();
|
||||
let in_flight = InFlightEvents::new(event_tx);
|
||||
let block_id = in_flight.start_text_block();
|
||||
in_flight.text_delta(block_id, "done".into());
|
||||
in_flight.text_done(block_id, "done".into());
|
||||
|
||||
let assistant_item = LoggedItem::Message {
|
||||
role: LoggedRole::Assistant,
|
||||
content: vec![LoggedContentPart::Text {
|
||||
text: "done".into(),
|
||||
}],
|
||||
};
|
||||
let assistant_entry = LogEntry::AssistantItem {
|
||||
ts: 1,
|
||||
item: assistant_item.clone(),
|
||||
};
|
||||
|
||||
in_flight.clear_for_committed_item_then(&assistant_item, || {
|
||||
sink.publish(assistant_entry);
|
||||
});
|
||||
|
||||
let in_flight_guard = in_flight.snapshot_guard();
|
||||
let (entries_snapshot, _entry_rx) = sink.subscribe_with_snapshot();
|
||||
let in_flight_snapshot = snapshot_from_guard(&in_flight_guard);
|
||||
|
||||
assert!(matches!(
|
||||
entries_snapshot.as_slice(),
|
||||
[LogEntry::AssistantItem { item, .. }] if item == &assistant_item
|
||||
));
|
||||
assert!(in_flight_snapshot.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn committed_item_clears_matching_in_flight_block() {
|
||||
let (event_tx, _) = broadcast::channel(16);
|
||||
let in_flight = InFlightEvents::new(event_tx);
|
||||
let block_id = in_flight.start_text_block();
|
||||
in_flight.text_delta(block_id, "done".into());
|
||||
in_flight.clear_for_committed_item_then(
|
||||
&LoggedItem::Message {
|
||||
role: session_store::LoggedRole::Assistant,
|
||||
content: vec![LoggedContentPart::Text {
|
||||
text: "done".into(),
|
||||
}],
|
||||
},
|
||||
|| (),
|
||||
);
|
||||
|
||||
let guard = in_flight.snapshot_guard();
|
||||
assert!(snapshot_from_guard(&guard).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn committed_reasoning_summary_clears_matching_in_flight_thinking_blocks() {
|
||||
let (event_tx, _) = broadcast::channel(16);
|
||||
let in_flight = InFlightEvents::new(event_tx);
|
||||
let first = in_flight.thinking_start();
|
||||
in_flight.thinking_delta(first, "summary A".into());
|
||||
in_flight.thinking_done(first, "".into());
|
||||
let second = in_flight.thinking_start();
|
||||
in_flight.thinking_delta(second, "summary B".into());
|
||||
in_flight.thinking_done(second, "".into());
|
||||
|
||||
in_flight.clear_for_committed_item_then(
|
||||
&LoggedItem::Reasoning {
|
||||
text: String::new(),
|
||||
summary: vec!["summary A".into(), "summary B".into()],
|
||||
encrypted_content: Some("opaque".into()),
|
||||
signature: None,
|
||||
},
|
||||
|| (),
|
||||
);
|
||||
|
||||
let guard = in_flight.snapshot_guard();
|
||||
assert!(snapshot_from_guard(&guard).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn committed_encrypted_only_reasoning_clears_empty_finished_thinking_block() {
|
||||
let (event_tx, _) = broadcast::channel(16);
|
||||
let in_flight = InFlightEvents::new(event_tx);
|
||||
let first = in_flight.thinking_start();
|
||||
in_flight.thinking_done(first, "".into());
|
||||
let second = in_flight.thinking_start();
|
||||
in_flight.thinking_delta(second, "still running".into());
|
||||
|
||||
in_flight.clear_for_committed_item_then(
|
||||
&LoggedItem::Reasoning {
|
||||
text: String::new(),
|
||||
summary: Vec::new(),
|
||||
encrypted_content: Some("opaque".into()),
|
||||
signature: None,
|
||||
},
|
||||
|| (),
|
||||
);
|
||||
|
||||
let guard = in_flight.snapshot_guard();
|
||||
assert_eq!(
|
||||
snapshot_from_guard(&guard).blocks,
|
||||
vec![InFlightBlock::Thinking {
|
||||
text: "still running".into(),
|
||||
finished: false,
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_omits_empty_finished_thinking_blocks() {
|
||||
let (event_tx, _) = broadcast::channel(16);
|
||||
let in_flight = InFlightEvents::new(event_tx);
|
||||
let empty_finished = in_flight.thinking_start();
|
||||
in_flight.thinking_done(empty_finished, "".into());
|
||||
let empty_running = in_flight.thinking_start();
|
||||
let visible_finished = in_flight.thinking_start();
|
||||
in_flight.thinking_delta(visible_finished, "visible".into());
|
||||
in_flight.thinking_done(visible_finished, "".into());
|
||||
|
||||
let guard = in_flight.snapshot_guard();
|
||||
assert_eq!(
|
||||
snapshot_from_guard(&guard).blocks,
|
||||
vec![
|
||||
InFlightBlock::Thinking {
|
||||
text: String::new(),
|
||||
finished: false,
|
||||
},
|
||||
InFlightBlock::Thinking {
|
||||
text: "visible".into(),
|
||||
finished: true,
|
||||
}
|
||||
]
|
||||
);
|
||||
assert_ne!(empty_running, empty_finished);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
//! Pre-run cleanup that fires when a Worker transitions out of `Paused`
|
||||
//! into a fresh turn via new user input.
|
||||
//!
|
||||
//! The previously in-flight turn is treated as finished. Any orphan
|
||||
//! `Item::ToolCall` (tool_use emitted by the LLM but whose tool did not
|
||||
//! run to completion before the pause) is closed with a synthetic
|
||||
//! `Item::ToolResult` so the next request is wire-valid under providers
|
||||
//! that require every `tool_use` to be followed by a matching
|
||||
//! `tool_result` (Anthropic). A short system note is then inserted so
|
||||
//! the LLM understands the prior work was cut short. Both side effects
|
||||
//! happen at the front of `Worker::run` when
|
||||
//! `worker.last_run_interrupted()` is set; see `Worker::apply_interrupt_prep`.
|
||||
|
||||
#[cfg(test)]
|
||||
use crate::prompt::catalog::PromptCatalog;
|
||||
use llm_engine::Item;
|
||||
|
||||
/// Build synthetic `Item::ToolResult` items for every unanswered
|
||||
/// `Item::ToolCall` in `history`, preserving order.
|
||||
pub(crate) fn orphan_tool_result_closures(history: &[Item], summary: &str) -> Vec<Item> {
|
||||
let mut answered: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
||||
for item in history {
|
||||
if let Item::ToolResult { call_id, .. } = item {
|
||||
answered.insert(call_id.as_str());
|
||||
}
|
||||
}
|
||||
let mut out = Vec::new();
|
||||
for item in history {
|
||||
if let Item::ToolCall { call_id, .. } = item {
|
||||
if !answered.contains(call_id.as_str()) {
|
||||
out.push(Item::tool_result(call_id.clone(), summary));
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Test-only helper to surface the canonical interrupt tool-result
|
||||
/// summary without round-tripping through a Worker — used by tests in
|
||||
/// this module that validate the closure logic.
|
||||
#[cfg(test)]
|
||||
fn interrupt_tool_result_summary() -> String {
|
||||
PromptCatalog::builtins_only()
|
||||
.unwrap()
|
||||
.interrupt_tool_result_summary()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn no_orphans_returns_empty() {
|
||||
let history = vec![Item::user_message("hi"), Item::assistant_message("hello")];
|
||||
let summary = interrupt_tool_result_summary();
|
||||
assert!(orphan_tool_result_closures(&history, &summary).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn paired_call_and_result_is_not_orphan() {
|
||||
let history = vec![
|
||||
Item::tool_call("c1", "Read", "{}"),
|
||||
Item::tool_result("c1", "ok"),
|
||||
];
|
||||
let summary = interrupt_tool_result_summary();
|
||||
assert!(orphan_tool_result_closures(&history, &summary).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unanswered_call_becomes_closure() {
|
||||
let history = vec![Item::tool_call("c1", "Read", "{}")];
|
||||
let summary = interrupt_tool_result_summary();
|
||||
let out = orphan_tool_result_closures(&history, &summary);
|
||||
assert_eq!(out.len(), 1);
|
||||
match &out[0] {
|
||||
Item::ToolResult {
|
||||
call_id,
|
||||
summary: got,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(call_id, "c1");
|
||||
assert_eq!(got, &summary);
|
||||
}
|
||||
other => panic!("expected ToolResult, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_orphans_are_closed_in_order() {
|
||||
let history = vec![
|
||||
Item::tool_call("c1", "Read", "{}"),
|
||||
Item::tool_call("c2", "Write", "{}"),
|
||||
Item::tool_result("c1", "ok"),
|
||||
Item::tool_call("c3", "Grep", "{}"),
|
||||
];
|
||||
let summary = interrupt_tool_result_summary();
|
||||
let out = orphan_tool_result_closures(&history, &summary);
|
||||
let ids: Vec<&str> = out
|
||||
.iter()
|
||||
.map(|i| match i {
|
||||
Item::ToolResult { call_id, .. } => call_id.as_str(),
|
||||
_ => unreachable!(),
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(ids, vec!["c2", "c3"]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
//! User-facing alert channel for Worker → client.
|
||||
//!
|
||||
//! Separate from `tracing` (which is for developer logs). Alerts
|
||||
//! are short human-readable messages the Worker layer wants a client to
|
||||
//! see — for example "compaction failed", "tool output truncated".
|
||||
//!
|
||||
//! Each alert is broadcast on the shared `Event` channel and
|
||||
//! also appended to an in-memory buffer so that clients connecting
|
||||
//! after the fact still see everything emitted during the session.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use protocol::{Alert, AlertLevel, AlertSource, Event};
|
||||
|
||||
/// Upper bound on buffered alerts. When exceeded, the oldest
|
||||
/// entries are discarded so a long-running session cannot leak
|
||||
/// memory through a pathological loop of recurring alerts
|
||||
/// (e.g. compaction failing every turn).
|
||||
const MAX_BUFFERED_ALERTS: usize = 512;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Alerter {
|
||||
inner: Arc<Inner>,
|
||||
}
|
||||
|
||||
struct Inner {
|
||||
event_tx: broadcast::Sender<Event>,
|
||||
buffer: Mutex<VecDeque<Alert>>,
|
||||
}
|
||||
|
||||
impl Alerter {
|
||||
pub fn new(event_tx: broadcast::Sender<Event>) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Inner {
|
||||
event_tx,
|
||||
buffer: Mutex::new(VecDeque::with_capacity(MAX_BUFFERED_ALERTS)),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Record and broadcast an alert.
|
||||
///
|
||||
/// The broadcast may have no subscribers (e.g. during Worker
|
||||
/// construction before any client has connected); the buffer
|
||||
/// guarantees the message is still delivered once a client
|
||||
/// attaches.
|
||||
///
|
||||
/// The buffer mutex is held across `broadcast::send` to make
|
||||
/// `subscribe_with_snapshot` race-free — a client that snapshots
|
||||
/// the buffer while holding the same lock sees every alert
|
||||
/// exactly once: older ones from the snapshot, newer ones from
|
||||
/// the freshly-subscribed receiver.
|
||||
pub fn alert(&self, level: AlertLevel, source: AlertSource, message: String) {
|
||||
let alert = Alert {
|
||||
level,
|
||||
source,
|
||||
message,
|
||||
timestamp_ms: now_ms(),
|
||||
};
|
||||
if let Ok(mut buf) = self.inner.buffer.lock() {
|
||||
if buf.len() >= MAX_BUFFERED_ALERTS {
|
||||
buf.pop_front();
|
||||
}
|
||||
buf.push_back(alert.clone());
|
||||
let _ = self.inner.event_tx.send(Event::Alert(alert));
|
||||
}
|
||||
}
|
||||
|
||||
/// Subscribe and atomically snapshot the current buffer.
|
||||
///
|
||||
/// The returned snapshot contains alerts emitted before
|
||||
/// this call; the receiver will deliver alerts emitted
|
||||
/// after. An alert cannot appear in both.
|
||||
pub fn subscribe_with_snapshot(&self) -> (Vec<Alert>, broadcast::Receiver<Event>) {
|
||||
let buf = self
|
||||
.inner
|
||||
.buffer
|
||||
.lock()
|
||||
.expect("alerter buffer mutex poisoned");
|
||||
let rx = self.inner.event_tx.subscribe();
|
||||
let snapshot: Vec<Alert> = buf.iter().cloned().collect();
|
||||
(snapshot, rx)
|
||||
}
|
||||
}
|
||||
|
||||
fn now_ms() -> i64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn alert_broadcasts_to_existing_subscriber() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(8);
|
||||
let alerter = Alerter::new(tx);
|
||||
let (_snapshot, mut rx) = alerter.subscribe_with_snapshot();
|
||||
|
||||
alerter.alert(
|
||||
AlertLevel::Warn,
|
||||
AlertSource::Compactor,
|
||||
"test message".into(),
|
||||
);
|
||||
|
||||
match rx.try_recv() {
|
||||
Ok(Event::Alert(a)) => assert_eq!(a.message, "test message"),
|
||||
other => panic!("unexpected event: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn late_subscriber_sees_earlier_alerts_via_snapshot() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(8);
|
||||
let alerter = Alerter::new(tx);
|
||||
|
||||
alerter.alert(AlertLevel::Error, AlertSource::Worker, "first".into());
|
||||
alerter.alert(AlertLevel::Warn, AlertSource::AgentsMd, "second".into());
|
||||
|
||||
let (snapshot, mut rx) = alerter.subscribe_with_snapshot();
|
||||
assert_eq!(snapshot.len(), 2);
|
||||
assert_eq!(snapshot[0].message, "first");
|
||||
assert_eq!(snapshot[1].message, "second");
|
||||
assert!(rx.try_recv().is_err()); // nothing pending on the receiver
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn buffer_discards_oldest_past_cap() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(1024);
|
||||
let alerter = Alerter::new(tx);
|
||||
|
||||
for i in 0..(MAX_BUFFERED_ALERTS + 50) {
|
||||
alerter.alert(AlertLevel::Warn, AlertSource::Engine, format!("msg-{i}"));
|
||||
}
|
||||
|
||||
let (snapshot, _rx) = alerter.subscribe_with_snapshot();
|
||||
assert_eq!(snapshot.len(), MAX_BUFFERED_ALERTS);
|
||||
// First 50 were evicted; the oldest remaining is msg-50.
|
||||
assert_eq!(snapshot.first().unwrap().message, "msg-50");
|
||||
let last = format!("msg-{}", MAX_BUFFERED_ALERTS + 49);
|
||||
assert_eq!(snapshot.last().unwrap().message, last);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subscribe_snapshot_and_live_do_not_overlap() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(8);
|
||||
let alerter = Alerter::new(tx);
|
||||
|
||||
alerter.alert(AlertLevel::Warn, AlertSource::Engine, "historic".into());
|
||||
let (snapshot, mut rx) = alerter.subscribe_with_snapshot();
|
||||
alerter.alert(AlertLevel::Error, AlertSource::Engine, "live".into());
|
||||
|
||||
assert_eq!(snapshot.len(), 1);
|
||||
assert_eq!(snapshot[0].message, "historic");
|
||||
match rx.try_recv() {
|
||||
Ok(Event::Alert(a)) => assert_eq!(a.message, "live"),
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
assert!(rx.try_recv().is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
//! `WorkerEvent` send / receive helpers.
|
||||
//!
|
||||
//! This module owns the parent-facing lifecycle-event primitive
|
||||
//! (`WorkerEvent`) that children fire upward on turn-end / error /
|
||||
//! shutdown / scope-sub-delegation. Three responsibilities live here:
|
||||
//!
|
||||
//! - **Send** a `Method::WorkerEvent` to the parent socket, fire-and-forget,
|
||||
//! logging failures without blocking the child.
|
||||
//! - **Render** agent-visible variants into human-readable strings for the
|
||||
//! parent's notification buffer. Control-plane-only variants may still have
|
||||
//! a renderer for diagnostics, but receive-side classification keeps them
|
||||
//! out of LLM history/context.
|
||||
//! - **Apply side effects** on the parent (registry / pod-registry
|
||||
//! updates) so that the receive path is idempotent and tolerant of
|
||||
//! out-of-order delivery.
|
||||
//!
|
||||
//! Transport is fire-and-forget — the ticket's decision is that
|
||||
//! callbacks are an optimisation and `ListWorkers` + `reclaim_stale` are
|
||||
//! the real fallback. This module is allowed to drop events on the
|
||||
//! floor (with a warn log) rather than retry.
|
||||
//!
|
||||
//! `apply_event_side_effects` takes its dependencies (registry, scope
|
||||
//! lock path, self identity) by reference so the caller owns lifetime
|
||||
//! and locking concerns.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use protocol::{Method, ScopeRule, WorkerEvent};
|
||||
|
||||
use crate::runtime::dir::SpawnedWorkerRecord;
|
||||
use crate::spawn::comm_tools::connect_and_send;
|
||||
use crate::spawn::registry::SpawnedWorkerRegistry;
|
||||
|
||||
/// Connect to `socket`, send a single `Method::WorkerEvent(event)`, and
|
||||
/// return. Used by children to report up to their parent.
|
||||
///
|
||||
/// This is a synchronous helper — callers that want fire-and-forget
|
||||
/// semantics should wrap the call in `tokio::spawn` themselves.
|
||||
pub async fn send_worker_event(socket: &Path, event: WorkerEvent) -> std::io::Result<()> {
|
||||
connect_and_send(socket, &Method::WorkerEvent(event)).await
|
||||
}
|
||||
|
||||
/// Spawn a fire-and-forget task that sends `event` to `socket`. If
|
||||
/// `socket` is `None`, no send happens (top-level Workers have no parent).
|
||||
/// Any send failure is logged at warn level but otherwise ignored —
|
||||
/// the parent is treated as best-effort.
|
||||
pub fn fire_and_forget(socket: Option<PathBuf>, event: WorkerEvent) {
|
||||
let Some(socket) = socket else { return };
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = send_worker_event(&socket, event).await {
|
||||
tracing::warn!(error = %e, socket = %socket.display(), "WorkerEvent send failed");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Render a variant into a one-line human-readable string.
|
||||
///
|
||||
/// 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 `ReadWorkerOutput` to fetch more
|
||||
/// detail if the event summary is not enough.
|
||||
pub fn render_event(event: &WorkerEvent) -> String {
|
||||
match event {
|
||||
WorkerEvent::TurnEnded { worker_name } => {
|
||||
format!("Worker `{worker_name}` finished a turn.")
|
||||
}
|
||||
WorkerEvent::Errored {
|
||||
worker_name,
|
||||
message,
|
||||
} => {
|
||||
format!("Worker `{worker_name}` reported an error: {message}")
|
||||
}
|
||||
WorkerEvent::ShutDown { worker_name } => {
|
||||
format!("Worker `{worker_name}` has stopped.")
|
||||
}
|
||||
WorkerEvent::ScopeSubDelegated {
|
||||
parent_worker,
|
||||
sub_worker,
|
||||
..
|
||||
} => {
|
||||
format!("Worker `{parent_worker}` spawned `{sub_worker}` and delegated scope to it.")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply the variant-specific side effect on the parent side.
|
||||
///
|
||||
/// All operations are idempotent so that out-of-order delivery (e.g.
|
||||
/// `TurnEnded` arriving after `ShutDown`) does not produce errors:
|
||||
///
|
||||
/// - `TurnEnded` / `Errored`: no system work; the LLM handles the
|
||||
/// semantic response.
|
||||
/// - `ShutDown`: remove the child from `spawned_workers.json`, Worker state,
|
||||
/// and reclaim its delegated scope/allocation. Missing entries are swallowed.
|
||||
/// - `ScopeSubDelegated`: register the grandchild locally and re-emit
|
||||
/// upward to our own parent if we have one. Duplicate grandchild
|
||||
/// entries (re-delivery) are swallowed.
|
||||
pub async fn apply_event_side_effects(
|
||||
event: &WorkerEvent,
|
||||
registry: &Arc<SpawnedWorkerRegistry>,
|
||||
self_name: &str,
|
||||
self_parent_socket: &Option<PathBuf>,
|
||||
) {
|
||||
match event {
|
||||
WorkerEvent::TurnEnded { .. } | WorkerEvent::Errored { .. } => {}
|
||||
|
||||
WorkerEvent::ShutDown { worker_name } => {
|
||||
if let Err(e) = registry.remove(worker_name).await {
|
||||
tracing::warn!(error = %e, worker = %worker_name, "registry remove on ShutDown failed");
|
||||
}
|
||||
}
|
||||
|
||||
WorkerEvent::ScopeSubDelegated {
|
||||
parent_worker,
|
||||
sub_worker,
|
||||
sub_socket,
|
||||
scope,
|
||||
} => {
|
||||
if registry.get(sub_worker).await.is_some() {
|
||||
return;
|
||||
}
|
||||
let callback_address = registry
|
||||
.get(parent_worker)
|
||||
.await
|
||||
.map(|r| r.socket_path)
|
||||
.unwrap_or_else(PathBuf::new);
|
||||
let record = SpawnedWorkerRecord {
|
||||
worker_name: sub_worker.clone(),
|
||||
socket_path: sub_socket.clone(),
|
||||
scope_delegated: scope.clone(),
|
||||
callback_address,
|
||||
};
|
||||
if let Err(e) = registry.add(record).await {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
sub_worker = %sub_worker,
|
||||
"registry add on ScopeSubDelegated failed"
|
||||
);
|
||||
}
|
||||
reemit_scope_sub_delegated(
|
||||
self_parent_socket,
|
||||
self_name,
|
||||
sub_worker.clone(),
|
||||
sub_socket.clone(),
|
||||
scope.clone(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn reemit_scope_sub_delegated(
|
||||
self_parent_socket: &Option<PathBuf>,
|
||||
self_name: &str,
|
||||
sub_worker: String,
|
||||
sub_socket: PathBuf,
|
||||
scope: Vec<ScopeRule>,
|
||||
) {
|
||||
let Some(parent_socket) = self_parent_socket.clone() else {
|
||||
return;
|
||||
};
|
||||
let event = WorkerEvent::ScopeSubDelegated {
|
||||
parent_worker: self_name.to_string(),
|
||||
sub_worker,
|
||||
sub_socket,
|
||||
scope,
|
||||
};
|
||||
fire_and_forget(Some(parent_socket), event);
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,6 @@
|
||||
pub mod alerter;
|
||||
pub mod event;
|
||||
pub mod server;
|
||||
|
||||
pub(crate) mod interceptor;
|
||||
pub(crate) mod notify_buffer;
|
||||
@@ -0,0 +1,198 @@
|
||||
//! Pending-notify buffer for `Method::Notify` and `Method::WorkerEvent`.
|
||||
//!
|
||||
//! Entries are queued here by the Controller (on receipt of the
|
||||
//! corresponding IPC method) and drained by
|
||||
//! `WorkerInterceptor::pending_history_appends`, which the Engine calls
|
||||
//! at the head of each turn loop iteration. The drain renders each
|
||||
//! pending entry into a typed `SystemItem` (with the `notify_wrapper`
|
||||
//! prompt applied), commits a `LogEntry::SystemItem` per entry through
|
||||
//! the session-log sink, and returns the corresponding
|
||||
//! `Item::system_message`s for the worker to append to its
|
||||
//! persistent history.
|
||||
//!
|
||||
//! This is the **single lane** for "system messages produced by Worker
|
||||
//! state that should land in the next LLM request": Notify,
|
||||
//! agent-visible WorkerEvent variants, and any future `<system-reminder>`
|
||||
//! injection all ride this queue.
|
||||
//! Per `tickets/notify-history-persist.md` and `AGENTS.md` (LLM
|
||||
//! context の加工原則), there is **no** "transient, history-skipping"
|
||||
//! lane — everything injected into a request is also committed to
|
||||
//! history so any LLM reaction has a visible trigger across turns,
|
||||
//! resume, and compaction, and so the Anthropic prompt cache prefix
|
||||
//! stays stable across requests.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use protocol::WorkerEvent;
|
||||
use session_store::SystemItem;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::prompt::catalog::{CatalogError, PromptCatalog};
|
||||
|
||||
/// Maximum queued pending entries. Oldest entries are dropped beyond this.
|
||||
const CAPACITY: usize = 128;
|
||||
|
||||
/// One pending entry awaiting drain into the next LLM request.
|
||||
///
|
||||
/// The buffer keeps the raw input shape so the drain step can decide
|
||||
/// the right `SystemItem` kind (and apply `notify_wrapper` to the
|
||||
/// rendered body) at the moment of commit, when the prompt catalog
|
||||
/// is available.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum PendingNotify {
|
||||
Notify { message: String },
|
||||
WorkerEvent { event: WorkerEvent },
|
||||
}
|
||||
|
||||
/// Shared, mutex-guarded buffer of pending entries.
|
||||
///
|
||||
/// Cloned between the Worker (producer) and WorkerInterceptor (consumer).
|
||||
#[derive(Clone, Default)]
|
||||
pub struct NotifyBuffer {
|
||||
inner: Arc<Mutex<VecDeque<PendingNotify>>>,
|
||||
}
|
||||
|
||||
impl NotifyBuffer {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Push a notify entry onto the queue. If the queue is full, the
|
||||
/// oldest entry is dropped and a `tracing::warn` is emitted — the
|
||||
/// caller should never hit this in normal operation.
|
||||
pub fn push_notify(&self, message: String) {
|
||||
self.push_entry(PendingNotify::Notify { message });
|
||||
}
|
||||
|
||||
/// Push a typed worker-event entry onto the queue.
|
||||
pub fn push_worker_event(&self, event: WorkerEvent) {
|
||||
self.push_entry(PendingNotify::WorkerEvent { event });
|
||||
}
|
||||
|
||||
fn push_entry(&self, entry: PendingNotify) {
|
||||
let mut q = self.inner.lock().expect("notify buffer poisoned");
|
||||
if q.len() >= CAPACITY {
|
||||
let dropped = q.pop_front();
|
||||
warn!(
|
||||
capacity = CAPACITY,
|
||||
dropped = ?dropped,
|
||||
"notify buffer overflow; dropped oldest"
|
||||
);
|
||||
}
|
||||
q.push_back(entry);
|
||||
}
|
||||
|
||||
/// Remove and return all pending entries in FIFO order.
|
||||
pub fn drain(&self) -> Vec<PendingNotify> {
|
||||
let mut q = self.inner.lock().expect("notify buffer poisoned");
|
||||
q.drain(..).collect()
|
||||
}
|
||||
|
||||
/// Number of pending entries. Primarily for tests.
|
||||
pub fn len(&self) -> usize {
|
||||
self.inner.lock().expect("notify buffer poisoned").len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
/// Render one pending entry into a typed `SystemItem`. The
|
||||
/// `notify_wrapper` prompt produces the LLM-context body for both
|
||||
/// `Notify` (raw message) and `WorkerEvent` (rendered event line).
|
||||
pub(crate) fn build_system_item(
|
||||
entry: &PendingNotify,
|
||||
prompts: &PromptCatalog,
|
||||
) -> Result<SystemItem, CatalogError> {
|
||||
match entry {
|
||||
PendingNotify::Notify { message } => {
|
||||
let body = prompts.notify_wrapper(message)?;
|
||||
Ok(SystemItem::Notification {
|
||||
message: message.clone(),
|
||||
body,
|
||||
})
|
||||
}
|
||||
PendingNotify::WorkerEvent { event } => {
|
||||
let rendered = session_store::render_worker_event(event);
|
||||
let body = prompts.notify_wrapper(&rendered)?;
|
||||
Ok(SystemItem::WorkerEvent {
|
||||
event: event.clone(),
|
||||
body,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn push_then_drain_preserves_order() {
|
||||
let buf = NotifyBuffer::new();
|
||||
buf.push_notify("one".into());
|
||||
buf.push_notify("two".into());
|
||||
let drained = buf.drain();
|
||||
assert_eq!(drained.len(), 2);
|
||||
match &drained[0] {
|
||||
PendingNotify::Notify { message } => assert_eq!(message, "one"),
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
assert!(buf.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capacity_drops_oldest() {
|
||||
let buf = NotifyBuffer::new();
|
||||
for i in 0..(CAPACITY + 5) {
|
||||
buf.push_notify(format!("msg{i}"));
|
||||
}
|
||||
let drained = buf.drain();
|
||||
assert_eq!(drained.len(), CAPACITY);
|
||||
match &drained[0] {
|
||||
PendingNotify::Notify { message } => assert_eq!(message, "msg5"),
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_system_item_for_notify_carries_wrapper_body() {
|
||||
let entry = PendingNotify::Notify {
|
||||
message: "hello".into(),
|
||||
};
|
||||
let catalog = PromptCatalog::builtins_only().unwrap();
|
||||
let item = build_system_item(&entry, &catalog).unwrap();
|
||||
match item {
|
||||
SystemItem::Notification { message, body } => {
|
||||
assert_eq!(message, "hello");
|
||||
assert!(body.contains("[Notification]"));
|
||||
assert!(body.contains("hello"));
|
||||
assert!(body.contains("not a blocking request"));
|
||||
}
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_system_item_for_worker_event_wraps_rendered_event_text() {
|
||||
let entry = PendingNotify::WorkerEvent {
|
||||
event: WorkerEvent::TurnEnded {
|
||||
worker_name: "child".into(),
|
||||
},
|
||||
};
|
||||
let catalog = PromptCatalog::builtins_only().unwrap();
|
||||
let item = build_system_item(&entry, &catalog).unwrap();
|
||||
match item {
|
||||
SystemItem::WorkerEvent { event, body } => {
|
||||
assert!(
|
||||
matches!(event, WorkerEvent::TurnEnded { ref worker_name } if worker_name == "child")
|
||||
);
|
||||
assert!(body.contains("[Notification]"));
|
||||
assert!(body.contains("`child`"));
|
||||
}
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,286 @@
|
||||
use std::io;
|
||||
use std::io::ErrorKind;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use protocol::stream::{JsonLineReader, JsonLineWriter};
|
||||
use tokio::net::UnixListener;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
use crate::controller::WorkerHandle;
|
||||
use crate::in_flight::snapshot_from_guard;
|
||||
use protocol::{Event, Method};
|
||||
|
||||
/// Unix socket server for Worker Protocol.
|
||||
///
|
||||
/// Listens on the Worker's runtime directory socket path.
|
||||
/// Each client connection gets bidirectional JSONL:
|
||||
/// - Client writes Method lines → forwarded to WorkerController
|
||||
/// - Worker events → written as Event lines to all connected clients
|
||||
pub struct SocketServer {
|
||||
_accept_task: JoinHandle<()>,
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
impl SocketServer {
|
||||
/// Start listening on the WorkerHandle's socket path.
|
||||
pub async fn start(handle: &WorkerHandle) -> Result<Self, io::Error> {
|
||||
let path = handle.runtime_dir.socket_path();
|
||||
|
||||
// Remove stale socket file if it exists
|
||||
let _ = tokio::fs::remove_file(&path).await;
|
||||
|
||||
let listener = UnixListener::bind(&path)?;
|
||||
let handle = handle.clone();
|
||||
|
||||
let _accept_task = tokio::spawn(async move {
|
||||
loop {
|
||||
match listener.accept().await {
|
||||
Ok((stream, _)) => {
|
||||
let handle = handle.clone();
|
||||
tokio::spawn(handle_connection(stream, handle));
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Self { _accept_task, path })
|
||||
}
|
||||
|
||||
/// The socket file path.
|
||||
pub fn path(&self) -> &std::path::Path {
|
||||
&self.path
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for SocketServer {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::remove_file(&self.path);
|
||||
}
|
||||
}
|
||||
|
||||
fn is_peer_disconnect_read_error(error: &io::Error) -> bool {
|
||||
matches!(
|
||||
error.kind(),
|
||||
ErrorKind::ConnectionReset
|
||||
| ErrorKind::ConnectionAborted
|
||||
| ErrorKind::BrokenPipe
|
||||
| ErrorKind::UnexpectedEof
|
||||
)
|
||||
}
|
||||
|
||||
fn live_entry_event(entry: session_store::LogEntry) -> Option<Event> {
|
||||
match entry {
|
||||
session_store::LogEntry::SegmentStart { .. } => {
|
||||
let value = serde_json::to_value(&entry).expect("LogEntry is Serialize");
|
||||
Some(Event::SegmentRotated { entry: value })
|
||||
}
|
||||
session_store::LogEntry::UserInput { segments, .. } => {
|
||||
Some(Event::UserMessage { segments })
|
||||
}
|
||||
session_store::LogEntry::SystemItem { item, .. } => {
|
||||
let value = serde_json::to_value(&item).expect("SystemItem is Serialize");
|
||||
Some(Event::SystemItem { item: value })
|
||||
}
|
||||
session_store::LogEntry::Invoke { trigger, .. } => {
|
||||
Some(Event::InvokeStart { kind: trigger })
|
||||
}
|
||||
other => {
|
||||
// `SegmentLogSink::is_live_relevant` keeps non-live-relevant
|
||||
// variants off the broadcast lane; reaching here means the two
|
||||
// are out of sync and we silently dropped a wire event. Log so a
|
||||
// future regression surfaces instead of vanishing.
|
||||
tracing::error!(
|
||||
entry_kind = ?std::mem::discriminant(&other),
|
||||
"session-log broadcast emitted a non-live-relevant entry; \
|
||||
sink filter and IPC dispatch are out of sync"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_connection(stream: tokio::net::UnixStream, handle: WorkerHandle) {
|
||||
let (reader, writer) = stream.into_split();
|
||||
let mut reader = JsonLineReader::new(reader);
|
||||
let mut writer = JsonLineWriter::new(writer);
|
||||
|
||||
// Hold the in-flight stream lock while taking the session-log mirror
|
||||
// snapshot. `LogEntry::AssistantItem` is mirror-only for live clients,
|
||||
// so a finalized assistant block must be observed either as an already
|
||||
// committed entry or as the still-present in-flight block. This lock
|
||||
// order matches `append_entry` (in-flight clear before sink publish) and
|
||||
// keeps the snapshot/live boundary gap-free.
|
||||
let (entries_snapshot, mut entry_rx, alert_snapshot, mut rx, in_flight) = {
|
||||
let in_flight_guard = handle.in_flight.snapshot_guard();
|
||||
let (entries_snapshot, entry_rx) = handle.sink.subscribe_with_snapshot();
|
||||
|
||||
// Atomically subscribe and snapshot buffered alerts so that warnings
|
||||
// emitted before this client connected are replayed exactly once.
|
||||
let (alert_snapshot, rx) = handle.alerter.subscribe_with_snapshot();
|
||||
let in_flight = snapshot_from_guard(&in_flight_guard);
|
||||
(entries_snapshot, entry_rx, alert_snapshot, rx, in_flight)
|
||||
};
|
||||
for alert in alert_snapshot {
|
||||
if writer.write(&Event::Alert(alert)).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Send the typed snapshot up front so late attachers can
|
||||
// reconstruct view state without an extra round trip.
|
||||
let snapshot_event = Event::Snapshot {
|
||||
entries: entries_snapshot
|
||||
.into_iter()
|
||||
.map(|e| serde_json::to_value(&e).expect("LogEntry is Serialize"))
|
||||
.collect(),
|
||||
greeting: handle.shared_state.greeting.clone(),
|
||||
status: handle.shared_state.get_status(),
|
||||
in_flight,
|
||||
};
|
||||
if writer.write(&snapshot_event).await.is_err() {
|
||||
return;
|
||||
}
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// Live session-log entries → dispatched as the role-specific
|
||||
// wire events. `SegmentLogSink` only broadcasts committed log
|
||||
// entries with live UI meaning; `UserInput` travels this lane so
|
||||
// the visible user line is ordered with `SegmentStart` rotation.
|
||||
entry = entry_rx.recv() => {
|
||||
match entry {
|
||||
Ok(entry) => {
|
||||
if let Some(event) = live_entry_event(entry) {
|
||||
if writer.write(&event).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
|
||||
// Slow client fell behind the broadcast buffer.
|
||||
// Drop the connection so the next reconnect
|
||||
// re-seeds the prefix via subscribe_with_snapshot.
|
||||
break;
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
}
|
||||
}
|
||||
// Broadcast events → this client
|
||||
event = rx.recv() => {
|
||||
match event {
|
||||
Ok(event) => {
|
||||
if writer.write(&event).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
// Client methods → handle or forward to controller
|
||||
method = reader.next::<Method>() => {
|
||||
match method {
|
||||
Ok(Some(Method::ListCompletions { kind, prefix })) => {
|
||||
let entries = match kind {
|
||||
protocol::CompletionKind::File => handle
|
||||
.shared_state
|
||||
.fs_view()
|
||||
.map(|view| view.list_file_completions(&prefix))
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|c| protocol::CompletionEntry {
|
||||
value: c.path,
|
||||
is_dir: c.is_dir,
|
||||
})
|
||||
.collect(),
|
||||
protocol::CompletionKind::Knowledge => handle
|
||||
.shared_state
|
||||
.list_knowledge_completions(&prefix)
|
||||
.into_iter()
|
||||
.map(|c| protocol::CompletionEntry {
|
||||
value: c.slug,
|
||||
is_dir: false,
|
||||
})
|
||||
.collect(),
|
||||
protocol::CompletionKind::Workflow => handle
|
||||
.shared_state
|
||||
.list_workflow_completions(&prefix)
|
||||
.into_iter()
|
||||
.map(|c| protocol::CompletionEntry {
|
||||
value: c.slug,
|
||||
is_dir: false,
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
if writer
|
||||
.write(&Event::Completions { kind, entries })
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Some(method)) => {
|
||||
let _ = handle.send(method).await;
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(e) if is_peer_disconnect_read_error(&e) => break,
|
||||
Err(e) => {
|
||||
if writer
|
||||
.write(&Event::Error {
|
||||
code: protocol::ErrorCode::InvalidRequest,
|
||||
message: format!("invalid method: {e}"),
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn peer_disconnect_read_errors_are_connection_close() {
|
||||
for kind in [
|
||||
ErrorKind::ConnectionReset,
|
||||
ErrorKind::ConnectionAborted,
|
||||
ErrorKind::BrokenPipe,
|
||||
ErrorKind::UnexpectedEof,
|
||||
] {
|
||||
let error = io::Error::new(kind, "peer disconnected");
|
||||
assert!(
|
||||
is_peer_disconnect_read_error(&error),
|
||||
"{kind:?} should be treated as a normal peer disconnect"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_data_is_not_peer_disconnect() {
|
||||
let error = io::Error::new(ErrorKind::InvalidData, "malformed method");
|
||||
assert!(!is_peer_disconnect_read_error(&error));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_input_log_entry_maps_to_user_message_event() {
|
||||
let segments = vec![protocol::Segment::text("hello from log")];
|
||||
let event = live_entry_event(session_store::LogEntry::UserInput {
|
||||
ts: session_store::segment_log::now_millis(),
|
||||
segments: segments.clone(),
|
||||
})
|
||||
.expect("UserInput must be live-relevant");
|
||||
|
||||
match event {
|
||||
Event::UserMessage { segments: echoed } => assert_eq!(echoed, segments),
|
||||
other => panic!("expected UserMessage, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
pub mod active_workflow;
|
||||
pub mod compact;
|
||||
pub mod controller;
|
||||
pub mod discovery;
|
||||
pub mod entrypoint;
|
||||
pub mod feature;
|
||||
pub mod fs_view;
|
||||
pub mod hook;
|
||||
pub(crate) mod in_flight;
|
||||
pub mod ipc;
|
||||
pub mod prompt;
|
||||
pub mod runtime;
|
||||
pub mod segment_log_sink;
|
||||
pub mod shared_state;
|
||||
mod shutdown_after_idle;
|
||||
pub mod spawn;
|
||||
pub mod workflow;
|
||||
|
||||
mod interrupt_prep;
|
||||
mod permission;
|
||||
mod ticket_event_notify;
|
||||
mod worker;
|
||||
|
||||
pub use compact::token_counter::{EstimateSource, SplitPoint, TokenEstimate};
|
||||
pub use controller::{ShutdownReceiver, WorkerController, WorkerHandle};
|
||||
pub use hook::{Hook, HookEventKind, HookRegistryBuilder};
|
||||
pub use ipc::alerter::Alerter;
|
||||
pub use ipc::server::SocketServer;
|
||||
pub use manifest::{
|
||||
AuthRef, ModelManifest, SchemeKind, Scope, WorkerManifest, WorkerManifestConfig,
|
||||
WorkerMetaConfig,
|
||||
};
|
||||
pub use prompt::catalog::{CatalogError, PromptCatalog, WorkerPrompt};
|
||||
pub use prompt::loader::PromptLoader;
|
||||
pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTemplate};
|
||||
pub use protocol::{ErrorCode, Event, Method, TurnResult, WorkerStatus};
|
||||
pub use provider::{ProviderError, build_client};
|
||||
pub use runtime::dir::RuntimeDir;
|
||||
pub use segment_log_sink::SegmentLogSink;
|
||||
pub use shared_state::WorkerSharedState;
|
||||
pub use worker::{Worker, WorkerError, WorkerRunResult, apply_worker_manifest};
|
||||
@@ -0,0 +1,221 @@
|
||||
use async_trait::async_trait;
|
||||
use llm_engine::llm_client::client::LlmClient;
|
||||
use manifest::{ToolPermissionAction, ToolPermissionConfig};
|
||||
use serde_json::Value;
|
||||
use session_store::Store;
|
||||
|
||||
use crate::Worker;
|
||||
use crate::hook::{Hook, HookPreToolAction, PreToolCall, ToolCallSummary};
|
||||
|
||||
/// Built-in manifest permission policy for `PreToolCall`.
|
||||
///
|
||||
/// This hook is registered by Worker before user hooks, so manifest-level deny
|
||||
/// rules fail closed before user extension code can approve a call.
|
||||
pub(crate) struct PermissionHook {
|
||||
config: ToolPermissionConfig,
|
||||
}
|
||||
|
||||
impl PermissionHook {
|
||||
pub(crate) fn new(config: ToolPermissionConfig) -> Self {
|
||||
Self { config }
|
||||
}
|
||||
|
||||
fn action_for(&self, input: &ToolCallSummary) -> ToolPermissionAction {
|
||||
let target = permission_target(&input.arguments);
|
||||
self.config
|
||||
.rules
|
||||
.iter()
|
||||
.find(|rule| {
|
||||
rule.tool.eq_ignore_ascii_case(&input.tool_name)
|
||||
&& wildcard_match(&rule.pattern, &target)
|
||||
})
|
||||
.map(|rule| rule.action)
|
||||
.unwrap_or(self.config.default_action)
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: LlmClient, St: Store> Worker<C, St> {
|
||||
pub(crate) fn apply_permissions_from_manifest(&mut self) {
|
||||
let Some(permissions) = self.manifest().permissions.clone() else {
|
||||
return;
|
||||
};
|
||||
self.add_pre_tool_call_hook(PermissionHook::new(permissions));
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreToolCall> for PermissionHook {
|
||||
async fn call(&self, input: &ToolCallSummary) -> HookPreToolAction {
|
||||
match self.action_for(input) {
|
||||
ToolPermissionAction::Allow => HookPreToolAction::Continue,
|
||||
ToolPermissionAction::Deny => HookPreToolAction::Deny(permission_denied_message(input)),
|
||||
ToolPermissionAction::Ask => {
|
||||
HookPreToolAction::Deny(permission_ask_unsupported_message(input))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn permission_denied_message(input: &ToolCallSummary) -> String {
|
||||
format!(
|
||||
"permission denied: tool `{}` arguments matched the manifest permission policy",
|
||||
input.tool_name
|
||||
)
|
||||
}
|
||||
|
||||
fn permission_ask_unsupported_message(input: &ToolCallSummary) -> String {
|
||||
format!(
|
||||
"permission ask unsupported: tool `{}` requires approval, but this runtime has no permission approval protocol; denied fail-closed",
|
||||
input.tool_name
|
||||
)
|
||||
}
|
||||
|
||||
fn permission_target(arguments: &Value) -> String {
|
||||
if let Value::Object(map) = arguments {
|
||||
for key in ["command", "file_path", "path", "pattern", "query", "url"] {
|
||||
if let Some(value) = map.get(key).and_then(Value::as_str) {
|
||||
return value.to_string();
|
||||
}
|
||||
}
|
||||
}
|
||||
serde_json::to_string(arguments).unwrap_or_else(|_| arguments.to_string())
|
||||
}
|
||||
|
||||
fn wildcard_match(pattern: &str, text: &str) -> bool {
|
||||
let pattern = pattern.as_bytes();
|
||||
let text = text.as_bytes();
|
||||
let (mut pi, mut ti) = (0usize, 0usize);
|
||||
let mut star: Option<usize> = None;
|
||||
let mut star_text = 0usize;
|
||||
|
||||
while ti < text.len() {
|
||||
if pi < pattern.len() && (pattern[pi] == b'?' || pattern[pi] == text[ti]) {
|
||||
pi += 1;
|
||||
ti += 1;
|
||||
} else if pi < pattern.len() && pattern[pi] == b'*' {
|
||||
star = Some(pi);
|
||||
pi += 1;
|
||||
star_text = ti;
|
||||
} else if let Some(star_pi) = star {
|
||||
pi = star_pi + 1;
|
||||
star_text += 1;
|
||||
ti = star_text;
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
while pi < pattern.len() && pattern[pi] == b'*' {
|
||||
pi += 1;
|
||||
}
|
||||
|
||||
pi == pattern.len()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::hook::HookPreToolAction;
|
||||
use manifest::ToolPermissionRule;
|
||||
|
||||
fn summary(tool_name: &str, arguments: Value) -> ToolCallSummary {
|
||||
ToolCallSummary {
|
||||
call_id: "call_1".into(),
|
||||
tool_name: tool_name.into(),
|
||||
arguments,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_matching_rule_wins_by_declaration_order() {
|
||||
let hook = PermissionHook::new(ToolPermissionConfig {
|
||||
default_action: ToolPermissionAction::Deny,
|
||||
rules: vec![
|
||||
ToolPermissionRule {
|
||||
tool: "bash".into(),
|
||||
pattern: "git *".into(),
|
||||
action: ToolPermissionAction::Allow,
|
||||
},
|
||||
ToolPermissionRule {
|
||||
tool: "Bash".into(),
|
||||
pattern: "git reset *".into(),
|
||||
action: ToolPermissionAction::Deny,
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
let input = summary("Bash", serde_json::json!({ "command": "git reset --hard" }));
|
||||
|
||||
assert_eq!(hook.action_for(&input), ToolPermissionAction::Allow);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_action_applies_when_no_rule_matches() {
|
||||
let hook = PermissionHook::new(ToolPermissionConfig {
|
||||
default_action: ToolPermissionAction::Deny,
|
||||
rules: Vec::new(),
|
||||
});
|
||||
|
||||
let input = summary("Read", serde_json::json!({ "file_path": "/tmp/a.txt" }));
|
||||
|
||||
assert_eq!(hook.action_for(&input), ToolPermissionAction::Deny);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deny_and_ask_fail_closed_as_public_deny_actions() {
|
||||
let deny = PermissionHook::new(ToolPermissionConfig {
|
||||
default_action: ToolPermissionAction::Deny,
|
||||
rules: Vec::new(),
|
||||
});
|
||||
let denied = deny
|
||||
.call(&summary(
|
||||
"Bash",
|
||||
serde_json::json!({ "command": "rm -rf target" }),
|
||||
))
|
||||
.await;
|
||||
match denied {
|
||||
HookPreToolAction::Deny(message) => {
|
||||
assert!(message.contains("permission denied"));
|
||||
assert!(message.contains("Bash"));
|
||||
}
|
||||
other => panic!("expected fail-closed deny action, got {other:?}"),
|
||||
}
|
||||
|
||||
let ask = PermissionHook::new(ToolPermissionConfig {
|
||||
default_action: ToolPermissionAction::Ask,
|
||||
rules: Vec::new(),
|
||||
});
|
||||
let asked = ask
|
||||
.call(&summary(
|
||||
"Bash",
|
||||
serde_json::json!({ "command": "git status" }),
|
||||
))
|
||||
.await;
|
||||
match asked {
|
||||
HookPreToolAction::Deny(message) => {
|
||||
assert!(message.contains("permission ask unsupported"));
|
||||
assert!(message.contains("denied fail-closed"));
|
||||
}
|
||||
other => panic!("expected ask fail-closed deny action, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn target_prefers_known_builtin_argument_fields() {
|
||||
assert_eq!(
|
||||
permission_target(&serde_json::json!({ "command": "rm -rf target" })),
|
||||
"rm -rf target"
|
||||
);
|
||||
assert_eq!(
|
||||
permission_target(&serde_json::json!({ "file_path": "/tmp/.env" })),
|
||||
"/tmp/.env"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wildcard_supports_star_and_question() {
|
||||
assert!(wildcard_match("rm *", "rm -rf target"));
|
||||
assert!(wildcard_match("file?.rs", "file1.rs"));
|
||||
assert!(!wildcard_match("rm *", "git status"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
//! `AGENTS.md` ingestion for system-prompt templates.
|
||||
//!
|
||||
//! Reads `AGENTS.md` directly under the Worker cwd and exposes its body
|
||||
//! to the template engine through `SystemPromptContext.agents_md`.
|
||||
//! Nested / parent-directory AGENTS.md files are intentionally ignored;
|
||||
//! subproject context is expressed by launching a Worker with that
|
||||
//! directory as cwd.
|
||||
//!
|
||||
//! No size cap is applied here — the whole file is read and embedded.
|
||||
//! System-prompt-size policing is the responsibility of a higher layer
|
||||
//! (Usage-driven warning after the first LLM round-trip).
|
||||
|
||||
use std::fs;
|
||||
use std::io::ErrorKind;
|
||||
use std::path::Path;
|
||||
|
||||
use tracing::warn;
|
||||
|
||||
/// Outcome of an `AGENTS.md` ingestion attempt.
|
||||
///
|
||||
/// `body` carries the text that should be handed to the template
|
||||
/// engine (if any); `warnings` are short human-readable messages that
|
||||
/// Worker forwards to the user-facing notification channel. The caller
|
||||
/// also gets `tracing::warn!` lines for the developer log.
|
||||
pub(crate) struct AgentsMdResult {
|
||||
pub body: Option<String>,
|
||||
pub warnings: Vec<String>,
|
||||
}
|
||||
|
||||
/// Read `AGENTS.md` from `cwd` if present. All non-fatal problems are
|
||||
/// both logged via `tracing::warn!` (developer-facing) and surfaced
|
||||
/// via `AgentsMdResult::warnings` (user-facing).
|
||||
///
|
||||
/// - Absent: `body = None`, no warning.
|
||||
/// - Non-UTF-8 or I/O error: `body = None`, warning.
|
||||
pub(crate) fn read_agents_md(cwd: &Path) -> AgentsMdResult {
|
||||
let path = cwd.join("AGENTS.md");
|
||||
let mut warnings = Vec::new();
|
||||
|
||||
match fs::read_to_string(&path) {
|
||||
Ok(body) => AgentsMdResult {
|
||||
body: Some(body),
|
||||
warnings,
|
||||
},
|
||||
Err(e) if e.kind() == ErrorKind::NotFound => AgentsMdResult {
|
||||
body: None,
|
||||
warnings,
|
||||
},
|
||||
Err(e) if e.kind() == ErrorKind::InvalidData => {
|
||||
warn!(path = %path.display(), error = %e, "AGENTS.md is not valid UTF-8");
|
||||
warnings.push(format!(
|
||||
"AGENTS.md ({}) is not valid UTF-8: {}",
|
||||
path.display(),
|
||||
e
|
||||
));
|
||||
AgentsMdResult {
|
||||
body: None,
|
||||
warnings,
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(path = %path.display(), error = %e, "failed to read AGENTS.md");
|
||||
warnings.push(format!(
|
||||
"failed to read AGENTS.md ({}): {}",
|
||||
path.display(),
|
||||
e
|
||||
));
|
||||
AgentsMdResult {
|
||||
body: None,
|
||||
warnings,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn absent_file_returns_none() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
assert!(read_agents_md(dir.path()).body.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reads_small_file_verbatim() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
fs::write(dir.path().join("AGENTS.md"), "# hello\nworld").unwrap();
|
||||
let result = read_agents_md(dir.path());
|
||||
assert_eq!(result.body.as_deref(), Some("# hello\nworld"));
|
||||
assert!(result.warnings.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reads_large_file_verbatim() {
|
||||
// Previously truncated at 64KB; now read whole. Size-policing
|
||||
// is deferred to the Usage-driven warning layer.
|
||||
let dir = TempDir::new().unwrap();
|
||||
let body = "a".repeat(128 * 1024);
|
||||
fs::write(dir.path().join("AGENTS.md"), &body).unwrap();
|
||||
let result = read_agents_md(dir.path());
|
||||
assert_eq!(result.body.as_ref().map(String::len), Some(128 * 1024));
|
||||
assert!(result.warnings.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_utf8_surfaces_warning() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
fs::write(dir.path().join("AGENTS.md"), [0xff, 0xfe, 0xfd]).unwrap();
|
||||
let result = read_agents_md(dir.path());
|
||||
assert!(result.body.is_none());
|
||||
assert_eq!(result.warnings.len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,771 @@
|
||||
//! Central catalog of Worker-level prompt strings.
|
||||
//!
|
||||
//! Prompts that Worker injects into a Engine (compaction system prompt,
|
||||
//! notification wrapper, interrupt notes, system-prompt trailing
|
||||
//! sections, AGENTS.md truncation notice, ...) are enumerated by
|
||||
//! [`WorkerPrompt`] and rendered through a single [`PromptCatalog`]. Direct
|
||||
//! `const &str` / `format!` authoring of these strings elsewhere in
|
||||
//! `crates/worker` is deliberately avoided — new injection points add a
|
||||
//! variant here, which forces a matching entry in
|
||||
//! `resources/prompts/internal.toml` (checked at build time) and keeps
|
||||
//! the "Worker tone" editable in one place.
|
||||
//!
|
||||
//! # Layering
|
||||
//!
|
||||
//! Values are merged key-wise from low priority to high:
|
||||
//!
|
||||
//! 1. **builtin** — `resources/prompts/internal.toml`, baked into the
|
||||
//! binary. Must cover every [`WorkerPrompt`] variant (build-time check).
|
||||
//! 2. **user** — `<config_dir>/prompts.toml`, when a caller supplies it.
|
||||
//! Optional.
|
||||
//! 3. **workspace** — `<project>/.yoi/prompts.toml`, when a caller
|
||||
//! supplies it. Optional.
|
||||
//! 4. **manifest pack** — `manifest.worker.prompt_pack`, an explicit path
|
||||
//! per-Worker. Optional.
|
||||
//!
|
||||
//! Unknown keys in layers 2–4 are logged via `tracing::warn!` and
|
||||
//! ignored (forward compatibility). Layer 1 is enforced at build time.
|
||||
//!
|
||||
//! # Template language
|
||||
//!
|
||||
//! All values are minijinja templates. `{% include "$prefix/..." %}`
|
||||
//! resolves through the same [`PromptLoader`] used by the system-prompt
|
||||
//! template, so long prompt bodies can be factored into `.md` files
|
||||
//! under `resources/prompts/...`, the user prompts library, or the
|
||||
//! workspace prompts library.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use minijinja::value::Value;
|
||||
use minijinja::{Environment, ErrorKind, UndefinedBehavior};
|
||||
use serde::Deserialize;
|
||||
use thiserror::Error;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::prompt::loader::PromptLoader;
|
||||
|
||||
// Generated by build.rs from `resources/prompts/internal.toml`.
|
||||
include!(concat!(env!("OUT_DIR"), "/internal_keys.rs"));
|
||||
|
||||
/// Source of the builtin pack. Baked in at compile time.
|
||||
const INTERNAL_TOML: &str = include_str!("../../../../resources/prompts/internal.toml");
|
||||
|
||||
/// Worker-level prompt injection point.
|
||||
///
|
||||
/// Adding a new variant also requires adding a matching key to
|
||||
/// `resources/prompts/internal.toml`; the build fails otherwise.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum WorkerPrompt {
|
||||
/// System prompt of the compaction (summary) Engine.
|
||||
CompactSystem,
|
||||
/// System prompt of the memory extract Engine.
|
||||
MemoryExtractSystem,
|
||||
/// System prompt of the memory consolidation (integration + tidy) Engine.
|
||||
MemoryConsolidationSystem,
|
||||
/// Wrapper around an incoming `Method::Notify` message injected into
|
||||
/// the next LLM request context as a transient system message.
|
||||
NotifyWrapper,
|
||||
/// Synthetic `Item::ToolResult` summary used to close out orphaned
|
||||
/// tool calls when a paused turn is interrupted by the user.
|
||||
InterruptToolResultSummary,
|
||||
/// System note prepended to the new turn after an interrupt.
|
||||
InterruptSystemNote,
|
||||
/// Trailing `## Working boundaries` section appended to every
|
||||
/// materialised system prompt.
|
||||
WorkingBoundariesSection,
|
||||
/// Trailing `## Project instructions (AGENTS.md)` section, appended
|
||||
/// after the scope summary when an AGENTS.md is present.
|
||||
AgentsMdSection,
|
||||
/// Trailing `## Resident memory summary` section, appended after the
|
||||
/// AGENTS.md section when memory is enabled, summary injection is enabled,
|
||||
/// and `memory/summary.md` has a valid non-empty body.
|
||||
ResidentMemorySummarySection,
|
||||
/// Trailing `## Resident knowledge` section, appended after the
|
||||
/// resident memory summary when memory is enabled, Knowledge resident
|
||||
/// injection is enabled, and at least one `knowledge/*` record advertises
|
||||
/// `model_invokation: true`.
|
||||
ResidentKnowledgeSection,
|
||||
/// Trailing `## Resident workflows` section, appended after resident
|
||||
/// knowledge when Workflow resident injection is enabled and at least one
|
||||
/// workflow advertises `model_invokation: true`.
|
||||
ResidentWorkflowsSection,
|
||||
/// Trailing Worker orchestration guidance, appended when registered tools
|
||||
/// include Worker-management capabilities.
|
||||
WorkerOrchestrationGuidanceSection,
|
||||
/// Weak Companion Notify payload for explicit Orchestrator Ticket events.
|
||||
TicketEventCompanionNotice,
|
||||
/// LLM-facing description for the SpawnWorker tool, including discovered
|
||||
/// profile selectors.
|
||||
SpawnWorkerToolDescription,
|
||||
}
|
||||
|
||||
impl WorkerPrompt {
|
||||
pub fn key(self) -> &'static str {
|
||||
match self {
|
||||
Self::CompactSystem => "compact_system",
|
||||
Self::MemoryExtractSystem => "memory_extract_system",
|
||||
Self::MemoryConsolidationSystem => "memory_consolidation_system",
|
||||
Self::NotifyWrapper => "notify_wrapper",
|
||||
Self::InterruptToolResultSummary => "interrupt_tool_result_summary",
|
||||
Self::InterruptSystemNote => "interrupt_system_note",
|
||||
Self::WorkingBoundariesSection => "working_boundaries_section",
|
||||
Self::AgentsMdSection => "agents_md_section",
|
||||
Self::ResidentMemorySummarySection => "resident_memory_summary_section",
|
||||
Self::ResidentKnowledgeSection => "resident_knowledge_section",
|
||||
Self::ResidentWorkflowsSection => "resident_workflows_section",
|
||||
Self::WorkerOrchestrationGuidanceSection => "worker_orchestration_guidance_section",
|
||||
Self::TicketEventCompanionNotice => "ticket_event_companion_notice",
|
||||
Self::SpawnWorkerToolDescription => "spawn_worker_tool_description",
|
||||
}
|
||||
}
|
||||
|
||||
/// All variants in declaration order. The associated `KEYS` slice
|
||||
/// mirrors this for const-eval coverage checks against
|
||||
/// `INTERNAL_KEYS` (generated by `build.rs`).
|
||||
pub const ALL: &'static [WorkerPrompt] = &[
|
||||
WorkerPrompt::CompactSystem,
|
||||
WorkerPrompt::MemoryExtractSystem,
|
||||
WorkerPrompt::MemoryConsolidationSystem,
|
||||
WorkerPrompt::NotifyWrapper,
|
||||
WorkerPrompt::InterruptToolResultSummary,
|
||||
WorkerPrompt::InterruptSystemNote,
|
||||
WorkerPrompt::WorkingBoundariesSection,
|
||||
WorkerPrompt::AgentsMdSection,
|
||||
WorkerPrompt::ResidentMemorySummarySection,
|
||||
WorkerPrompt::ResidentKnowledgeSection,
|
||||
WorkerPrompt::ResidentWorkflowsSection,
|
||||
WorkerPrompt::WorkerOrchestrationGuidanceSection,
|
||||
WorkerPrompt::TicketEventCompanionNotice,
|
||||
WorkerPrompt::SpawnWorkerToolDescription,
|
||||
];
|
||||
|
||||
pub const KEYS: &'static [&'static str] = &[
|
||||
"compact_system",
|
||||
"memory_extract_system",
|
||||
"memory_consolidation_system",
|
||||
"notify_wrapper",
|
||||
"interrupt_tool_result_summary",
|
||||
"interrupt_system_note",
|
||||
"working_boundaries_section",
|
||||
"agents_md_section",
|
||||
"resident_memory_summary_section",
|
||||
"resident_knowledge_section",
|
||||
"resident_workflows_section",
|
||||
"worker_orchestration_guidance_section",
|
||||
"ticket_event_companion_notice",
|
||||
"spawn_worker_tool_description",
|
||||
];
|
||||
}
|
||||
|
||||
// --- build-time bidirectional coverage check --------------------------------
|
||||
|
||||
const _: () = {
|
||||
// Every enum key must appear in the builtin TOML.
|
||||
let mut i = 0;
|
||||
while i < WorkerPrompt::KEYS.len() {
|
||||
if !const_slice_contains(INTERNAL_KEYS, WorkerPrompt::KEYS[i]) {
|
||||
panic!(
|
||||
"resources/prompts/internal.toml is missing a key declared by \
|
||||
WorkerPrompt — regenerate the TOML or remove the variant"
|
||||
);
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
// Every TOML key must correspond to an enum variant.
|
||||
let mut i = 0;
|
||||
while i < INTERNAL_KEYS.len() {
|
||||
if !const_slice_contains(WorkerPrompt::KEYS, INTERNAL_KEYS[i]) {
|
||||
panic!(
|
||||
"resources/prompts/internal.toml has a key not declared by \
|
||||
WorkerPrompt — add the variant or drop the key"
|
||||
);
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
};
|
||||
|
||||
const fn const_str_eq(a: &str, b: &str) -> bool {
|
||||
let a = a.as_bytes();
|
||||
let b = b.as_bytes();
|
||||
if a.len() != b.len() {
|
||||
return false;
|
||||
}
|
||||
let mut i = 0;
|
||||
while i < a.len() {
|
||||
if a[i] != b[i] {
|
||||
return false;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
const fn const_slice_contains(haystack: &[&str], needle: &str) -> bool {
|
||||
let mut i = 0;
|
||||
while i < haystack.len() {
|
||||
if const_str_eq(haystack[i], needle) {
|
||||
return true;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
// --- errors ----------------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum CatalogError {
|
||||
#[error("failed to read prompt pack {}: {source}", .path.display())]
|
||||
Io {
|
||||
path: PathBuf,
|
||||
#[source]
|
||||
source: std::io::Error,
|
||||
},
|
||||
#[error("failed to parse prompt pack {}: {source}", .path.display())]
|
||||
ParseToml {
|
||||
path: PathBuf,
|
||||
#[source]
|
||||
source: toml::de::Error,
|
||||
},
|
||||
#[error("failed to parse builtin prompt pack: {0}")]
|
||||
ParseBuiltin(#[source] toml::de::Error),
|
||||
#[error("failed to compile prompt template '{key}': {source}")]
|
||||
TemplateCompile {
|
||||
key: String,
|
||||
#[source]
|
||||
source: minijinja::Error,
|
||||
},
|
||||
#[error("failed to render prompt '{key}': {source}")]
|
||||
Render {
|
||||
key: String,
|
||||
#[source]
|
||||
source: minijinja::Error,
|
||||
},
|
||||
#[error("prompt key '{key}' is not registered in the catalog")]
|
||||
UnknownKey { key: String },
|
||||
}
|
||||
|
||||
// --- pack file shape -------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct PackFile {
|
||||
#[serde(default)]
|
||||
prompt: HashMap<String, String>,
|
||||
}
|
||||
|
||||
// --- catalog ---------------------------------------------------------------
|
||||
|
||||
/// Merged, compiled worker-prompt catalog.
|
||||
///
|
||||
/// Owns a `minijinja::Environment` with one template registered per
|
||||
/// [`WorkerPrompt`] key (after the 4-layer merge). Includes inside templates
|
||||
/// are resolved via a provided [`PromptLoader`], so values can pull from
|
||||
/// `$yoi` / `$user` / `$workspace`.
|
||||
pub struct PromptCatalog {
|
||||
env: Environment<'static>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for PromptCatalog {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("PromptCatalog").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl PromptCatalog {
|
||||
/// Builtin-only catalog. All `{% include %}` references must resolve
|
||||
/// through `$yoi` (user/workspace prefixes are unavailable).
|
||||
pub fn builtins_only() -> Result<Arc<Self>, CatalogError> {
|
||||
Self::load(&PromptLoader::builtins_only(), None)
|
||||
}
|
||||
|
||||
/// Load the catalog honouring the 4-layer overlay.
|
||||
///
|
||||
/// - Layer 1 (builtin): `INTERNAL_TOML` baked into the binary.
|
||||
/// - Layer 2 (user): `loader.user_pack_file()` if present.
|
||||
/// - Layer 3 (workspace): `loader.workspace_pack_file()` if present.
|
||||
/// - Layer 4 (manifest): `manifest_pack` as an absolute filesystem
|
||||
/// path (pre-resolved by profile/manifest resolution).
|
||||
pub fn load(
|
||||
loader: &PromptLoader,
|
||||
manifest_pack: Option<&Path>,
|
||||
) -> Result<Arc<Self>, CatalogError> {
|
||||
let mut merged = parse_builtin_pack()?;
|
||||
|
||||
if let Some(path) = loader.user_pack_file() {
|
||||
if path.is_file() {
|
||||
let pack = parse_pack_file(path)?;
|
||||
merge_into(&mut merged, pack, "user");
|
||||
}
|
||||
}
|
||||
if let Some(path) = loader.workspace_pack_file() {
|
||||
if path.is_file() {
|
||||
let pack = parse_pack_file(path)?;
|
||||
merge_into(&mut merged, pack, "workspace");
|
||||
}
|
||||
}
|
||||
if let Some(path) = manifest_pack {
|
||||
let pack = parse_pack_file(path)?;
|
||||
merge_into(&mut merged, pack, "manifest");
|
||||
}
|
||||
|
||||
build_catalog(merged, loader.clone()).map(Arc::new)
|
||||
}
|
||||
|
||||
/// Render a prompt by variant. `ctx` provides template variables; use
|
||||
/// [`Value::UNDEFINED`] (or a helper below) when the template takes
|
||||
/// no inputs.
|
||||
pub fn render(&self, prompt: WorkerPrompt, ctx: Value) -> Result<String, CatalogError> {
|
||||
let key = prompt.key();
|
||||
let tmpl = self
|
||||
.env
|
||||
.get_template(key)
|
||||
.map_err(|_| CatalogError::UnknownKey {
|
||||
key: key.to_string(),
|
||||
})?;
|
||||
tmpl.render(ctx).map_err(|source| CatalogError::Render {
|
||||
key: key.to_string(),
|
||||
source,
|
||||
})
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::CompactSystem` (no inputs).
|
||||
pub fn compact_system(&self) -> Result<String, CatalogError> {
|
||||
self.render(WorkerPrompt::CompactSystem, Value::UNDEFINED)
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::MemoryExtractSystem` with `{{ language }}`.
|
||||
pub fn memory_extract_system(&self, language: &str) -> Result<String, CatalogError> {
|
||||
self.render(
|
||||
WorkerPrompt::MemoryExtractSystem,
|
||||
single("language", language),
|
||||
)
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::MemoryConsolidationSystem` with `{{ language }}`.
|
||||
pub fn memory_consolidation_system(&self, language: &str) -> Result<String, CatalogError> {
|
||||
self.render(
|
||||
WorkerPrompt::MemoryConsolidationSystem,
|
||||
single("language", language),
|
||||
)
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::NotifyWrapper` with `{{ message }}`.
|
||||
pub fn notify_wrapper(&self, message: &str) -> Result<String, CatalogError> {
|
||||
self.render(WorkerPrompt::NotifyWrapper, single("message", message))
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::InterruptToolResultSummary` (no inputs).
|
||||
pub fn interrupt_tool_result_summary(&self) -> Result<String, CatalogError> {
|
||||
self.render(WorkerPrompt::InterruptToolResultSummary, Value::UNDEFINED)
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::InterruptSystemNote` (no inputs).
|
||||
pub fn interrupt_system_note(&self) -> Result<String, CatalogError> {
|
||||
self.render(WorkerPrompt::InterruptSystemNote, Value::UNDEFINED)
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::WorkingBoundariesSection` with `{{ scope_summary }}`.
|
||||
pub fn working_boundaries_section(&self, scope_summary: &str) -> Result<String, CatalogError> {
|
||||
self.render(
|
||||
WorkerPrompt::WorkingBoundariesSection,
|
||||
single("scope_summary", scope_summary),
|
||||
)
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::AgentsMdSection` with `{{ agents_md }}`.
|
||||
pub fn agents_md_section(&self, agents_md: &str) -> Result<String, CatalogError> {
|
||||
self.render(
|
||||
WorkerPrompt::AgentsMdSection,
|
||||
single("agents_md", agents_md),
|
||||
)
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::ResidentMemorySummarySection` with `{{ summary }}`.
|
||||
pub fn resident_memory_summary_section(&self, summary: &str) -> Result<String, CatalogError> {
|
||||
self.render(
|
||||
WorkerPrompt::ResidentMemorySummarySection,
|
||||
single("summary", summary),
|
||||
)
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::ResidentKnowledgeSection` with `{{ entries }}`
|
||||
/// (a pre-formatted list block authored by the caller).
|
||||
pub fn resident_knowledge_section(
|
||||
&self,
|
||||
entries: &str,
|
||||
knowledge_query_available: bool,
|
||||
memory_read_available: bool,
|
||||
) -> Result<String, CatalogError> {
|
||||
use std::collections::BTreeMap;
|
||||
let mut m: BTreeMap<&'static str, Value> = BTreeMap::new();
|
||||
m.insert("entries", Value::from(entries));
|
||||
m.insert(
|
||||
"knowledge_query_available",
|
||||
Value::from(knowledge_query_available),
|
||||
);
|
||||
m.insert("memory_read_available", Value::from(memory_read_available));
|
||||
self.render(WorkerPrompt::ResidentKnowledgeSection, Value::from(m))
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::ResidentWorkflowsSection` with `{{ entries }}`
|
||||
/// (a pre-formatted list block authored by the caller).
|
||||
pub fn resident_workflows_section(&self, entries: &str) -> Result<String, CatalogError> {
|
||||
self.render(
|
||||
WorkerPrompt::ResidentWorkflowsSection,
|
||||
single("entries", entries),
|
||||
)
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::WorkerOrchestrationGuidanceSection` (no inputs).
|
||||
pub fn worker_orchestration_guidance_section(&self) -> Result<String, CatalogError> {
|
||||
self.render(
|
||||
WorkerPrompt::WorkerOrchestrationGuidanceSection,
|
||||
Value::UNDEFINED,
|
||||
)
|
||||
}
|
||||
|
||||
/// Render `WorkerPrompt::SpawnWorkerToolDescription`.
|
||||
pub fn spawn_worker_tool_description(
|
||||
&self,
|
||||
available_profiles: &str,
|
||||
default_profile: &str,
|
||||
profile_diagnostic: &str,
|
||||
) -> Result<String, CatalogError> {
|
||||
use std::collections::BTreeMap;
|
||||
let mut m: BTreeMap<&'static str, Value> = BTreeMap::new();
|
||||
m.insert("available_profiles", Value::from(available_profiles));
|
||||
m.insert("default_profile", Value::from(default_profile));
|
||||
m.insert("profile_diagnostic", Value::from(profile_diagnostic));
|
||||
self.render(WorkerPrompt::SpawnWorkerToolDescription, Value::from(m))
|
||||
}
|
||||
}
|
||||
|
||||
fn single(key: &'static str, value: &str) -> Value {
|
||||
use std::collections::BTreeMap;
|
||||
let mut m: BTreeMap<&'static str, Value> = BTreeMap::new();
|
||||
m.insert(key, Value::from(value));
|
||||
Value::from(m)
|
||||
}
|
||||
|
||||
fn parse_builtin_pack() -> Result<HashMap<String, String>, CatalogError> {
|
||||
let parsed: PackFile = toml::from_str(INTERNAL_TOML).map_err(CatalogError::ParseBuiltin)?;
|
||||
Ok(parsed.prompt)
|
||||
}
|
||||
|
||||
fn parse_pack_file(path: &Path) -> Result<HashMap<String, String>, CatalogError> {
|
||||
let src = fs::read_to_string(path).map_err(|source| CatalogError::Io {
|
||||
path: path.to_path_buf(),
|
||||
source,
|
||||
})?;
|
||||
let parsed: PackFile = toml::from_str(&src).map_err(|source| CatalogError::ParseToml {
|
||||
path: path.to_path_buf(),
|
||||
source,
|
||||
})?;
|
||||
Ok(parsed.prompt)
|
||||
}
|
||||
|
||||
fn merge_into(
|
||||
base: &mut HashMap<String, String>,
|
||||
upper: HashMap<String, String>,
|
||||
origin: &'static str,
|
||||
) {
|
||||
for (k, v) in upper {
|
||||
if !WorkerPrompt::KEYS.iter().any(|declared| *declared == k) {
|
||||
warn!(
|
||||
origin = origin,
|
||||
key = %k,
|
||||
"unknown prompt pack key; ignoring"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
base.insert(k, v);
|
||||
}
|
||||
}
|
||||
|
||||
fn build_catalog(
|
||||
templates: HashMap<String, String>,
|
||||
loader: PromptLoader,
|
||||
) -> Result<PromptCatalog, CatalogError> {
|
||||
let mut env = Environment::new();
|
||||
env.set_undefined_behavior(UndefinedBehavior::Strict);
|
||||
|
||||
// Reuse the system-prompt-template resolver so `{% include
|
||||
// "$prefix/..." %}` inside a catalog value pulls from the same asset
|
||||
// namespaces.
|
||||
let loader_for_join = loader.clone();
|
||||
env.set_path_join_callback(move |name, parent| {
|
||||
let parent_ref = loader_for_join.parse_ref(parent, None).ok();
|
||||
match loader_for_join.parse_ref(name, parent_ref.as_ref()) {
|
||||
Ok(r) => r.to_qualified_string().into(),
|
||||
Err(_) => name.to_string().into(),
|
||||
}
|
||||
});
|
||||
|
||||
let loader_for_src = loader.clone();
|
||||
env.set_loader(move |name| {
|
||||
let reference = loader_for_src
|
||||
.parse_ref(name, None)
|
||||
.map_err(|e| minijinja::Error::new(ErrorKind::TemplateNotFound, e.to_string()))?;
|
||||
match loader_for_src.load(&reference) {
|
||||
Ok(src) => Ok(Some(src)),
|
||||
Err(e) => Err(minijinja::Error::new(
|
||||
ErrorKind::TemplateNotFound,
|
||||
e.to_string(),
|
||||
)),
|
||||
}
|
||||
});
|
||||
|
||||
for (k, v) in templates {
|
||||
env.add_template_owned(k.clone(), v)
|
||||
.map_err(|source| CatalogError::TemplateCompile {
|
||||
key: k.clone(),
|
||||
source,
|
||||
})?;
|
||||
}
|
||||
|
||||
Ok(PromptCatalog { env })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn loader_with_packs(
|
||||
user_dir: Option<PathBuf>,
|
||||
workspace_dir: Option<PathBuf>,
|
||||
user_pack: Option<PathBuf>,
|
||||
workspace_pack: Option<PathBuf>,
|
||||
) -> PromptLoader {
|
||||
PromptLoader::new(user_dir, workspace_dir).with_pack_files(user_pack, workspace_pack)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builtin_covers_every_variant() {
|
||||
let cat = PromptCatalog::builtins_only().unwrap();
|
||||
for p in WorkerPrompt::ALL {
|
||||
assert!(
|
||||
cat.env.get_template(p.key()).is_ok(),
|
||||
"builtin missing key: {}",
|
||||
p.key()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builtin_render_compact_system_includes_worker_instructions() {
|
||||
let cat = PromptCatalog::builtins_only().unwrap();
|
||||
let rendered = cat.compact_system().unwrap();
|
||||
assert!(rendered.contains("write_summary"));
|
||||
assert!(rendered.contains("mark_read_required"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internal_worker_prompts_do_not_include_default_memory_guidance() {
|
||||
let cat = PromptCatalog::builtins_only().unwrap();
|
||||
let compact = cat.compact_system().unwrap();
|
||||
let extract = cat.memory_extract_system("Japanese").unwrap();
|
||||
let consolidate = cat.memory_consolidation_system("Japanese").unwrap();
|
||||
for rendered in [compact, extract, consolidate] {
|
||||
assert!(!rendered.contains("### Memory and knowledge"));
|
||||
assert!(!rendered.contains("Do not query memory every turn"));
|
||||
assert!(!rendered.contains("Strong lookup triggers include"));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn memory_worker_prompts_include_language() {
|
||||
let cat = PromptCatalog::builtins_only().unwrap();
|
||||
let extract = cat.memory_extract_system("Japanese").unwrap();
|
||||
let consolidate = cat.memory_consolidation_system("Japanese").unwrap();
|
||||
assert!(extract.contains("`language`: `Japanese`"));
|
||||
assert!(consolidate.contains("`language`: `Japanese`"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notify_wrapper_interpolates_message() {
|
||||
let cat = PromptCatalog::builtins_only().unwrap();
|
||||
let out = cat.notify_wrapper("file changed").unwrap();
|
||||
assert!(out.contains("[Notification]"));
|
||||
assert!(out.contains("file changed"));
|
||||
assert!(out.contains("not a blocking request"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn working_boundaries_section_wraps_summary() {
|
||||
let cat = PromptCatalog::builtins_only().unwrap();
|
||||
let out = cat.working_boundaries_section("Readable: /a").unwrap();
|
||||
assert!(out.contains("## Working boundaries"));
|
||||
assert!(out.contains("Readable: /a"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agents_md_section_contains_marker() {
|
||||
let cat = PromptCatalog::builtins_only().unwrap();
|
||||
let out = cat.agents_md_section("PROJECT DOCS").unwrap();
|
||||
assert!(out.contains("## Project instructions (AGENTS.md)"));
|
||||
assert!(out.contains("PROJECT DOCS"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_pack_overrides_builtin() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let pack = tmp.path().join("prompts.toml");
|
||||
fs::write(
|
||||
&pack,
|
||||
r#"
|
||||
[prompt]
|
||||
interrupt_system_note = "[OVERRIDDEN]"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let loader = loader_with_packs(None, None, Some(pack), None);
|
||||
let cat = PromptCatalog::load(&loader, None).unwrap();
|
||||
assert_eq!(cat.interrupt_system_note().unwrap(), "[OVERRIDDEN]");
|
||||
// Other keys still come from the builtin.
|
||||
assert!(cat.notify_wrapper("x").unwrap().contains("[Notification]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_pack_wins_over_user_pack() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let user = tmp.path().join("user.toml");
|
||||
let ws = tmp.path().join("ws.toml");
|
||||
fs::write(
|
||||
&user,
|
||||
r#"
|
||||
[prompt]
|
||||
interrupt_system_note = "[USER]"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
fs::write(
|
||||
&ws,
|
||||
r#"
|
||||
[prompt]
|
||||
interrupt_system_note = "[WS]"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let loader = loader_with_packs(None, None, Some(user), Some(ws));
|
||||
let cat = PromptCatalog::load(&loader, None).unwrap();
|
||||
assert_eq!(cat.interrupt_system_note().unwrap(), "[WS]");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manifest_pack_wins_over_workspace_pack() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let ws = tmp.path().join("ws.toml");
|
||||
let mf = tmp.path().join("mf.toml");
|
||||
fs::write(
|
||||
&ws,
|
||||
r#"
|
||||
[prompt]
|
||||
interrupt_system_note = "[WS]"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
fs::write(
|
||||
&mf,
|
||||
r#"
|
||||
[prompt]
|
||||
interrupt_system_note = "[MF]"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let loader = loader_with_packs(None, None, None, Some(ws));
|
||||
let cat = PromptCatalog::load(&loader, Some(mf.as_path())).unwrap();
|
||||
assert_eq!(cat.interrupt_system_note().unwrap(), "[MF]");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_key_in_runtime_pack_is_ignored_with_warning() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let pack = tmp.path().join("p.toml");
|
||||
fs::write(
|
||||
&pack,
|
||||
r#"
|
||||
[prompt]
|
||||
interrupt_system_note = "[OK]"
|
||||
future_injection_point = "tolerated"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let loader = loader_with_packs(None, None, Some(pack), None);
|
||||
// Loads without error; the unknown key is dropped silently at
|
||||
// runtime (log warning is emitted via tracing).
|
||||
let cat = PromptCatalog::load(&loader, None).unwrap();
|
||||
assert_eq!(cat.interrupt_system_note().unwrap(), "[OK]");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manifest_pack_reads_from_absolute_path() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let pack = tmp.path().join("mine.toml");
|
||||
fs::write(
|
||||
&pack,
|
||||
r#"
|
||||
[prompt]
|
||||
interrupt_system_note = "[FROM-MANIFEST-PACK]"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let cat = PromptCatalog::load(&loader, Some(pack.as_path())).unwrap();
|
||||
assert_eq!(cat.interrupt_system_note().unwrap(), "[FROM-MANIFEST-PACK]");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn value_can_pull_long_text_via_include() {
|
||||
// A runtime pack that overrides `compact_system` with an
|
||||
// `{% include %}` into the same `$yoi` namespace — exercises
|
||||
// the template resolver path through all four layers.
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let pack = tmp.path().join("p.toml");
|
||||
fs::write(
|
||||
&pack,
|
||||
r#"
|
||||
[prompt]
|
||||
compact_system = "PREFIX\n{% include \"$yoi/internal/compact_system\" %}"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let loader = loader_with_packs(None, None, Some(pack), None);
|
||||
let cat = PromptCatalog::load(&loader, None).unwrap();
|
||||
let rendered = cat.compact_system().unwrap();
|
||||
assert!(rendered.starts_with("PREFIX\n"));
|
||||
assert!(rendered.contains("write_summary"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_orchestration_guidance_section_renders_resource_body() {
|
||||
let cat = PromptCatalog::builtins_only().unwrap();
|
||||
let rendered = cat.worker_orchestration_guidance_section().unwrap();
|
||||
assert!(rendered.contains("## Worker orchestration"));
|
||||
assert!(rendered.contains("spawned Worker notifications are background signals"));
|
||||
assert!(rendered.contains("does not need to keep a turn open"));
|
||||
assert!(rendered.contains("Do not use `sleep` or polling loops"));
|
||||
assert!(rendered.contains("worktree state, diff, and test results"));
|
||||
assert!(rendered.contains("not scheduler or auto-maintain authorization"));
|
||||
assert!(rendered.contains("bypass user/workflow authorization"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn spawn_worker_tool_description_renders_profile_block() {
|
||||
let cat = PromptCatalog::builtins_only().unwrap();
|
||||
let rendered = cat
|
||||
.spawn_worker_tool_description(
|
||||
"- `project:coder` — Coder\n- `project:reviewer` — Reviewer",
|
||||
"project:coder",
|
||||
"",
|
||||
)
|
||||
.unwrap();
|
||||
assert!(rendered.contains("Profile selection"));
|
||||
assert!(rendered.contains("Default profile: project:coder"));
|
||||
assert!(rendered.contains("`project:reviewer`"));
|
||||
assert!(rendered.contains("Special selector: inherit"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,425 @@
|
||||
//! Prefix-addressed prompt asset loader used by [`crate::SystemPromptTemplate`].
|
||||
//!
|
||||
//! Three prefixes address three physical libraries:
|
||||
//!
|
||||
//! | prefix | location |
|
||||
//! |--------------|---------------------------------------------------------|
|
||||
//! | `$yoi` | builtin, baked into the binary via `include_dir!` |
|
||||
//! | `$user` | `<config_dir>/prompts/` (resolved by `manifest::paths`) |
|
||||
//! | `$workspace` | `<project>/.yoi/prompts/` |
|
||||
//!
|
||||
//! A reference is `$<prefix>/<path>` where `<path>` is a `/`-separated
|
||||
//! name without the `.md` extension (e.g. `$yoi/common/header`).
|
||||
//! Unqualified names (no `$prefix/` at the front) are resolved relative
|
||||
//! to an optional current reference — typically the file that issued
|
||||
//! the `{% include %}` — so a prompt library can be authored as a
|
||||
//! self-contained directory.
|
||||
//!
|
||||
//! Missing files produce a [`LoaderError::NotFound`]; there is no
|
||||
//! fallthrough between layers.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use include_dir::{Dir, include_dir};
|
||||
use thiserror::Error;
|
||||
|
||||
static BUILTIN_PROMPTS: Dir<'static> = include_dir!("$CARGO_MANIFEST_DIR/../../resources/prompts");
|
||||
|
||||
const PREFIX_YOI: &str = "$yoi";
|
||||
const PREFIX_USER: &str = "$user";
|
||||
const PREFIX_WORKSPACE: &str = "$workspace";
|
||||
|
||||
/// Prefix-resolved reference to a prompt asset. Produced by
|
||||
/// [`PromptLoader::parse_ref`] from a user-supplied string such as
|
||||
/// `"$yoi/default"`.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PromptRef {
|
||||
prefix: Prefix,
|
||||
/// Relative path under the prefix root, without the `.md` extension.
|
||||
/// `/`-separated, never empty, never starts with `/`.
|
||||
path: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum Prefix {
|
||||
Yoi,
|
||||
User,
|
||||
Workspace,
|
||||
}
|
||||
|
||||
impl Prefix {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Yoi => PREFIX_YOI,
|
||||
Self::User => PREFIX_USER,
|
||||
Self::Workspace => PREFIX_WORKSPACE,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PromptRef {
|
||||
/// Produce a canonical `$prefix/path` string.
|
||||
pub fn to_qualified_string(&self) -> String {
|
||||
format!("{}/{}", self.prefix.as_str(), self.path)
|
||||
}
|
||||
|
||||
/// Directory portion (leading prefix segments minus the file name),
|
||||
/// joined with `/`. Returns an empty string when the ref points at
|
||||
/// a file directly under the prefix root.
|
||||
fn dir(&self) -> &str {
|
||||
match self.path.rsplit_once('/') {
|
||||
Some((dir, _)) => dir,
|
||||
None => "",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Errors produced when resolving a [`PromptRef`].
|
||||
#[derive(Debug, Error)]
|
||||
pub enum LoaderError {
|
||||
#[error("invalid prompt reference '{raw}': {reason}")]
|
||||
InvalidRef { raw: String, reason: String },
|
||||
#[error("unknown prompt prefix '{prefix}' in reference '{raw}'")]
|
||||
UnknownPrefix { raw: String, prefix: String },
|
||||
#[error(
|
||||
"unqualified prompt reference '{raw}' requires a current prefix \
|
||||
(include it from inside another template, or use an explicit \
|
||||
$prefix/path form)"
|
||||
)]
|
||||
UnqualifiedWithoutCurrent { raw: String },
|
||||
#[error("prompt prefix '{prefix}' is not configured for this loader")]
|
||||
PrefixNotConfigured { prefix: &'static str },
|
||||
#[error("prompt asset not found: '{}'", .reference.to_qualified_string())]
|
||||
NotFound { reference: PromptRef },
|
||||
#[error("failed to read prompt asset '{}': {source}", .reference.to_qualified_string())]
|
||||
Io {
|
||||
reference: PromptRef,
|
||||
#[source]
|
||||
source: std::io::Error,
|
||||
},
|
||||
}
|
||||
|
||||
/// Loader that resolves [`PromptRef`]s against the configured prompt
|
||||
/// libraries. Cheap to clone.
|
||||
///
|
||||
/// Also carries the auto-discovered `prompts.toml` pack file paths so
|
||||
/// [`crate::prompt::catalog::PromptCatalog`] can read the same user/workspace
|
||||
/// layers without a separate plumbing channel. These fields do not
|
||||
/// affect `$prefix` asset resolution — they are purely metadata
|
||||
/// consulted by the catalog loader.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptLoader {
|
||||
user_dir: Option<PathBuf>,
|
||||
workspace_dir: Option<PathBuf>,
|
||||
user_pack_file: Option<PathBuf>,
|
||||
workspace_pack_file: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl PromptLoader {
|
||||
/// Loader with only the builtin `$yoi` library available.
|
||||
/// `$user` / `$workspace` references fail with
|
||||
/// [`LoaderError::PrefixNotConfigured`].
|
||||
pub fn builtins_only() -> Self {
|
||||
Self {
|
||||
user_dir: None,
|
||||
workspace_dir: None,
|
||||
user_pack_file: None,
|
||||
workspace_pack_file: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Loader with optional user and workspace prompt directories.
|
||||
pub fn new(user_dir: Option<PathBuf>, workspace_dir: Option<PathBuf>) -> Self {
|
||||
Self {
|
||||
user_dir,
|
||||
workspace_dir,
|
||||
user_pack_file: None,
|
||||
workspace_pack_file: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Override pack file paths supplied by the caller's profile/manifest
|
||||
/// resolution context.
|
||||
pub fn with_pack_files(
|
||||
mut self,
|
||||
user_pack_file: Option<PathBuf>,
|
||||
workspace_pack_file: Option<PathBuf>,
|
||||
) -> Self {
|
||||
self.user_pack_file = user_pack_file;
|
||||
self.workspace_pack_file = workspace_pack_file;
|
||||
self
|
||||
}
|
||||
|
||||
/// Root of the `$user` prompt library, if configured.
|
||||
pub fn user_dir(&self) -> Option<&Path> {
|
||||
self.user_dir.as_deref()
|
||||
}
|
||||
|
||||
/// Root of the `$workspace` prompt library, if configured.
|
||||
pub fn workspace_dir(&self) -> Option<&Path> {
|
||||
self.workspace_dir.as_deref()
|
||||
}
|
||||
|
||||
/// Auto-discovered path to the user-layer `prompts.toml` pack, if any.
|
||||
pub fn user_pack_file(&self) -> Option<&Path> {
|
||||
self.user_pack_file.as_deref()
|
||||
}
|
||||
|
||||
/// Auto-discovered path to the workspace-layer `prompts.toml` pack, if any.
|
||||
pub fn workspace_pack_file(&self) -> Option<&Path> {
|
||||
self.workspace_pack_file.as_deref()
|
||||
}
|
||||
|
||||
/// Parse a string reference into a [`PromptRef`]. Unqualified
|
||||
/// references (no leading `$prefix/`) are resolved against
|
||||
/// `current`: the prefix is inherited, and the path is joined to
|
||||
/// the current ref's directory.
|
||||
pub fn parse_ref(
|
||||
&self,
|
||||
raw: &str,
|
||||
current: Option<&PromptRef>,
|
||||
) -> Result<PromptRef, LoaderError> {
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Err(LoaderError::InvalidRef {
|
||||
raw: raw.to_string(),
|
||||
reason: "reference must not be empty".into(),
|
||||
});
|
||||
}
|
||||
if let Some(prefix) = trimmed.strip_prefix('$') {
|
||||
let (prefix_name, rest) =
|
||||
prefix
|
||||
.split_once('/')
|
||||
.ok_or_else(|| LoaderError::InvalidRef {
|
||||
raw: raw.to_string(),
|
||||
reason: "prefix must be followed by '/'".into(),
|
||||
})?;
|
||||
let prefix = parse_prefix(raw, prefix_name)?;
|
||||
let path = normalize_path(raw, rest)?;
|
||||
Ok(PromptRef { prefix, path })
|
||||
} else {
|
||||
let Some(current) = current else {
|
||||
return Err(LoaderError::UnqualifiedWithoutCurrent {
|
||||
raw: raw.to_string(),
|
||||
});
|
||||
};
|
||||
let dir = current.dir();
|
||||
let joined = if dir.is_empty() {
|
||||
trimmed.to_string()
|
||||
} else {
|
||||
format!("{dir}/{trimmed}")
|
||||
};
|
||||
let path = normalize_path(raw, &joined)?;
|
||||
Ok(PromptRef {
|
||||
prefix: current.prefix,
|
||||
path,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve a [`PromptRef`] to its raw template source. Hard-errors
|
||||
/// when the prefix is not configured or the file does not exist.
|
||||
pub fn load(&self, reference: &PromptRef) -> Result<String, LoaderError> {
|
||||
match reference.prefix {
|
||||
Prefix::Yoi => load_from_include_dir(&BUILTIN_PROMPTS, reference),
|
||||
Prefix::User => match self.user_dir.as_deref() {
|
||||
Some(dir) => load_from_dir(dir, reference),
|
||||
None => Err(LoaderError::PrefixNotConfigured {
|
||||
prefix: PREFIX_USER,
|
||||
}),
|
||||
},
|
||||
Prefix::Workspace => match self.workspace_dir.as_deref() {
|
||||
Some(dir) => load_from_dir(dir, reference),
|
||||
None => Err(LoaderError::PrefixNotConfigured {
|
||||
prefix: PREFIX_WORKSPACE,
|
||||
}),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse `raw` against `current`, then load the resulting ref.
|
||||
/// Convenience wrapper for the minijinja loader hook.
|
||||
pub fn resolve(
|
||||
&self,
|
||||
raw: &str,
|
||||
current: Option<&PromptRef>,
|
||||
) -> Result<(PromptRef, String), LoaderError> {
|
||||
let reference = self.parse_ref(raw, current)?;
|
||||
let source = self.load(&reference)?;
|
||||
Ok((reference, source))
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_prefix(raw: &str, prefix_name: &str) -> Result<Prefix, LoaderError> {
|
||||
match prefix_name {
|
||||
"yoi" => Ok(Prefix::Yoi),
|
||||
"user" => Ok(Prefix::User),
|
||||
"workspace" => Ok(Prefix::Workspace),
|
||||
_ => Err(LoaderError::UnknownPrefix {
|
||||
raw: raw.to_string(),
|
||||
prefix: format!("${prefix_name}"),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_path(raw: &str, rest: &str) -> Result<String, LoaderError> {
|
||||
let cleaned = rest.trim_matches('/').trim();
|
||||
if cleaned.is_empty() {
|
||||
return Err(LoaderError::InvalidRef {
|
||||
raw: raw.to_string(),
|
||||
reason: "path component must not be empty".into(),
|
||||
});
|
||||
}
|
||||
if cleaned.split('/').any(|seg| seg == "." || seg == "..") {
|
||||
return Err(LoaderError::InvalidRef {
|
||||
raw: raw.to_string(),
|
||||
reason: "path must not contain '.' or '..' segments".into(),
|
||||
});
|
||||
}
|
||||
Ok(cleaned.to_string())
|
||||
}
|
||||
|
||||
fn load_from_dir(dir: &Path, reference: &PromptRef) -> Result<String, LoaderError> {
|
||||
let path = dir.join(format!("{}.md", reference.path));
|
||||
match std::fs::read_to_string(&path) {
|
||||
Ok(s) => Ok(s),
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Err(LoaderError::NotFound {
|
||||
reference: reference.clone(),
|
||||
}),
|
||||
Err(source) => Err(LoaderError::Io {
|
||||
reference: reference.clone(),
|
||||
source,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn load_from_include_dir(dir: &Dir<'static>, reference: &PromptRef) -> Result<String, LoaderError> {
|
||||
let path = format!("{}.md", reference.path);
|
||||
dir.get_file(&path)
|
||||
.and_then(|f| f.contents_utf8())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| LoaderError::NotFound {
|
||||
reference: reference.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn builtin_default_resolves() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let (r, source) = loader.resolve("$yoi/default", None).unwrap();
|
||||
assert_eq!(r.to_qualified_string(), "$yoi/default");
|
||||
assert!(!source.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builtin_ticket_role_instructions_resolve() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
for role in ["intake", "orchestrator", "coder", "reviewer"] {
|
||||
let (reference, source) = loader.resolve(&format!("$yoi/role/{role}"), None).unwrap();
|
||||
assert_eq!(reference.to_qualified_string(), format!("$yoi/role/{role}"));
|
||||
assert!(source.contains("first committed user message"));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builtin_subdirectory_lookup() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let (_, source) = loader.resolve("$yoi/common/tool-usage", None).unwrap();
|
||||
assert!(source.contains("tool"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_prefix_resolves() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let user_dir = tmp.path().to_path_buf();
|
||||
std::fs::write(user_dir.join("my.md"), "user-body").unwrap();
|
||||
let loader = PromptLoader::new(Some(user_dir), None);
|
||||
let (_, source) = loader.resolve("$user/my", None).unwrap();
|
||||
assert_eq!(source, "user-body");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_prefix_resolves() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let ws_dir = tmp.path().to_path_buf();
|
||||
std::fs::write(ws_dir.join("custom.md"), "ws-body").unwrap();
|
||||
let loader = PromptLoader::new(None, Some(ws_dir));
|
||||
let (_, source) = loader.resolve("$workspace/custom", None).unwrap();
|
||||
assert_eq!(source, "ws-body");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_file_is_hard_error() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let err = loader.resolve("$yoi/definitely-missing", None).unwrap_err();
|
||||
assert!(matches!(err, LoaderError::NotFound { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_prefix_not_configured_errors() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let err = loader.resolve("$user/my", None).unwrap_err();
|
||||
assert!(matches!(
|
||||
err,
|
||||
LoaderError::PrefixNotConfigured { prefix: "$user" }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_prefix_errors() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let err = loader.resolve("$bogus/x", None).unwrap_err();
|
||||
assert!(matches!(err, LoaderError::UnknownPrefix { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unqualified_ref_without_current_errors() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let err = loader.resolve("default", None).unwrap_err();
|
||||
assert!(matches!(err, LoaderError::UnqualifiedWithoutCurrent { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unqualified_ref_resolves_relative_to_current() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let current = loader.parse_ref("$yoi/common/tool-usage", None).unwrap();
|
||||
// Sibling lookup under the same prefix and directory.
|
||||
let sibling = loader.parse_ref("workspace", Some(¤t)).unwrap();
|
||||
assert_eq!(sibling.to_qualified_string(), "$yoi/common/workspace");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unqualified_ref_from_root_file_has_empty_dir() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let current = loader.parse_ref("$yoi/default", None).unwrap();
|
||||
let sibling = loader.parse_ref("other", Some(¤t)).unwrap();
|
||||
assert_eq!(sibling.to_qualified_string(), "$yoi/other");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_prefix_overrides_current() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let user_dir = tmp.path().to_path_buf();
|
||||
std::fs::write(user_dir.join("custom.md"), "user-body").unwrap();
|
||||
let loader = PromptLoader::new(Some(user_dir), None);
|
||||
|
||||
let current = loader.parse_ref("$yoi/default", None).unwrap();
|
||||
// Even with an $yoi-rooted current, an explicit $user
|
||||
// prefix must win.
|
||||
let (reference, source) = loader.resolve("$user/custom", Some(¤t)).unwrap();
|
||||
assert_eq!(reference.to_qualified_string(), "$user/custom");
|
||||
assert_eq!(source, "user-body");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn traversal_segments_rejected() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let err = loader.resolve("$yoi/../etc/passwd", None).unwrap_err();
|
||||
assert!(matches!(err, LoaderError::InvalidRef { .. }));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
pub(crate) mod agents_md;
|
||||
pub(crate) mod catalog;
|
||||
pub(crate) mod loader;
|
||||
pub(crate) mod system;
|
||||
@@ -0,0 +1,991 @@
|
||||
//! System prompt template machinery for the Worker layer.
|
||||
//!
|
||||
//! Manifests describe the system prompt body as a reference to a
|
||||
//! prompt asset (`worker.instruction`, see [`manifest::EngineManifest`]).
|
||||
//! [`SystemPromptTemplate`] resolves that reference through a
|
||||
//! [`PromptLoader`], parses the source as a minijinja template, and
|
||||
//! eagerly syntax-checks it at Worker construction. The final system
|
||||
//! prompt is materialised exactly once just before the first LLM turn:
|
||||
//! the rendered body is appended with a fixed trailing section carrying
|
||||
//! the Worker's `Scope` summary, (if present) the project's `AGENTS.md`
|
||||
//! contents, resident memory sections, and conditional Worker-orchestration
|
||||
//! guidance, then the whole string is handed to the Engine via
|
||||
//! `set_system_prompt`. Subsequent turns and compactions reuse that
|
||||
//! materialised string verbatim.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::{DateTime, SecondsFormat, Utc};
|
||||
use manifest::Scope;
|
||||
use memory::ResidentKnowledgeEntry;
|
||||
use minijinja::value::Value;
|
||||
use minijinja::{Environment, ErrorKind, UndefinedBehavior};
|
||||
use thiserror::Error;
|
||||
use workflow_crate::ResidentWorkflowEntry;
|
||||
|
||||
use crate::prompt::catalog::{CatalogError, PromptCatalog};
|
||||
use crate::prompt::loader::{LoaderError, PromptLoader, PromptRef};
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum SystemPromptError {
|
||||
#[error("failed to resolve instruction reference: {0}")]
|
||||
LoaderResolve(#[source] LoaderError),
|
||||
#[error("system prompt template parse error: {0}")]
|
||||
Parse(String),
|
||||
#[error("system prompt template render error: {0}")]
|
||||
Render(String),
|
||||
#[error("failed to render trailing section template: {0}")]
|
||||
Catalog(#[from] CatalogError),
|
||||
}
|
||||
|
||||
/// Parsed instruction template bound to a prompt loader.
|
||||
///
|
||||
/// Holds a minijinja Environment pre-populated with the instruction
|
||||
/// template registered under its fully-qualified name (`$prefix/path`).
|
||||
/// Includes are resolved via the loader using a path-join callback that
|
||||
/// tracks the including template's prefix and directory, so
|
||||
/// `{% include "sibling" %}` fragments work as expected.
|
||||
#[derive(Clone)]
|
||||
pub struct SystemPromptTemplate {
|
||||
env: Arc<Environment<'static>>,
|
||||
instruction_name: String,
|
||||
}
|
||||
|
||||
impl SystemPromptTemplate {
|
||||
/// Parse the instruction asset referenced by `instruction_ref`
|
||||
/// using the supplied [`PromptLoader`]. The reference is resolved
|
||||
/// at parse time so syntax errors surface immediately.
|
||||
pub fn parse(instruction_ref: &str, loader: PromptLoader) -> Result<Self, SystemPromptError> {
|
||||
let root_ref = loader
|
||||
.parse_ref(instruction_ref, None)
|
||||
.map_err(SystemPromptError::LoaderResolve)?;
|
||||
let source = loader
|
||||
.load(&root_ref)
|
||||
.map_err(SystemPromptError::LoaderResolve)?;
|
||||
let root_name = root_ref.to_qualified_string();
|
||||
|
||||
let mut env = Environment::new();
|
||||
env.set_undefined_behavior(UndefinedBehavior::Strict);
|
||||
|
||||
// Path-join callback: compute the target template name when a
|
||||
// template includes another by a possibly-unqualified string.
|
||||
// The joined name is then looked up via `set_loader` below.
|
||||
let loader_for_join = loader.clone();
|
||||
env.set_path_join_callback(move |name, parent| {
|
||||
let parent_ref = loader_for_join.parse_ref(parent, None).ok();
|
||||
match loader_for_join.parse_ref(name, parent_ref.as_ref()) {
|
||||
Ok(r) => r.to_qualified_string().into(),
|
||||
// Propagate the raw name on error so set_loader surfaces
|
||||
// a proper TemplateNotFound/LoaderError to the caller.
|
||||
Err(_) => name.to_string().into(),
|
||||
}
|
||||
});
|
||||
|
||||
let loader_for_src = loader.clone();
|
||||
env.set_loader(move |name| {
|
||||
let reference = loader_for_src
|
||||
.parse_ref(name, None)
|
||||
.map_err(|e| minijinja::Error::new(ErrorKind::TemplateNotFound, e.to_string()))?;
|
||||
match loader_for_src.load(&reference) {
|
||||
Ok(source) => Ok(Some(source)),
|
||||
Err(e) => Err(minijinja::Error::new(
|
||||
ErrorKind::TemplateNotFound,
|
||||
e.to_string(),
|
||||
)),
|
||||
}
|
||||
});
|
||||
|
||||
env.add_template_owned(root_name.clone(), source)
|
||||
.map_err(|e| SystemPromptError::Parse(e.to_string()))?;
|
||||
|
||||
Ok(Self {
|
||||
env: Arc::new(env),
|
||||
instruction_name: root_name,
|
||||
})
|
||||
}
|
||||
|
||||
/// Render the instruction body and append the fixed trailing
|
||||
/// section (scope summary + optional AGENTS.md). The trailing
|
||||
/// section is assembled in Rust so that authored templates cannot
|
||||
/// accidentally omit the scope boundary or the project instructions.
|
||||
pub fn render(&self, ctx: &SystemPromptContext<'_>) -> Result<String, SystemPromptError> {
|
||||
let tmpl = self
|
||||
.env
|
||||
.get_template(&self.instruction_name)
|
||||
.map_err(|e| SystemPromptError::Render(e.to_string()))?;
|
||||
let body = tmpl
|
||||
.render(ctx.to_minijinja_value())
|
||||
.map_err(|e| SystemPromptError::Render(e.to_string()))?;
|
||||
append_trailing_section(
|
||||
&body,
|
||||
ctx.prompts,
|
||||
ctx.scope,
|
||||
ctx.agents_md.as_deref(),
|
||||
ctx.resident_summary,
|
||||
ctx.resident_knowledge,
|
||||
ctx.resident_workflows,
|
||||
ToolCapabilities::from_tool_names(&ctx.tool_names),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for SystemPromptTemplate {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("SystemPromptTemplate")
|
||||
.field("instruction", &self.instruction_name)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
/// Inputs available to an instruction template at materialisation time.
|
||||
///
|
||||
/// Scope summary and AGENTS.md are deliberately **not** exposed to the
|
||||
/// template — they live in the Rust-owned trailing section so user
|
||||
/// templates cannot drop them on the floor.
|
||||
pub struct SystemPromptContext<'a> {
|
||||
pub now: DateTime<Utc>,
|
||||
pub cwd: &'a Path,
|
||||
/// Language policy exposed to instruction templates as `{{ language }}`.
|
||||
pub language: &'a str,
|
||||
pub scope: &'a Scope,
|
||||
pub tool_names: Vec<String>,
|
||||
/// Project-level instructions read from the nearest `AGENTS.md`.
|
||||
/// Not visible from the template; consumed by the trailing-section
|
||||
/// formatter in [`SystemPromptTemplate::render`].
|
||||
pub agents_md: Option<String>,
|
||||
/// The body of `<workspace>/.yoi/memory/summary.md`, with
|
||||
/// frontmatter stripped. `None` disables the resident summary section;
|
||||
/// empty strings are ignored by the trailing-section formatter.
|
||||
pub resident_summary: Option<&'a str>,
|
||||
/// Resident-injection candidates from `<workspace>/knowledge/*` whose
|
||||
/// frontmatter has `model_invokation: true`. `None` disables the
|
||||
/// section entirely (memory disabled, or a consolidation worker that opts
|
||||
/// out); `Some(&[])` also yields no section.
|
||||
pub resident_knowledge: Option<&'a [ResidentKnowledgeEntry]>,
|
||||
/// Resident workflow descriptions from `<workspace>/.yoi/workflow/*`
|
||||
/// whose frontmatter has `model_invokation: true`. `None` disables the
|
||||
/// section; consolidation workers opt out together with resident Knowledge.
|
||||
pub resident_workflows: Option<&'a [ResidentWorkflowEntry]>,
|
||||
/// Catalog used to render the fixed trailing section headers.
|
||||
/// Passed by reference so callers do not give up ownership across
|
||||
/// the short-lived render borrow.
|
||||
pub prompts: &'a PromptCatalog,
|
||||
}
|
||||
|
||||
impl<'a> SystemPromptContext<'a> {
|
||||
fn to_minijinja_value(&self) -> Value {
|
||||
let mut root: BTreeMap<String, Value> = BTreeMap::new();
|
||||
root.insert(
|
||||
"date".into(),
|
||||
Value::from(self.now.format("%Y-%m-%d").to_string()),
|
||||
);
|
||||
root.insert(
|
||||
"time".into(),
|
||||
Value::from(self.now.format("%H:%M:%S").to_string()),
|
||||
);
|
||||
root.insert(
|
||||
"datetime".into(),
|
||||
Value::from(self.now.to_rfc3339_opts(SecondsFormat::Secs, true)),
|
||||
);
|
||||
root.insert("cwd".into(), Value::from(self.cwd.display().to_string()));
|
||||
root.insert("language".into(), Value::from(self.language));
|
||||
root.insert(
|
||||
"tools".into(),
|
||||
Value::from(
|
||||
self.tool_names
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(Value::from)
|
||||
.collect::<Vec<_>>(),
|
||||
),
|
||||
);
|
||||
root.insert(
|
||||
"tool_capabilities".into(),
|
||||
ToolCapabilities::from_tool_names(&self.tool_names).to_minijinja_value(),
|
||||
);
|
||||
Value::from(root)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
|
||||
struct ToolCapabilities {
|
||||
memory_query: bool,
|
||||
knowledge_query: bool,
|
||||
memory_read: bool,
|
||||
memory_write: bool,
|
||||
memory_edit: bool,
|
||||
memory_delete: bool,
|
||||
worker_spawn: bool,
|
||||
worker_send: bool,
|
||||
worker_read_output: bool,
|
||||
worker_stop: bool,
|
||||
worker_list: bool,
|
||||
worker_restore: bool,
|
||||
}
|
||||
|
||||
impl ToolCapabilities {
|
||||
fn from_tool_names(names: &[String]) -> Self {
|
||||
let mut capabilities = Self::default();
|
||||
for name in names {
|
||||
match name.as_str() {
|
||||
"MemoryQuery" => capabilities.memory_query = true,
|
||||
"KnowledgeQuery" => capabilities.knowledge_query = true,
|
||||
"MemoryRead" => capabilities.memory_read = true,
|
||||
"MemoryWrite" => capabilities.memory_write = true,
|
||||
"MemoryEdit" => capabilities.memory_edit = true,
|
||||
"MemoryDelete" => capabilities.memory_delete = true,
|
||||
"SpawnWorker" => capabilities.worker_spawn = true,
|
||||
"SendToWorker" => capabilities.worker_send = true,
|
||||
"ReadWorkerOutput" => capabilities.worker_read_output = true,
|
||||
"StopWorker" => capabilities.worker_stop = true,
|
||||
"ListWorkers" => capabilities.worker_list = true,
|
||||
"RestoreWorker" => capabilities.worker_restore = true,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
capabilities
|
||||
}
|
||||
|
||||
fn memory_records(self) -> bool {
|
||||
self.memory_query
|
||||
|| self.memory_read
|
||||
|| self.memory_write
|
||||
|| self.memory_edit
|
||||
|| self.memory_delete
|
||||
}
|
||||
|
||||
fn memory_any(self) -> bool {
|
||||
self.memory_records() || self.knowledge_query
|
||||
}
|
||||
|
||||
fn memory_mutation(self) -> bool {
|
||||
self.memory_write || self.memory_edit || self.memory_delete
|
||||
}
|
||||
|
||||
fn worker_management(self) -> bool {
|
||||
self.worker_spawn
|
||||
|| self.worker_send
|
||||
|| self.worker_read_output
|
||||
|| self.worker_stop
|
||||
|| self.worker_list
|
||||
|| self.worker_restore
|
||||
}
|
||||
|
||||
fn to_minijinja_value(self) -> Value {
|
||||
let mut map: BTreeMap<&'static str, Value> = BTreeMap::new();
|
||||
map.insert("memory_any", Value::from(self.memory_any()));
|
||||
map.insert("memory_records", Value::from(self.memory_records()));
|
||||
map.insert("memory_query", Value::from(self.memory_query));
|
||||
map.insert("knowledge_query", Value::from(self.knowledge_query));
|
||||
map.insert("memory_read", Value::from(self.memory_read));
|
||||
map.insert("memory_write", Value::from(self.memory_write));
|
||||
map.insert("memory_edit", Value::from(self.memory_edit));
|
||||
map.insert("memory_delete", Value::from(self.memory_delete));
|
||||
map.insert("memory_mutation", Value::from(self.memory_mutation()));
|
||||
map.insert("worker_management", Value::from(self.worker_management()));
|
||||
Value::from(map)
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the final system prompt by appending the fixed trailing
|
||||
/// section to `body`. The Rust side owns the layout (blank-line
|
||||
/// separators, trailing-whitespace trim); each section's header + body
|
||||
/// comes from the prompt catalog (`WorkerPrompt::WorkingBoundariesSection`
|
||||
/// / `WorkerPrompt::AgentsMdSection`) so that wording can be overridden
|
||||
/// per-pack without touching this function.
|
||||
fn append_trailing_section(
|
||||
body: &str,
|
||||
prompts: &PromptCatalog,
|
||||
scope: &Scope,
|
||||
agents_md: Option<&str>,
|
||||
resident_summary: Option<&str>,
|
||||
resident_knowledge: Option<&[ResidentKnowledgeEntry]>,
|
||||
resident_workflows: Option<&[ResidentWorkflowEntry]>,
|
||||
tool_capabilities: ToolCapabilities,
|
||||
) -> Result<String, SystemPromptError> {
|
||||
let mut out = String::with_capacity(body.len() + 256);
|
||||
out.push_str(body);
|
||||
if !body.ends_with('\n') {
|
||||
out.push('\n');
|
||||
}
|
||||
out.push('\n');
|
||||
|
||||
let boundaries = prompts.working_boundaries_section(&scope.summary())?;
|
||||
out.push_str(boundaries.trim_end_matches(&['\n', ' '][..]));
|
||||
out.push('\n');
|
||||
if let Some(agents) = agents_md {
|
||||
out.push('\n');
|
||||
let section = prompts.agents_md_section(agents)?;
|
||||
out.push_str(section.trim_end_matches(&['\n', ' '][..]));
|
||||
out.push('\n');
|
||||
}
|
||||
if let Some(summary) = resident_summary {
|
||||
let summary = summary.trim_matches(&['\n', '\r'][..]);
|
||||
if !summary.trim().is_empty() {
|
||||
out.push('\n');
|
||||
let section = prompts.resident_memory_summary_section(summary)?;
|
||||
out.push_str(section.trim_end_matches(&['\n', ' '][..]));
|
||||
out.push('\n');
|
||||
}
|
||||
}
|
||||
if let Some(entries) = resident_knowledge {
|
||||
if !entries.is_empty() {
|
||||
out.push('\n');
|
||||
let formatted = format_resident_knowledge_entries(entries);
|
||||
let section = prompts.resident_knowledge_section(
|
||||
&formatted,
|
||||
tool_capabilities.knowledge_query,
|
||||
tool_capabilities.memory_read,
|
||||
)?;
|
||||
out.push_str(section.trim_end_matches(&['\n', ' '][..]));
|
||||
out.push('\n');
|
||||
}
|
||||
}
|
||||
if let Some(entries) = resident_workflows {
|
||||
if !entries.is_empty() {
|
||||
out.push('\n');
|
||||
let formatted = format_resident_workflow_entries(entries);
|
||||
let section = prompts.resident_workflows_section(&formatted)?;
|
||||
out.push_str(section.trim_end_matches(&['\n', ' '][..]));
|
||||
out.push('\n');
|
||||
}
|
||||
}
|
||||
if tool_capabilities.worker_management() {
|
||||
out.push('\n');
|
||||
let section = prompts.worker_orchestration_guidance_section()?;
|
||||
out.push_str(section.trim_end_matches(&['\n', ' '][..]));
|
||||
out.push('\n');
|
||||
}
|
||||
// Canonicalise the tail so the emitted prompt has a single form
|
||||
// regardless of how individual templates chose to end.
|
||||
while out.ends_with('\n') || out.ends_with(' ') {
|
||||
out.pop();
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// `- <slug>: <description>` per line. Description newlines are folded
|
||||
/// to spaces so a single entry stays on one row in the rendered prompt.
|
||||
fn format_resident_knowledge_entries(entries: &[ResidentKnowledgeEntry]) -> String {
|
||||
format_resident_entries(
|
||||
entries
|
||||
.iter()
|
||||
.map(|e| (e.slug.as_str(), e.description.as_str())),
|
||||
)
|
||||
}
|
||||
|
||||
fn format_resident_workflow_entries(entries: &[ResidentWorkflowEntry]) -> String {
|
||||
format_resident_entries(
|
||||
entries
|
||||
.iter()
|
||||
.map(|e| (e.slug.as_str(), e.description.as_str())),
|
||||
)
|
||||
}
|
||||
|
||||
fn format_resident_entries<'a>(entries: impl Iterator<Item = (&'a str, &'a str)>) -> String {
|
||||
let mut out = String::new();
|
||||
for (i, (slug, description)) in entries.enumerate() {
|
||||
if i > 0 {
|
||||
out.push('\n');
|
||||
}
|
||||
out.push_str("- ");
|
||||
out.push_str(slug);
|
||||
out.push_str(": ");
|
||||
for ch in description.chars() {
|
||||
if ch == '\n' || ch == '\r' {
|
||||
out.push(' ');
|
||||
} else {
|
||||
out.push(ch);
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Bridge used by [`Worker::ensure_system_prompt_materialized`] so tests
|
||||
/// can construct a synthetic context without going through a full Worker.
|
||||
#[doc(hidden)]
|
||||
pub fn __instruction_ref_for_tests(raw: &str, loader: &PromptLoader) -> Option<PromptRef> {
|
||||
loader.parse_ref(raw, None).ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use chrono::TimeZone;
|
||||
use manifest::{Permission, ScopeConfig, ScopeRule};
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn fixed_now() -> DateTime<Utc> {
|
||||
Utc.with_ymd_and_hms(2026, 4, 15, 9, 30, 0).unwrap()
|
||||
}
|
||||
|
||||
fn build_scope(dir: &Path) -> Scope {
|
||||
let cfg = ScopeConfig {
|
||||
allow: vec![ScopeRule {
|
||||
target: dir.to_path_buf(),
|
||||
permission: Permission::Write,
|
||||
recursive: true,
|
||||
}],
|
||||
deny: Vec::new(),
|
||||
};
|
||||
Scope::from_config(&cfg).unwrap()
|
||||
}
|
||||
|
||||
fn ctx<'a>(
|
||||
cwd: &'a Path,
|
||||
scope: &'a Scope,
|
||||
tools: Vec<String>,
|
||||
agents_md: Option<String>,
|
||||
) -> SystemPromptContext<'a> {
|
||||
SystemPromptContext {
|
||||
now: fixed_now(),
|
||||
cwd,
|
||||
language: manifest::defaults::WORKER_LANGUAGE,
|
||||
scope,
|
||||
tool_names: tools,
|
||||
agents_md,
|
||||
resident_summary: None,
|
||||
resident_knowledge: None,
|
||||
resident_workflows: None,
|
||||
prompts: test_prompts(),
|
||||
}
|
||||
}
|
||||
|
||||
fn ctx_with_summary<'a>(
|
||||
cwd: &'a Path,
|
||||
scope: &'a Scope,
|
||||
summary: Option<&'a str>,
|
||||
) -> SystemPromptContext<'a> {
|
||||
SystemPromptContext {
|
||||
now: fixed_now(),
|
||||
cwd,
|
||||
language: manifest::defaults::WORKER_LANGUAGE,
|
||||
scope,
|
||||
tool_names: Vec::new(),
|
||||
agents_md: None,
|
||||
resident_summary: summary,
|
||||
resident_knowledge: None,
|
||||
resident_workflows: None,
|
||||
prompts: test_prompts(),
|
||||
}
|
||||
}
|
||||
|
||||
fn ctx_with_resident<'a>(
|
||||
cwd: &'a Path,
|
||||
scope: &'a Scope,
|
||||
resident: &'a [ResidentKnowledgeEntry],
|
||||
) -> SystemPromptContext<'a> {
|
||||
SystemPromptContext {
|
||||
now: fixed_now(),
|
||||
cwd,
|
||||
language: manifest::defaults::WORKER_LANGUAGE,
|
||||
scope,
|
||||
tool_names: Vec::new(),
|
||||
agents_md: None,
|
||||
resident_summary: None,
|
||||
resident_knowledge: Some(resident),
|
||||
resident_workflows: None,
|
||||
prompts: test_prompts(),
|
||||
}
|
||||
}
|
||||
|
||||
fn ctx_with_resident_workflows<'a>(
|
||||
cwd: &'a Path,
|
||||
scope: &'a Scope,
|
||||
resident: &'a [ResidentWorkflowEntry],
|
||||
) -> SystemPromptContext<'a> {
|
||||
SystemPromptContext {
|
||||
now: fixed_now(),
|
||||
cwd,
|
||||
language: manifest::defaults::WORKER_LANGUAGE,
|
||||
scope,
|
||||
tool_names: Vec::new(),
|
||||
agents_md: None,
|
||||
resident_summary: None,
|
||||
resident_knowledge: None,
|
||||
resident_workflows: Some(resident),
|
||||
prompts: test_prompts(),
|
||||
}
|
||||
}
|
||||
|
||||
fn memory_tool_names() -> Vec<String> {
|
||||
[
|
||||
"MemoryQuery",
|
||||
"KnowledgeQuery",
|
||||
"MemoryRead",
|
||||
"MemoryWrite",
|
||||
"MemoryEdit",
|
||||
"MemoryDelete",
|
||||
]
|
||||
.into_iter()
|
||||
.map(String::from)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn worker_management_tool_names() -> Vec<String> {
|
||||
[
|
||||
"SpawnWorker",
|
||||
"SendToWorker",
|
||||
"ReadWorkerOutput",
|
||||
"StopWorker",
|
||||
"ListWorkers",
|
||||
"RestoreWorker",
|
||||
]
|
||||
.into_iter()
|
||||
.map(String::from)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Lazily-initialised builtin catalog shared across system-prompt
|
||||
/// tests, so every `ctx()` can hand out a `&'static PromptCatalog`
|
||||
/// reference without forcing test bodies to create one per call.
|
||||
fn test_prompts() -> &'static PromptCatalog {
|
||||
use std::sync::OnceLock;
|
||||
static CELL: OnceLock<Arc<PromptCatalog>> = OnceLock::new();
|
||||
CELL.get_or_init(|| PromptCatalog::builtins_only().unwrap())
|
||||
.as_ref()
|
||||
}
|
||||
|
||||
fn user_loader_with(file_name: &str, body: &str) -> (TempDir, PromptLoader) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
std::fs::write(tmp.path().join(file_name), body).unwrap();
|
||||
let loader = PromptLoader::new(Some(tmp.path().to_path_buf()), None);
|
||||
(tmp, loader)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn instruction_default_resolves_to_yoi_default() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let tmpl = SystemPromptTemplate::parse("$yoi/default", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx(dir.path(), &scope, memory_tool_names(), None))
|
||||
.unwrap();
|
||||
// Builtin default body must expose the tool and language policies.
|
||||
assert!(rendered.contains("### Memory and knowledge"));
|
||||
assert!(rendered.contains("small targeted `MemoryQuery` / `KnowledgeQuery`"));
|
||||
assert!(rendered.contains("Strong lookup triggers include"));
|
||||
assert!(rendered.contains("MemoryRead(kind=summary)"));
|
||||
assert!(rendered.contains("Do not query memory every turn"));
|
||||
assert!(rendered.contains("MemoryWrite"));
|
||||
assert!(rendered.contains("## Language"));
|
||||
assert!(rendered.contains("`language`: `match the user's language"));
|
||||
// Trailing section must be present.
|
||||
assert!(rendered.contains("## Working boundaries"));
|
||||
assert!(rendered.contains("Readable:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn instruction_default_omits_memory_guidance_without_memory_tools() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let tmpl = SystemPromptTemplate::parse("$yoi/default", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx(
|
||||
dir.path(),
|
||||
&scope,
|
||||
vec!["Read".into(), "Edit".into()],
|
||||
None,
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(!rendered.contains("### Memory and knowledge"));
|
||||
assert!(!rendered.contains("MemoryQuery"));
|
||||
assert!(!rendered.contains("KnowledgeQuery"));
|
||||
assert!(!rendered.contains("MemoryRead"));
|
||||
assert!(!rendered.contains("MemoryWrite"));
|
||||
assert!(!rendered.contains("MemoryEdit"));
|
||||
assert!(!rendered.contains("MemoryDelete"));
|
||||
assert!(rendered.contains("## Language"));
|
||||
assert!(rendered.contains("## Working boundaries"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn memory_guidance_names_only_available_memory_tools() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let tmpl = SystemPromptTemplate::parse("$yoi/default", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx(
|
||||
dir.path(),
|
||||
&scope,
|
||||
vec!["MemoryQuery".into(), "MemoryRead".into()],
|
||||
None,
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(rendered.contains("### Memory and knowledge"));
|
||||
assert!(rendered.contains("small targeted `MemoryQuery`"));
|
||||
assert!(rendered.contains("MemoryRead(kind=summary)"));
|
||||
assert!(!rendered.contains("KnowledgeQuery"));
|
||||
assert!(!rendered.contains("MemoryWrite"));
|
||||
assert!(!rendered.contains("MemoryEdit"));
|
||||
assert!(!rendered.contains("MemoryDelete"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_orchestration_guidance_is_included_for_worker_management_tools() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let tmpl = SystemPromptTemplate::parse("$yoi/default", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx(
|
||||
dir.path(),
|
||||
&scope,
|
||||
worker_management_tool_names(),
|
||||
None,
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(rendered.contains("## Worker orchestration"));
|
||||
assert!(rendered.contains("spawned Worker notifications are background signals"));
|
||||
assert!(rendered.contains("does not need to keep a turn open"));
|
||||
assert!(rendered.contains("Do not use `sleep` or polling loops"));
|
||||
assert!(rendered.contains("worktree state, diff, and test results"));
|
||||
assert!(rendered.contains("not scheduler or auto-maintain authorization"));
|
||||
assert!(rendered.contains("bypass user/workflow authorization"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_orchestration_guidance_is_omitted_without_worker_management_tools() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let tmpl = SystemPromptTemplate::parse("$yoi/default", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx(
|
||||
dir.path(),
|
||||
&scope,
|
||||
vec!["Read".into(), "Edit".into(), "MemoryRead".into()],
|
||||
None,
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert!(!rendered.contains("## Worker orchestration"));
|
||||
assert!(!rendered.contains("spawned Worker notifications are background signals"));
|
||||
assert!(!rendered.contains("does not need to keep a turn open"));
|
||||
assert!(!rendered.contains("Do not use `sleep` or polling loops"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn instruction_prefix_addressing_user() {
|
||||
let (_tmp, loader) = user_loader_with("greet.md", "HELLO from {{ cwd }}");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/greet", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl.render(&ctx(dir.path(), &scope, vec![], None)).unwrap();
|
||||
assert!(rendered.starts_with("HELLO from"));
|
||||
assert!(rendered.contains("## Working boundaries"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn instruction_prefix_addressing_workspace() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
std::fs::write(tmp.path().join("ws.md"), "WS {{ date }}").unwrap();
|
||||
let loader = PromptLoader::new(None, Some(tmp.path().to_path_buf()));
|
||||
let tmpl = SystemPromptTemplate::parse("$workspace/ws", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl.render(&ctx(dir.path(), &scope, vec![], None)).unwrap();
|
||||
assert!(rendered.starts_with("WS 2026-04-15"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn include_unqualified_resolves_relative_to_current_prefix() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
// parent.md and sibling.md both under the user root.
|
||||
std::fs::write(
|
||||
tmp.path().join("parent.md"),
|
||||
"PARENT\n{% include \"sibling\" %}",
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::write(tmp.path().join("sibling.md"), "SIBLING-BODY").unwrap();
|
||||
let loader = PromptLoader::new(Some(tmp.path().to_path_buf()), None);
|
||||
let tmpl = SystemPromptTemplate::parse("$user/parent", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl.render(&ctx(dir.path(), &scope, vec![], None)).unwrap();
|
||||
assert!(rendered.contains("PARENT"));
|
||||
assert!(rendered.contains("SIBLING-BODY"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn include_unqualified_from_subdirectory_resolves_in_same_dir() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
std::fs::create_dir(tmp.path().join("common")).unwrap();
|
||||
std::fs::write(
|
||||
tmp.path().join("common/header.md"),
|
||||
"HEADER\n{% include \"nested\" %}",
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::write(tmp.path().join("common/nested.md"), "NESTED-OK").unwrap();
|
||||
let loader = PromptLoader::new(Some(tmp.path().to_path_buf()), None);
|
||||
let tmpl = SystemPromptTemplate::parse("$user/common/header", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl.render(&ctx(dir.path(), &scope, vec![], None)).unwrap();
|
||||
assert!(rendered.contains("HEADER"));
|
||||
assert!(rendered.contains("NESTED-OK"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn include_explicit_prefix_overrides_relative() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
std::fs::write(
|
||||
tmp.path().join("root.md"),
|
||||
"U-ROOT\n{% include \"$yoi/common/tool-usage\" %}",
|
||||
)
|
||||
.unwrap();
|
||||
let loader = PromptLoader::new(Some(tmp.path().to_path_buf()), None);
|
||||
let tmpl = SystemPromptTemplate::parse("$user/root", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx(
|
||||
dir.path(),
|
||||
&scope,
|
||||
vec!["Read".into(), "Edit".into()],
|
||||
None,
|
||||
))
|
||||
.unwrap();
|
||||
assert!(rendered.contains("U-ROOT"));
|
||||
// Pulled in from the builtin tool-usage asset.
|
||||
assert!(rendered.contains("Read"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefix_with_missing_file_is_hard_error() {
|
||||
let loader = PromptLoader::builtins_only();
|
||||
let err = SystemPromptTemplate::parse("$yoi/definitely-missing", loader).unwrap_err();
|
||||
assert!(matches!(err, SystemPromptError::LoaderResolve(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fails_on_syntax_error() {
|
||||
let (_tmp, loader) = user_loader_with("broken.md", "{{ unclosed");
|
||||
let err = SystemPromptTemplate::parse("$user/broken", loader).unwrap_err();
|
||||
assert!(matches!(err, SystemPromptError::Parse(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_fails_on_undefined_variable() {
|
||||
let (_tmp, loader) = user_loader_with("ghost.md", "{{ ghost }}");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/ghost", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let err = tmpl
|
||||
.render(&ctx(dir.path(), &scope, vec![], None))
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, SystemPromptError::Render(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_substitutes_date_cwd_tools() {
|
||||
let (_tmp, loader) = user_loader_with(
|
||||
"vars.md",
|
||||
"date={{ date }} cwd={{ cwd }} tools={{ tools | join(',') }}",
|
||||
);
|
||||
let tmpl = SystemPromptTemplate::parse("$user/vars", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx(
|
||||
dir.path(),
|
||||
&scope,
|
||||
vec!["alpha".into(), "beta".into()],
|
||||
None,
|
||||
))
|
||||
.unwrap();
|
||||
assert!(rendered.contains("date=2026-04-15"));
|
||||
assert!(rendered.contains(&format!("cwd={}", dir.path().display())));
|
||||
assert!(rendered.contains("tools=alpha,beta"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_always_contains_scope_summary() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl.render(&ctx(dir.path(), &scope, vec![], None)).unwrap();
|
||||
assert!(rendered.contains("## Working boundaries"));
|
||||
assert!(rendered.contains("Readable:"));
|
||||
assert!(rendered.contains("Writable:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_contains_agents_md_when_present() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx(
|
||||
dir.path(),
|
||||
&scope,
|
||||
vec![],
|
||||
Some("PROJECT DOCS".into()),
|
||||
))
|
||||
.unwrap();
|
||||
assert!(rendered.contains("## Project instructions (AGENTS.md)"));
|
||||
assert!(rendered.contains("PROJECT DOCS"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_omits_agents_md_when_absent() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl.render(&ctx(dir.path(), &scope, vec![], None)).unwrap();
|
||||
assert!(!rendered.contains("AGENTS.md"));
|
||||
assert!(!rendered.contains("Project instructions"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_renders_resident_summary_body() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx_with_summary(
|
||||
dir.path(),
|
||||
&scope,
|
||||
Some("Persistent summary body"),
|
||||
))
|
||||
.unwrap();
|
||||
assert!(rendered.contains("## Resident memory summary"));
|
||||
assert!(rendered.contains("Persistent summary body"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_omits_resident_summary_when_none_or_empty() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx_with_summary(dir.path(), &scope, None))
|
||||
.unwrap();
|
||||
assert!(!rendered.contains("Resident memory summary"));
|
||||
|
||||
let rendered = tmpl
|
||||
.render(&ctx_with_summary(dir.path(), &scope, Some(" \n")))
|
||||
.unwrap();
|
||||
assert!(!rendered.contains("Resident memory summary"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_omits_resident_knowledge_when_none() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl.render(&ctx(dir.path(), &scope, vec![], None)).unwrap();
|
||||
assert!(!rendered.contains("Resident knowledge"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_omits_resident_knowledge_when_empty_slice() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let rendered = tmpl
|
||||
.render(&ctx_with_resident(dir.path(), &scope, &[]))
|
||||
.unwrap();
|
||||
assert!(!rendered.contains("Resident knowledge"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_renders_resident_knowledge_entries() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let entries = vec![
|
||||
ResidentKnowledgeEntry {
|
||||
slug: "alpha".into(),
|
||||
description: "first record".into(),
|
||||
},
|
||||
ResidentKnowledgeEntry {
|
||||
slug: "beta".into(),
|
||||
description: "second record\nwith newline".into(),
|
||||
},
|
||||
];
|
||||
let rendered = tmpl
|
||||
.render(&ctx_with_resident(dir.path(), &scope, &entries))
|
||||
.unwrap();
|
||||
assert!(rendered.contains("## Resident knowledge"));
|
||||
assert!(rendered.contains("- alpha: first record"));
|
||||
// Newline in description is folded to a space (one entry per line).
|
||||
assert!(rendered.contains("- beta: second record with newline"));
|
||||
assert!(!rendered.contains("KnowledgeQuery"));
|
||||
assert!(!rendered.contains("MemoryRead"));
|
||||
// Resident section sits *after* the working-boundaries header.
|
||||
let pos_boundaries = rendered.find("## Working boundaries").unwrap();
|
||||
let pos_resident = rendered.find("## Resident knowledge").unwrap();
|
||||
assert!(pos_resident > pos_boundaries);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_mentions_resident_knowledge_tools_when_available() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let entries = [ResidentKnowledgeEntry {
|
||||
slug: "alpha".into(),
|
||||
description: "first record".into(),
|
||||
}];
|
||||
let mut context = ctx_with_resident(dir.path(), &scope, &entries);
|
||||
context.tool_names = memory_tool_names();
|
||||
let rendered = tmpl.render(&context).unwrap();
|
||||
|
||||
assert!(rendered.contains("## Resident knowledge"));
|
||||
assert!(rendered.contains("KnowledgeQuery / MemoryRead"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_renders_resident_workflows() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let workflows = [ResidentWorkflowEntry {
|
||||
slug: "resident-flow".to_string(),
|
||||
description: "workflow resident desc\nwith newline".to_string(),
|
||||
}];
|
||||
let rendered = tmpl
|
||||
.render(&ctx_with_resident_workflows(dir.path(), &scope, &workflows))
|
||||
.unwrap();
|
||||
|
||||
assert!(rendered.contains("## Resident workflows"));
|
||||
assert!(rendered.contains("- resident-flow: workflow resident desc with newline"));
|
||||
let pos_boundaries = rendered.find("## Working boundaries").unwrap();
|
||||
let pos_resident = rendered.find("## Resident workflows").unwrap();
|
||||
assert!(pos_resident > pos_boundaries);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trailing_section_omits_empty_resident_workflows() {
|
||||
let (_tmp, loader) = user_loader_with("body.md", "BODY");
|
||||
let tmpl = SystemPromptTemplate::parse("$user/body", loader).unwrap();
|
||||
let dir = TempDir::new().unwrap();
|
||||
let scope = build_scope(dir.path());
|
||||
let workflows: [ResidentWorkflowEntry; 0] = [];
|
||||
let rendered = tmpl
|
||||
.render(&ctx_with_resident_workflows(dir.path(), &scope, &workflows))
|
||||
.unwrap();
|
||||
|
||||
assert!(!rendered.contains("Resident workflows"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
use std::io;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use manifest::{ScopeRule, paths};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::fs;
|
||||
|
||||
use crate::shared_state::WorkerSharedState;
|
||||
|
||||
/// One spawned-child record mirrored to `spawned_workers.json`.
|
||||
///
|
||||
/// Written by the spawner after registry changes so runtime-local tools
|
||||
/// have a materialised snapshot. Durable restore uses Worker state metadata;
|
||||
/// this file is not the authoritative source.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SpawnedWorkerRecord {
|
||||
/// Spawned Worker's identity.
|
||||
pub worker_name: String,
|
||||
/// Spawned Worker's Unix socket path.
|
||||
pub socket_path: PathBuf,
|
||||
/// Scope allow rules delegated to the spawned Worker.
|
||||
pub scope_delegated: Vec<ScopeRule>,
|
||||
/// Socket path the spawned Worker was told to use for callbacks
|
||||
/// (= this Worker's own socket when spawn happened).
|
||||
pub callback_address: PathBuf,
|
||||
}
|
||||
|
||||
/// Manages the Worker's runtime directory on tmpfs.
|
||||
///
|
||||
/// ```text
|
||||
/// <runtime_dir>/{worker_name}/
|
||||
/// ├── pid
|
||||
/// ├── status.json
|
||||
/// ├── manifest.toml
|
||||
/// ├── history.json
|
||||
/// └── sock (created by socket listener, not by RuntimeDir)
|
||||
/// ```
|
||||
///
|
||||
/// `<runtime_dir>` is resolved via [`manifest::paths::runtime_dir`].
|
||||
/// Files are written atomically (write tmp → rename).
|
||||
/// The directory is removed on drop.
|
||||
pub struct RuntimeDir {
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
impl RuntimeDir {
|
||||
/// Create the runtime directory and write the PID file.
|
||||
pub async fn create(base: &Path, worker_name: &str) -> Result<Self, io::Error> {
|
||||
let path = base.join(worker_name);
|
||||
fs::create_dir_all(&path).await?;
|
||||
|
||||
let pid = std::process::id().to_string();
|
||||
fs::write(path.join("pid"), pid.as_bytes()).await?;
|
||||
|
||||
Ok(Self { path })
|
||||
}
|
||||
|
||||
/// Create in the default base directory resolved via
|
||||
/// [`manifest::paths::runtime_dir`].
|
||||
pub async fn create_default(worker_name: &str) -> Result<Self, io::Error> {
|
||||
let base = default_base()?;
|
||||
Self::create(&base, worker_name).await
|
||||
}
|
||||
|
||||
/// Write status.json atomically.
|
||||
pub async fn write_status(&self, state: &WorkerSharedState) -> Result<(), io::Error> {
|
||||
let content = state.status_json();
|
||||
atomic_write(&self.path.join("status.json"), content.as_bytes()).await
|
||||
}
|
||||
|
||||
/// Write manifest.toml (typically once at startup).
|
||||
pub async fn write_manifest(&self, toml: &str) -> Result<(), io::Error> {
|
||||
atomic_write(&self.path.join("manifest.toml"), toml.as_bytes()).await
|
||||
}
|
||||
|
||||
/// Write `spawned_workers.json` atomically. The entries are the full
|
||||
/// set of spawned children known to this Worker — callers pass the
|
||||
/// replacement list, no incremental merge.
|
||||
pub async fn write_spawned_workers(
|
||||
&self,
|
||||
records: &[SpawnedWorkerRecord],
|
||||
) -> Result<(), io::Error> {
|
||||
let json = serde_json::to_vec_pretty(records).map_err(io::Error::other)?;
|
||||
atomic_write(&self.path.join("spawned_workers.json"), &json).await
|
||||
}
|
||||
|
||||
/// Path to this Worker's runtime directory.
|
||||
pub fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
|
||||
/// Path where the Unix socket should be created. External callers
|
||||
/// that only know the worker name (e.g. the TUI's attach flow)
|
||||
/// predict the same path via [`manifest::paths::worker_socket_path`].
|
||||
pub fn socket_path(&self) -> PathBuf {
|
||||
self.path.join("sock")
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for RuntimeDir {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::remove_dir_all(&self.path);
|
||||
}
|
||||
}
|
||||
|
||||
/// Atomic write: write to a temp file, then rename.
|
||||
async fn atomic_write(target: &Path, content: &[u8]) -> Result<(), io::Error> {
|
||||
let tmp = target.with_extension("tmp");
|
||||
fs::write(&tmp, content).await?;
|
||||
fs::rename(&tmp, target).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolve the default base directory for runtime data.
|
||||
///
|
||||
/// Thin wrapper over [`manifest::paths::runtime_dir`] that converts a
|
||||
/// missing-env situation into an `io::Error`.
|
||||
pub fn default_base() -> Result<PathBuf, io::Error> {
|
||||
paths::runtime_dir().ok_or_else(|| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::NotFound,
|
||||
"could not resolve runtime directory (no YOI_HOME / \
|
||||
YOI_RUNTIME_DIR / XDG_RUNTIME_DIR / HOME)",
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::shared_state::WorkerSharedState;
|
||||
use protocol::WorkerStatus;
|
||||
|
||||
fn test_state() -> WorkerSharedState {
|
||||
WorkerSharedState::new(
|
||||
"test-worker".into(),
|
||||
session_store::new_segment_id(),
|
||||
"[engine]\nname = \"test-worker\"".into(),
|
||||
protocol::Greeting {
|
||||
worker_name: "test-worker".into(),
|
||||
cwd: "/tmp".into(),
|
||||
provider: "anthropic".into(),
|
||||
model: "claude".into(),
|
||||
scope_summary: String::new(),
|
||||
tools: Vec::new(),
|
||||
context_window: 200_000,
|
||||
context_tokens: 0,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn creates_directory_and_pid() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let rt = RuntimeDir::create(tmp.path(), "my-worker").await.unwrap();
|
||||
|
||||
assert!(rt.path().join("pid").exists());
|
||||
let pid = std::fs::read_to_string(rt.path().join("pid")).unwrap();
|
||||
assert_eq!(pid, std::process::id().to_string());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_status_creates_file() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let rt = RuntimeDir::create(tmp.path(), "my-worker").await.unwrap();
|
||||
let state = test_state();
|
||||
|
||||
rt.write_status(&state).await.unwrap();
|
||||
|
||||
let content = std::fs::read_to_string(rt.path().join("status.json")).unwrap();
|
||||
let parsed: serde_json::Value = serde_json::from_str(&content).unwrap();
|
||||
assert_eq!(parsed["state"], "idle");
|
||||
assert_eq!(parsed["worker_name"], "test-worker");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_status_reflects_changes() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let rt = RuntimeDir::create(tmp.path(), "my-worker").await.unwrap();
|
||||
let state = test_state();
|
||||
|
||||
state.set_status(WorkerStatus::Running);
|
||||
rt.write_status(&state).await.unwrap();
|
||||
|
||||
let content = std::fs::read_to_string(rt.path().join("status.json")).unwrap();
|
||||
let parsed: serde_json::Value = serde_json::from_str(&content).unwrap();
|
||||
assert_eq!(parsed["state"], "running");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_manifest_creates_file() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let rt = RuntimeDir::create(tmp.path(), "my-worker").await.unwrap();
|
||||
|
||||
rt.write_manifest("[engine]\nname = \"test\"")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let content = std::fs::read_to_string(rt.path().join("manifest.toml")).unwrap();
|
||||
assert_eq!(content, "[engine]\nname = \"test\"");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_spawned_workers_creates_file() {
|
||||
use manifest::{Permission, ScopeRule};
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let rt = RuntimeDir::create(tmp.path(), "my-worker").await.unwrap();
|
||||
|
||||
let records = vec![SpawnedWorkerRecord {
|
||||
worker_name: "child".into(),
|
||||
socket_path: "/run/yoi/child/sock".into(),
|
||||
scope_delegated: vec![ScopeRule {
|
||||
target: "/tmp/work".into(),
|
||||
permission: Permission::Write,
|
||||
recursive: true,
|
||||
}],
|
||||
callback_address: "/run/yoi/my-worker/sock".into(),
|
||||
}];
|
||||
rt.write_spawned_workers(&records).await.unwrap();
|
||||
|
||||
let content = std::fs::read_to_string(rt.path().join("spawned_workers.json")).unwrap();
|
||||
let parsed: Vec<SpawnedWorkerRecord> = serde_json::from_str(&content).unwrap();
|
||||
assert_eq!(parsed.len(), 1);
|
||||
assert_eq!(parsed[0].worker_name, "child");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn socket_path() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let rt = RuntimeDir::create(tmp.path(), "my-worker").await.unwrap();
|
||||
assert_eq!(rt.socket_path(), rt.path().join("sock"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drop_removes_directory() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let dir_path;
|
||||
{
|
||||
let rt = RuntimeDir::create(tmp.path(), "my-worker").await.unwrap();
|
||||
dir_path = rt.path().to_owned();
|
||||
assert!(dir_path.exists());
|
||||
}
|
||||
assert!(!dir_path.exists());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod dir;
|
||||
pub use ::pod_registry;
|
||||
@@ -0,0 +1,375 @@
|
||||
//! Worker-side session-log mirror + broadcast.
|
||||
//!
|
||||
//! Owns the in-memory `Vec<LogEntry>` mirror that backs `Event::Snapshot`
|
||||
//! delivery to newly connected clients and the
|
||||
//! `broadcast::Sender<LogEntry>` that fans out per-entry commits to
|
||||
//! existing subscribers. Disk writes remain the responsibility of the
|
||||
//! Worker (which still owns the `Store` handle); the sink stays focused on
|
||||
//! the wire-side fan-out.
|
||||
//!
|
||||
//! Atomicity contract:
|
||||
//!
|
||||
//! 1. Worker writes the entry to disk via the `Store`.
|
||||
//! 2. Worker calls [`SegmentLogSink::publish`] which acquires the mirror
|
||||
//! mutex, pushes the entry, and fires `broadcast::send` — all under
|
||||
//! the same critical section.
|
||||
//!
|
||||
//! [`SegmentLogSink::subscribe_with_snapshot`] takes the same mutex,
|
||||
//! so the `(snapshot, receiver)` pair returned to a connecting client
|
||||
//! splits the entry sequence cleanly: every entry shows up in exactly
|
||||
//! one of `snapshot` or on `receiver`.
|
||||
//!
|
||||
//! Disk-write failures short-circuit before `publish`, so a failed
|
||||
//! entry never appears in the mirror or on the broadcast.
|
||||
|
||||
use std::sync::{Arc, Mutex as StdMutex};
|
||||
|
||||
use session_store::LogEntry;
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
/// Broadcast capacity for the live receiver. Slow subscribers that
|
||||
/// fall behind will see `RecvError::Lagged` and are expected to drop
|
||||
/// the connection so that the next reconnect's `subscribe_with_snapshot`
|
||||
/// re-seeds the prefix.
|
||||
const BROADCAST_CAPACITY: usize = 256;
|
||||
|
||||
/// In-memory mirror + broadcast fan-out for the active session log.
|
||||
///
|
||||
/// Clone is cheap (`Arc` clone) — the Worker hands one to the IPC layer
|
||||
/// for read-only `subscribe_with_snapshot` access and keeps one for
|
||||
/// its own write path.
|
||||
#[derive(Clone)]
|
||||
pub struct SegmentLogSink {
|
||||
inner: Arc<SinkInner>,
|
||||
}
|
||||
|
||||
struct SinkInner {
|
||||
/// Full session log mirror in commit order. Reset on session swap
|
||||
/// (compaction / fork) via [`SegmentLogSink::reset_with_initial`].
|
||||
mirror: StdMutex<Vec<LogEntry>>,
|
||||
/// Broadcast channel for live entry updates. The same `Sender`
|
||||
/// survives session swaps so existing subscribers keep their
|
||||
/// receiver — they observe the swap as a freshly broadcast
|
||||
/// `LogEntry::SegmentStart` and reset their view accordingly.
|
||||
broadcast_tx: broadcast::Sender<LogEntry>,
|
||||
}
|
||||
|
||||
impl SegmentLogSink {
|
||||
/// Create a fresh sink with an empty mirror. Used before any entry
|
||||
/// has been written (deferred SegmentStart) or as a placeholder in
|
||||
/// tests.
|
||||
pub fn new() -> Self {
|
||||
let (broadcast_tx, _) = broadcast::channel(BROADCAST_CAPACITY);
|
||||
Self {
|
||||
inner: Arc::new(SinkInner {
|
||||
mirror: StdMutex::new(Vec::new()),
|
||||
broadcast_tx,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a sink seeded with a prefix of entries already on disk.
|
||||
/// Used by restore / fork-at-restore code paths that materialise
|
||||
/// the existing log before the sink starts taking new commits.
|
||||
pub fn with_initial(entries: Vec<LogEntry>) -> Self {
|
||||
let (broadcast_tx, _) = broadcast::channel(BROADCAST_CAPACITY);
|
||||
Self {
|
||||
inner: Arc::new(SinkInner {
|
||||
mirror: StdMutex::new(entries),
|
||||
broadcast_tx,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Push `entry` to the mirror; selectively broadcast it.
|
||||
///
|
||||
/// MUST be called only after the Worker has successfully persisted the
|
||||
/// entry to the underlying `Store` — disk write is the gate. Failed
|
||||
/// disk writes must not call `publish`.
|
||||
///
|
||||
/// Live broadcast fires for committed session-log entries that
|
||||
/// socket clients must see in log order:
|
||||
/// - `LogEntry::SegmentStart` → `Event::SegmentRotated` on the wire.
|
||||
/// - `LogEntry::UserInput` → `Event::UserMessage`.
|
||||
/// - `LogEntry::SystemItem` → `Event::SystemItem`.
|
||||
/// - `LogEntry::Invoke` → `Event::InvokeStart`.
|
||||
/// Everything else (AssistantItem, ToolResult, TurnEnd,
|
||||
/// RunCompleted, RunErrored, PausedTurnAbandoned, LlmUsage, Extension,
|
||||
/// ConfigChanged) is reflected in the mirror so reconnect snapshots stay accurate,
|
||||
/// but is not sent live — the streaming events (TextDelta /
|
||||
/// ToolCallStart / ToolResult / TurnEnd / etc.) already provide
|
||||
/// that data, and re-broadcasting it as a typed entry would just
|
||||
/// double-render every block on the client side.
|
||||
pub fn publish(&self, entry: LogEntry) {
|
||||
let mut mirror = self
|
||||
.inner
|
||||
.mirror
|
||||
.lock()
|
||||
.expect("session log mirror mutex poisoned");
|
||||
mirror.push(entry.clone());
|
||||
if Self::is_live_relevant(&entry) {
|
||||
// SendError means there are zero subscribers; harmless. The
|
||||
// mirror lock is held across `send` so subscribers cannot
|
||||
// observe an inconsistent (snapshot, receiver) pair.
|
||||
let _ = self.inner.broadcast_tx.send(entry);
|
||||
}
|
||||
}
|
||||
|
||||
/// `true` for entry kinds that the IPC layer forwards to clients
|
||||
/// as a typed live event.
|
||||
fn is_live_relevant(entry: &LogEntry) -> bool {
|
||||
matches!(
|
||||
entry,
|
||||
LogEntry::SegmentStart { .. }
|
||||
| LogEntry::UserInput { .. }
|
||||
| LogEntry::SystemItem { .. }
|
||||
| LogEntry::Invoke { .. }
|
||||
)
|
||||
}
|
||||
|
||||
/// Atomically swap the mirror to `[initial]` and broadcast the new
|
||||
/// session-start entry. Used during compaction / fork: the new
|
||||
/// `LogEntry::SegmentStart` is the first entry of the replacement
|
||||
/// session, and existing subscribers transition by replaying it
|
||||
/// like any other live entry.
|
||||
///
|
||||
/// Existing snapshot prefixes seen by old subscribers stay valid
|
||||
/// for the prior session; the new `SegmentStart` on the broadcast
|
||||
/// is the signal to reset their derived view.
|
||||
pub fn reset_with_initial(&self, initial: LogEntry) {
|
||||
let mut mirror = self
|
||||
.inner
|
||||
.mirror
|
||||
.lock()
|
||||
.expect("session log mirror mutex poisoned");
|
||||
mirror.clear();
|
||||
mirror.push(initial.clone());
|
||||
let _ = self.inner.broadcast_tx.send(initial);
|
||||
}
|
||||
|
||||
/// Atomically swap the mirror to the supplied replacement-session prefix
|
||||
/// and broadcast the first entry as the live rotation signal. Entries after
|
||||
/// the first are already reflected in reconnect snapshots but are not
|
||||
/// broadcast live; this is intended for non-live extension state that must
|
||||
/// share the new segment prefix with SegmentStart.
|
||||
pub fn reset_with_initial_entries(&self, entries: Vec<LogEntry>) {
|
||||
let first = entries.first().cloned();
|
||||
let mut mirror = self
|
||||
.inner
|
||||
.mirror
|
||||
.lock()
|
||||
.expect("session log mirror mutex poisoned");
|
||||
*mirror = entries;
|
||||
if let Some(initial) = first {
|
||||
let _ = self.inner.broadcast_tx.send(initial);
|
||||
}
|
||||
}
|
||||
|
||||
/// Replace the mirror with the supplied prefix without broadcasting.
|
||||
///
|
||||
/// Used by restore paths that load a session's complete log into
|
||||
/// the mirror before any subscriber is connected. Callers that need
|
||||
/// to notify existing subscribers should use [`reset_with_initial`].
|
||||
pub fn replace_silent(&self, entries: Vec<LogEntry>) {
|
||||
let mut mirror = self
|
||||
.inner
|
||||
.mirror
|
||||
.lock()
|
||||
.expect("session log mirror mutex poisoned");
|
||||
*mirror = entries;
|
||||
}
|
||||
|
||||
/// Truncate the mirror without broadcasting.
|
||||
pub fn truncate_silent(&self, entries_len: usize) {
|
||||
let mut mirror = self
|
||||
.inner
|
||||
.mirror
|
||||
.lock()
|
||||
.expect("session log mirror mutex poisoned");
|
||||
mirror.truncate(entries_len);
|
||||
}
|
||||
|
||||
/// Atomically read the current mirror and subscribe to subsequent
|
||||
/// commits. The returned snapshot and receiver split the entry
|
||||
/// timeline into a duplicate-free, gap-free prefix/suffix pair.
|
||||
pub fn subscribe_with_snapshot(&self) -> (Vec<LogEntry>, broadcast::Receiver<LogEntry>) {
|
||||
let mirror = self
|
||||
.inner
|
||||
.mirror
|
||||
.lock()
|
||||
.expect("session log mirror mutex poisoned");
|
||||
let snapshot = mirror.clone();
|
||||
let rx = self.inner.broadcast_tx.subscribe();
|
||||
(snapshot, rx)
|
||||
}
|
||||
|
||||
/// Current entry count. Useful for tests / diagnostics.
|
||||
pub fn len(&self) -> usize {
|
||||
self.inner
|
||||
.mirror
|
||||
.lock()
|
||||
.expect("session log mirror mutex poisoned")
|
||||
.len()
|
||||
}
|
||||
|
||||
/// Whether the mirror is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SegmentLogSink {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use llm_engine::llm_client::RequestConfig;
|
||||
use session_store::segment_log::now_millis;
|
||||
|
||||
fn session_start() -> LogEntry {
|
||||
LogEntry::SegmentStart {
|
||||
ts: now_millis(),
|
||||
session_id: uuid::Uuid::nil(),
|
||||
system_prompt: None,
|
||||
config: RequestConfig::default(),
|
||||
history: vec![],
|
||||
forked_from: None,
|
||||
compacted_from: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn turn_end(n: usize) -> LogEntry {
|
||||
LogEntry::TurnEnd {
|
||||
ts: now_millis(),
|
||||
turn_count: n,
|
||||
}
|
||||
}
|
||||
|
||||
fn user_input(text: &str) -> LogEntry {
|
||||
LogEntry::UserInput {
|
||||
ts: now_millis(),
|
||||
segments: vec![protocol::Segment::Text {
|
||||
content: text.to_owned(),
|
||||
}],
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn publish_then_subscribe_returns_history_in_snapshot() {
|
||||
let sink = SegmentLogSink::new();
|
||||
sink.publish(session_start());
|
||||
sink.publish(turn_end(1));
|
||||
|
||||
let (snapshot, mut rx) = sink.subscribe_with_snapshot();
|
||||
assert_eq!(snapshot.len(), 2);
|
||||
assert!(matches!(snapshot[0], LogEntry::SegmentStart { .. }));
|
||||
assert!(matches!(
|
||||
snapshot[1],
|
||||
LogEntry::TurnEnd { turn_count: 1, .. }
|
||||
));
|
||||
assert!(rx.try_recv().is_err());
|
||||
}
|
||||
|
||||
fn notification_entry(text: &str) -> LogEntry {
|
||||
LogEntry::SystemItem {
|
||||
ts: now_millis(),
|
||||
item: session_store::SystemItem::Notification {
|
||||
message: text.to_owned(),
|
||||
body: format!("[Notification] {text}"),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subscribe_then_publish_delivers_only_live_relevant_entries() {
|
||||
let sink = SegmentLogSink::new();
|
||||
sink.publish(session_start());
|
||||
|
||||
let (snapshot, mut rx) = sink.subscribe_with_snapshot();
|
||||
assert_eq!(snapshot.len(), 1);
|
||||
|
||||
// TurnEnd is mirror-only — no live broadcast.
|
||||
sink.publish(turn_end(1));
|
||||
assert!(rx.try_recv().is_err(), "TurnEnd must not be broadcast live");
|
||||
|
||||
// UserInput is live-relevant because it is the persisted source
|
||||
// for Event::UserMessage.
|
||||
sink.publish(user_input("hi from log"));
|
||||
match rx.try_recv() {
|
||||
Ok(LogEntry::UserInput { segments, .. }) => {
|
||||
assert_eq!(segments.len(), 1);
|
||||
}
|
||||
other => panic!("expected UserInput, got {other:?}"),
|
||||
}
|
||||
|
||||
// SystemItem is live-relevant.
|
||||
sink.publish(notification_entry("hi"));
|
||||
match rx.try_recv() {
|
||||
Ok(LogEntry::SystemItem { .. }) => {}
|
||||
other => panic!("expected SystemItem, got {other:?}"),
|
||||
}
|
||||
|
||||
// Mirror still grew with all entries (snapshot completeness).
|
||||
let (after_snapshot, _) = sink.subscribe_with_snapshot();
|
||||
assert_eq!(after_snapshot.len(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_and_live_never_overlap() {
|
||||
let sink = SegmentLogSink::new();
|
||||
sink.publish(session_start());
|
||||
let (snapshot, mut rx) = sink.subscribe_with_snapshot();
|
||||
sink.publish(notification_entry("post-snapshot"));
|
||||
|
||||
assert_eq!(snapshot.len(), 1);
|
||||
match rx.try_recv() {
|
||||
Ok(LogEntry::SystemItem { .. }) => {}
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
assert!(rx.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_with_initial_clears_and_broadcasts() {
|
||||
let sink = SegmentLogSink::new();
|
||||
sink.publish(session_start());
|
||||
sink.publish(turn_end(1));
|
||||
|
||||
let (_pre_snapshot, mut rx) = sink.subscribe_with_snapshot();
|
||||
sink.reset_with_initial(session_start());
|
||||
|
||||
match rx.try_recv() {
|
||||
Ok(LogEntry::SegmentStart { .. }) => {}
|
||||
other => panic!("expected SegmentStart broadcast, got {other:?}"),
|
||||
}
|
||||
|
||||
let (post_snapshot, _) = sink.subscribe_with_snapshot();
|
||||
assert_eq!(post_snapshot.len(), 1);
|
||||
assert!(matches!(post_snapshot[0], LogEntry::SegmentStart { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_silent_does_not_broadcast() {
|
||||
let sink = SegmentLogSink::new();
|
||||
sink.publish(session_start());
|
||||
let (_pre_snapshot, mut rx) = sink.subscribe_with_snapshot();
|
||||
|
||||
sink.replace_silent(vec![session_start(), turn_end(1)]);
|
||||
|
||||
// No broadcast fired.
|
||||
assert!(rx.try_recv().is_err());
|
||||
let (post_snapshot, _) = sink.subscribe_with_snapshot();
|
||||
assert_eq!(post_snapshot.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_initial_seeds_the_mirror() {
|
||||
let sink = SegmentLogSink::with_initial(vec![session_start(), turn_end(1)]);
|
||||
let (snapshot, _) = sink.subscribe_with_snapshot();
|
||||
assert_eq!(snapshot.len(), 2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
use std::sync::{OnceLock, RwLock};
|
||||
|
||||
use protocol::WorkerStatus;
|
||||
use serde_json::json;
|
||||
use session_store::SegmentId;
|
||||
|
||||
use crate::fs_view::WorkerFsView;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct WorkflowCandidate {
|
||||
pub slug: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct KnowledgeCandidate {
|
||||
pub slug: String,
|
||||
}
|
||||
|
||||
/// Shared state between WorkerController and runtime directory.
|
||||
///
|
||||
/// Controller updates this in-memory; RuntimeDir writes the status
|
||||
/// snapshot to disk. Wrapped in `Arc` for sharing.
|
||||
///
|
||||
/// History and typed user-segment mirrors used to live here so the
|
||||
/// IPC layer could answer `Method::GetHistory`. Those reads now go
|
||||
/// directly through the session-log sink (`Event::Snapshot` +
|
||||
/// `Event::Entry`), so this struct holds only status, identity,
|
||||
/// greeting, and completion lookup hubs.
|
||||
pub struct WorkerSharedState {
|
||||
pub worker_name: String,
|
||||
pub segment_id: SegmentId,
|
||||
pub manifest_toml: String,
|
||||
pub greeting: protocol::Greeting,
|
||||
pub status: RwLock<WorkerStatus>,
|
||||
/// Worker-from-the-inside view of the filesystem. Set once in
|
||||
/// `WorkerController::start` after the `ScopedFs` is materialised, and
|
||||
/// read from the IPC server layer to answer `ListCompletions`
|
||||
/// queries without going through the controller. `None` until set
|
||||
/// (only relevant for unit tests that build a `WorkerSharedState`
|
||||
/// directly without spinning up a controller).
|
||||
fs_view: OnceLock<WorkerFsView>,
|
||||
workflows: OnceLock<Vec<WorkflowCandidate>>,
|
||||
knowledge: OnceLock<Vec<KnowledgeCandidate>>,
|
||||
}
|
||||
|
||||
impl WorkerSharedState {
|
||||
pub fn new(
|
||||
worker_name: String,
|
||||
segment_id: SegmentId,
|
||||
manifest_toml: String,
|
||||
greeting: protocol::Greeting,
|
||||
) -> Self {
|
||||
Self {
|
||||
worker_name,
|
||||
segment_id,
|
||||
manifest_toml,
|
||||
greeting,
|
||||
status: RwLock::new(WorkerStatus::Idle),
|
||||
fs_view: OnceLock::new(),
|
||||
workflows: OnceLock::new(),
|
||||
knowledge: OnceLock::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach the Worker's filesystem view. Called once during controller
|
||||
/// startup. Subsequent calls are silently ignored (`OnceLock`).
|
||||
pub fn set_fs_view(&self, view: WorkerFsView) {
|
||||
let _ = self.fs_view.set(view);
|
||||
}
|
||||
|
||||
/// Borrow the attached `WorkerFsView`, if any. Returns `None` for unit
|
||||
/// tests that didn't wire one up.
|
||||
pub fn fs_view(&self) -> Option<&WorkerFsView> {
|
||||
self.fs_view.get()
|
||||
}
|
||||
|
||||
pub fn set_workflows(&self, workflows: Vec<WorkflowCandidate>) {
|
||||
let _ = self.workflows.set(workflows);
|
||||
}
|
||||
|
||||
pub fn list_workflow_completions(&self, prefix: &str) -> Vec<WorkflowCandidate> {
|
||||
self.workflows
|
||||
.get()
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter(|candidate| candidate.slug.starts_with(prefix))
|
||||
.cloned()
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub fn set_knowledge(&self, knowledge: Vec<KnowledgeCandidate>) {
|
||||
let _ = self.knowledge.set(knowledge);
|
||||
}
|
||||
|
||||
pub fn list_knowledge_completions(&self, prefix: &str) -> Vec<KnowledgeCandidate> {
|
||||
self.knowledge
|
||||
.get()
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter(|candidate| candidate.slug.starts_with(prefix))
|
||||
.cloned()
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub fn set_status(&self, status: WorkerStatus) {
|
||||
if let Ok(mut s) = self.status.write() {
|
||||
*s = status;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_status(&self) -> WorkerStatus {
|
||||
self.status.read().map(|s| *s).unwrap_or(WorkerStatus::Idle)
|
||||
}
|
||||
|
||||
/// Serialize status as JSON.
|
||||
pub fn status_json(&self) -> String {
|
||||
let status = self.get_status();
|
||||
json!({
|
||||
"state": status,
|
||||
"segment_id": self.segment_id.to_string(),
|
||||
"worker_name": self.worker_name,
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn test_state() -> WorkerSharedState {
|
||||
WorkerSharedState::new(
|
||||
"test-worker".into(),
|
||||
session_store::new_segment_id(),
|
||||
"[engine]\nname = \"test-worker\"".into(),
|
||||
test_greeting(),
|
||||
)
|
||||
}
|
||||
|
||||
fn test_greeting() -> protocol::Greeting {
|
||||
protocol::Greeting {
|
||||
worker_name: "test-worker".into(),
|
||||
cwd: "/tmp".into(),
|
||||
provider: "anthropic".into(),
|
||||
model: "claude".into(),
|
||||
scope_summary: String::new(),
|
||||
tools: Vec::new(),
|
||||
context_window: 200_000,
|
||||
context_tokens: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn initial_status_is_idle() {
|
||||
let state = test_state();
|
||||
assert_eq!(state.get_status(), WorkerStatus::Idle);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_and_get_status() {
|
||||
let state = test_state();
|
||||
state.set_status(WorkerStatus::Running);
|
||||
assert_eq!(state.get_status(), WorkerStatus::Running);
|
||||
state.set_status(WorkerStatus::Paused);
|
||||
assert_eq!(state.get_status(), WorkerStatus::Paused);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn status_json_contains_fields() {
|
||||
let state = test_state();
|
||||
let json = state.status_json();
|
||||
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed["state"], "idle");
|
||||
assert_eq!(parsed["worker_name"], "test-worker");
|
||||
assert!(parsed["segment_id"].is_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn status_json_reflects_changes() {
|
||||
let state = test_state();
|
||||
state.set_status(WorkerStatus::Running);
|
||||
let json = state.status_json();
|
||||
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed["state"], "running");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn knowledge_completions_empty_when_unset() {
|
||||
let state = test_state();
|
||||
assert!(state.list_knowledge_completions("").is_empty());
|
||||
assert!(state.list_knowledge_completions("foo").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn knowledge_completions_filter_by_prefix() {
|
||||
let state = test_state();
|
||||
state.set_knowledge(vec![
|
||||
KnowledgeCandidate {
|
||||
slug: "alpha".into(),
|
||||
},
|
||||
KnowledgeCandidate {
|
||||
slug: "alphabet".into(),
|
||||
},
|
||||
KnowledgeCandidate {
|
||||
slug: "beta".into(),
|
||||
},
|
||||
]);
|
||||
let all = state.list_knowledge_completions("");
|
||||
assert_eq!(all.len(), 3);
|
||||
let alpha = state.list_knowledge_completions("alpha");
|
||||
assert_eq!(
|
||||
alpha.iter().map(|c| c.slug.as_str()).collect::<Vec<_>>(),
|
||||
vec!["alpha", "alphabet"]
|
||||
);
|
||||
assert!(state.list_knowledge_completions("zzz").is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,165 @@
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use protocol::WorkerStatus;
|
||||
use ticket::config::TicketRole;
|
||||
|
||||
use crate::hook::{Hook, HookPostToolAction, PostToolCall, ToolResultSummary};
|
||||
|
||||
const TICKET_INTAKE_READY_TOOL_NAME: &str = "TicketIntakeReady";
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct ShutdownAfterIdleRequest {
|
||||
requested: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl ShutdownAfterIdleRequest {
|
||||
pub(crate) fn request(&self) {
|
||||
self.requested.store(true, Ordering::Release);
|
||||
}
|
||||
|
||||
pub(crate) fn take(&self) -> bool {
|
||||
self.requested.swap(false, Ordering::AcqRel)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn is_requested(&self) -> bool {
|
||||
self.requested.load(Ordering::Acquire)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_ticket_intake_role(role: Option<&str>) -> bool {
|
||||
matches!(role.and_then(TicketRole::parse), Some(TicketRole::Intake))
|
||||
}
|
||||
|
||||
pub(crate) fn take_shutdown_request_after_status(
|
||||
shutdown_after_idle: &ShutdownAfterIdleRequest,
|
||||
status: WorkerStatus,
|
||||
) -> bool {
|
||||
status == WorkerStatus::Idle && shutdown_after_idle.take()
|
||||
}
|
||||
|
||||
pub(crate) struct TicketIntakeReadyShutdownHook {
|
||||
shutdown_after_idle: ShutdownAfterIdleRequest,
|
||||
eligible_ticket_intake_role: bool,
|
||||
}
|
||||
|
||||
impl TicketIntakeReadyShutdownHook {
|
||||
pub(crate) fn new(
|
||||
shutdown_after_idle: ShutdownAfterIdleRequest,
|
||||
eligible_ticket_intake_role: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
shutdown_after_idle,
|
||||
eligible_ticket_intake_role,
|
||||
}
|
||||
}
|
||||
|
||||
fn observe_tool_result(&self, info: &ToolResultSummary) {
|
||||
if self.eligible_ticket_intake_role
|
||||
&& info.tool_name == TICKET_INTAKE_READY_TOOL_NAME
|
||||
&& !info.is_error
|
||||
{
|
||||
self.shutdown_after_idle.request();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PostToolCall> for TicketIntakeReadyShutdownHook {
|
||||
async fn call(&self, info: &ToolResultSummary) -> HookPostToolAction {
|
||||
self.observe_tool_result(info);
|
||||
HookPostToolAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use llm_engine::tool::ToolOutput;
|
||||
|
||||
fn tool_result(name: &str, is_error: bool) -> ToolResultSummary {
|
||||
ToolResultSummary {
|
||||
call_id: "tool-1".to_string(),
|
||||
tool_name: name.to_string(),
|
||||
is_error,
|
||||
output: ToolOutput {
|
||||
summary: "result".to_string(),
|
||||
content: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn successful_ticket_intake_ready_schedules_shutdown_after_idle_for_intake_role() {
|
||||
let request = ShutdownAfterIdleRequest::default();
|
||||
let hook = TicketIntakeReadyShutdownHook::new(request.clone(), true);
|
||||
|
||||
hook.observe_tool_result(&tool_result(TICKET_INTAKE_READY_TOOL_NAME, false));
|
||||
|
||||
assert!(request.is_requested());
|
||||
assert!(request.take());
|
||||
assert!(!request.is_requested());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_ticket_intake_ready_does_not_schedule_shutdown_after_idle() {
|
||||
let request = ShutdownAfterIdleRequest::default();
|
||||
let hook = TicketIntakeReadyShutdownHook::new(request.clone(), true);
|
||||
|
||||
hook.observe_tool_result(&tool_result(TICKET_INTAKE_READY_TOOL_NAME, true));
|
||||
|
||||
assert!(!request.is_requested());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_intake_role_does_not_schedule_shutdown_after_idle() {
|
||||
let request = ShutdownAfterIdleRequest::default();
|
||||
let hook = TicketIntakeReadyShutdownHook::new(request.clone(), false);
|
||||
|
||||
hook.observe_tool_result(&tool_result(TICKET_INTAKE_READY_TOOL_NAME, false));
|
||||
|
||||
assert!(!request.is_requested());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn other_successful_tools_do_not_schedule_shutdown_after_idle() {
|
||||
let request = ShutdownAfterIdleRequest::default();
|
||||
let hook = TicketIntakeReadyShutdownHook::new(request.clone(), true);
|
||||
|
||||
hook.observe_tool_result(&tool_result("TicketShow", false));
|
||||
|
||||
assert!(!request.is_requested());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_ticket_intake_runtime_role_is_eligible() {
|
||||
assert!(is_ticket_intake_role(Some("intake")));
|
||||
assert!(!is_ticket_intake_role(Some("orchestrator")));
|
||||
assert!(!is_ticket_intake_role(Some("coder")));
|
||||
assert!(!is_ticket_intake_role(Some("reviewer")));
|
||||
assert!(!is_ticket_intake_role(Some("unknown")));
|
||||
assert!(!is_ticket_intake_role(None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shutdown_after_idle_is_taken_only_after_idle_status() {
|
||||
let request = ShutdownAfterIdleRequest::default();
|
||||
request.request();
|
||||
|
||||
assert!(!take_shutdown_request_after_status(
|
||||
&request,
|
||||
WorkerStatus::Running
|
||||
));
|
||||
assert!(request.is_requested());
|
||||
|
||||
assert!(take_shutdown_request_after_status(
|
||||
&request,
|
||||
WorkerStatus::Idle
|
||||
));
|
||||
assert!(!request.is_requested());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,704 @@
|
||||
//! Worker-to-Worker communication tools.
|
||||
//!
|
||||
//! Three tools in one module: `SendToWorker`, `ReadWorkerOutput`, `StopWorker`,
|
||||
//! all built on the same `SpawnedWorkerRegistry` handed in by
|
||||
//! the controller. Each operation is request-response: connect to the
|
||||
//! target's Unix socket, perform one method exchange, disconnect.
|
||||
//!
|
||||
//! These tools only touch Workers listed in the spawner's
|
||||
//! `SpawnedWorkerRegistry`; there is no machine-wide directory lookup, so
|
||||
//! the spawner can only reach its own descendants.
|
||||
|
||||
use std::path::Path;
|
||||
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::{ErrorCode, Event, InvokeKind, Method};
|
||||
use serde::Deserialize;
|
||||
use session_store::LogEntry;
|
||||
use tokio::net::UnixStream;
|
||||
|
||||
use crate::runtime::dir::SpawnedWorkerRecord;
|
||||
use crate::spawn::registry::SpawnedWorkerRegistry;
|
||||
|
||||
/// Timeout applied to each socket-level operation — connect, write,
|
||||
/// read. Kept short so a stuck child doesn't block the spawner's turn.
|
||||
const SOCKET_OP_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Shared input types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct NameInput {
|
||||
/// Name of a previously spawned Worker.
|
||||
name: String,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SendToWorker
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const SEND_TO_POD_DESCRIPTION: &str = "Send a text message to a previously spawned Worker. The spawned Worker \
|
||||
processes it as a user turn. Fails if the Worker is already executing a \
|
||||
turn — retry after it finishes. Does not wait for the turn to complete; \
|
||||
use `ReadWorkerOutput` to fetch results afterwards.";
|
||||
|
||||
#[derive(Debug, Deserialize, schemars::JsonSchema)]
|
||||
struct SendToWorkerInput {
|
||||
/// Target Worker name.
|
||||
name: String,
|
||||
/// Text delivered to the Worker as the next user message.
|
||||
message: String,
|
||||
}
|
||||
|
||||
struct SendToWorkerTool {
|
||||
registry: Arc<SpawnedWorkerRegistry>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SendToWorkerTool {
|
||||
async fn execute(
|
||||
&self,
|
||||
input_json: &str,
|
||||
_ctx: llm_engine::tool::ToolExecutionContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let input: SendToWorkerInput = serde_json::from_str(input_json)
|
||||
.map_err(|e| ToolError::InvalidArgument(format!("invalid SendToWorker input: {e}")))?;
|
||||
let record = self
|
||||
.registry
|
||||
.get(&input.name)
|
||||
.await
|
||||
.ok_or_else(|| unknown_worker_err(&input.name))?;
|
||||
|
||||
send_run_and_confirm(&record.socket_path, input.message)
|
||||
.await
|
||||
.map_err(|e| match e {
|
||||
SendRunError::AlreadyRunning => ToolError::ExecutionFailed(format!(
|
||||
"worker `{}` is already running a turn; wait for it to finish and retry",
|
||||
input.name
|
||||
)),
|
||||
SendRunError::Rejected { code, message } => ToolError::ExecutionFailed(format!(
|
||||
"worker `{}` rejected the run with {code:?}: {message}",
|
||||
input.name
|
||||
)),
|
||||
SendRunError::Io(msg) => {
|
||||
ToolError::ExecutionFailed(format!("send to `{}`: {msg}", input.name))
|
||||
}
|
||||
})?;
|
||||
|
||||
Ok(ToolOutput {
|
||||
summary: format!("sent message to `{}`", input.name),
|
||||
content: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn send_to_worker_tool(registry: Arc<SpawnedWorkerRegistry>) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = schemars::schema_for!(SendToWorkerInput);
|
||||
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
|
||||
let meta = ToolMeta::new("SendToWorker")
|
||||
.description(SEND_TO_POD_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(SendToWorkerTool {
|
||||
registry: registry.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ReadWorkerOutput
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const READ_POD_OUTPUT_DESCRIPTION: &str = "Fetch new assistant text from a spawned Worker since the last read. \
|
||||
Uses an internal cursor per-Worker so consecutive calls return only \
|
||||
newly-produced output. Returns the Worker's current status and the new \
|
||||
text, or reports `stopped` if the Worker can no longer be reached.";
|
||||
|
||||
struct ReadWorkerOutputTool {
|
||||
registry: Arc<SpawnedWorkerRegistry>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ReadWorkerOutputTool {
|
||||
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 ReadWorkerOutput input: {e}"))
|
||||
})?;
|
||||
let record = self
|
||||
.registry
|
||||
.get(&input.name)
|
||||
.await
|
||||
.ok_or_else(|| unknown_worker_err(&input.name))?;
|
||||
|
||||
let items = match fetch_history(&record.socket_path).await {
|
||||
Ok(items) => items,
|
||||
Err(_) => {
|
||||
return Ok(ToolOutput {
|
||||
summary: format!("worker `{}` is stopped (unreachable)", input.name),
|
||||
content: None,
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
let cursor = self.registry.cursor(&input.name).await;
|
||||
let new_items = if cursor >= items.len() {
|
||||
&[] as &[serde_json::Value]
|
||||
} else {
|
||||
&items[cursor..]
|
||||
};
|
||||
let new_text = extract_assistant_text(new_items);
|
||||
self.registry.set_cursor(&input.name, items.len()).await;
|
||||
|
||||
let summary = if new_text.is_empty() {
|
||||
format!("worker `{}` running; no new assistant text", input.name)
|
||||
} else {
|
||||
let lines = new_text.lines().count();
|
||||
format!(
|
||||
"worker `{}`: {lines} new line(s) of assistant text",
|
||||
input.name
|
||||
)
|
||||
};
|
||||
let content = if new_text.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(new_text)
|
||||
};
|
||||
Ok(ToolOutput { summary, content })
|
||||
}
|
||||
}
|
||||
|
||||
pub fn read_worker_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("ReadWorkerOutput")
|
||||
.description(READ_POD_OUTPUT_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(ReadWorkerOutputTool {
|
||||
registry: registry.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// StopWorker
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const STOP_POD_DESCRIPTION: &str = "Terminate a spawned Worker and reclaim the delegated scope. The Worker \
|
||||
receives `Shutdown`; its scope entry is released in the machine-wide \
|
||||
registry so the spawner can spawn a new Worker over the same paths.";
|
||||
|
||||
struct StopWorkerTool {
|
||||
registry: Arc<SpawnedWorkerRegistry>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for StopWorkerTool {
|
||||
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 StopWorker input: {e}")))?;
|
||||
let record = self
|
||||
.registry
|
||||
.get(&input.name)
|
||||
.await
|
||||
.ok_or_else(|| unknown_worker_err(&input.name))?;
|
||||
|
||||
// Best-effort Shutdown. The child's own `ScopeAllocationGuard`
|
||||
// releases its entry on clean exit; the parent reclaim below is the
|
||||
// authoritative operation for removing the child record and returning
|
||||
// delegated Write scope to the spawner.
|
||||
let _ = connect_and_send(&record.socket_path, &Method::Shutdown).await;
|
||||
|
||||
let scope_summary = summarize_scope(&record);
|
||||
|
||||
self.registry
|
||||
.remove(&record.worker_name)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
ToolError::ExecutionFailed(format!("update spawned worker registry: {e}"))
|
||||
})?;
|
||||
|
||||
Ok(ToolOutput {
|
||||
summary: format!(
|
||||
"stopped worker `{}`; reclaimed scope: {scope_summary}",
|
||||
record.worker_name
|
||||
),
|
||||
content: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn stop_worker_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("StopWorker")
|
||||
.description(STOP_POD_DESCRIPTION)
|
||||
.input_schema(schema_value);
|
||||
let tool: Arc<dyn Tool> = Arc::new(StopWorkerTool {
|
||||
registry: registry.clone(),
|
||||
});
|
||||
(meta, tool)
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn unknown_worker_err(name: &str) -> ToolError {
|
||||
ToolError::InvalidArgument(format!("no spawned worker named `{name}`"))
|
||||
}
|
||||
|
||||
fn summarize_scope(record: &SpawnedWorkerRecord) -> String {
|
||||
if record.scope_delegated.is_empty() {
|
||||
return "(none)".into();
|
||||
}
|
||||
let parts: Vec<String> = record
|
||||
.scope_delegated
|
||||
.iter()
|
||||
.map(|rule| {
|
||||
let perm = match rule.permission {
|
||||
manifest::Permission::Read => "read",
|
||||
manifest::Permission::Write => "write",
|
||||
};
|
||||
let recursive = if rule.recursive {
|
||||
""
|
||||
} else {
|
||||
" [non-recursive]"
|
||||
};
|
||||
format!("{perm}:{}{recursive}", rule.target.display())
|
||||
})
|
||||
.collect();
|
||||
parts.join(", ")
|
||||
}
|
||||
|
||||
/// Connect with a timeout, drain the server's connect-time snapshot,
|
||||
/// write one `Method` line, flush, and close.
|
||||
///
|
||||
/// The Worker socket protocol sends replayed alerts and an initial
|
||||
/// `Event::Snapshot` before it starts reading client methods. Send-only
|
||||
/// callers must consume that prefix; otherwise a large snapshot can block
|
||||
/// the server's writer before it reaches the method-read branch. Any
|
||||
/// socket error maps to an `io::Error`; the caller decides whether to
|
||||
/// surface it to the LLM or treat it as "worker stopped".
|
||||
pub(crate) async fn connect_and_send(socket: &Path, method: &Method) -> std::io::Result<()> {
|
||||
let stream = tokio::time::timeout(SOCKET_OP_TIMEOUT, UnixStream::connect(socket))
|
||||
.await
|
||||
.map_err(|_| std::io::Error::new(std::io::ErrorKind::TimedOut, "connect timed out"))??;
|
||||
let (r, w) = stream.into_split();
|
||||
let mut reader = JsonLineReader::new(r);
|
||||
let mut writer = JsonLineWriter::new(w);
|
||||
|
||||
drain_initial_snapshot(&mut reader).await?;
|
||||
|
||||
tokio::time::timeout(SOCKET_OP_TIMEOUT, writer.write(method))
|
||||
.await
|
||||
.map_err(|_| std::io::Error::new(std::io::ErrorKind::TimedOut, "write timed out"))??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drain_initial_snapshot<R>(reader: &mut JsonLineReader<R>) -> std::io::Result<()>
|
||||
where
|
||||
R: tokio::io::AsyncBufRead + Unpin,
|
||||
{
|
||||
loop {
|
||||
let event = tokio::time::timeout(SOCKET_OP_TIMEOUT, reader.next::<Event>())
|
||||
.await
|
||||
.map_err(|_| std::io::Error::new(std::io::ErrorKind::TimedOut, "read timed out"))??;
|
||||
match event {
|
||||
Some(Event::Snapshot { .. }) => return Ok(()),
|
||||
Some(_) => continue,
|
||||
None => {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::UnexpectedEof,
|
||||
"worker closed connection before Snapshot event",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Failure modes distinguished by `SendToWorker`.
|
||||
#[derive(Debug)]
|
||||
pub(crate) enum SendRunError {
|
||||
/// Target Worker responded with `Error { AlreadyRunning }` — the
|
||||
/// caller can retry once the current turn ends.
|
||||
AlreadyRunning,
|
||||
/// Target Worker explicitly rejected the run after delivery reached the
|
||||
/// controller.
|
||||
Rejected { code: ErrorCode, message: String },
|
||||
/// Transport, protocol, timeout, or unexpected EOF before acceptance
|
||||
/// evidence was observed.
|
||||
Io(String),
|
||||
}
|
||||
|
||||
/// Write `Method::Run` to the target and read back events until we see
|
||||
/// evidence that the controller accepted the run (`UserMessage`,
|
||||
/// `TurnStart`, or a user-send `InvokeStart`) or rejected it. The connect-time
|
||||
/// event prelude is drained before sending the method so large Snapshots and
|
||||
/// large Run payloads cannot block each other on the same socket. Times out
|
||||
/// per operation so a stuck Worker doesn't hang the tool.
|
||||
pub(crate) async fn send_run_and_confirm(socket: &Path, input: String) -> Result<(), SendRunError> {
|
||||
let stream = tokio::time::timeout(SOCKET_OP_TIMEOUT, UnixStream::connect(socket))
|
||||
.await
|
||||
.map_err(|_| SendRunError::Io("connect timed out".into()))?
|
||||
.map_err(|e| SendRunError::Io(format!("connect: {e}")))?;
|
||||
let (r, w) = stream.into_split();
|
||||
let mut writer = JsonLineWriter::new(w);
|
||||
let mut reader = JsonLineReader::new(r);
|
||||
|
||||
loop {
|
||||
let event = tokio::time::timeout(SOCKET_OP_TIMEOUT, reader.next::<Event>())
|
||||
.await
|
||||
.map_err(|_| SendRunError::Io("read initial Snapshot timed out".into()))?
|
||||
.map_err(|e| SendRunError::Io(format!("read initial Snapshot: {e}")))?;
|
||||
match event {
|
||||
Some(Event::Snapshot { .. }) => break,
|
||||
Some(Event::Alert(_)) => continue,
|
||||
Some(Event::Error {
|
||||
code: ErrorCode::AlreadyRunning,
|
||||
..
|
||||
}) => return Err(SendRunError::AlreadyRunning),
|
||||
Some(Event::Error { code, message }) => {
|
||||
return Err(SendRunError::Rejected { code, message });
|
||||
}
|
||||
Some(_) => continue,
|
||||
None => {
|
||||
return Err(SendRunError::Io(
|
||||
"connection closed before initial Snapshot".into(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tokio::time::timeout(
|
||||
SOCKET_OP_TIMEOUT,
|
||||
writer.write(&Method::Run {
|
||||
input: vec![protocol::Segment::text(input)],
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| SendRunError::Io("write timed out".into()))?
|
||||
.map_err(|e| SendRunError::Io(format!("write: {e}")))?;
|
||||
loop {
|
||||
let event = tokio::time::timeout(SOCKET_OP_TIMEOUT, reader.next::<Event>())
|
||||
.await
|
||||
.map_err(|_| SendRunError::Io("read response timed out".into()))?
|
||||
.map_err(|e| SendRunError::Io(format!("read response: {e}")))?;
|
||||
match event {
|
||||
Some(Event::Error {
|
||||
code: ErrorCode::AlreadyRunning,
|
||||
..
|
||||
}) => return Err(SendRunError::AlreadyRunning),
|
||||
Some(Event::Error { code, message }) => {
|
||||
return Err(SendRunError::Rejected { code, message });
|
||||
}
|
||||
Some(Event::InvokeStart {
|
||||
kind: InvokeKind::UserSend,
|
||||
})
|
||||
| Some(Event::UserMessage { .. })
|
||||
| Some(Event::TurnStart { .. }) => return Ok(()),
|
||||
// Other post-Snapshot events can race with the controller's
|
||||
// response; keep reading until the Run is accepted or rejected.
|
||||
Some(_) => continue,
|
||||
None => return Err(SendRunError::Io("connection closed before response".into())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Connect to a Worker's socket and read the connect-time `Event::Snapshot`.
|
||||
///
|
||||
/// Workers deliver the session-log mirror as the first non-Alert event on
|
||||
/// every new connection, so consuming it is sufficient — no explicit
|
||||
/// `GetHistory` method round trip. Returns the entries as raw JSON
|
||||
/// values; callers deserialize as `session_store::LogEntry` if they
|
||||
/// need typed access.
|
||||
async fn fetch_history(socket: &Path) -> std::io::Result<Vec<serde_json::Value>> {
|
||||
let stream = tokio::time::timeout(SOCKET_OP_TIMEOUT, UnixStream::connect(socket))
|
||||
.await
|
||||
.map_err(|_| std::io::Error::new(std::io::ErrorKind::TimedOut, "connect timed out"))??;
|
||||
let (r, _w) = stream.into_split();
|
||||
let mut reader = JsonLineReader::new(r);
|
||||
|
||||
loop {
|
||||
let event = tokio::time::timeout(SOCKET_OP_TIMEOUT, reader.next::<Event>())
|
||||
.await
|
||||
.map_err(|_| std::io::Error::new(std::io::ErrorKind::TimedOut, "read timed out"))??;
|
||||
match event {
|
||||
Some(Event::Snapshot { entries, .. }) => return Ok(entries),
|
||||
Some(_) => continue,
|
||||
None => {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::UnexpectedEof,
|
||||
"worker closed connection before Snapshot event",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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::*;
|
||||
|
||||
use protocol::{Alert, AlertLevel, AlertSource, Greeting, WorkerEvent, WorkerStatus};
|
||||
use tempfile::TempDir;
|
||||
use tokio::net::UnixListener;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
fn snapshot(entries: Vec<serde_json::Value>) -> Event {
|
||||
Event::Snapshot {
|
||||
entries,
|
||||
greeting: Greeting {
|
||||
worker_name: "server".into(),
|
||||
cwd: "/tmp".into(),
|
||||
provider: "test".into(),
|
||||
model: "test".into(),
|
||||
scope_summary: String::new(),
|
||||
tools: Vec::new(),
|
||||
context_window: 200_000,
|
||||
context_tokens: 0,
|
||||
},
|
||||
status: WorkerStatus::Idle,
|
||||
in_flight: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn serve_initial_events_then_method(
|
||||
listener: UnixListener,
|
||||
events: Vec<Event>,
|
||||
) -> JoinHandle<Option<Method>> {
|
||||
tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.ok()?;
|
||||
let (r, w) = stream.into_split();
|
||||
let mut reader = JsonLineReader::new(r);
|
||||
let mut writer = JsonLineWriter::new(w);
|
||||
for event in events {
|
||||
writer.write(&event).await.ok()?;
|
||||
}
|
||||
reader.next::<Method>().await.ok().flatten()
|
||||
})
|
||||
}
|
||||
|
||||
fn serve_initial_events_then_run_ack(
|
||||
listener: UnixListener,
|
||||
initial_events: Vec<Event>,
|
||||
ack: Event,
|
||||
) -> JoinHandle<Option<Method>> {
|
||||
tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.ok()?;
|
||||
let (r, w) = stream.into_split();
|
||||
let mut reader = JsonLineReader::new(r);
|
||||
let mut writer = JsonLineWriter::new(w);
|
||||
for event in initial_events {
|
||||
writer.write(&event).await.ok()?;
|
||||
}
|
||||
let method = reader.next::<Method>().await.ok().flatten()?;
|
||||
writer.write(&ack).await.ok()?;
|
||||
Some(method)
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_run_and_confirm_keeps_connection_open_until_user_message_ack() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let socket = tmp.path().join("worker.sock");
|
||||
let listener = UnixListener::bind(&socket).unwrap();
|
||||
let received = serve_initial_events_then_run_ack(
|
||||
listener,
|
||||
vec![
|
||||
Event::Alert(Alert {
|
||||
level: AlertLevel::Warn,
|
||||
source: AlertSource::Worker,
|
||||
message: "replayed alert".into(),
|
||||
timestamp_ms: 0,
|
||||
}),
|
||||
snapshot(Vec::new()),
|
||||
],
|
||||
Event::UserMessage {
|
||||
segments: vec![protocol::Segment::text("hello")],
|
||||
},
|
||||
);
|
||||
|
||||
send_run_and_confirm(&socket, "hello".into()).await.unwrap();
|
||||
|
||||
let method = received.await.unwrap().expect("expected method");
|
||||
match method {
|
||||
Method::Run { input } => {
|
||||
assert_eq!(protocol::Segment::flatten_to_text(&input), "hello");
|
||||
}
|
||||
other => panic!("expected Run, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_run_and_confirm_drains_alert_and_large_snapshot_before_large_run() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let socket = tmp.path().join("worker.sock");
|
||||
let listener = UnixListener::bind(&socket).unwrap();
|
||||
let large_snapshot_payload = "s".repeat(2 * 1024 * 1024);
|
||||
let large_run_payload = "r".repeat(2 * 1024 * 1024);
|
||||
let received = serve_initial_events_then_run_ack(
|
||||
listener,
|
||||
vec![
|
||||
Event::Alert(Alert {
|
||||
level: AlertLevel::Warn,
|
||||
source: AlertSource::Worker,
|
||||
message: "replayed alert".into(),
|
||||
timestamp_ms: 0,
|
||||
}),
|
||||
snapshot(vec![
|
||||
serde_json::json!({ "payload": large_snapshot_payload }),
|
||||
]),
|
||||
],
|
||||
Event::InvokeStart {
|
||||
kind: InvokeKind::UserSend,
|
||||
},
|
||||
);
|
||||
|
||||
send_run_and_confirm(&socket, large_run_payload.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let method = received.await.unwrap().expect("expected method");
|
||||
match method {
|
||||
Method::Run { input } => {
|
||||
assert_eq!(
|
||||
protocol::Segment::flatten_to_text(&input),
|
||||
large_run_payload
|
||||
);
|
||||
}
|
||||
other => panic!("expected Run, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_run_and_confirm_reports_already_running() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let socket = tmp.path().join("worker.sock");
|
||||
let listener = UnixListener::bind(&socket).unwrap();
|
||||
let received = serve_initial_events_then_run_ack(
|
||||
listener,
|
||||
vec![snapshot(Vec::new())],
|
||||
Event::Error {
|
||||
code: ErrorCode::AlreadyRunning,
|
||||
message: "busy".into(),
|
||||
},
|
||||
);
|
||||
|
||||
let err = send_run_and_confirm(&socket, "hello".into())
|
||||
.await
|
||||
.expect_err("expected AlreadyRunning");
|
||||
assert!(matches!(err, SendRunError::AlreadyRunning));
|
||||
assert!(matches!(received.await.unwrap(), Some(Method::Run { .. })));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connect_and_send_drains_initial_alert_and_snapshot_before_method() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let socket = tmp.path().join("worker.sock");
|
||||
let listener = UnixListener::bind(&socket).unwrap();
|
||||
let received = serve_initial_events_then_method(
|
||||
listener,
|
||||
vec![
|
||||
Event::Alert(Alert {
|
||||
level: AlertLevel::Warn,
|
||||
source: AlertSource::Worker,
|
||||
message: "replayed alert".into(),
|
||||
timestamp_ms: 0,
|
||||
}),
|
||||
snapshot(Vec::new()),
|
||||
],
|
||||
);
|
||||
|
||||
connect_and_send(&socket, &Method::Shutdown).await.unwrap();
|
||||
|
||||
let method = received.await.unwrap().expect("expected method");
|
||||
assert!(matches!(method, Method::Shutdown));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connect_and_send_delivers_method_after_large_initial_snapshot() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let socket = tmp.path().join("worker.sock");
|
||||
let listener = UnixListener::bind(&socket).unwrap();
|
||||
let large_payload = "x".repeat(2 * 1024 * 1024);
|
||||
let received = serve_initial_events_then_method(
|
||||
listener,
|
||||
vec![snapshot(vec![
|
||||
serde_json::json!({ "payload": large_payload }),
|
||||
])],
|
||||
);
|
||||
let expected = Method::WorkerEvent(WorkerEvent::TurnEnded {
|
||||
worker_name: "child".into(),
|
||||
});
|
||||
|
||||
connect_and_send(&socket, &expected).await.unwrap();
|
||||
|
||||
let method = received.await.unwrap().expect("expected method");
|
||||
match method {
|
||||
Method::WorkerEvent(WorkerEvent::TurnEnded { worker_name }) => {
|
||||
assert_eq!(worker_name, "child")
|
||||
}
|
||||
other => panic!("expected TurnEnded WorkerEvent, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
pub mod comm_tools;
|
||||
pub mod registry;
|
||||
pub mod tool;
|
||||
@@ -0,0 +1,451 @@
|
||||
//! Shared registry of Workers spawned by this Worker.
|
||||
//!
|
||||
//! `SpawnWorker` writes here; the worker-comm tools (`SendToWorker`,
|
||||
//! `ReadWorkerOutput`, `StopWorker`) read and mutate the same instance. Discovery
|
||||
//! tools consult this registry together with durable Worker state. Runtime
|
||||
//! write-through still materialises `spawned_workers.json`, but durable state lives
|
||||
//! in the spawner's Worker metadata.
|
||||
//!
|
||||
//! `ReadWorkerOutput` additionally owns a per-spawned-worker cursor here so
|
||||
//! two consecutive reads yield only new assistant text. The cursor is
|
||||
//! an item-index into the child's history; push-only history makes
|
||||
//! index stable across reads.
|
||||
//!
|
||||
//! Cursors intentionally do not persist; a restored registry starts with
|
||||
//! fresh read positions.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use manifest::{Permission, ScopeRule, SharedScope};
|
||||
use pod_store::{
|
||||
WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerSpawnedScopeRule,
|
||||
WorkerStoreError,
|
||||
};
|
||||
use tokio::net::UnixStream;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
|
||||
use crate::runtime::pod_registry;
|
||||
|
||||
type RegistryStateWriter = Arc<dyn Fn(&[SpawnedWorkerRecord]) -> io::Result<()> + Send + Sync>;
|
||||
type RegistryReclaimWriter = Arc<dyn Fn(&SpawnedWorkerRecord) -> io::Result<()> + Send + Sync>;
|
||||
|
||||
const RESTORE_REACHABILITY_TIMEOUT: Duration = Duration::from_millis(500);
|
||||
const REGISTRY_CLEANUP_TIMEOUT: Duration = Duration::from_secs(15);
|
||||
|
||||
pub struct SpawnedWorkerRegistry {
|
||||
records: Mutex<Vec<SpawnedWorkerRecord>>,
|
||||
cursors: Mutex<HashMap<String, usize>>,
|
||||
mutations: Mutex<()>,
|
||||
runtime_dir: Arc<RuntimeDir>,
|
||||
state_writer: Option<RegistryStateWriter>,
|
||||
reclaim_writer: Option<RegistryReclaimWriter>,
|
||||
parent_name: Option<String>,
|
||||
parent_scope: Option<SharedScope>,
|
||||
}
|
||||
|
||||
pub struct SpawnedWorkerRegistryLoad {
|
||||
pub registry: Arc<SpawnedWorkerRegistry>,
|
||||
pub reclaimed_unreachable: bool,
|
||||
}
|
||||
|
||||
impl SpawnedWorkerRegistry {
|
||||
pub fn new(runtime_dir: Arc<RuntimeDir>) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
records: Mutex::new(Vec::new()),
|
||||
cursors: Mutex::new(HashMap::new()),
|
||||
mutations: Mutex::new(()),
|
||||
runtime_dir,
|
||||
state_writer: None,
|
||||
reclaim_writer: None,
|
||||
parent_name: None,
|
||||
parent_scope: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a registry from the spawner's durable Worker state, pruning child
|
||||
/// records whose socket path is already gone. The surviving list is
|
||||
/// written through to both `spawned_workers.json` and Worker state so runtime
|
||||
/// and durable views start aligned.
|
||||
pub async fn load_from_worker_state<St>(
|
||||
runtime_dir: Arc<RuntimeDir>,
|
||||
store: St,
|
||||
worker_name: String,
|
||||
) -> io::Result<Arc<Self>>
|
||||
where
|
||||
St: WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
let loaded =
|
||||
Self::load_from_worker_state_with_reclaim(runtime_dir, store, worker_name, None)
|
||||
.await?;
|
||||
Ok(loaded.registry)
|
||||
}
|
||||
|
||||
pub async fn load_from_worker_state_with_reclaim<St>(
|
||||
runtime_dir: Arc<RuntimeDir>,
|
||||
store: St,
|
||||
worker_name: String,
|
||||
parent_scope: Option<SharedScope>,
|
||||
) -> io::Result<SpawnedWorkerRegistryLoad>
|
||||
where
|
||||
St: WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
let metadata = store
|
||||
.read_by_name(&worker_name)
|
||||
.map_err(store_error_to_io)?;
|
||||
let persisted_children = metadata
|
||||
.as_ref()
|
||||
.map(|m| m.spawned_children.clone())
|
||||
.unwrap_or_default();
|
||||
|
||||
let mut records = Vec::with_capacity(persisted_children.len());
|
||||
let mut pruned_records = Vec::new();
|
||||
for child in &persisted_children {
|
||||
let record = match record_from_worker_state(child) {
|
||||
Ok(record) => record,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
error = %err,
|
||||
worker = %child.worker_name,
|
||||
"dropping corrupt persisted spawned-worker record"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if is_reachable(&record.socket_path).await {
|
||||
records.push(record);
|
||||
} else {
|
||||
warn!(
|
||||
worker = %record.worker_name,
|
||||
socket = %record.socket_path.display(),
|
||||
"dropping unreachable persisted spawned-worker record"
|
||||
);
|
||||
pruned_records.push(record);
|
||||
}
|
||||
}
|
||||
|
||||
runtime_dir.write_spawned_workers(&records).await?;
|
||||
let state_writer = worker_state_writer(store.clone(), worker_name.clone());
|
||||
let reclaim_writer = worker_state_reclaim_writer(store.clone(), worker_name.clone());
|
||||
if metadata.is_none() {
|
||||
state_writer(&records)?;
|
||||
}
|
||||
|
||||
let mut reclaimed_unreachable = false;
|
||||
if !pruned_records.is_empty() {
|
||||
let reclaimed = pruned_records
|
||||
.iter()
|
||||
.map(|record| WorkerReclaimedChild {
|
||||
worker_name: record.worker_name.clone(),
|
||||
scope_delegated: record
|
||||
.scope_delegated
|
||||
.iter()
|
||||
.map(|rule| WorkerSpawnedScopeRule {
|
||||
target: rule.target.clone(),
|
||||
permission: match rule.permission {
|
||||
Permission::Read => "read".to_string(),
|
||||
Permission::Write => "write".to_string(),
|
||||
},
|
||||
recursive: rule.recursive,
|
||||
})
|
||||
.collect(),
|
||||
})
|
||||
.collect();
|
||||
store
|
||||
.reclaim_spawned_children(&worker_name, reclaimed)
|
||||
.map_err(store_error_to_io)?;
|
||||
reclaimed_unreachable = true;
|
||||
}
|
||||
if parent_scope.is_some() {
|
||||
for record in &pruned_records {
|
||||
reclaim_record(&worker_name, parent_scope.as_ref(), record)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(SpawnedWorkerRegistryLoad {
|
||||
registry: Arc::new(Self {
|
||||
records: Mutex::new(records),
|
||||
cursors: Mutex::new(HashMap::new()),
|
||||
mutations: Mutex::new(()),
|
||||
runtime_dir,
|
||||
state_writer: Some(state_writer),
|
||||
reclaim_writer: Some(reclaim_writer),
|
||||
parent_name: Some(worker_name),
|
||||
parent_scope,
|
||||
}),
|
||||
reclaimed_unreachable,
|
||||
})
|
||||
}
|
||||
|
||||
/// Append a new record and persist the full list. Returns an I/O
|
||||
/// error if either persisted write fails; the in-memory state is still
|
||||
/// updated in that case — the next successful write will reconcile.
|
||||
pub async fn add(&self, record: SpawnedWorkerRecord) -> io::Result<()> {
|
||||
let _mutation = self.mutations.lock().await;
|
||||
let snapshot = {
|
||||
let mut records = self.records.lock().await;
|
||||
records.push(record);
|
||||
records.clone()
|
||||
};
|
||||
self.persist_records(&snapshot).await
|
||||
}
|
||||
|
||||
/// Look up a record by worker name. Cloned so callers can drop the lock.
|
||||
pub async fn get(&self, worker_name: &str) -> Option<SpawnedWorkerRecord> {
|
||||
self.records
|
||||
.lock()
|
||||
.await
|
||||
.iter()
|
||||
.find(|r| r.worker_name == worker_name)
|
||||
.cloned()
|
||||
}
|
||||
|
||||
pub async fn list(&self) -> Vec<SpawnedWorkerRecord> {
|
||||
self.records.lock().await.clone()
|
||||
}
|
||||
|
||||
/// Remove the record for `worker_name`, persist, clear its cursor, and
|
||||
/// reclaim any delegated Write scope owned by that child. Returns the
|
||||
/// removed record (if any).
|
||||
pub async fn remove(&self, worker_name: &str) -> io::Result<Option<SpawnedWorkerRecord>> {
|
||||
let _mutation = self.mutations.lock().await;
|
||||
let (removed, snapshot) = {
|
||||
let mut records = self.records.lock().await;
|
||||
let idx = records.iter().position(|r| r.worker_name == worker_name);
|
||||
let removed = idx.map(|i| records.remove(i));
|
||||
let snapshot = records.clone();
|
||||
(removed, snapshot)
|
||||
};
|
||||
self.persist_records(&snapshot).await?;
|
||||
self.cursors.lock().await.remove(worker_name);
|
||||
if let Some(record) = &removed {
|
||||
self.reclaim_removed_record(record.clone()).await?;
|
||||
}
|
||||
Ok(removed)
|
||||
}
|
||||
|
||||
async fn reclaim_removed_record(&self, record: SpawnedWorkerRecord) -> io::Result<()> {
|
||||
let parent_name = self.parent_name.clone();
|
||||
let parent_scope = self.parent_scope.clone();
|
||||
let reclaim_writer = self.reclaim_writer.clone();
|
||||
let worker_name = record.worker_name.clone();
|
||||
let reclaim = tokio::task::spawn_blocking(move || {
|
||||
reclaim_removed_record_blocking(parent_name, parent_scope, reclaim_writer, record)
|
||||
});
|
||||
tokio::time::timeout(REGISTRY_CLEANUP_TIMEOUT, reclaim)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::TimedOut,
|
||||
format!("timed out reclaiming spawned worker `{worker_name}`"),
|
||||
)
|
||||
})?
|
||||
.map_err(|err| io::Error::other(format!("spawned-worker reclaim task failed: {err}")))?
|
||||
}
|
||||
|
||||
/// Read-only cursor lookup. Returns 0 when no cursor has been set.
|
||||
pub async fn cursor(&self, worker_name: &str) -> usize {
|
||||
self.cursors
|
||||
.lock()
|
||||
.await
|
||||
.get(worker_name)
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
pub async fn set_cursor(&self, worker_name: &str, cursor: usize) {
|
||||
self.cursors
|
||||
.lock()
|
||||
.await
|
||||
.insert(worker_name.to_string(), cursor);
|
||||
}
|
||||
|
||||
async fn persist_records(&self, records: &[SpawnedWorkerRecord]) -> io::Result<()> {
|
||||
self.runtime_dir.write_spawned_workers(records).await?;
|
||||
if let Some(write_state) = &self.state_writer {
|
||||
write_state(records)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn worker_state_writer<St>(store: St, worker_name: String) -> RegistryStateWriter
|
||||
where
|
||||
St: WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
Arc::new(move |records| {
|
||||
write_records_to_worker_state(&store, &worker_name, records).map_err(store_error_to_io)
|
||||
})
|
||||
}
|
||||
|
||||
fn worker_state_reclaim_writer<St>(store: St, worker_name: String) -> RegistryReclaimWriter
|
||||
where
|
||||
St: WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
{
|
||||
Arc::new(move |record| {
|
||||
let reclaimed = WorkerReclaimedChild {
|
||||
worker_name: record.worker_name.clone(),
|
||||
scope_delegated: record
|
||||
.scope_delegated
|
||||
.iter()
|
||||
.map(|rule| WorkerSpawnedScopeRule {
|
||||
target: rule.target.clone(),
|
||||
permission: match rule.permission {
|
||||
Permission::Read => "read".to_string(),
|
||||
Permission::Write => "write".to_string(),
|
||||
},
|
||||
recursive: rule.recursive,
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
store
|
||||
.reclaim_spawned_children(&worker_name, vec![reclaimed])
|
||||
.map(|_| ())
|
||||
.map_err(store_error_to_io)
|
||||
})
|
||||
}
|
||||
|
||||
fn reclaim_removed_record_blocking(
|
||||
parent_name: Option<String>,
|
||||
parent_scope: Option<SharedScope>,
|
||||
reclaim_writer: Option<RegistryReclaimWriter>,
|
||||
record: SpawnedWorkerRecord,
|
||||
) -> io::Result<()> {
|
||||
if let Some(parent_name) = parent_name {
|
||||
reclaim_record(&parent_name, parent_scope.as_ref(), &record)?;
|
||||
} else {
|
||||
release_child_allocation(&record.worker_name)?;
|
||||
}
|
||||
if let Some(write_reclaim) = reclaim_writer {
|
||||
write_reclaim(&record)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn reclaim_record(
|
||||
parent_name: &str,
|
||||
parent_scope: Option<&SharedScope>,
|
||||
record: &SpawnedWorkerRecord,
|
||||
) -> io::Result<()> {
|
||||
let write_rules = record
|
||||
.scope_delegated
|
||||
.iter()
|
||||
.filter(|rule| rule.permission == Permission::Write)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let lock_path = pod_registry::default_registry_path()
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::Other, err))?;
|
||||
let mut guard = pod_registry::LockFileGuard::open(&lock_path)
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::Other, err))?;
|
||||
pod_registry::reclaim_delegated_scope(
|
||||
&mut guard,
|
||||
parent_name,
|
||||
&record.worker_name,
|
||||
&record.scope_delegated,
|
||||
)
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::Other, err))?;
|
||||
|
||||
if let Some(scope) = parent_scope {
|
||||
scope
|
||||
.update(|current| current.with_removed_deny_rules(write_rules))
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn release_child_allocation(worker_name: &str) -> io::Result<()> {
|
||||
let lock_path = pod_registry::default_registry_path()
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::Other, err))?;
|
||||
let mut guard = pod_registry::LockFileGuard::open(&lock_path)
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::Other, err))?;
|
||||
match pod_registry::release_worker(&mut guard, worker_name) {
|
||||
Ok(()) | Err(pod_registry::ScopeLockError::UnknownWorker(_)) => Ok(()),
|
||||
Err(err) => Err(io::Error::new(io::ErrorKind::Other, err)),
|
||||
}
|
||||
}
|
||||
|
||||
fn write_records_to_worker_state<St>(
|
||||
store: &St,
|
||||
worker_name: &str,
|
||||
records: &[SpawnedWorkerRecord],
|
||||
) -> Result<(), WorkerStoreError>
|
||||
where
|
||||
St: WorkerMetadataStore,
|
||||
{
|
||||
let children = records
|
||||
.iter()
|
||||
.map(record_to_worker_state)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
store.set_spawned_children(worker_name, children)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn record_to_worker_state(
|
||||
record: &SpawnedWorkerRecord,
|
||||
) -> Result<WorkerSpawnedChild, serde_json::Error> {
|
||||
Ok(WorkerSpawnedChild {
|
||||
worker_name: record.worker_name.clone(),
|
||||
socket_path: record.socket_path.clone(),
|
||||
scope_delegated: record
|
||||
.scope_delegated
|
||||
.iter()
|
||||
.map(|rule| WorkerSpawnedScopeRule {
|
||||
target: rule.target.clone(),
|
||||
permission: match rule.permission {
|
||||
Permission::Read => "read".to_string(),
|
||||
Permission::Write => "write".to_string(),
|
||||
},
|
||||
recursive: rule.recursive,
|
||||
})
|
||||
.collect(),
|
||||
callback_address: record.callback_address.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn record_from_worker_state(
|
||||
child: &WorkerSpawnedChild,
|
||||
) -> Result<SpawnedWorkerRecord, serde_json::Error> {
|
||||
Ok(SpawnedWorkerRecord {
|
||||
worker_name: child.worker_name.clone(),
|
||||
socket_path: child.socket_path.clone(),
|
||||
scope_delegated: child
|
||||
.scope_delegated
|
||||
.iter()
|
||||
.map(|rule| {
|
||||
Ok(ScopeRule {
|
||||
target: rule.target.clone(),
|
||||
permission: match rule.permission.as_str() {
|
||||
"read" => Permission::Read,
|
||||
"write" => Permission::Write,
|
||||
other => {
|
||||
return Err(serde_json::Error::io(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
format!("invalid permission `{other}`"),
|
||||
)));
|
||||
}
|
||||
},
|
||||
recursive: rule.recursive,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?,
|
||||
callback_address: child.callback_address.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn store_error_to_io(error: WorkerStoreError) -> io::Error {
|
||||
io::Error::other(error)
|
||||
}
|
||||
|
||||
async fn is_reachable(socket: &Path) -> bool {
|
||||
tokio::time::timeout(RESTORE_REACHABILITY_TIMEOUT, UnixStream::connect(socket))
|
||||
.await
|
||||
.map(|result| result.is_ok())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,509 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use minijinja::Value as TemplateValue;
|
||||
use serde_json::Value;
|
||||
use std::collections::BTreeMap;
|
||||
use ticket::{LocalTicketBackend, TicketBackend, TicketIdOrSlug};
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use crate::discovery::{WeakNotifyDelivery, WorkerDiscovery};
|
||||
use crate::hook::{Hook, HookPostToolAction, PostToolCall, ToolResultSummary};
|
||||
use crate::prompt::catalog::{PromptCatalog, WorkerPrompt};
|
||||
use pod_store::WorkerMetadataStore;
|
||||
|
||||
const MAX_TITLE_CHARS: usize = 96;
|
||||
const MAX_SUMMARY_CHARS: usize = 160;
|
||||
const MAX_EVENT_KIND_CHARS: usize = 80;
|
||||
const MAX_MESSAGE_CHARS: usize = 768;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct TicketEventCompanionNotifyHook<
|
||||
St: WorkerMetadataStore + Clone + Send + Sync + 'static,
|
||||
> {
|
||||
backend: Arc<LocalTicketBackend>,
|
||||
discovery: WorkerDiscovery<St>,
|
||||
companion_worker_name: String,
|
||||
}
|
||||
|
||||
impl<St: WorkerMetadataStore + Clone + Send + Sync + 'static> TicketEventCompanionNotifyHook<St> {
|
||||
pub(crate) fn new(
|
||||
backend: LocalTicketBackend,
|
||||
discovery: WorkerDiscovery<St>,
|
||||
companion_worker_name: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
backend: Arc::new(backend),
|
||||
discovery,
|
||||
companion_worker_name: companion_worker_name.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<St: WorkerMetadataStore + Clone + Send + Sync + 'static> Hook<PostToolCall>
|
||||
for TicketEventCompanionNotifyHook<St>
|
||||
{
|
||||
async fn call(&self, summary: &ToolResultSummary) -> HookPostToolAction {
|
||||
let Some(notice) = build_ticket_event_notice(&self.backend, summary) else {
|
||||
return HookPostToolAction::Continue;
|
||||
};
|
||||
match self
|
||||
.discovery
|
||||
.ensure_existing_peer(&self.companion_worker_name)
|
||||
{
|
||||
Ok(Some(_)) => {
|
||||
debug!(
|
||||
ticket = %notice.ticket_id,
|
||||
event_kind = %notice.event_kind,
|
||||
companion = %self.companion_worker_name,
|
||||
"ensured Companion peer relationship before Ticket event notification"
|
||||
);
|
||||
}
|
||||
Ok(None) => {
|
||||
debug!(
|
||||
ticket = %notice.ticket_id,
|
||||
event_kind = %notice.event_kind,
|
||||
companion = %self.companion_worker_name,
|
||||
"skipping Companion peer registration because Companion metadata is missing"
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(
|
||||
ticket = %notice.ticket_id,
|
||||
event_kind = %notice.event_kind,
|
||||
companion = %self.companion_worker_name,
|
||||
error = %error,
|
||||
"failed to ensure Companion peer relationship before Ticket event notification"
|
||||
);
|
||||
}
|
||||
}
|
||||
let delivery = self
|
||||
.discovery
|
||||
.send_weak_notify_to_live_peer(&self.companion_worker_name, notice.message)
|
||||
.await;
|
||||
match delivery {
|
||||
WeakNotifyDelivery::Delivered => {
|
||||
debug!(
|
||||
ticket = %notice.ticket_id,
|
||||
event_kind = %notice.event_kind,
|
||||
companion = %self.companion_worker_name,
|
||||
"delivered weak Ticket event notification to Companion peer"
|
||||
);
|
||||
}
|
||||
skipped => {
|
||||
warn!(
|
||||
ticket = %notice.ticket_id,
|
||||
event_kind = %notice.event_kind,
|
||||
companion = %self.companion_worker_name,
|
||||
delivery = %skipped,
|
||||
"skipped weak Ticket event notification to Companion peer"
|
||||
);
|
||||
}
|
||||
}
|
||||
HookPostToolAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct TicketEventNotice {
|
||||
ticket_id: String,
|
||||
event_kind: String,
|
||||
message: String,
|
||||
}
|
||||
|
||||
fn build_ticket_event_notice(
|
||||
backend: &LocalTicketBackend,
|
||||
summary: &ToolResultSummary,
|
||||
) -> Option<TicketEventNotice> {
|
||||
if summary.is_error {
|
||||
return None;
|
||||
}
|
||||
let output = &summary.output;
|
||||
let content = output.content.as_deref()?;
|
||||
let content: Value = serde_json::from_str(content).ok()?;
|
||||
if !content.get("ok").and_then(Value::as_bool).unwrap_or(false) {
|
||||
return None;
|
||||
}
|
||||
|
||||
let event_kind = explicit_ticket_event_kind(summary.tool_name.as_str(), &content)?;
|
||||
let ticket_query = content.get("ticket").and_then(Value::as_str)?;
|
||||
let ticket = backend
|
||||
.show(TicketIdOrSlug::Query(ticket_query.to_string()))
|
||||
.ok()?;
|
||||
|
||||
let event_kind = sanitize_one_line(&event_kind, MAX_EVENT_KIND_CHARS);
|
||||
let ticket_id = ticket.meta.id.as_str();
|
||||
let title = sanitize_one_line(&ticket.meta.title, MAX_TITLE_CHARS);
|
||||
let state = ticket.meta.workflow_state.as_str();
|
||||
let output_summary = sanitize_one_line(&output.summary, MAX_SUMMARY_CHARS);
|
||||
let ref_path = event_ref_path(ticket_id, summary.tool_name.as_str());
|
||||
let message = render_ticket_event_notice_message(TicketEventNoticeValues {
|
||||
ticket_id,
|
||||
title: &title,
|
||||
state,
|
||||
event_kind: &event_kind,
|
||||
summary: &output_summary,
|
||||
ref_path: &ref_path,
|
||||
})?;
|
||||
|
||||
Some(TicketEventNotice {
|
||||
ticket_id: ticket_id.to_string(),
|
||||
event_kind,
|
||||
message: bound_chars(&message, MAX_MESSAGE_CHARS),
|
||||
})
|
||||
}
|
||||
|
||||
struct TicketEventNoticeValues<'a> {
|
||||
ticket_id: &'a str,
|
||||
title: &'a str,
|
||||
state: &'a str,
|
||||
event_kind: &'a str,
|
||||
summary: &'a str,
|
||||
ref_path: &'a str,
|
||||
}
|
||||
|
||||
fn render_ticket_event_notice_message(values: TicketEventNoticeValues<'_>) -> Option<String> {
|
||||
PromptCatalog::builtins_only()
|
||||
.ok()?
|
||||
.render(
|
||||
WorkerPrompt::TicketEventCompanionNotice,
|
||||
values.to_template(),
|
||||
)
|
||||
.ok()
|
||||
}
|
||||
|
||||
impl TicketEventNoticeValues<'_> {
|
||||
fn to_template(&self) -> TemplateValue {
|
||||
let mut values: BTreeMap<&'static str, TemplateValue> = BTreeMap::new();
|
||||
values.insert("ticket_id", TemplateValue::from(self.ticket_id));
|
||||
values.insert("title", TemplateValue::from(self.title));
|
||||
values.insert("state", TemplateValue::from(self.state));
|
||||
values.insert("event_kind", TemplateValue::from(self.event_kind));
|
||||
values.insert("summary", TemplateValue::from(self.summary));
|
||||
values.insert("ref_path", TemplateValue::from(self.ref_path));
|
||||
TemplateValue::from(values)
|
||||
}
|
||||
}
|
||||
|
||||
fn explicit_ticket_event_kind(tool_name: &str, content: &Value) -> Option<String> {
|
||||
match tool_name {
|
||||
"TicketComment" => content
|
||||
.get("event")
|
||||
.and_then(Value::as_str)
|
||||
.map(|event| format!("comment/{event}")),
|
||||
"TicketReview" => content
|
||||
.get("review")
|
||||
.and_then(Value::as_str)
|
||||
.map(|review| format!("review/{review}")),
|
||||
"TicketWorkflowState" => {
|
||||
let from = content.get("from").and_then(Value::as_str).unwrap_or("?");
|
||||
let to = content.get("to").and_then(Value::as_str).unwrap_or("?");
|
||||
Some(format!("state/{from}->{to}"))
|
||||
}
|
||||
"TicketIntakeReady" => Some("state/planning->ready".to_string()),
|
||||
"TicketClose" => Some("close/resolution".to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn event_ref_path(ticket_id: &str, tool_name: &str) -> String {
|
||||
let leaf = match tool_name {
|
||||
"TicketClose" => "resolution.md",
|
||||
"TicketIntakeReady" | "TicketWorkflowState" => "item.md",
|
||||
_ => "thread.md",
|
||||
};
|
||||
format!(".yoi/tickets/{ticket_id}/{leaf}")
|
||||
}
|
||||
|
||||
fn sanitize_one_line(input: &str, limit: usize) -> String {
|
||||
let collapsed = input.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
bound_chars(&collapsed, limit)
|
||||
}
|
||||
|
||||
fn bound_chars(input: &str, limit: usize) -> String {
|
||||
let mut out = String::new();
|
||||
for (idx, ch) in input.chars().filter(|ch| !ch.is_control()).enumerate() {
|
||||
if idx >= limit {
|
||||
out.push('…');
|
||||
break;
|
||||
}
|
||||
out.push(ch);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub(crate) fn companion_worker_name_for_workspace(
|
||||
workspace_root: &std::path::Path,
|
||||
) -> Option<String> {
|
||||
workspace_root
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.map(str::trim)
|
||||
.filter(|name| !name.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::WorkerStatus;
|
||||
use crate::runtime::dir::RuntimeDir;
|
||||
use crate::spawn::registry::SpawnedWorkerRegistry;
|
||||
use llm_engine::tool::ToolOutput;
|
||||
use pod_store::FsWorkerStore;
|
||||
use pod_store::WorkerMetadata;
|
||||
use protocol::stream::{JsonLineReader, JsonLineWriter};
|
||||
use protocol::{Event, Method};
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tempfile::tempdir;
|
||||
use ticket::NewTicket;
|
||||
use tokio::net::UnixListener;
|
||||
|
||||
fn create_backend_with_ticket(title: &str) -> (tempfile::TempDir, LocalTicketBackend, String) {
|
||||
let dir = tempdir().expect("tempdir");
|
||||
let backend = LocalTicketBackend::new(dir.path().to_path_buf());
|
||||
let mut input = NewTicket::new(title);
|
||||
input.body = ticket::MarkdownText::new("body");
|
||||
let ticket = backend.create(input).expect("create ticket");
|
||||
(dir, backend, ticket.id)
|
||||
}
|
||||
|
||||
fn tool_summary(tool_name: &str, output: ToolOutput) -> ToolResultSummary {
|
||||
ToolResultSummary {
|
||||
call_id: "test-call".to_string(),
|
||||
tool_name: tool_name.to_string(),
|
||||
output,
|
||||
is_error: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_bounded_event_scoped_notice_for_ticket_state_change() {
|
||||
let (_dir, backend, ticket_id) = create_backend_with_ticket(
|
||||
"A very long title that should be bounded but still identify the ticket precisely enough for Companion",
|
||||
);
|
||||
let output = ToolOutput {
|
||||
summary: "Changed ticket state from queued to inprogress with a deliberately long summary that should be bounded before entering the weak notification payload and should not contain large logs".into(),
|
||||
content: Some(
|
||||
json!({
|
||||
"ok": true,
|
||||
"ticket": ticket_id,
|
||||
"from": "queued",
|
||||
"to": "inprogress",
|
||||
})
|
||||
.to_string(),
|
||||
),
|
||||
};
|
||||
|
||||
let notice =
|
||||
build_ticket_event_notice(&backend, &tool_summary("TicketWorkflowState", output))
|
||||
.expect("notice");
|
||||
|
||||
assert_eq!(notice.ticket_id, ticket_id);
|
||||
assert_eq!(notice.event_kind, "state/queued->inprogress");
|
||||
assert!(notice.message.contains("auto_run=false"));
|
||||
assert!(notice.message.contains("event: state/queued->inprogress"));
|
||||
assert!(notice.message.contains("ref: .yoi/tickets/"));
|
||||
assert!(notice.message.chars().count() <= MAX_MESSAGE_CHARS + 1);
|
||||
|
||||
let expected = PromptCatalog::builtins_only()
|
||||
.expect("load prompt catalog")
|
||||
.render(
|
||||
WorkerPrompt::TicketEventCompanionNotice,
|
||||
TicketEventNoticeValues {
|
||||
ticket_id: ¬ice.ticket_id,
|
||||
title: &sanitize_one_line(
|
||||
"A very long title that should be bounded but still identify the ticket precisely enough for Companion",
|
||||
MAX_TITLE_CHARS,
|
||||
),
|
||||
state: "planning",
|
||||
event_kind: "state/queued->inprogress",
|
||||
summary: &sanitize_one_line(
|
||||
"Changed ticket state from queued to inprogress with a deliberately long summary that should be bounded before entering the weak notification payload and should not contain large logs",
|
||||
MAX_SUMMARY_CHARS,
|
||||
),
|
||||
ref_path: &format!(".yoi/tickets/{}/item.md", ticket_id),
|
||||
}
|
||||
.to_template(),
|
||||
)
|
||||
.expect("render prompt resource");
|
||||
assert_eq!(notice.message, bound_chars(&expected, MAX_MESSAGE_CHARS));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_passive_or_non_event_ticket_tools() {
|
||||
let (_dir, backend, ticket_id) = create_backend_with_ticket("Passive list test");
|
||||
let output = ToolOutput {
|
||||
summary: "Listed tickets".into(),
|
||||
content: Some(json!({"ok": true, "ticket": ticket_id}).to_string()),
|
||||
};
|
||||
|
||||
assert!(build_ticket_event_notice(&backend, &tool_summary("TicketList", output)).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notice_does_not_include_tool_content_body_or_error_details() {
|
||||
let (_dir, backend, ticket_id) = create_backend_with_ticket("Safe payload");
|
||||
let output = ToolOutput {
|
||||
summary: "Appended implementation_report to ticket".into(),
|
||||
content: Some(
|
||||
json!({
|
||||
"ok": true,
|
||||
"ticket": ticket_id,
|
||||
"event": "implementation_report",
|
||||
"body": "SECRET_TOKEN provider stack trace long diagnostic should not be copied",
|
||||
"error": "provider error details should not be copied"
|
||||
})
|
||||
.to_string(),
|
||||
),
|
||||
};
|
||||
|
||||
let notice = build_ticket_event_notice(&backend, &tool_summary("TicketComment", output))
|
||||
.expect("notice");
|
||||
|
||||
assert!(
|
||||
notice
|
||||
.message
|
||||
.contains("event: comment/implementation_report")
|
||||
);
|
||||
assert!(!notice.message.contains("SECRET_TOKEN"));
|
||||
assert!(!notice.message.contains("provider error details"));
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn ticket_event_hook_ensures_peer_and_delivers_weak_companion_notification() {
|
||||
let root = tempdir().expect("tempdir");
|
||||
let runtime_base = root.path().join("runtime");
|
||||
let store_dir = root.path().join("store");
|
||||
std::fs::create_dir_all(runtime_base.join("companion")).unwrap();
|
||||
let store = FsWorkerStore::new(&store_dir).unwrap();
|
||||
store
|
||||
.write(&WorkerMetadata {
|
||||
worker_name: "orchestrator".into(),
|
||||
active: None,
|
||||
workspace_root: None,
|
||||
spawned_children: Vec::new(),
|
||||
reclaimed_children: Vec::new(),
|
||||
peers: Vec::new(),
|
||||
resolved_manifest_snapshot: None,
|
||||
})
|
||||
.unwrap();
|
||||
store
|
||||
.write(&WorkerMetadata {
|
||||
worker_name: "companion".into(),
|
||||
active: None,
|
||||
workspace_root: None,
|
||||
spawned_children: Vec::new(),
|
||||
reclaimed_children: Vec::new(),
|
||||
peers: Vec::new(),
|
||||
resolved_manifest_snapshot: None,
|
||||
})
|
||||
.unwrap();
|
||||
let (_ticket_dir, backend, ticket_id) = create_backend_with_ticket("Companion event hook");
|
||||
let runtime_dir = Arc::new(
|
||||
RuntimeDir::create(&runtime_base, "orchestrator")
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let store_for_assert = store.clone();
|
||||
let hook = TicketEventCompanionNotifyHook::new(
|
||||
backend,
|
||||
WorkerDiscovery::new(
|
||||
store,
|
||||
"orchestrator".into(),
|
||||
runtime_base.clone(),
|
||||
root.path().to_path_buf(),
|
||||
SpawnedWorkerRegistry::new(runtime_dir),
|
||||
),
|
||||
"companion",
|
||||
);
|
||||
|
||||
let socket = runtime_base.join("companion").join("sock");
|
||||
let listener = UnixListener::bind(&socket).unwrap();
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel(1);
|
||||
let companion = tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.unwrap();
|
||||
let mut writer = JsonLineWriter::new(stream);
|
||||
writer
|
||||
.write(&Event::Snapshot {
|
||||
entries: Vec::new(),
|
||||
greeting: protocol::Greeting {
|
||||
worker_name: "companion".into(),
|
||||
cwd: "/tmp".into(),
|
||||
provider: "test".into(),
|
||||
model: "test".into(),
|
||||
scope_summary: String::new(),
|
||||
tools: Vec::new(),
|
||||
context_window: 0,
|
||||
context_tokens: 0,
|
||||
},
|
||||
status: WorkerStatus::Idle,
|
||||
in_flight: Default::default(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (stream, _) = listener.accept().await.unwrap();
|
||||
let (reader_half, writer_half) = stream.into_split();
|
||||
let mut reader = JsonLineReader::new(reader_half);
|
||||
let mut writer = JsonLineWriter::new(writer_half);
|
||||
writer
|
||||
.write(&Event::Snapshot {
|
||||
entries: Vec::new(),
|
||||
greeting: protocol::Greeting {
|
||||
worker_name: "companion".into(),
|
||||
cwd: "/tmp".into(),
|
||||
provider: "test".into(),
|
||||
model: "test".into(),
|
||||
scope_summary: String::new(),
|
||||
tools: Vec::new(),
|
||||
context_window: 0,
|
||||
context_tokens: 0,
|
||||
},
|
||||
status: WorkerStatus::Idle,
|
||||
in_flight: Default::default(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let method = reader.next::<Method>().await.unwrap().unwrap();
|
||||
if let Method::Notify { message, auto_run } = method {
|
||||
assert!(!auto_run);
|
||||
tx.send(message).await.unwrap();
|
||||
} else {
|
||||
panic!("expected Notify, got {method:?}");
|
||||
}
|
||||
});
|
||||
|
||||
let output = ToolOutput {
|
||||
summary: "Changed ticket state from queued to inprogress".into(),
|
||||
content: Some(
|
||||
json!({
|
||||
"ok": true,
|
||||
"ticket": ticket_id,
|
||||
"from": "queued",
|
||||
"to": "inprogress",
|
||||
})
|
||||
.to_string(),
|
||||
),
|
||||
};
|
||||
let action = hook
|
||||
.call(&tool_summary("TicketWorkflowState", output))
|
||||
.await;
|
||||
assert_eq!(action, HookPostToolAction::Continue);
|
||||
let message = rx.recv().await.unwrap();
|
||||
assert!(message.contains("event: state/queued->inprogress"));
|
||||
assert!(message.contains("title: Companion event hook"));
|
||||
let orchestrator = store_for_assert
|
||||
.read_by_name("orchestrator")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(orchestrator.peers.len(), 1);
|
||||
assert_eq!(orchestrator.peers[0].worker_name, "companion");
|
||||
let companion_metadata = store_for_assert.read_by_name("companion").unwrap().unwrap();
|
||||
assert_eq!(companion_metadata.peers.len(), 1);
|
||||
assert_eq!(companion_metadata.peers[0].worker_name, "orchestrator");
|
||||
companion.await.unwrap();
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,261 @@
|
||||
//! Worker-side Workflow resolver.
|
||||
//!
|
||||
//! Turns `Segment::WorkflowInvoke { slug }` into system-message attachments:
|
||||
//! dependency Knowledge bodies first, then the Workflow body. Resolution is
|
||||
//! strict for explicit user invocations: missing workflows, non-user-invocable
|
||||
//! workflows, and missing Knowledge requirements are returned as errors before
|
||||
//! the turn is handed to the Engine.
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use llm_engine::Item;
|
||||
use memory::WorkspaceLayout;
|
||||
use memory::schema::split_frontmatter;
|
||||
use workflow_crate::{Slug, WorkflowRegistry};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum WorkflowResolveError {
|
||||
InvalidSlug(workflow_crate::WorkflowLintError),
|
||||
NotFound {
|
||||
slug: String,
|
||||
},
|
||||
NotUserInvocable {
|
||||
slug: String,
|
||||
},
|
||||
KnowledgeNotFound {
|
||||
workflow: String,
|
||||
slug: String,
|
||||
},
|
||||
KnowledgeRead {
|
||||
workflow: String,
|
||||
slug: String,
|
||||
source: std::io::Error,
|
||||
},
|
||||
KnowledgeFrontmatter {
|
||||
workflow: String,
|
||||
slug: String,
|
||||
source: memory::LintError,
|
||||
},
|
||||
}
|
||||
|
||||
impl fmt::Display for WorkflowResolveError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Self::InvalidSlug(e) => write!(f, "invalid workflow slug: {e}"),
|
||||
Self::NotFound { slug } => write!(f, "workflow /{slug} is not registered"),
|
||||
Self::NotUserInvocable { slug } => {
|
||||
write!(f, "workflow /{slug} is not user-invocable")
|
||||
}
|
||||
Self::KnowledgeNotFound { workflow, slug } => write!(
|
||||
f,
|
||||
"workflow /{workflow} requires missing Knowledge slug `{slug}`"
|
||||
),
|
||||
Self::KnowledgeRead {
|
||||
workflow,
|
||||
slug,
|
||||
source,
|
||||
} => write!(
|
||||
f,
|
||||
"workflow /{workflow} could not read required Knowledge `{slug}`: {source}"
|
||||
),
|
||||
Self::KnowledgeFrontmatter {
|
||||
workflow,
|
||||
slug,
|
||||
source,
|
||||
} => write!(
|
||||
f,
|
||||
"workflow /{workflow} required Knowledge `{slug}` has invalid frontmatter: {source}"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for WorkflowResolveError {}
|
||||
|
||||
struct BuiltinKnowledgeResource {
|
||||
slug: &'static str,
|
||||
content: &'static str,
|
||||
}
|
||||
|
||||
const BUILTIN_KNOWLEDGE: &[BuiltinKnowledgeResource] = &[BuiltinKnowledgeResource {
|
||||
slug: "workflow-resource-boundary",
|
||||
content: include_str!("../../../../resources/knowledge/workflow-resource-boundary.md"),
|
||||
}];
|
||||
|
||||
fn builtin_knowledge(slug: &Slug) -> Option<&'static str> {
|
||||
BUILTIN_KNOWLEDGE
|
||||
.iter()
|
||||
.find(|resource| resource.slug == slug.as_str())
|
||||
.map(|resource| resource.content)
|
||||
}
|
||||
|
||||
fn read_required_knowledge(
|
||||
workflow: &Slug,
|
||||
layout: &WorkspaceLayout,
|
||||
req: &Slug,
|
||||
) -> Result<(String, &'static str), WorkflowResolveError> {
|
||||
let path = layout.knowledge_dir().join(format!("{req}.md"));
|
||||
match std::fs::read_to_string(&path) {
|
||||
Ok(raw) => Ok((raw, "workspace")),
|
||||
Err(source) if source.kind() == std::io::ErrorKind::NotFound => {
|
||||
if let Some(raw) = builtin_knowledge(req) {
|
||||
Ok((raw.to_string(), "builtin"))
|
||||
} else {
|
||||
Err(WorkflowResolveError::KnowledgeNotFound {
|
||||
workflow: workflow.to_string(),
|
||||
slug: req.to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
Err(source) => Err(WorkflowResolveError::KnowledgeRead {
|
||||
workflow: workflow.to_string(),
|
||||
slug: req.to_string(),
|
||||
source,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn resolve_workflow_invocation(
|
||||
registry: &WorkflowRegistry,
|
||||
layout: &WorkspaceLayout,
|
||||
raw_slug: &str,
|
||||
) -> Result<Vec<Item>, WorkflowResolveError> {
|
||||
let slug = Slug::parse(raw_slug.to_string())
|
||||
.map_err(|source| WorkflowResolveError::InvalidSlug(source.into()))?;
|
||||
let record = registry
|
||||
.get(&slug)
|
||||
.ok_or_else(|| WorkflowResolveError::NotFound {
|
||||
slug: raw_slug.to_string(),
|
||||
})?;
|
||||
if !record.user_invocable {
|
||||
return Err(WorkflowResolveError::NotUserInvocable {
|
||||
slug: raw_slug.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
let mut out = Vec::new();
|
||||
for req in &record.requires {
|
||||
let (raw, knowledge_source) = read_required_knowledge(&slug, layout, req)?;
|
||||
let (_yaml, body) = split_frontmatter(&raw).map_err(|source| {
|
||||
WorkflowResolveError::KnowledgeFrontmatter {
|
||||
workflow: slug.to_string(),
|
||||
slug: req.to_string(),
|
||||
source,
|
||||
}
|
||||
})?;
|
||||
out.push(Item::system_message(format!(
|
||||
"[Workflow /{} requires Knowledge #{} from {}]\n{}",
|
||||
slug,
|
||||
req,
|
||||
knowledge_source,
|
||||
body.trim_end()
|
||||
)));
|
||||
}
|
||||
out.push(Item::system_message(format!(
|
||||
"[Workflow /{} from {}]\n{}",
|
||||
slug,
|
||||
record.source.label(),
|
||||
record.body.trim_end()
|
||||
)));
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn write(path: &std::path::Path, content: &str) {
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent).unwrap();
|
||||
}
|
||||
std::fs::write(path, content).unwrap();
|
||||
}
|
||||
|
||||
fn setup() -> (TempDir, WorkspaceLayout, WorkflowRegistry) {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let layout = WorkspaceLayout::new(dir.path().to_path_buf());
|
||||
write(
|
||||
&dir.path().join(".yoi/knowledge/policy.md"),
|
||||
"---\ncreated_at: 2026-01-01T00:00:00Z\nupdated_at: 2026-01-01T00:00:00Z\nkind: policy\ndescription: p\nmodel_invokation: false\nuser_invocable: true\nlast_sources: []\n---\npolicy body\n",
|
||||
);
|
||||
write(
|
||||
&dir.path().join(".yoi/workflow/run-it.md"),
|
||||
"---\ndescription: run\nrequires: [policy]\n---\nworkflow body\n",
|
||||
);
|
||||
let registry = workflow_crate::load_workflows(&layout).unwrap();
|
||||
(dir, layout, registry)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_requires_before_workflow_body() {
|
||||
let (_dir, layout, registry) = setup();
|
||||
let items = resolve_workflow_invocation(®istry, &layout, "run-it").unwrap();
|
||||
assert_eq!(items.len(), 2);
|
||||
let first = format!("{:?}", items[0]);
|
||||
let second = format!("{:?}", items[1]);
|
||||
assert!(first.contains("Knowledge #policy"));
|
||||
assert!(first.contains("policy body"));
|
||||
assert!(second.contains("[Workflow /run-it from workspace workflow]"));
|
||||
assert!(second.contains("workflow body"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builtin_workflow_uses_builtin_required_knowledge_when_workspace_missing() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let layout = WorkspaceLayout::new(dir.path().to_path_buf());
|
||||
let registry = workflow_crate::load_workflows(&layout).unwrap();
|
||||
let items =
|
||||
resolve_workflow_invocation(®istry, &layout, "ticket-intake-workflow").unwrap();
|
||||
let first = format!("{:?}", items[0]);
|
||||
let second = format!("{:?}", items[1]);
|
||||
assert!(first.contains("Knowledge #workflow-resource-boundary from builtin"));
|
||||
assert!(first.contains("Builtin workflow resources live under"));
|
||||
assert!(second.contains("[Workflow /ticket-intake-workflow from builtin workflow]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_knowledge_overrides_builtin_required_knowledge() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let layout = WorkspaceLayout::new(dir.path().to_path_buf());
|
||||
write(
|
||||
&dir.path()
|
||||
.join(".yoi/knowledge/workflow-resource-boundary.md"),
|
||||
"---\ncreated_at: 2026-01-01T00:00:00Z\nupdated_at: 2026-01-01T00:00:00Z\nkind: policy\ndescription: p\nmodel_invokation: false\nuser_invocable: true\nlast_sources: []\n---\nworkspace override knowledge\n",
|
||||
);
|
||||
let registry = workflow_crate::load_workflows(&layout).unwrap();
|
||||
let items =
|
||||
resolve_workflow_invocation(®istry, &layout, "ticket-intake-workflow").unwrap();
|
||||
let first = format!("{:?}", items[0]);
|
||||
assert!(first.contains("Knowledge #workflow-resource-boundary from workspace"));
|
||||
assert!(first.contains("workspace override knowledge"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_invocable_false_errors() {
|
||||
let (dir, layout, _registry) = setup();
|
||||
write(
|
||||
&dir.path().join(".yoi/workflow/hidden.md"),
|
||||
"---\ndescription: hidden\nuser_invocable: false\n---\nbody\n",
|
||||
);
|
||||
let registry = workflow_crate::load_workflows(&layout).unwrap();
|
||||
let err = resolve_workflow_invocation(®istry, &layout, "hidden").unwrap_err();
|
||||
assert!(matches!(err, WorkflowResolveError::NotUserInvocable { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_required_knowledge_errors() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let layout = WorkspaceLayout::new(dir.path().to_path_buf());
|
||||
write(
|
||||
&dir.path().join(".yoi/workflow/bad.md"),
|
||||
"---\ndescription: bad\nrequires: [ghost]\n---\nbody\n",
|
||||
);
|
||||
let registry = workflow_crate::load_workflows(&layout).unwrap();
|
||||
let err = resolve_workflow_invocation(®istry, &layout, "bad").unwrap_err();
|
||||
assert!(matches!(
|
||||
err,
|
||||
WorkflowResolveError::KnowledgeNotFound { .. }
|
||||
));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user