kcode_agent_runtime/
lib.rs

1//! Provider-neutral agent 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, ModelContext, ReasoningEffort,
11    TokenUsage, ToolResult,
12};
13use kcode_intelligence_router::{
14    AgentContinuation, AgentProvider, Intelligence, ResolvedAgentModel, UsageReceipt,
15};
16use serde_json::{Value, json};
17use sha2::{Digest, Sha256};
18use uuid::Uuid;
19
20/// Default maximum provider calls for primary and delegated agent loops.
21pub const DEFAULT_ROUND_LIMIT: u64 = 250;
22const PROTOCOL_TOKEN_RESERVE: u64 = 4_096;
23const SUPERSEDED_TOOL_OUTPUT: &str = "[Tool output was displayed here, but has since been updated and now appears elsewhere in the context]";
24const REMOVED_TOOL_OUTPUT: &str = "[Tool output was displayed here, but is no longer current]";
25
26/// A boxed asynchronous host operation.
27pub type HostFuture<'a, T> = Pin<Box<dyn Future<Output = anyhow::Result<T>> + Send + 'a>>;
28
29/// One application tool call requested by a subagent.
30#[derive(Clone, Debug, PartialEq)]
31pub struct ToolCall {
32    /// Exact application tool name.
33    pub name: String,
34    /// Tool arguments.
35    pub arguments: Value,
36}
37
38/// One complete, typed audit fact selected by the subagent runtime.
39#[derive(Clone, Debug, PartialEq)]
40pub enum AuditEvent {
41    /// Immutable starting context and resolved provider capacity.
42    Started {
43        /// Parent operation that causally owns the subagent.
44        parent_operation_id: Uuid,
45        /// Caller-selected model identifier.
46        model: String,
47        /// Exact provider model.
48        provider_model: String,
49        /// Resolved provider transport.
50        provider: AgentProvider,
51        /// Total provider context window.
52        context_window_tokens: u64,
53        /// Maximum permitted input.
54        max_input_tokens: u64,
55        /// Ordered immutable starting context.
56        context: Vec<String>,
57        /// Focused subagent task.
58        task: String,
59        /// Opaque application metadata supplied at launch.
60        host: Value,
61    },
62    /// Exact input identity submitted for one provider round.
63    InferenceSubmitted {
64        /// Parent operation that causally owns the subagent.
65        parent_operation_id: Uuid,
66        /// One-based subagent round.
67        round: u64,
68        /// SHA-256 hash of the exact provider input.
69        manifest_hash: String,
70        /// Runtime input estimate including protocol reserve.
71        estimated_input_tokens: u64,
72    },
73    /// Application tool invocation requested by the provider.
74    ToolCall {
75        /// Parent operation that causally owns the subagent.
76        parent_operation_id: Uuid,
77        /// Exact application tool name.
78        name: String,
79        /// Exact tool arguments.
80        arguments: Value,
81    },
82    /// Complete application tool result retained for audit.
83    ToolResult {
84        /// Parent operation that causally owns the subagent.
85        parent_operation_id: Uuid,
86        /// Exact application tool name.
87        name: String,
88        /// Whether execution itself succeeded.
89        ok: bool,
90        /// Whether the resulting context projection fit.
91        projection_accepted: bool,
92        /// Exact application result.
93        result: String,
94    },
95    /// Canonical accounting for one completed or interrupted provider round.
96    ProviderReceipt {
97        /// Parent operation that causally owns the subagent.
98        parent_operation_id: Uuid,
99        /// One-based subagent round.
100        round: u64,
101        /// SHA-256 hash of the exact provider input.
102        manifest_hash: String,
103        /// Provider-native cumulative usage retained for exact audit detail.
104        usage: Option<kcode_codex_runtime_v2::TokenUsage>,
105        /// Canonical durable receipt written by the intelligence router.
106        receipt: Box<UsageReceipt>,
107    },
108    /// Final non-empty subagent answer.
109    Completed {
110        /// Parent operation that causally owns the subagent.
111        parent_operation_id: Uuid,
112        /// Caller-selected model identifier.
113        model: String,
114        /// Terminal response returned to the application.
115        response: String,
116    },
117}
118
119/// One replaceable state section rendered into every later context slice.
120#[derive(Clone, Debug, Eq, PartialEq)]
121pub struct StateUpdate {
122    /// Stable state identity. A later update with this key replaces the prior text.
123    pub key: String,
124    /// Current rendered state, or `None` to remove it.
125    pub text: Option<String>,
126}
127
128/// Result returned by the application after one tool or capture operation.
129#[derive(Clone, Debug, PartialEq)]
130pub struct ToolOutcome {
131    /// Exact result retained for audit.
132    pub text: String,
133    /// Whether the operation succeeded.
134    pub ok: bool,
135    /// Replaceable state made current by the operation.
136    pub state_updates: Vec<StateUpdate>,
137    /// State identities whose complete values this result displays.
138    ///
139    /// A later update to one of these identities replaces this retained result
140    /// with a concise supersession marker while preserving the exact audit
141    /// result.
142    pub displayed_state_keys: Vec<String>,
143    /// Opaque application token requesting a tool-free freeform output capture.
144    pub capture: Option<Value>,
145}
146
147impl ToolOutcome {
148    /// Constructs a simple successful result.
149    pub fn success(text: impl Into<String>) -> Self {
150        Self {
151            text: text.into(),
152            ok: true,
153            state_updates: Vec::new(),
154            displayed_state_keys: Vec::new(),
155            capture: None,
156        }
157    }
158
159    /// Constructs a simple failed result.
160    pub fn failure(text: impl Into<String>) -> Self {
161        Self {
162            text: text.into(),
163            ok: false,
164            state_updates: Vec::new(),
165            displayed_state_keys: Vec::new(),
166            capture: None,
167        }
168    }
169}
170
171/// Read-only capacity view supplied while a host evaluates a tool.
172#[derive(Clone)]
173pub struct ContextBudget {
174    projection: Projection,
175    max_input_tokens: u64,
176}
177
178impl ContextBudget {
179    /// Current estimated input tokens, including protocol reserve.
180    pub fn estimated_tokens(&self) -> u64 {
181        self.projection.estimated_tokens()
182    }
183
184    /// Maximum permitted input tokens.
185    pub fn max_input_tokens(&self) -> u64 {
186        self.max_input_tokens
187    }
188
189    /// Returns whether replacing one projected state would fit.
190    pub fn fits_state(&self, key: impl Into<String>, text: impl Into<String>) -> bool {
191        let mut projection = self.projection.clone();
192        projection.apply_updates(&[StateUpdate {
193            key: key.into(),
194            text: Some(text.into()),
195        }]);
196        projection.estimated_tokens() <= self.max_input_tokens
197    }
198}
199
200/// Application-owned behavior invoked by the generic subagent loop.
201pub trait Host: Send {
202    /// Renders the retained invocation text before execution.
203    fn render_tool_call(&mut self, call: &ToolCall) -> anyhow::Result<String>;
204
205    /// Executes one application tool.
206    fn execute_tool<'a>(
207        &'a mut self,
208        call: ToolCall,
209        operation_id: Uuid,
210        budget: ContextBudget,
211    ) -> HostFuture<'a, ToolOutcome>;
212
213    /// Completes an opaque freeform capture requested by a prior tool result.
214    fn complete_capture<'a>(
215        &'a mut self,
216        capture: Value,
217        contents: String,
218        budget: ContextBudget,
219    ) -> HostFuture<'a, ToolOutcome>;
220
221    /// Records one durable audit event selected by the runtime.
222    fn record(&mut self, event: AuditEvent) -> anyhow::Result<()>;
223}
224
225/// Inputs shared by every provider call in one primary session run.
226#[derive(Clone, Debug, Eq, PartialEq)]
227pub struct SessionRunRequest {
228    /// Stable user identifier used for router accounting.
229    pub user_id: String,
230    /// Top-level operation whose cancellation propagates to every provider call.
231    pub operation_id: Uuid,
232    /// Number of provider rounds already durably used by a restored session.
233    pub rounds_used: u64,
234    /// Maximum cumulative provider rounds permitted for the logical turn.
235    pub round_limit: u64,
236}
237
238/// Exact provider input prepared by the application for one primary round.
239#[derive(Clone, Debug, Eq, PartialEq)]
240pub struct PreparedRound {
241    /// Complete logical context retained for durable audit and capacity checks.
242    pub input: String,
243    /// New context submitted to the provider for this native turn.
244    pub provider_input: String,
245    /// Prior native Codex state when this turn resumes an affine thread.
246    pub continuation: Option<AgentContinuation>,
247    /// Exact requested model identifier.
248    pub model: String,
249    /// Provider-neutral reasoning effort.
250    pub reasoning_effort: String,
251    /// Session-specific explanation attached to the single application-tool bridge.
252    pub tool_description: String,
253    /// Optional complete provider-turn timeout.
254    pub timeout: Option<Duration>,
255}
256
257/// Result of preparing the next primary round.
258#[derive(Clone, Debug, Eq, PartialEq)]
259pub enum RoundPreparation {
260    /// Run the prepared provider round.
261    Run(PreparedRound),
262    /// Finish without starting another provider round.
263    Complete(Option<String>),
264}
265
266/// One durable provider-protocol fact emitted by the primary runtime.
267#[derive(Clone, Debug, PartialEq)]
268pub enum SessionEvent {
269    /// Exact input identity submitted for one provider round.
270    InferenceSubmitted {
271        /// Cumulative one-based provider round.
272        round: u64,
273        /// SHA-256 hash of the exact provider input.
274        manifest_hash: String,
275        /// Exact requested model identifier.
276        model: String,
277    },
278    /// Exact caller-controlled material submitted to the model.
279    ProviderInput {
280        /// Cumulative one-based provider round.
281        round: u64,
282        /// Normalized model input and structured material.
283        context: ModelContext,
284    },
285    /// Live usage for the current provider call.
286    UsageUpdated {
287        /// Cumulative one-based provider round.
288        round: u64,
289        /// Provider-native current-call usage, normalized from thread totals.
290        usage: TokenUsage,
291    },
292    /// Canonical receipt and final usage for a completed or interrupted call.
293    ProviderReceipt {
294        /// Cumulative one-based provider round.
295        round: u64,
296        /// Final provider-native current-call usage when available.
297        usage: Option<TokenUsage>,
298        /// Canonical durable receipt written by the intelligence router.
299        receipt: Box<UsageReceipt>,
300        /// Native Codex state after successful completion.
301        continuation: Option<AgentContinuation>,
302    },
303}
304
305/// Complete application handling of one primary-session tool call.
306#[derive(Clone, Debug, PartialEq)]
307pub struct SessionToolOutcome {
308    /// Exact result returned through the provider's native tool protocol.
309    pub text: String,
310    /// Whether tool execution succeeded.
311    pub ok: bool,
312    /// Opaque token requesting tool-free freeform output capture.
313    pub capture: Option<Value>,
314    /// Whether the logical session must stop after this tool response.
315    pub stop: bool,
316    /// Whether the logical session should finish after this provider round.
317    pub finish_after_round: bool,
318    /// Whether this tool already emitted the session's externally visible response.
319    pub emitted_response: bool,
320}
321
322impl SessionToolOutcome {
323    /// Constructs a simple successful tool result.
324    pub fn success(text: impl Into<String>) -> Self {
325        Self {
326            text: text.into(),
327            ok: true,
328            capture: None,
329            stop: false,
330            finish_after_round: false,
331            emitted_response: false,
332        }
333    }
334
335    /// Constructs a simple failed tool result.
336    pub fn failure(text: impl Into<String>) -> Self {
337        Self {
338            text: text.into(),
339            ok: false,
340            capture: None,
341            stop: false,
342            finish_after_round: false,
343            emitted_response: false,
344        }
345    }
346}
347
348/// Application preparation immediately before a primary provider turn resumes.
349#[derive(Clone, Debug, PartialEq)]
350pub enum ProviderResume {
351    /// Send the prepared native tool result and allow another inference.
352    Continue(SessionToolOutcome),
353    /// Finish the logical session without starting another inference.
354    Complete(Option<String>),
355    /// Account for the interrupted call and select another outer round without responding.
356    RestartFresh,
357}
358
359enum ProviderResumeAction {
360    Respond(SessionToolOutcome),
361    Interrupt(PostReceiptDisposition),
362}
363
364enum PostReceiptDisposition {
365    Return(Option<String>),
366    NextRound,
367}
368
369fn provider_resume_action(resume: ProviderResume) -> ProviderResumeAction {
370    match resume {
371        ProviderResume::Continue(outcome) => ProviderResumeAction::Respond(outcome),
372        ProviderResume::Complete(answer) => {
373            ProviderResumeAction::Interrupt(PostReceiptDisposition::Return(answer))
374        }
375        ProviderResume::RestartFresh => {
376            ProviderResumeAction::Interrupt(PostReceiptDisposition::NextRound)
377        }
378    }
379}
380
381/// Provider completion summarized for the application-owned session policy.
382#[derive(Clone, Debug, Eq, PartialEq)]
383pub struct RoundCompletion {
384    /// Terminal assistant text, which may be empty after tool use.
385    pub answer: String,
386    /// Whether the provider called at least one application tool.
387    pub used_tool: bool,
388    /// Whether a successful tool requested session completion after this round.
389    pub finish_requested: bool,
390    /// Whether a tool already emitted the externally visible response.
391    pub emitted_response: bool,
392}
393
394/// Application decision after completing a semantic primary-session transition.
395#[derive(Clone, Debug, Eq, PartialEq)]
396pub enum SessionControl {
397    /// Prepare another affine provider round.
398    Continue,
399    /// Finish the logical turn with the optional externally visible answer.
400    Complete(Option<String>),
401}
402
403/// Application-owned semantics invoked by the primary agent runtime.
404pub trait SessionHost: Send {
405    /// Prepares the complete context and runtime selection for one provider round.
406    fn prepare_round<'a>(&'a mut self, round: u64) -> HostFuture<'a, RoundPreparation>;
407
408    /// Records one provider-protocol fact and durably checkpoints its effects.
409    fn record<'a>(&'a mut self, event: SessionEvent) -> HostFuture<'a, ()>;
410
411    /// Executes and durably projects one parsed or invalid application tool call.
412    fn execute_tool<'a>(
413        &'a mut self,
414        call: anyhow::Result<ToolCall>,
415        provider_operation_id: Uuid,
416    ) -> HostFuture<'a, SessionToolOutcome>;
417
418    /// Reconciles application policy at the actual provider-resume boundary.
419    fn prepare_provider_resume<'a>(
420        &'a mut self,
421        outcome: SessionToolOutcome,
422    ) -> HostFuture<'a, ProviderResume> {
423        Box::pin(async move { Ok(ProviderResume::Continue(outcome)) })
424    }
425
426    /// Completes an opaque freeform capture requested by a prior tool result.
427    fn complete_capture<'a>(
428        &'a mut self,
429        capture: Value,
430        contents: String,
431    ) -> HostFuture<'a, SessionControl>;
432
433    /// Applies the terminal provider answer and decides whether another round is needed.
434    fn complete_round<'a>(
435        &'a mut self,
436        completion: RoundCompletion,
437    ) -> HostFuture<'a, SessionControl>;
438}
439
440/// Error returned when a restored primary turn exhausts its cumulative round budget.
441#[derive(Debug)]
442pub struct SessionRoundLimitError {
443    limit: u64,
444}
445
446impl std::fmt::Display for SessionRoundLimitError {
447    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
448        write!(
449            formatter,
450            "agent exceeded the {}-round tool-loop safety limit",
451            self.limit
452        )
453    }
454}
455
456impl std::error::Error for SessionRoundLimitError {}
457
458/// Returns whether an error is the primary runtime's round-limit signal.
459pub fn is_session_round_limit(error: &anyhow::Error) -> bool {
460    error.downcast_ref::<SessionRoundLimitError>().is_some()
461}
462
463/// Inputs for one complete subagent run.
464#[derive(Clone, Debug, PartialEq)]
465pub struct RunRequest {
466    /// Stable user identifier used for router accounting.
467    pub user_id: String,
468    /// Running parent operation whose cancellation propagates to each turn.
469    pub parent_operation_id: Uuid,
470    /// Exact requested model selector.
471    pub model: String,
472    /// Provider-neutral reasoning effort.
473    pub reasoning_effort: String,
474    /// Ordered immutable starting context sections.
475    pub context: Vec<String>,
476    /// Exact task presented after the context.
477    pub task: String,
478    /// Optional per-turn timeout.
479    pub timeout: Option<Duration>,
480    /// Additional application metadata included in the start audit event.
481    pub start_metadata: Value,
482}
483
484/// Completed subagent output.
485#[derive(Clone, Debug, Eq, PartialEq)]
486pub struct RunResult {
487    /// Final non-empty assistant answer.
488    pub answer: String,
489    /// Model used for every turn.
490    pub model: ResolvedAgentModel,
491}
492
493/// Cloneable provider-neutral subagent runtime.
494#[derive(Clone)]
495pub struct AgentRuntime {
496    intelligence: Intelligence,
497}
498
499impl AgentRuntime {
500    /// Constructs a runtime over the sole direct-model boundary.
501    pub fn new(intelligence: Intelligence) -> Self {
502        Self { intelligence }
503    }
504
505    /// Resolves a model without running an agent.
506    pub async fn resolve_model(&self, requested: &str) -> anyhow::Result<ResolvedAgentModel> {
507        self.intelligence
508            .resolve_agent_model(requested)
509            .await
510            .map_err(anyhow::Error::new)
511    }
512
513    /// Runs one primary logical turn across affine provider rounds until the host completes it.
514    pub async fn run_session<H: SessionHost>(
515        &self,
516        request: SessionRunRequest,
517        host: &mut H,
518    ) -> anyhow::Result<Option<String>> {
519        ensure!(
520            request.round_limit > 0,
521            "session round limit must be positive"
522        );
523        ensure!(
524            request.rounds_used <= request.round_limit,
525            "restored session round count exceeds its safety limit"
526        );
527        let user = self
528            .intelligence
529            .for_user(request.user_id)
530            .map_err(anyhow::Error::new)?;
531
532        'rounds: for round_index in request.rounds_used..request.round_limit {
533            let round = round_index + 1;
534            let prepared = match host.prepare_round(round).await? {
535                RoundPreparation::Run(prepared) => prepared,
536                RoundPreparation::Complete(answer) => return Ok(answer),
537            };
538            let manifest_hash = hex::encode(Sha256::digest(prepared.input.as_bytes()));
539            host.record(SessionEvent::InferenceSubmitted {
540                round,
541                manifest_hash: manifest_hash.clone(),
542                model: prepared.model.clone(),
543            })
544            .await?;
545
546            let mut provider_request = AgentRequest::new(prepared.provider_input, prepared.model);
547            provider_request.reasoning_effort = reasoning_effort(&prepared.reasoning_effort)?;
548            provider_request.tools = vec![ktool_definition(&prepared.tool_description)];
549            if let Some(timeout) = prepared.timeout {
550                provider_request.timeout = timeout;
551            }
552            let mut turn = match user
553                .start_agent_turn(
554                    request.operation_id,
555                    None,
556                    prepared.continuation,
557                    provider_request,
558                )
559                .await
560            {
561                Ok(turn) => turn,
562                Err(error) => {
563                    if let Some(receipt) = error.receipt().cloned() {
564                        host.record(SessionEvent::ProviderReceipt {
565                            round,
566                            usage: None,
567                            receipt: Box::new(receipt),
568                            continuation: None,
569                        })
570                        .await?;
571                    }
572                    return Err(anyhow::Error::new(error));
573                }
574            };
575            let mut used_tool = false;
576            let mut finish_requested = false;
577            let mut emitted_response = false;
578            let mut pending_capture: Option<Value> = None;
579            let completed = loop {
580                let event = match turn.next_event().await {
581                    Ok(Some(event)) => event,
582                    Ok(None) => {
583                        record_unavailable_session_turn(&mut turn, host, round).await?;
584                        anyhow::bail!("provider ended without a terminal turn event");
585                    }
586                    Err(error) => {
587                        if let Some(receipt) = error.receipt().cloned() {
588                            host.record(SessionEvent::ProviderReceipt {
589                                round,
590                                usage: None,
591                                receipt: Box::new(receipt),
592                                continuation: None,
593                            })
594                            .await?;
595                        }
596                        return Err(anyhow::Error::new(error));
597                    }
598                };
599                match event {
600                    AgentEvent::ProviderInput(_) => {}
601                    AgentEvent::ModelContextSubmitted(context) => {
602                        host.record(SessionEvent::ProviderInput { round, context })
603                            .await?;
604                    }
605                    AgentEvent::UsageUpdated(usage) => {
606                        host.record(SessionEvent::UsageUpdated { round, usage })
607                            .await?;
608                    }
609                    AgentEvent::ToolCall(native) => {
610                        used_tool = true;
611                        let call = parse_ktool_call(&native)
612                            .map_err(|error| anyhow::anyhow!("Invalid Ktool call: {error}"));
613                        let outcome = match host.execute_tool(call, request.operation_id).await {
614                            Ok(outcome) => outcome,
615                            Err(error) => {
616                                record_unavailable_session_turn(&mut turn, host, round).await?;
617                                return Err(error);
618                            }
619                        };
620                        let resume = match host.prepare_provider_resume(outcome).await {
621                            Ok(resume) => resume,
622                            Err(error) => {
623                                record_unavailable_session_turn(&mut turn, host, round).await?;
624                                return Err(error);
625                            }
626                        };
627                        let mut outcome = match provider_resume_action(resume) {
628                            ProviderResumeAction::Respond(outcome) => outcome,
629                            ProviderResumeAction::Interrupt(disposition) => {
630                                record_unavailable_session_turn(&mut turn, host, round).await?;
631                                match disposition {
632                                    PostReceiptDisposition::Return(answer) => return Ok(answer),
633                                    PostReceiptDisposition::NextRound => continue 'rounds,
634                                }
635                            }
636                        };
637                        finish_requested |= outcome.ok && outcome.finish_after_round;
638                        emitted_response |= outcome.ok && outcome.emitted_response;
639                        pending_capture = outcome.capture.take();
640                        let stop = outcome.stop;
641                        respond_session_or_record(
642                            &mut turn,
643                            host,
644                            round,
645                            &native.call_id,
646                            if outcome.ok {
647                                ToolResult::success(outcome.text)
648                            } else {
649                                ToolResult::failure(outcome.text)
650                            },
651                        )
652                        .await?;
653                        if stop {
654                            return Ok(None);
655                        }
656                    }
657                    AgentEvent::Completed(completed) => break completed,
658                }
659            };
660            let receipt = turn
661                .receipt()
662                .context("provider completed without a usage receipt")?
663                .clone();
664            let continuation = turn.continuation().cloned();
665            host.record(SessionEvent::ProviderReceipt {
666                round,
667                usage: completed.usage.clone(),
668                receipt: Box::new(receipt),
669                continuation,
670            })
671            .await?;
672
673            let control = if let Some(capture) = pending_capture {
674                host.complete_capture(capture, completed.answer).await?
675            } else {
676                host.complete_round(RoundCompletion {
677                    answer: completed.answer,
678                    used_tool,
679                    finish_requested,
680                    emitted_response,
681                })
682                .await?
683            };
684            match control {
685                SessionControl::Continue => {}
686                SessionControl::Complete(answer) => return Ok(answer),
687            }
688        }
689        Err(SessionRoundLimitError {
690            limit: request.round_limit,
691        }
692        .into())
693    }
694
695    /// Runs one fresh-context native subagent turn to a final non-empty answer.
696    pub async fn run<H: Host>(
697        &self,
698        request: RunRequest,
699        host: &mut H,
700    ) -> anyhow::Result<RunResult> {
701        let selected = self.resolve_model(&request.model).await?;
702        let reasoning_effort = reasoning_effort(&request.reasoning_effort)?;
703        let mut projection = Projection::new(request.context, request.task);
704        ensure_capacity(&projection, selected.max_input_tokens)?;
705        host.record(AuditEvent::Started {
706            parent_operation_id: request.parent_operation_id,
707            model: request.model.clone(),
708            provider_model: selected.provider_model.clone(),
709            provider: selected.provider,
710            context_window_tokens: selected.context_window_tokens,
711            max_input_tokens: selected.max_input_tokens,
712            context: projection.context.clone(),
713            task: projection.task.clone(),
714            host: request.start_metadata.clone(),
715        })?;
716        let user = self
717            .intelligence
718            .for_user(request.user_id)
719            .map_err(anyhow::Error::new)?;
720        for round in 0..1 {
721            let capturing = false;
722            ensure_capacity(&projection, selected.max_input_tokens)?;
723            let input = projection.render();
724            let manifest_hash = hex::encode(Sha256::digest(input.as_bytes()));
725            host.record(AuditEvent::InferenceSubmitted {
726                parent_operation_id: request.parent_operation_id,
727                round: round + 1,
728                manifest_hash: manifest_hash.clone(),
729                estimated_input_tokens: projection.estimated_tokens(),
730            })?;
731            let mut provider_request = AgentRequest::new(input, selected.requested_model.clone());
732            provider_request.reasoning_effort = reasoning_effort;
733            provider_request.ephemeral = true;
734            provider_request.tools = if capturing {
735                Vec::new()
736            } else {
737                vec![ktool_definition(
738                    "Call one available Ktool by its exact name.",
739                )]
740            };
741            if let Some(timeout) = request.timeout {
742                provider_request.timeout = timeout;
743            }
744            let child_operation_id = Uuid::new_v4();
745            let mut turn = match user
746                .start_agent_turn(
747                    child_operation_id,
748                    Some(request.parent_operation_id),
749                    None,
750                    provider_request,
751                )
752                .await
753            {
754                Ok(turn) => turn,
755                Err(error) => {
756                    if let Some(receipt) = error.receipt().cloned() {
757                        host.record(AuditEvent::ProviderReceipt {
758                            parent_operation_id: request.parent_operation_id,
759                            round: round + 1,
760                            manifest_hash,
761                            usage: None,
762                            receipt: Box::new(receipt),
763                        })?;
764                    }
765                    return Err(anyhow::Error::new(error));
766                }
767            };
768            let mut used_tool = false;
769            let mut pending_capture: Option<Value> = None;
770            let completed = loop {
771                let event = match turn.next_event().await {
772                    Ok(Some(event)) => event,
773                    Ok(None) => {
774                        let receipt = turn.finish_unavailable()?.clone();
775                        host.record(AuditEvent::ProviderReceipt {
776                            parent_operation_id: request.parent_operation_id,
777                            round: round + 1,
778                            manifest_hash: manifest_hash.clone(),
779                            usage: None,
780                            receipt: Box::new(receipt),
781                        })?;
782                        anyhow::bail!("subagent provider ended without a terminal turn event");
783                    }
784                    Err(error) => {
785                        if let Some(receipt) = error.receipt().cloned() {
786                            host.record(AuditEvent::ProviderReceipt {
787                                parent_operation_id: request.parent_operation_id,
788                                round: round + 1,
789                                manifest_hash: manifest_hash.clone(),
790                                usage: None,
791                                receipt: Box::new(receipt),
792                            })?;
793                        }
794                        return Err(anyhow::Error::new(error));
795                    }
796                };
797                match event {
798                    AgentEvent::ProviderInput(_) => {}
799                    AgentEvent::ModelContextSubmitted(_) => {}
800                    AgentEvent::UsageUpdated(_) => {}
801                    AgentEvent::ToolCall(native) => {
802                        used_tool = true;
803                        if capturing {
804                            respond_or_record(
805                                &mut turn,
806                                host,
807                                request.parent_operation_id,
808                                round + 1,
809                                &manifest_hash,
810                                &native.call_id,
811                                ToolResult::failure(
812                                    "No application tool is available while complete freeform output is being captured.",
813                                ),
814                            )
815                            .await?;
816                            continue;
817                        }
818                        if pending_capture.is_some() {
819                            respond_or_record(
820                                &mut turn,
821                                host,
822                                request.parent_operation_id,
823                                round + 1,
824                                &manifest_hash,
825                                &native.call_id,
826                                ToolResult::failure(
827                                    "A freeform output capture is pending; no other tool can run first.",
828                                ),
829                            )
830                            .await?;
831                            continue;
832                        }
833                        let call = match parse_ktool_call(&native) {
834                            Ok(call) => call,
835                            Err(error) => {
836                                let text = format!("Invalid application tool call: {error}");
837                                projection.push_history(format!("Ktool result:\n{text}"));
838                                respond_or_record(
839                                    &mut turn,
840                                    host,
841                                    request.parent_operation_id,
842                                    round + 1,
843                                    &manifest_hash,
844                                    &native.call_id,
845                                    ToolResult::failure(text),
846                                )
847                                .await?;
848                                continue;
849                            }
850                        };
851                        host.record(AuditEvent::ToolCall {
852                            parent_operation_id: request.parent_operation_id,
853                            name: call.name.clone(),
854                            arguments: call.arguments.clone(),
855                        })?;
856                        projection.push_history(format!(
857                            "Ktool call:\n{}",
858                            host.render_tool_call(&call)?
859                        ));
860                        let budget = ContextBudget {
861                            projection: projection.clone(),
862                            max_input_tokens: selected.max_input_tokens,
863                        };
864                        let mut outcome = host
865                            .execute_tool(call.clone(), child_operation_id, budget)
866                            .await
867                            .unwrap_or_else(|error| {
868                                ToolOutcome::failure(format!("{} failed: {error}", call.name))
869                            });
870                        let exact_result = outcome.text.clone();
871                        let initially_ok = outcome.ok;
872                        let mut provider_result = outcome.text.clone();
873                        let mut candidate = projection.clone();
874                        candidate.apply_updates(successful_state_updates(&outcome));
875                        candidate.push_history_displaying(
876                            format!("Ktool result:\n{provider_result}"),
877                            if initially_ok {
878                                outcome.displayed_state_keys.clone()
879                            } else {
880                                Default::default()
881                            },
882                        );
883                        let accepted = candidate.estimated_tokens() <= selected.max_input_tokens;
884                        if accepted {
885                            projection = candidate;
886                        } else {
887                            outcome.ok = false;
888                            outcome.capture = None;
889                            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();
890                            projection.push_history(format!("Ktool result:\n{provider_result}"));
891                        }
892                        host.record(AuditEvent::ToolResult {
893                            parent_operation_id: request.parent_operation_id,
894                            name: call.name.clone(),
895                            ok: initially_ok,
896                            projection_accepted: accepted,
897                            result: exact_result,
898                        })?;
899                        pending_capture = outcome.capture.take();
900                        respond_or_record(
901                            &mut turn,
902                            host,
903                            request.parent_operation_id,
904                            round + 1,
905                            &manifest_hash,
906                            &native.call_id,
907                            if outcome.ok {
908                                ToolResult::success(provider_result)
909                            } else {
910                                ToolResult::failure(provider_result)
911                            },
912                        )
913                        .await?;
914                    }
915                    AgentEvent::Completed(completed) => break completed,
916                }
917            };
918            let receipt = turn
919                .receipt()
920                .context("subagent provider completed without a usage receipt")?
921                .clone();
922            host.record(AuditEvent::ProviderReceipt {
923                parent_operation_id: request.parent_operation_id,
924                round: round + 1,
925                manifest_hash: manifest_hash.clone(),
926                usage: completed.usage.clone(),
927                receipt: Box::new(receipt),
928            })?;
929
930            if let Some(capture) = pending_capture {
931                ensure!(
932                    !completed.answer.trim().is_empty(),
933                    "subagent provider completed an output capture without contents"
934                );
935                let budget = ContextBudget {
936                    projection: projection.clone(),
937                    max_input_tokens: selected.max_input_tokens,
938                };
939                let outcome = host
940                    .complete_capture(capture, completed.answer, budget)
941                    .await?;
942                let mut candidate = projection.clone();
943                candidate.apply_updates(successful_state_updates(&outcome));
944                candidate.push_history(format!("Ktool result:\n{}", outcome.text));
945                ensure_capacity(&candidate, selected.max_input_tokens)?;
946                let answer = outcome.text.trim().to_owned();
947                ensure!(
948                    !answer.is_empty(),
949                    "subagent output capture completed without a result"
950                );
951                host.record(AuditEvent::Completed {
952                    parent_operation_id: request.parent_operation_id,
953                    model: request.model.clone(),
954                    response: answer.clone(),
955                })?;
956                return Ok(RunResult {
957                    answer,
958                    model: selected,
959                });
960            }
961            let answer = completed.answer.trim().to_owned();
962            if !answer.is_empty() {
963                host.record(AuditEvent::Completed {
964                    parent_operation_id: request.parent_operation_id,
965                    model: request.model.clone(),
966                    response: answer.clone(),
967                })?;
968                return Ok(RunResult {
969                    answer,
970                    model: selected,
971                });
972            }
973            ensure!(
974                !used_tool,
975                "subagent provider completed without a response or tool call"
976            );
977        }
978        anyhow::bail!("subagent provider completed without a terminal response")
979    }
980}
981
982async fn respond_or_record<H: Host>(
983    turn: &mut kcode_intelligence_router::AgentTurn,
984    host: &mut H,
985    parent_operation_id: Uuid,
986    round: u64,
987    manifest_hash: &str,
988    call_id: &str,
989    result: ToolResult,
990) -> anyhow::Result<()> {
991    if let Err(error) = turn.respond(call_id, result).await {
992        let receipt = turn.finish_unavailable()?.clone();
993        host.record(AuditEvent::ProviderReceipt {
994            parent_operation_id,
995            round,
996            manifest_hash: manifest_hash.into(),
997            usage: None,
998            receipt: Box::new(receipt),
999        })?;
1000        return Err(anyhow::Error::new(error));
1001    }
1002    Ok(())
1003}
1004
1005async fn respond_session_or_record<H: SessionHost>(
1006    turn: &mut kcode_intelligence_router::AgentTurn,
1007    host: &mut H,
1008    round: u64,
1009    call_id: &str,
1010    result: ToolResult,
1011) -> anyhow::Result<()> {
1012    if let Err(error) = turn.respond(call_id, result).await {
1013        record_unavailable_session_turn(turn, host, round).await?;
1014        return Err(anyhow::Error::new(error));
1015    }
1016    Ok(())
1017}
1018
1019async fn record_unavailable_session_turn<H: SessionHost>(
1020    turn: &mut kcode_intelligence_router::AgentTurn,
1021    host: &mut H,
1022    round: u64,
1023) -> anyhow::Result<()> {
1024    let receipt = turn.finish_unavailable()?.clone();
1025    host.record(SessionEvent::ProviderReceipt {
1026        round,
1027        usage: None,
1028        receipt: Box::new(receipt),
1029        continuation: None,
1030    })
1031    .await
1032}
1033
1034#[derive(Clone)]
1035struct Projection {
1036    context: Vec<String>,
1037    task: String,
1038    history: Vec<ProjectedHistory>,
1039    states: Vec<ProjectedState>,
1040}
1041
1042#[derive(Clone)]
1043struct ProjectedHistory {
1044    text: String,
1045    displayed_state_keys: Vec<String>,
1046}
1047
1048#[derive(Clone)]
1049struct ProjectedState {
1050    key: String,
1051    text: String,
1052}
1053
1054impl Projection {
1055    fn new(context: Vec<String>, task: String) -> Self {
1056        Self {
1057            context,
1058            task,
1059            history: Vec::new(),
1060            states: Vec::new(),
1061        }
1062    }
1063
1064    fn render(&self) -> String {
1065        self.context
1066            .iter()
1067            .map(String::as_str)
1068            .chain(std::iter::once(self.task.as_str()))
1069            .chain(self.history.iter().map(|entry| entry.text.as_str()))
1070            .chain(
1071                self.states
1072                    .iter()
1073                    .filter(|state| {
1074                        !self.history.iter().any(|entry| {
1075                            entry
1076                                .displayed_state_keys
1077                                .iter()
1078                                .any(|key| key == &state.key)
1079                        })
1080                    })
1081                    .map(|state| state.text.as_str()),
1082            )
1083            .filter(|section| !section.is_empty())
1084            .collect::<Vec<_>>()
1085            .join("\n\n")
1086    }
1087
1088    fn push_history(&mut self, text: impl Into<String>) {
1089        self.push_history_displaying(text, Vec::new());
1090    }
1091
1092    fn push_history_displaying(
1093        &mut self,
1094        text: impl Into<String>,
1095        displayed_state_keys: Vec<String>,
1096    ) {
1097        self.history.push(ProjectedHistory {
1098            text: text.into(),
1099            displayed_state_keys,
1100        });
1101    }
1102
1103    fn update_state(&mut self, key: String, text: Option<String>) {
1104        self.states.retain(|state| state.key != key);
1105        if let Some(text) = text {
1106            self.states.push(ProjectedState { key, text });
1107        }
1108    }
1109
1110    fn apply_updates(&mut self, updates: &[StateUpdate]) {
1111        for update in updates {
1112            let marker = if update.text.is_some() {
1113                SUPERSEDED_TOOL_OUTPUT
1114            } else {
1115                REMOVED_TOOL_OUTPUT
1116            };
1117            for entry in &mut self.history {
1118                if entry
1119                    .displayed_state_keys
1120                    .iter()
1121                    .any(|key| key == &update.key)
1122                {
1123                    entry.text = marker.into();
1124                    entry.displayed_state_keys.clear();
1125                }
1126            }
1127            self.update_state(update.key.clone(), update.text.clone());
1128        }
1129    }
1130
1131    fn estimated_tokens(&self) -> u64 {
1132        (self.render().chars().count() as u64)
1133            .div_ceil(4)
1134            .saturating_add(PROTOCOL_TOKEN_RESERVE)
1135    }
1136}
1137
1138fn successful_state_updates(outcome: &ToolOutcome) -> &[StateUpdate] {
1139    if outcome.ok {
1140        &outcome.state_updates
1141    } else {
1142        &[]
1143    }
1144}
1145
1146fn ensure_capacity(projection: &Projection, max_input_tokens: u64) -> anyhow::Result<()> {
1147    let estimated = projection.estimated_tokens();
1148    ensure!(
1149        estimated <= max_input_tokens,
1150        "subagent context requires approximately {estimated} input tokens, over the selected model's {max_input_tokens}-token input limit"
1151    );
1152    Ok(())
1153}
1154
1155fn ktool_definition(description: &str) -> DynamicTool {
1156    DynamicTool::new(
1157        "call_ktool",
1158        description,
1159        json!({
1160            "type": "object",
1161            "additionalProperties": false,
1162            "required": ["name", "arguments"],
1163            "properties": {
1164                "name": {"type": "string"},
1165                "arguments": {"type": "object"}
1166            }
1167        }),
1168    )
1169}
1170
1171fn parse_ktool_call(call: &DynamicToolCall) -> anyhow::Result<ToolCall> {
1172    ensure!(call.tool == "call_ktool", "unknown provider tool");
1173    let arguments = call
1174        .arguments
1175        .as_object()
1176        .context("call_ktool arguments must be an object")?;
1177    ensure!(
1178        arguments
1179            .keys()
1180            .all(|key| matches!(key.as_str(), "name" | "arguments")),
1181        "call_ktool contains unknown arguments"
1182    );
1183    let name = arguments
1184        .get("name")
1185        .and_then(Value::as_str)
1186        .map(str::trim)
1187        .filter(|name| !name.is_empty() && name.chars().count() <= 100)
1188        .context("call_ktool.name must be a non-empty bounded string")?
1189        .to_owned();
1190    let arguments = arguments
1191        .get("arguments")
1192        .filter(|value| value.is_object())
1193        .context("call_ktool.arguments must be an object")?
1194        .clone();
1195    Ok(ToolCall { name, arguments })
1196}
1197
1198fn reasoning_effort(value: &str) -> anyhow::Result<ReasoningEffort> {
1199    Ok(match value {
1200        "none" => ReasoningEffort::None,
1201        "minimal" => ReasoningEffort::Minimal,
1202        "low" => ReasoningEffort::Low,
1203        "medium" => ReasoningEffort::Medium,
1204        "high" => ReasoningEffort::High,
1205        "xhigh" => ReasoningEffort::XHigh,
1206        "max" => ReasoningEffort::Max,
1207        _ => anyhow::bail!("unsupported reasoning effort {value:?}"),
1208    })
1209}
1210
1211#[cfg(test)]
1212mod tests {
1213    use super::*;
1214
1215    #[test]
1216    fn projection_replaces_state_and_budget_accounts_for_reserve() {
1217        let mut projection = Projection::new(vec!["context".into()], "task".into());
1218        projection.update_state("file".into(), Some("old".into()));
1219        projection.update_state("file".into(), Some("new".into()));
1220        assert_eq!(projection.states.len(), 1);
1221        assert!(projection.render().contains("new"));
1222        assert!(!projection.render().contains("old"));
1223        assert!(projection.estimated_tokens() >= PROTOCOL_TOKEN_RESERVE);
1224    }
1225
1226    #[test]
1227    fn state_update_supersedes_earlier_display_and_moves_current_value_after_history() {
1228        let mut projection = Projection::new(vec!["context".into()], "task".into());
1229        projection.push_history("Ktool call: open");
1230        projection.push_history_displaying(
1231            "Ktool result:\ncomplete old source",
1232            vec!["tool-state:7".into()],
1233        );
1234        projection.push_history("Ktool call: write");
1235        projection.apply_updates(&[StateUpdate {
1236            key: "tool-state:7".into(),
1237            text: Some("complete new source".into()),
1238        }]);
1239        projection.push_history("Ktool result: write completed");
1240
1241        let rendered = projection.render();
1242        assert!(!rendered.contains("complete old source"));
1243        assert_eq!(rendered.matches(SUPERSEDED_TOOL_OUTPUT).count(), 1);
1244        assert_eq!(rendered.matches("complete new source").count(), 1);
1245        assert!(
1246            rendered.find(SUPERSEDED_TOOL_OUTPUT).unwrap()
1247                < rendered.find("Ktool call: write").unwrap()
1248        );
1249        assert!(
1250            rendered.find("Ktool result: write completed").unwrap()
1251                < rendered.find("complete new source").unwrap()
1252        );
1253    }
1254
1255    #[test]
1256    fn current_state_is_not_duplicated_when_the_latest_result_displays_it() {
1257        let mut projection = Projection::new(vec!["context".into()], "task".into());
1258        projection.apply_updates(&[StateUpdate {
1259            key: "tool-state:7".into(),
1260            text: Some("complete source".into()),
1261        }]);
1262        projection.push_history_displaying(
1263            "Ktool result:\ncomplete source",
1264            vec!["tool-state:7".into()],
1265        );
1266
1267        assert_eq!(projection.render().matches("complete source").count(), 1);
1268    }
1269
1270    #[test]
1271    fn state_update_only_supersedes_displays_with_the_same_identity() {
1272        let mut projection = Projection::new(vec!["context".into()], "task".into());
1273        projection.push_history_displaying("first output", vec!["first".into()]);
1274        projection.push_history_displaying("second output", vec!["second".into()]);
1275        projection.apply_updates(&[StateUpdate {
1276            key: "first".into(),
1277            text: Some("current first output".into()),
1278        }]);
1279
1280        let rendered = projection.render();
1281        assert!(
1282            !rendered
1283                .split("\n\n")
1284                .any(|section| section == "first output")
1285        );
1286        assert!(rendered.contains(SUPERSEDED_TOOL_OUTPUT));
1287        assert!(rendered.contains("second output"));
1288        assert!(rendered.contains("current first output"));
1289    }
1290
1291    #[test]
1292    fn removed_state_uses_a_truthful_supersession_marker() {
1293        let mut projection = Projection::new(vec!["context".into()], "task".into());
1294        projection.push_history_displaying("retired output", vec!["state".into()]);
1295        projection.apply_updates(&[StateUpdate {
1296            key: "state".into(),
1297            text: None,
1298        }]);
1299
1300        let rendered = projection.render();
1301        assert!(!rendered.contains("retired output"));
1302        assert!(rendered.contains(REMOVED_TOOL_OUTPUT));
1303    }
1304
1305    #[test]
1306    fn failed_outcomes_cannot_update_or_supersede_state() {
1307        let mut projection = Projection::new(vec!["context".into()], "task".into());
1308        projection.push_history_displaying("current output", vec!["state".into()]);
1309        let failed = ToolOutcome {
1310            text: "write failed".into(),
1311            ok: false,
1312            state_updates: vec![StateUpdate {
1313                key: "state".into(),
1314                text: Some("invalid update".into()),
1315            }],
1316            displayed_state_keys: Vec::new(),
1317            capture: None,
1318        };
1319
1320        projection.apply_updates(successful_state_updates(&failed));
1321
1322        let rendered = projection.render();
1323        assert!(rendered.contains("current output"));
1324        assert!(!rendered.contains("invalid update"));
1325        assert!(!rendered.contains(SUPERSEDED_TOOL_OUTPUT));
1326    }
1327
1328    #[test]
1329    fn large_stateful_result_is_rendered_exactly_once() {
1330        let loaded_node = format!(
1331            "Node body:\n{}\n\nFixed connections:\nfixed-node\n\nRecent connections:\nrecent-node",
1332            "x".repeat(2_000)
1333        );
1334        let outcome = ToolOutcome {
1335            text: loaded_node.clone(),
1336            ok: true,
1337            state_updates: vec![StateUpdate {
1338                key: "loaded-node".into(),
1339                text: Some(loaded_node.clone()),
1340            }],
1341            displayed_state_keys: vec!["loaded-node".into()],
1342            capture: None,
1343        };
1344        let mut projection =
1345            Projection::new(vec!["initial node description".into()], "task".into());
1346        let provider_result = outcome.text.clone();
1347        projection.apply_updates(successful_state_updates(&outcome));
1348        projection.push_history_displaying(
1349            format!("Ktool result:\n{provider_result}"),
1350            outcome.displayed_state_keys,
1351        );
1352
1353        let rendered = projection.render();
1354        assert_eq!(rendered.matches(loaded_node.as_str()).count(), 1);
1355        assert!(rendered.contains("Fixed connections:\nfixed-node"));
1356        assert!(rendered.contains("Recent connections:\nrecent-node"));
1357    }
1358
1359    #[test]
1360    fn native_tool_wrapper_is_strict() {
1361        let call = parse_ktool_call(&DynamicToolCall {
1362            call_id: "1".into(),
1363            tool: "call_ktool".into(),
1364            arguments: json!({"name": "Read", "arguments": {"id": 1}}),
1365        })
1366        .unwrap();
1367        assert_eq!(call.name, "Read");
1368        assert_eq!(call.arguments["id"], 1);
1369    }
1370
1371    #[test]
1372    fn restart_fresh_selects_interrupt_then_next_round_disposition() {
1373        let ProviderResumeAction::Interrupt(disposition) =
1374            provider_resume_action(ProviderResume::RestartFresh)
1375        else {
1376            panic!("restart-fresh must not select native response");
1377        };
1378        assert!(matches!(disposition, PostReceiptDisposition::NextRound));
1379    }
1380
1381    #[test]
1382    fn existing_provider_resume_actions_keep_their_prior_paths() {
1383        let outcome = SessionToolOutcome::success("durable result");
1384        assert!(matches!(
1385            provider_resume_action(ProviderResume::Continue(outcome.clone())),
1386            ProviderResumeAction::Respond(returned) if returned == outcome
1387        ));
1388        let ProviderResumeAction::Interrupt(disposition) =
1389            provider_resume_action(ProviderResume::Complete(Some("done".into())))
1390        else {
1391            panic!("complete must not select native response");
1392        };
1393        assert!(matches!(
1394            disposition,
1395            PostReceiptDisposition::Return(Some(answer)) if answer == "done"
1396        ));
1397    }
1398}
kcode_agent_runtime/lib.rs

kcode_agent_runtime/
lib.rs
