klieo-core 3.14.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.
//! - `compaction` — `maybe_compact` auto-wires `summarize_history` into the loop (Item C).
//! - `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.
//! - `structured` — [`run_structured`] bounded retry-with-feedback for
//!   structured-output parsing (see [`crate::response::parse_structured`]).
//! - `mod.rs` (this file) — [`run_steps`] + `run_loop` + shared `build_request`.

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

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 compaction::maybe_compact;
use guardrails::{check_post_llm, check_pre_llm};
use retry::complete_with_retry;
use step::{append_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, CaptureSinkProvider};
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(crate) use step::CompletionRecording;
pub use streaming::run_steps_streaming;
pub use structured::{run_structured, MAX_STRUCTURED_RETRIES};

/// 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. See
/// [`crate::runtime::run_structured`] for a bounded-retry variant with
/// different completion-recording semantics: it never calls this
/// function, precisely because every attempt here records `Completed`
/// unconditionally on a normal finish.
#[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> {
    record_run_entry(ctx).await?;

    run_loop(
        ctx,
        system_prompt,
        &thread,
        &opts,
        0,
        CompletionRecording::Record,
    )
    .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(())
}

/// Records the local parent → child link when the context names a spawning run.
///
/// Co-recorded with [`record_run_origin`] at run entry. Without it, a fan-out
/// renders as N disconnected roots: an observability view has no other way to
/// learn that one run created another in-process, since edges otherwise come
/// only from bus causation. Same best-effort-but-abort contract as its
/// neighbours — losing the edge silently is what this exists to prevent.
pub(crate) async fn record_run_parent(ctx: &AgentContext) -> Result<(), Error> {
    if let Some(parent_run) = ctx.parent_run() {
        ctx.episodic
            .record(
                ctx.run_id,
                Episode::SpawnedBy {
                    parent_run: parent_run.to_string(),
                },
            )
            .await?;
    }
    Ok(())
}

/// Record a run's lifecycle around work that never calls [`run_steps`].
///
/// An [`crate::agent::Agent`] impl is free to do its job without an LLM — a
/// recall, a tally, a lookup. Nothing then records an episode for it, and a
/// projected run with no `Completed` has no way to end: every observability
/// view shows it as permanently running, indistinguishable from a stage that
/// hung. Wrapping the work records `Started` (plus the same attribution/origin/
/// parent episodes `run_steps` records) and then `Completed` or `Failed`.
///
/// The result is returned untouched. A recording failure aborts before the work
/// runs, matching `run_steps`' own entry contract; a failure recording the
/// terminal episode is logged and swallowed, because the work has already
/// happened by then and losing the record must not turn a completed run into an
/// error.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::runtime::record_run_lifecycle;
/// use klieo_core::test_utils::fake_context;
/// use klieo_core::{Episode, EpisodicMemory};
///
/// let ctx = fake_context("tally");
/// let answer = record_run_lifecycle(&ctx, async { Ok::<_, klieo_core::error::Error>(41 + 1) })
///     .await
///     .unwrap();
/// assert_eq!(answer, 42);
/// let episodes = ctx.episodic.replay(ctx.run_id).await.unwrap();
/// assert!(matches!(episodes.last(), Some(Episode::Completed)));
/// # });
/// ```
pub async fn record_run_lifecycle<T, E, F>(ctx: &AgentContext, work: F) -> Result<T, E>
where
    F: std::future::Future<Output = Result<T, E>>,
    E: std::fmt::Display,
{
    if let Err(error) = record_run_entry(ctx).await {
        error!(%error, run_id = %ctx.run_id, "failed to record run entry");
    }
    let outcome = work.await;
    let terminal = match &outcome {
        Ok(_) => Episode::Completed,
        Err(error) => Episode::Failed {
            error: error.to_string(),
        },
    };
    if let Err(error) = ctx.episodic.record(ctx.run_id, terminal).await {
        error!(%error, run_id = %ctx.run_id, "failed to record run outcome");
    }
    outcome
}

/// The four episodes every run entry records: `Started`, plus tenant
/// attribution, cross-hop origin and local parent when the context carries them.
pub(crate) async fn record_run_entry(ctx: &AgentContext) -> Result<(), 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?;
    record_run_parent(ctx).await
}

/// Shared step loop behind [`run_steps`] and
/// [`crate::runtime::run_structured`]. `completion_recording` controls
/// whether a normal-finish step records `Episode::Completed` — see
/// [`CompletionRecording`] for why `run_structured` always passes
/// [`CompletionRecording::Suppress`] instead of the [`CompletionRecording::Record`]
/// every other caller uses.
pub(crate) async fn run_loop(
    ctx: &AgentContext,
    system_prompt: &str,
    thread: &ThreadId,
    opts: &RunOptions,
    mut step: u32,
    completion_recording: CompletionRecording,
) -> 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;

        maybe_compact(
            ctx,
            thread,
            opts.compaction.as_ref(),
            opts.max_history_tokens,
        )
        .await?;

        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?;

        // `RunOptions` first, then the context's own sink: an agent wrapper owns
        // its run options, while whoever built the context owns the run's
        // identity -- and a capture has to be filed under a `RunId`.
        if let Some(sink) = opts.capture_sink.as_ref().or(ctx.capture_sink()) {
            sink.record_llm_call(&req, &resp);
        }

        // Before the review gate, not after: the gate returns
        // `Error::Suspended` and never reaches the post-step skeleton, so
        // recording later loses this call from a suspended run entirely.
        let (provider, model) = step::split_provider_model(ctx.llm.name());
        step::record_llm_call(ctx, &resp.usage, latency_ms as u32, provider, model).await?;

        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,
        };
        match append_and_dispatch(ctx, thread, step, outcome, "blocking", completion_recording)
            .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 records_provider_and_model_from_the_client_name() {
        let mut ctx = fake_context("cost-live");
        ctx.llm = Arc::new(
            FakeLlmClient::new("openai:gpt-4o").with_steps(vec![FakeLlmStep::Text("ok".into())]),
        );
        let thread = ThreadId::new("t-cost");
        ctx.short_term
            .append(
                thread.clone(),
                Message {
                    role: crate::llm::Role::User,
                    content: "go".into(),
                    tool_calls: vec![],
                    tool_call_id: None,
                },
            )
            .await
            .unwrap();

        run_steps(&ctx, "sys", thread, RunOptions::default())
            .await
            .unwrap();

        let episodes = ctx.episodic.replay(ctx.run_id).await.unwrap();
        let (provider, model) = episodes
            .iter()
            .find_map(|e| match e {
                Episode::LlmCall {
                    provider, model, ..
                } => Some((provider.clone(), model.clone())),
                _ => None,
            })
            .expect("an LlmCall episode");
        assert_eq!(provider.as_deref(), Some("openai"));
        assert_eq!(model.as_deref(), Some("gpt-4o"));
    }

    #[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);
        assert_eq!(
            cp.messages.len(),
            2,
            "checkpoint must include the suspending turn's own assistant message, not just prior history"
        );
        let suspending_turn = cp.messages.last().unwrap();
        assert_eq!(suspending_turn.role, Role::Assistant);
        assert_eq!(suspending_turn.content, "hi");

        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:?}"),
        };
        assert_eq!(
            cp.messages.len(),
            1,
            "checkpoint must include the suspending turn's own assistant message even with no prior history"
        );
        assert_eq!(cp.messages[0].role, Role::Assistant);
        assert_eq!(cp.messages[0].content, "hi");

        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"
        );
    }
}