@teleporthq/teleport-plugin-next-workflows 0.43.58 → 0.43.60

package/__tests__/ai-chat-rag-sql-parameterization.test.ts

@@ -22,14 +22,34 @@ import { generateSqlValidatorCode } from '../src/sql-validator'
  */
 
 const SEMANTIC_QUERY = [
-  'SELECT d.id, d.content, d.search_content, d.keywords,',
-  '       1 - (d.embedding_json::vector <=> q.query_vector) AS similarity',
-  'FROM teleport_ai_chat_documents d',
-  "CROSS JOIN (SELECT NULLIF({{ queryEmbedding }}::text, '')::vector AS query_vector) q",
-  'WHERE d.embedding_json IS NOT NULL',
-  '  AND q.query_vector IS NOT NULL',
-  'ORDER BY d.embedding_json::vector <=> q.query_vector ASC',
-  'LIMIT 6',
+  'WITH corpus AS (',
+  '  SELECT count(*) AS total_docs,',
+  '         COALESCE(SUM(LENGTH(content)), 0) AS total_chars',
+  '  FROM teleport_ai_chat_documents',
+  ')',
+  'SELECT ranked.id, ranked.content, ranked.search_content, ranked.keywords, ranked.similarity',
+  'FROM (',
+  '  SELECT d.id, d.content, d.search_content, d.keywords,',
+  '         CASE',
+  '           WHEN q.query_vector IS NULL OR d.embedding_json IS NULL THEN NULL',
+  '           ELSE 1 - (d.embedding_json::vector <=> q.query_vector)',
+  '         END AS similarity,',
+  '         row_number() OVER (',
+  '           ORDER BY (q.query_vector IS NULL OR d.embedding_json IS NULL) ASC,',
+  '                    CASE',
+  '                      WHEN q.query_vector IS NULL OR d.embedding_json IS NULL THEN NULL',
+  '                      ELSE d.embedding_json::vector <=> q.query_vector',
+  '                    END ASC,',
+  '                    d.id ASC',
+  '         ) AS rn',
+  '  FROM teleport_ai_chat_documents d',
+  "  CROSS JOIN (SELECT NULLIF({{ queryEmbedding }}::text, '')::vector AS query_vector) q",
+  ') ranked',
+  'CROSS JOIN corpus',
+  'WHERE (corpus.total_docs <= 30 AND corpus.total_chars <= 24000)',
+  '   OR (ranked.similarity IS NOT NULL AND ranked.rn <= 6)',
+  'ORDER BY ranked.rn ASC',
+  'LIMIT 30',
 ].join('\n')
 
 const LEXICAL_QUERY = [

package/__tests__/ai-chat-tool-dispatch-segments.test.ts

@@ -0,0 +1,110 @@
+import { splitIntoSegments } from '../src/segment-splitter'
+import type { UIDLWorkflow, UIDLWorkflowNode, UIDLWorkflowEdge } from '@teleporthq/teleport-types'
+
+/**
+ * The AI Assistant Chat's action layer stands or falls on ONE property of the
+ * generated app: the tool must run BEFORE the answer is written, so the reply
+ * can report truthfully what happened.
+ *
+ * The tool runs inside a `general-custom-node`, which the splitter pins to the
+ * CLIENT, while both AI calls are SERVER nodes. So the walk has to produce
+ * `[server: decide] [client: run] [server: answer]`.
+ *
+ * The chat used to express this with two plain edges out of a parse node — one
+ * into the tool chain, one into the answer — and relied on the answer's PROMPT
+ * referencing the tool's summary to force the ordering, because
+ * `getTopologicalOrder` turns a `workflowContext` reference into a dependency
+ * edge. That worked, but only by accident of the data flow: remove the
+ * reference and the answer was written before the tool ran.
+ *
+ * The dispatch now lives inside its own custom node, so the ordering comes from
+ * a real EDGE. This test pins both halves: the three-segment shape, and that it
+ * survives even with no prompt reference at all.
+ */
+
+const ctx = (nodeId: string, path: string[]) => ({ type: 'workflowContext', nodeId, path })
+
+const node = (
+  id: string,
+  type: string,
+  stepNumber: number,
+  config: Record<string, unknown> = {}
+): UIDLWorkflowNode => ({ id, type, label: id, config, stepNumber } as unknown as UIDLWorkflowNode)
+
+const edge = (source: string, target: string, sourceHandle?: string): UIDLWorkflowEdge =>
+  ({ id: `${source}->${target}`, source, target, sourceHandle } as unknown as UIDLWorkflowEdge)
+
+/** The chat's shape from the merge script onward. */
+function buildChatWorkflow(options: { answerReferencesOutcome: boolean }): UIDLWorkflow {
+  const nodes: UIDLWorkflowNode[] = [
+    node('merge', 'general-custom-js', 1, { context: 'server', code: 'x' }),
+    node('decide', 'ai-custom-prompt', 2, {
+      jsonMode: true,
+      prompt: ['ctx:', ctx('merge', ['merge', 'context'])],
+    }),
+    node('run', 'general-custom-node', 3, {
+      customNodeId: 'custom_run_action',
+      parameters: [{ key: 'decisionText', value: ctx('decide', ['decide', 'response']) }],
+    }),
+    node('answer', 'ai-custom-prompt', 4, {
+      streaming: true,
+      prompt: [
+        'Context:',
+        ctx('merge', ['merge', 'context']),
+        ...(options.answerReferencesOutcome
+          ? ['\nOutcome: ', ctx('run', ['run', 'outcomeText'])]
+          : []),
+      ],
+    }),
+  ]
+
+  const edges: UIDLWorkflowEdge[] = [
+    edge('merge', 'decide'),
+    edge('decide', 'run'),
+    edge('run', 'answer'),
+  ]
+
+  return { name: 'send-chat-message', nodes, edges, trigger: { type: 'manual' } } as UIDLWorkflow
+}
+
+const segmentOf = (segments: ReturnType<typeof splitIntoSegments>, nodeId: string): number =>
+  segments.findIndex((segment) => segment.nodes.some((entry) => entry.id === nodeId))
+
+describe('AI chat tool dispatch — segment ordering', () => {
+  const segments = splitIntoSegments(buildChatWorkflow({ answerReferencesOutcome: true }))
+
+  it('runs the action on the CLIENT, where general-custom-node is dispatched', () => {
+    // The server handler for general-custom-node is a stub: placed in a server
+    // segment the nested custom node never runs and downstream references
+    // resolve to the raw marker object.
+    expect(segments[segmentOf(segments, 'run')].env).toBe('client')
+  })
+
+  it('keeps both AI calls on the server', () => {
+    expect(segments[segmentOf(segments, 'decide')].env).toBe('server')
+    expect(segments[segmentOf(segments, 'answer')].env).toBe('server')
+  })
+
+  it('orders the answer AFTER the action', () => {
+    const answerSegment = segmentOf(segments, 'answer')
+    expect(answerSegment).toBeGreaterThan(segmentOf(segments, 'decide'))
+    expect(answerSegment).toBeGreaterThan(segmentOf(segments, 'run'))
+  })
+
+  it('produces exactly the three-segment shape the design depends on', () => {
+    expect(segments.map((segment) => segment.env)).toEqual(['server', 'client', 'server'])
+  })
+
+  it('holds the ordering on the EDGE, not on the prompt reference', () => {
+    // The regression this replaces: with the old two-edges-out-of-parse shape,
+    // dropping the prompt reference let the answer overtake the tool. A real
+    // edge cannot be overtaken, so the shape is identical either way.
+    const withoutReference = splitIntoSegments(
+      buildChatWorkflow({ answerReferencesOutcome: false })
+    )
+    expect(withoutReference.map((segment) => segment.env)).toEqual(['server', 'client', 'server'])
+    expect(segmentOf(withoutReference, 'answer')).toBeGreaterThan(
+      segmentOf(withoutReference, 'run')
+    )
+  })
+})

package/__tests__/ai-custom-prompt-json-mode.test.ts

@@ -0,0 +1,78 @@
+import { aiCustomPrompt } from '../src/nodes/ai/ai-custom-prompt'
+
+/**
+ * `ai-custom-prompt` used to hard-code `jsonMode: false`, so a node whose only
+ * job is to emit a machine-readable decision (the AI Assistant Chat's tool
+ * router) had no way to ask the provider for a JSON-only reply. The flag is now
+ * forwarded from the node config.
+ *
+ * The generated provider utils are emitted as
+ * `var __ai_callProvider = typeof __ai_callProvider !== 'undefined' ? … : <def>`,
+ * so defining the name ahead of the handler source substitutes a spy for the
+ * real provider call.
+ */
+type ProviderCall = { jsonMode?: boolean; systemMessage?: string; userMessage?: string }
+type PromptHandler = (
+  config: Record<string, unknown>,
+  context: Record<string, unknown>,
+  streamCallback?: unknown
+) => Promise<Record<string, unknown>>
+
+function loadHandlerWithProviderSpy(): { handler: PromptHandler; calls: ProviderCall[] } {
+  const calls: ProviderCall[] = []
+  const source = aiCustomPrompt.generateHandler()
+  // eslint-disable-next-line no-new-func
+  const factory = new Function(
+    '__calls',
+    [
+      'var __ai_callProvider = function (params) {',
+      '  __calls.push(params);',
+      '  return Promise.resolve({ content: \'{"action":"none"}\', usage: { total_tokens: 3 } });',
+      '};',
+      source,
+      'return ai_custom_prompt;',
+    ].join('\n')
+  )
+  return { handler: factory(calls) as PromptHandler, calls }
+}
+
+describe('ai-custom-prompt jsonMode passthrough', () => {
+  it('asks the provider for JSON when the node config sets jsonMode', async () => {
+    const { handler, calls } = loadHandlerWithProviderSpy()
+
+    const result = await handler({ prompt: 'Decide.', token: 'sk-test', jsonMode: true }, {})
+
+    expect(calls).toHaveLength(1)
+    expect(calls[0].jsonMode).toBe(true)
+    expect(result.response).toBe('{"action":"none"}')
+  })
+
+  it('leaves JSON mode off for an ordinary prompt node', async () => {
+    const { handler, calls } = loadHandlerWithProviderSpy()
+
+    await handler({ prompt: 'Write a haiku.', token: 'sk-test' }, {})
+
+    expect(calls[0].jsonMode).toBe(false)
+  })
+
+  it('treats any non-true value as off, so a stale string config cannot enable it', async () => {
+    const { handler, calls } = loadHandlerWithProviderSpy()
+
+    await handler({ prompt: 'Hi.', token: 'sk-test', jsonMode: 'true' }, {})
+
+    expect(calls[0].jsonMode).toBe(false)
+  })
+
+  it('never sends jsonMode down the streaming path', async () => {
+    const { handler, calls } = loadHandlerWithProviderSpy()
+    const chunks: string[] = []
+
+    // No stream callback => the streaming branch is skipped; with one, the
+    // handler must not reach __ai_callProvider at all.
+    await handler({ prompt: 'Hi.', token: 'sk-test', jsonMode: true, streaming: true }, {}, () =>
+      chunks.push('x')
+    )
+
+    expect(calls).toHaveLength(0)
+  })
+})

package/__tests__/ai-select-database-data-guard.test.ts

@@ -0,0 +1,218 @@
+import { generateAiSqlSelectGuard } from '../src/nodes/ai/ai-sql-select-guard'
+
+/**
+ * The guard is the ONLY thing standing between an AI model's output and the
+ * database: the model can be prompt-injected into producing anything, so a
+ * miss here is a write/exfiltration primitive, and an over-match breaks the
+ * legitimate aggregate queries the node exists for (MIN/MAX/CASE…END).
+ * Exercised exactly as the generated app defines it — by evaluating the
+ * emitted source.
+ */
+function loadGuard() {
+  // eslint-disable-next-line no-new-func
+  return new Function(
+    `${generateAiSqlSelectGuard()}
+    return {
+      validate: __aisql_validateSelectQuery,
+      enforceLimit: __aisql_enforceLimit,
+      buildSystemPrompt: __aisql_buildSystemPrompt,
+    };`
+  )() as {
+    validate: (sql: string, allowedTables: string[]) => { valid: boolean; reason?: string }
+    enforceLimit: (sql: string, maxRows: number) => string
+    buildSystemPrompt: (tableSchemas: unknown, allowedTables: string[], maxRows: number) => string
+  }
+}
+
+const { validate, enforceLimit, buildSystemPrompt } = loadGuard()
+const ALLOWED = ['teleport_products', 'teleport_blog_posts']
+
+describe('accepts the queries the node exists for', () => {
+  it('plain SELECT over an allowed table', () => {
+    expect(
+      validate('SELECT "id", "name", "price" FROM "teleport_products" LIMIT 10', ALLOWED)
+    ).toEqual({
+      valid: true,
+    })
+  })
+
+  it('aggregates and superlatives (MIN/MAX/ORDER BY … LIMIT 1)', () => {
+    expect(
+      validate(
+        'SELECT "name", "price" FROM "teleport_products" ORDER BY "price" DESC LIMIT 1',
+        ALLOWED
+      ).valid
+    ).toBe(true)
+    expect(validate('SELECT MAX("price") FROM "teleport_products"', ALLOWED).valid).toBe(true)
+    expect(validate('SELECT COUNT(*) FROM "teleport_blog_posts";', ALLOWED).valid).toBe(true)
+  })
+
+  it('CASE … END (the forms-worker validator regression)', () => {
+    expect(
+      validate(
+        `SELECT "name", CASE WHEN "price" > 100 THEN 'premium' ELSE 'standard' END AS tier FROM "teleport_products"`,
+        ALLOWED
+      ).valid
+    ).toBe(true)
+  })
+
+  it('WITH … SELECT over allowed tables, CTE referenced by name', () => {
+    expect(
+      validate(
+        'WITH ranked AS (SELECT "name", "price" FROM "teleport_products") SELECT * FROM ranked ORDER BY "price" ASC LIMIT 1',
+        ALLOWED
+      ).valid
+    ).toBe(true)
+  })
+
+  it('JOIN with an alias between allowed tables', () => {
+    expect(
+      validate(
+        'SELECT p."name" FROM "teleport_products" p JOIN "teleport_blog_posts" b ON b."id" = p."id"',
+        ALLOWED
+      ).valid
+    ).toBe(true)
+  })
+
+  it('set-returning functions in FROM are not table refs', () => {
+    expect(
+      validate(
+        `SELECT kv.key FROM "teleport_products" p, jsonb_each_text(p."metadata") kv`,
+        ALLOWED
+      ).valid
+    ).toBe(true)
+  })
+
+  it('keywords inside string literals do not false-positive', () => {
+    expect(
+      validate(`SELECT "name" FROM "teleport_products" WHERE "status" = 'deleted'`, ALLOWED).valid
+    ).toBe(true)
+    expect(
+      validate(
+        `SELECT "name" FROM "teleport_products" WHERE "description" = 'DROP by the store'`,
+        ALLOWED
+      ).valid
+    ).toBe(true)
+  })
+})
+
+describe('rejects everything that is not a single read-only SELECT', () => {
+  const rejected = (sql: string, tables: string[] = ALLOWED) =>
+    validate(sql, tables).valid === false
+
+  it('writes and DDL', () => {
+    expect(rejected(`INSERT INTO "teleport_products" ("name") VALUES ('x')`)).toBe(true)
+    expect(rejected(`UPDATE "teleport_products" SET "price" = 0`)).toBe(true)
+    expect(rejected(`DELETE FROM "teleport_products"`)).toBe(true)
+    expect(rejected(`DROP TABLE "teleport_products"`)).toBe(true)
+    expect(rejected(`SELECT * INTO backup FROM "teleport_products"`)).toBe(true)
+  })
+
+  it('stacked statements and comments', () => {
+    expect(rejected(`SELECT 1; DELETE FROM "teleport_products"`)).toBe(true)
+    expect(rejected(`SELECT "name" FROM "teleport_products" -- hidden`)).toBe(true)
+    expect(rejected(`SELECT /* smuggle */ "name" FROM "teleport_products"`)).toBe(true)
+  })
+
+  it('system catalogs and pg_* objects', () => {
+    expect(rejected(`SELECT * FROM information_schema.tables`)).toBe(true)
+    expect(rejected(`SELECT * FROM pg_catalog.pg_tables`)).toBe(true)
+    expect(rejected(`SELECT * FROM pg_stat_activity`)).toBe(true)
+  })
+
+  it('tables outside the allowlist, in FROM, JOIN, and subqueries', () => {
+    expect(rejected(`SELECT * FROM "users"`)).toBe(true)
+    expect(rejected(`SELECT * FROM "teleport_products" JOIN "users" u ON u."id" = 1`)).toBe(true)
+    expect(
+      rejected(`SELECT * FROM "teleport_products" WHERE "id" IN (SELECT "id" FROM "users")`)
+    ).toBe(true)
+    expect(rejected(`SELECT * FROM otherschema."teleport_products"`)).toBe(true)
+  })
+
+  it('a CTE cannot launder a non-allowed table', () => {
+    expect(rejected(`WITH x AS (SELECT * FROM "users") SELECT * FROM x`)).toBe(true)
+  })
+
+  it('timing/file functions and locks', () => {
+    expect(rejected(`SELECT pg_sleep(10)`)).toBe(true)
+    expect(rejected(`SELECT current_setting('server_version')`)).toBe(true)
+    expect(rejected(`SELECT "name" FROM "teleport_products" FOR SHARE`)).toBe(true)
+    // FOR UPDATE falls to the UPDATE keyword rule
+    expect(rejected(`SELECT "name" FROM "teleport_products" FOR UPDATE`)).toBe(true)
+  })
+
+  it('transaction control and session state', () => {
+    expect(rejected(`BEGIN`)).toBe(true)
+    expect(rejected(`SET search_path TO public`)).toBe(true)
+  })
+
+  it('degenerate inputs', () => {
+    expect(rejected('')).toBe(true)
+    expect(rejected('   ')).toBe(true)
+    expect(rejected('SELECT 1'.padEnd(10001, '1'))).toBe(true)
+    expect(validate(null as unknown as string, ALLOWED).valid).toBe(false)
+  })
+})
+
+describe('enforceLimit', () => {
+  it('appends a LIMIT when the statement has none', () => {
+    expect(enforceLimit('SELECT * FROM "teleport_products"', 100)).toBe(
+      'SELECT * FROM "teleport_products" LIMIT 100'
+    )
+  })
+
+  it('clamps an oversized trailing LIMIT and keeps a compliant one', () => {
+    expect(enforceLimit('SELECT * FROM "teleport_products" LIMIT 5000', 100)).toBe(
+      'SELECT * FROM "teleport_products" LIMIT 100'
+    )
+    expect(enforceLimit('SELECT * FROM "teleport_products" LIMIT 5', 100)).toBe(
+      'SELECT * FROM "teleport_products" LIMIT 5'
+    )
+  })
+
+  it('a subquery LIMIT does not satisfy the outer statement', () => {
+    const sql = 'SELECT * FROM (SELECT * FROM "teleport_products" LIMIT 5) t WHERE t."price" > 0'
+    expect(enforceLimit(sql, 100)).toBe(sql + ' LIMIT 100')
+  })
+
+  it('strips a trailing semicolon and re-clamps maxRows itself', () => {
+    expect(enforceLimit('SELECT 1;', 999999)).toBe('SELECT 1 LIMIT 1000')
+    expect(enforceLimit('SELECT 1', 0)).toBe('SELECT 1 LIMIT 100')
+  })
+})
+
+describe('buildSystemPrompt', () => {
+  it('renders baked schemas and falls back for undescribed tables', () => {
+    const prompt = buildSystemPrompt(
+      [
+        {
+          table: 'teleport_products',
+          columns: [
+            { name: 'id', type: 'uuid', nullable: false },
+            { name: 'price', type: 'numeric' },
+          ],
+        },
+      ],
+      ALLOWED,
+      50
+    )
+    expect(prompt).toContain('TABLE "teleport_products" ("id" uuid NOT NULL, "price" numeric)')
+    expect(prompt).toContain('TABLE "teleport_blog_posts" (column names unknown')
+    expect(prompt).toContain('LIMIT of at most 50')
+    expect(prompt).toContain('needsQuery')
+  })
+
+  it('teaches the SQL habits the store chat needed in practice', () => {
+    const prompt = buildSystemPrompt([], ALLOWED, 50)
+    // "what price does the 054 product have?" — partial name, possibly living
+    // in a language-suffixed column.
+    expect(prompt).toContain('ILIKE')
+    expect(prompt).toContain('language-suffixed variants')
+    // "most expensive product" with several products at the same top price —
+    // ORDER BY + LIMIT 1 silently returned only one of them.
+    expect(prompt).toContain('MAX()/MIN() subquery')
+    expect(prompt).toContain('never ORDER BY with LIMIT 1')
+    // Rows must arrive identifiable and linkable.
+    expect(prompt).toContain('name/title/slug-like columns')
+  })
+})

package/__tests__/ai-select-database-data-handler.test.ts

@@ -0,0 +1,237 @@
+import { aiSelectDatabaseData } from '../src/nodes/ai/ai-select-database-data'
+import { generateSharedRuntimeUtilsCode } from '../src/executor-generator'
+
+/**
+ * The handler's result is serialized straight to the visitor's browser, so the
+ * contract under test is as much about what it must NEVER contain (the
+ * generated SQL, schema text, database error messages) as about what it
+ * returns. The AI call and fetch are mocked by pre-defining the guarded
+ * globals — the emitted `var X = typeof X !== 'undefined' ? X : …` pattern
+ * keeps whatever is already in scope.
+ */
+
+interface MockState {
+  aiResponses: string[]
+  aiCalls: Array<{ systemMessage: string; userMessage: string; jsonMode: boolean }>
+  fetchCalls: Array<{ url: string; body: { query: string; params: unknown[] } }>
+  fetchResponse: { ok: boolean; json: unknown }
+}
+
+function loadHandler(state: MockState) {
+  const mocks = `
+    var __ai_callProvider = async function (params) {
+      state.aiCalls.push({ systemMessage: params.systemMessage, userMessage: params.userMessage, jsonMode: params.jsonMode });
+      var next = state.aiResponses.shift();
+      if (next === '__throw__') { throw new Error('provider exploded'); }
+      return { content: next };
+    };
+  `
+  const source = aiSelectDatabaseData.generateServerHandler!()
+  const fetchMock = async (url: string, init: { body: string }) => {
+    state.fetchCalls.push({ url, body: JSON.parse(init.body) })
+    return {
+      ok: state.fetchResponse.ok,
+      json: async () => state.fetchResponse.json,
+    }
+  }
+  // eslint-disable-next-line no-new-func
+  return new Function('state', 'fetch', `${mocks}\n${source}\nreturn ai_select_database_data;`)(
+    state,
+    fetchMock
+  ) as (config: unknown, context: Record<string, unknown>) => Promise<any>
+}
+
+function makeState(overrides: Partial<MockState> = {}): MockState {
+  return {
+    aiResponses: ['{"needsQuery": true, "query": "SELECT \\"id\\" FROM \\"teleport_products\\""}'],
+    aiCalls: [],
+    fetchCalls: [],
+    fetchResponse: { ok: true, json: { rows: [{ id: '1' }, { id: '2' }] } },
+    ...overrides,
+  }
+}
+
+const BASE_CONFIG = {
+  dataSourceId: 'ds1',
+  allowedTables: ['teleport_products'],
+  prompt: 'What is the cheapest product?',
+  model: 'gpt-4o',
+  token: 'sk-test',
+  tableSchemas: [
+    { table: 'teleport_products', columns: [{ name: 'id', type: 'uuid', nullable: false }] },
+  ],
+}
+
+function loadIsFatalNodeResult(): (result: unknown) => boolean {
+  const utilsModule = { exports: {} as { isFatalNodeResult: (result: unknown) => boolean } }
+  // eslint-disable-next-line @typescript-eslint/no-implied-eval
+  new Function('module', 'exports', 'require', generateSharedRuntimeUtilsCode())(
+    utilsModule,
+    utilsModule.exports,
+    () => ({})
+  )
+  return utilsModule.exports.isFatalNodeResult
+}
+
+describe('happy path', () => {
+  it('generates, validates, executes, and returns only { executed, rowCount, rows, truncated }', async () => {
+    const state = makeState()
+    const handler = loadHandler(state)
+    const result = await handler(BASE_CONFIG, {})
+
+    expect(result).toEqual({
+      executed: true,
+      rowCount: 2,
+      rows: [{ id: '1' }, { id: '2' }],
+      // Two rows under a cap of 50: nothing was cut off, so a caller may
+      // report this as a complete result.
+      truncated: false,
+    })
+    expect(state.aiCalls).toHaveLength(1)
+    expect(state.aiCalls[0].jsonMode).toBe(true)
+    expect(state.aiCalls[0].systemMessage).toContain('TABLE "teleport_products"')
+    expect(state.fetchCalls).toHaveLength(1)
+    expect(state.fetchCalls[0].url).toBe('/api/data/ds1/raw-query')
+    // The executed SQL carries the enforced LIMIT
+    expect(state.fetchCalls[0].body.query).toBe('SELECT "id" FROM "teleport_products" LIMIT 100')
+  })
+
+  it('respects the AI deciding no query is needed, without touching the database', async () => {
+    const state = makeState({ aiResponses: ['{"needsQuery": false, "query": null}'] })
+    const handler = loadHandler(state)
+    const result = await handler(BASE_CONFIG, {})
+
+    expect(result).toEqual({ executed: false, rows: [], rowCount: 0, truncated: false })
+    expect(state.fetchCalls).toHaveLength(0)
+  })
+
+  it('retries once with feedback when the first SQL is rejected', async () => {
+    const state = makeState({
+      aiResponses: [
+        '{"needsQuery": true, "query": "DELETE FROM \\"teleport_products\\""}',
+        '{"needsQuery": true, "query": "SELECT \\"id\\" FROM \\"teleport_products\\" LIMIT 5"}',
+      ],
+    })
+    const handler = loadHandler(state)
+    const result = await handler(BASE_CONFIG, {})
+
+    expect(result.executed).toBe(true)
+    expect(state.aiCalls).toHaveLength(2)
+    expect(state.aiCalls[1].userMessage).toContain('was rejected')
+  })
+})
+
+describe('failure contract — nothing sensitive ever leaves the handler', () => {
+  const flatten = (value: unknown): string => JSON.stringify(value)
+
+  it('rejected SQL on both attempts fails closed with a fixed code', async () => {
+    const bad = '{"needsQuery": true, "query": "DELETE FROM \\"teleport_products\\""}'
+    const state = makeState({ aiResponses: [bad, bad] })
+    const handler = loadHandler(state)
+    const result = await handler(BASE_CONFIG, {})
+
+    expect(result).toEqual({
+      error: true,
+      message: 'AI database query failed (sql_rejected)',
+      code: 'sql_rejected',
+    })
+    expect(flatten(result)).not.toContain('DELETE')
+    expect(flatten(result)).not.toContain('teleport_products')
+  })
+
+  it('never forwards the data route error body (pg errors quote the SQL)', async () => {
+    const state = makeState({
+      fetchResponse: {
+        ok: false,
+        json: { error: 'syntax error at or near "SELECT \\"id\\" FROM teleport_products"' },
+      },
+    })
+    const handler = loadHandler(state)
+    const result = await handler(BASE_CONFIG, {})
+
+    expect(result.code).toBe('query_failed')
+    expect(flatten(result)).not.toContain('syntax error')
+    expect(flatten(result)).not.toContain('SELECT')
+  })
+
+  it('optional: true converts every failure into skipped data', async () => {
+    const state = makeState({ aiResponses: ['__throw__'] })
+    const handler = loadHandler(state)
+    const result = await handler({ ...BASE_CONFIG, optional: true }, {})
+
+    expect(result).toEqual({
+      executed: false,
+      rows: [],
+      rowCount: 0,
+      truncated: false,
+      skipped: true,
+      skipReason: 'provider_error',
+    })
+  })
+
+  it('missing configuration fails before any AI call', async () => {
+    const state = makeState()
+    const handler = loadHandler(state)
+    const result = await handler({ ...BASE_CONFIG, allowedTables: [] }, {})
+
+    expect(result.code).toBe('missing_configuration')
+    expect(state.aiCalls).toHaveLength(0)
+  })
+
+  it('the executor treats default failures as fatal and optional ones as data', async () => {
+    const isFatalNodeResult = loadIsFatalNodeResult()
+    const bad = '{"needsQuery": true, "query": "DELETE FROM \\"teleport_products\\""}'
+
+    const fatalState = makeState({ aiResponses: [bad, bad] })
+    const fatal = await loadHandler(fatalState)(BASE_CONFIG, {})
+    const optionalState = makeState({ aiResponses: [bad, bad] })
+    const optional = await loadHandler(optionalState)({ ...BASE_CONFIG, optional: true }, {})
+
+    expect(isFatalNodeResult(fatal)).toBe(true)
+    expect(isFatalNodeResult(optional)).toBe(false)
+  })
+})
+
+describe('registration', () => {
+  it('is a server node so its config (schemas, allowlist) is redacted from client bundles', () => {
+    expect(aiSelectDatabaseData.executionEnv).toBe('server')
+  })
+})
+
+describe('telling a capped result apart from a complete one', () => {
+  it('flags a full page of rows as truncated', async () => {
+    // The statement is capped, so a result that fills the cap is one page of a
+    // larger answer. Reporting its length as a total is what made the chat say
+    // "there are 50 products" for a catalogue of 220.
+    const rows = Array.from({ length: 50 }, (_, index) => ({ id: String(index) }))
+    const state = makeState({ fetchResponse: { ok: true, json: { rows } } })
+    const handler = loadHandler(state)
+
+    const result = (await handler({ ...BASE_CONFIG, maxRows: 50 }, {})) as {
+      truncated: boolean
+      rowCount: number
+    }
+
+    expect(result.truncated).toBe(true)
+    expect(result.rowCount).toBe(50)
+  })
+
+  it('does not flag an aggregate answer as truncated', async () => {
+    const state = makeState({
+      aiResponses: [
+        '{"needsQuery": true, "query": "SELECT COUNT(*) AS \\"total_count\\" FROM \\"teleport_products\\""}',
+      ],
+      fetchResponse: { ok: true, json: { rows: [{ total_count: 223 }] } },
+    })
+    const handler = loadHandler(state)
+
+    const result = (await handler({ ...BASE_CONFIG, maxRows: 50 }, {})) as {
+      truncated: boolean
+      rows: unknown[]
+    }
+
+    // One row carrying the real total — the only shape that can state one.
+    expect(result.truncated).toBe(false)
+    expect(result.rows).toEqual([{ total_count: 223 }])
+  })
+})