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

feat: add declarative Feature chat invocations and typed composers

Browse Source
This commit is contained in:
Hare
2026-10-05 18:44:41 +09:00
parent cc97987482
commit f7783b3d2a
69 changed files with 10065 additions and 657 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+59
View File
@@ -184,6 +184,7 @@ impl WorkerHandle {
&self,
kind: protocol::CompletionKind,
prefix: &str,
context: Option<&protocol::CompletionContext>,
) -> Vec<protocol::CompletionEntry> {
match kind {
protocol::CompletionKind::File => {
@@ -196,9 +197,62 @@ impl WorkerHandle {
.map(|candidate| protocol::CompletionEntry {
value: candidate.path,
is_dir: candidate.is_dir,
description: None,
usage: None,
invocation: None,
})
.collect()
}
protocol::CompletionKind::Feature => self
.shared_state
.feature_invocations()
.feature_completions(prefix),
protocol::CompletionKind::FeatureArgument => {
let Some(context) = context else {
return Vec::new();
};
let registry = self.shared_state.feature_invocations();
let Some(argument) = context.argument.as_deref() else {
return registry.argument_name_completions(&context.invocation, prefix);
};
let worker_file = registry
.descriptor(&context.invocation)
.and_then(|descriptor| {
descriptor
.arguments
.iter()
.find(|candidate| candidate.name == argument)
})
.is_some_and(|descriptor| {
matches!(
descriptor.completion,
protocol::InvocationCompletion::WorkerFile
)
});
if worker_file {
let Some(view) = self.shared_state.fs_view() else {
return Vec::new();
};
let candidates = view.list_file_completions(prefix).await;
if registry.descriptor(&context.invocation).is_none() {
return Vec::new();
}
return crate::feature::invocation::bounded_completion_entries(
candidates
.into_iter()
.map(|candidate| protocol::CompletionEntry {
value: candidate.path,
is_dir: candidate.is_dir,
description: None,
usage: None,
invocation: None,
}),
);
}
registry
.argument_completions(&context.invocation, argument, prefix)
.await
}
}
}
@@ -846,6 +900,7 @@ impl WorkerController {
manifest_toml.clone(),
greeting,
));
shared_state.set_feature_invocations(worker.feature_invocations());
let usage_state = shared_state.clone();
let usage_events = working_event_tx.clone();
worker.engine_mut().on_usage(move |event| {
@@ -1327,6 +1382,10 @@ where
let worker_enabled = feature_config.worker.enabled;
let sub_worker_enabled = feature_config.sub_worker.enabled;
let mut feature_registry = FeatureRegistryBuilder::new();
if !backend_job_profile {
feature_registry
.add_module(crate::feature::builtin::chat_invocation::AttachmentInvocationFeature);
}
if backend_job_profile {
let workspace_client = worker.workspace_client_handle();
if !workspace_client.is_available() || workspace_client.workspace_id().is_none() {
crates/worker/src/feature.rs
+56
View File
@@ -20,6 +20,7 @@ use agen::Engine;
use agen::llm_client::client::LlmClient;
use agen::state::Mutable;
use agen::tool::ToolDefinition;
use protocol::FeatureInvocationDescriptor;
use serde::{Deserialize, Serialize};
use thiserror::Error;
@@ -623,6 +624,9 @@ pub struct FeatureDescriptor {
pub provides_services: Vec<ServiceDeclaration>,
pub requires_services: Vec<ServiceRequirement>,
pub protocol_providers: Vec<ProtocolProviderDeclaration>,
/// Declarative, user-selected chat invocations contributed by this Feature.
/// Metadata is public; handlers are registered separately during install.
pub chat_invocations: Vec<FeatureInvocationDescriptor>,
}
impl FeatureDescriptor {
@@ -640,6 +644,7 @@ impl FeatureDescriptor {
provides_services: Vec::new(),
requires_services: Vec::new(),
protocol_providers: Vec::new(),
chat_invocations: Vec::new(),
}
}
@@ -682,6 +687,11 @@ impl FeatureDescriptor {
self.protocol_providers.push(provider);
self
}
pub fn with_chat_invocation(mut self, invocation: FeatureInvocationDescriptor) -> Self {
self.chat_invocations.push(invocation);
self
}
}
/// Feature module contribution boundary.
@@ -739,6 +749,7 @@ pub enum FeatureContributionKind {
BackgroundTask,
Service,
ProtocolProvider,
ChatInvocation,
Notification,
Alert,
Diagnostic,
@@ -765,6 +776,7 @@ pub struct FeatureInstallReport {
pub provided_services: Vec<ServiceDeclaration>,
pub resolved_service_requirements: Vec<ServiceRequirement>,
pub protocol_providers: Vec<ProtocolProviderLifecycleDiagnostic>,
pub installed_chat_invocations: Vec<FeatureInvocationDescriptor>,
pub skipped: Vec<SkippedContribution>,
pub diagnostics: Vec<FeatureDiagnostic>,
}
@@ -782,6 +794,7 @@ impl FeatureInstallReport {
provided_services: Vec::new(),
resolved_service_requirements: Vec::new(),
protocol_providers: Vec::new(),
installed_chat_invocations: Vec::new(),
skipped: Vec::new(),
diagnostics: Vec::new(),
}
@@ -794,6 +807,7 @@ impl FeatureInstallReport {
self.installed_instructions.clear();
self.declared_background_tasks.clear();
self.provided_services.clear();
self.installed_chat_invocations.clear();
}
fn mark_skipped(
@@ -819,6 +833,7 @@ struct FeatureContributionDeclarations {
provided_services: HashSet<(ServiceId, String)>,
required_services: HashSet<ServiceId>,
protocol_providers: HashSet<ProviderId>,
chat_invocations: Vec<FeatureInvocationDescriptor>,
}
impl FeatureContributionDeclarations {
@@ -859,6 +874,7 @@ impl FeatureContributionDeclarations {
.iter()
.map(|provider| provider.id.clone())
.collect(),
chat_invocations: descriptor.chat_invocations.clone(),
}
}
@@ -1527,6 +1543,7 @@ pub struct FeatureInstallContext<'a> {
hook_builder: &'a mut HookRegistryBuilder,
background_task_builder: &'a mut FeatureBackgroundTaskRegistryBuilder,
service_registry: &'a mut FeatureServiceRegistry,
invocation_registry: &'a mut FeatureInvocationRegistry,
report: &'a mut FeatureInstallReport,
}
@@ -1591,6 +1608,15 @@ impl FeatureInstallContext<'_> {
}
}
pub fn chat_invocations(&mut self) -> InvocationContributionRegistrar<'_> {
InvocationContributionRegistrar {
feature_id: self.feature_id,
declarations: &self.declarations.chat_invocations,
registry: self.invocation_registry,
installed: &mut self.report.installed_chat_invocations,
}
}
pub fn diagnostics(&mut self) -> FeatureDiagnosticSink<'_> {
FeatureDiagnosticSink {
report: self.report,
@@ -1604,6 +1630,7 @@ pub struct FeatureRegistryInstallReport {
pub reports: Vec<FeatureInstallReport>,
pub services: FeatureServiceRegistry,
pub background_tasks: FeatureBackgroundTaskRegistry,
pub chat_invocations: FeatureInvocationRegistry,
pub plan_error: Option<FeaturePlanError>,
}
@@ -2026,11 +2053,13 @@ impl FeatureRegistryBuilder {
reports,
services: FeatureServiceRegistry::default(),
background_tasks: FeatureBackgroundTaskRegistry::default(),
chat_invocations: FeatureInvocationRegistry::default(),
plan_error: Some(error),
};
}
};
let mut service_registry = FeatureServiceRegistry::default();
let mut invocation_registry = FeatureInvocationRegistry::default();
let mut background_task_builder = FeatureBackgroundTaskRegistryBuilder::default();
let mut reports = Vec::with_capacity(plan.ordered_indices.len());
let install_hook_checkpoint = hook_builder.checkpoint();
@@ -2055,6 +2084,7 @@ impl FeatureRegistryBuilder {
let hook_checkpoint = hook_builder.checkpoint();
let background_checkpoint = background_task_builder.checkpoint();
let service_checkpoint = service_registry.clone();
let invocation_checkpoint = invocation_registry.clone();
let tool_checkpoint = pending_tools.len();
let installed_tool_checkpoint = installed_tool_names.clone();
@@ -2102,6 +2132,7 @@ impl FeatureRegistryBuilder {
hook_builder,
background_task_builder: &mut background_task_builder,
service_registry: &mut service_registry,
invocation_registry: &mut invocation_registry,
report: &mut report,
};
module.install(&mut context)
@@ -2113,6 +2144,7 @@ impl FeatureRegistryBuilder {
hook_builder.rollback_to(hook_checkpoint);
background_task_builder.rollback_to(&background_checkpoint);
service_registry = service_checkpoint.clone();
invocation_registry = invocation_checkpoint.clone();
pending_tools.truncate(tool_checkpoint);
installed_tool_names = installed_tool_checkpoint.clone();
report.clear_installed_contributions();
@@ -2138,6 +2170,14 @@ impl FeatureRegistryBuilder {
)));
}
}
for invocation in &descriptor.chat_invocations {
if !report.installed_chat_invocations.contains(invocation) {
report.diagnostics.push(FeatureDiagnostic::error(format!(
"feature `{}` declared chat invocation `{}` but did not register an executable handler",
descriptor.id, invocation.identity
)));
}
}
if report
.diagnostics
.iter()
@@ -2146,6 +2186,7 @@ impl FeatureRegistryBuilder {
hook_builder.rollback_to(hook_checkpoint);
background_task_builder.rollback_to(&background_checkpoint);
service_registry = service_checkpoint.clone();
invocation_registry = invocation_checkpoint.clone();
pending_tools.truncate(tool_checkpoint);
installed_tool_names = installed_tool_checkpoint.clone();
report.clear_installed_contributions();
@@ -2176,6 +2217,7 @@ impl FeatureRegistryBuilder {
reports,
services: FeatureServiceRegistry::default(),
background_tasks: FeatureBackgroundTaskRegistry::default(),
chat_invocations: FeatureInvocationRegistry::default(),
plan_error: None,
}
} else {
@@ -2183,6 +2225,7 @@ impl FeatureRegistryBuilder {
reports,
services: service_registry,
background_tasks: background_task_builder.build(),
chat_invocations: invocation_registry,
plan_error: None,
}
}
@@ -2225,11 +2268,24 @@ pub enum FeatureInstallError {
first_feature: String,
duplicate_feature: String,
},
#[error("duplicate invocation identity `{identity}`")]
DuplicateInvocationIdentity { identity: String },
#[error("duplicate invocation name `{name}` from `{duplicate}`; first registered by `{first}`")]
DuplicateInvocationName {
name: String,
first: String,
duplicate: String,
},
#[error("feature install failed: {0}")]
Install(String),
}
pub mod background;
pub mod invocation;
pub use invocation::{
FeatureInvocationCompletionProvider, FeatureInvocationContext, FeatureInvocationHandler,
FeatureInvocationHandlerError, FeatureInvocationRegistry, InvocationContributionRegistrar,
};
pub mod builtin;
pub mod mcp;
pub(crate) mod session;
crates/worker/src/feature/builtin.rs
+1
View File
@@ -5,6 +5,7 @@
//! an external plugin-loading surface.
pub mod backend_job_result;
pub mod chat_invocation;
pub mod flow_transition;
pub mod manage_workdir;
pub mod manage_worker;
crates/worker/src/feature/builtin/chat_invocation.rs
+72
View File
@@ -0,0 +1,72 @@
//! Builtin metadata for client-local attachment staging.
//!
//! `/attach` uses the common invocation declaration and parser, but its path is
//! intentionally consumed by the client adapter. The Worker accepts only the
//! resulting immutable `UploadedFile` segment and never resolves a client path.
use async_trait::async_trait;
use protocol::{
FeatureInvocation, FeatureInvocationDescriptor, FeatureInvocationIdentity,
FeatureInvocationResult, FeatureInvocationSyntax, InvocationArgumentDescriptor,
InvocationArgumentType, InvocationClientAdapter, InvocationCompletion,
};
use crate::feature::{
FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureInvocationContext,
FeatureInvocationHandler, FeatureInvocationHandlerError, FeatureModule,
};
pub const ATTACH_FEATURE_ID: &str = "attachments";
pub const ATTACH_INVOCATION_ID: &str = "builtin:attachments/attach";
pub fn attach_invocation_descriptor() -> FeatureInvocationDescriptor {
FeatureInvocationDescriptor {
identity: FeatureInvocationIdentity(ATTACH_INVOCATION_ID.into()),
name: "attach".into(),
aliases: Vec::new(),
display_name: "Attach file".into(),
description: "Stage a client-local file as an editable attachment chip.".into(),
syntax: FeatureInvocationSyntax::Parenthesized,
arguments: vec![InvocationArgumentDescriptor {
name: "file".into(),
position: Some(0),
required: true,
value_type: InvocationArgumentType::ClientFile,
completion: InvocationCompletion::ClientFile,
description: Some("File available to this client".into()),
}],
client_adapter: Some(InvocationClientAdapter::Attachment),
}
}
pub struct AttachmentInvocationFeature;
impl FeatureModule for AttachmentInvocationFeature {
fn descriptor(&self) -> FeatureDescriptor {
FeatureDescriptor::builtin(ATTACH_FEATURE_ID, "Attachments")
.with_description("Client-local attachment staging and typed uploaded-file input")
.with_chat_invocation(attach_invocation_descriptor())
}
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
context.chat_invocations().register(
attach_invocation_descriptor(),
AttachmentMustBeStagedByClient,
)
}
}
struct AttachmentMustBeStagedByClient;
#[async_trait]
impl FeatureInvocationHandler for AttachmentMustBeStagedByClient {
async fn invoke(
&self,
_context: FeatureInvocationContext,
_invocation: &FeatureInvocation,
) -> Result<FeatureInvocationResult, FeatureInvocationHandlerError> {
Err(FeatureInvocationHandlerError::failed(
"attachment invocation reached the Worker without client-local upload staging",
))
}
}
crates/worker/src/feature/invocation.rs
+966
View File
@@ -0,0 +1,966 @@
use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use std::sync::Arc;
use async_trait::async_trait;
use protocol::{
CompletionEntry, FeatureInvocation, FeatureInvocationDescriptor, FeatureInvocationIdentity,
FeatureInvocationResult, InvocationCompletion, InvocationValue, validate_feature_invocation,
};
use super::{FeatureId, FeatureInstallError};
pub(crate) const MAX_INVOCATION_MESSAGE_BYTES: usize = 4 * 1024;
pub(crate) const MAX_INVOCATION_CONTEXT_BYTES: usize = 16 * 1024;
const MAX_COMPLETION_ENTRIES: usize = 64;
const MAX_COMPLETION_VALUE_BYTES: usize = 1024;
const MAX_COMPLETION_DESCRIPTION_BYTES: usize = 2 * 1024;
const MAX_COMPLETION_USAGE_BYTES: usize = 4 * 1024;
const MAX_INVOCATION_DESCRIPTOR_BYTES: usize =
protocol::invocation::MAX_FEATURE_INVOCATION_DESCRIPTOR_BYTES;
const MAX_COMPLETION_RESPONSE_BYTES: usize = 64 * 1024;
pub(crate) fn invocation_result_is_bounded(result: &FeatureInvocationResult) -> bool {
result.message.len() <= MAX_INVOCATION_MESSAGE_BYTES
&& result
.context
.as_ref()
.is_none_or(|context| context.len() <= MAX_INVOCATION_CONTEXT_BYTES)
}
fn truncate_display(text: &mut String, max_bytes: usize) {
let mut end = text.len().min(max_bytes);
while !text.is_char_boundary(end) {
end -= 1;
}
text.truncate(end);
}
/// Values are never truncated: doing so could silently select a different
/// operation. Display-only text may be shortened; oversized metadata is dropped.
pub(crate) fn bounded_completion_entries(
entries: impl IntoIterator<Item = CompletionEntry>,
) -> Vec<CompletionEntry> {
let mut result = Vec::new();
let mut bytes = 2; // JSON array delimiters.
for mut entry in entries.into_iter().take(MAX_COMPLETION_ENTRIES * 4) {
if entry.value.len() > MAX_COMPLETION_VALUE_BYTES {
continue;
}
if let Some(description) = &mut entry.description {
truncate_display(description, MAX_COMPLETION_DESCRIPTION_BYTES);
}
if let Some(usage) = &mut entry.usage {
truncate_display(usage, MAX_COMPLETION_USAGE_BYTES);
}
if entry.invocation.as_ref().is_some_and(|descriptor| {
serde_json::to_vec(descriptor).map_or(true, |encoded| {
encoded.len() > MAX_INVOCATION_DESCRIPTOR_BYTES
})
}) {
continue;
}
let Ok(encoded) = serde_json::to_vec(&entry) else {
continue;
};
let additional = encoded.len() + usize::from(!result.is_empty());
if bytes + additional > MAX_COMPLETION_RESPONSE_BYTES {
continue;
}
bytes += additional;
result.push(entry);
if result.len() == MAX_COMPLETION_ENTRIES {
break;
}
}
result
}
/// Host-owned execution context. `invocation_id` is the idempotency key a
/// handler must use when its operation has external side effects. Workspace,
/// credential, and permission authority remain captured by the handler itself.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FeatureInvocationContext {
pub invocation_id: String,
}
#[async_trait]
pub trait FeatureInvocationHandler: Send + Sync {
/// Side-effect-free runtime availability/permission check. Installation is
/// necessary but not sufficient: discovery and execution both consult this.
fn is_available(&self) -> bool {
true
}
/// Side-effect-free semantic and authority validation. Installed handlers
/// retain Workspace/credential/permission authority; request arguments never
/// supply it. The Worker calls this before persisting an execution start.
fn validate(
&self,
_invocation: &FeatureInvocation,
) -> Result<(), FeatureInvocationHandlerError> {
Ok(())
}
async fn invoke(
&self,
context: FeatureInvocationContext,
invocation: &FeatureInvocation,
) -> Result<FeatureInvocationResult, FeatureInvocationHandlerError>;
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("{message}")]
pub struct FeatureInvocationHandlerError {
pub message: String,
/// True only when the handler cannot determine whether its business-side
/// effect committed. The Worker never automatically retries this outcome.
pub outcome_unknown: bool,
}
impl FeatureInvocationHandlerError {
pub fn failed(message: impl Into<String>) -> Self {
Self {
message: message.into(),
outcome_unknown: false,
}
}
pub fn outcome_unknown(message: impl Into<String>) -> Self {
Self {
message: message.into(),
outcome_unknown: true,
}
}
}
#[async_trait]
pub trait FeatureInvocationCompletionProvider: Send + Sync {
/// The provider selector comes only from installed host metadata, never
/// from completion-request input. Providers retain their own scoped authority.
async fn complete(
&self,
completion: &InvocationCompletion,
argument: &str,
prefix: &str,
) -> Vec<CompletionEntry>;
}
#[derive(Clone)]
struct RegisteredInvocation {
descriptor: FeatureInvocationDescriptor,
handler: Arc<dyn FeatureInvocationHandler>,
completion_provider: Option<Arc<dyn FeatureInvocationCompletionProvider>>,
}
/// Installed invocation metadata and executable handlers. Equality/debug expose
/// only public descriptors; implementation objects and credentials never leak.
#[derive(Clone, Default)]
pub struct FeatureInvocationRegistry {
entries: BTreeMap<FeatureInvocationIdentity, RegisteredInvocation>,
names: BTreeMap<String, FeatureInvocationIdentity>,
}
impl fmt::Debug for FeatureInvocationRegistry {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("FeatureInvocationRegistry")
.field("descriptors", &self.descriptors())
.finish()
}
}
impl PartialEq for FeatureInvocationRegistry {
fn eq(&self, other: &Self) -> bool {
self.descriptors() == other.descriptors()
}
}
impl Eq for FeatureInvocationRegistry {}
impl FeatureInvocationRegistry {
pub fn descriptors(&self) -> Vec<FeatureInvocationDescriptor> {
self.entries
.values()
.filter(|entry| entry.handler.is_available())
.map(|entry| entry.descriptor.clone())
.collect()
}
pub fn descriptor(
&self,
identity: &FeatureInvocationIdentity,
) -> Option<&FeatureInvocationDescriptor> {
self.entries
.get(identity)
.filter(|entry| entry.handler.is_available())
.map(|entry| &entry.descriptor)
}
pub fn feature_completions(&self, prefix: &str) -> Vec<CompletionEntry> {
bounded_completion_entries(
self.entries
.values()
.filter(|entry| entry.handler.is_available())
.filter(|entry| {
entry.descriptor.name.starts_with(prefix)
|| entry
.descriptor
.aliases
.iter()
.any(|alias| alias.starts_with(prefix))
})
.map(|entry| CompletionEntry {
value: entry.descriptor.name.clone(),
is_dir: false,
description: Some(entry.descriptor.description.clone()),
usage: Some(entry.descriptor.usage()),
invocation: Some(entry.descriptor.clone()),
}),
)
}
pub fn argument_name_completions(
&self,
identity: &FeatureInvocationIdentity,
prefix: &str,
) -> Vec<CompletionEntry> {
bounded_completion_entries(
self.entries
.get(identity)
.filter(|entry| entry.handler.is_available())
.into_iter()
.flat_map(|entry| entry.descriptor.arguments.iter())
.filter(|argument| argument.name.starts_with(prefix))
.map(|argument| CompletionEntry {
value: format!("{}=", argument.name),
is_dir: false,
description: argument.description.clone(),
usage: None,
invocation: None,
}),
)
}
pub async fn argument_completions(
&self,
identity: &FeatureInvocationIdentity,
argument: &str,
prefix: &str,
) -> Vec<CompletionEntry> {
let Some(entry) = self
.entries
.get(identity)
.filter(|entry| entry.handler.is_available())
else {
return Vec::new();
};
let Some(descriptor) = entry
.descriptor
.arguments
.iter()
.find(|candidate| candidate.name == argument)
else {
return Vec::new();
};
let strings = match &descriptor.completion {
InvocationCompletion::None => match &descriptor.value_type {
protocol::InvocationArgumentType::Boolean => ["true", "false"]
.into_iter()
.filter(|value| value.starts_with(prefix))
.map(str::to_owned)
.collect(),
protocol::InvocationArgumentType::Enum { values } => values
.iter()
.filter(|value| value.starts_with(prefix))
.cloned()
.collect(),
_ => Vec::new(),
},
InvocationCompletion::WorkerFile | InvocationCompletion::ClientFile => Vec::new(),
InvocationCompletion::Static { values } => values
.iter()
.filter(|value| value.starts_with(prefix))
.cloned()
.collect(),
InvocationCompletion::Provider { .. } => {
return match &entry.completion_provider {
Some(provider) => {
let entries = provider
.complete(&descriptor.completion, argument, prefix)
.await;
if !entry.handler.is_available() {
return Vec::new();
}
bounded_completion_entries(entries.into_iter().map(|mut candidate| {
// Argument providers contribute values/display text,
// never new executable invocation metadata.
candidate.invocation = None;
candidate
}))
}
None => Vec::new(),
};
}
};
bounded_completion_entries(strings.into_iter().map(|value| CompletionEntry {
value,
is_dir: false,
description: descriptor.description.clone(),
usage: None,
invocation: None,
}))
}
pub(crate) fn validate_metadata(
&self,
invocation: &FeatureInvocation,
) -> Result<(), FeatureInvocationHandlerError> {
let entry = self
.entries
.get(&invocation.identity)
.filter(|entry| entry.handler.is_available())
.ok_or_else(|| {
FeatureInvocationHandlerError::failed(
"Feature invocation is not enabled or currently available",
)
})?;
if entry.descriptor.client_adapter.is_some() {
return Err(FeatureInvocationHandlerError::failed(
"client-local invocation must be staged by the client, not executed by the Worker",
));
}
validate_feature_invocation(&entry.descriptor, invocation).map_err(|error| {
let mut message = error.to_string();
truncate_display(&mut message, MAX_INVOCATION_MESSAGE_BYTES);
FeatureInvocationHandlerError::failed(message)
})
}
pub fn validate(
&self,
invocation: &FeatureInvocation,
) -> Result<(), FeatureInvocationHandlerError> {
self.validate_metadata(invocation)?;
self.entries[&invocation.identity]
.handler
.validate(invocation)
.map_err(|mut error| {
// Validation is side-effect-free; bound its diagnostic before
// any start, without claiming an uncertain business outcome.
truncate_display(&mut error.message, MAX_INVOCATION_MESSAGE_BYTES);
error
})
}
pub async fn invoke(
&self,
invocation: &FeatureInvocation,
) -> Result<FeatureInvocationResult, FeatureInvocationHandlerError> {
self.validate(invocation)?;
let entry = &self.entries[&invocation.identity];
let result = entry
.handler
.invoke(
FeatureInvocationContext {
invocation_id: invocation.invocation_id.clone(),
},
invocation,
)
.await
.map_err(|error| {
if error.message.len() > MAX_INVOCATION_MESSAGE_BYTES {
FeatureInvocationHandlerError::outcome_unknown(
"handler error exceeded the result byte limit; do not retry automatically",
)
} else {
error
}
})?;
if !invocation_result_is_bounded(&result) {
return Err(FeatureInvocationHandlerError::outcome_unknown(
"handler result exceeded the message/context byte limit; do not retry automatically",
));
}
if result.invocation_id != invocation.invocation_id
|| result.identity != invocation.identity
{
return Err(FeatureInvocationHandlerError::outcome_unknown(
"handler returned a result for a different invocation; do not retry automatically",
));
}
if result.status != protocol::FeatureInvocationStatus::Succeeded && result.context.is_some()
{
return Err(FeatureInvocationHandlerError::outcome_unknown(
"handler returned preparation context without a successful result",
));
}
if result.status == protocol::FeatureInvocationStatus::Succeeded
&& !entry.handler.is_available()
{
return Err(FeatureInvocationHandlerError::failed(
"handler reported success but the invocation is no longer available; preparation context withheld; do not retry automatically",
));
}
Ok(result)
}
pub(crate) fn register(
&mut self,
feature_id: FeatureId,
descriptor: FeatureInvocationDescriptor,
handler: Arc<dyn FeatureInvocationHandler>,
completion_provider: Option<Arc<dyn FeatureInvocationCompletionProvider>>,
) -> Result<(), FeatureInstallError> {
if descriptor.identity.0.len() > 256
|| serde_json::to_vec(&descriptor).map_or(true, |encoded| {
encoded.len() > MAX_INVOCATION_DESCRIPTOR_BYTES
})
{
return Err(FeatureInstallError::InvalidDescriptor(
"invocation metadata exceeds the host byte limit".into(),
));
}
descriptor
.validate()
.map_err(|error| FeatureInstallError::InvalidDescriptor(error.to_string()))?;
let expected_prefix = format!("{feature_id}/");
if !descriptor.identity.0.starts_with(&expected_prefix) {
return Err(FeatureInstallError::InvalidDescriptor(format!(
"invocation identity `{}` must be qualified by feature `{feature_id}`",
descriptor.identity
)));
}
if self.entries.contains_key(&descriptor.identity) {
return Err(FeatureInstallError::DuplicateInvocationIdentity {
identity: descriptor.identity.to_string(),
});
}
let public_names = std::iter::once(descriptor.name.as_str())
.chain(descriptor.aliases.iter().map(String::as_str))
.collect::<BTreeSet<_>>();
for name in &public_names {
if let Some(existing) = self.names.get(*name) {
return Err(FeatureInstallError::DuplicateInvocationName {
name: (*name).to_string(),
first: existing.to_string(),
duplicate: descriptor.identity.to_string(),
});
}
}
for name in public_names {
self.names
.insert(name.to_string(), descriptor.identity.clone());
}
self.entries.insert(
descriptor.identity.clone(),
RegisteredInvocation {
descriptor,
handler,
completion_provider,
},
);
Ok(())
}
pub fn argument_value<'a>(
invocation: &'a FeatureInvocation,
name: &str,
) -> Option<&'a InvocationValue> {
invocation
.arguments
.iter()
.find(|argument| argument.name == name)
.map(|argument| &argument.value)
}
}
pub struct InvocationContributionRegistrar<'a> {
pub(super) feature_id: &'a FeatureId,
pub(super) declarations: &'a [FeatureInvocationDescriptor],
pub(super) registry: &'a mut FeatureInvocationRegistry,
pub(super) installed: &'a mut Vec<FeatureInvocationDescriptor>,
}
impl InvocationContributionRegistrar<'_> {
pub fn register<H>(
&mut self,
descriptor: FeatureInvocationDescriptor,
handler: H,
) -> Result<(), FeatureInstallError>
where
H: FeatureInvocationHandler + 'static,
{
self.register_with_provider(descriptor, Arc::new(handler), None)
}
pub fn register_with_completion<H, P>(
&mut self,
descriptor: FeatureInvocationDescriptor,
handler: H,
provider: P,
) -> Result<(), FeatureInstallError>
where
H: FeatureInvocationHandler + 'static,
P: FeatureInvocationCompletionProvider + 'static,
{
self.register_with_provider(descriptor, Arc::new(handler), Some(Arc::new(provider)))
}
fn register_with_provider(
&mut self,
descriptor: FeatureInvocationDescriptor,
handler: Arc<dyn FeatureInvocationHandler>,
provider: Option<Arc<dyn FeatureInvocationCompletionProvider>>,
) -> Result<(), FeatureInstallError> {
if provider.is_none()
&& descriptor.arguments.iter().any(|argument| {
matches!(argument.completion, InvocationCompletion::Provider { .. })
})
{
return Err(FeatureInstallError::InvalidDescriptor(format!(
"invocation `{}` declares a completion provider but did not register one",
descriptor.identity
)));
}
if !self.declarations.contains(&descriptor) {
return Err(FeatureInstallError::UndeclaredContribution {
kind: super::FeatureContributionKind::ChatInvocation,
name: descriptor.identity.to_string(),
feature: self.feature_id.to_string(),
});
}
self.registry.register(
self.feature_id.clone(),
descriptor.clone(),
handler,
provider,
)?;
self.installed.push(descriptor);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use protocol::{
FeatureInvocationStatus, FeatureInvocationSyntax, InvocationArgumentDescriptor,
InvocationArgumentType, InvocationArgumentValue,
};
use std::sync::atomic::{AtomicUsize, Ordering};
fn descriptor(feature: &str, name: &str) -> FeatureInvocationDescriptor {
FeatureInvocationDescriptor {
identity: FeatureInvocationIdentity(format!("{feature}/{name}")),
name: name.into(),
aliases: Vec::new(),
display_name: name.into(),
description: format!("invoke {name}"),
syntax: FeatureInvocationSyntax::Parenthesized,
arguments: vec![InvocationArgumentDescriptor {
name: "value".into(),
position: Some(0),
required: true,
value_type: InvocationArgumentType::String,
completion: InvocationCompletion::Static {
values: vec!["alpha".into(), "beta".into()],
},
description: Some("selected value".into()),
}],
client_adapter: None,
}
}
struct CountingHandler(Arc<AtomicUsize>);
#[async_trait]
impl FeatureInvocationHandler for CountingHandler {
async fn invoke(
&self,
_context: FeatureInvocationContext,
invocation: &FeatureInvocation,
) -> Result<FeatureInvocationResult, FeatureInvocationHandlerError> {
self.0.fetch_add(1, Ordering::SeqCst);
Ok(FeatureInvocationResult {
invocation_id: invocation.invocation_id.clone(),
identity: invocation.identity.clone(),
status: FeatureInvocationStatus::Succeeded,
message: "done".into(),
context: Some("prepared".into()),
})
}
}
fn invocation(descriptor: &FeatureInvocationDescriptor) -> FeatureInvocation {
FeatureInvocation {
invocation_id: "invoke-1".into(),
identity: descriptor.identity.clone(),
name: descriptor.name.clone(),
arguments: vec![InvocationArgumentValue {
name: "value".into(),
value: InvocationValue::String("alpha".into()),
}],
}
}
#[tokio::test]
async fn installed_registry_validates_before_calling_handler_and_returns_context() {
let calls = Arc::new(AtomicUsize::new(0));
let descriptor = descriptor("builtin:test", "run");
let mut registry = FeatureInvocationRegistry::default();
registry
.register(
FeatureId::builtin("test"),
descriptor.clone(),
Arc::new(CountingHandler(calls.clone())),
None,
)
.unwrap();
let result = registry.invoke(&invocation(&descriptor)).await.unwrap();
assert_eq!(result.status, FeatureInvocationStatus::Succeeded);
assert_eq!(result.context.as_deref(), Some("prepared"));
assert_eq!(calls.load(Ordering::SeqCst), 1);
let mut malformed = invocation(&descriptor);
malformed.arguments.clear();
assert!(registry.invoke(&malformed).await.is_err());
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[test]
fn registration_rejects_public_name_collisions_across_source_qualified_features() {
let calls = Arc::new(AtomicUsize::new(0));
let mut registry = FeatureInvocationRegistry::default();
registry
.register(
FeatureId::builtin("first"),
descriptor("builtin:first", "run"),
Arc::new(CountingHandler(calls.clone())),
None,
)
.unwrap();
let error = registry
.register(
FeatureId::builtin("second"),
descriptor("builtin:second", "run"),
Arc::new(CountingHandler(calls)),
None,
)
.unwrap_err();
assert!(matches!(
error,
FeatureInstallError::DuplicateInvocationName { .. }
));
}
#[tokio::test]
async fn client_local_adapter_metadata_cannot_authorize_worker_execution() {
let calls = Arc::new(AtomicUsize::new(0));
let descriptor = crate::feature::builtin::chat_invocation::attach_invocation_descriptor();
let mut registry = FeatureInvocationRegistry::default();
registry
.register(
FeatureId::builtin("attachments"),
descriptor.clone(),
Arc::new(CountingHandler(calls.clone())),
None,
)
.unwrap();
let invocation = FeatureInvocation {
invocation_id: "client-path".into(),
identity: descriptor.identity,
name: descriptor.name,
arguments: vec![InvocationArgumentValue {
name: "file".into(),
value: InvocationValue::String("/client/private/file".into()),
}],
};
assert!(registry.invoke(&invocation).await.is_err());
assert_eq!(calls.load(Ordering::SeqCst), 0);
}
struct CountingCompletion(Arc<AtomicUsize>);
#[async_trait]
impl FeatureInvocationCompletionProvider for CountingCompletion {
async fn complete(
&self,
_completion: &InvocationCompletion,
_argument: &str,
_prefix: &str,
) -> Vec<CompletionEntry> {
self.0.fetch_add(1, Ordering::SeqCst);
Vec::new()
}
}
#[tokio::test]
async fn completion_dispatch_uses_installed_argument_metadata_not_provider_presence() {
let calls = Arc::new(AtomicUsize::new(0));
let provider_calls = Arc::new(AtomicUsize::new(0));
let descriptor = descriptor("builtin:test", "run");
let mut registry = FeatureInvocationRegistry::default();
registry
.register(
FeatureId::builtin("test"),
descriptor.clone(),
Arc::new(CountingHandler(calls.clone())),
Some(Arc::new(CountingCompletion(provider_calls.clone()))),
)
.unwrap();
assert_eq!(
registry
.argument_completions(&descriptor.identity, "value", "a")
.await[0]
.value,
"alpha"
);
assert!(
registry
.argument_completions(
&FeatureInvocationIdentity("foreign:feature/run".into()),
"value",
""
)
.await
.is_empty()
);
assert!(
registry
.argument_completions(&descriptor.identity, "credential", "")
.await
.is_empty()
);
assert_eq!(provider_calls.load(Ordering::SeqCst), 0);
assert_eq!(calls.load(Ordering::SeqCst), 0);
}
struct AvailabilityHandler {
available: Arc<std::sync::atomic::AtomicBool>,
calls: Arc<AtomicUsize>,
}
#[async_trait]
impl FeatureInvocationHandler for AvailabilityHandler {
fn is_available(&self) -> bool {
self.available.load(Ordering::SeqCst)
}
async fn invoke(
&self,
context: FeatureInvocationContext,
invocation: &FeatureInvocation,
) -> Result<FeatureInvocationResult, FeatureInvocationHandlerError> {
CountingHandler(self.calls.clone())
.invoke(context, invocation)
.await
}
}
struct TestProvider {
entries: Vec<CompletionEntry>,
revoke: Option<Arc<std::sync::atomic::AtomicBool>>,
}
#[async_trait]
impl FeatureInvocationCompletionProvider for TestProvider {
async fn complete(
&self,
_completion: &InvocationCompletion,
_argument: &str,
_prefix: &str,
) -> Vec<CompletionEntry> {
if let Some(available) = &self.revoke {
available.store(false, Ordering::SeqCst);
}
self.entries.clone()
}
}
#[tokio::test]
async fn provider_output_is_bounded_and_cannot_inject_invocation_metadata() {
let available = Arc::new(std::sync::atomic::AtomicBool::new(true));
let mut descriptor = descriptor("builtin:test", "run");
descriptor.arguments[0].completion = InvocationCompletion::Provider {
provider: "choices".into(),
};
let entries = (0..128)
.map(|index| CompletionEntry {
value: format!("value-{index}"),
is_dir: false,
description: Some("é".repeat(MAX_COMPLETION_DESCRIPTION_BYTES)),
usage: Some("usage".repeat(MAX_COMPLETION_USAGE_BYTES)),
invocation: Some(descriptor.clone()),
})
.chain([CompletionEntry {
value: "x".repeat(MAX_COMPLETION_VALUE_BYTES + 1),
is_dir: false,
description: None,
usage: None,
invocation: None,
}])
.collect();
let mut registry = FeatureInvocationRegistry::default();
registry
.register(
FeatureId::builtin("test"),
descriptor.clone(),
Arc::new(AvailabilityHandler {
available,
calls: Arc::new(AtomicUsize::new(0)),
}),
Some(Arc::new(TestProvider {
entries,
revoke: None,
})),
)
.unwrap();
let candidates = registry
.argument_completions(&descriptor.identity, "value", "")
.await;
assert!(!candidates.is_empty());
assert!(candidates.len() <= MAX_COMPLETION_ENTRIES);
assert!(serde_json::to_vec(&candidates).unwrap().len() <= MAX_COMPLETION_RESPONSE_BYTES);
assert!(candidates.iter().all(|candidate| {
candidate.value.len() <= MAX_COMPLETION_VALUE_BYTES
&& candidate.invocation.is_none()
&& candidate.description.as_ref().unwrap().len() <= MAX_COMPLETION_DESCRIPTION_BYTES
&& candidate.usage.as_ref().unwrap().len() <= MAX_COMPLETION_USAGE_BYTES
}));
let simple = (0..128).map(|index| CompletionEntry {
value: format!("value-{index}"),
is_dir: false,
description: None,
usage: None,
invocation: None,
});
assert_eq!(
bounded_completion_entries(simple).len(),
MAX_COMPLETION_ENTRIES
);
assert!(
bounded_completion_entries([CompletionEntry {
value: "x".repeat(MAX_COMPLETION_VALUE_BYTES + 1),
is_dir: false,
description: None,
usage: None,
invocation: None,
}])
.is_empty()
);
}
#[tokio::test]
async fn provider_revocation_during_completion_discards_its_output() {
let available = Arc::new(std::sync::atomic::AtomicBool::new(true));
let calls = Arc::new(AtomicUsize::new(0));
let mut descriptor = descriptor("builtin:test", "run");
descriptor.arguments[0].completion = InvocationCompletion::Provider {
provider: "choices".into(),
};
let mut registry = FeatureInvocationRegistry::default();
registry
.register(
FeatureId::builtin("test"),
descriptor.clone(),
Arc::new(AvailabilityHandler {
available: available.clone(),
calls: calls.clone(),
}),
Some(Arc::new(TestProvider {
entries: vec![CompletionEntry {
value: "private".into(),
is_dir: false,
description: None,
usage: None,
invocation: None,
}],
revoke: Some(available),
})),
)
.unwrap();
assert!(
registry
.argument_completions(&descriptor.identity, "value", "")
.await
.is_empty()
);
assert!(registry.descriptors().is_empty());
assert!(registry.invoke(&invocation(&descriptor)).await.is_err());
assert_eq!(calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn enum_and_boolean_completions_derive_from_type_schema_when_no_provider_declared() {
let calls = Arc::new(AtomicUsize::new(0));
for (value_type, prefix, expected) in [
(InvocationArgumentType::Boolean, "f", "false"),
(
InvocationArgumentType::Enum {
values: vec!["read".into(), "write".into()],
},
"w",
"write",
),
] {
let mut descriptor = descriptor("builtin:test", "run");
descriptor.arguments[0].value_type = value_type;
descriptor.arguments[0].completion = InvocationCompletion::None;
let mut registry = FeatureInvocationRegistry::default();
registry
.register(
FeatureId::builtin("test"),
descriptor.clone(),
Arc::new(CountingHandler(calls.clone())),
None,
)
.unwrap();
let candidates = registry
.argument_completions(&descriptor.identity, "value", prefix)
.await;
assert_eq!(candidates.len(), 1);
assert_eq!(candidates[0].value, expected);
}
assert_eq!(calls.load(Ordering::SeqCst), 0);
}
#[test]
fn oversized_installed_invocation_metadata_is_rejected() {
let mut descriptor = descriptor("builtin:test", "run");
descriptor.description = "x".repeat(MAX_INVOCATION_DESCRIPTOR_BYTES);
let mut registry = FeatureInvocationRegistry::default();
assert!(
registry
.register(
FeatureId::builtin("test"),
descriptor,
Arc::new(CountingHandler(Arc::new(AtomicUsize::new(0)))),
None,
)
.is_err()
);
assert!(registry.descriptors().is_empty());
}
#[tokio::test]
async fn metadata_and_argument_completion_are_side_effect_free() {
let calls = Arc::new(AtomicUsize::new(0));
let descriptor = descriptor("builtin:test", "run");
let mut registry = FeatureInvocationRegistry::default();
registry
.register(
FeatureId::builtin("test"),
descriptor.clone(),
Arc::new(CountingHandler(calls.clone())),
None,
)
.unwrap();
assert_eq!(registry.feature_completions("ru")[0].value, "run");
assert_eq!(
registry
.argument_completions(&descriptor.identity, "value", "a")
.await[0]
.value,
"alpha"
);
assert_eq!(calls.load(Ordering::SeqCst), 0);
}
}
crates/worker/src/ipc/interceptor.rs
+20 -2
View File
@@ -399,8 +399,23 @@ impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
.lock()
.expect("pending_attachments poisoned"),
);
// Preparations can commit before Engine receives the user prompt. The
// annotation queue alone does not materialize their model-visible items:
// return them through the normal prompt lifecycle, without recommitting.
// The annotator consumes their original IDs/provenance as Engine appends.
let mut prepared_items = self
.pending_committed_history
.lock()
.expect("pending committed history poisoned")
.iter()
.map(|entry| entry.item.clone())
.collect::<Vec<_>>();
Ok(if extras.is_empty() {
PromptAction::Continue
if prepared_items.is_empty() {
PromptAction::Continue
} else {
PromptAction::ContinueWith(prepared_items)
}
} else {
// Commit the typed system items first, then hand the
// matching `Item::system_message`s to the worker. Sync
@@ -410,7 +425,10 @@ impl Interceptor<SessionHistoryMetadata> for WorkerInterceptor {
self.attach_prompt_provenance(&mut extras);
let items: Vec<Item> = extras.iter().map(SystemItem::to_history_item).collect();
match self.commit_system_items(&extras) {
Ok(()) => PromptAction::ContinueWith(items),
Ok(()) => {
prepared_items.extend(items);
PromptAction::ContinueWith(prepared_items)
}
Err(error) => PromptAction::Cancel(format!("session persistence failed: {error}")),
}
})
crates/worker/src/ipc/protocol_session.rs
+16 -3
View File
@@ -121,9 +121,22 @@ pub async fn dispatch_worker_protocol_method(
method: Method,
) -> Option<Event> {
match method {
Method::ListCompletions { kind, prefix } => {
let entries = handle.completion_entries(kind, &prefix).await;
Some(Event::Completions { kind, entries })
Method::ListCompletions {
request_id,
kind,
prefix,
context,
} => {
let entries = handle
.completion_entries(kind, &prefix, context.as_ref())
.await;
Some(Event::Completions {
request_id,
kind,
prefix,
context,
entries,
})
}
method => {
let _ = handle.send(method).await;
crates/worker/src/shared_state.rs
+16
View File
@@ -48,6 +48,7 @@ pub struct WorkerSharedState {
/// `ListCompletions` queries without going through the controller. It is
/// unset only in unit tests that construct `WorkerSharedState` directly.
fs_view: OnceLock<WorkerFsView>,
feature_invocations: RwLock<crate::feature::FeatureInvocationRegistry>,
flow_transition_enabled: AtomicBool,
}
@@ -66,6 +67,7 @@ impl WorkerSharedState {
state: RwLock::new(WorkerStateSnapshot::initial()),
accepted_commands: RwLock::new(VecDeque::new()),
fs_view: OnceLock::new(),
feature_invocations: RwLock::new(crate::feature::FeatureInvocationRegistry::default()),
flow_transition_enabled: AtomicBool::new(false),
}
}
@@ -82,6 +84,20 @@ impl WorkerSharedState {
self.fs_view.get()
}
pub fn feature_invocations(&self) -> crate::feature::FeatureInvocationRegistry {
self.feature_invocations
.read()
.expect("feature invocation registry lock poisoned")
.clone()
}
pub fn set_feature_invocations(&self, registry: crate::feature::FeatureInvocationRegistry) {
*self
.feature_invocations
.write()
.expect("feature invocation registry lock poisoned") = registry;
}
pub fn enable_flow_transition(&self) {
self.flow_transition_enabled.store(true, Ordering::Release);
}
crates/worker/src/worker.rs
+889 -3
View File
File diff suppressed because it is too large Load Diff
crates/worker/tests/feature_invocation_test.rs
+581
View File
@@ -0,0 +1,581 @@
//! Installed host-feature execution, not lexical slash-command dispatch.
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use agen::Engine;
use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
use agen::llm_client::{ClientError, LlmClient, Request};
use async_trait::async_trait;
use futures::Stream;
use protocol::{
FeatureInvocation, FeatureInvocationDescriptor, FeatureInvocationIdentity,
FeatureInvocationResult, FeatureInvocationStatus, FeatureInvocationSyntax,
InvocationArgumentDescriptor, InvocationArgumentType, InvocationArgumentValue,
InvocationCompletion, InvocationValue, Segment,
};
use session_store::{CombinedStore, FsStore, FsWorkerStore, LogEntry, SystemItem};
use worker::feature::{
FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureInvocationContext,
FeatureInvocationHandler, FeatureInvocationHandlerError, FeatureModule, FeatureRegistryBuilder,
};
use worker::{Worker, WorkerFilesystemAuthority, WorkerManifest, WorkerWorkspaceContext};
#[derive(Clone)]
struct Trace {
calls: Arc<Mutex<Vec<String>>>,
requests: Arc<Mutex<Vec<Request>>>,
available: Arc<AtomicBool>,
}
impl Default for Trace {
fn default() -> Self {
Self {
calls: Arc::default(),
requests: Arc::default(),
available: Arc::new(AtomicBool::new(true)),
}
}
}
#[async_trait]
impl LlmClient for Trace {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(
&self,
request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
self.calls.lock().unwrap().push("llm".into());
self.requests.lock().unwrap().push(request);
Ok(Box::pin(futures::stream::iter(vec![
Ok(LlmEvent::text_block_start(0)),
Ok(LlmEvent::text_delta(0, "done")),
Ok(LlmEvent::text_block_stop(0, None)),
Ok(LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
})),
])))
}
}
fn descriptor() -> FeatureInvocationDescriptor {
FeatureInvocationDescriptor {
identity: FeatureInvocationIdentity("builtin:fake/prepare".into()),
name: "prepare".into(),
aliases: vec!["prep".into()],
display_name: "Prepare".into(),
description: "Prepare test context".into(),
syntax: FeatureInvocationSyntax::Parenthesized,
arguments: vec![InvocationArgumentDescriptor {
name: "value".into(),
position: Some(0),
required: true,
value_type: InvocationArgumentType::String,
completion: InvocationCompletion::None,
description: None,
}],
client_adapter: None,
}
}
#[derive(Clone)]
struct FakeFeature(Trace);
impl FeatureModule for FakeFeature {
fn descriptor(&self) -> FeatureDescriptor {
FeatureDescriptor::builtin("fake", "Fake").with_chat_invocation(descriptor())
}
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
context
.chat_invocations()
.register(descriptor(), self.clone())
}
}
#[async_trait]
impl FeatureInvocationHandler for FakeFeature {
fn is_available(&self) -> bool {
self.0.available.load(Ordering::SeqCst)
}
fn validate(
&self,
invocation: &FeatureInvocation,
) -> Result<(), FeatureInvocationHandlerError> {
if invocation.arguments[0].value == InvocationValue::String("denied".into()) {
return Err(FeatureInvocationHandlerError::failed(
"host permission denied",
));
}
Ok(())
}
async fn invoke(
&self,
context: FeatureInvocationContext,
invocation: &FeatureInvocation,
) -> Result<FeatureInvocationResult, FeatureInvocationHandlerError> {
assert_eq!(context.invocation_id, invocation.invocation_id);
let InvocationValue::String(value) = &invocation.arguments[0].value else {
panic!("validated string")
};
self.0.calls.lock().unwrap().push(value.clone());
match value.as_str() {
"fail" => {
return Err(FeatureInvocationHandlerError::failed(
"test preparation failed",
));
}
"unknown" => {
return Err(FeatureInvocationHandlerError::outcome_unknown(
"test effect uncertain",
));
}
"revoke-during-effect" => self.0.available.store(false, Ordering::SeqCst),
"large-error" => return Err(FeatureInvocationHandlerError::failed("x".repeat(4097))),
_ => {}
}
Ok(FeatureInvocationResult {
invocation_id: if value == "wrong" {
"wrong-id".into()
} else {
invocation.invocation_id.clone()
},
identity: invocation.identity.clone(),
status: FeatureInvocationStatus::Succeeded,
message: match value.as_str() {
"large-message" => "é".repeat(2049),
"limit-message" => "x".repeat(4096),
_ => "prepared".into(),
},
context: Some(match value.as_str() {
"large-context" => "é".repeat(8193),
"limit-context" => "x".repeat(16384),
_ => format!("trusted-preparation-{value}"),
}),
})
}
}
fn invoke(id: &str, value: &str) -> Segment {
Segment::FeatureInvoke {
invocation: FeatureInvocation {
invocation_id: id.into(),
identity: descriptor().identity,
name: "prepare".into(),
arguments: vec![InvocationArgumentValue {
name: "value".into(),
value: InvocationValue::String(value.into()),
}],
},
}
}
type TestStore = CombinedStore<FsStore, FsWorkerStore>;
async fn fixture(trace: Trace) -> (tempfile::TempDir, Worker<Trace, TestStore>) {
let dir = tempfile::tempdir().unwrap();
let manifest = WorkerManifest::from_toml(
r#"
[worker]
name = "invocation-test"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
max_tokens = 100
[scope]
allow = []
"#,
)
.unwrap();
let store = CombinedStore::new(
FsStore::new(dir.path().join("sessions")).unwrap(),
FsWorkerStore::new(dir.path().join("workers")).unwrap(),
);
let mut worker = Worker::new(
manifest,
Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(
trace.clone(),
),
store,
WorkerWorkspaceContext::no_workspace(),
WorkerFilesystemAuthority::None,
manifest::Scope::empty(),
)
.await
.unwrap();
let report =
worker.install_features(FeatureRegistryBuilder::new().with_module(FakeFeature(trace)));
assert!(!report.has_errors(), "{report:?}");
(dir, worker)
}
fn results(worker: &Worker<Trace, TestStore>) -> Vec<FeatureInvocationResult> {
worker
.sink()
.subscribe_with_snapshot()
.0
.into_iter()
.filter_map(|entry| match entry {
LogEntry::AnnotatedSystemItem { entry, .. } => match entry.item {
SystemItem::FeatureInvocationResult { result, .. } => Some(result),
_ => None,
},
_ => None,
})
.collect()
}
#[derive(Clone)]
struct TracePromptHook {
trace: Trace,
cancel: bool,
}
#[async_trait]
impl worker::hook::Hook<worker::hook::OnPromptSubmit> for TracePromptHook {
async fn call(
&self,
_input: &worker::hook::PromptSubmitInfo,
) -> Result<worker::hook::HookPromptAction, worker::hook::HookError> {
self.trace.calls.lock().unwrap().push("prompt-hook".into());
Ok(if self.cancel {
worker::hook::HookPromptAction::Cancel("test cancellation".into())
} else {
worker::hook::HookPromptAction::Continue
})
}
}
#[tokio::test]
async fn installed_feature_invocations_execute_in_order_before_model_and_context_is_visible() {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
worker.add_on_prompt_submit_hook(TracePromptHook {
trace: trace.clone(),
cancel: false,
});
worker
.run(vec![
invoke("one", "first"),
Segment::text("explain"),
invoke("two", "second"),
])
.await
.unwrap();
assert_eq!(
*trace.calls.lock().unwrap(),
["first", "second", "prompt-hook", "llm"]
);
let requests = trace.requests.lock().unwrap();
let request = format!("{:?}", requests[0]);
assert!(request.contains("trusted-preparation-first"));
assert!(request.contains("trusted-preparation-second"));
assert!(request.find("trusted-preparation-first") < request.find("trusted-preparation-second"));
assert_eq!(results(&worker).len(), 2);
// Preparation results must be model-visible in the same order as replay.
let history = worker.history();
let first = history
.iter()
.position(|item| {
item.as_text()
.is_some_and(|text| text.contains("trusted-preparation-first"))
})
.unwrap();
let second = history
.iter()
.position(|item| {
item.as_text()
.is_some_and(|text| text.contains("trusted-preparation-second"))
})
.unwrap();
assert!(first < second);
}
#[tokio::test]
async fn feature_failure_fences_later_invocations_model_resume_and_notifications() {
for value in ["fail", "unknown", "wrong"] {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
assert!(
worker
.run(vec![
invoke("one", "first"),
invoke("two", value),
invoke("three", "never")
])
.await
.is_err()
);
assert_eq!(*trace.calls.lock().unwrap(), ["first", value]);
let results = results(&worker);
assert_eq!(results.len(), 3);
assert_eq!(results[0].status, FeatureInvocationStatus::Succeeded);
assert_eq!(
results[1].status,
if value == "fail" {
FeatureInvocationStatus::Failed
} else {
FeatureInvocationStatus::OutcomeUnknown
}
);
assert_eq!(results[1].invocation_id, "two");
assert!(results[1].context.is_none());
assert_eq!(results[2].status, FeatureInvocationStatus::Failed);
assert!(worker.resume().await.is_err());
assert!(
worker
.run_for_notification(protocol::InvokeKind::Notify)
.await
.is_err()
);
assert!(trace.requests.lock().unwrap().is_empty());
}
}
#[tokio::test]
async fn feature_replay_reuses_result_and_rejects_id_payload_changes() {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
worker.run(vec![invoke("same", "first")]).await.unwrap();
worker.run(vec![invoke("same", "first")]).await.unwrap();
assert_eq!(*trace.calls.lock().unwrap(), ["first", "llm", "llm"]);
assert!(worker.run(vec![invoke("same", "different")]).await.is_err());
assert_eq!(*trace.calls.lock().unwrap(), ["first", "llm", "llm"]);
assert_eq!(
results(&worker).last().unwrap().status,
FeatureInvocationStatus::Failed
);
}
#[tokio::test]
async fn feature_success_replay_does_not_resurrect_disabled_installed_capability() {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
worker.run(vec![invoke("same", "first")]).await.unwrap();
assert!(
!worker
.install_features(FeatureRegistryBuilder::new())
.has_errors()
);
assert!(
worker
.feature_invocations()
.feature_completions("")
.is_empty()
);
assert!(worker.run(vec![invoke("same", "first")]).await.is_err());
assert_eq!(*trace.calls.lock().unwrap(), ["first", "llm"]);
assert_eq!(
results(&worker).last().unwrap().status,
FeatureInvocationStatus::Failed
);
}
#[tokio::test]
async fn feature_invocation_text_is_not_executable_and_duplicate_ids_are_rejected_before_effects() {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
worker
.run(vec![Segment::text("/prepare(first) quoted /prep(fail)")])
.await
.unwrap();
assert_eq!(*trace.calls.lock().unwrap(), ["llm"]);
assert!(results(&worker).is_empty());
assert!(
worker
.run(vec![invoke("same", "first"), invoke("same", "second")])
.await
.is_err()
);
assert_eq!(*trace.calls.lock().unwrap(), ["llm"]);
}
#[tokio::test]
async fn feature_handler_validation_rejects_semantics_and_forged_authority_before_effects() {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
assert!(worker.run(vec![invoke("denied", "denied")]).await.is_err());
let mut forged = invoke("forged", "first");
if let Segment::FeatureInvoke { invocation } = &mut forged {
invocation.arguments.push(InvocationArgumentValue {
name: "workspace_id".into(),
value: InvocationValue::String("foreign".into()),
});
}
assert!(worker.run(vec![forged]).await.is_err());
assert!(trace.calls.lock().unwrap().is_empty());
assert_eq!(results(&worker).len(), 2);
assert!(
results(&worker)
.iter()
.all(|result| result.status == FeatureInvocationStatus::Failed)
);
}
#[tokio::test]
async fn feature_prompt_hook_cancellation_does_not_rollback_durable_business_preparation() {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
worker.add_on_prompt_submit_hook(TracePromptHook {
trace: trace.clone(),
cancel: true,
});
let _ = worker.run(vec![invoke("one", "first")]).await;
assert_eq!(*trace.calls.lock().unwrap(), ["first", "prompt-hook"]);
assert_eq!(results(&worker).len(), 1);
assert!(worker.history().iter().any(|item| {
item.as_text()
.is_some_and(|text| text.contains("trusted-preparation-first"))
}));
let _ = worker.run(vec![invoke("one", "first")]).await;
assert_eq!(
*trace.calls.lock().unwrap(),
["first", "prompt-hook", "prompt-hook"]
);
}
#[tokio::test]
async fn feature_unknown_input_rejected_before_effects_commit_or_model() {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
let unknown: Segment = serde_json::from_value(serde_json::json!({
"kind": "future_side_effect", "args": { "target": "important" }
}))
.unwrap();
assert!(matches!(unknown, Segment::Unknown));
let before = worker.sink().subscribe_with_snapshot().0.len();
assert!(
worker
.run(vec![
invoke("one", "first"),
unknown,
Segment::text("explain")
])
.await
.is_err()
);
assert_eq!(worker.sink().subscribe_with_snapshot().0.len(), before);
assert!(trace.calls.lock().unwrap().is_empty());
}
#[tokio::test]
async fn feature_oversized_results_become_bounded_unknown_and_do_not_repeat_effects() {
for value in ["large-message", "large-context", "large-error"] {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
assert!(
worker
.run(vec![invoke("large", value), invoke("later", "never")])
.await
.is_err()
);
assert_eq!(*trace.calls.lock().unwrap(), [value]);
for result in results(&worker) {
assert!(result.message.len() <= 4096);
assert!(result.context.is_none());
}
assert_eq!(
results(&worker)[0].status,
FeatureInvocationStatus::OutcomeUnknown
);
assert!(worker.run(vec![invoke("large", value)]).await.is_err());
assert!(worker.resume().await.is_err());
assert_eq!(*trace.calls.lock().unwrap(), [value]);
}
for value in ["limit-message", "limit-context"] {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
worker.run(vec![invoke("limit", value)]).await.unwrap();
worker.run(vec![invoke("limit", value)]).await.unwrap();
assert_eq!(*trace.calls.lock().unwrap(), [value, "llm", "llm"]);
}
}
#[tokio::test]
async fn feature_dynamic_permission_revocation_hides_discovery_and_fences_execution_and_replay() {
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
let registry = worker.feature_invocations();
assert_eq!(registry.descriptors().len(), 1);
worker.run(vec![invoke("same", "first")]).await.unwrap();
trace.available.store(false, Ordering::SeqCst);
assert!(registry.descriptors().is_empty());
assert!(registry.descriptor(&descriptor().identity).is_none());
assert!(registry.feature_completions("").is_empty());
assert!(
registry
.argument_name_completions(&descriptor().identity, "")
.is_empty()
);
assert!(
registry
.argument_completions(&descriptor().identity, "value", "")
.await
.is_empty()
);
assert!(worker.run(vec![invoke("same", "first")]).await.is_err());
assert!(worker.run(vec![invoke("new", "second")]).await.is_err());
assert!(worker.resume().await.is_err());
assert_eq!(*trace.calls.lock().unwrap(), ["first", "llm"]);
trace.available.store(true, Ordering::SeqCst);
assert_eq!(registry.descriptors().len(), 1);
worker.run(vec![invoke("same", "first")]).await.unwrap();
assert_eq!(*trace.calls.lock().unwrap(), ["first", "llm", "llm"]);
}
#[tokio::test]
async fn feature_revocation_during_effect_withholds_preparation_without_retrying_committed_effect()
{
let trace = Trace::default();
let (_dir, mut worker) = fixture(trace.clone()).await;
assert!(
worker
.run(vec![invoke("same", "revoke-during-effect")])
.await
.is_err()
);
assert_eq!(*trace.calls.lock().unwrap(), ["revoke-during-effect"]);
let result = &results(&worker)[0];
assert_eq!(result.status, FeatureInvocationStatus::Failed);
assert!(result.context.is_none());
assert!(result.message.contains("reported success"));
trace.available.store(true, Ordering::SeqCst);
assert!(
worker
.run(vec![invoke("same", "revoke-during-effect")])
.await
.is_err()
);
assert_eq!(*trace.calls.lock().unwrap(), ["revoke-during-effect"]);
}
#[test]
fn notify_rejects_typed_side_effect_message_payloads() {
let methods = [
protocol::Method::Notify {
notification_request_id: "n-1".into(),
message: "/prepare(first)".into(),
},
protocol::Method::NotifyTracked {
notification_request_id: "n-1".into(),
message: "/prepare(first)".into(),
source: protocol::AuthenticatedInputSource::UntrustedWire,
},
];
for method in methods {
// Notify accepts an advisory string, not typed input or invocation
// objects. Ordinary text spelling an invocation remains advisory text.
let mut request = serde_json::to_value(method).unwrap();
assert!(serde_json::from_value::<protocol::Method>(request.clone()).is_ok());
let mut with_input = request.clone();
with_input["params"]["input"] = serde_json::json!([invoke("one", "first")]);
assert!(
serde_json::from_value::<protocol::Method>(with_input).is_err(),
"Notify must reject typed input rather than silently discard it"
);
request["params"]["message"] = serde_json::json!([invoke("one", "first")]);
assert!(serde_json::from_value::<protocol::Method>(request).is_err());
}
}