Skip to main content

offline_intelligence/api/
api_keys_api.rs

1//! API Keys Management Endpoints
2//!
3//! Provides REST endpoints for managing API keys:
4//! - Save/update HuggingFace and OpenRouter keys (encrypted with machine-specific key, stored in SQLite)
5//! - Retrieve keys (plaintext - decrypted using machine-specific key)
6//! - Delete keys
7//! - Mark keys as used with mode tracking
8//!
9//! Encryption/decryption is handled entirely inside `ApiKeysStore`.
10//! This layer only deals with plain-text values.
11
12use axum::{
13    extract::{Query, State},
14    http::StatusCode,
15    response::IntoResponse,
16    Json,
17};
18use serde::{Deserialize, Serialize};
19use tracing::{error, info, warn};
20
21use crate::{
22    memory_db::{ApiKeyType, ApiKeyRecord},
23    shared_state::UnifiedAppState,
24};
25
26// ─── Request / Response types ────────────────────────────────────────────────
27
28/// Request body for saving or updating an API key.
29#[derive(Debug, Deserialize)]
30pub struct SaveApiKeyRequest {
31    pub key_type: String, // "huggingface" or "openrouter"
32    pub value: String,    // Plain-text key value
33}
34
35/// Confirmation after saving an API key.
36#[derive(Debug, Serialize)]
37pub struct SaveApiKeyResponse {
38    pub success: bool,
39    pub message: String,
40}
41
42/// Response carrying a single API key value + metadata.
43#[derive(Debug, Serialize)]
44pub struct GetApiKeyResponse {
45    pub key_type: String,
46    pub value: Option<String>,
47    pub created_at: Option<String>,
48    pub last_used_at: Option<String>,
49    pub last_mode: Option<String>,
50    pub usage_count: Option<i64>,
51}
52
53/// Response carrying all API keys.
54#[derive(Debug, Serialize)]
55pub struct GetAllApiKeysResponse {
56    pub keys: Vec<GetApiKeyResponse>,
57}
58
59// ─── Helpers ─────────────────────────────────────────────────────────────────
60
61fn record_to_response(record: &ApiKeyRecord, value: Option<String>) -> GetApiKeyResponse {
62    GetApiKeyResponse {
63        key_type: record.key_type.clone(),
64        value,
65        created_at: Some(record.created_at.to_rfc3339()),
66        last_used_at: record.last_used_at.map(|dt| dt.to_rfc3339()),
67        last_mode: record.last_mode.clone(),
68        usage_count: Some(record.usage_count),
69    }
70}
71
72// ─── Handlers ────────────────────────────────────────────────────────────────
73
74/// `POST /api-keys` — save or update an API key.
75///
76/// The plaintext value is encrypted using machine-specific encryption and stored in SQLite.
77pub async fn save_api_key(
78    State(state): State<UnifiedAppState>,
79    Json(payload): Json<SaveApiKeyRequest>,
80) -> Result<impl IntoResponse, StatusCode> {
81    let key_type = ApiKeyType::from_str(&payload.key_type).ok_or(StatusCode::BAD_REQUEST)?;
82
83    info!("Saving API key for: {}", key_type.as_str());
84
85    match state
86        .shared_state
87        .database_pool
88        .api_keys
89        .save_key(key_type, &payload.value)
90    {
91        Ok(_) => Ok(Json(SaveApiKeyResponse {
92            success: true,
93            message: format!("API key saved for {}", payload.key_type),
94        })),
95        Err(e) => {
96            error!("Failed to save API key: {}", e);
97            Err(StatusCode::INTERNAL_SERVER_ERROR)
98        }
99    }
100}
101
102/// `GET /api-keys?key_type=<type>` — retrieve a single API key (plaintext).
103pub async fn get_api_key(
104    State(state): State<UnifiedAppState>,
105    Query(params): Query<std::collections::HashMap<String, String>>,
106) -> Result<impl IntoResponse, StatusCode> {
107    let key_type_str = params.get("key_type").ok_or(StatusCode::BAD_REQUEST)?;
108    let key_type = ApiKeyType::from_str(key_type_str).ok_or(StatusCode::BAD_REQUEST)?;
109
110    // Get plaintext (decrypts using machine-specific key).
111    let value = state
112        .shared_state
113        .database_pool
114        .api_keys
115        .get_key_plaintext(&key_type)
116        .map_err(|e| {
117            error!("Failed to retrieve API key '{}': {}", key_type_str, e);
118            StatusCode::INTERNAL_SERVER_ERROR
119        })?;
120
121    // Fetch metadata (created_at, last_used_at, …) for the response body.
122    let metadata = state
123        .shared_state
124        .database_pool
125        .api_keys
126        .get_key_metadata(&key_type)
127        .map_err(|e| {
128            error!("Failed to get API key metadata '{}': {}", key_type_str, e);
129            StatusCode::INTERNAL_SERVER_ERROR
130        })?;
131
132    Ok(Json(match metadata {
133        Some(record) => record_to_response(&record, value),
134        None => GetApiKeyResponse {
135            key_type: key_type_str.clone(),
136            value: None,
137            created_at: None,
138            last_used_at: None,
139            last_mode: None,
140            usage_count: None,
141        },
142    }))
143}
144
145/// `GET /api-keys/all` — retrieve all API keys with their plaintext values.
146pub async fn get_all_api_keys(
147    State(state): State<UnifiedAppState>,
148) -> Result<impl IntoResponse, StatusCode> {
149    match state
150        .shared_state
151        .database_pool
152        .api_keys
153        .get_all_keys_with_values()
154    {
155        Ok(entries) => {
156            let keys: Vec<GetApiKeyResponse> = entries
157                .into_iter()
158                .map(|(record, value)| record_to_response(&record, Some(value)))
159                .collect();
160            Ok(Json(GetAllApiKeysResponse { keys }))
161        }
162        Err(e) => {
163            error!("Failed to retrieve all API keys: {}", e);
164            Err(StatusCode::INTERNAL_SERVER_ERROR)
165        }
166    }
167}
168
169/// `DELETE /api-keys?key_type=<type>` — delete an API key from storage.
170pub async fn delete_api_key(
171    State(state): State<UnifiedAppState>,
172    Query(params): Query<std::collections::HashMap<String, String>>,
173) -> Result<impl IntoResponse, StatusCode> {
174    let key_type_str = params.get("key_type").ok_or(StatusCode::BAD_REQUEST)?;
175    let key_type = ApiKeyType::from_str(key_type_str).ok_or(StatusCode::BAD_REQUEST)?;
176
177    match state
178        .shared_state
179        .database_pool
180        .api_keys
181        .delete_key(key_type)
182    {
183        Ok(true) => {
184            info!("API key deleted for: {}", key_type_str);
185            Ok(Json(serde_json::json!({
186                "success": true,
187                "message": format!("API key deleted for {}", key_type_str)
188            })))
189        }
190        Ok(false) => {
191            warn!("API key not found for deletion: {}", key_type_str);
192            Ok(Json(serde_json::json!({
193                "success": false,
194                "message": format!("API key not found for {}", key_type_str)
195            })))
196        }
197        Err(e) => {
198            error!("Failed to delete API key: {}", e);
199            Err(StatusCode::INTERNAL_SERVER_ERROR)
200        }
201    }
202}
203
204/// `POST /api-keys/mark-used` — record that a key was used with a given mode.
205#[derive(Debug, Deserialize)]
206pub struct MarkKeyUsedRequest {
207    pub key_type: String,
208    pub mode: String, // "offline" or "online"
209}
210
211pub async fn mark_key_used(
212    State(state): State<UnifiedAppState>,
213    Json(payload): Json<MarkKeyUsedRequest>,
214) -> Result<impl IntoResponse, StatusCode> {
215    let key_type = ApiKeyType::from_str(&payload.key_type).ok_or(StatusCode::BAD_REQUEST)?;
216
217    match state
218        .shared_state
219        .database_pool
220        .api_keys
221        .mark_used(key_type, &payload.mode)
222    {
223        Ok(_) => Ok(Json(serde_json::json!({
224            "success": true,
225            "message": "Key usage recorded"
226        }))),
227        Err(e) => {
228            error!("Failed to mark key as used: {}", e);
229            Err(StatusCode::INTERNAL_SERVER_ERROR)
230        }
231    }
232}
233
234/// Request body for verifying an API key.
235#[derive(Debug, Deserialize)]
236pub struct VerifyApiKeyRequest {
237    pub key_type: String,
238    pub api_key: String,
239}
240
241/// Response for API key verification.
242#[derive(Debug, Serialize)]
243pub struct VerifyApiKeyResponse {
244    pub valid: bool,
245    pub message: String,
246}
247
248/// `POST /api-keys/verify` — verify an API key by calling the provider's API.
249///
250/// - For OpenRouter: Calls GET https://openrouter.ai/api/v1/key
251/// - For HuggingFace: Calls GET https://huggingface.co/api/whoami
252pub async fn verify_api_key(
253    State(state): State<UnifiedAppState>,
254    Json(payload): Json<VerifyApiKeyRequest>,
255) -> Result<impl IntoResponse, StatusCode> {
256    let key_type = ApiKeyType::from_str(&payload.key_type).ok_or(StatusCode::BAD_REQUEST)?;
257
258    if payload.api_key.trim().is_empty() {
259        return Ok(Json(VerifyApiKeyResponse {
260            valid: false,
261            message: "API key cannot be empty".to_string(),
262        }));
263    }
264
265    let client = reqwest::Client::new();
266    let http_client = state.http_client.clone();
267
268    match key_type {
269        ApiKeyType::OpenRouter => {
270            let url = "https://openrouter.ai/api/v1/key";
271            match http_client
272                .get(url)
273                .header("Authorization", format!("Bearer {}", payload.api_key))
274                .send()
275                .await
276            {
277                Ok(resp) => {
278                    if resp.status().is_success() {
279                        info!("OpenRouter API key verified successfully, saving to database");
280                        
281                        // Save the verified key to the database
282                        if let Err(e) = state.shared_state.database_pool.api_keys.save_key(key_type, &payload.api_key) {
283                            error!("Failed to save verified OpenRouter API key: {}", e);
284                            return Ok(Json(VerifyApiKeyResponse {
285                                valid: false,
286                                message: "Key verified but failed to save. Please try again.".to_string(),
287                            }));
288                        }
289                        
290                        Ok(Json(VerifyApiKeyResponse {
291                            valid: true,
292                            message: "OpenRouter API key is valid".to_string(),
293                        }))
294                    } else if resp.status() == StatusCode::UNAUTHORIZED || resp.status() == StatusCode::FORBIDDEN {
295                        info!("OpenRouter API key verification failed: invalid credentials");
296                        Ok(Json(VerifyApiKeyResponse {
297                            valid: false,
298                            message: "Invalid OpenRouter API key. Please check and try again.".to_string(),
299                        }))
300                    } else {
301                        let status = resp.status();
302                        error!("OpenRouter API key verification returned unexpected status: {}", status);
303                        Ok(Json(VerifyApiKeyResponse {
304                            valid: false,
305                            message: format!("OpenRouter API error: {}", status),
306                        }))
307                    }
308                }
309                Err(e) => {
310                    error!("Failed to verify OpenRouter API key: {}", e);
311                    Ok(Json(VerifyApiKeyResponse {
312                        valid: false,
313                        message: "Failed to connect to OpenRouter API. Please check your internet connection.".to_string(),
314                    }))
315                }
316            }
317        }
318        ApiKeyType::HuggingFace => {
319            let url = "https://huggingface.co/api/whoami";
320            match http_client
321                .get(url)
322                .header("Authorization", format!("Bearer {}", payload.api_key))
323                .send()
324                .await
325            {
326                Ok(resp) => {
327                    if resp.status().is_success() {
328                        info!("HuggingFace token verified successfully, saving to database");
329                        
330                        // Save the verified token to the database
331                        if let Err(e) = state.shared_state.database_pool.api_keys.save_key(key_type, &payload.api_key) {
332                            error!("Failed to save verified HuggingFace token: {}", e);
333                            return Ok(Json(VerifyApiKeyResponse {
334                                valid: false,
335                                message: "Token verified but failed to save. Please try again.".to_string(),
336                            }));
337                        }
338                        
339                        Ok(Json(VerifyApiKeyResponse {
340                            valid: true,
341                            message: "HuggingFace token is valid".to_string(),
342                        }))
343                    } else if resp.status() == StatusCode::UNAUTHORIZED || resp.status() == StatusCode::FORBIDDEN {
344                        info!("HuggingFace token verification failed: invalid credentials");
345                        Ok(Json(VerifyApiKeyResponse {
346                            valid: false,
347                            message: "Invalid HuggingFace token. Please check and try again.".to_string(),
348                        }))
349                    } else {
350                        let status = resp.status();
351                        error!("HuggingFace API key verification returned unexpected status: {}", status);
352                        Ok(Json(VerifyApiKeyResponse {
353                            valid: false,
354                            message: format!("HuggingFace API error: {}", status),
355                        }))
356                    }
357                }
358                Err(e) => {
359                    error!("Failed to verify HuggingFace token: {}", e);
360                    Ok(Json(VerifyApiKeyResponse {
361                        valid: false,
362                        message: "Failed to connect to HuggingFace API. Please check your internet connection.".to_string(),
363                    }))
364                }
365            }
366        }
367    }
368}