storkit/server/src/http/anthropic.rs

use crate::http::context::{AppContext, OpenApiResult, bad_request};
use crate::llm::chat;
use crate::store::StoreOps;
use poem_openapi::{Object, OpenApi, Tags, payload::Json};
use reqwest::header::{HeaderMap, HeaderValue};
use serde::Deserialize;
use std::sync::Arc;

const ANTHROPIC_MODELS_URL: &str = "https://api.anthropic.com/v1/models";
const ANTHROPIC_VERSION: &str = "2023-06-01";
const KEY_ANTHROPIC_API_KEY: &str = "anthropic_api_key";

#[derive(Deserialize)]
struct AnthropicModelsResponse {
    data: Vec<AnthropicModelInfo>,
}

#[derive(Deserialize)]
struct AnthropicModelInfo {
    id: String,
}

fn get_anthropic_api_key(ctx: &AppContext) -> Result<String, String> {
    match ctx.store.get(KEY_ANTHROPIC_API_KEY) {
        Some(value) => {
            if let Some(key) = value.as_str() {
                if key.is_empty() {
                    Err("Anthropic API key is empty. Please set your API key.".to_string())
                } else {
                    Ok(key.to_string())
                }
            } else {
                Err("Stored API key is not a string".to_string())
            }
        }
        None => Err("Anthropic API key not found. Please set your API key.".to_string()),
    }
}

#[derive(Deserialize, Object)]
struct ApiKeyPayload {
    api_key: String,
}

#[derive(Tags)]
enum AnthropicTags {
    Anthropic,
}

pub struct AnthropicApi {
    ctx: Arc<AppContext>,
}

impl AnthropicApi {
    pub fn new(ctx: Arc<AppContext>) -> Self {
        Self { ctx }
    }
}

#[OpenApi(tag = "AnthropicTags::Anthropic")]
impl AnthropicApi {
    /// Check whether an Anthropic API key is stored.
    ///
    /// Returns `true` if a non-empty key is present, otherwise `false`.
    #[oai(path = "/anthropic/key/exists", method = "get")]
    async fn get_anthropic_api_key_exists(&self) -> OpenApiResult<Json<bool>> {
        let exists =
            chat::get_anthropic_api_key_exists(self.ctx.store.as_ref()).map_err(bad_request)?;
        Ok(Json(exists))
    }

    /// Store or update the Anthropic API key used for requests.
    ///
    /// Returns `true` when the key is saved successfully.
    #[oai(path = "/anthropic/key", method = "post")]
    async fn set_anthropic_api_key(
        &self,
        payload: Json<ApiKeyPayload>,
    ) -> OpenApiResult<Json<bool>> {
        chat::set_anthropic_api_key(self.ctx.store.as_ref(), payload.0.api_key)
            .map_err(bad_request)?;
        Ok(Json(true))
    }

    /// List available Anthropic models.
    #[oai(path = "/anthropic/models", method = "get")]
    async fn list_anthropic_models(&self) -> OpenApiResult<Json<Vec<String>>> {
        self.list_anthropic_models_from(ANTHROPIC_MODELS_URL).await
    }
}

impl AnthropicApi {
    async fn list_anthropic_models_from(
        &self,
        url: &str,
    ) -> OpenApiResult<Json<Vec<String>>> {
        let api_key = get_anthropic_api_key(self.ctx.as_ref()).map_err(bad_request)?;
        let client = reqwest::Client::new();
        let mut headers = HeaderMap::new();
        headers.insert(
            "x-api-key",
            HeaderValue::from_str(&api_key).map_err(|e| bad_request(e.to_string()))?,
        );
        headers.insert(
            "anthropic-version",
            HeaderValue::from_static(ANTHROPIC_VERSION),
        );

        let response = client
            .get(url)
            .headers(headers)
            .send()
            .await
            .map_err(|e| bad_request(e.to_string()))?;

        if !response.status().is_success() {
            let status = response.status();
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(bad_request(format!(
                "Anthropic API error {status}: {error_text}"
            )));
        }

        let body = response
            .json::<AnthropicModelsResponse>()
            .await
            .map_err(|e| bad_request(e.to_string()))?;
        let models = body.data.into_iter().map(|m| m.id).collect();

        Ok(Json(models))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::http::context::AppContext;
    use serde_json::json;
    use std::sync::Arc;
    use tempfile::TempDir;

    fn test_ctx(dir: &TempDir) -> AppContext {
        AppContext::new_test(dir.path().to_path_buf())
    }

    fn make_api(dir: &TempDir) -> AnthropicApi {
        AnthropicApi::new(Arc::new(test_ctx(dir)))
    }

    // -- get_anthropic_api_key (private helper) --

    #[test]
    fn get_api_key_returns_err_when_not_set() {
        let dir = TempDir::new().unwrap();
        let ctx = test_ctx(&dir);
        let result = get_anthropic_api_key(&ctx);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("not found"));
    }

    #[test]
    fn get_api_key_returns_err_when_empty() {
        let dir = TempDir::new().unwrap();
        let ctx = test_ctx(&dir);
        ctx.store.set(KEY_ANTHROPIC_API_KEY, json!(""));
        let result = get_anthropic_api_key(&ctx);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("empty"));
    }

    #[test]
    fn get_api_key_returns_err_when_not_string() {
        let dir = TempDir::new().unwrap();
        let ctx = test_ctx(&dir);
        ctx.store.set(KEY_ANTHROPIC_API_KEY, json!(12345));
        let result = get_anthropic_api_key(&ctx);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("not a string"));
    }

    #[test]
    fn get_api_key_returns_key_when_set() {
        let dir = TempDir::new().unwrap();
        let ctx = test_ctx(&dir);
        ctx.store.set(KEY_ANTHROPIC_API_KEY, json!("sk-ant-test123"));
        let result = get_anthropic_api_key(&ctx);
        assert_eq!(result.unwrap(), "sk-ant-test123");
    }

    // -- get_anthropic_api_key_exists endpoint --

    #[tokio::test]
    async fn key_exists_returns_false_when_not_set() {
        let dir = TempDir::new().unwrap();
        let api = make_api(&dir);
        let result = api.get_anthropic_api_key_exists().await.unwrap();
        assert!(!result.0);
    }

    #[tokio::test]
    async fn key_exists_returns_true_when_set() {
        let dir = TempDir::new().unwrap();
        let ctx = AppContext::new_test(dir.path().to_path_buf());
        ctx.store.set(KEY_ANTHROPIC_API_KEY, json!("sk-ant-test123"));
        let api = AnthropicApi::new(Arc::new(ctx));
        let result = api.get_anthropic_api_key_exists().await.unwrap();
        assert!(result.0);
    }

    // -- set_anthropic_api_key endpoint --

    #[tokio::test]
    async fn set_api_key_returns_true() {
        let dir = TempDir::new().unwrap();
        let api = make_api(&dir);
        let payload = Json(ApiKeyPayload {
            api_key: "sk-ant-test123".to_string(),
        });
        let result = api.set_anthropic_api_key(payload).await.unwrap();
        assert!(result.0);
    }

    #[tokio::test]
    async fn set_then_exists_returns_true() {
        let dir = TempDir::new().unwrap();
        let ctx = Arc::new(AppContext::new_test(dir.path().to_path_buf()));
        let api = AnthropicApi::new(ctx);
        api.set_anthropic_api_key(Json(ApiKeyPayload {
            api_key: "sk-ant-test123".to_string(),
        }))
        .await
        .unwrap();
        let result = api.get_anthropic_api_key_exists().await.unwrap();
        assert!(result.0);
    }

    // -- list_anthropic_models endpoint --

    #[tokio::test]
    async fn list_models_fails_when_no_key() {
        let dir = TempDir::new().unwrap();
        let api = make_api(&dir);
        let result = api.list_anthropic_models().await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn list_models_fails_with_invalid_header_value() {
        let dir = TempDir::new().unwrap();
        let ctx = AppContext::new_test(dir.path().to_path_buf());
        // A header value containing a newline is invalid
        ctx.store
            .set(KEY_ANTHROPIC_API_KEY, json!("bad\nvalue"));
        let api = AnthropicApi::new(Arc::new(ctx));
        let result = api.list_anthropic_models_from("http://127.0.0.1:1").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn list_models_fails_when_server_unreachable() {
        let dir = TempDir::new().unwrap();
        let ctx = AppContext::new_test(dir.path().to_path_buf());
        ctx.store
            .set(KEY_ANTHROPIC_API_KEY, json!("sk-ant-test123"));
        let api = AnthropicApi::new(Arc::new(ctx));
        // Port 1 is reserved and should immediately refuse the connection
        let result = api.list_anthropic_models_from("http://127.0.0.1:1").await;
        assert!(result.is_err());
    }

    #[test]
    fn new_creates_api_instance() {
        let dir = TempDir::new().unwrap();
        let _api = make_api(&dir);
    }
}