klieo-core 2.2.0

//! Runtime loop driving an [`crate::agent::Agent`] through LLM ↔ tool
//! exchanges.
//!
//! This is the narrow primitive every higher layer (graph, actor,
//! spec-aggregate) builds on. It does one thing: alternate LLM
//! completions with tool dispatch, persisting episodic events, until
//! the LLM emits a final text reply, hits `max_steps`, or is cancelled.
//!
//! Module map:
//! - `checkpoint` — `persist_checkpoint` / `resume_from_checkpoint` behaviour.
//! - `dispatch` — shared tool-dispatch helper (deduped between drivers).
//! - `guardrails` — `check_pre_llm` / `check_post_llm` dispatch.
//! - `options` — [`RunOptions`] tunables + defaults.
//! - `retry` — bounded exponential-backoff retry policy + constants.
//! - `review` — [`ReviewPolicy`] suspend gate + [`NeverReview`] default.
//! - `streaming` — [`run_steps_streaming`] and its loop driver.
//! - `streaming_forward` — chunk-pump helper for the streaming loop.
//! - `mod.rs` (this file) — [`run_steps`] + `run_loop` + shared `build_request`.

mod capture_sink;
mod checkpoint;
mod dispatch;
mod guardrails;
mod options;
mod retry;
mod review;
mod step;
mod streaming;
mod streaming_forward;

use crate::agent::{AgentContext, AgentEvent};
use crate::error::{Error, LlmError};
use crate::ids::ThreadId;
use crate::llm::{ChatRequest, Message, Role};
use crate::memory::Episode;
use std::time::Instant;
use tracing::{error, info, instrument};

use guardrails::{check_post_llm, check_pre_llm};
use retry::complete_with_retry;
use step::{record_and_dispatch, StepDisposition, StepOutcome};

/// Best-effort send of an [`AgentEvent`] to `ctx.progress`.
/// No-op when the channel is unset; dropped receivers are
/// silently ignored (broadcast semantics — zero subscribers is
/// legitimate).
pub(crate) fn emit(ctx: &AgentContext, event: AgentEvent) {
    if let Some(tx) = &ctx.progress {
        let _ = tx.send(event);
    }
}

/// Sanitise a runtime error into a stable wire-safe reason
/// string for `AgentEvent::Failed`. Matches the policy
/// `streaming::terminal_chunk_for` already uses on the
/// streaming path.
fn sanitised_reason(err: &Error) -> String {
    match err {
        Error::Llm(LlmError::RateLimit { .. }) => "llm rate-limited".into(),
        Error::Llm(LlmError::Unauthorized) => "llm auth failed".into(),
        Error::Llm(LlmError::Server(_)) => "llm server error".into(),
        Error::Llm(LlmError::BadRequest(_)) => "llm bad request".into(),
        Error::Llm(LlmError::Network(_)) => "llm network error".into(),
        Error::Llm(LlmError::Timeout) => "llm timeout".into(),
        Error::Llm(LlmError::Decoding(_)) => "llm decoding error".into(),
        Error::Llm(LlmError::Unsupported(_)) => "llm unsupported capability".into(),
        Error::Llm(LlmError::Cancelled) => "cancelled".into(),
        Error::Cancelled => "cancelled".into(),
        Error::MaxStepsExceeded { .. } => "max steps exceeded".into(),
        Error::Bus(_) => "bus error".into(),
        Error::Memory(_) => "memory error".into(),
        Error::Tool(_) => "tool error".into(),
        Error::Refused { .. } => "guardrail rejected".into(),
        Error::Handoff { .. } => "guardrail handoff".into(),
        Error::Suspended { .. } => "run suspended".into(),
        _ => "internal error".into(),
    }
}

pub use crate::checkpoint::{ApprovalDecision, RunCheckpoint, CHECKPOINT_BUCKET};
pub use capture_sink::CaptureSink;
pub use checkpoint::{
    gc_checkpoints, resume_from_checkpoint, spawn_checkpoint_gc, CheckpointGcHandle,
};
pub(crate) use dispatch::dispatch_tool_calls;
pub use options::RunOptions;
pub use review::{NeverReview, ReviewPolicy};
pub use streaming::run_steps_streaming;

/// Drive the agent's LLM/tool loop until completion. Caller is
/// responsible for appending the user message before invoking — this
/// function only consumes / extends short-term memory.
///
/// Returns the assistant's final text response.
///
/// **Episode logging.** Records `Episode::Started` on entry, `LlmCall`
/// per cycle, `ToolCall` per dispatched tool, and `Completed` on
/// success. **Does NOT record `Episode::Failed` on errors** — failures
/// propagate via the `Result::Err` return and the caller decides how
/// to log them. This differs from [`run_steps_streaming`], which
/// always records a terminal `Failed` episode because errors cannot
/// flow back through a returned stream.
#[instrument(level = "debug", skip(ctx, system_prompt), fields(run_id = %ctx.run_id))]
pub async fn run_steps(
    ctx: &AgentContext,
    system_prompt: &str,
    thread: ThreadId,
    opts: RunOptions,
) -> Result<String, Error> {
    ctx.episodic
        .record(
            ctx.run_id,
            Episode::Started {
                agent: ctx.agent_name.clone(),
            },
        )
        .await?;
    record_run_attribution(ctx).await?;
    record_run_origin(ctx).await?;

    run_loop(ctx, system_prompt, &thread, &opts, 0).await
}

/// Stamps the derived non-PII tenant label onto the audit trail when
/// the caller installed one via [`AgentContext::with_tenant_label`].
/// Same error-handling contract as the adjacent `Episode::Started`
/// record — a memory-record failure aborts run entry rather than
/// silently losing attribution.
pub(crate) async fn record_run_attribution(ctx: &AgentContext) -> Result<(), Error> {
    if let Some(label) = ctx.tenant_label.as_ref() {
        ctx.episodic
            .record(
                ctx.run_id,
                Episode::RunAttributed {
                    tenant_label: label.clone(),
                },
            )
            .await?;
    }
    Ok(())
}

/// Records the cross-hop provenance origin onto the audit trail when an
/// authenticated caller installed a parent-chain anchor via
/// [`AgentContext::with_parent_anchor`]. Co-recorded with
/// [`record_run_attribution`] so the (unverified) parent claim is
/// attributable to the authenticated principal. Same error contract as
/// the adjacent `Episode::Started` record — a memory-record failure
/// aborts run entry rather than silently losing the origin link.
pub(crate) async fn record_run_origin(ctx: &AgentContext) -> Result<(), Error> {
    if let Some(anchor) = ctx.parent_anchor.as_ref() {
        ctx.episodic
            .record(
                ctx.run_id,
                Episode::RunOrigin {
                    parent_anchor: anchor.clone(),
                },
            )
            .await?;
    }
    Ok(())
}

pub(crate) async fn run_loop(
    ctx: &AgentContext,
    system_prompt: &str,
    thread: &ThreadId,
    opts: &RunOptions,
    mut step: u32,
) -> Result<String, Error> {
    loop {
        if ctx.cancel.is_cancelled() {
            error!(run_id = %ctx.run_id, thread_id = %thread, operation = "run_loop", "cancelled");
            emit(
                ctx,
                AgentEvent::Failed {
                    reason: "cancelled".into(),
                },
            );
            return Err(Error::Cancelled);
        }
        if step >= opts.max_steps {
            emit(
                ctx,
                AgentEvent::Failed {
                    reason: format!("max steps exceeded ({})", opts.max_steps),
                },
            );
            return Err(Error::MaxStepsExceeded {
                steps: opts.max_steps,
            });
        }
        step += 1;

        let req = build_request(ctx, system_prompt, thread, opts.max_history_tokens).await?;

        check_pre_llm(&opts.guardrails, &req).await?;

        let llm_started = Instant::now();
        emit(ctx, AgentEvent::LlmCallStarted);
        let resp = match complete_with_retry(ctx.llm.as_ref(), &ctx.cancel, req.clone()).await {
            Ok(r) => r,
            Err(e) => {
                let latency_ms = llm_started.elapsed().as_millis() as u64;
                emit(
                    ctx,
                    AgentEvent::LlmCallCompleted {
                        tokens: 0,
                        latency_ms,
                    },
                );
                emit(
                    ctx,
                    AgentEvent::Failed {
                        reason: sanitised_reason(&e),
                    },
                );
                return Err(e);
            }
        };
        let latency_ms = llm_started.elapsed().as_millis() as u64;
        let tokens = resp.usage.prompt_tokens + resp.usage.completion_tokens;
        emit(ctx, AgentEvent::LlmCallCompleted { tokens, latency_ms });

        check_post_llm(&opts.guardrails, &req, &resp).await?;

        if let Some(sink) = &opts.capture_sink {
            sink.record_llm_call(&req, &resp);
        }

        if let Some(reason) = opts
            .review_policy
            .should_pause_for_approval(step, &resp.message)
            .await?
        {
            let checkpoint = checkpoint::build_suspend_checkpoint(
                ctx,
                thread,
                step,
                &resp.message,
                resp.finish_reason,
                opts.max_history_tokens,
            )
            .await?;
            if let Some(bucket) = &opts.checkpoint_kv_bucket {
                checkpoint::persist_checkpoint(ctx, bucket, &checkpoint).await?;
            }
            info!(run_id = %ctx.run_id, thread_id = %thread, step, %reason, "run suspended for human review");
            emit(
                ctx,
                AgentEvent::Suspended {
                    reason: reason.clone(),
                },
            );
            return Err(Error::Suspended {
                checkpoint: Box::new(checkpoint),
                reason,
            });
        }

        let outcome = StepOutcome {
            message: resp.message,
            finish_reason: resp.finish_reason,
            usage: resp.usage,
            latency_ms: latency_ms as u32,
        };
        match record_and_dispatch(ctx, thread, step, outcome, "blocking").await? {
            StepDisposition::Done(content) => {
                emit(ctx, AgentEvent::Completed);
                return Ok(content);
            }
            StepDisposition::Continue => continue,
        }
    }
}

/// Build a single [`ChatRequest`] from short-term memory + the system
/// prompt + the active tool catalogue. Shared by [`run_steps`] and
/// [`run_steps_streaming`] so prompt-assembly stays identical between
/// paths.
pub(super) async fn build_request(
    ctx: &AgentContext,
    system_prompt: &str,
    thread: &ThreadId,
    max_history_tokens: usize,
) -> Result<ChatRequest, Error> {
    let history = ctx
        .short_term
        .load(thread.clone(), max_history_tokens)
        .await?;
    let mut messages = Vec::with_capacity(history.len() + 1);
    if !system_prompt.is_empty() {
        messages.push(Message {
            role: Role::System,
            content: system_prompt.into(),
            tool_calls: vec![],
            tool_call_id: None,
        });
    }
    messages.extend(history);

    Ok(ChatRequest {
        messages,
        tools: ctx.tools.catalogue(),
        ..ChatRequest::new(vec![])
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::Error;
    use crate::llm::{ChatRequest, ChatResponse, FinishReason, Message};
    use crate::test_utils::{fake_context, fake_kv, FakeLlmClient, FakeLlmStep};
    use async_trait::async_trait;
    use std::sync::{Arc, Mutex};

    #[tokio::test]
    async fn capture_sink_records_each_successful_llm_call() {
        #[derive(Default)]
        struct RecordingSink {
            calls: Mutex<Vec<(String, FinishReason)>>,
        }
        impl CaptureSink for RecordingSink {
            fn record_llm_call(&self, _request: &ChatRequest, response: &ChatResponse) {
                self.calls
                    .lock()
                    .unwrap()
                    .push((response.message.content.clone(), response.finish_reason));
            }
        }

        let mut ctx = fake_context("capture-test");
        ctx.llm = Arc::new(
            FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("answer".into())]),
        );
        let thread = ThreadId::new("t-capture");
        ctx.short_term
            .append(
                thread.clone(),
                Message {
                    role: crate::llm::Role::User,
                    content: "go".into(),
                    tool_calls: vec![],
                    tool_call_id: None,
                },
            )
            .await
            .unwrap();

        let sink = Arc::new(RecordingSink::default());
        let opts = RunOptions::default().with_capture_sink(sink.clone());
        let out = run_steps(&ctx, "sys", thread, opts).await.unwrap();
        assert_eq!(out, "answer");

        let calls = sink.calls.lock().unwrap();
        assert_eq!(calls.len(), 1, "one record per successful LLM call");
        assert_eq!(calls[0].0, "answer");
        assert_eq!(calls[0].1, FinishReason::Stop);
    }

    #[tokio::test]
    async fn capture_sink_not_called_when_llm_errors() {
        #[derive(Default)]
        struct CountingSink {
            calls: Mutex<usize>,
        }
        impl CaptureSink for CountingSink {
            fn record_llm_call(&self, _r: &ChatRequest, _resp: &ChatResponse) {
                *self.calls.lock().unwrap() += 1;
            }
        }

        let mut ctx = fake_context("capture-err");
        // Empty script → the first LLM call errors (script exhausted).
        ctx.llm = Arc::new(FakeLlmClient::new("fake"));
        let thread = ThreadId::new("t-err");
        ctx.short_term
            .append(
                thread.clone(),
                Message {
                    role: crate::llm::Role::User,
                    content: "go".into(),
                    tool_calls: vec![],
                    tool_call_id: None,
                },
            )
            .await
            .unwrap();

        let sink = Arc::new(CountingSink::default());
        let opts = RunOptions::default().with_capture_sink(sink.clone());
        let result = run_steps(&ctx, "sys", thread, opts).await;
        assert!(result.is_err(), "empty script makes the LLM call fail");
        assert_eq!(
            *sink.calls.lock().unwrap(),
            0,
            "sink must not fire on the error path"
        );
    }

    #[tokio::test]
    async fn run_suspends_when_policy_pauses_and_persists_checkpoint() {
        struct PauseFirst;

        #[async_trait]
        impl ReviewPolicy for PauseFirst {
            async fn should_pause_for_approval(
                &self,
                step: u32,
                _m: &Message,
            ) -> Result<Option<String>, Error> {
                Ok((step == 1).then(|| "manual review".to_string()))
            }
        }

        let mut ctx = fake_context("gate-test");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("hi".into())]));
        ctx.kv = fake_kv();
        let (progress_tx, mut progress_rx) = tokio::sync::broadcast::channel(8);
        ctx.progress = Some(progress_tx);

        let thread = ThreadId::new("t-gate");
        ctx.short_term
            .append(
                thread.clone(),
                Message {
                    role: crate::llm::Role::User,
                    content: "go".into(),
                    tool_calls: vec![],
                    tool_call_id: None,
                },
            )
            .await
            .unwrap();

        let opts = RunOptions::default()
            .with_review_policy(Arc::new(PauseFirst))
            .with_checkpoint_bucket(CHECKPOINT_BUCKET);

        let err = run_steps(&ctx, "sys", thread.clone(), opts)
            .await
            .unwrap_err();

        let cp = match err {
            Error::Suspended { checkpoint, reason } => {
                assert_eq!(reason, "manual review");
                checkpoint
            }
            other => panic!("expected Suspended, got {other:?}"),
        };
        assert_eq!(cp.step_index, 1);

        let mut suspended_reasons = Vec::new();
        while let Ok(event) = progress_rx.try_recv() {
            if let AgentEvent::Suspended { reason } = event {
                suspended_reasons.push(reason);
            }
        }
        assert_eq!(
            suspended_reasons,
            vec!["manual review".to_string()],
            "exactly one Suspended event must reach subscribers, carrying the reason"
        );

        let stored = ctx
            .kv
            .get(CHECKPOINT_BUCKET, &cp.run_id.to_string())
            .await
            .unwrap()
            .expect("checkpoint must be persisted to kv");
        let persisted: RunCheckpoint = serde_json::from_slice(&stored.value).unwrap();
        assert_eq!(persisted.run_id, cp.run_id);
        assert_eq!(persisted.step_index, 1);
        assert_eq!(persisted.thread_id, cp.thread_id);
    }

    #[tokio::test]
    async fn run_suspends_without_persisting_when_no_bucket_configured() {
        struct PauseFirst;

        #[async_trait]
        impl ReviewPolicy for PauseFirst {
            async fn should_pause_for_approval(
                &self,
                step: u32,
                _m: &Message,
            ) -> Result<Option<String>, Error> {
                Ok((step == 1).then(|| "manual review".to_string()))
            }
        }

        let mut ctx = fake_context("gate-no-bucket");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("hi".into())]));
        ctx.kv = fake_kv();

        let thread = ThreadId::new("t-no-bucket");

        let opts = RunOptions::default().with_review_policy(Arc::new(PauseFirst));

        let err = run_steps(&ctx, "sys", thread.clone(), opts)
            .await
            .unwrap_err();

        let cp = match err {
            Error::Suspended { checkpoint, reason } => {
                assert_eq!(reason, "manual review");
                checkpoint
            }
            other => panic!("expected Suspended, got {other:?}"),
        };

        let stored = ctx
            .kv
            .get(CHECKPOINT_BUCKET, &cp.run_id.to_string())
            .await
            .unwrap();
        assert!(
            stored.is_none(),
            "no checkpoint_kv_bucket means the checkpoint travels only in Error::Suspended, not kv"
        );
    }
}