buzz/crates/buzz-agent/tests/databricks_oauth.rs

//! Integration tests for the PKCE OAuth token source.
//!
//! No browser dance — we cover the silent-refresh and cache-hit paths
//! against a stubbed OIDC server (axum). The interactive browser flow is
//! exercised manually via the `buzz-agent auth databricks` subcommand
//! (see `lib.rs::auth_subcommand`).
//!
//! The second test module (further down) is an ACP-level envelope
//! regression: it spawns the real `buzz-agent` binary with
//! `DATABRICKS_TOKEN` set and a stub HTTP server, then asserts the wire
//! shape we send to Databricks.

use std::net::SocketAddr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};

use axum::extract::Form;
use axum::{routing::get, routing::post, Json, Router};
use buzz_agent::auth::{PkceOAuthConfig, PkceOAuthTokenSource, TokenSource};
use serde::Deserialize;
use serde_json::json;
use sha2::{Digest, Sha256};
use tempfile::TempDir;

#[derive(Deserialize)]
struct TokenForm {
    grant_type: String,
    #[allow(dead_code)]
    refresh_token: Option<String>,
}

/// Boot a stub OIDC server that:
///   - serves discovery at `/.well-known/oauth-authorization-server`
///   - issues a fresh access token for every `refresh_token` request
///   - counts how many refresh hits it gets
async fn spawn_oidc() -> (String, Arc<AtomicU64>) {
    let counter = Arc::new(AtomicU64::new(0));
    let counter_for_token = counter.clone();

    // Bind first so we know our own base URL before building the router.
    let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))
        .await
        .unwrap();
    let addr = listener.local_addr().unwrap();
    let base = format!("http://{addr}");
    let base_for_discovery = base.clone();

    let app = Router::new()
        .route(
            "/.well-known/oauth-authorization-server",
            get(move || {
                let base = base_for_discovery.clone();
                async move {
                    Json(json!({
                        "authorization_endpoint": format!("{base}/authorize"),
                        "token_endpoint": format!("{base}/token"),
                    }))
                }
            }),
        )
        .route(
            "/token",
            post(move |Form(form): Form<TokenForm>| {
                let counter = counter_for_token.clone();
                async move {
                    let n = counter.fetch_add(1, Ordering::SeqCst) + 1;
                    assert_eq!(form.grant_type, "refresh_token");
                    Json(json!({
                        "access_token": format!("fresh-token-{n}"),
                        "refresh_token": "rotated-refresh",
                        "expires_in": 3600,
                    }))
                }
            }),
        );

    tokio::spawn(async move {
        let _ = axum::serve(listener, app).await;
    });
    (base, counter)
}

/// Cache key construction matches the auth module: sha256(discovery|client|scopes).
fn cache_path_for(cache_dir: &std::path::Path, cfg: &PkceOAuthConfig) -> std::path::PathBuf {
    use sha2::Digest;
    let mut h = sha2::Sha256::new();
    h.update(cfg.discovery_url.as_bytes());
    h.update(b"|");
    h.update(cfg.client_id.as_bytes());
    h.update(b"|");
    h.update(cfg.scopes.join(",").as_bytes());
    let hash = hex::encode(h.finalize());
    cache_dir
        .join(&cfg.cache_namespace)
        .join(format!("{hash}.json"))
}

/// Write a token file the engine should pick up on construction.
fn seed_cache(path: &std::path::Path, body: serde_json::Value) {
    std::fs::create_dir_all(path.parent().unwrap()).unwrap();
    std::fs::write(path, serde_json::to_vec(&body).unwrap()).unwrap();
}

#[tokio::test]
async fn cache_hit_short_circuits_network() {
    let tmp = TempDir::new().unwrap();

    let (base, refresh_counter) = spawn_oidc().await;
    let cfg = PkceOAuthConfig {
        discovery_url: format!("{base}/.well-known/oauth-authorization-server"),
        client_id: "test-client".into(),
        scopes: vec!["a".into(), "b".into()],
        cache_namespace: "databricks".into(),
        cache_dir_override: Some(tmp.path().to_path_buf()),
    };

    // Seed an unexpired token in the cache.
    let future = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs()
        + 3600;
    let path = cache_path_for(tmp.path(), &cfg);
    seed_cache(
        &path,
        json!({
            "access_token": "cached-token",
            "refresh_token": "rt",
            "expires_at": future,
        }),
    );

    let src = PkceOAuthTokenSource::new(cfg).unwrap();
    let bearer = src.bearer().await.unwrap();
    assert_eq!(bearer, "cached-token");
    assert_eq!(
        refresh_counter.load(Ordering::SeqCst),
        0,
        "no refresh should fire"
    );
}

#[tokio::test]
async fn expired_cache_silently_refreshes() {
    let tmp = TempDir::new().unwrap();

    let (base, refresh_counter) = spawn_oidc().await;
    let cfg = PkceOAuthConfig {
        discovery_url: format!("{base}/.well-known/oauth-authorization-server"),
        client_id: "test-client".into(),
        scopes: vec!["a".into()],
        cache_namespace: "databricks".into(),
        cache_dir_override: Some(tmp.path().to_path_buf()),
    };

    // Seed an already-expired token with a refresh_token.
    let path = cache_path_for(tmp.path(), &cfg);
    seed_cache(
        &path,
        json!({
            "access_token": "stale",
            "refresh_token": "valid-refresh",
            "expires_at": 1u64, // way in the past
        }),
    );

    let src = PkceOAuthTokenSource::new(cfg).unwrap();
    let bearer = src.bearer().await.unwrap();
    assert_eq!(bearer, "fresh-token-1");
    assert_eq!(refresh_counter.load(Ordering::SeqCst), 1);

    // A second call should hit the in-memory cache and skip the network.
    let bearer2 = src.bearer().await.unwrap();
    assert_eq!(bearer2, "fresh-token-1");
    assert_eq!(refresh_counter.load(Ordering::SeqCst), 1);
}

#[tokio::test]
async fn refreshed_token_is_persisted_to_disk() {
    let tmp = TempDir::new().unwrap();

    let (base, _) = spawn_oidc().await;
    let cfg = PkceOAuthConfig {
        discovery_url: format!("{base}/.well-known/oauth-authorization-server"),
        client_id: "test-client".into(),
        scopes: vec!["a".into()],
        cache_namespace: "databricks".into(),
        cache_dir_override: Some(tmp.path().to_path_buf()),
    };

    let path = cache_path_for(tmp.path(), &cfg);
    seed_cache(
        &path,
        json!({
            "access_token": "stale",
            "refresh_token": "valid-refresh",
            "expires_at": 1u64,
        }),
    );

    let src = PkceOAuthTokenSource::new(cfg).unwrap();
    let _ = src.bearer().await.unwrap();

    let on_disk: serde_json::Value =
        serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
    assert_eq!(on_disk["access_token"], "fresh-token-1");
    assert_eq!(on_disk["refresh_token"], "rotated-refresh");
    assert!(on_disk["expires_at"].is_u64());
}

#[tokio::test]
async fn refresh_now_runs_grant_on_unexpired_rejected_token() {
    let tmp = TempDir::new().unwrap();

    let (base, refresh_counter) = spawn_oidc().await;
    let cfg = PkceOAuthConfig {
        discovery_url: format!("{base}/.well-known/oauth-authorization-server"),
        client_id: "test-client".into(),
        scopes: vec!["a".into()],
        cache_namespace: "databricks".into(),
        cache_dir_override: Some(tmp.path().to_path_buf()),
    };

    // The exact 401 case this whole change exists to fix: a token that is
    // still locally *unexpired* but the server rejected it (skew, revocation,
    // a node that never saw it). is_expired() says "keep it", so a clock-based
    // gate would no-op and the agent would die. refresh_now() must instead key
    // off identity — the cached token equals the rejected one — and run the
    // grant anyway. The stub never serves a browser flow, so a fresh token
    // here proves the refresh-token grant ran, not the interactive path.
    let future = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs()
        + 3600;
    let path = cache_path_for(tmp.path(), &cfg);
    seed_cache(
        &path,
        json!({
            "access_token": "rejected",
            "refresh_token": "valid-refresh",
            "expires_at": future,
        }),
    );

    let src = PkceOAuthTokenSource::new(cfg).unwrap();
    let bearer = src.refresh_now("rejected").await.unwrap();
    assert_eq!(bearer, "fresh-token-1", "grant ran despite local freshness");
    assert_eq!(refresh_counter.load(Ordering::SeqCst), 1, "grant ran once");

    // The refresh token was preserved (rotated, not discarded): the saved
    // token still carries one, so a future 401 can refresh again instead of
    // falling to the browser flow. This is the property defect #1 broke.
    let on_disk: serde_json::Value =
        serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
    assert_eq!(on_disk["access_token"], "fresh-token-1");
    assert_eq!(on_disk["refresh_token"], "rotated-refresh");
}

#[tokio::test]
async fn refresh_now_coalesces_when_another_caller_already_refreshed() {
    let tmp = TempDir::new().unwrap();

    let (base, refresh_counter) = spawn_oidc().await;
    let cfg = PkceOAuthConfig {
        discovery_url: format!("{base}/.well-known/oauth-authorization-server"),
        client_id: "test-client".into(),
        scopes: vec!["a".into()],
        cache_namespace: "databricks".into(),
        cache_dir_override: Some(tmp.path().to_path_buf()),
    };

    // A concurrent caller already replaced the rejected token: the cached
    // token differs from the one we hold. Coalesce by identity — return the
    // new token without burning a second grant, so N concurrent 401s on the
    // same stale token collapse onto one refresh. Note the cached token is
    // *unexpired* here too, so this proves coalescing keys off identity, not
    // the clock (which agrees with both the old and new token).
    let future = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs()
        + 3600;
    let path = cache_path_for(tmp.path(), &cfg);
    seed_cache(
        &path,
        json!({
            "access_token": "already-refreshed",
            "refresh_token": "rt",
            "expires_at": future,
        }),
    );

    let src = PkceOAuthTokenSource::new(cfg).unwrap();
    let bearer = src.refresh_now("the-rejected-one").await.unwrap();
    assert_eq!(bearer, "already-refreshed");
    assert_eq!(
        refresh_counter.load(Ordering::SeqCst),
        0,
        "no grant when a sibling already refreshed the rejected token"
    );
}

#[tokio::test]
async fn refresh_now_without_refresh_token_is_terminal() {
    let tmp = TempDir::new().unwrap();

    let (base, refresh_counter) = spawn_oidc().await;
    let cfg = PkceOAuthConfig {
        discovery_url: format!("{base}/.well-known/oauth-authorization-server"),
        client_id: "test-client".into(),
        scopes: vec!["a".into()],
        cache_namespace: "databricks".into(),
        cache_dir_override: Some(tmp.path().to_path_buf()),
    };

    // The rejected token is still the cached one and there's no refresh token
    // to fall back on. refresh_now() must fail terminally (LlmAuth) rather
    // than open a browser — the headless hang this whole change exists to
    // prevent.
    let path = cache_path_for(tmp.path(), &cfg);
    seed_cache(
        &path,
        json!({
            "access_token": "rejected",
            "refresh_token": serde_json::Value::Null,
            "expires_at": 1u64,
        }),
    );

    let src = PkceOAuthTokenSource::new(cfg).unwrap();
    let err = src.refresh_now("rejected").await.unwrap_err();
    // `types::AgentError` isn't a public path; match on its Display, which
    // prefixes `LlmAuth` variants with "llm auth:". A terminal LlmAuth (not
    // a browser hang) is the whole point of this path.
    let msg = err.to_string();
    assert!(
        msg.starts_with("llm auth:"),
        "expected terminal LlmAuth, got: {msg}"
    );
    assert_eq!(
        refresh_counter.load(Ordering::SeqCst),
        0,
        "no grant attempted"
    );
}

// ACP-level envelope regression test.
//
// Boots the real buzz-agent binary with `DATABRICKS_TOKEN` set (so the
// OAuth dance is skipped) pointed at a stub HTTP server that captures every
// inbound request. Asserts the wire-level shape Databricks model serving
// requires: path is `/serving-endpoints/<model>/invocations`, Authorization
// is `Bearer <token>`, and the JSON body has *no* top-level `"model"`. This
// locks in the DRY envelope behavior so a refactor of `post_openai` can't
// silently break Databricks.

use std::collections::VecDeque;
use std::process::Stdio;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::net::TcpListener;
use tokio::sync::Mutex;

#[derive(Clone, Debug)]
struct CapturedRequest {
    path: String,
    authorization: Option<String>,
    body: serde_json::Value,
}

async fn spawn_capturing_server(
    responses: Vec<serde_json::Value>,
) -> (String, Arc<Mutex<Vec<CapturedRequest>>>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let url = format!("http://{}", listener.local_addr().unwrap());
    let queue = Arc::new(Mutex::new(VecDeque::from(responses)));
    let captured = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
    let cap_for_task = captured.clone();
    tokio::spawn(async move {
        loop {
            let (mut sock, _) = match listener.accept().await {
                Ok(p) => p,
                Err(_) => return,
            };
            let queue = queue.clone();
            let captured = cap_for_task.clone();
            tokio::spawn(async move {
                let mut buf = Vec::new();
                let mut tmp = [0u8; 8192];
                while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
                    match sock.read(&mut tmp).await {
                        Ok(0) | Err(_) => return,
                        Ok(n) => buf.extend_from_slice(&tmp[..n]),
                    }
                    if buf.len() > 4_000_000 {
                        return;
                    }
                }
                let header_end = buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4;
                let header_str = String::from_utf8_lossy(&buf[..header_end]).to_string();
                let (request_line, rest) = header_str.split_once('\n').unwrap_or(("", ""));
                let path = request_line
                    .split_whitespace()
                    .nth(1)
                    .unwrap_or("")
                    .to_string();
                let mut authorization = None;
                let mut body_len = 0usize;
                for line in rest.lines() {
                    // Split case-insensitively on the colon but keep the value's case intact.
                    let Some((name, value)) = line.split_once(':') else {
                        continue;
                    };
                    let value = value.trim().trim_end_matches('\r').to_string();
                    match name.trim().to_ascii_lowercase().as_str() {
                        "authorization" => authorization = Some(value),
                        "content-length" => body_len = value.parse().unwrap_or(0),
                        _ => {}
                    }
                }
                while buf.len() < header_end + body_len {
                    match sock.read(&mut tmp).await {
                        Ok(0) | Err(_) => return,
                        Ok(n) => buf.extend_from_slice(&tmp[..n]),
                    }
                }
                let body: serde_json::Value =
                    serde_json::from_slice(&buf[header_end..header_end + body_len])
                        .unwrap_or(json!(null));
                captured.lock().await.push(CapturedRequest {
                    path,
                    authorization,
                    body,
                });
                let body = queue
                    .lock()
                    .await
                    .pop_front()
                    .unwrap_or_else(|| json!({ "error": "no canned response" }));
                let body_s = serde_json::to_string(&body).unwrap();
                let resp = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                    body_s.len(),
                    body_s,
                );
                let _ = sock.write_all(resp.as_bytes()).await;
                let _ = sock.shutdown().await;
            });
        }
    });
    (url, captured)
}

struct AgentHarness {
    child: tokio::process::Child,
    stdin: tokio::process::ChildStdin,
    stdout: BufReader<tokio::process::ChildStdout>,
    next_id: i64,
    _home: Option<TempDir>,
}

impl Drop for AgentHarness {
    fn drop(&mut self) {
        let _ = self.child.start_kill();
    }
}

impl AgentHarness {
    async fn spawn_provider(provider: &str, base_url: &str, model: &str) -> Self {
        Self::spawn_provider_with_options(provider, base_url, model, 1, Some("test-bearer")).await
    }

    async fn spawn_oauth_provider(
        provider: &str,
        base_url: &str,
        model: &str,
        max_sessions: usize,
    ) -> Self {
        Self::spawn_provider_with_options(provider, base_url, model, max_sessions, None).await
    }

    async fn spawn_provider_with_max_sessions(
        provider: &str,
        base_url: &str,
        model: &str,
        max_sessions: usize,
    ) -> Self {
        Self::spawn_provider_with_options(
            provider,
            base_url,
            model,
            max_sessions,
            Some("test-bearer"),
        )
        .await
    }

    async fn spawn_provider_with_options(
        provider: &str,
        base_url: &str,
        model: &str,
        max_sessions: usize,
        token: Option<&str>,
    ) -> Self {
        let bin = env!("CARGO_BIN_EXE_buzz-agent");
        let home = token
            .is_none()
            .then(|| TempDir::new().expect("create isolated OAuth home"));
        let mut cmd = tokio::process::Command::new(bin);
        cmd.env("BUZZ_AGENT_PROVIDER", provider)
            .env("DATABRICKS_HOST", base_url)
            .env("DATABRICKS_MODEL", model)
            .env_remove("DATABRICKS_TOKEN")
            .env("BUZZ_AGENT_LLM_TIMEOUT_SECS", "5")
            .env("BUZZ_AGENT_TOOL_TIMEOUT_SECS", "5")
            .env("BUZZ_AGENT_MAX_ROUNDS", "2")
            .env("BUZZ_AGENT_MAX_SESSIONS", max_sessions.to_string())
            .env("BUZZ_AGENT_MCP_INIT_TIMEOUT_SECS", "2")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .kill_on_drop(true);
        if let Some(token) = token {
            cmd.env("DATABRICKS_TOKEN", token);
        }
        if let Some(home) = &home {
            cmd.env("HOME", home.path());
        }
        let mut child = cmd.spawn().expect("spawn buzz-agent");
        let stdin = child.stdin.take().unwrap();
        let stdout = BufReader::new(child.stdout.take().unwrap());
        Self {
            child,
            stdin,
            stdout,
            next_id: 1,
            _home: home,
        }
    }

    fn oauth_home(&self) -> &std::path::Path {
        self._home
            .as_ref()
            .expect("harness was not started in OAuth mode")
            .path()
    }

    async fn send(&mut self, method: &str, params: serde_json::Value) -> i64 {
        let id = self.next_id;
        self.next_id += 1;
        let mut s = serde_json::to_string(
            &json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }),
        )
        .unwrap();
        s.push('\n');
        self.stdin.write_all(s.as_bytes()).await.unwrap();
        self.stdin.flush().await.unwrap();
        id
    }

    async fn recv_for(&mut self, want_id: i64) -> serde_json::Value {
        loop {
            let mut line = String::new();
            let n = tokio::time::timeout(Duration::from_secs(15), self.stdout.read_line(&mut line))
                .await
                .expect("recv timeout")
                .expect("read line");
            assert!(n > 0, "agent EOF");
            let v: serde_json::Value = serde_json::from_str(&line).expect("non-JSON line");
            if v.get("id") == Some(&json!(want_id)) {
                return v;
            }
        }
    }
}

async fn run_single_prompt(provider: &str, base: &str, model: &str) {
    let mut h = AgentHarness::spawn_provider(provider, base, model).await;
    h.send(
        "initialize",
        json!({ "protocolVersion": 1, "clientCapabilities": {} }),
    )
    .await;
    h.recv_for(1).await;
    h.send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
        .await;
    let r = h.recv_for(2).await;
    let sid = r["result"]["sessionId"].as_str().unwrap().to_string();
    h.send(
        "session/prompt",
        json!({ "sessionId": sid, "prompt": [{ "type": "text", "text": "say ok" }] }),
    )
    .await;
    let _ = h.recv_for(3).await;
}

async fn run_captured_prompt(
    provider: &str,
    model: &str,
    llm_canned: Vec<serde_json::Value>,
) -> CapturedRequest {
    // session/new triggers model catalog discovery against the same stub server.
    // Prepend a minimal valid discovery response (empty endpoints list) so the
    // discovery call is served cleanly and the LLM canned responses follow.
    // Legacy Databricks discovery hits /api/2.0/serving-endpoints;
    // Databricks v2 discovery hits /api/ai-gateway/v2/endpoints.
    let discovery_resp = json!({ "endpoints": [], "next_page_token": null });
    let mut all_canned = vec![discovery_resp];
    all_canned.extend(llm_canned);

    let (base, captured) = spawn_capturing_server(all_canned).await;
    run_single_prompt(provider, &base, model).await;

    // Filter out discovery requests — keep only the LLM invocation(s).
    // Discovery paths: /api/2.0/serving-endpoints, /api/ai-gateway/v2/endpoints*
    // LLM paths: /serving-endpoints/*, /ai-gateway/anthropic/*, /ai-gateway/openai/*, /ai-gateway/mlflow/*
    let reqs = captured.lock().await;
    let llm_reqs: Vec<_> = reqs
        .iter()
        .filter(|r| {
            !r.path.starts_with("/api/2.0/serving-endpoints")
                && !r.path.starts_with("/api/ai-gateway/v2/endpoints")
        })
        .collect();
    assert_eq!(llm_reqs.len(), 1, "expected exactly one LLM request");
    llm_reqs[0].clone()
}

#[tokio::test]
async fn databricks_envelope_routes_through_serving_endpoints_and_strips_model() {
    // One canned chat-completions-shaped response → assistant says "ok"
    // with end_turn so the agent loop exits cleanly.
    let canned = vec![json!({
        "id": "x",
        "object": "chat.completion",
        "choices": [{
            "index": 0,
            "message": { "role": "assistant", "content": "ok" },
            "finish_reason": "stop"
        }]
    })];
    let model = "goose-claude-4-6-sonnet";
    let req = run_captured_prompt("databricks", model, canned).await;

    assert_eq!(
        req.path.as_str(),
        format!("/serving-endpoints/{model}/invocations"),
        "Databricks must route to serving-endpoints/{{model}}/invocations"
    );
    assert_eq!(
        req.authorization.as_deref(),
        Some("Bearer test-bearer"),
        "Authorization must be the static DATABRICKS_TOKEN as a Bearer"
    );
    assert!(
        req.body.get("model").is_none(),
        "request body must NOT include `model` (Databricks rejects it): {:?}",
        req.body
    );
    // Sanity: the rest of the chat envelope should still be there.
    assert!(
        req.body
            .get("messages")
            .and_then(|v| v.as_array())
            .is_some(),
        "request body should keep the chat `messages` field"
    );
}

#[tokio::test]
async fn databricks_v2_gpt5_routes_through_ai_gateway_responses() {
    let canned = vec![json!({
        "status": "completed",
        "output": [{
            "type": "message",
            "content": [{ "type": "output_text", "text": "ok" }]
        }]
    })];
    let model = "databricks-gpt-5-5";
    let req = run_captured_prompt("databricks_v2", model, canned).await;

    assert_eq!(
        req.path.as_str(),
        "/ai-gateway/openai/v1/responses",
        "Databricks v2 GPT-5 models must route through AI Gateway Responses"
    );
    assert_eq!(
        req.authorization.as_deref(),
        Some("Bearer test-bearer"),
        "Authorization must be the static DATABRICKS_TOKEN as a Bearer"
    );
    assert_eq!(req.body["model"], model);
    assert!(
        req.body.get("input").and_then(|v| v.as_array()).is_some(),
        "Responses request body should keep `input`: {:?}",
        req.body
    );
}

#[tokio::test]
async fn databricks_v2_claude_routes_through_ai_gateway_anthropic_messages() {
    let canned = vec![json!({
        "stop_reason": "end_turn",
        "content": [{ "type": "text", "text": "ok" }]
    })];
    let model = "databricks-claude-opus-4-7";
    let req = run_captured_prompt("databricks_v2", model, canned).await;

    assert_eq!(
        req.path.as_str(),
        "/ai-gateway/anthropic/v1/messages",
        "Databricks v2 Claude models must route through AI Gateway Anthropic Messages"
    );
    assert_eq!(
        req.authorization.as_deref(),
        Some("Bearer test-bearer"),
        "Authorization must be the static DATABRICKS_TOKEN as a Bearer"
    );
    assert_eq!(req.body["model"], model);
    assert!(
        req.body
            .get("messages")
            .and_then(|v| v.as_array())
            .is_some(),
        "Anthropic request body should keep `messages`: {:?}",
        req.body
    );
}

#[tokio::test]
async fn databricks_v2_other_models_route_through_ai_gateway_mlflow_chat() {
    let canned = vec![json!({
        "id": "x",
        "object": "chat.completion",
        "choices": [{
            "index": 0,
            "message": { "role": "assistant", "content": "ok" },
            "finish_reason": "stop"
        }]
    })];
    let model = "custom-chat-model";
    let req = run_captured_prompt("databricks_v2", model, canned).await;

    assert_eq!(
        req.path.as_str(),
        "/ai-gateway/mlflow/v1/chat/completions",
        "Databricks v2 fallback models must route through AI Gateway MLflow Chat"
    );
    assert_eq!(
        req.authorization.as_deref(),
        Some("Bearer test-bearer"),
        "Authorization must be the static DATABRICKS_TOKEN as a Bearer"
    );
    assert_eq!(req.body["model"], model);
    assert!(
        req.body
            .get("messages")
            .and_then(|v| v.as_array())
            .is_some(),
        "Chat request body should keep `messages`: {:?}",
        req.body
    );
}

// ---------- session/set_model integration tests ----------

/// Helper: run initialize + session/new + optional set_model + session/prompt on a
/// freshly-spawned harness against the given stub server base URL. Returns the
/// session/prompt response and the session ID so callers can also call set_model.
async fn run_with_set_model(
    provider: &str,
    base: &str,
    initial_model: &str,
    switch_to_model: Option<&str>,
) -> (String, serde_json::Value) {
    let mut h = AgentHarness::spawn_provider(provider, base, initial_model).await;
    h.send(
        "initialize",
        json!({ "protocolVersion": 1, "clientCapabilities": {} }),
    )
    .await;
    h.recv_for(1).await;
    h.send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
        .await;
    let r = h.recv_for(2).await;
    let sid = r["result"]["sessionId"].as_str().unwrap().to_string();

    if let Some(new_model) = switch_to_model {
        h.send(
            "session/set_model",
            json!({ "sessionId": sid, "modelId": new_model }),
        )
        .await;
        let set_r = h.recv_for(3).await;
        // Verify the response carries the expected modelId.
        assert_eq!(
            set_r["result"]["modelId"],
            json!(new_model),
            "set_model response must echo the new modelId"
        );
        h.send(
            "session/prompt",
            json!({ "sessionId": sid, "prompt": [{ "type": "text", "text": "say ok" }] }),
        )
        .await;
        let prompt_r = h.recv_for(4).await;
        (sid, prompt_r)
    } else {
        h.send(
            "session/prompt",
            json!({ "sessionId": sid, "prompt": [{ "type": "text", "text": "say ok" }] }),
        )
        .await;
        let prompt_r = h.recv_for(3).await;
        (sid, prompt_r)
    }
}

/// After session/set_model switches from model A to model B, the next
/// session/prompt must route to B's Databricks serving-endpoint URL and
/// strip the `model` field from the body (legacy Databricks behaviour).
#[tokio::test]
async fn session_set_model_switches_databricks_legacy_route() {
    let initial_model = "initial-model";
    let switched_model = "switched-model";

    // Two canned responses: one for the discovery call (session/new),
    // one for the LLM call after the switch.
    let canned = vec![
        json!({ "endpoints": [], "next_page_token": null }), // discovery
        json!({                                               // LLM response
            "id": "x",
            "object": "chat.completion",
            "choices": [{
                "index": 0,
                "message": { "role": "assistant", "content": "ok" },
                "finish_reason": "stop"
            }]
        }),
    ];
    let (base, captured) = spawn_capturing_server(canned).await;
    run_with_set_model("databricks", &base, initial_model, Some(switched_model)).await;

    let reqs = captured.lock().await;
    let llm_reqs: Vec<_> = reqs
        .iter()
        .filter(|r| {
            !r.path.starts_with("/api/2.0/serving-endpoints")
                && !r.path.starts_with("/api/ai-gateway/v2/endpoints")
        })
        .collect();
    assert_eq!(
        llm_reqs.len(),
        1,
        "expected exactly one LLM request after switch"
    );
    let req = &llm_reqs[0];

    // The request must go to the SWITCHED model endpoint, not the initial one.
    assert_eq!(
        req.path.as_str(),
        format!("/serving-endpoints/{switched_model}/invocations"),
        "Databricks legacy must route to the switched model endpoint"
    );
    // The body must not include `model` (Databricks rejects it).
    assert!(
        req.body.get("model").is_none(),
        "request body must NOT include `model` after switch: {:?}",
        req.body
    );
}

/// After session/set_model switches a Databricks v2 session from a GPT-5 model
/// (OpenAI Responses route) to a Claude model (Anthropic Messages route),
/// the next prompt must hit the Anthropic AI Gateway path.
#[tokio::test]
async fn session_set_model_switches_databricks_v2_route() {
    let initial_model = "databricks-gpt-5-5"; // → OpenAI Responses
    let switched_model = "databricks-claude-opus-4-7"; // → Anthropic Messages

    let canned = vec![
        json!({ "endpoints": [], "next_page_token": null }), // discovery (v2: /api/ai-gateway/v2/endpoints)
        json!({                                               // LLM response (Anthropic Messages shape)
            "stop_reason": "end_turn",
            "content": [{ "type": "text", "text": "ok" }]
        }),
    ];
    let (base, captured) = spawn_capturing_server(canned).await;
    run_with_set_model("databricks_v2", &base, initial_model, Some(switched_model)).await;

    let reqs = captured.lock().await;
    let llm_reqs: Vec<_> = reqs
        .iter()
        .filter(|r| {
            !r.path.starts_with("/api/2.0/serving-endpoints")
                && !r.path.starts_with("/api/ai-gateway/v2/endpoints")
        })
        .collect();
    assert_eq!(
        llm_reqs.len(),
        1,
        "expected exactly one LLM request after v2 route switch"
    );
    let req = &llm_reqs[0];

    assert_eq!(
        req.path.as_str(),
        "/ai-gateway/anthropic/v1/messages",
        "After switching to a Claude model, Databricks v2 must route to Anthropic Messages"
    );
    assert_eq!(
        req.body["model"],
        json!(switched_model),
        "body must carry the switched model ID"
    );
}

/// session/set_model with an unknown session ID must return an invalid_params
/// error without touching any LLM endpoint.
#[tokio::test]
async fn session_set_model_unknown_session_returns_error() {
    // Spawn with a single discovery canned response; no LLM response needed.
    let canned = vec![json!({ "endpoints": [], "next_page_token": null })];
    let (base, _captured) = spawn_capturing_server(canned).await;

    let mut h = AgentHarness::spawn_provider("databricks", &base, "some-model").await;
    h.send(
        "initialize",
        json!({ "protocolVersion": 1, "clientCapabilities": {} }),
    )
    .await;
    h.recv_for(1).await;
    h.send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
        .await;
    h.recv_for(2).await;

    // Call set_model with a bogus session ID.
    h.send(
        "session/set_model",
        json!({ "sessionId": "nonexistent-session-id", "modelId": "new-model" }),
    )
    .await;
    let r = h.recv_for(3).await;

    assert!(
        r.get("error").is_some(),
        "set_model with unknown session must return an error: {:?}",
        r
    );
    let msg = r["error"]["message"].as_str().unwrap_or("");
    assert!(
        msg.contains("unknown session"),
        "error message must mention unknown session, got: {msg}"
    );
}

/// session/set_model with an empty modelId must return an invalid_params error.
#[tokio::test]
async fn session_set_model_empty_model_id_returns_error() {
    let canned = vec![json!({ "endpoints": [], "next_page_token": null })];
    let (base, _captured) = spawn_capturing_server(canned).await;

    let mut h = AgentHarness::spawn_provider("databricks", &base, "some-model").await;
    h.send(
        "initialize",
        json!({ "protocolVersion": 1, "clientCapabilities": {} }),
    )
    .await;
    h.recv_for(1).await;
    h.send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
        .await;
    let r = h.recv_for(2).await;
    let sid = r["result"]["sessionId"].as_str().unwrap().to_string();

    // Empty string modelId.
    h.send(
        "session/set_model",
        json!({ "sessionId": sid, "modelId": "" }),
    )
    .await;
    let r = h.recv_for(3).await;

    assert!(
        r.get("error").is_some(),
        "set_model with empty modelId must return an error: {:?}",
        r
    );
    let msg = r["error"]["message"].as_str().unwrap_or("");
    assert!(
        msg.contains("modelId"),
        "error message must mention modelId, got: {msg}"
    );
}

#[tokio::test]
async fn model_discovery_surfaces_rejected_static_token_as_auth_failure() {
    use axum::http::StatusCode;
    use buzz_agent::config::{Config, Provider};
    use buzz_agent::discover_databricks_models;

    let requests = Arc::new(AtomicU64::new(0));
    let requests_for_route = requests.clone();
    let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))
        .await
        .unwrap();
    let host = format!("http://{}", listener.local_addr().unwrap());
    let app = Router::new().route(
        "/api/ai-gateway/v2/endpoints",
        get(move || {
            let requests = requests_for_route.clone();
            async move {
                requests.fetch_add(1, Ordering::SeqCst);
                (StatusCode::UNAUTHORIZED, "rejected bearer rejected")
            }
        }),
    );
    tokio::spawn(async move {
        let _ = axum::serve(listener, app).await;
    });

    let cfg = Config::for_discovery(Provider::DatabricksV2, "rejected".into(), host);
    let error = discover_databricks_models(&cfg).await.unwrap_err();

    assert!(
        error.to_string().starts_with("llm auth:"),
        "401 must retain auth semantics: {error}"
    );
    assert!(
        !error.to_string().contains("rejected bearer"),
        "auth errors must not propagate provider bodies that may echo credentials: {error}"
    );
    assert_eq!(
        requests.load(Ordering::SeqCst),
        1,
        "a static token cannot refresh, so discovery must not issue a duplicate request"
    );
}

fn databricks_oauth_cache_path(home: &std::path::Path, host: &str) -> std::path::PathBuf {
    let discovery_url = format!(
        "{}/oidc/.well-known/oauth-authorization-server",
        host.trim_end_matches('/')
    );
    let mut hasher = Sha256::new();
    hasher.update(discovery_url.as_bytes());
    hasher.update(b"|");
    hasher.update(b"databricks-cli");
    hasher.update(b"|");
    hasher.update(b"all-apis,offline_access");
    let hash = hex::encode(hasher.finalize());
    home.join(".config")
        .join("buzz-agent")
        .join("oauth")
        .join("databricks")
        .join(format!("{hash}.json"))
}

fn write_cached_oauth_token(home: &std::path::Path, host: &str, access_token: &str) {
    let path = databricks_oauth_cache_path(home, host);
    std::fs::create_dir_all(path.parent().unwrap()).unwrap();
    std::fs::write(
        path,
        serde_json::to_vec(&json!({
            "access_token": access_token,
            "refresh_token": null,
            "expires_at": SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs()
                + 3600,
        }))
        .unwrap(),
    )
    .unwrap();
}

#[tokio::test]
async fn oauth_missing_token_uses_configured_model_then_retries_discovery() {
    let attempts = Arc::new(AtomicU64::new(0));
    let attempts_for_route = attempts.clone();
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let host = format!("http://{}", listener.local_addr().unwrap());
    let app = Router::new().route(
        "/api/ai-gateway/v2/endpoints",
        get(move || {
            let attempts = attempts_for_route.clone();
            async move {
                attempts.fetch_add(1, Ordering::SeqCst);
                Json(json!({
                    "endpoints": [{"name": "authenticated-model"}],
                    "next_page_token": null,
                }))
            }
        }),
    );
    tokio::spawn(async move {
        let _ = axum::serve(listener, app).await;
    });

    let configured_model = "  configured-model  ";
    let mut h =
        AgentHarness::spawn_oauth_provider("databricks_v2", &host, configured_model, 2).await;
    let initialize = h
        .send(
            "initialize",
            json!({ "protocolVersion": 1, "clientCapabilities": {} }),
        )
        .await;
    assert!(h.recv_for(initialize).await.get("result").is_some());

    let first = h
        .send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
        .await;
    let first_response = h.recv_for(first).await;
    assert!(
        first_response["result"]["sessionId"].is_string(),
        "missing OAuth token blocked session creation: {first_response}"
    );
    assert_eq!(
        first_response["result"]["models"]["availableModels"],
        json!([{"modelId": "configured-model", "name": "configured-model"}])
    );
    assert_eq!(attempts.load(Ordering::SeqCst), 0);

    write_cached_oauth_token(h.oauth_home(), &host, "cached-bearer");

    let second = h
        .send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
        .await;
    let second_response = h.recv_for(second).await;
    assert!(
        second_response["result"]["sessionId"].is_string(),
        "later authenticated session failed: {second_response}"
    );
    assert_eq!(
        second_response["result"]["models"]["availableModels"],
        json!([{"modelId": "authenticated-model", "name": "authenticated-model"}])
    );
    assert_eq!(
        attempts.load(Ordering::SeqCst),
        1,
        "OAuth fallback was cached instead of retrying discovery"
    );
}

#[tokio::test]
async fn non_auth_discovery_failure_uses_configured_model_without_caching_fallback() {
    use axum::http::StatusCode;

    let attempts = Arc::new(AtomicU64::new(0));
    let attempts_for_route = attempts.clone();
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let host = format!("http://{}", listener.local_addr().unwrap());
    let app = Router::new().route(
        "/api/ai-gateway/v2/endpoints",
        get(move || {
            let attempts = attempts_for_route.clone();
            async move {
                attempts.fetch_add(1, Ordering::SeqCst);
                (StatusCode::SERVICE_UNAVAILABLE, "catalog unavailable")
            }
        }),
    );
    tokio::spawn(async move {
        let _ = axum::serve(listener, app).await;
    });

    let configured_model = "  configured-model  ";
    let normalized_configured_model = configured_model.trim();
    let mut h =
        AgentHarness::spawn_provider_with_max_sessions("databricks_v2", &host, configured_model, 2)
            .await;
    let initialize = h
        .send(
            "initialize",
            json!({ "protocolVersion": 1, "clientCapabilities": {} }),
        )
        .await;
    assert!(h.recv_for(initialize).await.get("result").is_some());

    for expected_attempts in 1..=2 {
        let request = h
            .send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
            .await;
        let response = h.recv_for(request).await;
        assert!(
            response["result"]["sessionId"].is_string(),
            "non-auth catalog failure blocked session creation: {response}"
        );
        assert_eq!(
            response["result"]["models"]["availableModels"],
            json!([{"modelId": normalized_configured_model, "name": normalized_configured_model}])
        );
        assert_eq!(attempts.load(Ordering::SeqCst), expected_attempts);
    }
}

#[tokio::test]
async fn rejected_static_token_does_not_consume_capacity_or_spawn_mcp() {
    use axum::http::StatusCode;

    let attempts = Arc::new(AtomicU64::new(0));
    let attempts_for_route = attempts.clone();
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let host = format!("http://{}", listener.local_addr().unwrap());
    let app = Router::new().route(
        "/api/ai-gateway/v2/endpoints",
        get(move || {
            let attempts = attempts_for_route.clone();
            async move {
                if attempts.fetch_add(1, Ordering::SeqCst) == 0 {
                    Err((StatusCode::UNAUTHORIZED, "rejected"))
                } else {
                    Ok(Json(json!({
                        "endpoints": [{"name": "discovered-model"}],
                        "next_page_token": null,
                    })))
                }
            }
        }),
    );
    tokio::spawn(async move {
        let _ = axum::serve(listener, app).await;
    });

    let mut h = AgentHarness::spawn_provider("databricks_v2", &host, "discovered-model").await;
    let initialize = h
        .send(
            "initialize",
            json!({ "protocolVersion": 1, "clientCapabilities": {} }),
        )
        .await;
    assert!(h.recv_for(initialize).await.get("result").is_some());

    let pid_dir = TempDir::new().unwrap();
    let pid_file = pid_dir.path().join("mcp.pid");
    let fake_mcp = env!("CARGO_BIN_EXE_fake-mcp");
    let mcp_servers = json!([{
        "name": "must-not-spawn",
        "command": fake_mcp,
        "args": [],
        "env": [{
            "name": "FAKE_MCP_PID_FILE",
            "value": pid_file.to_string_lossy(),
        }],
    }]);

    let failed = h
        .send(
            "session/new",
            json!({ "cwd": "/tmp", "mcpServers": mcp_servers }),
        )
        .await;
    let failed_response = h.recv_for(failed).await;
    assert!(failed_response.get("error").is_some(), "{failed_response}");
    assert!(
        failed_response["error"]["message"]
            .as_str()
            .unwrap_or_default()
            .contains("llm auth"),
        "rejected static token did not retain auth semantics: {failed_response}"
    );
    tokio::time::sleep(Duration::from_millis(100)).await;
    assert!(
        !pid_file.exists(),
        "MCP process spawned before failed discovery was resolved"
    );

    let retry = h
        .send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
        .await;
    let retry_response = h.recv_for(retry).await;
    assert!(
        retry_response["result"]["sessionId"].is_string(),
        "failed discovery consumed the sole session slot: {retry_response}"
    );
    assert_eq!(attempts.load(Ordering::SeqCst), 2);
}