maestro/src/bridge/a2a/request-handler-auth-gate.test.ts

// @vitest-environment node
/**
 * AuthGatedRequestHandler — unit + E2E tests.
 *
 * What is tested:
 *   - getTask / cancelTask / resubscribe reject UnauthenticatedUser (or missing context)
 *     with A2AError code -32600 (INVALID_REQUEST).
 *   - Authenticated contexts pass the auth gate and reach the underlying store
 *     (which throws taskNotFound -32001, not the auth error).
 *   - E2E via supertest: unauthenticated `tasks/get` HTTP request → proper JSON-RPC
 *     error response (not 500).
 *
 * Test approach: unit-test AuthGatedRequestHandler directly with InMemoryTaskStore
 * and a stub executor (fastest, no HTTP round-trip needed for unit assertions).
 * One E2E supertest case covers the full HTTP path.
 */
import { describe, it, expect, vi } from 'vitest';
import express from 'express';
import request from 'supertest';
import {
  A2AError,
  DefaultRequestHandler,
  InMemoryTaskStore,
  ServerCallContext,
  UnauthenticatedUser,
} from '@a2a-js/sdk/server';
import { Repository } from '../../db/repository.js';
import { buildBaseCard } from './agent-card.js';
import { AuthGatedRequestHandler, mountA2aRequestHandler } from './request-handler.js';
import { A2aAuthenticatedUser } from './user-builder.js';
import type { A2aPrincipal } from './token-auth.js';
import { A2aLimiter, resolveA2aLimits } from './limiter.js';

// ─── fixtures ─────────────────────────────────────────────────────────────────

const BASE_DEPS = { baseUrl: 'https://m', issuer: 'https://m/oidc', version: '1.0.0' };

const FAKE_PRINCIPAL: A2aPrincipal = {
  actingUserId: 'u1', clientId: 'cli1', grantId: 'g1',
  delegation: {
    id: 'd1', userId: 'u1', clientId: 'cli1', grantId: 'g1',
    grantedSpaceIds: ['s1'], grantedSkills: ['research'],
    expiresAt: null, revokedAt: null,
  },
};

/** Helper — construct AuthGatedRequestHandler with an in-memory store. */
function makeHandler() {
  const store = new InMemoryTaskStore();
  const fakeExecutor = { execute: async () => {} } as any;
  const baseCard = buildBaseCard(BASE_DEPS);
  return new AuthGatedRequestHandler(baseCard, store, fakeExecutor);
}

const UNAUTH_CTX = new ServerCallContext(undefined, new UnauthenticatedUser());
const AUTH_CTX   = new ServerCallContext(undefined, new A2aAuthenticatedUser(FAKE_PRINCIPAL));

// ─── getTask ──────────────────────────────────────────────────────────────────

describe('AuthGatedRequestHandler.getTask', () => {
  it('rejects UnauthenticatedUser with A2AError', async () => {
    const h = makeHandler();
    await expect(h.getTask({ id: 'x' }, UNAUTH_CTX)).rejects.toBeInstanceOf(A2AError);
  });

  it('rejected error carries INVALID_REQUEST code (-32600)', async () => {
    const h = makeHandler();
    await expect(h.getTask({ id: 'x' }, UNAUTH_CTX))
      .rejects.toMatchObject({ code: -32600 });
  });

  it('fail-closed: undefined context also rejects', async () => {
    const h = makeHandler();
    await expect(h.getTask({ id: 'x' }, undefined)).rejects.toMatchObject({ code: -32600 });
  });

  it('authenticated context passes auth gate (taskNotFound, not auth error)', async () => {
    const h = makeHandler();
    const err = await h.getTask({ id: 'nonexistent' }, AUTH_CTX).catch(e => e);
    expect(err).toBeInstanceOf(A2AError);
    // -32001 = TASK_NOT_FOUND (auth gate was cleared, underlying store was reached)
    expect(err.code).not.toBe(-32600);
  });
});

// ─── cancelTask ───────────────────────────────────────────────────────────────

describe('AuthGatedRequestHandler.cancelTask', () => {
  it('rejects UnauthenticatedUser with INVALID_REQUEST (-32600)', async () => {
    const h = makeHandler();
    await expect(h.cancelTask({ id: 'x' }, UNAUTH_CTX))
      .rejects.toMatchObject({ code: -32600 });
  });

  it('fail-closed: undefined context rejects', async () => {
    const h = makeHandler();
    await expect(h.cancelTask({ id: 'x' }, undefined)).rejects.toMatchObject({ code: -32600 });
  });

  it('authenticated context passes auth gate (taskNotFound, not auth error)', async () => {
    const h = makeHandler();
    const err = await h.cancelTask({ id: 'nonexistent' }, AUTH_CTX).catch(e => e);
    expect(err).toBeInstanceOf(A2AError);
    expect(err.code).not.toBe(-32600);
  });
});

// ─── resubscribe ──────────────────────────────────────────────────────────────

describe('AuthGatedRequestHandler.resubscribe', () => {
  it('rejects UnauthenticatedUser synchronously (throws before returning generator)', () => {
    const h = makeHandler();
    expect(() => h.resubscribe({ id: 'x' }, UNAUTH_CTX)).toThrow(A2AError);
  });

  it('thrown error carries INVALID_REQUEST code (-32600)', () => {
    const h = makeHandler();
    let caught: unknown;
    try { h.resubscribe({ id: 'x' }, UNAUTH_CTX); } catch (e) { caught = e; }
    expect(caught).toBeInstanceOf(A2AError);
    expect((caught as A2AError).code).toBe(-32600);
  });

  it('fail-closed: undefined context throws synchronously', () => {
    const h = makeHandler();
    expect(() => h.resubscribe({ id: 'x' }, undefined)).toThrow(A2AError);
  });

  it('authenticated context returns an AsyncGenerator without throwing', () => {
    const h = makeHandler();
    let gen: ReturnType<typeof h.resubscribe> | undefined;
    expect(() => { gen = h.resubscribe({ id: 'x' }, AUTH_CTX); }).not.toThrow();
    expect(gen).toBeDefined();
    expect(typeof gen![Symbol.asyncIterator]).toBe('function');
    gen!.return?.(undefined); // clean up
  });
});

// ─── Rate limiting + resubscribe slot cap (Task 7) ────────────────────────────

describe('AuthGatedRequestHandler — rate limiting', () => {
  function makeHandlerWithLimiter(limiter: A2aLimiter) {
    const store = new InMemoryTaskStore();
    const fakeExecutor = { execute: async () => {} } as any;
    const baseCard = buildBaseCard(BASE_DEPS);
    return new AuthGatedRequestHandler(baseCard, store, fakeExecutor, undefined, limiter);
  }

  it('getTask: second call past rate=1 throws A2AError (rate limit exceeded)', async () => {
    const limiter = new A2aLimiter(resolveA2aLimits({ ratePerMinute: 1 }));
    const h = makeHandlerWithLimiter(limiter);
    // First call: auth OK, 1 rate token consumed → falls through to taskNotFound (not rate error)
    await h.getTask({ id: 'x' }, AUTH_CTX).catch(() => {});
    // Second call: no token left → rate-limit error
    const err = await h.getTask({ id: 'x' }, AUTH_CTX).catch(e => e);
    expect(err).toBeInstanceOf(A2AError);
    expect(err.code).toBe(-32600);
    expect(err.message).toMatch(/rate limit/);
  });

  it('getTask: unauthenticated throws auth error even when budget is fully exhausted', async () => {
    const limiter = new A2aLimiter(resolveA2aLimits({ ratePerMinute: 1 }));
    // Pre-drain the bucket for grant g1 (the FAKE_PRINCIPAL's grantId)
    limiter.tryConsumeRate('g1');
    const h = makeHandlerWithLimiter(limiter);
    // Auth check must fire before rate check; message must be 'unauthorized', not 'rate limit'
    const err = await h.getTask({ id: 'x' }, UNAUTH_CTX).catch(e => e);
    expect(err).toBeInstanceOf(A2AError);
    expect(err.message).toMatch(/unauthorized/);
  });

  it('cancelTask: second call past rate=1 throws A2AError (rate limit exceeded)', async () => {
    const limiter = new A2aLimiter(resolveA2aLimits({ ratePerMinute: 1 }));
    const h = makeHandlerWithLimiter(limiter);
    // First call: auth OK, 1 rate token consumed → falls through to taskNotFound (not rate error)
    await h.cancelTask({ id: 'x' }, AUTH_CTX).catch(() => {});
    // Second call: no token left → rate-limit error
    const err = await h.cancelTask({ id: 'x' }, AUTH_CTX).catch(e => e);
    expect(err).toBeInstanceOf(A2AError);
    expect(err.code).toBe(-32600);
    expect(err.message).toMatch(/rate limit/);
  });
});

describe('AuthGatedRequestHandler — resubscribe slot cap (Task 7)', () => {
  function makeHandlerWithLimiter(limiter: A2aLimiter) {
    const store = new InMemoryTaskStore();
    const fakeExecutor = { execute: async () => {} } as any;
    const baseCard = buildBaseCard(BASE_DEPS);
    return new AuthGatedRequestHandler(baseCard, store, fakeExecutor, undefined, limiter);
  }

  it('resubscribe: second call throws when slot cap (maxConcurrentResubscribe=1) is reached', () => {
    const limiter = new A2aLimiter(resolveA2aLimits({ maxConcurrentResubscribe: 1, ratePerMinute: 100 }));
    const h = makeHandlerWithLimiter(limiter);
    // First resubscribe: slot acquired
    const gen1 = h.resubscribe({ id: 'x' }, AUTH_CTX);
    expect(gen1).toBeDefined();
    // Second resubscribe: slot full → synchronous throw
    expect(() => h.resubscribe({ id: 'x' }, AUTH_CTX)).toThrow(A2AError);
    gen1.return?.(undefined); // clean up (async; we don't await)
  });

  it('resubscribe: slot released via finally when generator is iterated and exhausted', async () => {
    const limiter = new A2aLimiter(resolveA2aLimits({ maxConcurrentResubscribe: 1, ratePerMinute: 100 }));
    const h = makeHandlerWithLimiter(limiter);
    // Acquire the slot
    const gen1 = h.resubscribe({ id: 'x' }, AUTH_CTX);
    // Start the generator — inner store will throw (task not found), triggering finally
    await gen1.next().catch(() => {});
    // Slot should now be released; a new resubscribe must not throw
    let gen2: AsyncGenerator<any> | undefined;
    expect(() => { gen2 = h.resubscribe({ id: 'x' }, AUTH_CTX); }).not.toThrow();
    expect(gen2).toBeDefined();
    await gen2!.next().catch(() => {}); // clean up
  });
});

// ─── ingress governance: sendMessage / sendMessageStream (2C-3 review #2) ─────
// payload + rate are enforced here, BEFORE execute() runs, so a denial throws a JSON-RPC
// error and persists nothing (the SDK only writes an a2a_tasks row once execute() publishes).

describe('AuthGatedRequestHandler — ingress governance (send/stream)', () => {
  const MSG = {
    message: { kind: 'message', role: 'user', messageId: 'm1', parts: [{ kind: 'text', text: 'hello world' }] },
  } as any;

  function makeSendHandler(limiter: A2aLimiter) {
    const store = new InMemoryTaskStore();
    const execute = vi.fn(async () => {});
    const fakeExecutor = { execute } as any;
    const baseCard = buildBaseCard(BASE_DEPS);
    const h = new AuthGatedRequestHandler(baseCard, store, fakeExecutor, undefined, limiter);
    return { h, execute };
  }

  it('sendMessage: oversized payload → throws -32600 and execute() is never invoked', async () => {
    const limiter = new A2aLimiter(resolveA2aLimits({ maxPayloadBytes: 1, ratePerMinute: 100 }));
    const { h, execute } = makeSendHandler(limiter);
    await expect(h.sendMessage(MSG, AUTH_CTX)).rejects.toMatchObject({ code: -32600 });
    expect(execute).not.toHaveBeenCalled();
  });

  it('sendMessage: rate bucket drained → throws -32600 and execute() is never invoked', async () => {
    const limiter = new A2aLimiter(resolveA2aLimits({ ratePerMinute: 1, maxPayloadBytes: 65536 }));
    limiter.tryConsumeRate('g1'); // drain the single token (AUTH_CTX principal grantId = g1)
    const { h, execute } = makeSendHandler(limiter);
    await expect(h.sendMessage(MSG, AUTH_CTX)).rejects.toMatchObject({ code: -32600 });
    expect(execute).not.toHaveBeenCalled();
  });

  it('sendMessage: unauthenticated → throws -32600 (no failed Task persisted, no execute)', async () => {
    const limiter = new A2aLimiter(resolveA2aLimits({ ratePerMinute: 100 }));
    const { h, execute } = makeSendHandler(limiter);
    await expect(h.sendMessage(MSG, UNAUTH_CTX)).rejects.toMatchObject({ code: -32600 });
    expect(execute).not.toHaveBeenCalled();
  });

  it('sendMessageStream: oversized payload → SYNCHRONOUS throw before the generator (no SSE, no execute)', () => {
    const limiter = new A2aLimiter(resolveA2aLimits({ maxPayloadBytes: 1, ratePerMinute: 100 }));
    const { h, execute } = makeSendHandler(limiter);
    expect(() => h.sendMessageStream(MSG, AUTH_CTX)).toThrow(A2AError);
    expect(execute).not.toHaveBeenCalled();
  });

  it('sendMessageStream: rate drained → synchronous throw before the generator', () => {
    const limiter = new A2aLimiter(resolveA2aLimits({ ratePerMinute: 1, maxPayloadBytes: 65536 }));
    limiter.tryConsumeRate('g1');
    const { h, execute } = makeSendHandler(limiter);
    expect(() => h.sendMessageStream(MSG, AUTH_CTX)).toThrow(A2AError);
    expect(execute).not.toHaveBeenCalled();
  });

  it('within limits: gate passes and consumes a rate token (second call is throttled)', () => {
    // sendMessageStream returns the generator synchronously; the base body (and execute())
    // only runs on the first .next(), which would block on a quiet stub — so we assert the
    // gate side effects instead: the first call passes and drains the single token, the
    // second call is rejected, proving the first consumed it.
    const limiter = new A2aLimiter(resolveA2aLimits({ ratePerMinute: 1, maxPayloadBytes: 65536 }));
    const { h } = makeSendHandler(limiter);
    let gen: AsyncGenerator<any> | undefined;
    expect(() => { gen = h.sendMessageStream(MSG, AUTH_CTX); }).not.toThrow();
    expect(gen).toBeDefined();
    void gen!.return?.(undefined); // unstarted generator → completes without running the body
    expect(() => h.sendMessageStream(MSG, AUTH_CTX)).toThrow(A2AError);
  });
});

// ─── maxStreamSeconds fires for a quiet stream (Task 7 review fix) ───────────

describe('AuthGatedRequestHandler — maxStreamSeconds wall-clock race', () => {
  it('quiet stream (no events) completes and releases slot at deadline', async () => {
    // Use a fractional maxStreamSeconds (0.05 = 50 ms) for a fast test.
    // A2aLimiter accepts raw limits directly — bypassing resolveA2aLimits integer≥1 constraint
    // is intentional here so we don't need to wait a full second.
    const limiter = new A2aLimiter({
      ratePerMinute: 100,
      maxConcurrentTasks: 5,
      maxPayloadBytes: 65536,
      maxStreamSeconds: 0.05, // 50 ms deadline
      skillBudgetPerHour: 100,
      maxConcurrentResubscribe: 1,
    });
    const store = new InMemoryTaskStore();
    const fakeExecutor = { execute: async () => {} } as any;
    const baseCard = buildBaseCard(BASE_DEPS);
    const h = new AuthGatedRequestHandler(baseCard, store, fakeExecutor, undefined, limiter);

    // Stub super.resubscribe with a generator that never emits events (quiet live stream).
    // This is the exact scenario the timer-race fixes: the old for-await loop would
    // hold the slot indefinitely; the new Promise.race releases it at the deadline.
    const neverYielding: AsyncGenerator<any, void, undefined> = {
      [Symbol.asyncIterator]() { return this; },
      next: () => new Promise<never>(() => {}), // blocks indefinitely
      return: async (v: any) => ({ done: true as const, value: v }),
      throw: async (e: any): Promise<never> => { throw e; },
    };
    const spy = vi.spyOn(DefaultRequestHandler.prototype, 'resubscribe').mockReturnValue(
      neverYielding as any,
    );

    try {
      // Slot cap=1: acquire it
      const gen = h.resubscribe({ id: 'x' }, AUTH_CTX);
      // Kick the generator — suspends inside the 50 ms Promise.race
      const firstNext = gen.next();
      // Wait 120 ms (real timers) — the 50 ms deadline fires, race resolves with 'timeout'
      await new Promise(res => setTimeout(res, 120));
      // Generator must have completed via the timeout path
      const result = await firstNext;
      expect(result.done).toBe(true);
      // Slot must be released — second resubscribe must not throw
      let gen2: AsyncGenerator<any> | undefined;
      expect(() => { gen2 = h.resubscribe({ id: 'x' }, AUTH_CTX); }).not.toThrow();
      expect(gen2).toBeDefined();
      await gen2!.return?.(undefined); // cleanup
    } finally {
      spy.mockRestore();
    }
  }, 2000); // generous timeout for the real-timer wait

  it('quiet REAL async generator: deadline still releases slot even though inner .return() never settles', async () => {
    // Regression guard for the follow-up fix: a REAL async generator suspended inside its
    // body (not at a yield) queues a .return() behind the still-pending .next(), so the
    // return never settles for a quiet stream. The old `await it.return()` would hang the
    // wrapper forever (never reaching finally → slot leak). The stub in the test above hid
    // this because its hand-rolled `return` resolved immediately. Here we use a genuine
    // async generator to reproduce the queuing semantics.
    const limiter = new A2aLimiter({
      ratePerMinute: 100,
      maxConcurrentTasks: 5,
      maxPayloadBytes: 65536,
      maxStreamSeconds: 0.05, // 50 ms deadline
      skillBudgetPerHour: 100,
      maxConcurrentResubscribe: 1,
    });
    const store = new InMemoryTaskStore();
    const fakeExecutor = { execute: async () => {} } as any;
    const baseCard = buildBaseCard(BASE_DEPS);
    const h = new AuthGatedRequestHandler(baseCard, store, fakeExecutor, undefined, limiter);

    // A real async generator that suspends forever inside its body. Calling .return() while
    // the .next() promise is pending queues the return; it will not settle for this stream.
    async function* realQuiet(): AsyncGenerator<any, void, undefined> {
      await new Promise<void>(() => {}); // never resolves
      yield 1; // unreachable
    }
    const spy = vi.spyOn(DefaultRequestHandler.prototype, 'resubscribe').mockReturnValue(
      realQuiet() as any,
    );

    try {
      const gen = h.resubscribe({ id: 'x' }, AUTH_CTX);
      const firstNext = gen.next();
      await new Promise(res => setTimeout(res, 120)); // 50 ms deadline fires
      // Must complete via the timeout path despite the inner return never settling.
      const result = await firstNext;
      expect(result.done).toBe(true);
      // Slot released → a fresh resubscribe must not throw.
      let gen2: AsyncGenerator<any> | undefined;
      expect(() => { gen2 = h.resubscribe({ id: 'x' }, AUTH_CTX); }).not.toThrow();
      await gen2!.return?.(undefined);
    } finally {
      spy.mockRestore();
    }
  }, 2000);
});

// ─── fail-closed: authenticated but missing grantId is denied when metering is active ──

describe('AuthGatedRequestHandler — fail-closed on missing grantId', () => {
  const NO_GRANT_PRINCIPAL = { ...FAKE_PRINCIPAL, grantId: undefined } as any;
  const NO_GRANT_CTX = new ServerCallContext(undefined, new A2aAuthenticatedUser(NO_GRANT_PRINCIPAL));

  function makeHandlerWithLimiter() {
    const store = new InMemoryTaskStore();
    const fakeExecutor = { execute: async () => {} } as any;
    const baseCard = buildBaseCard(BASE_DEPS);
    const limiter = new A2aLimiter(resolveA2aLimits({ ratePerMinute: 100, maxConcurrentResubscribe: 5 }));
    return new AuthGatedRequestHandler(baseCard, store, fakeExecutor, undefined, limiter);
  }

  it('getTask denies an authenticated caller with no grantId (limiter active)', async () => {
    const h = makeHandlerWithLimiter();
    await expect(h.getTask({ id: 'x' }, NO_GRANT_CTX)).rejects.toMatchObject({ code: -32600 });
  });

  it('cancelTask denies an authenticated caller with no grantId (limiter active)', async () => {
    const h = makeHandlerWithLimiter();
    await expect(h.cancelTask({ id: 'x' }, NO_GRANT_CTX)).rejects.toMatchObject({ code: -32600 });
  });

  it('resubscribe denies an authenticated caller with no grantId (limiter active)', () => {
    const h = makeHandlerWithLimiter();
    expect(() => h.resubscribe({ id: 'x' }, NO_GRANT_CTX)).toThrow(A2AError);
  });

  it('without a limiter, a missing grantId is unmetered (reaches the store, not denied)', async () => {
    // No limiter → metering off → must not fail-closed. Store throws taskNotFound (-32001),
    // proving the auth+grant gate was passed rather than denied for the missing grant.
    const h = makeHandler();
    await expect(h.getTask({ id: 'x' }, NO_GRANT_CTX)).rejects.toMatchObject({ code: -32001 });
  });
});

// ─── E2E: tasks/get via HTTP → JSON-RPC error when unauthenticated ───────────

describe('mountA2aRequestHandler – tasks/get auth gate (E2E)', () => {
  function seedRepo(repo: Repository) {
    (repo as any).db.prepare(
      "INSERT INTO users (id, email, name, avatar_url, role, status, created_at, updated_at) " +
      "VALUES ('u1','u1@test','User One',NULL,'user','active','2026-01-01','2026-01-01')",
    ).run();
  }

  function fakeProvider() {
    return { AccessToken: { async find(_v: string) { return undefined; } } } as any;
  }

  it('unauthenticated tasks/get returns JSON-RPC error (not 500, code=-32600)', async () => {
    const repo = new Repository(':memory:');
    seedRepo(repo);
    const app = express();
    mountA2aRequestHandler(app, { provider: fakeProvider(), repo, ...BASE_DEPS });

    const res = await request(app)
      .post('/a2a')
      .set('Content-Type', 'application/json')
      .send({ jsonrpc: '2.0', id: 99, method: 'tasks/get', params: { id: 'some-task-id' } });

    expect(res.status).not.toBe(500);
    expect(res.headers['content-type']).toMatch(/json/);
    expect(res.body.jsonrpc).toBe('2.0');
    expect(res.body).toHaveProperty('error');
    expect(res.body.error.code).toBe(-32600); // INVALID_REQUEST = unauthorized
  });
});