use super::{
message::MAX_MESSAGE_BYTES,
model::{
AgentDescriptor, AgentId, AgentStatus, AgentUpdate, MessageId, MessagePriority,
MessagePurpose, agent_prompt,
},
runtime::{AgentDirectoryEntry, AgentSummary, OutputContract, Registry, forward_events},
};
use nanocodex::{
Model, Tool,
tools::{
ToolsBuilder,
contract::{ToolContext, ToolDefinition, ToolInput, ToolOutput, ToolResult, async_trait},
},
};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::{
io,
sync::{Arc, Weak},
time::Duration,
};
use tokio::sync::oneshot;
const DEFAULT_WAIT_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_WAIT_TIMEOUT: Duration = Duration::from_secs(300);
const SPAWN_AGENT_TOOL: &str = "spawn_agent";
const SUBMIT_RESULT_TOOL: &str = "submit_result";
const SEND_AGENT_MESSAGE_TOOL: &str = "send_agent_message";
const LIST_AGENTS_TOOL: &str = "list_agents";
const WAIT_AGENT_TOOL: &str = "wait_agent";
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct AgentTask {
role: String,
task: String,
model: SubagentModel,
output_schema: Value,
}
#[derive(Clone, Copy, Deserialize)]
#[serde(rename_all = "snake_case")]
enum SubagentModel {
Selected,
Luna,
}
impl SubagentModel {
fn resolve(self, selected: Model, allow_luna: bool) -> Result<Model, io::Error> {
match self {
Self::Selected => Ok(selected),
Self::Luna if allow_luna => Ok(Model::Luna),
Self::Luna => Err(io::Error::new(
io::ErrorKind::InvalidInput,
"Luna subagents are disabled by `subagents.allow_luna`",
)),
}
}
}
#[derive(Serialize)]
struct AgentStartReport {
agent_id: AgentId,
model: Model,
role: String,
status: AgentStatus,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct WaitTask {
agent_ids: Vec<AgentId>,
#[serde(default)]
timeout_ms: Option<u64>,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct TargetAgent {
agent_id: AgentId,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct DirectoryTask {
#[serde(default)]
include_completed: bool,
#[serde(default)]
include_self: bool,
}
#[derive(Serialize)]
struct AgentDirectory {
agents: Vec<AgentDirectoryEntry>,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct SendMessageTask {
agent_id: AgentId,
message: String,
#[serde(default)]
priority: MessagePriority,
#[serde(default)]
purpose: MessagePurpose,
#[serde(default)]
in_reply_to: Option<MessageId>,
}
#[derive(Serialize)]
struct WaitReport {
agents: Vec<AgentSummary>,
timed_out: bool,
}
#[derive(Serialize)]
struct LifecycleReport {
agents: Vec<AgentSummary>,
}
fn json_output(value: &impl Serialize) -> ToolResult {
Ok(ToolOutput::from_json(serde_json::to_value(value)?, true))
}
struct SpawnAgent {
registry: Weak<Registry>,
selected_model: Model,
allow_luna: bool,
}
#[async_trait]
impl Tool for SpawnAgent {
fn definition(&self) -> ToolDefinition {
let (models, model_description) = if self.allow_luna {
(
json!(["selected", "luna"]),
"Use `selected` for the session's selected model or `luna` when low latency matters more than reasoning capability.",
)
} else {
(
json!(["selected"]),
"Use `selected` for the session's selected model.",
)
};
ToolDefinition::function(
SPAWN_AGENT_TOOL,
"Starts a reusable clean-room subagent without inherited conversation history and immediately returns its ID.",
json!({
"type": "object",
"properties": {
"role": {
"type": "string",
"description": "A short role describing the subagent's specialty."
},
"task": {
"type": "string",
"description": "A complete, focused task for the subagent."
},
"model": {
"type": "string",
"enum": models,
"description": model_description
},
"output_schema": {
"description": "The JSON Schema that every successful result from this agent must satisfy. Use an object with one string field for a free-form report."
}
},
"required": ["role", "task", "model", "output_schema"],
"additionalProperties": false
}),
)
.with_output_schema(spawn_agent_output_schema())
}
async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
let AgentTask {
role,
task,
model,
output_schema,
} = input.decode_json()?;
let model = model.resolve(self.selected_model, self.allow_luna)?;
let contract = OutputContract::compile(&output_schema)?;
let registry = self
.registry
.upgrade()
.ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
let capacity = registry.reserve_turn()?;
let reservation = registry.reserve(context.session_id()).await?;
let id = reservation.id;
let (child, events) = registry.spawn_agent(model)?;
let session_id = child.session_id().to_string();
let descriptor = AgentDescriptor {
id,
session_id,
model,
role: role.clone(),
task: task.clone(),
parent: reservation.parent,
};
let (start_events, events_ready) = oneshot::channel();
let event_task = forward_events(
reservation.root_session_id.clone(),
id,
events,
events_ready,
Arc::downgrade(®istry),
registry.updates.clone(),
);
registry
.insert(
reservation.root_session_id.clone(),
descriptor.clone(),
child,
event_task,
contract,
)
.await?;
registry.send(&reservation.root_session_id, AgentUpdate::Added(descriptor));
let _ = start_events.send(());
registry
.launch_initial_turn(
&reservation.root_session_id,
id,
agent_prompt(id, &task),
capacity,
)
.await?;
json_output(&AgentStartReport {
agent_id: id,
model,
role,
status: AgentStatus::Running,
})
}
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct SubmitResultArgs {
turn_token: u64,
output: Value,
}
struct SubmitResult {
registry: Weak<Registry>,
}
#[async_trait]
impl Tool for SubmitResult {
fn definition(&self) -> ToolDefinition {
ToolDefinition::function(
SUBMIT_RESULT_TOOL,
"Submits the current subagent turn's final JSON output. Call exactly once with a value matching the output schema in the task prompt. Invalid values can be corrected and retried.",
json!({
"type": "object",
"properties": {
"output": {
"description": "The final JSON value required by this agent's output schema."
},
"turn_token": {
"type": "integer",
"minimum": 1,
"description": "The current turn token stated in the task prompt."
}
},
"required": ["turn_token", "output"],
"additionalProperties": false
}),
)
.with_output_schema(json!({
"type": "object",
"properties": {
"accepted": { "type": "boolean", "const": true }
},
"required": ["accepted"],
"additionalProperties": false
}))
}
async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
let SubmitResultArgs { turn_token, output } = input.decode_json()?;
let registry = self
.registry
.upgrade()
.ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
registry
.submit_result(context.session_id(), turn_token, output)
.await?;
Ok(ToolOutput::from_json(json!({ "accepted": true }), true))
}
}
struct SendAgentMessage {
registry: Weak<Registry>,
}
#[async_trait]
impl Tool for SendAgentMessage {
fn definition(&self) -> ToolDefinition {
ToolDefinition::function(
SEND_AGENT_MESSAGE_TOOL,
"Sends a bounded directed message to any other agent in the same task tree. Deferred messages start an idle agent or queue behind its active turn. If a send is queued, do not wait for it inside the current turn; finish the turn so queued messages can be delivered. Urgent messages steer a running agent at its next safe model boundary. Delegate messages replace the recipient's assigned task, retain its output schema, and require management authority.",
json!({
"type": "object",
"properties": {
"agent_id": {
"type": "integer",
"minimum": 1,
"description": "The recipient from list_agents. Any non-closing agent in the same task tree can receive coordination messages."
},
"message": {
"type": "string",
"minLength": 1,
"maxLength": MAX_MESSAGE_BYTES,
"description": "The focused message body. The runtime enforces a 2048-byte UTF-8 limit."
},
"priority": {
"type": "string",
"enum": ["deferred", "urgent"],
"default": "deferred",
"description": "Urgent steers an active turn; deferred preserves turn boundaries. A queued deferred send requires the current turn to finish before delivery."
},
"purpose": {
"type": "string",
"enum": ["delegate", "coordinate", "finding", "question", "reply"],
"default": "coordinate",
"description": "A typed coordination intent. Delegate is restricted to agents the sender can manage."
},
"in_reply_to": {
"type": "integer",
"minimum": 1,
"description": "A message ID from the same two-party thread. Replies must reverse the original direction."
}
},
"required": ["agent_id", "message"],
"additionalProperties": false
}),
)
}
async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
let SendMessageTask {
agent_id,
message,
priority,
purpose,
in_reply_to,
} = input.decode_json()?;
let registry = self
.registry
.upgrade()
.ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
let receipt = registry
.send_message(
context.session_id(),
agent_id,
priority,
purpose,
in_reply_to,
message,
)
.await?;
json_output(&receipt)
}
}
struct ListAgents {
registry: Weak<Registry>,
}
#[async_trait]
impl Tool for ListAgents {
fn definition(&self) -> ToolDefinition {
ToolDefinition::function(
LIST_AGENTS_TOOL,
"Lists a compact directory of agents in the same task tree. Active recipients are returned by default; completed agents can be included when a follow-up message is needed.",
json!({
"type": "object",
"properties": {
"include_completed": {
"type": "boolean",
"default": false,
"description": "Includes completed, interrupted, failed, and closed agents."
},
"include_self": {
"type": "boolean",
"default": false,
"description": "Includes the calling agent for topology inspection. Self-messaging remains unavailable."
}
},
"additionalProperties": false
}),
)
}
async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
let DirectoryTask {
include_completed,
include_self,
} = input.decode_json()?;
let registry = self
.registry
.upgrade()
.ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
json_output(&AgentDirectory {
agents: registry
.directory(context.session_id(), include_completed, include_self)
.await,
})
}
}
struct WaitAgent {
registry: Weak<Registry>,
}
#[async_trait]
impl Tool for WaitAgent {
fn definition(&self) -> ToolDefinition {
ToolDefinition::function(
WAIT_AGENT_TOOL,
"Waits until any requested subagent reaches a terminal status and returns current statuses and reports. Use one call with multiple IDs instead of polling the workspace.",
json!({
"type": "object",
"properties": {
"agent_ids": {
"type": "array",
"items": { "type": "integer", "minimum": 1 },
"minItems": 1,
"description": "Agent IDs returned by spawn_agent. Waiting returns when any one becomes terminal."
},
"timeout_ms": {
"type": "integer",
"minimum": 1,
"maximum": 300000,
"description": "Bounded wait in milliseconds. Defaults to 30000."
}
},
"required": ["agent_ids"],
"additionalProperties": false
}),
)
.with_output_schema(wait_agent_output_schema())
}
async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
let WaitTask {
agent_ids,
timeout_ms,
} = input.decode_json()?;
let registry = self
.registry
.upgrade()
.ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
let duration = timeout_ms
.map(Duration::from_millis)
.unwrap_or(DEFAULT_WAIT_TIMEOUT)
.min(MAX_WAIT_TIMEOUT);
let (agents, timed_out) = registry
.wait(context.session_id(), &agent_ids, duration)
.await?;
json_output(&WaitReport { agents, timed_out })
}
}
#[derive(Clone, Copy)]
enum LifecycleOperation {
Interrupt,
Close,
}
struct ChangeAgentLifecycle {
registry: Weak<Registry>,
operation: LifecycleOperation,
}
impl ChangeAgentLifecycle {
fn tool_name(&self) -> &'static str {
match self.operation {
LifecycleOperation::Interrupt => "interrupt_agent",
LifecycleOperation::Close => "close_agent",
}
}
}
#[async_trait]
impl Tool for ChangeAgentLifecycle {
fn definition(&self) -> ToolDefinition {
let description = match self.operation {
LifecycleOperation::Interrupt => {
"Interrupts an agent's active turn and every active descendant, waits for their model and tool resources to stop, and keeps the sessions reusable."
}
LifecycleOperation::Close => {
"Closes an agent and its entire descendant subtree, waiting for active model and tool resources to stop before returning. Closed agents remain inspectable but are not reusable."
}
};
ToolDefinition::function(
self.tool_name(),
description,
json!({
"type": "object",
"properties": {
"agent_id": {
"type": "integer",
"minimum": 1,
"description": "The root of the subagent subtree to stop."
}
},
"required": ["agent_id"],
"additionalProperties": false
}),
)
}
async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
let TargetAgent { agent_id } = input.decode_json()?;
let registry = self
.registry
.upgrade()
.ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
let agents = match self.operation {
LifecycleOperation::Interrupt => {
registry.interrupt(context.session_id(), agent_id).await?
}
LifecycleOperation::Close => registry.close(context.session_id(), agent_id).await?,
};
json_output(&LifecycleReport { agents })
}
}
impl super::runtime::WeakSubagents {
pub fn install_tools(
&self,
tools: ToolsBuilder,
selected_model: Model,
allow_luna: bool,
) -> ToolsBuilder {
let registry = self.registry.clone();
tools
.tool(SpawnAgent {
registry: registry.clone(),
selected_model,
allow_luna,
})
.tool(SubmitResult {
registry: registry.clone(),
})
.tool(SendAgentMessage {
registry: registry.clone(),
})
.tool(ListAgents {
registry: registry.clone(),
})
.tool(WaitAgent {
registry: registry.clone(),
})
.tool(ChangeAgentLifecycle {
registry: registry.clone(),
operation: LifecycleOperation::Interrupt,
})
.tool(ChangeAgentLifecycle {
registry,
operation: LifecycleOperation::Close,
})
}
}
fn spawn_agent_output_schema() -> Value {
json!({
"type": "object",
"properties": {
"agent_id": { "type": "integer" },
"model": { "type": "string" },
"role": { "type": "string" },
"status": {
"type": "object",
"properties": { "state": { "type": "string", "const": "running" } },
"required": ["state"],
"additionalProperties": false
}
},
"required": ["agent_id", "model", "role", "status"],
"additionalProperties": false
})
}
fn wait_agent_output_schema() -> Value {
json!({
"type": "object",
"properties": {
"agents": {
"type": "array",
"items": {
"type": "object",
"properties": {
"agent_id": { "type": "integer" },
"model": { "type": "string" },
"role": { "type": "string" },
"task": { "type": "string" },
"parent_agent_id": { "type": ["integer", "null"] },
"status": agent_status_schema(),
"last_output": {}
},
"required": ["agent_id", "model", "role", "task", "parent_agent_id", "status"],
"additionalProperties": false
}
},
"timed_out": { "type": "boolean" }
},
"required": ["agents", "timed_out"],
"additionalProperties": false
})
}
fn agent_status_schema() -> Value {
let state_only = ["pending", "running", "interrupted", "closing", "closed"].map(|state| {
json!({
"type": "object",
"properties": { "state": { "type": "string", "const": state } },
"required": ["state"],
"additionalProperties": false
})
});
let mut variants = state_only.into_iter().collect::<Vec<_>>();
variants.push(json!({
"type": "object",
"properties": {
"state": { "type": "string", "const": "completed" },
"output": {}
},
"required": ["state", "output"],
"additionalProperties": false
}));
variants.push(json!({
"type": "object",
"properties": {
"state": { "type": "string", "const": "failed" },
"error": { "type": "string" }
},
"required": ["state", "error"],
"additionalProperties": false
}));
json!({ "oneOf": variants })
}
#[cfg(test)]
mod tests {
use super::{SendAgentMessage, SpawnAgent, SubagentModel, SubmitResult, WaitAgent};
use crate::runtime::Registry;
use nanocodex::{Model, Tool};
use serde_json::json;
use std::sync::Weak;
#[test]
fn spawn_agent_requires_an_explicit_bounded_model_choice() {
let definition = SpawnAgent {
registry: Weak::<Registry>::new(),
selected_model: Model::Terra,
allow_luna: true,
}
.definition();
let parameters = definition.parameters().unwrap().as_value();
let output = definition.output_schema().unwrap();
assert_eq!(
parameters["properties"]["model"]["enum"],
json!(["selected", "luna"])
);
assert!(
parameters["required"]
.as_array()
.unwrap()
.contains(&json!("model"))
);
assert_eq!(
SubagentModel::Selected.resolve(Model::Terra, true).unwrap(),
Model::Terra
);
assert_eq!(
SubagentModel::Luna.resolve(Model::Terra, true).unwrap(),
Model::Luna
);
assert!(
output.as_value()["required"]
.as_array()
.unwrap()
.contains(&json!("model"))
);
}
#[test]
fn spawn_agent_excludes_and_rejects_luna_when_disabled() {
let definition = SpawnAgent {
registry: Weak::<Registry>::new(),
selected_model: Model::Terra,
allow_luna: false,
}
.definition();
let parameters = definition.parameters().unwrap().as_value();
assert_eq!(
parameters["properties"]["model"]["enum"],
json!(["selected"])
);
assert!(
SubagentModel::Luna
.resolve(Model::Terra, false)
.unwrap_err()
.to_string()
.contains("subagents.allow_luna")
);
}
#[test]
fn send_message_definition_names_deferred_delivery_and_queued_waiting() {
let definition = SendAgentMessage {
registry: Weak::<Registry>::new(),
}
.definition();
let priority = &definition.parameters().unwrap().as_value()["properties"]["priority"];
assert_eq!(priority["enum"], json!(["deferred", "urgent"]));
assert_eq!(priority["default"], json!("deferred"));
assert!(definition.description().contains("do not wait"));
assert!(definition.description().contains("finish the turn"));
}
#[test]
fn submit_result_requires_the_turn_token_and_one_output_value() {
let definition = SubmitResult {
registry: Weak::<Registry>::new(),
}
.definition();
let parameters = definition.parameters().unwrap().as_value();
assert_eq!(parameters["required"], json!(["turn_token", "output"]));
assert_eq!(parameters["additionalProperties"], json!(false));
assert_eq!(parameters["properties"].as_object().unwrap().len(), 2);
}
#[test]
fn wait_agent_only_refers_to_clean_spawns() {
let definition = WaitAgent {
registry: Weak::<Registry>::new(),
}
.definition();
let description =
&definition.parameters().unwrap().as_value()["properties"]["agent_ids"]["description"];
let output = definition.output_schema().unwrap();
let agent = &output.as_value()["properties"]["agents"]["items"];
assert!(description.as_str().unwrap().contains("spawn_agent"));
assert!(!description.as_str().unwrap().contains("fork_agent"));
assert_eq!(agent["properties"]["model"], json!({ "type": "string" }));
assert!(
agent["required"]
.as_array()
.unwrap()
.contains(&json!("model"))
);
}
}