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}
356
357/// Provider completion summarized for the application-owned session policy.
358#[derive(Clone, Debug, Eq, PartialEq)]
359pub struct RoundCompletion {
360    /// Terminal assistant text, which may be empty after tool use.
361    pub answer: String,
362    /// Whether the provider called at least one application tool.
363    pub used_tool: bool,
364    /// Whether a successful tool requested session completion after this round.
365    pub finish_requested: bool,
366    /// Whether a tool already emitted the externally visible response.
367    pub emitted_response: bool,
368}
369
370/// Application decision after completing a semantic primary-session transition.
371#[derive(Clone, Debug, Eq, PartialEq)]
372pub enum SessionControl {
373    /// Prepare another affine provider round.
374    Continue,
375    /// Finish the logical turn with the optional externally visible answer.
376    Complete(Option<String>),
377}
378
379/// Application-owned semantics invoked by the primary agent runtime.
380pub trait SessionHost: Send {
381    /// Prepares the complete context and runtime selection for one provider round.
382    fn prepare_round<'a>(&'a mut self, round: u64) -> HostFuture<'a, RoundPreparation>;
383
384    /// Records one provider-protocol fact and durably checkpoints its effects.
385    fn record<'a>(&'a mut self, event: SessionEvent) -> HostFuture<'a, ()>;
386
387    /// Executes and durably projects one parsed or invalid application tool call.
388    fn execute_tool<'a>(
389        &'a mut self,
390        call: anyhow::Result<ToolCall>,
391        provider_operation_id: Uuid,
392    ) -> HostFuture<'a, SessionToolOutcome>;
393
394    /// Reconciles application policy at the actual provider-resume boundary.
395    fn prepare_provider_resume<'a>(
396        &'a mut self,
397        outcome: SessionToolOutcome,
398    ) -> HostFuture<'a, ProviderResume> {
399        Box::pin(async move { Ok(ProviderResume::Continue(outcome)) })
400    }
401
402    /// Completes an opaque freeform capture requested by a prior tool result.
403    fn complete_capture<'a>(
404        &'a mut self,
405        capture: Value,
406        contents: String,
407    ) -> HostFuture<'a, SessionControl>;
408
409    /// Applies the terminal provider answer and decides whether another round is needed.
410    fn complete_round<'a>(
411        &'a mut self,
412        completion: RoundCompletion,
413    ) -> HostFuture<'a, SessionControl>;
414}
415
416/// Error returned when a restored primary turn exhausts its cumulative round budget.
417#[derive(Debug)]
418pub struct SessionRoundLimitError {
419    limit: u64,
420}
421
422impl std::fmt::Display for SessionRoundLimitError {
423    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
424        write!(
425            formatter,
426            "agent exceeded the {}-round tool-loop safety limit",
427            self.limit
428        )
429    }
430}
431
432impl std::error::Error for SessionRoundLimitError {}
433
434/// Returns whether an error is the primary runtime's round-limit signal.
435pub fn is_session_round_limit(error: &anyhow::Error) -> bool {
436    error.downcast_ref::<SessionRoundLimitError>().is_some()
437}
438
439/// Inputs for one complete subagent run.
440#[derive(Clone, Debug, PartialEq)]
441pub struct RunRequest {
442    /// Stable user identifier used for router accounting.
443    pub user_id: String,
444    /// Running parent operation whose cancellation propagates to each turn.
445    pub parent_operation_id: Uuid,
446    /// Exact requested model selector.
447    pub model: String,
448    /// Provider-neutral reasoning effort.
449    pub reasoning_effort: String,
450    /// Ordered immutable context sections.
451    pub context: Vec<String>,
452    /// Exact task presented after the context.
453    pub task: String,
454    /// Optional per-turn timeout.
455    pub timeout: Option<Duration>,
456    /// Additional application metadata included in the start audit event.
457    pub start_metadata: Value,
458}
459
460/// Completed subagent output.
461#[derive(Clone, Debug, Eq, PartialEq)]
462pub struct RunResult {
463    /// Final non-empty assistant answer.
464    pub answer: String,
465    /// Model used for every turn.
466    pub model: ResolvedAgentModel,
467}
468
469/// Cloneable provider-neutral subagent runtime.
470#[derive(Clone)]
471pub struct AgentRuntime {
472    intelligence: Intelligence,
473}
474
475impl AgentRuntime {
476    /// Constructs a runtime over the sole direct-model boundary.
477    pub fn new(intelligence: Intelligence) -> Self {
478        Self { intelligence }
479    }
480
481    /// Resolves a model without running an agent.
482    pub async fn resolve_model(&self, requested: &str) -> anyhow::Result<ResolvedAgentModel> {
483        self.intelligence
484            .resolve_agent_model(requested)
485            .await
486            .map_err(anyhow::Error::new)
487    }
488
489    /// Runs one primary logical turn across affine provider rounds until the host completes it.
490    pub async fn run_session<H: SessionHost>(
491        &self,
492        request: SessionRunRequest,
493        host: &mut H,
494    ) -> anyhow::Result<Option<String>> {
495        ensure!(
496            request.round_limit > 0,
497            "session round limit must be positive"
498        );
499        ensure!(
500            request.rounds_used <= request.round_limit,
501            "restored session round count exceeds its safety limit"
502        );
503        let user = self
504            .intelligence
505            .for_user(request.user_id)
506            .map_err(anyhow::Error::new)?;
507
508        for round_index in request.rounds_used..request.round_limit {
509            let round = round_index + 1;
510            let prepared = match host.prepare_round(round).await? {
511                RoundPreparation::Run(prepared) => prepared,
512                RoundPreparation::Complete(answer) => return Ok(answer),
513            };
514            let manifest_hash = hex::encode(Sha256::digest(prepared.input.as_bytes()));
515            host.record(SessionEvent::InferenceSubmitted {
516                round,
517                manifest_hash: manifest_hash.clone(),
518                model: prepared.model.clone(),
519            })
520            .await?;
521
522            let mut provider_request = AgentRequest::new(prepared.provider_input, prepared.model);
523            provider_request.reasoning_effort = reasoning_effort(&prepared.reasoning_effort)?;
524            provider_request.tools = vec![ktool_definition(&prepared.tool_description)];
525            if let Some(timeout) = prepared.timeout {
526                provider_request.timeout = timeout;
527            }
528            let mut turn = match user
529                .start_agent_turn(
530                    request.operation_id,
531                    None,
532                    prepared.continuation,
533                    provider_request,
534                )
535                .await
536            {
537                Ok(turn) => turn,
538                Err(error) => {
539                    if let Some(receipt) = error.receipt().cloned() {
540                        host.record(SessionEvent::ProviderReceipt {
541                            round,
542                            usage: None,
543                            receipt: Box::new(receipt),
544                            continuation: None,
545                        })
546                        .await?;
547                    }
548                    return Err(anyhow::Error::new(error));
549                }
550            };
551            let mut used_tool = false;
552            let mut finish_requested = false;
553            let mut emitted_response = false;
554            let mut pending_capture: Option<Value> = None;
555            let completed = loop {
556                let event = match turn.next_event().await {
557                    Ok(Some(event)) => event,
558                    Ok(None) => {
559                        record_unavailable_session_turn(&mut turn, host, round).await?;
560                        anyhow::bail!("provider ended without a terminal turn event");
561                    }
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                match event {
576                    AgentEvent::ProviderInput(_) => {}
577                    AgentEvent::ModelContextSubmitted(context) => {
578                        host.record(SessionEvent::ProviderInput { round, context })
579                            .await?;
580                    }
581                    AgentEvent::UsageUpdated(usage) => {
582                        host.record(SessionEvent::UsageUpdated { round, usage })
583                            .await?;
584                    }
585                    AgentEvent::ToolCall(native) => {
586                        used_tool = true;
587                        let call = parse_ktool_call(&native)
588                            .map_err(|error| anyhow::anyhow!("Invalid Ktool call: {error}"));
589                        let outcome = match host.execute_tool(call, request.operation_id).await {
590                            Ok(outcome) => outcome,
591                            Err(error) => {
592                                record_unavailable_session_turn(&mut turn, host, round).await?;
593                                return Err(error);
594                            }
595                        };
596                        let mut outcome = match host.prepare_provider_resume(outcome).await {
597                            Ok(ProviderResume::Continue(outcome)) => outcome,
598                            Ok(ProviderResume::Complete(answer)) => {
599                                record_unavailable_session_turn(&mut turn, host, round).await?;
600                                return Ok(answer);
601                            }
602                            Err(error) => {
603                                record_unavailable_session_turn(&mut turn, host, round).await?;
604                                return Err(error);
605                            }
606                        };
607                        finish_requested |= outcome.ok && outcome.finish_after_round;
608                        emitted_response |= outcome.ok && outcome.emitted_response;
609                        pending_capture = outcome.capture.take();
610                        let stop = outcome.stop;
611                        respond_session_or_record(
612                            &mut turn,
613                            host,
614                            round,
615                            &native.call_id,
616                            if outcome.ok {
617                                ToolResult::success(outcome.text)
618                            } else {
619                                ToolResult::failure(outcome.text)
620                            },
621                        )
622                        .await?;
623                        if stop {
624                            return Ok(None);
625                        }
626                    }
627                    AgentEvent::Completed(completed) => break completed,
628                }
629            };
630            let receipt = turn
631                .receipt()
632                .context("provider completed without a usage receipt")?
633                .clone();
634            let continuation = turn.continuation().cloned();
635            host.record(SessionEvent::ProviderReceipt {
636                round,
637                usage: completed.usage.clone(),
638                receipt: Box::new(receipt),
639                continuation,
640            })
641            .await?;
642
643            let control = if let Some(capture) = pending_capture {
644                host.complete_capture(capture, completed.answer).await?
645            } else {
646                host.complete_round(RoundCompletion {
647                    answer: completed.answer,
648                    used_tool,
649                    finish_requested,
650                    emitted_response,
651                })
652                .await?
653            };
654            match control {
655                SessionControl::Continue => {}
656                SessionControl::Complete(answer) => return Ok(answer),
657            }
658        }
659        Err(SessionRoundLimitError {
660            limit: request.round_limit,
661        }
662        .into())
663    }
664
665    /// Runs one fresh-context native subagent turn to a final non-empty answer.
666    pub async fn run<H: Host>(
667        &self,
668        request: RunRequest,
669        host: &mut H,
670    ) -> anyhow::Result<RunResult> {
671        let selected = self.resolve_model(&request.model).await?;
672        let reasoning_effort = reasoning_effort(&request.reasoning_effort)?;
673        let mut projection = Projection::new(request.context, request.task);
674        ensure_capacity(&projection, selected.max_input_tokens)?;
675        host.record(AuditEvent::Started {
676            parent_operation_id: request.parent_operation_id,
677            model: request.model.clone(),
678            provider_model: selected.provider_model.clone(),
679            provider: selected.provider,
680            context_window_tokens: selected.context_window_tokens,
681            max_input_tokens: selected.max_input_tokens,
682            context: projection.context.clone(),
683            task: projection.task.clone(),
684            host: request.start_metadata.clone(),
685        })?;
686        let user = self
687            .intelligence
688            .for_user(request.user_id)
689            .map_err(anyhow::Error::new)?;
690        for round in 0..1 {
691            let capturing = false;
692            ensure_capacity(&projection, selected.max_input_tokens)?;
693            let input = projection.render();
694            let manifest_hash = hex::encode(Sha256::digest(input.as_bytes()));
695            host.record(AuditEvent::InferenceSubmitted {
696                parent_operation_id: request.parent_operation_id,
697                round: round + 1,
698                manifest_hash: manifest_hash.clone(),
699                estimated_input_tokens: projection.estimated_tokens(),
700            })?;
701            let mut provider_request = AgentRequest::new(input, selected.requested_model.clone());
702            provider_request.reasoning_effort = reasoning_effort;
703            provider_request.ephemeral = true;
704            provider_request.tools = if capturing {
705                Vec::new()
706            } else {
707                vec![ktool_definition(
708                    "Call one available Ktool by its exact name.",
709                )]
710            };
711            if let Some(timeout) = request.timeout {
712                provider_request.timeout = timeout;
713            }
714            let child_operation_id = Uuid::new_v4();
715            let mut turn = match user
716                .start_agent_turn(
717                    child_operation_id,
718                    Some(request.parent_operation_id),
719                    None,
720                    provider_request,
721                )
722                .await
723            {
724                Ok(turn) => turn,
725                Err(error) => {
726                    if let Some(receipt) = error.receipt().cloned() {
727                        host.record(AuditEvent::ProviderReceipt {
728                            parent_operation_id: request.parent_operation_id,
729                            round: round + 1,
730                            manifest_hash,
731                            usage: None,
732                            receipt: Box::new(receipt),
733                        })?;
734                    }
735                    return Err(anyhow::Error::new(error));
736                }
737            };
738            let mut used_tool = false;
739            let mut pending_capture: Option<Value> = None;
740            let completed = loop {
741                let event = match turn.next_event().await {
742                    Ok(Some(event)) => event,
743                    Ok(None) => {
744                        let receipt = turn.finish_unavailable()?.clone();
745                        host.record(AuditEvent::ProviderReceipt {
746                            parent_operation_id: request.parent_operation_id,
747                            round: round + 1,
748                            manifest_hash: manifest_hash.clone(),
749                            usage: None,
750                            receipt: Box::new(receipt),
751                        })?;
752                        anyhow::bail!("subagent provider ended without a terminal turn event");
753                    }
754                    Err(error) => {
755                        if let Some(receipt) = error.receipt().cloned() {
756                            host.record(AuditEvent::ProviderReceipt {
757                                parent_operation_id: request.parent_operation_id,
758                                round: round + 1,
759                                manifest_hash: manifest_hash.clone(),
760                                usage: None,
761                                receipt: Box::new(receipt),
762                            })?;
763                        }
764                        return Err(anyhow::Error::new(error));
765                    }
766                };
767                match event {
768                    AgentEvent::ProviderInput(_) => {}
769                    AgentEvent::ModelContextSubmitted(_) => {}
770                    AgentEvent::UsageUpdated(_) => {}
771                    AgentEvent::ToolCall(native) => {
772                        used_tool = true;
773                        if capturing {
774                            respond_or_record(
775                                &mut turn,
776                                host,
777                                request.parent_operation_id,
778                                round + 1,
779                                &manifest_hash,
780                                &native.call_id,
781                                ToolResult::failure(
782                                    "No application tool is available while complete freeform output is being captured.",
783                                ),
784                            )
785                            .await?;
786                            continue;
787                        }
788                        if pending_capture.is_some() {
789                            respond_or_record(
790                                &mut turn,
791                                host,
792                                request.parent_operation_id,
793                                round + 1,
794                                &manifest_hash,
795                                &native.call_id,
796                                ToolResult::failure(
797                                    "A freeform output capture is pending; no other tool can run first.",
798                                ),
799                            )
800                            .await?;
801                            continue;
802                        }
803                        let call = match parse_ktool_call(&native) {
804                            Ok(call) => call,
805                            Err(error) => {
806                                let text = format!("Invalid application tool call: {error}");
807                                projection.push_history(format!("Ktool result:\n{text}"));
808                                respond_or_record(
809                                    &mut turn,
810                                    host,
811                                    request.parent_operation_id,
812                                    round + 1,
813                                    &manifest_hash,
814                                    &native.call_id,
815                                    ToolResult::failure(text),
816                                )
817                                .await?;
818                                continue;
819                            }
820                        };
821                        host.record(AuditEvent::ToolCall {
822                            parent_operation_id: request.parent_operation_id,
823                            name: call.name.clone(),
824                            arguments: call.arguments.clone(),
825                        })?;
826                        projection.push_history(format!(
827                            "Ktool call:\n{}",
828                            host.render_tool_call(&call)?
829                        ));
830                        let budget = ContextBudget {
831                            projection: projection.clone(),
832                            max_input_tokens: selected.max_input_tokens,
833                        };
834                        let mut outcome = host
835                            .execute_tool(call.clone(), child_operation_id, budget)
836                            .await
837                            .unwrap_or_else(|error| {
838                                ToolOutcome::failure(format!("{} failed: {error}", call.name))
839                            });
840                        let exact_result = outcome.text.clone();
841                        let initially_ok = outcome.ok;
842                        let mut provider_result = outcome.text.clone();
843                        let mut candidate = projection.clone();
844                        candidate.apply_updates(successful_state_updates(&outcome));
845                        candidate.push_history_displaying(
846                            format!("Ktool result:\n{provider_result}"),
847                            if initially_ok {
848                                outcome.displayed_state_keys.clone()
849                            } else {
850                                Default::default()
851                            },
852                        );
853                        let accepted = candidate.estimated_tokens() <= selected.max_input_tokens;
854                        if accepted {
855                            projection = candidate;
856                        } else {
857                            outcome.ok = false;
858                            outcome.capture = None;
859                            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();
860                            projection.push_history(format!("Ktool result:\n{provider_result}"));
861                        }
862                        host.record(AuditEvent::ToolResult {
863                            parent_operation_id: request.parent_operation_id,
864                            name: call.name.clone(),
865                            ok: initially_ok,
866                            projection_accepted: accepted,
867                            result: exact_result,
868                        })?;
869                        pending_capture = outcome.capture.take();
870                        respond_or_record(
871                            &mut turn,
872                            host,
873                            request.parent_operation_id,
874                            round + 1,
875                            &manifest_hash,
876                            &native.call_id,
877                            if outcome.ok {
878                                ToolResult::success(provider_result)
879                            } else {
880                                ToolResult::failure(provider_result)
881                            },
882                        )
883                        .await?;
884                    }
885                    AgentEvent::Completed(completed) => break completed,
886                }
887            };
888            let receipt = turn
889                .receipt()
890                .context("subagent provider completed without a usage receipt")?
891                .clone();
892            host.record(AuditEvent::ProviderReceipt {
893                parent_operation_id: request.parent_operation_id,
894                round: round + 1,
895                manifest_hash: manifest_hash.clone(),
896                usage: completed.usage.clone(),
897                receipt: Box::new(receipt),
898            })?;
899
900            if let Some(capture) = pending_capture {
901                ensure!(
902                    !completed.answer.trim().is_empty(),
903                    "subagent provider completed an output capture without contents"
904                );
905                let budget = ContextBudget {
906                    projection: projection.clone(),
907                    max_input_tokens: selected.max_input_tokens,
908                };
909                let outcome = host
910                    .complete_capture(capture, completed.answer, budget)
911                    .await?;
912                let mut candidate = projection.clone();
913                candidate.apply_updates(successful_state_updates(&outcome));
914                candidate.push_history(format!("Ktool result:\n{}", outcome.text));
915                ensure_capacity(&candidate, selected.max_input_tokens)?;
916                let answer = outcome.text.trim().to_owned();
917                ensure!(
918                    !answer.is_empty(),
919                    "subagent output capture completed without a result"
920                );
921                host.record(AuditEvent::Completed {
922                    parent_operation_id: request.parent_operation_id,
923                    model: request.model.clone(),
924                    response: answer.clone(),
925                })?;
926                return Ok(RunResult {
927                    answer,
928                    model: selected,
929                });
930            }
931            let answer = completed.answer.trim().to_owned();
932            if !answer.is_empty() {
933                host.record(AuditEvent::Completed {
934                    parent_operation_id: request.parent_operation_id,
935                    model: request.model.clone(),
936                    response: answer.clone(),
937                })?;
938                return Ok(RunResult {
939                    answer,
940                    model: selected,
941                });
942            }
943            ensure!(
944                !used_tool,
945                "subagent provider completed without a response or tool call"
946            );
947        }
948        anyhow::bail!("subagent provider completed without a terminal response")
949    }
950}
951
952async fn respond_or_record<H: Host>(
953    turn: &mut kcode_intelligence_router::AgentTurn,
954    host: &mut H,
955    parent_operation_id: Uuid,
956    round: u64,
957    manifest_hash: &str,
958    call_id: &str,
959    result: ToolResult,
960) -> anyhow::Result<()> {
961    if let Err(error) = turn.respond(call_id, result).await {
962        let receipt = turn.finish_unavailable()?.clone();
963        host.record(AuditEvent::ProviderReceipt {
964            parent_operation_id,
965            round,
966            manifest_hash: manifest_hash.into(),
967            usage: None,
968            receipt: Box::new(receipt),
969        })?;
970        return Err(anyhow::Error::new(error));
971    }
972    Ok(())
973}
974
975async fn respond_session_or_record<H: SessionHost>(
976    turn: &mut kcode_intelligence_router::AgentTurn,
977    host: &mut H,
978    round: u64,
979    call_id: &str,
980    result: ToolResult,
981) -> anyhow::Result<()> {
982    if let Err(error) = turn.respond(call_id, result).await {
983        record_unavailable_session_turn(turn, host, round).await?;
984        return Err(anyhow::Error::new(error));
985    }
986    Ok(())
987}
988
989async fn record_unavailable_session_turn<H: SessionHost>(
990    turn: &mut kcode_intelligence_router::AgentTurn,
991    host: &mut H,
992    round: u64,
993) -> anyhow::Result<()> {
994    let receipt = turn.finish_unavailable()?.clone();
995    host.record(SessionEvent::ProviderReceipt {
996        round,
997        usage: None,
998        receipt: Box::new(receipt),
999        continuation: None,
1000    })
1001    .await
1002}
1003
1004#[derive(Clone)]
1005struct Projection {
1006    context: Vec<String>,
1007    task: String,
1008    history: Vec<ProjectedHistory>,
1009    states: Vec<ProjectedState>,
1010}
1011
1012#[derive(Clone)]
1013struct ProjectedHistory {
1014    text: String,
1015    displayed_state_keys: Vec<String>,
1016}
1017
1018#[derive(Clone)]
1019struct ProjectedState {
1020    key: String,
1021    text: String,
1022}
1023
1024impl Projection {
1025    fn new(context: Vec<String>, task: String) -> Self {
1026        Self {
1027            context,
1028            task,
1029            history: Vec::new(),
1030            states: Vec::new(),
1031        }
1032    }
1033
1034    fn render(&self) -> String {
1035        self.context
1036            .iter()
1037            .map(String::as_str)
1038            .chain(std::iter::once(self.task.as_str()))
1039            .chain(self.history.iter().map(|entry| entry.text.as_str()))
1040            .chain(
1041                self.states
1042                    .iter()
1043                    .filter(|state| {
1044                        !self.history.iter().any(|entry| {
1045                            entry
1046                                .displayed_state_keys
1047                                .iter()
1048                                .any(|key| key == &state.key)
1049                        })
1050                    })
1051                    .map(|state| state.text.as_str()),
1052            )
1053            .filter(|section| !section.is_empty())
1054            .collect::<Vec<_>>()
1055            .join("\n\n")
1056    }
1057
1058    fn push_history(&mut self, text: impl Into<String>) {
1059        self.push_history_displaying(text, Vec::new());
1060    }
1061
1062    fn push_history_displaying(
1063        &mut self,
1064        text: impl Into<String>,
1065        displayed_state_keys: Vec<String>,
1066    ) {
1067        self.history.push(ProjectedHistory {
1068            text: text.into(),
1069            displayed_state_keys,
1070        });
1071    }
1072
1073    fn update_state(&mut self, key: String, text: Option<String>) {
1074        self.states.retain(|state| state.key != key);
1075        if let Some(text) = text {
1076            self.states.push(ProjectedState { key, text });
1077        }
1078    }
1079
1080    fn apply_updates(&mut self, updates: &[StateUpdate]) {
1081        for update in updates {
1082            let marker = if update.text.is_some() {
1083                SUPERSEDED_TOOL_OUTPUT
1084            } else {
1085                REMOVED_TOOL_OUTPUT
1086            };
1087            for entry in &mut self.history {
1088                if entry
1089                    .displayed_state_keys
1090                    .iter()
1091                    .any(|key| key == &update.key)
1092                {
1093                    entry.text = marker.into();
1094                    entry.displayed_state_keys.clear();
1095                }
1096            }
1097            self.update_state(update.key.clone(), update.text.clone());
1098        }
1099    }
1100
1101    fn estimated_tokens(&self) -> u64 {
1102        (self.render().chars().count() as u64)
1103            .div_ceil(4)
1104            .saturating_add(PROTOCOL_TOKEN_RESERVE)
1105    }
1106}
1107
1108fn successful_state_updates(outcome: &ToolOutcome) -> &[StateUpdate] {
1109    if outcome.ok {
1110        &outcome.state_updates
1111    } else {
1112        &[]
1113    }
1114}
1115
1116fn ensure_capacity(projection: &Projection, max_input_tokens: u64) -> anyhow::Result<()> {
1117    let estimated = projection.estimated_tokens();
1118    ensure!(
1119        estimated <= max_input_tokens,
1120        "subagent context requires approximately {estimated} input tokens, over the selected model's {max_input_tokens}-token input limit"
1121    );
1122    Ok(())
1123}
1124
1125fn ktool_definition(description: &str) -> DynamicTool {
1126    DynamicTool::new(
1127        "call_ktool",
1128        description,
1129        json!({
1130            "type": "object",
1131            "additionalProperties": false,
1132            "required": ["name", "arguments"],
1133            "properties": {
1134                "name": {"type": "string"},
1135                "arguments": {"type": "object"}
1136            }
1137        }),
1138    )
1139}
1140
1141fn parse_ktool_call(call: &DynamicToolCall) -> anyhow::Result<ToolCall> {
1142    ensure!(call.tool == "call_ktool", "unknown provider tool");
1143    let arguments = call
1144        .arguments
1145        .as_object()
1146        .context("call_ktool arguments must be an object")?;
1147    ensure!(
1148        arguments
1149            .keys()
1150            .all(|key| matches!(key.as_str(), "name" | "arguments")),
1151        "call_ktool contains unknown arguments"
1152    );
1153    let name = arguments
1154        .get("name")
1155        .and_then(Value::as_str)
1156        .map(str::trim)
1157        .filter(|name| !name.is_empty() && name.chars().count() <= 100)
1158        .context("call_ktool.name must be a non-empty bounded string")?
1159        .to_owned();
1160    let arguments = arguments
1161        .get("arguments")
1162        .filter(|value| value.is_object())
1163        .context("call_ktool.arguments must be an object")?
1164        .clone();
1165    Ok(ToolCall { name, arguments })
1166}
1167
1168fn reasoning_effort(value: &str) -> anyhow::Result<ReasoningEffort> {
1169    Ok(match value {
1170        "none" => ReasoningEffort::None,
1171        "minimal" => ReasoningEffort::Minimal,
1172        "low" => ReasoningEffort::Low,
1173        "medium" => ReasoningEffort::Medium,
1174        "high" => ReasoningEffort::High,
1175        "xhigh" => ReasoningEffort::XHigh,
1176        "max" => ReasoningEffort::Max,
1177        _ => anyhow::bail!("unsupported reasoning effort {value:?}"),
1178    })
1179}
1180
1181#[cfg(test)]
1182mod tests {
1183    use super::*;
1184
1185    #[test]
1186    fn projection_replaces_state_and_budget_accounts_for_reserve() {
1187        let mut projection = Projection::new(vec!["context".into()], "task".into());
1188        projection.update_state("file".into(), Some("old".into()));
1189        projection.update_state("file".into(), Some("new".into()));
1190        assert_eq!(projection.states.len(), 1);
1191        assert!(projection.render().contains("new"));
1192        assert!(!projection.render().contains("old"));
1193        assert!(projection.estimated_tokens() >= PROTOCOL_TOKEN_RESERVE);
1194    }
1195
1196    #[test]
1197    fn state_update_supersedes_earlier_display_and_moves_current_value_after_history() {
1198        let mut projection = Projection::new(vec!["context".into()], "task".into());
1199        projection.push_history("Ktool call: open");
1200        projection.push_history_displaying(
1201            "Ktool result:\ncomplete old source",
1202            vec!["tool-state:7".into()],
1203        );
1204        projection.push_history("Ktool call: write");
1205        projection.apply_updates(&[StateUpdate {
1206            key: "tool-state:7".into(),
1207            text: Some("complete new source".into()),
1208        }]);
1209        projection.push_history("Ktool result: write completed");
1210
1211        let rendered = projection.render();
1212        assert!(!rendered.contains("complete old source"));
1213        assert_eq!(rendered.matches(SUPERSEDED_TOOL_OUTPUT).count(), 1);
1214        assert_eq!(rendered.matches("complete new source").count(), 1);
1215        assert!(
1216            rendered.find(SUPERSEDED_TOOL_OUTPUT).unwrap()
1217                < rendered.find("Ktool call: write").unwrap()
1218        );
1219        assert!(
1220            rendered.find("Ktool result: write completed").unwrap()
1221                < rendered.find("complete new source").unwrap()
1222        );
1223    }
1224
1225    #[test]
1226    fn current_state_is_not_duplicated_when_the_latest_result_displays_it() {
1227        let mut projection = Projection::new(vec!["context".into()], "task".into());
1228        projection.apply_updates(&[StateUpdate {
1229            key: "tool-state:7".into(),
1230            text: Some("complete source".into()),
1231        }]);
1232        projection.push_history_displaying(
1233            "Ktool result:\ncomplete source",
1234            vec!["tool-state:7".into()],
1235        );
1236
1237        assert_eq!(projection.render().matches("complete source").count(), 1);
1238    }
1239
1240    #[test]
1241    fn state_update_only_supersedes_displays_with_the_same_identity() {
1242        let mut projection = Projection::new(vec!["context".into()], "task".into());
1243        projection.push_history_displaying("first output", vec!["first".into()]);
1244        projection.push_history_displaying("second output", vec!["second".into()]);
1245        projection.apply_updates(&[StateUpdate {
1246            key: "first".into(),
1247            text: Some("current first output".into()),
1248        }]);
1249
1250        let rendered = projection.render();
1251        assert!(
1252            !rendered
1253                .split("\n\n")
1254                .any(|section| section == "first output")
1255        );
1256        assert!(rendered.contains(SUPERSEDED_TOOL_OUTPUT));
1257        assert!(rendered.contains("second output"));
1258        assert!(rendered.contains("current first output"));
1259    }
1260
1261    #[test]
1262    fn removed_state_uses_a_truthful_supersession_marker() {
1263        let mut projection = Projection::new(vec!["context".into()], "task".into());
1264        projection.push_history_displaying("retired output", vec!["state".into()]);
1265        projection.apply_updates(&[StateUpdate {
1266            key: "state".into(),
1267            text: None,
1268        }]);
1269
1270        let rendered = projection.render();
1271        assert!(!rendered.contains("retired output"));
1272        assert!(rendered.contains(REMOVED_TOOL_OUTPUT));
1273    }
1274
1275    #[test]
1276    fn failed_outcomes_cannot_update_or_supersede_state() {
1277        let mut projection = Projection::new(vec!["context".into()], "task".into());
1278        projection.push_history_displaying("current output", vec!["state".into()]);
1279        let failed = ToolOutcome {
1280            text: "write failed".into(),
1281            ok: false,
1282            state_updates: vec![StateUpdate {
1283                key: "state".into(),
1284                text: Some("invalid update".into()),
1285            }],
1286            displayed_state_keys: Vec::new(),
1287            capture: None,
1288        };
1289
1290        projection.apply_updates(successful_state_updates(&failed));
1291
1292        let rendered = projection.render();
1293        assert!(rendered.contains("current output"));
1294        assert!(!rendered.contains("invalid update"));
1295        assert!(!rendered.contains(SUPERSEDED_TOOL_OUTPUT));
1296    }
1297
1298    #[test]
1299    fn large_stateful_result_is_rendered_exactly_once() {
1300        let loaded_node = format!(
1301            "Node body:\n{}\n\nFixed connections:\nfixed-node\n\nRecent connections:\nrecent-node",
1302            "x".repeat(2_000)
1303        );
1304        let outcome = ToolOutcome {
1305            text: loaded_node.clone(),
1306            ok: true,
1307            state_updates: vec![StateUpdate {
1308                key: "loaded-node".into(),
1309                text: Some(loaded_node.clone()),
1310            }],
1311            displayed_state_keys: vec!["loaded-node".into()],
1312            capture: None,
1313        };
1314        let mut projection =
1315            Projection::new(vec!["initial node description".into()], "task".into());
1316        let provider_result = outcome.text.clone();
1317        projection.apply_updates(successful_state_updates(&outcome));
1318        projection.push_history_displaying(
1319            format!("Ktool result:\n{provider_result}"),
1320            outcome.displayed_state_keys,
1321        );
1322
1323        let rendered = projection.render();
1324        assert_eq!(rendered.matches(loaded_node.as_str()).count(), 1);
1325        assert!(rendered.contains("Fixed connections:\nfixed-node"));
1326        assert!(rendered.contains("Recent connections:\nrecent-node"));
1327    }
1328
1329    #[test]
1330    fn native_tool_wrapper_is_strict() {
1331        let call = parse_ktool_call(&DynamicToolCall {
1332            call_id: "1".into(),
1333            tool: "call_ktool".into(),
1334            arguments: json!({"name": "Read", "arguments": {"id": 1}}),
1335        })
1336        .unwrap();
1337        assert_eq!(call.name, "Read");
1338        assert_eq!(call.arguments["id"], 1);
1339    }
1340}
kcode_agent_runtime/lib.rs

kcode_agent_runtime/
lib.rs
