fegger
9823761e03
Add `cancelStream` method to allow aborting active requests. Store the current `AbortController` on the client instance and reset it when the stream completes or is cancelled. Update `ChatView` to call `cancelStream` on close. Add comprehensive tests for stream cancellation scenarios, including aborting active requests, handling cancellation when no stream is active, clearing the controller after normal completion, and allowing new streams after cancellation.
663 lines
20 KiB
TypeScript
Executable File
663 lines
20 KiB
TypeScript
Executable File
import { OllamaClient } from '../src/ollama-client';
|
|
import { OllamaMessage, OllamaTool } from '../src/types';
|
|
|
|
describe('OllamaClient', () => {
|
|
let client: OllamaClient;
|
|
let mockFetch: jest.Mock;
|
|
|
|
function createMockReader(data: string) {
|
|
const encoder = new TextEncoder();
|
|
const encoded = encoder.encode(data);
|
|
let called = false;
|
|
return {
|
|
read: () => {
|
|
if (!called) {
|
|
called = true;
|
|
return Promise.resolve({ done: false, value: encoded });
|
|
}
|
|
return Promise.resolve({ done: true, value: new Uint8Array(0) });
|
|
},
|
|
releaseLock: jest.fn(),
|
|
};
|
|
}
|
|
|
|
const mockMessages: OllamaMessage[] = [
|
|
{ role: 'system', content: 'You are helpful.' },
|
|
{ role: 'user', content: 'Hello' },
|
|
];
|
|
|
|
const mockTools: OllamaTool[] = [
|
|
{
|
|
type: 'function',
|
|
function: {
|
|
name: 'test_tool',
|
|
description: 'A test tool',
|
|
parameters: {
|
|
type: 'object',
|
|
properties: { input: { type: 'string' } },
|
|
required: ['input'],
|
|
},
|
|
},
|
|
},
|
|
];
|
|
|
|
beforeEach(() => {
|
|
mockFetch = jest.fn();
|
|
client = new OllamaClient('http://localhost:11434', 'llama3', mockFetch);
|
|
});
|
|
|
|
afterEach(() => {
|
|
jest.clearAllMocks();
|
|
// Removed cancelStream call as we now use local controllers
|
|
});
|
|
|
|
describe('chat (non-streaming)', () => {
|
|
it('should send a non-streaming request and return the response', async () => {
|
|
const mockResponse = {
|
|
ok: true,
|
|
json: () => Promise.resolve({ message: { content: 'Hello back!' } }),
|
|
};
|
|
mockFetch.mockResolvedValue(mockResponse);
|
|
|
|
const result = await client.chat(mockMessages, mockTools);
|
|
|
|
expect(result.content).toBe('Hello back!');
|
|
expect(mockFetch).toHaveBeenCalledWith(
|
|
'http://localhost:11434/api/chat',
|
|
expect.objectContaining({
|
|
method: 'POST',
|
|
body: JSON.stringify({
|
|
model: 'llama3',
|
|
messages: mockMessages,
|
|
tools: mockTools,
|
|
stream: false,
|
|
}),
|
|
})
|
|
);
|
|
});
|
|
|
|
it('should throw on non-OK response', async () => {
|
|
mockFetch.mockResolvedValue({ ok: false, status: 500 });
|
|
|
|
await expect(client.chat(mockMessages, mockTools)).rejects.toThrow('Ollama API error: 500');
|
|
});
|
|
|
|
it('should handle missing message content gracefully', async () => {
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: () => Promise.resolve({}),
|
|
});
|
|
|
|
const result = await client.chat(mockMessages, mockTools);
|
|
expect(result.content).toBe('');
|
|
expect(result.tool_calls).toEqual([]);
|
|
});
|
|
|
|
it('should include abort signal in fetch options', async () => {
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: () => Promise.resolve({ message: { content: 'ok' } }),
|
|
});
|
|
|
|
await client.chat(mockMessages, mockTools);
|
|
|
|
const fetchOptions = mockFetch.mock.calls[0][1];
|
|
expect(fetchOptions.signal).toBeInstanceOf(AbortSignal);
|
|
});
|
|
|
|
it('should forward tool_calls from response when present', async () => {
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
json: () =>
|
|
Promise.resolve({
|
|
message: {
|
|
content: 'result',
|
|
tool_calls: [{ function: { name: 'create_file', arguments: '{}' } }],
|
|
},
|
|
}),
|
|
});
|
|
|
|
const result = await client.chat(mockMessages, mockTools);
|
|
expect(result.tool_calls).toEqual([{ function: { name: 'create_file', arguments: '{}' } }]);
|
|
});
|
|
});
|
|
|
|
describe('streamChat', () => {
|
|
it('should send a streaming request and yield chunks', async () => {
|
|
const streamData = [
|
|
JSON.stringify({ message: { content: 'He' } }),
|
|
JSON.stringify({ message: { content: 'llo' } }),
|
|
JSON.stringify({ message: { content: '!' } }),
|
|
'',
|
|
].join('\n');
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
const chunks: string[] = [];
|
|
|
|
for await (const chunk of stream) {
|
|
chunks.push(chunk.content);
|
|
}
|
|
|
|
expect(chunks).toEqual(['He', 'llo', '!']);
|
|
expect(mockReader.releaseLock).toHaveBeenCalled();
|
|
});
|
|
|
|
it('should skip malformed JSON chunks and log a warning', async () => {
|
|
const streamData = [
|
|
JSON.stringify({ message: { content: 'valid' } }),
|
|
'this is not json',
|
|
JSON.stringify({ message: { content: 'also valid' } }),
|
|
'',
|
|
].join('\n');
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
const chunks: string[] = [];
|
|
|
|
for await (const chunk of stream) {
|
|
chunks.push(chunk.content);
|
|
}
|
|
|
|
expect(chunks).toEqual(['valid', 'also valid']);
|
|
expect(consoleWarnSpy).toHaveBeenCalledWith(
|
|
expect.stringContaining(
|
|
'[ollama-client] WARN: Skipped malformed chunk: this is not json... - Unexpected token \'h\', "this is not json" is not valid JSON'
|
|
)
|
|
);
|
|
consoleWarnSpy.mockRestore();
|
|
});
|
|
|
|
it('should throw when too many chunks are malformed', async () => {
|
|
const streamData = Array(51).fill('invalid json').join('\n') + '\n';
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
|
|
await expect(
|
|
(async () => {
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
})()
|
|
).rejects.toThrow(/malformed/);
|
|
});
|
|
|
|
it('should throw on non-OK response', async () => {
|
|
mockFetch.mockResolvedValue({ ok: false, status: 404 });
|
|
|
|
await expect(client.streamChatAsPromise(mockMessages, mockTools)).rejects.toThrow(
|
|
'Ollama API error: 404'
|
|
);
|
|
});
|
|
|
|
it('should throw when response has no body', async () => {
|
|
mockFetch.mockResolvedValue({ ok: true, body: undefined });
|
|
|
|
await expect(client.streamChatAsPromise(mockMessages, mockTools)).rejects.toThrow(
|
|
'No response body'
|
|
);
|
|
});
|
|
|
|
it('should throw on invalid content type', async () => {
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: {
|
|
getReader: () => ({
|
|
read: () => Promise.resolve({ done: true, value: new Uint8Array(0) }),
|
|
}),
|
|
},
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'text/html' : null),
|
|
},
|
|
});
|
|
|
|
await expect(client.streamChatAsPromise(mockMessages, mockTools)).rejects.toThrow(
|
|
'Invalid response format'
|
|
);
|
|
});
|
|
|
|
it('should propagate Ollama error messages from the stream', async () => {
|
|
const streamData = JSON.stringify({ error: 'model not found' }) + '\n';
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
|
|
await expect(
|
|
(async () => {
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
})()
|
|
).rejects.toThrow('Ollama error: model not found');
|
|
});
|
|
|
|
it('should yield tool_calls when present in streamed response', async () => {
|
|
const streamData = [
|
|
JSON.stringify({
|
|
message: {
|
|
content: '',
|
|
tool_calls: [{ function: { name: 'create_file', arguments: '{"path":"a.md"}' } }],
|
|
},
|
|
}),
|
|
'',
|
|
].join('\n');
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
let lastChunk: any;
|
|
|
|
for await (const chunk of stream) {
|
|
lastChunk = chunk;
|
|
}
|
|
|
|
expect(lastChunk.tool_calls).toEqual([
|
|
{ function: { name: 'create_file', arguments: '{"path":"a.md"}' } },
|
|
]);
|
|
});
|
|
|
|
it('should default tool_calls to empty array when not present', async () => {
|
|
const streamData = JSON.stringify({ message: { content: 'hello' } }) + '\n';
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
let lastChunk: any;
|
|
|
|
for await (const chunk of stream) {
|
|
lastChunk = chunk;
|
|
}
|
|
|
|
expect(lastChunk.tool_calls).toEqual([]);
|
|
});
|
|
|
|
it('should send correct request body with stream:true', async () => {
|
|
const streamData = JSON.stringify({ message: { content: 'ok' } }) + '\n';
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
|
|
expect(mockFetch).toHaveBeenCalledWith(
|
|
'http://localhost:11434/api/chat',
|
|
expect.objectContaining({
|
|
method: 'POST',
|
|
body: JSON.stringify({
|
|
model: 'llama3',
|
|
messages: mockMessages,
|
|
tools: mockTools,
|
|
stream: true,
|
|
}),
|
|
signal: expect.any(AbortSignal),
|
|
})
|
|
);
|
|
});
|
|
});
|
|
|
|
describe('streamChat with retry logic', () => {
|
|
it('should retry on 5xx errors and eventually succeed', async () => {
|
|
let callCount = 0;
|
|
mockFetch.mockImplementation(async () => {
|
|
callCount++;
|
|
if (callCount === 1) {
|
|
return { ok: false, status: 500 };
|
|
}
|
|
if (callCount === 2) {
|
|
return { ok: false, status: 502 };
|
|
}
|
|
return {
|
|
ok: true,
|
|
body: {
|
|
getReader: () => ({
|
|
read: () => Promise.resolve({ done: true, value: new Uint8Array(0) }),
|
|
releaseLock: () => {},
|
|
}),
|
|
},
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
};
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
const chunks: string[] = [];
|
|
|
|
for await (const _ of stream) {
|
|
chunks.push('chunk');
|
|
}
|
|
|
|
expect(callCount).toBe(3);
|
|
expect(chunks.length).toBe(0);
|
|
});
|
|
|
|
it('should give up after maxRetries attempts', async () => {
|
|
mockFetch.mockResolvedValue({ ok: false, status: 500 });
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
|
|
await expect(
|
|
(async () => {
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
})()
|
|
).rejects.toThrow('Ollama API error: 500');
|
|
});
|
|
|
|
it('should not retry on 4xx errors', async () => {
|
|
mockFetch.mockResolvedValue({ ok: false, status: 404 });
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
|
|
await expect(
|
|
(async () => {
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
})()
|
|
).rejects.toThrow('Ollama API error: 404');
|
|
});
|
|
});
|
|
|
|
describe('cancelStream', () => {
|
|
it('should abort an active streaming request when cancelled before fetch resolves', async () => {
|
|
let capturedSignal: AbortSignal | undefined;
|
|
|
|
mockFetch.mockImplementation((_url, options?: any) => {
|
|
capturedSignal = options?.signal;
|
|
|
|
return Promise.race([
|
|
// Simulate slow network response
|
|
new Promise(() => {
|
|
// Never resolves on its own - relies on abort
|
|
}),
|
|
// Reject when signal is aborted (like real fetch does)
|
|
new Promise<never>((_, reject) => {
|
|
if (capturedSignal?.aborted) {
|
|
reject(new DOMException('The operation was aborted.', 'AbortError'));
|
|
return;
|
|
}
|
|
capturedSignal?.addEventListener('abort', () => {
|
|
reject(new DOMException('The operation was aborted.', 'AbortError'));
|
|
});
|
|
}),
|
|
]);
|
|
});
|
|
|
|
// Consume the stream in an async function so we can await the rejection
|
|
const consumeStream = async () => {
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
};
|
|
|
|
// Start consuming - this triggers the fetch
|
|
const consumePromise = consumeStream();
|
|
|
|
// Wait a tick for fetch to be invoked
|
|
await new Promise((resolve) => setTimeout(resolve, 10));
|
|
|
|
// Verify controller is set after fetch starts
|
|
expect(client['currentStreamController']).not.toBeNull();
|
|
|
|
// Cancel the stream - this aborts the controller and triggers fetch rejection
|
|
client.cancelStream();
|
|
|
|
// Verify signal was aborted
|
|
expect(capturedSignal?.aborted).toBe(true);
|
|
|
|
// Verify controller was cleared by cancelStream
|
|
expect(client['currentStreamController']).toBeNull();
|
|
|
|
// The stream consumption must reject with an abort error
|
|
await expect(consumePromise).rejects.toThrow('The operation was aborted.');
|
|
});
|
|
|
|
it('should handle cancel when no active stream', () => {
|
|
// Calling cancelStream with no active stream must not throw
|
|
expect(() => client.cancelStream()).not.toThrow();
|
|
expect(client['currentStreamController']).toBeNull();
|
|
});
|
|
|
|
it('should clear the controller after stream completes normally', async () => {
|
|
const streamData = JSON.stringify({ message: { content: 'done' } }) + '\n';
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
|
|
// Controller should be cleared after normal completion
|
|
expect(client['currentStreamController']).toBeNull();
|
|
});
|
|
|
|
it('should allow a new stream after cancelling a previous one', async () => {
|
|
// First fetch: pending and abortable
|
|
mockFetch.mockImplementationOnce((_url, options?: any) => {
|
|
const signal = options?.signal;
|
|
|
|
return Promise.race([
|
|
new Promise(() => {
|
|
// Never resolves on its own
|
|
}),
|
|
new Promise<never>((_, reject) => {
|
|
if (signal?.aborted) {
|
|
reject(new DOMException('The operation was aborted.', 'AbortError'));
|
|
return;
|
|
}
|
|
signal?.addEventListener('abort', () => {
|
|
reject(new DOMException('The operation was aborted.', 'AbortError'));
|
|
});
|
|
}),
|
|
]);
|
|
});
|
|
|
|
// Second fetch: resolves immediately with valid stream
|
|
mockFetch.mockResolvedValueOnce({
|
|
ok: true,
|
|
body: {
|
|
getReader: () => ({
|
|
read: () => Promise.resolve({ done: true, value: new Uint8Array(0) }),
|
|
releaseLock: () => {},
|
|
}),
|
|
},
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
// Start first stream and consume it
|
|
const consumeFirst = async () => {
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
for await (const _ of stream) {
|
|
/* consume */
|
|
}
|
|
};
|
|
|
|
const firstStreamPromise = consumeFirst();
|
|
await new Promise((resolve) => setTimeout(resolve, 10));
|
|
|
|
// Cancel first stream
|
|
client.cancelStream();
|
|
await expect(firstStreamPromise).rejects.toThrow('The operation was aborted.');
|
|
|
|
// Controller is cleared, can start a new stream
|
|
expect(client['currentStreamController']).toBeNull();
|
|
|
|
// Start second stream - should succeed independently
|
|
const stream2 = await client.streamChat(mockMessages, mockTools);
|
|
for await (const _ of stream2) {
|
|
/* consume */
|
|
}
|
|
|
|
// Second stream completes and clears controller
|
|
expect(client['currentStreamController']).toBeNull();
|
|
});
|
|
});
|
|
|
|
describe('streamChat final buffer parsing', () => {
|
|
it('should parse final buffer content when stream ends with partial line', async () => {
|
|
const streamData = [
|
|
JSON.stringify({ message: { content: 'First' } }),
|
|
JSON.stringify({ message: { content: 'Second' } }),
|
|
JSON.stringify({ message: { content: 'Third' } }),
|
|
'', // Final line should be empty to signal end
|
|
].join('\n');
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
const chunks: string[] = [];
|
|
|
|
for await (const chunk of stream) {
|
|
chunks.push(chunk.content);
|
|
}
|
|
|
|
expect(chunks).toEqual(['First', 'Second', 'Third']);
|
|
expect(mockReader.releaseLock).toHaveBeenCalled();
|
|
});
|
|
|
|
it('should handle malformed final buffer content gracefully', async () => {
|
|
const streamData = [
|
|
JSON.stringify({ message: { content: 'Valid' } }),
|
|
'malformed json',
|
|
'', // Final line should be empty to signal end
|
|
].join('\n');
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
const chunks: string[] = [];
|
|
|
|
for await (const chunk of stream) {
|
|
chunks.push(chunk.content);
|
|
}
|
|
|
|
expect(chunks).toEqual(['Valid']);
|
|
expect(consoleWarnSpy).toHaveBeenCalledWith(
|
|
expect.stringContaining('Skipped malformed chunk')
|
|
);
|
|
consoleWarnSpy.mockRestore();
|
|
});
|
|
|
|
it('should parse final buffer content even when it contains message data', async () => {
|
|
const streamData = [
|
|
JSON.stringify({ message: { content: 'First' } }),
|
|
'', // Final line should be empty to signal end
|
|
JSON.stringify({ message: { content: 'Final' } }),
|
|
].join('\n');
|
|
|
|
const mockReader = createMockReader(streamData);
|
|
|
|
mockFetch.mockResolvedValue({
|
|
ok: true,
|
|
body: { getReader: () => mockReader },
|
|
headers: {
|
|
get: (name: string) => (name === 'content-type' ? 'application/x-ndjson' : null),
|
|
},
|
|
});
|
|
|
|
const stream = await client.streamChat(mockMessages, mockTools);
|
|
const chunks: string[] = [];
|
|
|
|
for await (const chunk of stream) {
|
|
chunks.push(chunk.content);
|
|
}
|
|
|
|
expect(chunks).toEqual(['First', 'Final']);
|
|
expect(mockReader.releaseLock).toHaveBeenCalled();
|
|
});
|
|
});
|
|
});
|