Skip to content
Merged
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
121 changes: 63 additions & 58 deletions bin/memayu/src/doctor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,13 @@
//!
//! Runs a battery of non-destructive checks against the effective config:
//! config validity, storage reachability/schema, and LLM/embedder connectivity
//! (key validity + model availability). Exits 0 when everything is healthy and
//! 1 when problems are found.
//! (key validity + real test-call reachability). Exits 0 when everything is
//! healthy and 1 when problems are found.
//!
//! Nothing here mutates state: storage is inspected, not opened-for-write, and
//! providers are probed via `GET /models` rather than issuing completions.
//! providers are probed with a real (non-mutating) completion/embedding call
//! rather than a `GET /models` listing, so proxy gateways that don't expose
//! `/models` are not misreported.

use console::style;
use memayu_config::{Config, StorageBackend};
Expand Down Expand Up @@ -60,13 +62,7 @@ pub async fn cmd_doctor(config: &Config) -> i32 {
// ── 3. Providers ──
if config.extraction_mode != ExtractionMode::Raw {
println!("\n{} Checking LLM provider", section("llm"));
let (passed, warned) = check_provider(
"LLM",
&config.llm.base_url,
config.llm.api_key.as_deref(),
&config.llm.model,
)
.await;
let (passed, warned) = check_provider("LLM", &config.llm, ProviderKind::Llm).await;
ok &= passed;
warnings += usize::from(warned);
}
Expand All @@ -86,13 +82,8 @@ pub async fn cmd_doctor(config: &Config) -> i32 {
);
println!(" model dir: {}", memayu_config::model_dir().display());
} else {
let (passed, warned) = check_provider(
"embedder",
&config.embedder.base_url,
config.embedder.api_key.as_deref(),
&config.embedder.model,
)
.await;
let (passed, warned) =
check_provider("embedder", &config.embedder, ProviderKind::Embedder).await;
ok &= passed;
warnings += usize::from(warned);
}
Expand Down Expand Up @@ -274,56 +265,70 @@ async fn inspect_storage(config: &Config) -> Result<StorageInspection, String> {
}
}

/// Probe a provider's `GET /models` endpoint, print the outcome, and return
/// `(passed, warned)` so the caller can fold it into the overall result.
/// Probe a provider with a real test call (an embedding for the embedder, a
/// completion for the LLM) rather than a `GET /models` listing, print the
/// outcome, and return `(passed, warned)` so the caller can fold it into the
/// overall result. A proxy that never exposes `/models` or renames the model
/// still reports success as long as the real call works (issue #48).
async fn check_provider(
name: &str,
base_url: &str,
api_key: Option<&str>,
model: &str,
cfg: &memayu_config::ProviderConfig,
kind: ProviderKind,
) -> (bool, bool) {
match memayu_llm_client::check_models(base_url, api_key, model).await {
memayu_llm_client::ModelsCheck::Ok { model_available } => {
if model_available {
println!(
" {} reachable, key accepted, model \"{model}\" advertised",
style("✔").green().bold()
);
(true, false)
} else {
println!(
" {} reachable, key accepted, but \"{model}\" not in /models list",
style("⚠").yellow().bold()
);
(true, true)
match kind {
ProviderKind::Llm => {
let provider = memayu_llm_client::HttpLlmProvider::new(cfg.clone());
match provider.probe().await {
Ok(()) => {
println!(
" {} reachable, key accepted, model \"{}\" answered a test completion",
style("✔").green().bold(),
cfg.model
);
(true, false)
}
Err(e) => {
println!(
" {} {} failed a test completion: {}",
style("✘").red().bold(),
name,
truncate(e)
);
(false, false)
}
}
}
memayu_llm_client::ModelsCheck::Unauthorized => {
println!(
" {} API key rejected (401) for {name}",
style("✘").red().bold()
);
(false, false)
}
memayu_llm_client::ModelsCheck::Http { status, detail } => {
println!(
" {} {name} returned HTTP {status}: {}",
style("✘").red().bold(),
truncate(detail)
);
(false, false)
}
memayu_llm_client::ModelsCheck::Unreachable(msg) => {
println!(
" {} cannot reach {name} at {base_url}: {}",
style("✘").red().bold(),
truncate(msg)
);
(false, false)
ProviderKind::Embedder => {
let provider = memayu_llm_client::HttpEmbedderProvider::new(cfg.clone());
match provider.probe().await {
Ok(dim) => {
println!(
" {} reachable, key accepted, model \"{}\" returned a {dim}-dim embedding",
style("✔").green().bold(),
cfg.model
);
(true, false)
}
Err(e) => {
println!(
" {} {} failed a test embedding: {}",
style("✘").red().bold(),
name,
truncate(e)
);
(false, false)
}
}
}
}
}

/// Which kind of provider [`check_provider`] should exercise.
enum ProviderKind {
Llm,
Embedder,
}

fn truncate(s: String) -> String {
const MAX: usize = 160;
if s.len() <= MAX {
Expand Down
136 changes: 100 additions & 36 deletions bin/memayu/tests/doctor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,60 +9,88 @@ use std::sync::atomic::{AtomicU16, Ordering};

const EMB: [f32; 3] = [0.1, 0.2, 0.3];

/// How the mock `/models` endpoint should behave.
/// How the mock provider should answer real test calls.
#[derive(Clone, Copy)]
enum ModelsBehavior {
/// 200 with the configured model advertised.
enum ProviderBehavior {
/// `/embeddings` and `/chat/completions` both succeed. There is no
/// `/models` endpoint at all (proxy scenario).
Good,
/// 200 but the configured model is NOT advertised.
ModelMissing,
/// 401 (bad API key).
/// 401 on real calls (bad API key).
Unauthorized,
/// 500 on real calls (provider misconfigured).
ServerError,
}

/// Bind an HTTP mock that answers `/models` (per `behavior`) and `/embeddings`
/// with a fixed embedding. Returns the bound port.
fn start_mock(behavior: ModelsBehavior) -> u16 {
/// Bind an HTTP mock that has NO `/models` endpoint and answers the real test
/// calls (`POST /embeddings`, `POST /chat/completions`) according to
/// `behavior`. Returns the bound port.
fn start_mock(behavior: ProviderBehavior) -> u16 {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock");
let port = listener.local_addr().unwrap().port();
std::thread::spawn(move || {
for stream in listener.incoming() {
let Ok(mut stream) = stream else { break };
// Read the full request: headers plus Content-Length bytes of body.
let mut buf = Vec::new();
let mut tmp = [0u8; 1024];
let mut tmp = [0u8; 2048];
loop {
match stream.read(&mut tmp) {
Ok(0) => break,
Ok(n) => {
buf.extend_from_slice(&tmp[..n]);
if buf.windows(4).any(|w| w == b"\r\n\r\n") {
break;
if let Some(header_end) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
let header = String::from_utf8_lossy(&buf[..header_end]).into_owned();
let content_length = header
.lines()
.find_map(|line| {
let mut parts = line.splitn(2, ':');
if parts
.next()
.map(|k| k.trim().eq_ignore_ascii_case("content-length"))
== Some(true)
{
parts.next()?.trim().parse::<usize>().ok()
} else {
None
}
})
.unwrap_or(0);
if buf.len() >= header_end + 4 + content_length {
break;
}
}
}
Err(_) => break,
}
}

let head = String::from_utf8_lossy(&buf);
let (status, body) = if head.starts_with("GET /models") {
// No /models endpoint: the proxy exposes only the real call
// surface. This is exactly what issue #48 exercises.
("404 Not Found", r#"{"error":"not found"}"#.to_string())
} else {
match behavior {
ModelsBehavior::Good => (
"200 OK",
r#"{"data":[{"id":"test-embed"},{"id":"other-model"}]}"#.to_string(),
),
ModelsBehavior::ModelMissing => {
("200 OK", r#"{"data":[{"id":"other-model"}]}"#.to_string())
ProviderBehavior::Good => {
if head.starts_with("POST /embeddings") {
(
"200 OK",
format!(
r#"{{"data":[{{"embedding":[{:.6},{:.6},{:.6}]}}]}}"#,
EMB[0], EMB[1], EMB[2]
),
)
} else {
// /chat/completions
(
"200 OK",
r#"{"choices":[{"message":{"content":"{\"decision\":\"add\",\"memory_id\":null,\"content\":\"ok\"}"}}]}"#.to_string(),
)
}
}
ModelsBehavior::Unauthorized => ("401 Unauthorized", String::new()),
ProviderBehavior::Unauthorized => ("401 Unauthorized", String::new()),
ProviderBehavior::ServerError => ("500 Internal Server Error", String::new()),
}
} else {
// /embeddings
(
"200 OK",
format!(
r#"{{"data":[{{"embedding":[{:.6},{:.6},{:.6}]}}]}}"#,
EMB[0], EMB[1], EMB[2]
),
)
};
let resp = format!(
"HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
Expand Down Expand Up @@ -139,7 +167,7 @@ fn prime_schema(port: u16, db: &std::path::Path) {

#[test]
fn doctor_fresh_machine_exits_zero() {
let port = start_mock(ModelsBehavior::Good);
let port = start_mock(ProviderBehavior::Good);
let db = temp_db_path(); // does not exist yet
let (code, stdout, _) = run_doctor(port, &db, &[]);
assert_eq!(code, 0, "fresh machine should be healthy");
Expand All @@ -148,7 +176,7 @@ fn doctor_fresh_machine_exits_zero() {

#[test]
fn doctor_used_store_exits_zero() {
let port = start_mock(ModelsBehavior::Good);
let port = start_mock(ProviderBehavior::Good);
let db = temp_db_path();
prime_schema(port, &db);
let (code, stdout, _) = run_doctor(port, &db, &[]);
Expand All @@ -158,20 +186,56 @@ fn doctor_used_store_exits_zero() {

#[test]
fn doctor_bad_key_exits_one() {
let port = start_mock(ModelsBehavior::Unauthorized);
let port = start_mock(ProviderBehavior::Unauthorized);
let db = temp_db_path();
let (code, stdout, _) = run_doctor(port, &db, &[]);
assert_eq!(code, 1, "rejected key must fail doctor");
assert!(stdout.contains("401"), "stdout: {stdout}");
}

#[test]
fn doctor_model_missing_is_warning_but_healthy() {
let port = start_mock(ModelsBehavior::ModelMissing);
fn doctor_proxy_without_models_endpoint_is_healthy() {
// The mock has no /models endpoint at all (a proxy like 9Router), yet the
// real embed test call succeeds. Doctor must report success (issue #48).
let port = start_mock(ProviderBehavior::Good);
let db = temp_db_path();
let (code, stdout, _) = run_doctor(port, &db, &[]);
assert_eq!(code, 0, "real test call must succeed through a proxy");
assert!(
stdout.contains("returned a 3-dim embedding"),
"stdout: {stdout}"
);
}

#[test]
fn doctor_server_error_exits_one() {
let port = start_mock(ProviderBehavior::ServerError);
let db = temp_db_path();
let (code, stdout, _) = run_doctor(port, &db, &[]);
assert_eq!(code, 0, "unadvertised model is only a warning");
assert!(stdout.contains("not in /models list"), "stdout: {stdout}");
assert_eq!(code, 1, "misconfigured provider must still fail doctor");
assert!(stdout.contains("500"), "stdout: {stdout}");
}

#[test]
fn doctor_llm_proxy_without_models_endpoint_is_healthy() {
// In llm mode doctor also exercises the LLM with a real completion. The
// mock has no /models endpoint, but both real calls succeed.
let port = start_mock(ProviderBehavior::Good);
let db = temp_db_path();
let (code, stdout, _) = run_doctor(
port,
&db,
&[
("MEMAYU_EXTRACTION_MODE", "llm"),
("MEMAYU_LLM_BASE_URL", &format!("http://127.0.0.1:{port}")),
("MEMAYU_LLM_MODEL", "proxy-llm-model"),
],
);
assert_eq!(code, 0, "LLM test completion must succeed through a proxy");
assert!(
stdout.contains("answered a test completion"),
"stdout: {stdout}"
);
}

#[test]
Expand All @@ -185,7 +249,7 @@ fn doctor_missing_config_exits_one() {

#[test]
fn doctor_dimension_mismatch_fails() {
let port = start_mock(ModelsBehavior::Good);
let port = start_mock(ProviderBehavior::Good);
let db = temp_db_path();
prime_schema(port, &db); // stores dimension 3
// Run doctor with a conflicting configured dimension.
Expand Down
11 changes: 11 additions & 0 deletions crates/memayu-llm-client/src/embedder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,17 @@ impl HttpEmbedderProvider {
)
.await
}

/// Issue a minimal, real embedding call and return the vector dimension.
/// A success here means the endpoint is reachable, the key is accepted, and
/// the configured model answers — the same test the setup/dashboard flows
/// use (PRD-01 §4.2) — without consulting any `/models` listing.
pub async fn probe(&self) -> Result<usize, String> {
self.embed("dimension probe")
.await
.map(|v| v.len())
.map_err(|e| e.to_string())
}
}

fn retry_after_ms(headers: &reqwest::header::HeaderMap) -> Option<u64> {
Expand Down
Loading
Loading