code-mode.ts

| import type { AgentContext } from "agents"; | | import { McpAgent } from "agents/mcp"; | | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; | | | | import { Effect, Redacted } from "effect"; | | import { | | CloudflareAccessConfig, | | CloudflareAccessInvalidError, | | CloudflareAccessJwtVerifier, | | CloudflareAccessKeyResolver, | | CloudflareAccessMisconfiguredError, | | CloudflareAccessRequiredError, | | CloudflareAccessUnsupportedIdentityError, | | type CloudflareAccessPrincipal, | | } from "@mulroy/cloudflare-access"; | | | | import { createYnabCodeModeServerFromEnv } from "./ynab-api.ts"; | | | | function getYnabMcpRequestLogContext(request: Request): Record<string, string | null> { | | return { | | accept: request.headers.get("accept"), | | cfMcpMethod: request.headers.get("cf-mcp-method"), | | cfRay: request.headers.get("cf-ray"), | | cfWorker: request.headers.get("cf-worker"), | | contentLength: request.headers.get("content-length"), | | method: request.method, | | sessionId: request.headers.get("mcp-session-id"), | | upgrade: request.headers.get("upgrade"), | | userAgent: request.headers.get("user-agent"), | | }; | | } | | | | async function requireYnabMcpAccess( | | request: Request, | | env: Env, | | ): Promise<CloudflareAccessPrincipal | Response> { | | try { | | const principal = await Effect.runPromise( | | Effect.gen(function* () { | | const verifier = yield* CloudflareAccessJwtVerifier; | | return yield* verifier.verify( | | Redacted.make(request.headers.get("cf-access-jwt-assertion") ?? ""), | | ); | | }).pipe( | | Effect.provide(CloudflareAccessJwtVerifier.layer), | | Effect.provide(CloudflareAccessKeyResolver.layerRemoteJwks), | | Effect.provide(CloudflareAccessConfig.layerFromEnv(env)), | | ), | | ); | | | | return principal; | | } catch (error) { | | if ( | | error instanceof CloudflareAccessRequiredError || | | error instanceof CloudflareAccessInvalidError || | | error instanceof CloudflareAccessUnsupportedIdentityError || | | error instanceof CloudflareAccessMisconfiguredError | | ) { | | return Response.json( | | { error: error.error, message: error.message }, | | { status: error.status }, | | ); | | } | | | | return Response.json( | | { | | error: "cloudflare_access_invalid", | | message: "Unexpected Cloudflare Access verification failure", | | }, | | { status: 500 }, | | ); | | } | | } | | | | export class YnabMcpAgent extends McpAgent { | | private readonly runtimeEnv: Env; | | | | server = new McpServer({ | | name: "for-ynab-api", | | version: "1.0.0", | | }); | | | | constructor(ctx: AgentContext, env: Env) { | | super(ctx, env); | | this.runtimeEnv = env; | | } | | | | async init(): Promise<void> { | | console.log("ynab mcp agent init", { | | agentName: this.name, | | sessionId: this.getSessionId(), | | transportType: this.getTransportType(), | | }); | | | | this.server = createYnabCodeModeServerFromEnv(this.runtimeEnv); | | } | | | | async onConnect(...args: Parameters<McpAgent["onConnect"]>): Promise<void> { | | const [connection, context] = args; | | | | console.log("ynab mcp agent websocket connect", { | | activeConnectionCount: Array.from(this.getConnections()).length, | | agentName: this.name, | | cfMcpMethod: context.request.headers.get("cf-mcp-method"), | | connectionId: connection.id, | | requestSessionId: context.request.headers.get("mcp-session-id"), | | sessionId: this.getSessionId(), | | transportType: this.getTransportType(), | | upgrade: context.request.headers.get("upgrade"), | | }); | | | | await super.onConnect(...args); | | } | | | | async onClose(...args: Parameters<McpAgent["onClose"]>): Promise<void> { | | const [connection, code, reason, wasClean] = args; | | | | console.log("ynab mcp agent websocket close", { | | activeConnectionCount: Array.from(this.getConnections()).length, | | agentName: this.name, | | code, | | connectionId: connection.id, | | reason, | | sessionId: this.getSessionId(), | | transportType: this.getTransportType(), | | wasClean, | | }); | | | | await super.onClose(...args); | | } | | } | | | | const mcpHandler = YnabMcpAgent.serve("/s/ynab", { binding: "YnabMcpAgent" }); | | | | export default { | | async fetch(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> { | | const url = new URL(request.url); | | | | if (url.pathname === "/health") { | | return Response.json({ ok: true }); | | } | | | | if (url.pathname.startsWith("/s/ynab")) { | | const requestLogContext = getYnabMcpRequestLogContext(request); | | | | console.log("ynab mcp request received", requestLogContext); | | | | const result = await requireYnabMcpAccess(request, env); | | | | if (result instanceof Response) { | | console.warn("ynab mcp access denied", { | | ...requestLogContext, | | status: result.status, | | }); | | | | return result; | | } | | | | console.log("MCP access granted", { | | ...requestLogContext, | | principalKind: result.kind, | | }); | | } | | | | const handlerStartedAt = Date.now(); | | const response = await mcpHandler.fetch(request, env, ctx); | | | | if (url.pathname.startsWith("/s/ynab")) { | | console.log("ynab mcp handler response", { | | durationMs: Date.now() - handlerStartedAt, | | method: request.method, | | requestSessionId: request.headers.get("mcp-session-id"), | | responseContentType: response.headers.get("content-type"), | | responseSessionId: response.headers.get("mcp-session-id"), | | status: response.status, | | }); | | } | | | | return response; | | }, | | }; |