feat: gate built-in tools by profile features

This commit is contained in:
2026-06-09 21:05:38 +09:00
parent 41133e0cd5
commit f0f6cc92d8
19 changed files with 833 additions and 109 deletions
+70 -34
View File
@@ -4,6 +4,7 @@ use std::sync::atomic::Ordering;
use llm_worker::WorkerError;
use llm_worker::llm_client::client::LlmClient;
use manifest::TicketFeatureAccessConfig;
use pod_store::PodMetadataStore;
use session_store::Store;
use tokio::sync::{broadcast, mpsc, oneshot};
@@ -223,7 +224,7 @@ impl PodController {
runtime_dir.socket_path(),
runtime_base.to_path_buf(),
spawned_registry.clone(),
);
)?;
// Intake role Pods self-terminate only after a successful
// TicketIntakeReady turn has fully settled back to Idle. The request
@@ -499,7 +500,7 @@ fn register_pod_tools<C, St>(
spawner_socket: PathBuf,
runtime_base: PathBuf,
spawned_registry: Arc<SpawnedPodRegistry>,
) -> tools::ScopedFs
) -> std::io::Result<tools::ScopedFs>
where
C: LlmClient + Clone + 'static,
St: Store + PodMetadataStore + Clone + 'static,
@@ -513,6 +514,7 @@ where
let session_id_for_usage = pod.segment_id().to_string();
let memory_config = pod.manifest().memory.clone();
let web_config = pod.manifest().web.clone();
let feature_config = pod.manifest().feature.clone();
let spawner_name = pod.manifest().pod.name.clone();
let spawner_manifest = pod.manifest().clone();
let prompts = pod.prompts().clone();
@@ -534,24 +536,47 @@ where
fs,
tracker.clone(),
bash_output_dir,
web_config,
));
if feature_config.web.enabled {
pod.worker_mut()
.register_tools(tools::web_builtin_tools(web_config));
}
let mut feature_registry = FeatureRegistryBuilder::new();
feature_registry.add_module(task_feature);
feature_registry.add_module(crate::feature::builtin::ticket_tools_feature(
&workspace_root,
));
if feature_config.task.enabled {
feature_registry.add_module(task_feature);
}
if feature_config.ticket.enabled || feature_config.ticket_orchestration.enabled {
let ticket_access = match feature_config.ticket.access {
TicketFeatureAccessConfig::ReadOnly => {
crate::feature::builtin::ticket::TicketFeatureAccess::ReadOnly
}
TicketFeatureAccessConfig::Lifecycle => {
crate::feature::builtin::ticket::TicketFeatureAccess::Lifecycle
}
};
feature_registry.add_module(
crate::feature::builtin::ticket::ticket_tools_feature_with_options(
&workspace_root,
feature_config.ticket.enabled.then_some(ticket_access),
feature_config.ticket_orchestration.enabled,
),
);
}
let _feature_install_report = pod.install_features(feature_registry);
let worker = pod.worker_mut();
// Memory subsystem opt-in. When `[memory]` is present in the
// manifest, register the memory-specific Read/Write/Edit tools that
// target `<workspace>/memory/` and `<workspace>/knowledge/` with
// their built-in linter. Companion deny rules on the generic CRUD
// scope were already applied during `Pod::from_manifest`.
if let Some(mem) = memory_config.as_ref() {
// Memory tools require both explicit feature exposure and memory storage
// configuration. This keeps resident-memory config separate from the
// model-visible Memory*/Knowledge* tool surface.
if feature_config.memory.enabled {
let mem = memory_config.as_ref().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"[feature.memory].enabled = true requires a [memory] configuration section",
)
})?;
let layout = memory::WorkspaceLayout::resolve(mem, &workspace_root);
let query_cfg = memory::tool::QueryConfig::from(mem);
worker.register_tool(memory::tool::read_tool_with_usage(
@@ -567,28 +592,39 @@ where
// Pod-orchestration tools (SpawnPod + the four comm tools) share
// the Pod-scoped `SpawnedPodRegistry` (also consumed by the main
// loop's `PodEvent` handler).
worker.register_tool(spawn_pod_tool(
spawner_name.clone(),
spawner_socket,
runtime_base.clone(),
workspace_root.clone(),
pwd.clone(),
spawned_registry.clone(),
self_parent_socket,
spawner_manifest,
scope_handle,
prompts,
));
worker.register_tool(send_to_pod_tool(spawned_registry.clone()));
worker.register_tool(read_pod_output_tool(spawned_registry.clone()));
worker.register_tool(stop_pod_tool(spawned_registry.clone()));
let discovery = PodDiscovery::new(pod_store, spawner_name, runtime_base, pwd, spawned_registry);
worker.register_tool(list_pods_tool(discovery.clone()));
worker.register_tool(restore_pod_tool(discovery.clone()));
worker.register_tool(send_to_peer_pod_tool(discovery));
// loop's `PodEvent` handler). Expose them only behind the explicit
// profile feature and require delegation authority up front so enabling
// the surface cannot imply broad child scope by accident.
if feature_config.pod_management.enabled {
if spawner_manifest.delegation_scope.allow.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"[feature.pod_management].enabled = true requires non-empty [[delegation_scope.allow]]",
));
}
worker.register_tool(spawn_pod_tool(
spawner_name.clone(),
spawner_socket,
runtime_base.clone(),
workspace_root.clone(),
pwd.clone(),
spawned_registry.clone(),
self_parent_socket,
spawner_manifest,
scope_handle,
prompts,
));
worker.register_tool(send_to_pod_tool(spawned_registry.clone()));
worker.register_tool(read_pod_output_tool(spawned_registry.clone()));
worker.register_tool(stop_pod_tool(spawned_registry.clone()));
let discovery =
PodDiscovery::new(pod_store, spawner_name, runtime_base, pwd, spawned_registry);
worker.register_tool(list_pods_tool(discovery.clone()));
worker.register_tool(restore_pod_tool(discovery.clone()));
worker.register_tool(send_to_peer_pod_tool(discovery));
}
pod.attach_tracker(tracker);
fs_for_view
Ok(fs_for_view)
}
/// Idle/Paused event loop. Each iteration either fires a staged
+1
View File
@@ -10,4 +10,5 @@ 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,
};
+119 -15
View File
@@ -9,7 +9,11 @@ use std::path::{Path, PathBuf};
use ticket::{
LocalTicketBackend,
config::{DEFAULT_TICKET_BACKEND_RELATIVE_PATH, TicketConfig},
tool::{TICKET_READ_ONLY_TOOL_NAMES, TICKET_TOOL_NAMES, ticket_tools},
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_tools,
},
};
use crate::feature::{
@@ -32,10 +36,17 @@ pub enum TicketFeatureAccess {
}
impl TicketFeatureAccess {
pub fn tool_names(self) -> &'static [&'static str] {
pub fn base_tool_names(self) -> &'static [&'static str] {
match self {
Self::ReadOnly => &TICKET_READ_ONLY_TOOL_NAMES,
Self::Lifecycle => &TICKET_TOOL_NAMES,
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,
}
}
}
@@ -46,6 +57,8 @@ pub struct TicketFeature {
record_language: Option<String>,
config_error: Option<String>,
access: TicketFeatureAccess,
include_base_tools: bool,
include_orchestration_tools: bool,
}
impl TicketFeature {
@@ -54,11 +67,21 @@ impl TicketFeature {
}
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: access.unwrap_or(TicketFeatureAccess::Lifecycle),
include_base_tools: access.is_some(),
include_orchestration_tools,
}
}
@@ -69,22 +92,36 @@ impl TicketFeature {
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_access(backend_root, access);
let mut feature =
Self::new_with_options(backend_root, access, include_orchestration_tools);
feature.record_language = record_language;
feature
}
Err(error) => Self {
backend_root: workspace.join(DEFAULT_TICKET_BACKEND_RELATIVE_PATH),
record_language: None,
config_error: Some(error.to_string()),
access,
},
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,
}
}
}
}
@@ -96,6 +133,23 @@ impl TicketFeature {
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 authority(&self) -> HostAuthority {
HostAuthority::TicketBackend {
root: self.backend_root.display().to_string(),
@@ -122,7 +176,8 @@ impl FeatureModule for TicketFeature {
self.authority(),
AUTHORITY_REASON,
));
for name in self.access.tool_names() {
let enabled_tool_names = self.enabled_tool_names();
for name in &enabled_tool_names {
descriptor = descriptor.with_tool(ToolDeclaration::new(*name, tool_description(name)));
}
descriptor
@@ -152,12 +207,15 @@ impl FeatureModule for TicketFeature {
let authority = self.authority();
let backend = LocalTicketBackend::new(usable_root)
.with_record_language(self.record_language.as_deref());
let allowed_tool_names = self.access.tool_names();
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.contains(&name.as_str()) {
if !allowed_tool_names
.iter()
.any(|allowed| *allowed == name.as_str())
{
continue;
}
tools.register(
@@ -211,12 +269,24 @@ pub fn ticket_tools_feature_with_access(
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();
@@ -269,6 +339,40 @@ mod tests {
assert_eq!(descriptor.requested_host_authorities.len(), 1);
}
#[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();
+1
View File
@@ -775,6 +775,7 @@ fn manifest_to_reusable_config(manifest: &PodManifest) -> PodManifestConfig {
default_action: Some(p.default_action),
rules: p.rules.clone(),
}),
feature: manifest.feature.clone().into(),
compaction: manifest
.compaction
.as_ref()
+121
View File
@@ -215,6 +215,127 @@ async fn wait_for_status(handle: &PodHandle, status: PodStatus) {
// ---------------------------------------------------------------------------
fn request_tool_names(request: &Request) -> Vec<String> {
let mut names = request
.tools
.iter()
.map(|tool| tool.name.clone())
.collect::<Vec<_>>();
names.sort();
names
}
async fn wait_for_captured_request(client: &MockClient) -> Request {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
let requests = client.captured_requests();
if let Some(request) = requests.into_iter().next() {
return request;
}
assert!(
tokio::time::Instant::now() < deadline,
"timed out waiting for captured LLM request"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
}
#[tokio::test]
async fn feature_flags_default_to_core_tool_surface_only() {
let client = MockClient::new(simple_text_events());
let client_for_assert = client.clone();
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
handle.send(Method::run_text("Hello")).await.unwrap();
wait_for_status(&handle, PodStatus::Idle).await;
let request = wait_for_captured_request(&client_for_assert).await;
let names = request_tool_names(&request);
assert_eq!(names, vec!["Bash", "Edit", "Glob", "Grep", "Read", "Write"]);
assert!(!names.iter().any(|name| name == "TaskCreate"));
assert!(!names.iter().any(|name| name == "WebSearch"));
assert!(!names.iter().any(|name| name == "SpawnPod"));
}
#[tokio::test]
async fn enabled_task_and_web_features_register_their_tools() {
let manifest = r#"
[pod]
name = "feature-test-pod"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[worker]
max_tokens = 100
[feature.task]
enabled = true
[feature.web]
enabled = true
[web]
enabled = false
[[scope.allow]]
target = "./"
permission = "write"
"#;
let client = MockClient::new(simple_text_events());
let client_for_assert = client.clone();
let pod = make_pod_with_pwd_and_manifest(client, manifest).await.0;
let handle = spawn_controller(pod).await;
handle.send(Method::run_text("Hello")).await.unwrap();
wait_for_status(&handle, PodStatus::Idle).await;
let request = wait_for_captured_request(&client_for_assert).await;
let names = request_tool_names(&request);
assert!(names.iter().any(|name| name == "TaskCreate"));
assert!(names.iter().any(|name| name == "TaskUpdate"));
assert!(names.iter().any(|name| name == "WebSearch"));
assert!(names.iter().any(|name| name == "WebFetch"));
assert!(!names.iter().any(|name| name == "SpawnPod"));
assert!(!names.iter().any(|name| name == "MemoryRead"));
}
#[tokio::test]
async fn pod_management_feature_requires_delegation_scope() {
let manifest = r#"
[pod]
name = "pod-management-feature-test"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[worker]
max_tokens = 100
[feature.pod_management]
enabled = true
[[scope.allow]]
target = "./"
permission = "write"
"#;
let client = MockClient::new(simple_text_events());
let pod = make_pod_with_pwd_and_manifest(client, manifest).await.0;
let tmp = tempfile::tempdir().unwrap();
let result = PodController::spawn(pod, tmp.path()).await;
assert!(result.is_err());
let message = result.err().unwrap().to_string();
assert!(
message.contains("[feature.pod_management].enabled = true requires non-empty"),
"unexpected error: {message}"
);
}
#[tokio::test]
async fn run_end_returns_to_idle_without_busy_status() {
let client = MockClient::new(simple_text_events());