polyc_llm/

preflight.rs

//! Startup preflight: verify a backend actually delivers native tool calling and
//! schema-constrained structured output, against the live endpoint.
//!
//! A self-hosted OpenAI-compatible server only delivers these if its serving
//! runtime is configured for them (a tool-call parser + a grammar/constrained
//! decoding backend). When it isn't, tool calls arrive as plain text and the
//! response schema is ignored — and the harness's tool loop silently no-ops.
//! This module turns that silent degradation into an explicit, logged (or fatal)
//! signal by actually exercising both capabilities once at startup.
//!
//! Each probe yields a tri-state [`ProbeOutcome`]: `Supported` (verified),
//! `Unsupported` (the backend ran but didn't honor the constraint), or `Errored`
//! (the call itself failed — transport, cold start, timeout). The distinction
//! matters at the call site: a definitive `Unsupported` is a capability verdict
//! worth failing startup over, whereas an `Errored` probe is "couldn't verify"
//! and must not crash-loop a backend that is merely warming up.

use futures::StreamExt;

use crate::{Chunk, CompletionRequest, DynProvider, JsonSchema, Message, ToolChoice, ToolSpec};

/// The result of one capability probe.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProbeOutcome {
    /// The backend ran the probe and honored the constraint.
    Supported,
    /// The backend ran the probe but did not honor the constraint (no tool call
    /// emitted under `tool_choice = Required`, or output that violated the schema).
    /// A definitive capability verdict.
    Unsupported,
    /// The probe call itself failed (transport error, cold start, timeout). NOT a
    /// capability verdict — the backend may be fine once warm.
    Errored,
}

/// What the live endpoint actually did when probed.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PreflightReport {
    /// Outcome of the forced-tool-call probe.
    pub native_tool_calling: ProbeOutcome,
    /// Outcome of the schema-constrained structured-output probe.
    pub structured_output: ProbeOutcome,
    /// Human-readable notes (parse failures, transport errors) for logging.
    pub notes: Vec<String>,
}

impl PreflightReport {
    /// Both capabilities verified — the backend is a full peer.
    #[must_use]
    pub fn ok(&self) -> bool {
        self.native_tool_calling == ProbeOutcome::Supported
            && self.structured_output == ProbeOutcome::Supported
    }

    /// A probe gave a definitive *negative* capability verdict (ran, didn't honor
    /// the constraint). This — not a transport error — is what justifies failing
    /// startup under a strict policy.
    #[must_use]
    pub fn has_unsupported(&self) -> bool {
        self.native_tool_calling == ProbeOutcome::Unsupported
            || self.structured_output == ProbeOutcome::Unsupported
    }
}

/// The schema used by both probes: a one-field object the smallest model can
/// satisfy, with `additionalProperties: false` so a grammar backend has
/// something to enforce.
fn probe_schema() -> serde_json::Value {
    serde_json::json!({
        "type": "object",
        "properties": { "ok": { "type": "boolean" } },
        "required": ["ok"],
        "additionalProperties": false,
    })
}

/// Probe a backend's real capabilities with two live calls.
///
/// One forces a tool call; the other requests structured output and checks the
/// reply against the full probe schema. Never panics.
pub async fn preflight(provider: &DynProvider, model: &str) -> PreflightReport {
    let mut notes = Vec::new();
    let native_tool_calling = match check_tool_call(provider, model).await {
        Ok(true) => ProbeOutcome::Supported,
        Ok(false) => {
            notes.push(
                "tool-call probe: no tool call emitted under tool_choice=required".to_owned(),
            );
            ProbeOutcome::Unsupported
        }
        Err(e) => {
            notes.push(format!("tool-call probe errored: {e}"));
            ProbeOutcome::Errored
        }
    };
    let structured_output = match check_structured_output(provider, model).await {
        Ok(true) => ProbeOutcome::Supported,
        Ok(false) => {
            notes.push("structured-output probe: reply did not conform to the schema".to_owned());
            ProbeOutcome::Unsupported
        }
        Err(e) => {
            notes.push(format!("structured-output probe errored: {e}"));
            ProbeOutcome::Errored
        }
    };
    PreflightReport {
        native_tool_calling,
        structured_output,
        notes,
    }
}

/// Force a tool call and report whether the backend emitted one.
async fn check_tool_call(provider: &DynProvider, model: &str) -> Result<bool, String> {
    let mut req = CompletionRequest::new(model);
    req.max_tokens = Some(256);
    req.tools = vec![ToolSpec::new(
        "preflight_probe",
        "A connectivity probe. Call it with ok=true.",
        probe_schema(),
    )];
    req.tool_choice = ToolChoice::Required;
    req.messages = vec![Message::user(
        "Call the preflight_probe tool with ok set to true.",
    )];

    let mut stream = provider.complete(req).await.map_err(|e| e.to_string())?;
    while let Some(item) = stream.next().await {
        let chunk = item.map_err(|e| e.to_string())?;
        if matches!(chunk, Chunk::ToolCallStart { .. }) {
            return Ok(true);
        }
    }
    Ok(false)
}

/// Request structured output and report whether the reply conforms to the schema.
///
/// Checks the reply against the FULL probe schema (exact shape — not merely "has
/// the key"), so a backend that returns prose, extra keys, or a wrong-typed field
/// fails. This verifies the observable property callers actually depend on:
/// schema-conforming structured output when requested. It is deliberately a
/// black-box check — distinguishing true token-level grammar enforcement from a
/// model that simply complied is not reliably observable from outside (real
/// backends vary, and a prose-pulling "adversarial" prompt false-fails endpoints
/// that serve structured output fine in practice), so the probe asserts
/// conformance-on-request rather than enforcement-against-any-prompt.
async fn check_structured_output(provider: &DynProvider, model: &str) -> Result<bool, String> {
    let mut req = CompletionRequest::new(model);
    req.max_tokens = Some(256);
    req.response_format = Some(JsonSchema(probe_schema()));
    req.messages = vec![Message::user(
        "Reply with a JSON object that sets \"ok\" to true.",
    )];

    let mut stream = provider.complete(req).await.map_err(|e| e.to_string())?;
    let mut text = String::new();
    while let Some(item) = stream.next().await {
        if let Chunk::TextDelta(t) = item.map_err(|e| e.to_string())? {
            text.push_str(&t);
        }
    }
    Ok(json_matches_probe(&text))
}

/// True if `text` (possibly wrapped in a markdown code fence) is a JSON object
/// that *conforms to the probe schema*: exactly one key `ok`, with a boolean
/// value, and nothing else.
///
/// Checking the full shape — not merely that `ok` is present — is what makes this
/// a test of schema-*constrained* decoding: a server that ignored
/// `additionalProperties: false` (extra keys), the `ok` type, or returned prose
/// fails, exactly as it should. Lenient on a surrounding markdown fence only.
fn json_matches_probe(text: &str) -> bool {
    let trimmed = strip_code_fence(text.trim());
    let Ok(serde_json::Value::Object(map)) = serde_json::from_str::<serde_json::Value>(trimmed)
    else {
        return false;
    };
    map.len() == 1 && matches!(map.get("ok"), Some(serde_json::Value::Bool(_)))
}

/// Strip a leading/trailing ```` ```json ```` … ```` ``` ```` fence if present.
///
/// Public so callers parsing a JSON-in-a-fence model reply (e.g. e2e tests) share
/// one fence-handling implementation instead of re-deriving it.
#[must_use]
pub fn strip_code_fence(s: &str) -> &str {
    let s = s
        .strip_prefix("```json")
        .or_else(|| s.strip_prefix("```"))
        .unwrap_or(s);
    s.trim().trim_end_matches("```").trim()
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]
    use super::*;
    use crate::{StopReason, Usage, error::DummyError, into_dyn};
    use futures::stream;

    #[test]
    fn json_matches_probe_accepts_plain_and_fenced() {
        assert!(json_matches_probe(r#"{"ok": true}"#));
        assert!(json_matches_probe("```json\n{\"ok\": false}\n```"));
        assert!(json_matches_probe("```\n{\"ok\": true}\n```"));
    }

    #[test]
    fn json_matches_probe_rejects_prose_and_wrong_shape() {
        assert!(!json_matches_probe("The sky is blue."));
        assert!(!json_matches_probe(r#"{"status": "fine"}"#));
        assert!(!json_matches_probe(""));
        // Schema-constrained: extra keys (additionalProperties) and a wrong-typed
        // `ok` must fail — a server that merely returned JSON-with-`ok` is not
        // proof of constrained decoding.
        assert!(!json_matches_probe(r#"{"ok": true, "extra": 1}"#));
        assert!(!json_matches_probe(r#"{"ok": "yes"}"#));
    }

    /// Fully capable: emits a tool call and schema-conforming JSON.
    struct CapableProvider;
    #[async_trait::async_trait]
    impl crate::LlmProvider for CapableProvider {
        type Error = DummyError;
        async fn complete(
            &self,
            req: CompletionRequest,
        ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
        {
            let chunks = if req.tool_choice == ToolChoice::Required {
                vec![
                    Ok(Chunk::tool_call_start("c1", "preflight_probe")),
                    Ok(Chunk::tool_call_args_delta("c1", "{\"ok\":true}")),
                    Ok(Chunk::tool_call_end("c1")),
                    Ok(Chunk::Stop(StopReason::ToolUse)),
                ]
            } else {
                vec![
                    Ok(Chunk::text_delta("{\"ok\": true}")),
                    Ok(Chunk::Usage(Usage {
                        input_tokens: 1,
                        output_tokens: 1,
                        ..Default::default()
                    })),
                    Ok(Chunk::Stop(StopReason::EndTurn)),
                ]
            };
            Ok(stream::iter(chunks).boxed())
        }
    }

    /// Ran fine but honored neither constraint: no tool call, prose reply.
    struct DegradedProvider;
    #[async_trait::async_trait]
    impl crate::LlmProvider for DegradedProvider {
        type Error = DummyError;
        async fn complete(
            &self,
            _req: CompletionRequest,
        ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
        {
            Ok(stream::iter(vec![
                Ok(Chunk::text_delta("The sky is blue.")),
                Ok(Chunk::Stop(StopReason::EndTurn)),
            ])
            .boxed())
        }
    }

    /// Unreachable: every call fails before streaming (transport/cold start).
    struct ErroringProvider;
    #[async_trait::async_trait]
    impl crate::LlmProvider for ErroringProvider {
        type Error = DummyError;
        async fn complete(
            &self,
            _req: CompletionRequest,
        ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
        {
            Err(DummyError::Other("connection refused".to_owned()))
        }
    }

    #[tokio::test]
    async fn preflight_passes_a_capable_backend() {
        let p = into_dyn(CapableProvider);
        let report = preflight(&*p, "m").await;
        assert!(report.ok(), "{report:?}");
        assert!(!report.has_unsupported());
    }

    #[tokio::test]
    async fn preflight_flags_a_degraded_backend_as_unsupported() {
        let p = into_dyn(DegradedProvider);
        let report = preflight(&*p, "m").await;
        assert!(!report.ok());
        assert!(
            report.has_unsupported(),
            "degraded backend is a capability verdict"
        );
        assert_eq!(report.native_tool_calling, ProbeOutcome::Unsupported);
        assert_eq!(report.structured_output, ProbeOutcome::Unsupported);
    }

    #[tokio::test]
    async fn preflight_marks_transport_failure_errored_not_unsupported() {
        // A cold/unreachable backend must NOT read as a definitive capability
        // verdict — otherwise strict mode crash-loops a backend that's merely
        // warming up.
        let p = into_dyn(ErroringProvider);
        let report = preflight(&*p, "m").await;
        assert!(!report.ok());
        assert!(
            !report.has_unsupported(),
            "transport error is not 'unsupported'"
        );
        assert_eq!(report.native_tool_calling, ProbeOutcome::Errored);
        assert_eq!(report.structured_output, ProbeOutcome::Errored);
    }
}