use axum::extract::Path; use axum::http::{HeaderMap, StatusCode}; use axum::routing::{patch, post}; use axum::{Extension, Json, Router}; use serde_json::{Value, json}; use sqlx::PgPool; use crate::model::{ ContextScope, MemoryCreateRequest, MemoryDeleteRequest, MemoryDeleteResponse, MemoryResponse, MemorySearchRequest, MemorySearchResponse, MemoryUpdateRequest, }; use crate::repository::ContextRepository; type ApiError = (StatusCode, Json); #[derive(Clone)] pub struct ContextRouteState { pub db: PgPool, } impl ContextRouteState { pub fn new(db: PgPool) -> Self { Self { db } } } pub fn routes(db: PgPool) -> Router where S: Clone + Send + Sync + 'static, { Router::new() .route("/context/memories", post(create_memory)) .route("/context/memories/search", post(search_memories)) .route( "/context/memories/{id}", patch(update_memory).delete(delete_memory), ) .layer(Extension(ContextRouteState::new(db))) } async fn create_memory( headers: HeaderMap, Extension(state): Extension, Json(request): Json, ) -> Result, ApiError> { authorize(&headers)?; let scope = validate_scope(request.identity.scope())?; let content = required_text(&request.content, "content")?; let metadata = request.metadata.unwrap_or_else(|| json!({})).to_string(); let memory = ContextRepository::create_memory(&state.db, &scope, content, metadata) .await .map_err(api_internal_error)?; Ok(Json(MemoryResponse { success: true, memory, })) } async fn search_memories( headers: HeaderMap, Extension(state): Extension, Json(request): Json, ) -> Result, ApiError> { authorize(&headers)?; let scope = validate_scope(request.identity.scope())?; let limit = request.limit.unwrap_or(20).clamp(1, 50); let query = request .query .as_deref() .map(str::trim) .filter(|query| !query.is_empty()); let memories = ContextRepository::search_memories(&state.db, &scope.scope_key, query, limit) .await .map_err(api_internal_error)?; Ok(Json(MemorySearchResponse { success: true, memories, })) } async fn update_memory( headers: HeaderMap, Extension(state): Extension, Path(id): Path, Json(request): Json, ) -> Result, ApiError> { authorize(&headers)?; let scope = validate_scope(request.identity.scope())?; let content = required_text(&request.content, "content")?; let metadata = request.metadata.map(|value| value.to_string()); let memory = ContextRepository::update_memory(&state.db, &scope.scope_key, id, content, metadata) .await .map_err(api_internal_error)? .ok_or_else(|| api_not_found("长期记忆不存在"))?; Ok(Json(MemoryResponse { success: true, memory, })) } async fn delete_memory( headers: HeaderMap, Extension(state): Extension, Path(id): Path, Json(request): Json, ) -> Result, ApiError> { authorize(&headers)?; let scope = validate_scope(request.identity.scope())?; let deleted = ContextRepository::delete_memory(&state.db, &scope.scope_key, id) .await .map_err(api_internal_error)?; Ok(Json(MemoryDeleteResponse { success: true, deleted, })) } fn authorize(headers: &HeaderMap) -> Result<(), ApiError> { let expected = context_api_key()?; let token = headers .get("x-ias-context-token") .and_then(|value| value.to_str().ok()) .map(str::trim) .filter(|value| !value.is_empty()) .or_else(|| bearer_token(headers)); if token.is_some_and(|token| token == expected) { Ok(()) } else { Err(api_unauthorized("context 内部接口鉴权失败")) } } fn bearer_token(headers: &HeaderMap) -> Option<&str> { let value = headers.get("authorization")?.to_str().ok()?.trim(); value .strip_prefix("Bearer ") .or_else(|| value.strip_prefix("bearer ")) .map(str::trim) .filter(|value| !value.is_empty()) } fn context_api_key() -> Result { std::env::var("IA_CONTEXT_API_KEY") .or_else(|_| std::env::var("IA_CONTEXT_TOKEN")) .or_else(|_| std::env::var("ROUTER_API_KEY")) .map(|value| value.trim().to_string()) .ok() .filter(|value| !value.is_empty()) .ok_or_else(|| api_unavailable("未配置 context 内部访问密钥")) } fn validate_scope(scope: ContextScope) -> Result { if scope.account_id.trim().is_empty() { return Err(api_bad_request("account_id 不能为空")); } if scope.from_user_id.trim().is_empty() { return Err(api_bad_request("from_user_id 不能为空")); } Ok(scope) } fn required_text(value: &str, field: &str) -> Result { let value = value.trim(); if value.is_empty() { return Err(api_bad_request(format!("{field} 不能为空"))); } Ok(value.to_string()) } fn api_bad_request(message: impl Into) -> ApiError { ( StatusCode::BAD_REQUEST, Json(json!({ "success": false, "message": message.into(), })), ) } fn api_not_found(message: impl Into) -> ApiError { ( StatusCode::NOT_FOUND, Json(json!({ "success": false, "message": message.into(), })), ) } fn api_unauthorized(message: impl Into) -> ApiError { ( StatusCode::UNAUTHORIZED, Json(json!({ "success": false, "message": message.into(), })), ) } fn api_unavailable(message: impl Into) -> ApiError { ( StatusCode::SERVICE_UNAVAILABLE, Json(json!({ "success": false, "message": message.into(), })), ) } fn api_internal_error(error: anyhow::Error) -> ApiError { eprintln!("context 接口失败: {error}"); ( StatusCode::INTERNAL_SERVER_ERROR, Json(json!({ "success": false, "message": "context 接口失败", })), ) }