1use super::{
2 message::MAX_MESSAGE_BYTES,
3 model::{
4 AgentDescriptor, AgentId, AgentStatus, AgentUpdate, MessageId, MessagePriority,
5 MessagePurpose, agent_prompt,
6 },
7 runtime::{AgentDirectoryEntry, AgentSummary, OutputContract, Registry, forward_events},
8};
9use nanocodex::{
10 Model, Tool,
11 tools::{
12 ToolsBuilder,
13 contract::{ToolContext, ToolDefinition, ToolInput, ToolOutput, ToolResult, async_trait},
14 },
15};
16use serde::{Deserialize, Serialize};
17use serde_json::{Value, json};
18use std::{
19 io,
20 sync::{Arc, Weak},
21 time::Duration,
22};
23use tokio::sync::oneshot;
24
25const DEFAULT_WAIT_TIMEOUT: Duration = Duration::from_secs(30);
26const MAX_WAIT_TIMEOUT: Duration = Duration::from_secs(300);
27const SPAWN_AGENT_TOOL: &str = "spawn_agent";
28const SUBMIT_RESULT_TOOL: &str = "submit_result";
29const SEND_AGENT_MESSAGE_TOOL: &str = "send_agent_message";
30const LIST_AGENTS_TOOL: &str = "list_agents";
31const WAIT_AGENT_TOOL: &str = "wait_agent";
32
33#[derive(Deserialize)]
34#[serde(deny_unknown_fields)]
35struct AgentTask {
36 role: String,
37 task: String,
38 model: SubagentModel,
39 output_schema: Value,
40}
41
42#[derive(Clone, Copy, Deserialize)]
43#[serde(rename_all = "snake_case")]
44enum SubagentModel {
45 Selected,
46 Luna,
47}
48
49impl SubagentModel {
50 fn resolve(self, selected: Model, allow_luna: bool) -> Result<Model, io::Error> {
51 match self {
52 Self::Selected => Ok(selected),
53 Self::Luna if allow_luna => Ok(Model::Luna),
54 Self::Luna => Err(io::Error::new(
55 io::ErrorKind::InvalidInput,
56 "Luna subagents are disabled by `subagents.allow_luna`",
57 )),
58 }
59 }
60}
61
62#[derive(Serialize)]
63struct AgentStartReport {
64 agent_id: AgentId,
65 model: Model,
66 role: String,
67 status: AgentStatus,
68}
69
70#[derive(Deserialize)]
71#[serde(deny_unknown_fields)]
72struct WaitTask {
73 agent_ids: Vec<AgentId>,
74 #[serde(default)]
75 timeout_ms: Option<u64>,
76}
77
78#[derive(Deserialize)]
79#[serde(deny_unknown_fields)]
80struct TargetAgent {
81 agent_id: AgentId,
82}
83
84#[derive(Deserialize)]
85#[serde(deny_unknown_fields)]
86struct DirectoryTask {
87 #[serde(default)]
88 include_completed: bool,
89 #[serde(default)]
90 include_self: bool,
91}
92
93#[derive(Serialize)]
94struct AgentDirectory {
95 agents: Vec<AgentDirectoryEntry>,
96}
97
98#[derive(Deserialize)]
99#[serde(deny_unknown_fields)]
100struct SendMessageTask {
101 agent_id: AgentId,
102 message: String,
103 #[serde(default)]
104 priority: MessagePriority,
105 #[serde(default)]
106 purpose: MessagePurpose,
107 #[serde(default)]
108 in_reply_to: Option<MessageId>,
109}
110
111#[derive(Serialize)]
112struct WaitReport {
113 agents: Vec<AgentSummary>,
114 timed_out: bool,
115}
116
117#[derive(Serialize)]
118struct LifecycleReport {
119 agents: Vec<AgentSummary>,
120}
121
122fn json_output(value: &impl Serialize) -> ToolResult {
123 Ok(ToolOutput::from_json(serde_json::to_value(value)?, true))
124}
125
126struct SpawnAgent {
127 registry: Weak<Registry>,
128 selected_model: Model,
129 allow_luna: bool,
130}
131
132#[async_trait]
133impl Tool for SpawnAgent {
134 fn definition(&self) -> ToolDefinition {
135 let (models, model_description) = if self.allow_luna {
136 (
137 json!(["selected", "luna"]),
138 "Use `selected` for the session's selected model or `luna` when low latency matters more than reasoning capability.",
139 )
140 } else {
141 (
142 json!(["selected"]),
143 "Use `selected` for the session's selected model.",
144 )
145 };
146 ToolDefinition::function(
147 SPAWN_AGENT_TOOL,
148 "Starts a reusable clean-room subagent without inherited conversation history and immediately returns its ID.",
149 json!({
150 "type": "object",
151 "properties": {
152 "role": {
153 "type": "string",
154 "description": "A short role describing the subagent's specialty."
155 },
156 "task": {
157 "type": "string",
158 "description": "A complete, focused task for the subagent."
159 },
160 "model": {
161 "type": "string",
162 "enum": models,
163 "description": model_description
164 },
165 "output_schema": {
166 "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."
167 }
168 },
169 "required": ["role", "task", "model", "output_schema"],
170 "additionalProperties": false
171 }),
172 )
173 .with_output_schema(spawn_agent_output_schema())
174 }
175
176 async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
177 let AgentTask {
178 role,
179 task,
180 model,
181 output_schema,
182 } = input.decode_json()?;
183 let model = model.resolve(self.selected_model, self.allow_luna)?;
184 let contract = OutputContract::compile(&output_schema)?;
185 let registry = self
186 .registry
187 .upgrade()
188 .ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
189 let capacity = registry.reserve_turn()?;
190 let reservation = registry.reserve(context.session_id()).await?;
191 let id = reservation.id;
192 let (child, events) = registry.spawn_agent(model)?;
193 let session_id = child.session_id().to_string();
194 let descriptor = AgentDescriptor {
195 id,
196 session_id,
197 model,
198 role: role.clone(),
199 task: task.clone(),
200 parent: reservation.parent,
201 };
202 let (start_events, events_ready) = oneshot::channel();
203 let event_task = forward_events(
204 reservation.root_session_id.clone(),
205 id,
206 events,
207 events_ready,
208 Arc::downgrade(®istry),
209 registry.updates.clone(),
210 );
211 registry
212 .insert(
213 reservation.root_session_id.clone(),
214 descriptor.clone(),
215 child,
216 event_task,
217 contract,
218 )
219 .await?;
220 registry.send(&reservation.root_session_id, AgentUpdate::Added(descriptor));
221 let _ = start_events.send(());
222
223 registry
224 .launch_initial_turn(
225 &reservation.root_session_id,
226 id,
227 agent_prompt(id, &task),
228 capacity,
229 )
230 .await?;
231 json_output(&AgentStartReport {
232 agent_id: id,
233 model,
234 role,
235 status: AgentStatus::Running,
236 })
237 }
238}
239
240#[derive(Deserialize)]
241#[serde(deny_unknown_fields)]
242struct SubmitResultArgs {
243 turn_token: u64,
244 output: Value,
245}
246
247struct SubmitResult {
248 registry: Weak<Registry>,
249}
250
251#[async_trait]
252impl Tool for SubmitResult {
253 fn definition(&self) -> ToolDefinition {
254 ToolDefinition::function(
255 SUBMIT_RESULT_TOOL,
256 "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.",
257 json!({
258 "type": "object",
259 "properties": {
260 "output": {
261 "description": "The final JSON value required by this agent's output schema."
262 },
263 "turn_token": {
264 "type": "integer",
265 "minimum": 1,
266 "description": "The current turn token stated in the task prompt."
267 }
268 },
269 "required": ["turn_token", "output"],
270 "additionalProperties": false
271 }),
272 )
273 .with_output_schema(json!({
274 "type": "object",
275 "properties": {
276 "accepted": { "type": "boolean", "const": true }
277 },
278 "required": ["accepted"],
279 "additionalProperties": false
280 }))
281 }
282
283 async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
284 let SubmitResultArgs { turn_token, output } = input.decode_json()?;
285 let registry = self
286 .registry
287 .upgrade()
288 .ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
289 registry
290 .submit_result(context.session_id(), turn_token, output)
291 .await?;
292 Ok(ToolOutput::from_json(json!({ "accepted": true }), true))
293 }
294}
295
296struct SendAgentMessage {
297 registry: Weak<Registry>,
298}
299
300#[async_trait]
301impl Tool for SendAgentMessage {
302 fn definition(&self) -> ToolDefinition {
303 ToolDefinition::function(
304 SEND_AGENT_MESSAGE_TOOL,
305 "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.",
306 json!({
307 "type": "object",
308 "properties": {
309 "agent_id": {
310 "type": "integer",
311 "minimum": 1,
312 "description": "The recipient from list_agents. Any non-closing agent in the same task tree can receive coordination messages."
313 },
314 "message": {
315 "type": "string",
316 "minLength": 1,
317 "maxLength": MAX_MESSAGE_BYTES,
318 "description": "The focused message body. The runtime enforces a 2048-byte UTF-8 limit."
319 },
320 "priority": {
321 "type": "string",
322 "enum": ["deferred", "urgent"],
323 "default": "deferred",
324 "description": "Urgent steers an active turn; deferred preserves turn boundaries. A queued deferred send requires the current turn to finish before delivery."
325 },
326 "purpose": {
327 "type": "string",
328 "enum": ["delegate", "coordinate", "finding", "question", "reply"],
329 "default": "coordinate",
330 "description": "A typed coordination intent. Delegate is restricted to agents the sender can manage."
331 },
332 "in_reply_to": {
333 "type": "integer",
334 "minimum": 1,
335 "description": "A message ID from the same two-party thread. Replies must reverse the original direction."
336 }
337 },
338 "required": ["agent_id", "message"],
339 "additionalProperties": false
340 }),
341 )
342 }
343
344 async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
345 let SendMessageTask {
346 agent_id,
347 message,
348 priority,
349 purpose,
350 in_reply_to,
351 } = input.decode_json()?;
352 let registry = self
353 .registry
354 .upgrade()
355 .ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
356 let receipt = registry
357 .send_message(
358 context.session_id(),
359 agent_id,
360 priority,
361 purpose,
362 in_reply_to,
363 message,
364 )
365 .await?;
366 json_output(&receipt)
367 }
368}
369
370struct ListAgents {
371 registry: Weak<Registry>,
372}
373
374#[async_trait]
375impl Tool for ListAgents {
376 fn definition(&self) -> ToolDefinition {
377 ToolDefinition::function(
378 LIST_AGENTS_TOOL,
379 "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.",
380 json!({
381 "type": "object",
382 "properties": {
383 "include_completed": {
384 "type": "boolean",
385 "default": false,
386 "description": "Includes completed, interrupted, failed, and closed agents."
387 },
388 "include_self": {
389 "type": "boolean",
390 "default": false,
391 "description": "Includes the calling agent for topology inspection. Self-messaging remains unavailable."
392 }
393 },
394 "additionalProperties": false
395 }),
396 )
397 }
398
399 async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
400 let DirectoryTask {
401 include_completed,
402 include_self,
403 } = input.decode_json()?;
404 let registry = self
405 .registry
406 .upgrade()
407 .ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
408 json_output(&AgentDirectory {
409 agents: registry
410 .directory(context.session_id(), include_completed, include_self)
411 .await,
412 })
413 }
414}
415
416struct WaitAgent {
417 registry: Weak<Registry>,
418}
419
420#[async_trait]
421impl Tool for WaitAgent {
422 fn definition(&self) -> ToolDefinition {
423 ToolDefinition::function(
424 WAIT_AGENT_TOOL,
425 "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.",
426 json!({
427 "type": "object",
428 "properties": {
429 "agent_ids": {
430 "type": "array",
431 "items": { "type": "integer", "minimum": 1 },
432 "minItems": 1,
433 "description": "Agent IDs returned by spawn_agent. Waiting returns when any one becomes terminal."
434 },
435 "timeout_ms": {
436 "type": "integer",
437 "minimum": 1,
438 "maximum": 300000,
439 "description": "Bounded wait in milliseconds. Defaults to 30000."
440 }
441 },
442 "required": ["agent_ids"],
443 "additionalProperties": false
444 }),
445 )
446 .with_output_schema(wait_agent_output_schema())
447 }
448
449 async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
450 let WaitTask {
451 agent_ids,
452 timeout_ms,
453 } = input.decode_json()?;
454 let registry = self
455 .registry
456 .upgrade()
457 .ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
458 let duration = timeout_ms
459 .map(Duration::from_millis)
460 .unwrap_or(DEFAULT_WAIT_TIMEOUT)
461 .min(MAX_WAIT_TIMEOUT);
462 let (agents, timed_out) = registry
463 .wait(context.session_id(), &agent_ids, duration)
464 .await?;
465 json_output(&WaitReport { agents, timed_out })
466 }
467}
468
469#[derive(Clone, Copy)]
470enum LifecycleOperation {
471 Interrupt,
472 Close,
473}
474
475struct ChangeAgentLifecycle {
476 registry: Weak<Registry>,
477 operation: LifecycleOperation,
478}
479
480impl ChangeAgentLifecycle {
481 fn tool_name(&self) -> &'static str {
482 match self.operation {
483 LifecycleOperation::Interrupt => "interrupt_agent",
484 LifecycleOperation::Close => "close_agent",
485 }
486 }
487}
488
489#[async_trait]
490impl Tool for ChangeAgentLifecycle {
491 fn definition(&self) -> ToolDefinition {
492 let description = match self.operation {
493 LifecycleOperation::Interrupt => {
494 "Interrupts an agent's active turn and every active descendant, waits for their model and tool resources to stop, and keeps the sessions reusable."
495 }
496 LifecycleOperation::Close => {
497 "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."
498 }
499 };
500 ToolDefinition::function(
501 self.tool_name(),
502 description,
503 json!({
504 "type": "object",
505 "properties": {
506 "agent_id": {
507 "type": "integer",
508 "minimum": 1,
509 "description": "The root of the subagent subtree to stop."
510 }
511 },
512 "required": ["agent_id"],
513 "additionalProperties": false
514 }),
515 )
516 }
517
518 async fn execute(&self, input: ToolInput, context: ToolContext<'_>) -> ToolResult {
519 let TargetAgent { agent_id } = input.decode_json()?;
520 let registry = self
521 .registry
522 .upgrade()
523 .ok_or_else(|| std::io::Error::other("subagent runtime is closed"))?;
524 let agents = match self.operation {
525 LifecycleOperation::Interrupt => {
526 registry.interrupt(context.session_id(), agent_id).await?
527 }
528 LifecycleOperation::Close => registry.close(context.session_id(), agent_id).await?,
529 };
530 json_output(&LifecycleReport { agents })
531 }
532}
533
534impl super::runtime::WeakSubagents {
535 pub fn install_tools(
540 &self,
541 tools: ToolsBuilder,
542 selected_model: Model,
543 allow_luna: bool,
544 ) -> ToolsBuilder {
545 let registry = self.registry.clone();
546 tools
547 .tool(SpawnAgent {
548 registry: registry.clone(),
549 selected_model,
550 allow_luna,
551 })
552 .tool(SubmitResult {
553 registry: registry.clone(),
554 })
555 .tool(SendAgentMessage {
556 registry: registry.clone(),
557 })
558 .tool(ListAgents {
559 registry: registry.clone(),
560 })
561 .tool(WaitAgent {
562 registry: registry.clone(),
563 })
564 .tool(ChangeAgentLifecycle {
565 registry: registry.clone(),
566 operation: LifecycleOperation::Interrupt,
567 })
568 .tool(ChangeAgentLifecycle {
569 registry,
570 operation: LifecycleOperation::Close,
571 })
572 }
573}
574
575fn spawn_agent_output_schema() -> Value {
576 json!({
577 "type": "object",
578 "properties": {
579 "agent_id": { "type": "integer" },
580 "model": { "type": "string" },
581 "role": { "type": "string" },
582 "status": {
583 "type": "object",
584 "properties": { "state": { "type": "string", "const": "running" } },
585 "required": ["state"],
586 "additionalProperties": false
587 }
588 },
589 "required": ["agent_id", "model", "role", "status"],
590 "additionalProperties": false
591 })
592}
593
594fn wait_agent_output_schema() -> Value {
595 json!({
596 "type": "object",
597 "properties": {
598 "agents": {
599 "type": "array",
600 "items": {
601 "type": "object",
602 "properties": {
603 "agent_id": { "type": "integer" },
604 "model": { "type": "string" },
605 "role": { "type": "string" },
606 "task": { "type": "string" },
607 "parent_agent_id": { "type": ["integer", "null"] },
608 "status": agent_status_schema(),
609 "last_output": {}
610 },
611 "required": ["agent_id", "model", "role", "task", "parent_agent_id", "status"],
612 "additionalProperties": false
613 }
614 },
615 "timed_out": { "type": "boolean" }
616 },
617 "required": ["agents", "timed_out"],
618 "additionalProperties": false
619 })
620}
621
622fn agent_status_schema() -> Value {
623 let state_only = ["pending", "running", "interrupted", "closing", "closed"].map(|state| {
624 json!({
625 "type": "object",
626 "properties": { "state": { "type": "string", "const": state } },
627 "required": ["state"],
628 "additionalProperties": false
629 })
630 });
631 let mut variants = state_only.into_iter().collect::<Vec<_>>();
632 variants.push(json!({
633 "type": "object",
634 "properties": {
635 "state": { "type": "string", "const": "completed" },
636 "output": {}
637 },
638 "required": ["state", "output"],
639 "additionalProperties": false
640 }));
641 variants.push(json!({
642 "type": "object",
643 "properties": {
644 "state": { "type": "string", "const": "failed" },
645 "error": { "type": "string" }
646 },
647 "required": ["state", "error"],
648 "additionalProperties": false
649 }));
650 json!({ "oneOf": variants })
651}
652
653#[cfg(test)]
654mod tests {
655 use super::{SendAgentMessage, SpawnAgent, SubagentModel, SubmitResult, WaitAgent};
656 use crate::runtime::Registry;
657 use nanocodex::{Model, Tool};
658 use serde_json::json;
659 use std::sync::Weak;
660
661 #[test]
662 fn spawn_agent_requires_an_explicit_bounded_model_choice() {
663 let definition = SpawnAgent {
664 registry: Weak::<Registry>::new(),
665 selected_model: Model::Terra,
666 allow_luna: true,
667 }
668 .definition();
669 let parameters = definition.parameters().unwrap().as_value();
670 let output = definition.output_schema().unwrap();
671
672 assert_eq!(
673 parameters["properties"]["model"]["enum"],
674 json!(["selected", "luna"])
675 );
676 assert!(
677 parameters["required"]
678 .as_array()
679 .unwrap()
680 .contains(&json!("model"))
681 );
682 assert_eq!(
683 SubagentModel::Selected.resolve(Model::Terra, true).unwrap(),
684 Model::Terra
685 );
686 assert_eq!(
687 SubagentModel::Luna.resolve(Model::Terra, true).unwrap(),
688 Model::Luna
689 );
690 assert!(
691 output.as_value()["required"]
692 .as_array()
693 .unwrap()
694 .contains(&json!("model"))
695 );
696 }
697
698 #[test]
699 fn spawn_agent_excludes_and_rejects_luna_when_disabled() {
700 let definition = SpawnAgent {
701 registry: Weak::<Registry>::new(),
702 selected_model: Model::Terra,
703 allow_luna: false,
704 }
705 .definition();
706 let parameters = definition.parameters().unwrap().as_value();
707
708 assert_eq!(
709 parameters["properties"]["model"]["enum"],
710 json!(["selected"])
711 );
712 assert!(
713 SubagentModel::Luna
714 .resolve(Model::Terra, false)
715 .unwrap_err()
716 .to_string()
717 .contains("subagents.allow_luna")
718 );
719 }
720
721 #[test]
722 fn send_message_definition_names_deferred_delivery_and_queued_waiting() {
723 let definition = SendAgentMessage {
724 registry: Weak::<Registry>::new(),
725 }
726 .definition();
727 let priority = &definition.parameters().unwrap().as_value()["properties"]["priority"];
728
729 assert_eq!(priority["enum"], json!(["deferred", "urgent"]));
730 assert_eq!(priority["default"], json!("deferred"));
731 assert!(definition.description().contains("do not wait"));
732 assert!(definition.description().contains("finish the turn"));
733 }
734
735 #[test]
736 fn submit_result_requires_the_turn_token_and_one_output_value() {
737 let definition = SubmitResult {
738 registry: Weak::<Registry>::new(),
739 }
740 .definition();
741 let parameters = definition.parameters().unwrap().as_value();
742
743 assert_eq!(parameters["required"], json!(["turn_token", "output"]));
744 assert_eq!(parameters["additionalProperties"], json!(false));
745 assert_eq!(parameters["properties"].as_object().unwrap().len(), 2);
746 }
747
748 #[test]
749 fn wait_agent_only_refers_to_clean_spawns() {
750 let definition = WaitAgent {
751 registry: Weak::<Registry>::new(),
752 }
753 .definition();
754 let description =
755 &definition.parameters().unwrap().as_value()["properties"]["agent_ids"]["description"];
756 let output = definition.output_schema().unwrap();
757 let agent = &output.as_value()["properties"]["agents"]["items"];
758
759 assert!(description.as_str().unwrap().contains("spawn_agent"));
760 assert!(!description.as_str().unwrap().contains("fork_agent"));
761 assert_eq!(agent["properties"]["model"], json!({ "type": "string" }));
762 assert!(
763 agent["required"]
764 .as_array()
765 .unwrap()
766 .contains(&json!("model"))
767 );
768 }
769}