path: root/libpanto/src/auth.zig

//! Provider authentication: named auth sessions resolved into request-ready
//! credentials before a provider stream opens.
//!
//! A provider names the auth session it uses (`auth = "<name>"`); core
//! resolves that session into a `ResolvedCredential` (a bearer/api-key value,
//! an optional dynamic `base_url`, and optional auth-derived request headers).
//! This covers three families with one configuration shape:
//!
//!   - `api_key`     — static key from a literal or an environment variable.
//!   - `oauth_device` — OAuth 2.0 device-authorization flow. Two dialects:
//!       * `token` — standard device flow (GitHub Copilot). The poll endpoint
//!         returns the OAuth token response directly.
//!       * `codex` — OpenAI Codex device flow. The poll endpoint returns an
//!         authorization code plus a server-generated PKCE verifier; core
//!         exchanges those at `token_url` for `{access,refresh,id}_token`.
//!
//! This module owns the *mechanism* (config types, the persisted token shape,
//! and — added incrementally — the HTTP flows and refresh lifecycle). It is
//! UI- and filesystem-policy-agnostic: token storage takes a directory path,
//! and interactive device-code prompts are delivered through a caller-supplied
//! `Presenter`. That keeps libpanto reusable while the embedder owns where
//! `$PANTO_HOME` is and how a device code is shown to the user.

const std = @import("std");
const builtin = @import("builtin");
const Io = std.Io;
const config_mod = @import("config.zig");
const http = @import("http_helper.zig");

/// A single request header. Re-exported from `config.zig` so callers can
/// build auth-derived headers and provider `extra_headers` from one type.
pub const Header = config_mod.Header;

// ===========================================================================
// Auth configuration (parsed from `[auth.<name>]`)
// ===========================================================================

pub const AuthType = enum { api_key, oauth_device };

/// Device-flow completion shape. See the module doc.
pub const DeviceDialect = enum { token, codex };

/// Wire encoding for the `token` dialect's device-code / poll request bodies.
pub const TokenRequestFormat = enum { form, json };

/// Static API-key auth. `key` is the resolved credential — a literal, or the
/// result of `${env:VAR}`-style substitution performed by the embedder. An
/// absent/empty key means unresolved; providers using such a session are
/// omitted from the active config (export the key and they reappear).
pub const ApiKeyAuth = struct {
    key: ?[]const u8 = null,
};

/// A secondary token exchange run after the durable OAuth token is obtained
/// (GitHub Copilot's `/copilot_internal/v2/token`). The exchange result —
/// not the OAuth token — becomes the request credential, and may also
/// override the provider `base_url`.
pub const ExchangeConfig = struct {
    method: []const u8 = "GET",
    url: []const u8,
    /// Which stored token to send as the exchange request's bearer.
    bearer: BearerSource = .oauth_access_token,
    /// Dotted JSON path to the credential in the exchange response.
    token_json_path: []const u8 = "token",
    /// Dotted JSON path to the credential's unix-seconds expiry, if any.
    expires_at_json_path: ?[]const u8 = null,
    /// Dotted JSON path to a dynamic `base_url` in the exchange response.
    base_url_json_path: ?[]const u8 = null,

    pub const BearerSource = enum { oauth_access_token };
};

/// OAuth 2.0 device-authorization auth.
pub const OAuthDeviceAuth = struct {
    dialect: DeviceDialect = .token,
    client_id: []const u8,
    /// Endpoint that issues the device/user code.
    device_code_url: []const u8,
    /// `token` dialect: poll-for-token and refresh endpoint.
    /// `codex` dialect: authorization-code exchange and refresh endpoint.
    token_url: []const u8,
    /// `codex` dialect: endpoint polled for the authorization code.
    device_poll_url: ?[]const u8 = null,
    /// Browser URL shown to the user (codex). For the `token` dialect the
    /// verification URI comes from the device-code response.
    verification_url: ?[]const u8 = null,
    scope: ?[]const u8 = null,
    /// Request encoding for the `token` dialect device-code / poll bodies.
    token_request_format: TokenRequestFormat = .form,
    /// `codex` dialect: redirect URI sent in the code exchange.
    redirect_uri: ?[]const u8 = null,
    /// `codex` dialect: JWT claim (on the id_token) holding the account id,
    /// e.g. `https://api.openai.com/auth`. When set, the resolver extracts an
    /// account id and emits it as a header (see provider wiring).
    account_id_jwt_claim: ?[]const u8 = null,
    /// Optional post-login token exchange (Copilot).
    exchange: ?ExchangeConfig = null,
};

/// One named auth session's configuration. The union tag mirrors the TOML
/// `type` field.
pub const AuthConfig = union(AuthType) {
    api_key: ApiKeyAuth,
    oauth_device: OAuthDeviceAuth,
};

// ===========================================================================
// Persisted token state (`$PANTO_HOME/auth/<name>.json`)
// ===========================================================================

/// Durable auth state for one session. Only the fields relevant to the auth
/// type are populated. Treat the on-disk file like a password.
pub const TokenSet = struct {
    /// `"api_key"` | `"oauth_device"` — records which family wrote the file.
    type: []const u8 = "oauth_device",
    /// Durable OAuth access token (the `ghu_...` for Copilot; the JWT access
    /// token for Codex).
    access_token: ?[]const u8 = null,
    refresh_token: ?[]const u8 = null,
    id_token: ?[]const u8 = null,
    /// Unix-seconds expiry of `access_token` (when known; OAuth `expires_in`
    /// or a JWT `exp`). Null for tokens with no intrinsic expiry (Copilot's
    /// durable `ghu_` token).
    expires_at: ?i64 = null,
    /// Account id derived from the id_token (Codex `chatgpt-account-id`).
    account_id: ?[]const u8 = null,
    /// Result of the secondary exchange (Copilot short-lived API token).
    exchange: ?ExchangeToken = null,

    pub const ExchangeToken = struct {
        token: []const u8,
        /// Unix-seconds expiry of the exchanged token.
        expires_at: ?i64 = null,
        /// Dynamic base_url returned by the exchange, if any.
        base_url: ?[]const u8 = null,
    };
};

// ===========================================================================
// Resolved credential (handed to the provider for one turn)
// ===========================================================================

/// The request-ready output of resolving an auth session: the secret to place
/// in the provider's auth header, plus any dynamic base_url and auth-derived
/// headers. The embedder injects these into the active `ProviderConfig`.
pub const ResolvedCredential = struct {
    /// Value for the provider auth header (the bearer / x-api-key value).
    api_key: []const u8,
    /// Overrides the provider's configured `base_url` when present (e.g.
    /// Copilot's `endpoints.api`).
    base_url_override: ?[]const u8 = null,
    /// Auth-derived headers (e.g. `chatgpt-account-id`) to merge onto the
    /// provider's request.
    extra_headers: []const Header = &.{},
};

// ===========================================================================
// Token storage (`$PANTO_HOME/auth/<name>.json`)
// ===========================================================================

/// A parsed `TokenSet` plus the arena owning its strings. Call `deinit`.
pub const ParsedTokenSet = std.json.Parsed(TokenSet);

/// Owner-only file permissions for token files (POSIX 0o600). On Windows the
/// platform `Permissions` enum has no mode, so fall back to its default.
fn tokenFilePermissions() Io.File.Permissions {
    if (builtin.os.tag == .windows) return .default_file;
    return Io.File.Permissions.fromMode(0o600);
}

/// Load and parse `<auth_dir>/<name>.json`. Returns null if the file does not
/// exist. The caller owns the result and must `deinit` it.
pub fn loadTokenSet(
    allocator: std.mem.Allocator,
    io: Io,
    auth_dir: []const u8,
    name: []const u8,
) !?ParsedTokenSet {
    const fname = try std.fmt.allocPrint(allocator, "{s}.json", .{name});
    defer allocator.free(fname);

    var dir = Io.Dir.cwd().openDir(io, auth_dir, .{}) catch |err| switch (err) {
        error.FileNotFound => return null,
        else => return err,
    };
    defer dir.close(io);

    const bytes = dir.readFileAlloc(io, fname, allocator, .limited(1 << 20)) catch |err| switch (err) {
        error.FileNotFound => return null,
        else => return err,
    };
    defer allocator.free(bytes);

    return try std.json.parseFromSlice(TokenSet, allocator, bytes, .{
        .ignore_unknown_fields = true,
        .allocate = .alloc_always,
    });
}

/// Serialize and write `ts` to `<auth_dir>/<name>.json` with owner-only
/// permissions, creating `auth_dir` if needed. Null optional fields are
/// omitted to keep the file tidy.
pub fn saveTokenSet(
    allocator: std.mem.Allocator,
    io: Io,
    auth_dir: []const u8,
    name: []const u8,
    ts: TokenSet,
) !void {
    Io.Dir.cwd().createDirPath(io, auth_dir) catch |err| switch (err) {
        error.PathAlreadyExists => {},
        else => return err,
    };

    const json = try std.json.Stringify.valueAlloc(allocator, ts, .{
        .emit_null_optional_fields = false,
        .whitespace = .indent_2,
    });
    defer allocator.free(json);

    const fname = try std.fmt.allocPrint(allocator, "{s}.json", .{name});
    defer allocator.free(fname);

    var dir = try Io.Dir.cwd().openDir(io, auth_dir, .{});
    defer dir.close(io);
    try dir.writeFile(io, .{
        .sub_path = fname,
        .data = json,
        .flags = .{ .permissions = tokenFilePermissions() },
    });
}

/// Delete `<auth_dir>/<name>.json`. Returns true if a file was removed, false
/// if it did not exist (idempotent logout).
pub fn deleteTokenSet(
    allocator: std.mem.Allocator,
    io: Io,
    auth_dir: []const u8,
    name: []const u8,
) !bool {
    const fname = try std.fmt.allocPrint(allocator, "{s}.json", .{name});
    defer allocator.free(fname);
    var dir = Io.Dir.cwd().openDir(io, auth_dir, .{}) catch |err| switch (err) {
        error.FileNotFound => return false,
        else => return err,
    };
    defer dir.close(io);
    dir.deleteFile(io, fname) catch |err| switch (err) {
        error.FileNotFound => return false,
        else => return err,
    };
    return true;
}

// ===========================================================================
// OAuth device flow
// ===========================================================================

pub const AuthError = error{
    /// No persisted token and no interactive presenter to run a login.
    AuthLoginRequired,
    /// The OAuth endpoint returned an error (denied, expired, malformed).
    AuthFlowFailed,
    /// The device-flow poll exceeded its deadline.
    AuthLoginTimeout,
    /// The configured auth could not be resolved into a credential.
    AuthUnresolved,
};

/// The device-code prompt shown to the user. The embedder renders this
/// (a verification URL + a short user code) however suits its UI.
pub const DeviceCodePrompt = struct {
    verification_uri: []const u8,
    user_code: []const u8,
    /// Seconds until the code expires (0 if the endpoint didn't say).
    expires_in: i64 = 0,
};

/// Caller-supplied UI hook for interactive device login. `onDeviceCode` is
/// called once with the prompt; `onStatus` (optional) receives progress text.
pub const Presenter = struct {
    ptr: *anyopaque,
    vtable: *const VTable,

    pub const VTable = struct {
        on_device_code: *const fn (ptr: *anyopaque, prompt: DeviceCodePrompt) void,
        on_status: ?*const fn (ptr: *anyopaque, msg: []const u8) void = null,
    };

    fn deviceCode(self: Presenter, prompt: DeviceCodePrompt) void {
        self.vtable.on_device_code(self.ptr, prompt);
    }
    fn status(self: Presenter, msg: []const u8) void {
        if (self.vtable.on_status) |f| f(self.ptr, msg);
    }
};

/// The OAuth token-endpoint response, fields borrowed from an arena.
pub const OAuthTokens = struct {
    access_token: []const u8,
    refresh_token: ?[]const u8 = null,
    id_token: ?[]const u8 = null,
    /// Seconds-to-live from `expires_in`, if present.
    expires_in: ?i64 = null,
};

/// Identifiers needed to poll a started device authorization.
const DeviceAuth = struct {
    /// `token` dialect poll parameter.
    device_code: ?[]const u8 = null,
    /// `codex` dialect poll parameter.
    device_auth_id: ?[]const u8 = null,
    user_code: []const u8,
    verification_uri: []const u8,
    interval_secs: u32 = 5,
    expires_in: i64 = 0,
};

/// `application/x-www-form-urlencoded` body from name/value pairs. Values are
/// percent-encoded (unreserved chars pass through). Allocated in `arena`.
fn formEncode(arena: std.mem.Allocator, pairs: []const [2][]const u8) ![]u8 {
    var out: std.ArrayList(u8) = .empty;
    for (pairs, 0..) |kv, i| {
        if (i != 0) try out.append(arena, '&');
        try percentEncode(arena, &out, kv[0]);
        try out.append(arena, '=');
        try percentEncode(arena, &out, kv[1]);
    }
    return out.toOwnedSlice(arena);
}

fn percentEncode(arena: std.mem.Allocator, out: *std.ArrayList(u8), s: []const u8) !void {
    for (s) |c| {
        const unreserved = (c >= 'A' and c <= 'Z') or (c >= 'a' and c <= 'z') or
            (c >= '0' and c <= '9') or c == '-' or c == '.' or c == '_' or c == '~';
        if (unreserved) {
            try out.append(arena, c);
        } else {
            try out.print(arena, "%{X:0>2}", .{c});
        }
    }
}

/// Build the device-code / token request body in the dialect's encoding.
fn encodeBody(
    arena: std.mem.Allocator,
    format: TokenRequestFormat,
    pairs: []const [2][]const u8,
) !struct { body: []u8, content_type: []const u8 } {
    switch (format) {
        .form => return .{ .body = try formEncode(arena, pairs), .content_type = "application/x-www-form-urlencoded" },
        .json => {
            var aw: std.Io.Writer.Allocating = .init(arena);
            var s: std.json.Stringify = .{ .writer = &aw.writer };
            try s.beginObject();
            for (pairs) |kv| {
                try s.objectField(kv[0]);
                try s.write(kv[1]);
            }
            try s.endObject();
            return .{ .body = aw.written(), .content_type = "application/json" };
        },
    }
}

/// Request a device/user code. `headers` are the client-identity headers
/// (the provider's `extra_headers`) sent on every auth HTTP call. Allocates
/// the result into `arena`.
fn requestDeviceAuth(
    arena: std.mem.Allocator,
    client: *std.http.Client,
    oauth: OAuthDeviceAuth,
    headers: []const Header,
) !DeviceAuth {
    const enc = switch (oauth.dialect) {
        // GitHub: `client_id` (+ optional scope), in the configured encoding.
        .token => try encodeBody(arena, oauth.token_request_format, blk: {
            if (oauth.scope) |sc| {
                break :blk &[_][2][]const u8{ .{ "client_id", oauth.client_id }, .{ "scope", sc } };
            } else break :blk &[_][2][]const u8{.{ "client_id", oauth.client_id }};
        }),
        // Codex: JSON `{client_id}`.
        .codex => try encodeBody(arena, .json, &.{.{ "client_id", oauth.client_id }}),
    };

    const resp = try http.request(arena, client, .POST, oauth.device_code_url, .{
        .headers = headers,
        .body = enc.body,
        .content_type = enc.content_type,
    });
    if (!resp.ok()) return AuthError.AuthFlowFailed;

    const parsed = std.json.parseFromSlice(std.json.Value, arena, resp.body, .{}) catch
        return AuthError.AuthFlowFailed;
    const v = parsed.value;

    const user_code = http.jsonStringAtPath(v, "user_code") orelse
        http.jsonStringAtPath(v, "usercode") orelse return AuthError.AuthFlowFailed;
    const interval: u32 = blk: {
        if (http.jsonIntAtPath(v, "interval")) |iv| break :blk @intCast(@max(iv, 1));
        // Codex returns interval as a string.
        if (http.jsonStringAtPath(v, "interval")) |is| {
            if (std.fmt.parseInt(i64, is, 10) catch null) |iv| break :blk @intCast(@max(iv, 1));
        }
        break :blk 5;
    };
    const verification_uri = switch (oauth.dialect) {
        .token => http.jsonStringAtPath(v, "verification_uri") orelse
            http.jsonStringAtPath(v, "verification_uri_complete") orelse return AuthError.AuthFlowFailed,
        .codex => oauth.verification_url orelse return AuthError.AuthFlowFailed,
    };
    return .{
        .device_code = http.jsonStringAtPath(v, "device_code"),
        .device_auth_id = http.jsonStringAtPath(v, "device_auth_id"),
        .user_code = user_code,
        .verification_uri = verification_uri,
        .interval_secs = interval,
        .expires_in = http.jsonIntAtPath(v, "expires_in") orelse 0,
    };
}

const PollResult = union(enum) {
    pending,
    done: OAuthTokens,
};

/// Poll once for completion. Allocates any returned tokens into `arena`.
fn pollOnce(
    arena: std.mem.Allocator,
    client: *std.http.Client,
    oauth: OAuthDeviceAuth,
    da: DeviceAuth,
    headers: []const Header,
) !PollResult {
    switch (oauth.dialect) {
        .token => {
            const enc = try encodeBody(arena, oauth.token_request_format, &.{
                .{ "client_id", oauth.client_id },
                .{ "device_code", da.device_code orelse return AuthError.AuthFlowFailed },
                .{ "grant_type", "urn:ietf:params:oauth:grant-type:device_code" },
            });
            const resp = try http.request(arena, client, .POST, oauth.token_url, .{
                .headers = headers,
                .body = enc.body,
                .content_type = enc.content_type,
            });
            const parsed = std.json.parseFromSlice(std.json.Value, arena, resp.body, .{}) catch
                return AuthError.AuthFlowFailed;
            const v = parsed.value;
            if (http.jsonStringAtPath(v, "error")) |e| {
                if (std.mem.eql(u8, e, "authorization_pending") or std.mem.eql(u8, e, "slow_down"))
                    return .pending;
                return AuthError.AuthFlowFailed;
            }
            const tokens = parseOAuthTokens(v) orelse return AuthError.AuthFlowFailed;
            return .{ .done = tokens };
        },
        .codex => {
            const poll_url = oauth.device_poll_url orelse return AuthError.AuthFlowFailed;
            const enc = try encodeBody(arena, .json, &.{
                .{ "device_auth_id", da.device_auth_id orelse return AuthError.AuthFlowFailed },
                .{ "user_code", da.user_code },
            });
            const resp = try http.request(arena, client, .POST, poll_url, .{
                .headers = headers,
                .body = enc.body,
                .content_type = enc.content_type,
            });
            // 403/404 => still pending (reference behavior).
            if (resp.status == 403 or resp.status == 404) return .pending;
            if (!resp.ok()) return AuthError.AuthFlowFailed;
            const parsed = std.json.parseFromSlice(std.json.Value, arena, resp.body, .{}) catch
                return AuthError.AuthFlowFailed;
            const code = http.jsonStringAtPath(parsed.value, "authorization_code") orelse return .pending;
            const verifier = http.jsonStringAtPath(parsed.value, "code_verifier") orelse return AuthError.AuthFlowFailed;
            // Exchange the authorization code + server-generated PKCE verifier.
            const tokens = try exchangeCode(arena, client, oauth, code, verifier, headers);
            return .{ .done = tokens };
        },
    }
}

/// Codex authorization-code exchange at the token endpoint.
fn exchangeCode(
    arena: std.mem.Allocator,
    client: *std.http.Client,
    oauth: OAuthDeviceAuth,
    code: []const u8,
    verifier: []const u8,
    headers: []const Header,
) !OAuthTokens {
    const redirect = oauth.redirect_uri orelse "https://auth.openai.com/deviceauth/callback";
    const enc = try encodeBody(arena, .form, &.{
        .{ "grant_type", "authorization_code" },
        .{ "client_id", oauth.client_id },
        .{ "code", code },
        .{ "code_verifier", verifier },
        .{ "redirect_uri", redirect },
    });
    const resp = try http.request(arena, client, .POST, oauth.token_url, .{
        .headers = headers,
        .body = enc.body,
        .content_type = enc.content_type,
    });
    if (!resp.ok()) return AuthError.AuthFlowFailed;
    const parsed = std.json.parseFromSlice(std.json.Value, arena, resp.body, .{}) catch
        return AuthError.AuthFlowFailed;
    return parseOAuthTokens(parsed.value) orelse AuthError.AuthFlowFailed;
}

/// Refresh an access token. Allocates into `arena`.
pub fn refreshTokens(
    arena: std.mem.Allocator,
    client: *std.http.Client,
    oauth: OAuthDeviceAuth,
    refresh_token: []const u8,
    headers: []const Header,
) !OAuthTokens {
    const enc = try encodeBody(arena, .form, &.{
        .{ "grant_type", "refresh_token" },
        .{ "refresh_token", refresh_token },
        .{ "client_id", oauth.client_id },
    });
    const resp = try http.request(arena, client, .POST, oauth.token_url, .{
        .headers = headers,
        .body = enc.body,
        .content_type = enc.content_type,
    });
    if (!resp.ok()) return AuthError.AuthFlowFailed;
    const parsed = std.json.parseFromSlice(std.json.Value, arena, resp.body, .{}) catch
        return AuthError.AuthFlowFailed;
    return parseOAuthTokens(parsed.value) orelse AuthError.AuthFlowFailed;
}

/// Pluck an `OAuthTokens` out of a parsed token response. Null if no
/// `access_token`. Borrows from the parsed value.
fn parseOAuthTokens(v: std.json.Value) ?OAuthTokens {
    const access = http.jsonStringAtPath(v, "access_token") orelse return null;
    return .{
        .access_token = access,
        .refresh_token = http.jsonStringAtPath(v, "refresh_token"),
        .id_token = http.jsonStringAtPath(v, "id_token"),
        .expires_in = http.jsonIntAtPath(v, "expires_in"),
    };
}

/// Run the full interactive device login: request a code, present it, then
/// poll until authorized (or the deadline passes). Allocates the resulting
/// tokens into `arena`.
pub fn login(
    arena: std.mem.Allocator,
    io: Io,
    client: *std.http.Client,
    oauth: OAuthDeviceAuth,
    presenter: Presenter,
    headers: []const Header,
) !OAuthTokens {
    const da = try requestDeviceAuth(arena, client, oauth, headers);
    presenter.deviceCode(.{
        .verification_uri = da.verification_uri,
        .user_code = da.user_code,
        .expires_in = da.expires_in,
    });
    presenter.status("waiting for authorization…");

    const start_ns = std.Io.Clock.now(.real, io).nanoseconds;
    const deadline_ns = start_ns + 15 * std.time.ns_per_min;
    while (true) {
        io.sleep(.fromSeconds(@intCast(da.interval_secs)), .real) catch {};
        const result = try pollOnce(arena, client, oauth, da, headers);
        switch (result) {
            .pending => {},
            .done => |toks| return toks,
        }
        if (std.Io.Clock.now(.real, io).nanoseconds >= deadline_ns) return AuthError.AuthLoginTimeout;
    }
}

// ===========================================================================
// Secondary token exchange (Copilot)
// ===========================================================================

/// Run the configured exchange (Copilot's `/v2/token`) using `bearer` as the
/// request bearer. Returns the exchanged token + optional expiry/base_url,
/// allocated into `arena`.
pub fn runExchange(
    arena: std.mem.Allocator,
    client: *std.http.Client,
    exchange: ExchangeConfig,
    bearer: []const u8,
    headers: []const Header,
) !TokenSet.ExchangeToken {
    const auth_value = try std.fmt.allocPrint(arena, "Bearer {s}", .{bearer});
    var hdrs: std.ArrayList(Header) = .empty;
    try hdrs.append(arena, .{ .name = "authorization", .value = auth_value });
    for (headers) |h| try hdrs.append(arena, h);

    const method: http.Method = if (std.ascii.eqlIgnoreCase(exchange.method, "POST")) .POST else .GET;
    const resp = try http.request(arena, client, method, exchange.url, .{ .headers = hdrs.items });
    if (!resp.ok()) return AuthError.AuthFlowFailed;

    const parsed = std.json.parseFromSlice(std.json.Value, arena, resp.body, .{}) catch
        return AuthError.AuthFlowFailed;
    const v = parsed.value;
    const token = http.jsonStringAtPath(v, exchange.token_json_path) orelse return AuthError.AuthFlowFailed;
    return .{
        .token = token,
        .expires_at = if (exchange.expires_at_json_path) |p| http.jsonIntAtPath(v, p) else null,
        .base_url = if (exchange.base_url_json_path) |p| http.jsonStringAtPath(v, p) else null,
    };
}

// ===========================================================================
// JWT + credential building (pure)
// ===========================================================================

/// Decode a JWT's payload (the middle segment) into a parsed JSON value.
/// Returns null if the token is malformed. Caller owns the result.
pub fn decodeJwtPayload(allocator: std.mem.Allocator, token: []const u8) ?std.json.Parsed(std.json.Value) {
    var it = std.mem.splitScalar(u8, token, '.');
    _ = it.next() orelse return null; // header
    const payload_b64 = it.next() orelse return null;
    if (it.next() == null) return null; // require a signature segment

    const dec = std.base64.url_safe_no_pad.Decoder;
    const n = dec.calcSizeForSlice(payload_b64) catch return null;
    const buf = allocator.alloc(u8, n) catch return null;
    defer allocator.free(buf);
    dec.decode(buf, payload_b64) catch return null;
    return std.json.parseFromSlice(std.json.Value, allocator, buf, .{}) catch null;
}

/// The `exp` (unix-seconds expiry) claim of a JWT access/id token, or null.
pub fn parseJwtExp(allocator: std.mem.Allocator, token: []const u8) ?i64 {
    var parsed = decodeJwtPayload(allocator, token) orelse return null;
    defer parsed.deinit();
    return http.jsonIntAtPath(parsed.value, "exp");
}

/// Extract the ChatGPT account id from an id_token. `claim_path` is the JWT
/// claim holding the auth object (e.g. `https://api.openai.com/auth`); the
/// account id is its `chatgpt_account_id` field. Returns an owned copy.
pub fn extractAccountId(
    allocator: std.mem.Allocator,
    id_token: []const u8,
    claim_path: []const u8,
) ?[]u8 {
    var parsed = decodeJwtPayload(allocator, id_token) orelse return null;
    defer parsed.deinit();
    // `claim_path` is a single literal claim key (e.g.
    // `https://api.openai.com/auth`) that itself contains dots, so it is
    // looked up directly rather than walked as a dotted path.
    const claim = switch (parsed.value) {
        .object => |o| o.get(claim_path) orelse return null,
        else => return null,
    };
    const account = switch (claim) {
        .object => |o| o.get("chatgpt_account_id") orelse return null,
        else => return null,
    };
    return switch (account) {
        .string => |s| allocator.dupe(u8, s) catch null,
        else => null,
    };
}

/// True when an access token is missing or within `margin` seconds of expiry.
pub fn needsRefresh(ts: TokenSet, now_unix: i64, margin: i64) bool {
    if (ts.access_token == null) return false; // nothing to refresh (login path)
    const exp = ts.expires_at orelse return false; // no intrinsic expiry (e.g. ghu_)
    return exp - now_unix <= margin;
}

/// True when an exchange is configured but the stored exchange token is
/// missing or within `margin` seconds of expiry.
pub fn needsExchange(oauth: OAuthDeviceAuth, ts: TokenSet, now_unix: i64, margin: i64) bool {
    if (oauth.exchange == null) return false;
    const ex = ts.exchange orelse return true;
    const exp = ex.expires_at orelse return false; // no expiry => assume durable
    return exp - now_unix <= margin;
}

/// Build the request-ready credential from a (refreshed/exchanged) token set.
/// The exchanged token wins over the access token; a codex account id becomes
/// a `chatgpt-account-id` header. All strings are duped into `arena`.
pub fn buildCredential(
    arena: std.mem.Allocator,
    oauth: OAuthDeviceAuth,
    ts: TokenSet,
) !ResolvedCredential {
    var base_url_override: ?[]const u8 = null;
    const secret: []const u8 = blk: {
        if (ts.exchange) |ex| {
            if (ex.base_url) |bu| base_url_override = try arena.dupe(u8, bu);
            break :blk ex.token;
        }
        break :blk ts.access_token orelse return AuthError.AuthUnresolved;
    };

    var headers: std.ArrayList(Header) = .empty;
    if (oauth.account_id_jwt_claim != null) {
        if (ts.account_id) |aid| {
            try headers.append(arena, .{
                .name = "chatgpt-account-id",
                .value = try arena.dupe(u8, aid),
            });
        }
    }

    return .{
        .api_key = try arena.dupe(u8, secret),
        .base_url_override = base_url_override,
        .extra_headers = try headers.toOwnedSlice(arena),
    };
}

/// Assemble a persisted `TokenSet` from a fresh OAuth token response.
/// `expires_at` is derived from `expires_in` (preferred) or the access
/// token's JWT `exp`. A codex account id is extracted when configured. All
/// strings are duped into `arena`.
pub fn tokensToTokenSet(
    arena: std.mem.Allocator,
    oauth: OAuthDeviceAuth,
    tokens: OAuthTokens,
    now_unix: i64,
) !TokenSet {
    const access = try arena.dupe(u8, tokens.access_token);
    const expires_at: ?i64 = blk: {
        if (tokens.expires_in) |e| break :blk now_unix + e;
        break :blk parseJwtExp(arena, access);
    };
    var account_id: ?[]const u8 = null;
    if (oauth.account_id_jwt_claim) |claim| {
        if (tokens.id_token) |idt| {
            account_id = extractAccountId(arena, idt, claim);
        }
    }
    return .{
        .type = "oauth_device",
        .access_token = access,
        .refresh_token = if (tokens.refresh_token) |r| try arena.dupe(u8, r) else null,
        .id_token = if (tokens.id_token) |i| try arena.dupe(u8, i) else null,
        .expires_at = expires_at,
        .account_id = account_id,
    };
}

const t = std.testing;

test "token storage: save, load, delete round-trip" {
    const io = t.io;
    var tmp = t.tmpDir(.{});
    defer tmp.cleanup();
    var path_buf: [std.fs.max_path_bytes]u8 = undefined;
    const n = try tmp.dir.realPath(io, &path_buf);
    const auth_dir = try std.fmt.allocPrint(t.allocator, "{s}/auth", .{path_buf[0..n]});
    defer t.allocator.free(auth_dir);

    // Absent => null.
    try t.expect((try loadTokenSet(t.allocator, io, auth_dir, "ghost")) == null);

    const ts: TokenSet = .{
        .type = "oauth_device",
        .access_token = "ghu_abc",
        .refresh_token = "rt_xyz",
        .expires_at = 1700000000,
        .exchange = .{ .token = "tkn", .expires_at = 1700001800, .base_url = "https://api.x" },
    };
    try saveTokenSet(t.allocator, io, auth_dir, "github_copilot", ts);

    var loaded = (try loadTokenSet(t.allocator, io, auth_dir, "github_copilot")).?;
    defer loaded.deinit();
    try t.expectEqualStrings("ghu_abc", loaded.value.access_token.?);
    try t.expectEqualStrings("rt_xyz", loaded.value.refresh_token.?);
    try t.expectEqual(@as(?i64, 1700000000), loaded.value.expires_at);
    try t.expect(loaded.value.id_token == null);
    try t.expectEqualStrings("tkn", loaded.value.exchange.?.token);
    try t.expectEqualStrings("https://api.x", loaded.value.exchange.?.base_url.?);

    try t.expect(try deleteTokenSet(t.allocator, io, auth_dir, "github_copilot"));
    try t.expect(!try deleteTokenSet(t.allocator, io, auth_dir, "github_copilot"));
    try t.expect((try loadTokenSet(t.allocator, io, auth_dir, "github_copilot")) == null);
}

test "AuthConfig: api_key variant" {
    const a: AuthConfig = .{ .api_key = .{ .key = "sk-resolved" } };
    try t.expectEqual(AuthType.api_key, @as(AuthType, a));
    try t.expectEqualStrings("sk-resolved", a.api_key.key.?);
}

test "AuthConfig: oauth_device defaults" {
    const a: AuthConfig = .{ .oauth_device = .{
        .client_id = "Iv1.x",
        .device_code_url = "https://github.com/login/device/code",
        .token_url = "https://github.com/login/oauth/access_token",
    } };
    try t.expectEqual(DeviceDialect.token, a.oauth_device.dialect);
    try t.expectEqual(TokenRequestFormat.form, a.oauth_device.token_request_format);
    try t.expect(a.oauth_device.exchange == null);
}

test "TokenSet: defaults" {
    const ts: TokenSet = .{};
    try t.expectEqualStrings("oauth_device", ts.type);
    try t.expect(ts.access_token == null);
    try t.expect(ts.exchange == null);
}

/// Build a `header.payload.sig` JWT whose payload is `payload_json`,
/// base64url-no-pad encoded. Allocated in `arena`.
fn makeJwt(arena: std.mem.Allocator, payload_json: []const u8) ![]u8 {
    const enc = std.base64.url_safe_no_pad.Encoder;
    const header = "{\"alg\":\"none\"}";
    var hbuf: [64]u8 = undefined;
    const h = enc.encode(&hbuf, header);
    const pbuf = try arena.alloc(u8, enc.calcSize(payload_json.len));
    const p = enc.encode(pbuf, payload_json);
    return std.fmt.allocPrint(arena, "{s}.{s}.sig", .{ h, p });
}

test "parseJwtExp + extractAccountId" {
    var aa = std.heap.ArenaAllocator.init(t.allocator);
    defer aa.deinit();
    const a = aa.allocator();
    const jwt = try makeJwt(a,
        \\{"exp":1700000000,"https://api.openai.com/auth":{"chatgpt_account_id":"acct_123"}}
    );
    try t.expectEqual(@as(?i64, 1700000000), parseJwtExp(a, jwt));
    const aid = extractAccountId(a, jwt, "https://api.openai.com/auth").?;
    try t.expectEqualStrings("acct_123", aid);
    // Wrong claim path => null.
    try t.expect(extractAccountId(a, jwt, "nope") == null);
    // Malformed token => null.
    try t.expect(parseJwtExp(a, "not-a-jwt") == null);
}

test "parseOAuthTokens: required access_token, optional rest" {
    var parsed = try std.json.parseFromSlice(std.json.Value, t.allocator,
        \\{"access_token":"at","refresh_token":"rt","id_token":"it","expires_in":1800}
    , .{});
    defer parsed.deinit();
    const toks = parseOAuthTokens(parsed.value).?;
    try t.expectEqualStrings("at", toks.access_token);
    try t.expectEqualStrings("rt", toks.refresh_token.?);
    try t.expectEqual(@as(?i64, 1800), toks.expires_in);

    var p2 = try std.json.parseFromSlice(std.json.Value, t.allocator, "{\"error\":\"x\"}", .{});
    defer p2.deinit();
    try t.expect(parseOAuthTokens(p2.value) == null);
}

test "formEncode: percent-encodes reserved characters" {
    var aa = std.heap.ArenaAllocator.init(t.allocator);
    defer aa.deinit();
    const body = try formEncode(aa.allocator(), &.{
        .{ "grant_type", "urn:ietf:params:oauth:grant-type:device_code" },
        .{ "client_id", "Iv1.abc" },
    });
    try t.expectEqualStrings(
        "grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code&client_id=Iv1.abc",
        body,
    );
}

test "needsRefresh / needsExchange predicates" {
    const now: i64 = 1000;
    // Durable token (no expiry) never needs refresh.
    try t.expect(!needsRefresh(.{ .access_token = "x" }, now, 120));
    // Within margin => refresh.
    try t.expect(needsRefresh(.{ .access_token = "x", .expires_at = 1100 }, now, 120));
    // Comfortably fresh => no.
    try t.expect(!needsRefresh(.{ .access_token = "x", .expires_at = 5000 }, now, 120));

    const oauth: OAuthDeviceAuth = .{
        .client_id = "c",
        .device_code_url = "u",
        .token_url = "u",
        .exchange = .{ .url = "https://x" },
    };
    // Exchange configured but none stored => needs exchange.
    try t.expect(needsExchange(oauth, .{ .access_token = "x" }, now, 120));
    // Stored exchange near expiry => needs exchange.
    try t.expect(needsExchange(oauth, .{ .exchange = .{ .token = "t", .expires_at = 1050 } }, now, 120));
    // Fresh exchange => no.
    try t.expect(!needsExchange(oauth, .{ .exchange = .{ .token = "t", .expires_at = 9000 } }, now, 120));
    // No exchange configured => never.
    const no_ex: OAuthDeviceAuth = .{ .client_id = "c", .device_code_url = "u", .token_url = "u" };
    try t.expect(!needsExchange(no_ex, .{ .access_token = "x" }, now, 120));
}

test "buildCredential: copilot exchange token + dynamic base_url" {
    var aa = std.heap.ArenaAllocator.init(t.allocator);
    defer aa.deinit();
    const oauth: OAuthDeviceAuth = .{
        .client_id = "c",
        .device_code_url = "u",
        .token_url = "u",
        .exchange = .{ .url = "https://x" },
    };
    const ts: TokenSet = .{
        .access_token = "ghu_durable",
        .exchange = .{ .token = "tid", .base_url = "https://api.individual.githubcopilot.com" },
    };
    const cred = try buildCredential(aa.allocator(), oauth, ts);
    try t.expectEqualStrings("tid", cred.api_key);
    try t.expectEqualStrings("https://api.individual.githubcopilot.com", cred.base_url_override.?);
    try t.expectEqual(@as(usize, 0), cred.extra_headers.len);
}

test "buildCredential: codex access token + account-id header" {
    var aa = std.heap.ArenaAllocator.init(t.allocator);
    defer aa.deinit();
    const oauth: OAuthDeviceAuth = .{
        .dialect = .codex,
        .client_id = "c",
        .device_code_url = "u",
        .token_url = "u",
        .device_poll_url = "p",
        .account_id_jwt_claim = "https://api.openai.com/auth",
    };
    const ts: TokenSet = .{ .access_token = "jwt-access", .account_id = "acct_9" };
    const cred = try buildCredential(aa.allocator(), oauth, ts);
    try t.expectEqualStrings("jwt-access", cred.api_key);
    try t.expect(cred.base_url_override == null);
    try t.expectEqual(@as(usize, 1), cred.extra_headers.len);
    try t.expectEqualStrings("chatgpt-account-id", cred.extra_headers[0].name);
    try t.expectEqualStrings("acct_9", cred.extra_headers[0].value);
}

test "tokensToTokenSet: derives expires_at from expires_in and extracts account_id" {
    var aa = std.heap.ArenaAllocator.init(t.allocator);
    defer aa.deinit();
    const a = aa.allocator();
    const idt = try makeJwt(a,
        \\{"https://api.openai.com/auth":{"chatgpt_account_id":"acct_77"}}
    );
    const oauth: OAuthDeviceAuth = .{
        .dialect = .codex,
        .client_id = "c",
        .device_code_url = "u",
        .token_url = "u",
        .device_poll_url = "p",
        .account_id_jwt_claim = "https://api.openai.com/auth",
    };
    const ts = try tokensToTokenSet(a, oauth, .{
        .access_token = "at",
        .refresh_token = "rt",
        .id_token = idt,
        .expires_in = 1800,
    }, 1000);
    try t.expectEqual(@as(?i64, 2800), ts.expires_at);
    try t.expectEqualStrings("rt", ts.refresh_token.?);
    try t.expectEqualStrings("acct_77", ts.account_id.?);
}