diff --git a/src/main/ipc/pty/ipc/spawn-commit-persist.ts b/src/main/ipc/pty/ipc/spawn-commit-persist.ts index 03b7ef9e0bd..c2e630390e1 100644 --- a/src/main/ipc/pty/ipc/spawn-commit-persist.ts +++ b/src/main/ipc/pty/ipc/spawn-commit-persist.ts @@ -3,6 +3,7 @@ import { closeStartupQueryAuthorityForPty } from '../provider/registry' import { createTerminalSessionStateSaveFailureMessage } from '../../../../shared/terminal-session-state-save-failure' import { recordCodexPaneAccountForSpawn } from '../host-env/codex-home' import { persistAdmittedStablePaneBinding } from '../pane/stable-owner' +import { persistPaneBinding } from '../pane/pane-binding-partition' import { swapReplacedPaneBinding } from '../pane/pane-owner-replacement' import { pendingByPaneKey, @@ -25,6 +26,7 @@ export async function persistPtyIpcSpawnCommit(ctx: PtyIpcSpawnState): Promise

, + binding: Parameters[0], + connectionId: string | null | undefined +): Promise { + return ORCA_LAYOUT_OWNER_IS_MAIN && connectionId + ? store.persistPtyBinding(binding, toSshExecutionHostId(connectionId)) + : store.persistPtyBinding(binding) +} diff --git a/src/main/ipc/pty/pane/pane-owner-replacement.ts b/src/main/ipc/pty/pane/pane-owner-replacement.ts index b3aaef0ea5b..89b742fb468 100644 --- a/src/main/ipc/pty/pane/pane-owner-replacement.ts +++ b/src/main/ipc/pty/pane/pane-owner-replacement.ts @@ -4,6 +4,7 @@ import type { Store } from '../../../persistence' import type { PersistPtyBindingArgs } from '../../../persistence/loading-store/pty-binding-persistence' import type { OrcaRuntimeService } from '../../../runtime/orca-runtime' import { resolveStablePaneOwner, type StablePaneOwner } from './stable-owner' +import { ORCA_LAYOUT_OWNER_IS_MAIN } from '../../../../shared/layout-owner-switch' /** The owner a replacing spawn stopped. Its binding stays until the replacement's bind swaps it. */ export type ReplacedPaneOwner = { @@ -69,10 +70,9 @@ export function swapReplacedPaneBinding( replaced: ReplacedPaneOwner ): () => PersistPtyBindingArgs | null { return () => { - const bound = - store.getWorkspaceSession().terminalLayoutsByTabId?.[binding.tabId]?.ptyIdsByLeafId?.[ - binding.leafId - ] + const bound = store.getWorkspaceSession( + ORCA_LAYOUT_OWNER_IS_MAIN ? replaced.pane?.hostId : undefined + ).terminalLayoutsByTabId?.[binding.tabId]?.ptyIdsByLeafId?.[binding.leafId] // Unbound is ok: the renderer's pre-connect clear can withdraw it (split tab). return bound === undefined || bound === replaced.ptyId ? binding : null } diff --git a/src/main/ipc/pty/pane/stable-owner.ts b/src/main/ipc/pty/pane/stable-owner.ts index f190e5ab2cd..fb1b0ea1035 100644 --- a/src/main/ipc/pty/pane/stable-owner.ts +++ b/src/main/ipc/pty/pane/stable-owner.ts @@ -20,6 +20,7 @@ import { import { ptyIncarnationById } from '../provider/ownership-state' import { releaseAbsentPaneOwner } from './absent-pane-owner' import { spawnCommitBindingOrigin } from '../../../persistence/loading-store/pty-binding-span' +import { persistPaneBinding } from './pane-binding-partition' export type StablePaneOwner = { handle?: string @@ -197,22 +198,27 @@ export async function persistAdmittedStablePaneBinding(args: { owner: StablePaneOwner | null result: PtySpawnResult worktreeId: string | undefined + connectionId: string | null | undefined startupCwd: string | undefined }): Promise { const expectedBinding = stablePanePersistenceFence(args.owner) if (!args.store || !args.owner || !args.worktreeId || !expectedBinding) { return false } - const persisted = await args.store.persistPtyBinding({ - worktreeId: args.worktreeId, - tabId: args.owner.tabId, - leafId: args.owner.leafId, - ptyId: args.result.id, - ...(args.result.incarnationId ? { incarnationId: args.result.incarnationId } : {}), - ...(args.startupCwd ? { startupCwd: args.startupCwd } : {}), - expectedBinding, - origin: spawnCommitBindingOrigin(args.result) - }) + const persisted = await persistPaneBinding( + args.store, + { + worktreeId: args.worktreeId, + tabId: args.owner.tabId, + leafId: args.owner.leafId, + ptyId: args.result.id, + ...(args.result.incarnationId ? { incarnationId: args.result.incarnationId } : {}), + ...(args.startupCwd ? { startupCwd: args.startupCwd } : {}), + expectedBinding, + origin: spawnCommitBindingOrigin(args.result) + }, + args.connectionId + ) if (persisted === false) { throw new Error('terminal_pane_owner_changed') } diff --git a/src/main/ipc/pty/runtime/spawn-commit.ts b/src/main/ipc/pty/runtime/spawn-commit.ts index b72c9b2a79b..d42281dc482 100644 --- a/src/main/ipc/pty/runtime/spawn-commit.ts +++ b/src/main/ipc/pty/runtime/spawn-commit.ts @@ -23,6 +23,7 @@ import { seedResumeLaunchIdentityFromSpawn } from '../pane/resume-launch-identit import { seedHeadlessTerminalFromSpawnResult } from '../pane/terminal-spawn-restore' import { recordPtySpawnTelemetry } from '../pane/spawn-telemetry' import { persistAdmittedStablePaneBinding } from '../pane/stable-owner' +import { persistPaneBinding } from '../pane/pane-binding-partition' import { isNativeWindowsLocalPtySpawn, markNativeWindowsConptyPty @@ -66,6 +67,7 @@ async function commitReservedRuntimePtySpawn(ctx: RuntimePtySpawnState) { owner: ctx.stablePaneOwner, result: ctx.result, worktreeId: ctx.hostSessionBinding?.worktreeId, + connectionId: args.connectionId, startupCwd: ctx.cwd }) } catch (error) { @@ -152,7 +154,7 @@ async function commitReservedRuntimePtySpawn(ctx: RuntimePtySpawnState) { ...(placement ? { placement } : {}), origin: spawnCommitBindingOrigin(ctx.result, expectedSourceBinding) } - const persisted = await store.persistPtyBinding(binding) + const persisted = await persistPaneBinding(store, binding, args.connectionId) if (persisted === false) { throw new Error('terminal_split_source_not_found') } diff --git a/src/main/persistence/loading-store/pty-spawn-ssh-binding-partition.test.ts b/src/main/persistence/loading-store/pty-spawn-ssh-binding-partition.test.ts new file mode 100644 index 00000000000..bd7881cd516 --- /dev/null +++ b/src/main/persistence/loading-store/pty-spawn-ssh-binding-partition.test.ts @@ -0,0 +1,119 @@ +import { expect, it, vi } from 'vitest' +import { admitRuntimeSpawnPane, fixture } from './profile-state-delayed-authority-fixture' +import { commitRuntimePtySpawn } from '../../ipc/pty/runtime/spawn-commit' +import { createRuntimePtySpawnState } from '../../ipc/pty/runtime/spawn-state' +import type { PtyRuntimeControllerDeps } from '../../ipc/pty/runtime/controller-deps' +import { OrcaRuntimeService } from '../../runtime/orca-runtime' +import { persistPtyIpcSpawnCommit } from '../../ipc/pty/ipc/spawn-commit-persist' +import { createPtyIpcSpawnState } from '../../ipc/pty/ipc/spawn-state' +import type { PtySpawnIpcDeps } from '../../ipc/pty/ipc/spawn-types' +import { getDefaultWorkspaceSession } from '../../../shared/constants' +import type { WorkspaceSessionState } from '../../../shared/workspace-session-state-types' + +vi.mock('../../../shared/layout-owner-switch', () => ({ ORCA_LAYOUT_OWNER_IS_MAIN: true })) +vi.mock('../../telemetry/client', () => ({ track: vi.fn() })) +vi.mock('../../telemetry/cohort-classifier', () => ({ + getCohortAtEmit: () => ({ nth_repo_added: 2 }) +})) +vi.mock('../../ssh/ssh-config-parser', () => ({ + loadUserSshConfig: () => ({ hosts: [] }), + sshConfigHostsToTargets: () => [] +})) + +const SSH_HOST_ID = 'ssh:ssh-1' +const pane = { + worktreeId: 'repo-ssh::/remote/wt', + tabId: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', + leafId: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb' +} + +function sshSession(ptyId: string | null): WorkspaceSessionState { + return { + ...getDefaultWorkspaceSession(), + tabsByWorktree: { + [pane.worktreeId]: [ + { + id: pane.tabId, + worktreeId: pane.worktreeId, + title: 'Terminal 1', + customTitle: null, + color: null, + sortOrder: 0, + createdAt: 1, + ptyId + } + ] + }, + terminalLayoutsByTabId: { + [pane.tabId]: { + root: { type: 'leaf', leafId: pane.leafId }, + activeLeafId: pane.leafId, + expandedLeafId: null, + ...(ptyId ? { ptyIdsByLeafId: { [pane.leafId]: ptyId } } : {}) + } + } + } +} + +function boundPtyId(session: WorkspaceSessionState): string | undefined { + return session.terminalLayoutsByTabId[pane.tabId]?.ptyIdsByLeafId?.[pane.leafId] +} + +// Main keeps its own layout over a window save, so a binding only a save carries would be lost. +it.each([ + { spawn: 'fresh', sessionId: undefined }, + { spawn: 'reattach', sessionId: 'ssh:ssh-1@@remote-pty' } +])('binds a window-spawned SSH pane ($spawn) in its SSH partition', async ({ sessionId }) => { + const { store } = await fixture() + store.setWorkspaceSession(sshSession(null), SSH_HOST_ID) + const local = store.getWorkspaceSession('local') + const ptyId = 'ssh:ssh-1@@remote-pty' + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: only the commit's persist step runs; its store is real. + const deps = { runtime: new OrcaRuntimeService(store), store } as PtySpawnIpcDeps + const ctx = createPtyIpcSpawnState(deps, { + ...pane, + connectionId: 'ssh-1', + ...(sessionId ? { sessionId } : {}), + cols: 80, + rows: 24 + }) + ctx.result = { id: ptyId, incarnationId: 'cccccccc-cccc-4ccc-8ccc-cccccccccccc' } + ctx.stablePaneOwner = null + ctx.validatedLeafId = pane.leafId + + await persistPtyIpcSpawnCommit(ctx) + + expect(boundPtyId(store.getWorkspaceSession(SSH_HOST_ID))).toBe(ptyId) + expect(store.getWorkspaceSession('local').tabsByWorktree).toEqual(local.tabsByWorktree) +}) + +it('binds a runtime-started SSH pane in the SSH partition it was admitted to', async () => { + const { store } = await fixture() + const runtime = new OrcaRuntimeService(store) + const ptyId = 'ssh:ssh-1@@runtime-pty' + const incarnationId = 'dddddddd-dddd-4ddd-8ddd-dddddddddddd' + runtime.onPtySpawned(ptyId, incarnationId) + runtime.beginPtyRegistration(ptyId, incarnationId) + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: only commit runs; its runtime/store are real and preflight-only dependencies are unreachable. + const deps = { + runtime, + store, + options: {}, + sendPtySpawnedToRenderer: vi.fn() + } as PtyRuntimeControllerDeps + const ctx = createRuntimePtySpawnState(deps, { + ...pane, + connectionId: 'ssh-1', + cols: 80, + rows: 24 + }) + ctx.result = { id: ptyId, incarnationId } + ctx.stablePaneOwner = null + await admitRuntimeSpawnPane(store, pane, 'ssh-1') + ctx.hostSessionBinding = { store, ...pane } + ctx.metadataLeafId = pane.leafId + + await commitRuntimePtySpawn(ctx) + + expect(boundPtyId(store.getWorkspaceSession(SSH_HOST_ID))).toBe(ptyId) +})