|
1 | | -import OAuthProvider from '@cloudflare/workers-oauth-provider' |
2 | | -import { McpAgent } from 'agents/mcp' |
3 | | - |
4 | | -import { AccountManager } from '@repo/mcp-common/src/account-manager' |
5 | | -import { handleApiTokenMode, isApiTokenRequest } from '@repo/mcp-common/src/api-token-mode' |
6 | | -import { |
7 | | - createAuthHandlers, |
8 | | - handleTokenExchangeCallback, |
9 | | -} from '@repo/mcp-common/src/cloudflare-oauth-handler' |
10 | | -import { getEnv } from '@repo/mcp-common/src/env' |
11 | | -import { getProps } from '@repo/mcp-common/src/get-props' |
| 1 | +import { createAuthenticatedMcpApp } from '@repo/mcp-common/src/mcp-app' |
12 | 2 | import { RequiredScopes } from '@repo/mcp-common/src/scopes' |
13 | | -import { CloudflareMCPServer } from '@repo/mcp-common/src/server' |
14 | 3 |
|
15 | | -import { MetricsTracker } from '../../../packages/mcp-observability/src' |
16 | 4 | import { registerAIGatewayTools } from './tools/ai-gateway.tools' |
17 | 5 |
|
18 | | -import type { AuthProps } from '@repo/mcp-common/src/cloudflare-oauth-handler' |
19 | 6 | import type { Env } from './ai-gateway.context' |
20 | 7 |
|
21 | | -const env = getEnv<Env>() |
22 | | - |
23 | | -const metrics = new MetricsTracker(env.MCP_METRICS, { |
24 | | - name: env.MCP_SERVER_NAME, |
25 | | - version: env.MCP_SERVER_VERSION, |
26 | | -}) |
27 | | - |
28 | | -// Context from the auth process, encrypted & stored in the auth token |
29 | | -// and provided to the DurableMCP as this.props |
30 | | -type Props = AuthProps |
31 | | -type State = Record<string, never> |
32 | | - |
33 | | -export class AIGatewayMCP extends McpAgent<Env, State, Props> { |
34 | | - _server: CloudflareMCPServer | undefined |
35 | | - set server(server: CloudflareMCPServer) { |
36 | | - this._server = server |
37 | | - } |
38 | | - get server(): CloudflareMCPServer { |
39 | | - if (!this._server) { |
40 | | - throw new Error('Tried to access server before it was initialized') |
41 | | - } |
42 | | - |
43 | | - return this._server |
44 | | - } |
45 | | - |
46 | | - constructor(ctx: DurableObjectState, env: Env) { |
47 | | - super(ctx, env) |
48 | | - } |
49 | | - |
50 | | - async init() { |
51 | | - // TODO: Probably we'll want to track account tokens usage through an account identifier at some point |
52 | | - const props = getProps(this) |
53 | | - const userId = props.type === 'user_token' ? props.user.id : undefined |
54 | | - const accountManager = new AccountManager(props) |
55 | | - |
56 | | - this.server = new CloudflareMCPServer({ |
57 | | - userId, |
58 | | - wae: this.env.MCP_METRICS, |
59 | | - serverInfo: { |
60 | | - name: this.env.MCP_SERVER_NAME, |
61 | | - version: this.env.MCP_SERVER_VERSION, |
62 | | - }, |
63 | | - accountManager, |
64 | | - options: { instructions: accountManager.instructionsSuffix() }, |
65 | | - }) |
66 | | - |
67 | | - // Register Cloudflare Log Push tools |
68 | | - registerAIGatewayTools(this) |
69 | | - } |
70 | | -} |
71 | | - |
72 | 8 | const AIGatewayScopes = { |
73 | 9 | ...RequiredScopes, |
74 | 10 | 'account:read': 'See your account info such as account details, analytics, and memberships.', |
75 | 11 | 'aig:read': 'Grants read level access to AI Gateway.', |
76 | 12 | } as const |
77 | 13 |
|
78 | | -export default { |
79 | | - fetch: async (req: Request, env: Env, ctx: ExecutionContext) => { |
80 | | - if (await isApiTokenRequest(req, env)) { |
81 | | - return await handleApiTokenMode(AIGatewayMCP, req, env, ctx) |
82 | | - } |
| 14 | +const app = createAuthenticatedMcpApp<Env>({ |
| 15 | + serviceHostnames: ['ai-gateway-staging.mcp.cloudflare.com', 'ai-gateway.mcp.cloudflare.com'], |
| 16 | + scopes: AIGatewayScopes, |
| 17 | + register: registerAIGatewayTools, |
| 18 | +}) |
| 19 | + |
| 20 | +export const mcpHandler = app.mcpHandler |
83 | 21 |
|
84 | | - return new OAuthProvider({ |
85 | | - apiHandlers: { |
86 | | - '/mcp': AIGatewayMCP.serve('/mcp'), |
87 | | - '/sse': AIGatewayMCP.serveSSE('/sse'), |
88 | | - }, |
89 | | - defaultHandler: createAuthHandlers({ scopes: AIGatewayScopes, metrics }), |
90 | | - authorizeEndpoint: '/oauth/authorize', |
91 | | - tokenEndpoint: '/token', |
92 | | - tokenExchangeCallback: (options) => |
93 | | - handleTokenExchangeCallback( |
94 | | - options, |
95 | | - env.CLOUDFLARE_CLIENT_ID, |
96 | | - env.CLOUDFLARE_CLIENT_SECRET |
97 | | - ), |
98 | | - // Cloudflare access token TTL |
99 | | - accessTokenTTL: 3600, |
100 | | - refreshTokenTTL: 2592000, // 30 days |
101 | | - // TODO: Remove after 2026-05-01 — all pre-0.4.0 grants will have expired by then |
102 | | - resourceMatchOriginOnly: true, |
103 | | - clientRegistrationEndpoint: '/register', |
104 | | - }).fetch(req, env, ctx) |
105 | | - }, |
106 | | -} |
| 22 | +export default app.worker |
0 commit comments