1#![deny(missing_docs)]
4#![forbid(unsafe_code)]
5
6use std::{future::Future, pin::Pin, time::Duration};
7
8use anyhow::{Context, ensure};
9use kcode_codex_runtime_v2::{
10 AgentEvent, AgentRequest, DynamicTool, DynamicToolCall, ReasoningEffort, ToolResult,
11};
12use kcode_intelligence_router::{AgentProvider, Intelligence, ResolvedAgentModel, UsageReceipt};
13use serde_json::{Value, json};
14use sha2::{Digest, Sha256};
15use uuid::Uuid;
16
17const DEFAULT_ROUND_LIMIT: u64 = 100;
18const PROTOCOL_TOKEN_RESERVE: u64 = 4_096;
19const INLINE_TOOL_RESULT_CHARACTERS: usize = 1_000;
20
21pub type HostFuture<'a, T> = Pin<Box<dyn Future<Output = anyhow::Result<T>> + Send + 'a>>;
23
24#[derive(Clone, Debug, PartialEq)]
26pub struct ToolCall {
27 pub name: String,
29 pub arguments: Value,
31}
32
33#[derive(Clone, Debug, PartialEq)]
35pub enum AuditEvent {
36 Started {
38 parent_operation_id: Uuid,
40 model: String,
42 provider_model: String,
44 provider: AgentProvider,
46 context_window_tokens: u64,
48 max_input_tokens: u64,
50 context: Vec<String>,
52 task: String,
54 host: Value,
56 },
57 InferenceSubmitted {
59 parent_operation_id: Uuid,
61 round: u64,
63 manifest_hash: String,
65 estimated_input_tokens: u64,
67 },
68 ToolCall {
70 parent_operation_id: Uuid,
72 name: String,
74 arguments: Value,
76 },
77 ToolResult {
79 parent_operation_id: Uuid,
81 name: String,
83 ok: bool,
85 projection_accepted: bool,
87 result: String,
89 },
90 ProviderReceipt {
92 parent_operation_id: Uuid,
94 round: u64,
96 manifest_hash: String,
98 usage: Option<kcode_codex_runtime_v2::TokenUsage>,
100 receipt: Box<UsageReceipt>,
102 },
103 Completed {
105 parent_operation_id: Uuid,
107 model: String,
109 response: String,
111 },
112}
113
114#[derive(Clone, Debug, Eq, PartialEq)]
116pub struct StateUpdate {
117 pub key: String,
119 pub text: Option<String>,
121}
122
123#[derive(Clone, Debug, PartialEq)]
125pub struct ToolOutcome {
126 pub text: String,
128 pub ok: bool,
130 pub state_updates: Vec<StateUpdate>,
132 pub capture: Option<Value>,
134}
135
136impl ToolOutcome {
137 pub fn success(text: impl Into<String>) -> Self {
139 Self {
140 text: text.into(),
141 ok: true,
142 state_updates: Vec::new(),
143 capture: None,
144 }
145 }
146
147 pub fn failure(text: impl Into<String>) -> Self {
149 Self {
150 text: text.into(),
151 ok: false,
152 state_updates: Vec::new(),
153 capture: None,
154 }
155 }
156}
157
158#[derive(Clone)]
160pub struct ContextBudget {
161 projection: Projection,
162 max_input_tokens: u64,
163}
164
165impl ContextBudget {
166 pub fn estimated_tokens(&self) -> u64 {
168 self.projection.estimated_tokens()
169 }
170
171 pub fn max_input_tokens(&self) -> u64 {
173 self.max_input_tokens
174 }
175
176 pub fn fits_state(&self, key: impl Into<String>, text: impl Into<String>) -> bool {
178 let mut projection = self.projection.clone();
179 projection.update_state(key.into(), Some(text.into()));
180 projection.estimated_tokens() <= self.max_input_tokens
181 }
182}
183
184pub trait Host: Send {
186 fn render_tool_call(&mut self, call: &ToolCall) -> anyhow::Result<String>;
188
189 fn execute_tool<'a>(
191 &'a mut self,
192 call: ToolCall,
193 operation_id: Uuid,
194 budget: ContextBudget,
195 ) -> HostFuture<'a, ToolOutcome>;
196
197 fn complete_capture<'a>(
199 &'a mut self,
200 capture: Value,
201 contents: String,
202 budget: ContextBudget,
203 ) -> HostFuture<'a, ToolOutcome>;
204
205 fn record(&mut self, event: AuditEvent) -> anyhow::Result<()>;
207}
208
209#[derive(Clone, Debug, PartialEq)]
211pub struct RunRequest {
212 pub user_id: String,
214 pub parent_operation_id: Uuid,
216 pub model: String,
218 pub reasoning_effort: String,
220 pub context: Vec<String>,
222 pub task: String,
224 pub timeout: Option<Duration>,
226 pub start_metadata: Value,
228}
229
230#[derive(Clone, Debug, Eq, PartialEq)]
232pub struct RunResult {
233 pub answer: String,
235 pub model: ResolvedAgentModel,
237}
238
239#[derive(Clone)]
241pub struct AgentRuntime {
242 intelligence: Intelligence,
243 round_limit: u64,
244}
245
246impl AgentRuntime {
247 pub fn new(intelligence: Intelligence) -> Self {
249 Self {
250 intelligence,
251 round_limit: DEFAULT_ROUND_LIMIT,
252 }
253 }
254
255 pub async fn resolve_model(&self, requested: &str) -> anyhow::Result<ResolvedAgentModel> {
257 self.intelligence
258 .resolve_agent_model(requested)
259 .await
260 .map_err(anyhow::Error::new)
261 }
262
263 pub async fn run<H: Host>(
265 &self,
266 request: RunRequest,
267 host: &mut H,
268 ) -> anyhow::Result<RunResult> {
269 let selected = self.resolve_model(&request.model).await?;
270 let reasoning_effort = reasoning_effort(&request.reasoning_effort)?;
271 let mut projection = Projection::new(request.context, request.task);
272 ensure_capacity(&projection, selected.max_input_tokens)?;
273 host.record(AuditEvent::Started {
274 parent_operation_id: request.parent_operation_id,
275 model: request.model.clone(),
276 provider_model: selected.provider_model.clone(),
277 provider: selected.provider,
278 context_window_tokens: selected.context_window_tokens,
279 max_input_tokens: selected.max_input_tokens,
280 context: projection.context.clone(),
281 task: projection.task.clone(),
282 host: request.start_metadata.clone(),
283 })?;
284 let user = self
285 .intelligence
286 .for_user(request.user_id)
287 .map_err(anyhow::Error::new)?;
288 let mut deferred_capture: Option<Value> = None;
289
290 for round in 0..self.round_limit {
291 let capturing = deferred_capture.is_some();
292 ensure_capacity(&projection, selected.max_input_tokens)?;
293 let input = projection.render();
294 let manifest_hash = hex::encode(Sha256::digest(input.as_bytes()));
295 host.record(AuditEvent::InferenceSubmitted {
296 parent_operation_id: request.parent_operation_id,
297 round: round + 1,
298 manifest_hash: manifest_hash.clone(),
299 estimated_input_tokens: projection.estimated_tokens(),
300 })?;
301 let mut provider_request = AgentRequest::new(input, selected.requested_model.clone());
302 provider_request.reasoning_effort = reasoning_effort;
303 provider_request.ephemeral = true;
304 provider_request.tools = if capturing {
305 Vec::new()
306 } else {
307 vec![ktool_definition()]
308 };
309 if let Some(timeout) = request.timeout {
310 provider_request.timeout = timeout;
311 }
312 let child_operation_id = Uuid::new_v4();
313 let mut turn = match user
314 .start_agent_turn(
315 child_operation_id,
316 Some(request.parent_operation_id),
317 provider_request,
318 )
319 .await
320 {
321 Ok(turn) => turn,
322 Err(error) => {
323 if let Some(receipt) = error.receipt().cloned() {
324 host.record(AuditEvent::ProviderReceipt {
325 parent_operation_id: request.parent_operation_id,
326 round: round + 1,
327 manifest_hash,
328 usage: None,
329 receipt: Box::new(receipt),
330 })?;
331 }
332 return Err(anyhow::Error::new(error));
333 }
334 };
335 let mut used_tool = false;
336 let mut pending_capture: Option<Value> = None;
337 let mut requires_rerender = false;
338 let completed = loop {
339 let event = match turn.next_event().await {
340 Ok(Some(event)) => event,
341 Ok(None) => {
342 let receipt = turn.finish_unavailable()?.clone();
343 host.record(AuditEvent::ProviderReceipt {
344 parent_operation_id: request.parent_operation_id,
345 round: round + 1,
346 manifest_hash: manifest_hash.clone(),
347 usage: None,
348 receipt: Box::new(receipt),
349 })?;
350 anyhow::bail!("subagent provider ended without a terminal turn event");
351 }
352 Err(error) => {
353 if let Some(receipt) = error.receipt().cloned() {
354 host.record(AuditEvent::ProviderReceipt {
355 parent_operation_id: request.parent_operation_id,
356 round: round + 1,
357 manifest_hash: manifest_hash.clone(),
358 usage: None,
359 receipt: Box::new(receipt),
360 })?;
361 }
362 return Err(anyhow::Error::new(error));
363 }
364 };
365 match event {
366 AgentEvent::ProviderInput(_) => {}
367 AgentEvent::UsageUpdated(_) => {}
368 AgentEvent::ToolCall(native) => {
369 used_tool = true;
370 if capturing {
371 respond_or_record(
372 &mut turn,
373 host,
374 request.parent_operation_id,
375 round + 1,
376 &manifest_hash,
377 &native.call_id,
378 ToolResult::failure(
379 "No application tool is available while complete freeform output is being captured.",
380 ),
381 )
382 .await?;
383 continue;
384 }
385 if pending_capture.is_some() {
386 respond_or_record(
387 &mut turn,
388 host,
389 request.parent_operation_id,
390 round + 1,
391 &manifest_hash,
392 &native.call_id,
393 ToolResult::failure(
394 "A freeform output capture is pending; no other tool can run first.",
395 ),
396 )
397 .await?;
398 continue;
399 }
400 if requires_rerender {
401 respond_or_record(
402 &mut turn,
403 host,
404 request.parent_operation_id,
405 round + 1,
406 &manifest_hash,
407 &native.call_id,
408 ToolResult::failure(
409 "A state update is waiting to be re-rendered. End this slice before calling another tool.",
410 ),
411 )
412 .await?;
413 continue;
414 }
415 let call = match parse_ktool_call(&native) {
416 Ok(call) => call,
417 Err(error) => {
418 let text = format!("Invalid application tool call: {error}");
419 projection.push_history(format!("Ktool result:\n{text}"));
420 respond_or_record(
421 &mut turn,
422 host,
423 request.parent_operation_id,
424 round + 1,
425 &manifest_hash,
426 &native.call_id,
427 ToolResult::failure(text),
428 )
429 .await?;
430 continue;
431 }
432 };
433 host.record(AuditEvent::ToolCall {
434 parent_operation_id: request.parent_operation_id,
435 name: call.name.clone(),
436 arguments: call.arguments.clone(),
437 })?;
438 projection.push_history(format!(
439 "Ktool call:\n{}",
440 host.render_tool_call(&call)?
441 ));
442 let budget = ContextBudget {
443 projection: projection.clone(),
444 max_input_tokens: selected.max_input_tokens,
445 };
446 let mut outcome = host
447 .execute_tool(call.clone(), child_operation_id, budget)
448 .await
449 .unwrap_or_else(|error| {
450 ToolOutcome::failure(format!("{} failed: {error}", call.name))
451 });
452 let exact_result = outcome.text.clone();
453 let initially_ok = outcome.ok;
454 let mut provider_result =
455 compact_tool_result(&outcome.text, &outcome.state_updates);
456 let mut candidate = projection.clone();
457 candidate.apply_updates(&outcome.state_updates);
458 candidate.push_history(format!("Ktool result:\n{provider_result}"));
459 let accepted = candidate.estimated_tokens() <= selected.max_input_tokens;
460 if accepted {
461 projection = candidate;
462 requires_rerender = !outcome.state_updates.is_empty();
463 } else {
464 outcome.ok = false;
465 outcome.capture = None;
466 provider_result = "The tool ran, but its result or updated state could not fit in the subagent context. Do not retry it; report the capacity failure to Kennedy.".into();
467 projection.push_history(format!("Ktool result:\n{provider_result}"));
468 }
469 host.record(AuditEvent::ToolResult {
470 parent_operation_id: request.parent_operation_id,
471 name: call.name.clone(),
472 ok: initially_ok,
473 projection_accepted: accepted,
474 result: exact_result,
475 })?;
476 pending_capture = outcome.capture.take();
477 respond_or_record(
478 &mut turn,
479 host,
480 request.parent_operation_id,
481 round + 1,
482 &manifest_hash,
483 &native.call_id,
484 if outcome.ok {
485 ToolResult::success(provider_result)
486 } else {
487 ToolResult::failure(provider_result)
488 },
489 )
490 .await?;
491 }
492 AgentEvent::Completed(completed) => break completed,
493 }
494 };
495 let receipt = turn
496 .receipt()
497 .context("subagent provider completed without a usage receipt")?
498 .clone();
499 host.record(AuditEvent::ProviderReceipt {
500 parent_operation_id: request.parent_operation_id,
501 round: round + 1,
502 manifest_hash: manifest_hash.clone(),
503 usage: completed.usage.clone(),
504 receipt: Box::new(receipt),
505 })?;
506
507 let capture = deferred_capture.take().or(pending_capture);
508 if let Some(capture) = capture {
509 if !capturing && completed.answer.is_empty() {
510 deferred_capture = Some(capture);
511 continue;
512 }
513 let budget = ContextBudget {
514 projection: projection.clone(),
515 max_input_tokens: selected.max_input_tokens,
516 };
517 let outcome = host
518 .complete_capture(capture, completed.answer, budget)
519 .await?;
520 let mut candidate = projection.clone();
521 candidate.apply_updates(&outcome.state_updates);
522 candidate.push_history(format!("Ktool result:\n{}", outcome.text));
523 ensure_capacity(&candidate, selected.max_input_tokens)?;
524 projection = candidate;
525 continue;
526 }
527 if requires_rerender {
528 let draft = completed.answer.trim();
529 if !draft.is_empty() {
530 projection.push_history(format!(
531 "Assistant draft produced before the state refresh:\n{draft}"
532 ));
533 }
534 continue;
535 }
536 let answer = completed.answer.trim().to_owned();
537 if !answer.is_empty() {
538 host.record(AuditEvent::Completed {
539 parent_operation_id: request.parent_operation_id,
540 model: request.model.clone(),
541 response: answer.clone(),
542 })?;
543 return Ok(RunResult {
544 answer,
545 model: selected,
546 });
547 }
548 ensure!(
549 used_tool,
550 "subagent provider completed without a response or tool call"
551 );
552 }
553 anyhow::bail!(
554 "subagent exceeded the {}-round tool-loop safety limit",
555 self.round_limit
556 )
557 }
558}
559
560async fn respond_or_record<H: Host>(
561 turn: &mut kcode_intelligence_router::AgentTurn,
562 host: &mut H,
563 parent_operation_id: Uuid,
564 round: u64,
565 manifest_hash: &str,
566 call_id: &str,
567 result: ToolResult,
568) -> anyhow::Result<()> {
569 if let Err(error) = turn.respond(call_id, result).await {
570 let receipt = turn.finish_unavailable()?.clone();
571 host.record(AuditEvent::ProviderReceipt {
572 parent_operation_id,
573 round,
574 manifest_hash: manifest_hash.into(),
575 usage: None,
576 receipt: Box::new(receipt),
577 })?;
578 return Err(anyhow::Error::new(error));
579 }
580 Ok(())
581}
582
583#[derive(Clone)]
584struct Projection {
585 context: Vec<String>,
586 task: String,
587 history: Vec<String>,
588 states: Vec<ProjectedState>,
589}
590
591#[derive(Clone)]
592struct ProjectedState {
593 key: String,
594 text: String,
595}
596
597impl Projection {
598 fn new(context: Vec<String>, task: String) -> Self {
599 Self {
600 context,
601 task,
602 history: Vec::new(),
603 states: Vec::new(),
604 }
605 }
606
607 fn render(&self) -> String {
608 self.context
609 .iter()
610 .map(String::as_str)
611 .chain(std::iter::once(self.task.as_str()))
612 .chain(self.history.iter().map(String::as_str))
613 .chain(self.states.iter().map(|state| state.text.as_str()))
614 .filter(|section| !section.is_empty())
615 .collect::<Vec<_>>()
616 .join("\n\n")
617 }
618
619 fn push_history(&mut self, text: impl Into<String>) {
620 self.history.push(text.into());
621 }
622
623 fn update_state(&mut self, key: String, text: Option<String>) {
624 self.states.retain(|state| state.key != key);
625 if let Some(text) = text {
626 self.states.push(ProjectedState { key, text });
627 }
628 }
629
630 fn apply_updates(&mut self, updates: &[StateUpdate]) {
631 for update in updates {
632 self.update_state(update.key.clone(), update.text.clone());
633 }
634 }
635
636 fn estimated_tokens(&self) -> u64 {
637 (self.render().chars().count() as u64)
638 .div_ceil(4)
639 .saturating_add(PROTOCOL_TOKEN_RESERVE)
640 }
641}
642
643fn compact_tool_result(text: &str, states: &[StateUpdate]) -> String {
644 if states.is_empty() {
645 return text.to_owned();
646 }
647 let result = if text.chars().count() <= INLINE_TOOL_RESULT_CHARACTERS {
648 text
649 } else {
650 "Tool completed successfully."
651 };
652 format!(
653 "{result}\n\nThe updated state will be rendered in the next fresh context slice; end this slice now."
654 )
655}
656
657fn ensure_capacity(projection: &Projection, max_input_tokens: u64) -> anyhow::Result<()> {
658 let estimated = projection.estimated_tokens();
659 ensure!(
660 estimated <= max_input_tokens,
661 "subagent context requires approximately {estimated} input tokens, over the selected model's {max_input_tokens}-token input limit"
662 );
663 Ok(())
664}
665
666fn ktool_definition() -> DynamicTool {
667 DynamicTool::new(
668 "call_ktool",
669 "Call one available Ktool by its exact name.",
670 json!({
671 "type": "object",
672 "additionalProperties": false,
673 "required": ["name", "arguments"],
674 "properties": {
675 "name": {"type": "string"},
676 "arguments": {"type": "object"}
677 }
678 }),
679 )
680}
681
682fn parse_ktool_call(call: &DynamicToolCall) -> anyhow::Result<ToolCall> {
683 ensure!(call.tool == "call_ktool", "unknown provider tool");
684 let arguments = call
685 .arguments
686 .as_object()
687 .context("call_ktool arguments must be an object")?;
688 ensure!(
689 arguments
690 .keys()
691 .all(|key| matches!(key.as_str(), "name" | "arguments")),
692 "call_ktool contains unknown arguments"
693 );
694 let name = arguments
695 .get("name")
696 .and_then(Value::as_str)
697 .map(str::trim)
698 .filter(|name| !name.is_empty() && name.chars().count() <= 100)
699 .context("call_ktool.name must be a non-empty bounded string")?
700 .to_owned();
701 let arguments = arguments
702 .get("arguments")
703 .filter(|value| value.is_object())
704 .context("call_ktool.arguments must be an object")?
705 .clone();
706 Ok(ToolCall { name, arguments })
707}
708
709fn reasoning_effort(value: &str) -> anyhow::Result<ReasoningEffort> {
710 Ok(match value {
711 "none" => ReasoningEffort::None,
712 "minimal" => ReasoningEffort::Minimal,
713 "low" => ReasoningEffort::Low,
714 "medium" => ReasoningEffort::Medium,
715 "high" => ReasoningEffort::High,
716 "xhigh" => ReasoningEffort::XHigh,
717 "max" => ReasoningEffort::Max,
718 _ => anyhow::bail!("unsupported reasoning effort {value:?}"),
719 })
720}
721
722#[cfg(test)]
723mod tests {
724 use super::*;
725
726 #[test]
727 fn projection_replaces_state_and_budget_accounts_for_reserve() {
728 let mut projection = Projection::new(vec!["context".into()], "task".into());
729 projection.update_state("file".into(), Some("old".into()));
730 projection.update_state("file".into(), Some("new".into()));
731 assert_eq!(projection.states.len(), 1);
732 assert!(projection.render().contains("new"));
733 assert!(!projection.render().contains("old"));
734 assert!(projection.estimated_tokens() >= PROTOCOL_TOKEN_RESERVE);
735 }
736
737 #[test]
738 fn state_changes_compact_large_tool_results() {
739 let compacted = compact_tool_result(
740 &"x".repeat(INLINE_TOOL_RESULT_CHARACTERS + 1),
741 &[StateUpdate {
742 key: "state".into(),
743 text: Some("current".into()),
744 }],
745 );
746 assert!(compacted.starts_with("Tool completed successfully."));
747 assert!(compacted.contains("fresh context slice"));
748 }
749
750 #[test]
751 fn native_tool_wrapper_is_strict() {
752 let call = parse_ktool_call(&DynamicToolCall {
753 call_id: "1".into(),
754 tool: "call_ktool".into(),
755 arguments: json!({"name": "Read", "arguments": {"id": 1}}),
756 })
757 .unwrap();
758 assert_eq!(call.name, "Read");
759 assert_eq!(call.arguments["id"], 1);
760 }
761}