klieo-core 3.0.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,
    /// Provider name for this call, derived from the active
    /// [`crate::llm::LlmClient::name`]; `None` when the name has no
    /// `provider:model` shape. Recorded so a live run prices without a sink.
    pub provider: Option<String>,
    /// None when the client name lacks a `provider:model` shape.
    pub model: Option<String>,
}

/// Split an [`crate::llm::LlmClient::name`] of the conventional
/// `"provider:model"` shape into its provider and model id (on the FIRST colon,
/// so a colon inside the model tag survives). Returns `(None, None)` when the
/// name lacks that shape, so cost projection degrades to unpriced rather than
/// mispricing. Price-key aliasing (e.g. `gemini`→`google`) is the price table's
/// concern, not this layer's.
pub(crate) fn split_provider_model(name: &str) -> (Option<String>, Option<String>) {
    match name.split_once(':') {
        Some((provider, model)) if !provider.is_empty() && !model.is_empty() => {
            (Some(provider.to_string()), Some(model.to_string()))
        }
        _ => (None, None),
    }
}

/// 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,
        provider,
        model,
    } = outcome;

    ctx.episodic
        .record(
            ctx.run_id,
            Episode::LlmCall {
                tokens: usage.prompt_tokens + usage.completion_tokens,
                latency_ms,
                provider,
                model,
                prompt_tokens: Some(usage.prompt_tokens),
                completion_tokens: Some(usage.completion_tokens),
            },
        )
        .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(),
        ))),
    }
}

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

    #[test]
    fn splits_provider_and_model_from_a_conventional_name() {
        assert_eq!(
            split_provider_model("openai:gpt-4o"),
            (Some("openai".into()), Some("gpt-4o".into()))
        );
    }

    #[test]
    fn keeps_a_colon_in_the_model_half() {
        // Ollama names carry a colon inside the model tag.
        assert_eq!(
            split_provider_model("ollama:qwen2.5:14b"),
            (Some("ollama".into()), Some("qwen2.5:14b".into()))
        );
    }

    #[test]
    fn keeps_the_raw_provider_prefix_aliasing_is_the_price_tables_job() {
        assert_eq!(
            split_provider_model("gemini:gemini-2.0-flash"),
            (Some("gemini".into()), Some("gemini-2.0-flash".into()))
        );
    }

    #[test]
    fn yields_none_when_the_name_has_no_provider_model_shape() {
        assert_eq!(split_provider_model("fake"), (None, None));
        assert_eq!(split_provider_model(":m"), (None, None));
        assert_eq!(split_provider_model("p:"), (None, None));
    }
}