obsidian_ollama/tests/ollama-client.test.ts

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');
    }, 15000);

    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(
        'Model "llama3" not found. Run \`ollama pull llama3\` first.'
      );
    });

    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);
    }, 15000);

    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');
    }, 15000);

    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('Model "llama3" not found. Run \`ollama pull llama3\` first.');
    });
  });

  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('cancelStream race condition', () => {
    it('should not clear new stream controller when old stream finally block executes', async () => {
      // Race condition scenario:
      // 1. Stream A starts → controllerA assigned to currentStreamController
      // 2. Stream B starts → controllerB replaces controllerA
      // 3. Stream A's fetch rejects (abort) → finally block runs
      // 4. Guard (this.currentStreamController === controllerA) is false
      // 5. controllerB survives

      // First fetch call: pending promise with manually triggerable abort
      let fireFirstAbort: (() => void) | undefined;
      mockFetch.mockImplementationOnce((_url, options?: any) => {
        const signal = options?.signal;
        return 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'));
            },
            { once: true }
          );
          // Capture abort trigger for manual control
          fireFirstAbort = () => {
            signal?.dispatchEvent(new CustomEvent('abort'));
          };
        });
      });

      // Second fetch call: also pending, so controller stays assigned
      // We don't need it to complete - just need to verify controller survives abort
      mockFetch.mockImplementationOnce((_url, options?: any) => {
        const signal = options?.signal;
        return 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'));
            },
            { once: true }
          );
        });
      });

      // Start first stream and consume it (triggers fetch + controller assignment)
      const stream1 = await client.streamChat(mockMessages, mockTools);
      const consumeFirst = async () => {
        for await (const _ of stream1) {
          /* consume */
        }
      };
      const firstPromise = consumeFirst();

      // Wait for fetch to be triggered (controller should be set)
      await new Promise((resolve) => setTimeout(resolve, 10));
      const firstController = client['currentStreamController'];
      expect(firstController).not.toBeNull();

      // Start second stream and consume it (triggers fetch + replaces controller)
      const stream2 = await client.streamChat(mockMessages, mockTools);
      const consumeSecond = async () => {
        for await (const _ of stream2) {
          /* consume */
        }
      };
      const secondPromise = consumeSecond();

      // Wait for second stream fetch to trigger and assign its controller
      await new Promise((resolve) => setTimeout(resolve, 10));

      const secondController = client['currentStreamController'];
      expect(secondController).not.toBeNull();
      expect(secondController).not.toBe(firstController);

      // Defer abort to next microtask so Jest associates rejection with expectation
      // Then immediately await the rejection
      queueMicrotask(() => fireFirstAbort?.());
      await expect(firstPromise).rejects.toThrow('The operation was aborted.');

      // Second stream's controller should still be intact
      // (not cleared by first stream's finally block due to guard)
      expect(client['currentStreamController']).toBe(secondController);

      // Clean up second stream so it doesn't leak
      secondController?.abort();
      await secondPromise.catch(() => {});
    });
  });

  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();
    });
  });
});