Merge pull request #1532 from kirill-markin/loop-claude/lean-agent-sql-write-payload

Shrink agent SQL write payloads and stop post-commit write failures

Author: kirill-markin

apps/backend/src/aiTools/agentSql.ts

@@ -15,8 +15,10 @@ import { executeSqlReadBatch, executeSqlReadStatement } from "./agentSql/readExe
 import {
   isSqlMutationStatement,
   isSqlReadStatement,
+  type AgentSqlBatchExecutionResult,
   type AgentSqlContext,
   type AgentSqlExecutionResult,
+  type AgentSqlMutationExecutionResult,
   type AgentSqlPayload,
   type AgentSqlSinglePayload,
 } from "./agentSql/shared";
@@ -85,15 +87,18 @@ function toStatementSqls(sql: string, statements: ReadonlyArray<ParsedSqlStateme
 }
 
 /**
- * Single source of truth for the agent SQL result-size budget shared by both
- * the MCP surface (`sql_query` / `sql_execute`) and the REST surface
- * (`POST /agent/sql/query` / `POST /agent/sql/execute`).
+ * Read-path result-size budget shared by the MCP surface (`sql_query`) and the
+ * REST surface (`POST /agent/sql/query`).
  *
  * The agent envelope serializes `result.data`, so the budget is measured
  * against the serialized `data` payload. On overflow we fail with an actionable
  * error (matching the repo's "clear, actionable errors / no silent fallbacks"
  * principle) instead of returning a payload that exceeds the directory's
- * tool-result token limit. The remedies are concrete: narrow the result set.
+ * tool-result token limit. Nothing is committed on a read, and the remedies are
+ * concrete: narrow the result set.
+ *
+ * Writes must never reach this: their transaction is already committed when the
+ * size is measured, so they drop rows instead (see the reducers below).
  */
 function assertSqlResultWithinSizeBudget<T extends AgentSqlExecutionResult>(result: T): T {
   const serializedLength = JSON.stringify(result.data).length;
@@ -108,6 +113,57 @@ function assertSqlResultWithinSizeBudget<T extends AgentSqlExecutionResult>(resu
   return result;
 }
 
+/**
+ * Appended to the instructions of a committed write whose payload had to shrink,
+ * so the model knows the rows are missing by design and how to get them.
+ */
+const OMITTED_MUTATION_ROWS_INSTRUCTION =
+  "The affected rows were omitted from this result because the payload exceeded the result-size budget. The write itself succeeded, so do not repeat it: run a follow-up SELECT query if you need the affected rows.";
+
+/**
+ * Shrinks an oversized committed write result instead of rejecting it.
+ *
+ * `executeSqlMutationStatement` and `executeSqlMutationBatch` return after their
+ * transaction committed, so answering with `QUERY_RESULT_TOO_LARGE` would report
+ * a successful write as a failure and invite the caller to retry it and
+ * duplicate the data. Only the returned rows are dropped; the counts, the echoed
+ * SQL, and the atomicity contract stay intact.
+ */
+function reduceMutationResultToSizeBudget(
+  result: AgentSqlMutationExecutionResult,
+): AgentSqlMutationExecutionResult {
+  if (JSON.stringify(result.data).length <= MAX_SQL_RESULT_CHARS) {
+    return result;
+  }
+
+  return {
+    data: {
+      ...result.data,
+      rows: [],
+    },
+    instructions: `${result.instructions} ${OMITTED_MUTATION_ROWS_INSTRUCTION}`,
+  };
+}
+
+function reduceBatchMutationResultToSizeBudget(
+  result: AgentSqlBatchExecutionResult,
+): AgentSqlBatchExecutionResult {
+  if (JSON.stringify(result.data).length <= MAX_SQL_RESULT_CHARS) {
+    return result;
+  }
+
+  return {
+    data: {
+      ...result.data,
+      statements: result.data.statements.map((statement) => ({
+        ...statement,
+        rows: [],
+      })),
+    },
+    instructions: `${result.instructions} ${OMITTED_MUTATION_ROWS_INSTRUCTION}`,
+  };
+}
+
 function getAgentSqlFingerprint(sql: string): string {
   return createHash("sha256")
     .update(sql)
@@ -320,12 +376,12 @@ export async function runSqlExecute(
 
     if (statements.every(isSqlMutationStatement)) {
       if (statements.length === 1) {
-        return assertSqlResultWithinSizeBudget(
+        return reduceMutationResultToSizeBudget(
           await executeSqlMutationStatement(dependencies, context, sql, statements[0]),
         );
       }
 
-      return assertSqlResultWithinSizeBudget(
+      return reduceBatchMutationResultToSizeBudget(
         await executeSqlMutationBatch(dependencies, context, sql, statements, statementSqls),
       );
     }

apps/backend/src/aiTools/agentSql/batchMutation.ts

@@ -31,13 +31,13 @@ import {
   buildCreateDeckInput,
   makeBatchNormalizedSql,
   requireSqlMutationTargetIds,
+  toCardIdRows,
   toCardRow,
-  toCreatedCardRows,
-  toCreatedDeckRows,
+  toDeckIdRows,
   toDeckRow,
   wrapBatchExecutionError,
+  type AgentSqlBatchExecutionResult,
   type AgentSqlContext,
-  type AgentSqlExecutionResult,
   type AgentSqlMutationAssignment,
   type AgentSqlMutationStatement,
   type AgentSqlSinglePayload,
@@ -233,7 +233,7 @@ export async function executeSqlMutationBatch(
   sql: string,
   statements: ReadonlyArray<AgentSqlMutationStatement>,
   statementSqls: ReadonlyArray<string>,
-): Promise<AgentSqlExecutionResult> {
+): Promise<AgentSqlBatchExecutionResult> {
   const replicaId = await dependencies.ensureAgentSyncReplica(
     context.workspaceId,
     context.userId,
@@ -273,9 +273,7 @@ export async function executeSqlMutationBatch(
           payloads.push({
             statementType: "insert",
             resource: "cards",
-            sql: statementSql,
-            normalizedSql: statement.normalizedSql,
-            rows: toCreatedCardRows(createdCards),
+            rows: toCardIdRows(createdCards),
             affectedCount: createdCards.length,
           });
           continue;
@@ -309,9 +307,7 @@ export async function executeSqlMutationBatch(
           payloads.push({
             statementType: "insert",
             resource: "decks",
-            sql: statementSql,
-            normalizedSql: statement.normalizedSql,
-            rows: toCreatedDeckRows(createdDecks),
+            rows: toDeckIdRows(createdDecks),
             affectedCount: createdDecks.length,
           });
           continue;
@@ -351,9 +347,7 @@ export async function executeSqlMutationBatch(
           payloads.push({
             statementType: "update",
             resource: "cards",
-            sql: statementSql,
-            normalizedSql: statement.normalizedSql,
-            rows: toCreatedCardRows(updatedCards),
+            rows: toCardIdRows(updatedCards),
             affectedCount: updatedCards.length,
           });
           continue;
@@ -385,9 +379,7 @@ export async function executeSqlMutationBatch(
           payloads.push({
             statementType: "update",
             resource: "decks",
-            sql: statementSql,
-            normalizedSql: statement.normalizedSql,
-            rows: toCreatedDeckRows(updatedDecks),
+            rows: toDeckIdRows(updatedDecks),
             affectedCount: updatedDecks.length,
           });
           continue;
@@ -411,8 +403,6 @@ export async function executeSqlMutationBatch(
           payloads.push({
             statementType: "delete",
             resource: "cards",
-            sql: statementSql,
-            normalizedSql: statement.normalizedSql,
             rows: [],
             affectedCount: targetIds.length,
           });
@@ -436,8 +426,6 @@ export async function executeSqlMutationBatch(
         payloads.push({
           statementType: "delete",
           resource: "decks",
-          sql: statementSql,
-          normalizedSql: statement.normalizedSql,
           rows: [],
           affectedCount: targetIds.length,
         });

apps/backend/src/aiTools/agentSql/readExecution.ts

@@ -26,6 +26,7 @@ import {
   type AgentSqlExecutionResult,
   type AgentSqlReadExecutionResult,
   type AgentSqlReadStatement,
+  type AgentSqlReadStatementPayload,
   type AgentSqlSinglePayload,
 } from "./shared";
 
@@ -140,12 +141,19 @@ export async function loadSelectRows(
   return collectReviewEventRows(dependencies, context.userId, context.workspaceId);
 }
 
-export async function executeSqlReadStatement(
+/**
+ * Runs one read statement into a batch-entry payload, which carries no echo of
+ * the submitted SQL. `executeSqlReadStatement` adds that echo for the
+ * single-statement top-level payload.
+ */
+async function executeReadStatementPayload(
   dependencies: AgentToolOperationDependencies,
   context: AgentSqlContext,
-  sql: string,
   statement: AgentSqlReadStatement,
-): Promise<AgentSqlReadExecutionResult> {
+): Promise<Readonly<{
+  data: AgentSqlReadStatementPayload;
+  instructions: string;
+}>> {
   if (statement.type === "show_tables") {
     const rows = getSqlResourceDescriptors()
       .filter((descriptor) => statement.likePattern === null || likePatternToRegExp(statement.likePattern).test(descriptor.resourceName))
@@ -159,8 +167,6 @@ export async function executeSqlReadStatement(
       data: {
         statementType: "show_tables",
         resource: null,
-        sql,
-        normalizedSql: statement.normalizedSql,
         rows,
         rowCount: rows.length,
         limit: null,
@@ -186,8 +192,6 @@ export async function executeSqlReadStatement(
       data: {
         statementType: "describe",
         resource: statement.resourceName,
-        sql,
-        normalizedSql: statement.normalizedSql,
         rows,
         rowCount: rows.length,
         limit: null,
@@ -204,8 +208,6 @@ export async function executeSqlReadStatement(
     data: {
       statementType: "select",
       resource: statement.source.resourceName,
-      sql,
-      normalizedSql: statement.normalizedSql,
       rows: result.rows,
       rowCount: result.rowCount,
       limit: result.limit,
@@ -216,6 +218,24 @@ export async function executeSqlReadStatement(
   };
 }
 
+export async function executeSqlReadStatement(
+  dependencies: AgentToolOperationDependencies,
+  context: AgentSqlContext,
+  sql: string,
+  statement: AgentSqlReadStatement,
+): Promise<AgentSqlReadExecutionResult> {
+  const result = await executeReadStatementPayload(dependencies, context, statement);
+
+  return {
+    data: {
+      ...result.data,
+      sql,
+      normalizedSql: statement.normalizedSql,
+    },
+    instructions: result.instructions,
+  };
+}
+
 export async function executeSqlReadBatch(
   dependencies: AgentToolOperationDependencies,
   context: AgentSqlContext,
@@ -227,12 +247,7 @@ export async function executeSqlReadBatch(
 
   for (const [index, statement] of statements.entries()) {
     try {
-      const result = await executeSqlReadStatement(
-        dependencies,
-        context,
-        statementSqls[index] ?? statement.normalizedSql,
-        statement,
-      );
+      const result = await executeReadStatementPayload(dependencies, context, statement);
       payloads.push(result.data);
     } catch (error) {
       wrapBatchExecutionError(error, index, statementSqls[index] ?? statement.normalizedSql);