klieo-mcp-server 2.2.0

//! MCP `sampling/createMessage` wire-format types.
//!
//! Sampling is the server-initiated path that asks the MCP client to
//! run an LLM completion on the server's behalf. The shape of the
//! request/response envelope is fixed by the MCP spec revision the
//! crate advertises ([`crate::MCP_PROTOCOL_VERSION`]); this module
//! captures it as plain Rust types with `serde` renames mapping
//! the snake_case Rust fields onto the camelCase JSON the wire uses.
//!
//! ## Text-content-only scope
//!
//! `SamplingContent` carries the `text` variant only. Image, audio, and
//! resource content blocks are part of the MCP spec but intentionally
//! out of scope here — a content blob tagged `image` / `audio` /
//! `resource` will fail to deserialise with serde's "unknown variant"
//! error, surfacing as an `UnsupportedContent` failure at the call
//! site rather than silently dropping non-text payloads. See ADR-027
//! for the rationale.
//!
//! ## Wire-format renames
//!
//! Field names on the JSON wire follow MCP's camelCase convention
//! (`maxTokens`, `systemPrompt`, `modelPreferences`, `stopSequences`,
//! `costPriority`, `speedPriority`, `intelligencePriority`,
//! `stopReason`). Rust fields use snake_case and rely on
//! `#[serde(rename = "...")]` to bridge the two. Optional fields use
//! `skip_serializing_if = "Option::is_none"` so the serialised
//! payload omits absent values rather than emitting `null`s the spec
//! does not require.

use serde::{Deserialize, Serialize};

/// One `sampling/createMessage` request envelope sent from a klieo MCP
/// server to its client. Carries the message history plus optional
/// model preferences, a system prompt, sampling controls, and an upper
/// bound on the number of tokens the client may generate.
#[non_exhaustive]
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SamplingRequest {
    /// Conversation history the client should complete. Ordered
    /// oldest-first; the client's reply continues from the last entry.
    pub messages: Vec<SamplingMessage>,

    /// Optional hints to the client about which model family / cost
    /// / latency / intelligence profile to use. Clients MAY ignore
    /// these and pick their own model.
    #[serde(rename = "modelPreferences", skip_serializing_if = "Option::is_none")]
    pub model_preferences: Option<ModelPreferences>,

    /// Optional system prompt to prepend to the conversation. When
    /// `None`, the client uses its own default (or no system prompt).
    #[serde(rename = "systemPrompt", skip_serializing_if = "Option::is_none")]
    pub system_prompt: Option<String>,

    /// Upper bound on the number of tokens the client may generate.
    /// Required by the MCP spec; clients MUST honour it.
    #[serde(rename = "maxTokens")]
    pub max_tokens: u32,

    /// Optional sampling temperature. Range and default are
    /// client-defined; common practice is 0.0–2.0 with 1.0 as default.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,

    /// Optional list of stop sequences. Generation halts as soon as
    /// the model emits any of these strings.
    #[serde(rename = "stopSequences", skip_serializing_if = "Option::is_none")]
    pub stop_sequences: Option<Vec<String>>,
}

impl SamplingRequest {
    /// Constructs a new [`SamplingRequest`] with required fields.
    /// Optional fields default to `None`.
    pub fn new(messages: Vec<SamplingMessage>, max_tokens: u32) -> Self {
        Self {
            messages,
            model_preferences: None,
            system_prompt: None,
            max_tokens,
            temperature: None,
            stop_sequences: None,
        }
    }
}

/// One message in a [`SamplingRequest::messages`] history. The MCP
/// spec defines `role` as the free-form string `"user"` or
/// `"assistant"`; richer typing is intentionally deferred to the
/// transport boundary.
#[non_exhaustive]
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SamplingMessage {
    /// Speaker role. Wire values: `"user"` or `"assistant"`.
    pub role: String,

    /// Message payload. Only [`SamplingContent::Text`] is supported
    /// in this cluster; see the module-level note on unsupported
    /// content variants.
    pub content: SamplingContent,
}

impl SamplingMessage {
    /// Constructs a new [`SamplingMessage`] with the given role and content.
    ///
    /// Required for external callers after `#[non_exhaustive]` prevents
    /// struct-literal construction outside the crate.
    pub fn new(role: impl Into<String>, content: SamplingContent) -> Self {
        Self {
            role: role.into(),
            content,
        }
    }
}

/// Discriminated content payload for a [`SamplingMessage`] or
/// [`SamplingResponse`]. Serialised with `{"type":"<variant>", ...}`
/// per the MCP wire format.
///
/// Only the `text` variant is implemented in this cluster. An
/// incoming blob tagged `"image"`, `"audio"`, or `"resource"` will
/// fail to deserialise with serde's "unknown variant" error, which
/// the call site surfaces as an unsupported-content failure. See
/// ADR-027 for the scope decision.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "lowercase")]
#[non_exhaustive]
pub enum SamplingContent {
    /// Plain UTF-8 text payload.
    Text {
        /// The text content of the message.
        text: String,
    },
}

/// Optional hints from the server to the client about which model to
/// use for a [`SamplingRequest`]. All fields are advisory — the
/// client may pick any model it has access to and ignore these.
///
/// Priority fields are normalised to the range `0.0..=1.0`; the
/// client weighs them against its own catalogue when choosing a model.
#[non_exhaustive]
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct ModelPreferences {
    /// Optional ordered list of model-name hints. The first hint
    /// whose name the client recognises wins.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub hints: Option<Vec<ModelHint>>,

    /// Optional cost-priority weight (`0.0` = cost-insensitive, `1.0`
    /// = cheapest model preferred).
    #[serde(rename = "costPriority", skip_serializing_if = "Option::is_none")]
    pub cost_priority: Option<f32>,

    /// Optional speed-priority weight (`0.0` = latency-insensitive,
    /// `1.0` = fastest model preferred).
    #[serde(rename = "speedPriority", skip_serializing_if = "Option::is_none")]
    pub speed_priority: Option<f32>,

    /// Optional intelligence-priority weight (`0.0` = simplest model
    /// acceptable, `1.0` = most capable model preferred).
    #[serde(
        rename = "intelligencePriority",
        skip_serializing_if = "Option::is_none"
    )]
    pub intelligence_priority: Option<f32>,
}

/// One model-name hint inside [`ModelPreferences::hints`]. The wire
/// shape is `{"name": "..."}` per the MCP spec.
#[non_exhaustive]
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ModelHint {
    /// Model family / name string the server is hinting at (e.g.
    /// `"claude-3-5-sonnet"`, `"gpt-4o"`).
    pub name: String,
}

/// The `sampling/createMessage` response envelope returned by an MCP
/// client. Carries the generated message, the model the client
/// actually used, and an optional `stopReason` describing why
/// generation halted.
#[non_exhaustive]
#[derive(Clone, Debug, Deserialize)]
pub struct SamplingResponse {
    /// Speaker role of the generated message. Typically
    /// `"assistant"`.
    pub role: String,

    /// Generated content payload.
    pub content: SamplingContent,

    /// Identifier of the model the client used to generate
    /// `content`. May differ from any hint the server supplied.
    pub model: String,

    /// Optional reason generation stopped (e.g. `"endTurn"`,
    /// `"maxTokens"`, `"stopSequence"`). Wire field is
    /// camelCase `stopReason`.
    #[serde(rename = "stopReason")]
    pub stop_reason: Option<String>,
}

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

    #[test]
    fn sampling_request_serialises_camelcase_field_names() {
        let req = SamplingRequest {
            messages: vec![SamplingMessage {
                role: "user".into(),
                content: SamplingContent::Text {
                    text: "hello".into(),
                },
            }],
            model_preferences: Some(ModelPreferences {
                hints: Some(vec![ModelHint {
                    name: "claude-3-5-sonnet".into(),
                }]),
                cost_priority: Some(0.2),
                speed_priority: Some(0.5),
                intelligence_priority: Some(0.9),
            }),
            system_prompt: Some("be concise".into()),
            max_tokens: 256,
            temperature: Some(0.7),
            stop_sequences: Some(vec!["STOP".into()]),
        };
        let json = serde_json::to_value(&req).expect("serialises");
        assert!(json.get("maxTokens").is_some(), "maxTokens key present");
        assert!(json.get("max_tokens").is_none(), "no snake_case leak");
        assert!(
            json.get("systemPrompt").is_some(),
            "systemPrompt key present"
        );
        assert!(
            json.get("modelPreferences").is_some(),
            "modelPreferences key present"
        );
        assert!(
            json.get("stopSequences").is_some(),
            "stopSequences key present"
        );
        let prefs = &json["modelPreferences"];
        assert!(prefs.get("costPriority").is_some(), "costPriority present");
        assert!(
            prefs.get("speedPriority").is_some(),
            "speedPriority present"
        );
        assert!(
            prefs.get("intelligencePriority").is_some(),
            "intelligencePriority present"
        );

        // Absent optional fields are omitted, not serialised as null.
        let minimal = SamplingRequest {
            messages: vec![],
            model_preferences: None,
            system_prompt: None,
            max_tokens: 16,
            temperature: None,
            stop_sequences: None,
        };
        let minimal_json = serde_json::to_value(&minimal).expect("serialises");
        assert!(minimal_json.get("systemPrompt").is_none());
        assert!(minimal_json.get("modelPreferences").is_none());
        assert!(minimal_json.get("stopSequences").is_none());
        assert!(minimal_json.get("temperature").is_none());
    }

    #[test]
    fn sampling_response_deserialises_text_content() {
        let fixture = r#"{"role":"assistant","content":{"type":"text","text":"hi"},"model":"test","stopReason":"endTurn"}"#;
        let resp: SamplingResponse = serde_json::from_str(fixture).expect("parses");
        assert_eq!(resp.role, "assistant");
        assert_eq!(resp.model, "test");
        assert_eq!(resp.stop_reason.as_deref(), Some("endTurn"));
        match resp.content {
            SamplingContent::Text { text } => assert_eq!(text, "hi"),
        }
    }

    #[test]
    fn sampling_content_image_fails_to_deserialise() {
        let fixture = r#"{"type":"image","data":"...","mimeType":"image/png"}"#;
        let result: Result<SamplingContent, _> = serde_json::from_str(fixture);
        assert!(
            result.is_err(),
            "image-tagged content must surface as unknown-variant error"
        );
    }

    #[test]
    fn sampling_request_new_defaults_optional_fields_to_none() {
        let msg = SamplingMessage::new("user", SamplingContent::Text { text: "hi".into() });
        let req = SamplingRequest::new(vec![msg], 128);
        assert_eq!(req.messages.len(), 1);
        assert_eq!(req.max_tokens, 128);
        assert!(req.model_preferences.is_none());
        assert!(req.system_prompt.is_none());
        assert!(req.temperature.is_none());
        assert!(req.stop_sequences.is_none());
    }

    #[test]
    fn sampling_message_new_accepts_str_and_string_role() {
        let m1 = SamplingMessage::new("user", SamplingContent::Text { text: "a".into() });
        let m2 = SamplingMessage::new(
            String::from("assistant"),
            SamplingContent::Text { text: "b".into() },
        );
        assert_eq!(m1.role, "user");
        assert_eq!(m2.role, "assistant");
        assert!(matches!(m1.content, SamplingContent::Text { .. }));
    }
}