klieo-core 0.6.0

//! Shared post-step skeleton for both [`super::run_steps`] (blocking)
//! and [`super::streaming::run_steps_streaming`] (streaming).
//!
//! Both loops differ only in how they produce a per-step assistant
//! message:
//!
//! - blocking: `complete_with_retry` returns one [`ChatResponse`].
//! - streaming: `forward_chunks` reconstructs `(Message, FinishReason)`
//!   from a chunk pump (usage unavailable today, so passed as zero).
//!
//! Everything after that point — episodic `LlmCall` record, short-term
//! append, the four-arm `FinishReason` match — was previously copied
//! twice. This module owns that shared tail; the two loops orchestrate
//! their provider-specific call and then hand the result here.

use crate::agent::AgentContext;
use crate::error::{Error, LlmError};
use crate::ids::ThreadId;
use crate::llm::{FinishReason, Message, Usage};
use crate::memory::Episode;
use tracing::{debug, info};

use super::dispatch::dispatch_tool_calls;

/// Provider-call result handed to [`record_and_dispatch`]. Holds
/// exactly what the post-step skeleton needs: the assistant message,
/// the terminal `finish_reason`, usage (zero for streaming until
/// usage plumbing lands), and observed provider latency.
pub(crate) struct StepOutcome {
    pub message: Message,
    pub finish_reason: FinishReason,
    pub usage: Usage,
    pub latency_ms: u32,
}

/// Tells the caller loop whether to terminate or continue after the
/// post-step skeleton has run.
///
/// - `Done(content)` — loop should exit. `content` is the final
///   assistant text. The streaming caller discards it because chunks
///   were already forwarded; the blocking caller returns it.
/// - `Continue` — a tool-call cycle was just dispatched; the loop
///   should build the next request and call the provider again.
pub(crate) enum StepDisposition {
    Done(String),
    Continue,
}

/// Persist one step's outcome and decide whether to continue. Shared
/// between blocking and streaming loops.
///
/// Steps performed (identical to the inline code both loops previously
/// carried):
///
/// 1. Record `Episode::LlmCall { tokens, latency_ms }`.
/// 2. Append the assistant message to short-term memory.
/// 3. Dispatch on `finish_reason`:
///    - `Stop` / `Length` → record `Completed`, return `Done`.
///    - `ToolCalls` → dispatch tools via [`dispatch_tool_calls`] under
///      `dispatch_label`, return `Continue`.
///    - `ContentFilter` → return `Err(Llm(BadRequest))`.
///    - `Error` → return `Err(Llm(Server))`.
///
/// `dispatch_label` is the tracing-mode string the dispatcher tags its
/// spans with: `"blocking"` for [`super::run_steps`], `"streaming"`
/// for the streaming driver. The label flows through unchanged so
/// observability stays distinguishable between paths.
pub(crate) async fn record_and_dispatch(
    ctx: &AgentContext,
    thread: &ThreadId,
    step: u32,
    outcome: StepOutcome,
    dispatch_label: &'static str,
) -> Result<StepDisposition, Error> {
    let StepOutcome {
        message,
        finish_reason,
        usage,
        latency_ms,
    } = outcome;

    ctx.episodic
        .record(
            ctx.run_id,
            Episode::LlmCall {
                tokens: usage.prompt_tokens + usage.completion_tokens,
                latency_ms,
            },
        )
        .await?;

    ctx.short_term
        .append(thread.clone(), message.clone())
        .await?;

    match finish_reason {
        FinishReason::Stop | FinishReason::Length => {
            info!(
                step,
                content_len = message.content.len(),
                mode = dispatch_label,
                "agent finished"
            );
            ctx.episodic.record(ctx.run_id, Episode::Completed).await?;
            Ok(StepDisposition::Done(message.content))
        }
        FinishReason::ToolCalls => {
            debug!(
                step,
                count = message.tool_calls.len(),
                mode = dispatch_label,
                "dispatching tools"
            );
            dispatch_tool_calls(ctx, thread, &message.tool_calls, dispatch_label).await?;
            Ok(StepDisposition::Continue)
        }
        FinishReason::ContentFilter => Err(Error::Llm(LlmError::BadRequest(
            "content-filter triggered".into(),
        ))),
        FinishReason::Error => Err(Error::Llm(LlmError::Server(
            "provider returned Error finish_reason".into(),
        ))),
    }
}