@kentcdodds/ai

import { readNamedSetting, writeNamedSetting } from './legacy-value.ts'

const defaultModel = '@cf/mistralai/mistral-small-3.1-24b-instruct'
const cloudflareApiToken = '{{secret:cloudflareApiToken}}'
const gatewayName = 'kody'

/**
 * Fetch resolves live secret placeholders in JSON bodies. Chat text (PR
 * prose, docs examples) must use the inert angle-bracket form so the
 * gateway does not look up a documentation example name or inject real
 * credentials into the model request. The Authorization header still uses
 * the live cloudflareApiToken placeholder above.
 */
const liveSecretPlaceholder =
	/\{\{secret:([a-zA-Z0-9._-]+)(\|scope=(?:session|package|user))?\}\}/g
const liveBasicAuthPlaceholder =
	/\{\{secret-basic:username=([a-zA-Z0-9._-]+),password=([a-zA-Z0-9._-]+)(\|scope=(?:session|package|user))?\}\}/g

function neutralizeSecretPlaceholdersInText(value: string) {
	return value
		.replace(liveSecretPlaceholder, '{{secret:<$1>$2}}')
		.replace(
			liveBasicAuthPlaceholder,
			'{{secret-basic:username=<$1>,password=<$2>$3}}',
		)
}

function neutralizeModelMessage(message: ModelStepMessage): ModelStepMessage {
	if (typeof message.content !== 'string') return message
	return {
		...message,
		content: neutralizeSecretPlaceholdersInText(message.content),
	}
}

export type ModelStepMessage = {
	role: 'system' | 'user' | 'assistant' | 'tool'
	content?: string | null
	tool_calls?: unknown[]
	tool_call_id?: string
	name?: string
}

export type ModelStepInput = {
	messages: ModelStepMessage[]
	system?: string | { content: string; cache?: 'prefix' }
	modelId?: string
	/** Alias for `modelId` (lower precedence). Matches the user-value name hosts already know. */
	cloudflareAiModel?: string
	conversationId?: string
	/**
	 * Completion token budget (default 4096). Reasoning models spend hidden
	 * reasoning tokens from this same budget, so hosts expecting long visible
	 * output from them should raise it and check `finishReason === 'length'`.
	 */
	maxTokens?: number
	/** When false, omit tool definitions (final answer pass). Defaults to true. */
	useTools?: boolean
	/**
	 * Custom OpenAI-style tool definitions for host-side tool loops. When set
	 * (and `useTools` is not false) these replace the default Kody
	 * search/execute tools. The host executes returned `toolCalls` itself and
	 * replays them as assistant `tool_calls` + `tool` messages.
	 */
	tools?: unknown[]
}

export type ModelStepToolCall = {
	toolCallId: string
	toolName: string
	input: Record<string, unknown>
}

export type ModelStepResult = {
	ok: boolean
	error?: string
	text: string
	reasoningText: string
	finishReason: string
	toolCalls: ModelStepToolCall[]
	conversationId: string
}

function clean(value: unknown) {
	return String(value ?? '').trim()
}

function stringify(value: unknown) {
	try {
		return JSON.stringify(value, null, 2)
	} catch {
		return String(value)
	}
}

async function readUserValue(name: string) {
	return await readNamedSetting(name)
}

async function resolveAccountId() {
	const accountId = await readUserValue('cloudflareAccountId')
	if (!accountId) throw new Error('cloudflareAccountId is required in this package\'s storage.')
	return accountId
}

async function resolveModelId(input: ModelStepInput) {
	if (input.modelId) return input.modelId
	if (input.cloudflareAiModel) return input.cloudflareAiModel
	return (await readUserValue('cloudflareAiModel').catch(() => '')) || defaultModel
}

async function cloudflareApi<T>(
	accountId: string,
	path: string,
	init: RequestInit = {},
): Promise<{ ok: boolean; status: number; data: T | null; text: string }> {
	const response = await fetch('https://api.cloudflare.com/client/v4/accounts/' + accountId + path, {
		...init,
		headers: {
			Authorization: 'Bearer ' + cloudflareApiToken,
			'Content-Type': 'application/json',
			...(init.headers || {}),
		},
	})
	const text = await response.text()
	let data: T | null = null
	try {
		data = text ? (JSON.parse(text) as T) : null
	} catch {
		data = null
	}
	return { ok: response.ok, status: response.status, data, text }
}

async function resolveGateway(accountId: string): Promise<string> {
	const existing = await readUserValue('cloudflareAiGatewayId')
	if (existing) return existing

	const list = await cloudflareApi<{
		success?: boolean
		result?: Array<{ id?: string }>
		errors?: Array<{ message?: string }>
	}>(accountId, '/ai-gateway/gateways')

	if (!list.ok) {
		const message =
			list.data?.errors?.map((error) => error.message).filter(Boolean).join('; ') ||
			'HTTP ' + list.status
		throw new Error(
			'Could not list AI Gateways (permission or API error): ' +
				message +
				'. Save cloudflareAiGatewayId in this package\'s storage or grant AI Gateway permissions to cloudflareApiToken.',
		)
	}

	const gateways = Array.isArray(list.data?.result) ? list.data!.result! : []
	const match = gateways.find((gateway) => gateway.id === gatewayName)
	if (match?.id) {
		await writeNamedSetting('cloudflareAiGatewayId', match.id)
		return match.id
	}

	const created = await cloudflareApi<{
		success?: boolean
		result?: { id?: string }
		errors?: Array<{ message?: string }>
	}>(accountId, '/ai-gateway/gateways', {
		method: 'POST',
		body: JSON.stringify({ id: gatewayName, cache_ttl: 0, collect_logs: true }),
	})

	if (created.ok && created.data?.result?.id) {
		const gatewayId = created.data.result.id
		await writeNamedSetting('cloudflareAiGatewayId', gatewayId)
		return gatewayId
	}

	const message =
		created.data?.errors?.map((error) => error.message).filter(Boolean).join('; ') ||
		'HTTP ' + created.status
	throw new Error(
		'Could not create AI Gateway "' +
			gatewayName +
			'": ' +
			message +
			'. Save cloudflareAiGatewayId in this package\'s storage or grant AI Gateway permissions to cloudflareApiToken.',
	)
}

/** Tool schemas hosts should expose when they execute Kody search/execute themselves. */
export function kodyAgentToolDefinitions() {
	return [
		{
			type: 'function',
			function: {
				name: 'search',
				description:
					'Discover Kody capabilities, saved packages, persisted values, integrations, and secret metadata (not secret values). Use a natural-language `query` for ranked matches, or pass `entity` (`"{id}:{type}"` or an array of refs) for exact detail/snippet for one or more hits before you execute.',
				parameters: {
					type: 'object',
					properties: {
						query: {
							type: 'string',
							description:
								'Natural language description of what you need, or an exact package UUID / kody id / hosted package URL.',
						},
						entity: {
							description:
								'Exact entity ref like "name:capability", "uuid:package", "name:integration", "user:name:value", or "name:secret", or an array of 1–10 refs.',
							anyOf: [{ type: 'string' }, { type: 'array', items: { type: 'string' } }],
						},
						limit: {
							type: 'number',
							description: 'Max ranked matches (default 15).',
						},
					},
				},
			},
		},
		{
			type: 'function',
			function: {
				name: 'execute',
				description:
					'Run one complete ESM module in the Kody sandbox. Required shape: `export default async function main(input = {}) { ... }`. Prefer `import { kody } from "kody:runtime"` and `kody:@scope/package[/export]` from search/entity detail. Project/slim large payloads before returning. Pass optional `params` as main\'s first argument.',
				parameters: {
					type: 'object',
					properties: {
						code: {
							type: 'string',
							description: 'Full ESM module string with a default export function.',
						},
						params: {
							type: 'object',
							description: 'Optional JSON passed as the first argument to main.',
						},
					},
					required: ['code'],
				},
			},
		},
	]
}

function normalizeSystem(system: ModelStepInput['system']) {
	if (typeof system === 'string') return system
	if (system && typeof system.content === 'string') return system.content
	return ''
}

function sanitizeMessagesForModel(messages: ModelStepMessage[]) {
	const out: ModelStepMessage[] = []
	for (const message of messages) {
		const last = out[out.length - 1]
		if (message.role === 'user' && last?.role === 'tool') {
			out.push({
				role: 'assistant',
				content: 'Tool calls finished. I will continue from those results.',
			})
		}
		const next: ModelStepMessage = {
			role: message.role,
			content: message.content == null ? '' : message.content,
		}
		if (message.tool_calls) next.tool_calls = message.tool_calls
		if (message.tool_call_id) next.tool_call_id = message.tool_call_id
		if (message.name) next.name = message.name
		out.push(next)
	}
	return out
}

function parseArguments(value: unknown): Record<string, unknown> {
	if (!value) return {}
	if (typeof value === 'object' && !Array.isArray(value)) return value as Record<string, unknown>
	if (typeof value !== 'string') return {}
	try {
		const parsed = JSON.parse(value)
		return parsed && typeof parsed === 'object' && !Array.isArray(parsed)
			? (parsed as Record<string, unknown>)
			: {}
	} catch {
		return {}
	}
}

function choice(payload: unknown) {
	const choices =
		payload && typeof payload === 'object' ? (payload as { choices?: unknown[] }).choices : null
	return Array.isArray(choices) && choices[0] && typeof choices[0] === 'object'
		? (choices[0] as Record<string, unknown>)
		: null
}

function message(payload: unknown) {
	const msg = choice(payload)?.message
	return msg && typeof msg === 'object' ? (msg as Record<string, unknown>) : null
}

function textFrom(payload: unknown) {
	const msg = message(payload)
	if (typeof msg?.content === 'string') return msg.content
	return ''
}

function reasoningFrom(payload: unknown) {
	const msg = message(payload)
	if (typeof msg?.reasoning_content === 'string') return msg.reasoning_content
	if (typeof msg?.reasoning === 'string') return msg.reasoning
	return ''
}

function finishReasonFrom(payload: unknown) {
	const reason = choice(payload)?.finish_reason
	return typeof reason === 'string' ? reason : 'stop'
}

function toolCallsFrom(payload: unknown) {
	const msg = message(payload)
	return Array.isArray(msg?.tool_calls) ? msg.tool_calls : []
}

function toolId(call: unknown, index: number) {
	const record = call && typeof call === 'object' ? (call as Record<string, unknown>) : {}
	return typeof record.id === 'string' && record.id ? record.id : 'tool_' + index
}

function toolName(call: unknown) {
	const record = call && typeof call === 'object' ? (call as Record<string, unknown>) : {}
	const fn = record.function
	if (fn && typeof fn === 'object' && typeof (fn as { name?: string }).name === 'string') {
		return (fn as { name: string }).name
	}
	return typeof record.name === 'string' ? record.name : ''
}

function toolArgs(call: unknown) {
	const record = call && typeof call === 'object' ? (call as Record<string, unknown>) : {}
	const fn = record.function
	const args =
		fn && typeof fn === 'object' ? (fn as { arguments?: unknown }).arguments : record.arguments
	return parseArguments(args)
}

/**
 * One model completion only. Does not execute Kody tools.
 * Host packages should run `kody.search` / `kody.execute` in their own package context.
 */
export async function runModelStep(input: ModelStepInput = { messages: [] }): Promise<ModelStepResult> {
	const conversationId = input.conversationId || crypto.randomUUID()
	try {
		if (!Array.isArray(input.messages) || input.messages.length === 0) {
			throw new Error('messages must include at least one message.')
		}
		const accountId = await resolveAccountId()
		const gatewayId = await resolveGateway(accountId)
		const modelId = await resolveModelId(input)
		const system = normalizeSystem(input.system)
		const messages: ModelStepMessage[] = []
		if (system) {
			messages.push({
				role: 'system',
				content: neutralizeSecretPlaceholdersInText(system),
			})
		}
		for (const message of input.messages) {
			messages.push(
				neutralizeModelMessage({
					role: message.role || 'user',
					content: message.content ?? '',
					tool_calls: message.tool_calls,
					tool_call_id: message.tool_call_id,
					name: message.name,
				}),
			)
		}

		const requestedMaxTokens = Number(input.maxTokens)
		const maxTokens =
			Number.isInteger(requestedMaxTokens) && requestedMaxTokens > 0
				? Math.min(requestedMaxTokens, 32768)
				: 4096
		const body: Record<string, unknown> = {
			model: modelId,
			messages: sanitizeMessagesForModel(messages),
			max_tokens: maxTokens,
		}
		if (input.useTools !== false) {
			body.tools =
				Array.isArray(input.tools) && input.tools.length > 0
					? input.tools
					: kodyAgentToolDefinitions()
			body.tool_choice = 'auto'
		}

		const response = await fetch(
			'https://gateway.ai.cloudflare.com/v1/' +
				accountId +
				'/' +
				gatewayId +
				'/workers-ai/v1/chat/completions',
			{
				method: 'POST',
				headers: {
					Authorization: 'Bearer ' + cloudflareApiToken,
					'Content-Type': 'application/json',
				},
				body: JSON.stringify(body),
			},
		)
		const raw = await response.text()
		let payload: unknown = raw
		try {
			payload = raw ? JSON.parse(raw) : null
		} catch {
			payload = raw
		}
		if (!response.ok) {
			throw new Error(
				'Workers AI chat request failed: ' + response.status + ' ' + stringify(payload),
			)
		}
		const calls = toolCallsFrom(payload)
		return {
			ok: true,
			text: textFrom(payload),
			reasoningText: reasoningFrom(payload),
			finishReason: finishReasonFrom(payload),
			toolCalls: calls.map((call, index) => ({
				toolCallId: toolId(call, index),
				toolName: toolName(call),
				input: toolArgs(call),
			})),
			conversationId,
		}
	} catch (error) {
		return {
			ok: false,
			error: error instanceof Error ? error.message : String(error),
			text: '',
			reasoningText: '',
			finishReason: 'error',
			toolCalls: [],
			conversationId,
		}
	}
}

export default runModelStep