codexia 0.3.2

//! HTTP server state, router wiring, and endpoint tests.

use crate::{
    codex::client::CodexClient, error::Result, openai::response::ModelList, status::StatusClient,
    token::TokenManager,
};
use axum::{
    Router,
    routing::{get, post},
};
use reqwest::Client;
use std::{net::SocketAddr, sync::Arc};
use tokio::net::TcpListener;

mod auth;
mod handlers;
mod status_response;

pub use auth::authorize;

/// Shared application state used by all HTTP route handlers.
#[derive(Clone)]
pub struct AppState {
    token_manager: TokenManager,
    codex: CodexClient,
    status: StatusClient,
    api_key: Option<Arc<str>>,
    models: ModelList,
}

impl AppState {
    /// Builds application state from the configured token manager, upstream client, and model list.
    #[must_use]
    pub fn new(
        token_manager: TokenManager,
        codex: CodexClient,
        api_key: Option<String>,
        models: ModelList,
    ) -> Self {
        let status = StatusClient::new(Client::new(), codex.base_url().to_owned());
        Self {
            token_manager,
            status,
            codex,
            api_key: api_key.map(Arc::from),
            models,
        }
    }
}

/// Binds the local listener and serves the Axum router until shutdown.
///
/// # Errors
///
/// Returns an error when binding the socket or serving the router fails.
pub async fn serve(addr: SocketAddr, state: AppState) -> Result<()> {
    let listener = TcpListener::bind(addr).await?;
    axum::serve(listener, router(state)).await?;
    Ok(())
}

/// Builds the HTTP router for all supported endpoints.
pub fn router(state: AppState) -> Router {
    // Keep route registration centralized so the server surface stays easy to audit.
    Router::new()
        .route("/health", get(handlers::health))
        .route("/v1/auth/refresh", post(handlers::manual_refresh))
        .route("/v1/status", get(handlers::status))
        .route("/v1/models", get(handlers::models))
        .route("/v1/messages", post(handlers::messages))
        .route(
            "/v1/messages/count_tokens",
            post(handlers::count_message_tokens),
        )
        .route("/v1/chat/completions", post(handlers::chat_completions))
        .with_state(state)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        Error,
        config::{AuthStore, Credentials, now_unix},
        oauth::CodexOAuthClient,
        openai::response::ModelList,
        testsupport::TempDir,
    };
    use axum::{
        Json,
        body::to_bytes,
        extract::Form,
        http::{HeaderMap, HeaderValue, StatusCode},
        routing::{get, post},
    };
    use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
    use reqwest::Client;
    use serde_json::{Value, json};
    use std::collections::HashMap;
    use tokio::net::TcpListener;

    fn jwt_with_account_id(account_id: &str) -> String {
        let encoded = URL_SAFE_NO_PAD.encode(
            serde_json::to_vec(&json!({
                "https://api.openai.com/auth": { "chatgpt_account_id": account_id }
            }))
            .unwrap(),
        );
        format!("header.{encoded}.sig")
    }

    async fn refresh_handler(Form(form): Form<HashMap<String, String>>) -> Json<Value> {
        assert_eq!(form.get("refresh_token").unwrap(), "old_refresh");
        Json(json!({
            "access_token": jwt_with_account_id("acc_refreshed"),
            "refresh_token": "new_refresh",
            "expires_in": 3600
        }))
    }

    async fn spawn_refresh_server() -> String {
        let app = Router::new().route("/token", post(refresh_handler));
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = format!("http://{}/token", listener.local_addr().unwrap());
        tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        url
    }

    async fn account_handler() -> Json<Value> {
        Json(json!({
            "accounts": {
                "default": {
                    "account": {
                        "name": "Personal",
                        "structure": "personal"
                    },
                    "entitlement": {
                        "subscription_plan": "chatgptplus",
                        "has_active_subscription": true,
                        "expires_at": "2026-05-01T00:00:00Z"
                    }
                }
            }
        }))
    }

    async fn usage_handler() -> Json<Value> {
        Json(json!({
            "email": "test@example.com",
            "plan_type": "pro",
            "rate_limit": {
                "primary_window": {
                    "used_percent": 10,
                    "reset_at": "2026-04-27T12:00:00Z"
                },
                "secondary_window": {
                    "remaining_percent": 90,
                    "reset_at": "2026-05-01T00:00:00Z"
                }
            },
            "credits": { "balance": 1 }
        }))
    }

    async fn spawn_status_server() -> String {
        let app = Router::new()
            .route("/accounts/check/v4-2023-04-27", get(account_handler))
            .route("/wham/usage", get(usage_handler));
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = format!("http://{}", listener.local_addr().unwrap());
        tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        url
    }

    async fn codex_complete_handler() -> impl axum::response::IntoResponse {
        (
            [(axum::http::header::CONTENT_TYPE, "text/event-stream")],
            "data: {\"type\":\"response.completed\",\"response\":{\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"OK\"}]}],\"usage\":{\"input_tokens\":12,\"output_tokens\":5,\"total_tokens\":17}}}\n\n",
        )
    }

    async fn codex_tool_call_handler() -> impl axum::response::IntoResponse {
        (
            [(axum::http::header::CONTENT_TYPE, "text/event-stream")],
            "data: {\"type\":\"response.completed\",\"response\":{\"output\":[{\"type\":\"function_call\",\"call_id\":\"call_1\",\"name\":\"lookup\",\"arguments\":\"{\\\"q\\\":\\\"x\\\"}\"}],\"usage\":{\"input_tokens\":7,\"output_tokens\":3,\"total_tokens\":10}}}\n\n",
        )
    }

    async fn spawn_codex_server(tool_call: bool) -> String {
        let app = if tool_call {
            Router::new().route("/codex/responses", post(codex_tool_call_handler))
        } else {
            Router::new().route("/codex/responses", post(codex_complete_handler))
        };
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let url = format!("http://{}", listener.local_addr().unwrap());
        tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        url
    }

    fn test_state(store: AuthStore, token_url: String, api_key: Option<String>) -> AppState {
        let http = Client::new();
        AppState::new(
            TokenManager::new(
                store,
                CodexOAuthClient::new_with_token_url(http.clone(), token_url),
            ),
            CodexClient::new(http, "http://codex.invalid"),
            api_key,
            ModelList::from_ids(["gpt-test"]),
        )
    }

    #[test]
    fn authorizes_when_no_api_key_is_configured() {
        assert!(authorize(&HeaderMap::new(), None).is_ok());
    }

    #[test]
    fn authorizes_bearer_token() {
        let mut headers = HeaderMap::new();
        headers.insert(
            axum::http::header::AUTHORIZATION,
            HeaderValue::from_static("Bearer secret"),
        );

        assert!(authorize(&headers, Some("secret")).is_ok());
    }

    #[test]
    fn rejects_wrong_api_key() {
        let mut headers = HeaderMap::new();
        headers.insert("x-api-key", HeaderValue::from_static("wrong"));

        assert!(matches!(
            authorize(&headers, Some("secret")),
            Err(Error::Unauthorized)
        ));
    }

    #[tokio::test]
    async fn manual_refresh_refreshes_credentials_without_returning_tokens() {
        let dir = TempDir::new().unwrap();
        let store = AuthStore::new(dir.path().join("auth.json"));
        store
            .save(&Credentials {
                access_token: "old_access".into(),
                refresh_token: "old_refresh".into(),
                expires_at: now_unix() + 600,
                account_id: "acc_old".into(),
            })
            .unwrap();
        let state = test_state(
            store.clone(),
            spawn_refresh_server().await,
            Some("secret".into()),
        );
        let mut headers = HeaderMap::new();
        headers.insert("x-api-key", HeaderValue::from_static("secret"));

        let response = handlers::manual_refresh(axum::extract::State(state), headers).await;

        assert_eq!(response.status(), StatusCode::OK);
        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
        let value = serde_json::from_slice::<Value>(&body).unwrap();
        assert_eq!(value["account_id"], "acc_refreshed");
        assert!(value.get("expires_at").is_some());
        assert!(value.get("access_token").is_none());
        assert!(value.get("refresh_token").is_none());

        let saved = store.load().unwrap().unwrap();
        assert_eq!(saved.refresh_token, "new_refresh");
        assert_eq!(saved.account_id, "acc_refreshed");
    }

    #[tokio::test]
    async fn manual_refresh_requires_api_key_when_configured() {
        let dir = TempDir::new().unwrap();
        let store = AuthStore::new(dir.path().join("auth.json"));
        let state = test_state(store, spawn_refresh_server().await, Some("secret".into()));

        let response =
            handlers::manual_refresh(axum::extract::State(state), HeaderMap::new()).await;

        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
    }

    #[tokio::test]
    async fn status_returns_account_and_rate_limit_snapshot() {
        let dir = TempDir::new().unwrap();
        let store = AuthStore::new(dir.path().join("auth.json"));
        store
            .save(&Credentials {
                access_token: "access".into(),
                refresh_token: "refresh".into(),
                expires_at: now_unix() + 600,
                account_id: "acc_old".into(),
            })
            .unwrap();

        let status_base_url = spawn_status_server().await;
        let http = Client::new();
        let state = AppState::new(
            TokenManager::new(
                store,
                CodexOAuthClient::new_with_token_url(http.clone(), spawn_refresh_server().await),
            ),
            CodexClient::new(http, status_base_url),
            Some("secret".into()),
            ModelList::from_ids(["gpt-test"]),
        );
        let mut headers = HeaderMap::new();
        headers.insert("x-api-key", HeaderValue::from_static("secret"));

        let response = handlers::status(axum::extract::State(state), headers).await;

        assert_eq!(response.status(), StatusCode::OK);
        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
        let value = serde_json::from_slice::<Value>(&body).unwrap();
        assert_eq!(value["account_id"], "acc_old");
        assert_eq!(value["account"]["plan"], "chatgptplus");
        assert!(value["account"]["subscription_expires_at_local"].is_string());
        assert!(value["account"]["subscription_remaining_seconds"].is_number());
        assert_eq!(value["rate_limits"][0]["name"], "5h");
        assert_eq!(value["rate_limits"][0]["remaining_percent"], 90.0);
        assert!(value["rate_limits"][0]["reset_at_local"].is_string());
        assert!(value["rate_limits"][0]["reset_in_seconds"].is_number());
        assert!(value["token"]["expires_at_local"].is_string());
        assert!(value["warnings"].as_array().unwrap().is_empty());
    }

    #[tokio::test]
    async fn messages_returns_anthropic_message_response() {
        let dir = TempDir::new().unwrap();
        let store = AuthStore::new(dir.path().join("auth.json"));
        store
            .save(&Credentials {
                access_token: "access".into(),
                refresh_token: "refresh".into(),
                expires_at: now_unix() + 600,
                account_id: "acc_old".into(),
            })
            .unwrap();
        let http = Client::new();
        let state = AppState::new(
            TokenManager::new(
                store,
                CodexOAuthClient::new_with_token_url(http.clone(), spawn_refresh_server().await),
            ),
            CodexClient::new(http, spawn_codex_server(false).await),
            Some("secret".into()),
            ModelList::from_ids(["gpt-5.5"]),
        );
        let mut headers = HeaderMap::new();
        headers.insert("x-api-key", HeaderValue::from_static("secret"));

        let response = handlers::messages(
            axum::extract::State(state),
            headers,
            Json(
                serde_json::from_value(json!({
                    "model": "gpt-5.5",
                    "max_tokens": 128,
                    "messages": [{"role": "user", "content": "hello"}]
                }))
                .unwrap(),
            ),
        )
        .await;

        assert_eq!(response.status(), StatusCode::OK);
        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
        let value = serde_json::from_slice::<Value>(&body).unwrap();
        assert_eq!(value["type"], "message");
        assert_eq!(value["role"], "assistant");
        assert_eq!(value["content"][0]["type"], "text");
        assert_eq!(value["content"][0]["text"], "OK");
        assert_eq!(value["stop_reason"], "end_turn");
        assert_eq!(value["usage"]["input_tokens"], 12);
    }

    #[tokio::test]
    async fn messages_returns_tool_use_blocks() {
        let dir = TempDir::new().unwrap();
        let store = AuthStore::new(dir.path().join("auth.json"));
        store
            .save(&Credentials {
                access_token: "access".into(),
                refresh_token: "refresh".into(),
                expires_at: now_unix() + 600,
                account_id: "acc_old".into(),
            })
            .unwrap();
        let http = Client::new();
        let state = AppState::new(
            TokenManager::new(
                store,
                CodexOAuthClient::new_with_token_url(http.clone(), spawn_refresh_server().await),
            ),
            CodexClient::new(http, spawn_codex_server(true).await),
            Some("secret".into()),
            ModelList::from_ids(["gpt-5.5"]),
        );
        let mut headers = HeaderMap::new();
        headers.insert("x-api-key", HeaderValue::from_static("secret"));

        let response = handlers::messages(
            axum::extract::State(state),
            headers,
            Json(
                serde_json::from_value(json!({
                    "model": "gpt-5.5",
                    "max_tokens": 128,
                    "messages": [{"role": "user", "content": "hello"}]
                }))
                .unwrap(),
            ),
        )
        .await;

        assert_eq!(response.status(), StatusCode::OK);
        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
        let value = serde_json::from_slice::<Value>(&body).unwrap();
        assert_eq!(value["content"][0]["type"], "tool_use");
        assert_eq!(value["content"][0]["id"], "call_1");
        assert_eq!(value["content"][0]["name"], "lookup");
        assert_eq!(value["stop_reason"], "tool_use");
    }

    #[tokio::test]
    async fn count_message_tokens_returns_estimate() {
        let dir = TempDir::new().unwrap();
        let store = AuthStore::new(dir.path().join("auth.json"));
        let state = test_state(store, spawn_refresh_server().await, Some("secret".into()));
        let mut headers = HeaderMap::new();
        headers.insert("x-api-key", HeaderValue::from_static("secret"));

        let response = handlers::count_message_tokens(
            axum::extract::State(state),
            headers,
            Json(
                serde_json::from_value(json!({
                    "model": "gpt-5.5",
                    "messages": [{"role": "user", "content": "Reply with the single word OK"}]
                }))
                .unwrap(),
            ),
        )
        .await;

        assert_eq!(response.status(), StatusCode::OK);
        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
        let value = serde_json::from_slice::<Value>(&body).unwrap();
        assert!(value["input_tokens"].as_u64().unwrap() > 0);
    }
}