sync: update from private repo (edc775f2)

src/engine/tools/ssh-console.ts

@@ -6,10 +6,12 @@
  * (piece membership, access decision, enabled / abuse / host-key state)
  * via the exported `preflight` helper in ssh.ts. After that:
  *
- *   - SshConsoleEnsure: find-or-open the session keyed by (localTaskId).
- *     If a session already exists for a different connection on this task,
- *     close it with reason 'connection_change' and open a fresh one. Apply
- *     per-connection session cap (evict oldest with 'session_cap_evict').
+ *   - SshConsoleEnsure: find-or-open the session keyed by
+ *     (localTaskId, connectionId). A task can hold multiple concurrent
+ *     sessions, one per connection: opening a DIFFERENT connection ADDS a
+ *     session (it does not close the others). `force_replace` only restarts
+ *     the SAME connection's session (close 'connection_change' + reopen).
+ *     Enforces per-task, per-user, and per-connection session caps.
  *
  *   - SshConsoleSend: deny-list check the input lines, then forward the
  *     bytes to the session (which writes them straight to the PTY — same
@@ -63,7 +65,7 @@ export interface OpenConsoleParams {
   userId: string;
   /** For the grant check (access resolver matches against this piece). */
   pieceName: string;
-  /** Piece allowed_ssh_connections (['*'] for user-initiated). */
+  /** Workspace/space-scoped SSH connection IDs (['*'] for user-initiated console). */
   allowedConnections: string[];
   cols?: number;
   rows?: number;
@@ -72,6 +74,15 @@ export interface OpenConsoleParams {
   initiator: 'agent' | 'user';
   /** Optional job id for the audit rows (agent path passes ctx.jobId). */
   jobId?: string | null;
+  /**
+   * Per-space isolation (spec §11): the space of the job/task opening the
+   * console. Gates access to space-scoped connections in the SAME way SshExec
+   * does via ctx.spaceId. MUST be the already-resolved space id (personal-WS
+   * NULL→personal-space handled by the caller, mirroring the worker). When
+   * omitted, resolveAccess treats the caller as a no-space (legacy) job and
+   * denies every space-scoped connection with `space_mismatch`.
+   */
+  spaceId?: string | null;
 }
 
 export interface OpenConsoleResult {
@@ -86,7 +97,8 @@ export interface OpenConsoleResult {
     | 'no_grant'
     | 'host_key_not_verified'
     | 'host_key_mismatch'
-    | 'connection_change_requires_force'
+    | 'task_session_cap'
+    | 'user_session_cap'
     | 'abuse_locked'
     | 'connection_disabled'
     | 'connection_not_found'
@@ -104,6 +116,46 @@ export interface OpenConsoleResult {
   session?: ConsoleSession;
 }
 
+// ──────────────────────────────────────────────────────────────────────
+// Per-task open lock (TOCTOU guard for the find-or-open core)
+// ──────────────────────────────────────────────────────────────────────
+//
+// Opening a console session is a check-then-act sequence: look up any
+// existing (task, connection) session, enforce the per-task / per-user caps,
+// open the SSH channel, then register. Two concurrent opens for the SAME task
+// must not both pass the cap check, and must not both register a fresh session
+// for the same (task, connection) — the later register() would silently
+// overwrite (and orphan) the first, leaking its SSH client + PTY. We serialize
+// the whole existence-check → cap-check → channel-open → register sequence per
+// task with a minimal chained-promise mutex (opens are infrequent, so the
+// per-task serialization cost is negligible).
+const taskOpenLocks = new Map<string, Promise<void>>();
+
+async function withTaskLock<T>(taskId: string, fn: () => Promise<T>): Promise<T> {
+  const prev = taskOpenLocks.get(taskId) ?? Promise.resolve();
+  let release!: () => void;
+  const gate = new Promise<void>((r) => {
+    release = r;
+  });
+  const chained = prev.then(() => gate);
+  taskOpenLocks.set(taskId, chained);
+  // Wait for the previous holder to release. Swallow its rejection — each
+  // holder is independent, so one open failing must not wedge the queue.
+  await prev.catch(() => {});
+  try {
+    return await fn();
+  } finally {
+    release();
+    // If nobody queued behind us, drop the map entry so it can't grow
+    // unbounded across tasks. Safe under JS's single-threaded model: any
+    // later waiter's synchronous set() has already replaced this entry
+    // before this finally runs.
+    if (taskOpenLocks.get(taskId) === chained) {
+      taskOpenLocks.delete(taskId);
+    }
+  }
+}
+
 /**
  * Find-or-open a ConsoleSession bound to (taskId, connectionId). Runs the
  * full preflight (piece membership, access decision, enabled / abuse /
@@ -111,6 +163,13 @@ export interface OpenConsoleResult {
  * registers the ConsoleSession, enforces the per-connection cap, and writes
  * the `ssh.console.open` audit (with `initiator` in the detail).
  *
+ * Add-not-replace: a task may hold several sessions at once, one per
+ * connection. Opening a NEW connection adds a session; it never closes the
+ * task's other connections. Only `force_replace` on the SAME connection
+ * restarts that one shell. The whole sequence is serialized per task
+ * (`withTaskLock`) so concurrent opens can't exceed the caps or orphan a
+ * session by double-registering the same (task, connection).
+ *
  * Returns a structured OpenConsoleResult. On the happy path / reuse,
  * `ok: true` with `session` set; on any gate failure, `ok: false` with a
  * structured `error` code + `message`.
@@ -119,7 +178,6 @@ export async function openConsoleSession(
   deps: OpenConsoleDeps,
   params: OpenConsoleParams,
 ): Promise<OpenConsoleResult> {
-  const { sub } = deps;
   const connectionId = params.connectionId;
   if (!connectionId) {
     return { ok: false, error: 'missing_connection_id', message: 'SshConsoleEnsure: connection_id is required.' };
@@ -133,13 +191,30 @@ export async function openConsoleSession(
     };
   }
 
-  // If a session already exists for this task, branch on whether it's the
-  // same connection. Same → reuse. Different → reject by default (so a
-  // single LLM connection_id slip can't kill the user's live shell), opt
-  // into the swap with force_replace=true.
-  const existing = sub.sessionRegistry.get(localTaskId);
+  // Serialize the existence-check → cap-check → channel-open → register for
+  // this task. See withTaskLock above for why (TOCTOU on caps + register).
+  return withTaskLock(localTaskId, () => openConsoleSessionLocked(deps, params, connectionId, localTaskId));
+}
+
+async function openConsoleSessionLocked(
+  deps: OpenConsoleDeps,
+  params: OpenConsoleParams,
+  connectionId: string,
+  localTaskId: string,
+): Promise<OpenConsoleResult> {
+  const { sub } = deps;
+
+  // Existence check is per (task, connection) now. Same connection already
+  // live → reuse (or restart on force_replace). A DIFFERENT connection is
+  // simply a new session below — opening it never touches other connections.
+  const existing = sub.sessionRegistry.get(localTaskId, connectionId);
   if (existing) {
-    if (existing.connectionId === connectionId) {
+    if (params.forceReplace === true) {
+      // Same-connection restart: close ONLY this connection's session, then
+      // fall through to open a fresh one. Sibling sessions on other
+      // connections of this task are left untouched.
+      await sub.sessionRegistry.closeSession(localTaskId, connectionId, 'connection_change');
+    } else {
       return {
         ok: true,
         alreadyActive: true,
@@ -149,22 +224,40 @@ export async function openConsoleSession(
         session: existing,
       };
     }
-    const forceReplace = params.forceReplace === true;
-    if (!forceReplace) {
-      const ageSec = Math.max(0, Math.floor((Date.now() - existing.startedAt) / 1000));
-      const idleSec = Math.max(0, Math.floor((Date.now() - existing.lastActivityAt) / 1000));
+  }
+
+  // Per-task cap: how many live sessions this task already holds across all
+  // its connections. At/over the cap → refuse, closing nothing.
+  const maxSessionsPerTask = sub.config.console.maxSessionsPerTask;
+  if (
+    typeof maxSessionsPerTask === 'number' &&
+    sub.sessionRegistry.listForTask(localTaskId).length >= maxSessionsPerTask
+  ) {
+    return {
+      ok: false,
+      error: 'task_session_cap',
+      message:
+        `SshConsoleEnsure: this task already has the maximum of ${maxSessionsPerTask} open console ` +
+        `sessions. Close one (or reuse an existing connection) before opening another.`,
+    };
+  }
+
+  // Per-user cap: how many live sessions this initiator owns across ALL
+  // tasks/connections. 0 = unlimited. At/over the cap → refuse.
+  const maxSessionsPerUser = sub.config.console.maxSessionsPerUser;
+  if (typeof maxSessionsPerUser === 'number' && maxSessionsPerUser > 0) {
+    const ownedByUser = sub.sessionRegistry
+      .listAll()
+      .filter((s) => s.startedByUserId === params.userId).length;
+    if (ownedByUser >= maxSessionsPerUser) {
       return {
         ok: false,
-        error: 'connection_change_requires_force',
-        connectionId: existing.connectionId,
+        error: 'user_session_cap',
         message:
-          `SshConsoleEnsure: this task already has an active session on connection ${existing.connectionId} ` +
-          `(age=${ageSec}s, last_activity=${idleSec}s ago). ` +
-          `Use connection_id="${existing.connectionId}" to continue working in the existing shell, ` +
-          `or pass force_replace=true to close it and open a new session on ${connectionId}.`,
+          `SshConsoleEnsure: you already have the maximum of ${maxSessionsPerUser} open console ` +
+          `sessions across all tasks. Close one before opening another.`,
       };
     }
-    await sub.sessionRegistry.closeForTask(localTaskId, 'connection_change');
   }
 
   // Build the minimal ToolContext-shaped object the shared `preflight`
@@ -180,6 +273,11 @@ export async function openConsoleSession(
     jobId: params.jobId ?? undefined,
     pieceName: params.pieceName,
     allowedSshConnections: params.allowedConnections,
+    // Per-space isolation (spec §11): thread the task/job space so the shared
+    // preflight's resolveAccess gates space-scoped connections identically to
+    // SshExec. Without this every workspace-registered connection is denied
+    // with `space_mismatch`.
+    spaceId: params.spaceId ?? undefined,
   } as ToolContext;
 
   // Full 12-step preflight (same path as SshExec).
@@ -334,10 +432,12 @@ export async function openConsoleSession(
     'success',
   );
 
-  // Enforce the per-connection session cap (evict oldest).
+  // Enforce the per-connection session cap (evict oldest). Close each evicted
+  // session by its exact (task, connection) key — NOT closeForTask, which
+  // would collaterally kill other connections the evicted task holds.
   const evict = sub.sessionRegistry.enforceCap(connectionId);
   for (const e of evict) {
-    sub.sessionRegistry.closeForTask(e.localTaskId, 'session_cap_evict').catch((err) =>
+    sub.sessionRegistry.closeSession(e.localTaskId, e.connectionId, 'session_cap_evict').catch((err) =>
       logger.warn(`[ssh-console] evict close error: ${(err as Error).message}`),
     );
   }
@@ -345,6 +445,180 @@ export async function openConsoleSession(
   return { ok: true, alreadyActive: false, connectionId, cols, rows, session };
 }
 
+// ──────────────────────────────────────────────────────────────────────
+// Agent current-connection pointer + shared session routing (Task 4)
+// ──────────────────────────────────────────────────────────────────────
+//
+// Task 3 made openConsoleSession add-not-replace, so a task can now hold
+// several live console sessions at once (one per connection). Send / Run /
+// Snapshot need to resolve an omitted connection_id to the RIGHT one rather
+// than an arbitrary "the" active session. We track, per task, which
+// connection the agent last acted on (Ensure/Send/Run/Snapshot) and use
+// that as the default when connection_id is omitted and there's more than
+// one live session.
+//
+// Populated by ensureSessionInternal (the agent-facing wrapper) — NOT by
+// the shared openConsoleSession core, which is also used by the
+// human-initiated HTTP console endpoint. Attributing a human's console
+// open to "the agent's current connection" would be wrong.
+const agentCurrentConnection = new Map<string, string>();
+
+function markAgentConnection(taskId: string, connectionId: string): void {
+  agentCurrentConnection.set(taskId, connectionId);
+}
+
+/** Test-only: clear pointer state between test cases (module-scope Map
+ * would otherwise leak taskId reuse — e.g. 'task-1' — across test files). */
+export function __resetAgentConnectionPointer(): void {
+  agentCurrentConnection.clear();
+}
+
+export type ResolveSessionResult =
+  | { ok: true; session: ConsoleSession; connectionId: string }
+  | { ok: false; error: 'not_found'; connectionId: string }
+  | { ok: false; error: 'ambiguous'; connectionIds: string[] }
+  | { ok: false; error: 'none' };
+
+/**
+ * Shared routing for SshConsoleSend / SshConsoleRun / SshConsoleSnapshot
+ * when connection_id may be omitted:
+ *
+ *   1. connection_id given → exact (taskId, connectionId) lookup. No live
+ *      session there → 'not_found'. We do NOT auto-open here: silently
+ *      opening a fresh session for an explicit-but-wrong id would mask a
+ *      typo'd connection_id instead of surfacing it.
+ *   2. else the agent's current-connection pointer, if set and still live.
+ *      A pointer whose session has since closed is treated as unset (and
+ *      dropped) rather than resolved.
+ *   3. else, if the task has exactly one live session, that one
+ *      (unambiguous single-session case — most tasks).
+ *   4. else, if the task has more than one live session, 'ambiguous' — the
+ *      agent must disambiguate with connection_id.
+ *   5. else 'none' — no live session at all for this task.
+ */
+export function resolveTargetSession(
+  sub: SshSubsystem,
+  taskId: string,
+  connectionIdMaybe: string | undefined,
+): ResolveSessionResult {
+  if (connectionIdMaybe) {
+    const session = sub.sessionRegistry.get(taskId, connectionIdMaybe);
+    if (!session) return { ok: false, error: 'not_found', connectionId: connectionIdMaybe };
+    return { ok: true, session, connectionId: connectionIdMaybe };
+  }
+
+  const pointerConnId = agentCurrentConnection.get(taskId);
+  if (pointerConnId) {
+    const session = sub.sessionRegistry.get(taskId, pointerConnId);
+    if (session) return { ok: true, session, connectionId: pointerConnId };
+    // Stale pointer (session closed since) — drop it and fall through to
+    // the listForTask-based resolution below.
+    agentCurrentConnection.delete(taskId);
+  }
+
+  const sessions = sub.sessionRegistry.listForTask(taskId);
+  if (sessions.length === 1) {
+    const only = sessions[0]!;
+    return { ok: true, session: only, connectionId: only.connectionId };
+  }
+  if (sessions.length > 1) {
+    return { ok: false, error: 'ambiguous', connectionIds: sessions.map((s) => s.connectionId) };
+  }
+  return { ok: false, error: 'none' };
+}
+
+/**
+ * Screen-auto-inject lookup consumed by buildGlobalPreamble (prompt.ts) via
+ * __setActiveSessionLookup — shows the agent every one of its live console
+ * screens at the top of each movement's system prompt (one block per open
+ * connection). Returns ALL live sessions for the task, unfiltered; prompt.ts
+ * is responsible for scoping the result down to the current movement's
+ * `allowedSshConnections` (Task 5) — this function must not do that
+ * filtering itself, since it has no visibility into which movement is
+ * asking.
+ */
+export function listAgentConsoleSessions(sub: SshSubsystem, taskId: string): ConsoleSession[] {
+  return sub.sessionRegistry.listForTask(taskId);
+}
+
+/**
+ * Manually close a single (task, connection) console session on behalf of a
+ * human actor (task owner/admin or the session's own starter — the HTTP route
+ * enforces that gate). Closes via the registry with reason `user_close` AND
+ * records an audit row naming the ACTOR.
+ *
+ * Why the extra audit row: ConsoleSession.close() already writes an
+ * `ssh.console.close` audit, but it attributes the action to
+ * `startedByUserId` — correct for idle/duration/transport teardown, WRONG for
+ * a manual close by a different user (e.g. an admin killing someone's stray
+ * session). This row captures "who actually pressed close" with the original
+ * starter preserved in the detail for forensics.
+ *
+ * Returns `{ ok: false, error: 'not_found' }` when no live session exists for
+ * the pair (raced with idle sweep / another close) so the caller can 404.
+ */
+export async function closeConsoleSessionByUser(
+  sub: SshSubsystem,
+  args: { taskId: string; connectionId: string; actorUserId: string },
+): Promise<{ ok: boolean; error?: 'not_found' }> {
+  const session = sub.sessionRegistry.get(args.taskId, args.connectionId);
+  if (!session) return { ok: false, error: 'not_found' };
+  sub.auditRepo.beginAndComplete(
+    {
+      action: 'ssh.console.close',
+      connectionId: args.connectionId,
+      ownerId: session.ownerId,
+      actingUserId: args.actorUserId,
+      detail: {
+        reason: 'user_close',
+        initiator: 'user',
+        manual: true,
+        started_by_user_id: session.startedByUserId,
+      },
+    },
+    'success',
+  );
+  await sub.sessionRegistry.closeSession(args.taskId, args.connectionId, 'user_close');
+  return { ok: true };
+}
+
+/** No-session-at-all wording differs slightly per tool (kept close to the
+ * pre-Task-4 text so existing agent guidance / tests stay stable). */
+function noActiveSessionMessage(toolName: string): string {
+  if (toolName === 'SshConsoleSnapshot') {
+    return `${toolName}: no live session for this task. Call SshConsoleEnsure({connection_id}) first.`;
+  }
+  return (
+    `${toolName}: connection_id is required when this task has no active session. ` +
+    `Call SshListConnections() to find the right UUID, then SshConsoleEnsure({connection_id}) to open one.`
+  );
+}
+
+function notFoundMessage(toolName: string, connectionId: string): string {
+  if (toolName === 'SshConsoleSnapshot') {
+    return `${toolName}: no live session for connection ${connectionId} on this task. Call SshConsoleEnsure({connection_id: "${connectionId}"}) first.`;
+  }
+  return `${toolName}: no session for connection ${connectionId}; call SshConsoleEnsure({connection_id: "${connectionId}"}) first.`;
+}
+
+function ambiguousMessage(toolName: string, connectionIds: string[]): string {
+  return `${toolName}: connection_id required (multiple sessions open: ${connectionIds.join(', ')}).`;
+}
+
+function resolveErrorToToolResult(
+  toolName: string,
+  resolved: Exclude<ResolveSessionResult, { ok: true }>,
+): ToolResult {
+  switch (resolved.error) {
+    case 'not_found':
+      return err(notFoundMessage(toolName, resolved.connectionId));
+    case 'ambiguous':
+      return err(ambiguousMessage(toolName, resolved.connectionIds));
+    case 'none':
+      return err(noActiveSessionMessage(toolName));
+  }
+}
+
 // ──────────────────────────────────────────────────────────────────────
 // Tool definitions
 // ──────────────────────────────────────────────────────────────────────
@@ -355,7 +629,7 @@ export const TOOL_DEFS: Record<string, ToolDef> = {
     function: {
       name: 'SshConsoleEnsure',
       description:
-        'SSH console セッションを確保します（無ければ open、有れば再利用）。既存セッションが別 connection_id にある場合はデフォルトでエラー、force_replace=true で強制置換。詳細は ReadToolDoc({ name: "SshConsoleEnsure" })。',
+        'SSH console セッションを確保します（無ければ open、有れば再利用）。接続ごとに1セッションで、別の connection_id を指定すると既存を閉じずに新しいセッションを追加します（タスクあたりの上限あり）。force_replace=true は同じ接続のセッションを閉じて張り直す用途。詳細は ReadToolDoc({ name: "SshConsoleEnsure" })。',
       parameters: {
         type: 'object',
         properties: {
@@ -364,7 +638,7 @@ export const TOOL_DEFS: Record<string, ToolDef> = {
           rows: { type: 'number' },
           force_replace: {
             type: 'boolean',
-            description: '既存セッションが別 connection_id にあるとき、true なら旧セッションを閉じて新規 open。default=false (mismatch ならエラー返却)。',
+            description: '同じ connection_id の既存セッションを閉じて新規 open し直す（詰まり時の再接続用）。default=false。別の接続を開く場合は不要——自動で追加される。',
           },
         },
         required: ['connection_id'],
@@ -376,11 +650,11 @@ export const TOOL_DEFS: Record<string, ToolDef> = {
     function: {
       name: 'SshConsoleSend',
       description:
-        'console に入力を送る。通常のシェルコマンド実行には SshConsoleRun を使うこと（blocking + exit_code 取得）。SshConsoleSend は対話操作 (vim/REPL/sudo/TUI) と本当の中断専用。printable な shell コマンドには server が自動で末尾に "\\n" を付加して実行する (例: "ls -la" でも実行される)。TUI 操作 / sudo password / control byte (Ctrl-C 等) は raw のまま送られるので "\\n" を含めるかどうかは呼び出し側次第。connection_id はタスクに active session があれば省略可。詳細は ReadToolDoc({ name: "SshConsoleSend" })。',
+        'console に入力を送る。通常のシェルコマンド実行には SshConsoleRun を使うこと（blocking + exit_code 取得）。SshConsoleSend は対話操作 (vim/REPL/sudo/TUI) と本当の中断専用。printable な shell コマンドには server が自動で末尾に "\\n" を付加して実行する (例: "ls -la" でも実行される)。TUI 操作 / sudo password / control byte (Ctrl-C 等) は raw のまま送られるので "\\n" を含めるかどうかは呼び出し側次第。connection_id は開いているセッションが1つなら省略可（複数なら必須）。詳細は ReadToolDoc({ name: "SshConsoleSend" })。',
       parameters: {
         type: 'object',
         properties: {
-          connection_id: { type: 'string', description: '省略時はこのタスクの active session を自動採用。明示する場合は active session の id と一致する必要あり (mismatch はエラー)。' },
+          connection_id: { type: 'string', description: '省略時は、開いているセッションが1つならそれを自動採用（複数開いている場合は connection_id 必須）。明示するとその接続のセッションに送る。' },
           input: { type: 'string' },
           wait_ms: { type: 'number' },
           confirm_interrupt: {
@@ -401,7 +675,7 @@ export const TOOL_DEFS: Record<string, ToolDef> = {
       parameters: {
         type: 'object',
         properties: {
-          connection_id: { type: 'string', description: '省略時はこのタスクの active session を自動採用。' },
+          connection_id: { type: 'string', description: '省略時は、開いているセッションが1つならそれを自動採用（複数開いている場合は connection_id 必須）。' },
           kind: { type: 'string', enum: ['screen', 'scrollback'] },
           max_bytes: { type: 'number' },
         },
@@ -419,7 +693,7 @@ export const TOOL_DEFS: Record<string, ToolDef> = {
         type: 'object',
         properties: {
           command: { type: 'string', description: '実行するシェルコマンド。' },
-          connection_id: { type: 'string', description: '省略時はこのタスクの active session を自動採用。' },
+          connection_id: { type: 'string', description: '省略時は、開いているセッションが1つならそれを自動採用（複数開いている場合は connection_id 必須）。' },
           timeout_ms: { type: 'number', description: 'タイムアウト（ms）。デフォルト 120000、最大 600000。タイムアウト時はコマンドを kill しない。' },
           idle_ms: { type: 'number', description: '出力が途切れた後のアイドル判定時間（ms）。0=無効（デフォルト）。' },
         },
@@ -509,12 +783,22 @@ async function ensureSessionInternal(
       forceReplace: input.force_replace === true,
       initiator: 'agent',
       jobId: ctx.jobId ?? undefined,
+      // Per-space isolation (spec §11): the worker already resolved ctx.spaceId
+      // (personal-WS NULL→personal-space). Pass it so the console gate matches
+      // SshExec — otherwise workspace-registered connections fail space_mismatch.
+      spaceId: ctx.spaceId ?? null,
     },
   );
 
   if (!result.ok || !result.session) {
     return err(result.message ?? 'SshConsoleEnsure: failed to open session.');
   }
+  // Record this as the agent's current connection for this task — Send /
+  // Run / Snapshot resolve an omitted connection_id to this pointer when
+  // the task has more than one live session. ctx.taskId is guaranteed
+  // non-empty here (openConsoleSession already rejected an empty taskId
+  // with 'missing_task_context' before result.ok could be true).
+  markAgentConnection(ctx.taskId as string, result.session.connectionId);
   return { opened: !result.alreadyActive, session: result.session };
 }
 
@@ -589,28 +873,18 @@ async function sendInput(
     return err('SshConsoleSend: this tool requires a local task context (ctx.taskId).');
   }
 
-  // Resolve connection_id: explicit arg or active session's id.
-  // - omitted + active session → use active session's connection_id (recommended)
-  // - omitted + no active session → instruct the agent to discover via SshListConnections
-  // - explicit + matches active → fine
-  // - explicit + mismatches active → reject with surface of the actual id (don't auto-swap)
-  let connectionId = typeof input.connection_id === 'string' ? input.connection_id : '';
-  const existingSession = sub.sessionRegistry.get(localTaskId);
-  if (!connectionId) {
-    if (!existingSession) {
-      return err(
-        'SshConsoleSend: connection_id is required when this task has no active session. ' +
-          'Call SshListConnections() to find the right UUID, then SshConsoleEnsure({connection_id}) to open one.',
-      );
-    }
-    connectionId = existingSession.connectionId;
-  } else if (existingSession && existingSession.connectionId !== connectionId) {
-    return err(
-      `SshConsoleSend: this task has an active session on connection ${existingSession.connectionId}, not ${connectionId}. ` +
-        `Either omit connection_id (uses the active one), pass connection_id="${existingSession.connectionId}", ` +
-        `or call SshConsoleEnsure({connection_id: "${connectionId}", force_replace: true}) to swap intentionally.`,
-    );
+  // Resolve the target session: explicit connection_id (exact lookup, no
+  // auto-open), else the agent's current-connection pointer, else the
+  // task's sole live session, else an ambiguity/no-session error. See
+  // resolveTargetSession for the full precedence (Task 4: a task can now
+  // hold several live sessions at once).
+  const connectionIdInput =
+    typeof input.connection_id === 'string' && input.connection_id ? input.connection_id : undefined;
+  const resolved = resolveTargetSession(sub, localTaskId, connectionIdInput);
+  if (!resolved.ok) {
+    return resolveErrorToToolResult('SshConsoleSend', resolved);
   }
+  const { session, connectionId } = resolved;
 
   // Many LLMs serialize tool args correctly so "\n" in input means an
   // actual LF byte. But some surface escape sequences *literally*
@@ -644,18 +918,9 @@ async function sendInput(
     );
   }
 
-  // Find-or-open the session. We already resolved/validated connectionId
-  // above (either matches the existing session or there is none), so
-  // ensureSessionInternal will always reuse the live session or open a
-  // brand-new one without ever hitting its mismatch-reject path.
-  const ensured = await ensureSessionInternal({ connection_id: connectionId }, ctx, sub);
-  if ('isError' in ensured) return ensured;
-  const session = ensured.session;
-
-  // Re-run preflight to get the canonical connection for deny-list eval +
-  // audit logging. (ensureSessionInternal already did the same call but
-  // discarded the result; we want the latest snapshot in case admin
-  // policy changed between Ensure and Send.)
+  // Preflight to get the canonical connection for deny-list eval + audit
+  // logging (latest snapshot in case admin policy changed since the
+  // session was opened).
   const pre = preflight({
     toolName: 'SshExec',
     connectionId,
@@ -741,6 +1006,11 @@ async function sendInput(
 
   // (auto-newline was applied earlier — agent self-correction is no longer needed)
 
+  // Record this as the agent's current connection now that the write
+  // actually happened (the interrupt-guard early return above never
+  // reaches here, since it didn't write anything).
+  markAgentConnection(localTaskId, connectionId);
+
   sub.auditRepo.beginAndComplete(
     {
       action: 'ssh.console.send',
@@ -799,26 +1069,17 @@ async function snapshot(
   }
 
   // Snapshot does NOT auto-Ensure — if there is no live session, the LLM
-  // should call SshConsoleEnsure / Send first.
-  const session = sub.sessionRegistry.get(localTaskId);
-  if (!session) {
-    return err(
-      'SshConsoleSnapshot: no live session for this task. Call SshConsoleEnsure({connection_id}) first.',
-    );
-  }
-
-  // Resolve connection_id: explicit arg or the active session's id.
-  // Mismatch (explicit id ≠ active session id) → reject with surface of
-  // the actual id so the LLM can self-correct.
-  let connectionId = typeof input.connection_id === 'string' ? input.connection_id : '';
-  if (!connectionId) {
-    connectionId = session.connectionId;
-  } else if (session.connectionId !== connectionId) {
-    return err(
-      `SshConsoleSnapshot: this task has an active session on connection ${session.connectionId}, not ${connectionId}. ` +
-        `Either omit connection_id (uses the active one) or pass connection_id="${session.connectionId}".`,
-    );
+  // should call SshConsoleEnsure / Send first. resolveTargetSession routes
+  // an omitted connection_id to the agent's current-connection pointer, the
+  // task's sole live session, or a clear ambiguity/no-session error (Task 4:
+  // a task can now hold several live sessions at once).
+  const connectionIdInput =
+    typeof input.connection_id === 'string' && input.connection_id ? input.connection_id : undefined;
+  const resolved = resolveTargetSession(sub, localTaskId, connectionIdInput);
+  if (!resolved.ok) {
+    return resolveErrorToToolResult('SshConsoleSnapshot', resolved);
   }
+  const { session, connectionId } = resolved;
 
   // Preflight (for audit + access check; we still want to see read attempts
   // in the audit log).
@@ -832,6 +1093,10 @@ async function snapshot(
   if (!pre.ok) return pre.error;
   const { connection, actingUserId, pieceName } = pre;
 
+  // Record this as the agent's current connection now that access has been
+  // confirmed (the operation cannot fail past this point).
+  markAgentConnection(localTaskId, connectionId);
+
   const kind = input.kind === 'scrollback' ? 'scrollback' : 'screen';
 
   if (kind === 'screen') {
@@ -952,28 +1217,17 @@ async function run(
     return err('SshConsoleRun: this tool requires a local task context (ctx.taskId).');
   }
 
-  // Resolve connection_id: explicit arg or active session's id.
-  let connectionId = typeof input.connection_id === 'string' ? input.connection_id : '';
-  const existingSession = sub.sessionRegistry.get(localTaskId);
-  if (!connectionId) {
-    if (!existingSession) {
-      return err(
-        'SshConsoleRun: connection_id is required when this task has no active session. ' +
-          'Call SshListConnections() to find the right UUID, then SshConsoleEnsure({connection_id}) to open one.',
-      );
-    }
-    connectionId = existingSession.connectionId;
-  } else if (existingSession && existingSession.connectionId !== connectionId) {
-    return err(
-      `SshConsoleRun: this task has an active session on connection ${existingSession.connectionId}, not ${connectionId}. ` +
-        `Either omit connection_id (uses the active one) or pass connection_id="${existingSession.connectionId}".`,
-    );
+  // Resolve the target session: explicit connection_id (exact lookup, no
+  // auto-open), else the agent's current-connection pointer, else the
+  // task's sole live session, else an ambiguity/no-session error (Task 4:
+  // a task can now hold several live sessions at once).
+  const connectionIdInput =
+    typeof input.connection_id === 'string' && input.connection_id ? input.connection_id : undefined;
+  const resolved = resolveTargetSession(sub, localTaskId, connectionIdInput);
+  if (!resolved.ok) {
+    return resolveErrorToToolResult('SshConsoleRun', resolved);
   }
-
-  // Find-or-open session.
-  const ensured = await ensureSessionInternal({ connection_id: connectionId }, ctx, sub);
-  if ('isError' in ensured) return ensured;
-  const session = ensured.session;
+  const { session, connectionId } = resolved;
 
   // Preflight for audit + access check.
   const pre = preflight({
@@ -986,6 +1240,11 @@ async function run(
   if (!pre.ok) return pre.error;
   const { connection, actingUserId, pieceName } = pre;
 
+  // Record this as the agent's current connection now that access has been
+  // confirmed (the run below can still time out / be interrupted, but it
+  // did act on this connection).
+  markAgentConnection(localTaskId, connectionId);
+
   // Deny-check the command.
   const denyResult = checkConsoleInput(
     command + '\n',