From 4049bc77d6177fecc2e7307e32b5f221187da472 Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:25:44 -0700 Subject: [PATCH] feat(orcad): gate agent launches on readiness and refuse callers that leave while waiting (#35710) * feat(orcad): gate agent launches on accounts and cancel feature waits for departed callers requires: may read the parsed params, so terminal.create and session.tabs.createTerminal wait on accounts only when they start an agent; plain shells stay ungated. createAgentSession, ensureAgentSession, agentSession.create, agent.launch (unless it reuses a terminal), git.generateCommitMessage and the worker-launching orchestration methods also wait on accounts. The feature wait races the request signal and the connection close and refuses without running the handler. * feat(orcad): gate terminal.ensureAgentSession only for explicit launches An attach to an existing session is restore latency and needs no fresh account state; only kind 'explicit' may start an agent. * test: reset orcad feature readiness after every unit test A startOrcad test that fails partway left latches closed, so later tests in the same worker could wait 10 s on a gated method. --- config/scripts/ci-unit-files.mjs | 3 +- .../vitest-orcad-feature-readiness-reset.ts | 7 ++ .../rpc/agent-launch-feature-requirement.ts | 45 ++++++++ src/main/runtime/rpc/core.ts | 10 +- .../rpc/dispatcher-feature-readiness.test.ts | 106 +++++++++++++++++- .../runtime/rpc/dispatcher-stream-options.ts | 2 + src/main/runtime/rpc/dispatcher.ts | 12 +- src/main/runtime/rpc/methods/agent-launch.ts | 3 + src/main/runtime/rpc/methods/agent-session.ts | 3 + .../git-commit-message-generation-methods.ts | 1 + .../orchestration/federation/federation.ts | 1 + .../methods/orchestration/worker/workers.ts | 1 + src/main/runtime/rpc/methods/session-tabs.ts | 2 + .../rpc/methods/structured-agent-session.ts | 1 + .../terminal/terminal-lifecycle-methods.ts | 2 + .../orcad-feature-readiness-manifest.test.ts | 66 ++++++++++- .../rpc/orcad-feature-readiness.test.ts | 17 ++- .../runtime/rpc/orcad-feature-readiness.ts | 80 +++++++++---- ...me-rpc-local-disconnect-forwarding.test.ts | 50 +++++++++ .../runtime-rpc-request-admission.ts | 1 + 20 files changed, 374 insertions(+), 39 deletions(-) create mode 100644 config/scripts/vitest-orcad-feature-readiness-reset.ts create mode 100644 src/main/runtime/rpc/agent-launch-feature-requirement.ts diff --git a/config/scripts/ci-unit-files.mjs b/config/scripts/ci-unit-files.mjs index 294560a1487..d2688470c4b 100644 --- a/config/scripts/ci-unit-files.mjs +++ b/config/scripts/ci-unit-files.mjs @@ -18,7 +18,8 @@ export const UNIT_SETUP_FILES = [ 'config/scripts/vitest-bun-node-builtins.ts', 'config/scripts/vitest-dom-setup.ts', 'config/scripts/vitest-host-ports-setup.ts', - 'config/scripts/vitest-caller-identity-env-setup.ts' + 'config/scripts/vitest-caller-identity-env-setup.ts', + 'config/scripts/vitest-orcad-feature-readiness-reset.ts' ] export const UNIT_EXCLUDE = [ diff --git a/config/scripts/vitest-orcad-feature-readiness-reset.ts b/config/scripts/vitest-orcad-feature-readiness-reset.ts new file mode 100644 index 00000000000..d13f3db9db4 --- /dev/null +++ b/config/scripts/vitest-orcad-feature-readiness-reset.ts @@ -0,0 +1,7 @@ +// A startOrcad test that fails partway must not leave latches closed: later tests in the worker would wait 10 s. +import { afterEach } from 'vitest' +import { _resetOrcadFeatureReadinessForTests } from '../../src/main/runtime/rpc/orcad-feature-readiness' + +afterEach(() => { + _resetOrcadFeatureReadinessForTests() +}) diff --git a/src/main/runtime/rpc/agent-launch-feature-requirement.ts b/src/main/runtime/rpc/agent-launch-feature-requirement.ts new file mode 100644 index 00000000000..8f43263ebfb --- /dev/null +++ b/src/main/runtime/rpc/agent-launch-feature-requirement.ts @@ -0,0 +1,45 @@ +import { recognizeAgentLaunchCommand } from '../../../shared/agent-command-line-intent' +import type { OrcadFeature } from './orcad-feature-readiness' + +const ACCOUNTS: readonly OrcadFeature[] = ['accounts'] +const NONE: readonly OrcadFeature[] = [] + +type TerminalCreateLaunchFields = { + command?: string + agent?: unknown + launchAgent?: unknown + launchConfig?: unknown + resumeProviderSession?: unknown + sourceChatFork?: unknown +} + +/** + * A terminal create that starts an agent needs its account home; a plain shell needs nothing. + * An agent behind a bare command the recognizer does not know stays ungated (only Claude/Codex use accounts). + */ +export function terminalCreateFeatureRequirement( + params: TerminalCreateLaunchFields +): readonly OrcadFeature[] { + const agentLaunch = + params.agent !== undefined || + params.launchAgent !== undefined || + params.launchConfig !== undefined || + params.resumeProviderSession !== undefined || + params.sourceChatFork !== undefined || + recognizeAgentLaunchCommand(params.command) !== null + return agentLaunch ? ACCOUNTS : NONE +} + +/** Reusing a live terminal starts no process, so it needs no account home. */ +export function agentLaunchFeatureRequirement(params: { + reuseTerminal?: unknown +}): readonly OrcadFeature[] { + return params.reuseTerminal ? NONE : ACCOUNTS +} + +/** Only an explicit ensure may start an agent; attaching to an existing session is restore latency. */ +export function ensureAgentSessionFeatureRequirement(params: { + kind?: unknown +}): readonly OrcadFeature[] { + return params.kind === 'explicit' ? ACCOUNTS : NONE +} diff --git a/src/main/runtime/rpc/core.ts b/src/main/runtime/rpc/core.ts index 738f5bb6526..349f5d3ae90 100644 --- a/src/main/runtime/rpc/core.ts +++ b/src/main/runtime/rpc/core.ts @@ -23,7 +23,7 @@ import type { RpcCallerIdentity } from './rpc-caller-identity' import type { OrcadPeerDesktopsResult } from '../../../shared/orcad-peer-desktops-contract' import type { RpcBinarySender } from './rpc-binary-sender' import type { RpcMethodPermission } from './rpc-method-permission' -import type { OrcadFeature } from './orcad-feature-readiness' +import type { OrcadFeatureRequirement } from './orcad-feature-readiness' import type { RegisterDeviceKeyParams, RegisterDeviceKeyResult @@ -180,7 +180,7 @@ export type RpcTypedMethod> readonly handler: (params: RpcParsedParams, ctx: RpcContext) => TResult } @@ -195,7 +195,7 @@ export type RpcTypedStreamingMethod> readonly stream: true readonly handler: ( params: RpcParsedParams, @@ -214,7 +214,7 @@ type RpcMethodWithParams

= { readonly name: string readonly permission: RpcMethodPermission readonly params: ZodType | null - readonly requires?: readonly OrcadFeature[] + readonly requires?: OrcadFeatureRequirement

readonly handler: (params: P, ctx: RpcContext) => unknown } @@ -222,7 +222,7 @@ type RpcStreamingMethodWithParams

= { readonly name: string readonly permission: RpcMethodPermission readonly params: ZodType | null - readonly requires?: readonly OrcadFeature[] + readonly requires?: OrcadFeatureRequirement

readonly stream: true readonly handler: (params: P, ctx: RpcContext, emit: (result: unknown) => void) => Promise } diff --git a/src/main/runtime/rpc/dispatcher-feature-readiness.test.ts b/src/main/runtime/rpc/dispatcher-feature-readiness.test.ts index bfb9974cbfe..c84e0a5aee6 100644 --- a/src/main/runtime/rpc/dispatcher-feature-readiness.test.ts +++ b/src/main/runtime/rpc/dispatcher-feature-readiness.test.ts @@ -4,13 +4,19 @@ import { z } from 'zod' import { RpcDispatcher } from './dispatcher' import { defineMethod, defineStreamingMethod, type RpcRequest } from './core' import type { OrcaRuntimeService } from '../orca-runtime' +import { ensureAgentSessionFeatureRequirement } from './agent-launch-feature-requirement' import { _resetOrcadFeatureReadinessForTests, installOrcadFeatureReadiness, ORCAD_FEATURE_WAIT_MS } from './orcad-feature-readiness' -const request = (method: string): RpcRequest => ({ id: 'r1', authToken: 't', method, params: {} }) +const request = (method: string, params: unknown = {}): RpcRequest => ({ + id: 'r1', + authToken: 't', + method, + params +}) function dispatcher(): { dispatcher: RpcDispatcher; ran: string[] } { const ran: string[] = [] @@ -39,6 +45,26 @@ function dispatcher(): { dispatcher: RpcDispatcher; ran: string[] } { return {} } }), + defineMethod({ + name: 'terminal.create', + requires: (params) => (params.command === 'claude' ? ['accounts'] : []), + permission: 'workspace', + params: z.object({ command: z.string().optional() }), + handler: (params) => { + ran.push(`create ${params.command ?? 'shell'}`) + return {} + } + }), + defineMethod({ + name: 'terminal.ensureAgentSession', + requires: ensureAgentSessionFeatureRequirement, + permission: 'workspace', + params: z.object({ kind: z.enum(['automatic', 'explicit']) }), + handler: (params) => { + ran.push(`ensure ${params.kind}`) + return {} + } + }), defineStreamingMethod({ name: 'accounts.subscribe', requires: ['accounts'], @@ -104,4 +130,82 @@ describe('RPC methods that require an orcad startup phase', () => { expect(ran).toEqual(['subscribe']) expect(replies).toHaveLength(1) }) + + it('gate a terminal create only when its params start an agent', async () => { + const readiness = installOrcadFeatureReadiness() + const { dispatcher: rpc, ran } = dispatcher() + const agent = rpc.dispatch(request('terminal.create', { command: 'claude' })) + await expect(rpc.dispatch(request('terminal.create', {}))).resolves.toMatchObject({ ok: true }) + expect(ran).toEqual(['create shell']) + readiness.open('accounts') + await expect(agent).resolves.toMatchObject({ ok: true }) + expect(ran).toEqual(['create shell', 'create claude']) + }) + + // A handler run for a caller that left leaks its close cleanup, or runs a retried mutation twice. + it('refuse a unary call whose request is aborted while it waits, never running it', async () => { + vi.useFakeTimers() + const readiness = installOrcadFeatureReadiness() + const { dispatcher: rpc, ran } = dispatcher() + const aborted = new AbortController() + let reply: unknown = 'pending' + void rpc.dispatch(request('aiVault.searchStatus'), { signal: aborted.signal }).then((r) => { + reply = r + }) + await vi.advanceTimersByTimeAsync(1) + aborted.abort() + await vi.advanceTimersByTimeAsync(1) + expect(reply).toMatchObject({ ok: false, error: { code: 'runtime_unavailable' } }) + readiness.open('sessionSearch') + await vi.advanceTimersByTimeAsync(1) + expect(ran).toEqual([]) + }) + + it('refuse a stream whose connection closes while it waits, never running its handler', async () => { + vi.useFakeTimers() + const readiness = installOrcadFeatureReadiness() + const { dispatcher: rpc, ran } = dispatcher() + const closed = new AbortController() + const replies: string[] = [] + void rpc.dispatchStreaming(request('accounts.subscribe'), (reply) => replies.push(reply), { + connectionSignal: closed.signal + }) + await vi.advanceTimersByTimeAsync(1) + closed.abort() + await vi.advanceTimersByTimeAsync(1) + expect(replies.map((reply) => JSON.parse(reply))).toMatchObject([ + { ok: false, error: { code: 'runtime_unavailable' } } + ]) + readiness.open('accounts') + await vi.advanceTimersByTimeAsync(1) + expect(ran).toEqual([]) + }) + + it('refuse at once when the caller already left, without waiting out the bound', async () => { + vi.useFakeTimers() + installOrcadFeatureReadiness() + const { dispatcher: rpc, ran } = dispatcher() + let reply: unknown = 'pending' + void rpc + .dispatch(request('aiVault.searchStatus'), { connectionSignal: AbortSignal.abort() }) + .then((r) => { + reply = r + }) + await vi.advanceTimersByTimeAsync(1) + expect(reply).toMatchObject({ ok: false, error: { code: 'runtime_unavailable' } }) + expect(ran).toEqual([]) + }) + + it('run an attach-only ensureAgentSession at once while every latch is closed', async () => { + const readiness = installOrcadFeatureReadiness() + const { dispatcher: rpc, ran } = dispatcher() + const explicit = rpc.dispatch(request('terminal.ensureAgentSession', { kind: 'explicit' })) + await expect( + rpc.dispatch(request('terminal.ensureAgentSession', { kind: 'automatic' })) + ).resolves.toMatchObject({ ok: true }) + expect(ran).toEqual(['ensure automatic']) + readiness.open('accounts') + await expect(explicit).resolves.toMatchObject({ ok: true }) + expect(ran).toEqual(['ensure automatic', 'ensure explicit']) + }) }) diff --git a/src/main/runtime/rpc/dispatcher-stream-options.ts b/src/main/runtime/rpc/dispatcher-stream-options.ts index bf454638ae1..349c589acc0 100644 --- a/src/main/runtime/rpc/dispatcher-stream-options.ts +++ b/src/main/runtime/rpc/dispatcher-stream-options.ts @@ -28,6 +28,8 @@ export type RpcDispatchStreamingOptions = Pick< > & { /** What the transport proved the caller may do; absent means this host's owner. */ callerScope?: RpcCallerScope + /** Aborts when the caller's connection closes; only for the dispatcher's own waits, never handlers. */ + connectionSignal?: AbortSignal } export function rpcContextFromTransport( diff --git a/src/main/runtime/rpc/dispatcher.ts b/src/main/runtime/rpc/dispatcher.ts index 5f24e240e1f..d6c08bacfde 100644 --- a/src/main/runtime/rpc/dispatcher.ts +++ b/src/main/runtime/rpc/dispatcher.ts @@ -33,7 +33,7 @@ import { recordRuntimeFeatureInteraction } from './runtime-feature-interaction' import type { RpcCallerScope } from './rpc-caller-scope' import { runWithActiveRpcCaller } from './rpc-active-caller' import { admitRpcRequest } from './rpc-request-admission-prologue' -import { rpcMethodFeatureWait } from './orcad-feature-readiness' +import { resolveOrcadFeatureRequirement, rpcMethodFeatureWait } from './orcad-feature-readiness' export type DispatcherOptions = { runtime: OrcaRuntimeService @@ -84,7 +84,10 @@ export class RpcDispatcher { emulatorProbe(`rpc ${request.method}`, request.params) } try { - const featureWait = rpcMethodFeatureWait(admission.method.requires) + const featureWait = rpcMethodFeatureWait( + resolveOrcadFeatureRequirement(admission.method.requires, admission.params), + [options?.signal, options?.connectionSignal] + ) if (featureWait) { await featureWait } @@ -118,7 +121,10 @@ export class RpcDispatcher { return } const { method, request, params, caller } = admission - const featureWait = rpcMethodFeatureWait(method.requires) + const featureWait = rpcMethodFeatureWait( + resolveOrcadFeatureRequirement(method.requires, params), + [options?.signal, options?.connectionSignal] + ) if (featureWait) { try { await featureWait diff --git a/src/main/runtime/rpc/methods/agent-launch.ts b/src/main/runtime/rpc/methods/agent-launch.ts index ca7e0340499..adda5fbdcb1 100644 --- a/src/main/runtime/rpc/methods/agent-launch.ts +++ b/src/main/runtime/rpc/methods/agent-launch.ts @@ -3,6 +3,7 @@ * operationId adds durable at-most-once execution and recorded replay; unknown outcomes refuse, while stranded-surface recovery remains separate. */ import { AGENT_LAUNCH_RUNTIME_CAPABILITY } from '../../../../shared/agent-launch-runtime-capability' +import { agentLaunchFeatureRequirement } from '../agent-launch-feature-requirement' import { AgentLaunchTabClosedError } from '../../../../shared/agent-launch-tab-closed' import { computeAgentLaunchFingerprint } from '../../../../shared/agent-launch-operation' import type { AgentLaunchResult } from '../../../../shared/agent-launch-intent' @@ -265,6 +266,7 @@ function runReplaySafeAgentLaunch( export const AGENT_LAUNCH_METHODS = [ defineMethod({ name: 'agent.launchReplay', + requires: agentLaunchFeatureRequirement, permission: 'workspace', params: AgentLaunchReplay, handler: async (params, context): Promise => { @@ -296,6 +298,7 @@ export const AGENT_LAUNCH_METHODS = [ }), defineMethod({ name: 'agent.launch', + requires: agentLaunchFeatureRequirement, permission: 'workspace', params: AgentLaunch, handler: async (params, context): Promise => { diff --git a/src/main/runtime/rpc/methods/agent-session.ts b/src/main/runtime/rpc/methods/agent-session.ts index 7ab5942df7b..17d8f30e9f8 100644 --- a/src/main/runtime/rpc/methods/agent-session.ts +++ b/src/main/runtime/rpc/methods/agent-session.ts @@ -1,3 +1,4 @@ +import { ensureAgentSessionFeatureRequirement } from '../agent-launch-feature-requirement' import { terminalImageViewerSupported } from '../../terminal-image-viewer-env' import type { RuntimeAgentSessionRpcCaller } from '../../../../shared/agent-session-host-authority' import { @@ -43,6 +44,7 @@ function assertOperationTimestampWithinFutureSkew(clientOperationId: string): vo export const AGENT_SESSION_METHODS = [ defineMethod({ name: 'terminal.ensureAgentSession', + requires: ensureAgentSessionFeatureRequirement, permission: 'workspace', params: EnsureAgentSessionParams, handler: (params, context) => @@ -53,6 +55,7 @@ export const AGENT_SESSION_METHODS = [ }), defineMethod({ name: 'terminal.createAgentSession', + requires: ['accounts'], permission: 'workspace', params: CreateAgentSessionParams, handler: (params, context) => { diff --git a/src/main/runtime/rpc/methods/git-commit-message-generation-methods.ts b/src/main/runtime/rpc/methods/git-commit-message-generation-methods.ts index 5b3f063336b..cb85e3fe8f9 100644 --- a/src/main/runtime/rpc/methods/git-commit-message-generation-methods.ts +++ b/src/main/runtime/rpc/methods/git-commit-message-generation-methods.ts @@ -64,6 +64,7 @@ function buildCommitMessageGenerationOverride(params: { export const GIT_COMMIT_MESSAGE_GENERATION_METHODS = [ defineMethod({ name: 'git.generateCommitMessage', + requires: ['accounts'], permission: 'workspace', params: GitGenerateCommitMessage, handler: async (params, { runtime }) => { diff --git a/src/main/runtime/rpc/methods/orchestration/federation/federation.ts b/src/main/runtime/rpc/methods/orchestration/federation/federation.ts index 3b2c07ca406..8eea4554751 100644 --- a/src/main/runtime/rpc/methods/orchestration/federation/federation.ts +++ b/src/main/runtime/rpc/methods/orchestration/federation/federation.ts @@ -28,6 +28,7 @@ import { prepareFederatedAttachmentAuthority } from './federation-attachment-aut export const ORCHESTRATION_FEDERATION_ATTACH_METHODS = [ defineMethod({ name: 'orchestration.federationAttachStart', + requires: ['accounts'], permission: 'workspace', params: FederationAttachStartParams, handler: async (params, { runtime, orchestrationMutation }) => { diff --git a/src/main/runtime/rpc/methods/orchestration/worker/workers.ts b/src/main/runtime/rpc/methods/orchestration/worker/workers.ts index eb5ef4aca8f..349e18e7a5d 100644 --- a/src/main/runtime/rpc/methods/orchestration/worker/workers.ts +++ b/src/main/runtime/rpc/methods/orchestration/worker/workers.ts @@ -20,6 +20,7 @@ import { ORCHESTRATION_SESSION_CALLER_ERROR_CODES as SESSION_CODES } from '../.. export const ORCHESTRATION_WORKER_START_METHODS = [ defineMethod({ name: 'orchestration.workerStart', + requires: ['accounts'], permission: 'workspace', params: WorkerStartParams, handler: async ( diff --git a/src/main/runtime/rpc/methods/session-tabs.ts b/src/main/runtime/rpc/methods/session-tabs.ts index 1ff40616529..7a6dc2494ec 100644 --- a/src/main/runtime/rpc/methods/session-tabs.ts +++ b/src/main/runtime/rpc/methods/session-tabs.ts @@ -1,3 +1,4 @@ +import { terminalCreateFeatureRequirement } from '../agent-launch-feature-requirement' import { terminalImageViewerSupported } from '../../terminal-image-viewer-env' import { resolveRuntimeNavigationTarget } from '../../../../shared/runtime-navigation' import { getExplicitWorktreeIdSelector } from '../../runtime-worktree-selection' @@ -53,6 +54,7 @@ export const SESSION_TAB_METHODS = [ ...SESSION_TAB_CLOSE_METHODS, defineMethod({ name: 'session.tabs.createTerminal', + requires: terminalCreateFeatureRequirement, permission: 'workspace', params: CreateTerminalTab, handler: async (params, context) => { diff --git a/src/main/runtime/rpc/methods/structured-agent-session.ts b/src/main/runtime/rpc/methods/structured-agent-session.ts index 261d64c1a86..de829e50725 100644 --- a/src/main/runtime/rpc/methods/structured-agent-session.ts +++ b/src/main/runtime/rpc/methods/structured-agent-session.ts @@ -116,6 +116,7 @@ export const STRUCTURED_AGENT_SESSION_METHODS = [ }), defineMethod({ name: 'agentSession.create', + requires: ['accounts'], permission: 'workspace', params: CreateParams, handler: ORCA_LAYOUT_OWNER_IS_MAIN diff --git a/src/main/runtime/rpc/methods/terminal/terminal-lifecycle-methods.ts b/src/main/runtime/rpc/methods/terminal/terminal-lifecycle-methods.ts index 0047032dd0c..d8c0773a7ff 100644 --- a/src/main/runtime/rpc/methods/terminal/terminal-lifecycle-methods.ts +++ b/src/main/runtime/rpc/methods/terminal/terminal-lifecycle-methods.ts @@ -1,3 +1,4 @@ +import { terminalCreateFeatureRequirement } from '../../agent-launch-feature-requirement' import { terminalImageViewerSupported } from '../../../terminal-image-viewer-env' import { TERMINAL_WORKSPACE_STOP_METHODS } from './terminal-workspace-stop-methods' import { ORCA_SIDEBAR_INPUTS_IS_MAIN } from '../../../../../shared/sidebar-inputs-switch' @@ -34,6 +35,7 @@ export const TERMINAL_LIFECYCLE_METHODS = [ }), defineMethod({ name: 'terminal.create', + requires: terminalCreateFeatureRequirement, permission: 'workspace', params: TerminalCreateParams, handler: async (params, context) => { diff --git a/src/main/runtime/rpc/orcad-feature-readiness-manifest.test.ts b/src/main/runtime/rpc/orcad-feature-readiness-manifest.test.ts index 5b11d3d7d8a..9192d7cbe8f 100644 --- a/src/main/runtime/rpc/orcad-feature-readiness-manifest.test.ts +++ b/src/main/runtime/rpc/orcad-feature-readiness-manifest.test.ts @@ -2,11 +2,33 @@ import { readFileSync } from 'node:fs' import { join } from 'node:path' import { describe, expect, it } from 'vitest' import { ALL_RPC_METHODS } from './methods' +import { buildRegistry } from './core' import { ACCOUNT_METHODS } from './methods/accounts' -import type { OrcadFeature } from './orcad-feature-readiness' +import { resolveOrcadFeatureRequirement, type OrcadFeature } from './orcad-feature-readiness' -const requiresOf = (name: string): readonly OrcadFeature[] | undefined => - ALL_RPC_METHODS.find((method) => method.name === name)?.requires +// Every method's declared needs for these params; param-level ones read them. +const REGISTRY = buildRegistry(ALL_RPC_METHODS) +const requiresOf = (name: string, params: unknown = {}): readonly OrcadFeature[] => { + const method = REGISTRY.get(name) + expect(method, name).toBeDefined() + return resolveOrcadFeatureRequirement(method?.requires, params) +} + +// Launches that start an agent (and so read its account home); a terminal create is one only with an agent. +const AGENT_LAUNCHES = [ + 'terminal.createAgentSession', + 'agentSession.create', + 'git.generateCommitMessage', + 'orchestration.workerStart', + 'orchestration.federationAttachStart' +] +const PARAM_GATED_LAUNCHES = [ + 'terminal.ensureAgentSession', + 'terminal.create', + 'session.tabs.createTerminal', + 'agent.launch', + 'agent.launchReplay' +] describe('which RPC methods wait for an orcad startup phase', () => { it('gates the account services, search index, orchestration and orphan adoption', () => { @@ -20,6 +42,9 @@ describe('which RPC methods wait for an orcad startup phase', () => { ) { expected.push('accounts') } + if (AGENT_LAUNCHES.includes(method.name)) { + expected.push('accounts') + } if (method.name.startsWith('orchestration.')) { expected.push('orchestration') } @@ -29,13 +54,44 @@ describe('which RPC methods wait for an orcad startup phase', () => { if (method.name === 'terminal.adoptOrphans') { expected.push('unshownPlacement') } - expect({ name: method.name, requires: method.requires ?? [] }).toEqual({ + if (PARAM_GATED_LAUNCHES.includes(method.name)) { + expect(typeof method.requires, method.name).toBe('function') + continue + } + expect({ name: method.name, requires: requiresOf(method.name) }).toEqual({ name: method.name, requires: expected }) } }) + it('gates a terminal create only when it starts an agent; plain shells stay ungated', () => { + for (const name of ['terminal.create', 'session.tabs.createTerminal']) { + expect(requiresOf(name, {}), name).toEqual([]) + expect(requiresOf(name, { command: 'ls -la' }), name).toEqual([]) + expect(requiresOf(name, { command: 'claude --resume' }), name).toEqual(['accounts']) + expect(requiresOf(name, { command: 'codex' }), name).toEqual(['accounts']) + expect(requiresOf(name, { launchAgent: 'claude' }), name).toEqual(['accounts']) + expect(requiresOf(name, { launchConfig: { agentArgs: '', agentEnv: {} } }), name).toEqual([ + 'accounts' + ]) + expect(requiresOf(name, { resumeProviderSession: { id: 's' } }), name).toEqual(['accounts']) + } + expect(requiresOf('session.tabs.createTerminal', { agent: 'claude' })).toEqual(['accounts']) + }) + + it('gates only an explicit ensureAgentSession; an attach is restore latency', () => { + expect(requiresOf('terminal.ensureAgentSession', { kind: 'explicit' })).toEqual(['accounts']) + expect(requiresOf('terminal.ensureAgentSession', { kind: 'automatic' })).toEqual([]) + }) + + it('gates agent.launch unless it reuses a live terminal', () => { + for (const name of ['agent.launch', 'agent.launchReplay']) { + expect(requiresOf(name, {}), name).toEqual(['accounts']) + expect(requiresOf(name, { reuseTerminal: { handle: 'h' } }), name).toEqual([]) + } + }) + it('never gates what every client calls at connect', () => { for (const name of [ 'status.get', @@ -46,7 +102,7 @@ describe('which RPC methods wait for an orcad startup phase', () => { 'accounts.setHostLaunchEnv' ]) { expect(ALL_RPC_METHODS.some((method) => method.name === name)).toBe(true) - expect(requiresOf(name)).toBeUndefined() + expect(requiresOf(name)).toEqual([]) } }) diff --git a/src/main/runtime/rpc/orcad-feature-readiness.test.ts b/src/main/runtime/rpc/orcad-feature-readiness.test.ts index f9300bbdea5..cfa8a9a0e3b 100644 --- a/src/main/runtime/rpc/orcad-feature-readiness.test.ts +++ b/src/main/runtime/rpc/orcad-feature-readiness.test.ts @@ -1,19 +1,18 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import { - _resetOrcadFeatureReadinessForTests, installOrcadFeatureReadiness, ORCAD_FEATURES, rpcMethodFeatureWait } from './orcad-feature-readiness' +// No local readiness reset: the shared vitest setup must restore it (asserted below). afterEach(() => { - _resetOrcadFeatureReadinessForTests() vi.useRealTimers() }) describe('OrcadFeatureReadiness', () => { it('adds no await for open or undeclared features, so ungated dispatch starts synchronously', () => { - expect(rpcMethodFeatureWait(undefined)).toBeNull() + expect(rpcMethodFeatureWait([])).toBeNull() expect(rpcMethodFeatureWait(['accounts'])).toBeNull() const readiness = installOrcadFeatureReadiness() expect(rpcMethodFeatureWait([])).toBeNull() @@ -39,10 +38,20 @@ describe('OrcadFeatureReadiness', () => { it('refuses with runtime_unavailable after the bound and leaves no timer behind', async () => { vi.useFakeTimers() installOrcadFeatureReadiness() - const wait = rpcMethodFeatureWait(['sessionSearch'], 50) + const wait = rpcMethodFeatureWait(['sessionSearch'], [], 50) const refused = expect(wait).rejects.toThrow('runtime_unavailable') await vi.advanceTimersByTimeAsync(50) await refused expect(vi.getTimerCount()).toBe(0) }) + + // The shared setup resets after every test; these two run in order in one worker. + it('may leave latches closed when a test stops partway', () => { + installOrcadFeatureReadiness() + expect(rpcMethodFeatureWait(['accounts'])).not.toBeNull() + }) + + it('finds every latch open again in the next test', () => { + expect(rpcMethodFeatureWait(['accounts'])).toBeNull() + }) }) diff --git a/src/main/runtime/rpc/orcad-feature-readiness.ts b/src/main/runtime/rpc/orcad-feature-readiness.ts index 4f6a8745940..e3291f79d89 100644 --- a/src/main/runtime/rpc/orcad-feature-readiness.ts +++ b/src/main/runtime/rpc/orcad-feature-readiness.ts @@ -37,20 +37,40 @@ export class OrcadFeatureReadiness { this.pending.delete(feature) } - /** True once every feature is open; false if any is still closed after `timeoutMs`. */ - async waitFor(features: readonly OrcadFeature[], timeoutMs: number): Promise { + /** + * 'ready' once every feature is open; 'timeout' if any is still closed after `timeoutMs`; + * 'aborted' as soon as any signal aborts (the client went away). + */ + async waitFor( + features: readonly OrcadFeature[], + timeoutMs: number, + signals: readonly (AbortSignal | undefined)[] = [] + ): Promise<'ready' | 'timeout' | 'aborted'> { + const live = signals.filter((signal): signal is AbortSignal => signal !== undefined) + if (live.some((signal) => signal.aborted)) { + return 'aborted' + } const waits = features.flatMap((feature) => this.pending.get(feature)?.opened ?? []) if (waits.length === 0) { - return true + return 'ready' } let timer: ReturnType | undefined - const timedOut = new Promise((resolve) => { - timer = setTimeout(() => resolve(false), timeoutMs) + const cleanups: (() => void)[] = [] + const settled = new Promise<'timeout' | 'aborted'>((resolve) => { + timer = setTimeout(() => resolve('timeout'), timeoutMs) + for (const signal of live) { + const onAbort = (): void => resolve('aborted') + signal.addEventListener('abort', onAbort, { once: true }) + cleanups.push(() => signal.removeEventListener('abort', onAbort)) + } }) try { - return await Promise.race([Promise.all(waits).then(() => true), timedOut]) + return await Promise.race([Promise.all(waits).then(() => 'ready' as const), settled]) } finally { clearTimeout(timer) + for (const cleanup of cleanups) { + cleanup() + } } } } @@ -72,31 +92,51 @@ export function _resetOrcadFeatureReadinessForTests(): void { current = new OrcadFeatureReadiness() } +/** A method's startup needs: fixed, or read from its parsed params (an agent launch, not a shell). */ +export type OrcadFeatureRequirement

= + | readonly OrcadFeature[] + | ((params: P) => readonly OrcadFeature[]) + +export function resolveOrcadFeatureRequirement

( + requires: OrcadFeatureRequirement

| undefined, + params: P +): readonly OrcadFeature[] { + return typeof requires === 'function' ? requires(params) : (requires ?? []) +} + /** * The wait for a method's required features, or null when nothing is closed, so ungated dispatch - * keeps its synchronous start. Refuses with the existing retryable `runtime_unavailable`. + * keeps its synchronous start. Refuses with the existing retryable `runtime_unavailable`, also when + * the request or its connection ends while waiting: the handler must never run for a caller that + * left (a streaming handler's close cleanup would never fire; a retried mutation would run twice). */ export function rpcMethodFeatureWait( - requires: readonly OrcadFeature[] | undefined, + requires: readonly OrcadFeature[], + signals: readonly (AbortSignal | undefined)[] = [], timeoutMs = ORCAD_FEATURE_WAIT_MS ): Promise | null { - if (!requires?.some((feature) => !current.isOpen(feature))) { + if (!requires.some((feature) => !current.isOpen(feature))) { return null } - return current.waitFor(requires, timeoutMs).then((ready) => { - if (!ready) { + return current.waitFor(requires, timeoutMs, signals).then((outcome) => { + if (outcome !== 'ready') { throw new Error('runtime_unavailable') } }) } -/** Declares `feature` on each method of a family, e.g. every `orchestration.*` method. */ -export function requiringOrcadFeature( - feature: OrcadFeature, - methods: readonly M[] -): M[] { - return methods.map((method) => ({ - ...method, - requires: [...(method.requires ?? []), feature] - })) +/** Adds `feature` to each method of a family, e.g. every `orchestration.*` method. */ +export function requiringOrcadFeature< + M extends { readonly requires?: OrcadFeatureRequirement } +>(feature: OrcadFeature, methods: readonly M[]): M[] { + return methods.map((method) => { + const own = method.requires + return { + ...method, + requires: + typeof own === 'function' + ? (params: never) => [...own(params), feature] + : [...(own ?? []), feature] + } + }) } diff --git a/src/main/runtime/runtime-rpc-local-disconnect-forwarding.test.ts b/src/main/runtime/runtime-rpc-local-disconnect-forwarding.test.ts index 504b51c6201..f0217946451 100644 --- a/src/main/runtime/runtime-rpc-local-disconnect-forwarding.test.ts +++ b/src/main/runtime/runtime-rpc-local-disconnect-forwarding.test.ts @@ -6,6 +6,10 @@ import { expect, it, vi } from 'vitest' import { OrcaRuntimeService } from './orca-runtime' import { defineMethod } from './rpc/core' import { OrcaRuntimeRpcServer } from './runtime-rpc' +import { + _resetOrcadFeatureReadinessForTests, + installOrcadFeatureReadiness +} from './rpc/orcad-feature-readiness' it.each(['browser.eval', 'linear.commentEdit', 'linear.commentDelete'] as const)( 'forwards local disconnect to %s without long-poll keepalives', @@ -48,3 +52,49 @@ it.each(['browser.eval', 'linear.commentEdit', 'linear.commentDelete'] as const) } } ) + +// A short RPC gets no handler signal, but a caller that disconnects mid feature-wait is never run. +it('refuses a gated short RPC whose local caller disconnects while it waits', async () => { + const userDataPath = mkdtempSync(join(tmpdir(), 'local-disconnect-feature-wait-')) + const ran = vi.fn() + const readiness = installOrcadFeatureReadiness() + const server = new OrcaRuntimeRpcServer({ + runtime: new OrcaRuntimeService(), + userDataPath, + methods: [ + defineMethod({ + name: 'aiVault.searchStatus', + requires: ['sessionSearch'], + permission: 'workspace', + params: null, + handler: () => { + ran() + return {} + } + }) + ] + }) + const disconnect = new AbortController() + try { + const pending = server['handleMessage']( + JSON.stringify({ + id: 'req-1', + method: 'aiVault.searchStatus', + authToken: server['authToken'] + }), + { signal: disconnect.signal, startKeepalive: vi.fn() } + ) + disconnect.abort() + await expect(pending).resolves.toMatchObject({ + ok: false, + error: { code: 'runtime_unavailable' } + }) + readiness.open('sessionSearch') + await Promise.resolve() + expect(ran).not.toHaveBeenCalled() + } finally { + _resetOrcadFeatureReadinessForTests() + await server.stop() + rmSync(userDataPath, { recursive: true, force: true }) + } +}) diff --git a/src/main/runtime/runtime-rpc/runtime-rpc-request-admission.ts b/src/main/runtime/runtime-rpc/runtime-rpc-request-admission.ts index e4c7999a4d0..e676227c9e9 100644 --- a/src/main/runtime/runtime-rpc/runtime-rpc-request-admission.ts +++ b/src/main/runtime/runtime-rpc/runtime-rpc-request-admission.ts @@ -48,6 +48,7 @@ export class RuntimeRpcRequestAdmission extends RuntimeRpcBinaryRouting { longPoll || DISCONNECT_FORWARDED_METHODS.has(request.method) ? context?.signal : undefined, + connectionSignal: context?.signal, callerScope }) } finally {