notebook-ai-extension/src/__tests__/opencode_client.spec.ts

import { ServerConnection } from '@jupyterlab/services';

import {
  callOpenCodeEdit,
  callOpenCodeProviders,
  callOpenCodeReplyPermission,
  callOpenCodeReplyQuestion,
  subscribeOpenCodeEvents
} from '../api/opencode_client';
import type { OpenCodeRequest } from '../types';

jest.mock('@jupyterlab/services', () => {
  const actual = jest.requireActual('@jupyterlab/services');
  return {
    ...actual,
    ServerConnection: {
      ...actual.ServerConnection,
      makeRequest: jest.fn(),
    },
  };
});

const mockedMakeRequest = ServerConnection.makeRequest as jest.MockedFunction<
  typeof ServerConnection.makeRequest
>;

function mockServerSettings(): ServerConnection.ISettings {
  return {
    baseUrl: 'http://localhost:8888',
    wsUrl: 'ws://localhost:8888',
    token: 'test',
    init: { headers: { 'X-Test': '1' } },
    fetch: {} as unknown as typeof fetch,
    RequestClass: {} as unknown as typeof Request,
    Headers: {} as unknown as typeof Headers,
    appendToken: false,
    pageUrl: '',
    settings: {} as unknown as ServerConnection.ISettings,
    displayName: 'JupyterLab',
    handleError: () => undefined,
    requestHeaders: {},
    wsHeaders: {},
    userSettings: {},
  } as unknown as ServerConnection.ISettings;
}

function mockResponse(overrides: {
  ok: boolean;
  status: number;
  statusText: string;
  text: string;
}): Response {
  return {
    ok: overrides.ok,
    status: overrides.status,
    statusText: overrides.statusText,
    text: async () => overrides.text,
    headers: {} as unknown as Headers,
    json: async () => JSON.parse(overrides.text),
    url: '',
    redirected: false,
    type: 'basic',
    body: null,
    bodyUsed: false,
    arrayBuffer: async () => new ArrayBuffer(0),
    blob: async () => new Blob(),
    formData: async () => new FormData(),
    clone: () => mockResponse(overrides),
  } as unknown as Response;
}

const sampleRequest: OpenCodeRequest = {
  prompt: 'add type hints',
  context: {
    notebookPath: 'foo.ipynb',
    cellId: 'cell-1',
    language: 'python',
    cellIndex: 0,
    totalCells: 1,
    source: 'def foo(x): return x',
    previousCode: null,
    error: null,
  },
};

describe('callOpenCodeEdit', () => {
  beforeEach(() => {
    mockedMakeRequest.mockReset();
  });

  it('POSTs to /opencode-bridge/edit and returns async {ok, sessionId, notebookPath}', async () => {
    const respBody = {
      ok: true,
      sessionId: 'sid',
      notebookPath: 'foo.ipynb',
    };
    mockedMakeRequest.mockResolvedValue(
      mockResponse({ ok: true, status: 200, statusText: 'OK', text: JSON.stringify(respBody) })
    );

    const settings = mockServerSettings();
    const resp = await callOpenCodeEdit(sampleRequest, settings);

    expect(mockedMakeRequest).toHaveBeenCalledWith(
      'http://localhost:8888/opencode-bridge/edit',
      expect.objectContaining({
        method: 'POST',
        body: JSON.stringify(sampleRequest),
        headers: { 'Content-Type': 'application/json' },
      }),
      settings
    );
    expect(resp.ok).toBe(true);
    if (resp.ok) {
      // Async response has no `markdown` field — the LLM reply arrives
      // via the SSE stream, not embedded in this response.
      expect(resp.sessionId).toBe('sid');
      expect((resp as any).markdown).toBeUndefined();
    }
  });

  it('returns failure response when server returns ok: false', async () => {
    mockedMakeRequest.mockResolvedValue(
      mockResponse({
        ok: false,
        status: 502,
        statusText: 'Bad Gateway',
        text: JSON.stringify({ ok: false, error: 'opencode down' }),
      })
    );

    const resp = await callOpenCodeEdit(sampleRequest, mockServerSettings());
    expect(resp.ok).toBe(false);
    if (!resp.ok) {
      expect(resp.error).toBe('opencode down');
    }
  });

  it('throws on network error', async () => {
    mockedMakeRequest.mockRejectedValue(new Error('ECONNREFUSED'));
    await expect(
      callOpenCodeEdit(sampleRequest, mockServerSettings())
    ).rejects.toThrow(/network error/);
  });

  it('throws on empty response body', async () => {
    mockedMakeRequest.mockResolvedValue(
      mockResponse({ ok: true, status: 200, statusText: 'OK', text: '' })
    );
    await expect(
      callOpenCodeEdit(sampleRequest, mockServerSettings())
    ).rejects.toThrow(/empty response/);
  });
});

describe('callOpenCodeProviders', () => {
  beforeEach(() => {
    mockedMakeRequest.mockReset();
  });

  it('GETs /opencode-bridge/providers and returns parsed JSON', async () => {
    const respBody = {
      providers: [
        { id: 'anthropic', models: [{ id: 'claude-sonnet-4-20250514' }] },
        { id: 'openai', models: [{ id: 'gpt-4' }, { id: 'gpt-3.5-turbo' }] }
      ]
    };
    mockedMakeRequest.mockResolvedValue({
      ok: true,
      status: 200,
      statusText: 'OK',
      text: async () => JSON.stringify(respBody),
      json: async () => respBody,
      headers: {} as any,
      url: '',
      redirected: false,
      type: 'basic',
      body: null,
      bodyUsed: false,
      arrayBuffer: async () => new ArrayBuffer(0),
      blob: async () => new Blob(),
      formData: async () => new FormData(),
      clone: function () { return this; }
    } as any);

    const settings = mockServerSettings();
    const resp = await callOpenCodeProviders(settings);

    expect(mockedMakeRequest).toHaveBeenCalledWith(
      'http://localhost:8888/opencode-bridge/providers',
      expect.objectContaining({ method: 'GET' }),
      settings
    );
    expect(resp.providers).toHaveLength(2);
    expect(resp.providers[0].id).toBe('anthropic');
    expect(resp.providers[1].models).toHaveLength(2);
  });

  it('throws on non-2xx response', async () => {
    mockedMakeRequest.mockResolvedValue({
      ok: false,
      status: 502,
      statusText: 'Bad Gateway',
      text: async () => '',
      json: async () => null,
      headers: {} as any,
      url: '',
      redirected: false,
      type: 'basic',
      body: null,
      bodyUsed: false,
      arrayBuffer: async () => new ArrayBuffer(0),
      blob: async () => new Blob(),
      formData: async () => new FormData(),
      clone: function () { return this; }
    } as any);

    await expect(
      callOpenCodeProviders(mockServerSettings())
    ).rejects.toThrow(/502/);
  });

  it('throws on network error', async () => {
    mockedMakeRequest.mockRejectedValue(new Error('ECONNREFUSED'));
    await expect(
      callOpenCodeProviders(mockServerSettings())
    ).rejects.toThrow(/network error/);
  });
});

describe('callOpenCodeReplyPermission', () => {
  beforeEach(() => {
    mockedMakeRequest.mockReset();
  });

  it('POSTs to /opencode-bridge/permissions/:permId with {response}', async () => {
    mockedMakeRequest.mockResolvedValue(
      mockResponse({
        ok: true,
        status: 200,
        statusText: 'OK',
        text: JSON.stringify({ ok: true, permissionId: 'perm-1', response: 'once' })
      })
    );

    const r = await callOpenCodeReplyPermission(
      's1',
      'perm-1',
      'once',
      mockServerSettings()
    );
    expect(r).toEqual({ ok: true, permissionId: 'perm-1', response: 'once' });
    expect(mockedMakeRequest).toHaveBeenCalledWith(
      'http://localhost:8888/opencode-bridge/permissions/perm-1?session=s1',
      expect.objectContaining({
        method: 'POST',
        body: JSON.stringify({ response: 'once' })
      }),
      expect.anything()
    );
  });

  it('throws on non-2xx', async () => {
    mockedMakeRequest.mockResolvedValue(
      mockResponse({
        ok: false,
        status: 400,
        statusText: 'Bad Request',
        text: 'bad response'
      })
    );
    await expect(
      callOpenCodeReplyPermission('s1', 'p', 'reject', mockServerSettings())
    ).rejects.toThrow(/400/);
  });
});

describe('callOpenCodeReplyQuestion', () => {
  beforeEach(() => {
    mockedMakeRequest.mockReset();
  });

  it('POSTs to /opencode-bridge/questions/:qId/reply with {answer}', async () => {
    mockedMakeRequest.mockResolvedValue(
      mockResponse({
        ok: true,
        status: 200,
        statusText: 'OK',
        text: JSON.stringify({ ok: true, questionId: 'q-1' })
      })
    );

    const r = await callOpenCodeReplyQuestion(
      's1',
      'q-1',
      'use 3.12',
      mockServerSettings()
    );
    expect(r).toEqual({ ok: true, questionId: 'q-1' });
    expect(mockedMakeRequest).toHaveBeenCalledWith(
      'http://localhost:8888/opencode-bridge/questions/q-1/reply?session=s1',
      expect.objectContaining({
        method: 'POST',
        body: JSON.stringify({ answer: 'use 3.12' })
      }),
      expect.anything()
    );
  });
});

/**
 * subscribeOpenCodeEvents uses the global `fetch` (not ServerConnection),
 * so we mock fetch instead of ServerConnection.makeRequest.
 */
describe('subscribeOpenCodeEvents', () => {
  let originalFetch: typeof fetch;
  let mockReader: { read: jest.Mock };
  let capturedUrl: string | undefined;
  let capturedInit: RequestInit | undefined;

  beforeEach(() => {
    originalFetch = globalThis.fetch;
    mockReader = {
      read: jest
        .fn()
        // First call returns a chunk; second returns done.
        .mockResolvedValueOnce({
          value: new TextEncoder().encode(
            'data: {"type":"session.next.text.delta","properties":{"sessionID":"sid","delta":"hello"}}\n\n' +
            'data: {"type":"session.idle","properties":{"sessionID":"sid"}}\n\n'
          ),
          done: false
        })
        .mockResolvedValueOnce({ value: undefined, done: true })
    };
    globalThis.fetch = jest.fn(async (url: any, init?: RequestInit) => {
      capturedUrl = String(url);
      capturedInit = init;
      return {
        ok: true,
        status: 200,
        statusText: 'OK',
        body: {
          getReader: () => mockReader
        }
      } as unknown as Response;
    }) as unknown as typeof fetch;
  });

  afterEach(() => {
    globalThis.fetch = originalFetch;
  });

  it('GETs /opencode-bridge/events?session=<sid> and dispatches parsed events', async () => {
    const onEvent = jest.fn();
    const settings = mockServerSettings();
    const sub = subscribeOpenCodeEvents('sid', settings, { onEvent });

    // Let the async fetch + reader resolve.
    await new Promise(r => setTimeout(r, 10));
    sub.close();

    expect(capturedUrl).toBe(
      'http://localhost:8888/opencode-bridge/events?session=sid'
    );
    expect(capturedInit?.method).toBe('GET');
    expect((capturedInit?.headers as any).Accept).toBe('text/event-stream');
    expect((capturedInit?.headers as any)['X-XSRFToken']).toBe('test');

    expect(onEvent).toHaveBeenCalledTimes(2);
    expect(onEvent.mock.calls[0][0].type).toBe('session.next.text.delta');
    expect(onEvent.mock.calls[0][0].properties.delta).toBe('hello');
    expect(onEvent.mock.calls[1][0].type).toBe('session.idle');
  });

  it('close() is idempotent and aborts the in-flight fetch', async () => {
    const onEvent = jest.fn();
    const sub = subscribeOpenCodeEvents('sid', mockServerSettings(), { onEvent });
    // Let the async fetch actually run and capture the signal.
    await new Promise(r => setTimeout(r, 5));
    sub.close();
    sub.close(); // second call must not throw
    const signal = capturedInit?.signal as AbortSignal | undefined;
    expect(signal).toBeDefined();
    expect(signal?.aborted).toBe(true);
  });
});