From 6e4f5f33aff4a2a00545efc8de00c866cc43972b Mon Sep 17 00:00:00 2001 From: Andrew Beal Date: Thu, 29 Jan 2026 15:04:43 +0000 Subject: [PATCH] Create new tests for Orchestration, Planning and Execution agents. --- __tests__/Services/ExecutionAgent.test.ts | 917 ++++++++++++++++++ .../Services/MultiAgentIntegration.test.ts | 487 ++++++++++ __tests__/Services/OrchestrationAgent.test.ts | 114 +++ __tests__/Services/PlanningAgent.test.ts | 911 +++++++++++++++++ 4 files changed, 2429 insertions(+) create mode 100644 __tests__/Services/ExecutionAgent.test.ts create mode 100644 __tests__/Services/MultiAgentIntegration.test.ts create mode 100644 __tests__/Services/OrchestrationAgent.test.ts create mode 100644 __tests__/Services/PlanningAgent.test.ts diff --git a/__tests__/Services/ExecutionAgent.test.ts b/__tests__/Services/ExecutionAgent.test.ts new file mode 100644 index 0000000..9cb677b --- /dev/null +++ b/__tests__/Services/ExecutionAgent.test.ts @@ -0,0 +1,917 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { ExecutionAgent } from '../../Services/AIServices/ExecutionAgent'; +import { RegisterSingleton, DeregisterAllServices } from '../../Services/DependencyService'; +import { Services } from '../../Services/Services'; +import { AIFunction } from '../../Enums/AIFunction'; +import { AIFunctionCall } from '../../AIClasses/AIFunctionCall'; +import { AIFunctionResponse } from '../../AIClasses/FunctionDefinitions/AIFunctionResponse'; +import type { ExecutionStep } from '../../Types/ExecutionStep'; + +/** + * UNIT TESTS - ExecutionAgent + * + * Tests the execution agent that: + * - Executes individual steps via CompleteTask + * - Calls AI functions during execution + * - Retries when CompleteTask is not called + * - Respects max depth limits + * - Uses step context during execution + */ + +describe('ExecutionAgent - Unit Tests', () => { + let service: ExecutionAgent; + let mockAI: any; + let mockPrompt: any; + let mockAIFunctionService: any; + + const createMockCallbacks = () => ({ + onSubmit: vi.fn(), + onStreamingUpdate: vi.fn(), + onThoughtUpdate: vi.fn(), + onPlanningStarted: vi.fn(), + onPlanningFinished: vi.fn(), + onUserQuestion: vi.fn().mockResolvedValue('User answer'), + onPlanUpdate: vi.fn(), + onPlanStepUpdate: vi.fn(), + onPlanReset: vi.fn(), + onComplete: vi.fn(), + onCancel: vi.fn() + }); + + beforeEach(() => { + // Mock IPrompt + mockPrompt = { + systemInstruction: vi.fn().mockReturnValue('System instruction'), + userInstruction: vi.fn().mockResolvedValue('User instruction'), + executionInstruction: vi.fn().mockReturnValue('Execution instruction') + }; + RegisterSingleton(Services.IPrompt, mockPrompt); + + // Mock AIFunctionService + mockAIFunctionService = { + performAIFunction: vi.fn().mockResolvedValue( + new AIFunctionResponse(AIFunction.SearchVaultFiles, { results: [] }, 'test-tool-id') + ) + }; + RegisterSingleton(Services.AIFunctionService, mockAIFunctionService); + + // Mock IAIClass + mockAI = { + systemPrompt: '', + userInstruction: '', + toolDefinitions: [], + agentType: undefined, + streamRequest: vi.fn() + }; + RegisterSingleton(Services.IAIClass, mockAI); + + // Create service + service = new ExecutionAgent(); + }); + + afterEach(() => { + DeregisterAllServices(); + vi.restoreAllMocks(); + }); + + describe('CompleteTask success', () => { + it('should complete task successfully', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Search for files', + instruction: 'Search the vault for markdown files' + }; + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Task completed', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Found 10 markdown files' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(result?.success).toBe(true); + expect(result?.description).toBe('Found 10 markdown files'); + }); + + it('should complete task with context', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Process files', + instruction: 'Process the found files' + }; + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Task completed', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Processed 5 files', + context: 'File names: note1.md, note2.md, note3.md, note4.md, note5.md' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(result?.success).toBe(true); + expect(result?.context).toContain('File names'); + }); + + it('should handle multiple steps independently', async () => { + const callbacks = createMockCallbacks(); + + const step1: ExecutionStep = { + description: 'Step 1', + instruction: 'First task' + }; + + const step2: ExecutionStep = { + description: 'Step 2', + instruction: 'Second task' + }; + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Completed' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result1 = await service.runExecutionAgent(step1, callbacks); + + // Create new agent for second step + const service2 = new ExecutionAgent(); + service2.resolveAIProvider(); + const result2 = await service2.runExecutionAgent(step2, callbacks); + + expect(result1).toBeDefined(); + expect(result2).toBeDefined(); + expect(result1?.success).toBe(true); + expect(result2?.success).toBe(true); + }); + }); + + describe('CompleteTask failure', () => { + it('should handle task failure', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Delete files', + instruction: 'Delete temporary files' + }; + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Task failed', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: false, + description: 'No temporary files found' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(result?.success).toBe(false); + expect(result?.description).toBe('No temporary files found'); + }); + + it('should handle task failure with error details', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Write file', + instruction: 'Write content to file' + }; + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Failed', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: false, + description: 'Permission denied: cannot write to /protected/file.md', + context: 'The target directory is read-only' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(result?.success).toBe(false); + expect(result?.description).toContain('Permission denied'); + expect(result?.context).toContain('read-only'); + }); + }); + + describe('Function execution', () => { + it('should call SearchVaultFiles during execution', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Find files', + instruction: 'Search for files with tag #important' + }; + + let functionCallCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + functionCallCount++; + if (functionCallCount === 1) { + // Call search function + yield { + content: 'Searching', + functionCall: new AIFunctionCall( + AIFunction.SearchVaultFiles, + { + search_terms: ['#important'], + user_message: 'Searching for tagged files' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + // Complete task + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Found 3 files with #important tag' + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(mockAIFunctionService.performAIFunction).toHaveBeenCalled(); + expect(result?.success).toBe(true); + expect(functionCallCount).toBe(2); + }); + + it('should call ReadVaultFiles during execution', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Read files', + instruction: 'Read contents of note.md' + }; + + let functionCallCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + functionCallCount++; + if (functionCallCount === 1) { + yield { + content: 'Reading', + functionCall: new AIFunctionCall( + AIFunction.ReadVaultFiles, + { + file_paths: ['note.md'], + user_message: 'Reading file' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'File read successfully' + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(mockAIFunctionService.performAIFunction).toHaveBeenCalled(); + expect(result?.success).toBe(true); + }); + + it('should call WriteVaultFile during execution', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Write file', + instruction: 'Create new file with content' + }; + + let functionCallCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + functionCallCount++; + if (functionCallCount === 1) { + yield { + content: 'Writing', + functionCall: new AIFunctionCall( + AIFunction.WriteVaultFile, + { + file_path: 'new-note.md', + content: '# New Note\n\nContent here', + user_message: 'Creating file' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'File created successfully' + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(mockAIFunctionService.performAIFunction).toHaveBeenCalled(); + expect(result?.success).toBe(true); + }); + + it('should call multiple functions in sequence', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Update files', + instruction: 'Search, read, and update files' + }; + + let functionCallCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + functionCallCount++; + if (functionCallCount === 1) { + yield { + content: 'Searching', + functionCall: new AIFunctionCall( + AIFunction.SearchVaultFiles, + { + search_terms: ['TODO'], + user_message: 'Searching' + }, + 'tool-1' + ), + isComplete: true + }; + } else if (functionCallCount === 2) { + yield { + content: 'Reading', + functionCall: new AIFunctionCall( + AIFunction.ReadVaultFiles, + { + file_paths: ['found.md'], + user_message: 'Reading' + }, + 'tool-2' + ), + isComplete: true + }; + } else if (functionCallCount === 3) { + yield { + content: 'Writing', + functionCall: new AIFunctionCall( + AIFunction.PatchVaultFile, + { + file_path: 'found.md', + patch_operations: [], + user_message: 'Updating' + }, + 'tool-3' + ), + isComplete: true + }; + } else { + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Updated 1 file' + }, + 'tool-4' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(mockAIFunctionService.performAIFunction).toHaveBeenCalledTimes(3); + expect(result?.success).toBe(true); + }); + }); + + describe('Recursive retry', () => { + it('should retry when CompleteTask is not called', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Process files', + instruction: 'Process files' + }; + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount === 1) { + // No CompleteTask + yield { + content: 'Working on it...', + isComplete: true + }; + } else { + // Complete on retry + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Completed' + }, + 'tool-1' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(attemptCount).toBe(2); + }); + + it('should retry multiple times before completing', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Complex task', + instruction: 'Do complex task' + }; + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount < 3) { + // No completion on first two attempts + yield { + content: 'Still processing...', + isComplete: true + }; + } else { + // Complete on third attempt + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Finally completed' + }, + 'tool-1' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(attemptCount).toBe(3); + }); + + it('should maintain conversation across retries', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Task', + instruction: 'Do task' + }; + + let attemptCount = 0; + let conversationGrows = false; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + + // Check if conversation has retry message + if (attemptCount > 1) { + const conversation = mockAI.streamRequest.mock.calls[mockAI.streamRequest.mock.calls.length - 1][0]; + conversationGrows = conversation.contents.length > 2; // Initial + retries + } + + if (attemptCount < 2) { + yield { + content: 'Not done', + isComplete: true + }; + } else { + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Completed' + }, + 'tool-1' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + await service.runExecutionAgent(step, callbacks); + + expect(conversationGrows).toBe(true); + }); + }); + + describe('Max depth protection', () => { + it('should return undefined when max execution depth is exceeded', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Impossible task', + instruction: 'Task that never completes' + }; + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + // Never complete task, force retries + yield { + content: 'Still working...', + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeUndefined(); + expect(attemptCount).toBe(3); // MAX_AGENT_DEPTH + }); + + it('should complete just before max depth', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Almost failing task', + instruction: 'Task that completes at last moment' + }; + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount < 3) { + yield { + content: 'Retrying...', + isComplete: true + }; + } else { + // Complete on third attempt (max depth) + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Barely made it' + }, + 'tool-1' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(result?.success).toBe(true); + expect(attemptCount).toBe(3); + }); + }); + + describe('Step context usage', () => { + it('should include step context in execution', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Process files', + instruction: 'Process the files', + context: 'Files found: note1.md, note2.md, note3.md' + }; + + let contextWasIncluded = false; + + mockAI.streamRequest.mockImplementation(async function* () { + const conversation = mockAI.streamRequest.mock.calls[mockAI.streamRequest.mock.calls.length - 1][0]; + const firstMessage = conversation.contents[0]; + contextWasIncluded = firstMessage.content.includes('note1.md'); + + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Processed files from context' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(contextWasIncluded).toBe(true); + }); + + it('should work without context', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Simple task', + instruction: 'Do simple task' + // No context + }; + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Task completed without context' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(result?.success).toBe(true); + }); + + it('should handle long context', async () => { + const callbacks = createMockCallbacks(); + const longContext = 'File list: ' + Array.from({ length: 100 }, (_, i) => `file${i}.md`).join(', '); + const step: ExecutionStep = { + description: 'Process many files', + instruction: 'Process all files', + context: longContext + }; + + let contextLength = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + const conversation = mockAI.streamRequest.mock.calls[mockAI.streamRequest.mock.calls.length - 1][0]; + const firstMessage = conversation.contents[0]; + contextLength = firstMessage.content.length; + + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Processed all files' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(contextLength).toBeGreaterThan(longContext.length); + }); + }); + + describe('Invalid CompleteTask arguments', () => { + it('should reject CompleteTask with invalid arguments', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Task', + instruction: 'Do task' + }; + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount === 1) { + // Invalid: missing description + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true + // Missing description field + }, + 'tool-1' + ), + isComplete: true + }; + } else { + // Valid completion + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Valid completion' + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(attemptCount).toBeGreaterThan(1); + }); + + it('should reject CompleteTask with wrong type', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Task', + instruction: 'Do task' + }; + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount === 1) { + // Invalid: success is not boolean + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: 'yes', + description: 'Done' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + // Valid + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Done properly' + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(attemptCount).toBeGreaterThan(1); + }); + }); + + describe('Thought updates', () => { + it('should update thought callback during function calls', async () => { + const callbacks = createMockCallbacks(); + const step: ExecutionStep = { + description: 'Task', + instruction: 'Do task with updates' + }; + + let functionCallCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + functionCallCount++; + if (functionCallCount === 1) { + yield { + content: 'Searching', + functionCall: new AIFunctionCall( + AIFunction.SearchVaultFiles, + { + search_terms: ['test'], + user_message: 'Searching for files' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Completed' + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + await service.runExecutionAgent(step, callbacks); + + expect(callbacks.onThoughtUpdate).toHaveBeenCalledWith('Searching for files'); + }); + }); +}); diff --git a/__tests__/Services/MultiAgentIntegration.test.ts b/__tests__/Services/MultiAgentIntegration.test.ts new file mode 100644 index 0000000..c692216 --- /dev/null +++ b/__tests__/Services/MultiAgentIntegration.test.ts @@ -0,0 +1,487 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { MainAgent } from '../../Services/AIServices/MainAgent'; +import { PlanningAgent } from '../../Services/AIServices/PlanningAgent'; +import { ExecutionAgent } from '../../Services/AIServices/ExecutionAgent'; +import { OrchestrationAgent } from '../../Services/AIServices/OrchestrationAgent'; +import { RegisterSingleton, DeregisterAllServices } from '../../Services/DependencyService'; +import { Services } from '../../Services/Services'; +import { Conversation } from '../../Conversations/Conversation'; +import { ConversationContent } from '../../Conversations/ConversationContent'; +import { Role } from '../../Enums/Role'; +import { AIFunction } from '../../Enums/AIFunction'; +import { AIFunctionCall } from '../../AIClasses/AIFunctionCall'; +import { AIFunctionResponse } from '../../AIClasses/FunctionDefinitions/AIFunctionResponse'; +import type { ExecutionStep } from '../../Types/ExecutionStep'; + +/** + * INTEGRATION TESTS - Multi-Agent Architecture + * + * Tests the multi-agent components independently: + * - MainAgent initialization and setup + * - PlanningAgent plan creation + * - ExecutionAgent step execution + * - OrchestrationAgent setup + * + * Note: Full end-to-end workflow tests would require extensive mocking of nested + * agent instantiation, which is not practical given the current architecture. + * These tests verify that each agent can be initialized and has the correct interface. + */ + +describe('Multi-Agent Integration Tests', () => { + let mockAI: any; + let mockPrompt: any; + let mockAIFunctionService: any; + + const createMockCallbacks = () => ({ + onSubmit: vi.fn(), + onStreamingUpdate: vi.fn(), + onThoughtUpdate: vi.fn(), + onPlanningStarted: vi.fn(), + onPlanningFinished: vi.fn(), + onUserQuestion: vi.fn().mockResolvedValue('User answer'), + onPlanUpdate: vi.fn(), + onPlanStepUpdate: vi.fn(), + onPlanReset: vi.fn(), + onComplete: vi.fn(), + onCancel: vi.fn() + }); + + beforeEach(() => { + // Mock IPrompt + mockPrompt = { + systemInstruction: vi.fn().mockReturnValue('System instruction'), + userInstruction: vi.fn().mockResolvedValue('User instruction'), + planningInstruction: vi.fn().mockReturnValue('Planning instruction'), + orchestrationInstruction: vi.fn().mockReturnValue('Orchestration instruction'), + executionInstruction: vi.fn().mockReturnValue('Execution instruction') + }; + RegisterSingleton(Services.IPrompt, mockPrompt); + + // Mock AIFunctionService + mockAIFunctionService = { + performAIFunction: vi.fn().mockResolvedValue( + new AIFunctionResponse(AIFunction.SearchVaultFiles, { results: ['file1.md', 'file2.md'] }, 'test-tool-id') + ) + }; + RegisterSingleton(Services.AIFunctionService, mockAIFunctionService); + + // Mock IAIClass + mockAI = { + systemPrompt: '', + userInstruction: '', + toolDefinitions: [], + agentType: undefined, + streamRequest: vi.fn() + }; + RegisterSingleton(Services.IAIClass, mockAI); + }); + + afterEach(() => { + DeregisterAllServices(); + vi.restoreAllMocks(); + }); + + describe('MainAgent Integration', () => { + it('should initialize MainAgent with all services', () => { + const service = new MainAgent(); + expect(service).toBeDefined(); + expect(service.runMainAgent).toBeDefined(); + }); + + it('should resolve AI provider', () => { + const service = new MainAgent(); + service.resolveAIProvider(); + expect(service).toBeDefined(); + }); + + it('should throw error if running without resolved provider', async () => { + const service = new MainAgent(); + const conversation = new Conversation(); + const callbacks = createMockCallbacks(); + + await expect(service.runMainAgent(conversation, true, false, callbacks)) + .rejects.toThrow('Error: No AI provider has been set!'); + }); + }); + + describe('PlanningAgent Integration', () => { + it('should initialize PlanningAgent', () => { + const service = new PlanningAgent(); + expect(service).toBeDefined(); + expect(service.runPlanningAgent).toBeDefined(); + }); + + it('should create a simple plan', async () => { + const service = new PlanningAgent(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create a plan' + })); + const callbacks = createMockCallbacks(); + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Test step', + instruction: 'Test instruction' + } + ] + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(result).toBeDefined(); + expect(result?.executionSteps).toHaveLength(1); + expect(result?.isReplan).toBe(false); + }); + + it('should mark plan as replan when requested', async () => { + const service = new PlanningAgent(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Revise plan' + })); + const callbacks = createMockCallbacks(); + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Replan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Revised step', + instruction: 'New approach' + } + ] + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, true); + + expect(result).toBeDefined(); + expect(result?.isReplan).toBe(true); + }); + }); + + describe('ExecutionAgent Integration', () => { + it('should initialize ExecutionAgent', () => { + const service = new ExecutionAgent(); + expect(service).toBeDefined(); + expect(service.runExecutionAgent).toBeDefined(); + }); + + it('should execute a step successfully', async () => { + const service = new ExecutionAgent(); + const step: ExecutionStep = { + description: 'Test step', + instruction: 'Do something' + }; + const callbacks = createMockCallbacks(); + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Task completed' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(result?.success).toBe(true); + expect(result?.description).toBe('Task completed'); + }); + + it('should handle step failure', async () => { + const service = new ExecutionAgent(); + const step: ExecutionStep = { + description: 'Test step', + instruction: 'Do something' + }; + const callbacks = createMockCallbacks(); + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Failed', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: false, + description: 'Task failed' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(result?.success).toBe(false); + expect(result?.description).toBe('Task failed'); + }); + + it('should include context in execution', async () => { + const service = new ExecutionAgent(); + const step: ExecutionStep = { + description: 'Test step', + instruction: 'Process data', + context: 'Previous result: 42' + }; + const callbacks = createMockCallbacks(); + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Processed with context', + context: 'Used previous result' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(result).toBeDefined(); + expect(result?.success).toBe(true); + expect(result?.context).toContain('Used previous result'); + }); + }); + + describe('OrchestrationAgent Integration', () => { + it('should initialize OrchestrationAgent', () => { + const service = new OrchestrationAgent(); + expect(service).toBeDefined(); + expect(service.runPlannedWorkflow).toBeDefined(); + }); + + it('should have correct method interface', () => { + const service = new OrchestrationAgent(); + expect(typeof service.runPlannedWorkflow).toBe('function'); + }); + }); + + describe('Agent Depth Limits', () => { + it('PlanningAgent should respect max depth', async () => { + const service = new PlanningAgent(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create plan' + })); + const callbacks = createMockCallbacks(); + + let attemptCount = 0; + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + // Never submit plan, force retries + yield { + content: 'Still thinking...', + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + // Should return undefined after max depth + expect(result).toBeUndefined(); + expect(attemptCount).toBe(3); // MAX_AGENT_DEPTH + }); + + it('ExecutionAgent should respect max depth', async () => { + const service = new ExecutionAgent(); + const step: ExecutionStep = { + description: 'Test step', + instruction: 'Do something' + }; + const callbacks = createMockCallbacks(); + + let attemptCount = 0; + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + // Never complete task, force retries + yield { + content: 'Still working...', + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + // Should return undefined after max depth + expect(result).toBeUndefined(); + expect(attemptCount).toBe(3); // MAX_AGENT_DEPTH + }); + }); + + describe('Agent Interoperability', () => { + it('should allow PlanningAgent and ExecutionAgent to run independently', async () => { + const planningAgent = new PlanningAgent(); + const executionAgent = new ExecutionAgent(); + + expect(planningAgent).toBeDefined(); + expect(executionAgent).toBeDefined(); + expect(planningAgent).not.toBe(executionAgent); + }); + + it('should create independent conversation contexts', () => { + const conversation1 = new Conversation(); + const conversation2 = new Conversation(); + + conversation1.contents.push(new ConversationContent({ + role: Role.User, + content: 'Request 1' + })); + + conversation2.contents.push(new ConversationContent({ + role: Role.User, + content: 'Request 2' + })); + + expect(conversation1.contents).toHaveLength(1); + expect(conversation2.contents).toHaveLength(1); + expect(conversation1.contents[0].content).not.toBe(conversation2.contents[0].content); + }); + }); + + describe('Agent Tool Restrictions', () => { + it('should allow PlanningAgent to use read-only tools', async () => { + const service = new PlanningAgent(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Plan based on vault content' + })); + const callbacks = createMockCallbacks(); + + let searchCalled = false; + mockAI.streamRequest.mockImplementation(async function* () { + if (!searchCalled) { + searchCalled = true; + yield { + content: 'Searching', + functionCall: new AIFunctionCall( + AIFunction.SearchVaultFiles, + { + search_terms: ['test'], + user_message: 'Searching' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Process files', + instruction: 'Process found files' + } + ] + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(mockAIFunctionService.performAIFunction).toHaveBeenCalled(); + expect(result).toBeDefined(); + }); + + it('should allow ExecutionAgent to use all vault tools', async () => { + const service = new ExecutionAgent(); + const step: ExecutionStep = { + description: 'Search and process', + instruction: 'Find and read files' + }; + const callbacks = createMockCallbacks(); + + let functionCallCount = 0; + mockAI.streamRequest.mockImplementation(async function* () { + functionCallCount++; + if (functionCallCount === 1) { + yield { + content: 'Searching', + functionCall: new AIFunctionCall( + AIFunction.SearchVaultFiles, + { + search_terms: ['test'], + user_message: 'Searching' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + yield { + content: 'Done', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Found and processed files' + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runExecutionAgent(step, callbacks); + + expect(mockAIFunctionService.performAIFunction).toHaveBeenCalled(); + expect(result?.success).toBe(true); + }); + }); +}); diff --git a/__tests__/Services/OrchestrationAgent.test.ts b/__tests__/Services/OrchestrationAgent.test.ts new file mode 100644 index 0000000..770639a --- /dev/null +++ b/__tests__/Services/OrchestrationAgent.test.ts @@ -0,0 +1,114 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { OrchestrationAgent } from '../../Services/AIServices/OrchestrationAgent'; +import { RegisterSingleton, DeregisterAllServices } from '../../Services/DependencyService'; +import { Services } from '../../Services/Services'; +import { AIFunction } from '../../Enums/AIFunction'; +import { AIFunctionResponse } from '../../AIClasses/FunctionDefinitions/AIFunctionResponse'; + +/** + * UNIT TESTS - OrchestrationAgent + * + * Tests the orchestration agent's basic functionality: + * - Service initialization + * - Agent configuration + * + * Note: Full workflow integration tests with nested agents are in MultiAgentIntegration.test.ts + * These tests focus on the OrchestrationAgent's setup and configuration. + */ + +describe('OrchestrationAgent - Unit Tests', () => { + let service: OrchestrationAgent; + let mockAI: any; + let mockPrompt: any; + let mockAIFunctionService: any; + + beforeEach(() => { + // Mock IPrompt + mockPrompt = { + systemInstruction: vi.fn().mockReturnValue('System instruction'), + userInstruction: vi.fn().mockResolvedValue('User instruction'), + planningInstruction: vi.fn().mockReturnValue('Planning instruction'), + orchestrationInstruction: vi.fn().mockReturnValue('Orchestration instruction'), + executionInstruction: vi.fn().mockReturnValue('Execution instruction') + }; + RegisterSingleton(Services.IPrompt, mockPrompt); + + // Mock AIFunctionService + mockAIFunctionService = { + performAIFunction: vi.fn().mockResolvedValue( + new AIFunctionResponse(AIFunction.SearchVaultFiles, { results: [] }, 'test-tool-id') + ) + }; + RegisterSingleton(Services.AIFunctionService, mockAIFunctionService); + + // Mock IAIClass + mockAI = { + systemPrompt: '', + userInstruction: '', + toolDefinitions: [], + agentType: undefined, + streamRequest: vi.fn() + }; + RegisterSingleton(Services.IAIClass, mockAI); + + // Create service + service = new OrchestrationAgent(); + }); + + afterEach(() => { + DeregisterAllServices(); + vi.restoreAllMocks(); + }); + + describe('constructor', () => { + it('should initialize with all required services', () => { + const testService = new OrchestrationAgent(); + + expect(testService).toBeDefined(); + }); + }); + + describe('resolveAIProvider', () => { + it('should resolve AI provider successfully', () => { + service.resolveAIProvider(); + + expect(service).toBeDefined(); + }); + + it('should set up AI provider from dependency service', () => { + service.resolveAIProvider(); + + // Verify the service can be used + expect(service).toHaveProperty('runPlannedWorkflow'); + }); + }); + + describe('service methods', () => { + it('should have runPlannedWorkflow method', () => { + expect(service.runPlannedWorkflow).toBeDefined(); + expect(typeof service.runPlannedWorkflow).toBe('function'); + }); + + it('should throw error if AI provider not resolved before running workflow', async () => { + const callbacks = { + onSubmit: vi.fn(), + onStreamingUpdate: vi.fn(), + onThoughtUpdate: vi.fn(), + onPlanningStarted: vi.fn(), + onPlanningFinished: vi.fn(), + onUserQuestion: vi.fn(), + onPlanUpdate: vi.fn(), + onPlanStepUpdate: vi.fn(), + onPlanReset: vi.fn(), + onComplete: vi.fn(), + onCancel: vi.fn() + }; + + // Note: runPlannedWorkflow spawns new agents internally which get + // their own AI provider from the dependency service, so this test + // verifies the service is properly constructed + expect(service.runPlannedWorkflow).toBeDefined(); + expect(typeof service.runPlannedWorkflow).toBe('function'); + }); + }); +}); diff --git a/__tests__/Services/PlanningAgent.test.ts b/__tests__/Services/PlanningAgent.test.ts new file mode 100644 index 0000000..b98e2dc --- /dev/null +++ b/__tests__/Services/PlanningAgent.test.ts @@ -0,0 +1,911 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { PlanningAgent } from '../../Services/AIServices/PlanningAgent'; +import { RegisterSingleton, DeregisterAllServices } from '../../Services/DependencyService'; +import { Services } from '../../Services/Services'; +import { Conversation } from '../../Conversations/Conversation'; +import { ConversationContent } from '../../Conversations/ConversationContent'; +import { Role } from '../../Enums/Role'; +import { AIFunction } from '../../Enums/AIFunction'; +import { AIFunctionCall } from '../../AIClasses/AIFunctionCall'; +import { AIFunctionResponse } from '../../AIClasses/FunctionDefinitions/AIFunctionResponse'; + +/** + * UNIT TESTS - PlanningAgent + * + * Tests the planning agent that: + * - Creates execution plans via SubmitPlan + * - Asks user questions during planning + * - Enforces tool restrictions (only planning tools allowed) + * - Handles replan vs initial plan scenarios + * - Respects max depth limits + */ + +describe('PlanningAgent - Unit Tests', () => { + let service: PlanningAgent; + let mockAI: any; + let mockPrompt: any; + let mockAIFunctionService: any; + + const createMockCallbacks = () => ({ + onSubmit: vi.fn(), + onStreamingUpdate: vi.fn(), + onThoughtUpdate: vi.fn(), + onPlanningStarted: vi.fn(), + onPlanningFinished: vi.fn(), + onUserQuestion: vi.fn().mockResolvedValue('User answer'), + onPlanUpdate: vi.fn(), + onPlanStepUpdate: vi.fn(), + onPlanReset: vi.fn(), + onComplete: vi.fn(), + onCancel: vi.fn() + }); + + beforeEach(() => { + // Mock IPrompt + mockPrompt = { + systemInstruction: vi.fn().mockReturnValue('System instruction'), + userInstruction: vi.fn().mockResolvedValue('User instruction'), + planningInstruction: vi.fn().mockReturnValue('Planning instruction') + }; + RegisterSingleton(Services.IPrompt, mockPrompt); + + // Mock AIFunctionService + mockAIFunctionService = { + performAIFunction: vi.fn().mockResolvedValue( + new AIFunctionResponse(AIFunction.SearchVaultFiles, { results: [] }, 'test-tool-id') + ) + }; + RegisterSingleton(Services.AIFunctionService, mockAIFunctionService); + + // Mock IAIClass + mockAI = { + systemPrompt: '', + userInstruction: '', + toolDefinitions: [], + agentType: undefined, + streamRequest: vi.fn() + }; + RegisterSingleton(Services.IAIClass, mockAI); + + // Create service + service = new PlanningAgent(); + }); + + afterEach(() => { + DeregisterAllServices(); + vi.restoreAllMocks(); + }); + + describe('SubmitPlan - Valid plan submission', () => { + it('should submit a valid plan with multiple steps', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create a plan to organize notes' + })); + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Creating plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Search for notes', + instruction: 'Search the vault for all notes' + }, + { + description: 'Categorize notes', + instruction: 'Organize notes by topic' + }, + { + description: 'Create index', + instruction: 'Create an index file' + } + ], + user_message: 'Plan created' + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(result).toBeDefined(); + expect(result?.executionSteps).toHaveLength(3); + expect(result?.isReplan).toBe(false); + expect(result?.executionSteps[0].description).toBe('Search for notes'); + expect(result?.executionSteps[1].description).toBe('Categorize notes'); + expect(result?.executionSteps[2].description).toBe('Create index'); + }); + + it('should submit a plan with single step', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Simple task' + })); + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Complete task', + instruction: 'Do the single task' + } + ] + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(result).toBeDefined(); + expect(result?.executionSteps).toHaveLength(1); + }); + + it('should mark plan as replan when isReplan is true', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Previous plan failed, create new plan' + })); + + mockAI.streamRequest.mockImplementation(async function* () { + yield { + content: 'Replan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Revised approach', + instruction: 'New strategy' + } + ] + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, true); + + expect(result).toBeDefined(); + expect(result?.isReplan).toBe(true); + }); + }); + + describe('SubmitPlan validation', () => { + it('should reject invalid plan structure', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create plan' + })); + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount === 1) { + // Invalid: missing required fields + yield { + content: 'Invalid plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + // Missing instruction + description: 'Incomplete step' + } + ] + }, + 'tool-1' + ), + isComplete: true + }; + } else { + // Valid plan + yield { + content: 'Valid plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Complete step', + instruction: 'Full instruction' + } + ] + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(result).toBeDefined(); + expect(attemptCount).toBeGreaterThan(1); + }); + + it('should retry when no plan is submitted', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create plan' + })); + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount === 1) { + // No plan submitted + yield { + content: 'Thinking about the plan...', + isComplete: true + }; + } else { + // Submit plan on retry + yield { + content: 'Plan ready', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Step 1', + instruction: 'Instruction 1' + } + ] + }, + 'tool-1' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(result).toBeDefined(); + expect(attemptCount).toBe(2); + }); + }); + + describe('AskUserQuestionPlanning', () => { + it('should ask user question during planning', async () => { + const callbacks = createMockCallbacks(); + callbacks.onUserQuestion.mockResolvedValue('Notes about AI'); + + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Search for relevant notes' + })); + + let functionCallCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + functionCallCount++; + if (functionCallCount === 1) { + // Ask user question + yield { + content: 'Question', + functionCall: new AIFunctionCall( + AIFunction.AskUserQuestionPlanning, + { + question: 'What topic should I search for?', + user_message: 'Asking user for clarification' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + // Submit plan with user's answer incorporated + yield { + content: 'Plan with answer', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Search for AI notes', + instruction: 'Search for notes about AI' + } + ] + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(callbacks.onUserQuestion).toHaveBeenCalledWith('What topic should I search for?'); + expect(result).toBeDefined(); + expect(functionCallCount).toBe(2); + }); + + it('should handle multiple user questions in planning', async () => { + const callbacks = createMockCallbacks(); + callbacks.onUserQuestion + .mockResolvedValueOnce('markdown') + .mockResolvedValueOnce('daily notes'); + + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Organize files' + })); + + let functionCallCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + functionCallCount++; + if (functionCallCount === 1) { + yield { + content: 'Q1', + functionCall: new AIFunctionCall( + AIFunction.AskUserQuestionPlanning, + { + question: 'What file type?', + user_message: 'Asking about file type' + }, + 'tool-1' + ), + isComplete: true + }; + } else if (functionCallCount === 2) { + yield { + content: 'Q2', + functionCall: new AIFunctionCall( + AIFunction.AskUserQuestionPlanning, + { + question: 'Which folder?', + user_message: 'Asking about folder' + }, + 'tool-2' + ), + isComplete: true + }; + } else { + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Organize markdown files in daily notes', + instruction: 'Process files' + } + ] + }, + 'tool-3' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(callbacks.onUserQuestion).toHaveBeenCalledTimes(2); + expect(result).toBeDefined(); + }); + + it('should reject invalid AskUserQuestionPlanning arguments', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create plan' + })); + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount === 1) { + // Invalid: missing question + yield { + content: 'Invalid', + functionCall: new AIFunctionCall( + AIFunction.AskUserQuestionPlanning, + { + // Missing question field + user_message: 'Message' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + // Valid plan + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Step', + instruction: 'Instruction' + } + ] + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + await service.runPlanningAgent(conversation, callbacks, false); + + expect(attemptCount).toBeGreaterThan(1); + }); + }); + + describe('Tool restrictions', () => { + it('should allow planning tools (SearchVaultFiles)', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Plan to find files' + })); + + let searchCalled = false; + + mockAI.streamRequest.mockImplementation(async function* () { + if (!searchCalled) { + searchCalled = true; + yield { + content: 'Searching', + functionCall: new AIFunctionCall( + AIFunction.SearchVaultFiles, + { + search_terms: ['test'], + user_message: 'Searching' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Process files', + instruction: 'Process found files' + } + ] + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(mockAIFunctionService.performAIFunction).toHaveBeenCalled(); + expect(result).toBeDefined(); + }); + + it('should allow planning tools (ReadVaultFiles)', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Plan based on file content' + })); + + let readCalled = false; + + mockAI.streamRequest.mockImplementation(async function* () { + if (!readCalled) { + readCalled = true; + yield { + content: 'Reading', + functionCall: new AIFunctionCall( + AIFunction.ReadVaultFiles, + { + file_paths: ['test.md'], + user_message: 'Reading' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Update file', + instruction: 'Update based on content' + } + ] + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(mockAIFunctionService.performAIFunction).toHaveBeenCalled(); + expect(result).toBeDefined(); + }); + + it('should deny destructive tools (WriteVaultFile)', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create plan' + })); + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount === 1) { + // Try to write file (not allowed in planning) + yield { + content: 'Writing', + functionCall: new AIFunctionCall( + AIFunction.WriteVaultFile, + { + file_path: 'test.md', + content: 'content' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + // Submit valid plan + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Write file', + instruction: 'Write file in execution' + } + ] + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + // WriteVaultFile should be denied, not executed + expect(mockAIFunctionService.performAIFunction).not.toHaveBeenCalledWith( + expect.objectContaining({ name: AIFunction.WriteVaultFile }) + ); + expect(result).toBeDefined(); + expect(attemptCount).toBeGreaterThan(1); + }); + + it('should deny execution tools (CompleteTask)', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create plan' + })); + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount === 1) { + // Try to complete task (execution agent tool) + yield { + content: 'Completing', + functionCall: new AIFunctionCall( + AIFunction.CompleteTask, + { + success: true, + description: 'Done' + }, + 'tool-1' + ), + isComplete: true + }; + } else { + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Step', + instruction: 'Instruction' + } + ] + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + await service.runPlanningAgent(conversation, callbacks, false); + + expect(attemptCount).toBeGreaterThan(1); + }); + + it('should deny orchestration tools (CompleteStep)', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create plan' + })); + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount === 1) { + // Try to complete step (orchestration agent tool) + yield { + content: 'Completing', + functionCall: new AIFunctionCall( + AIFunction.CompleteStep, + { + confirm_completion: true + }, + 'tool-1' + ), + isComplete: true + }; + } else { + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Step', + instruction: 'Instruction' + } + ] + }, + 'tool-2' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + await service.runPlanningAgent(conversation, callbacks, false); + + expect(attemptCount).toBeGreaterThan(1); + }); + }); + + describe('Max depth protection', () => { + it('should return undefined when max planning depth is exceeded', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create plan' + })); + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + // Never submit a plan, forcing retries + yield { + content: 'Still thinking...', + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(result).toBeUndefined(); + expect(attemptCount).toBe(3); // MAX_AGENT_DEPTH + }); + + it('should track depth across retries', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create plan' + })); + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount < 3) { + // No plan on first two attempts + yield { + content: 'Working on it', + isComplete: true + }; + } else { + // Submit plan on third attempt (just before max depth) + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Final step', + instruction: 'Final instruction' + } + ] + }, + 'tool-1' + ), + isComplete: true + }; + } + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(result).toBeDefined(); + expect(attemptCount).toBe(3); + }); + }); + + describe('Empty plan handling', () => { + it('should accept empty steps array as valid plan', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Create plan' + })); + + mockAI.streamRequest.mockImplementation(async function* () { + // Empty plan is valid according to schema + yield { + content: 'Empty plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [] + }, + 'tool-1' + ), + isComplete: true + }; + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + // Empty plan should be accepted + expect(result).toBeDefined(); + expect(result?.executionSteps).toEqual([]); + }); + }); + + describe('Conversation structure', () => { + it('should maintain conversation across retries', async () => { + const callbacks = createMockCallbacks(); + const conversation = new Conversation(); + conversation.contents.push(new ConversationContent({ + role: Role.User, + content: 'Initial request' + })); + + let attemptCount = 0; + + mockAI.streamRequest.mockImplementation(async function* () { + attemptCount++; + if (attemptCount === 1) { + // No plan + yield { + content: 'Need more info', + isComplete: true + }; + } else { + // Check that retry message was added + const lastUserMessage = conversation.contents + .filter(c => c.role === Role.User) + .pop(); + const hasRetryMessage = lastUserMessage?.content?.includes('submit'); + + if (hasRetryMessage) { + yield { + content: 'Plan', + functionCall: new AIFunctionCall( + AIFunction.SubmitPlan, + { + steps: [ + { + description: 'Step', + instruction: 'Instruction' + } + ] + }, + 'tool-1' + ), + isComplete: true + }; + } + } + }); + + service.resolveAIProvider(); + const result = await service.runPlanningAgent(conversation, callbacks, false); + + expect(result).toBeDefined(); + expect(conversation.contents.length).toBeGreaterThan(1); + }); + }); +});