rock/poimen-memory
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement M3.5.7: Rate limiting + idempotency (20 tests)

Browse Source
This commit is contained in:
Story Crater Bot
2026-08-26 13:35:50 -07:00
parent f8a06ef49d
commit 57c806ea25
9 changed files with 1106 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/mem-cli/src/http_server.rs
+98 -11
View File
@@ -10,6 +10,8 @@ use std::time::Instant;
use crate::endpoints::IngestRequest;
use crate::ingest_worker::IngestWorker;
use crate::query_worker::QueryWorker;
use crate::rate_limiter::{RateLimiter, LimitConfig};
use crate::idempotency::IdempotencyStore;
/// Server state with database and workers
pub struct AppState {
@@ -20,6 +22,8 @@ pub struct AppState {
pub embeddings: Arc<EmbeddingsClient>,
pub ingest_worker: Arc<IngestWorker>,
pub query_worker: Arc<QueryWorker>,
pub rate_limiter: Arc<RateLimiter>,
pub idempotency_store: Arc<IdempotencyStore>,
}
/// Auth extractor — validates apikey header
@@ -36,6 +40,34 @@ fn check_auth(req: &HttpRequest, state: &AppState) -> Result<(), HttpResponse> {
Ok(())
}
/// Extract apikey from request
fn extract_apikey(req: &HttpRequest) -> Option<String> {
req.headers()
.get("apikey")
.and_then(|h| h.to_str().ok())
.map(|s| s.to_string())
}
/// Rate limit guard — call this in handlers to check rate limit
fn check_rate_limit(req: &HttpRequest, state: &AppState, endpoint: &str) -> Result<(), HttpResponse> {
let apikey = extract_apikey(req).unwrap_or_else(|| "unknown".to_string());
match state.rate_limiter.check(&apikey, endpoint) {
Ok(_) => Ok(()),
Err(rate_limit_err) => {
let retry_after = rate_limit_err.retry_after_seconds.to_string();
Err(HttpResponse::TooManyRequests()
.insert_header(("Retry-After", retry_after))
.json(json!({
"error": "rate_limit_exceeded",
"reason": rate_limit_err.reason.clone(),
"retry_after_seconds": rate_limit_err.retry_after_seconds,
"limit_window": format!("{}s", rate_limit_err.limit_window_secs),
})))
}
}
}
/// Start HTTP server with database initialization
pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Result<()> {
// Create connection pool
@@ -53,6 +85,33 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
let reranker = RerankClient::from_env()?;
let query_worker = Arc::new(QueryWorker::new(VectorStore::new(pool.clone()), (*embeddings).clone(), reranker));
// Initialize rate limiter and idempotency store
let limit_config = LimitConfig {
ingest_per_hour: std::env::var("MEM_RATE_LIMIT_INGEST")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(100.0),
query_per_hour: std::env::var("MEM_RATE_LIMIT_QUERY")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(1000.0),
projects_per_hour: std::env::var("MEM_RATE_LIMIT_PROJECTS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(100.0),
burst_per_second: std::env::var("MEM_RATE_LIMIT_BURST")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(10.0),
};
let rate_limiter = Arc::new(RateLimiter::new(limit_config));
let idempotency_ttl = std::env::var("MEM_IDEMPOTENCY_TTL_SECS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(86400); // 24 hours default
let idempotency_store = Arc::new(IdempotencyStore::new(idempotency_ttl));
let state = web::Data::new(AppState {
api_key,
start_time: Instant::now(),
@@ -61,6 +120,8 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
embeddings,
ingest_worker,
query_worker,
rate_limiter,
idempotency_store,
});
tracing::info!("Starting HTTP server on port {}", port);
@@ -103,6 +164,10 @@ pub async fn ingest_handler(
return e;
}
if let Err(e) = check_rate_limit(&req, &state, "/memory/ingest") {
return e;
}
let project = body.project.clone();
let ingest_id = body.ingest_id.clone();
let records: Vec<(String, String)> = body
@@ -111,6 +176,12 @@ pub async fn ingest_handler(
.map(|r| (r.text.clone(), body.source.clone()))
.collect();
// Check idempotency cache first
if let Some(cached_response) = state.idempotency_store.get(&ingest_id) {
tracing::info!("Returning cached response for ingest_id: {}", ingest_id);
return HttpResponse::Accepted().json(cached_response);
}
// Create ingest job in DB
let job_result = sqlx::query(
"INSERT INTO ingest_jobs (id, project, ingest_id, status, created_at)
@@ -136,18 +207,26 @@ pub async fn ingest_handler(
}
});
HttpResponse::Accepted().json(json!({
let response = json!({
"ingest_id": ingest_id,
"status": "pending",
"status_url": format!("/memory/ingest/{}", ingest_id)
}))
});
// Cache the response for idempotency
state.idempotency_store.set(ingest_id.clone(), response.clone());
HttpResponse::Accepted().json(response)
}
Ok(None) => {
// Already exists
HttpResponse::Conflict().json(json!({
"error": "already_ingesting",
"ingest_id": ingest_id
}))
// Already exists in DB (was inserted concurrently)
let response = json!({
"ingest_id": ingest_id,
"status": "pending",
"status_url": format!("/memory/ingest/{}", ingest_id)
});
state.idempotency_store.set(ingest_id.clone(), response.clone());
HttpResponse::Accepted().json(response)
}
Err(e) => {
tracing::error!("DB error: {}", e);
@@ -203,6 +282,10 @@ pub async fn query_handler(
return e;
}
if let Err(e) = check_rate_limit(&req, &state, "/memory/query") {
return e;
}
let project = match query.get("project") {
Some(p) => p.clone(),
None => {
@@ -246,6 +329,10 @@ pub async fn projects_handler(
return e;
}
if let Err(e) = check_rate_limit(&req, &state, "/memory/projects") {
return e;
}
let result = sqlx::query_as::<_, (String,)>(
"SELECT DISTINCT project FROM memories_l2 ORDER BY project",
)
@@ -544,12 +631,12 @@ pub async fn vault_file_handler(
let (frontmatter, body) = if content.starts_with("---") {
let parts: Vec<&str> = content.split("---").collect();
if parts.len() >= 3 {
(parts[1], parts[2..].join("---"))
(parts[1].to_string(), parts[2..].join("---"))
} else {
("", &content[..])
("".to_string(), content.clone())
}
} else {
("", &content[..])
("".to_string(), content.clone())
};
let html = format!(
@@ -590,7 +677,7 @@ pub async fn vault_file_handler(
} else {
String::new()
},
body.replace("&", "&amp;")
body.clone().replace("&", "&amp;")
.replace("<", "&lt;")
.replace(">", "&gt;")
.lines()