agenticSystem/mcp-server/auth/apiClient.js

import axios from "axios";
import { McpError, ErrorCode } from "@modelcontextprotocol/sdk/types.js";
import { sessionApiClients, getSessionCredentials, setCredentials, findRoleByToken } from "./credentials.js";
import { assertSafeCmsTarget } from "../utils/cmsTargetSafety.js";
import { ensureFreshSessionCredentials, refreshSessionCredentials } from "./sessionRefresh.js";
const DEFAULT_ROLE = 'developer';

export const isAcaiTokenFailure = (error) => {
    const status = error?.response?.status;
    const responseText = JSON.stringify(error?.response?.data || "");
    return status === 401
        || (status === 403 && /token|jwt|auth|unauthoriz|expired|no v[aá]lido/i.test(responseText));
};

export const runWithTokenRefreshRetry = async (
    sessionId,
    operation,
    { refreshCredentials = refreshSessionCredentials } = {},
) => {
    try {
        return await operation();
    } catch (error) {
        if (!isAcaiTokenFailure(error) || error?.config?._acaiTokenRetry) throw error;
        await refreshCredentials(sessionId, { force: true });
        return operation();
    }
};

/**
 * Check if session is configured with valid credentials
 */
export const ensureConfigured = async (sessionId) => {
    const creds = await getSessionCredentials(sessionId);
    if (!creds.token || !creds.web_url || !creds.api_web_url) {
        throw new McpError(
            ErrorCode.InvalidRequest,
            "Acai site not configured for safe local execution. Use the project MCP config/select_project flow so ACAI_API_WEB_URL points to the local environment."
        );
    }

    try {
        assertSafeCmsTarget(creds, "apiClient");
    } catch (error) {
        throw new McpError(
            ErrorCode.InvalidRequest,
            error.message
        );
    }
};

/**
 * Rebuild API client for a session
 */
export const rebuildApiClient = async (sessionId, { refreshCredentials = refreshSessionCredentials } = {}) => {
    const creds = await getSessionCredentials(sessionId);
    if (!creds.token || !creds.web_url || !creds.api_web_url) {
        return null;
    }
    assertSafeCmsTarget(creds, "apiClient");
    const client = axios.create({
        baseURL: creds.api_web_url,
        headers: {
            "Content-Type": "application/json",
            "X-Acai-Token": creds.token,
            ...(creds.forge_host ? { Host: creds.forge_host } : {}),
        },
    });

    // Always resolve the latest session token. The client outlives individual
    // JWTs, so capturing `creds` here would keep sending a rotated token.
    client.interceptors.request.use(async (config) => {
        const latest = await getSessionCredentials(sessionId);
        if (latest.token) {
            config.headers["X-Acai-Token"] = latest.token;
        }
        return config;
    });

    client.interceptors.response.use(undefined, async (error) => {
        const original = error.config;
        if (!isAcaiTokenFailure(error) || !original || original._acaiTokenRetry) {
            throw error;
        }

        original._acaiTokenRetry = true;
        const fresh = await refreshCredentials(sessionId, { force: true });
        original.headers = original.headers || {};
        original.headers["X-Acai-Token"] = fresh.token;
        return client.request(original);
    });

    sessionApiClients.set(sessionId, client);
    return client;
};

/**
 * Get or create API client for a session
 * @param {string} sessionId - The session ID
 */
export const getApiClient = async (sessionId) => {
    const creds = await getSessionCredentials(sessionId);

    console.error(`[API Client] getApiClient called for session ${sessionId}`);
    console.error(`[API Client] Current creds: token=${!!creds.token}, web_url=${creds.web_url}`);

    await ensureConfigured(sessionId);
    let client = sessionApiClients.get(sessionId);
    if (!client) {
        console.error(`[API Client] No cached client, rebuilding...`);
        client = await rebuildApiClient(sessionId);
    }
    if (!client) {
        throw new McpError(
            ErrorCode.InvalidRequest,
            "Unable to create API client. Verify credentials and try again."
        );
    }
    console.error(`[API Client] Returning client for session ${sessionId}`);
    return client;
};

/**
 * Set credentials and rebuild API client
 * @param {Object} credentials - The credentials object
 * @param {string} sessionId - The session ID
 * @param {string} mcpSessionId - Optional MCP-Session-Id for persistence across SSE reconnections
 */
export const setCredentialsAndRebuild = async (credentials, sessionId, mcpSessionId = null) => {
    await setCredentials(credentials, sessionId, mcpSessionId);
    await rebuildApiClient(sessionId);
};

/**
 * Wrapper for authenticated handlers
 * Supports both session-based auth and inline credentials (stateless mode).
 *
 * If args contains acaiToken + acaiWebsite, these are used directly,
 * allowing Claude to send credentials with each request.
 */
export const withAuth = (handler) => {
    return async (args, extra) => {
        const sessionId = extra?.sessionId || "_default";

        console.error(`[withAuth] Called with sessionId: ${sessionId}`);

        // Check for inline credentials (stateless mode)
        const inlineCredentials = {
            acaiToken: args.acaiToken,
            acaiWebsite: args.acaiWebsite,
            acaiTokenHash: args.acaiTokenHash
        };

        const hasInlineCredentials = inlineCredentials.acaiToken && inlineCredentials.acaiWebsite;

        if (hasInlineCredentials) {
            // Lookup role by token before storing credentials
            const role = findRoleByToken(inlineCredentials.acaiToken) || DEFAULT_ROLE;
            console.error(`[withAuth] Using INLINE credentials: website=${inlineCredentials.acaiWebsite}, role=${role}`);

            // Temporarily store inline credentials in session for this request
            await setCredentials({
                token: inlineCredentials.acaiToken,
                website: inlineCredentials.acaiWebsite,
                web_url: `https://${inlineCredentials.acaiWebsite}`,
                api_web_url: null,
                forge_host: null,
                tokenHash: inlineCredentials.acaiTokenHash || null,
                profileName: 'inline',
                role: role
            }, sessionId);
        }

        await ensureFreshSessionCredentials(sessionId);
        console.error(`[withAuth] Getting API client for session ${sessionId}...`);
        await getApiClient(sessionId);
        console.error(`[withAuth] API client ready, calling handler...`);

        const handlerExtra = {
            ...extra,
            sessionId,
            inlineCredentials: hasInlineCredentials ? inlineCredentials : null,
        };
        return runWithTokenRefreshRetry(
            sessionId,
            () => handler(args, handlerExtra),
        );
    };
};