diff --git a/src/main/pi/agent-status-extension-source.test.ts b/src/main/pi/agent-status-extension-source.test.ts index 7df089f13b9..a32bc441667 100644 --- a/src/main/pi/agent-status-extension-source.test.ts +++ b/src/main/pi/agent-status-extension-source.test.ts @@ -195,15 +195,16 @@ describe('getPiAgentStatusExtensionSource', () => { it('tracks persistent OMP sessions and clears ephemeral session ids', async () => { const harness = createHarness({ kind: 'omp' }) let sessionId = 'omp-session-8' - const sessionManager = { getSessionId: () => sessionId, getSessionFile: () => '/tmp/s' } + let sessionFile: string | undefined = '/tmp/s' + const sessionManager = { getSessionId: () => sessionId, getSessionFile: () => sessionFile } await harness.callHook('agent_start', undefined, { sessionManager }) sessionId = 'omp-session-9' await harness.callHook('before_agent_start', { prompt: 'hi' }, { sessionManager }) await vi.waitFor(() => expect(harness.fetchMock).toHaveBeenCalledTimes(2)) - await harness.callHook('agent_end', undefined, { - sessionManager: { getSessionId: () => 'omp-ephemeral' } - }) + sessionId = 'omp-ephemeral' + sessionFile = undefined + await harness.callHook('agent_end', undefined, { sessionManager }) await vi.waitFor(() => expect(harness.fetchMock).toHaveBeenCalledTimes(3)) expect( @@ -236,21 +237,17 @@ describe('getPiAgentStatusExtensionSource', () => { ) }) - await harness.callHook('agent_start', undefined, { - sessionManager: { - getSessionId: () => 'omp-session-8', - getSessionFile: () => '/tmp/omp-session-8.jsonl' - } - }) + let sessionId = 'omp-session-8' + const sessionManager = { + getSessionId: () => sessionId, + getSessionFile: () => '/tmp/session.jsonl' + } + await harness.callHook('agent_start', undefined, { sessionManager }) + sessionId = 'omp-session-9' await harness.callHook( 'message_end', { message: { role: 'assistant', content: 'done' } }, - { - sessionManager: { - getSessionId: () => 'omp-session-9', - getSessionFile: () => '/tmp/omp-session-9.jsonl' - } - } + { sessionManager } ) await harness.callHook('message_end', { message: { role: 'user', content: 'next' } }, {}) diff --git a/src/main/pi/agent-status-handler-source.ts b/src/main/pi/agent-status-handler-source.ts index e6c2353aa0a..e2baa3e7716 100644 --- a/src/main/pi/agent-status-handler-source.ts +++ b/src/main/pi/agent-status-handler-source.ts @@ -1,5 +1,6 @@ import { getPiPrefillHandlerSourceLines } from './prefill-extension-source' import type { PiAgentKind } from '../../shared/pi-agent-kind' +import { getOmpSessionOwnerHandlerSourceLines } from './omp-session-status-owner-source' import { getPiAgentStatusUiPromptHandlerSourceLines } from './agent-status-ui-prompt-source' // Why: keep the generated handler registrations separate from hook transport; @@ -8,7 +9,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] const sessionStartHandler = kind !== 'omp' ? [ - " pi.on('session_start', (event, ctx) => {", + " onStatus('session_start', (event, ctx) => {", ' updateSessionMetadata(ctx)', ...(kind === 'pi' ? [' piUiPromptDepth = 0'] : []), ' // Why: /reload re-registers the active session, but it is not a', @@ -40,7 +41,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] kind === 'prime-agent' ? [] : [ - ` pi.on('tool_approval_requested', (event${ctxParam}) => {`, + ` onStatus('tool_approval_requested', (event${ctxParam}) => {`, ...captureSessionMetadata, ' if (!isOmpRuntime()) return', " post('tool_approval_requested', {", @@ -50,7 +51,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' })', ' })', '', - ` pi.on('tool_approval_resolved', (event${ctxParam}) => {`, + ` onStatus('tool_approval_resolved', (event${ctxParam}) => {`, ...captureSessionMetadata, ' if (!isOmpRuntime()) return', " post('tool_approval_resolved', {", @@ -129,14 +130,24 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' })' ] : []), + ...getOmpSessionOwnerHandlerSourceLines(), ...sessionStartHandler, + ...(kind === 'omp' ? getPiPrefillHandlerSourceLines('omp') : []), + ` onStatus('before_agent_start', (event${ctxParam}) => {`, + ...sessionStartHandler, +<<<<<<< HEAD ...(kind === 'omp' ? getPiPrefillHandlerSourceLines('omp') : []), ` pi.on('before_agent_start', (event${ctxParam}) => {`, +||||||| parent of 2d948aac594 (fix(omp): fence pane status to the root session manager) + ` pi.on('before_agent_start', (event${ctxParam}) => {`, +======= + ` onStatus('before_agent_start', (event${ctxParam}) => {`, +>>>>>>> 2d948aac594 (fix(omp): fence pane status to the root session manager) ...captureSessionMetadata, " post('before_agent_start', { prompt: event.prompt ?? '' })", ' })', '', - ` pi.on('agent_start', (${bareCtxParams}) => {`, + ` onStatus('agent_start', (${bareCtxParams}) => {`, ...captureSessionMetadata, ' clearPendingAgentEndCheck()', ' agentEndReported = false', @@ -146,7 +157,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] " post('agent_start')", ' })', '', - ` pi.on('tool_execution_start', (event${ctxParam}) => {`, + ` onStatus('tool_execution_start', (event${ctxParam}) => {`, ...captureSessionMetadata, " post('tool_execution_start', {", ' tool_name: event.toolName,', @@ -154,7 +165,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' })', ' })', '', - ` pi.on('tool_call', (event${ctxParam}) => {`, + ` onStatus('tool_call', (event${ctxParam}) => {`, ...captureSessionMetadata, " post('tool_call', {", ' tool_name: event.toolName,', @@ -162,7 +173,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' })', ' })', '', - ` pi.on('tool_execution_end', (event${ctxParam}) => {`, + ` onStatus('tool_execution_end', (event${ctxParam}) => {`, ...captureSessionMetadata, " post('tool_execution_end', {", ' tool_name: event.toolName,', @@ -175,7 +186,7 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' // so the dashboard preview reflects the most recent reply even before', ' // agent_end fires. message_end is the right hook because pi guarantees', ' // it fires after the message is finalized (post-streaming).', - ` pi.on('message_end', (event${ctxParam}) => {`, + ` onStatus('message_end', (event${ctxParam}) => {`, ...captureSessionMetadata, " if (event.message?.role !== 'assistant') return", ' const text = extractAssistantText(event.message)', @@ -234,14 +245,14 @@ export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] ' agentEndIdleRecheckMs = Math.min(agentEndIdleRecheckMs * 2, AGENT_END_IDLE_RECHECK_MAX_MS)', ' }', '', - ` pi.on('agent_settled', (${bareCtxParams}) => {`, + ` onStatus('agent_settled', (${bareCtxParams}) => {`, ...captureSessionMetadata, ' agentSettledSupported = true', ' clearPendingAgentEndCheck()', ' postAgentEndOnce()', ' })', '', - " pi.on('agent_end', (event, ctx) => {", + " onStatus('agent_end', (event, ctx) => {", ...captureSessionMetadata, ' if (event?.willContinue === true) {', ' clearPendingAgentEndCheck()', diff --git a/src/main/pi/agent-status-ui-prompt-source.ts b/src/main/pi/agent-status-ui-prompt-source.ts index b641a5b3328..e71642ca357 100644 --- a/src/main/pi/agent-status-ui-prompt-source.ts +++ b/src/main/pi/agent-status-ui-prompt-source.ts @@ -7,14 +7,14 @@ export function getPiAgentStatusUiPromptHandlerSourceLines(kind: PiAgentKind): s } return [ - " pi.on('ui_prompt_start', () => {", + " onStatus('ui_prompt_start', () => {", ' if (isOmpRuntime()) return', ' piUiPromptDepth++', ' if (piUiPromptDepth > 1) return', " post('ui_prompt_start')", ' })', '', - " pi.on('ui_prompt_end', (_event, ctx) => {", + " onStatus('ui_prompt_end', (_event, ctx) => {", ' if (isOmpRuntime() || piUiPromptDepth === 0) return', ' piUiPromptDepth--', ' if (piUiPromptDepth > 0) return', @@ -31,9 +31,9 @@ export function getPiAgentStatusUiPromptHandlerSourceLines(kind: PiAgentKind): s " post('ui_prompt_end', { is_idle: isIdle })", ' })', '', - " pi.on('session_shutdown', () => {", + " onStatus('session_shutdown', () => {", ' resetPostQueue()', - ' clearPendingAgentEndCheck()', + ' clearPendingAgentEndCheck()' ' if (isOmpRuntime()) return', ' // Why: pi tears an open dialog down through resetExtensionUI without resolving its', ' // promise, so a replaced session never emits the matching ui_prompt_end and the wait', diff --git a/src/main/pi/omp-session-status-owner-source.ts b/src/main/pi/omp-session-status-owner-source.ts new file mode 100644 index 00000000000..167fcbc9a01 --- /dev/null +++ b/src/main/pi/omp-session-status-owner-source.ts @@ -0,0 +1,33 @@ +// OMP loads the same extension in each in-process task session. +export function getOmpSessionOwnerHandlerSourceLines(): string[] { + return [ + ' // SessionManager survives reload/new/resume; task children own a different instance.', + ' function ownsSessionStatus(ctx): boolean {', + ' if (!isOmpRuntime()) return true', + ' const manager = ctx?.sessionManager', + " if (!manager || typeof manager !== 'object') return true", + ' // Keep ownership through module reload and shutdown while child sessions drain.', + " const key = Symbol.for('orca.omp.status-session-owners')", + ' let owners = Reflect.get(globalThis, key)', + ' if (!(owners instanceof Map)) {', + ' owners = new Map()', + ' Reflect.set(globalThis, key, owners)', + ' }', + ' const pane = JSON.stringify([process.env.ORCA_PANE_KEY, process.env.ORCA_AGENT_LAUNCH_TOKEN])', + ' const owner = owners.get(pane)', + ' if (owner) return owner === manager', + ' owners.set(pane, manager)', + ' return true', + ' }', + '', + ' function onStatus(name, handler): void {', + ' pi.on(name, (event, ctx) => {', + ' if (!ownsSessionStatus(ctx)) return', + ' return handler(event, ctx)', + ' })', + ' }', + '', + " onStatus('session_start', () => {})", + '' + ] +} diff --git a/src/main/pi/omp-session-status-owner.test.ts b/src/main/pi/omp-session-status-owner.test.ts new file mode 100644 index 00000000000..c3c14db1da4 --- /dev/null +++ b/src/main/pi/omp-session-status-owner.test.ts @@ -0,0 +1,102 @@ +import { describe, expect, it } from 'vitest' +import { createAgentStatusExtensionHarness } from './agent-status-extension-test-harness' + +const settle = async (): Promise => { + for (let i = 0; i < 10; i++) { + await Promise.resolve() + } +} + +describe('OMP session status ownership', () => { + it.each(['omp', 'pi'] as const)( + 'fences child callbacks before they change %s pane metadata', + async (kind) => { + const harness = createAgentStatusExtensionHarness({ kind, argv: ['bun', '/opt/omp/bin/omp'] }) + const root = { + sessionManager: { getSessionId: () => 'root', getSessionFile: () => '/root.jsonl' } + } + await harness.callHook('session_start', {}, root) + await settle() + harness.fetchMock.mockClear() + const rootHandlers = { ...harness.handlers } + harness.reload() + const child = { + sessionManager: { getSessionId: () => 'child', getSessionFile: () => '/child.jsonl' } + } + for (const name of [ + 'session_start', + 'before_agent_start', + 'agent_start', + 'tool_call', + 'tool_execution_start', + 'tool_execution_end', + 'tool_approval_requested', + 'tool_approval_resolved', + 'message_end', + 'agent_end', + 'agent_settled' + ]) { + await harness.callHook( + name, + { message: { role: 'assistant', content: 'child answer' } }, + child + ) + await settle() + } + expect(harness.fetchMock).not.toHaveBeenCalled() + await rootHandlers.agent_start({}, root) + await settle() + await rootHandlers.agent_end({}, root) + await settle() + const bodies = harness.fetchMock.mock.calls.map((call) => JSON.parse(call[1].body)) + expect(bodies.map((body) => body.payload.session_id)).toEqual(['root', 'root']) + expect(bodies.at(-1).payload.hook_event_name).toBe('agent_end') + } + ) + + it('preserves a headless owner through reload, new, and resume', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'omp' }) + let sessionId = 'initial' + const root = { + hasUI: false, + sessionManager: { getSessionId: () => sessionId, getSessionFile: () => '/session.jsonl' } + } + await harness.callHook('session_start', {}, root) + for (const next of ['initial', 'new', 'resumed']) { + sessionId = next + harness.reload() + await harness.callHook('session_start', { reason: 'reload' }, root) + await harness.callHook('agent_start', {}, root) + await settle() + } + expect( + harness.fetchMock.mock.calls.map((call) => JSON.parse(call[1].body).payload.session_id) + ).toEqual(['initial', 'new', 'resumed']) + }) + it('isolates separately launched panes in the same host process', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'omp' }) + const parent = { + sessionManager: { getSessionId: () => 'parent', getSessionFile: () => '/parent.jsonl' } + } + const separate = { + sessionManager: { getSessionId: () => 'separate', getSessionFile: () => '/separate.jsonl' } + } + await harness.callHook('session_start', {}, parent) + harness.processEnv.ORCA_PANE_KEY = 'pane-2' + harness.processEnv.ORCA_AGENT_LAUNCH_TOKEN = 'launch-2' + harness.reload() + await harness.callHook('agent_start', {}, separate) + await settle() + expect(JSON.parse(harness.fetchMock.mock.calls[0][1].body).payload.session_id).toBe('separate') + }) + it('keeps reporting for legacy callbacks without a session manager', async () => { + const harness = createAgentStatusExtensionHarness({ kind: 'omp' }) + await harness.callHook('agent_start') + await settle() + await harness.callHook('agent_end') + await settle() + expect( + harness.fetchMock.mock.calls.map((call) => JSON.parse(call[1].body).payload.hook_event_name) + ).toEqual(['agent_start', 'agent_end']) + }) +}) diff --git a/tests/tools/omp-child-session-runtime-smoke.mjs b/tests/tools/omp-child-session-runtime-smoke.mjs new file mode 100644 index 00000000000..17bdda1a2df --- /dev/null +++ b/tests/tools/omp-child-session-runtime-smoke.mjs @@ -0,0 +1,97 @@ +// Run with Bun and a read-only OMP checkout path as the first argument. +import assert from 'node:assert/strict' +import { mkdtemp, writeFile, rm } from 'node:fs/promises' +import { createServer } from 'node:http' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' +import { pathToFileURL } from 'node:url' +import { build } from 'esbuild' + +const reference = process.argv[2] +assert.ok(reference, 'Pass the read-only oh-my-pi source checkout path') +const source = (path) => + pathToFileURL(join(resolve(reference), 'packages/coding-agent/src', path)).href +const { loadExtensions } = await import(source('extensibility/extensions/loader.ts')) +const { EventBus } = await import(source('utils/event-bus.ts')) +const { SessionManager } = await import(source('session/session-manager.ts')) +const scratch = await mkdtemp(join(tmpdir(), 'orca-omp-child-status-')) +const posts = [] +const server = createServer(async (request, response) => { + let body = '' + for await (const chunk of request) { + body += chunk + } + posts.push(JSON.parse(body).payload) + response.writeHead(200).end() +}) +await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) +try { + await build({ + entryPoints: ['src/main/pi/agent-status-extension-source.ts'], + bundle: true, + platform: 'node', + format: 'esm', + outfile: join(scratch, 'generator.mjs') + }) + const { getPiAgentStatusExtensionSource } = await import( + pathToFileURL(join(scratch, 'generator.mjs')).href + ) + const extensionPath = join(scratch, 'orca-agent-status.ts') + await writeFile(extensionPath, getPiAgentStatusExtensionSource('omp')) + process.env.ORCA_PANE_KEY = 'test-parent-pane' + process.env.ORCA_AGENT_HOOK_PORT = String(server.address().port) + process.env.ORCA_AGENT_HOOK_TOKEN = 'test-token' + delete process.env.ORCA_AGENT_HOOK_ENDPOINT + delete process.env.ORCA_PI_STATUS_OWNED + const load = async () => { + const result = await loadExtensions([extensionPath], scratch, new EventBus()) + assert.deepEqual(result.errors, []) + return result.extensions[0] + } + const rootManager = SessionManager.inMemory(scratch) + const childManager = SessionManager.inMemory(scratch) + const emit = async (extension, type, manager) => { + for (const handler of extension.handlers.get(type) ?? []) { + await handler({ type }, { sessionManager: manager, hasUI: false }) + } + await new Promise((resolve) => setTimeout(resolve, 40)) + } + const root = await load() + await emit(root, 'session_start', rootManager) + await emit(root, 'agent_start', rootManager) + const child = await load() + assert.notEqual(root, child) + await emit(child, 'session_start', childManager) + await emit(child, 'agent_start', childManager) + await emit(child, 'agent_end', childManager) + assert.deepEqual( + posts.map((post) => post.hook_event_name), + ['agent_start'] + ) + await emit(root, 'agent_end', rootManager) + await writeFile(extensionPath, `${getPiAgentStatusExtensionSource('omp')}\n// Reloaded module\n`) + const reloaded = await load() + await emit(reloaded, 'session_start', rootManager) + const previousId = rootManager.getSessionId() + await rootManager.newSession() + assert.notEqual(rootManager.getSessionId(), previousId) + await emit(reloaded, 'agent_start', rootManager) + await emit(reloaded, 'agent_end', rootManager) + assert.deepEqual( + posts.map((post) => post.hook_event_name), + ['agent_start', 'agent_end', 'agent_start', 'agent_end'] + ) + console.log( + JSON.stringify({ + platform: process.platform, + posts: posts.map((post) => post.hook_event_name), + distinctManagers: rootManager !== childManager, + scope: + 'Actual OMP loader, EventBus and in-memory SessionManager; synthetic lifecycle callbacks; real native HTTP' + }) + ) +} finally { + server.closeAllConnections() + await new Promise((resolve) => server.close(resolve)) + await rm(scratch, { recursive: true, force: true }) +}