kcode-agent-runtime 0.1.0

//! Provider-neutral subagent loops over `kcode-intelligence-router`.

#![deny(missing_docs)]
#![forbid(unsafe_code)]

use std::{future::Future, pin::Pin, time::Duration};

use anyhow::{Context, ensure};
use kcode_codex_runtime_v2::{
    AgentEvent, AgentRequest, DynamicTool, DynamicToolCall, ReasoningEffort, ToolResult,
};
use kcode_intelligence_router::{Intelligence, ResolvedAgentModel};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use uuid::Uuid;

const DEFAULT_ROUND_LIMIT: u64 = 100;
const PROTOCOL_TOKEN_RESERVE: u64 = 4_096;
const INLINE_TOOL_RESULT_CHARACTERS: usize = 1_000;

/// A boxed asynchronous host operation.
pub type HostFuture<'a, T> = Pin<Box<dyn Future<Output = anyhow::Result<T>> + Send + 'a>>;

/// One application tool call requested by a subagent.
#[derive(Clone, Debug, PartialEq)]
pub struct ToolCall {
    /// Exact application tool name.
    pub name: String,
    /// Tool arguments.
    pub arguments: Value,
}

/// One replaceable state section rendered into every later context slice.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StateUpdate {
    /// Stable state identity. A later update with this key replaces the prior text.
    pub key: String,
    /// Current rendered state, or `None` to remove it.
    pub text: Option<String>,
}

/// Result returned by the application after one tool or capture operation.
#[derive(Clone, Debug, PartialEq)]
pub struct ToolOutcome {
    /// Exact result retained for audit.
    pub text: String,
    /// Whether the operation succeeded.
    pub ok: bool,
    /// Replaceable state made current by the operation.
    pub state_updates: Vec<StateUpdate>,
    /// Opaque application token requesting a tool-free freeform output capture.
    pub capture: Option<Value>,
}

impl ToolOutcome {
    /// Constructs a simple successful result.
    pub fn success(text: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            ok: true,
            state_updates: Vec::new(),
            capture: None,
        }
    }

    /// Constructs a simple failed result.
    pub fn failure(text: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            ok: false,
            state_updates: Vec::new(),
            capture: None,
        }
    }
}

/// Read-only capacity view supplied while a host evaluates a tool.
#[derive(Clone)]
pub struct ContextBudget {
    projection: Projection,
    max_input_tokens: u64,
}

impl ContextBudget {
    /// Current estimated input tokens, including protocol reserve.
    pub fn estimated_tokens(&self) -> u64 {
        self.projection.estimated_tokens()
    }

    /// Maximum permitted input tokens.
    pub fn max_input_tokens(&self) -> u64 {
        self.max_input_tokens
    }

    /// Returns whether replacing one projected state would fit.
    pub fn fits_state(&self, key: impl Into<String>, text: impl Into<String>) -> bool {
        let mut projection = self.projection.clone();
        projection.update_state(key.into(), Some(text.into()));
        projection.estimated_tokens() <= self.max_input_tokens
    }
}

/// Application-owned behavior invoked by the generic subagent loop.
pub trait Host: Send {
    /// Renders the retained invocation text before execution.
    fn render_tool_call(&mut self, call: &ToolCall) -> anyhow::Result<String>;

    /// Executes one application tool.
    fn execute_tool<'a>(
        &'a mut self,
        call: ToolCall,
        operation_id: Uuid,
        budget: ContextBudget,
    ) -> HostFuture<'a, ToolOutcome>;

    /// Completes an opaque freeform capture requested by a prior tool result.
    fn complete_capture<'a>(
        &'a mut self,
        capture: Value,
        contents: String,
        budget: ContextBudget,
    ) -> HostFuture<'a, ToolOutcome>;

    /// Records one durable audit event selected by the runtime.
    fn record(&mut self, label: &str, value: Value) -> anyhow::Result<()>;
}

/// Inputs for one complete subagent run.
#[derive(Clone, Debug, PartialEq)]
pub struct RunRequest {
    /// Stable user identifier used for router accounting.
    pub user_id: String,
    /// Running parent operation whose cancellation propagates to each turn.
    pub parent_operation_id: Uuid,
    /// Exact requested model selector.
    pub model: String,
    /// Provider-neutral reasoning effort.
    pub reasoning_effort: String,
    /// Ordered immutable context sections.
    pub context: Vec<String>,
    /// Exact task presented after the context.
    pub task: String,
    /// Optional per-turn timeout.
    pub timeout: Option<Duration>,
    /// Additional application metadata included in the start audit event.
    pub start_metadata: Value,
}

/// Completed subagent output.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct RunResult {
    /// Final non-empty assistant answer.
    pub answer: String,
    /// Model used for every turn.
    pub model: ResolvedAgentModel,
}

/// Cloneable provider-neutral subagent runtime.
#[derive(Clone)]
pub struct AgentRuntime {
    intelligence: Intelligence,
    round_limit: u64,
}

impl AgentRuntime {
    /// Constructs a runtime over the sole direct-model boundary.
    pub fn new(intelligence: Intelligence) -> Self {
        Self {
            intelligence,
            round_limit: DEFAULT_ROUND_LIMIT,
        }
    }

    /// Resolves a model without running an agent.
    pub async fn resolve_model(&self, requested: &str) -> anyhow::Result<ResolvedAgentModel> {
        self.intelligence
            .resolve_agent_model(requested)
            .await
            .map_err(anyhow::Error::new)
    }

    /// Runs one fresh-context subagent to a final non-empty answer.
    pub async fn run<H: Host>(
        &self,
        request: RunRequest,
        host: &mut H,
    ) -> anyhow::Result<RunResult> {
        let selected = self.resolve_model(&request.model).await?;
        let reasoning_effort = reasoning_effort(&request.reasoning_effort)?;
        let mut projection = Projection::new(request.context, request.task);
        ensure_capacity(&projection, selected.max_input_tokens)?;
        host.record(
            "subagent_started",
            json!({
                "model": request.model,
                "providerModel": selected.provider_model,
                "provider": format!("{:?}", selected.provider),
                "contextWindowTokens": selected.context_window_tokens,
                "maxInputTokens": selected.max_input_tokens,
                "context": projection.context,
                "task": projection.task,
                "host": request.start_metadata,
            }),
        )?;
        let user = self
            .intelligence
            .for_user(request.user_id)
            .map_err(anyhow::Error::new)?;
        let mut deferred_capture: Option<Value> = None;

        for round in 0..self.round_limit {
            let capturing = deferred_capture.is_some();
            ensure_capacity(&projection, selected.max_input_tokens)?;
            let input = projection.render();
            let manifest_hash = hex::encode(Sha256::digest(input.as_bytes()));
            host.record(
                "subagent_inference_submitted",
                json!({
                    "round": round + 1,
                    "manifestHash": manifest_hash,
                    "estimatedInputTokens": projection.estimated_tokens(),
                }),
            )?;
            let mut provider_request = AgentRequest::new(input, selected.requested_model.clone());
            provider_request.reasoning_effort = reasoning_effort;
            provider_request.ephemeral = true;
            provider_request.tools = if capturing {
                Vec::new()
            } else {
                vec![ktool_definition()]
            };
            if let Some(timeout) = request.timeout {
                provider_request.timeout = timeout;
            }
            let child_operation_id = Uuid::new_v4();
            let mut turn = user
                .start_agent_turn(
                    child_operation_id,
                    Some(request.parent_operation_id),
                    provider_request,
                )
                .await
                .map_err(anyhow::Error::new)?;
            let mut used_tool = false;
            let mut pending_capture: Option<Value> = None;
            let mut requires_rerender = false;
            let completed = loop {
                let event = turn
                    .next_event()
                    .await
                    .map_err(anyhow::Error::new)?
                    .context("subagent provider ended without a terminal turn event")?;
                match event {
                    AgentEvent::ProviderInput(_) => {}
                    AgentEvent::ToolCall(native) => {
                        used_tool = true;
                        if capturing {
                            turn.respond(
                                &native.call_id,
                                ToolResult::failure(
                                    "No application tool is available while complete freeform output is being captured.",
                                ),
                            )
                            .await
                            .map_err(anyhow::Error::new)?;
                            continue;
                        }
                        if pending_capture.is_some() {
                            turn.respond(
                                &native.call_id,
                                ToolResult::failure(
                                    "A freeform output capture is pending; no other tool can run first.",
                                ),
                            )
                            .await
                            .map_err(anyhow::Error::new)?;
                            continue;
                        }
                        if requires_rerender {
                            turn.respond(
                                &native.call_id,
                                ToolResult::failure(
                                    "A state update is waiting to be re-rendered. End this slice before calling another tool.",
                                ),
                            )
                            .await
                            .map_err(anyhow::Error::new)?;
                            continue;
                        }
                        let call = match parse_ktool_call(&native) {
                            Ok(call) => call,
                            Err(error) => {
                                let text = format!("Invalid application tool call: {error}");
                                projection.push_history(format!("Ktool result:\n{text}"));
                                turn.respond(&native.call_id, ToolResult::failure(text))
                                    .await
                                    .map_err(anyhow::Error::new)?;
                                continue;
                            }
                        };
                        host.record(
                            "subagent_tool_call",
                            json!({"name": call.name, "arguments": call.arguments}),
                        )?;
                        projection.push_history(format!(
                            "Ktool call:\n{}",
                            host.render_tool_call(&call)?
                        ));
                        let budget = ContextBudget {
                            projection: projection.clone(),
                            max_input_tokens: selected.max_input_tokens,
                        };
                        let mut outcome = host
                            .execute_tool(call.clone(), child_operation_id, budget)
                            .await
                            .unwrap_or_else(|error| {
                                ToolOutcome::failure(format!("{} failed: {error}", call.name))
                            });
                        let exact_result = outcome.text.clone();
                        let initially_ok = outcome.ok;
                        let mut provider_result =
                            compact_tool_result(&outcome.text, &outcome.state_updates);
                        let mut candidate = projection.clone();
                        candidate.apply_updates(&outcome.state_updates);
                        candidate.push_history(format!("Ktool result:\n{provider_result}"));
                        let accepted = candidate.estimated_tokens() <= selected.max_input_tokens;
                        if accepted {
                            projection = candidate;
                            requires_rerender = !outcome.state_updates.is_empty();
                        } else {
                            outcome.ok = false;
                            outcome.capture = None;
                            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();
                            projection.push_history(format!("Ktool result:\n{provider_result}"));
                        }
                        host.record(
                            "subagent_tool_result",
                            json!({
                                "name": call.name,
                                "ok": initially_ok,
                                "projectionAccepted": accepted,
                                "result": exact_result,
                            }),
                        )?;
                        pending_capture = outcome.capture.take();
                        turn.respond(
                            &native.call_id,
                            if outcome.ok {
                                ToolResult::success(provider_result)
                            } else {
                                ToolResult::failure(provider_result)
                            },
                        )
                        .await
                        .map_err(anyhow::Error::new)?;
                    }
                    AgentEvent::Completed(completed) => break completed,
                }
            };
            host.record(
                "subagent_provider_receipt",
                json!({
                    "round": round + 1,
                    "usage": completed.usage.as_ref().map(|usage| json!({
                        "inputTokens": usage.input_tokens,
                        "outputTokens": usage.output_tokens,
                        "cachedInputTokens": usage.cached_input_tokens,
                        "reasoningOutputTokens": usage.reasoning_output_tokens,
                        "lastInputTokens": usage.last_input_tokens,
                        "lastOutputTokens": usage.last_output_tokens,
                    })),
                }),
            )?;

            let capture = deferred_capture.take().or(pending_capture);
            if let Some(capture) = capture {
                if !capturing && completed.answer.is_empty() {
                    deferred_capture = Some(capture);
                    continue;
                }
                let budget = ContextBudget {
                    projection: projection.clone(),
                    max_input_tokens: selected.max_input_tokens,
                };
                let outcome = host
                    .complete_capture(capture, completed.answer, budget)
                    .await?;
                let mut candidate = projection.clone();
                candidate.apply_updates(&outcome.state_updates);
                candidate.push_history(format!("Ktool result:\n{}", outcome.text));
                ensure_capacity(&candidate, selected.max_input_tokens)?;
                projection = candidate;
                continue;
            }
            if requires_rerender {
                let draft = completed.answer.trim();
                if !draft.is_empty() {
                    projection.push_history(format!(
                        "Assistant draft produced before the state refresh:\n{draft}"
                    ));
                }
                continue;
            }
            let answer = completed.answer.trim().to_owned();
            if !answer.is_empty() {
                host.record(
                    "subagent_completed",
                    json!({"model": request.model, "response": answer}),
                )?;
                return Ok(RunResult {
                    answer,
                    model: selected,
                });
            }
            ensure!(
                used_tool,
                "subagent provider completed without a response or tool call"
            );
        }
        anyhow::bail!(
            "subagent exceeded the {}-round tool-loop safety limit",
            self.round_limit
        )
    }
}

#[derive(Clone)]
struct Projection {
    context: Vec<String>,
    task: String,
    history: Vec<String>,
    states: Vec<ProjectedState>,
}

#[derive(Clone)]
struct ProjectedState {
    key: String,
    text: String,
}

impl Projection {
    fn new(context: Vec<String>, task: String) -> Self {
        Self {
            context,
            task,
            history: Vec::new(),
            states: Vec::new(),
        }
    }

    fn render(&self) -> String {
        self.context
            .iter()
            .map(String::as_str)
            .chain(std::iter::once(self.task.as_str()))
            .chain(self.history.iter().map(String::as_str))
            .chain(self.states.iter().map(|state| state.text.as_str()))
            .filter(|section| !section.is_empty())
            .collect::<Vec<_>>()
            .join("\n\n")
    }

    fn push_history(&mut self, text: impl Into<String>) {
        self.history.push(text.into());
    }

    fn update_state(&mut self, key: String, text: Option<String>) {
        self.states.retain(|state| state.key != key);
        if let Some(text) = text {
            self.states.push(ProjectedState { key, text });
        }
    }

    fn apply_updates(&mut self, updates: &[StateUpdate]) {
        for update in updates {
            self.update_state(update.key.clone(), update.text.clone());
        }
    }

    fn estimated_tokens(&self) -> u64 {
        (self.render().chars().count() as u64)
            .div_ceil(4)
            .saturating_add(PROTOCOL_TOKEN_RESERVE)
    }
}

fn compact_tool_result(text: &str, states: &[StateUpdate]) -> String {
    if states.is_empty() {
        return text.to_owned();
    }
    let result = if text.chars().count() <= INLINE_TOOL_RESULT_CHARACTERS {
        text
    } else {
        "Tool completed successfully."
    };
    format!(
        "{result}\n\nThe updated state will be rendered in the next fresh context slice; end this slice now."
    )
}

fn ensure_capacity(projection: &Projection, max_input_tokens: u64) -> anyhow::Result<()> {
    let estimated = projection.estimated_tokens();
    ensure!(
        estimated <= max_input_tokens,
        "subagent context requires approximately {estimated} input tokens, over the selected model's {max_input_tokens}-token input limit"
    );
    Ok(())
}

fn ktool_definition() -> DynamicTool {
    DynamicTool::new(
        "call_ktool",
        "Call one available Ktool by its exact name.",
        json!({
            "type": "object",
            "additionalProperties": false,
            "required": ["name", "arguments"],
            "properties": {
                "name": {"type": "string"},
                "arguments": {"type": "object"}
            }
        }),
    )
}

fn parse_ktool_call(call: &DynamicToolCall) -> anyhow::Result<ToolCall> {
    ensure!(call.tool == "call_ktool", "unknown provider tool");
    let arguments = call
        .arguments
        .as_object()
        .context("call_ktool arguments must be an object")?;
    ensure!(
        arguments
            .keys()
            .all(|key| matches!(key.as_str(), "name" | "arguments")),
        "call_ktool contains unknown arguments"
    );
    let name = arguments
        .get("name")
        .and_then(Value::as_str)
        .map(str::trim)
        .filter(|name| !name.is_empty() && name.chars().count() <= 100)
        .context("call_ktool.name must be a non-empty bounded string")?
        .to_owned();
    let arguments = arguments
        .get("arguments")
        .filter(|value| value.is_object())
        .context("call_ktool.arguments must be an object")?
        .clone();
    Ok(ToolCall { name, arguments })
}

fn reasoning_effort(value: &str) -> anyhow::Result<ReasoningEffort> {
    Ok(match value {
        "none" => ReasoningEffort::None,
        "minimal" => ReasoningEffort::Minimal,
        "low" => ReasoningEffort::Low,
        "medium" => ReasoningEffort::Medium,
        "high" => ReasoningEffort::High,
        "xhigh" => ReasoningEffort::XHigh,
        "max" => ReasoningEffort::Max,
        _ => anyhow::bail!("unsupported reasoning effort {value:?}"),
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn projection_replaces_state_and_budget_accounts_for_reserve() {
        let mut projection = Projection::new(vec!["context".into()], "task".into());
        projection.update_state("file".into(), Some("old".into()));
        projection.update_state("file".into(), Some("new".into()));
        assert_eq!(projection.states.len(), 1);
        assert!(projection.render().contains("new"));
        assert!(!projection.render().contains("old"));
        assert!(projection.estimated_tokens() >= PROTOCOL_TOKEN_RESERVE);
    }

    #[test]
    fn state_changes_compact_large_tool_results() {
        let compacted = compact_tool_result(
            &"x".repeat(INLINE_TOOL_RESULT_CHARACTERS + 1),
            &[StateUpdate {
                key: "state".into(),
                text: Some("current".into()),
            }],
        );
        assert!(compacted.starts_with("Tool completed successfully."));
        assert!(compacted.contains("fresh context slice"));
    }

    #[test]
    fn native_tool_wrapper_is_strict() {
        let call = parse_ktool_call(&DynamicToolCall {
            call_id: "1".into(),
            tool: "call_ktool".into(),
            arguments: json!({"name": "Read", "arguments": {"id": 1}}),
        })
        .unwrap();
        assert_eq!(call.name, "Read");
        assert_eq!(call.arguments["id"], 1);
    }
}