Add a docs server based on the AI Assistant Vectorize index

apps/docs-autorag/src/index.ts

@@ -1,5 +1,6 @@
 import { McpAgent } from 'agents/mcp'
 
+import { createApiHandler } from '@repo/mcp-common/src/api-handler'
 import { getEnv } from '@repo/mcp-common/src/env'
 import { CloudflareMCPServer } from '@repo/mcp-common/src/server'
 
@@ -35,4 +36,4 @@ export class CloudflareDocumentationMCP extends McpAgent<Env, State, Props> {
 	}
 }
 
-export default CloudflareDocumentationMCP.mount('/sse')
+export default createApiHandler(CloudflareDocumentationMCP)

apps/docs-autorag/wrangler.jsonc

@@ -49,7 +49,7 @@
 		"staging": {
 			"name": "mcp-cloudflare-docs-autorag-staging",
 			"account_id": "6702657b6aa048cf3081ff3ff3c9c52f",
-			"routes": [{ "pattern": "docs-staging.mcp.cloudflare.com", "custom_domain": true }],
+			"routes": [{ "pattern": "docs-autorag-staging.mcp.cloudflare.com", "custom_domain": true }],
 			"durable_objects": {
 				"bindings": [
 					{
@@ -75,7 +75,7 @@
 		"production": {
 			"name": "mcp-cloudflare-docs-autorag-production",
 			"account_id": "6702657b6aa048cf3081ff3ff3c9c52f",
-			"routes": [{ "pattern": "docs.mcp.cloudflare.com", "custom_domain": true }],
+			"routes": [{ "pattern": "docs-autorag.mcp.cloudflare.com", "custom_domain": true }],
 			"durable_objects": {
 				"bindings": [
 					{

apps/docs-vectorize/.eslintrc.cjs

@@ -0,0 +1,5 @@
+/** @type {import("eslint").Linter.Config} */
+module.exports = {
+	root: true,
+	extends: ['@repo/eslint-config/default.cjs'],
+}

apps/docs-vectorize/README.md

@@ -0,0 +1,19 @@
+# Model Context Protocol (MCP) Server + Cloudflare Documentation (via Autorag)
+
+This is a [Model Context Protocol (MCP)](https://modelcontextprotocol.io/introduction) server that supports remote MCP connections. It connects to a Vectorize DB (in this case, indexed w/ the Cloudflare docs)
+
+The Cloudflare account this worker is deployed on already has this Vectorize DB setup and indexed.
+
+## Running locally
+
+```
+pnpm run start
+```
+
+Then connect to the server via remote MCP at `http://localhost:8976/sse`
+
+## Deploying
+
+```
+pnpm run deploy --env [ENVIRONMENT]
+```

apps/docs-vectorize/package.json

@@ -0,0 +1,34 @@
+{
+	"name": "docs-vectorize",
+	"version": "0.0.1",
+	"private": true,
+	"scripts": {
+		"check:lint": "run-eslint-workers",
+		"check:types": "run-tsc",
+		"deploy": "run-wrangler-deploy",
+		"dev": "wrangler dev --experimental-vectorize-bind-to-prod",
+		"start": "npm run dev",
+		"types": "wrangler types --include-env=false",
+		"test": "vitest run"
+	},
+	"dependencies": {
+		"@cloudflare/workers-oauth-provider": "0.0.3",
+		"@hono/zod-validator": "0.4.3",
+		"@modelcontextprotocol/sdk": "1.10.2",
+		"@repo/mcp-common": "workspace:*",
+		"@repo/mcp-observability": "workspace:*",
+		"agents": "0.0.67",
+		"cloudflare": "4.2.0",
+		"hono": "4.7.6",
+		"mime": "4.0.6",
+		"zod": "3.24.2"
+	},
+	"devDependencies": {
+		"@cloudflare/vitest-pool-workers": "0.8.14",
+		"@types/node": "22.14.1",
+		"prettier": "3.5.3",
+		"typescript": "5.5.4",
+		"vitest": "3.0.9",
+		"wrangler": "4.10.0"
+	}
+}

apps/docs-vectorize/src/context.ts

@@ -0,0 +1,11 @@
+import type { CloudflareDocumentationMCP } from './index'
+
+export interface Env {
+	ENVIRONMENT: 'development' | 'staging' | 'production'
+	MCP_SERVER_NAME: string
+	MCP_SERVER_VERSION: string
+	MCP_OBJECT: DurableObjectNamespace<CloudflareDocumentationMCP>
+	MCP_METRICS: AnalyticsEngineDataset
+	AI: Ai
+	VECTORIZE: VectorizeIndex
+}

apps/docs-vectorize/src/index.ts

@@ -0,0 +1,39 @@
+import { McpAgent } from 'agents/mcp'
+
+import { createApiHandler } from '@repo/mcp-common/src/api-handler'
+import { getEnv } from '@repo/mcp-common/src/env'
+import { CloudflareMCPServer } from '@repo/mcp-common/src/server'
+
+import { registerDocsTools } from './tools/docs'
+
+import type { Env } from './context'
+
+const env = getEnv<Env>()
+
+// The docs MCP server isn't stateful, so we don't have state/props
+export type Props = never
+
+export type State = never
+
+export class CloudflareDocumentationMCP extends McpAgent<Env, State, Props> {
+	server = new CloudflareMCPServer({
+		wae: env.MCP_METRICS,
+		serverInfo: {
+			name: env.MCP_SERVER_NAME,
+			version: env.MCP_SERVER_VERSION,
+		},
+	})
+
+	constructor(
+		public ctx: DurableObjectState,
+		public env: Env
+	) {
+		super(ctx, env)
+	}
+
+	async init() {
+		registerDocsTools(this)
+	}
+}
+
+export default createApiHandler(CloudflareDocumentationMCP)

apps/docs-vectorize/src/tools/docs.ts

@@ -0,0 +1,113 @@
+import { z } from 'zod'
+
+import type { CloudflareDocumentationMCP } from '../index'
+
+// Always return 10 results for simplicity, don't make it configurable
+const TOP_K = 10
+
+/**
+ * Registers the docs search tool with the MCP server
+ * @param agent The MCP server instance
+ */
+export function registerDocsTools(agent: CloudflareDocumentationMCP) {
+	// Register the worker logs analysis tool by worker name
+	agent.server.tool(
+		'search_cloudflare_documentation',
+		`Search the Cloudflare documentation.
+
+		This tool should be used to answer any question about Cloudflare products or features, including:
+		- Workers, Pages, R2, Images, Stream, D1, Durable Objects, KV, Workflows, Hyperdrive, Queues
+		- AutoRAG, Workers AI, Vectorize, AI Gateway, Browser Rendering
+		- Zero Trust, Access, Tunnel, Gateway, Browser Isolation, WARP, DDOS, Magic Transit, Magic WAN
+		- CDN, Cache, DNS, Zaraz, Argo, Rulesets, Terraform, Account and Billing
+
+		Results are returned as semantically similar chunks to the query.
+		`,
+		{
+			query: z.string(),
+		},
+		async ({ query }) => {
+			const results = await queryVectorize(agent.env.AI, agent.env.VECTORIZE, query, TOP_K)
+			const resultsAsXml = results
+				.map((result) => {
+					return `<result>
+<url>${result.url}</url>
+<text>
+${result.text}
+</text>
+</result>`
+				})
+				.join('\n')
+			return {
+				content: [{ type: 'text', text: resultsAsXml }],
+			}
+		}
+	)
+}
+
+async function queryVectorize(ai: Ai, vectorizeIndex: VectorizeIndex, query: string, topK: number) {
+	// Recommendation from: https://huggingface.co/BAAI/bge-base-en-v1.5#model-list
+	const [queryEmbedding] = await getEmbeddings(ai, [
+		'Represent this sentence for searching relevant passages: ' + query,
+	])
+
+	const { matches } = await vectorizeIndex.query(queryEmbedding, {
+		topK,
+		returnMetadata: 'all',
+		returnValues: false,
+	})
+
+	return matches.map((match, _i) => ({
+		similarity: Math.min(match.score, 1),
+		id: match.id,
+		url: sourceToUrl(String(match.metadata?.filePath ?? '')),
+		text: String(match.metadata?.text ?? ''),
+	}))
+}
+
+const TOP_DIR = 'src/content/docs'
+function sourceToUrl(path: string) {
+	const prefix = `${TOP_DIR}/`
+	return (
+		'https://developers.cloudflare.com/' +
+		(path.startsWith(prefix) ? path.slice(prefix.length) : path)
+			.replace(/index\.mdx$/, '')
+			.replace(/\.mdx$/, '')
+	)
+}
+
+async function getEmbeddings(ai: Ai, strings: string[]) {
+	const response = await doWithRetries(() =>
+		ai.run('@cf/baai/bge-base-en-v1.5', {
+			text: strings,
+			// @ts-expect-error pooling not in types yet
+			pooling: 'cls',
+		})
+	)
+
+	return response.data
+}
+
+/**
+ * @template T
+ * @param {() => Promise<T>} action
+ */
+async function doWithRetries<T>(action: () => Promise<T>) {
+	const NUM_RETRIES = 10
+	const INIT_RETRY_MS = 50
+	for (let i = 0; i <= NUM_RETRIES; i++) {
+		try {
+			return await action()
+		} catch (e) {
+			// TODO: distinguish between user errors (4xx) and system errors (5xx)
+			console.error(e)
+			if (i === NUM_RETRIES) {
+				throw e
+			}
+			// Exponential backoff with full jitter
+			await scheduler.wait(Math.random() * INIT_RETRY_MS * Math.pow(2, i))
+		}
+	}
+	// Should never reach here – last loop iteration should return
+	throw new Error('An unknown error occurred')
+}

apps/docs-vectorize/tsconfig.json

@@ -0,0 +1,4 @@
+{
+	"extends": "@repo/typescript-config/workers.json",
+	"include": ["*/**.ts", "./vitest.config.ts"]
+}

apps/docs-vectorize/vitest.config.ts

@@ -0,0 +1,24 @@
+import { defineWorkersConfig } from '@cloudflare/vitest-pool-workers/config'
+
+import type { Env } from './src/context'
+
+export interface TestEnv extends Env {
+	CLOUDFLARE_MOCK_ACCOUNT_ID: string
+	CLOUDFLARE_MOCK_API_TOKEN: string
+}
+
+export default defineWorkersConfig({
+	test: {
+		poolOptions: {
+			workers: {
+				wrangler: { configPath: `${__dirname}/wrangler.jsonc` },
+				miniflare: {
+					bindings: {
+						CLOUDFLARE_MOCK_ACCOUNT_ID: 'mock-account-id',
+						CLOUDFLARE_MOCK_API_TOKEN: 'mock-api-token',
+					} satisfies Partial<TestEnv>,
+				},
+			},
+		},
+	},
+})