diff --git a/frontend/internal-packages/agent/README.md b/frontend/internal-packages/agent/README.md index 5bb80cc57e..d0f8d999af 100644 --- a/frontend/internal-packages/agent/README.md +++ b/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; diff --git a/frontend/internal-packages/agent/src/constants.ts b/frontend/internal-packages/agent/src/constants.ts index c4848a8594..55b4544c50 100644 --- a/frontend/internal-packages/agent/src/constants.ts +++ b/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 diff --git a/frontend/internal-packages/agent/src/createGraph.integration.test.ts b/frontend/internal-packages/agent/src/createGraph.integration.test.ts index 4e709eff56..02fd82caab 100644 --- a/frontend/internal-packages/agent/src/createGraph.integration.test.ts +++ b/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' @@ -35,6 +36,7 @@ describe('createGraph Integration', () => { // Act const streamEvents = graph.streamEvents(initialState, { ...config, + recursionLimit: DEFAULT_RECURSION_LIMIT, streamMode: 'messages', version: 'v2', subgraphs: true, diff --git a/frontend/internal-packages/agent/src/createGraph.test.ts b/frontend/internal-packages/agent/src/createGraph.test.ts index d986653245..70da169a4d 100644 --- a/frontend/internal-packages/agent/src/createGraph.test.ts +++ b/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; diff --git a/frontend/internal-packages/agent/src/createGraph.ts b/frontend/internal-packages/agent/src/createGraph.ts index 1c83a71ebd..67bd14851c 100644 --- a/frontend/internal-packages/agent/src/createGraph.ts +++ b/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') diff --git a/frontend/internal-packages/agent/src/db-agent/utils/convertAnalyzedRequirementsToPrompt.test.ts b/frontend/internal-packages/agent/src/db-agent/utils/convertAnalyzedRequirementsToPrompt.test.ts index 0103a0c246..3f4156a5ad 100644 --- a/frontend/internal-packages/agent/src/db-agent/utils/convertAnalyzedRequirementsToPrompt.test.ts +++ b/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,11 +251,17 @@ describe('convertAnalyzedRequirementsToPrompt', () => { Design a user management system with authentication - ## Test Cases" + ## Requirements + + + + ## Schema Issues to Fix + + 1. Non-existent testcase issue" `) }) - it('should filter testcases without showing IDs or issue details', () => { + it('should filter testcases and show issue details in Schema Issues section', () => { const userInput = 'Design a user management system with authentication' const schemaIssues = [ { testcaseId: '4', description: 'User table structure issue' }, @@ -265,7 +279,8 @@ describe('convertAnalyzedRequirementsToPrompt', () => { ) expect(result).toContain('Create new user') expect(result).not.toContain('[4]') - expect(result).not.toContain('User table structure issue') + expect(result).toContain('## Schema Issues to Fix') + expect(result).toContain('User table structure issue') }) }) }) diff --git a/frontend/internal-packages/agent/src/db-agent/utils/convertAnalyzedRequirementsToPrompt.ts b/frontend/internal-packages/agent/src/db-agent/utils/convertAnalyzedRequirementsToPrompt.ts index 4c06fc5a2f..aaad806f2d 100644 --- a/frontend/internal-packages/agent/src/db-agent/utils/convertAnalyzedRequirementsToPrompt.ts +++ b/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() } diff --git a/frontend/internal-packages/agent/src/lead-agent/classifyMessage/index.test.ts b/frontend/internal-packages/agent/src/lead-agent/classifyMessage/index.test.ts new file mode 100644 index 0000000000..91d3c23c06 --- /dev/null +++ b/frontend/internal-packages/agent/src/lead-agent/classifyMessage/index.test.ts @@ -0,0 +1,123 @@ +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 => ({ + messages: [], + analyzedRequirements: { + goal: '', + testcases: {}, + }, + schemaIssues: [], + schemaData: { tables: {}, enums: {}, extensions: {} }, + 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' }) + }) +}) diff --git a/frontend/internal-packages/agent/src/lead-agent/classifyMessage/index.ts b/frontend/internal-packages/agent/src/lead-agent/classifyMessage/index.ts index 61adeac92b..75dd60385e 100644 --- a/frontend/internal-packages/agent/src/lead-agent/classifyMessage/index.ts +++ b/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 { - // 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, diff --git a/frontend/internal-packages/agent/src/lead-agent/utils/workflowStatus.ts b/frontend/internal-packages/agent/src/lead-agent/utils/workflowStatus.ts index 80082f4a4a..91d6afb1ee 100644 --- a/frontend/internal-packages/agent/src/lead-agent/utils/workflowStatus.ts +++ b/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) +} diff --git a/frontend/internal-packages/agent/src/qa-agent/shared/qaAgentAnnotation.ts b/frontend/internal-packages/agent/src/qa-agent/shared/qaAgentAnnotation.ts index 897d90a4cb..93e2f98d8b 100644 --- a/frontend/internal-packages/agent/src/qa-agent/shared/qaAgentAnnotation.ts +++ b/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' /** @@ -26,7 +26,7 @@ export const qaAgentAnnotation = Annotation.Root({ }), }), designSessionId: Annotation, - schemaIssues: schemaIssuesAnnotation, + schemaIssues: qaSchemaIssuesAnnotation, generatedSqls: generatedSqlsAnnotation, next: Annotation({ reducer: (x, y) => y ?? x ?? END, diff --git a/frontend/internal-packages/agent/src/qa-agent/testcaseGeneration/testcaseAnnotation.ts b/frontend/internal-packages/agent/src/qa-agent/testcaseGeneration/testcaseAnnotation.ts index a0794cd454..9653c0880a 100644 --- a/frontend/internal-packages/agent/src/qa-agent/testcaseGeneration/testcaseAnnotation.ts +++ b/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>({ +/** + * 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>({ reducer: (prev, next) => prev.concat(next), default: () => [], }) @@ -27,6 +34,6 @@ export const testcaseAnnotation = Annotation.Root({ currentTestcase: Annotation, schemaData: Annotation, goal: Annotation, - schemaIssues: schemaIssuesAnnotation, + schemaIssues: qaSchemaIssuesAnnotation, generatedSqls: generatedSqlsAnnotation, }) diff --git a/frontend/internal-packages/agent/src/workflowAnnotation.ts b/frontend/internal-packages/agent/src/workflowAnnotation.ts index c1da7c5a8f..df7d9e2791 100644 --- a/frontend/internal-packages/agent/src/workflowAnnotation.ts +++ b/frontend/internal-packages/agent/src/workflowAnnotation.ts @@ -1,7 +1,7 @@ import { Annotation, END, MessagesAnnotation } from '@langchain/langgraph' import type { AnalyzedRequirements } from '@liam-hq/artifact' import type { Schema } from '@liam-hq/schema' -import { schemaIssuesAnnotation } from './qa-agent/testcaseGeneration/testcaseAnnotation' +import { workflowSchemaIssuesAnnotation } from './workflowSchemaIssuesAnnotation' export const workflowAnnotation = Annotation.Root({ ...MessagesAnnotation.spec, @@ -16,7 +16,7 @@ export const workflowAnnotation = Annotation.Root({ organizationId: Annotation, userId: Annotation, designSessionId: Annotation, - schemaIssues: schemaIssuesAnnotation, + schemaIssues: workflowSchemaIssuesAnnotation, next: Annotation({ reducer: (x, y) => y ?? x ?? END, diff --git a/frontend/internal-packages/agent/src/workflowSchemaIssuesAnnotation.test.ts b/frontend/internal-packages/agent/src/workflowSchemaIssuesAnnotation.test.ts new file mode 100644 index 0000000000..0d312954e2 --- /dev/null +++ b/frontend/internal-packages/agent/src/workflowSchemaIssuesAnnotation.test.ts @@ -0,0 +1,79 @@ +import { StateGraph } from '@langchain/langgraph' +import { describe, expect, it } from 'vitest' +import type { WorkflowState } from './types' +import { workflowAnnotation } from './workflowAnnotation' + +describe('workflowSchemaIssuesAnnotation', () => { + it('should replace schemaIssues (not concat)', async () => { + const graph = new StateGraph(workflowAnnotation) + .addNode('step1', (_state: WorkflowState) => ({ + schemaIssues: [{ testcaseId: 'test-1', description: 'Issue 1' }], + })) + .addNode('step2', (_state: WorkflowState) => ({ + schemaIssues: [{ testcaseId: 'test-2', description: 'Issue 2' }], + })) + .addEdge('__start__', 'step1') + .addEdge('step1', 'step2') + .addEdge('step2', '__end__') + .compile() + + const result = await graph.invoke({ + messages: [], + schemaData: { tables: {}, enums: {}, extensions: {} }, + }) + + // Should be replaced (only step2's issue), not concatenated + expect(result.schemaIssues).toEqual([ + { testcaseId: 'test-2', description: 'Issue 2' }, + ]) + }) + + it('should clear schemaIssues when set to empty array', async () => { + const graph = new StateGraph(workflowAnnotation) + .addNode('step1', (_state: WorkflowState) => ({ + schemaIssues: [ + { testcaseId: 'test-1', description: 'Issue 1' }, + { testcaseId: 'test-2', description: 'Issue 2' }, + ], + })) + .addNode('step2', (_state: WorkflowState) => ({ + schemaIssues: [], + })) + .addEdge('__start__', 'step1') + .addEdge('step1', 'step2') + .addEdge('step2', '__end__') + .compile() + + const result = await graph.invoke({ + messages: [], + schemaData: { tables: {}, enums: {}, extensions: {} }, + }) + + // Should be cleared + expect(result.schemaIssues).toEqual([]) + }) + + it('should keep schemaIssues when not updated', async () => { + const graph = new StateGraph(workflowAnnotation) + .addNode('step1', (_state: WorkflowState) => ({ + schemaIssues: [{ testcaseId: 'test-1', description: 'Issue 1' }], + })) + .addNode('step2', (_state: WorkflowState) => ({ + // Not updating schemaIssues + })) + .addEdge('__start__', 'step1') + .addEdge('step1', 'step2') + .addEdge('step2', '__end__') + .compile() + + const result = await graph.invoke({ + messages: [], + schemaData: { tables: {}, enums: {}, extensions: {} }, + }) + + // Should keep step1's value + expect(result.schemaIssues).toEqual([ + { testcaseId: 'test-1', description: 'Issue 1' }, + ]) + }) +}) diff --git a/frontend/internal-packages/agent/src/workflowSchemaIssuesAnnotation.ts b/frontend/internal-packages/agent/src/workflowSchemaIssuesAnnotation.ts new file mode 100644 index 0000000000..b8de853676 --- /dev/null +++ b/frontend/internal-packages/agent/src/workflowSchemaIssuesAnnotation.ts @@ -0,0 +1,22 @@ +import { Annotation } from '@langchain/langgraph' + +export type SchemaIssue = { + testcaseId: string + description: string +} + +/** + * Schema issues annotation for workflow-level state management. + * + * Uses a replacement reducer instead of concat because: + * - Workflow needs to clear issues after DB agent processing + * - Setting schemaIssues: [] should actually clear the array + * - With concat reducer, [] would be concatenated (prev.concat([]) = prev) + * + * This is different from QA agent's annotation which uses concat + * for parallel issue collection. + */ +export const workflowSchemaIssuesAnnotation = Annotation>({ + reducer: (prev, next) => next ?? prev, + default: () => [], +})