bDS/tests/renderer/python/PythonRuntimeManager.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { PythonRuntimeManager } from '../../../src/renderer/python/PythonRuntimeManager';
import { createMacroRenderOptions } from '../../../src/renderer/python/macroRenderOptions';

class MockWorker {
  onmessage: ((event: MessageEvent) => void) | null = null;
  onerror: ((event: ErrorEvent) => void) | null = null;
  terminated = false;
  postedMessages: unknown[] = [];

  postMessage(message: unknown): void {
    this.postedMessages.push(message);
  }

  terminate(): void {
    this.terminated = true;
  }

  emitMessage(data: unknown): void {
    this.onmessage?.({ data } as MessageEvent);
  }

  emitError(error: Error): void {
    this.onerror?.({ error } as ErrorEvent);
  }
}

describe('PythonRuntimeManager', () => {
  beforeEach(() => {
    vi.useFakeTimers();
  });

  afterEach(() => {
    vi.useRealTimers();
  });

  it('initializes worker once and resolves on ready signal', async () => {
    const createdWorkers: MockWorker[] = [];
    const manager = new PythonRuntimeManager(() => {
      const worker = new MockWorker();
      createdWorkers.push(worker);
      return worker as unknown as Worker;
    });

    const initA = manager.initialize();
    const initB = manager.initialize();

    expect(initA).toBe(initB);
    expect(createdWorkers).toHaveLength(1);

    createdWorkers[0].emitMessage({ type: 'ready' });
    await expect(initA).resolves.toBeUndefined();
    expect(manager.isReady()).toBe(true);
  });

  it('rejects when worker emits an error before ready', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitError(new Error('bootstrap failed'));

    await expect(initPromise).rejects.toThrow('bootstrap failed');
    expect(manager.isReady()).toBe(false);
  });

  it('executes code and returns stdout and result', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.execute('print("hello")\n1 + 1');
    await Promise.resolve();
    const request = worker.postedMessages[0] as { type: string; requestId: string; code: string };

    worker.emitMessage({ type: 'stdout', requestId: request.requestId, chunk: 'hello\n' });
    worker.emitMessage({ type: 'runResult', requestId: request.requestId, result: '2' });

    await expect(runPromise).resolves.toEqual({ result: '2', stdout: 'hello\n' });
  });

  it('forwards compile cache key in execute request options', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.execute('value = 1', { cacheKey: 'script-1:3' });
    await Promise.resolve();
    const request = worker.postedMessages[0] as { requestId: string; cacheKey?: string };
    expect(request.cacheKey).toBe('script-1:3');

    worker.emitMessage({ type: 'runResult', requestId: request.requestId, result: '' });
    await expect(runPromise).resolves.toEqual({ result: '', stdout: '' });
  });

  it('forwards selected entrypoint in execute request options', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.execute('def helper():\n    return 42', { entrypoint: 'helper' });
    await Promise.resolve();
    const request = worker.postedMessages[0] as { requestId: string; entrypoint?: string };
    expect(request.entrypoint).toBe('helper');

    worker.emitMessage({ type: 'runResult', requestId: request.requestId, result: '42' });
    await expect(runPromise).resolves.toEqual({ result: '42', stdout: '' });
  });

  it('inspects script and returns available function names', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const inspectPromise = manager.inspectEntrypoints('def render(context):\n    return {}\n\ndef helper():\n    return 1');
    await Promise.resolve();

    const request = worker.postedMessages[0] as { type: string; requestId: string; code: string };
    expect(request.type).toBe('inspectEntrypoints');

    worker.emitMessage({ type: 'entrypoints', requestId: request.requestId, entrypoints: ['render', 'helper'] });
    await expect(inspectPromise).resolves.toEqual(['render', 'helper']);
  });

  it('checks syntax and returns structured syntax errors', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const syntaxPromise = manager.syntaxCheck('def broken(:\n  pass');
    await Promise.resolve();

    const request = worker.postedMessages[0] as { type: string; requestId: string; code: string };
    expect(request.type).toBe('syntaxCheck');

    worker.emitMessage({
      type: 'syntaxResult',
      requestId: request.requestId,
      errors: [
        {
          line: 1,
          column: 11,
          endLine: 1,
          endColumn: 12,
          message: 'invalid syntax',
        },
      ],
    });

    await expect(syntaxPromise).resolves.toEqual({
      errors: [
        {
          line: 1,
          column: 11,
          endLine: 1,
          endColumn: 12,
          message: 'invalid syntax',
        },
      ],
    });
  });

  it('rejects when runtime returns run error', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.execute('raise RuntimeError("boom")');
    await Promise.resolve();
    const request = worker.postedMessages[0] as { requestId: string };
    worker.emitMessage({ type: 'runError', requestId: request.requestId, error: 'boom' });

    await expect(runPromise).rejects.toThrow('boom');
  });

  it('queues concurrent execute calls and sends the next request only after completion', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const firstRun = manager.execute('1 + 1');
    const secondRun = manager.execute('2 + 2');
    await Promise.resolve();

    expect(worker.postedMessages).toHaveLength(1);

    const firstRequest = worker.postedMessages[0] as { requestId: string };
    worker.emitMessage({ type: 'runResult', requestId: firstRequest.requestId, result: '2' });
    await expect(firstRun).resolves.toEqual({ result: '2', stdout: '' });

    await Promise.resolve();
    expect(worker.postedMessages).toHaveLength(2);

    const secondRequest = worker.postedMessages[1] as { requestId: string };
    worker.emitMessage({ type: 'runResult', requestId: secondRequest.requestId, result: '4' });
    await expect(secondRun).resolves.toEqual({ result: '4', stdout: '' });
  });

  it('terminates timed out worker and recovers with a new worker', async () => {
    const workers: MockWorker[] = [];
    const manager = new PythonRuntimeManager(() => {
      const worker = new MockWorker();
      workers.push(worker);
      return worker as unknown as Worker;
    });

    const firstInit = manager.initialize();
    workers[0].emitMessage({ type: 'ready' });
    await firstInit;

    const timedOutRun = manager.execute('while True: pass', { timeoutMs: 100 });
    await Promise.resolve();
    vi.advanceTimersByTime(101);
    await expect(timedOutRun).rejects.toThrow('timed out');
    expect(workers[0].terminated).toBe(true);
    expect(manager.isReady()).toBe(false);

    const secondInit = manager.initialize();
    workers[1].emitMessage({ type: 'ready' });
    await secondInit;

    const secondRun = manager.execute('40 + 2');
    await Promise.resolve();
    const request = workers[1].postedMessages[0] as { requestId: string };
    workers[1].emitMessage({ type: 'runResult', requestId: request.requestId, result: '42' });

    await expect(secondRun).resolves.toEqual({ result: '42', stdout: '' });
  });

  it('rejects macro execution when ABI context is invalid on caller side', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.renderMacroV1('def render(context):\n  return {"html": "<p>ok</p>"}', {
      env: {
        isPreview: 'yes',
      },
    });

    await expect(runPromise).rejects.toThrow('Invalid macro context');
    expect(worker.postedMessages).toHaveLength(0);
  });

  it('returns validated macro result and stdout', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.renderMacroV1('def render(context):\n  return {"html": "<p>ok</p>"}', {
      env: {
        isPreview: true,
      },
      params: {
        title: 'Hello',
      },
    });

    await Promise.resolve();
    const request = worker.postedMessages[0] as { requestId: string };
    worker.emitMessage({ type: 'stdout', requestId: request.requestId, chunk: 'rendering\n' });
    worker.emitMessage({ type: 'macroResult', requestId: request.requestId, result: { html: '<p>ok</p>' } });

    await expect(runPromise).resolves.toEqual({ result: { html: '<p>ok</p>' }, stdout: 'rendering\n' });
  });

  it('accepts optional env hook and source metadata for macro execution', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.renderMacroV1('def render(context):\n  return {"html": "<p>ok</p>"}', {
      env: {
        isPreview: true,
        hook: 'post:render',
        source: {
          kind: 'post',
          id: 'post-1',
        },
      },
    });

    await Promise.resolve();
    const request = worker.postedMessages[0] as {
      requestId: string;
      context: { env: { hook?: string; source?: { kind: string; id?: string } } };
    };
    expect(request.context.env.hook).toBe('post:render');
    expect(request.context.env.source).toEqual({ kind: 'post', id: 'post-1' });

    worker.emitMessage({ type: 'macroResult', requestId: request.requestId, result: { html: '<p>ok</p>' } });
    await expect(runPromise).resolves.toEqual({ result: { html: '<p>ok</p>' }, stdout: '' });
  });

  it('injects env hook and source metadata from macro execution options', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.renderMacroV1(
      'def render(context):\n  return {"html": "<p>ok</p>"}',
      {
        env: {
          isPreview: true,
        },
      },
      createMacroRenderOptions({
        hook: 'preview:macro',
        source: {
          kind: 'post',
          id: 'post-77',
        },
      })
    );

    await Promise.resolve();
    const request = worker.postedMessages[0] as {
      requestId: string;
      context: { env: { hook?: string; source?: { kind: string; id?: string } } };
    };

    expect(request.context.env.hook).toBe('preview:macro');
    expect(request.context.env.source).toEqual({ kind: 'post', id: 'post-77' });

    worker.emitMessage({ type: 'macroResult', requestId: request.requestId, result: { html: '<p>ok</p>' } });
    await expect(runPromise).resolves.toEqual({ result: { html: '<p>ok</p>' }, stdout: '' });
  });

  it('preserves explicit env hook and source over macro execution options', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.renderMacroV1(
      'def render(context):\n  return {"html": "<p>ok</p>"}',
      {
        env: {
          isPreview: true,
          hook: 'explicit:hook',
          source: {
            kind: 'page',
            id: 'page-9',
          },
        },
      },
      createMacroRenderOptions({
        hook: 'preview:macro',
        source: {
          kind: 'post',
          id: 'post-77',
        },
      })
    );

    await Promise.resolve();
    const request = worker.postedMessages[0] as {
      requestId: string;
      context: { env: { hook?: string; source?: { kind: string; id?: string } } };
    };

    expect(request.context.env.hook).toBe('explicit:hook');
    expect(request.context.env.source).toEqual({ kind: 'page', id: 'page-9' });

    worker.emitMessage({ type: 'macroResult', requestId: request.requestId, result: { html: '<p>ok</p>' } });
    await expect(runPromise).resolves.toEqual({ result: { html: '<p>ok</p>' }, stdout: '' });
  });

  it('rejects macro execution when worker result violates ABI schema', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.renderMacroV1('def render(context):\n  return {"html": "<p>ok</p>"}', {
      env: {
        isPreview: true,
      },
    });

    await Promise.resolve();
    const request = worker.postedMessages[0] as { requestId: string };
    worker.emitMessage({ type: 'macroResult', requestId: request.requestId, result: { html: 42 } });

    await expect(runPromise).rejects.toThrow('Invalid macro result');
  });

  it('handles worker apiCall by invoking host bridge and returning apiResult', async () => {
    const worker = new MockWorker();
    const invokeApiCall = vi.fn().mockResolvedValue({ id: 'post-1', title: 'Hello' });
    const manager = new PythonRuntimeManager(
      () => worker as unknown as Worker,
      {
        invokeApiCall,
      }
    );

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.execute('print("hello")');
    await Promise.resolve();

    const runRequest = worker.postedMessages[0] as { requestId: string };
    worker.emitMessage({
      type: 'apiCall',
      requestId: runRequest.requestId,
      callId: 'call-1',
      method: 'posts.get',
      args: { postId: 'post-1' },
    });

    await Promise.resolve();

    expect(invokeApiCall).toHaveBeenCalledWith('posts.get', { postId: 'post-1' });
    expect(worker.postedMessages[1]).toEqual({
      type: 'apiResult',
      requestId: runRequest.requestId,
      callId: 'call-1',
      ok: true,
      result: { id: 'post-1', title: 'Hello' },
    });

    worker.emitMessage({ type: 'runResult', requestId: runRequest.requestId, result: 'done' });
    await expect(runPromise).resolves.toEqual({ result: 'done', stdout: '' });
  });

  it('returns apiResult with error when host bridge invocation fails', async () => {
    const worker = new MockWorker();
    const invokeApiCall = vi.fn().mockRejectedValue(new Error('unknown api method'));
    const manager = new PythonRuntimeManager(
      () => worker as unknown as Worker,
      {
        invokeApiCall,
      }
    );

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.execute('print("hello")');
    await Promise.resolve();

    const runRequest = worker.postedMessages[0] as { requestId: string };
    worker.emitMessage({
      type: 'apiCall',
      requestId: runRequest.requestId,
      callId: 'call-2',
      method: 'posts.nonExisting',
      args: {},
    });

    await Promise.resolve();

    expect(worker.postedMessages[1]).toEqual({
      type: 'apiResult',
      requestId: runRequest.requestId,
      callId: 'call-2',
      ok: false,
      error: 'unknown api method',
    });

    worker.emitMessage({ type: 'runResult', requestId: runRequest.requestId, result: 'done' });
    await expect(runPromise).resolves.toEqual({ result: 'done', stdout: '' });
  });

  it('does not time out when timeoutMs is 0', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.execute('long_running()', { timeoutMs: 0 });
    await Promise.resolve();

    // Advance time well past any default timeout — script must still be pending
    vi.advanceTimersByTime(60_000);
    expect(worker.terminated).toBe(false);

    const request = worker.postedMessages[0] as { requestId: string };
    worker.emitMessage({ type: 'runResult', requestId: request.requestId, result: 'done' });

    await expect(runPromise).resolves.toEqual({ result: 'done', stdout: '' });
  });

  it('queued inspectEntrypoints with timeoutMs 0 does not kill running execute', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    // Start a long-running execute with no timeout
    const runPromise = manager.execute('long_running()', { timeoutMs: 0 });
    await Promise.resolve();

    // Queue inspectEntrypoints (default timeout) while execute is running
    const inspectPromise = manager.inspectEntrypoints('def render(): pass');
    await Promise.resolve();

    // Advance past the default 5000ms timeout
    vi.advanceTimersByTime(6000);

    // Worker must still be alive — the queued inspect must not kill it
    expect(worker.terminated).toBe(false);

    // Finish the execute
    const runRequest = worker.postedMessages[0] as { requestId: string };
    worker.emitMessage({ type: 'runResult', requestId: runRequest.requestId, result: 'done' });
    await expect(runPromise).resolves.toEqual({ result: 'done', stdout: '' });

    // Now the inspect request dispatches — respond to it
    await Promise.resolve();
    const inspectRequest = worker.postedMessages[1] as { requestId: string };
    worker.emitMessage({ type: 'entrypoints', requestId: inspectRequest.requestId, entrypoints: ['render'] });
    await expect(inspectPromise).resolves.toEqual(['render']);
  });

  it('calls onStdout callback for each stdout chunk during execution', async () => {
    const worker = new MockWorker();
    const manager = new PythonRuntimeManager(() => worker as unknown as Worker);

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const stdoutChunks: string[] = [];
    const runPromise = manager.execute('print("a")\nprint("b")', {
      onStdout: (chunk) => { stdoutChunks.push(chunk); },
    });
    await Promise.resolve();

    const request = worker.postedMessages[0] as { requestId: string };
    worker.emitMessage({ type: 'stdout', requestId: request.requestId, chunk: 'a\n' });
    worker.emitMessage({ type: 'stdout', requestId: request.requestId, chunk: 'b\n' });
    worker.emitMessage({ type: 'runResult', requestId: request.requestId, result: '' });

    const result = await runPromise;
    expect(stdoutChunks).toEqual(['a\n', 'b\n']);
    expect(result.stdout).toBe('a\nb\n');
  });

  it('calls onToast handler when worker sends a toast message', async () => {
    const worker = new MockWorker();
    const toasts: Array<{ message: string; toastType?: string }> = [];
    const manager = new PythonRuntimeManager(
      () => worker as unknown as Worker,
      {
        onToast: (message, toastType) => { toasts.push({ message, toastType }); },
      }
    );

    const initPromise = manager.initialize();
    worker.emitMessage({ type: 'ready' });
    await initPromise;

    const runPromise = manager.execute('toast("hello")');
    await Promise.resolve();

    const request = worker.postedMessages[0] as { requestId: string };
    worker.emitMessage({ type: 'toast', message: 'hello', toastType: 'success' });
    worker.emitMessage({ type: 'toast', message: 'oops', toastType: 'error' });
    worker.emitMessage({ type: 'toast', message: 'note' });

    expect(toasts).toEqual([
      { message: 'hello', toastType: 'success' },
      { message: 'oops', toastType: 'error' },
      { message: 'note', toastType: undefined },
    ]);

    worker.emitMessage({ type: 'runResult', requestId: request.requestId, result: '' });
    await expect(runPromise).resolves.toEqual({ result: '', stdout: '' });
  });
});