maestro/src/engine/tools/ssh-console.test.ts
2026-06-03 05:08:00 +00:00

557 lines
18 KiB
TypeScript

/**
* 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 } from './ssh-console.js';
import { setSshSubsystem, 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),
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() };
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,
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));
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('rejects mismatching connection_id by default and surfaces the active id', async () => {
const { sub, registry } = mkStubSubsystem();
registry.get.mockReturnValue({
localTaskId: 'task-1',
connectionId: 'conn-OLD',
isClosed: false,
startedAt: Date.now() - 60_000,
lastActivityAt: Date.now() - 5_000,
});
setSshSubsystem(sub);
const res = await executeTool('SshConsoleEnsure', { connection_id: 'conn-NEW' }, mkCtx());
expect(res?.isError).toBe(true);
expect(res?.output).toContain('conn-OLD');
expect(res?.output).toContain('force_replace');
expect(registry.closeForTask).not.toHaveBeenCalled();
});
it('closes old session and opens new when force_replace=true', async () => {
const { sub, registry } = mkStubSubsystem();
registry.get.mockReturnValue({
localTaskId: 'task-1',
connectionId: 'conn-OLD',
isClosed: false,
startedAt: Date.now() - 60_000,
lastActivityAt: Date.now() - 5_000,
});
setSshSubsystem(sub);
const res = await executeTool(
'SshConsoleEnsure',
{ connection_id: 'conn-NEW', force_replace: true },
mkCtx(),
);
expect(registry.closeForTask).toHaveBeenCalledWith('task-1', 'connection_change');
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);
});
});
describe('SshConsoleSend', () => {
beforeEach(() => setSshSubsystem(null));
it('writes input to session and returns screen snapshot', async () => {
const { sub, registry } = mkStubSubsystem();
const writes: Buffer[] = [];
const fakeSession = {
localTaskId: 'task-1',
connectionId: 'conn-1',
cols: 80,
rows: 24,
isClosed: false,
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 fakeSession = {
localTaskId: 'task-1', connectionId: 'conn-1', cols: 80, rows: 24,
isClosed: false,
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 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);
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[] = [];
registry.get.mockReturnValue({
localTaskId: 'task-1',
connectionId: 'conn-1',
cols: 80, rows: 24,
isClosed: false,
write: (b: Buffer) => writes.push(b),
snapshotScreen: () => ({ cols: 80, rows: 24, text: 'ok', cursor: { x: 0, y: 0 } }),
totalOutputBytes: 0,
} as any);
setSshSubsystem(sub);
const res = await executeTool('SshConsoleSend', { input: 'whoami\n' }, mkCtx());
expect(res?.isError).toBe(false);
expect(writes[0]!.toString()).toBe('whoami\n');
});
it('rejects mismatching connection_id and surfaces the active id', async () => {
const { sub, registry } = mkStubSubsystem();
registry.get.mockReturnValue({
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,
} as any);
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-ACTIVE');
expect(res?.output).toContain('force_replace');
});
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');
});
});
describe('SshConsoleSnapshot', () => {
beforeEach(() => setSshSubsystem(null));
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();
registry.get.mockReturnValue({
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 }),
} as any);
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('rejects mismatching connection_id and surfaces the active id', async () => {
const { sub, registry } = mkStubSubsystem();
registry.get.mockReturnValue({
localTaskId: 'task-1',
connectionId: 'conn-ACTIVE',
cols: 80, rows: 24,
isClosed: false,
snapshotScreen: () => ({ cols: 80, rows: 24, text: '', cursor: { x: 0, y: 0 } }),
} as any);
setSshSubsystem(sub);
const res = await executeTool(
'SshConsoleSnapshot',
{ connection_id: 'conn-WRONG' },
mkCtx(),
);
expect(res?.isError).toBe(true);
expect(res?.output).toContain('conn-ACTIVE');
});
});