@oceanum/oceanumlab 4.6.0

package/src/__tests__/chatRouter.spec.ts

@@ -0,0 +1,315 @@
+import { ChatRouter } from '../chatRouter';
+import type { INotebookSnapshot } from '../notebookContext';
+
+const settings = {
+  get: () => ({ composite: 'a-token' })
+} as unknown as ConstructorParameters<typeof ChatRouter>[0];
+
+const snapshot = (
+  cells: INotebookSnapshot['cells'],
+  selected: INotebookSnapshot['selected'] = null
+): INotebookSnapshot => ({ cells, selected });
+
+let sent: Record<string, unknown> | undefined;
+let sentHeaders: Record<string, string> | undefined;
+
+interface IRequestInit {
+  body: string;
+  headers: Record<string, string>;
+  signal?: AbortSignal;
+}
+
+/** Answer every request with what `response` builds from it. */
+function answering(response: (init: IRequestInit) => unknown): void {
+  (globalThis as any).fetch = jest.fn(
+    async (_url: string, init: IRequestInit) => {
+      sent = JSON.parse(init.body);
+      sentHeaders = init.headers;
+      return response(init);
+    }
+  );
+}
+
+beforeEach(() => {
+  sent = undefined;
+  sentHeaders = undefined;
+  answering(() => ({
+    ok: true,
+    status: 200,
+    headers: { get: () => 'application/json' },
+    json: async (): Promise<unknown> => ({ message: 'ok', blocks: [] })
+  }));
+});
+
+const frame = (event: string, data: unknown): string =>
+  `event: ${event}\ndata: ${JSON.stringify(data)}\n\n`;
+
+/**
+ * A 200 whose body streams `frames` and ends -- or, with `hang`, then stays
+ * open until `signal` aborts, failing the read in flight as fetch does.
+ */
+function streamed(frames: string[], signal?: AbortSignal, hang = false) {
+  const encoder = new TextEncoder();
+  const chunks = frames.map(f => encoder.encode(f));
+  let index = 0;
+  return {
+    ok: true,
+    status: 200,
+    headers: { get: () => 'text/event-stream' },
+    json: async (): Promise<unknown> => {
+      throw new Error('a stream is not JSON');
+    },
+    body: {
+      getReader: () => ({
+        read: (): Promise<{ done: boolean; value?: Uint8Array }> => {
+          if (index < chunks.length) {
+            return Promise.resolve({ done: false, value: chunks[index++] });
+          }
+          if (!hang) {
+            return Promise.resolve({ done: true, value: undefined });
+          }
+          return new Promise((_ok, fail) => {
+            signal?.addEventListener('abort', () =>
+              fail(new DOMException('Aborted', 'AbortError'))
+            );
+          });
+        },
+        releaseLock: (): void => undefined
+      })
+    }
+  };
+}
+
+describe('ChatRouter: the notebook sent with a request', () => {
+  it('is the conversation notebook, code and markdown', async () => {
+    const router = new ChatRouter(settings, async () =>
+      snapshot([
+        { kind: 'code', source: 'x = 1' },
+        { kind: 'markdown', source: '## Notes' },
+        { kind: 'code', source: 'y = x' }
+      ])
+    );
+
+    await router.route('what does this do?');
+
+    expect(sent?.notebookCells).toEqual([
+      'x = 1',
+      '# %% [markdown]\n# ## Notes',
+      'y = x'
+    ]);
+  });
+
+  it('carries the selected cell of that notebook, which a code answer may replace', async () => {
+    const router = new ChatRouter(settings, async () =>
+      snapshot([{ kind: 'code', source: 'df.head()' }], {
+        source: 'df.head()',
+        isCode: true
+      })
+    );
+
+    const result = await router.route('fix this');
+
+    expect(sent?.codeContext).toBe('df.head()');
+    expect(result.hasCodeCellSelected).toBe(true);
+  });
+
+  it('asks for the notebook on every request, so it follows the conversation', async () => {
+    let current = snapshot([{ kind: 'code', source: 'a = 1' }]);
+    const router = new ChatRouter(settings, async () => current);
+
+    await router.route('one');
+    expect(sent?.notebookCells).toEqual(['a = 1']);
+
+    current = snapshot([{ kind: 'code', source: 'b = 2' }]);
+    await router.route('two');
+    expect(sent?.notebookCells).toEqual(['b = 2']);
+  });
+
+  it('sends no notebook when there is none', async () => {
+    const router = new ChatRouter(settings, async () => null);
+
+    await router.route('hello');
+
+    expect(sent?.notebookCells).toBeUndefined();
+    expect(sent?.codeContext).toBeUndefined();
+    expect(sent?.context).toBeUndefined();
+  });
+
+  it('still sends the question when the notebook cannot be read', async () => {
+    // Context is best-effort; a notebook that could not be read must not
+    // cost the user the answer.
+    const router = new ChatRouter(settings, async () => {
+      throw new Error('could not read');
+    });
+
+    await router.route('hello');
+
+    expect(sent?.prompt).toBe('hello');
+    expect(sent?.notebookCells).toBeUndefined();
+  });
+});
+
+describe('ChatRouter: the streamed answer', () => {
+  const noNotebook = async (): Promise<null> => null;
+
+  it('asks for the stream only when someone is listening', async () => {
+    // Without a listener the request is byte-identical to what it always was.
+    const router = new ChatRouter(settings, noNotebook);
+
+    await router.route('hello');
+    expect(sentHeaders?.Accept).toBeUndefined();
+
+    await router.route('hello', [], undefined, () => undefined);
+    expect(sentHeaders?.Accept).toBe('text/event-stream');
+  });
+
+  it('reports each phase, and takes the answer from done', async () => {
+    // The events before `done` are for the user to look at. Assembled from
+    // them, the answer would come out short whenever a frame went missing.
+    answering(() =>
+      streamed([
+        frame('status', { phase: 'generating' }),
+        ': keepalive\n\n',
+        frame('status', { phase: 'tool', tool: 'search_catalog' }),
+        frame('message', { message: 'partial' }),
+        frame('done', {
+          message: 'The whole answer.',
+          blocks: [{ type: 'code', content: 'x = 1' }]
+        })
+      ])
+    );
+    const progress: unknown[] = [];
+    const router = new ChatRouter(settings, noNotebook);
+
+    const result = await router.route('find waves', [], undefined, p =>
+      progress.push(p)
+    );
+
+    expect(progress).toEqual([
+      { phase: 'generating' },
+      { phase: 'tool', tool: 'search_catalog' }
+    ]);
+    expect(result.response).toEqual({
+      message: 'The whole answer.',
+      blocks: [{ type: 'code', content: 'x = 1' }]
+    });
+  });
+
+  it('reads JSON when the server answers with JSON anyway', async () => {
+    // An older backend, or a proxy that buffered the stream away.
+    const router = new ChatRouter(settings, noNotebook);
+
+    const result = await router.route('hello', [], undefined, () => undefined);
+
+    expect(result.response).toEqual({ message: 'ok', blocks: [] });
+  });
+
+  it('turns an error event into the error the chat shows', async () => {
+    // The status was already 200, so the failure can only arrive as an event.
+    answering(() =>
+      streamed([
+        frame('status', { phase: 'generating' }),
+        frame('error', {
+          detail: 'Internal error',
+          status_code: 500,
+          code: 'internal_error'
+        })
+      ])
+    );
+    const router = new ChatRouter(settings, noNotebook);
+
+    await expect(
+      router.route('hi', [], undefined, () => undefined)
+    ).rejects.toMatchObject({
+      name: 'ChatRouterError',
+      message: 'Backend error: Internal error',
+      statusCode: 500
+    });
+  });
+
+  it('fails a stream that ends without its answer', async () => {
+    // A dropped connection: a failure, not a shorter answer.
+    answering(() => streamed([frame('status', { phase: 'generating' })]));
+    const router = new ChatRouter(settings, noNotebook);
+
+    await expect(
+      router.route('hi', [], undefined, () => undefined)
+    ).rejects.toThrow('The response ended before it was complete.');
+  });
+
+  it('reports a connection lost mid-answer as an error, not silence', async () => {
+    // A redeploy or a network drop fails the read in flight with a bare
+    // TypeError. Unwrapped, it reached the command as unexpected, and the chat
+    // showed an empty answer.
+    let reads = 0;
+    answering(() => ({
+      ...streamed([]),
+      body: {
+        getReader: () => ({
+          read: (): Promise<{ done: boolean; value?: Uint8Array }> => {
+            reads += 1;
+            return reads === 1
+              ? Promise.resolve({
+                  done: false,
+                  value: new TextEncoder().encode(
+                    frame('status', { phase: 'generating' })
+                  )
+                })
+              : Promise.reject(new TypeError('terminated'));
+          },
+          releaseLock: (): void => undefined
+        })
+      }
+    }));
+    const router = new ChatRouter(settings, noNotebook);
+
+    await expect(
+      router.route('hi', [], undefined, () => undefined)
+    ).rejects.toMatchObject({
+      name: 'ChatRouterError',
+      message:
+        'The connection to Oceanum AI was lost before the response was complete.'
+    });
+  });
+
+  it('streams the observe round too', async () => {
+    answering(() =>
+      streamed([
+        frame('status', { phase: 'interpreting' }),
+        frame('done', { message: 'Read it.', blocks: [] })
+      ])
+    );
+    const progress: unknown[] = [];
+    const router = new ChatRouter(settings, noNotebook);
+
+    const response = await router.observe(
+      'hi',
+      [],
+      [{ code: 'x', status: 'ok', stdout: '', error: null, message: 'm' }],
+      undefined,
+      p => progress.push(p)
+    );
+
+    expect(sentHeaders?.Accept).toBe('text/event-stream');
+    expect(progress).toEqual([{ phase: 'interpreting' }]);
+    expect(response).toEqual({ message: 'Read it.', blocks: [] });
+  });
+
+  it('stops in the middle of the stream when the request is aborted', async () => {
+    const controller = new AbortController();
+    answering(init =>
+      streamed([frame('status', { phase: 'generating' })], init.signal, true)
+    );
+    const progress: unknown[] = [];
+    const router = new ChatRouter(settings, noNotebook);
+
+    const pending = router.route('hi', [], controller.signal, p =>
+      progress.push(p)
+    );
+    await new Promise(resolve => setTimeout(resolve, 0));
+    controller.abort();
+
+    await expect(pending).rejects.toMatchObject({ name: 'AbortError' });
+    expect(progress).toEqual([{ phase: 'generating' }]);
+  });
+});

package/src/__tests__/conversationPin.spec.ts

@@ -0,0 +1,236 @@
+import type { NotebookPanel } from '@jupyterlab/notebook';
+import { ConversationPin, INotebookHost } from '../conversationPin';
+
+type Slot = (sender: unknown, path: string) => void;
+
+interface IFakePanel {
+  id: string;
+  isDisposed: boolean;
+  title: { label: string };
+  context: {
+    path: string;
+    pathChanged: {
+      slots: Array<[Slot, unknown]>;
+      connect(slot: Slot, thisArg: unknown): void;
+      disconnect(slot: Slot, thisArg: unknown): void;
+    };
+  };
+}
+
+const panel = (path: string): IFakePanel => {
+  const slots: Array<[Slot, unknown]> = [];
+  return {
+    id: `id:${path}`,
+    isDisposed: false,
+    title: { label: path.split('/').pop() ?? path },
+    context: {
+      path,
+      pathChanged: {
+        slots,
+        connect: (slot, thisArg) => slots.push([slot, thisArg]),
+        disconnect: (slot, thisArg) => {
+          const at = slots.findIndex(([s, t]) => s === slot && t === thisArg);
+          if (at >= 0) {
+            slots.splice(at, 1);
+          }
+        }
+      }
+    }
+  };
+};
+
+/** JupyterLab renaming or moving an open notebook. */
+const rename = (p: IFakePanel, path: string): void => {
+  p.context.path = path;
+  p.title.label = path.split('/').pop() ?? path;
+  for (const [slot, thisArg] of [...p.context.pathChanged.slots]) {
+    slot.call(thisArg, p.context, path);
+  }
+};
+
+const asPanel = (p: IFakePanel | null) => p as unknown as NotebookPanel | null;
+
+describe('ConversationPin', () => {
+  let active: IFakePanel | null;
+  let created: number;
+  let onDisk: Set<string>;
+  let host: jest.Mocked<INotebookHost>;
+  let pin: ConversationPin;
+
+  beforeEach(() => {
+    active = null;
+    created = 0;
+    onDisk = new Set();
+    host = {
+      activeTab: jest.fn(() => asPanel(active)),
+      // Like notebook:create-new: the new notebook becomes the active tab.
+      create: jest.fn(async () => {
+        created += 1;
+        active = panel(`Untitled${created}.ipynb`);
+        return asPanel(active);
+      }),
+      reopen: jest.fn(async (path: string) =>
+        asPanel(onDisk.has(path) ? panel(path) : null)
+      ),
+      activate: jest.fn()
+    };
+    pin = new ConversationPin(host);
+  });
+
+  it('pins the notebook in the active tab when a conversation starts', async () => {
+    active = panel('work/a.ipynb');
+    await expect(pin.start()).resolves.toBe('a.ipynb');
+    expect(host.create).not.toHaveBeenCalled();
+  });
+
+  it('creates a notebook when the active tab is not one, and pins it', async () => {
+    await expect(pin.start()).resolves.toBe('Untitled1.ipynb');
+    expect(host.create).toHaveBeenCalledTimes(1);
+    await expect(pin.notebook()).resolves.toBe(active);
+  });
+
+  it('starts an unstarted conversation at its first message, the same way', async () => {
+    await expect(pin.ensure()).resolves.toBe('Untitled1.ipynb');
+    await expect(pin.ensure()).resolves.toBe('Untitled1.ipynb');
+    expect(host.create).toHaveBeenCalledTimes(1);
+  });
+
+  it('keeps the pin for the whole conversation, whatever tab is active', async () => {
+    const a = panel('a.ipynb');
+    active = a;
+    await pin.start();
+
+    active = panel('b.ipynb');
+    await expect(pin.ensure()).resolves.toBe('a.ipynb');
+    await expect(pin.current()).resolves.toBe(a);
+    await expect(pin.notebook()).resolves.toBe(a);
+  });
+
+  it('creates one notebook for two New chats in quick succession', async () => {
+    // The second must find the notebook the first created as the active tab.
+    const [first, second] = await Promise.all([pin.start(), pin.start()]);
+    expect(host.create).toHaveBeenCalledTimes(1);
+    expect(first).toBe('Untitled1.ipynb');
+    expect(second).toBe('Untitled1.ipynb');
+  });
+
+  it('does not reopen a closed notebook just to be asked about', async () => {
+    // Re-opening the tab is for an answer that has cells to place.
+    const a = panel('work/a.ipynb');
+    onDisk.add('work/a.ipynb');
+    active = a;
+    await pin.start();
+
+    a.isDisposed = true;
+    await expect(pin.current()).resolves.toBeNull();
+    expect(pin.path()).toBe('work/a.ipynb');
+    expect(host.reopen).not.toHaveBeenCalled();
+  });
+
+  it('opens a closed notebook again, from where it was, to place an answer', async () => {
+    const a = panel('work/a.ipynb');
+    onDisk.add('work/a.ipynb');
+    active = a;
+    await pin.start();
+
+    a.isDisposed = true;
+    active = null;
+    const reopened = await pin.notebook();
+
+    expect(host.reopen).toHaveBeenCalledWith('work/a.ipynb');
+    expect(reopened).not.toBe(a);
+    expect(pin.name()).toBe('a.ipynb');
+    expect(host.create).not.toHaveBeenCalled();
+  });
+
+  it('reports the last known name while the notebook is closed', async () => {
+    const a = panel('work/a.ipynb');
+    active = a;
+    await pin.start();
+    a.isDisposed = true;
+    expect(pin.name()).toBe('a.ipynb');
+  });
+
+  it('follows a rename made while the notebook was open', async () => {
+    // Renamed and then closed with no request in between: found under the
+    // new name, not replaced by a new notebook.
+    const a = panel('work/a.ipynb');
+    onDisk.add('work/renamed.ipynb');
+    active = a;
+    await pin.start();
+
+    rename(a, 'work/renamed.ipynb');
+    a.isDisposed = true;
+    await pin.notebook();
+
+    expect(host.reopen).toHaveBeenCalledWith('work/renamed.ipynb');
+    expect(host.create).not.toHaveBeenCalled();
+  });
+
+  it('stops following a notebook once the conversation has moved on', async () => {
+    const a = panel('work/a.ipynb');
+    active = a;
+    await pin.start();
+    active = panel('work/b.ipynb');
+    await pin.start();
+
+    rename(a, 'work/elsewhere.ipynb');
+    expect(pin.path()).toBe('work/b.ipynb');
+  });
+
+  it('replaces a notebook that cannot be opened again with a new one', async () => {
+    const a = panel('work/gone.ipynb');
+    active = a;
+    await pin.start();
+
+    a.isDisposed = true;
+    active = null;
+    const replacement = await pin.notebook();
+
+    expect(host.create).toHaveBeenCalledTimes(1);
+    expect(replacement).toBe(active);
+    expect(pin.name()).toBe('Untitled1.ipynb');
+  });
+
+  it('brings the notebook to the front when shown, opening it first if closed', async () => {
+    const a = panel('work/a.ipynb');
+    onDisk.add('work/a.ipynb');
+    active = a;
+    await pin.start();
+    active = panel('b.ipynb');
+
+    await pin.show();
+    expect(host.activate).toHaveBeenLastCalledWith(a);
+
+    a.isDisposed = true;
+    const reopened = await pin.show();
+    expect(host.activate).toHaveBeenLastCalledWith(reopened);
+  });
+
+  it('retries a notebook that could not be created', async () => {
+    host.create.mockResolvedValueOnce(null);
+    await expect(pin.start()).resolves.toBeNull();
+
+    // Read `active` only after notebook() has run: it is the notebook the
+    // retry created.
+    const created = await pin.notebook();
+    expect(created).not.toBeNull();
+    expect(created).toBe(active);
+    expect(host.create).toHaveBeenCalledTimes(2);
+  });
+
+  it('keeps working after a step fails', async () => {
+    // One failure is its caller's to report; it must not block the queue.
+    const a = panel('work/a.ipynb');
+    onDisk.add('work/a.ipynb');
+    active = a;
+    await pin.start();
+    a.isDisposed = true;
+
+    host.reopen.mockRejectedValueOnce(new Error('server went away'));
+    await expect(pin.notebook()).rejects.toThrow('server went away');
+
+    await expect(pin.notebook()).resolves.not.toBeNull();
+    await expect(pin.ensure()).resolves.toBe('a.ipynb');
+  });
+});

package/src/__tests__/kernelHandoff.spec.ts

@@ -0,0 +1,106 @@
+import { CodeCell } from '@jupyterlab/cells';
+import type { NotebookPanel } from '@jupyterlab/notebook';
+import { NotebookActions } from '@jupyterlab/notebook';
+import { KernelHandoff } from '../kernelHandoff';
+
+// The real modules pull in ES-module UI code that jest cannot load. Stand-ins
+// with just what placing a code block without running it touches.
+jest.mock('@jupyterlab/notebook', () => ({
+  NotebookActions: { insertBelow: jest.fn(), changeCellType: jest.fn() },
+  NotebookPanel: class NotebookPanel {}
+}));
+jest.mock('@jupyterlab/cells', () => ({ CodeCell: class CodeCell {} }));
+
+/** A code cell holding `source`; `text()` reads what it holds now. */
+function codeCell(source: string): { cell: CodeCell; text: () => string } {
+  let text = source;
+  const cell = Object.create(CodeCell.prototype) as CodeCell;
+  Object.defineProperty(cell, 'isDisposed', { value: false });
+  Object.defineProperty(cell, 'model', {
+    value: {
+      sharedModel: {
+        setSource: (value: string) => {
+          text = value;
+        }
+      }
+    }
+  });
+  return { cell, text: () => text };
+}
+
+/** A notebook panel with `selected` selected; inserting selects `inserted`. */
+function panelWith(selected: CodeCell, inserted: CodeCell): NotebookPanel {
+  const content = { activeCell: selected };
+  (NotebookActions.insertBelow as jest.Mock).mockImplementation(() => {
+    content.activeCell = inserted;
+  });
+  return { content, sessionContext: {} } as unknown as NotebookPanel;
+}
+
+const answer = {
+  message: 'Fixed.',
+  blocks: [{ type: 'code' as const, content: 'x = 2' }]
+};
+
+describe('KernelHandoff: which notebook it touches', () => {
+  it('leaves the tabs alone when an answer has nothing to place', async () => {
+    // The conversation's notebook is re-opened for an answer with cells to
+    // place, not for one that is only words.
+    const target = jest.fn(async (): Promise<null> => null);
+    const handoff = new KernelHandoff(target, () => null);
+
+    const placed = await handoff.inject({ message: 'Just words.', blocks: [] });
+
+    expect(target).not.toHaveBeenCalled();
+    expect(placed.message).toBe('Just words.');
+  });
+
+  it('asks for the conversation notebook once when there are cells to place', async () => {
+    const target = jest.fn(async (): Promise<null> => null);
+    const handoff = new KernelHandoff(target, () => null);
+
+    const placed = await handoff.inject({
+      message: 'Here.',
+      blocks: [{ type: 'code', content: 'x = 1' }]
+    });
+
+    expect(target).toHaveBeenCalledTimes(1);
+    expect(placed.message).toContain('(Could not open a notebook)');
+    expect(placed.runs).toEqual([]);
+  });
+});
+
+describe('KernelHandoff: replacing the selected cell', () => {
+  it('replaces it in the notebook the request read it from', async () => {
+    const selected = codeCell('x = 1');
+    const inserted = codeCell('');
+    const panel = panelWith(selected.cell, inserted.cell);
+    const handoff = new KernelHandoff(
+      async () => panel,
+      () => panel
+    );
+
+    await handoff.inject(answer, { replaceCodeCell: true });
+
+    expect(selected.text()).toBe('x = 2');
+    expect(inserted.text()).toBe('');
+  });
+
+  it('never replaces a cell in a notebook opened again since', async () => {
+    // Closed while the answer was on its way and opened again: its selected
+    // cell is not the one the user chose, but some other work of theirs.
+    const selected = codeCell('users_latest_work()');
+    const inserted = codeCell('');
+    const reopened = panelWith(selected.cell, inserted.cell);
+    const readFrom = {} as NotebookPanel;
+    const handoff = new KernelHandoff(
+      async () => reopened,
+      () => readFrom
+    );
+
+    await handoff.inject(answer, { replaceCodeCell: true });
+
+    expect(selected.text()).toBe('users_latest_work()');
+    expect(inserted.text()).toBe('x = 2');
+  });
+});