code-mode.ts

import { DynamicWorkerExecutor } from "@cloudflare/codemode";
import { openApiMcpServer } from "@cloudflare/codemode/mcp";
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";

import {
  createFetchYnabApiClient,
  createYnabOpenApiRequestHandler,
  type YnabApiClient,
} from "./ynab-client.ts";
import ynabOpenApi from "./ynab-openapi.json";

export interface CreateYnabCodeModeServerOptions {
  client: YnabApiClient;
  executor: DynamicWorkerExecutor;
}

export function createYnabCodeModeServer(options: CreateYnabCodeModeServerOptions): McpServer {
  return openApiMcpServer({
    spec: ynabOpenApi,
    executor: options.executor,
    request: createYnabOpenApiRequestHandler(options.client),
  });
}

export function createYnabCodeModeServerFromEnv(env: Env): McpServer {
  const executor = new DynamicWorkerExecutor({
    loader: env.LOADER,
  });
  const client = createFetchYnabApiClient({
    accessToken: env.YNAB_ACCESS_TOKEN,
    fetch: globalThis.fetch.bind(globalThis),
  });

  return createYnabCodeModeServer({ client, executor });
}

ynab-mcp-agent.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;
  },
};