@deepagents/text2sql 0.17.1 → 0.19.0

package/dist/index.js

@@ -519,7 +519,8 @@ import {
   Bash,
   MountableFs,
   OverlayFs,
-  defineCommand
+  defineCommand,
+  parse
 } from "just-bash";
 import { AsyncLocalStorage } from "node:async_hooks";
 import * as path from "node:path";
@@ -650,6 +651,300 @@ function createSqlCommand(adapter, metaStore) {
     }
   });
 }
+var BLOCKED_DB_CLIENT_COMMANDS = /* @__PURE__ */ new Set([
+  "psql",
+  "sqlite3",
+  "mysql",
+  "duckdb"
+]);
+var BLOCKED_RAW_SQL_COMMANDS = /* @__PURE__ */ new Set(["select", "with"]);
+var ALLOWED_SQL_PROXY_SUBCOMMANDS = /* @__PURE__ */ new Set(["run", "validate"]);
+var SQL_PROXY_ENFORCEMENT_MESSAGE = [
+  "Direct database querying through bash is blocked.",
+  "Use SQL proxy commands in this order:",
+  '1) sql validate "SELECT ..."',
+  '2) sql run "SELECT ..."'
+].join("\n");
+function cloneInspectionContext(context) {
+  return {
+    functionDefinitions: new Map(context.functionDefinitions),
+    callStack: new Set(context.callStack)
+  };
+}
+function asStaticWordText(word) {
+  if (!word) {
+    return null;
+  }
+  return asStaticWordPartText(
+    word.parts
+  );
+}
+function asStaticWordPartText(parts) {
+  let text = "";
+  for (const part of parts) {
+    const type = part.type;
+    if (type === "Literal" || type === "SingleQuoted" || type === "Escaped") {
+      if (typeof part.value !== "string") {
+        return null;
+      }
+      text += part.value;
+      continue;
+    }
+    if (type === "DoubleQuoted") {
+      if (!Array.isArray(part.parts)) {
+        return null;
+      }
+      const inner = asStaticWordPartText(
+        part.parts
+      );
+      if (inner == null) {
+        return null;
+      }
+      text += inner;
+      continue;
+    }
+    return null;
+  }
+  return text;
+}
+function isScriptNode(value) {
+  if (typeof value !== "object" || value === null) {
+    return false;
+  }
+  const node = value;
+  return node.type === "Script" && Array.isArray(node.statements);
+}
+function scriptContainsBlockedCommand(script, context) {
+  return statementsContainBlockedCommand(script.statements, context);
+}
+function statementsContainBlockedCommand(statements, context) {
+  for (const statement of statements) {
+    if (statementContainsBlockedCommand(statement, context)) {
+      return true;
+    }
+  }
+  return false;
+}
+function statementContainsBlockedCommand(statement, context) {
+  for (const pipeline of statement.pipelines) {
+    if (pipelineContainsBlockedCommand(pipeline, context)) {
+      return true;
+    }
+  }
+  return false;
+}
+function pipelineContainsBlockedCommand(pipeline, context) {
+  for (const command of pipeline.commands) {
+    if (command.type === "FunctionDef") {
+      context.functionDefinitions.set(command.name, command);
+      continue;
+    }
+    if (commandContainsBlockedCommand(command, context)) {
+      return true;
+    }
+  }
+  return false;
+}
+function stringCommandContainsBlockedCommand(command, context) {
+  let script;
+  try {
+    script = parse(command);
+  } catch {
+    return false;
+  }
+  return scriptContainsBlockedCommand(script, cloneInspectionContext(context));
+}
+function wordContainsBlockedCommand(word, context) {
+  if (!word) {
+    return false;
+  }
+  return wordPartContainsBlockedCommand(
+    word.parts,
+    context
+  );
+}
+function wordPartContainsBlockedCommand(parts, context) {
+  for (const part of parts) {
+    if (partContainsBlockedCommand(part, context)) {
+      return true;
+    }
+  }
+  return false;
+}
+function partContainsBlockedCommand(node, context) {
+  const type = node.type;
+  if (type === "CommandSubstitution" || type === "ProcessSubstitution") {
+    if (isScriptNode(node.body)) {
+      return scriptContainsBlockedCommand(
+        node.body,
+        cloneInspectionContext(context)
+      );
+    }
+    return false;
+  }
+  if (type === "ArithCommandSubst" && typeof node.command === "string") {
+    return stringCommandContainsBlockedCommand(node.command, context);
+  }
+  for (const value of Object.values(node)) {
+    if (Array.isArray(value)) {
+      for (const item of value) {
+        if (typeof item === "object" && item !== null) {
+          if (partContainsBlockedCommand(item, context)) {
+            return true;
+          }
+        }
+      }
+      continue;
+    }
+    if (typeof value === "object" && value !== null) {
+      if (partContainsBlockedCommand(value, context)) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+function functionInvocationContainsBlockedCommand(functionName, context) {
+  const definition = context.functionDefinitions.get(functionName);
+  if (!definition) {
+    return false;
+  }
+  if (context.callStack.has(functionName)) {
+    return false;
+  }
+  const invocationContext = cloneInspectionContext(context);
+  invocationContext.callStack.add(functionName);
+  return commandContainsBlockedCommand(definition.body, invocationContext);
+}
+function commandContainsBlockedCommand(command, context) {
+  switch (command.type) {
+    case "SimpleCommand":
+      return isBlockedSimpleCommand(command, context);
+    case "If":
+      return command.clauses.some(
+        (clause) => statementsContainBlockedCommand(
+          clause.condition,
+          cloneInspectionContext(context)
+        ) || statementsContainBlockedCommand(
+          clause.body,
+          cloneInspectionContext(context)
+        )
+      ) || (command.elseBody ? statementsContainBlockedCommand(
+        command.elseBody,
+        cloneInspectionContext(context)
+      ) : false);
+    case "For":
+    case "CStyleFor":
+      return statementsContainBlockedCommand(
+        command.body,
+        cloneInspectionContext(context)
+      );
+    case "While":
+    case "Until":
+      return statementsContainBlockedCommand(
+        command.condition,
+        cloneInspectionContext(context)
+      ) || statementsContainBlockedCommand(
+        command.body,
+        cloneInspectionContext(context)
+      );
+    case "Case":
+      return command.items.some(
+        (item) => statementsContainBlockedCommand(
+          item.body,
+          cloneInspectionContext(context)
+        )
+      );
+    case "Subshell":
+    case "Group":
+      return statementsContainBlockedCommand(
+        command.body,
+        cloneInspectionContext(context)
+      );
+    case "FunctionDef":
+      return false;
+    case "ArithmeticCommand":
+    case "ConditionalCommand":
+      return false;
+    default: {
+      const _unreachable = command;
+      return _unreachable;
+    }
+  }
+}
+function isBlockedSimpleCommand(command, context) {
+  if (wordContainsBlockedCommand(command.name, context)) {
+    return true;
+  }
+  if (command.args.some((arg) => wordContainsBlockedCommand(arg, context))) {
+    return true;
+  }
+  if (command.assignments.some(
+    (assignment) => wordContainsBlockedCommand(assignment.value, context) || (assignment.array?.some(
+      (value) => wordContainsBlockedCommand(value, context)
+    ) ?? false)
+  )) {
+    return true;
+  }
+  if (command.redirections.some((redirection) => {
+    if (redirection.target.type === "Word") {
+      return wordContainsBlockedCommand(
+        redirection.target,
+        context
+      );
+    }
+    if (redirection.target.type === "HereDoc" && redirection.target.content) {
+      return wordContainsBlockedCommand(redirection.target.content, context);
+    }
+    return false;
+  })) {
+    return true;
+  }
+  const commandName = asStaticWordText(command.name);
+  if (!commandName) {
+    return false;
+  }
+  const normalizedName = path.posix.basename(commandName).toLowerCase();
+  if (BLOCKED_DB_CLIENT_COMMANDS.has(normalizedName)) {
+    return true;
+  }
+  if (BLOCKED_RAW_SQL_COMMANDS.has(normalizedName)) {
+    return true;
+  }
+  if (normalizedName === "sql") {
+    const subcommand = asStaticWordText(command.args[0])?.toLowerCase();
+    return !subcommand || !ALLOWED_SQL_PROXY_SUBCOMMANDS.has(subcommand);
+  }
+  if (functionInvocationContainsBlockedCommand(commandName, context)) {
+    return true;
+  }
+  return false;
+}
+function getSqlProxyEnforcementResult(command) {
+  const trimmed = command.trim();
+  if (!trimmed) {
+    return null;
+  }
+  let script;
+  try {
+    script = parse(trimmed);
+  } catch {
+    return null;
+  }
+  const blocked = scriptContainsBlockedCommand(script, {
+    functionDefinitions: /* @__PURE__ */ new Map(),
+    callStack: /* @__PURE__ */ new Set()
+  });
+  if (!blocked) {
+    return null;
+  }
+  return {
+    stdout: "",
+    stderr: `${SQL_PROXY_ENFORCEMENT_MESSAGE}
+`,
+    exitCode: 1
+  };
+}
 async function createResultTools(options) {
   const { adapter, skillMounts, filesystem: baseFs } = options;
   const metaStore = new AsyncLocalStorage();
@@ -685,6 +980,16 @@ async function createResultTools(options) {
       return { result };
     }
   });
+  const guardedSandbox = {
+    ...sandbox,
+    executeCommand: async (command) => {
+      const blockedResult = getSqlProxyEnforcementResult(command);
+      if (blockedResult) {
+        return blockedResult;
+      }
+      return sandbox.executeCommand(command);
+    }
+  };
   const bash = tool2({
     ...tools2.bash,
     inputSchema: z3.object({
@@ -696,6 +1001,10 @@ async function createResultTools(options) {
       if (!execute) {
         throw new Error("bash tool execution is not available");
       }
+      const blockedResult = getSqlProxyEnforcementResult(command);
+      if (blockedResult) {
+        return blockedResult;
+      }
       return metaStore.run({}, async () => {
         const result = await execute({ command }, execOptions);
         const meta = metaStore.getStore()?.value;
@@ -703,12 +1012,15 @@ async function createResultTools(options) {
       });
     },
     toModelOutput: ({ output }) => {
+      if (typeof output !== "object" || output === null) {
+        return { type: "json", value: output };
+      }
       const { meta, ...rest } = output;
       return { type: "json", value: rest };
     }
   });
   return {
-    sandbox,
+    sandbox: guardedSandbox,
     tools: {
       ...tools2,
       bash
@@ -779,8 +1091,7 @@ async function toSql(options) {
           role: "You are an expert SQL query generator. You translate natural language questions into precise, efficient SQL queries based on the provided database schema.",
           objective: "Translate natural language questions into precise, efficient SQL queries"
         }),
-        ...options.instructions,
-        ...options.schemaFragments
+        ...options.fragments
       );
       if (errors.length) {
         context.set(
@@ -801,19 +1112,21 @@ async function toSql(options) {
       const sqlOutput = structuredOutput2({
         model,
         context,
-        schema: z4.union([
-          z4.object({
-            sql: z4.string().describe("The SQL query that answers the question"),
-            reasoning: z4.string().optional().describe("The reasoning steps taken to generate the SQL")
-          }),
-          z4.object({
-            error: z4.string().describe(
-              "Error message explaining why the question cannot be answered with the given schema"
-            )
-          })
-        ])
+        schema: z4.object({
+          result: z4.union([
+            z4.object({
+              sql: z4.string().describe("The SQL query that answers the question"),
+              reasoning: z4.string().describe("The reasoning steps taken to generate the SQL")
+            }),
+            z4.object({
+              error: z4.string().describe(
+                "Error message explaining why the question cannot be answered with the given schema"
+              )
+            })
+          ])
+        })
       });
-      const output = await sqlOutput.generate();
+      const { result: output } = await sqlOutput.generate();
       if ("error" in output) {
         throw new UnanswerableSQLError(output.error);
       }
@@ -843,6 +1156,16 @@ function formatErrorMessage(error) {
   }
   return error.message;
 }
+function isModelUnavailableError(error) {
+  if (!APICallError.isInstance(error)) {
+    return false;
+  }
+  const message2 = error.message.toLowerCase();
+  const responseBody = (error.responseBody ?? "").toLowerCase();
+  const is404ModelError = error.statusCode === 404 && (message2.includes("model") || responseBody.includes("model_not_found"));
+  const errorCode = typeof error.data === "object" && error.data !== null && "error" in error.data && typeof error.data.error === "object" && error.data.error !== null && "code" in error.data.error && typeof error.data.error.code === "string" ? error.data.error.code.toLowerCase() : void 0;
+  return is404ModelError || errorCode === "model_not_found" || responseBody.includes('"code":"model_not_found"') || message2.includes("model") && message2.includes("does not exist or you do not have access to it");
+}
 async function withRetry(computation, options = { retries: 3 }) {
   const errors = [];
   let attempts = 0;
@@ -856,6 +1179,9 @@ async function withRetry(computation, options = { retries: 3 }) {
         if (UnanswerableSQLError.isInstance(context.error)) {
           return false;
         }
+        if (isModelUnavailableError(context.error)) {
+          return false;
+        }
         if (SQLValidationError.isInstance(context.error)) {
           return true;
         }
@@ -4302,8 +4628,7 @@ var Text2Sql = class {
     const result = await toSql({
       input,
       adapter: this.#config.adapter,
-      schemaFragments,
-      instructions: [],
+      fragments: schemaFragments,
       model: this.#config.model
     });
     return result.sql;