maestro/src/engine/agent-loop-console.test.ts

/**
 * Phase 9 / Task 24: job-crossing screen injection.
 *
 * Phase 4 unit tests in agent-loop.test.ts already verify the *one-shot*
 * `buildSystemPrompt` console injection logic. This file is a focused
 * regression for the multi-iteration / job-crossing property: when the
 * orchestrator runs multiple ReAct iterations (or even multiple jobs
 * within the same local task), each new system prompt must re-read the
 * live PTY screen rather than caching a stale snapshot. The plan
 * (docs/superpowers/plans/2026-05-13-ssh-console.md §9.2) calls this out
 * because the orchestrator's auto-context-trim path can rebuild prompts
 * mid-task without an explicit Send/Snapshot tool call.
 */
import { describe, it, expect, afterEach } from 'vitest';
import { buildSystemPrompt, __setActiveSessionLookup, type Movement } from './agent-loop.js';

function makeConsoleMovement(allowedTools: string[], allowedSshConnections: string[] = ['*']): Movement {
  return {
    name: 'm',
    edit: false,
    persona: 'p',
    instruction: 'i',
    allowedTools,
    allowedSshConnections,
    rules: [{ condition: 'done', next: 'COMPLETE' }],
    defaultNext: 'COMPLETE',
  };
}

describe('console session lookup across jobs', () => {
  afterEach(() => {
    __setActiveSessionLookup(null);
  });

  it('agent-loop reads injected screen on every iteration for the same task', () => {
    const fakeSession = {
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      snapshotScreen: () => ({
        cols: 80,
        rows: 24,
        text: '$ pwd\n/var/log\n$ ',
        cursor: { x: 0, y: 0 },
      }),
    };
    __setActiveSessionLookup((_t: string) => [fakeSession]);

    // Build the prompt twice (simulating two ReAct iterations against the
    // same task: e.g. the engine reruns buildSystemPrompt after a
    // context-window trim, or for a fresh movement visit on revisit). Both
    // prompts must contain the live screen tail.
    const p1 = buildSystemPrompt(
      makeConsoleMovement(['SshConsoleSend']),
      1, // visitCount
      5, // maxVisits
      [], // tools
      null, // missionBrief
      undefined, // userId
      undefined, // userFolderRoot
      undefined, // workspacePath
      't1', // taskId
    );
    const p2 = buildSystemPrompt(
      makeConsoleMovement(['SshConsoleSend']),
      2,
      5,
      [],
      null,
      undefined,
      undefined,
      undefined,
      't1',
    );
    expect(p1).toContain('/var/log');
    expect(p2).toContain('/var/log');
    expect(p1).toContain('Console screen');
    expect(p2).toContain('Console screen');
  });

  it('does NOT inject the console screen when the movement declared no SSH connections (shared_tools leak guard)', () => {
    // Regression: a console tool can reach allowedTools via a piece-level
    // shared_tools union. Injecting live PTY content into a movement that
    // never declared `allowed_ssh_connections` would leak terminal output.
    __setActiveSessionLookup((_t: string) => [{
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      snapshotScreen: () => ({ cols: 80, rows: 24, text: 'secret-host$ cat /etc/shadow\n', cursor: { x: 0, y: 0 } }),
    }]);
    const args = [1, 5, [], null, undefined, undefined, undefined, 't1'] as const;
    // No connections declared (empty array = explicit deny) → no injection.
    const denied = buildSystemPrompt(makeConsoleMovement(['SshConsoleSnapshot'], []), ...args);
    // Field omitted entirely → no injection. Built inline so the helper's
    // default ['*'] does not apply.
    const omittedMovement: Movement = {
      name: 'm', edit: false, persona: 'p', instruction: 'i',
      allowedTools: ['SshConsoleSnapshot'],
      rules: [{ condition: 'done', next: 'COMPLETE' }], defaultNext: 'COMPLETE',
    };
    const omitted = buildSystemPrompt(omittedMovement, ...args);
    // Declared connection → injection still works.
    const allowed = buildSystemPrompt(makeConsoleMovement(['SshConsoleSnapshot'], ['*']), ...args);
    expect(denied).not.toContain('secret-host');
    expect(denied).not.toContain('Console screen');
    expect(omitted).not.toContain('secret-host');
    expect(allowed).toContain('secret-host');
  });

  it('updates the injected screen when the session text changes between calls', () => {
    // Simulates a real shell session: the AI's first iteration sees the
    // login banner; a subsequent SshConsoleSend mutates the screen; the
    // *next* iteration must see the post-Send screen, not the cached one.
    let screen = '$ ';
    __setActiveSessionLookup((_t: string) => [{
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      snapshotScreen: () => ({ cols: 80, rows: 24, text: screen, cursor: { x: 0, y: 0 } }),
    }]);

    const before = buildSystemPrompt(
      makeConsoleMovement(['SshConsoleSnapshot']),
      1,
      5,
      [],
      null,
      undefined,
      undefined,
      undefined,
      't1',
    );
    expect(before).toContain('Console screen');
    expect(before).not.toContain('hello-world');

    screen = '$ echo hello-world\nhello-world\n$ ';
    const after = buildSystemPrompt(
      makeConsoleMovement(['SshConsoleSnapshot']),
      2,
      5,
      [],
      null,
      undefined,
      undefined,
      undefined,
      't1',
    );
    expect(after).toContain('hello-world');
  });

  it('lookup receives the taskId so distinct tasks resolve to distinct sessions', () => {
    const lookups: string[] = [];
    __setActiveSessionLookup((tid: string) => {
      lookups.push(tid);
      return [{
        connectionId: 'conn-1',
        cols: 80,
        rows: 24,
        snapshotScreen: () => ({ cols: 80, rows: 24, text: `screen-for-${tid}`, cursor: { x: 0, y: 0 } }),
      }];
    });

    const pA = buildSystemPrompt(
      makeConsoleMovement(['SshConsoleSend']),
      1,
      5,
      [],
      null,
      undefined,
      undefined,
      undefined,
      'task-A',
    );
    const pB = buildSystemPrompt(
      makeConsoleMovement(['SshConsoleSend']),
      1,
      5,
      [],
      null,
      undefined,
      undefined,
      undefined,
      'task-B',
    );
    expect(lookups).toEqual(['task-A', 'task-B']);
    expect(pA).toContain('screen-for-task-A');
    expect(pB).toContain('screen-for-task-B');
    expect(pA).not.toContain('screen-for-task-B');
    expect(pB).not.toContain('screen-for-task-A');
  });
});