255 lines
12 KiB
TypeScript
255 lines
12 KiB
TypeScript
/**
|
||
* mountA2aRequestHandler — DefaultRequestHandler + jsonRpcHandler を /a2a にマウント。
|
||
*
|
||
* Plan 2A の base card route (`/.well-known/agent-card.json`) と
|
||
* extended card route (`/a2a/agent-card/extended`) はそのまま残す。
|
||
* SDK の DefaultRequestHandler が `agent/getAuthenticatedExtendedCard` RPC を別途提供するが、
|
||
* Plan 2A の自前 REST route と競合しない（path が異なる）。
|
||
* 両者は同じ `buildExtendedCard` + `computeEffectiveScope` を呼ぶため実装ドリフトは起きない。
|
||
*/
|
||
import express from 'express';
|
||
import type { Express } from 'express';
|
||
import { A2AError, DefaultRequestHandler, type ServerCallContext } from '@a2a-js/sdk/server';
|
||
import { jsonRpcHandler } from '@a2a-js/sdk/server/express';
|
||
import type Provider from 'oidc-provider';
|
||
import type { Repository } from '../../db/repository.js';
|
||
import { logger } from '../../logger.js';
|
||
import { SqliteA2aTaskStore } from './task-store.js';
|
||
import { MaestroA2aExecutor } from './executor.js';
|
||
import { buildBaseCard, buildExtendedCard } from './agent-card.js';
|
||
import { computeEffectiveScope } from './delegation.js';
|
||
import { makeA2aUserBuilder } from './user-builder.js';
|
||
import type { A2aPrincipal } from './token-auth.js';
|
||
import { A2aLimiter, resolveA2aLimits, type A2aLimits } from './limiter.js';
|
||
|
||
export interface MountA2aRequestHandlerDeps {
|
||
provider: Provider;
|
||
repo: Repository;
|
||
baseUrl: string;
|
||
issuer: string;
|
||
version: string;
|
||
limits?: Partial<A2aLimits>;
|
||
}
|
||
|
||
type CardDeps = { baseUrl: string; issuer: string; version: string };
|
||
|
||
/**
|
||
* DefaultRequestHandler のサブクラス。
|
||
* `tasks/get`・`tasks/resubscribe`・`tasks/cancel` の 3 RPC に認証ゲートを追加する。
|
||
*
|
||
* 委任を取り消されたクライアント（または未認証クライアント）は、
|
||
* すでに知っているタスク ID を使ってもタスクの状態・成果物を読み取れない（fail-closed）。
|
||
* context.user が存在しない・isAuthenticated !== true の場合は A2AError.invalidRequest を投げ、
|
||
* JsonRpcTransportHandler がそれを -32600 の JSON-RPC エラーレスポンスに変換する。
|
||
*
|
||
* sendMessage / sendMessageStream はゲートしない：executor.execute() が
|
||
* principal のチェックを行い、unauthorized ならタスクを failed 状態で終了させる（既存の保護）。
|
||
*
|
||
* Task 7: limiter が注入された場合、認証チェックの後に rate-limit と
|
||
* resubscribe スロットキャップを適用する（auth → rate → slot の順序を保証）。
|
||
*/
|
||
export class AuthGatedRequestHandler extends DefaultRequestHandler {
|
||
private limiter?: A2aLimiter;
|
||
|
||
constructor(
|
||
agentCard: ConstructorParameters<typeof DefaultRequestHandler>[0],
|
||
taskStore: ConstructorParameters<typeof DefaultRequestHandler>[1],
|
||
executor: ConstructorParameters<typeof DefaultRequestHandler>[2],
|
||
extendedCardProvider?: ConstructorParameters<typeof DefaultRequestHandler>[6],
|
||
limiter?: A2aLimiter,
|
||
) {
|
||
super(agentCard, taskStore, executor, undefined, undefined, undefined, extendedCardProvider);
|
||
this.limiter = limiter;
|
||
}
|
||
|
||
override async getTask(params: Parameters<DefaultRequestHandler['getTask']>[0], context?: ServerCallContext) {
|
||
this.requireAuthenticated(context);
|
||
const grantId = this.meteredGrantId(context);
|
||
if (this.limiter && !this.limiter.tryConsumeRate(grantId!)) {
|
||
throw A2AError.invalidRequest('rate limit exceeded');
|
||
}
|
||
return super.getTask(params, context);
|
||
}
|
||
|
||
override async cancelTask(params: Parameters<DefaultRequestHandler['cancelTask']>[0], context?: ServerCallContext) {
|
||
this.requireAuthenticated(context);
|
||
const grantId = this.meteredGrantId(context);
|
||
if (this.limiter && !this.limiter.tryConsumeRate(grantId!)) {
|
||
throw A2AError.invalidRequest('rate limit exceeded');
|
||
}
|
||
return super.cancelTask(params, context);
|
||
}
|
||
|
||
/**
|
||
* 認証ゲートを同期的に実行してから super.resubscribe() を返す。
|
||
* 同期 throw にすることで、JsonRpcTransportHandler がジェネレータを取得する前に
|
||
* エラーを捕捉できる（SSE ヘッダ送信前に正規の JSON-RPC エラーレスポンスを返せる）。
|
||
*
|
||
* limiter 注入時:
|
||
* 1. auth check (sync throw)
|
||
* 2. rate check (sync throw)
|
||
* 3. slot acquire (sync throw if cap reached)
|
||
* 4. wrap inner generator — finally releases slot on any exit path
|
||
*/
|
||
override resubscribe(params: Parameters<DefaultRequestHandler['resubscribe']>[0], context?: ServerCallContext) {
|
||
// 1. Auth — must be first and synchronous
|
||
this.requireAuthenticated(context);
|
||
// Fail-closed: with a limiter active, an authenticated caller missing a grantId is denied.
|
||
const grantId = this.meteredGrantId(context);
|
||
|
||
// No limiter → unmetered (e.g. test without limiter / metering off)
|
||
if (!this.limiter) {
|
||
return super.resubscribe(params, context);
|
||
}
|
||
|
||
// 2. Rate limit (limiter present ⟹ grantId is defined, else meteredGrantId threw)
|
||
if (!this.limiter.tryConsumeRate(grantId!)) {
|
||
throw A2AError.invalidRequest('rate limit exceeded');
|
||
}
|
||
|
||
// 3. Slot cap — synchronous acquire before returning generator
|
||
if (!this.limiter.tryAcquireResubscribe(grantId!)) {
|
||
throw A2AError.invalidRequest('too many concurrent resubscribe streams');
|
||
}
|
||
|
||
// 4. Wrap inner generator with a real wall-clock race so the stream-duration cap
|
||
// fires even when the inner stream emits no events (quiet-stream problem).
|
||
// Release slot in finally on every exit path.
|
||
const inner = super.resubscribe(params, context);
|
||
const limiter = this.limiter;
|
||
const maxMs = this.limiter.limits.maxStreamSeconds * 1000;
|
||
async function* wrapped() {
|
||
// Idempotent release: guards against the JS generator edge case where a consumer
|
||
// calls .return() before the first .next() — in that case the generator body's
|
||
// finally block does not execute. In practice the SDK SSE path always calls
|
||
// .next() before .return(), so this window is unreachable in production.
|
||
let released = false;
|
||
function release() {
|
||
if (!released) { released = true; limiter.releaseResubscribe(grantId!); }
|
||
}
|
||
const deadline = Date.now() + maxMs;
|
||
const it = inner[Symbol.asyncIterator]();
|
||
try {
|
||
while (true) {
|
||
const remaining = deadline - Date.now();
|
||
// Fire-and-forget the inner .return(): a `.return()` issued while a prior
|
||
// `.next()` is still pending is queued behind it and does not settle until
|
||
// that next() settles — which never happens for a quiet stream (the exact
|
||
// case this cap targets). Awaiting it here would suspend `wrapped()` forever,
|
||
// skip the `finally` below, and leak the resubscribe slot. Instead we drop the
|
||
// slot immediately via `return` (→ finally → release) and let the inner
|
||
// generator wind down on its own.
|
||
if (remaining <= 0) { void it.return?.(undefined)?.catch(() => {}); return; }
|
||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||
const timeout = new Promise<'timeout'>(res => {
|
||
timer = setTimeout(() => res('timeout'), remaining);
|
||
});
|
||
const step = await Promise.race([it.next(), timeout]).finally(() => {
|
||
if (timer !== undefined) clearTimeout(timer);
|
||
});
|
||
if (step === 'timeout') { void it.return?.(undefined)?.catch(() => {}); return; }
|
||
if (step.done) return;
|
||
yield step.value;
|
||
}
|
||
} finally {
|
||
release();
|
||
}
|
||
}
|
||
return wrapped() as ReturnType<DefaultRequestHandler['resubscribe']>;
|
||
}
|
||
|
||
private extractGrantId(context?: ServerCallContext): string | undefined {
|
||
const p = (context?.user as any)?.a2aPrincipal as A2aPrincipal | undefined;
|
||
return p?.grantId;
|
||
}
|
||
|
||
/**
|
||
* Resolve the grantId for metering. When a limiter is active (a2a metering enabled),
|
||
* an authenticated caller MUST carry a grantId — a missing one is treated as deny
|
||
* (fail-closed), never as "skip the limits". Without a limiter (tests / metering off)
|
||
* the caller is simply unmetered and undefined is returned.
|
||
*/
|
||
private meteredGrantId(context?: ServerCallContext): string | undefined {
|
||
const grantId = this.extractGrantId(context);
|
||
if (this.limiter && !grantId) {
|
||
throw A2AError.invalidRequest('unauthorized: missing delegation grant');
|
||
}
|
||
return grantId;
|
||
}
|
||
|
||
private requireAuthenticated(context?: ServerCallContext): void {
|
||
const user = context?.user;
|
||
if (!user || user.isAuthenticated !== true) {
|
||
throw A2AError.invalidRequest(
|
||
'unauthorized: caller not authenticated or delegation has been revoked',
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
/**
|
||
* Extended card provider（DefaultRequestHandler 7th 引数）。
|
||
* ServerCallContext.user.a2aPrincipal が存在すれば scope 計算後の extended card を返す。
|
||
* principal がない（未認証）または scope が null（unknown/inactive user）の場合は
|
||
* base card を返す（スコープ情報の漏洩を防ぐ fail-closed 設計）。
|
||
*
|
||
* テスト可能にするため named export にする。
|
||
*/
|
||
export function makeExtendedCardProvider(
|
||
repo: Repository,
|
||
baseCard: ReturnType<typeof buildBaseCard>,
|
||
deps: CardDeps,
|
||
) {
|
||
return async (context?: ServerCallContext) => {
|
||
const principal = (context?.user as any)?.a2aPrincipal as A2aPrincipal | undefined;
|
||
if (!principal) return baseCard; // 未認証 → base card（no scope leak）
|
||
try {
|
||
const scope = await computeEffectiveScope(repo, principal);
|
||
if (!scope) return baseCard; // unknown/inactive user → base card
|
||
return buildExtendedCard(principal, {
|
||
baseUrl: deps.baseUrl,
|
||
issuer: deps.issuer,
|
||
version: deps.version,
|
||
userVisibleSpaceIds: scope.userVisibleSpaceIds,
|
||
allowlistsBySpace: scope.allowlistsBySpace,
|
||
});
|
||
} catch (err) {
|
||
logger.warn(`[a2a] extended card provider failed: ${err}`);
|
||
return baseCard;
|
||
}
|
||
};
|
||
}
|
||
|
||
/**
|
||
* DefaultRequestHandler + jsonRpcHandler を Express app の /a2a にマウントする。
|
||
* server.ts の `if (a2aCfg.enabled)` ブロック内から呼ぶこと。
|
||
* /a2a は /api 配下でないため requireAuth ブランケットに掛からない。
|
||
* 認証は makeA2aUserBuilder（UserBuilder）が担う。
|
||
*/
|
||
export function mountA2aRequestHandler(app: Express, deps: MountA2aRequestHandlerDeps): void {
|
||
const limits = resolveA2aLimits(deps.limits);
|
||
const limiter = new A2aLimiter(limits);
|
||
const taskStore = new SqliteA2aTaskStore(deps.repo);
|
||
const executor = new MaestroA2aExecutor(deps.repo, { limiter });
|
||
const cardDeps: CardDeps = {
|
||
baseUrl: deps.baseUrl,
|
||
issuer: deps.issuer,
|
||
version: deps.version,
|
||
};
|
||
const baseCard = buildBaseCard(cardDeps);
|
||
const extendedCardProvider = makeExtendedCardProvider(deps.repo, baseCard, cardDeps);
|
||
const handler = new AuthGatedRequestHandler(
|
||
baseCard,
|
||
taskStore,
|
||
executor,
|
||
extendedCardProvider, // 4th arg: forwarded to super's 7th slot
|
||
limiter, // 5th arg: shared rate/stream limiter (Task 7)
|
||
);
|
||
const userBuilder = makeA2aUserBuilder(deps.provider, deps.repo);
|
||
// body-size limit: must precede jsonRpcHandler so oversized bodies are rejected
|
||
// before the SDK parser sees them. body-parser skips re-parsing when req.body is
|
||
// already set, so the SDK's own express.json() inside jsonRpcHandler won't double-parse.
|
||
app.use('/a2a', express.json({ limit: limits.maxPayloadBytes }));
|
||
app.use('/a2a', jsonRpcHandler({ requestHandler: handler, userBuilder }));
|
||
}
|