Taskmaster

import { jest } from '@jest/globals'; import { z } from 'zod'; // Mock the AI SDK module const mockGenerateObject = jest.fn(); const mockStreamObject = jest.fn(); const mockZodSchema = jest.fn((schema) => ({ _zodSchema: schema })); jest.unstable_mockModule('ai', () => ({ generateObject: mockGenerateObject, streamObject: mockStreamObject, zodSchema: mockZodSchema, generateText: jest.fn(), streamText: jest.fn(), JSONParseError: class JSONParseError extends Error {}, NoObjectGeneratedError: class NoObjectGeneratedError extends Error { static isInstance(error) { return error instanceof NoObjectGeneratedError; } } })); // Mock the gemini-cli SDK module const mockModel = jest.fn((modelId) => ({ modelId, type: 'gemini-cli-model' })); mockModel.languageModel = mockModel; mockModel.chat = mockModel; jest.unstable_mockModule('ai-sdk-provider-gemini-cli', () => ({ createGeminiProvider: jest.fn(() => mockModel) })); // Mock utilities jest.unstable_mockModule('../../../scripts/modules/utils.js', () => ({ log: jest.fn(), findProjectRoot: jest.fn(() => '/mock/project'), resolveEnvVariable: jest.fn((key) => process.env[key]) })); jest.unstable_mockModule('../../../scripts/modules/config-manager.js', () => ({ isProxyEnabled: jest.fn(() => false), getAnonymousTelemetryEnabled: jest.fn(() => true), setSuppressConfigWarnings: jest.fn(), isConfigWarningSuppressed: jest.fn(() => false) })); // Import after mocking const { GeminiCliProvider } = await import( '../../../src/ai-providers/gemini-cli.js' ); describe('GeminiCliProvider Structured Output Integration', () => { let provider; // Sample Zod schema for testing const testSchema = z.object({ title: z.string(), description: z.string(), priority: z.enum(['low', 'medium', 'high']) }); beforeEach(() => { provider = new GeminiCliProvider(); jest.clearAllMocks(); // Reset mock implementations mockGenerateObject.mockReset(); mockStreamObject.mockReset(); }); describe('generateObject', () => { it('should forward schema to AI SDK generateObject', async () => { // Setup mock response mockGenerateObject.mockResolvedValue({ object: { title: 'Test Task', description: 'A test task description', priority: 'high' }, usage: { promptTokens: 100, completionTokens: 50, totalTokens: 150 } }); const params = { modelId: 'gemini-2.5-pro', messages: [ { role: 'system', content: 'You are a helpful assistant.' }, { role: 'user', content: 'Create a task for testing.' } ], schema: testSchema, objectName: 'task', maxTokens: 1000 }; const result = await provider.generateObject(params); // Verify generateObject was called expect(mockGenerateObject).toHaveBeenCalledTimes(1); // Verify schema was passed through const callArgs = mockGenerateObject.mock.calls[0][0]; expect(callArgs.schema).toBe(testSchema); // Verify result is returned correctly expect(result.object).toEqual({ title: 'Test Task', description: 'A test task description', priority: 'high' }); }); it('should use mode "auto" since needsExplicitJsonSchema is false', async () => { mockGenerateObject.mockResolvedValue({ object: { title: 'Test', description: 'Test', priority: 'low' }, usage: { promptTokens: 10, completionTokens: 10, totalTokens: 20 } }); await provider.generateObject({ modelId: 'gemini-2.5-flash', messages: [{ role: 'user', content: 'Test' }], schema: testSchema, objectName: 'task', maxTokens: 500 }); const callArgs = mockGenerateObject.mock.calls[0][0]; // Mode should be 'auto' because needsExplicitJsonSchema is false expect(callArgs.mode).toBe('auto'); // Verify the provider flag is correctly set expect(provider.needsExplicitJsonSchema).toBe(false); }); it('should pass schemaName and schemaDescription', async () => { mockGenerateObject.mockResolvedValue({ object: { title: 'Test', description: 'Test', priority: 'medium' }, usage: { promptTokens: 10, completionTokens: 10, totalTokens: 20 } }); await provider.generateObject({ modelId: 'gemini-2.5-pro', messages: [{ role: 'user', content: 'Test' }], schema: testSchema, objectName: 'myCustomObject', maxTokens: 500 }); const callArgs = mockGenerateObject.mock.calls[0][0]; expect(callArgs.schemaName).toBe('myCustomObject'); expect(callArgs.schemaDescription).toBe( 'Generate a valid JSON object for myCustomObject' ); }); it('should return usage statistics from SDK response', async () => { mockGenerateObject.mockResolvedValue({ object: { title: 'Test', description: 'Test', priority: 'high' }, usage: { promptTokens: 250, completionTokens: 100, totalTokens: 350 } }); const result = await provider.generateObject({ modelId: 'gemini-2.5-pro', messages: [{ role: 'user', content: 'Test' }], schema: testSchema, objectName: 'task', maxTokens: 1000 }); expect(result.usage).toEqual({ inputTokens: 250, outputTokens: 100, totalTokens: 350 }); }); it('should not include temperature when supportsTemperature is false', async () => { mockGenerateObject.mockResolvedValue({ object: { title: 'Test', description: 'Test', priority: 'low' }, usage: { promptTokens: 10, completionTokens: 10, totalTokens: 20 } }); // GeminiCliProvider sets supportsTemperature = false expect(provider.supportsTemperature).toBe(false); await provider.generateObject({ modelId: 'gemini-2.5-pro', messages: [{ role: 'user', content: 'Test' }], schema: testSchema, objectName: 'task', maxTokens: 500, temperature: 0.7 // This should be ignored }); const callArgs = mockGenerateObject.mock.calls[0][0]; expect(callArgs.temperature).toBeUndefined(); }); }); describe('streamObject', () => { it('should forward schema to AI SDK streamObject with zodSchema wrapper', async () => { const mockStreamResult = { partialObjectStream: { [Symbol.asyncIterator]: async function* () { yield { title: 'Test' }; yield { title: 'Test', description: 'Description' }; yield { title: 'Test', description: 'Description', priority: 'high' }; } } }; mockStreamObject.mockResolvedValue(mockStreamResult); const params = { modelId: 'gemini-2.5-pro', messages: [{ role: 'user', content: 'Stream a task' }], schema: testSchema, maxTokens: 1000 }; const result = await provider.streamObject(params); // Verify streamObject was called expect(mockStreamObject).toHaveBeenCalledTimes(1); // Verify zodSchema wrapper was used expect(mockZodSchema).toHaveBeenCalledWith(testSchema); // Verify the wrapped schema was passed const callArgs = mockStreamObject.mock.calls[0][0]; expect(callArgs.schema).toEqual({ _zodSchema: testSchema }); // Verify stream result is returned expect(result).toBe(mockStreamResult); }); it('should use default mode "auto" for streamObject', async () => { mockStreamObject.mockResolvedValue({ partialObjectStream: {} }); await provider.streamObject({ modelId: 'gemini-2.5-flash', messages: [{ role: 'user', content: 'Test' }], schema: testSchema, maxTokens: 500 }); const callArgs = mockStreamObject.mock.calls[0][0]; expect(callArgs.mode).toBe('auto'); }); it('should pass maxOutputTokens to streamObject', async () => { mockStreamObject.mockResolvedValue({ partialObjectStream: {} }); await provider.streamObject({ modelId: 'gemini-2.5-pro', messages: [{ role: 'user', content: 'Test' }], schema: testSchema, maxTokens: 2000 }); const callArgs = mockStreamObject.mock.calls[0][0]; expect(callArgs.maxOutputTokens).toBe(2000); }); }); describe('SDK integration', () => { it('should create model with correct modelId', async () => { mockGenerateObject.mockResolvedValue({ object: { title: 'Test', description: 'Test', priority: 'low' }, usage: { totalTokens: 10 } }); await provider.generateObject({ modelId: 'gemini-2.5-pro', messages: [{ role: 'user', content: 'Test' }], schema: testSchema, objectName: 'task' }); const callArgs = mockGenerateObject.mock.calls[0][0]; // The model should be the result of calling mockModel with the modelId expect(callArgs.model).toEqual({ modelId: 'gemini-2.5-pro', type: 'gemini-cli-model' }); }); it('should work with gemini-2.5-flash model', async () => { mockGenerateObject.mockResolvedValue({ object: { title: 'Fast', description: 'Quick response', priority: 'medium' }, usage: { totalTokens: 20 } }); const result = await provider.generateObject({ modelId: 'gemini-2.5-flash', messages: [{ role: 'user', content: 'Quick task' }], schema: testSchema, objectName: 'task' }); expect(result.object.title).toBe('Fast'); const callArgs = mockGenerateObject.mock.calls[0][0]; expect(callArgs.model.modelId).toBe('gemini-2.5-flash'); }); it('should work with gemini-3-pro-preview model', async () => { mockGenerateObject.mockResolvedValue({ object: { title: 'Preview', description: 'Latest model', priority: 'high' }, usage: { totalTokens: 30 } }); const result = await provider.generateObject({ modelId: 'gemini-3-pro-preview', messages: [{ role: 'user', content: 'Test preview' }], schema: testSchema, objectName: 'task' }); expect(result.object.title).toBe('Preview'); const callArgs = mockGenerateObject.mock.calls[0][0]; expect(callArgs.model.modelId).toBe('gemini-3-pro-preview'); }); }); });