windmill/frontend/src/lib/components/copilot/chat/datatableTools.ts

import { z } from 'zod'
import { WorkspaceService, type CompletedJob } from '$lib/gen'
import type { DataTableTables } from '$lib/gen/types.gen'
import { runScript } from '$lib/components/jobs/utils'
import { datatableReference } from '$lib/components/dbTypes'
import {
	createToolDef,
	executeTestRun,
	type ToolDisplayMessage,
	type ChatJobResultFormat
} from './shared'
import { NONE, RUN_PREVIEW, type SessionTool } from './sessionCapabilities'

/**
 * Workspace-scoped datatable tools, with no app whitelist and no creation policy.
 *
 * Datatables are workspace-level managed PostgreSQL databases. The backend
 * endpoints used here (`list_datatable_tables`, `get_datatable_table_schema`)
 * and SQL execution (`datatable://<name>`) need no app context: the server
 * decides what the caller reaches, as the datatable role they name or its
 * default. This is the unrestricted counterpart to the app-mode datatable
 * tools in `app/core.ts`, which additionally filter by the app's whitelist.
 */

// ============= Utility =============

/** Memoize a factory function - the factory is only called once, on first access */
const memo = <T>(factory: () => T): (() => T) => {
	let cached: T | undefined
	return () => (cached ??= factory())
}

// ============= Pure workspace-scoped operations =============

/** List the datatables configured in the workspace, with their schema/table names: all of them as
 * their default role, or only `datatableName`, as `role` when one is given. */
export async function listDatatables(
	workspace: string,
	datatableName?: string,
	role?: string
): Promise<DataTableTables[]> {
	if (datatableName === undefined) return await WorkspaceService.listDataTableTables({ workspace })
	return await WorkspaceService.listDataTableTables({
		workspace,
		datatableName,
		...(role !== undefined && { roleFor: datatableName, role })
	})
}

/** Get the columns (column_name -> compact_type) of one datatable table. */
export async function getDatatableColumns(
	workspace: string,
	datatableName: string,
	schemaName: string,
	tableName: string,
	role?: string
): Promise<Record<string, string>> {
	const schema = await WorkspaceService.getDataTableTableSchema({
		workspace,
		datatableName,
		schemaName,
		tableName,
		role
	})
	return schema.columns
}

// ============= Error helpers =============

/**
 * The backend returns "datatable <name> not found" when no datatable with that
 * name is configured in the workspace settings. That is a hard, blocking
 * prerequisite (not a transient failure), so we surface an explicit, actionable
 * message instead of the raw internal error.
 */
function isDatatableNotConfiguredError(error: string | null | undefined): boolean {
	return typeof error === 'string' && /datatable\s+\S+\s+not found/i.test(error)
}

function datatableNotConfiguredMessage(datatableName: string): string {
	return (
		`Datatable "${datatableName}" is not configured in this workspace, so this operation cannot run. ` +
		`Datatables are not created by SQL — they must be set up first by the user in the workspace settings ` +
		`(Workspace settings → Data Tables) before any table can be queried or created. ` +
		`This is a required, blocking prerequisite: do not retry on this or another name. ` +
		`Tell the user they need to configure a datatable (e.g. named "${datatableName}") in their workspace settings, then try again.`
	)
}

const NO_DATATABLES_CONFIGURED_MESSAGE =
	'No datatables are configured in this workspace. Datatable operations (querying data, creating or altering tables) are blocked until a datatable exists. ' +
	'Datatables are not created by SQL — the user must set one up in the workspace settings (Workspace settings → Data Tables) first. ' +
	'Do not call exec_datatable_sql or assume a "main" datatable exists; instead tell the user this is a required prerequisite and ask them to configure a datatable in their workspace settings.'

// ============= Tool definitions =============

// The same rule the server applies to `-- role <name>`; a name it would refuse fails here instead.
const getRoleSchema = memo(() =>
	z
		.string()
		.regex(/^[A-Za-z0-9_-]{1,63}$/)
		.optional()
		.describe(
			"The datatable role to connect as, when the code you are working on uses one (an app's `data.roles` entry, or the `role` it passes to wmill.datatable). Omit for the datatable's default role."
		)
)

const getListDatatablesSchema = memo(() =>
	z.object({
		datatable_name: z
			.string()
			.optional()
			.describe('List only this datatable. Required with `role`.'),
		role: getRoleSchema()
	})
)
const getListDatatablesToolDef = memo(() =>
	createToolDef(
		getListDatatablesSchema(),
		'list_datatables',
		'List datatables configured in the workspace with schema and table names only. Does not include column definitions. Use this directly for table-list or available-tables summaries. Only call get_datatable_table_schema when column names/types are required.'
	)
)

const getGetDatatableTableSchemaSchema = memo(() =>
	z.object({
		datatable_name: z.string().describe('The datatable name to inspect, e.g. "main".'),
		schema_name: z.string().describe('The schema name, e.g. "public".'),
		table_name: z.string().describe('The table name to inspect.'),
		role: getRoleSchema()
	})
)
const getGetDatatableTableSchemaToolDef = memo(() =>
	createToolDef(
		getGetDatatableTableSchemaSchema(),
		'get_datatable_table_schema',
		'Get column definitions for one datatable table. Do not call this for row counts or table-list summaries; list_datatables is enough for those.'
	)
)

const getExecDatatableSqlSchema = memo(() =>
	z.object({
		datatable_name: z
			.string()
			.describe(
				'The name of the datatable to query (e.g., "main"). Must be one of the datatables configured in the workspace.'
			),
		sql: z
			.string()
			.describe(
				'The SQL query to execute. Supports SELECT, INSERT, UPDATE, DELETE, CREATE TABLE, ALTER TABLE, DROP TABLE, etc. For SELECT queries, results are returned as an array of objects. A newly created table will appear in list_datatables automatically.'
			),
		role: getRoleSchema(),
		background: z
			.boolean()
			.optional()
			.describe(
				'Run in the background without waiting. Set true for queries you expect to be slow (large scans, heavy migrations) — you will be notified when it finishes. Leave unset for normal queries, which wait briefly and only background automatically if slow.'
			),
		wait_seconds: z
			.number()
			.optional()
			.describe(
				'How many seconds to wait for the query in-turn before it detaches into the background jobs tray. Defaults to 15. Raise it (capped at 120) for a query you expect to finish in ~30–60s and want the result in this same turn. Ignored when background is true.'
			)
	})
)
const getExecDatatableSqlToolDef = memo(() =>
	createToolDef(
		getExecDatatableSqlSchema(),
		'exec_datatable_sql',
		'Execute a SQL query on a workspace datatable. Use this to explore data, test queries, create/alter tables, or make changes. Creating a table is a normal CREATE TABLE statement — no registration step is needed.'
	)
)

/** Maximum rows returned to the model for a SELECT query. */
const MAX_ROWS = 100

/**
 * Shape a finished datatable SQL job into the model-visible result + tool card:
 * SELECT rows capped at MAX_ROWS, DDL/DML reported as zero rows, and the "datatable
 * not configured" backend error rewritten as the actionable blocking message.
 *
 * Pure and keyed only on the datatable name so it can run from BOTH the inline
 * completion path (executeTestRun.formatCompletion) AND a detached/rehydrated job
 * reconstructed from its persisted ChatJobResultFormat (see formatChatJobCompletion).
 */
export function formatDatatableSqlCompletion(
	job: CompletedJob,
	datatableName: string
): { llmText: string; card: Partial<ToolDisplayMessage> } {
	if (job.success) {
		// Successful runs always carry a `result` array (empty for DDL/DML),
		// so SELECT rows and zero-row statements share one reporting path.
		const rows = Array.isArray(job.result) ? (job.result as Record<string, any>[]) : []
		const rowCount = rows.length
		const payload =
			rowCount > MAX_ROWS
				? {
						success: true,
						rowCount,
						result: rows.slice(0, MAX_ROWS),
						note: `Showing first ${MAX_ROWS} of ${rowCount} rows`
					}
				: { success: true, rowCount, result: rows }
		return {
			llmText: JSON.stringify(payload, null, 2),
			card: {
				content: `Query returned ${rowCount} row(s)`,
				result: JSON.stringify(job.result, null, 2)
			}
		}
	}
	const raw =
		(job.result as any)?.error?.message ??
		(typeof job.result === 'string' ? job.result : JSON.stringify(job.result)) ??
		// JSON.stringify(undefined) is `undefined` (not a string), so a failed
		// job with no result would otherwise yield an empty error message.
		'Unknown error'
	const errorMsg = isDatatableNotConfiguredError(raw)
		? datatableNotConfiguredMessage(datatableName)
		: raw
	return {
		llmText: JSON.stringify({ success: false, error: errorMsg }),
		card: { content: `Error: ${errorMsg}`, error: errorMsg }
	}
}

/**
 * Reconstruct a tool's terminal formatter from the serializable descriptor stored
 * on a ChatJob. Lets a background job that detached (and may have survived a reload,
 * losing any in-memory closure) still report through its launching tool's result
 * contract. Dispatches on `kind`; extend as more tools gain persisted formatters.
 */
export function formatChatJobCompletion(
	job: CompletedJob,
	resultFormat: ChatJobResultFormat
): { llmText: string; card: Partial<ToolDisplayMessage> } {
	switch (resultFormat.kind) {
		case 'datatable':
			return formatDatatableSqlCompletion(job, resultFormat.datatableName)
	}
}

/**
 * The unrestricted workspace datatable tools, for registration in global mode.
 * Helper-free: each tool reads `workspace` directly from the tool call params.
 */
export function getDatatableTools(): SessionTool<{}>[] {
	return [
		{
			requires: NONE,
			def: getListDatatablesToolDef(),
			planModeSafe: true,
			fn: async ({ args, workspace, toolId, toolCallbacks }) => {
				toolCallbacks.setToolStatus(toolId, { content: 'Listing datatables...' })
				try {
					const parsedArgs = getListDatatablesSchema().parse(args ?? {})
					if (parsedArgs.role !== undefined && parsedArgs.datatable_name === undefined) {
						throw new Error('`role` needs `datatable_name`, the datatable it is a role of')
					}
					const metadata = await listDatatables(
						workspace,
						parsedArgs.datatable_name,
						parsedArgs.role
					)
					if (metadata.length === 0) {
						toolCallbacks.setToolStatus(toolId, {
							content: 'No datatables configured — set one up in workspace settings'
						})
						return NO_DATATABLES_CONFIGURED_MESSAGE
					}
					const totalTables = metadata.reduce(
						(acc, datatable) =>
							acc +
							Object.values(datatable.schemas).reduce((sum, tables) => sum + tables.length, 0),
						0
					)
					toolCallbacks.setToolStatus(toolId, {
						content: `Listed ${metadata.length} datatable(s) with ${totalTables} table(s)`
					})
					// Only what the model acts on: the roles it may pass, not the creation privileges
					// the manager's UI gates on.
					return JSON.stringify(
						metadata.map((d) => ({
							datatable_name: d.datatable_name,
							schemas: d.schemas,
							...(d.error && { error: d.error }),
							...(d.migrations_enabled && { migrations_enabled: true }),
							...(d.permissioned && { usable_roles: d.usable_roles, default_role: d.default_role })
						})),
						null,
						2
					)
				} catch (e) {
					const errorMsg = `Error listing datatables: ${e instanceof Error ? e.message : String(e)}`
					toolCallbacks.setToolStatus(toolId, { content: errorMsg, error: errorMsg })
					return errorMsg
				}
			}
		},
		{
			requires: NONE,
			def: getGetDatatableTableSchemaToolDef(),
			planModeSafe: true,
			fn: async ({ args, workspace, toolId, toolCallbacks }) => {
				const parsedArgs = getGetDatatableTableSchemaSchema().parse(args)
				toolCallbacks.setToolStatus(toolId, {
					content: `Getting schema for ${parsedArgs.datatable_name}.${parsedArgs.schema_name}.${parsedArgs.table_name}...`
				})
				try {
					const columns = await getDatatableColumns(
						workspace,
						parsedArgs.datatable_name,
						parsedArgs.schema_name,
						parsedArgs.table_name,
						parsedArgs.role
					)
					toolCallbacks.setToolStatus(toolId, {
						content: `Retrieved schema for ${parsedArgs.schema_name}.${parsedArgs.table_name}`
					})
					return JSON.stringify(
						{
							datatable_name: parsedArgs.datatable_name,
							schema_name: parsedArgs.schema_name,
							table_name: parsedArgs.table_name,
							columns
						},
						null,
						2
					)
				} catch (e) {
					const raw = e instanceof Error ? e.message : String(e)
					const errorMsg = isDatatableNotConfiguredError(raw)
						? datatableNotConfiguredMessage(parsedArgs.datatable_name)
						: `Error getting table schema: ${raw}`
					toolCallbacks.setToolStatus(toolId, { content: errorMsg, error: errorMsg })
					return errorMsg
				}
			}
		},
		{
			// `runScript` below goes through /jobs/run/preview, which jobs.rs refuses operators.
			requires: RUN_PREVIEW,
			def: getExecDatatableSqlToolDef(),
			requiresConfirmation: true,
			confirmAction: 'run',
			confirmationMessage: 'Execute SQL on datatable',
			queuedLabel: (args) => `Execute SQL on "${args?.datatable_name ?? 'datatable'}"`,
			showDetails: true,
			fn: async ({ args, workspace, toolId, toolCallbacks }) => {
				const parsedArgs = getExecDatatableSqlSchema().parse(args)
				const name = parsedArgs.datatable_name
				// Route through executeTestRun so a slow query detaches into the jobs
				// tray (and honors `background`) instead of blocking the chat turn,
				// while keeping the datatable-specific result/error shaping.
				return executeTestRun({
					jobStarter: () =>
						runScript({
							workspace,
							requestBody: {
								language: 'postgresql',
								content: parsedArgs.sql,
								args: { database: datatableReference(name, parsedArgs.role) }
							}
						}),
					workspace,
					toolCallbacks,
					toolId,
					contextName: 'script',
					label: `SQL · ${name}`,
					background: parsedArgs.background,
					detachAfterMs:
						parsedArgs.wait_seconds == null
							? undefined
							: Math.max(0, parsedArgs.wait_seconds) * 1000,
					startMessage: `Executing SQL on "${name}"...`,
					runningMessage: `SQL running on "${name}"...`,
					// Inline path shapes the result here; the descriptor mirrors it so a
					// detached/rehydrated completion reconstructs the same contract.
					formatCompletion: (job) => formatDatatableSqlCompletion(job, name),
					resultFormat: { kind: 'datatable', datatableName: name }
				})
			}
		}
	]
}