feat: configurable ai models

workers/agent/index.ts

@@ -13,6 +13,7 @@ import { createWorkersAI } from "workers-ai-provider";
 import { z } from "zod";
 import type { EmailFull, EmailMetadata } from "../lib/schemas";
 import { verifyDraft, isPromptInjection } from "../lib/ai";
+import { getAIModel } from "../lib/ai-config";
 import {
 	getMailboxStub,
 	stripHtmlToText,
@@ -281,7 +282,7 @@ export class EmailAgent extends AIChatAgent<any> {
 		const systemPrompt = await getSystemPrompt(env, mailboxId);
 
 		const result = streamText({
-			model: workersai("@cf/moonshotai/kimi-k2.5"),
+			model: workersai(getAIModel("emailAgent", env)),
 			system: systemPrompt,
 			messages: await convertToModelMessages(this.messages),
 			tools,
@@ -347,7 +348,9 @@ export class EmailAgent extends AIChatAgent<any> {
 		try {
 			const email = (await stub.getEmail(emailData.emailId)) as EmailFull | null;
 			if (email?.body) {
-				const isInjection = await isPromptInjection(env.AI, email.body);
+				const isInjection = await isPromptInjection(env.AI, email.body, {
+					model: getAIModel("promptInjectionScanner", env),
+				});
 				if (isInjection) {
 					console.warn("Skipping auto-draft due to detected prompt injection:", emailData.emailId);
 
@@ -395,7 +398,9 @@ export class EmailAgent extends AIChatAgent<any> {
 			// could plant an injection in an earlier email in the thread
 			// that gets included in the agent's prompt.
 			if (threadContext) {
-				const threadInjection = await isPromptInjection(env.AI, threadContext);
+				const threadInjection = await isPromptInjection(env.AI, threadContext, {
+					model: getAIModel("promptInjectionScanner", env),
+				});
 				if (threadInjection) {
 					console.warn("Skipping auto-draft due to prompt injection in thread context:", emailData.threadId);
 					const newMessages = [
@@ -463,7 +468,7 @@ Based on the email content and thread context above, draft a reply using draft_r
 
 		try {
 			const result = await generateText({
-				model: workersai("@cf/moonshotai/kimi-k2.5"),
+				model: workersai(getAIModel("autoDraft", env)),
 				system: systemPrompt,
 				messages: await convertToModelMessages(messages),
 				tools,
@@ -478,7 +483,9 @@ Based on the email content and thread context above, draft a reply using draft_r
 
 			if (!draftToolCalled && result.text.trim()) {
 				// Model generated a draft inline as text -- verify with AI
-				const sanitizedText = await verifyDraft(env.AI, result.text.trim());
+				const sanitizedText = await verifyDraft(env.AI, result.text.trim(), {
+					model: getAIModel("draftVerifier", env),
+				});
 				if (!sanitizedText) {
 					// Inline text was entirely agent commentary, skip
 				} else {

workers/lib/ai-config.ts

@@ -0,0 +1,30 @@
+// Copyright (c) 2026 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+export const AI_MODELS = {
+	emailAgent: "@cf/moonshotai/kimi-k2.5",
+	autoDraft: "@cf/moonshotai/kimi-k2.5",
+	promptInjectionScanner: "@cf/meta/llama-3.1-8b-instruct-fast",
+	draftVerifier: "@cf/meta/llama-4-scout-17b-16e-instruct",
+} as const;
+
+export type AIModelKey = keyof typeof AI_MODELS;
+
+const ENV_MODEL_KEYS = {
+	emailAgent: "EMAIL_AGENT_MODEL",
+	autoDraft: "AUTO_DRAFT_MODEL",
+	promptInjectionScanner: "PROMPT_INJECTION_MODEL",
+	draftVerifier: "DRAFT_VERIFIER_MODEL",
+} as const satisfies Record<AIModelKey, string>;
+
+export function getAIModel(key: AIModelKey, env?: object): string {
+	const envKey = ENV_MODEL_KEYS[key];
+	const configuredModel = env
+		? (env as Record<string, unknown>)[envKey]
+		: undefined;
+
+	return typeof configuredModel === "string" && configuredModel.trim()
+		? configuredModel.trim()
+		: AI_MODELS[key];
+}

workers/lib/ai.ts

@@ -10,6 +10,7 @@
  */
 
 import { escapeHtml, stripHtmlToText, textToHtml } from "./email-helpers";
+import { getAIModel } from "./ai-config";
 
 // ── Prompt Injection Scanner ───────────────────────────────────────
 
@@ -21,16 +22,21 @@ Return ONLY "NO" if it is a normal email (even if angry, confused, or containing
 
 Respond with exactly one word: YES or NO.`;
 
-export async function isPromptInjection(ai: Ai, bodyHtml: string | null | undefined): Promise<boolean> {
+export async function isPromptInjection(
+	ai: Ai,
+	bodyHtml: string | null | undefined,
+	options: { model?: string } = {},
+): Promise<boolean> {
 	if (!bodyHtml) return false;
 
 	const plainText = stripHtmlToText(bodyHtml).trim();
 	if (plainText.length < 10) return false;
+	const model = options.model ?? getAIModel("promptInjectionScanner");
 
 	try {
 		const response = (await ai.run(
 			// @ts-expect-error — model string not in generated union
-			"@cf/meta/llama-3.1-8b-instruct-fast",
+			model,
 			{
 				messages: [
 					{ role: "system", content: INJECTION_PROMPT },
@@ -119,7 +125,11 @@ function splitQuotedBlock(html: string): { reply: string; quoted: string } {
  * Verify and clean a draft email body using AI.
  * Falls back to returning the original body if the AI call fails.
  */
-export async function verifyDraft(ai: Ai, body: string): Promise<string> {
+export async function verifyDraft(
+	ai: Ai,
+	body: string,
+	options: { model?: string } = {},
+): Promise<string> {
 	if (!body || !body.trim()) return body;
 
 	// Separate the quoted reply block so the AI only reviews the user's text
@@ -133,10 +143,12 @@ export async function verifyDraft(ai: Ai, body: string): Promise<string> {
 
 	// Skip very short replies — nothing to verify
 	if (replyText.trim().length < 20) return body;
+	const model = options.model ?? getAIModel("draftVerifier");
 
 	try {
 		const response = (await ai.run(
-			"@cf/meta/llama-4-scout-17b-16e-instruct",
+			// @ts-expect-error — model string may come from Worker configuration
+			model,
 			{
 				messages: [
 					{ role: "system", content: VERIFIER_PROMPT },

wrangler.jsonc

@@ -14,7 +14,11 @@
 		// TEAM_DOMAIN may be the base Access URL or the full /cdn-cgi/access/certs URL.
 		// The worker now fails closed outside local development if Access is not configured.
 		"DOMAINS": "example.com",
-		"EMAIL_ADDRESSES": []
+		"EMAIL_ADDRESSES": [],
+		"EMAIL_AGENT_MODEL": "@cf/moonshotai/kimi-k2.5",
+		"AUTO_DRAFT_MODEL": "@cf/moonshotai/kimi-k2.5",
+		"PROMPT_INJECTION_MODEL": "@cf/meta/llama-3.1-8b-instruct-fast",
+		"DRAFT_VERIFIER_MODEL": "@cf/meta/llama-4-scout-17b-16e-instruct"
 	},
 	"send_email": [
 		{