Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,15 @@ impl ConversationHandler {
Self { store }
}

/// Verifies that the underlying conversation store is available.
///
/// # Errors
///
/// Returns `ExecutorError` when the database health query fails or times out.
pub async fn check_health(&self) -> ExecutorResult<()> {
self.store.check_health().await.map_err(ExecutorError::Storage)
}

/// Gets an existing conversation or creates one.
///
/// Reads `conversation_id` from `ctx.original_request`.
Expand Down
26 changes: 26 additions & 0 deletions crates/agentic-server-core/src/storage/conversation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,15 @@

use std::convert::TryFrom;
use std::sync::Arc;
use std::time::Duration;

use super::models::{conversation, item, response};
use super::pool::DbPool;
use super::types::{ConversationData, InOutItem, ResponseMetadata, StorageError, StoreResult};
use crate::utils::common::{serialize_to_string, uuid7_str};

const DATABASE_HEALTH_TIMEOUT_SECONDS: u64 = 2;

/// Conversation storage operations.
#[derive(Clone, Debug)]
pub struct ConversationStore {
Expand Down Expand Up @@ -36,6 +39,29 @@ impl ConversationStore {
self.pool.as_deref().ok_or(StorageError::NotConfigured)
}

/// Verifies that configured storage can execute a bounded query.
///
/// Disabled storage has no required database dependency and is considered healthy.
///
/// # Errors
///
/// Returns an error when the query fails or exceeds the health-check deadline.
pub async fn check_health(&self) -> StoreResult<()> {
let Some(pool) = self.pool.as_deref() else {
return Ok(());
};
let query = sqlx::query("SELECT 1").execute(pool);
match tokio::time::timeout(Duration::from_secs(DATABASE_HEALTH_TIMEOUT_SECONDS), query).await {
Ok(result) => {
result?;
Ok(())
}
Err(_) => Err(StorageError::HealthCheckTimeout {
timeout_seconds: DATABASE_HEALTH_TIMEOUT_SECONDS,
}),
}
}

/// Creates a new conversation.
///
/// # Errors
Expand Down
4 changes: 4 additions & 0 deletions crates/agentic-server-core/src/storage/types/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,10 @@ pub enum StorageError {
#[error("storage not configured or disabled")]
NotConfigured,

/// Database health check exceeded its deadline.
#[error("database health check timed out after {timeout_seconds} seconds")]
HealthCheckTimeout { timeout_seconds: u64 },

/// Serialization or deserialization of data failed.
///
/// Wraps `serde_json::Error` and automatically converts from it via `#[from]`.
Expand Down
1 change: 1 addition & 0 deletions crates/agentic-server/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ criterion.workspace = true
futures.workspace = true
reqwest = { workspace = true, features = ["json"] }
serde_json.workspace = true
serde_yml = "0.0.12"
tokio = { workspace = true, features = ["test-util"] }
tokio-tungstenite.workspace = true
uuid = { version = "1", features = ["v7"] }
Expand Down
1 change: 1 addition & 0 deletions crates/agentic-server/benches/gateway_bench.rs
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,7 @@ async fn spawn_gateway(llm_url: &str) -> (Arc<reqwest::Client>, String) {
shutdown_token: CancellationToken::new(),
websocket_tracker: WebSocketTracker::default(),
llm_api_base: config.llm_api_base,
skip_llm_ready_check: config.skip_llm_ready_check,
openai_api_key: config.openai_api_key,
};

Expand Down
1 change: 1 addition & 0 deletions crates/agentic-server/benches/proxy_bench.rs
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,7 @@ async fn spawn_gateway(config: Config) -> String {
shutdown_token: CancellationToken::new(),
websocket_tracker: WebSocketTracker::default(),
llm_api_base: config.llm_api_base,
skip_llm_ready_check: config.skip_llm_ready_check,
openai_api_key: config.openai_api_key,
};
let server_config = ServerConfig::from_env();
Expand Down
2 changes: 2 additions & 0 deletions crates/agentic-server/src/app.rs
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,8 @@ pub struct AppState {
pub websocket_tracker: WebSocketTracker,
/// vLLM base URL — used by the `/ready` health probe.
pub llm_api_base: String,
/// Whether `/ready` should omit the upstream health check.
pub skip_llm_ready_check: bool,
/// Server-configured API key; used as fallback when the request carries no
/// `Authorization` header on the executor path.
pub openai_api_key: Option<String>,
Expand Down
54 changes: 40 additions & 14 deletions crates/agentic-server/src/handler/http/models.rs
Original file line number Diff line number Diff line change
Expand Up @@ -105,28 +105,54 @@ pub async fn health() -> impl IntoResponse {
StatusCode::OK
}

pub async fn ready(State(state): State<AppState>) -> impl IntoResponse {
async fn upstream_is_ready(state: &AppState) -> bool {
let base = state.llm_api_base.trim_end_matches('/');
let url = format!("{base}/health");

let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(2))
.build();

let Ok(client) = client else {
return StatusCode::SERVICE_UNAVAILABLE;
};
let mut request = state.exec_ctx.client.get(&url);
if let Some(key) = state
.openai_api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
{
request = request.bearer_auth(key);
}

match client.get(&url).send().await {
Ok(resp) if resp.status().is_success() => StatusCode::OK,
Ok(resp) => {
match tokio::time::timeout(std::time::Duration::from_secs(2), request.send()).await {
Ok(Ok(resp)) if resp.status().is_success() => true,
Ok(Ok(resp)) => {
warn!("LLM backend not ready: {}", resp.status());
StatusCode::SERVICE_UNAVAILABLE
false
}
Err(e) => {
Ok(Err(e)) => {
warn!("LLM backend unreachable: {e}");
StatusCode::SERVICE_UNAVAILABLE
false
}
Err(_) => {
warn!("LLM backend readiness check timed out");
false
}
}
}

async fn configured_upstream_is_ready(state: &AppState) -> bool {
state.skip_llm_ready_check || upstream_is_ready(state).await
}

pub async fn ready(State(state): State<AppState>) -> impl IntoResponse {
let (upstream_ready, database_result) = tokio::join!(
configured_upstream_is_ready(&state),
state.exec_ctx.conv_handler.check_health()
);
if let Err(error) = &database_result {
warn!("database not ready: {error}");
}

if upstream_ready && database_result.is_ok() {
StatusCode::OK
} else {
StatusCode::SERVICE_UNAVAILABLE
}
}

Expand Down
1 change: 1 addition & 0 deletions crates/agentic-server/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ async fn build_state(config: &Config, shutdown_token: CancellationToken) -> Resu
shutdown_token,
websocket_tracker: WebSocketTracker::default(),
llm_api_base: config.llm_api_base.clone(),
skip_llm_ready_check: config.skip_llm_ready_check,
openai_api_key: config.openai_api_key.clone(),
})
}
Expand Down
7 changes: 6 additions & 1 deletion crates/agentic-server/tests/common/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,12 @@ pub fn test_config(llm_url: &str) -> Config {
}

pub fn test_state(config: &Config) -> AppState {
test_state_with_conversation_store(config, ConversationStore::disabled())
}

pub fn test_state_with_conversation_store(config: &Config, conversation_store: ConversationStore) -> AppState {
let exec_ctx = ExecutionContext::new(
ConversationHandler::new(ConversationStore::disabled()),
ConversationHandler::new(conversation_store),
ResponseHandler::new(ResponseStore::disabled()),
Arc::new(reqwest::Client::new()),
config.llm_api_base.clone(),
Expand All @@ -40,6 +44,7 @@ pub fn test_state(config: &Config) -> AppState {
shutdown_token: CancellationToken::new(),
websocket_tracker: WebSocketTracker::default(),
llm_api_base: config.llm_api_base.clone(),
skip_llm_ready_check: config.skip_llm_ready_check,
openai_api_key: config.openai_api_key.clone(),
}
}
Expand Down
65 changes: 64 additions & 1 deletion crates/agentic-server/tests/health_test.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,14 @@
mod common;

use agentic_core::config::Config;
use common::{spawn_gateway, spawn_mock_llm, test_config, test_state};
use agentic_core::storage::{ConversationStore, create_pool};
use axum::Router;
use axum::http::HeaderMap;
use axum::response::IntoResponse;
use axum::routing::get;
use common::{spawn_gateway, spawn_mock_llm, test_config, test_state, test_state_with_conversation_store};
use http::StatusCode;
use tokio::net::TcpListener;

fn test_config_no_key(llm_url: &str) -> Config {
Config {
Expand Down Expand Up @@ -45,3 +52,59 @@ async fn test_ready_returns_503_when_llm_unreachable() {
let resp = reqwest::get(format!("{gw_url}/ready")).await.unwrap();
assert_eq!(resp.status(), 503);
}

#[tokio::test]
async fn test_ready_skips_upstream_when_configured() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let dead_addr = listener.local_addr().unwrap();
drop(listener);
let config = Config {
skip_llm_ready_check: true,
..test_config_no_key(&format!("http://{dead_addr}"))
};
let (gw_url, _gateway) = spawn_gateway(test_state(&config)).await;

let resp = reqwest::get(format!("{gw_url}/ready")).await.unwrap();

assert_eq!(resp.status(), 200);
}

#[tokio::test]
async fn test_ready_authenticates_upstream_health_check() {
let app = Router::new().route(
"/health",
get(|headers: HeaderMap| async move {
if headers.get("authorization").and_then(|value| value.to_str().ok()) == Some("Bearer test-key") {
StatusCode::OK.into_response()
} else {
StatusCode::UNAUTHORIZED.into_response()
}
}),
);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let upstream = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });

let (gw_url, gateway) = spawn_gateway(test_state(&test_config(&format!("http://{addr}")))).await;
let resp = reqwest::get(format!("{gw_url}/ready")).await.unwrap();

assert_eq!(resp.status(), 200);
upstream.abort();
gateway.abort();
}

#[tokio::test]
async fn test_ready_returns_503_when_database_is_unreachable() {
let (llm_url, upstream) = spawn_mock_llm().await;
let pool = create_pool(Some("sqlite::memory:")).await.unwrap();
let store = ConversationStore::new(pool.clone());
pool.close().await;
let state = test_state_with_conversation_store(&test_config(&llm_url), store);
let (gw_url, gateway) = spawn_gateway(state).await;

let resp = reqwest::get(format!("{gw_url}/ready")).await.unwrap();

assert_eq!(resp.status(), 503);
upstream.abort();
gateway.abort();
}
Loading