steep/crates/steep-server/src/open_router.rs

//! Minimal OpenRouter chat-completions client.
//!
//! Steep only needs a non-streaming `POST /chat/completions` with optional
//! tool calling, so the wire types are implemented in-tree instead of pulling a
//! third-party SDK (and its own `reqwest` version) into the workspace.
//!
//! Only the response fields Steep consumes are modelled: `content`,
//! `tool_calls` and `usage.cost`. Unknown fields are ignored.

use std::time::Duration;

use anyhow::Context;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tracing::instrument;

/// OpenRouter API root, version prefix included.
const BASE_URL: &str = "https://openrouter.ai/api/v1";

/// What the model is allowed to do on a given turn.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolChoice {
    /// The model decides whether to call a tool.
    Auto,
    /// The model must answer, without calling any tool. The tools stay declared,
    /// so the conversation remains the same shape as on the other turns.
    None,
}

impl ToolChoice {
    fn as_str(self) -> &'static str {
        match self {
            Self::Auto => "auto",
            Self::None => "none",
        }
    }
}

/// Result of a completion that may contain tool calls.
pub struct ToolChatResult {
    pub message: Option<String>,
    pub tool_calls: Vec<ToolCall>,
    pub cost: Option<f64>,
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    User,
    Assistant,
    Tool,
}

/// A single turn of the conversation.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Message {
    pub role: Role,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

impl Message {
    pub fn new(role: Role, content: impl Into<String>) -> Self {
        Self {
            role,
            content: Some(content.into()),
            tool_calls: None,
            tool_call_id: None,
        }
    }

    /// Assistant turn requesting tool calls. Models often answer with tool
    /// calls but no text, in which case the content is sent as `null`.
    pub fn assistant_with_tool_calls(
        content: impl Into<String>,
        tool_calls: Vec<ToolCall>,
    ) -> Self {
        let content = content.into();

        Self {
            role: Role::Assistant,
            content: (!content.is_empty()).then_some(content),
            tool_calls: Some(tool_calls),
            tool_call_id: None,
        }
    }

    /// Result of a tool call, linked to the request by `tool_call_id`.
    pub fn tool_response(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            role: Role::Tool,
            content: Some(content.into()),
            tool_calls: None,
            tool_call_id: Some(tool_call_id.into()),
        }
    }
}

/// A tool the model may call.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Tool {
    #[serde(rename = "type", default = "function_type")]
    pub kind: String,
    pub function: FunctionDefinition,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct FunctionDefinition {
    pub name: String,
    pub description: String,
    /// JSON schema describing the accepted arguments.
    pub parameters: Value,
}

impl Tool {
    pub fn new(name: impl Into<String>, description: impl Into<String>, parameters: Value) -> Self {
        Self {
            kind: function_type(),
            function: FunctionDefinition {
                name: name.into(),
                description: description.into(),
                parameters,
            },
        }
    }
}

/// A tool call requested by the model.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ToolCall {
    pub id: String,
    #[serde(rename = "type", default = "function_type")]
    pub kind: String,
    pub function: FunctionCall,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct FunctionCall {
    pub name: String,
    /// Arguments as a JSON-encoded string. Kept verbatim so that echoing the
    /// call back into the conversation does not re-encode or corrupt it.
    pub arguments: String,
}

impl ToolCall {
    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn name(&self) -> &str {
        &self.function.name
    }

    pub fn arguments_json(&self) -> &str {
        &self.function.arguments
    }
}

fn function_type() -> String {
    String::from("function")
}

#[derive(Serialize)]
struct ChatRequest<'a> {
    model: &'a str,
    messages: &'a [Message],
    reasoning: Reasoning,
    tools: &'a [Tool],
    tool_choice: &'a str,
}

/// OpenRouter reasoning config; Steep always asks for reasoning.
#[derive(Serialize)]
struct Reasoning {
    enabled: bool,
}

#[derive(Deserialize)]
struct ChatResponse {
    #[serde(default)]
    choices: Vec<Choice>,
    #[serde(default)]
    usage: Option<Usage>,
}

#[derive(Deserialize)]
struct Choice {
    message: ResponseMessage,
}

#[derive(Deserialize)]
struct ResponseMessage {
    #[serde(default)]
    content: Option<Content>,
    #[serde(default)]
    tool_calls: Option<Vec<ToolCall>>,
}

/// Message content is a plain string for most models and a list of typed parts
/// for others; both are collapsed to text.
#[derive(Deserialize)]
#[serde(untagged)]
enum Content {
    Text(String),
    Parts(Vec<ContentPart>),
}

#[derive(Deserialize)]
struct ContentPart {
    #[serde(default)]
    text: Option<String>,
}

impl Content {
    fn into_text(self) -> String {
        match self {
            Self::Text(text) => text,
            Self::Parts(parts) => parts
                .into_iter()
                .filter_map(|part| part.text)
                .collect::<Vec<_>>()
                .join(""),
        }
    }
}

#[derive(Deserialize)]
struct Usage {
    #[serde(default)]
    cost: Option<f64>,
}

/// Error payload returned by OpenRouter for a failed request.
#[derive(Deserialize)]
struct ErrorResponse {
    error: ErrorDetail,
}

#[derive(Deserialize)]
struct ErrorDetail {
    message: String,
}

#[derive(Clone)]
pub struct OpenRouterClient {
    client: reqwest::Client,
    api_key: String,
    model: String,
}

impl OpenRouterClient {
    pub fn new(token: &str, model: &str, timeout: Duration) -> anyhow::Result<Self> {
        Ok(Self {
            client: reqwest::Client::builder().timeout(timeout).build()?,
            api_key: String::from(token),
            model: String::from(model),
        })
    }

    /// The model every completion is requested from.
    pub fn model(&self) -> &str {
        &self.model
    }

    /// Sends a conversation with tool definitions and returns either a final
    /// message or the tool calls requested by the model.
    ///
    /// The conversation is borrowed rather than taken by value: the caller keeps
    /// appending to it between iterations, and a copy per iteration would be pure
    /// waste.
    #[instrument(skip(self, messages, tools), err)]
    pub async fn chat_with_tools(
        &self,
        messages: &[Message],
        tools: &[Tool],
        tool_choice: ToolChoice,
    ) -> anyhow::Result<ToolChatResult> {
        let response = self.complete(messages, tools, tool_choice).await?;

        let cost = response.usage.and_then(|usage| usage.cost);
        let message = response
            .choices
            .into_iter()
            .next()
            .ok_or_else(|| anyhow::anyhow!("No choices in response"))?
            .message;

        Ok(ToolChatResult {
            message: message.content.map(Content::into_text),
            tool_calls: message.tool_calls.unwrap_or_default(),
            cost,
        })
    }

    async fn complete(
        &self,
        messages: &[Message],
        tools: &[Tool],
        tool_choice: ToolChoice,
    ) -> anyhow::Result<ChatResponse> {
        let request = ChatRequest {
            model: &self.model,
            messages,
            reasoning: Reasoning { enabled: true },
            tools,
            tool_choice: tool_choice.as_str(),
        };

        let response = self
            .client
            .post(format!("{BASE_URL}/chat/completions"))
            .bearer_auth(&self.api_key)
            .json(&request)
            .send()
            .await
            .context("failed to reach OpenRouter")?;

        let status = response.status();
        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            anyhow::bail!("OpenRouter returned {status}: {}", error_message(&body));
        }

        response
            .json::<ChatResponse>()
            .await
            .context("invalid OpenRouter response")
    }
}

/// Extracts the message from an OpenRouter error body, falling back to the raw
/// body when it is not the expected JSON shape.
fn error_message(body: &str) -> String {
    serde_json::from_str::<ErrorResponse>(body)
        .map(|response| response.error.message)
        .unwrap_or_else(|_| body.trim().to_string())
}

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

    /// Builds a tool call the way the API returns one, so the fixture also
    /// covers deserialization.
    fn tool_call(id: &str, name: &str, arguments: &str) -> ToolCall {
        serde_json::from_value(json!({
            "id": id,
            "type": "function",
            "function": { "name": name, "arguments": arguments }
        }))
        .unwrap()
    }

    #[test]
    fn tool_call_arguments_are_kept_verbatim() {
        let call = tool_call("call_1", "grep", r#"{"pattern":"fn main"}"#);
        let message = Message::assistant_with_tool_calls("", vec![call]);

        let serialized = serde_json::to_value(&message).unwrap();

        assert_eq!(serialized["content"], Value::Null);
        assert_eq!(serialized["tool_calls"][0]["type"], "function");
        assert_eq!(serialized["tool_calls"][0]["function"]["name"], "grep");
        assert_eq!(
            serialized["tool_calls"][0]["function"]["arguments"],
            r#"{"pattern":"fn main"}"#
        );
    }

    #[test]
    fn tool_response_carries_role_and_tool_call_id() {
        let message = Message::tool_response("call_1", "src/main.rs");
        let serialized = serde_json::to_value(&message).unwrap();

        assert_eq!(serialized["role"], "tool");
        assert_eq!(serialized["tool_call_id"], "call_1");
        assert_eq!(serialized["content"], "src/main.rs");
    }

    #[test]
    fn request_serializes_tool_choice_and_reasoning() {
        let messages = [Message::new(Role::User, "hi")];
        let tools = [Tool::new("ls", "List files", json!({"type": "object"}))];

        let request = serde_json::to_value(ChatRequest {
            model: "some/model",
            messages: &messages,
            reasoning: Reasoning { enabled: true },
            tools: &tools,
            tool_choice: ToolChoice::Auto.as_str(),
        })
        .unwrap();

        assert_eq!(request["model"], "some/model");
        assert_eq!(request["reasoning"]["enabled"], true);
        assert_eq!(request["tool_choice"], "auto");
        assert_eq!(request["tools"][0]["function"]["name"], "ls");
        assert_eq!(request["messages"][0]["role"], "user");
    }

    #[test]
    fn a_forced_answer_turn_disables_the_tools() {
        assert_eq!(ToolChoice::None.as_str(), "none");
        assert_eq!(ToolChoice::Auto.as_str(), "auto");
    }

    #[test]
    fn response_parses_tool_calls_and_cost() {
        let response: ChatResponse = serde_json::from_value(json!({
            "choices": [{
                "message": {
                    "role": "assistant",
                    "content": null,
                    "tool_calls": [{
                        "id": "call_1",
                        "type": "function",
                        "function": { "name": "ls", "arguments": "{\"path\":\"src\"}" }
                    }]
                }
            }],
            "usage": { "cost": 0.0021 }
        }))
        .unwrap();

        assert_eq!(response.usage.and_then(|usage| usage.cost), Some(0.0021));

        let message = response.choices.into_iter().next().unwrap().message;
        assert!(message.content.is_none());

        let calls = message.tool_calls.unwrap();
        assert_eq!(calls[0].name(), "ls");
        assert_eq!(calls[0].id(), "call_1");
        assert_eq!(calls[0].arguments_json(), r#"{"path":"src"}"#);
    }

    #[test]
    fn response_parses_parts_content() {
        let response: ChatResponse = serde_json::from_value(json!({
            "choices": [{
                "message": {
                    "role": "assistant",
                    "content": [
                        { "type": "text", "text": "hello " },
                        { "type": "text", "text": "world" }
                    ]
                }
            }]
        }))
        .unwrap();

        assert!(response.usage.is_none());

        let message = response.choices.into_iter().next().unwrap().message;
        assert_eq!(message.content.unwrap().into_text(), "hello world");
    }

    #[test]
    fn error_message_prefers_api_message() {
        let body = r#"{"error":{"message":"No auth credentials found","code":401}}"#;
        assert_eq!(error_message(body), "No auth credentials found");
    }

    #[test]
    fn error_message_falls_back_to_raw_body() {
        assert_eq!(
            error_message("  <html>bad gateway</html>  "),
            "<html>bad gateway</html>"
        );
    }
}