Skip to content

Commit 445be72

Browse files
rifkybujanarifkybujana
committed
feat: add support for working directory in GenerationService and GenerateHandler, including tests for workspace context management
1 parent 7de1a17 commit 445be72

5 files changed

Lines changed: 253 additions & 1 deletion

File tree

‎packages/cli/src/api/handlers/generate.handler.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ export class GenerateHandler implements RouteHandler {
3333
const model = body?.model;
3434
const apiKey = body?.api_key;
3535
const baseUrl = body?.base_url;
36+
const workingDirectory = body?.working_directory;
3637

3738
if (!input) {
3839
return HttpUtils.sendJson(
@@ -58,6 +59,7 @@ export class GenerateHandler implements RouteHandler {
5859
model,
5960
apiKey,
6061
baseUrl,
62+
workingDirectory,
6163
);
6264
} else {
6365
await this.handleNonStreamingResponse(
@@ -70,6 +72,7 @@ export class GenerateHandler implements RouteHandler {
7072
model,
7173
apiKey,
7274
baseUrl,
75+
workingDirectory,
7376
);
7477
}
7578
} catch (e) {
@@ -92,6 +95,7 @@ export class GenerateHandler implements RouteHandler {
9295
model?: string,
9396
apiKey?: string,
9497
baseUrl?: string,
98+
workingDirectory?: string,
9599
): Promise<void> {
96100
HttpUtils.setupSseHeaders(res, enableCors);
97101

@@ -135,6 +139,7 @@ export class GenerateHandler implements RouteHandler {
135139
model,
136140
apiKey,
137141
baseUrl,
142+
workingDirectory,
138143
},
139144
);
140145

@@ -154,6 +159,7 @@ export class GenerateHandler implements RouteHandler {
154159
model?: string,
155160
apiKey?: string,
156161
baseUrl?: string,
162+
workingDirectory?: string,
157163
): Promise<void> {
158164
const { finalText, transcript, history: updatedHistory } =
159165
await this.generationService.generateResponse(
@@ -165,6 +171,7 @@ export class GenerateHandler implements RouteHandler {
165171
model,
166172
apiKey,
167173
baseUrl,
174+
workingDirectory,
168175
},
169176
);
170177

‎packages/cli/src/api/services/generation.service.test.ts‎

Lines changed: 211 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,211 @@
1+
/**
2+
* @license
3+
* Copyright 2025 Kolosal Inc.
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest';
8+
import type { Config, WorkspaceContext } from '@kolosal-ai/kolosal-ai-core';
9+
import { ApprovalMode } from '@kolosal-ai/kolosal-ai-core';
10+
import { GenerationService } from '../services/generation.service.js';
11+
12+
// Mock the core module
13+
vi.mock('@kolosal-ai/kolosal-ai-core', async () => {
14+
const actual = await vi.importActual('@kolosal-ai/kolosal-ai-core');
15+
return {
16+
...actual,
17+
WorkspaceContext: vi.fn().mockImplementation((directory: string) => ({
18+
directory,
19+
getDirectories: vi.fn().mockReturnValue([directory]),
20+
})),
21+
};
22+
});
23+
24+
describe('GenerationService - Working Directory', () => {
25+
let mockConfig: Partial<Config>;
26+
let generationService: GenerationService;
27+
let originalWorkspaceContext: Partial<WorkspaceContext>;
28+
29+
beforeEach(() => {
30+
// Create original workspace context
31+
originalWorkspaceContext = {
32+
getDirectories: vi.fn().mockReturnValue(['/original/path']),
33+
};
34+
35+
// Create mock config
36+
mockConfig = {
37+
getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT),
38+
setApprovalMode: vi.fn(),
39+
getWorkspaceContext: vi.fn().mockReturnValue(originalWorkspaceContext),
40+
setWorkspaceContext: vi.fn(),
41+
getGeminiClient: vi.fn().mockReturnValue({
42+
isInitialized: vi.fn().mockReturnValue(true),
43+
setHistory: vi.fn(),
44+
sendMessageStream: vi.fn().mockReturnValue([]),
45+
}),
46+
};
47+
48+
generationService = new GenerationService(mockConfig as Config);
49+
});
50+
51+
afterEach(() => {
52+
vi.clearAllMocks();
53+
});
54+
55+
it('should not change workspace context when no working directory is provided', async () => {
56+
const abortController = new AbortController();
57+
58+
// Mock the generation flow to avoid actual AI calls
59+
const mockSendMessageStream = vi.fn().mockReturnValue([]);
60+
mockConfig.getGeminiClient = vi.fn().mockReturnValue({
61+
isInitialized: vi.fn().mockReturnValue(true),
62+
setHistory: vi.fn(),
63+
sendMessageStream: mockSendMessageStream,
64+
});
65+
66+
try {
67+
await generationService.generateResponse(
68+
'test input',
69+
'test-prompt-id',
70+
abortController.signal,
71+
{
72+
// No working directory provided
73+
}
74+
);
75+
} catch {
76+
// Ignore errors for this test - we're only testing workspace context behavior
77+
}
78+
79+
// Workspace context should not be changed
80+
expect(mockConfig.setWorkspaceContext).not.toHaveBeenCalled();
81+
});
82+
83+
it('should temporarily change workspace context when working directory is provided', async () => {
84+
const abortController = new AbortController();
85+
const testWorkingDirectory = '/test/working/directory';
86+
87+
// Mock the generation flow to avoid actual AI calls
88+
const mockSendMessageStream = vi.fn().mockReturnValue([]);
89+
mockConfig.getGeminiClient = vi.fn().mockReturnValue({
90+
isInitialized: vi.fn().mockReturnValue(true),
91+
setHistory: vi.fn(),
92+
sendMessageStream: mockSendMessageStream,
93+
});
94+
95+
try {
96+
await generationService.generateResponse(
97+
'test input',
98+
'test-prompt-id',
99+
abortController.signal,
100+
{
101+
workingDirectory: testWorkingDirectory,
102+
}
103+
);
104+
} catch {
105+
// Ignore errors for this test - we're only testing workspace context behavior
106+
}
107+
108+
// Workspace context should be changed and then restored
109+
expect(mockConfig.setWorkspaceContext).toHaveBeenCalledTimes(2);
110+
111+
// First call should set the new workspace context
112+
const firstCall = (mockConfig.setWorkspaceContext as any).mock.calls[0];
113+
expect(firstCall[0]).toEqual(expect.objectContaining({
114+
getDirectories: expect.any(Function),
115+
}));
116+
117+
// Second call should restore the original workspace context
118+
const secondCall = (mockConfig.setWorkspaceContext as any).mock.calls[1];
119+
expect(secondCall[0]).toBe(originalWorkspaceContext);
120+
});
121+
122+
it('should restore original workspace context even if generation fails', async () => {
123+
const abortController = new AbortController();
124+
const testWorkingDirectory = '/test/working/directory';
125+
126+
// Mock the generation to throw an error
127+
mockConfig.getGeminiClient = vi.fn().mockReturnValue({
128+
isInitialized: vi.fn().mockReturnValue(true),
129+
setHistory: vi.fn(),
130+
sendMessageStream: vi.fn().mockRejectedValue(new Error('Generation failed')),
131+
});
132+
133+
try {
134+
await generationService.generateResponse(
135+
'test input',
136+
'test-prompt-id',
137+
abortController.signal,
138+
{
139+
workingDirectory: testWorkingDirectory,
140+
}
141+
);
142+
} catch {
143+
// Expected to fail
144+
}
145+
146+
// Workspace context should still be restored even after error
147+
expect(mockConfig.setWorkspaceContext).toHaveBeenCalledTimes(2);
148+
149+
// Second call should restore the original workspace context
150+
const secondCall = (mockConfig.setWorkspaceContext as any).mock.calls[1];
151+
expect(secondCall[0]).toBe(originalWorkspaceContext);
152+
});
153+
154+
it('should handle WorkspaceContext creation failure gracefully', async () => {
155+
const abortController = new AbortController();
156+
const testWorkingDirectory = '/test/working/directory';
157+
158+
// Mock the generation flow
159+
const mockSendMessageStream = vi.fn().mockReturnValue([]);
160+
mockConfig.getGeminiClient = vi.fn().mockReturnValue({
161+
isInitialized: vi.fn().mockReturnValue(true),
162+
setHistory: vi.fn(),
163+
sendMessageStream: mockSendMessageStream,
164+
});
165+
166+
// This test verifies that the code handles directory setup gracefully
167+
// Since WorkspaceContext is mocked to succeed, we just verify the workspace context was set
168+
try {
169+
await generationService.generateResponse(
170+
'test input',
171+
'test-prompt-id',
172+
abortController.signal,
173+
{
174+
workingDirectory: testWorkingDirectory,
175+
}
176+
);
177+
} catch {
178+
// Ignore errors for this test
179+
}
180+
181+
// The workspace context should be changed even with mocked implementation
182+
expect(mockConfig.setWorkspaceContext).toHaveBeenCalledTimes(2);
183+
});
184+
185+
it('should set approval mode to YOLO and restore it afterwards', async () => {
186+
const abortController = new AbortController();
187+
188+
// Mock the generation flow
189+
const mockSendMessageStream = vi.fn().mockReturnValue([]);
190+
mockConfig.getGeminiClient = vi.fn().mockReturnValue({
191+
isInitialized: vi.fn().mockReturnValue(true),
192+
setHistory: vi.fn(),
193+
sendMessageStream: mockSendMessageStream,
194+
});
195+
196+
try {
197+
await generationService.generateResponse(
198+
'test input',
199+
'test-prompt-id',
200+
abortController.signal,
201+
{}
202+
);
203+
} catch {
204+
// Ignore errors for this test
205+
}
206+
207+
// Should set to YOLO mode and then restore original
208+
expect(mockConfig.setApprovalMode).toHaveBeenCalledWith(ApprovalMode.YOLO);
209+
expect(mockConfig.setApprovalMode).toHaveBeenLastCalledWith(ApprovalMode.DEFAULT);
210+
});
211+
});

‎packages/cli/src/api/services/generation.service.ts‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,38 @@ export class GenerationService {
2929
model?: string;
3030
apiKey?: string;
3131
baseUrl?: string;
32+
workingDirectory?: string;
3233
} = {},
3334
): Promise<GenerationResult> {
34-
const { onContentChunk, onEvent, conversationHistory, model, apiKey, baseUrl } = options;
35+
const { onContentChunk, onEvent, conversationHistory, model, apiKey, baseUrl, workingDirectory } = options;
3536

3637
// Set approval mode to YOLO for API requests to auto-approve tool calls
3738
const originalApprovalMode = this.config.getApprovalMode();
3839
this.config.setApprovalMode(ApprovalMode.YOLO);
3940

41+
// Store original workspace context if we need to temporarily change it
42+
let originalWorkspaceContext: any = null;
43+
let tempWorkspaceContext: any = null;
44+
45+
if (workingDirectory) {
46+
try {
47+
// Import the WorkspaceContext class
48+
const { WorkspaceContext } = await import('@kolosal-ai/kolosal-ai-core');
49+
50+
// Store the original workspace context
51+
originalWorkspaceContext = this.config.getWorkspaceContext();
52+
53+
// Create a temporary workspace context with the specified working directory
54+
tempWorkspaceContext = new WorkspaceContext(workingDirectory, []);
55+
56+
// Temporarily replace the workspace context in the config
57+
this.config.setWorkspaceContext(tempWorkspaceContext);
58+
} catch (error) {
59+
console.warn('[API] Failed to set working directory:', error);
60+
// Continue with the original workspace context
61+
}
62+
}
63+
4064
try {
4165
let geminiClient = this.config.getGeminiClient();
4266

@@ -84,6 +108,11 @@ export class GenerationService {
84108
onEvent,
85109
);
86110
} finally {
111+
// Restore original workspace context if it was changed
112+
if (originalWorkspaceContext && workingDirectory) {
113+
this.config.setWorkspaceContext(originalWorkspaceContext);
114+
}
115+
87116
// Restore original approval mode
88117
this.config.setApprovalMode(originalApprovalMode);
89118
}

‎packages/cli/src/api/types/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ export interface GenerateRequest {
4747
model?: string;
4848
api_key?: string;
4949
base_url?: string;
50+
working_directory?: string;
5051
}
5152

5253
export interface GenerateResponse {

‎packages/core/src/config/config.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -657,6 +657,10 @@ export class Config {
657657
return this.workspaceContext;
658658
}
659659

660+
setWorkspaceContext(workspaceContext: WorkspaceContext): void {
661+
this.workspaceContext = workspaceContext;
662+
}
663+
660664
getToolRegistry(): ToolRegistry {
661665
return this.toolRegistry;
662666
}

0 commit comments

Comments
 (0)