test/controller.test.js

import test from 'node:test';
import assert from 'node:assert/strict';
import { WorkspaceController } from '../src/controller.js';
import { ContextOverflowError, DeviceLostError } from '../runtime/minicpm5.js';

const deferred = () => {
  let resolve;
  let reject;
  const promise = new Promise((yes, no) => { resolve = yes; reject = no; });
  return { promise, resolve, reject };
};

function answer(content = 'Done.', { calls, reasoning, finish = calls ? 'tool_calls' : 'stop' } = {}) {
  return {
    finish_reason: finish,
    message: {
      role: 'assistant', content,
      ...(reasoning ? { reasoning_content: reasoning } : {}),
      ...(calls ? { tool_calls: calls } : {}),
    },
    usage: { prompt_tokens: 80, cached_tokens: 40, completion_tokens: 12, reasoning_tokens: 3 },
    timings: { ttft_ms: 15, decode_ms: 50, tokens_per_second: 220 },
  };
}

function waitForAbort(signal) {
  if (signal.aborted) return Promise.resolve();
  return new Promise(resolve => signal.addEventListener('abort', resolve, { once: true }));
}

function modelWithSteps(steps) {
  const loss = deferred();
  const disposed = deferred();
  const requests = [];
  const model = {
    lost: loss.promise,
    lose: loss.resolve,
    disposed: disposed.promise,
    disposeCount: 0,
    sessionCount: 0,
    sessionDisposeCount: 0,
    contextLength: 4096,
    requests,
    deviceInfo: () => ({ vendor: 'test', architecture: 'mock-gpu' }),
    dispose() { this.disposeCount += 1; disposed.resolve(); },
    create_session() {
      this.sessionCount += 1;
      return {
        dispose() { model.sessionDisposeCount += 1; },
        generate(history, options) {
          const step = steps[requests.length];
          assert.ok(step, 'unexpected extra generation');
          requests.push({ history, options });
          const completion = deferred();
          completion.promise.catch(() => {});
          return {
            result: completion.promise,
            async *[Symbol.asyncIterator]() {
              try {
                for (const event of step.events ?? []) yield event;
                if (step.waitForAbort) await waitForAbort(options.signal);
                if (step.error) throw step.error;
                const result = step.result ?? step;
                completion.resolve(result);
              } catch (error) {
                completion.reject(error);
                throw error;
              }
            },
          };
        },
      };
    },
  };
  return model;
}

function harness({ models = [], saved, checkAvailability, fromPretrained, tools = [], save } = {}) {
  const events = [];
  const snapshots = [];
  const waiters = [];
  const calls = { support: 0, load: 0 };
  const files = {};
  const workspace = {
    tools,
    snapshot: async () => ({ ...files }),
    serialize: async () => Object.entries(files).map(([path, content]) => ({ path, content })),
    write: async (path, content) => { files[path] = content; },
    read: async path => files[path],
    bash: { exec: async command => ({ stdout: command, stderr: '', exitCode: 0 }) },
  };
  const store = {
    async save(snapshot) {
      snapshots.push(structuredClone(snapshot));
      if (save) await save(snapshot);
    },
  };
  const Model = {
    async checkAvailability(...args) {
      calls.support += 1;
      return checkAvailability ? checkAvailability(...args) : { ok: true };
    },
    async from_pretrained(...args) {
      const index = calls.load++;
      if (fromPretrained) return fromPretrained(index, ...args);
      assert.ok(models[index], 'unexpected extra model load');
      return models[index];
    },
  };
  const controller = new WorkspaceController({ Model, workspace, store, saved, emit(event) {
    // Match postMessage/IndexedDB boundaries; source objects continue changing.
    events.push(structuredClone(event));
    for (const waiter of [...waiters]) {
      if (waiter.test(event)) {
        waiters.splice(waiters.indexOf(waiter), 1);
        waiter.resolve(event);
      }
    }
  } });
  const until = predicate => {
    const existing = events.find(predicate);
    if (existing) return Promise.resolve(existing);
    return new Promise(resolve => waiters.push({ test: predicate, resolve }));
  };
  return { controller, events, snapshots, calls, files, until };
}

const agentEvents = events => events.filter(event => event.type === 'agent_event').map(event => event.event);

test('native agent replies and tools render in the existing terminal protocol and persist separately', async () => {
  const call = { id: 'write-1', type: 'function', function: {
    name: 'write', arguments: { path: '/workspace/result.txt', content: '42' },
  } };
  const first = answer('Writing the result.', { calls: [call], reasoning: 'I should save it.' });
  const last = answer('Saved result.txt.');
  const model = modelWithSteps([
    { result: first, events: [
      { type: 'prefill', done: 32, total: 64 },
      { type: 'reasoning', text: 'I should save it.' },
      { type: 'text', text: 'Writing the result.' },
    ] },
    { result: last, events: [{ type: 'text', text: 'Saved result.txt.' }] },
  ]);
  const h = harness({ models: [model] });
  h.controller.workspace.tools.push({
    name: 'write', parameters: { type: 'object', properties: {
      path: { type: 'string' }, content: { type: 'string' },
    } },
    async execute({ path, content }) {
      await h.controller.workspace.write(path, content);
      return { content: 'File saved.', details: { path } };
    },
  });
  await h.controller.dispatch('load');
  const ui = await h.controller.dispatch('prompt', { text: 'Save 42.' });
  const native = h.snapshots.at(-1).messages;

  assert.equal(h.files['/workspace/result.txt'], '42');
  assert.deepEqual(native.map(message => message.role), ['user', 'assistant', 'tool', 'assistant']);
  assert.deepEqual(native[1], first.message, 'canonical reasoning and tool calls must remain restorable');
  assert.deepEqual(native[2], { role: 'tool', tool_call_id: 'write-1', name: 'write', content: 'File saved.' });
  assert.deepEqual(ui.messages.map(message => message.role), ['user', 'assistant', 'toolResult', 'assistant']);
  assert.deepEqual(ui.messages[1].content.map(block => block.type), ['thinking', 'text', 'toolCall']);
  assert.equal(ui.messages[1].stopReason, 'toolUse');
  assert.equal(ui.messages[2].content[0].text, 'File saved.');
  assert.equal(ui.messages[3].content[0].text, 'Saved result.txt.');
  assert.equal(ui.messages[3].usage.cacheRead, 40);
  assert.equal(ui.messages[3].usage.output, 12);
  assert.ok(native.every(message => !Object.hasOwn(message, 'timestamp')));
  assert.ok(!native.some(message => message.role === 'system'), 'instructions are supplied fresh on restore');
  assert.deepEqual(model.requests[1].history.slice(-2).map(message => message.role), ['assistant', 'tool']);
  const start = agentEvents(h.events).find(event => event.type === 'tool_execution_start');
  assert.equal(start.toolName, 'write');
  assert.deepEqual(start.args, call.function.arguments);
  const end = agentEvents(h.events).find(event => event.type === 'tool_execution_end');
  assert.equal(end.result.content[0].text, 'File saved.');
  assert.equal(end.isError, false);
  assert.ok(h.events.some(event => event.type === 'inference_activity' && event.done === 32 && event.total === 64));
  const metrics = h.events.filter(event => event.type === 'generation');
  assert.equal(metrics.at(-1).tokensPerSecond, 220, 'show runtime token timing, not character throughput');
  assert.equal(metrics.at(-1).outputTokens, 12);
  assert.equal(h.events.find(event => event.type === 'loaded').contextLength, 4096);
  assert.equal(agentEvents(h.events).at(-1).type, 'agent_end');
  assert.equal(h.events.at(-1).type, 'idle');
});

test('restoring a saved chat sends native history to the model and display history to the UI', async () => {
  const native = [
    { role: 'user', content: 'Earlier question' },
    { role: 'assistant', content: 'Earlier answer', reasoning_content: 'Earlier reasoning' },
  ];
  const display = [
    { role: 'user', content: 'Earlier question', timestamp: 1 },
    { role: 'assistant', content: [{ type: 'text', text: 'Earlier answer' }], timestamp: 2, stopReason: 'stop' },
  ];
  const model = modelWithSteps([answer('Follow-up answer')]);
  const h = harness({ models: [model], saved: { messages: native, displayMessages: display } });
  assert.deepEqual((await h.controller.uiSnapshot()).messages, display);
  await h.controller.dispatch('load');
  await h.controller.dispatch('prompt', { text: 'Follow-up' });

  assert.deepEqual(model.requests[0].history.slice(1, 3), native);
  assert.equal(model.requests[0].history.at(-1).content, 'Follow-up');
  assert.equal(h.snapshots.at(-1).messages[1].reasoning_content, 'Earlier reasoning');
  const exported = await h.controller.dispatch('snapshot');
  assert.equal(exported.messages[1].content, 'Earlier answer');
  assert.equal(exported.displayMessages[1].content[0].text, 'Earlier answer');
});

test('new chat clears model/display history and reuses the loaded model and session', async () => {
  const model = modelWithSteps([answer('First answer'), answer('Fresh answer')]);
  const h = harness({ models: [model] });
  await h.controller.dispatch('load');
  await h.controller.dispatch('prompt', { text: 'First question' });
  const cleared = await h.controller.dispatch('new_chat');
  assert.deepEqual(cleared.messages, []);
  assert.deepEqual(h.snapshots.at(-1).messages, []);
  await h.controller.dispatch('prompt', { text: 'Fresh question' });

  assert.equal(h.calls.load, 1);
  assert.equal(model.sessionCount, 1);
  assert.deepEqual(model.requests[1].history.map(message => message.role), ['system', 'user']);
  assert.equal(model.requests[1].history.at(-1).content, 'Fresh question');
});

test('stop settles generation, preserves its visible partial, and excludes it from future model history', async () => {
  const partial = answer('Interrupted text', { finish: 'cancelled' });
  const model = modelWithSteps([
    { result: partial, events: [{ type: 'text', text: 'Interrupted text' }], waitForAbort: true },
    answer('Next answer'),
  ]);
  const h = harness({ models: [model] });
  await h.controller.dispatch('load');
  const pending = h.controller.dispatch('prompt', { text: 'Interrupted question' });
  await h.until(event => event.type === 'agent_event' && event.event.type === 'message_start');
  assert.equal(await h.controller.dispatch('stop'), true);
  const ui = await pending;

  assert.equal(ui.messages.at(-1).content[0].text, 'Interrupted text');
  assert.equal(ui.messages.at(-1).stopReason, 'aborted');
  assert.deepEqual(h.snapshots.at(-1).messages, []);
  assert.ok(h.events.some(event => event.type === 'notice' && /omitted from the model history/.test(event.text)));
  await h.controller.dispatch('prompt', { text: 'Next question' });
  assert.deepEqual(model.requests[1].history.map(message => message.role), ['system', 'user']);
  assert.equal(model.requests[1].history.at(-1).content, 'Next question');
  assert.equal(h.controller.busy, false);
});

test('busy operations reject extra prompts, edits, resets, and export while stop remains available', async () => {
  const model = modelWithSteps([{ result: answer('', { finish: 'cancelled' }), waitForAbort: true }]);
  const h = harness({ models: [model] });
  await h.controller.dispatch('load');
  const pending = h.controller.dispatch('prompt', { text: 'Wait' });
  await h.until(event => event.type === 'agent_event' && event.event.type === 'message_start');
  for (const action of ['prompt', 'write', 'new_chat', 'snapshot', 'load']) {
    await assert.rejects(h.controller.dispatch(action, { text: 'Other', path: '/x', content: 'bad' }), /current operation/);
  }
  assert.deepEqual(h.files, {});
  await h.controller.dispatch('stop');
  await pending;
  assert.equal(h.controller.busy, false);
  assert.equal(h.controller.operation, null);
});

test('a failed model download can be retried and never leaves the controller busy', async () => {
  const model = modelWithSteps([answer('Ready')]);
  const h = harness({ fromPretrained: async index => {
    if (index === 0) throw new Error('Download interrupted');
    return model;
  } });
  await assert.rejects(h.controller.dispatch('load'), /Download interrupted/);
  assert.equal(h.controller.busy, false);
  assert.equal(h.controller.model, null);
  assert.equal(h.controller.agent, null);
  await h.controller.dispatch('load');
  await h.controller.dispatch('prompt', { text: 'Hello' });
  assert.equal(h.calls.load, 2);
  assert.equal(h.events.filter(event => event.type === 'loaded').length, 1);
});

test('unsupported GPUs and cache-only loading fail before any weight download', async () => {
  const unsupported = harness({ checkAvailability: async () => ({ ok: false, reason: 'Insufficient GPU limits' }) });
  await assert.rejects(unsupported.controller.dispatch('load'), /Insufficient GPU limits/);
  assert.equal(unsupported.calls.load, 0);
  assert.equal(unsupported.controller.busy, false);

  const cacheOnly = harness();
  await assert.rejects(cacheOnly.controller.dispatch('load', { cachedOnly: true }), error => error.code === 'MODEL_CACHE_MISS');
  assert.equal(cacheOnly.calls.support, 0);
  assert.equal(cacheOnly.calls.load, 0);
  assert.equal(await cacheOnly.controller.dispatch('cache_status'), false);
});

test('GPU loss releases the old session/model and allows a reload with the saved conversation', async () => {
  const oldModel = modelWithSteps([answer('Before loss')]);
  const newModel = modelWithSteps([answer('After reload')]);
  const h = harness({ models: [oldModel, newModel] });
  await h.controller.dispatch('load');
  await h.controller.dispatch('prompt', { text: 'Remember this' });
  oldModel.lose(new DeviceLostError('unknown', 'Adapter removed'));
  await h.until(event => event.type === 'fatal' && /GPU device lost/.test(event.error));
  await oldModel.disposed;

  assert.equal(oldModel.disposeCount, 1);
  assert.equal(oldModel.sessionDisposeCount, 1);
  assert.equal(h.controller.model, null);
  await h.controller.dispatch('load');
  await h.controller.dispatch('prompt', { text: 'Continue' });
  assert.equal(newModel.requests[0].history[1].content, 'Remember this');
  assert.equal(newModel.requests[0].history[2].content, 'Before loss');
  assert.equal(h.controller.deviceLoss, null);
  assert.equal(h.calls.load, 2);
});

test('GPU loss during a reply ends the operation and retains the last completed conversation', async () => {
  const deviceError = new DeviceLostError('unknown', 'Adapter disconnected');
  const oldModel = modelWithSteps([
    answer('Completed answer'),
    { error: deviceError, events: [{ type: 'text', text: 'Interrupted' }], waitForAbort: true },
  ]);
  const newModel = modelWithSteps([answer('Recovered answer')]);
  const h = harness({ models: [oldModel, newModel] });
  await h.controller.dispatch('load');
  await h.controller.dispatch('prompt', { text: 'Completed question' });
  h.events.length = 0;
  const pending = h.controller.dispatch('prompt', { text: 'Interrupted question' });
  const rejected = assert.rejects(pending, error => error === deviceError);
  await h.until(event => event.type === 'agent_event' && event.event.type === 'message_start');
  oldModel.lose(deviceError);
  await rejected;
  await oldModel.disposed;

  assert.equal(h.controller.busy, false);
  assert.equal(h.snapshots.at(-1).displayMessages.at(-1).stopReason, 'aborted');
  assert.deepEqual(h.snapshots.at(-1).messages.map(message => message.content), ['Completed question', 'Completed answer']);
  assert.equal(agentEvents(h.events).at(-1).type, 'agent_end');
  await h.controller.dispatch('load');
  await h.controller.dispatch('prompt', { text: 'Continue after loss' });
  assert.deepEqual(newModel.requests[0].history.slice(1).map(message => message.content), [
    'Completed question', 'Completed answer', 'Continue after loss',
  ]);
});

test('context overflow and generation failures end the visible reply without corrupting canonical history', async () => {
  const overflow = new ContextOverflowError(9000, 4096);
  const model = modelWithSteps([
    { error: overflow },
    { error: new Error('Shader execution failed'), events: [{ type: 'text', text: 'Partial' }] },
    answer('Recovered'),
  ]);
  const h = harness({ models: [model] });
  await h.controller.dispatch('load');
  await assert.rejects(h.controller.dispatch('prompt', { text: 'Too long' }), error => error === overflow);
  assert.ok(h.events.some(event => event.type === 'notice' && /9000 tokens.*4096/.test(event.text)));
  assert.equal(h.snapshots.at(-1).displayMessages.at(-1).stopReason, 'error');
  assert.deepEqual(h.snapshots.at(-1).messages, []);
  await assert.rejects(h.controller.dispatch('prompt', { text: 'Fails' }), /Shader execution failed/);
  const visible = h.snapshots.at(-1).displayMessages.at(-1);
  assert.equal(visible.stopReason, 'error');
  assert.equal(visible.content[0].text, 'Partial');
  assert.equal(visible.errorMessage, 'Shader execution failed');
  assert.deepEqual(h.snapshots.at(-1).messages, []);
  await h.controller.dispatch('prompt', { text: 'Try again' });
  assert.deepEqual(model.requests[2].history.map(message => message.role), ['system', 'user']);
  assert.equal(h.snapshots.at(-1).messages.at(-1).content, 'Recovered');
  assert.equal(h.controller.busy, false);
});

test('load progress distinguishes byte download counts from tensors and compilation', async () => {
  const model = modelWithSteps([]);
  const h = harness({ fromPretrained: async (_index, _id, options) => {
    options.onProgress({ status: 'init', message: 'Requesting GPU' });
    options.onProgress({ status: 'weights', kind: 'bytes', loaded: 512, total: 1024, fraction: 0.5 });
    options.onProgress({ status: 'weights', kind: 'tensors', loaded: 10, total: 20, message: 'Uploading' });
    options.onProgress({ status: 'weights', kind: 'tensors', message: 'Compiling kernels' });
    options.onProgress({ status: 'ready', fraction: 1, message: 'Ready' });
    return model;
  } });
  await h.controller.dispatch('load');
  const progress = h.events.filter(event => event.type === 'load_progress');
  assert.deepEqual(progress.filter(event => event.phase === 'download'), [
    { type: 'load_progress', phase: 'download', loaded: 512, total: 1024 },
  ]);
  assert.equal(progress.find(event => event.message === 'Uploading').phase, 'compile');
  assert.equal(progress.find(event => event.message === 'Compiling kernels').phase, 'compile');
  assert.equal(h.events.filter(event => event.type === 'loaded').length, 1);
  assert.ok(h.events.findIndex(event => event.type === 'loaded') > h.events.findIndex(event => event.message === 'Compiling kernels'));
});

test('a persistence failure keeps the answer usable and reports that changes remain in memory', async () => {
  const model = modelWithSteps([answer('Still available')]);
  const h = harness({ models: [model], save: async () => { throw new Error('Storage quota exceeded'); } });
  await h.controller.dispatch('load');
  const ui = await h.controller.dispatch('prompt', { text: 'Hello' });
  assert.equal(ui.messages.at(-1).content[0].text, 'Still available');
  assert.ok(h.events.some(event => event.type === 'persistence_error' && /in memory.*Storage quota/.test(event.error)));
  assert.equal(h.controller.history.at(-1).content, 'Still available');
  assert.equal(h.controller.busy, false);
});