/**
 * Unit tests for the SSH Console tools (SshConsoleEnsure / Send / Snapshot).
 *
 * Strategy: stub the full SshSubsystem rather than booting repos — these
 * tests focus on the orchestration logic (find-or-open, deny check,
 * snapshot routing). The 12-step preflight is already covered by the
 * SshExec tests in ssh.test.ts.
 */
import { describe, it, expect, beforeEach, vi } from 'vitest';
import {
  executeTool,
  TOOL_DEFS,
  unescapeAiInput,
  openConsoleSession,
  resolveTargetSession,
  __resetAgentConnectionPointer,
  type OpenConsoleParams,
} from './ssh-console.js';
import { setSshSubsystem, preflight, type SshSubsystem } from './ssh.js';
import type { ToolContext } from './core.js';

describe('unescapeAiInput', () => {
  it('passes real LF through unchanged', () => {
    expect(unescapeAiInput('ls -la\n')).toBe('ls -la\n');
  });

  it('converts literal 2-char "\\\\n" to real LF', () => {
    // source '\\n' = 2 chars (backslash + n); should become real LF.
    expect(unescapeAiInput('ls -la\\n')).toBe('ls -la\n');
  });

  it('converts literal 2-char "\\\\r" to real CR', () => {
    expect(unescapeAiInput('uptime\\r')).toBe('uptime\r');
  });

  it('converts literal "\\\\t" and "\\\\0"', () => {
    expect(unescapeAiInput('a\\tb')).toBe('a\tb');
    expect(unescapeAiInput('a\\0b')).toBe('a\0b');
  });

  it('converts \\xHH hex escapes', () => {
    expect(unescapeAiInput('\\x03')).toBe('\x03');           // Ctrl-C
    expect(unescapeAiInput('\\x1b:q!\\n')).toBe('\x1b:q!\n'); // Esc + vim exit
  });

  it('handles double-backslash correctly', () => {
    // source '\\\\' = 2 chars (\\); should become single backslash.
    expect(unescapeAiInput('a\\\\b')).toBe('a\\b');
  });

  it('preserves unknown escapes literally', () => {
    expect(unescapeAiInput('foo\\q')).toBe('foo\\q');
  });

  it('does not touch a string with no backslashes', () => {
    expect(unescapeAiInput('hello world')).toBe('hello world');
    expect(unescapeAiInput('')).toBe('');
  });
});

function mkConn(overrides: Partial<{
  commandDenyPatterns: string | null;
  commandAllowPatterns: string | null;
}> = {}) {
  return {
    id: 'conn-1',
    ownerId: 'u1',
    label: 'test',
    host: 'localhost',
    port: 22,
    username: 'me',
    privateKeyEnc: Buffer.alloc(0),
    passphraseEnc: null,
    keyVersion: 1,
    keyFingerprint: 'fp-key',
    hostKeyType: 'ssh-ed25519',
    hostKeyB64: 'aaa',
    hostKeyFingerprint: 'fp',
    hostKeyRecordedAt: '2026-01-01',
    hostKeyVerifiedAt: '2026-01-01',
    hostKeyPending: false,
    hostKeyPendingB64: null,
    hostKeyPendingFingerprint: null,
    hostKeyPendingToken: null,
    hostKeyPendingSource: null,
    commandDenyPatterns: overrides.commandDenyPatterns ?? null,
    commandAllowPatterns: overrides.commandAllowPatterns ?? null,
    remotePathPrefix: '/',
    allowRemoteUnrestricted: true,
    allowPrivateAddresses: true,
    enabled: true,
    disabledByAdmin: false,
    disabledByAdminReason: null,
    disabledByAdminAt: null,
    disabledByAdminUserId: null,
    createdAt: '2026-01-01',
    updatedAt: '2026-01-01',
  };
}

function mkStubSubsystem() {
  const audit = {
    beginAndComplete: vi.fn().mockReturnValue(1),
    begin: vi.fn().mockReturnValue(1),
    complete: vi.fn(),
    listAuditRows: vi.fn(),
    pruneOlderThan: vi.fn(),
    promotePendingToAborted: vi.fn(),
  };
  const registry = {
    get: vi.fn().mockReturnValue(null),
    register: vi.fn(),
    enforceCap: vi.fn().mockReturnValue([]),
    closeForTask: vi.fn().mockResolvedValue(undefined),
    closeSession: vi.fn().mockResolvedValue(undefined),
    listForTask: vi.fn().mockReturnValue([]),
    listAll: vi.fn().mockReturnValue([]),
    listForConnection: vi.fn().mockReturnValue([]),
    sweep: vi.fn(),
    startSweepTimer: vi.fn(),
    stopSweepTimer: vi.fn(),
    shutdown: vi.fn(),
  };
  const connectionRepo = {
    resolveConnection: vi.fn().mockReturnValue(mkConn()),
  };
  const abuseRepo = {
    isLocked: vi.fn().mockReturnValue({ locked: false }),
    checkAndRecordFailure: vi.fn(),
    recordSuccess: vi.fn(),
  };
  const accessResolver = { resolveAccess: vi.fn().mockReturnValue({ allowed: true }) };
  const channel = {
    write: vi.fn(),
    end: vi.fn(),
    setWindow: vi.fn(),
    on: vi.fn(),
  };
  const client = { end: vi.fn(), on: vi.fn() };
  const openShellChannel = vi.fn().mockResolvedValue({
    channel,
    client,
    hostFingerprint: 'SHA256:fake',
  });
  const sub = {
    connectionRepo,
    auditRepo: audit,
    abuseRepo,
    accessResolver,
    sessionRegistry: registry,
    openShellChannel,
    getUserAccess: () => ({ isAdmin: false, orgIds: [] }),
    decryptKeyMaterial: () => Buffer.alloc(0),
    decryptPassphrase: () => null,
    sshExec: vi.fn(),
    sshUpload: vi.fn(),
    sshDownload: vi.fn(),
    maintenance: { isActive: () => false, snapshot: () => ({ active: false }), enter: () => {}, exit: () => {} } as SshSubsystem['maintenance'],
    config: {
      enabled: true,
      allowPrivateAddresses: true,
      callTimeoutSeconds: 30,
      maxOutputBytes: 1024,
      maxUploadSizeMb: 10,
      maxDownloadSizeMb: 10,
      auditRetentionDays: 90,
      adminBypassesGrants: true,
      abuseWindowMinutes: 10,
      abuseFailureThreshold: 5,
      abuseLockMinutes: 30,
      console: {
        enabled: true,
        idleTimeoutSeconds: 60,
        maxSessionDurationSeconds: 600,
        scrollbackBytes: 4096,
        maxSessionsPerConnection: 3,
        maxSessionsPerTask: 5,
        maxSessionsPerUser: 20,
        maxInputBytesPerSend: 1024,
        autoInjectScreenLines: 24,
        defaultCols: 80,
        defaultRows: 24,
      },
    },
  } as unknown as SshSubsystem;
  return { sub, audit, registry, openShellChannel, connectionRepo, channel };
}

function mkCtx(overrides: Partial<ToolContext> = {}): ToolContext {
  return {
    workspacePath: '/tmp',
    editAllowed: true,
    taskId: 'task-1',
    userId: 'u1',
    ownerId: 'u1',
    jobId: 'j1',
    pieceName: 'p',
    allowedSshConnections: ['*'],
    ...overrides,
  };
}

describe('SshConsoleEnsure', () => {
  beforeEach(() => {
    setSshSubsystem(null);
    __resetAgentConnectionPointer();
  });

  it('is registered in TOOL_DEFS', () => {
    expect(TOOL_DEFS.SshConsoleEnsure).toBeDefined();
    expect(TOOL_DEFS.SshConsoleSend).toBeDefined();
    expect(TOOL_DEFS.SshConsoleSnapshot).toBeDefined();
  });

  it('opens new session when none exists', async () => {
    const { sub, registry, openShellChannel } = mkStubSubsystem();
    setSshSubsystem(sub);
    const res = await executeTool('SshConsoleEnsure', { connection_id: 'conn-1' }, mkCtx());
    expect(res?.isError).toBe(false);
    expect(openShellChannel).toHaveBeenCalled();
    expect(registry.register).toHaveBeenCalled();
  });

  it('reuses existing session for same task + connection', async () => {
    const { sub, registry, openShellChannel } = mkStubSubsystem();
    const existing = {
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      isClosed: false,
    };
    registry.get.mockReturnValue(existing);
    setSshSubsystem(sub);
    const res = await executeTool('SshConsoleEnsure', { connection_id: 'conn-1' }, mkCtx());
    expect(res?.isError).toBe(false);
    expect(openShellChannel).not.toHaveBeenCalled();
  });

  it('opening a DIFFERENT connection adds a session and does NOT close the existing one', async () => {
    // Task 3: add-not-replace. conn-OLD is active; opening conn-NEW just
    // opens a second session. No force_replace needed, nothing closed.
    const { sub, registry, openShellChannel } = mkStubSubsystem();
    const old = {
      localTaskId: 'task-1',
      connectionId: 'conn-OLD',
      isClosed: false,
      startedByUserId: 'u1',
      startedAt: Date.now() - 60_000,
      lastActivityAt: Date.now() - 5_000,
    };
    // Key-aware: only conn-OLD has an existing session.
    registry.get.mockImplementation((_t: string, c: string) => (c === 'conn-OLD' ? old : null));
    registry.listForTask.mockReturnValue([old]);
    setSshSubsystem(sub);
    const res = await executeTool('SshConsoleEnsure', { connection_id: 'conn-NEW' }, mkCtx());
    expect(res?.isError).toBe(false);
    expect(openShellChannel).toHaveBeenCalledTimes(1);
    expect(registry.closeForTask).not.toHaveBeenCalled();
    expect(registry.closeSession).not.toHaveBeenCalled();
  });

  it('force_replace on the SAME connection closes just that connection then reopens', async () => {
    const { sub, registry, openShellChannel } = mkStubSubsystem();
    const existing = {
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      isClosed: false,
      startedByUserId: 'u1',
      startedAt: Date.now() - 60_000,
      lastActivityAt: Date.now() - 5_000,
    };
    registry.get.mockImplementation((_t: string, c: string) => (c === 'conn-1' ? existing : null));
    setSshSubsystem(sub);
    const res = await executeTool(
      'SshConsoleEnsure',
      { connection_id: 'conn-1', force_replace: true },
      mkCtx(),
    );
    expect(registry.closeSession).toHaveBeenCalledWith('task-1', 'conn-1', 'connection_change');
    expect(registry.closeForTask).not.toHaveBeenCalled();
    expect(openShellChannel).toHaveBeenCalledTimes(1);
    expect(res?.isError).toBe(false);
  });

  it('rejects when console.enabled is false', async () => {
    const { sub } = mkStubSubsystem();
    sub.config.console.enabled = false;
    setSshSubsystem(sub);
    const res = await executeTool('SshConsoleEnsure', { connection_id: 'conn-1' }, mkCtx());
    expect(res?.isError).toBe(true);
  });
});

// ──────────────────────────────────────────────────────────────────────
// Task 3: openConsoleSession add-not-replace + task/user caps + TOCTOU lock
// ──────────────────────────────────────────────────────────────────────

/** A minimal stateful registry mirroring SessionRegistry's (task, conn) keying
 *  so we can assert listForTask length, add-not-replace, and concurrency. */
function mkStatefulRegistry(opts: { maxSessionsPerConnection?: number } = {}) {
  const byTask = new Map<string, Map<string, any>>();
  const reg = {
    get: vi.fn((taskId: string, connId: string) => {
      const s = byTask.get(taskId)?.get(connId);
      return s && !s.isClosed ? s : null;
    }),
    listForTask: vi.fn((taskId: string) =>
      [...(byTask.get(taskId)?.values() ?? [])].filter((s: any) => !s.isClosed)),
    listAll: vi.fn(() => {
      const all: any[] = [];
      for (const inner of byTask.values()) for (const s of inner.values()) if (!s.isClosed) all.push(s);
      return all;
    }),
    register: vi.fn((session: any) => {
      let inner = byTask.get(session.localTaskId);
      if (!inner) { inner = new Map(); byTask.set(session.localTaskId, inner); }
      inner.set(session.connectionId, session);
    }),
    closeSession: vi.fn(async (taskId: string, connId: string) => {
      const inner = byTask.get(taskId);
      if (inner?.has(connId)) {
        inner.delete(connId);
        if (inner.size === 0) byTask.delete(taskId);
      }
    }),
    closeForTask: vi.fn(async () => {}),
    enforceCap: vi.fn((connId: string) => {
      const max = opts.maxSessionsPerConnection ?? 999;
      const sessions: any[] = [];
      for (const inner of byTask.values())
        for (const s of inner.values()) if (!s.isClosed && s.connectionId === connId) sessions.push(s);
      const over = sessions.length - max;
      if (over <= 0) return [];
      sessions.sort((a, b) => a.startedAt - b.startedAt);
      return sessions.slice(0, over);
    }),
    listForConnection: vi.fn(() => []),
    sweep: vi.fn(), startSweepTimer: vi.fn(), stopSweepTimer: vi.fn(), shutdown: vi.fn(),
  };
  return { reg, byTask };
}

function mkPreSession(taskId: string, connId: string, userId = 'u1') {
  const now = Date.now();
  return {
    localTaskId: taskId, connectionId: connId, startedByUserId: userId,
    cols: 80, rows: 24, isClosed: false, startedAt: now, lastActivityAt: now,
  };
}

function mkOpenParams(overrides: Partial<OpenConsoleParams> = {}): OpenConsoleParams {
  return {
    taskId: 'task-1',
    connectionId: 'conn-1',
    ownerId: 'u1',
    userId: 'u1',
    pieceName: 'p',
    allowedConnections: ['*'],
    initiator: 'agent',
    jobId: 'j1',
    spaceId: null,
    ...overrides,
  };
}

describe('openConsoleSession (Task 3 multi-session core)', () => {
  beforeEach(() => {
    setSshSubsystem(null);
    __resetAgentConnectionPointer();
  });

  it('(a) opening connection B while A is active leaves A in listForTask', async () => {
    const { sub } = mkStubSubsystem();
    const { reg, byTask } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;
    // Seed an active session on conn-A.
    reg.register(mkPreSession('task-1', 'conn-A'));

    const res = await openConsoleSession(
      { sub, preflight },
      mkOpenParams({ connectionId: 'conn-B' }),
    );
    expect(res.ok).toBe(true);
    expect(res.alreadyActive).toBe(false);
    const live = byTask.get('task-1')!;
    expect(live.has('conn-A')).toBe(true);
    expect(live.has('conn-B')).toBe(true);
    expect(reg.listForTask('task-1').length).toBe(2);
    expect(reg.closeSession).not.toHaveBeenCalled();
  });

  it('(b) opening the SAME connection twice reuses (alreadyActive)', async () => {
    const { sub, openShellChannel } = mkStubSubsystem();
    const { reg } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;

    const first = await openConsoleSession({ sub, preflight }, mkOpenParams({ connectionId: 'conn-1' }));
    expect(first.ok).toBe(true);
    expect(first.alreadyActive).toBe(false);
    expect(openShellChannel).toHaveBeenCalledTimes(1);

    const second = await openConsoleSession({ sub, preflight }, mkOpenParams({ connectionId: 'conn-1' }));
    expect(second.ok).toBe(true);
    expect(second.alreadyActive).toBe(true);
    // No second channel open — the live session was reused.
    expect(openShellChannel).toHaveBeenCalledTimes(1);
    expect(reg.listForTask('task-1').length).toBe(1);
  });

  it('(c) force_replace on same connection closes only that connection', async () => {
    const { sub, openShellChannel } = mkStubSubsystem();
    const { reg } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;
    reg.register(mkPreSession('task-1', 'conn-1'));
    reg.register(mkPreSession('task-1', 'conn-2')); // a sibling that must survive

    const res = await openConsoleSession(
      { sub, preflight },
      mkOpenParams({ connectionId: 'conn-1', forceReplace: true }),
    );
    expect(res.ok).toBe(true);
    expect(reg.closeSession).toHaveBeenCalledWith('task-1', 'conn-1', 'connection_change');
    // Only conn-1 was closed; conn-2 sibling remains.
    const calls = reg.closeSession.mock.calls.filter((c: any[]) => c[2] === 'connection_change');
    expect(calls.every((c: any[]) => c[1] === 'conn-1')).toBe(true);
    expect(reg.get('task-1', 'conn-2')).not.toBeNull();
    expect(openShellChannel).toHaveBeenCalledTimes(1);
  });

  it('(d) at task cap, a new connection returns task_session_cap and closes nothing', async () => {
    const { sub, openShellChannel } = mkStubSubsystem();
    sub.config.console.maxSessionsPerTask = 1;
    const { reg } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;
    reg.register(mkPreSession('task-1', 'conn-A'));

    const res = await openConsoleSession(
      { sub, preflight },
      mkOpenParams({ connectionId: 'conn-B' }),
    );
    expect(res.ok).toBe(false);
    expect(res.error).toBe('task_session_cap');
    expect(openShellChannel).not.toHaveBeenCalled();
    expect(reg.closeSession).not.toHaveBeenCalled();
    expect(reg.listForTask('task-1').length).toBe(1);
  });

  it('(d2) at per-user cap, a new session returns user_session_cap', async () => {
    const { sub, openShellChannel } = mkStubSubsystem();
    sub.config.console.maxSessionsPerUser = 1;
    const { reg } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;
    // u1 already owns a session on a DIFFERENT task.
    reg.register(mkPreSession('task-OTHER', 'conn-X', 'u1'));

    const res = await openConsoleSession(
      { sub, preflight },
      mkOpenParams({ connectionId: 'conn-1', userId: 'u1' }),
    );
    expect(res.ok).toBe(false);
    expect(res.error).toBe('user_session_cap');
    expect(openShellChannel).not.toHaveBeenCalled();
  });

  it('(d3) per-user cap of 0 means unlimited', async () => {
    const { sub } = mkStubSubsystem();
    sub.config.console.maxSessionsPerUser = 0;
    const { reg } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;
    for (let i = 0; i < 3; i++) reg.register(mkPreSession(`task-${i}`, `conn-${i}`, 'u1'));

    const res = await openConsoleSession(
      { sub, preflight },
      mkOpenParams({ connectionId: 'conn-1' }),
    );
    expect(res.ok).toBe(true);
  });

  it('(e) two concurrent opens for the same new (task,conn) register exactly one session', async () => {
    const { sub, openShellChannel } = mkStubSubsystem();
    const { reg, byTask } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;
    // Make channel-open slow so both calls would interleave without the lock.
    const orig = openShellChannel;
    (sub as any).openShellChannel = vi.fn(async (args: any) => {
      await new Promise((r) => setTimeout(r, 25));
      return orig(args);
    });

    const [r1, r2] = await Promise.all([
      openConsoleSession({ sub, preflight }, mkOpenParams({ connectionId: 'conn-1' })),
      openConsoleSession({ sub, preflight }, mkOpenParams({ connectionId: 'conn-1' })),
    ]);
    expect(r1.ok).toBe(true);
    expect(r2.ok).toBe(true);
    // Exactly one channel opened, exactly one registered session — no orphan.
    expect((sub as any).openShellChannel).toHaveBeenCalledTimes(1);
    expect(byTask.get('task-1')!.size).toBe(1);
    // One of them opened, the other reused.
    expect([r1.alreadyActive, r2.alreadyActive].sort()).toEqual([false, true]);
  });
});

describe('SshConsoleSend', () => {
  beforeEach(() => {
    setSshSubsystem(null);
    __resetAgentConnectionPointer();
  });

  it('writes input to session and returns screen snapshot', async () => {
    const { sub, registry } = mkStubSubsystem();
    const writes: Buffer[] = [];
    const now = Date.now();
    const fakeSession = {
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      isClosed: false,
      idleMs: () => 10_000,
      lastAiInputAt: now - 10_000,
      startedAt: now - 20_000,
      write: (b: Buffer) => writes.push(b),
      snapshotScreen: () => ({
        cols: 80,
        rows: 24,
        text: 'prompt$ ls\n',
        cursor: { x: 0, y: 1 },
      }),
      totalOutputBytes: 100,
    };
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);
    const res = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-1', input: 'ls\n', wait_ms: 50 },
      mkCtx(),
    );
    expect(res?.isError).toBe(false);
    expect(writes[0]!.toString()).toBe('ls\n');
    const parsed = JSON.parse(res!.output);
    expect(parsed.bytes_sent).toBe(3);
    expect(parsed.screen_after).toContain('prompt');
    expect(parsed.warning).toBeUndefined();
  });

  it('auto-appends \\n when input is printable without line terminator', async () => {
    const { sub, registry, connectionRepo } = mkStubSubsystem();
    connectionRepo.resolveConnection.mockReturnValue(mkConn());
    const now = Date.now();
    const fakeSession = {
      localTaskId: 'task-1', connectionId: 'conn-1', cols: 80, rows: 24,
      isClosed: false,
      idleMs: () => 10_000,
      lastAiInputAt: now - 10_000,
      startedAt: now - 20_000,
      write: vi.fn(),
      snapshotScreen: () => ({ cols: 80, rows: 24, text: 'prompt$ ', cursor: { x: 0, y: 0 } }),
      totalOutputBytes: 0,
    };
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);
    const res = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-1', input: 'ls -la' }, // no newline
      mkCtx(),
    );
    expect(res?.isError).toBe(false);
    const parsed = JSON.parse(res!.output);
    expect(parsed.auto_newline_appended).toBe(true);
    expect(parsed.warning).toBeUndefined();
    expect(parsed.bytes_sent).toBe(7); // 'ls -la' (6) + '\n' (1)
    // Verify the bytes actually written to PTY include the appended newline
    expect(fakeSession.write).toHaveBeenCalledTimes(1);
    const writtenBuf = fakeSession.write.mock.calls[0][0] as Buffer;
    expect(writtenBuf.toString('utf8')).toBe('ls -la\n');
  });

  it('does NOT auto-append newline for control bytes (Ctrl-C / Ctrl-D / Esc / Tab)', async () => {
    const { sub, registry, connectionRepo } = mkStubSubsystem();
    connectionRepo.resolveConnection.mockReturnValue(mkConn());
    const now = Date.now();
    const fakeSession = {
      localTaskId: 'task-1', connectionId: 'conn-1', cols: 80, rows: 24,
      isClosed: false,
      // Large idleMs so \x03 passes the guard (session appears idle, no block)
      idleMs: () => 10_000,
      lastAiInputAt: now - 10_000,
      startedAt: now - 20_000,
      write: vi.fn(),
      snapshotScreen: () => ({ cols: 80, rows: 24, text: '', cursor: { x: 0, y: 0 } }),
      totalOutputBytes: 0,
    };
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);
    for (const input of ['\x03', '\x04', '\x1b:q!', '\t']) {
      const res = await executeTool('SshConsoleSend', { connection_id: 'conn-1', input }, mkCtx());
      const parsed = JSON.parse(res!.output);
      expect(parsed.auto_newline_appended, `for ${JSON.stringify(input)}`).toBeFalsy();
    }
  });

  it('rejects when deny-list line hit', async () => {
    const { sub, registry, connectionRepo } = mkStubSubsystem();
    connectionRepo.resolveConnection.mockReturnValue(
      mkConn({ commandDenyPatterns: '^rm -rf /\\b' }),
    );
    const fakeSession = {
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      isClosed: false,
      write: vi.fn(),
      snapshotScreen: () => ({
        cols: 80,
        rows: 24,
        text: '',
        cursor: { x: 0, y: 0 },
      }),
      totalOutputBytes: 0,
    };
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);
    const res = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-1', input: 'rm -rf /\n' },
      mkCtx(),
    );
    expect(res?.isError).toBe(true);
    expect(fakeSession.write).not.toHaveBeenCalled();
  });

  it('rejects input over max_input_bytes_per_send', async () => {
    const { sub, registry } = mkStubSubsystem();
    sub.config.console.maxInputBytesPerSend = 4;
    registry.get.mockReturnValue({
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      isClosed: false,
      write: vi.fn(),
      snapshotScreen: () => ({
        cols: 80,
        rows: 24,
        text: '',
        cursor: { x: 0, y: 0 },
      }),
      totalOutputBytes: 0,
    } as any);
    setSshSubsystem(sub);
    const res = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-1', input: '12345' },
      mkCtx(),
    );
    expect(res?.isError).toBe(true);
  });

  it('omitting connection_id uses the active session', async () => {
    const { sub, registry, connectionRepo } = mkStubSubsystem();
    connectionRepo.resolveConnection.mockReturnValue(mkConn());
    const writes: Buffer[] = [];
    const now = Date.now();
    const fakeSession = {
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      cols: 80, rows: 24,
      isClosed: false,
      idleMs: () => 10_000,
      lastAiInputAt: now - 10_000,
      startedAt: now - 20_000,
      write: (b: Buffer) => writes.push(b),
      snapshotScreen: () => ({ cols: 80, rows: 24, text: 'ok', cursor: { x: 0, y: 0 } }),
      totalOutputBytes: 0,
    };
    // Task 4: with connection_id omitted and no pointer set yet,
    // resolveTargetSession falls back to the task's sole live session —
    // registry.get is keyed by (task, connection), so listForTask must
    // reflect that same single session too.
    registry.get.mockReturnValue(fakeSession as any);
    registry.listForTask.mockReturnValue([fakeSession]);
    setSshSubsystem(sub);
    const res = await executeTool('SshConsoleSend', { input: 'whoami\n' }, mkCtx());
    expect(res?.isError).toBe(false);
    expect(writes[0]!.toString()).toBe('whoami\n');
  });

  it('explicit connection_id for a non-existent session errors and does not fall back to another live session', async () => {
    // Task 4: registry is now (task, connection)-keyed — an explicit
    // connection_id does an exact lookup, not a "does this match the one
    // active session" comparison. A live session on a DIFFERENT connection
    // must not be silently used.
    const { sub, registry } = mkStubSubsystem();
    registry.get.mockImplementation((_taskId: string, connId: string) =>
      connId === 'conn-ACTIVE'
        ? {
            localTaskId: 'task-1',
            connectionId: 'conn-ACTIVE',
            cols: 80, rows: 24,
            isClosed: false,
            write: vi.fn(),
            snapshotScreen: () => ({ cols: 80, rows: 24, text: '', cursor: { x: 0, y: 0 } }),
            totalOutputBytes: 0,
          }
        : null,
    );
    setSshSubsystem(sub);
    const res = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-WRONG', input: 'ls\n' },
      mkCtx(),
    );
    expect(res?.isError).toBe(true);
    expect(res?.output).toContain('conn-WRONG');
    expect(res?.output).toContain('SshConsoleEnsure');
  });

  it('errors when no active session and no connection_id provided', async () => {
    const { sub, registry } = mkStubSubsystem();
    registry.get.mockReturnValue(null);
    setSshSubsystem(sub);
    const res = await executeTool('SshConsoleSend', { input: 'ls\n' }, mkCtx());
    expect(res?.isError).toBe(true);
    expect(res?.output).toContain('SshListConnections');
  });

  // ── Task 4: interrupt guard tests ──────────────────────────────────

  /** Build a fake session for Send that includes liveness getters needed by Task 4. */
  function mkSendSession(overrides: {
    idleMs?: number;
    lastAiInputAt?: number;
    startedAt?: number;
    screenText?: string;
  } = {}) {
    const now = Date.now();
    return {
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      isClosed: false,
      totalOutputBytes: 0,
      idleMs: vi.fn().mockReturnValue(overrides.idleMs ?? 500),
      lastAiInputAt: overrides.lastAiInputAt ?? now - 1000,
      startedAt: overrides.startedAt ?? now - 2000,
      write: vi.fn(),
      snapshotScreen: () => ({
        cols: 80,
        rows: 24,
        text: overrides.screenText ?? 'prompt$ ',
        cursor: { x: 0, y: 0 },
      }),
    };
  }

  it('blocks \\x03 when session is recently active and confirm_interrupt is absent', async () => {
    const { sub, registry } = mkStubSubsystem();
    const now = Date.now();
    // Small idleMs (recent output) → guard fires
    const fakeSession = mkSendSession({
      idleMs: 500,               // < GUARD_IDLE_MS (3000) → session alive
      lastAiInputAt: now - 1000, // < GUARD_RECENT_INPUT_MS (5000)
      startedAt: now - 2000,     // < GUARD_SESSION_AGE_MS (5000)
    });
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);

    const res = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-1', input: '\x03', wait_ms: 0 },
      mkCtx(),
    );

    expect(res?.isError).toBe(false);
    const data = JSON.parse(res!.output);
    expect(data.ok).toBe(false);
    expect(data.interrupt_blocked).toBe(true);
    expect(data.require_confirm).toBe(true);
    expect(data.reason).toMatch(/running/i);
    expect(data.idle_ms).toBeDefined();
    expect(data.hint).toContain('confirm_interrupt');
    // session.write must NOT have been called
    expect(fakeSession.write).not.toHaveBeenCalled();
  });

  it('allows \\x03 when confirm_interrupt=true even if session is recently active', async () => {
    const { sub, registry } = mkStubSubsystem();
    const now = Date.now();
    const fakeSession = mkSendSession({
      idleMs: 500,
      lastAiInputAt: now - 1000,
      startedAt: now - 2000,
    });
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);

    const res = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-1', input: '\x03', wait_ms: 0, confirm_interrupt: true },
      mkCtx(),
    );

    expect(res?.isError).toBe(false);
    const data = JSON.parse(res!.output);
    expect(data.ok).toBe(true);
    expect(data.interrupt_blocked).toBeUndefined();
    // session.write MUST have been called with the \x03 byte
    expect(fakeSession.write).toHaveBeenCalledTimes(1);
    const writtenBuf = fakeSession.write.mock.calls[0][0] as Buffer;
    expect(writtenBuf.includes(0x03)).toBe(true);
  });

  it('allows \\x03 without confirm when session is clearly idle (large idleMs + old timestamps)', async () => {
    const { sub, registry } = mkStubSubsystem();
    const now = Date.now();
    const fakeSession = mkSendSession({
      idleMs: 10_000,             // > GUARD_IDLE_MS (3000)
      lastAiInputAt: now - 9_000, // > GUARD_RECENT_INPUT_MS (5000)
      startedAt: now - 20_000,    // > GUARD_SESSION_AGE_MS (5000)
    });
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);

    const res = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-1', input: '\x03', wait_ms: 0 },
      mkCtx(),
    );

    expect(res?.isError).toBe(false);
    const data = JSON.parse(res!.output);
    expect(data.ok).toBe(true);
    expect(data.interrupt_blocked).toBeUndefined();
    expect(fakeSession.write).toHaveBeenCalledTimes(1);
  });

  it('normal printable Send result includes idle_ms field', async () => {
    const { sub, registry } = mkStubSubsystem();
    const fakeSession = mkSendSession({ idleMs: 42 });
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);

    const res = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-1', input: 'ls\n', wait_ms: 0 },
      mkCtx(),
    );

    expect(res?.isError).toBe(false);
    const data = JSON.parse(res!.output);
    expect(data.ok).toBe(true);
    expect(typeof data.idle_ms).toBe('number');
    expect(data.maybe_waiting_for_input).toBeDefined();
  });
});

describe('SshConsoleSnapshot', () => {
  beforeEach(() => {
    setSshSubsystem(null);
    __resetAgentConnectionPointer();
  });

  it('returns screen when kind=screen', async () => {
    const { sub, registry } = mkStubSubsystem();
    registry.get.mockReturnValue({
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      isClosed: false,
      snapshotScreen: () => ({ cols: 80, rows: 24, text: 'screen', cursor: { x: 1, y: 2 } }),
      snapshotScrollback: () => ({ text: 'scroll', byteCount: 6, truncated: false }),
    } as any);
    setSshSubsystem(sub);
    const res = await executeTool('SshConsoleSnapshot', { connection_id: 'conn-1' }, mkCtx());
    expect(res?.isError).toBe(false);
    const data = JSON.parse(res!.output);
    expect(data.kind).toBe('screen');
    expect(data.text).toBe('screen');
    expect(data.cursor).toEqual({ x: 1, y: 2 });
  });

  it('returns scrollback when kind=scrollback', async () => {
    const { sub, registry } = mkStubSubsystem();
    registry.get.mockReturnValue({
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      isClosed: false,
      snapshotScreen: () => ({ cols: 80, rows: 24, text: '', cursor: { x: 0, y: 0 } }),
      snapshotScrollback: (_opts: { maxBytes: number }) => ({
        text: 'tail',
        byteCount: 9999,
        truncated: true,
      }),
    } as any);
    setSshSubsystem(sub);
    const res = await executeTool(
      'SshConsoleSnapshot',
      { connection_id: 'conn-1', kind: 'scrollback', max_bytes: 4 },
      mkCtx(),
    );
    expect(res?.isError).toBe(false);
    const data = JSON.parse(res!.output);
    expect(data.kind).toBe('scrollback');
    expect(data.text).toBe('tail');
    expect(data.truncated).toBe(true);
  });

  it('returns error when no active session', async () => {
    const { sub, registry } = mkStubSubsystem();
    registry.get.mockReturnValue(null);
    setSshSubsystem(sub);
    const res = await executeTool('SshConsoleSnapshot', { connection_id: 'conn-1' }, mkCtx());
    expect(res?.isError).toBe(true);
    expect(res?.output).toContain('no live session');
  });

  it('omitting connection_id uses the active session', async () => {
    const { sub, registry } = mkStubSubsystem();
    const fakeSession = {
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      cols: 80, rows: 24,
      isClosed: false,
      snapshotScreen: () => ({ cols: 80, rows: 24, text: 'auto', cursor: { x: 0, y: 0 } }),
      snapshotScrollback: () => ({ text: '', byteCount: 0, truncated: false }),
    };
    // Task 4: omitted connection_id + no pointer yet falls back to the
    // task's sole live session via listForTask, not just registry.get.
    registry.get.mockReturnValue(fakeSession as any);
    registry.listForTask.mockReturnValue([fakeSession]);
    setSshSubsystem(sub);
    const res = await executeTool('SshConsoleSnapshot', {}, mkCtx());
    expect(res?.isError).toBe(false);
    const data = JSON.parse(res!.output);
    expect(data.text).toBe('auto');
  });

  it('explicit connection_id for a non-existent session errors and does not fall back to another live session', async () => {
    // Task 4: exact (task, connection) lookup — a live session on a
    // DIFFERENT connection must not be silently used.
    const { sub, registry } = mkStubSubsystem();
    registry.get.mockImplementation((_taskId: string, connId: string) =>
      connId === 'conn-ACTIVE'
        ? {
            localTaskId: 'task-1',
            connectionId: 'conn-ACTIVE',
            cols: 80, rows: 24,
            isClosed: false,
            snapshotScreen: () => ({ cols: 80, rows: 24, text: '', cursor: { x: 0, y: 0 } }),
          }
        : null,
    );
    setSshSubsystem(sub);
    const res = await executeTool(
      'SshConsoleSnapshot',
      { connection_id: 'conn-WRONG' },
      mkCtx(),
    );
    expect(res?.isError).toBe(true);
    expect(res?.output).toContain('conn-WRONG');
    expect(res?.output).toContain('SshConsoleEnsure');
  });
});

describe('SshConsoleRun', () => {
  beforeEach(() => {
    setSshSubsystem(null);
    __resetAgentConnectionPointer();
  });

  /** Build a fake session that supports onOutput, write (spy), scrollbackBytes, snapshotScreen,
   *  lastHumanInputAt, lastAiInputAt, idleMs, startedAt — matching the ConsoleSession public surface. */
  function mkFakeSession() {
    const outputListeners: Set<(chunk: Buffer) => void> = new Set();
    let _scrollback = Buffer.alloc(0);
    let _lastHumanInputAt = 0;
    let _lastAiInputAt = 0;
    let _startedAt = Date.now();

    const session = {
      localTaskId: 'task-1',
      connectionId: 'conn-1',
      cols: 80,
      rows: 24,
      isClosed: false,
      startedAt: _startedAt,
      get lastHumanInputAt() { return _lastHumanInputAt; },
      get lastAiInputAt() { return _lastAiInputAt; },
      idleMs: vi.fn().mockReturnValue(50),
      totalOutputBytes: 0,
      /** Spy-able write: when source='human', bump humanInputAt so the run loop detects it. */
      write: vi.fn((buf: Buffer, source: 'ai' | 'human') => {
        if (source === 'human') _lastHumanInputAt = Date.now();
        else _lastAiInputAt = Date.now();
        _scrollback = Buffer.concat([_scrollback, buf]);
      }),
      onOutput: (listener: (chunk: Buffer) => void) => {
        outputListeners.add(listener);
        return () => { outputListeners.delete(listener); };
      },
      scrollbackBytes: () => _scrollback,
      snapshotScreen: () => ({ cols: 80, rows: 24, text: 'prompt$ ', cursor: { x: 0, y: 0 } }),
      /** Test helper: push a chunk to all output listeners + append to scrollback. */
      _pushOutput(chunk: Buffer) {
        _scrollback = Buffer.concat([_scrollback, chunk]);
        for (const l of outputListeners) l(chunk);
      },
    };
    return session;
  }

  it('SshConsoleRun is registered in TOOL_DEFS', () => {
    expect(TOOL_DEFS.SshConsoleRun).toBeDefined();
  });

  it('resolves done+exit_code=0 when marker is pushed to session output', async () => {
    const { sub, registry } = mkStubSubsystem();
    const fakeSession = mkFakeSession();
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);

    // Start the run — it will block waiting for output.
    const runPromise = executeTool(
      'SshConsoleRun',
      { connection_id: 'conn-1', command: 'echo hello', timeout_ms: 5000 },
      mkCtx(),
    );

    // Let the run() function start and subscribe to onOutput.
    await new Promise((r) => setImmediate(r));

    // Extract the nonce from the command text written to session.write (the markerCommand).
    // write spy: first call is the markerCommand text (after deny-check etc.)
    // Find the write call for source='ai' that contains the marker.
    let nonce: string | undefined;
    for (const call of fakeSession.write.mock.calls) {
      const text = (call[0] as Buffer).toString('utf8');
      const m = /__MAESTRO_DONE_([0-9a-f]+):/.exec(text);
      if (m) { nonce = m[1]; break; }
    }
    expect(nonce).toBeDefined();

    // Push some output then the marker line.
    fakeSession._pushOutput(Buffer.from('hello\n'));
    fakeSession._pushOutput(Buffer.from(`__MAESTRO_DONE_${nonce!}:0__\n`));

    const res = await runPromise;
    expect(res?.isError).toBe(false);
    const data = JSON.parse(res!.output);
    expect(data.ok).toBe(true);
    expect(data.done).toBe(true);
    expect(data.exit_code).toBe(0);
    expect(data.marker_found).toBe(true);
    expect(data.output).toContain('hello');
  });

  it('returns timed_out=true on timeout and does NOT write Ctrl-C (\\x03)', async () => {
    const { sub, registry } = mkStubSubsystem();
    const fakeSession = mkFakeSession();
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);

    const res = await executeTool(
      'SshConsoleRun',
      { connection_id: 'conn-1', command: 'sleep 99', timeout_ms: 50 },
      mkCtx(),
    );

    expect(res?.isError).toBe(false);
    const data = JSON.parse(res!.output);
    expect(data.done).toBe(false);
    expect(data.timed_out).toBe(true);

    // Ctrl-C must NEVER be sent on timeout.
    const allWritten = fakeSession.write.mock.calls
      .map((c) => (c[0] as Buffer).toString('binary'))
      .join('');
    expect(allWritten).not.toContain('\x03');
  });

  it('returns interrupted_by=human when human writes during the wait', async () => {
    const { sub, registry } = mkStubSubsystem();
    const fakeSession = mkFakeSession();
    registry.get.mockReturnValue(fakeSession);
    setSshSubsystem(sub);

    const runPromise = executeTool(
      'SshConsoleRun',
      { connection_id: 'conn-1', command: 'sleep 10', timeout_ms: 5000 },
      mkCtx(),
    );

    // Let the run loop subscribe and start the interval.
    await new Promise((r) => setImmediate(r));

    // Simulate a human keystroke — this updates lastHumanInputAt on the fake session.
    fakeSession.write(Buffer.from('q'), 'human');

    // Wait for the interval (100ms) to fire and detect the human input.
    await new Promise((r) => setTimeout(r, 200));

    const res = await runPromise;
    expect(res?.isError).toBe(false);
    const data = JSON.parse(res!.output);
    expect(data.done).toBe(false);
    expect(data.interrupted_by).toBe('human');
    expect(data.timed_out).toBe(false);
  });
});

// ──────────────────────────────────────────────────────────────────────
// Task 4: agent current-connection pointer + resolveTargetSession routing
// ──────────────────────────────────────────────────────────────────────
//
// Task 3 made openConsoleSession add-not-replace, so a task can hold several
// live sessions at once. These tests cover the routing an omitted
// connection_id gets in that world: the agent's current-connection pointer
// (set by the most recent successful Ensure/Send/Run/Snapshot), falling
// back to the task's sole session when unambiguous, and a clear error
// otherwise. Uses the same (task, connection)-keyed stateful registry stub
// as the Task 3 tests above (mkStatefulRegistry / mkPreSession).

/** A session stub with everything SshConsoleSend needs (write spy,
 *  snapshotScreen, liveness fields) AND the registry keying fields
 *  (localTaskId/connectionId/isClosed) so mkStatefulRegistry can hold it. */
function mkRoutableSession(taskId: string, connId: string) {
  const now = Date.now();
  return {
    localTaskId: taskId,
    connectionId: connId,
    cols: 80,
    rows: 24,
    isClosed: false,
    startedAt: now - 20_000,
    totalOutputBytes: 0,
    idleMs: vi.fn().mockReturnValue(10_000),
    lastAiInputAt: now - 10_000,
    write: vi.fn(),
    snapshotScreen: () => ({ cols: 80, rows: 24, text: 'ok', cursor: { x: 0, y: 0 } }),
  };
}

describe('resolveTargetSession (Task 4 routing)', () => {
  beforeEach(() => {
    setSshSubsystem(null);
    __resetAgentConnectionPointer();
  });

  it('(a) two sessions, no connection_id: a Send after the agent used connection A routes to A', async () => {
    const { sub } = mkStubSubsystem();
    const { reg } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;
    const sessA = mkRoutableSession('task-1', 'conn-A');
    const sessB = mkRoutableSession('task-1', 'conn-B');
    reg.register(sessA);
    reg.register(sessB);
    setSshSubsystem(sub);

    // The agent explicitly acts on connection A first — this sets the
    // per-task current-connection pointer to conn-A.
    const first = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-A', input: 'ls\n' },
      mkCtx(),
    );
    expect(first?.isError).toBe(false);
    expect(sessA.write).toHaveBeenCalledTimes(1);

    // A later Send with connection_id omitted must route to the pointer
    // (conn-A), not error out just because the task now has 2 sessions.
    const second = await executeTool('SshConsoleSend', { input: 'pwd\n' }, mkCtx());
    expect(second?.isError).toBe(false);
    expect(sessA.write).toHaveBeenCalledTimes(2);
    expect(sessB.write).not.toHaveBeenCalled();
  });

  it('(b) two sessions, no pointer set yet, no connection_id → ambiguity error', async () => {
    const { sub } = mkStubSubsystem();
    const { reg } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;
    reg.register(mkRoutableSession('task-1', 'conn-A'));
    reg.register(mkRoutableSession('task-1', 'conn-B'));
    setSshSubsystem(sub);

    const res = await executeTool('SshConsoleSend', { input: 'ls\n' }, mkCtx());
    expect(res?.isError).toBe(true);
    expect(res?.output).toContain('connection_id required (multiple sessions open');
  });

  it('(c) one session, omitted connection_id → routes to the sole session', async () => {
    const { sub } = mkStubSubsystem();
    const { reg } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;
    const sess = mkRoutableSession('task-1', 'conn-solo');
    reg.register(sess);
    setSshSubsystem(sub);

    const res = await executeTool('SshConsoleSend', { input: 'ls\n' }, mkCtx());
    expect(res?.isError).toBe(false);
    expect(sess.write).toHaveBeenCalledTimes(1);
  });

  it('(d) explicit connection_id for a non-existent session → clear error (no fallback to another live session)', async () => {
    const { sub } = mkStubSubsystem();
    const { reg } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;
    reg.register(mkRoutableSession('task-1', 'conn-real'));
    setSshSubsystem(sub);

    const res = await executeTool(
      'SshConsoleSend',
      { connection_id: 'conn-ghost', input: 'ls\n' },
      mkCtx(),
    );
    expect(res?.isError).toBe(true);
    expect(res?.output).toContain('conn-ghost');
    expect(res?.output).toContain('SshConsoleEnsure');
  });

  // ── Direct unit coverage of resolveTargetSession itself ──────────────

  it('unit: falls back through pointer → sole session → ambiguous → none in precedence order', () => {
    const { sub } = mkStubSubsystem();
    const { reg } = mkStatefulRegistry();
    (sub as any).sessionRegistry = reg;

    // none: no sessions at all.
    expect(resolveTargetSession(sub, 'task-1', undefined)).toEqual({ ok: false, error: 'none' });

    // sole session: unambiguous even without a pointer.
    const solo = mkRoutableSession('task-1', 'conn-solo');
    reg.register(solo);
    expect(resolveTargetSession(sub, 'task-1', undefined)).toEqual({
      ok: true,
      session: solo,
      connectionId: 'conn-solo',
    });

    // ambiguous: a second session with no pointer set.
    const second = mkRoutableSession('task-1', 'conn-second');
    reg.register(second);
    const ambiguous = resolveTargetSession(sub, 'task-1', undefined);
    expect(ambiguous.ok).toBe(false);
    if (!ambiguous.ok) expect(ambiguous.error).toBe('ambiguous');

    // explicit connection_id for a session that doesn't exist → not_found,
    // independent of how many OTHER sessions are live.
    const notFound = resolveTargetSession(sub, 'task-1', 'conn-missing');
    expect(notFound).toEqual({ ok: false, error: 'not_found', connectionId: 'conn-missing' });

    // explicit connection_id for a session that DOES exist always resolves,
    // even with multiple sessions live and no pointer set.
    expect(resolveTargetSession(sub, 'task-1', 'conn-second')).toEqual({
      ok: true,
      session: second,
      connectionId: 'conn-second',
    });
  });
});
