[AXON-1796] Add agent model selection to Rovo Dev chat (#1660)

Author: marcomura

CHANGELOG.md

@@ -1,5 +1,12 @@
 ### [Report an Issue](https://github.com/atlassian/atlascode/issues)
 
+
+## What's new in 4.0.22
+
+### Features
+
+- Rovo Dev: Agent model selection both via /models command, and dedicated drop-down menu
+
 ## What's new in 4.0.21
 
 ### Bug fixes

src/rovo-dev/client/responseParser.test.ts

@@ -700,5 +700,79 @@ describe('RovoDevResponseParser', () => {
             expect(toolReturnEvent.toolCallMessage).toBeDefined();
             expect(toolReturnEvent.toolCallMessage.tool_name).toBe('get_weather');
         });
+
+        describe('models responses', () => {
+            it('should parse a models response with model change message', () => {
+                const input =
+                    'event: models\ndata: {"model_name": "GPT-4", "model_id": "gpt-4", "message": "Agent model changed to GPT-4"}\n\n';
+
+                const results = Array.from(parser.parse(input));
+
+                expect(results).toHaveLength(1);
+                expect(results[0]).toEqual({
+                    event_kind: 'models',
+                    data: {
+                        model_name: 'GPT-4',
+                        model_id: 'gpt-4',
+                        message: 'Agent model changed to GPT-4',
+                    },
+                });
+            });
+
+            it('should parse a models response with available models list', () => {
+                const input = `event: models\ndata: {"models": [{"name": "GPT-4", "model_id": "gpt-4", "description": "Most capable", "credit_multiplier": "1.5"}, {"name": "GPT-3.5", "model_id": "gpt-3.5-turbo", "description": "Fast", "credit_multiplier": "1.0"}]}\n\n`;
+
+                const results = Array.from(parser.parse(input));
+
+                expect(results).toHaveLength(1);
+                expect(results[0].event_kind).toBe('models');
+                const modelsResponse = results[0] as any;
+                expect(modelsResponse.data.models).toHaveLength(2);
+                expect(modelsResponse.data.models[0]).toEqual({
+                    name: 'GPT-4',
+                    model_id: 'gpt-4',
+                    description: 'Most capable',
+                    credit_multiplier: '1.5',
+                });
+            });
+
+            it('should parse models response with empty models array', () => {
+                const input = 'event: models\ndata: {"models": []}\n\n';
+
+                const results = Array.from(parser.parse(input));
+
+                expect(results).toHaveLength(1);
+                expect(results[0].event_kind).toBe('models');
+                const modelsResponse = results[0] as any;
+                expect(modelsResponse.data.models).toEqual([]);
+            });
+
+            it('should handle models event as a non-bufferable generic event', () => {
+                const allResults: RovoDevResponse[] = [];
+
+                // Parse a models event
+                let input = 'event: models\ndata: {"model_id": "gpt-4", "message": "Model changed"}\n\n';
+                allResults.push(...Array.from(parser.parse(input)));
+
+                // Parse another generic event to ensure models is not buffered
+                input = 'event: status\ndata: {}\n\n';
+                allResults.push(...Array.from(parser.parse(input)));
+
+                expect(allResults).toHaveLength(2);
+                expect(allResults[0].event_kind).toBe('models');
+                expect(allResults[1].event_kind).toBe('status');
+            });
+
+            it('should parse models event with both message and models array', () => {
+                const input = `event: models\ndata: {"message": "Available models", "models": [{"name": "GPT-4", "model_id": "gpt-4", "description": "Test", "credit_multiplier": "1.0"}]}\n\n`;
+
+                const results = Array.from(parser.parse(input));
+
+                expect(results).toHaveLength(1);
+                const modelsResponse = results[0] as any;
+                expect(modelsResponse.data.message).toBe('Available models');
+                expect(modelsResponse.data.models).toHaveLength(1);
+            });
+        });
     });
 });

src/rovo-dev/client/responseParser.ts

@@ -3,6 +3,7 @@
 import {
     RovoDevClearResponse,
     RovoDevExceptionResponse,
+    RovoDevModelsResponse,
     RovoDevOnCallToolStartResponse,
     RovoDevParsingError,
     RovoDevPromptsResponse,
@@ -190,6 +191,7 @@ type RovoDevSingleChunk =
     | RovoDevStatusChunk
     | RovoDevUsageChunk
     | RovoDevPromptsChunk
+    | RovoDevModelsResponse
     | RovoDevCloseChunk
     | RovoDevReplayEndChunk
     | RovoDevRequestUsageChunk;
@@ -536,6 +538,7 @@ export class RovoDevResponseParser {
             case 'status':
             case 'usage':
             case 'prompts':
+            case 'models':
                 return buffer
                     ? generateError(Error(`Rovo Dev parser error: ${chunk.event_kind} seem to be split`))
                     : chunk;

src/rovo-dev/client/responseParserInterfaces.ts

@@ -137,6 +137,21 @@ export interface RovoDevPromptsResponse {
     };
 }
 
+export interface RovoDevModelsResponse {
+    event_kind: 'models';
+    data: {
+        model_name?: string;
+        model_id?: string;
+        message?: string;
+        models?: {
+            name: string;
+            model_id: string;
+            description: string;
+            credit_multiplier: string;
+        }[];
+    };
+}
+
 export type RovoDevResponse =
     | RovoDevParsingError
     | RovoDevUserPromptResponse
@@ -152,6 +167,7 @@ export type RovoDevResponse =
     | RovoDevStatusResponse
     | RovoDevUsageResponse
     | RovoDevPromptsResponse
+    | RovoDevModelsResponse
     | RovoDevCloseResponse
     | RovoDevReplayEndResponse
     | RovoDevIgnoredResponse;

src/rovo-dev/client/rovoDevApiClient.test.ts

@@ -1874,4 +1874,219 @@ describe('RovoDevApiClient', () => {
             );
         });
     });
+
+    describe('getAgentModel method', () => {
+        it('should return current agent model successfully', async () => {
+            const mockGetAgentModelResponse = {
+                model_name: 'GPT-4',
+                model_id: 'gpt-4',
+                credit_multiplier: '1.5',
+                message: 'Current model is GPT-4',
+            };
+            const mockResponse = {
+                status: 200,
+                json: jest.fn().mockResolvedValue(mockGetAgentModelResponse),
+                headers: mockStandardResponseHeaders(),
+            } as unknown as Response;
+
+            mockFetch.mockResolvedValue(mockResponse);
+
+            const result = await client.getAgentModel();
+
+            expect(mockFetch).toHaveBeenCalledWith('http://localhost:8080/v3/agent-model', {
+                method: 'GET',
+                headers: {
+                    accept: 'text/event-stream',
+                    'Content-Type': 'application/json',
+                    Authorization: 'Bearer sessionToken',
+                },
+                body: undefined,
+            });
+            expect(result).toEqual(mockGetAgentModelResponse);
+            expect(result.model_id).toBe('gpt-4');
+            expect(result.model_name).toBe('GPT-4');
+            expect(result.credit_multiplier).toBe('1.5');
+        });
+
+        it('should throw error when API call fails', async () => {
+            const mockResponse = {
+                status: 500,
+                statusText: 'Internal Server Error',
+                headers: mockStandardResponseHeaders(),
+            } as Response;
+
+            mockFetch.mockResolvedValue(mockResponse);
+
+            await expect(client.getAgentModel()).rejects.toThrow("Failed to fetch '/v3/agent-model API: HTTP 500");
+        });
+    });
+
+    describe('setAgentModel method', () => {
+        it('should set agent model successfully', async () => {
+            const mockSetAgentModelResponse = {
+                model_name: 'GPT-4',
+                model_id: 'gpt-4',
+                message: 'Agent model set to GPT-4',
+            };
+            const mockResponse = {
+                status: 200,
+                json: jest.fn().mockResolvedValue(mockSetAgentModelResponse),
+                headers: mockStandardResponseHeaders(),
+            } as unknown as Response;
+
+            mockFetch.mockResolvedValue(mockResponse);
+
+            const result = await client.setAgentModel('gpt-4');
+
+            expect(mockFetch).toHaveBeenCalledWith('http://localhost:8080/v3/agent-model', {
+                method: 'PUT',
+                headers: {
+                    accept: 'text/event-stream',
+                    'Content-Type': 'application/json',
+                    Authorization: 'Bearer sessionToken',
+                },
+                body: JSON.stringify({ model_id: 'gpt-4' }),
+            });
+            expect(result.model_id).toBe('gpt-4');
+            expect(result.message).toBe('Agent model set to GPT-4');
+        });
+
+        it('should handle different model IDs', async () => {
+            const mockSetAgentModelResponse = {
+                model_name: 'Claude 3',
+                model_id: 'claude-3',
+                message: 'Agent model set to Claude 3',
+            };
+            const mockResponse = {
+                status: 200,
+                json: jest.fn().mockResolvedValue(mockSetAgentModelResponse),
+                headers: mockStandardResponseHeaders(),
+            } as unknown as Response;
+
+            mockFetch.mockResolvedValue(mockResponse);
+
+            const result = await client.setAgentModel('claude-3');
+
+            expect(mockFetch).toHaveBeenCalledWith('http://localhost:8080/v3/agent-model', {
+                method: 'PUT',
+                headers: {
+                    accept: 'text/event-stream',
+                    'Content-Type': 'application/json',
+                    Authorization: 'Bearer sessionToken',
+                },
+                body: JSON.stringify({ model_id: 'claude-3' }),
+            });
+            expect(result.model_id).toBe('claude-3');
+        });
+
+        it('should throw error when API call fails', async () => {
+            const mockResponse = {
+                status: 500,
+                statusText: 'Internal Server Error',
+                headers: mockStandardResponseHeaders(),
+            } as Response;
+
+            mockFetch.mockResolvedValue(mockResponse);
+
+            await expect(client.setAgentModel('gpt-4')).rejects.toThrow(
+                "Failed to fetch '/v3/agent-model API: HTTP 500",
+            );
+        });
+
+        it('should throw error when API returns 400 (invalid model)', async () => {
+            const mockResponse = {
+                status: 400,
+                statusText: 'Bad Request',
+                headers: mockStandardResponseHeaders(),
+            } as Response;
+
+            mockFetch.mockResolvedValue(mockResponse);
+
+            await expect(client.setAgentModel('invalid-model')).rejects.toThrow(
+                "Failed to fetch '/v3/agent-model API: HTTP 400",
+            );
+        });
+    });
+
+    describe('getAvailableAgentModels method', () => {
+        it('should return list of available agent models successfully', async () => {
+            const mockAvailableModelsResponse = {
+                models: [
+                    {
+                        name: 'GPT-4',
+                        model_id: 'gpt-4',
+                        description: 'Most capable model',
+                        credit_multiplier: '1.5',
+                    },
+                    {
+                        name: 'GPT-3.5 Turbo',
+                        model_id: 'gpt-3.5-turbo',
+                        description: 'Fast and efficient',
+                        credit_multiplier: '1.0',
+                    },
+                    {
+                        name: 'Claude 3',
+                        model_id: 'claude-3',
+                        description: 'Anthropic model',
+                        credit_multiplier: '2.0',
+                    },
+                ],
+            };
+            const mockResponse = {
+                status: 200,
+                json: jest.fn().mockResolvedValue(mockAvailableModelsResponse),
+                headers: mockStandardResponseHeaders(),
+            } as unknown as Response;
+
+            mockFetch.mockResolvedValue(mockResponse);
+
+            const result = await client.getAvailableAgentModels();
+
+            expect(mockFetch).toHaveBeenCalledWith('http://localhost:8080/v3/agent-models', {
+                method: 'GET',
+                headers: {
+                    accept: 'text/event-stream',
+                    'Content-Type': 'application/json',
+                    Authorization: 'Bearer sessionToken',
+                },
+                body: undefined,
+            });
+            expect(result).toEqual(mockAvailableModelsResponse);
+            expect(result.models).toHaveLength(3);
+            expect(result.models[0].model_id).toBe('gpt-4');
+            expect(result.models[0].name).toBe('GPT-4');
+            expect(result.models[0].credit_multiplier).toBe('1.5');
+        });
+
+        it('should handle empty models list', async () => {
+            const mockAvailableModelsResponse = {
+                models: [],
+            };
+            const mockResponse = {
+                status: 200,
+                json: jest.fn().mockResolvedValue(mockAvailableModelsResponse),
+                headers: mockStandardResponseHeaders(),
+            } as unknown as Response;
+
+            mockFetch.mockResolvedValue(mockResponse);
+
+            const result = await client.getAvailableAgentModels();
+
+            expect(result.models).toHaveLength(0);
+        });
+
+        it('should throw error when API call fails', async () => {
+            const mockResponse = {
+                status: 500,
+                statusText: 'Internal Server Error',
+                headers: mockStandardResponseHeaders(),
+            } as Response;
+
+            mockFetch.mockResolvedValue(mockResponse);
+
+            await expect(client.getAvailableAgentModels()).rejects.toThrow(
+                "Failed to fetch '/v3/agent-models API: HTTP 500",
+            );
+        });
+    });
 });