feat: implement DB Agent routing in LangGraph workflow

frontend/internal-packages/agent/README.md

@@ -21,6 +21,7 @@ graph TD;
 	__start__ -.-> validateInitialSchema;
 	__start__ -.-> leadAgent;
 	leadAgent -.-> pmAgent;
+	leadAgent -.-> dbAgent;
 	leadAgent -.-> __end__;
 	classDef default fill:#f2f0ff,line-height:1.2;
 	classDef first fill-opacity:0;

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

@@ -5,15 +5,18 @@
  * Important: Node retries do NOT count toward this limit. The limit only
  * applies to transitions between nodes.
  *
- * The workflow has 8 nodes:
- * - Normal execution: 9 transitions (START → 8 nodes → END)
- * - With error loops: May have additional transitions when errors occur
- *   (e.g., validateSchema → designSchema)
+ * TEMPORARY LIMITATION (set to 10):
+ * Due to issues with schemaDesignTool, the DB Agent cannot resolve schema issues
+ * on the second and subsequent attempts, causing infinite loops between:
+ * leadAgent → dbAgent → qaAgent → leadAgent (when schemaIssues exist)
  *
- * Setting this to 50 ensures:
- * - Complete workflow execution under normal conditions
- * - Ample headroom for complex error handling loops and retries
- * - Protection against infinite loops while allowing for complex workflows
- * - Sufficient capacity for finding optimal workflow patterns
+ * Current behavior with limit=10:
+ * - Allows multiple iterations of: PM Agent → DB Agent → QA Agent → Lead Agent → DB Agent
+ * - The workflow will fail after 10 loops if issues persist
+ * - Provides more opportunities for the DB Agent to refine the schema
+ *
+ * TODO: Increase this limit after fixing schemaDesignTool to properly handle
+ * schema modifications (e.g., unique constraint issues, JSON patch errors)
+ * See: route06/liam-internal#5642
  */
-export const DEFAULT_RECURSION_LIMIT = 50
+export const DEFAULT_RECURSION_LIMIT = 10

frontend/internal-packages/agent/src/createGraph.integration.test.ts

@@ -6,6 +6,7 @@ import {
   getTestConfig,
   outputStreamEvents,
 } from '../test-utils/workflowTestHelpers'
+import { DEFAULT_RECURSION_LIMIT } from './constants'
 import { createGraph } from './createGraph'
 import type { WorkflowState } from './types'
 
@@ -36,6 +37,7 @@ describe('createGraph Integration', () => {
     // Act
     const streamEvents = graph.streamEvents(initialState, {
       ...config,
+      recursionLimit: DEFAULT_RECURSION_LIMIT,
       streamMode: 'messages',
       version: 'v2',
       subgraphs: true,

frontend/internal-packages/agent/src/createGraph.test.ts

@@ -20,6 +20,7 @@ graph TD;
 	__start__ -.-> validateInitialSchema;
 	__start__ -.-> leadAgent;
 	leadAgent -.-> pmAgent;
+	leadAgent -.-> dbAgent;
 	leadAgent -.-> __end__;
 	classDef default fill:#f2f0ff,line-height:1.2;
 	classDef first fill-opacity:0;

frontend/internal-packages/agent/src/createGraph.ts

@@ -36,7 +36,8 @@ export const createGraph = (checkpointer?: BaseCheckpointSaver) => {
     const modifiedState = { ...state, messages: [], prompt }
     const output = await dbAgentSubgraph.invoke(modifiedState, config)
 
-    return { ...state, ...output }
+    // Clear schemaIssues after DB agent processing to prevent infinite loops
+    return { ...state, ...output, schemaIssues: [] }
   }
 
   const callQaAgent = async (state: WorkflowState, config: RunnableConfig) => {
@@ -92,6 +93,7 @@ export const createGraph = (checkpointer?: BaseCheckpointSaver) => {
 
     .addConditionalEdges('leadAgent', (state) => state.next, {
       pmAgent: 'pmAgent',
+      dbAgent: 'dbAgent',
       [END]: END,
     })
     .addEdge('pmAgent', 'dbAgent')

frontend/internal-packages/agent/src/db-agent/utils/convertAnalyzedRequirementsToPrompt.test.ts

@@ -71,7 +71,7 @@ describe('convertAnalyzedRequirementsToPrompt', () => {
 
       Design a user management system with authentication
 
-      ## Test Cases
+      ## Requirements
 
       - authentication: User login (SELECT), User logout (UPDATE), Password reset (UPDATE)
       - userManagement: Create new user (INSERT), Update user info (UPDATE), Delete user (DELETE)"
@@ -96,7 +96,7 @@ describe('convertAnalyzedRequirementsToPrompt', () => {
 
       Build a simple system
 
-      ## Test Cases"
+      ## Requirements"
     `)
   })
 
@@ -128,7 +128,7 @@ describe('convertAnalyzedRequirementsToPrompt', () => {
 
       Add a basic feature
 
-      ## Test Cases
+      ## Requirements
 
       - basic: Basic feature test (INSERT)"
     `)
@@ -156,9 +156,13 @@ describe('convertAnalyzedRequirementsToPrompt', () => {
 
         Design a user management system with authentication
 
-        ## Test Cases
+        ## Requirements
 
-        - authentication: User logout (UPDATE)"
+        - authentication: User logout (UPDATE)
+
+        ## Schema Issues to Fix
+
+        1. Missing logout table"
       `)
     })
 
@@ -185,7 +189,7 @@ describe('convertAnalyzedRequirementsToPrompt', () => {
 
         Design a user management system with authentication
 
-        ## Test Cases
+        ## Requirements
 
         - authentication: User login (SELECT), User logout (UPDATE), Password reset (UPDATE)
         - userManagement: Create new user (INSERT), Update user info (UPDATE), Delete user (DELETE)"
@@ -213,9 +217,13 @@ describe('convertAnalyzedRequirementsToPrompt', () => {
 
         Design a user management system with authentication
 
-        ## Test Cases
+        ## Requirements
+
+        - authentication: User login (SELECT)
+
+        ## Schema Issues to Fix
 
-        - authentication: User login (SELECT)"
+        1. Login form missing"
       `)
     })
 
@@ -243,7 +251,12 @@ describe('convertAnalyzedRequirementsToPrompt', () => {
 
         Design a user management system with authentication
 
-        ## Test Cases"
+        ## Requirements
+
+
+        ## Schema Issues to Fix
+
+        1. Non-existent testcase issue"
       `)
     })
 

frontend/internal-packages/agent/src/db-agent/utils/convertAnalyzedRequirementsToPrompt.ts

@@ -37,6 +37,11 @@ export const convertRequirementsToPrompt = (
     )
     .join('\n')
 
+  const schemaIssuesSection =
+    schemaIssues && schemaIssues.length > 0
+      ? `\n\n## Schema Issues to Fix\n\n${schemaIssues.map((issue, index) => `${index + 1}. ${issue.description}`).join('\n')}`
+      : ''
+
   return `## Session Goal
 
 ${requirements.goal}
@@ -45,7 +50,7 @@ ${requirements.goal}
 
 ${userInput}
 
-## Test Cases
+## Requirements
 
-${testCasesSection}`.trim()
+${testCasesSection}${schemaIssuesSection}`.trim()
 }

frontend/internal-packages/agent/src/lead-agent/classifyMessage/index.test.ts

@@ -0,0 +1,124 @@
+import { END } from '@langchain/langgraph'
+import { describe, expect, it } from 'vitest'
+import type { WorkflowState } from '../../types'
+import { classifyMessage } from './index'
+
+const createMockState = (
+  overrides: Partial<WorkflowState> = {},
+): WorkflowState => ({
+  messages: [],
+  analyzedRequirements: {
+    goal: '',
+    testcases: {},
+  },
+  schemaIssues: [],
+  schemaData: { tables: {}, enums: {}, extensions: {} },
+  buildingSchemaId: 'test-schema-id',
+  organizationId: 'test-org-id',
+  userId: 'test-user-id',
+  designSessionId: 'test-session-id',
+  next: END,
+  ...overrides,
+})
+
+describe('classifyMessage', () => {
+  it('should route to dbAgent when schema issues exist', async () => {
+    const state = createMockState({
+      schemaIssues: [
+        {
+          testcaseId: 'test-1',
+          description: 'Missing foreign key constraint',
+        },
+      ],
+    })
+
+    const result = await classifyMessage(state)
+
+    expect(result.goto).toEqual([END])
+    expect(result.update).toEqual({ next: 'dbAgent' })
+  })
+
+  it('should route to dbAgent when multiple schema issues exist', async () => {
+    const state = createMockState({
+      schemaIssues: [
+        {
+          testcaseId: 'test-1',
+          description: 'Missing foreign key constraint',
+        },
+        { testcaseId: 'test-2', description: 'Invalid table name' },
+      ],
+    })
+
+    const result = await classifyMessage(state)
+
+    expect(result.goto).toEqual([END])
+    expect(result.update).toEqual({ next: 'dbAgent' })
+  })
+
+  it('should prioritize schema issues over QA completion', async () => {
+    const state = createMockState({
+      schemaIssues: [
+        { testcaseId: 'test-1', description: 'Schema issue found' },
+      ],
+      analyzedRequirements: {
+        goal: 'Test goal',
+        testcases: {
+          functional: [
+            {
+              id: 'test-1',
+              title: 'Test case',
+              type: 'SELECT',
+              sql: 'SELECT 1',
+              testResults: [],
+            },
+          ],
+        },
+      },
+    })
+
+    const result = await classifyMessage(state)
+
+    expect(result.goto).toEqual([END])
+    expect(result.update).toEqual({ next: 'dbAgent' })
+  })
+
+  it('should route to summarizeWorkflow when QA completed and no schema issues', async () => {
+    const state = createMockState({
+      schemaIssues: [],
+      analyzedRequirements: {
+        goal: 'Test goal',
+        testcases: {
+          functional: [
+            {
+              id: 'test-1',
+              title: 'Test case',
+              type: 'SELECT',
+              sql: 'SELECT 1',
+              testResults: [],
+            },
+          ],
+        },
+      },
+    })
+
+    const result = await classifyMessage(state)
+
+    expect(result.goto).toEqual(['summarizeWorkflow'])
+    expect(result.update).toBeUndefined()
+  })
+
+  it('should route to pmAgent when no schema issues and QA not completed', async () => {
+    const state = createMockState({
+      schemaIssues: [],
+      analyzedRequirements: {
+        goal: 'Test goal',
+        testcases: {},
+      },
+    })
+
+    const result = await classifyMessage(state)
+
+    expect(result.goto).toEqual([END])
+    expect(result.update).toEqual({ next: 'pmAgent' })
+  })
+})

frontend/internal-packages/agent/src/lead-agent/classifyMessage/index.ts

@@ -1,14 +1,22 @@
 import { Command, END } from '@langchain/langgraph'
 import type { WorkflowState } from '../../types'
-import { isQACompleted } from '../utils/workflowStatus'
+import { isQACompleted, shouldRouteDBAgent } from '../utils/workflowStatus'
 
 export async function classifyMessage(state: WorkflowState): Promise<Command> {
-  // 1. Check if QA is completed first (highest priority)
+  // 1. Check if DB Agent routing is needed (highest priority)
+  if (shouldRouteDBAgent(state)) {
+    return new Command({
+      update: { next: 'dbAgent' },
+      goto: END,
+    })
+  }
+
+  // 2. Check if QA is completed (second priority)
   if (isQACompleted(state)) {
     return new Command({ goto: 'summarizeWorkflow' })
   }
 
-  // 2. Direct routing to pmAgent
+  // 3. Direct routing to pmAgent (default)
   return new Command({
     update: { next: 'pmAgent' },
     goto: END,

frontend/internal-packages/agent/src/lead-agent/utils/workflowStatus.ts

@@ -14,3 +14,11 @@ export function isQACompleted(state: WorkflowState): boolean {
 
   return false
 }
+
+function hasSchemaIssues(state: WorkflowState): boolean {
+  return state.schemaIssues.length > 0
+}
+
+export function shouldRouteDBAgent(state: WorkflowState): boolean {
+  return hasSchemaIssues(state)
+}

frontend/internal-packages/agent/src/qa-agent/shared/qaAgentAnnotation.ts

@@ -3,7 +3,7 @@ import type { AnalyzedRequirements } from '@liam-hq/artifact'
 import type { Schema } from '@liam-hq/schema'
 import {
   generatedSqlsAnnotation,
-  schemaIssuesAnnotation,
+  qaSchemaIssuesAnnotation,
 } from '../testcaseGeneration/testcaseAnnotation'
 
 /**
@@ -27,7 +27,7 @@ export const qaAgentAnnotation = Annotation.Root({
   }),
   designSessionId: Annotation<string>,
   buildingSchemaId: Annotation<string>,
-  schemaIssues: schemaIssuesAnnotation,
+  schemaIssues: qaSchemaIssuesAnnotation,
   generatedSqls: generatedSqlsAnnotation,
   next: Annotation<string>({
     reducer: (x, y) => y ?? x ?? END,

frontend/internal-packages/agent/src/qa-agent/testcaseGeneration/testcaseAnnotation.ts

@@ -1,13 +1,20 @@
 import { Annotation, MessagesAnnotation } from '@langchain/langgraph'
 import type { Schema } from '@liam-hq/schema'
+import type { SchemaIssue } from '../../workflowSchemaIssuesAnnotation'
 import type { TestCaseData } from '../distributeRequirements'
 
-type SchemaIssue = {
-  testcaseId: string
-  description: string
-}
-
-export const schemaIssuesAnnotation = Annotation<Array<SchemaIssue>>({
+/**
+ * Schema issues annotation for QA agent's parallel processing.
+ *
+ * Uses concat reducer because:
+ * - Multiple testcase nodes run in parallel
+ * - Each node may discover different schema issues
+ * - All issues must be collected together
+ *
+ * This is different from workflow-level annotation which uses replacement
+ * for clearing issues after DB agent processing.
+ */
+export const qaSchemaIssuesAnnotation = Annotation<Array<SchemaIssue>>({
   reducer: (prev, next) => prev.concat(next),
   default: () => [],
 })
@@ -27,6 +34,6 @@ export const testcaseAnnotation = Annotation.Root({
   currentTestcase: Annotation<TestCaseData>,
   schemaData: Annotation<Schema>,
   goal: Annotation<string>,
-  schemaIssues: schemaIssuesAnnotation,
+  schemaIssues: qaSchemaIssuesAnnotation,
   generatedSqls: generatedSqlsAnnotation,
 })