summaryrefslogtreecommitdiff
path: root/packages/server/src/tools
diff options
context:
space:
mode:
Diffstat (limited to 'packages/server/src/tools')
-rw-r--r--packages/server/src/tools/edit.test.ts23
-rw-r--r--packages/server/src/tools/edit.ts10
-rw-r--r--packages/server/src/tools/shell.ts1
-rw-r--r--packages/server/src/tools/tools.ts3
-rw-r--r--packages/server/src/tools/write-file.test.ts76
-rw-r--r--packages/server/src/tools/write-file.ts8
6 files changed, 94 insertions, 27 deletions
diff --git a/packages/server/src/tools/edit.test.ts b/packages/server/src/tools/edit.test.ts
index 88216d53..08d0860d 100644
--- a/packages/server/src/tools/edit.test.ts
+++ b/packages/server/src/tools/edit.test.ts
@@ -223,7 +223,9 @@ describe('EditTool', () => {
old_string: 'old',
new_string: 'new',
};
- expect(await tool.shouldConfirmExecute(params)).toBe(false);
+ expect(
+ await tool.shouldConfirmExecute(params, new AbortController().signal),
+ ).toBe(false);
});
it('should request confirmation for valid edit', async () => {
@@ -235,7 +237,10 @@ describe('EditTool', () => {
};
// ensureCorrectEdit will be called by shouldConfirmExecute
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 1 });
- const confirmation = await tool.shouldConfirmExecute(params);
+ const confirmation = await tool.shouldConfirmExecute(
+ params,
+ new AbortController().signal,
+ );
expect(confirmation).toEqual(
expect.objectContaining({
title: `Confirm Edit: ${testFile}`,
@@ -253,7 +258,9 @@ describe('EditTool', () => {
new_string: 'new',
};
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 0 });
- expect(await tool.shouldConfirmExecute(params)).toBe(false);
+ expect(
+ await tool.shouldConfirmExecute(params, new AbortController().signal),
+ ).toBe(false);
});
it('should return false if multiple occurrences of old_string are found (ensureCorrectEdit returns > 1)', async () => {
@@ -264,7 +271,9 @@ describe('EditTool', () => {
new_string: 'new',
};
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 2 });
- expect(await tool.shouldConfirmExecute(params)).toBe(false);
+ expect(
+ await tool.shouldConfirmExecute(params, new AbortController().signal),
+ ).toBe(false);
});
it('should request confirmation for creating a new file (empty old_string)', async () => {
@@ -279,7 +288,10 @@ describe('EditTool', () => {
// as shouldConfirmExecute handles this for diff generation.
// If it is called, it should return 0 occurrences for a new file.
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 0 });
- const confirmation = await tool.shouldConfirmExecute(params);
+ const confirmation = await tool.shouldConfirmExecute(
+ params,
+ new AbortController().signal,
+ );
expect(confirmation).toEqual(
expect.objectContaining({
title: `Confirm Edit: ${newFileName}`,
@@ -328,6 +340,7 @@ describe('EditTool', () => {
const confirmation = (await tool.shouldConfirmExecute(
params,
+ new AbortController().signal,
)) as FileDiff;
expect(mockCalled).toBe(true); // Check if the mock implementation was run
diff --git a/packages/server/src/tools/edit.ts b/packages/server/src/tools/edit.ts
index 781483ae..d85c89b0 100644
--- a/packages/server/src/tools/edit.ts
+++ b/packages/server/src/tools/edit.ts
@@ -174,7 +174,10 @@ Expectation for parameters:
* @returns An object describing the potential edit outcome
* @throws File system errors if reading the file fails unexpectedly (e.g., permissions)
*/
- private async calculateEdit(params: EditToolParams): Promise<CalculatedEdit> {
+ private async calculateEdit(
+ params: EditToolParams,
+ abortSignal: AbortSignal,
+ ): Promise<CalculatedEdit> {
const expectedReplacements = 1;
let currentContent: string | null = null;
let fileExists = false;
@@ -210,6 +213,7 @@ Expectation for parameters:
currentContent,
params,
this.client,
+ abortSignal,
);
finalOldString = correctedEdit.params.old_string;
finalNewString = correctedEdit.params.new_string;
@@ -262,6 +266,7 @@ Expectation for parameters:
*/
async shouldConfirmExecute(
params: EditToolParams,
+ abortSignal: AbortSignal,
): Promise<ToolCallConfirmationDetails | false> {
if (this.config.getAlwaysSkipModificationConfirmation()) {
return false;
@@ -300,6 +305,7 @@ Expectation for parameters:
currentContent,
params,
this.client,
+ abortSignal,
);
finalOldString = correctedEdit.params.old_string;
finalNewString = correctedEdit.params.new_string;
@@ -376,7 +382,7 @@ Expectation for parameters:
let editData: CalculatedEdit;
try {
- editData = await this.calculateEdit(params);
+ editData = await this.calculateEdit(params, _signal);
} catch (error) {
const errorMsg = error instanceof Error ? error.message : String(error);
return {
diff --git a/packages/server/src/tools/shell.ts b/packages/server/src/tools/shell.ts
index 6ee36b8d..a708c93f 100644
--- a/packages/server/src/tools/shell.ts
+++ b/packages/server/src/tools/shell.ts
@@ -98,6 +98,7 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
async shouldConfirmExecute(
params: ShellToolParams,
+ _abortSignal: AbortSignal,
): Promise<ToolCallConfirmationDetails | false> {
if (this.validateToolParams(params)) {
return false; // skip confirmation, execute call will fail immediately
diff --git a/packages/server/src/tools/tools.ts b/packages/server/src/tools/tools.ts
index c57bbd39..8ec11bf0 100644
--- a/packages/server/src/tools/tools.ts
+++ b/packages/server/src/tools/tools.ts
@@ -57,6 +57,7 @@ export interface Tool<
*/
shouldConfirmExecute(
params: TParams,
+ abortSignal: AbortSignal,
): Promise<ToolCallConfirmationDetails | false>;
/**
@@ -137,6 +138,8 @@ export abstract class BaseTool<
shouldConfirmExecute(
// eslint-disable-next-line @typescript-eslint/no-unused-vars
params: TParams,
+ // eslint-disable-next-line @typescript-eslint/no-unused-vars
+ abortSignal: AbortSignal,
): Promise<ToolCallConfirmationDetails | false> {
return Promise.resolve(false);
}
diff --git a/packages/server/src/tools/write-file.test.ts b/packages/server/src/tools/write-file.test.ts
index 83e75a33..3fd97c9e 100644
--- a/packages/server/src/tools/write-file.test.ts
+++ b/packages/server/src/tools/write-file.test.ts
@@ -110,18 +110,32 @@ describe('WriteFileTool', () => {
// Default mock implementations that return valid structures
mockEnsureCorrectEdit.mockImplementation(
async (
- currentContent: string,
+ _currentContent: string,
params: EditToolParams,
_client: GeminiClient,
- ): Promise<CorrectedEditResult> =>
- Promise.resolve({
+ signal?: AbortSignal, // Make AbortSignal optional to match usage
+ ): Promise<CorrectedEditResult> => {
+ if (signal?.aborted) {
+ return Promise.reject(new Error('Aborted'));
+ }
+ return Promise.resolve({
params: { ...params, new_string: params.new_string ?? '' },
occurrences: 1,
- }),
+ });
+ },
);
mockEnsureCorrectFileContent.mockImplementation(
- async (content: string, _client: GeminiClient): Promise<string> =>
- Promise.resolve(content ?? ''),
+ async (
+ content: string,
+ _client: GeminiClient,
+ signal?: AbortSignal,
+ ): Promise<string> => {
+ // Make AbortSignal optional
+ if (signal?.aborted) {
+ return Promise.reject(new Error('Aborted'));
+ }
+ return Promise.resolve(content ?? '');
+ },
);
});
@@ -181,6 +195,7 @@ describe('WriteFileTool', () => {
const filePath = path.join(rootDir, 'new_corrected_file.txt');
const proposedContent = 'Proposed new content.';
const correctedContent = 'Corrected new content.';
+ const abortSignal = new AbortController().signal;
// Ensure the mock is set for this specific test case if needed, or rely on beforeEach
mockEnsureCorrectFileContent.mockResolvedValue(correctedContent);
@@ -188,11 +203,13 @@ describe('WriteFileTool', () => {
const result = await tool._getCorrectedFileContent(
filePath,
proposedContent,
+ abortSignal,
);
expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
proposedContent,
mockGeminiClientInstance,
+ abortSignal,
);
expect(mockEnsureCorrectEdit).not.toHaveBeenCalled();
expect(result.correctedContent).toBe(correctedContent);
@@ -206,6 +223,7 @@ describe('WriteFileTool', () => {
const originalContent = 'Original existing content.';
const proposedContent = 'Proposed replacement content.';
const correctedProposedContent = 'Corrected replacement content.';
+ const abortSignal = new AbortController().signal;
fs.writeFileSync(filePath, originalContent, 'utf8');
// Ensure this mock is active and returns the correct structure
@@ -222,6 +240,7 @@ describe('WriteFileTool', () => {
const result = await tool._getCorrectedFileContent(
filePath,
proposedContent,
+ abortSignal,
);
expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
@@ -232,6 +251,7 @@ describe('WriteFileTool', () => {
file_path: filePath,
},
mockGeminiClientInstance,
+ abortSignal,
);
expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
expect(result.correctedContent).toBe(correctedProposedContent);
@@ -243,6 +263,7 @@ describe('WriteFileTool', () => {
it('should return error if reading an existing file fails (e.g. permissions)', async () => {
const filePath = path.join(rootDir, 'unreadable_file.txt');
const proposedContent = 'some content';
+ const abortSignal = new AbortController().signal;
fs.writeFileSync(filePath, 'content', { mode: 0o000 });
const readError = new Error('Permission denied');
@@ -255,6 +276,7 @@ describe('WriteFileTool', () => {
const result = await tool._getCorrectedFileContent(
filePath,
proposedContent,
+ abortSignal,
);
expect(fs.readFileSync).toHaveBeenCalledWith(filePath, 'utf8');
@@ -274,16 +296,17 @@ describe('WriteFileTool', () => {
});
describe('shouldConfirmExecute', () => {
+ const abortSignal = new AbortController().signal;
it('should return false if params are invalid (relative path)', async () => {
const params = { file_path: 'relative.txt', content: 'test' };
- const confirmation = await tool.shouldConfirmExecute(params);
+ const confirmation = await tool.shouldConfirmExecute(params, abortSignal);
expect(confirmation).toBe(false);
});
it('should return false if params are invalid (outside root)', async () => {
const outsidePath = path.resolve(tempDir, 'outside-root.txt');
const params = { file_path: outsidePath, content: 'test' };
- const confirmation = await tool.shouldConfirmExecute(params);
+ const confirmation = await tool.shouldConfirmExecute(params, abortSignal);
expect(confirmation).toBe(false);
});
@@ -298,7 +321,7 @@ describe('WriteFileTool', () => {
throw readError;
});
- const confirmation = await tool.shouldConfirmExecute(params);
+ const confirmation = await tool.shouldConfirmExecute(params, abortSignal);
expect(confirmation).toBe(false);
vi.spyOn(fs, 'readFileSync').mockImplementation(originalReadFileSync);
@@ -314,11 +337,13 @@ describe('WriteFileTool', () => {
const params = { file_path: filePath, content: proposedContent };
const confirmation = (await tool.shouldConfirmExecute(
params,
+ abortSignal,
)) as ToolEditConfirmationDetails;
expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
proposedContent,
mockGeminiClientInstance,
+ abortSignal,
);
expect(confirmation).toEqual(
expect.objectContaining({
@@ -343,7 +368,6 @@ describe('WriteFileTool', () => {
'Corrected replacement for confirmation.';
fs.writeFileSync(filePath, originalContent, 'utf8');
- // Ensure this mock is active and returns the correct structure
mockEnsureCorrectEdit.mockResolvedValue({
params: {
file_path: filePath,
@@ -356,6 +380,7 @@ describe('WriteFileTool', () => {
const params = { file_path: filePath, content: proposedContent };
const confirmation = (await tool.shouldConfirmExecute(
params,
+ abortSignal,
)) as ToolEditConfirmationDetails;
expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
@@ -366,6 +391,7 @@ describe('WriteFileTool', () => {
file_path: filePath,
},
mockGeminiClientInstance,
+ abortSignal,
);
expect(confirmation).toEqual(
expect.objectContaining({
@@ -381,9 +407,10 @@ describe('WriteFileTool', () => {
});
describe('execute', () => {
+ const abortSignal = new AbortController().signal;
it('should return error if params are invalid (relative path)', async () => {
const params = { file_path: 'relative.txt', content: 'test' };
- const result = await tool.execute(params, new AbortController().signal);
+ const result = await tool.execute(params, abortSignal);
expect(result.llmContent).toMatch(/Error: Invalid parameters provided/);
expect(result.returnDisplay).toMatch(/Error: File path must be absolute/);
});
@@ -391,7 +418,7 @@ describe('WriteFileTool', () => {
it('should return error if params are invalid (path outside root)', async () => {
const outsidePath = path.resolve(tempDir, 'outside-root.txt');
const params = { file_path: outsidePath, content: 'test' };
- const result = await tool.execute(params, new AbortController().signal);
+ const result = await tool.execute(params, abortSignal);
expect(result.llmContent).toMatch(/Error: Invalid parameters provided/);
expect(result.returnDisplay).toMatch(
/Error: File path must be within the root directory/,
@@ -409,7 +436,7 @@ describe('WriteFileTool', () => {
throw readError;
});
- const result = await tool.execute(params, new AbortController().signal);
+ const result = await tool.execute(params, abortSignal);
expect(result.llmContent).toMatch(/Error checking existing file/);
expect(result.returnDisplay).toMatch(
/Error checking existing file: Simulated read error for execute/,
@@ -427,16 +454,20 @@ describe('WriteFileTool', () => {
const params = { file_path: filePath, content: proposedContent };
- const confirmDetails = await tool.shouldConfirmExecute(params);
+ const confirmDetails = await tool.shouldConfirmExecute(
+ params,
+ abortSignal,
+ );
if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
}
- const result = await tool.execute(params, new AbortController().signal);
+ const result = await tool.execute(params, abortSignal);
expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
proposedContent,
mockGeminiClientInstance,
+ abortSignal,
);
expect(result.llmContent).toMatch(
/Successfully created and wrote to new file/,
@@ -477,12 +508,15 @@ describe('WriteFileTool', () => {
const params = { file_path: filePath, content: proposedContent };
- const confirmDetails = await tool.shouldConfirmExecute(params);
+ const confirmDetails = await tool.shouldConfirmExecute(
+ params,
+ abortSignal,
+ );
if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
}
- const result = await tool.execute(params, new AbortController().signal);
+ const result = await tool.execute(params, abortSignal);
expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
initialContent,
@@ -492,6 +526,7 @@ describe('WriteFileTool', () => {
file_path: filePath,
},
mockGeminiClientInstance,
+ abortSignal,
);
expect(result.llmContent).toMatch(/Successfully overwrote file/);
expect(fs.readFileSync(filePath, 'utf8')).toBe(correctedProposedContent);
@@ -513,12 +548,15 @@ describe('WriteFileTool', () => {
const params = { file_path: filePath, content };
// Simulate confirmation if your logic requires it before execute, or remove if not needed for this path
- const confirmDetails = await tool.shouldConfirmExecute(params);
+ const confirmDetails = await tool.shouldConfirmExecute(
+ params,
+ abortSignal,
+ );
if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
}
- await tool.execute(params, new AbortController().signal);
+ await tool.execute(params, abortSignal);
expect(fs.existsSync(dirPath)).toBe(true);
expect(fs.statSync(dirPath).isDirectory()).toBe(true);
diff --git a/packages/server/src/tools/write-file.ts b/packages/server/src/tools/write-file.ts
index 60646cc2..2285c819 100644
--- a/packages/server/src/tools/write-file.ts
+++ b/packages/server/src/tools/write-file.ts
@@ -141,6 +141,7 @@ export class WriteFileTool extends BaseTool<WriteFileToolParams, ToolResult> {
*/
async shouldConfirmExecute(
params: WriteFileToolParams,
+ abortSignal: AbortSignal,
): Promise<ToolCallConfirmationDetails | false> {
if (this.config.getAlwaysSkipModificationConfirmation()) {
return false;
@@ -154,6 +155,7 @@ export class WriteFileTool extends BaseTool<WriteFileToolParams, ToolResult> {
const correctedContentResult = await this._getCorrectedFileContent(
params.file_path,
params.content,
+ abortSignal,
);
if (correctedContentResult.error) {
@@ -193,7 +195,7 @@ export class WriteFileTool extends BaseTool<WriteFileToolParams, ToolResult> {
async execute(
params: WriteFileToolParams,
- _signal: AbortSignal,
+ abortSignal: AbortSignal,
): Promise<ToolResult> {
const validationError = this.validateToolParams(params);
if (validationError) {
@@ -206,6 +208,7 @@ export class WriteFileTool extends BaseTool<WriteFileToolParams, ToolResult> {
const correctedContentResult = await this._getCorrectedFileContent(
params.file_path,
params.content,
+ abortSignal,
);
if (correctedContentResult.error) {
@@ -277,6 +280,7 @@ export class WriteFileTool extends BaseTool<WriteFileToolParams, ToolResult> {
private async _getCorrectedFileContent(
filePath: string,
proposedContent: string,
+ abortSignal: AbortSignal,
): Promise<GetCorrectedFileContentResult> {
let originalContent = '';
let fileExists = false;
@@ -316,6 +320,7 @@ export class WriteFileTool extends BaseTool<WriteFileToolParams, ToolResult> {
file_path: filePath,
},
this.client,
+ abortSignal,
);
correctedContent = correctedParams.new_string;
} else {
@@ -323,6 +328,7 @@ export class WriteFileTool extends BaseTool<WriteFileToolParams, ToolResult> {
correctedContent = await ensureCorrectFileContent(
proposedContent,
this.client,
+ abortSignal,
);
}
return { originalContent, correctedContent, fileExists };