[PR 10] Add validateSchemaNode Error Handling for DML

frontend/internal-packages/agent/src/chat/workflow/constants.ts

@@ -17,4 +17,9 @@ export const WORKFLOW_RETRY_CONFIG = {
    * by going back to the designSchema node with the error information
    */
   MAX_DDL_EXECUTION_RETRIES: 1,
+  /**
+   * Maximum number of retries for DML execution failures
+   * When DML execution fails with retryable errors, the workflow will retry
+   */
+  MAX_DML_EXECUTION_RETRIES: 1,
 } as const

frontend/internal-packages/agent/src/chat/workflow/nodes/validateSchemaNode.test.ts

@@ -231,4 +231,157 @@ describe('validateSchemaNode', () => {
     expect(executeQuery).not.toHaveBeenCalled()
     expect(result).toEqual(state)
   })
+
+  it('should handle missing DML statements', async () => {
+    const state = createMockState({
+      dmlStatements: undefined,
+    })
+
+    const repositories = createMockRepositories()
+    const result = await validateSchemaNode(state, {
+      configurable: { repositories, logger: mockLogger },
+    })
+
+    expect(executeQuery).not.toHaveBeenCalled()
+    expect(result).toEqual(state)
+  })
+
+  it('should retry execution on retryable errors', async () => {
+    const mockResults: SqlResult[] = [
+      {
+        success: false,
+        sql: 'INSERT INTO users VALUES (1);',
+        result: { error: 'foreign key constraint violation' },
+        id: 'result-1',
+        metadata: {
+          executionTime: 5,
+          timestamp: new Date().toISOString(),
+        },
+      },
+    ]
+
+    vi.mocked(executeQuery).mockResolvedValue(mockResults)
+
+    const state = createMockState({
+      dmlStatements: 'INSERT INTO users VALUES (1);',
+    })
+
+    const repositories = createMockRepositories()
+    const result = await validateSchemaNode(state, {
+      configurable: { repositories, logger: mockLogger },
+    })
+
+    expect(result.shouldRetryDmlExecution).toBe(true)
+    expect(result.dmlRetryReason).toContain('foreign key constraint violation')
+    expect(result.retryCount['dmlExecutionRetry']).toBe(1)
+    expect(mockLogger.log).toHaveBeenCalledWith(
+      '[validateSchemaNode] Execution failed with retryable error, scheduling retry',
+    )
+  })
+
+  it('should not retry execution on non-retryable errors', async () => {
+    const mockResults: SqlResult[] = [
+      {
+        success: false,
+        sql: 'INSERT INTO users VALUES;',
+        result: { error: 'syntax error at or near "VALUES"' },
+        id: 'result-1',
+        metadata: {
+          executionTime: 2,
+          timestamp: new Date().toISOString(),
+        },
+      },
+    ]
+
+    vi.mocked(executeQuery).mockResolvedValue(mockResults)
+
+    const state = createMockState({
+      dmlStatements: 'INSERT INTO users VALUES;',
+    })
+
+    const repositories = createMockRepositories()
+    const result = await validateSchemaNode(state, {
+      configurable: { repositories, logger: mockLogger },
+    })
+
+    expect(result.shouldRetryDmlExecution).toBeUndefined()
+    expect(result.dmlExecutionErrors).toContain('syntax error')
+    expect(mockLogger.log).toHaveBeenCalledWith(
+      '[validateSchemaNode] Completed with errors',
+    )
+  })
+
+  it('should not retry after max retries exceeded', async () => {
+    const mockResults: SqlResult[] = [
+      {
+        success: false,
+        sql: 'INSERT INTO users VALUES (1);',
+        result: { error: 'foreign key constraint violation' },
+        id: 'result-1',
+        metadata: {
+          executionTime: 5,
+          timestamp: new Date().toISOString(),
+        },
+      },
+    ]
+
+    vi.mocked(executeQuery).mockResolvedValue(mockResults)
+
+    const state = createMockState({
+      dmlStatements: 'INSERT INTO users VALUES (1);',
+      retryCount: { dmlExecutionRetry: 1 }, // Already at max retries
+    })
+
+    const repositories = createMockRepositories()
+    const result = await validateSchemaNode(state, {
+      configurable: { repositories, logger: mockLogger },
+    })
+
+    expect(result.shouldRetryDmlExecution).toBeUndefined()
+    expect(result.dmlExecutionErrors).toContain(
+      'foreign key constraint violation',
+    )
+    expect(result.retryCount['dmlExecutionRetry']).toBe(1)
+  })
+
+  it('should handle mixed success and failure in execution', async () => {
+    const mockResults: SqlResult[] = [
+      {
+        success: true,
+        sql: 'INSERT INTO categories VALUES (1, "test");',
+        result: { rows: [], columns: [] },
+        id: 'result-1',
+        metadata: {
+          executionTime: 5,
+          timestamp: new Date().toISOString(),
+          affectedRows: 1,
+        },
+      },
+      {
+        success: false,
+        sql: 'INSERT INTO products VALUES (1, 999);',
+        result: { error: 'foreign key constraint violation on category_id' },
+        id: 'result-2',
+        metadata: {
+          executionTime: 3,
+          timestamp: new Date().toISOString(),
+        },
+      },
+    ]
+
+    vi.mocked(executeQuery).mockResolvedValue(mockResults)
+
+    const state = createMockState({
+      dmlStatements:
+        'INSERT INTO categories VALUES (1, "test"); INSERT INTO products VALUES (1, 999);',
+    })
+
+    const repositories = createMockRepositories()
+    const result = await validateSchemaNode(state, {
+      configurable: { repositories, logger: mockLogger },
+    })
+
+    expect(result.shouldRetryDmlExecution).toBe(true)
+    expect(result.dmlRetryReason).toContain('foreign key constraint violation')
+  })
 })

frontend/internal-packages/agent/src/chat/workflow/nodes/validateSchemaNode.ts

@@ -1,9 +1,27 @@
 import type { RunnableConfig } from '@langchain/core/runnables'
 import { executeQuery } from '@liam-hq/pglite-server'
 import type { SqlResult } from '@liam-hq/pglite-server/src/types'
+import { WORKFLOW_RETRY_CONFIG } from '../constants'
 import { getConfigurable } from '../shared/getConfigurable'
 import type { WorkflowState } from '../types'
 
+const NODE_NAME = 'validateSchemaNode'
+
+// Error patterns that might benefit from retry
+const RETRYABLE_ERROR_PATTERNS = [
+  'constraint violation',
+  'foreign key',
+  'duplicate key',
+  'deadlock',
+  'lock timeout',
+]
+
+function isRetryableError(errorMessage: string): boolean {
+  const lowerError = errorMessage.toLowerCase()
+  return RETRYABLE_ERROR_PATTERNS.some((pattern) =>
+    lowerError.includes(pattern),
+  )
+}
 /**
  * Validate Schema Node - Combined DDL/DML Execution & Validation
  * Executes DDL and DML together in a single query to validate schema with test data
@@ -54,6 +72,39 @@ export async function validateSchemaNode(
       )
       .join('; ')
 
+    configurableResult.value.logger.error(
+      `[${NODE_NAME}] Execution failed: ${errorMessages}`,
+    )
+
+    // Check if errors are retryable
+    const hasRetryableError = isRetryableError(errorMessages)
+    const currentRetryCount = state.retryCount?.['dmlExecutionRetry'] || 0
+
+    if (
+      hasRetryableError &&
+      currentRetryCount < WORKFLOW_RETRY_CONFIG.MAX_DML_EXECUTION_RETRIES
+    ) {
+      configurableResult.value.logger.log(
+        `[${NODE_NAME}] Execution failed with retryable error, scheduling retry`,
+      )
+      configurableResult.value.logger.log(
+        `[${NODE_NAME}] Completed with retry scheduled`,
+      )
+
+      return {
+        ...state,
+        dmlExecutionErrors: errorMessages,
+        shouldRetryDmlExecution: true,
+        dmlRetryReason: errorMessages,
+        retryCount: {
+          ...state.retryCount,
+          dmlExecutionRetry: currentRetryCount + 1,
+        },
+      }
+    }
+
+    // Non-retryable error or max retries exceeded
+    configurableResult.value.logger.log(`[${NODE_NAME}] Completed with errors`)
     return {
       ...state,
       dmlExecutionErrors: errorMessages,

frontend/internal-packages/agent/src/chat/workflow/types.ts

@@ -33,6 +33,8 @@ export type WorkflowState = {
   // DML execution results
   dmlExecutionSuccessful?: boolean | undefined
   dmlExecutionErrors?: string | undefined
+  shouldRetryDmlExecution?: boolean | undefined
+  dmlRetryReason?: string | undefined
 
   // Schema update fields
   buildingSchemaId: string