Building Production MCP Servers on Cloudflare Workers, KV & Durable Objects

// src/worker.ts
import { DurableObjectNamespace } from '@cloudflare/workers-types';

interface Env {
  AGENT_SESSION: DurableObjectNamespace;
  TOOL_CACHE: KVNamespace;
  // Cloudflare Access configuration, e.g., AUDIENCE_TAG
  AUDIENCE_TAG: string;
}

export default {
  async fetch(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> {
    const url = new URL(request.url);

    // mTLS and Cloudflare Access validation
    const clientCert = request.headers.get('CF-Client-Cert');
    if (!clientCert) {
      return new Response('mTLS certificate required', { status: 401 });
    }
    // Further validation of clientCert against expected CAs or subject DNs
    // For Cloudflare Access, validate JWT from CF-Access-Jwt-Assertion header
    const jwt = request.headers.get('CF-Access-Jwt-Assertion');
    if (!jwt || !await validateCloudflareAccessJWT(jwt, env.AUDIENCE_TAG)) {
      return new Response('Unauthorized', { status: 401 });
    }

    if (url.pathname.startsWith('/agent/stream/')) {
      const agentId = url.pathname.split('/')[3];
      if (!agentId) {
        return new Response('Agent ID required', { status: 400 });
      }

      // Get Durable Object for the agent session
      const id = env.AGENT_SESSION.idFromName(agentId);
      const stub = env.AGENT_SESSION.get(id);

      // Forward the request to the Durable Object
      return stub.fetch(request);
    }

    // Handle JSON-RPC 2.0 or other requests
    if (url.pathname.startsWith('/agent/rpc/')) {
      const agentId = url.pathname.split('/')[3];
      if (!agentId) {
        return new Response('Agent ID required', { status: 400 });
      }

      const id = env.AGENT_SESSION.idFromName(agentId);
      const stub = env.AGENT_SESSION.get(id);
      return stub.fetch(request);
    }

    return new Response('Not Found', { status: 404 });
  },
};

// Placeholder for Cloudflare Access JWT validation
// In a production environment, use a library or Cloudflare's own validation endpoint
async function validateCloudflareAccessJWT(jwt: string, audienceTag: string): Promise<boolean> {
  // Example: Fetch JWKS from Cloudflare Access and verify signature and audience
  // This is a simplified example. Production code requires robust JWT validation.
  try {
    const response = await fetch(`https://your-team.cloudflareaccess.com/cdn-cgi/access/certs`);
    const jwks = await response.json();
    // ... perform JWT validation using a library like 'jose' ...
    // Check 'aud' claim against audienceTag
    return true; // Placeholder
  } catch (e) {
    console.error('JWT validation failed:', e);
    return false;
  }
}

// src/durable_object.ts
import { DurableObjectState, DurableObjectEnv } from '@cloudflare/workers-types';

interface AgentSessionState {
  // Example: current model context, tool states, etc.
  context: string[];
  toolResults: Record<string, any>;
}

export class AgentSession {
  state: DurableObjectState;
  env: DurableObjectEnv;
  private sessions: WebSocket[] = []; // For WebSocket-based streaming (alternative to SSE)
  private sseControllers: ReadableStreamDefaultController<Uint8Array>[] = []; // For SSE

  constructor(state: DurableObjectState, env: DurableObjectEnv) {
    this.state = state;
    this.env = env;
    this.state.blockConcurrencyWhile(async () => {
      // Initialize state from storage if available
      const storedState = await this.state.storage.get<AgentSessionState>('agentState');
      if (storedState) {
        // Restore state
      } else {
        // Initialize new state
        await this.state.storage.put('agentState', { context: [], toolResults: {} });
      }
    });
  }

  async fetch(request: Request): Promise<Response> {
    const url = new URL(request.url);

    if (url.pathname.endsWith('/stream')) {
      // Handle SSE connection
      const { readable, writable } = new TransformStream();
      const writer = writable.getWriter();

      // SSE headers
      const headers = {
        'Content-Type': 'text/event-stream',
        'Cache-Control': 'no-cache',
        'Connection': 'keep-alive',
      };

      // Create a controller to push data
      const controller = new ReadableStreamDefaultController<Uint8Array>();
      this.sseControllers.push(controller);

      // Send initial connection message
      await writer.write(new TextEncoder().encode('event: connected\ndata: {}\n\n'));

      // Clean up on disconnect
      request.signal.addEventListener('abort', () => {
        const index = this.sseControllers.indexOf(controller);
        if (index > -1) {
          this.sseControllers.splice(index, 1);
        }
        writer.close();
      });

      return new Response(new ReadableStream({
        start(c) {
          Object.assign(controller, c); // Assign the actual controller
        },
        cancel() {
          const index = this.sseControllers.indexOf(controller);
          if (index > -1) {
            this.sseControllers.splice(index, 1);
          }
        }
      }), { headers });

    } else if (url.pathname.endsWith('/rpc')) {
      // Handle JSON-RPC 2.0 requests
      if (request.method !== 'POST') {
        return new Response('Method Not Allowed', { status: 405 });
      }

      try {
        const rpcRequest = await request.json() as { jsonrpc: string; method: string; params?: any; id?: string | number; };

        if (rpcRequest.jsonrpc !== '2.0') {
          return this.jsonRpcErrorResponse(-32600, 'Invalid Request', rpcRequest.id);
        }

        let rpcResponse: any;
        switch (rpcRequest.method) {
          case 'agent.executeTool':
            rpcResponse = await this.executeTool(rpcRequest.params, rpcRequest.id);
            break;
          case 'agent.getContext':
            rpcResponse = await this.getContext(rpcRequest.id);
            break;
          case 'agent.sendMessage':
            rpcResponse = await this.sendMessage(rpcRequest.params, rpcRequest.id);
            break;
          default:
            rpcResponse = this.jsonRpcErrorResponse(-32601, 'Method not found', rpcRequest.id);
        }

        // Push updates to SSE clients if applicable
        if (rpcRequest.method === 'agent.sendMessage') {
          this.broadcastSSE('agentMessage', rpcResponse.result);
        }

        return new Response(JSON.stringify(rpcResponse), {
          headers: { 'Content-Type': 'application/json' },
        });

      } catch (e: any) {
        console.error('JSON-RPC error:', e);
        return this.jsonRpcErrorResponse(-32700, 'Parse error', null);
      }
    }

    return new Response('Not Found', { status: 404 });
  }

  private async executeTool(params: any, id?: string | number): Promise<any> {
    const { toolName, args } = params;
    const cacheKey = `tool:${toolName}:${JSON.stringify(args)}`;

    // Check KV cache first
    const cachedResult = await this.env.TOOL_CACHE.get(cacheKey);
    if (cachedResult) {
      this.broadcastSSE('toolCacheHit', { toolName, args, result: JSON.parse(cachedResult) });
      return { jsonrpc: '2.0', result: JSON.parse(cachedResult), id };
    }

    // Simulate tool execution
    console.log(`Executing tool: ${toolName} with args:`, args);
    await new Promise(resolve => setTimeout(resolve, 500)); // Simulate work

    const result = {
      output: `Result of ${toolName} with args ${JSON.stringify(args)}`,
      timestamp: Date.now(),
    };

    // Store result in KV cache (e.g., for 1 hour)
    await this.env.TOOL_CACHE.put(cacheKey, JSON.stringify(result), { expirationTtl: 3600 });

    // Update Durable Object state
    const agentState = await this.state.storage.get<AgentSessionState>('agentState');
    if (agentState) {
      agentState.toolResults[toolName] = result;
      await this.state.storage.put('agentState', agentState);
    }

    this.broadcastSSE('toolExecuted', { toolName, args, result });
    return { jsonrpc: '2.0', result, id };
  }

  private async getContext(id?: string | number): Promise<any> {
    const agentState = await this.state.storage.get<AgentSessionState>('agentState');
    return { jsonrpc: '2.0', result: agentState?.context || [], id };
  }

  private async sendMessage(params: any, id?: string | number): Promise<any> {
    const { message } = params;
    // Simulate model processing and response
    console.log('Agent received message:', message);
    await new Promise(resolve => setTimeout(resolve, 200));

    const modelResponse = `Echo: ${message}`;

    // Update Durable Object state
    const agentState = await this.state.storage.get<AgentSessionState>('agentState');
    if (agentState) {
      agentState.context.push(`User: ${message}`);
      agentState.context.push(`Model: ${modelResponse}`);
      await this.state.storage.put('agentState', agentState);
    }

    this.broadcastSSE('modelResponse', { message: modelResponse });
    return { jsonrpc: '2.0', result: { response: modelResponse }, id };
  }

  private jsonRpcErrorResponse(code: number, message: string, id: string | number | null): any {
    return {
      jsonrpc: '2.0',
      error: { code, message },
      id,
    };
  }

  private broadcastSSE(event: string, data: any) {
    const encoder = new TextEncoder();
    const message = `event: ${event}\ndata: ${JSON.stringify(data)}\n\n`;
    const encodedMessage = encoder.encode(message);

    this.sseControllers.forEach(controller => {
      try {
        controller.enqueue(encodedMessage);
      } catch (e) {
        console.error('Failed to enqueue SSE message:', e);
        // Consider removing disconnected controllers here
      }
    });
  }
}

// wrangler.toml
name = "mcp-worker"
main = "src/worker.ts"
compatibility_date = "2024-01-01"
compatibility_flags = ["nodejs_compat"]

[[kv_namespaces]]
binding = "TOOL_CACHE"
id = "YOUR_KV_NAMESPACE_ID" # Replace with your KV namespace ID

[[durable_objects.bindings]]
name = "AGENT_SESSION"
class_name = "AgentSession"
script_name = "mcp-worker" # Refers to the worker itself, where the DO class is defined

// Inside AgentSession.executeTool method
const cacheKey = `tool:${toolName}:${JSON.stringify(args)}`;
const cachedResult = await this.env.TOOL_CACHE.get(cacheKey);
if (cachedResult) {
  // ... return cached result ...
}
// ... execute tool, then cache result ...
await this.env.TOOL_CACHE.put(cacheKey, JSON.stringify(result), { expirationTtl: 3600 });

// src/worker.ts (excerpt)
// mTLS and Cloudflare Access validation
const clientCert = request.headers.get('CF-Client-Cert');
if (!clientCert) {
  return new Response('mTLS certificate required', { status: 401 });
}
// Further validation of clientCert against expected CAs or subject DNs
const jwt = request.headers.get('CF-Access-Jwt-Assertion');
if (!jwt || !await validateCloudflareAccessJWT(jwt, env.AUDIENCE_TAG)) {
  return new Response('Unauthorized', { status: 401 });
}