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 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,
}

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 {
        let bin = env!("CARGO_BIN_EXE_buzz-agent");
        let mut cmd = tokio::process::Command::new(bin);
        cmd.env("BUZZ_AGENT_PROVIDER", provider)
            .env("DATABRICKS_HOST", base_url)
            .env("DATABRICKS_MODEL", model)
            .env("DATABRICKS_TOKEN", "test-bearer")
            .env("BUZZ_AGENT_LLM_TIMEOUT_SECS", "5")
            .env("BUZZ_AGENT_TOOL_TIMEOUT_SECS", "5")
            .env("BUZZ_AGENT_MAX_ROUNDS", "2")
            .env("BUZZ_AGENT_MCP_INIT_TIMEOUT_SECS", "2")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .kill_on_drop(true);
        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,
        }
    }

    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}"
    );
}