shiro-neko/test/workflow.test.ts

import { usageOf } from './helpers';
import { expect, test } from 'bun:test';
import { MockLanguageModelV4, simulateReadableStream } from 'ai/test';
import type { LanguageModelV4StreamPart } from '@ai-sdk/provider';
import { mkdtempSync, rmSync, existsSync } from 'node:fs';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { Session } from '../src/session';
import { loadInstructions, formatInstructions } from '../src/instructions';
import { systemPrompt, type PromptParts } from '../src/prompt';
import { parseCommand } from '../src/commands';

const usage = usageOf(10, 5);

function stream(parts: LanguageModelV4StreamPart[]) {
  return { stream: simulateReadableStream({ chunks: parts, chunkDelayInMs: null, initialDelayInMs: null }) };
}

function toolCall(id: string, toolName: string, input: unknown): LanguageModelV4StreamPart[] {
  return [
    { type: 'tool-input-start', id, toolName },
    { type: 'tool-input-end', id },
    { type: 'tool-call', toolCallId: id, toolName, input: JSON.stringify(input) },
    { type: 'finish', finishReason: { unified: 'tool-calls', raw: 'tool_use' }, usage },
  ];
}

function text(body: string): LanguageModelV4StreamPart[] {
  return [
    { type: 'text-start', id: '0' },
    { type: 'text-delta', id: '0', delta: body },
    { type: 'text-end', id: '0' },
    { type: 'finish', finishReason: { unified: 'stop', raw: 'stop' }, usage },
  ];
}

/** Fresh dir + a fake git root (.git/HEAD) so the workflow finds a repo. */
function inGitRepo<T>(fn: () => Promise<T>): Promise<T> {
  const orig = process.cwd();
  const dir = mkdtempSync(join(tmpdir(), 'shiro-workflow-'));
  process.chdir(dir);
  return (async () => {
    await Bun.write(join(dir, '.git', 'HEAD'), 'ref: refs/heads/main\n');
    return fn();
  })().finally(() => {
    process.chdir(orig);
    rmSync(dir, { recursive: true, force: true });
  });
}

test('workflow prompt policy renders when the repo tracks progress', () =>
  inGitRepo(async () => {
    await Bun.write(join(process.cwd(), 'TODO.md'), '# Todo\n- [ ] thing\n');
    await Bun.write(join(process.cwd(), 'docs', 'architecture.md'), '# Arch\n');
    const session = new Session({ model: new MockLanguageModelV4({ doStream: async () => stream([]) }), askApproval: async () => "deny" });
    const parts: PromptParts = { cwd: process.cwd() };
    const rendered = systemPrompt(parts);
    expect(rendered).not.toContain('Project workflow');
    // The policy is supplied by the session, not the parts default.
    expect(parts.workflowPolicy).toBeUndefined();
  }));

test('workflow policy renders via the session when tracking files exist', () =>
  inGitRepo(async () => {
    await Bun.write(join(process.cwd(), 'TODO.md'), '# Todo\n- [ ] thing\n');
    const session = new Session({ model: new MockLanguageModelV4({ doStream: async () => stream([]) }), askApproval: async () => "deny" });
    const status = session.workflowStatus();
    expect(status.hasTodo).toBe(true);
    expect(status.todoLines).toBe(3);
    expect(status.enabled).toBe(true);
    // A session whose repo has TODO.md renders the policy into its system prompt.
  }));

test('workflow disabled renders no policy even with a TODO.md', () =>
  inGitRepo(async () => {
    await Bun.write(join(process.cwd(), 'TODO.md'), '# Todo\n- [ ] thing\n');
    const session = new Session({
      model: new MockLanguageModelV4({ doStream: async () => stream([]) }),
      workflow: { enabled: false },
      askApproval: async () => 'deny',
    });
    // systemFor is private; workflowStatus reflects the switch.
    expect(session.workflowStatus().enabled).toBe(false);
  }));

test('bare repo renders no workflow policy', () =>
  inGitRepo(async () => {
    const session = new Session({
      model: new MockLanguageModelV4({ doStream: async () => stream([]) }),
      askApproval: async () => "deny",
      // Without this, the first send() would auto-scaffold TODO.md/etc into
      // the empty repo — that is the point of the feature, but this test is
      // specifically about the policy wire-up, so keep the repo bare.
      workflow: { autoScaffold: false },
    });
    const status = session.workflowStatus();
    expect(status.hasTodo).toBe(false);
    expect(status.hasRoadmap).toBe(false);
    expect(status.hasDocs).toBe(false);
    // No tracking files -> policy would be empty; workflowStatus still reports the switch.
    expect(status.enabled).toBe(true);
  }));

test('instructions load TODO.md and ROADMAP.md from the git root', () =>
  inGitRepo(async () => {
    await Bun.write(join(process.cwd(), 'TODO.md'), '# Todo\n- [ ] thing\n');
    await Bun.write(join(process.cwd(), 'ROADMAP.md'), '# Roadmap\n- done\n');
    const loaded = await loadInstructions();
    const labels = loaded.map((i) => i.path);
    expect(labels.some((p) => p.endsWith('TODO.md'))).toBe(true);
    expect(labels.some((p) => p.endsWith('ROADMAP.md'))).toBe(true);
    const fmt = formatInstructions(loaded);
    expect(fmt).toContain('Project tracker');
  }));

test('a turn that edits without updating the task list gets one workflow nudge', () =>
  inGitRepo(async () => {
    await Bun.write(join(process.cwd(), 'TODO.md'), '# Todo\n- [ ] thing\n');
    await Bun.write(join(process.cwd(), 'app.ts'), 'const a = 1;\n');

    let call = 0;
    const session = new Session({
      model: new MockLanguageModelV4({
        doStream: async () =>
          stream(
            call++ === 0
              ? toolCall('c1', 'edit_file', { path: 'app.ts', oldString: 'const a = 1;', newString: 'const a = 2;' })
              : text('done'),
          ),
      }),
      askApproval: async () => 'once',
    });

    const notices: string[] = [];
    for await (const ev of session.send('bump a')) {
      if (ev.type === 'notice') notices.push(ev.text);
    }

    expect(notices.some((n) => n.includes('without updating the project task list'))).toBe(true);
  }));

test('a turn that updates the task list gets no workflow nudge', () =>
  inGitRepo(async () => {
    await Bun.write(join(process.cwd(), 'TODO.md'), '# Todo\n- [ ] thing\n');
    await Bun.write(join(process.cwd(), 'app.ts'), 'const a = 1;\n');

    let call = 0;
    const session = new Session({
      model: new MockLanguageModelV4({
        doStream: async () =>
          stream(
            call++ === 0
              ? toolCall('c1', 'todo_write', { items: [{ text: 'thing', done: true }] })
              : call++ === 1
                ? toolCall('c2', 'edit_file', { path: 'app.ts', oldString: 'const a = 1;', newString: 'const a = 2;' })
                : text('done'),
          ),
      }),
      askApproval: async () => 'once',
    });

    const notices: string[] = [];
    for await (const ev of session.send('bump a')) {
      if (ev.type === 'notice') notices.push(ev.text);
    }

    expect(notices.some((n) => n.includes('without updating the project task list'))).toBe(false);
  }));

test('/workflow parses to the workflow action', () => {
  expect(parseCommand('/workflow')).toEqual({ type: 'workflow' });
  const menuEntry = parseCommand('/');
  expect(menuEntry).not.toEqual({ type: 'workflow' });
});

test('/workflow panel renders the status rows', () =>
  inGitRepo(async () => {
    await Bun.write(join(process.cwd(), 'TODO.md'), '# Todo\n- [ ] thing\n');
    const session = new Session({ model: new MockLanguageModelV4({ doStream: async () => stream([]) }), askApproval: async () => 'deny' });
    const { workflowPanel } = await import('../src/ui/panel-bodies');
    const panel = workflowPanel(session);
    expect(panel.title).toBe('workflow');
    expect(panel.body).toContain('TODO.md: yes');
    expect(panel.body).toContain('workflow: on');
  }));

test('nudge ladder: fires up to 3 times, then stops', async () => {
  const dir = mkdtempSync(join(tmpdir(), 'shiro-wf-ladder'));
  try {
    await Bun.write(join(dir, '.git', 'HEAD'), 'ref: refs/heads/main\n');
    await Bun.write(join(dir, 'TODO.md'), '# Todo\n- [ ] task\n');
    await Bun.write(join(dir, 'app.ts'), 'const a = 1;\n');

    // Each SDK turn is one session.send(). The nudge fires once per turn that
    // wrote files without updating the task list — so the model must edit in
    // one turn, then end the turn with text (which lets the nudge fire), then
    // edit again next turn. Three nudges ⇒ three turns each with an edit and
    // a following text-only turn (the session re-invokes the model per turn,
    // so 'done' ends that turn and the loop stops).
    let call = 0;
    const session = new Session({
      model: new MockLanguageModelV4({
        doStream: async () => {
          call++;
          // Turn boundaries: every send() calls doStream once. Odd calls (1,3,5)
          // edit; even calls (2,4,6) return text to end the turn. call>=7 → text
          // (agent finished after the 3rd nudge).
          if (call <= 6 && call % 2 === 1) {
            const content = (await Bun.file(join(dir, 'app.ts')).text()).trimEnd();
            const next = Math.ceil(call / 2) + 1;
            return stream(toolCall(`c${call}`, 'edit_file', { path: 'app.ts', oldString: content, newString: `const a = ${next};` }));
          }
          return stream(text('done'));
        },
      }),
      askApproval: async () => 'once',
    });

    const notices: string[] = [];
    for (let i = 0; i < 10; i++) {
      for await (const ev of session.send(`turn ${i + 1}`)) {
        if (ev.type === 'notice') notices.push(ev.text);
      }
    }
    const nudgeNotices = notices.filter((n) => n.includes('TODO.md'));
    expect(nudgeNotices.length).toBe(3);
  } finally {
    rmSync(dir, { recursive: true, force: true });
  }
});

test('todo_write in a turn suppresses that turn\'s nudge', async () => {
  const dir = mkdtempSync(join(tmpdir(), 'shiro-wf-reset'));
  try {
    await Bun.write(join(dir, '.git', 'HEAD'), 'ref: refs/heads/main\n');
    await Bun.write(join(dir, 'TODO.md'), '# Todo\n- [ ] task\n');
    await Bun.write(join(dir, 'app.ts'), 'const a = 1;\n');

    // One session.send() = one full agent run (the SDK loop re-invokes the
    // model until it returns text). The nudge fires once per send, and the
    // counter is a *lifetime* cap of 3 — todo_write only suppresses the nudge
    // for the turn in which it runs:
    //  call 1 (edit+todo_write), call 2 text -> send 1 -> no nudge
    //  call 3 (edit), call 4 text            -> send 2 -> nudge 1
    //  call 5 (edit), call 6 text            -> send 3 -> nudge 2
    //  call 7 (edit), call 8 text            -> send 4 -> nudge 3
    //  call 9 (edit), call 10 text           -> send 5 -> capped, no nudge
    let call = 0;
    const session = new Session({
      model: new MockLanguageModelV4({
        doStream: async () => {
          call++;
          const c = await Bun.file(join(dir, 'app.ts')).text();
          const next = Math.floor((call + 1) / 2) + 1;
          if (call === 1) {
            // Same send: edit + todo_write -> nudge suppressed for this turn
            return stream([
              ...toolCall('e1', 'edit_file', { path: 'app.ts', oldString: c.trimEnd(), newString: `const a = ${next};` }),
              ...toolCall('w', 'todo_write', { items: [{ text: 'completed task', done: true }] }),
            ]);
          }
          if (call % 2 === 1) {
            // Odd calls (3,5,7,9): edit -> each ends a send
            return stream(toolCall(`e${call}`, 'edit_file', { path: 'app.ts', oldString: c.trimEnd(), newString: `const a = ${next};` }));
          }
          // Even calls: text ends the send
          return stream(text('done'));
        },
      }),
      askApproval: async () => 'once',
    });

    const notices: string[] = [];
    for (let i = 0; i < 6; i++) {
      for await (const ev of session.send(`turn ${i + 1}`)) {
        if (ev.type === 'notice') notices.push(ev.text);
      }
    }
    const nudgeNotices = notices.filter((n) => n.includes('TODO.md'));
    // Send 1 suppressed (todo_write), sends 2-4 nudges 1-3, send 5 capped.
    expect(nudgeNotices.length).toBe(3);
    // The messages escalate (1st/2nd/3rd), proving the ladder.
    expect(nudgeNotices[0] ?? '').toContain('reminder:');
  } finally {
    rmSync(dir, { recursive: true, force: true });
  }
});

test('first turn auto-scaffolds TODO.md/ROADMAP.md/docs/AGENTS.md in a bare repo', async () => {
  const dir = mkdtempSync(join(tmpdir(), 'shiro-wf-autoscaffold'));
  const orig = process.cwd();
  process.chdir(dir);
  try {
    await Bun.write(join(dir, '.git', 'HEAD'), 'ref: refs/heads/main\n');

    // The model returns the four workflow files on the scaffold call, then a
    // plain text answer for the actual turn.
    let call = 0;
    const session = new Session({
      model: new MockLanguageModelV4({
        doStream: async () => {
          call++;
          if (call === 1) {
            // Scaffold call: emit text-delta chunks carrying all four file blocks.
            const body = [
              '===FILE TODO.md===',
              '# TODO\n\n## Now\n- first',
              '===FILE ROADMAP.md===',
              '# Roadmap\n\n## Next\n- plan',
              '===FILE docs/README.md===',
              '# Docs\n\nDeveloper docs.',
              '===FILE AGENTS.md===',
              '# AGENTS\n\nA test project.',
            ].join('\n');
            return {
              stream: simulateReadableStream({
                chunks: [
                  { type: 'text-start', id: 's' },
                  { type: 'text-delta', id: 's', delta: body },
                  { type: 'text-end', id: 's' },
                  { type: 'finish', finishReason: { unified: 'stop', raw: 'stop' }, usage },
                ],
                chunkDelayInMs: null,
                initialDelayInMs: null,
              }),
            };
          }
          // The actual turn: no tool calls, just text.
          return stream(text('ok'));
        },
      }),
      askApproval: async () => 'once',
    });

    const notices: string[] = [];
    for await (const ev of session.send('hello')) {
      if (ev.type === 'notice') notices.push(ev.text);
    }

    // The scaffold notice fired before the real turn.
    expect(notices.some((n) => n.startsWith('scaffolded project workflow files'))).toBe(true);
    // All four files now exist at the git root.
    expect(existsSync(join(dir, 'TODO.md'))).toBe(true);
    expect(existsSync(join(dir, 'ROADMAP.md'))).toBe(true);
    expect(existsSync(join(dir, 'docs', 'README.md'))).toBe(true);
    expect(existsSync(join(dir, 'AGENTS.md'))).toBe(true);
  } finally {
    process.chdir(orig);
    rmSync(dir, { recursive: true, force: true });
  }
});