stack/packages/mosaic/src/fleet/tmux-runtime-transport.test.ts

import { describe, expect, it, vi } from 'vitest';
import type { CommandResult, CommandRunner, FleetRoster } from '../commands/fleet.js';
import { FleetTmuxRuntimeTransport } from './tmux-runtime-transport.js';
import type { FleetRuntimeTransportError } from './tmux-runtime-transport.js';

const roster: FleetRoster = {
  version: 1,
  transport: 'tmux',
  tmux: { socketName: 'tess-fleet', holderSession: '_holder' },
  defaults: { workingDirectory: '~/src' },
  runtimes: { codex: { resetCommand: '/clear' } },
  agents: [{ name: 'coder0', runtime: 'codex', className: 'code' }],
};

function commandResult(stdout = '', exitCode = 0, stderr = ''): CommandResult {
  return { stdout, stderr, exitCode };
}

describe('FleetTmuxRuntimeTransport security boundary', (): void => {
  it('rejects a prefix target before it invokes tmux', async (): Promise<void> => {
    const runner = vi.fn<CommandRunner>(async (): Promise<CommandResult> => commandResult());
    const transport = new FleetTmuxRuntimeTransport({
      rosterLoader: async (): Promise<FleetRoster> => roster,
      runner,
      mosaicHome: '/mosaic',
    });

    await expect(transport.verifySession('coder')).rejects.toMatchObject({
      code: 'not_found',
    } satisfies Partial<FleetRuntimeTransportError>);
    expect(runner).not.toHaveBeenCalled();
  });

  it('uses the roster socket and exact pane target while verifying peer identity', async (): Promise<void> => {
    const runner = vi.fn<CommandRunner>(
      async (): Promise<CommandResult> => commandResult('111 codex 0 0 0 0\n'),
    );
    const transport = new FleetTmuxRuntimeTransport({
      rosterLoader: async (): Promise<FleetRoster> => roster,
      runner,
      mosaicHome: '/mosaic',
    });

    await expect(transport.verifySession('coder0')).resolves.toEqual({
      id: 'coder0',
      runtimeId: 'codex',
      socketName: 'tess-fleet',
    });
    expect(runner).toHaveBeenCalledWith('tmux', [
      '-L',
      'tess-fleet',
      'list-panes',
      '-t',
      '=coder0:0.0',
      '-F',
      '#{pane_pid} #{pane_current_command} #{pane_dead} #{pane_activity} #{window_activity} #{session_activity}',
    ]);
  });

  it('denies a live pane whose runtime does not match the roster identity', async (): Promise<void> => {
    const runner = vi.fn<CommandRunner>(
      async (): Promise<CommandResult> => commandResult('111 python3 0 0 0 0\n'),
    );
    const transport = new FleetTmuxRuntimeTransport({
      rosterLoader: async (): Promise<FleetRoster> => roster,
      runner,
      mosaicHome: '/mosaic',
    });

    await expect(transport.verifySession('coder0')).rejects.toMatchObject({
      code: 'forbidden',
    } satisfies Partial<FleetRuntimeTransportError>);
    expect(runner).toHaveBeenCalledTimes(1);
  });

  it('surfaces a runtime identity mismatch instead of hiding it from fleet status', async (): Promise<void> => {
    const rosterWithMismatchedPeer: FleetRoster = {
      ...roster,
      agents: [...roster.agents, { name: 'coder1', runtime: 'codex', className: 'code' }],
    };
    const runner = vi.fn<CommandRunner>(
      async (_command: string, args: string[]): Promise<CommandResult> =>
        commandResult(
          args.includes('=coder1:0.0') ? '222 python3 0 0 0 0\n' : '111 codex 0 0 0 0\n',
        ),
    );
    const transport = new FleetTmuxRuntimeTransport({
      rosterLoader: async (): Promise<FleetRoster> => rosterWithMismatchedPeer,
      runner,
      mosaicHome: '/mosaic',
    });

    await expect(transport.listSessions()).rejects.toMatchObject({
      code: 'forbidden',
    } satisfies Partial<FleetRuntimeTransportError>);
  });

  it('denies sending when the roster socket does not contain the exact target', async (): Promise<void> => {
    const runner = vi.fn<CommandRunner>(
      async (): Promise<CommandResult> => commandResult('', 1, "can't find session: coder0"),
    );
    const transport = new FleetTmuxRuntimeTransport({
      rosterLoader: async (): Promise<FleetRoster> => roster,
      runner,
      mosaicHome: '/mosaic',
    });

    await expect(transport.sendMessage('coder0', 'hello', 'tess')).rejects.toMatchObject({
      code: 'unavailable',
    } satisfies Partial<FleetRuntimeTransportError>);
    expect(runner).toHaveBeenCalledTimes(1);
  });

  it('sends only through the maintained sender after exact target and identity verification', async (): Promise<void> => {
    const runner = vi
      .fn<CommandRunner>()
      .mockResolvedValueOnce(commandResult('111 codex 0 0 0 0\n'))
      .mockResolvedValueOnce(commandResult());
    const transport = new FleetTmuxRuntimeTransport({
      rosterLoader: async (): Promise<FleetRoster> => roster,
      runner,
      mosaicHome: '/mosaic',
    });

    await transport.sendMessage('coder0', 'hello', 'tess');

    expect(runner).toHaveBeenNthCalledWith(2, '/mosaic/tools/tmux/agent-send.sh', [
      '-L',
      'tess-fleet',
      '-S',
      'tess',
      '-s',
      'coder0',
      '-m',
      'hello',
    ]);
  });

  it('terminates only the exact roster target after identity verification', async (): Promise<void> => {
    const runner = vi
      .fn<CommandRunner>()
      .mockResolvedValueOnce(commandResult('111 codex 0 0 0 0\n'))
      .mockResolvedValueOnce(commandResult());
    const transport = new FleetTmuxRuntimeTransport({
      rosterLoader: async (): Promise<FleetRoster> => roster,
      runner,
      mosaicHome: '/mosaic',
    });

    await transport.terminate('coder0');

    expect(runner).toHaveBeenNthCalledWith(2, 'tmux', [
      '-L',
      'tess-fleet',
      'kill-session',
      '-t',
      '=coder0',
    ]);
  });
});