responses-proxy 0.1.0

package/dist/health-websocket-manager.js

@@ -0,0 +1,174 @@
+import fastifyWebsocket from "@fastify/websocket";
+export class HealthWebSocketManager {
+    healthService;
+    broadcastInterval;
+    connections = new Set();
+    updateInterval = null;
+    initialized = false;
+    constructor(healthService, broadcastInterval = 5000 // 5 seconds
+    ) {
+        this.healthService = healthService;
+        this.broadcastInterval = broadcastInterval;
+    }
+    /**
+     * Initialize WebSocket support on Fastify server
+     */
+    async initialize(app) {
+        if (this.initialized) {
+            return;
+        }
+        this.initialized = true;
+        // Register WebSocket plugin
+        await app.register(fastifyWebsocket);
+        // Health updates WebSocket endpoint
+        await app.register(async (fastify) => {
+            fastify.get('/ws/health', { websocket: true }, (connection, req) => {
+                console.log('New health WebSocket connection established');
+                // Add connection to active connections
+                this.connections.add(connection);
+                // Send initial health summary
+                const summary = this.healthService.getHealthSummary();
+                const summaryMessage = {
+                    type: 'health_summary',
+                    summary,
+                    timestamp: Date.now()
+                };
+                connection.send(JSON.stringify(summaryMessage));
+                // Send current health data for all providers
+                const allHealth = this.healthService.getAllProviderHealth();
+                for (const [providerId, metrics] of allHealth) {
+                    const updateMessage = {
+                        type: 'health_update',
+                        providerId,
+                        metrics,
+                        timestamp: Date.now()
+                    };
+                    connection.send(JSON.stringify(updateMessage));
+                }
+                // Handle connection close
+                connection.on('close', () => {
+                    console.log('Health WebSocket connection closed');
+                    this.connections.delete(connection);
+                });
+                // Handle connection error
+                connection.on('error', (error) => {
+                    console.error('Health WebSocket error:', error);
+                    this.connections.delete(connection);
+                });
+            });
+        });
+        // Start broadcasting health updates
+        this.startBroadcasting();
+    }
+    /**
+     * Start broadcasting health updates to all connected clients
+     */
+    startBroadcasting() {
+        if (this.updateInterval) {
+            clearInterval(this.updateInterval);
+        }
+        this.updateInterval = setInterval(() => {
+            this.broadcastHealthUpdates();
+        }, this.broadcastInterval);
+    }
+    /**
+     * Stop broadcasting health updates
+     */
+    stopBroadcasting() {
+        if (this.updateInterval) {
+            clearInterval(this.updateInterval);
+            this.updateInterval = null;
+        }
+    }
+    /**
+     * Broadcast health updates to all connected clients
+     */
+    broadcastHealthUpdates() {
+        if (this.connections.size === 0) {
+            return;
+        }
+        try {
+            // Get current health data
+            const allHealth = this.healthService.getAllProviderHealth();
+            const summary = this.healthService.getHealthSummary();
+            // Broadcast summary
+            const summaryMessage = {
+                type: 'health_summary',
+                summary,
+                timestamp: Date.now()
+            };
+            this.broadcast(summaryMessage);
+            // Broadcast individual provider updates
+            for (const [providerId, metrics] of allHealth) {
+                const updateMessage = {
+                    type: 'health_update',
+                    providerId,
+                    metrics,
+                    timestamp: Date.now()
+                };
+                this.broadcast(updateMessage);
+            }
+        }
+        catch (error) {
+            console.error('Error broadcasting health updates:', error);
+        }
+    }
+    /**
+     * Broadcast a message to all connected clients
+     */
+    broadcast(message) {
+        const messageStr = JSON.stringify(message);
+        const deadConnections = [];
+        for (const connection of this.connections) {
+            try {
+                if (connection.readyState === 1) { // WebSocket.OPEN
+                    connection.send(messageStr);
+                }
+                else {
+                    deadConnections.push(connection);
+                }
+            }
+            catch (error) {
+                console.error('Error sending WebSocket message:', error);
+                deadConnections.push(connection);
+            }
+        }
+        // Clean up dead connections
+        for (const deadConnection of deadConnections) {
+            this.connections.delete(deadConnection);
+        }
+    }
+    /**
+     * Broadcast immediate health update for a specific provider
+     */
+    broadcastProviderUpdate(providerId, metrics) {
+        const message = {
+            type: 'health_update',
+            providerId,
+            metrics,
+            timestamp: Date.now()
+        };
+        this.broadcast(message);
+    }
+    /**
+     * Get current connection count
+     */
+    getConnectionCount() {
+        return this.connections.size;
+    }
+    /**
+     * Close all connections and stop broadcasting
+     */
+    shutdown() {
+        this.stopBroadcasting();
+        for (const connection of this.connections) {
+            try {
+                connection.close();
+            }
+            catch (error) {
+                console.error('Error closing WebSocket connection:', error);
+            }
+        }
+        this.connections.clear();
+    }
+}

package/dist/http-rate-limit.js

@@ -0,0 +1,36 @@
+export class InMemoryHttpRateLimiter {
+    buckets = new Map();
+    consume(key, options) {
+        const nowMs = options.nowMs ?? Date.now();
+        this.prune(nowMs, options.windowMs);
+        const existing = this.buckets.get(key);
+        if (!existing || nowMs - existing.windowStartedAt >= options.windowMs) {
+            this.buckets.set(key, { windowStartedAt: nowMs, hits: 1 });
+            return {
+                allowed: true,
+                remaining: Math.max(options.maxRequests - 1, 0),
+                retryAfterMs: 0,
+            };
+        }
+        if (existing.hits >= options.maxRequests) {
+            return {
+                allowed: false,
+                remaining: 0,
+                retryAfterMs: Math.max(options.windowMs - (nowMs - existing.windowStartedAt), 0),
+            };
+        }
+        existing.hits += 1;
+        return {
+            allowed: true,
+            remaining: Math.max(options.maxRequests - existing.hits, 0),
+            retryAfterMs: 0,
+        };
+    }
+    prune(nowMs, windowMs) {
+        for (const [key, bucket] of this.buckets) {
+            if (nowMs - bucket.windowStartedAt >= windowMs * 2) {
+                this.buckets.delete(key);
+            }
+        }
+    }
+}

package/dist/http-rate-limit.test.js

@@ -0,0 +1,62 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { InMemoryHttpRateLimiter } from "./http-rate-limit.js";
+test("rate limiter allows requests until bucket limit is reached", () => {
+    const limiter = new InMemoryHttpRateLimiter();
+    const first = limiter.consume("responses:token-a", {
+        windowMs: 1_000,
+        maxRequests: 2,
+        nowMs: 10_000,
+    });
+    assert.equal(first.allowed, true);
+    assert.equal(first.remaining, 1);
+    assert.equal(first.retryAfterMs, 0);
+    const second = limiter.consume("responses:token-a", {
+        windowMs: 1_000,
+        maxRequests: 2,
+        nowMs: 10_100,
+    });
+    assert.equal(second.allowed, true);
+    assert.equal(second.remaining, 0);
+    const third = limiter.consume("responses:token-a", {
+        windowMs: 1_000,
+        maxRequests: 2,
+        nowMs: 10_200,
+    });
+    assert.equal(third.allowed, false);
+    assert.equal(third.remaining, 0);
+    assert.equal(third.retryAfterMs, 800);
+});
+test("rate limiter resets bucket after window expires", () => {
+    const limiter = new InMemoryHttpRateLimiter();
+    assert.equal(limiter.consume("health:ip-a", {
+        windowMs: 1_000,
+        maxRequests: 1,
+        nowMs: 5_000,
+    }).allowed, true);
+    assert.equal(limiter.consume("health:ip-a", {
+        windowMs: 1_000,
+        maxRequests: 1,
+        nowMs: 5_500,
+    }).allowed, false);
+    const reset = limiter.consume("health:ip-a", {
+        windowMs: 1_000,
+        maxRequests: 1,
+        nowMs: 6_000,
+    });
+    assert.equal(reset.allowed, true);
+    assert.equal(reset.remaining, 0);
+});
+test("rate limiter isolates keys", () => {
+    const limiter = new InMemoryHttpRateLimiter();
+    assert.equal(limiter.consume("responses:token-a", {
+        windowMs: 1_000,
+        maxRequests: 1,
+        nowMs: 1,
+    }).allowed, true);
+    assert.equal(limiter.consume("responses:token-b", {
+        windowMs: 1_000,
+        maxRequests: 1,
+        nowMs: 2,
+    }).allowed, true);
+});

package/dist/kiro-auth.js

@@ -0,0 +1,136 @@
+export class KiroAuthError extends Error {
+    statusCode = 409;
+    body;
+    constructor(code, message) {
+        super(message);
+        this.body = { type: "authentication_error", code, message };
+    }
+}
+const refreshLocks = new Map();
+const accountCursors = new Map();
+/**
+ * Refreshes a Kiro/CodeWhisperer token using the AWS SSO-OIDC endpoint for IDC
+ * accounts (clientId + clientSecret present), or the Kiro social refresh
+ * endpoint otherwise. Mirrors the flow used by the 9router desktop app.
+ */
+export async function refreshKiroToken(account, options = { defaultRegion: "us-east-1" }) {
+    const fetchImpl = options.fetchImpl ?? fetch;
+    const { clientId, clientSecret, region } = account.providerSpecificData;
+    const resolvedRegion = region?.trim() || options.defaultRegion;
+    let payload;
+    if (clientId && clientSecret) {
+        const url = `https://oidc.${resolvedRegion}.amazonaws.com/token`;
+        const response = await fetchImpl(url, {
+            method: "POST",
+            headers: { "Content-Type": "application/json" },
+            body: JSON.stringify({
+                clientId,
+                clientSecret,
+                refreshToken: account.refreshToken,
+                grantType: "refresh_token",
+            }),
+        });
+        if (!response.ok) {
+            const text = await response.text().catch(() => "");
+            throw new KiroAuthError("KIRO_TOKEN_REFRESH_FAILED", `Token refresh failed (${response.status}): ${text}`);
+        }
+        payload = (await response.json());
+    }
+    else {
+        const response = await fetchImpl("https://prod.us-east-1.auth.desktop.kiro.dev/refreshToken", {
+            method: "POST",
+            headers: { "Content-Type": "application/json" },
+            body: JSON.stringify({ refreshToken: account.refreshToken }),
+        });
+        if (!response.ok) {
+            const text = await response.text().catch(() => "");
+            throw new KiroAuthError("KIRO_TOKEN_REFRESH_FAILED", `Token refresh failed (${response.status}): ${text}`);
+        }
+        payload = (await response.json());
+    }
+    if (!payload.accessToken) {
+        throw new KiroAuthError("KIRO_TOKEN_REFRESH_FAILED", "Token refresh did not return an accessToken");
+    }
+    const expiresIn = typeof payload.expiresIn === "number" && payload.expiresIn > 0 ? payload.expiresIn : 3600;
+    return {
+        accessToken: payload.accessToken,
+        refreshToken: payload.refreshToken || account.refreshToken,
+        expiresAt: new Date(Date.now() + expiresIn * 1000).toISOString(),
+        expiresIn,
+    };
+}
+/**
+ * Selects an account (optionally pinned), refreshes its token if it is within
+ * the refresh lead window, persists the new token back to 9router's DB, and
+ * returns the credentials to use for a CodeWhisperer request.
+ */
+export async function resolveKiroCredentials(args) {
+    const account = args.accountId?.trim()
+        ? args.store.getAccount(args.accountId.trim())
+        : selectNextAccount(args.poolKey ?? "kiro", args.store, args.rotationMode ?? "round_robin");
+    if (!account || !account.isActive) {
+        throw new KiroAuthError("KIRO_ACCOUNT_UNAVAILABLE", "No connected Kiro accounts are available. Sign in through 9router first.");
+    }
+    const fresh = await ensureFreshAccount(account, args);
+    return {
+        accountId: fresh.id,
+        accessToken: fresh.accessToken,
+        profileArn: fresh.providerSpecificData.profileArn ?? null,
+        region: fresh.providerSpecificData.region?.trim() || args.defaultRegion,
+    };
+}
+async function ensureFreshAccount(account, args) {
+    const now = args.now ?? new Date();
+    const leadMs = args.refreshLeadSeconds * 1000;
+    const expiresAtMs = account.expiresAt ? Date.parse(account.expiresAt) : NaN;
+    const stillValid = Number.isFinite(expiresAtMs) && expiresAtMs - now.getTime() > leadMs;
+    if (stillValid && account.accessToken) {
+        return account;
+    }
+    const existingLock = refreshLocks.get(account.id);
+    if (existingLock) {
+        return existingLock;
+    }
+    const refreshPromise = (async () => {
+        const update = await refreshKiroToken(account, {
+            defaultRegion: args.defaultRegion,
+            fetchImpl: args.fetchImpl,
+        });
+        const persisted = args.store.updateTokens(account.id, update, now);
+        // Always overlay the fresh token onto whatever account we have. When write-back
+        // is disabled, `persisted` is the stale DB row, so without this overlay the
+        // refreshed token would be discarded and an expired token returned.
+        const base = persisted ?? account;
+        return {
+            ...base,
+            accessToken: update.accessToken,
+            refreshToken: update.refreshToken ?? base.refreshToken,
+            expiresAt: update.expiresAt,
+            expiresIn: update.expiresIn ?? base.expiresIn,
+        };
+    })().finally(() => {
+        refreshLocks.delete(account.id);
+    });
+    refreshLocks.set(account.id, refreshPromise);
+    return refreshPromise;
+}
+function selectNextAccount(poolKey, store, rotationMode) {
+    const accounts = store.listAvailableAccounts();
+    if (!accounts.length) {
+        return undefined;
+    }
+    if (rotationMode === "first_available") {
+        return accounts[0];
+    }
+    if (rotationMode === "random") {
+        return accounts[Math.floor(Math.random() * accounts.length)];
+    }
+    const cursor = accountCursors.get(poolKey) ?? 0;
+    const account = accounts[cursor % accounts.length];
+    accountCursors.set(poolKey, (cursor + 1) % accounts.length);
+    return account;
+}
+export function __resetKiroAuthStateForTests() {
+    refreshLocks.clear();
+    accountCursors.clear();
+}

package/dist/kiro-auth.test.js

@@ -0,0 +1,234 @@
+import assert from "node:assert/strict";
+import { beforeEach, test } from "node:test";
+import { KiroAuthError, __resetKiroAuthStateForTests, refreshKiroToken, resolveKiroCredentials, } from "./kiro-auth.js";
+function makeAccount(overrides = {}) {
+    return {
+        id: "acct-a",
+        name: "Account A",
+        priority: 1,
+        isActive: true,
+        accessToken: "access-a",
+        refreshToken: "refresh-a",
+        expiresAt: "2026-06-01T00:00:00.000Z",
+        expiresIn: 3600,
+        providerSpecificData: {
+            profileArn: "arn:keep",
+            clientId: null,
+            clientSecret: null,
+            region: "us-east-1",
+            authMethod: null,
+            startUrl: null,
+        },
+        raw: {},
+        createdAt: "2026-05-01T00:00:00.000Z",
+        updatedAt: "2026-05-01T00:00:00.000Z",
+        ...overrides,
+    };
+}
+class FakeStore {
+    accounts = new Map();
+    updates = [];
+    constructor(accounts) {
+        for (const account of accounts) {
+            this.accounts.set(account.id, account);
+        }
+    }
+    getAccount(id) {
+        return this.accounts.get(id);
+    }
+    listAvailableAccounts() {
+        return [...this.accounts.values()].filter((a) => a.isActive);
+    }
+    updateTokens(id, update, now = new Date()) {
+        this.updates.push({ id, update });
+        const current = this.accounts.get(id);
+        if (!current) {
+            return undefined;
+        }
+        const next = {
+            ...current,
+            accessToken: update.accessToken,
+            refreshToken: update.refreshToken ?? current.refreshToken,
+            expiresAt: update.expiresAt,
+            expiresIn: update.expiresIn ?? current.expiresIn,
+            updatedAt: now.toISOString(),
+        };
+        this.accounts.set(id, next);
+        return next;
+    }
+    asStore() {
+        return this;
+    }
+}
+function mockFetch(handler) {
+    const calls = [];
+    const fetchImpl = (async (input, init) => {
+        const url = typeof input === "string" ? input : input.toString();
+        calls.push({ url, init });
+        const { status = 200, body } = handler(url, init);
+        return new Response(JSON.stringify(body), {
+            status,
+            headers: { "Content-Type": "application/json" },
+        });
+    });
+    return { fetchImpl, calls };
+}
+beforeEach(() => {
+    __resetKiroAuthStateForTests();
+});
+test("refreshKiroToken uses the social endpoint when no IDC client credentials", async () => {
+    const { fetchImpl, calls } = mockFetch(() => ({
+        body: { accessToken: "new-access", refreshToken: "new-refresh", expiresIn: 1800 },
+    }));
+    const update = await refreshKiroToken(makeAccount(), { defaultRegion: "us-east-1", fetchImpl });
+    assert.equal(calls.length, 1);
+    assert.equal(calls[0].url, "https://prod.us-east-1.auth.desktop.kiro.dev/refreshToken");
+    assert.equal(JSON.parse(calls[0].init?.body).refreshToken, "refresh-a");
+    assert.equal(update.accessToken, "new-access");
+    assert.equal(update.refreshToken, "new-refresh");
+    assert.equal(update.expiresIn, 1800);
+});
+test("refreshKiroToken uses the regional OIDC endpoint for IDC accounts", async () => {
+    const { fetchImpl, calls } = mockFetch(() => ({
+        body: { accessToken: "idc-access", refreshToken: "idc-refresh", expiresIn: 3600 },
+    }));
+    const account = makeAccount({
+        providerSpecificData: {
+            profileArn: null,
+            clientId: "client-x",
+            clientSecret: "secret-x",
+            region: "eu-west-1",
+            authMethod: "idc",
+            startUrl: null,
+        },
+    });
+    await refreshKiroToken(account, { defaultRegion: "us-east-1", fetchImpl });
+    assert.equal(calls[0].url, "https://oidc.eu-west-1.amazonaws.com/token");
+    const body = JSON.parse(calls[0].init?.body);
+    assert.equal(body.clientId, "client-x");
+    assert.equal(body.clientSecret, "secret-x");
+    assert.equal(body.grantType, "refresh_token");
+});
+test("refreshKiroToken throws KiroAuthError on a non-ok response", async () => {
+    const { fetchImpl } = mockFetch(() => ({ status: 400, body: { error: "bad" } }));
+    await assert.rejects(() => refreshKiroToken(makeAccount(), { defaultRegion: "us-east-1", fetchImpl }), (error) => error instanceof KiroAuthError && error.body.code === "KIRO_TOKEN_REFRESH_FAILED");
+});
+test("refreshKiroToken throws when the response omits an accessToken", async () => {
+    const { fetchImpl } = mockFetch(() => ({ body: { refreshToken: "r" } }));
+    await assert.rejects(() => refreshKiroToken(makeAccount(), { defaultRegion: "us-east-1", fetchImpl }), (error) => error instanceof KiroAuthError);
+});
+test("resolveKiroCredentials returns a still-valid account without refreshing", async () => {
+    const store = new FakeStore([makeAccount()]);
+    const { fetchImpl, calls } = mockFetch(() => ({ body: {} }));
+    const creds = await resolveKiroCredentials({
+        store: store.asStore(),
+        accountId: "acct-a",
+        defaultRegion: "us-east-1",
+        refreshLeadSeconds: 120,
+        fetchImpl,
+        now: new Date("2026-05-15T00:00:00.000Z"),
+    });
+    assert.equal(calls.length, 0);
+    assert.equal(creds.accessToken, "access-a");
+    assert.equal(creds.profileArn, "arn:keep");
+    assert.equal(creds.region, "us-east-1");
+});
+test("resolveKiroCredentials refreshes when inside the lead window and persists the new token", async () => {
+    const store = new FakeStore([
+        makeAccount({ expiresAt: "2026-05-15T00:01:00.000Z" }),
+    ]);
+    const { fetchImpl, calls } = mockFetch(() => ({
+        body: { accessToken: "fresh-access", refreshToken: "fresh-refresh", expiresIn: 3600 },
+    }));
+    const creds = await resolveKiroCredentials({
+        store: store.asStore(),
+        accountId: "acct-a",
+        defaultRegion: "us-east-1",
+        refreshLeadSeconds: 120,
+        fetchImpl,
+        now: new Date("2026-05-15T00:00:00.000Z"),
+    });
+    assert.equal(calls.length, 1);
+    assert.equal(creds.accessToken, "fresh-access");
+    assert.equal(store.updates.length, 1);
+    assert.equal(store.getAccount("acct-a")?.accessToken, "fresh-access");
+});
+test("resolveKiroCredentials rotates round-robin across available accounts", async () => {
+    const store = new FakeStore([
+        makeAccount({ id: "acct-a", priority: 1 }),
+        makeAccount({ id: "acct-b", priority: 2 }),
+    ]);
+    const { fetchImpl } = mockFetch(() => ({ body: {} }));
+    const common = {
+        store: store.asStore(),
+        defaultRegion: "us-east-1",
+        refreshLeadSeconds: 120,
+        fetchImpl,
+        now: new Date("2026-05-15T00:00:00.000Z"),
+    };
+    const first = await resolveKiroCredentials(common);
+    const second = await resolveKiroCredentials(common);
+    const third = await resolveKiroCredentials(common);
+    assert.equal(first.accountId, "acct-a");
+    assert.equal(second.accountId, "acct-b");
+    assert.equal(third.accountId, "acct-a");
+});
+test("resolveKiroCredentials returns the fresh token even when write-back is disabled", async () => {
+    // Simulate writeBack:false: updateTokens persists nothing and returns the stale row.
+    class NoWriteBackStore extends FakeStore {
+        updateTokens(id) {
+            return this.getAccount(id);
+        }
+    }
+    const store = new NoWriteBackStore([makeAccount({ expiresAt: "2026-05-15T00:01:00.000Z" })]);
+    const { fetchImpl, calls } = mockFetch(() => ({
+        body: { accessToken: "fresh-access", refreshToken: "fresh-refresh", expiresIn: 3600 },
+    }));
+    const creds = await resolveKiroCredentials({
+        store: store.asStore(),
+        accountId: "acct-a",
+        defaultRegion: "us-east-1",
+        refreshLeadSeconds: 120,
+        fetchImpl,
+        now: new Date("2026-05-15T00:00:00.000Z"),
+    });
+    assert.equal(calls.length, 1);
+    // The refreshed token must be used, not the stale one the store returned.
+    assert.equal(creds.accessToken, "fresh-access");
+    // And the store row is intentionally left unchanged (no write-back).
+    assert.equal(store.getAccount("acct-a")?.accessToken, "access-a");
+});
+test("resolveKiroCredentials throws KIRO_ACCOUNT_UNAVAILABLE when no accounts are connected", async () => {
+    const store = new FakeStore([]);
+    const { fetchImpl } = mockFetch(() => ({ body: {} }));
+    await assert.rejects(() => resolveKiroCredentials({
+        store: store.asStore(),
+        defaultRegion: "us-east-1",
+        refreshLeadSeconds: 120,
+        fetchImpl,
+    }), (error) => error instanceof KiroAuthError && error.body.code === "KIRO_ACCOUNT_UNAVAILABLE");
+});
+test("concurrent refreshes for the same account share a single refresh call", async () => {
+    const store = new FakeStore([makeAccount({ expiresAt: "2026-05-15T00:01:00.000Z" })]);
+    let calls = 0;
+    const fetchImpl = (async () => {
+        calls += 1;
+        await new Promise((resolve) => setTimeout(resolve, 5));
+        return new Response(JSON.stringify({ accessToken: "fresh", refreshToken: "fresh-r", expiresIn: 3600 }), { status: 200, headers: { "Content-Type": "application/json" } });
+    });
+    const args = {
+        store: store.asStore(),
+        accountId: "acct-a",
+        defaultRegion: "us-east-1",
+        refreshLeadSeconds: 120,
+        fetchImpl,
+        now: new Date("2026-05-15T00:00:00.000Z"),
+    };
+    const [a, b] = await Promise.all([
+        resolveKiroCredentials(args),
+        resolveKiroCredentials(args),
+    ]);
+    assert.equal(calls, 1);
+    assert.equal(a.accessToken, "fresh");
+    assert.equal(b.accessToken, "fresh");
+});