Implement stream cancellation in OllamaClient
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.
This commit is contained in:
+140
-6
@@ -423,15 +423,149 @@ describe('OllamaClient', () => {
|
||||
});
|
||||
|
||||
describe('cancelStream', () => {
|
||||
it('should abort the current request', () => {
|
||||
// Removed cancelStream call as we now use local controllers
|
||||
// expect(client['abortController']).toBeNull(); // Removed as abortController no longer exists
|
||||
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', () => {
|
||||
// Removed cancelStream test as we now use local controllers
|
||||
// expect(() => client.cancelStream()).not.toThrow(); // Removed as cancelStream no longer exists
|
||||
// expect(client['abortController']).toBeNull(); // Removed as abortController no longer exists
|
||||
// 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();
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
Reference in New Issue
Block a user