mirror of
https://github.com/stablyai/orca.git
synced 2026-10-27 16:02:42 +00:00
feat(ssh): agent-status over SSH (#1702)
Co-authored-by: Orca <help@stably.ai>
This commit is contained in:
Brennan Benson
Orca
@@ -0,0 +1,323 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
import type { Store } from '../persistence'
|
||||
import type { SshPortForwardManager } from './ssh-port-forward'
|
||||
import type { SshConnection } from './ssh-connection'
|
||||
import type { MultiplexerTransport } from './ssh-channel-multiplexer'
|
||||
import type { AgentHookRelayEnvelope } from '../../shared/agent-hook-relay'
|
||||
import { RelayDispatcher } from '../../relay/dispatcher'
|
||||
import {
|
||||
AGENT_HOOK_NOTIFICATION_METHOD,
|
||||
AGENT_HOOK_REQUEST_REPLAY_METHOD,
|
||||
ORCA_FEATURE_REMOTE_AGENT_HOOKS_ENV,
|
||||
REMOTE_AGENT_HOOK_ENV
|
||||
} from '../../shared/agent-hook-relay'
|
||||
import { agentHookServer, _internals as agentHookInternals } from '../agent-hooks/server'
|
||||
import { getSshPtyProvider } from '../ipc/pty'
|
||||
|
||||
vi.mock('./ssh-relay-deploy', () => ({
|
||||
deployAndLaunchRelay: vi.fn()
|
||||
}))
|
||||
|
||||
const { deployAndLaunchRelay } = await import('./ssh-relay-deploy')
|
||||
const { SshRelaySession } = await import('./ssh-relay-session')
|
||||
|
||||
type CapturedStatus = {
|
||||
paneKey: string
|
||||
tabId?: string
|
||||
worktreeId?: string
|
||||
connectionId: string | null
|
||||
payload: {
|
||||
state: string
|
||||
prompt: string
|
||||
agentType?: string
|
||||
toolName?: string
|
||||
}
|
||||
}
|
||||
|
||||
type FakeRelay = {
|
||||
transport: MultiplexerTransport
|
||||
dispatcher: RelayDispatcher
|
||||
ptySpawnRequests: Record<string, unknown>[]
|
||||
replayEnvelopes: AgentHookRelayEnvelope[]
|
||||
notifyAgentHook: (envelope: AgentHookRelayEnvelope | Record<string, unknown>) => void
|
||||
dispose: () => void
|
||||
}
|
||||
|
||||
// Why: mock below SSH at the relay transport boundary so CI covers session,
|
||||
// mux, provider, and hook-ingest wiring without relying on a local sshd.
|
||||
function createFakeRelay(): FakeRelay {
|
||||
let relayFeed: ((data: Buffer) => void) | null = null
|
||||
const clientDataCallbacks: ((data: Buffer) => void)[] = []
|
||||
const clientCloseCallbacks: (() => void)[] = []
|
||||
const ptySpawnRequests: Record<string, unknown>[] = []
|
||||
const replayEnvelopes: AgentHookRelayEnvelope[] = []
|
||||
|
||||
const transport: MultiplexerTransport = {
|
||||
write: (data) => {
|
||||
setImmediate(() => relayFeed?.(data))
|
||||
},
|
||||
onData: (cb) => {
|
||||
clientDataCallbacks.push(cb)
|
||||
},
|
||||
onClose: (cb) => {
|
||||
clientCloseCallbacks.push(cb)
|
||||
},
|
||||
close: () => {
|
||||
for (const cb of clientCloseCallbacks) {
|
||||
cb()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const dispatcher = new RelayDispatcher((data) => {
|
||||
setImmediate(() => {
|
||||
for (const cb of clientDataCallbacks) {
|
||||
cb(data)
|
||||
}
|
||||
})
|
||||
})
|
||||
relayFeed = (data) => dispatcher.feed(data)
|
||||
|
||||
dispatcher.onRequest('session.resolveHome', async (params) => ({
|
||||
resolvedPath: params.path === '~' ? '/home/orca' : params.path
|
||||
}))
|
||||
dispatcher.onRequest('git.listWorktrees', async () => [])
|
||||
dispatcher.onRequest('ports.detect', async () => ({ ports: [], platform: 'linux' }))
|
||||
dispatcher.onRequest('pty.spawn', async (params) => {
|
||||
ptySpawnRequests.push(params)
|
||||
return { id: `remote-pty-${ptySpawnRequests.length}` }
|
||||
})
|
||||
dispatcher.onRequest(AGENT_HOOK_REQUEST_REPLAY_METHOD, async () => {
|
||||
// Why: relay replay must arrive after Orca wires its listener and before
|
||||
// the request resolves, matching the real relay ordering contract.
|
||||
for (const envelope of replayEnvelopes) {
|
||||
dispatcher.notify(
|
||||
AGENT_HOOK_NOTIFICATION_METHOD,
|
||||
envelope as unknown as Record<string, unknown>
|
||||
)
|
||||
}
|
||||
return { replayed: replayEnvelopes.length }
|
||||
})
|
||||
|
||||
return {
|
||||
transport,
|
||||
dispatcher,
|
||||
ptySpawnRequests,
|
||||
replayEnvelopes,
|
||||
notifyAgentHook: (envelope) => {
|
||||
dispatcher.notify(AGENT_HOOK_NOTIFICATION_METHOD, envelope as Record<string, unknown>)
|
||||
},
|
||||
dispose: () => dispatcher.dispose()
|
||||
}
|
||||
}
|
||||
|
||||
function createSession(targetId: string): InstanceType<typeof SshRelaySession> {
|
||||
const store = {
|
||||
getRepos: vi.fn().mockReturnValue([]),
|
||||
getSshRemotePtyLeases: vi.fn().mockReturnValue([]),
|
||||
markSshRemotePtyLease: vi.fn(),
|
||||
markSshRemotePtyLeases: vi.fn()
|
||||
} as unknown as Store
|
||||
const portForwardManager = {
|
||||
removeAllForwards: vi.fn().mockResolvedValue(undefined)
|
||||
} as unknown as SshPortForwardManager
|
||||
const getMainWindow = vi.fn().mockReturnValue({
|
||||
isDestroyed: () => false,
|
||||
webContents: { send: vi.fn() }
|
||||
})
|
||||
return new SshRelaySession(targetId, getMainWindow, store, portForwardManager)
|
||||
}
|
||||
|
||||
async function waitForStatusCount(events: CapturedStatus[], count: number): Promise<void> {
|
||||
await vi.waitFor(() => expect(events).toHaveLength(count), { timeout: 1500 })
|
||||
}
|
||||
|
||||
function captureAgentStatuses(events: CapturedStatus[]): void {
|
||||
agentHookServer.setListener((event) => {
|
||||
events.push({
|
||||
paneKey: event.paneKey,
|
||||
tabId: event.tabId,
|
||||
worktreeId: event.worktreeId,
|
||||
connectionId: event.connectionId,
|
||||
payload: {
|
||||
state: event.payload.state,
|
||||
prompt: event.payload.prompt,
|
||||
agentType: event.payload.agentType,
|
||||
toolName: event.payload.toolName
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
function makeEnvelope(overrides: Partial<AgentHookRelayEnvelope> = {}): AgentHookRelayEnvelope {
|
||||
return {
|
||||
source: 'codex',
|
||||
paneKey: 'tab-ssh:0',
|
||||
tabId: 'tab-ssh',
|
||||
worktreeId: 'wt-ssh',
|
||||
connectionId: null,
|
||||
env: REMOTE_AGENT_HOOK_ENV,
|
||||
version: '1',
|
||||
payload: {
|
||||
state: 'working',
|
||||
prompt: 'remote prompt',
|
||||
agentType: 'codex'
|
||||
},
|
||||
...overrides
|
||||
}
|
||||
}
|
||||
|
||||
describe('SshRelaySession agent hooks over a fake relay transport', () => {
|
||||
let previousRemoteHooksFlag: string | undefined
|
||||
let warnSpy: ReturnType<typeof vi.spyOn>
|
||||
let session: InstanceType<typeof SshRelaySession> | null = null
|
||||
let relay: FakeRelay | null = null
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
previousRemoteHooksFlag = process.env[ORCA_FEATURE_REMOTE_AGENT_HOOKS_ENV]
|
||||
process.env[ORCA_FEATURE_REMOTE_AGENT_HOOKS_ENV] = '1'
|
||||
warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {})
|
||||
agentHookServer.setListener(null)
|
||||
agentHookInternals.resetCachesForTests()
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
session?.dispose()
|
||||
relay?.dispose()
|
||||
session = null
|
||||
relay = null
|
||||
agentHookServer.setListener(null)
|
||||
agentHookInternals.resetCachesForTests()
|
||||
warnSpy.mockRestore()
|
||||
if (previousRemoteHooksFlag === undefined) {
|
||||
delete process.env[ORCA_FEATURE_REMOTE_AGENT_HOOKS_ENV]
|
||||
} else {
|
||||
process.env[ORCA_FEATURE_REMOTE_AGENT_HOOKS_ENV] = previousRemoteHooksFlag
|
||||
}
|
||||
})
|
||||
|
||||
it('establishes through a fake relay, spawns a remote PTY, and forwards agent status', async () => {
|
||||
relay = createFakeRelay()
|
||||
vi.mocked(deployAndLaunchRelay).mockResolvedValue({
|
||||
transport: relay.transport,
|
||||
platform: 'linux-x64'
|
||||
})
|
||||
const events: CapturedStatus[] = []
|
||||
captureAgentStatuses(events)
|
||||
|
||||
session = createSession('conn-fake')
|
||||
await session.establish({} as SshConnection)
|
||||
|
||||
const provider = getSshPtyProvider('conn-fake')
|
||||
expect(provider).toBeDefined()
|
||||
const spawn = await provider!.spawn({
|
||||
cols: 120,
|
||||
rows: 40,
|
||||
cwd: '/home/orca/project',
|
||||
env: {
|
||||
ORCA_PANE_KEY: 'tab-ssh:0',
|
||||
ORCA_TAB_ID: 'tab-ssh',
|
||||
ORCA_WORKTREE_ID: 'wt-ssh'
|
||||
}
|
||||
})
|
||||
|
||||
expect(spawn.id).toBe('remote-pty-1')
|
||||
expect(relay.ptySpawnRequests).toHaveLength(1)
|
||||
expect(relay.ptySpawnRequests[0]).toMatchObject({
|
||||
cwd: '/home/orca/project',
|
||||
env: {
|
||||
ORCA_PANE_KEY: 'tab-ssh:0',
|
||||
ORCA_TAB_ID: 'tab-ssh',
|
||||
ORCA_WORKTREE_ID: 'wt-ssh'
|
||||
}
|
||||
})
|
||||
|
||||
relay.notifyAgentHook(makeEnvelope())
|
||||
|
||||
await waitForStatusCount(events, 1)
|
||||
expect(events[0]).toEqual({
|
||||
paneKey: 'tab-ssh:0',
|
||||
tabId: 'tab-ssh',
|
||||
worktreeId: 'wt-ssh',
|
||||
connectionId: 'conn-fake',
|
||||
payload: {
|
||||
state: 'working',
|
||||
prompt: 'remote prompt',
|
||||
agentType: 'codex',
|
||||
toolName: undefined
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
it('asks the fake relay for cached hook replay after the session wires its listener', async () => {
|
||||
relay = createFakeRelay()
|
||||
relay.replayEnvelopes.push(
|
||||
makeEnvelope({
|
||||
paneKey: 'tab-replay:0',
|
||||
tabId: 'tab-replay',
|
||||
worktreeId: 'wt-replay',
|
||||
payload: {
|
||||
state: 'waiting',
|
||||
prompt: 'cached remote prompt',
|
||||
agentType: 'claude',
|
||||
toolName: 'Bash'
|
||||
}
|
||||
})
|
||||
)
|
||||
vi.mocked(deployAndLaunchRelay).mockResolvedValue({
|
||||
transport: relay.transport,
|
||||
platform: 'linux-x64'
|
||||
})
|
||||
const events: CapturedStatus[] = []
|
||||
captureAgentStatuses(events)
|
||||
|
||||
session = createSession('conn-replay')
|
||||
await session.establish({} as SshConnection)
|
||||
|
||||
await waitForStatusCount(events, 1)
|
||||
expect(events[0]).toMatchObject({
|
||||
paneKey: 'tab-replay:0',
|
||||
tabId: 'tab-replay',
|
||||
worktreeId: 'wt-replay',
|
||||
connectionId: 'conn-replay',
|
||||
payload: {
|
||||
state: 'waiting',
|
||||
prompt: 'cached remote prompt',
|
||||
agentType: 'claude',
|
||||
toolName: 'Bash'
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
it('drops malformed remote hook notifications at Orca main before caching', async () => {
|
||||
relay = createFakeRelay()
|
||||
vi.mocked(deployAndLaunchRelay).mockResolvedValue({
|
||||
transport: relay.transport,
|
||||
platform: 'linux-x64'
|
||||
})
|
||||
const events: CapturedStatus[] = []
|
||||
captureAgentStatuses(events)
|
||||
|
||||
session = createSession('conn-validate')
|
||||
await session.establish({} as SshConnection)
|
||||
|
||||
relay.notifyAgentHook({
|
||||
source: 'codex',
|
||||
paneKey: 'tab-bad:0',
|
||||
connectionId: null,
|
||||
env: REMOTE_AGENT_HOOK_ENV,
|
||||
version: '1',
|
||||
payload: {
|
||||
state: 'not-a-real-state',
|
||||
prompt: 'should not be cached',
|
||||
agentType: 'codex'
|
||||
}
|
||||
})
|
||||
|
||||
await new Promise((resolve) => setImmediate(resolve))
|
||||
expect(events).toHaveLength(0)
|
||||
expect(agentHookServer.getStatusSnapshot()).toEqual([])
|
||||
})
|
||||
})
|
||||
@@ -16,6 +16,12 @@ import { SshChannelMultiplexer } from './ssh-channel-multiplexer'
|
||||
import { SshPtyProvider, isSshPtyNotFoundError } from '../providers/ssh-pty-provider'
|
||||
import { SshFilesystemProvider } from '../providers/ssh-filesystem-provider'
|
||||
import { SshGitProvider } from '../providers/ssh-git-provider'
|
||||
import { agentHookServer } from '../agent-hooks/server'
|
||||
import {
|
||||
AGENT_HOOK_NOTIFICATION_METHOD,
|
||||
AGENT_HOOK_REQUEST_REPLAY_METHOD,
|
||||
isRemoteAgentHooksEnabled
|
||||
} from '../../shared/agent-hook-relay'
|
||||
import {
|
||||
registerSshPtyProvider,
|
||||
unregisterSshPtyProvider,
|
||||
@@ -409,9 +415,81 @@ export class SshRelaySession {
|
||||
registerSshGitProvider(this.targetId, gitProvider)
|
||||
|
||||
this.wireUpPtyEvents(ptyProvider)
|
||||
this.wireUpAgentHookEvents(mux)
|
||||
return true
|
||||
}
|
||||
|
||||
// Why: route the relay's `agent.hook` JSON-RPC notification into Orca's
|
||||
// shared `agentHookServer` via `ingestRemote`. The wire envelope carries
|
||||
// `connectionId: null` (the relay does not know Orca's local handle); we
|
||||
// stamp the real value here from `this.targetId` so the renderer can drop
|
||||
// in-flight events for connections that have torn down. After wiring is
|
||||
// in place we kick off a request-driven replay so any cached payload from
|
||||
// before the channel was up survives the reconnect — see §5 Path 3.
|
||||
//
|
||||
// The Orca-side mux's `notificationHandlers` is a flat array — each
|
||||
// handler must filter by method name itself.
|
||||
private wireUpAgentHookEvents(mux: SshChannelMultiplexer): void {
|
||||
if (!isRemoteAgentHooksEnabled()) {
|
||||
return
|
||||
}
|
||||
mux.onNotification((method, params) => {
|
||||
if (method !== AGENT_HOOK_NOTIFICATION_METHOD) {
|
||||
return
|
||||
}
|
||||
const envelope = params as {
|
||||
paneKey?: unknown
|
||||
tabId?: unknown
|
||||
worktreeId?: unknown
|
||||
env?: unknown
|
||||
version?: unknown
|
||||
payload?: unknown
|
||||
}
|
||||
if (typeof envelope.paneKey !== 'string') {
|
||||
return
|
||||
}
|
||||
// Why: forward env/version verbatim so Orca's warn-once cross-build /
|
||||
// dev-vs-prod diagnostics fire on remote events the same as on local
|
||||
// ones — see docs/design/agent-status-over-ssh.md §3 ("Replay /
|
||||
// version mismatch") and the relay's wire envelope at
|
||||
// src/shared/agent-hook-relay.ts.
|
||||
agentHookServer.ingestRemote(
|
||||
{
|
||||
paneKey: envelope.paneKey,
|
||||
tabId: typeof envelope.tabId === 'string' ? envelope.tabId : undefined,
|
||||
worktreeId: typeof envelope.worktreeId === 'string' ? envelope.worktreeId : undefined,
|
||||
env: typeof envelope.env === 'string' ? envelope.env : undefined,
|
||||
version: typeof envelope.version === 'string' ? envelope.version : undefined,
|
||||
payload: envelope.payload
|
||||
},
|
||||
this.targetId
|
||||
)
|
||||
})
|
||||
|
||||
// Why: ask the relay to replay every cached paneKey it remembers. Issued
|
||||
// *after* the handler is wired so the request-driven replay shape
|
||||
// strictly trails our subscription on the dispatcher's single write
|
||||
// callback. Best-effort: a relay that does not know the method
|
||||
// (e.g. older relay binary) returns -32601, which we swallow.
|
||||
void mux.request(AGENT_HOOK_REQUEST_REPLAY_METHOD).catch((err) => {
|
||||
const code = (err as { code?: unknown })?.code
|
||||
if (code === -32601) {
|
||||
return
|
||||
}
|
||||
// Why: a normal disconnect/teardown rejects the in-flight request with
|
||||
// "Multiplexer disposed"; suppress the warn for that path so reconnect
|
||||
// cycles aren't noisy.
|
||||
if (mux.isDisposed()) {
|
||||
return
|
||||
}
|
||||
console.warn(
|
||||
`[ssh-relay-session] agent_hook.requestReplay failed for ${this.targetId}: ${
|
||||
err instanceof Error ? err.message : String(err)
|
||||
}`
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
private teardownProviders(reason: 'shutdown' | 'connection_lost'): void {
|
||||
this.muxDisposeCleanup?.()
|
||||
this.muxDisposeCleanup = null
|
||||
|
||||
Reference in New Issue
Block a user