rig-core 0.45.0

//! Execution-local correlation and facts from provider request boundaries.
//!
//! ```
//! use rig_core::observe::scrub_diagnostic;
//!
//! assert_eq!(scrub_diagnostic("Bearer secret", &[]), "[redacted]");
//! ```

use std::sync::{Arc, Mutex};

use serde::{Deserialize, Serialize};

use super::{Action, Emitter, Observation, Stage, Subject, Witness};

/// Bounds diagnostic text and removes control characters. Known credentials
/// or credential markers redact the whole message; oversized messages are
/// replaced rather than truncated to a potentially sensitive prefix.
pub fn scrub_diagnostic(value: &str, secrets: &[String]) -> String {
    scrub::text(value, secrets)
}

/// The secrets a URL carries for diagnostic redaction: userinfo and known
/// credential query parameters, origin-form request URIs included. The
/// returned values are secrets: keep them runtime-only and never write
/// them into a diagnostic or an artifact.
pub fn diagnostic_url_secrets(url: &str) -> Vec<String> {
    scrub::url_secrets(url)
}

/// One fact about a provider operation, independently of its effect record.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AdapterObservation {
    /// Caller-owned logical operation identity within this execution.
    pub operation: String,
    /// One-based HTTP send ordinal; absent only when correlation is exhausted.
    pub attempt: Option<u64>,
    /// One-based host dispatch attempt, when explicitly supplied by the host.
    /// Independent of the HTTP send ordinal: a dispatch may send more than once.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub host_attempt: Option<std::num::NonZeroU64>,
    /// What the owning boundary observed.
    pub event: AdapterEvent,
    /// Scrubbed volatile diagnostics, excluded from semantic comparison.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub analysis: Option<AdapterAnalysis>,
}

/// Metadata from the provider boundary. Request/response bodies are not included.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "event", rename_all = "snake_case")]
pub enum AdapterEvent {
    /// Fields actually present in a provider payload; absent fields are not updates.
    Provider {
        /// The provider's verdict and selected model, independent of transport closure.
        verdict: AdapterVerdict,
    },
    /// A provider error envelope, including envelopes carried under HTTP 200.
    ErrorEnvelope {
        /// Scrubbed original envelope fields, not an inferred HTTP status.
        error: AdapterErrorEnvelope,
    },
    /// The attempt's normalized usage, observed once when its reply ends.
    Usage {
        /// The counters of the response's usage.
        usage: AdapterUsage,
    },
    /// A request is about to be sent.
    Started {
        /// HTTP method.
        method: String,
        /// Provider-declared route template, never a credential-bearing URI.
        route: String,
    },
    /// The transport returned response headers.
    Response {
        /// Actual HTTP response status, including successful responses.
        status: u16,
    },
    /// The body reached EOF, independently of any provider terminal verdict.
    TransportEof {
        /// Number of complete frames, excluding provider-recognized analysis-only frames.
        after: usize,
        /// Undelimited SSE bytes remaining at EOF; zero means no partial frame.
        partial_bytes: usize,
    },
    /// The owning request boundary closed this attempt.
    Finished {
        /// Whether the attempt completed, failed, or was dropped.
        ending: AdapterEnding,
    },
    /// A complete transport frame could not be decoded; the driver may continue.
    Corrupt {
        /// One-based ordinal counting known, unknown and corrupt frames,
        /// excluding provider-recognized analysis-only frames.
        frame: usize,
    },
    /// The context cannot assign another unique send ordinal.
    IdentityExhausted,
}

/// Sparse provider metadata: a present field supersedes its previous value.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AdapterVerdict {
    /// Provider finish code, including unknown codes after scrubbing.
    pub finish_reason: Option<String>,
    /// Provider prompt-block/refusal reason, when reported.
    pub block_reason: Option<String>,
    /// Provider detail accompanying a finish/block reason.
    pub detail: Option<String>,
    /// Actual provider model/version, not the requested alias.
    pub model: Option<String>,
}

/// Analysis-only response metadata. Never use identifiers as operation join keys.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AdapterAnalysis {
    /// Provider response identifier, scrubbed and bounded.
    pub response_id: Option<String>,
    /// Allowlisted response headers. None means the transport supplied no map;
    /// an empty map means it supplied no allowlisted values.
    pub headers: Option<std::collections::BTreeMap<String, String>>,
}

/// Original provider error-envelope fields, separate from Rig retry classification.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AdapterErrorEnvelope {
    /// Provider code in its scalar wire spelling; unknown codes stay observable.
    pub code: Option<String>,
    /// Provider status name, when present.
    pub status: Option<String>,
    /// Bounded, scrubbed provider message, when present.
    pub message: Option<String>,
}

/// Error-envelope projection with optional code, message, and `type`/`status`.
/// Codes remain JSON values until emission validates their scalar shape.
#[derive(Default)]
pub struct ObservedError {
    pub code: Option<serde_json::Value>,
    pub kind: Option<String>,
    pub message: Option<String>,
}

impl ObservedError {
    /// The error envelope a provider document carries, located as its
    /// [`crate::provider_response::ProviderResponseError`] reads it: an
    /// error event's own `type` is its tag, not a kind, and an `error`
    /// string is the message alone.
    pub fn of(doc: &serde_json::Value) -> Option<Self> {
        let (error, nested) = crate::provider_response::located_error(doc)?;
        let text = |key: &str| error.get(key).and_then(serde_json::Value::as_str);
        let kind = nested.then(|| text("type").or_else(|| text("status")));
        Some(Self {
            code: error.get("code").filter(|code| !code.is_null()).cloned(),
            kind: kind.flatten().map(str::to_owned),
            message: error
                .as_str()
                .or_else(|| text("message"))
                .map(str::to_owned),
        })
    }

    /// Emit this envelope as the attempt's error-envelope fact, scrubbed.
    pub fn emit(self, sink: &mut ObservationSink<'_>) {
        let code = self.code.map(|code| match code {
            serde_json::Value::String(code) => sink.scrub(&code),
            serde_json::Value::Number(code) => code.to_string(),
            _ => "[invalid]".to_owned(),
        });
        sink.error(AdapterErrorEnvelope {
            code,
            status: self.kind.map(|value| sink.scrub(&value)),
            message: self.message.map(|value| sink.scrub(&value)),
        });
    }
}

/// The counters of the [`Usage`](crate::completion::Usage) the decoder
/// normalized into the reply's end, observed once when the reply ends; an
/// attempt that fails before its end observes none. Unreported is `None`.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AdapterUsage {
    /// Every input token: uncached, cached, cache writes and hosted-tool prompt.
    pub input_tokens: Option<u64>,
    /// Every output token, reasoning included.
    pub output_tokens: Option<u64>,
    /// Always `None`; kept for serde compatibility. Read the response's
    /// [`Usage::total_tokens`](crate::completion::Usage::total_tokens).
    pub total_tokens: Option<u64>,
    /// The part of the input read from a provider-managed cache.
    pub cached_input_tokens: Option<u64>,
    /// The part of the output spent on internal reasoning.
    pub reasoning_tokens: Option<u64>,
    /// The part of the input a provider's hosted tools added to the prompt.
    pub tool_input_tokens: Option<u64>,
}

impl From<&crate::completion::Usage> for AdapterUsage {
    fn from(usage: &crate::completion::Usage) -> Self {
        // Exhaustive: a new `Usage` counter must be mapped or ignored here.
        let crate::completion::Usage {
            input_tokens,
            output_tokens,
            total_tokens: _,
            cached_input_tokens,
            cache_creation_input_tokens: _,
            tool_use_prompt_tokens,
            reasoning_tokens,
            cost: _,
        } = *usage;
        Self {
            input_tokens,
            output_tokens,
            total_tokens: None,
            cached_input_tokens,
            reasoning_tokens,
            tool_input_tokens: tool_use_prompt_tokens,
        }
    }
}

/// Boundary identified by the adapter's original typed error, not its message.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AdapterErrorBoundary {
    /// Request construction or validation failed.
    Request,
    /// The provider returned a failure response or error envelope.
    ProviderResponse,
    /// Response decoding or validation failed.
    Decode,
    /// A typed transport termination was reported.
    Transport,
    /// An erased client error does not establish a more specific boundary.
    #[default]
    Unknown,
}

impl AdapterErrorBoundary {
    pub(crate) fn from_http(error: &crate::http_client::Error) -> Self {
        use crate::http_client::Error as H;
        match error {
            H::Protocol(_) | H::InvalidHeaderValue(_) | H::NoHeaders | H::Unbuilt(_) => {
                Self::Request
            }
            H::InvalidContentType(_) => Self::Decode,
            H::StreamEnded => Self::Transport,
            H::InvalidStatusCodeWithDetails { .. } => Self::ProviderResponse,
            _ => Self::Unknown,
        }
    }
}

/// How a provider attempt closed, distinct from the run's ending.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "ending", rename_all = "snake_case")]
pub enum AdapterEnding {
    /// The request produced a decoded response.
    Decoded,
    /// The request failed; provider payloads remain in the existing error report.
    Error {
        /// Known boundary before error conversion erases transport subtypes.
        #[serde(default)]
        boundary: AdapterErrorBoundary,
        /// Stable Rig error classification.
        kind: String,
        /// HTTP status reported by the failure, when available.
        status: Option<u16>,
        /// Retryability under the existing policy, not a decision to retry.
        retryable: bool,
    },
    /// A provider terminal record was decoded from the stream.
    Terminal,
    /// Transport EOF without a provider terminal record.
    Eof {
        /// Number of complete frames, excluding provider-recognized analysis-only frames.
        after: usize,
    },
    /// EOF left an undelimited SSE event; bytes are never retained in the fact.
    PartialFrame {
        /// Raw bytes since the last blank event delimiter, saturated at usize::MAX.
        byte_count: usize,
        /// Complete frames previously delivered, excluding analysis-only frames.
        after: usize,
    },
    /// The future or stream was dropped before its boundary closed.
    Dropped,
}

/// A caller-supplied witness and logical operation identity.
///
/// Clone this context for retries of the same operation. Create a new context
/// with a distinct identity for a different call, including identical parallel
/// requests. Identity is execution-local; it does not assert replay equivalence.
/// This handle is runtime-only and must never be serialized into provider data.
#[derive(Clone)]
pub struct AdapterContext {
    inner: Arc<AdapterContextInner>,
}

struct AdapterContextInner {
    sink: Arc<dyn Witness + Send + Sync>,
    subject: Subject,
    operation: String,
    next: Arc<Mutex<Option<u64>>>,
    host_attempt: Option<std::num::NonZeroU64>,
}

impl std::fmt::Debug for AdapterContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // Caller-supplied identity can contain sensitive data; do not log it.
        f.debug_struct("AdapterContext").finish_non_exhaustive()
    }
}

impl AdapterContext {
    /// Caller-owned identity of this logical operation within the execution.
    pub fn operation(&self) -> &str {
        &self.inner.operation
    }

    /// Bind a logical operation to its witness and existing subject.
    /// The caller must use a non-sensitive, execution-unique operation identity.
    pub fn new(
        sink: Arc<dyn Witness + Send + Sync>,
        subject: Subject,
        operation: impl Into<String>,
    ) -> Self {
        Self {
            inner: Arc::new(AdapterContextInner {
                sink,
                subject,
                operation: operation.into(),
                next: Arc::new(Mutex::new(Some(1))),
                host_attempt: None,
            }),
        }
    }

    /// Bind an explicit host attempt to its current dispatch subject.
    ///
    /// The logical operation, witness and HTTP send counter remain shared.
    /// Existing clones retain their own subject and host ordinal, so an older
    /// in-flight attempt cannot be relabeled by a concurrent retry. The host
    /// owns ordinal allocation and must start a new context for a changed
    /// logical request rather than reusing this operation identity.
    pub fn for_host_attempt(&self, subject: Subject, attempt: std::num::NonZeroU64) -> Self {
        Self {
            inner: Arc::new(AdapterContextInner {
                sink: self.inner.sink.clone(),
                subject,
                operation: self.inner.operation.clone(),
                next: self.inner.next.clone(),
                host_attempt: Some(attempt),
            }),
        }
    }

    /// Begins a send and captures request credentials for diagnostic scrubbing.
    /// `route` must be a credential-free provider template without base-URL
    /// prefixes or query data. Returns `None` when attempt IDs are exhausted.
    pub(crate) fn attempt_for<B>(
        &self,
        request: &http::Request<B>,
        route: &str,
    ) -> Option<AdapterAttempt> {
        let mut attempt = self.begin(request.method(), route)?;
        attempt.secrets = scrub::request_secrets(request);
        Some(attempt)
    }

    fn emit(&self, attempt: Option<u64>, event: AdapterEvent) {
        self.emit_with_analysis(attempt, event, None);
    }

    fn emit_with_analysis(
        &self,
        attempt: Option<u64>,
        event: AdapterEvent,
        analysis: Option<AdapterAnalysis>,
    ) {
        self.inner.sink.observe(Observation::new(
            self.inner.subject.clone(),
            Stage::Handler,
            Emitter::named("rig-core/adapter"),
            Action::Adapter {
                observation: AdapterObservation {
                    operation: self.inner.operation.clone(),
                    attempt,
                    host_attempt: self.inner.host_attempt,
                    event,
                    analysis,
                },
            },
        ));
    }

    /// Allocates a send ordinal and emits its start under a credential-free
    /// route template. Emits exhaustion at the last ordinal; later calls return
    /// `None` without reusing identities or affecting provider execution.
    pub(crate) fn begin(&self, method: &http::Method, route: &str) -> Option<AdapterAttempt> {
        let number = {
            let mut next = self
                .inner
                .next
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            let number = (*next)?;
            *next = number.checked_add(1);
            number
        };
        self.emit(
            Some(number),
            AdapterEvent::Started {
                method: method.to_string(),
                route: route.to_owned(),
            },
        );
        if number == u64::MAX {
            self.emit(None, AdapterEvent::IdentityExhausted);
        }
        Some(AdapterAttempt {
            context: self.clone(),
            number,
            closed: false,
            response_seen: false,
            unframed: 0,
            secrets: Vec::new(),
            pending_response_id: None,
            error_boundary: None,
        })
    }
}

/// The request owns this guard until completion or cancellation.
pub(crate) struct AdapterAttempt {
    context: AdapterContext,
    number: u64,
    closed: bool,
    response_seen: bool,
    /// Bytes the transport's framer had not framed when the body ended.
    unframed: usize,
    secrets: Vec<String>,
    pending_response_id: Option<String>,
    error_boundary: Option<AdapterErrorBoundary>,
}

impl AdapterAttempt {
    pub(crate) fn text(&self, text: &str) -> String {
        scrub::text(text, &self.secrets)
    }

    fn emit_with_analysis(&self, event: AdapterEvent, analysis: AdapterAnalysis) {
        let analysis = (analysis != AdapterAnalysis::default()).then_some(analysis);
        self.context
            .emit_with_analysis(Some(self.number), event, analysis);
    }

    fn emit(&self, event: AdapterEvent) {
        self.context.emit(Some(self.number), event);
    }

    /// Project a payload's facts through the projector that understands it.
    pub(crate) fn project(&mut self, project: impl FnOnce(&mut ObservationSink<'_>)) {
        project(&mut ObservationSink { attempt: self });
    }

    pub(crate) fn provider(&mut self, verdict: AdapterVerdict, response_id: Option<String>) {
        if response_id.is_some() {
            self.pending_response_id = response_id;
        }
        // An ID-only payload must not manufacture a semantic provider event.
        // Retain at most one ID until the next verdict or the attempt closure.
        if verdict != AdapterVerdict::default() {
            let response_id = self.pending_response_id.take();
            self.emit_with_analysis(
                AdapterEvent::Provider { verdict },
                AdapterAnalysis {
                    response_id,
                    ..AdapterAnalysis::default()
                },
            );
        }
    }

    pub(crate) fn response_with_headers(
        &mut self,
        status: http::StatusCode,
        headers: Option<&http::HeaderMap>,
    ) {
        if !self.response_seen {
            self.response_seen = true;
            self.emit_with_analysis(
                AdapterEvent::Response {
                    status: status.as_u16(),
                },
                AdapterAnalysis {
                    headers: headers.map(|h| scrub::headers(h, &self.secrets)),
                    ..AdapterAnalysis::default()
                },
            );
        }
    }

    pub(crate) fn finish(&mut self, ending: AdapterEnding) {
        if !self.closed {
            self.closed = true;
            let response_id = self.pending_response_id.take();
            self.emit_with_analysis(
                AdapterEvent::Finished { ending },
                AdapterAnalysis {
                    response_id,
                    ..AdapterAnalysis::default()
                },
            );
        }
    }
}

impl Drop for AdapterAttempt {
    fn drop(&mut self) {
        self.finish(AdapterEnding::Dropped);
    }
}

/// Where a payload projector writes its observation facts.
///
/// A projector reads verdicts, ids and error envelopes off a raw payload
/// before normalization discards them. Text it forwards must go through
/// [`Self::scrub`]: a payload can echo credentials.
///
/// Usage is not a projector's to report: the driver observes it once, from
/// the usage the decoder normalized into the reply's end.
///
/// ```compile_fail
/// use rig_core::wire::{AdapterEvent, AdapterUsage, ObservationSink};
///
/// // A projector cannot parse usage a second time off the raw payload.
/// fn project(sink: &mut ObservationSink<'_>) {
///     sink.emit(AdapterEvent::Usage {
///         usage: AdapterUsage::default(),
///     });
/// }
/// ```
pub struct ObservationSink<'a> {
    attempt: &'a mut AdapterAttempt,
}

impl ObservationSink<'_> {
    /// Record the provider's error envelope.
    pub fn error(&mut self, error: AdapterErrorEnvelope) {
        self.attempt.emit(AdapterEvent::ErrorEnvelope { error });
    }

    /// Record the provider's verdict, and the response id it named.
    pub fn provider(&mut self, verdict: AdapterVerdict, response_id: Option<String>) {
        self.attempt.provider(verdict, response_id);
    }

    /// Bound and redact diagnostic text from the payload.
    pub fn scrub(&self, value: &str) -> String {
        self.attempt.text(value)
    }
}

/// Shared by the SSE transport and frame driver, which own different boundaries.
#[derive(Clone, Default)]
pub(crate) struct AdapterSlot(Arc<Mutex<Option<AdapterAttempt>>>);

impl AdapterSlot {
    pub(crate) fn transport_eof(&self, after: usize) {
        if let Some(attempt) = self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .as_ref()
        {
            attempt.emit(AdapterEvent::TransportEof {
                after,
                partial_bytes: attempt.unframed,
            });
        }
    }

    /// Project a reply payload's facts through the projector that
    /// understands it.
    pub(crate) fn project(&self, project: impl FnOnce(&mut ObservationSink<'_>)) {
        if let Some(attempt) = self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .as_mut()
        {
            attempt.project(project);
        }
    }

    /// Close the attempt as its reply ended: the one place usage is
    /// observed, once, from the decoder's end, just before `Terminal`.
    pub(crate) fn terminal(&self, usage: Option<&crate::completion::Usage>) {
        let mut attempt = self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(attempt) = attempt.as_mut().filter(|attempt| !attempt.closed) {
            let usage = usage.map(AdapterUsage::from);
            if let Some(usage) = usage.filter(|usage| *usage != AdapterUsage::default()) {
                attempt.emit(AdapterEvent::Usage { usage });
            }
            attempt.finish(AdapterEnding::Terminal);
        }
    }

    /// Install the attempt this send's facts belong to.
    pub(crate) fn install(&self, attempt: Option<AdapterAttempt>) {
        *self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner) = attempt;
    }

    pub(crate) fn response(&self, status: http::StatusCode) {
        self.response_with_headers(status, None);
    }

    pub(crate) fn response_with_headers(
        &self,
        status: http::StatusCode,
        headers: Option<&http::HeaderMap>,
    ) {
        if let Some(attempt) = self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .as_mut()
        {
            attempt.response_with_headers(status, headers);
        }
    }

    fn finish(&self, ending: AdapterEnding) {
        if let Some(attempt) = self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .as_mut()
        {
            attempt.finish(ending);
        }
    }

    /// Preserve information before the SSE transport converts the owned error.
    pub(crate) fn error_boundary(&self, boundary: AdapterErrorBoundary) {
        if let Some(attempt) = self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .as_mut()
        {
            attempt.error_boundary = Some(boundary);
        }
    }

    pub(crate) fn fail(&self, error: &crate::error::ProviderError) {
        if let Some(status) = error.provider_response_status() {
            self.response(status);
        }
        let report = error.report();
        let boundary = self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .as_ref()
            .and_then(|attempt| attempt.error_boundary)
            .unwrap_or_else(|| error.boundary());
        self.finish(AdapterEnding::Error {
            boundary,
            kind: report.kind.code().to_owned(),
            status: report.http_status,
            retryable: report.is_retryable(),
        });
    }

    /// Record the bytes the transport's framer had not framed when the
    /// reply's body ended.
    pub(crate) fn unframed(&self, bytes: usize) {
        if let Some(attempt) = self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .as_mut()
        {
            attempt.unframed = bytes;
        }
    }

    pub(crate) fn eof(&self, after: usize) {
        if let Some(attempt) = self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .as_mut()
        {
            let byte_count = attempt.unframed;
            let ending = if byte_count == 0 {
                AdapterEnding::Eof { after }
            } else {
                AdapterEnding::PartialFrame { byte_count, after }
            };
            attempt.finish(ending);
        }
    }

    pub(crate) fn corrupt(&self, frame: usize) {
        if let Some(attempt) = self
            .0
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .as_ref()
        {
            attempt
                .context
                .emit(Some(attempt.number), AdapterEvent::Corrupt { frame });
        }
    }
}

mod scrub;
#[cfg(test)]
mod tests;