feat(agent-harness): authorize scoped MCP gateway connections

apps/web/src/lib/agent-harness/mcp.test.ts

@@ -0,0 +1,283 @@
+import { beforeEach, expect, it, jest } from '@jest/globals';
+import { GatewayError } from '@kilocode/mcp-gateway';
+import { TRPCError } from '@trpc/server';
+import type * as McpModule from './mcp';
+
+const invocation = {
+  conversationId: '11111111-1111-4111-8111-111111111111',
+  operationId: '22222222-2222-4222-8222-222222222222',
+  request: { name: 'mcp.discover', arguments: {} },
+};
+let organizationId: string | null,
+  unavailable: boolean,
+  denied: boolean,
+  configurationVersion: number,
+  mintVersion: number,
+  gatewayBaseUrl: string,
+  canonicalOverride: string | undefined,
+  failure: Error | undefined,
+  failureStage:
+    | 'authorization'
+    | 'configuration'
+    | 'availability'
+    | 'route'
+    | 'mint'
+    | 'verification';
+const config = (id: string) => ({
+  configId: id,
+  canonicalUrl: `${gatewayBaseUrl}/${id}`,
+  registryMetadata: { providerSecret: 'provider-secret' },
+});
+function throwIfFailed(stage: typeof failureStage) {
+  if (failure && failureStage === stage) throw failure;
+}
+jest.mock('./authorization', () => ({
+  harnessInputDigest: (args: unknown) => JSON.stringify(args),
+  authorizeHarnessCapability: async (token: string, scope: Record<string, unknown>) => {
+    throwIfFailed('authorization');
+    if (
+      denied ||
+      token !== 'capability' ||
+      scope.audience !== 'agent-harness:operations' ||
+      scope.conversationId !== invocation.conversationId ||
+      scope.dispatchId !== invocation.operationId ||
+      scope.operation !== invocation.request.name ||
+      scope.definitionVersion !== '1' ||
+      JSON.stringify(scope.target) !== '{"kind":"backend"}' ||
+      scope.inputDigest !== JSON.stringify(invocation.request.arguments)
+    )
+      throw new TRPCError({ code: 'FORBIDDEN' });
+    return { authority: { userId: 'oauth/github:owner', organizationId } };
+  },
+}));
+jest.mock('@/lib/mcp-gateway/services', () => ({
+  createGatewayServices: () => {
+    throwIfFailed('configuration');
+    return {
+      availableService: {
+        listAvailableConfigs: async (
+          user: string,
+          context: { type: string; organizationId?: string }
+        ) => {
+          throwIfFailed('availability');
+          return unavailable || user !== 'oauth/github:owner'
+            ? []
+            : [
+                config(
+                  context.type === 'personal'
+                    ? 'personal'
+                    : (context.organizationId ?? 'wrong-context')
+                ),
+              ];
+        },
+      },
+      routeService: {
+        resolveResource: async (url: string) => {
+          throwIfFailed('route');
+          return {
+            route: { configId: url.split('/').at(-1) },
+            resolved: {
+              config: { config_version: configurationVersion, encrypted_secret: 'provider-secret' },
+            },
+          };
+        },
+        canonicalUrl: (route: { configId: string }) =>
+          canonicalOverride ?? config(route.configId).canonicalUrl,
+      },
+      tokenService: {
+        mintDerivedConnectToken: async ({
+          route,
+          userId,
+          executionContext,
+        }: {
+          route: { configId: string };
+          userId: string;
+          executionContext: { type: string; organizationId?: string };
+        }) => {
+          throwIfFailed('mint');
+          if (
+            route.configId !== (organizationId ?? 'personal') ||
+            userId !== 'oauth/github:owner' ||
+            executionContext.type !== (organizationId === null ? 'personal' : 'organization') ||
+            executionContext.organizationId !== (organizationId ?? undefined)
+          )
+            throw new Error('Wrong authority');
+          return { token: 'derived-only' };
+        },
+        verifyUserInfoToken: async (token: string) => {
+          throwIfFailed('verification');
+          if (token !== 'derived-only') throw new Error('Wrong token');
+          return { config_version: mintVersion };
+        },
+      },
+    };
+  },
+}));
+const { authorizeHarnessMcp } = jest.requireActual<typeof McpModule>('./mcp');
+beforeEach(() => {
+  organizationId = null;
+  unavailable = false;
+  denied = false;
+  configurationVersion = 1;
+  mintVersion = 1;
+  gatewayBaseUrl = 'https://gateway.example';
+  canonicalOverride = undefined;
+  failure = undefined;
+  failureStage = 'mint';
+  invocation.request = { name: 'mcp.discover', arguments: {} };
+});
+
+it.each([null, 'organization'])('uses only current conversation scope: %s', async scope => {
+  organizationId = scope;
+  configurationVersion = mintVersion = 7;
+  const result = await authorizeHarnessMcp('capability', invocation);
+  expect(result).toEqual([
+    {
+      serverId: scope ?? 'personal',
+      configurationVersion: '7',
+      url: config(scope ?? 'personal').canonicalUrl,
+      authorization: 'Bearer derived-only',
+    },
+  ]);
+  expect(JSON.stringify(result)).not.toContain('provider-secret');
+});
+it('returns an honest empty set and rejects substituted authority', async () => {
+  unavailable = true;
+  expect(await authorizeHarnessMcp('capability', invocation)).toEqual([]);
+  for (const patch of [
+    { userId: 'another-owner' },
+    { organizationId: 'another-context' },
+    { 'provider-secret': 'untrusted' },
+    { request: { name: 'kilo.organizations', arguments: {} } },
+  ]) {
+    await expect(
+      authorizeHarnessMcp('capability', { ...invocation, ...patch })
+    ).rejects.toMatchObject({
+      code: 'BAD_REQUEST',
+      message: 'invalid_input',
+      cause: undefined,
+    });
+  }
+  denied = true;
+  await expect(authorizeHarnessMcp('capability', invocation)).rejects.toMatchObject({
+    code: 'FORBIDDEN',
+    message: 'access_revoked',
+  });
+});
+it.each([
+  ['capability', { conversationId: invocation.operationId }],
+  ['capability', { operationId: invocation.conversationId }],
+  ['forged', {}],
+] as const)('rejects a forged capability: %s %j', async (token, patch) => {
+  await expect(authorizeHarnessMcp(token, { ...invocation, ...patch })).rejects.toMatchObject({
+    code: 'FORBIDDEN',
+    message: 'access_revoked',
+  });
+});
+it.each([
+  ['forbidden', 'reauthorization_required'],
+  ['invalid_grant', 'reauthorization_required'],
+  ['temporarily_unavailable', 'unavailable_server'],
+  ['not_found', 'unavailable_server'],
+  ['invalid_request', 'unsafe_destination'],
+  ['access_denied', 'access_revoked'],
+] as const)('sanitizes gateway %s without merging recovery states', async (code, message) => {
+  failure = new GatewayError(code, 'provider-secret', 403);
+  await expect(authorizeHarnessMcp('capability', invocation)).rejects.toMatchObject({
+    message,
+    cause: undefined,
+  });
+});
+it.each(['FORBIDDEN', 'UNAUTHORIZED'] as const)('sanitizes %s denial', async code => {
+  failureStage = 'authorization';
+  failure = new TRPCError({
+    code,
+    message: 'unavailable_server',
+    cause: new Error('provider-secret'),
+  });
+  await expect(authorizeHarnessMcp('capability', invocation)).rejects.toMatchObject({
+    code: 'FORBIDDEN',
+    message: 'access_revoked',
+    cause: undefined,
+  });
+});
+it.each([
+  'authorization',
+  'configuration',
+  'availability',
+  'route',
+  'mint',
+  'verification',
+] as const)('removes sensitive failures from %s', async stage => {
+  failureStage = stage;
+  for (const message of ['provider-secret', 'unavailable_server']) {
+    failure = new TRPCError({
+      code: 'SERVICE_UNAVAILABLE',
+      message,
+      cause: new Error('provider-secret'),
+    });
+    await expect(authorizeHarnessMcp('capability', invocation)).rejects.toMatchObject({
+      code: 'SERVICE_UNAVAILABLE',
+      message: 'unavailable_server',
+      cause: undefined,
+    });
+  }
+});
+it('reports missing gateway configuration rather than an empty discovery', async () => {
+  failureStage = 'configuration';
+  failure = new Error('provider-secret');
+  await expect(authorizeHarnessMcp('capability', invocation)).rejects.toMatchObject({
+    code: 'SERVICE_UNAVAILABLE',
+    message: 'unavailable_server',
+    cause: undefined,
+  });
+});
+it.each([
+  'http://gateway.example',
+  'https://provider-secret@gateway.example',
+  'https://user:provider-secret@gateway.example',
+])('refuses unsafe gateway destinations: %s', async url => {
+  gatewayBaseUrl = url;
+  await expect(authorizeHarnessMcp('capability', invocation)).rejects.toMatchObject({
+    code: 'BAD_REQUEST',
+    message: 'unsafe_destination',
+    cause: undefined,
+  });
+});
+it('refuses a destination that differs from the available canonical route', async () => {
+  canonicalOverride = 'https://other.example/personal';
+  await expect(authorizeHarnessMcp('capability', invocation)).rejects.toMatchObject({
+    message: 'unsafe_destination',
+  });
+});
+it('rejects configuration changes while minting authorization', async () => {
+  mintVersion = 2;
+  await expect(authorizeHarnessMcp('capability', invocation)).rejects.toMatchObject({
+    message: 'definition_changed',
+  });
+});
+it.each([
+  [null, 'personal', '1', null],
+  ['organization', 'organization', '1', null],
+  ['organization', 'personal', '1', 'unavailable_server'],
+  [null, 'other', '1', 'unavailable_server'],
+  [null, 'personal', '2', 'definition_changed'],
+])(
+  'checks the scoped server and version before a call: %s/%s/%s',
+  async (scope, serverId, configurationVersion, message) => {
+    organizationId = scope;
+    invocation.request = {
+      name: 'mcp.call',
+      arguments: {
+        serverId,
+        configurationVersion,
+        name: 'remote',
+        definitionVersion: 'immutable',
+        arguments: {},
+      },
+    };
+    const result = authorizeHarnessMcp('capability', invocation);
+    if (message) await expect(result).rejects.toMatchObject({ message });
+    else expect(await result).toHaveLength(1);
+  }
+);

apps/web/src/lib/agent-harness/mcp.ts

@@ -0,0 +1,113 @@
+import 'server-only';
+import { z } from 'zod';
+import { TRPCError } from '@trpc/server';
+import { GatewayError } from '@kilocode/mcp-gateway';
+import { toolDefinitions, ToolRequestSchema } from '@kilocode/agent-harness/tools';
+import { createGatewayServices } from '@/lib/mcp-gateway/services';
+import { authorizeHarnessCapability, harnessInputDigest } from './authorization';
+
+const Invocation = z.strictObject({
+  conversationId: z.uuid(),
+  operationId: z.uuid(),
+  request: ToolRequestSchema,
+});
+
+/** Internal server-to-server access only. Never project these ephemeral tokens into client or model state. */
+export async function authorizeHarnessMcp(token: string, input: unknown) {
+  const parsed = Invocation.safeParse(input);
+  if (!parsed.success) throw new TRPCError({ code: 'BAD_REQUEST', message: 'invalid_input' });
+  const { conversationId, operationId, request } = parsed.data;
+  if (request.name !== 'mcp.discover' && request.name !== 'mcp.call') {
+    throw new TRPCError({ code: 'BAD_REQUEST', message: 'invalid_input' });
+  }
+  const definition = toolDefinitions.find(tool => tool.name === request.name);
+  if (!definition) throw new TRPCError({ code: 'BAD_REQUEST', message: 'invalid_input' });
+  try {
+    const { authority } = await authorizeHarnessCapability(token, {
+      audience: 'agent-harness:operations',
+      conversationId,
+      operation: request.name,
+      definitionVersion: definition.version,
+      inputDigest: harnessInputDigest(request.arguments),
+      dispatchId: operationId,
+      target: { kind: 'backend' },
+    });
+    const executionContext =
+      authority.organizationId === null
+        ? { type: 'personal' as const }
+        : { type: 'organization' as const, organizationId: authority.organizationId };
+    const gateway = createGatewayServices();
+    const available = await gateway.availableService.listAvailableConfigs(
+      authority.userId,
+      executionContext
+    );
+    const selected =
+      request.name === 'mcp.discover'
+        ? available
+        : available.filter(config => config.configId === request.arguments.serverId);
+    if (request.name === 'mcp.call' && selected.length !== 1) {
+      throw new TRPCError({ code: 'NOT_FOUND', message: 'unavailable_server' });
+    }
+    const connections = [];
+    for (const config of selected) {
+      const { route, resolved } = await gateway.routeService.resolveResource(config.canonicalUrl);
+      const url = gateway.routeService.canonicalUrl(route);
+      const destination = new URL(url);
+      if (
+        url !== config.canonicalUrl ||
+        destination.protocol !== 'https:' ||
+        destination.username ||
+        destination.password ||
+        route.configId !== config.configId
+      ) {
+        throw new TRPCError({ code: 'BAD_REQUEST', message: 'unsafe_destination' });
+      }
+      const configurationVersion = String(resolved.config.config_version);
+      if (
+        request.name === 'mcp.call' &&
+        request.arguments.configurationVersion !== configurationVersion
+      ) {
+        throw new TRPCError({ code: 'PRECONDITION_FAILED', message: 'definition_changed' });
+      }
+      const access = await gateway.tokenService.mintDerivedConnectToken({
+        route,
+        userId: authority.userId,
+        executionContext,
+      });
+      const claims = await gateway.tokenService.verifyUserInfoToken(access.token);
+      if (String(claims.config_version) !== configurationVersion) {
+        throw new TRPCError({ code: 'PRECONDITION_FAILED', message: 'definition_changed' });
+      }
+      connections.push({
+        serverId: config.configId,
+        configurationVersion,
+        url,
+        authorization: `Bearer ${access.token}`,
+      });
+    }
+    return connections;
+  } catch (error) {
+    if (
+      (error instanceof GatewayError && error.code === 'access_denied') ||
+      (error instanceof TRPCError && (error.code === 'FORBIDDEN' || error.code === 'UNAUTHORIZED'))
+    ) {
+      throw new TRPCError({ code: 'FORBIDDEN', message: 'access_revoked' });
+    }
+    if (
+      error instanceof TRPCError &&
+      ['unavailable_server', 'unsafe_destination', 'definition_changed'].includes(error.message)
+    ) {
+      throw new TRPCError({ code: error.code, message: error.message });
+    }
+    if (error instanceof GatewayError && error.code === 'invalid_request') {
+      throw new TRPCError({ code: 'BAD_REQUEST', message: 'unsafe_destination' });
+    }
+    const reauthorize =
+      error instanceof GatewayError &&
+      (error.code === 'forbidden' || error.code === 'invalid_grant');
+    throw new TRPCError({
+      code: reauthorize ? 'PRECONDITION_FAILED' : 'SERVICE_UNAVAILABLE',
+      message: reauthorize ? 'reauthorization_required' : 'unavailable_server',
+    });
+  }
+}