Skip to content

Commit 271aaf7

Browse files
feat(utils): add tool call support
Add Ollama tool calling with 6 built-in tools (read_file, write_file, run_shell, list_dir, grep_search, view_range), smart approval by default, and Shift+Tab toggle for full auto mode. 1. Update Message Types **File:** `src/utils/ollama.ts` - Extend `Message` interface to include `tool_calls` and `tool_results` - Add `ToolCall` and `ToolResult` interfaces 2. Add Tool Definitions **File:** `src/utils/tools.ts` (new) - Define 6 tool schemas for Ollama API (name, description, parameters) - Export `TOOLS` array for passing to `client.chat()` 3. Add Tool Execution **File:** `src/utils/tools.ts` - Implement tool handlers using Node.js built-ins: - `read_file`: `fs.readFileSync` - `write_file`: `fs.writeFileSync` - `run_shell`: `child_process.exec` with `promisify` (requires approval) - `list_dir`: `fs.readdirSync` - `grep_search`: Custom implementation with `fs` + regex - `view_range`: `fs.readFileSync` with line slicing - Each returns `{ content: string, error?: string }` 4. Update Chat Streaming **File:** `src/utils/ollama.ts` - Modify `streamChat()` to accept tools parameter - Handle tool call chunks from Ollama response - Yield tool calls separately from content chunks - Add `executeChat()` function for full round-trip (chat → tool call → result → continue) 5. Add Tool Approval UI State **File:** `src/components/Chat.tsx` - Add state: `autoExecuteMode` (boolean), `pendingToolCalls` (array) - Add keyboard handler for Shift+Tab to toggle mode - Show mode indicator in UI (similar to other TUI) 6. Handle Tool Approval Flow **File:** `src/components/Chat.tsx` - When tool call received: - If read-only tool: auto-execute - If write/modify tool and not auto mode: pause stream, show approval prompt - If auto mode: execute without prompt - Execute tools via `executeTool()` from `src/utils/tools.ts` - Send tool results back to Ollama to continue conversation 7. Render Tool Calls and Results (TUI-style inline) **File:** `src/components/Messages.tsx` - Tool calls appear as collapsed/dimmed chips between messages (e.g., "📁 read_file: src/utils.ts") - Show brief description with args, expandable for full details - Use gray/dim color, smaller text, minimal visual weight - Tool results collapsed by default, expandable for full output - Group tool call + result as single inline unit, not separate messages 8. Add Approval Prompt Component **File:** `src/components/ToolApproval.tsx` (new) - Simple Y/n prompt for pending tool calls - Show tool name, arguments, and potential impact - Handle keyboard input (Y = execute, n = skip with message to LLM) 9. Update Tests **Files:** `src/utils/ollama.test.ts`, `src/components/Chat.test.tsx` - Mock tool calls in stream response - Test tool execution and approval flow - Test auto mode toggle - Test tool result rendering in Messages
1 parent a529983 commit 271aaf7

13 files changed

Lines changed: 1584 additions & 48 deletions

src/components/Chat.test.tsx

Lines changed: 313 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -64,10 +64,15 @@ import { Chat } from './Chat';
6464
vi.mock('../utils', () => ({
6565
ollama: {
6666
streamChat: vi.fn().mockImplementation(function* () {
67-
yield 'Mocked';
68-
yield ' response';
67+
yield { type: 'content', content: 'Mocked' };
68+
yield { type: 'content', content: ' response' };
6969
}),
7070
},
71+
tools: {
72+
TOOLS: [],
73+
TOOLS_REQUIRING_APPROVAL: new Set(),
74+
executeTool: vi.fn(),
75+
},
7176
}));
7277

7378
async function typeText(
@@ -133,7 +138,9 @@ describe('Chat', () => {
133138
const frame = lastFrame() ?? '';
134139
const lines = frame
135140
.split('\n')
136-
.filter((line) => line.trim() && !line.includes('>'));
141+
.filter(
142+
(line) => line.trim() && !line.includes('>') && !line.includes('Mode:'),
143+
);
137144
expect(lines).toHaveLength(0);
138145
});
139146

@@ -179,10 +186,311 @@ describe('Chat', () => {
179186
expect(vi.mocked(streamChat)).toHaveBeenLastCalledWith(
180187
expect.any(Array),
181188
'llama3',
189+
expect.any(Array),
182190
);
183191
});
184192
});
185193

194+
describe('Chat keyboard shortcuts', () => {
195+
beforeEach(() => {
196+
mockState.clear();
197+
});
198+
199+
it('toggles autoExecute mode with Shift+Tab', async () => {
200+
const chat = <Chat model="gemma4" onCommand={vi.fn()} />;
201+
const { lastFrame, stdin } = render(chat);
202+
203+
// Initial state should show "Smart" mode
204+
expect(lastFrame()).toContain('Mode: Smart');
205+
206+
// Send Shift+Tab escape sequence
207+
stdin.write('\x1B[Z');
208+
await tick();
209+
210+
// After Shift+Tab, should show "Auto" mode
211+
expect(lastFrame()).toContain('Mode: Auto');
212+
213+
// Toggle back with another Shift+Tab
214+
stdin.write('\x1B[Z');
215+
await tick();
216+
217+
expect(lastFrame()).toContain('Mode: Smart');
218+
});
219+
});
220+
221+
describe('Chat with tool calls', () => {
222+
beforeEach(() => {
223+
mockState.clear();
224+
});
225+
226+
it('shows tool approval when tool requires approval', async () => {
227+
const { ollama, tools } = await import('../utils');
228+
const { streamChat } = ollama;
229+
230+
vi.mocked(streamChat).mockImplementationOnce(async function* () {
231+
await Promise.resolve();
232+
yield {
233+
type: 'tool_calls',
234+
tool_calls: [
235+
{
236+
function: {
237+
name: 'write_file',
238+
arguments: { path: '/test.txt', content: 'hello' },
239+
},
240+
},
241+
],
242+
};
243+
});
244+
245+
// Set write_file as requiring approval
246+
vi.spyOn(tools.TOOLS_REQUIRING_APPROVAL, 'has').mockReturnValue(true);
247+
248+
const chat = <Chat model="gemma4" onCommand={vi.fn()} />;
249+
const { lastFrame, rerender } = render(chat);
250+
251+
await typeText(rerender, 'write a file', chat);
252+
submitInput('write a file');
253+
rerender(chat);
254+
await waitForStream();
255+
rerender(chat);
256+
257+
expect(lastFrame()).toContain('Tool requires approval');
258+
});
259+
260+
it('auto-executes tool that does not require approval', async () => {
261+
const { ollama, tools } = await import('../utils');
262+
const { streamChat } = ollama;
263+
264+
vi.mocked(streamChat).mockImplementationOnce(async function* () {
265+
await Promise.resolve();
266+
yield {
267+
type: 'tool_calls',
268+
tool_calls: [
269+
{
270+
function: {
271+
name: 'read_file',
272+
arguments: { path: '/test.txt' },
273+
},
274+
},
275+
],
276+
};
277+
});
278+
279+
// Mock executeTool to return content (auto-executed since read_file doesn't require approval)
280+
const mockExecute = vi.fn().mockResolvedValue({
281+
content: 'file contents',
282+
});
283+
vi.mocked(tools.executeTool).mockImplementation(mockExecute);
284+
285+
// read_file does not require approval
286+
vi.spyOn(tools.TOOLS_REQUIRING_APPROVAL, 'has').mockReturnValue(false);
287+
288+
const chat = <Chat model="gemma4" onCommand={vi.fn()} />;
289+
const { rerender } = render(chat);
290+
291+
await typeText(rerender, 'read file', chat);
292+
submitInput('read file');
293+
rerender(chat);
294+
await waitForStream();
295+
296+
expect(mockExecute).toHaveBeenCalledWith('read_file', {
297+
path: '/test.txt',
298+
});
299+
});
300+
301+
it('handles tool result error', async () => {
302+
const { ollama, tools } = await import('../utils');
303+
const { streamChat } = ollama;
304+
305+
vi.mocked(streamChat).mockImplementationOnce(async function* () {
306+
await Promise.resolve();
307+
yield {
308+
type: 'tool_calls',
309+
tool_calls: [
310+
{
311+
function: {
312+
name: 'read_file',
313+
arguments: { path: '/missing.txt' },
314+
},
315+
},
316+
],
317+
};
318+
});
319+
320+
// Use local mock implementation for executeTool
321+
const mockExecute = vi.fn().mockResolvedValue({
322+
content: '',
323+
error: 'File not found',
324+
});
325+
vi.mocked(tools.executeTool).mockImplementation(mockExecute);
326+
327+
// read_file does not require approval
328+
vi.spyOn(tools.TOOLS_REQUIRING_APPROVAL, 'has').mockReturnValue(false);
329+
330+
const chat = <Chat model="gemma4" onCommand={vi.fn()} />;
331+
const { lastFrame, rerender } = render(chat);
332+
333+
await typeText(rerender, 'read file', chat);
334+
submitInput('read file');
335+
rerender(chat);
336+
await waitForStream();
337+
338+
// The tool result message should contain the error
339+
expect(lastFrame()).toContain('File not found');
340+
});
341+
342+
it('handles tool approval rejection', async () => {
343+
const { ollama, tools } = await import('../utils');
344+
const { streamChat } = ollama;
345+
346+
vi.mocked(streamChat).mockImplementationOnce(async function* () {
347+
await Promise.resolve();
348+
yield {
349+
type: 'tool_calls',
350+
tool_calls: [
351+
{
352+
function: {
353+
name: 'write_file',
354+
arguments: { path: '/test.txt', content: 'hello' },
355+
},
356+
},
357+
],
358+
};
359+
});
360+
361+
vi.spyOn(tools.TOOLS_REQUIRING_APPROVAL, 'has').mockReturnValue(true);
362+
363+
const chat = <Chat model="gemma4" onCommand={vi.fn()} />;
364+
const { lastFrame, rerender, stdin } = render(chat);
365+
366+
await typeText(rerender, 'write a file', chat);
367+
submitInput('write a file');
368+
rerender(chat);
369+
await waitForStream();
370+
rerender(chat);
371+
372+
// Verify approval prompt is shown
373+
expect(lastFrame()).toContain('Tool requires approval');
374+
375+
// Reject the tool (move to No with right arrow, then Enter)
376+
stdin.write('\x1B[C'); // Right arrow
377+
await tick();
378+
stdin.write('\r'); // Enter
379+
await tick();
380+
rerender(chat);
381+
382+
// Should show rejection message
383+
expect(lastFrame()).toContain('declined');
384+
});
385+
386+
it('handles tool approval acceptance', async () => {
387+
const { ollama, tools } = await import('../utils');
388+
const { streamChat } = ollama;
389+
390+
vi.mocked(streamChat).mockImplementationOnce(async function* () {
391+
await Promise.resolve();
392+
yield {
393+
type: 'tool_calls',
394+
tool_calls: [
395+
{
396+
function: {
397+
name: 'write_file',
398+
arguments: { path: '/test.txt', content: 'hello' },
399+
},
400+
},
401+
],
402+
};
403+
});
404+
405+
// Second call after tool execution
406+
vi.mocked(streamChat).mockImplementationOnce(async function* () {
407+
await Promise.resolve();
408+
yield { type: 'content', content: 'Done' };
409+
});
410+
411+
const mockExecute = vi.fn().mockResolvedValue({
412+
content: 'File written successfully',
413+
});
414+
vi.mocked(tools.executeTool).mockImplementation(mockExecute);
415+
416+
vi.spyOn(tools.TOOLS_REQUIRING_APPROVAL, 'has').mockReturnValue(true);
417+
418+
const chat = <Chat model="gemma4" onCommand={vi.fn()} />;
419+
const { lastFrame, rerender, stdin } = render(chat);
420+
421+
await typeText(rerender, 'write a file', chat);
422+
submitInput('write a file');
423+
rerender(chat);
424+
await waitForStream();
425+
rerender(chat);
426+
427+
// Verify approval prompt is shown
428+
expect(lastFrame()).toContain('Tool requires approval');
429+
430+
// Approve the tool by pressing Enter (yes is default)
431+
stdin.write('\r'); // Enter
432+
await tick();
433+
rerender(chat);
434+
435+
// Should have called executeTool
436+
expect(mockExecute).toHaveBeenCalledWith('write_file', {
437+
path: '/test.txt',
438+
content: 'hello',
439+
});
440+
});
441+
442+
it('handles tool result with error in approval flow', async () => {
443+
const { ollama, tools } = await import('../utils');
444+
const { streamChat } = ollama;
445+
446+
vi.mocked(streamChat).mockImplementationOnce(async function* () {
447+
await Promise.resolve();
448+
yield {
449+
type: 'tool_calls',
450+
tool_calls: [
451+
{
452+
function: {
453+
name: 'write_file',
454+
arguments: { path: '/test.txt', content: 'hello' },
455+
},
456+
},
457+
],
458+
};
459+
});
460+
461+
// Second call after tool execution (with error)
462+
vi.mocked(streamChat).mockImplementationOnce(async function* () {
463+
await Promise.resolve();
464+
yield { type: 'content', content: 'Error handled' };
465+
});
466+
467+
const mockExecute = vi.fn().mockResolvedValue({
468+
content: '',
469+
error: 'Permission denied',
470+
});
471+
vi.mocked(tools.executeTool).mockImplementation(mockExecute);
472+
473+
vi.spyOn(tools.TOOLS_REQUIRING_APPROVAL, 'has').mockReturnValue(true);
474+
475+
const chat = <Chat model="gemma4" onCommand={vi.fn()} />;
476+
const { rerender, stdin } = render(chat);
477+
478+
await typeText(rerender, 'write a file', chat);
479+
submitInput('write a file');
480+
rerender(chat);
481+
await waitForStream();
482+
rerender(chat);
483+
484+
// Approve the tool by pressing Enter
485+
stdin.write('\r');
486+
await tick();
487+
rerender(chat);
488+
489+
// Should have called executeTool
490+
expect(mockExecute).toHaveBeenCalled();
491+
});
492+
});
493+
186494
describe('Chat with error', () => {
187495
beforeEach(() => {
188496
mockState.clear();
@@ -193,7 +501,7 @@ describe('Chat with error', () => {
193501
const { streamChat } = ollama;
194502
vi.mocked(streamChat).mockImplementationOnce(async function* () {
195503
await Promise.resolve();
196-
yield '';
504+
yield { type: 'content', content: '' };
197505
throw new Error('Connection failed');
198506
});
199507

@@ -213,7 +521,7 @@ describe('Chat with error', () => {
213521
const { streamChat } = ollama;
214522
vi.mocked(streamChat).mockImplementationOnce(async function* () {
215523
await Promise.resolve();
216-
yield '';
524+
yield { type: 'content', content: '' };
217525
// eslint-disable-next-line @typescript-eslint/only-throw-error
218526
throw { toString: () => 'Custom error' };
219527
});

0 commit comments

Comments
 (0)