klieo-core 3.10.0

//! Bounded retry-with-feedback loop for structured-output parsing.
//!
//! [`crate::runtime::run_steps`] returns raw text;
//! [`crate::response::parse_structured`] recovers a typed `T` from it but
//! is itself permanent (a malformed reply will fail identically on a
//! second parse of the SAME string). This module composes the two: on a
//! parse failure, it appends a feedback message describing *why* the
//! reply was rejected to the same thread and asks the model again,
//! bounded by [`MAX_STRUCTURED_RETRIES`].
//!
//! This is deliberately the opposite mechanism from
//! [`crate::agent::SimpleAgent::run_with_thread`]'s fresh-thread retry
//! independence: here the model must see its own failed attempt plus the
//! validation error to have any chance of fixing it.
//!
//! # Why not a [`crate::guardrail::Guardrail`]?
//!
//! Guardrails inspect the raw `ChatResponse` *before* text extraction
//! (`super::guardrails::check_post_llm`); [`parse_structured`] operates
//! on the already-extracted text content *after* that point. A guardrail
//! has no seam to reach into a parser that hasn't run yet, so this retry
//! is composed directly on [`crate::runtime::run_steps`] / `run_loop`
//! instead — the same primitives every other runtime driver is built on
//! — rather than off `Guardrail::post_llm`.

use super::{record_run_attribution, record_run_origin, run_loop, CompletionRecording, RunOptions};
use crate::agent::AgentContext;
use crate::error::Error;
use crate::ids::ThreadId;
use crate::llm::{Message, Role};
use crate::memory::Episode;
use crate::response::{parse_structured, KlieoResponse};

/// Maximum number of feedback-and-retry rounds after the first parse
/// failure.
///
/// Each round calls `run_loop` directly rather than issuing one bare
/// LLM call: `run_loop` iterates up to `opts.max_steps` LLM calls per
/// round whenever the model makes tool calls before producing a final
/// reply (the `StepDisposition::Continue` path in `runtime::step`), and
/// every one of those calls is itself subject to
/// `super::retry::MAX_LLM_RETRIES` transport retries. Worst case this
/// function issues `(MAX_STRUCTURED_RETRIES + 1) * opts.max_steps *
/// (MAX_LLM_RETRIES + 1)` LLM calls for one logical request — with the
/// crate defaults (`max_steps = 16`, `MAX_LLM_RETRIES = 3`) that is up
/// to `3 * 16 * 4 = 192` calls, not the 12 a naive one-call-per-round
/// reading would suggest. Kept small and separate from the
/// transport-retry budget on purpose — this is a real, user-visible cost
/// multiplier, not a free safety net.
pub const MAX_STRUCTURED_RETRIES: u32 = 2;

/// Run the agent loop and parse its final reply as `T`, retrying with
/// validation-error feedback (same thread) up to [`MAX_STRUCTURED_RETRIES`]
/// times on a parse failure.
///
/// Same calling contract as [`crate::runtime::run_steps`]: the caller
/// must append the user's message to `thread` before invoking — this
/// function only consumes / extends short-term memory, it does not seed
/// the initial turn.
///
/// # Episode bookkeeping
///
/// Performs the same entry preamble [`crate::runtime::run_steps`] does —
/// records `Episode::Started` plus run attribution/origin — directly,
/// exactly once, then drives every attempt (the first one AND every
/// retry) through `run_loop` with its `Episode::Completed` recording
/// suppressed (see `CompletionRecording` in `runtime::step`). This function records
/// `Episode::Completed` itself, exactly once, only at the point it is
/// about to return `Ok(T)`. An attempt whose reply [`parse_structured`]
/// rejects never reaches that point, so it is never counted as
/// completed; exhausting [`MAX_STRUCTURED_RETRIES`] returns `Err` with
/// no `Completed` recorded at all — mirroring
/// [`crate::runtime::run_steps`]'s existing "does not record
/// `Episode::Failed` on error, caller decides" convention (this function
/// adds no compensating `Failed` episode either).
///
/// [`crate::runtime::run_steps`] is deliberately never called here, not
/// even for the first attempt: it always drives `run_loop` with
/// completion recording ON, which would independently mark a
/// first-attempt reply that later fails to parse as "completed" in the
/// audit trail.
///
/// Compaction is forced off only for the retry sub-calls (the initial
/// attempt keeps whatever `opts.compaction` the caller configured) so the
/// feedback message injected before a retry cannot be summarised away
/// before the model reads it.
///
/// # Not provider-enforced
///
/// This is prompt-based structured output with client-side parse/retry,
/// not provider-enforced schema validation: `T::json_schema()` is never
/// wired into `ChatRequest.response_format` (today's
/// [`crate::runtime::run_steps`] / `build_request` expose no such hook —
/// adding one is out of scope here). Callers relying on this for hard
/// schema guarantees should expect the retry path to trigger more often
/// than it would against a provider that natively enforces the schema.
pub async fn run_structured<T: KlieoResponse>(
    ctx: &AgentContext,
    system_prompt: &str,
    thread: ThreadId,
    opts: RunOptions,
) -> Result<T, 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?;

    let content = run_loop(
        ctx,
        system_prompt,
        &thread,
        &opts,
        0,
        CompletionRecording::Suppress,
    )
    .await?;
    let mut last_err = match parse_structured::<T>(&content) {
        Ok(v) => return record_completed(ctx, v).await,
        Err(e) => e,
    };

    let retry_opts = opts.without_compaction();
    for _ in 0..MAX_STRUCTURED_RETRIES {
        let feedback = format!(
            "Your previous reply could not be parsed: {last_err}. \
             Reply again with ONLY valid JSON matching the required schema."
        );
        ctx.short_term
            .append(
                thread.clone(),
                Message {
                    role: Role::User,
                    content: feedback,
                    tool_calls: vec![],
                    tool_call_id: None,
                },
            )
            .await?;
        let content = run_loop(
            ctx,
            system_prompt,
            &thread,
            &retry_opts,
            0,
            CompletionRecording::Suppress,
        )
        .await?;
        match parse_structured::<T>(&content) {
            Ok(v) => return record_completed(ctx, v).await,
            Err(e) => last_err = e,
        }
    }
    Err(last_err)
}

/// Records the single terminal `Episode::Completed` for a
/// [`run_structured`] call, at the one point it is about to return
/// `Ok` — see "Episode bookkeeping" on [`run_structured`] for why every
/// individual attempt suppresses this instead of recording it inline.
async fn record_completed<T>(ctx: &AgentContext, value: T) -> Result<T, Error> {
    ctx.episodic.record(ctx.run_id, Episode::Completed).await?;
    Ok(value)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::summarize::SummarizeOptions;
    use crate::test_utils::{fake_context, FakeLlmClient, FakeLlmStep};
    use serde::Deserialize;
    use serde_json::json;
    use std::sync::Arc;

    #[derive(Debug, Deserialize, PartialEq)]
    struct Greeting {
        greeting: String,
    }

    impl KlieoResponse for Greeting {
        fn json_schema() -> serde_json::Value {
            json!({
                "type": "object",
                "properties": { "greeting": { "type": "string" } },
                "required": ["greeting"]
            })
        }
    }

    #[tokio::test]
    async fn malformed_first_reply_retries_with_feedback_then_succeeds() {
        let mut ctx = fake_context("structured-retry");
        ctx.llm = Arc::new(FakeLlmClient::new("fake").with_steps(vec![
            FakeLlmStep::Text("not json at all".into()),
            FakeLlmStep::Text(r#"{"greeting":"hi"}"#.into()),
        ]));
        let thread = ThreadId::new("t-structured");

        let out: Greeting = run_structured(&ctx, "sys", thread.clone(), RunOptions::default())
            .await
            .unwrap();
        assert_eq!(
            out,
            Greeting {
                greeting: "hi".into()
            }
        );

        let history = ctx.short_term.load(thread, 1024).await.unwrap();
        let feedback_msgs: Vec<_> = history
            .iter()
            .filter(|m| m.content.contains("could not be parsed"))
            .collect();
        assert_eq!(
            feedback_msgs.len(),
            1,
            "exactly one feedback message injected for the one bad reply"
        );
    }

    #[tokio::test]
    async fn retry_sub_call_disables_compaction_so_the_feedback_survives() {
        let mut ctx = fake_context("structured-compaction-off");
        // Exactly two scripted replies: the malformed first attempt, and
        // the well-formed retry reply. If the retry sub-call used the
        // caller's `opts` (compaction ON) instead of `retry_opts`
        // (compaction forced off), `maybe_compact` would fire before the
        // retry's own completion call — the by-then-non-empty history
        // (bad reply + injected feedback) trivially exceeds the
        // pathologically low `trigger_token_budget` below. That
        // summarization is itself an LLM call, so it would consume this
        // fake client's one remaining scripted step instead of the real
        // retry reply, and the run would fail with a "script exhausted"
        // error instead of succeeding.
        ctx.llm = Arc::new(FakeLlmClient::new("fake").with_steps(vec![
            FakeLlmStep::Text("not json at all".into()),
            FakeLlmStep::Text(r#"{"greeting":"hi"}"#.into()),
        ]));
        let thread = ThreadId::new("t-compaction-off");

        // `max_history_tokens` is pinned to the window this pathological
        // trigger implies, because `maybe_compact` rejects a pair where the
        // request path could restore more history than compaction can see.
        // The fake LLM ignores request content, so the tiny load budget does
        // not weaken what this test asserts: that only two scripted steps are
        // consumed, i.e. no third call for a summarization.
        let compaction = SummarizeOptions {
            trigger_token_budget: 1,
            keep_recent_messages: 0,
            ..SummarizeOptions::default()
        };
        let history_budget = crate::summarize::visible_window_tokens(&compaction);
        let opts = RunOptions::default()
            .with_compaction(compaction)
            .with_max_history_tokens(history_budget);

        let out: Greeting = run_structured(&ctx, "sys", thread, opts).await.unwrap();
        assert_eq!(
            out,
            Greeting {
                greeting: "hi".into()
            },
            "retry must reach the real scripted reply, not a summarizer call \
             consuming it first"
        );
    }

    #[tokio::test]
    async fn valid_first_reply_needs_no_retry() {
        let mut ctx = fake_context("structured-no-retry");
        ctx.llm = Arc::new(
            FakeLlmClient::new("fake")
                .with_steps(vec![FakeLlmStep::Text(r#"{"greeting":"hi"}"#.into())]),
        );
        let thread = ThreadId::new("t-clean");

        let out: Greeting = run_structured(&ctx, "sys", thread.clone(), RunOptions::default())
            .await
            .unwrap();
        assert_eq!(
            out,
            Greeting {
                greeting: "hi".into()
            }
        );

        let history = ctx.short_term.load(thread, 1024).await.unwrap();
        assert_eq!(
            history.len(),
            1,
            "only the assistant reply — no feedback message"
        );
    }

    #[tokio::test]
    async fn exhausting_all_retries_returns_the_last_parse_error() {
        let mut ctx = fake_context("structured-exhausted");
        // Each fixture fails at a distinct serde_json position (verified
        // empirically) so the assertion below can pin the LAST error
        // specifically — a regression that kept the FIRST-captured error
        // instead of overwriting `last_err` on each retry would return
        // "line 1 col 1" here and fail this test.
        ctx.llm = Arc::new(FakeLlmClient::new("fake").with_steps(vec![
            FakeLlmStep::Text("bad".into()),                 // line 1 col 1
            FakeLlmStep::Text("{\"greeting\": bad}".into()), // line 1 col 14
            FakeLlmStep::Text("{\"greeting\": \"ok\" extra}".into()), // line 1 col 19
        ]));
        let thread = ThreadId::new("t-exhausted");

        let err = run_structured::<Greeting>(&ctx, "sys", thread, RunOptions::default())
            .await
            .unwrap_err();
        assert!(matches!(err, Error::BadResponse(_)));
        assert!(
            err.to_string().contains("line 1 col 19"),
            "expected the LAST fixture's parse error (line 1 col 19), got: {err}"
        );
        assert_eq!(
            MAX_STRUCTURED_RETRIES, 2,
            "three scripted replies must exactly exhaust a budget of 1 initial + 2 retries"
        );
    }

    #[tokio::test]
    async fn retry_then_succeed_records_exactly_one_completed_episode() {
        let mut ctx = fake_context("structured-episode-success");
        ctx.llm = Arc::new(FakeLlmClient::new("fake").with_steps(vec![
            FakeLlmStep::Text("not json at all".into()),
            FakeLlmStep::Text(r#"{"greeting":"hi"}"#.into()),
        ]));
        let thread = ThreadId::new("t-episode-success");

        let _out: Greeting = run_structured(&ctx, "sys", thread, RunOptions::default())
            .await
            .unwrap();

        let episodes = ctx.episodic.replay(ctx.run_id).await.unwrap();
        let completed_count = episodes
            .iter()
            .filter(|e| matches!(e, Episode::Completed))
            .count();
        assert_eq!(
            completed_count, 1,
            "one logical call that retried once then succeeded must record \
             Episode::Completed exactly once, not once per attempt"
        );
    }

    #[tokio::test]
    async fn exhausting_all_retries_records_zero_completed_episodes() {
        let mut ctx = fake_context("structured-episode-exhausted");
        // Same distinct-error fixtures as
        // `exhausting_all_retries_returns_the_last_parse_error` — three
        // rejected attempts, none of which should ever reach the single
        // `record_completed` call site.
        ctx.llm = Arc::new(FakeLlmClient::new("fake").with_steps(vec![
            FakeLlmStep::Text("bad".into()),
            FakeLlmStep::Text("{\"greeting\": bad}".into()),
            FakeLlmStep::Text("{\"greeting\": \"ok\" extra}".into()),
        ]));
        let thread = ThreadId::new("t-episode-exhausted");

        let result = run_structured::<Greeting>(&ctx, "sys", thread, RunOptions::default()).await;
        assert!(result.is_err(), "all three scripted replies are malformed");

        let episodes = ctx.episodic.replay(ctx.run_id).await.unwrap();
        let completed_count = episodes
            .iter()
            .filter(|e| matches!(e, Episode::Completed))
            .count();
        assert_eq!(
            completed_count, 0,
            "a call that exhausts all retries and returns Err must record no \
             Episode::Completed at all — the audit trail must not claim \
             success for a structurally failed call"
        );
    }
}