Create new tests for Orchestration, Planning and Execution agents.

__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');
+		});
+	});
+});

__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);
+		});
+	});
+});

__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');
+		});
+	});
+});

__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);
+		});
+	});
+});