1use crate::agentloop::action::{SelfHandler, ToolClass};
18use crate::agentloop::stop::{Outcome, TerminalStatus};
19use crate::intel::client::IntelClient;
20use crate::mcp::client::McpClient;
21use crate::obs::log::Logger;
22use crate::subagent::protocol::ALLOWED_TOOLS_ROLE;
23use crate::supervisor::budget::Budget;
24use crate::wire::intel::{Message, Request, ToolDef, Usage};
25use serde_json::{Value, json};
26use std::collections::HashMap;
27use std::fmt;
28use std::sync::Arc;
29use std::sync::atomic::{AtomicBool, Ordering};
30use std::time::Instant;
31
32const PER_CALL_MAX_TOKENS: u32 = 4096;
34
35const SYSTEM_PROMPT: &str = "You are agentd, an autonomous agent. Accomplish the user's \
36instruction by calling the available tools and reasoning over their results. Call a tool when you \
37need information or need to act. When the task is complete, reply with your final answer and do \
38NOT call a tool. If the task cannot be done, say so plainly. Be concise and factual.";
39
40#[derive(Debug)]
44pub enum LoopAbort {
45 Intel(String),
47 Mcp(String),
49}
50
51impl fmt::Display for LoopAbort {
52 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53 match self {
54 LoopAbort::Intel(m) => write!(f, "intelligence: {m}"),
55 LoopAbort::Mcp(m) => write!(f, "mcp: {m}"),
56 }
57 }
58}
59
60pub struct LoopInput {
64 pub instruction: String,
65 pub output_contract: Option<String>,
66 pub seed: Vec<(String, String)>,
69 pub model: String,
70 pub max_steps: u32,
71 pub max_tokens: u64,
72 pub deadline: Instant,
73 pub cancel: Option<Arc<AtomicBool>>,
77}
78
79pub struct Session<'a> {
86 servers: &'a [McpClient],
87 tools: Vec<ToolDef>,
88 tool_to_server: HashMap<String, usize>,
89 resources: ResourceCatalogue,
90 model: String,
91 messages: Vec<Message>,
92 allowed: Vec<Vec<String>>,
99}
100
101impl<'a> Session<'a> {
102 pub fn prepare(
108 servers: &'a [McpClient],
109 input: &LoopInput,
110 self_handler: &mut dyn SelfHandler,
111 ) -> Result<Session<'a>, LoopAbort> {
112 let allowed = seed_grants(&input.seed);
113 let (mut tools, mut tool_to_server) = build_catalogue(servers)?;
114 let code = crate::tools::defs();
118 if !code.is_empty() {
119 tools.retain(|t| !crate::tools::is_registered(&t.name));
120 tools.extend(code);
121 }
122 tools.extend(self_handler.tools());
123 let resources = collect_resources(servers);
124 if !resources.owner.is_empty() || self_handler.serves_self_resources() {
127 tools.push(resource_read_tool_def());
128 }
129 narrow_catalogue(&allowed, &mut tools, &mut tool_to_server);
135 let mut messages = vec![Message::system(system_prompt(
136 input.output_contract.as_deref(),
137 ))];
138 if let Some(note) = resources.catalogue_note() {
139 messages.push(Message::system(note));
140 }
141 for (role, content) in &input.seed {
142 if role == ALLOWED_TOOLS_ROLE {
145 continue;
146 }
147 messages.push(seed_message(role, content));
148 }
149 messages.push(Message::user(&input.instruction));
150 Ok(Session {
151 servers,
152 tools,
153 tool_to_server,
154 resources,
155 model: input.model.clone(),
156 messages,
157 allowed,
158 })
159 }
160
161 pub fn refresh_tools(&mut self, self_handler: &mut dyn SelfHandler) -> Result<(), LoopAbort> {
169 let (mut tools, mut tool_to_server) = build_catalogue(self.servers)?;
170 let code = crate::tools::defs();
172 if !code.is_empty() {
173 tools.retain(|t| !crate::tools::is_registered(&t.name));
174 tools.extend(code);
175 }
176 tools.extend(self_handler.tools());
177 if !self.resources.owner.is_empty() || self_handler.serves_self_resources() {
178 tools.push(resource_read_tool_def());
179 }
180 narrow_catalogue(&self.allowed, &mut tools, &mut tool_to_server);
184 self.tools = tools;
185 self.tool_to_server = tool_to_server;
186 Ok(())
187 }
188
189 pub fn tools_len(&self) -> usize {
191 self.tools.len()
192 }
193
194 pub fn tool_class(&self, name: &str) -> ToolClass {
206 if crate::tools::is_registered(name) {
211 ToolClass::Code
212 } else if self.tool_to_server.contains_key(name) {
213 ToolClass::Mcp
214 } else {
215 ToolClass::SelfControl
216 }
217 }
218
219 pub fn tool_permitted(&self, name: &str) -> bool {
225 grant_permits(&self.allowed, name)
226 }
227
228 pub fn deliver(&mut self, content: &str) {
232 self.messages.push(Message::user(content));
233 }
234
235 pub fn set_model(&mut self, model: &str) {
240 self.model = model.to_string();
241 }
242
243 pub fn model(&self) -> &str {
247 &self.model
248 }
249
250 pub fn transcript_len(&self) -> usize {
255 self.messages.len()
256 }
257
258 pub fn truncate_transcript(&mut self, len: usize) {
263 if len < self.messages.len() {
264 self.messages.truncate(len);
265 }
266 }
267
268 pub fn run_turn(
281 &mut self,
282 intel: &IntelClient,
283 self_handler: &mut dyn SelfHandler,
284 log: &Logger,
285 budget: &mut Budget,
286 cancel: Option<&Arc<AtomicBool>>,
287 ) -> Result<(Outcome, Usage), LoopAbort> {
288 let mut last_text: Option<String> = None;
289 let run_start = crate::obs::otel::now_unix_nanos();
293 let mut run_span = crate::obs::otel::run_begin(log.ctx().trace_id.as_deref(), run_start);
294 let (mut tok_in, mut tok_out) = (0u64, 0u64);
295
296 log.info(
297 "loop.start",
298 json!({"tools": self.tools.len(), "servers": self.servers.len(), "resources": self.resources.owner.len(), "max_steps": budget.max_steps()}),
299 );
300
301 loop {
302 if cancel.is_some_and(|c| c.load(Ordering::Relaxed)) {
303 log.warn(
304 "loop.final",
305 json!({"status": "cancelled", "steps": budget.steps()}),
306 );
307 run_span.finish(&self.model, tok_in, tok_out, false);
308 return Ok((
309 Outcome {
310 status: TerminalStatus::Cancelled,
311 partial: last_text.is_some(),
312 result: json!(last_text.unwrap_or_default()),
313 scheduled: self_handler.take_scheduled(),
314 subscriptions: self_handler.take_subscriptions(),
315 },
316 Usage {
317 input_tokens: tok_in,
318 output_tokens: tok_out,
319 },
320 ));
321 }
322 if let Some(status) = budget.exceeded() {
323 log.warn("loop.final", json!({"status": status.as_str(), "steps": budget.steps(), "tokens": budget.tokens()}));
324 run_span.finish(&self.model, tok_in, tok_out, false);
325 return Ok((
326 Outcome {
327 status,
328 partial: last_text.is_some(),
329 result: json!(last_text.unwrap_or_default()),
330 scheduled: self_handler.take_scheduled(),
331 subscriptions: self_handler.take_subscriptions(),
332 },
333 Usage {
334 input_tokens: tok_in,
335 output_tokens: tok_out,
336 },
337 ));
338 }
339
340 log.debug(
344 "loop.step",
345 json!({"step": budget.steps(), "tokens": budget.tokens(), "messages": self.messages.len()}),
346 );
347
348 let req = Request {
349 model: self.model.clone(),
350 messages: self.messages.clone(),
351 tools: self.tools.clone(),
352 max_tokens: PER_CALL_MAX_TOKENS,
353 temperature: Some(0.0),
354 };
355
356 log.debug(
357 "intel.call",
358 json!({"step": budget.steps(), "messages": self.messages.len()}),
359 );
360 let chat_start = crate::obs::otel::now_unix_nanos();
361 let resp = intel
362 .complete(&req)
363 .map_err(|e| LoopAbort::Intel(e.to_string()))?;
364 budget.record_usage(resp.usage);
365 budget.record_step();
366 tok_in += resp.usage.input_tokens;
367 tok_out += resp.usage.output_tokens;
368 run_span.record_chat(
369 &self.model,
370 resp.usage.input_tokens,
371 resp.usage.output_tokens,
372 true,
373 chat_start,
374 );
375 log.debug(
376 "intel.result",
377 json!({"tool_calls": resp.tool_calls.len(), "tokens_in": resp.usage.input_tokens, "tokens_out": resp.usage.output_tokens}),
378 );
379
380 if resp.wants_tools() {
381 if let Some(t) = resp.text.as_deref().filter(|t| !t.is_empty()) {
382 last_text = Some(t.to_string());
383 }
384 let tool_calls = resp.tool_calls.clone();
385 self.messages.push(Message::Assistant {
386 text: resp.text,
387 tool_calls: tool_calls.clone(),
388 });
389
390 for tc in &tool_calls {
391 let mut call = json!({"tool": tc.name, "id": tc.id});
392 if log.content_capture() {
396 call["args"] = json!(truncate_for_log(&tc.arguments.to_string()));
397 }
398 log.info("tool.call", call);
399 let tool_start = crate::obs::otel::now_unix_nanos();
400 let (content, is_error) = if !self.tool_permitted(&tc.name) {
401 (
406 format!(
407 "error: tool '{}' is not in this subagent's allowed tools",
408 tc.name
409 ),
410 true,
411 )
412 } else if tc.name == "resource.read" {
413 let uri = tc
417 .arguments
418 .get("uri")
419 .and_then(Value::as_str)
420 .unwrap_or("")
421 .trim();
422 if uri.starts_with("agentd://") || uri.starts_with("agent://") {
423 self_handler.read_resource(uri).unwrap_or_else(|| {
424 (format!("unknown agentd resource: {uri}"), true)
425 })
426 } else {
427 read_resource_tool(self.servers, &self.resources.owner, &tc.arguments)
428 }
429 } else {
430 match self_handler.handle(&tc.name, &tc.arguments) {
431 Some(r) => r, None => match crate::tools::dispatch(&tc.name, &tc.arguments) {
435 Some(r) => r,
436 None => dispatch_tool(
437 self.servers,
438 &self.tool_to_server,
439 &tc.name,
440 &tc.arguments,
441 ),
442 },
443 }
444 };
445 run_span.record_tool(&tc.name, !is_error, tool_start);
446 let mut result =
447 json!({"tool": tc.name, "is_error": is_error, "bytes": content.len()});
448 if log.content_capture() {
449 result["content"] = json!(truncate_for_log(&content));
450 }
451 log.info("tool.result", result);
452 self.messages
453 .push(Message::tool_result(&tc.id, content, is_error));
454 }
455 continue;
456 }
457
458 let text = resp.text.clone().or(last_text).unwrap_or_default();
462 self.messages.push(Message::Assistant {
463 text: Some(text.clone()),
464 tool_calls: Vec::new(),
465 });
466 log.info(
467 "loop.final",
468 json!({"status": "completed", "steps": budget.steps(), "tokens": budget.tokens()}),
469 );
470 run_span.finish(&self.model, tok_in, tok_out, true);
471 return Ok((
472 Outcome {
473 status: TerminalStatus::Completed,
474 partial: false,
475 result: json!(text),
476 scheduled: self_handler.take_scheduled(),
477 subscriptions: self_handler.take_subscriptions(),
478 },
479 Usage {
480 input_tokens: tok_in,
481 output_tokens: tok_out,
482 },
483 ));
484 }
485 }
486}
487
488pub fn run_loop(
499 intel: &IntelClient,
500 servers: &[McpClient],
501 input: &LoopInput,
502 self_handler: &mut dyn SelfHandler,
503 log: &Logger,
504) -> Result<(Outcome, Usage), LoopAbort> {
505 let mut session = Session::prepare(servers, input, self_handler)?;
506 let mut budget = Budget::new(input.max_steps, input.max_tokens, input.deadline);
507 session.run_turn(intel, self_handler, log, &mut budget, input.cancel.as_ref())
508}
509
510const CONTENT_LOG_CAP: usize = 4096;
514
515fn truncate_for_log(s: &str) -> String {
518 if s.chars().count() <= CONTENT_LOG_CAP {
519 return s.to_string();
520 }
521 let mut t: String = s.chars().take(CONTENT_LOG_CAP).collect();
522 t.push_str(&format!(
523 "…(+{} more bytes)",
524 s.len().saturating_sub(t.len())
525 ));
526 t
527}
528
529fn build_catalogue(
533 servers: &[McpClient],
534) -> Result<(Vec<ToolDef>, HashMap<String, usize>), LoopAbort> {
535 let mut tools = Vec::new();
536 let mut routing = HashMap::new();
537 for (i, server) in servers.iter().enumerate() {
538 let listed = server
539 .list_tools()
540 .map_err(|e| LoopAbort::Mcp(e.to_string()))?;
541 for t in listed {
542 routing.entry(t.name.clone()).or_insert(i);
543 tools.push(ToolDef {
544 name: t.name,
545 description: t.description.unwrap_or_default(),
546 input_schema: t.input_schema,
547 });
548 }
549 }
550 Ok((tools, routing))
551}
552
553fn seed_grants(seed: &[(String, String)]) -> Vec<Vec<String>> {
557 seed.iter()
558 .filter(|(role, _)| role == ALLOWED_TOOLS_ROLE)
559 .map(|(_, content)| crate::subagent::protocol::parse_allowed_tools(content))
560 .collect()
561}
562
563fn grant_permits(grants: &[Vec<String>], name: &str) -> bool {
567 grants
568 .iter()
569 .all(|g| g.iter().any(|p| crate::registry::pattern_matches(p, name)))
570}
571
572fn narrow_catalogue(
576 grants: &[Vec<String>],
577 tools: &mut Vec<ToolDef>,
578 routing: &mut HashMap<String, usize>,
579) {
580 if grants.is_empty() {
581 return;
582 }
583 tools.retain(|t| grant_permits(grants, &t.name));
584 routing.retain(|name, _| grant_permits(grants, name));
585}
586
587fn dispatch_tool(
591 servers: &[McpClient],
592 routing: &HashMap<String, usize>,
593 name: &str,
594 arguments: &Value,
595) -> (String, bool) {
596 match routing.get(name) {
597 Some(&i) => match servers[i].call_tool(name, Some(arguments.clone())) {
598 Ok(res) => (res.text(), res.is_error()),
599 Err(e) => (format!("tool transport error: {e}"), true),
600 },
601 None => (format!("error: no such tool '{name}'"), true),
602 }
603}
604
605fn system_prompt(contract: Option<&str>) -> String {
608 match contract {
609 Some(c) if !c.is_empty() => format!("{SYSTEM_PROMPT}\n\nOutput contract:\n{c}"),
610 _ => SYSTEM_PROMPT.to_string(),
611 }
612}
613
614fn seed_message(role: &str, content: &str) -> Message {
617 match role {
618 "system" => Message::system(content),
619 "assistant" => Message::Assistant {
620 text: Some(content.to_string()),
621 tool_calls: Vec::new(),
622 },
623 _ => Message::user(content),
624 }
625}
626
627const RESOURCE_CAP: usize = 50;
630
631struct ResourceCatalogue {
634 owner: HashMap<String, usize>,
635 entries: Vec<(String, String)>, truncated: bool,
637}
638
639impl ResourceCatalogue {
640 fn catalogue_note(&self) -> Option<String> {
642 if self.entries.is_empty() {
643 return None;
644 }
645 let mut s = String::from(
646 "Available MCP resources — read the current content of any with the resource.read tool:\n",
647 );
648 for (uri, label) in &self.entries {
649 if label.is_empty() {
650 s.push_str(&format!("- {uri}\n"));
651 } else {
652 s.push_str(&format!("- {uri} — {label}\n"));
653 }
654 }
655 if self.truncated {
656 s.push_str(&format!(
657 "(… more than {RESOURCE_CAP} resources; list truncated)\n"
658 ));
659 }
660 Some(s)
661 }
662}
663
664fn collect_resources(servers: &[McpClient]) -> ResourceCatalogue {
667 let mut owner = HashMap::new();
668 let mut entries = Vec::new();
669 let mut truncated = false;
670 'outer: for (i, s) in servers.iter().enumerate() {
671 let Ok(list) = s.list_resources() else {
672 continue;
673 };
674 for r in list {
675 if entries.len() >= RESOURCE_CAP {
676 truncated = true;
677 break 'outer;
678 }
679 if !owner.contains_key(&r.uri) {
680 let label = r.title.or(r.name).or(r.description).unwrap_or_default();
681 owner.insert(r.uri.clone(), i);
682 entries.push((r.uri, label));
683 }
684 }
685 }
686 ResourceCatalogue {
687 owner,
688 entries,
689 truncated,
690 }
691}
692
693fn resource_read_tool_def() -> ToolDef {
694 ToolDef {
695 name: "resource.read".into(),
696 description: "Read the current content of an available MCP resource by its uri (see the \
697 resource catalogue). Use this to pull a resource's body when you need it."
698 .into(),
699 input_schema: json!({
700 "type": "object",
701 "properties": {"uri": {"type": "string", "description": "the resource uri to read"}},
702 "required": ["uri"]
703 }),
704 }
705}
706
707fn read_resource_tool(
710 servers: &[McpClient],
711 owner: &HashMap<String, usize>,
712 args: &Value,
713) -> (String, bool) {
714 let uri = args.get("uri").and_then(Value::as_str).unwrap_or("").trim();
715 if uri.is_empty() {
716 return ("error: resource.read requires a 'uri'".into(), true);
717 }
718 let candidates: Vec<usize> = match owner.get(uri) {
719 Some(i) => vec![*i],
720 None => (0..servers.len()).collect(), };
722 for i in candidates {
723 if let Ok(r) = servers[i].read_resource(uri) {
724 return (r.text(), false);
725 }
726 }
727 (format!("resource.read: no server could read '{uri}'"), true)
728}
729
730#[cfg(test)]
731mod tests {
732 use super::*;
733
734 #[test]
735 fn resource_catalogue_note_lists_uris() {
736 let c = ResourceCatalogue {
737 owner: HashMap::new(),
738 entries: vec![
739 ("file:///a.json".into(), "inbox".into()),
740 ("db://orders".into(), String::new()),
741 ],
742 truncated: false,
743 };
744 let note = c.catalogue_note().unwrap();
745 assert!(note.contains("resource.read"));
746 assert!(note.contains("file:///a.json — inbox"));
747 assert!(note.contains("- db://orders\n"));
748 }
749
750 #[test]
751 fn empty_catalogue_is_no_note() {
752 let c = ResourceCatalogue {
753 owner: HashMap::new(),
754 entries: vec![],
755 truncated: false,
756 };
757 assert!(c.catalogue_note().is_none());
758 }
759
760 #[test]
761 fn resource_read_rejects_missing_uri() {
762 let (msg, err) = read_resource_tool(&[], &HashMap::new(), &json!({}));
763 assert!(err);
764 assert!(msg.contains("uri"));
765 }
766
767 #[test]
768 fn resource_read_no_server_is_an_error_observation() {
769 let (msg, err) = read_resource_tool(&[], &HashMap::new(), &json!({"uri": "file:///x"}));
770 assert!(err);
771 assert!(msg.contains("file:///x"));
772 }
773
774 #[test]
775 fn system_prompt_appends_contract() {
776 let p = system_prompt(Some("Return JSON."));
777 assert!(p.contains("Output contract:"));
778 assert!(p.contains("Return JSON."));
779 assert_eq!(system_prompt(None), SYSTEM_PROMPT);
780 }
781
782 #[test]
783 fn a_code_registered_tool_classifies_code_and_wins_a_name_collision() {
784 let _guard = crate::tools::test_registry_guard();
788 crate::tools::register(crate::tools::CodeTool::new(
789 "runner.code_tool",
790 "a native tool",
791 json!({"type": "object"}),
792 |_| Ok(json!("native")),
793 ))
794 .expect("register");
795 let mut tool_to_server = HashMap::new();
796 tool_to_server.insert("runner.code_tool".to_string(), 0usize);
798 let sess = Session {
799 servers: &[],
800 tools: vec![],
801 tool_to_server,
802 resources: ResourceCatalogue {
803 owner: HashMap::new(),
804 entries: vec![],
805 truncated: false,
806 },
807 model: "m".into(),
808 messages: vec![],
809 allowed: Vec::new(),
810 };
811 assert_eq!(
812 sess.tool_class("runner.code_tool"),
813 ToolClass::Code,
814 "code wins the collision — a server cannot steal a registered tool's calls"
815 );
816 let (content, is_err) =
818 crate::tools::dispatch("runner.code_tool", &json!({})).expect("code tool dispatches");
819 assert!(!is_err);
820 assert_eq!(content, "native");
821 assert!(crate::tools::unregister("runner.code_tool"));
822 }
823
824 #[test]
825 fn catalogue_partitions_into_mcp_and_self_control_classes() {
826 use crate::agentloop::action::SELF_CONTROL_TOOLS;
827 let mcp = ["db.query", "http.get"];
833 let mut tool_to_server = HashMap::new();
834 let mut tools: Vec<ToolDef> = Vec::new();
835 for n in mcp {
836 tool_to_server.insert(n.to_string(), 0usize);
837 tools.push(ToolDef {
838 name: n.into(),
839 description: String::new(),
840 input_schema: json!({}),
841 });
842 }
843 for n in SELF_CONTROL_TOOLS {
846 tools.push(ToolDef {
847 name: (*n).into(),
848 description: String::new(),
849 input_schema: json!({}),
850 });
851 }
852 let sess = Session {
853 servers: &[],
854 tools,
855 tool_to_server,
856 resources: ResourceCatalogue {
857 owner: HashMap::new(),
858 entries: vec![],
859 truncated: false,
860 },
861 model: "m".into(),
862 messages: vec![],
863 allowed: Vec::new(),
864 };
865 for n in mcp {
867 assert_eq!(sess.tool_class(n), ToolClass::Mcp, "{n} is an MCP tool");
868 }
869 for n in SELF_CONTROL_TOOLS {
870 assert_eq!(
871 sess.tool_class(n),
872 ToolClass::SelfControl,
873 "{n} is self/control"
874 );
875 }
876 let (mut n_mcp, mut n_self, mut n_code) = (0usize, 0usize, 0usize);
879 for t in &sess.tools {
880 match sess.tool_class(&t.name) {
881 ToolClass::Mcp => n_mcp += 1,
882 ToolClass::SelfControl => n_self += 1,
883 ToolClass::Code => n_code += 1,
884 }
885 }
886 assert_eq!(n_code, 0, "no code tools registered here");
887 assert_eq!(n_mcp, mcp.len(), "every MCP tool classified");
888 assert_eq!(
889 n_self,
890 SELF_CONTROL_TOOLS.len(),
891 "every self tool classified"
892 );
893 for bad in [
895 "exec", "shell", "bash", "sh", "command", "system", "eval", "run",
896 ] {
897 assert!(
898 !SELF_CONTROL_TOOLS.contains(&bad),
899 "no local-exec self-tool: {bad}"
900 );
901 }
902 }
903
904 #[test]
905 fn dispatch_unknown_tool_is_error_observation() {
906 let routing = HashMap::new();
907 let (content, is_error) = dispatch_tool(&[], &routing, "ghost", &Value::Null);
908 assert!(is_error);
909 assert!(content.contains("ghost"));
910 }
911
912 #[test]
913 fn loop_abort_display() {
914 assert!(LoopAbort::Intel("down".into()).to_string().contains("down"));
915 }
916
917 #[test]
918 fn truncate_for_log_caps_and_marks() {
919 let short = "{\"a\":1}";
920 assert_eq!(truncate_for_log(short), short); let big = "x".repeat(CONTENT_LOG_CAP + 500);
922 let out = truncate_for_log(&big);
923 assert!(out.len() < big.len());
924 assert!(
925 out.contains("more bytes"),
926 "truncation is marked: {}",
927 &out[out.len() - 32..]
928 );
929 let multi = "é".repeat(CONTENT_LOG_CAP + 10);
931 let _ = truncate_for_log(&multi);
932 }
933
934 #[test]
944 fn refresh_tools_picks_up_a_changed_handler_catalogue() {
945 let _guard = crate::tools::test_registry_guard();
952 struct GrowingHandler {
953 grown: bool,
954 }
955 impl SelfHandler for GrowingHandler {
956 fn tools(&self) -> Vec<ToolDef> {
957 let mut t = vec![ToolDef {
958 name: "alpha".into(),
959 description: String::new(),
960 input_schema: Value::Null,
961 }];
962 if self.grown {
963 t.push(ToolDef {
964 name: "beta".into(),
965 description: String::new(),
966 input_schema: Value::Null,
967 });
968 }
969 t
970 }
971 fn handle(&mut self, _name: &str, _args: &Value) -> Option<(String, bool)> {
972 None
973 }
974 }
975 let input = LoopInput {
976 instruction: "x".into(),
977 output_contract: None,
978 seed: Vec::new(),
979 model: "m".into(),
980 max_steps: 5,
981 max_tokens: 1000,
982 deadline: std::time::Instant::now() + std::time::Duration::from_secs(5),
983 cancel: None,
984 };
985 let mut handler = GrowingHandler { grown: false };
986 let mut session = Session::prepare(&[], &input, &mut handler).unwrap();
987 let before = session.tools_len();
988 let transcript = session.transcript_len();
989 handler.grown = true;
990 session.refresh_tools(&mut handler).unwrap();
991 assert_eq!(session.tools_len(), before + 1, "the new tool is live");
992 assert_eq!(session.transcript_len(), transcript, "transcript untouched");
993 assert_eq!(session.tool_class("beta"), ToolClass::SelfControl);
995 }
996
997 #[test]
998 fn a_seed_grant_narrows_the_catalogue_the_dispatch_and_nothing_else() {
999 let _guard = crate::tools::test_registry_guard();
1004 struct TwoTools;
1005 impl SelfHandler for TwoTools {
1006 fn tools(&self) -> Vec<ToolDef> {
1007 ["alpha", "beta"]
1008 .into_iter()
1009 .map(|n| ToolDef {
1010 name: n.into(),
1011 description: String::new(),
1012 input_schema: Value::Null,
1013 })
1014 .collect()
1015 }
1016 fn handle(&mut self, _name: &str, _args: &Value) -> Option<(String, bool)> {
1017 Some(("served".into(), false))
1018 }
1019 }
1020 let grant = LoopInput {
1021 instruction: "x".into(),
1022 output_contract: None,
1023 seed: vec![
1024 (
1025 crate::subagent::protocol::ALLOWED_TOOLS_ROLE.to_string(),
1026 "[\"alpha\"]".to_string(),
1027 ),
1028 ("user".to_string(), "a real seed message".to_string()),
1029 ],
1030 model: "m".into(),
1031 max_steps: 5,
1032 max_tokens: 1000,
1033 deadline: std::time::Instant::now() + std::time::Duration::from_secs(5),
1034 cancel: None,
1035 };
1036 let mut handler = TwoTools;
1037 let narrowed = Session::prepare(&[], &grant, &mut handler).unwrap();
1038 assert_eq!(narrowed.tools_len(), 1, "only the granted tool is offered");
1039 assert!(narrowed.tool_permitted("alpha"));
1040 assert!(
1041 !narrowed.tool_permitted("beta"),
1042 "a filtered-out tool is refused at dispatch, not served"
1043 );
1044 assert_eq!(narrowed.transcript_len(), 3);
1047
1048 let mut plain = grant;
1050 plain.seed.remove(0);
1051 let wide = Session::prepare(&[], &plain, &mut handler).unwrap();
1052 assert_eq!(wide.tools_len(), 2);
1053 assert!(wide.tool_permitted("beta"));
1054 }
1055
1056 #[cfg(unix)]
1057 mod usage_producer {
1058 use super::*;
1059 use crate::intel::client::IntelClient;
1060 use crate::obs::log::{Comp, Level, LogCtx, Logger};
1061 use std::time::{Duration, Instant};
1062
1063 struct NoopHandler;
1067 impl SelfHandler for NoopHandler {
1068 fn tools(&self) -> Vec<ToolDef> {
1069 Vec::new()
1070 }
1071 fn handle(&mut self, _name: &str, _args: &Value) -> Option<(String, bool)> {
1072 None
1073 }
1074 }
1075
1076 fn test_log() -> Logger {
1077 Logger::new(
1078 LogCtx {
1079 run_id: "r".into(),
1080 agent_id: "0".into(),
1081 agent_path: "0".into(),
1082 comp: Comp::Agent,
1083 pid: 0,
1084 trace_id: None,
1085 },
1086 Level::Error, )
1088 }
1089
1090 fn start_mock_llm(addr_file: &std::path::Path, script: &'static str) -> String {
1095 let s = addr_file.to_str().unwrap().to_string();
1096 std::thread::spawn(move || {
1097 crate::intel::mock::run(&s, script);
1098 });
1099 let deadline = Instant::now() + Duration::from_secs(3);
1100 while !addr_file.exists() {
1101 assert!(Instant::now() < deadline, "mock-llm never announced");
1102 std::thread::sleep(Duration::from_millis(10));
1103 }
1104 let addr = std::fs::read_to_string(addr_file).expect("read mock-llm addr-file");
1105 format!("http://{}", addr.trim())
1106 }
1107
1108 fn input(instruction: &str) -> LoopInput {
1109 LoopInput {
1110 instruction: instruction.into(),
1111 output_contract: None,
1112 seed: Vec::new(),
1113 model: "mock".into(),
1114 max_steps: 8,
1115 max_tokens: 100_000,
1116 deadline: Instant::now() + Duration::from_secs(10),
1117 cancel: None,
1118 }
1119 }
1120
1121 #[test]
1122 fn run_turn_returns_the_turns_token_usage() {
1123 let dir = tempfile::tempdir().unwrap();
1127 let sock = dir.path().join("llm.addr");
1128 let url = start_mock_llm(&sock, "final");
1129
1130 let intel = IntelClient::from_parts(&url, None).unwrap();
1131 let inp = input("do the thing");
1132 let mut handler = NoopHandler;
1133 let mut session = Session::prepare(&[], &inp, &mut handler).unwrap();
1134 let mut budget = Budget::new(inp.max_steps, inp.max_tokens, inp.deadline);
1135
1136 let (outcome, usage) = session
1137 .run_turn(&intel, &mut handler, &test_log(), &mut budget, None)
1138 .expect("turn runs against the mock LLM");
1139
1140 assert_eq!(outcome.status, TerminalStatus::Completed);
1141 assert_eq!(
1144 usage.input_tokens, 11,
1145 "input tokens surfaced from the model"
1146 );
1147 assert_eq!(
1148 usage.output_tokens, 5,
1149 "output tokens surfaced from the model"
1150 );
1151 assert!(usage.total() > 0, "the rolled-up Usage is non-zero");
1152 }
1153
1154 #[test]
1155 fn run_loop_returns_the_runs_total_token_usage() {
1156 let dir = tempfile::tempdir().unwrap();
1161 let sock = dir.path().join("llm.addr");
1162 let url = start_mock_llm(&sock, "read");
1163
1164 let intel = IntelClient::from_parts(&url, None).unwrap();
1165 let inp = input("read the resource");
1166 let mut handler = NoopHandler;
1167
1168 let (outcome, usage) =
1169 run_loop(&intel, &[], &inp, &mut handler, &test_log()).expect("one-shot run");
1170
1171 assert_eq!(outcome.status, TerminalStatus::Completed);
1172 assert_eq!(usage.input_tokens, 22, "summed input over both model calls");
1175 assert_eq!(
1176 usage.output_tokens, 12,
1177 "summed output over both model calls"
1178 );
1179 }
1180 }
1181}