solo_api/llm/

sampling.rs

// SPDX-License-Identifier: Apache-2.0

//! [`SamplingLlmClient`] — `LlmClient` impl backed by an MCP client's
//! `sampling/createMessage` capability.
//!
//! Per the v0.9.0 design (`docs/dev-log/0098-v0.9.0-implementation-plan.md`
//! §6 "Sampling-backed LLM client" / MAJOR 1 + MAJOR 3 resolutions):
//!
//!   * Steward holds an `Arc<dyn LlmClient>`. When `LlmConfig::McpSampling`
//!     is configured, the Steward's `LlmClient` is a `SamplingLlmClient`
//!     constructed at MCP `initialize` time (when the live peer becomes
//!     available — the `TenantHandle::steward_slot` LATE-population path).
//!
//!   * `SamplingLlmClient::complete()` translates the workspace's
//!     `Message` → `rmcp::SamplingMessage`, calls
//!     `peer.create_message(params).await`, extracts the assistant's
//!     text from the returned `CreateMessageResult`, and emits a
//!     per-call `AuditOperation::LlmSamplingCall` row through the
//!     tenant's `WriteHandle` (lesson #30: sync in writer-actor tx
//!     for ACID).
//!
//!   * **Privacy invariant**: the audit `details_json` carries metadata
//!     only — model hint, message count, max_tokens, duration_ms,
//!     total prompt character count, output character count. **The raw
//!     prompt content MUST NOT appear in the audit row**. Pinned by
//!     [`tests::audit_row_omits_raw_prompt_text`].
//!
//!   * Error paths land structured audit rows:
//!     - Client refusal → `result = "forbidden"`,
//!       `details_json.reason = "client_refused"`.
//!     - Timeout → `result = "error"`,
//!       `details_json.reason = "timeout"`.
//!     - Other transport / malformed-response → `result = "error"`,
//!       `details_json.reason = <category>`.
//!
//!   * Per-call rate-limit / coalescing is **deferred to v0.9.0 P4**
//!     (`SamplingCoordinator`). P2 wires the per-call path only.

use std::sync::Arc;
use std::time::{Duration, Instant};

use async_trait::async_trait;
use rmcp::model::{
    CreateMessageRequestParams, CreateMessageResult, ModelHint,
    ModelPreferences, Role as RmcpRole, SamplingMessage,
    SamplingMessageContent,
};
use rmcp::service::{Peer, RoleServer, ServiceError};
use solo_core::{Error as CoreError, LlmClient, Message, Result as CoreResult, Role};
use solo_storage::{AuditEvent, AuditOperation, AuditResult, WriteHandle};

/// Default per-call timeout. Drives the bounded wait around
/// `peer.create_message`; if the client refuses or stalls, the caller
/// sees a structured timeout error instead of an indefinite hang.
///
/// 30 seconds matches the consolidate-timer's cadence margins: an
/// LLM call slower than this would already starve the Steward batch
/// in P4's coordinator. Configurable per-construct via
/// [`SamplingLlmClient::with_timeout`].
pub const DEFAULT_SAMPLING_TIMEOUT: Duration = Duration::from_secs(30);

/// Default max_tokens for sampling completions. Matches
/// `solo-steward::StewardConfig::default().abstraction_max_tokens`
/// so the wire shape is identical to what the Steward would have
/// requested from any other backend.
const DEFAULT_SAMPLING_MAX_TOKENS: u32 = 512;

/// Error surface for [`SamplingClient::create_message`]. Combines the
/// real rmcp `ServiceError` (when wrapping a live `Peer<RoleServer>`)
/// with [`super::super::test_support::fake_mcp_client::FakeSamplingError`]
/// (when driving the fixture from tests).
#[derive(Debug)]
pub enum SamplingError {
    /// Forwarded from `rmcp::Peer::create_message`.
    Service(ServiceError),
    /// Routed from [`super::super::test_support::fake_mcp_client::
    /// FakeSamplingError`] in test paths.
    #[cfg(any(test, feature = "test-support"))]
    Fake(crate::test_support::FakeSamplingError),
}

impl std::fmt::Display for SamplingError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Service(e) => write!(f, "{e}"),
            #[cfg(any(test, feature = "test-support"))]
            Self::Fake(e) => write!(f, "{e}"),
        }
    }
}

impl std::error::Error for SamplingError {}

impl SamplingError {
    /// Classifier used by [`SamplingLlmClient::complete`] to map the
    /// transport-level error to an audit-row category + a Solo
    /// [`CoreError`] variant.
    ///
    /// `(category_for_audit, treat_as_forbidden)` — `forbidden` becomes
    /// `AuditResult::Forbidden` + `CoreError::Forbidden`; everything
    /// else is `AuditResult::Error` + `CoreError::Llm`.
    pub fn classify(&self) -> (&'static str, bool) {
        match self {
            Self::Service(_) => ("transport_error", false),
            #[cfg(any(test, feature = "test-support"))]
            Self::Fake(e) => match e {
                crate::test_support::FakeSamplingError::Refused { .. } => {
                    ("client_refused", true)
                }
                crate::test_support::FakeSamplingError::Transport { .. } => {
                    ("transport_error", false)
                }
                crate::test_support::FakeSamplingError::MalformedResponse {
                    ..
                } => ("malformed_response", false),
            },
        }
    }
}

/// Trait abstracting the `sampling/createMessage` RPC. The production
/// impl wraps `Arc<Peer<RoleServer>>`; the test impl is
/// [`super::super::test_support::fake_mcp_client::FakeMcpClient`].
///
/// Separating the trait from the concrete `Peer<RoleServer>` is the
/// way around rmcp's `Peer` having private constructors — we can't
/// build a fake `Peer` for tests, so we inject behind a trait.
#[async_trait]
pub trait SamplingClient: Send + Sync {
    async fn create_message(
        &self,
        params: CreateMessageRequestParams,
    ) -> Result<CreateMessageResult, SamplingError>;
}

/// Production wrapper around `rmcp::Peer<RoleServer>`. The Peer is
/// cheap to clone (internally `Arc`-backed) and stays valid for the
/// lifetime of the MCP session.
pub struct PeerSamplingClient {
    peer: Peer<RoleServer>,
}

impl PeerSamplingClient {
    pub fn new(peer: Peer<RoleServer>) -> Self {
        Self { peer }
    }
}

#[async_trait]
impl SamplingClient for PeerSamplingClient {
    async fn create_message(
        &self,
        params: CreateMessageRequestParams,
    ) -> Result<CreateMessageResult, SamplingError> {
        self.peer
            .create_message(params)
            .await
            .map_err(SamplingError::Service)
    }
}

/// `LlmClient` impl whose `complete()` calls back via the connected
/// MCP client's sampling capability.
///
/// Construct via [`SamplingLlmClient::new`] (production path: wraps a
/// real `Peer<RoleServer>`) or [`SamplingLlmClient::with_sampling_client`]
/// (test path: takes the abstracted [`SamplingClient`] trait object
/// directly so [`super::super::test_support::fake_mcp_client::
/// FakeMcpClient`] can drive it).
///
/// Cheap to clone — every field is `Arc`-shared.
#[derive(Clone)]
pub struct SamplingLlmClient {
    /// The RPC channel back to the MCP client.
    sampling_client: Arc<dyn SamplingClient>,
    /// Per-tenant `WriteHandle` for the synchronous audit emit. Routes
    /// through the writer-actor's mpsc so the
    /// `AuditOperation::LlmSamplingCall` INSERT lands in a dedicated
    /// `BEGIN IMMEDIATE` transaction on the writer's connection.
    write_handle: WriteHandle,
    /// Cached audit `principal_subject` for the MCP session. Resolved
    /// at session init time (see `mcp::resolve_mcp_principal`); `None`
    /// for unauthenticated stdio sessions.
    audit_principal: Option<String>,
    /// `max_tokens` value sent on every `sampling/createMessage`.
    /// Defaults to [`DEFAULT_SAMPLING_MAX_TOKENS`]; configurable via
    /// [`Self::with_max_tokens`].
    max_tokens: u32,
    /// Bounded wait on `create_message`. See
    /// [`DEFAULT_SAMPLING_TIMEOUT`].
    timeout: Duration,
}

impl SamplingLlmClient {
    /// Build a client wrapping a real `Peer<RoleServer>`. Production
    /// path — called from
    /// [`crate::mcp::SoloMcpServer::populate_sampling_steward`] when an MCP
    /// session reaches `initialize` with a sampling-capable peer.
    pub fn new(
        peer: Peer<RoleServer>,
        write_handle: WriteHandle,
        audit_principal: Option<String>,
    ) -> Self {
        Self::with_sampling_client(
            Arc::new(PeerSamplingClient::new(peer)),
            write_handle,
            audit_principal,
        )
    }

    /// Test-friendly constructor accepting any [`SamplingClient`]
    /// implementation. Pair with
    /// [`super::super::test_support::fake_mcp_client::FakeMcpClient`]
    /// in tests.
    pub fn with_sampling_client(
        sampling_client: Arc<dyn SamplingClient>,
        write_handle: WriteHandle,
        audit_principal: Option<String>,
    ) -> Self {
        Self {
            sampling_client,
            write_handle,
            audit_principal,
            max_tokens: DEFAULT_SAMPLING_MAX_TOKENS,
            timeout: DEFAULT_SAMPLING_TIMEOUT,
        }
    }

    /// Override the per-call `max_tokens` cap.
    pub fn with_max_tokens(mut self, n: u32) -> Self {
        self.max_tokens = n.max(1);
        self
    }

    /// Override the per-call timeout.
    pub fn with_timeout(mut self, t: Duration) -> Self {
        self.timeout = t;
        self
    }

    /// Build the `CreateMessageRequestParams` from Solo's `Message`
    /// vec. Splits out `Role::System` into the `system_prompt` field
    /// (rmcp's `SamplingMessage::role` is only User / Assistant) and
    /// hints the user's MCP client toward a Claude-class model.
    fn build_request(&self, messages: &[Message]) -> CreateMessageRequestParams {
        // Split system messages out of the conversation history; the
        // sampling protocol carries the system prompt as a top-level
        // field rather than inline.
        let mut system_parts: Vec<String> = Vec::new();
        let mut samp_messages: Vec<SamplingMessage> = Vec::new();
        for m in messages {
            match m.role {
                Role::System => system_parts.push(m.content.clone()),
                Role::User => {
                    samp_messages.push(SamplingMessage::user_text(&m.content));
                }
                Role::Assistant => {
                    samp_messages
                        .push(SamplingMessage::assistant_text(&m.content));
                }
            }
        }
        // rmcp 1.7's struct literals are non-exhaustive across crate
        // boundaries; build via the typed constructors + builders.
        let preferences = ModelPreferences::new()
            .with_hints(vec![ModelHint::new("claude")])
            .with_intelligence_priority(0.7)
            .with_speed_priority(0.3)
            .with_cost_priority(0.4);
        let mut params =
            CreateMessageRequestParams::new(samp_messages, self.max_tokens)
                .with_model_preferences(preferences);
        if !system_parts.is_empty() {
            params = params.with_system_prompt(system_parts.join("\n\n"));
        }
        params
    }

    /// Build the audit `AuditEvent` carrying ONLY metadata. No raw
    /// prompt content lands in `details_json`.
    ///
    /// Pinned by [`tests::audit_row_omits_raw_prompt_text`].
    fn audit_event(
        &self,
        params: &CreateMessageRequestParams,
        outcome: SamplingOutcome,
    ) -> AuditEvent {
        let prompt_chars: usize = params
            .messages
            .iter()
            .flat_map(|m| m.content.iter())
            .filter_map(|c| c.as_text().map(|t| t.text.len()))
            .sum::<usize>()
            + params
                .system_prompt
                .as_ref()
                .map(|s| s.len())
                .unwrap_or(0);
        // ~4 chars per token for the rough English-text estimate used
        // by `solo doctor --check-llm` and Anthropic's docs. Recorded
        // for operator capacity-planning.
        let input_tokens_est = (prompt_chars / 4) as u64;
        let model_hint = params
            .model_preferences
            .as_ref()
            .and_then(|p| p.hints.as_ref())
            .and_then(|h| h.first())
            .and_then(|h| h.name.clone())
            .unwrap_or_else(|| "(none)".to_string());

        let mut details = serde_json::Map::new();
        details.insert(
            "model_hint".to_string(),
            serde_json::Value::String(model_hint),
        );
        details.insert(
            "messages_count".to_string(),
            serde_json::Value::Number(params.messages.len().into()),
        );
        details.insert(
            "max_tokens".to_string(),
            serde_json::Value::Number(params.max_tokens.into()),
        );
        details.insert(
            "prompt_chars".to_string(),
            serde_json::Value::Number(prompt_chars.into()),
        );
        details.insert(
            "input_tokens_est".to_string(),
            serde_json::Value::Number(input_tokens_est.into()),
        );

        let result = match &outcome {
            SamplingOutcome::Ok {
                duration_ms,
                model,
                output_chars,
            } => {
                let output_tokens_est = (*output_chars / 4) as u64;
                details.insert(
                    "duration_ms".to_string(),
                    serde_json::Value::Number((*duration_ms).into()),
                );
                details.insert(
                    "model".to_string(),
                    serde_json::Value::String(model.clone()),
                );
                details.insert(
                    "output_chars".to_string(),
                    serde_json::Value::Number((*output_chars).into()),
                );
                details.insert(
                    "output_tokens_est".to_string(),
                    serde_json::Value::Number(output_tokens_est.into()),
                );
                AuditResult::Ok
            }
            SamplingOutcome::Forbidden {
                reason,
                duration_ms,
            } => {
                details.insert(
                    "duration_ms".to_string(),
                    serde_json::Value::Number((*duration_ms).into()),
                );
                details.insert(
                    "reason".to_string(),
                    serde_json::Value::String(reason.to_string()),
                );
                AuditResult::Forbidden
            }
            SamplingOutcome::Error {
                reason,
                duration_ms,
            } => {
                details.insert(
                    "duration_ms".to_string(),
                    serde_json::Value::Number((*duration_ms).into()),
                );
                details.insert(
                    "reason".to_string(),
                    serde_json::Value::String(reason.to_string()),
                );
                AuditResult::Error
            }
        };

        AuditEvent {
            ts_ms: chrono::Utc::now().timestamp_millis(),
            principal_subject: self.audit_principal.clone(),
            operation: AuditOperation::LlmSamplingCall,
            target_id: None,
            result,
            details: Some(serde_json::Value::Object(details)),
        }
    }
}

/// Internal outcome category for the audit-row builder.
enum SamplingOutcome {
    Ok {
        duration_ms: u64,
        model: String,
        output_chars: usize,
    },
    Forbidden {
        reason: &'static str,
        duration_ms: u64,
    },
    Error {
        reason: &'static str,
        duration_ms: u64,
    },
}

#[async_trait]
impl LlmClient for SamplingLlmClient {
    fn name(&self) -> &str {
        "mcp-sampling"
    }

    async fn complete(&self, messages: &[Message]) -> CoreResult<Message> {
        let params = self.build_request(messages);
        let start = Instant::now();

        // Bounded wait on `peer.create_message`. The fold of (rmcp
        // ServiceError | FakeError | tokio timeout) into the
        // `Outcome` enum keeps the audit path single-sourced.
        let rpc = tokio::time::timeout(
            self.timeout,
            self.sampling_client.create_message(params.clone()),
        )
        .await;
        let duration_ms = start.elapsed().as_millis().min(u128::from(u64::MAX))
            as u64;

        let (core_result, outcome): (CoreResult<Message>, SamplingOutcome) =
            match rpc {
                Ok(Ok(result)) => {
                    match extract_text(&result) {
                        Ok(text) => {
                            let output_chars = text.len();
                            let outcome = SamplingOutcome::Ok {
                                duration_ms,
                                model: result.model.clone(),
                                output_chars,
                            };
                            (Ok(Message::assistant(text)), outcome)
                        }
                        Err(reason) => (
                            Err(CoreError::llm(format!(
                                "mcp sampling: malformed response: {reason}"
                            ))),
                            SamplingOutcome::Error {
                                reason: "malformed_response",
                                duration_ms,
                            },
                        ),
                    }
                }
                Ok(Err(e)) => {
                    let (category, is_forbidden) = e.classify();
                    let outcome = if is_forbidden {
                        SamplingOutcome::Forbidden {
                            reason: category,
                            duration_ms,
                        }
                    } else {
                        SamplingOutcome::Error {
                            reason: category,
                            duration_ms,
                        }
                    };
                    let err = if is_forbidden {
                        CoreError::forbidden(format!("mcp sampling: {e}"))
                    } else {
                        CoreError::llm(format!("mcp sampling: {e}"))
                    };
                    (Err(err), outcome)
                }
                Err(_elapsed) => (
                    Err(CoreError::llm(format!(
                        "mcp sampling: timeout after {}ms",
                        duration_ms
                    ))),
                    SamplingOutcome::Error {
                        reason: "timeout",
                        duration_ms,
                    },
                ),
            };

        // Synchronous audit emit through the writer-actor (lesson #30).
        // Failure to land the audit row is operator-visible: the
        // sampling call's caller sees the storage error and can decide
        // whether to abort (we DO abort here — without the audit row
        // we have no record of the call).
        let event = self.audit_event(&params, outcome);
        if let Err(audit_err) = self.write_handle.emit_llm_sampling_audit(event).await
        {
            // If the RPC itself succeeded but the audit failed, treat
            // the call as failed. The audit row IS the persisted
            // record of the call; an undocumented sampling call is
            // not acceptable.
            return Err(CoreError::storage(format!(
                "mcp sampling: audit emit failed: {audit_err}"
            )));
        }

        core_result
    }
}

/// Pull the assistant's text out of the rmcp result. Walks every text
/// content block in the message (the spec allows either a single
/// `SamplingContent::Single` or a `SamplingContent::Multiple`) and
/// concatenates them with newlines. Returns `Err(reason)` if no text
/// blocks were present — the malformed-response path.
fn extract_text(result: &CreateMessageResult) -> Result<String, &'static str> {
    if result.message.role != RmcpRole::Assistant {
        return Err("response role was not Assistant");
    }
    let mut out = String::new();
    for content in result.message.content.iter() {
        if let SamplingMessageContent::Text(text) = content {
            if !out.is_empty() {
                out.push('\n');
            }
            out.push_str(&text.text);
        }
    }
    if out.is_empty() {
        Err("no text content blocks")
    } else {
        Ok(out)
    }
}

/// v0.9.0 P2: build a sampling-backed `Arc<Steward>` for a tenant that
/// has resolved `LlmConfig::McpSampling` and just attached an MCP
/// session.
///
/// Called from [`crate::mcp::SoloMcpServer::populate_sampling_steward`] at
/// MCP `initialize` time once the peer's sampling capability is
/// confirmed. The returned `Arc<Steward>` is written into
/// `tenant.steward_slot()` so the writer-actor + consolidate timer
/// can read a populated slot on their next tick.
///
/// v0.9.0 P5 (M3 wiring): the live `PeerSamplingClient` is now wrapped
/// in a [`super::SamplingCoordinator`] before being handed to
/// `SamplingLlmClient`. Concurrent `complete()` calls within the
/// coalesce window collapse into one `peer.create_message` RPC and the
/// response is demultiplexed back per-task — matching the
/// `[sampling] coalesce_window_ms` / `coalesce_max_requests` config the
/// operator wrote in `solo.config.toml`. Per-call audit emit semantics
/// are unchanged: every logical request still lands one
/// `AuditOperation::LlmSamplingCall` row, no raw prompt content escapes
/// to the audit row.
///
/// Edge case (clamping): the `[sampling]` block accepts values that
/// effectively disable batching — `coalesce_max_requests = 1` and / or
/// `coalesce_window_ms = 0` reduce the coordinator to pass-through (one
/// inner call per submission). The coordinator's
/// [`super::SamplingCoordinator::with_settings`] clamps `max_batch` to
/// `max(1)` so a zero value still produces a single-element flush
/// immediately rather than panicking or deadlocking.
pub fn build_sampling_steward(
    peer: Peer<RoleServer>,
    write_handle: WriteHandle,
    audit_principal: Option<String>,
    steward_config: solo_steward::StewardConfig,
    sampling_config: solo_storage::SamplingConfig,
) -> Arc<solo_steward::Steward> {
    let inner: Arc<dyn SamplingClient> = Arc::new(PeerSamplingClient::new(peer));
    let coordinator: Arc<dyn SamplingClient> = super::SamplingCoordinator::with_settings(
        inner,
        std::time::Duration::from_millis(sampling_config.coalesce_window_ms),
        sampling_config.coalesce_max_requests as usize,
    );
    let client = SamplingLlmClient::with_sampling_client(
        coordinator,
        write_handle,
        audit_principal,
    )
    .with_max_tokens(steward_config.abstraction_max_tokens.min(65_536) as u32);
    Arc::new(solo_steward::Steward::new(Arc::new(client), steward_config))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_support::{FakeMcpClient, FakeResponse, FakeSamplingError};
    use rmcp::model::CreateMessageResult;
    use solo_core::TenantId;
    use solo_storage::{
        EmbedderConfig, HnswParams, InitParams, KeyMaterial, StubEmbedder,
        TenantHandle, TenantRegistry, TenantRegistryParams, init,
        open_sqlcipher,
    };
    use std::path::PathBuf;
    use std::sync::Arc;
    use tempfile::TempDir;
    use zeroize::Zeroizing;

    const TEST_PASSPHRASE: &str = "v0.9.0-p2-sampling-tests";

    /// Bootstrap a per-tenant `TenantHandle` whose writer-actor accepts
    /// the new `WriteCommand::EmitLlmSamplingAudit` variant.
    ///
    /// Mirrors the v0.8.x test discipline (see
    /// `crates/solo-storage/src/tenants/handle_registry_tests.rs`'s
    /// `fresh_init_dir`): build a real tenant DB on disk via the same
    /// `init()` helper users invoke, wrap in a `TenantRegistry`, and
    /// surface the `WriteHandle` for direct `SamplingLlmClient`
    /// wiring.
    struct Harness {
        _tmp: TempDir,
        _registry: Arc<TenantRegistry>,
        _tenant: Arc<TenantHandle>,
        write_handle: solo_storage::WriteHandle,
        db_path: PathBuf,
        key: KeyMaterial,
    }

    async fn harness() -> Harness {
        let tmp = TempDir::new().expect("tempdir");
        let data_dir = tmp.path().to_path_buf();
        let _ = init(InitParams {
            data_dir: data_dir.clone(),
            passphrase: Zeroizing::new(TEST_PASSPHRASE.into()),
            force: false,
            embedder: EmbedderConfig {
                name: "stub".into(),
                version: "v1".into(),
                dim: 32,
                dtype: "f32".into(),
            },
        })
        .expect("init");

        let cfg = solo_storage::SoloConfig::read(
            &data_dir.join("solo.config.toml"),
        )
        .expect("read cfg");
        let key = KeyMaterial::derive(
            TEST_PASSPHRASE,
            &cfg.salt_bytes().expect("salt"),
        )
        .expect("derive key");

        let embedder: Arc<dyn solo_core::Embedder> =
            Arc::new(StubEmbedder::new("stub", "v1", 32));
        let registry = Arc::new(
            TenantRegistry::open(TenantRegistryParams {
                data_dir: data_dir.clone(),
                key: key.clone(),
                embedder: embedder.clone(),
                hnsw_params: HnswParams::default(),
                steward: None,
                runtime_handle: Some(tokio::runtime::Handle::current()),
                steward_factory: None,
                triples_batch_signal: None,
            })
            .expect("open registry"),
        );

        let tenant_id = TenantId::default_tenant();
        let tenant = registry
            .get_or_open(&tenant_id)
            .await
            .expect("get_or_open default tenant");
        let write_handle = tenant.write().clone();
        let db_path = tenant.db_path().to_path_buf();

        Harness {
            _tmp: tmp,
            _registry: registry,
            _tenant: tenant,
            write_handle,
            db_path,
            key,
        }
    }

    /// Helper: count the `audit_events` rows whose `operation` is the
    /// given string. Opens a fresh connection to the tenant DB so we
    /// avoid contention with the writer-actor's own connection.
    fn count_audit_rows(db_path: &std::path::Path, key: &KeyMaterial, op: &str) -> i64 {
        let conn = open_sqlcipher(db_path, key).expect("open db");
        conn.query_row(
            "SELECT COUNT(*) FROM audit_events WHERE operation = ?",
            rusqlite::params![op],
            |r| r.get(0),
        )
        .expect("count")
    }

    /// Helper: load the most-recent `llm.sampling_call` audit row and
    /// return `(result, principal_subject, details_json)`.
    fn latest_sampling_audit_details(
        db_path: &std::path::Path,
        key: &KeyMaterial,
    ) -> (String, Option<String>, serde_json::Value) {
        let conn = open_sqlcipher(db_path, key).expect("open db");
        let (result, principal, details_str): (String, Option<String>, Option<String>) = conn
            .query_row(
                "SELECT result, principal_subject, details_json
                 FROM audit_events
                 WHERE operation = 'llm.sampling_call'
                 ORDER BY ts_ms DESC, rowid DESC
                 LIMIT 1",
                [],
                |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
            )
            .expect("query");
        let details: serde_json::Value =
            serde_json::from_str(&details_str.expect("details_json present"))
                .expect("parse details");
        (result, principal, details)
    }

    /// Happy path: a successful `create_message` round-trip returns
    /// the assistant text wrapped in a `Message::assistant`, and lands
    /// exactly one `llm.sampling_call` audit row with `result = 'ok'`.
    #[tokio::test]
    async fn sampling_complete_happy_path_returns_text() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text("derived theme")));
        let client = SamplingLlmClient::with_sampling_client(
            fake.clone(),
            h.write_handle.clone(),
            Some("alice".into()),
        );
        let messages = vec![Message::user("summarise these episodes")];
        let result = client.complete(&messages).await.expect("ok");
        assert_eq!(result.role, Role::Assistant);
        assert_eq!(result.content, "derived theme");

        // Exactly one audit row landed.
        assert_eq!(
            count_audit_rows(&h.db_path, &h.key, "llm.sampling_call"),
            1
        );
        let (result_str, principal, details) =
            latest_sampling_audit_details(&h.db_path, &h.key);
        assert_eq!(result_str, "ok");
        assert_eq!(principal.as_deref(), Some("alice"));
        assert_eq!(details["model_hint"], "claude");
        assert_eq!(details["model"], "fake-claude");
        assert_eq!(details["messages_count"], 1);
        assert_eq!(details["max_tokens"], 512);
    }

    /// Privacy invariant: the audit row's `details_json` MUST NOT
    /// contain the raw prompt content. Pinned by string inspection of
    /// the persisted JSON.
    #[tokio::test]
    async fn audit_row_omits_raw_prompt_text() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text("ok")));
        let client = SamplingLlmClient::with_sampling_client(
            fake,
            h.write_handle.clone(),
            None,
        );
        let secret = "THE-USER-ID-IS-bobby-1234";
        let messages = vec![
            Message::system("you are a friendly assistant"),
            Message::user(secret),
        ];
        client.complete(&messages).await.expect("ok");

        let (_, _, details) =
            latest_sampling_audit_details(&h.db_path, &h.key);
        let serialised =
            serde_json::to_string(&details).expect("serialise details");
        assert!(
            !serialised.contains(secret),
            "audit details must not carry raw prompt content; was: {serialised}"
        );
        assert!(
            !serialised.contains("you are a friendly assistant"),
            "audit details must not carry system prompt; was: {serialised}"
        );
        // Metadata IS present, even though the prompt is not.
        assert_eq!(details["messages_count"], 1);
        assert!(details["prompt_chars"].as_u64().unwrap() > 0);
    }

    /// Client refusal: maps to `CoreError::Forbidden` + audit row
    /// `result = 'forbidden'` + `details_json.reason = 'client_refused'`.
    #[tokio::test]
    async fn client_refusal_returns_forbidden_and_audits() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text("ignored")));
        fake.reject_with("user dismissed approval");
        let client = SamplingLlmClient::with_sampling_client(
            fake,
            h.write_handle.clone(),
            Some("alice".into()),
        );
        let err = client
            .complete(&[Message::user("anything")])
            .await
            .unwrap_err();
        match err {
            CoreError::Forbidden(_) => {}
            other => panic!("expected Forbidden, got {other:?}"),
        }
        let (result_str, _, details) =
            latest_sampling_audit_details(&h.db_path, &h.key);
        assert_eq!(result_str, "forbidden");
        assert_eq!(details["reason"], "client_refused");
    }

    /// Timeout: tokio::time::timeout fires before the fake's `Slow`
    /// response resolves; client returns `CoreError::Llm` + audit row
    /// `result = 'error'` + `details_json.reason = 'timeout'`.
    ///
    /// Real wall-clock: 80ms slow response vs 30ms client timeout.
    /// Margin is loose enough for slow CI without making the test
    /// drag.
    #[tokio::test]
    async fn timeout_returns_error_with_timeout_reason() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::slow(
            "late",
            Duration::from_millis(800),
        )));
        let client = SamplingLlmClient::with_sampling_client(
            fake,
            h.write_handle.clone(),
            None,
        )
        .with_timeout(Duration::from_millis(30));
        let err = client
            .complete(&[Message::user("hello")])
            .await
            .unwrap_err();
        match err {
            CoreError::Llm(msg) => assert!(msg.contains("timeout")),
            other => panic!("expected Llm, got {other:?}"),
        }
        let (result_str, _, details) =
            latest_sampling_audit_details(&h.db_path, &h.key);
        assert_eq!(result_str, "error");
        assert_eq!(details["reason"], "timeout");
    }

    /// Malformed response: the fake returns a result with zero text
    /// content blocks; client surfaces `CoreError::Llm` + audit row
    /// `result = 'error'` + `details_json.reason = 'malformed_response'`.
    #[tokio::test]
    async fn malformed_response_returns_error_with_reason() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::EmptyContent));
        let client = SamplingLlmClient::with_sampling_client(
            fake,
            h.write_handle.clone(),
            None,
        );
        let err = client
            .complete(&[Message::user("hi")])
            .await
            .unwrap_err();
        assert!(matches!(err, CoreError::Llm(_)));
        let (result_str, _, details) =
            latest_sampling_audit_details(&h.db_path, &h.key);
        assert_eq!(result_str, "error");
        assert_eq!(details["reason"], "malformed_response");
    }

    /// `principal_subject = None` works — audit row still emits with
    /// NULL.
    #[tokio::test]
    async fn no_principal_emits_audit_with_null_principal() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text("ok")));
        let client = SamplingLlmClient::with_sampling_client(
            fake,
            h.write_handle.clone(),
            None,
        );
        client.complete(&[Message::user("hi")]).await.expect("ok");
        let (_, principal, _) =
            latest_sampling_audit_details(&h.db_path, &h.key);
        assert_eq!(principal, None);
    }

    /// Concurrency: 8 parallel `complete()` calls land 8 audit rows.
    /// Audit IDs (autoincrement rowid) must be distinct — verifies the
    /// writer-actor serialises the per-call audit emit (no
    /// interleaving / dropped rows).
    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn parallel_completes_serialise_audit_rows() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text("ok")));
        let client = SamplingLlmClient::with_sampling_client(
            fake.clone(),
            h.write_handle.clone(),
            Some("alice".into()),
        );
        let mut futs = Vec::new();
        for _ in 0..8 {
            let c = client.clone();
            futs.push(tokio::spawn(async move {
                c.complete(&[Message::user("hi")]).await
            }));
        }
        for f in futs {
            f.await.expect("join").expect("ok");
        }
        assert_eq!(
            count_audit_rows(&h.db_path, &h.key, "llm.sampling_call"),
            8,
            "8 parallel calls must land 8 audit rows"
        );

        // Each was a separate request to the fake.
        assert_eq!(fake.record_requests().len(), 8);
    }

    /// `complete` translates the workspace's `Message::system` into the
    /// `system_prompt` top-level field; user/assistant roles map to
    /// rmcp's `SamplingMessage::user_text` / `assistant_text`.
    #[tokio::test]
    async fn build_request_splits_system_from_messages() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text("ok")));
        let client = SamplingLlmClient::with_sampling_client(
            fake.clone(),
            h.write_handle.clone(),
            None,
        );
        client
            .complete(&[
                Message::system("be terse"),
                Message::user("question"),
                Message::assistant("answer"),
            ])
            .await
            .expect("ok");
        let recorded = fake.record_requests();
        assert_eq!(recorded.len(), 1);
        let req = &recorded[0];
        assert_eq!(
            req.system_prompt.as_deref(),
            Some("be terse"),
            "Role::System must map to system_prompt"
        );
        assert_eq!(req.messages.len(), 2);
        // The remaining two messages are the user + assistant turns.
        assert_eq!(req.messages[0].role, RmcpRole::User);
        assert_eq!(req.messages[1].role, RmcpRole::Assistant);
    }

    /// `model_preferences` carries the `claude` hint per plan §6.
    /// Pins the wire shape so a future change is a conscious decision.
    #[tokio::test]
    async fn build_request_includes_claude_model_hint() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text("ok")));
        let client = SamplingLlmClient::with_sampling_client(
            fake.clone(),
            h.write_handle.clone(),
            None,
        );
        client
            .complete(&[Message::user("hi")])
            .await
            .expect("ok");
        let recorded = fake.record_requests();
        let prefs = recorded[0].model_preferences.as_ref().expect("prefs");
        let hint = prefs
            .hints
            .as_ref()
            .and_then(|h| h.first())
            .and_then(|h| h.name.clone())
            .expect("hint name");
        assert_eq!(hint, "claude");
    }

    /// `with_max_tokens(n)` propagates to the request's
    /// `max_tokens` field.
    #[tokio::test]
    async fn with_max_tokens_overrides_default() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text("ok")));
        let client = SamplingLlmClient::with_sampling_client(
            fake.clone(),
            h.write_handle.clone(),
            None,
        )
        .with_max_tokens(2048);
        client
            .complete(&[Message::user("hi")])
            .await
            .expect("ok");
        let recorded = fake.record_requests();
        assert_eq!(recorded[0].max_tokens, 2048);
    }

    /// Reconfiguring the fake mid-test produces distinct audit rows
    /// for each call (positive then negative).
    #[tokio::test]
    async fn reconfigurable_fake_distinguishes_audit_rows() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text("ok")));
        let client = SamplingLlmClient::with_sampling_client(
            fake.clone(),
            h.write_handle.clone(),
            Some("alice".into()),
        );

        client.complete(&[Message::user("a")]).await.expect("ok");
        fake.reject_with("user said no");
        let _ = client.complete(&[Message::user("b")]).await;

        let conn = open_sqlcipher(&h.db_path, &h.key).expect("open");
        let mut stmt = conn
            .prepare(
                "SELECT result FROM audit_events WHERE operation = 'llm.sampling_call' ORDER BY ts_ms ASC, rowid ASC",
            )
            .expect("prepare");
        let rows: Vec<String> = stmt
            .query_map([], |r| r.get::<_, String>(0))
            .expect("query")
            .map(|r| r.expect("row"))
            .collect();
        assert_eq!(rows, vec!["ok".to_string(), "forbidden".to_string()]);
    }

    /// `extract_text` walks single-block content.
    #[test]
    fn extract_text_pulls_text_from_single_block() {
        let result = CreateMessageResult::new(
            SamplingMessage::assistant_text("hello"),
            "fake".into(),
        );
        assert_eq!(extract_text(&result).unwrap(), "hello");
    }

    /// `extract_text` rejects an empty-content response.
    #[test]
    fn extract_text_rejects_empty_content() {
        let result = CreateMessageResult::new(
            SamplingMessage::new_multiple(RmcpRole::Assistant, Vec::new()),
            "fake".into(),
        );
        assert!(extract_text(&result).is_err());
    }

    /// `extract_text` rejects a User-role response (impossible per
    /// spec but pinning the defensive check).
    #[test]
    fn extract_text_rejects_non_assistant_role() {
        let result = CreateMessageResult::new(
            SamplingMessage::user_text("hello"),
            "fake".into(),
        );
        assert!(extract_text(&result).is_err());
    }

    /// `SamplingError::classify` maps each fake variant to the right
    /// audit category.
    #[test]
    fn sampling_error_classify_maps_fake_variants() {
        let refused = SamplingError::Fake(FakeSamplingError::Refused {
            reason: "x".into(),
        });
        let (cat, forb) = refused.classify();
        assert_eq!(cat, "client_refused");
        assert!(forb);

        let transport = SamplingError::Fake(FakeSamplingError::Transport {
            message: "x".into(),
        });
        let (cat, forb) = transport.classify();
        assert_eq!(cat, "transport_error");
        assert!(!forb);

        let malformed =
            SamplingError::Fake(FakeSamplingError::MalformedResponse {
                message: "x".into(),
            });
        let (cat, forb) = malformed.classify();
        assert_eq!(cat, "malformed_response");
        assert!(!forb);
    }

    // -------- v0.9.0 P5a (M3 wiring) — SamplingCoordinator integration --------
    //
    // These tests pin the contract that `build_sampling_steward` wraps the
    // live peer in a `SamplingCoordinator` before handing it to
    // `SamplingLlmClient`. They cannot call `build_sampling_steward`
    // directly (it takes a real `Peer<RoleServer>` whose constructors are
    // private inside rmcp), but they exercise the **exact same wiring
    // shape** by substituting `FakeMcpClient` for `PeerSamplingClient`.
    // The production code path is:
    //
    //     PeerSamplingClient -> SamplingCoordinator -> SamplingLlmClient
    //
    // The tested shape is:
    //
    //     FakeMcpClient      -> SamplingCoordinator -> SamplingLlmClient
    //
    // Only the leaf `SamplingClient` impl differs; the
    // `SamplingClient` trait is the same Arc-of-dyn in both paths.

    /// SamplingCoordinator wrapping a `FakeMcpClient` and feeding
    /// `SamplingLlmClient::with_sampling_client` is the same Arc-of-dyn
    /// shape `build_sampling_steward` constructs at MCP-initialize
    /// time. Single-element flushes pass through unwrapped, so a lone
    /// `complete()` call still emits one audit row and produces the
    /// expected text.
    #[tokio::test]
    async fn sampling_llm_client_uses_coordinator_in_production_path() {
        let h = harness().await;
        let fake: Arc<dyn SamplingClient> =
            Arc::new(FakeMcpClient::new(FakeResponse::text("ok")));
        let coord: Arc<dyn SamplingClient> =
            super::super::SamplingCoordinator::with_settings(
                fake.clone(),
                Duration::from_millis(50),
                10,
            );
        let client = SamplingLlmClient::with_sampling_client(
            coord,
            h.write_handle.clone(),
            Some("alice".into()),
        );
        let result = client
            .complete(&[Message::user("test")])
            .await
            .expect("ok");
        assert_eq!(result.role, Role::Assistant);
        assert_eq!(result.content, "ok");
        // Single audit row landed — per-call audit semantics
        // unchanged by the coordinator wrap.
        assert_eq!(
            count_audit_rows(&h.db_path, &h.key, "llm.sampling_call"),
            1,
            "one logical call → one audit row, even through coordinator"
        );
    }

    /// End-to-end batching pin: N concurrent `complete()` calls within
    /// the coalesce window resolve as ONE inner `create_message` RPC
    /// on the underlying `FakeMcpClient`, but N audit rows still land
    /// (one per logical call — the privacy + audit invariants from P2
    /// hold).
    ///
    /// This is the v0.9.0 release notes' "⌈N/M⌉ peer.create_message
    /// calls per coalesce window" claim, exercised through the same
    /// trait-object chain that `build_sampling_steward` constructs in
    /// production.
    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn coordinator_coalesces_concurrent_calls_into_one_inner_rpc() {
        // Coalesced JSON response for 5 tasks — matches the
        // `[{task_index, response}]` shape `flush_batch` demuxes
        // multi-element batches into.
        let response = serde_json::to_string(&(0..5)
            .map(|i| serde_json::json!({
                "task_index": i,
                "response": format!("response-{i}"),
            }))
            .collect::<Vec<_>>())
            .unwrap();

        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text(&response)));
        let coord: Arc<dyn SamplingClient> =
            super::super::SamplingCoordinator::with_settings(
                fake.clone(),
                // Wide window so all 5 submissions land in one batch.
                Duration::from_secs(5),
                10,
            );
        let client = SamplingLlmClient::with_sampling_client(
            coord,
            h.write_handle.clone(),
            Some("alice".into()),
        );

        // Fire 5 concurrent `complete()` calls; the coordinator should
        // coalesce them into ONE `FakeMcpClient::create_message` call.
        let mut futs = Vec::new();
        for i in 0..5 {
            let c = client.clone();
            futs.push(tokio::spawn(async move {
                c.complete(&[Message::user(format!("task-{i}"))]).await
            }));
        }
        for f in futs {
            f.await.expect("join").expect("ok");
        }

        // EXACTLY one inner RPC.
        assert_eq!(
            fake.record_requests().len(),
            1,
            "5 logical calls within window must coalesce to 1 inner RPC"
        );
        // BUT 5 audit rows — per-logical-call audit invariant preserved.
        assert_eq!(
            count_audit_rows(&h.db_path, &h.key, "llm.sampling_call"),
            5,
            "5 logical calls → 5 audit rows (coordinator doesn't merge audits)"
        );
    }

    /// Edge case: `coalesce_max_requests = 1` reduces the coordinator
    /// to pass-through (each submit flushes a 1-element batch
    /// immediately). With max_batch=1 and a wide window, 3 concurrent
    /// calls land 3 inner RPCs — coordinator is operating as if no
    /// batching were configured.
    ///
    /// Pins the brief's documented edge-case: zero / one-valued config
    /// reduces to pass-through, never panics or deadlocks. Mirrors
    /// `SamplingCoordinator::with_settings`'s `max_batch.max(1)`
    /// clamping for the `coalesce_max_requests = 0` case.
    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn coordinator_max_batch_one_acts_as_passthrough() {
        let h = harness().await;
        let fake = Arc::new(FakeMcpClient::new(FakeResponse::text("ok")));
        let coord: Arc<dyn SamplingClient> =
            super::super::SamplingCoordinator::with_settings(
                fake.clone(),
                Duration::from_secs(5),
                // max_batch=1 → every submission flushes immediately as
                // a 1-element batch; pass-through behaviour.
                1,
            );
        let client = SamplingLlmClient::with_sampling_client(
            coord,
            h.write_handle.clone(),
            None,
        );
        let mut futs = Vec::new();
        for _ in 0..3 {
            let c = client.clone();
            futs.push(tokio::spawn(async move {
                c.complete(&[Message::user("hi")]).await
            }));
        }
        for f in futs {
            f.await.expect("join").expect("ok");
        }
        // 3 logical calls → 3 inner RPCs (no coalescing).
        assert_eq!(
            fake.record_requests().len(),
            3,
            "max_batch=1 must pass through every submission as its own RPC"
        );
        assert_eq!(
            count_audit_rows(&h.db_path, &h.key, "llm.sampling_call"),
            3
        );
    }
}