kcode_agent_runtime/
lib.rs

1//! Provider-neutral subagent loops over `kcode-intelligence-router`.
2
3#![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::{Intelligence, ResolvedAgentModel};
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
21/// A boxed asynchronous host operation.
22pub type HostFuture<'a, T> = Pin<Box<dyn Future<Output = anyhow::Result<T>> + Send + 'a>>;
23
24/// One application tool call requested by a subagent.
25#[derive(Clone, Debug, PartialEq)]
26pub struct ToolCall {
27    /// Exact application tool name.
28    pub name: String,
29    /// Tool arguments.
30    pub arguments: Value,
31}
32
33/// One replaceable state section rendered into every later context slice.
34#[derive(Clone, Debug, Eq, PartialEq)]
35pub struct StateUpdate {
36    /// Stable state identity. A later update with this key replaces the prior text.
37    pub key: String,
38    /// Current rendered state, or `None` to remove it.
39    pub text: Option<String>,
40}
41
42/// Result returned by the application after one tool or capture operation.
43#[derive(Clone, Debug, PartialEq)]
44pub struct ToolOutcome {
45    /// Exact result retained for audit.
46    pub text: String,
47    /// Whether the operation succeeded.
48    pub ok: bool,
49    /// Replaceable state made current by the operation.
50    pub state_updates: Vec<StateUpdate>,
51    /// Opaque application token requesting a tool-free freeform output capture.
52    pub capture: Option<Value>,
53}
54
55impl ToolOutcome {
56    /// Constructs a simple successful result.
57    pub fn success(text: impl Into<String>) -> Self {
58        Self {
59            text: text.into(),
60            ok: true,
61            state_updates: Vec::new(),
62            capture: None,
63        }
64    }
65
66    /// Constructs a simple failed result.
67    pub fn failure(text: impl Into<String>) -> Self {
68        Self {
69            text: text.into(),
70            ok: false,
71            state_updates: Vec::new(),
72            capture: None,
73        }
74    }
75}
76
77/// Read-only capacity view supplied while a host evaluates a tool.
78#[derive(Clone)]
79pub struct ContextBudget {
80    projection: Projection,
81    max_input_tokens: u64,
82}
83
84impl ContextBudget {
85    /// Current estimated input tokens, including protocol reserve.
86    pub fn estimated_tokens(&self) -> u64 {
87        self.projection.estimated_tokens()
88    }
89
90    /// Maximum permitted input tokens.
91    pub fn max_input_tokens(&self) -> u64 {
92        self.max_input_tokens
93    }
94
95    /// Returns whether replacing one projected state would fit.
96    pub fn fits_state(&self, key: impl Into<String>, text: impl Into<String>) -> bool {
97        let mut projection = self.projection.clone();
98        projection.update_state(key.into(), Some(text.into()));
99        projection.estimated_tokens() <= self.max_input_tokens
100    }
101}
102
103/// Application-owned behavior invoked by the generic subagent loop.
104pub trait Host: Send {
105    /// Renders the retained invocation text before execution.
106    fn render_tool_call(&mut self, call: &ToolCall) -> anyhow::Result<String>;
107
108    /// Executes one application tool.
109    fn execute_tool<'a>(
110        &'a mut self,
111        call: ToolCall,
112        operation_id: Uuid,
113        budget: ContextBudget,
114    ) -> HostFuture<'a, ToolOutcome>;
115
116    /// Completes an opaque freeform capture requested by a prior tool result.
117    fn complete_capture<'a>(
118        &'a mut self,
119        capture: Value,
120        contents: String,
121        budget: ContextBudget,
122    ) -> HostFuture<'a, ToolOutcome>;
123
124    /// Records one durable audit event selected by the runtime.
125    fn record(&mut self, label: &str, value: Value) -> anyhow::Result<()>;
126}
127
128/// Inputs for one complete subagent run.
129#[derive(Clone, Debug, PartialEq)]
130pub struct RunRequest {
131    /// Stable user identifier used for router accounting.
132    pub user_id: String,
133    /// Running parent operation whose cancellation propagates to each turn.
134    pub parent_operation_id: Uuid,
135    /// Exact requested model selector.
136    pub model: String,
137    /// Provider-neutral reasoning effort.
138    pub reasoning_effort: String,
139    /// Ordered immutable context sections.
140    pub context: Vec<String>,
141    /// Exact task presented after the context.
142    pub task: String,
143    /// Optional per-turn timeout.
144    pub timeout: Option<Duration>,
145    /// Additional application metadata included in the start audit event.
146    pub start_metadata: Value,
147}
148
149/// Completed subagent output.
150#[derive(Clone, Debug, Eq, PartialEq)]
151pub struct RunResult {
152    /// Final non-empty assistant answer.
153    pub answer: String,
154    /// Model used for every turn.
155    pub model: ResolvedAgentModel,
156}
157
158/// Cloneable provider-neutral subagent runtime.
159#[derive(Clone)]
160pub struct AgentRuntime {
161    intelligence: Intelligence,
162    round_limit: u64,
163}
164
165impl AgentRuntime {
166    /// Constructs a runtime over the sole direct-model boundary.
167    pub fn new(intelligence: Intelligence) -> Self {
168        Self {
169            intelligence,
170            round_limit: DEFAULT_ROUND_LIMIT,
171        }
172    }
173
174    /// Resolves a model without running an agent.
175    pub async fn resolve_model(&self, requested: &str) -> anyhow::Result<ResolvedAgentModel> {
176        self.intelligence
177            .resolve_agent_model(requested)
178            .await
179            .map_err(anyhow::Error::new)
180    }
181
182    /// Runs one fresh-context subagent to a final non-empty answer.
183    pub async fn run<H: Host>(
184        &self,
185        request: RunRequest,
186        host: &mut H,
187    ) -> anyhow::Result<RunResult> {
188        let selected = self.resolve_model(&request.model).await?;
189        let reasoning_effort = reasoning_effort(&request.reasoning_effort)?;
190        let mut projection = Projection::new(request.context, request.task);
191        ensure_capacity(&projection, selected.max_input_tokens)?;
192        host.record(
193            "subagent_started",
194            json!({
195                "model": request.model,
196                "providerModel": selected.provider_model,
197                "provider": format!("{:?}", selected.provider),
198                "contextWindowTokens": selected.context_window_tokens,
199                "maxInputTokens": selected.max_input_tokens,
200                "context": projection.context,
201                "task": projection.task,
202                "host": request.start_metadata,
203            }),
204        )?;
205        let user = self
206            .intelligence
207            .for_user(request.user_id)
208            .map_err(anyhow::Error::new)?;
209        let mut deferred_capture: Option<Value> = None;
210
211        for round in 0..self.round_limit {
212            let capturing = deferred_capture.is_some();
213            ensure_capacity(&projection, selected.max_input_tokens)?;
214            let input = projection.render();
215            let manifest_hash = hex::encode(Sha256::digest(input.as_bytes()));
216            host.record(
217                "subagent_inference_submitted",
218                json!({
219                    "round": round + 1,
220                    "manifestHash": manifest_hash,
221                    "estimatedInputTokens": projection.estimated_tokens(),
222                }),
223            )?;
224            let mut provider_request = AgentRequest::new(input, selected.requested_model.clone());
225            provider_request.reasoning_effort = reasoning_effort;
226            provider_request.ephemeral = true;
227            provider_request.tools = if capturing {
228                Vec::new()
229            } else {
230                vec![ktool_definition()]
231            };
232            if let Some(timeout) = request.timeout {
233                provider_request.timeout = timeout;
234            }
235            let child_operation_id = Uuid::new_v4();
236            let mut turn = user
237                .start_agent_turn(
238                    child_operation_id,
239                    Some(request.parent_operation_id),
240                    provider_request,
241                )
242                .await
243                .map_err(anyhow::Error::new)?;
244            let mut used_tool = false;
245            let mut pending_capture: Option<Value> = None;
246            let mut requires_rerender = false;
247            let completed = loop {
248                let event = turn
249                    .next_event()
250                    .await
251                    .map_err(anyhow::Error::new)?
252                    .context("subagent provider ended without a terminal turn event")?;
253                match event {
254                    AgentEvent::ProviderInput(_) => {}
255                    AgentEvent::ToolCall(native) => {
256                        used_tool = true;
257                        if capturing {
258                            turn.respond(
259                                &native.call_id,
260                                ToolResult::failure(
261                                    "No application tool is available while complete freeform output is being captured.",
262                                ),
263                            )
264                            .await
265                            .map_err(anyhow::Error::new)?;
266                            continue;
267                        }
268                        if pending_capture.is_some() {
269                            turn.respond(
270                                &native.call_id,
271                                ToolResult::failure(
272                                    "A freeform output capture is pending; no other tool can run first.",
273                                ),
274                            )
275                            .await
276                            .map_err(anyhow::Error::new)?;
277                            continue;
278                        }
279                        if requires_rerender {
280                            turn.respond(
281                                &native.call_id,
282                                ToolResult::failure(
283                                    "A state update is waiting to be re-rendered. End this slice before calling another tool.",
284                                ),
285                            )
286                            .await
287                            .map_err(anyhow::Error::new)?;
288                            continue;
289                        }
290                        let call = match parse_ktool_call(&native) {
291                            Ok(call) => call,
292                            Err(error) => {
293                                let text = format!("Invalid application tool call: {error}");
294                                projection.push_history(format!("Ktool result:\n{text}"));
295                                turn.respond(&native.call_id, ToolResult::failure(text))
296                                    .await
297                                    .map_err(anyhow::Error::new)?;
298                                continue;
299                            }
300                        };
301                        host.record(
302                            "subagent_tool_call",
303                            json!({"name": call.name, "arguments": call.arguments}),
304                        )?;
305                        projection.push_history(format!(
306                            "Ktool call:\n{}",
307                            host.render_tool_call(&call)?
308                        ));
309                        let budget = ContextBudget {
310                            projection: projection.clone(),
311                            max_input_tokens: selected.max_input_tokens,
312                        };
313                        let mut outcome = host
314                            .execute_tool(call.clone(), child_operation_id, budget)
315                            .await
316                            .unwrap_or_else(|error| {
317                                ToolOutcome::failure(format!("{} failed: {error}", call.name))
318                            });
319                        let exact_result = outcome.text.clone();
320                        let initially_ok = outcome.ok;
321                        let mut provider_result =
322                            compact_tool_result(&outcome.text, &outcome.state_updates);
323                        let mut candidate = projection.clone();
324                        candidate.apply_updates(&outcome.state_updates);
325                        candidate.push_history(format!("Ktool result:\n{provider_result}"));
326                        let accepted = candidate.estimated_tokens() <= selected.max_input_tokens;
327                        if accepted {
328                            projection = candidate;
329                            requires_rerender = !outcome.state_updates.is_empty();
330                        } else {
331                            outcome.ok = false;
332                            outcome.capture = None;
333                            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();
334                            projection.push_history(format!("Ktool result:\n{provider_result}"));
335                        }
336                        host.record(
337                            "subagent_tool_result",
338                            json!({
339                                "name": call.name,
340                                "ok": initially_ok,
341                                "projectionAccepted": accepted,
342                                "result": exact_result,
343                            }),
344                        )?;
345                        pending_capture = outcome.capture.take();
346                        turn.respond(
347                            &native.call_id,
348                            if outcome.ok {
349                                ToolResult::success(provider_result)
350                            } else {
351                                ToolResult::failure(provider_result)
352                            },
353                        )
354                        .await
355                        .map_err(anyhow::Error::new)?;
356                    }
357                    AgentEvent::Completed(completed) => break completed,
358                }
359            };
360            host.record(
361                "subagent_provider_receipt",
362                json!({
363                    "round": round + 1,
364                    "usage": completed.usage.as_ref().map(|usage| json!({
365                        "inputTokens": usage.input_tokens,
366                        "outputTokens": usage.output_tokens,
367                        "cachedInputTokens": usage.cached_input_tokens,
368                        "reasoningOutputTokens": usage.reasoning_output_tokens,
369                        "lastInputTokens": usage.last_input_tokens,
370                        "lastOutputTokens": usage.last_output_tokens,
371                    })),
372                }),
373            )?;
374
375            let capture = deferred_capture.take().or(pending_capture);
376            if let Some(capture) = capture {
377                if !capturing && completed.answer.is_empty() {
378                    deferred_capture = Some(capture);
379                    continue;
380                }
381                let budget = ContextBudget {
382                    projection: projection.clone(),
383                    max_input_tokens: selected.max_input_tokens,
384                };
385                let outcome = host
386                    .complete_capture(capture, completed.answer, budget)
387                    .await?;
388                let mut candidate = projection.clone();
389                candidate.apply_updates(&outcome.state_updates);
390                candidate.push_history(format!("Ktool result:\n{}", outcome.text));
391                ensure_capacity(&candidate, selected.max_input_tokens)?;
392                projection = candidate;
393                continue;
394            }
395            if requires_rerender {
396                let draft = completed.answer.trim();
397                if !draft.is_empty() {
398                    projection.push_history(format!(
399                        "Assistant draft produced before the state refresh:\n{draft}"
400                    ));
401                }
402                continue;
403            }
404            let answer = completed.answer.trim().to_owned();
405            if !answer.is_empty() {
406                host.record(
407                    "subagent_completed",
408                    json!({"model": request.model, "response": answer}),
409                )?;
410                return Ok(RunResult {
411                    answer,
412                    model: selected,
413                });
414            }
415            ensure!(
416                used_tool,
417                "subagent provider completed without a response or tool call"
418            );
419        }
420        anyhow::bail!(
421            "subagent exceeded the {}-round tool-loop safety limit",
422            self.round_limit
423        )
424    }
425}
426
427#[derive(Clone)]
428struct Projection {
429    context: Vec<String>,
430    task: String,
431    history: Vec<String>,
432    states: Vec<ProjectedState>,
433}
434
435#[derive(Clone)]
436struct ProjectedState {
437    key: String,
438    text: String,
439}
440
441impl Projection {
442    fn new(context: Vec<String>, task: String) -> Self {
443        Self {
444            context,
445            task,
446            history: Vec::new(),
447            states: Vec::new(),
448        }
449    }
450
451    fn render(&self) -> String {
452        self.context
453            .iter()
454            .map(String::as_str)
455            .chain(std::iter::once(self.task.as_str()))
456            .chain(self.history.iter().map(String::as_str))
457            .chain(self.states.iter().map(|state| state.text.as_str()))
458            .filter(|section| !section.is_empty())
459            .collect::<Vec<_>>()
460            .join("\n\n")
461    }
462
463    fn push_history(&mut self, text: impl Into<String>) {
464        self.history.push(text.into());
465    }
466
467    fn update_state(&mut self, key: String, text: Option<String>) {
468        self.states.retain(|state| state.key != key);
469        if let Some(text) = text {
470            self.states.push(ProjectedState { key, text });
471        }
472    }
473
474    fn apply_updates(&mut self, updates: &[StateUpdate]) {
475        for update in updates {
476            self.update_state(update.key.clone(), update.text.clone());
477        }
478    }
479
480    fn estimated_tokens(&self) -> u64 {
481        (self.render().chars().count() as u64)
482            .div_ceil(4)
483            .saturating_add(PROTOCOL_TOKEN_RESERVE)
484    }
485}
486
487fn compact_tool_result(text: &str, states: &[StateUpdate]) -> String {
488    if states.is_empty() {
489        return text.to_owned();
490    }
491    let result = if text.chars().count() <= INLINE_TOOL_RESULT_CHARACTERS {
492        text
493    } else {
494        "Tool completed successfully."
495    };
496    format!(
497        "{result}\n\nThe updated state will be rendered in the next fresh context slice; end this slice now."
498    )
499}
500
501fn ensure_capacity(projection: &Projection, max_input_tokens: u64) -> anyhow::Result<()> {
502    let estimated = projection.estimated_tokens();
503    ensure!(
504        estimated <= max_input_tokens,
505        "subagent context requires approximately {estimated} input tokens, over the selected model's {max_input_tokens}-token input limit"
506    );
507    Ok(())
508}
509
510fn ktool_definition() -> DynamicTool {
511    DynamicTool::new(
512        "call_ktool",
513        "Call one available Ktool by its exact name.",
514        json!({
515            "type": "object",
516            "additionalProperties": false,
517            "required": ["name", "arguments"],
518            "properties": {
519                "name": {"type": "string"},
520                "arguments": {"type": "object"}
521            }
522        }),
523    )
524}
525
526fn parse_ktool_call(call: &DynamicToolCall) -> anyhow::Result<ToolCall> {
527    ensure!(call.tool == "call_ktool", "unknown provider tool");
528    let arguments = call
529        .arguments
530        .as_object()
531        .context("call_ktool arguments must be an object")?;
532    ensure!(
533        arguments
534            .keys()
535            .all(|key| matches!(key.as_str(), "name" | "arguments")),
536        "call_ktool contains unknown arguments"
537    );
538    let name = arguments
539        .get("name")
540        .and_then(Value::as_str)
541        .map(str::trim)
542        .filter(|name| !name.is_empty() && name.chars().count() <= 100)
543        .context("call_ktool.name must be a non-empty bounded string")?
544        .to_owned();
545    let arguments = arguments
546        .get("arguments")
547        .filter(|value| value.is_object())
548        .context("call_ktool.arguments must be an object")?
549        .clone();
550    Ok(ToolCall { name, arguments })
551}
552
553fn reasoning_effort(value: &str) -> anyhow::Result<ReasoningEffort> {
554    Ok(match value {
555        "none" => ReasoningEffort::None,
556        "minimal" => ReasoningEffort::Minimal,
557        "low" => ReasoningEffort::Low,
558        "medium" => ReasoningEffort::Medium,
559        "high" => ReasoningEffort::High,
560        "xhigh" => ReasoningEffort::XHigh,
561        "max" => ReasoningEffort::Max,
562        _ => anyhow::bail!("unsupported reasoning effort {value:?}"),
563    })
564}
565
566#[cfg(test)]
567mod tests {
568    use super::*;
569
570    #[test]
571    fn projection_replaces_state_and_budget_accounts_for_reserve() {
572        let mut projection = Projection::new(vec!["context".into()], "task".into());
573        projection.update_state("file".into(), Some("old".into()));
574        projection.update_state("file".into(), Some("new".into()));
575        assert_eq!(projection.states.len(), 1);
576        assert!(projection.render().contains("new"));
577        assert!(!projection.render().contains("old"));
578        assert!(projection.estimated_tokens() >= PROTOCOL_TOKEN_RESERVE);
579    }
580
581    #[test]
582    fn state_changes_compact_large_tool_results() {
583        let compacted = compact_tool_result(
584            &"x".repeat(INLINE_TOOL_RESULT_CHARACTERS + 1),
585            &[StateUpdate {
586                key: "state".into(),
587                text: Some("current".into()),
588            }],
589        );
590        assert!(compacted.starts_with("Tool completed successfully."));
591        assert!(compacted.contains("fresh context slice"));
592    }
593
594    #[test]
595    fn native_tool_wrapper_is_strict() {
596        let call = parse_ktool_call(&DynamicToolCall {
597            call_id: "1".into(),
598            tool: "call_ktool".into(),
599            arguments: json!({"name": "Read", "arguments": {"id": 1}}),
600        })
601        .unwrap();
602        assert_eq!(call.name, "Read");
603        assert_eq!(call.arguments["id"], 1);
604    }
605}
kcode_agent_runtime/lib.rs

kcode_agent_runtime/
lib.rs
