orca/tests/e2e/cross-version-wire/structured-agent-session-surface-execution.ts

import { availableRevealWorkspace } from '../../../src/main/runtime/session-tabs/structured-workspace-authority.test-fixture'
import { expect, vi } from 'vitest'
import type { structuredHostStub } from './structured-agent-session-host-fixture'
import { AGENT_SESSION_ITEM_DETAIL_RUNTIME_CAPABILITY } from '../../../src/shared/agent-session/agent-session-runtime-capabilities'
import type { OrcaRuntimeService } from '../../../src/main/runtime/orca-runtime'
import type { StructuredSessionSurfaceRuntime } from '../../../src/main/runtime/agent-session/structured-agent-session-close'
import { RuntimeSubscriptionRegistry } from '../../../src/main/runtime/client/runtime-subscription-registry'
import { resetStructuredAgentSessionClosesBeforeCreateForTests } from '../../../src/main/runtime/rpc/methods/structured-sessions/creation/structured-agent-session-close-before-create'
import { AGENT_SESSION_CREATE_MESSAGE_RUNTIME_CAPABILITY } from '../../../src/shared/agent-session/agent-session-create-capabilities'
import {
  attachParams,
  ATTENTION_READ,
  paramsFor,
  SESSION,
  WORKSPACE,
  STRUCTURED_CALLS
} from './structured-agent-session-surface-manifest'
import type {
  AgentSessionWireBuild,
  RpcClientIdentity,
  RpcReply
} from './versioned-agent-session-wire'

export function runtimeStub(overrides: Record<string, unknown> = {}) {
  // A fresh runtime remembers no chat closed before its create; scenarios reuse one session id.
  resetStructuredAgentSessionClosesBeforeCreateForTests()
  const subscriptions = new RuntimeSubscriptionRegistry()
  let raw: ReturnType<
    StructuredSessionSurfaceRuntime['currentLocalStructuredSessionPresentation']
  > = {
    worktreeId: WORKSPACE,
    tabId: `tab-${SESSION}`
  }
  let saved = [{ ...raw }]
  const presentation = {
    currentLocalStructuredSessionPresentation: vi.fn<
      StructuredSessionSurfaceRuntime['currentLocalStructuredSessionPresentation']
    >((id) => (id === SESSION ? raw : null)),
    captureStructuredAgentSessionPresentation: vi.fn<
      StructuredSessionSurfaceRuntime['captureStructuredAgentSessionPresentation']
    >((id) => (id === SESSION ? { raw, saved: [...saved] } : { raw: null, saved: [] })),
    retireStructuredAgentSessionTabFromSnapshot: vi.fn<
      StructuredSessionSurfaceRuntime['retireStructuredAgentSessionTabFromSnapshot']
    >((id, tabId, worktreeId) => {
      if (id !== SESSION || !raw || raw.tabId !== tabId || raw.worktreeId !== worktreeId) {
        return false
      }
      raw = null
      return true
    }),
    removeCapturedStructuredSessionSavedTabs: vi.fn<
      StructuredSessionSurfaceRuntime['removeCapturedStructuredSessionSavedTabs']
    >((id, targets) => {
      if (id === SESSION && raw === null) {
        saved = saved.filter(
          (tab) =>
            !targets.some(
              (target) => target.tabId === tab.tabId && target.worktreeId === tab.worktreeId
            )
        )
      }
    })
  } satisfies StructuredSessionSurfaceRuntime
  return {
    ...presentation,
    getRuntimeId: () => 'runtime-1',
    getClientSettings: () => ({ experimentalStructuredNativeChat: true }),
    ensureStructuredAgentSessionHost: async () => undefined,
    prepareStructuredAgentSessionRevealWorkspace: vi.fn(availableRevealWorkspace),
    getStructuredAgentSessionCreateSupport: async () => ({ supported: true }),
    structuredAgentSessionLaunchSeedOptions: () => undefined,
    resolveStructuredAgentSessionCreateIntent: async () => {
      const {
        envelope: _envelope,
        providerHandle: _providerHandle,
        ...resolved
      } = attachParams(null)
      return resolved
    },
    publishStructuredAgentSessionTab: vi.fn<OrcaRuntimeService['publishStructuredAgentSessionTab']>(
      async (input) =>
        input.capture?.({
          tableId: raw?.tabId ?? null,
          presentation: { raw, saved: [...saved] }
        })
    ),
    subscribeStructuredAgentAvailability: (read: (agents: unknown[]) => void) => {
      read([])
      return () => {}
    },
    registerSubscriptionCleanup: subscriptions.register.bind(subscriptions),
    registerOwnedSubscriptionCleanup: subscriptions.registerOwned.bind(subscriptions),
    cleanupSubscription: subscriptions.cleanup.bind(subscriptions),
    cleanupSubscriptionsByPrefix: subscriptions.cleanupByPrefix.bind(subscriptions),
    ...overrides
  }
}

// Collect every streaming reply: refusals must arrive explicitly, never as silence.
export async function callBuild<Result = unknown>(
  build: AgentSessionWireBuild,
  method: string,
  params: unknown,
  client: RpcClientIdentity,
  runtime: unknown = runtimeStub(),
  onReply?: (reply: RpcReply<Result>) => void
): Promise<RpcReply<Result>[]> {
  const replies: RpcReply<Result>[] = []
  await build.createDispatcher(runtime).dispatchStreaming(
    { id: `request-${method}`, authToken: 'cross-version-token', method, params },
    (raw) => {
      // The build serializes its own reply; each surface case checks its result/error shape.
      const reply: RpcReply<Result> = JSON.parse(raw)
      replies.push(reply)
      onReply?.(reply)
    },
    client
  )
  return replies
}

// Require every declared method to reach its host and return its declared result in each build pairing.
export async function expectDeclaredSurfaceExecutes(
  build: AgentSessionWireBuild,
  hostCalls: ReturnType<typeof structuredHostStub>,
  clientCapabilities: readonly string[]
): Promise<void> {
  const retirement = vi.fn<OrcaRuntimeService['acknowledgeNotificationAttention']>()
  const runtime = runtimeStub({ acknowledgeNotificationAttention: retirement })
  for (const { method, hostMethod, result, clientCapability } of STRUCTURED_CALLS) {
    const atRestCreate =
      method === 'agentSession.create' &&
      build.capabilities.includes(AGENT_SESSION_CREATE_MESSAGE_RUNTIME_CAPABILITY) &&
      clientCapabilities.includes(AGENT_SESSION_CREATE_MESSAGE_RUNTIME_CAPABILITY)
    const expectedHostMethod = atRestCreate ? 'create' : hostMethod
    const otherCreateMethod =
      method === 'agentSession.create' ? (atRestCreate ? 'attach' : 'create') : null
    // Shared host methods need per-call deltas; an earlier call must not satisfy this assertion.
    const before = expectedHostMethod ? hostCalls[expectedHostMethod].mock.calls.length : 0
    const otherBefore = otherCreateMethod ? hostCalls[otherCreateMethod].mock.calls.length : 0
    if (method === 'agentSession.readVisual') {
      hostCalls.getConversation.mockReturnValueOnce(null)
    }
    const replies = await callBuild(
      build,
      method,
      paramsFor(method),
      { clientKind: 'runtime', clientCapabilities },
      runtime
    )
    if (clientCapability && !clientCapabilities.includes(clientCapability)) {
      expect(replies).toMatchObject([
        {
          ok: false,
          error: { message: expect.stringContaining('structured_agent_session_unsupported') }
        }
      ])
      if (expectedHostMethod) {
        expect(hostCalls[expectedHostMethod].mock.calls.length).toBe(before)
      }
      continue
    }
    if (expectedHostMethod) {
      expect(
        hostCalls[expectedHostMethod].mock.calls.length - before,
        `${build.label}: ${method} did not reach the host: ${JSON.stringify(replies)}`
      ).toBe(1)
    }
    if (otherCreateMethod) {
      expect(hostCalls[otherCreateMethod].mock.calls.length - otherBefore).toBe(0)
    }
    if (method === 'agentSession.close') {
      const target = { worktreeId: WORKSPACE, tabId: `tab-${SESSION}` }
      expect(runtime.captureStructuredAgentSessionPresentation).toHaveBeenCalledExactlyOnceWith(
        SESSION
      )
      expect(runtime.retireStructuredAgentSessionTabFromSnapshot).toHaveBeenCalledExactlyOnceWith(
        SESSION,
        target.tabId,
        target.worktreeId
      )
      expect(runtime.removeCapturedStructuredSessionSavedTabs).toHaveBeenCalledExactlyOnceWith(
        SESSION,
        [target]
      )
      expect(hostCalls.getSessionTabId(SESSION)).toBeNull()
      expect(hostCalls.notifySessionTabHidden).toHaveBeenCalledExactlyOnceWith(SESSION)
      expect(runtime.captureStructuredAgentSessionPresentation(SESSION)).toEqual({
        raw: null,
        saved: []
      })
    }
    if (method === 'agentSession.acknowledgeAttention') {
      expect(retirement).toHaveBeenCalledExactlyOnceWith({
        reads: [
          {
            ...ATTENTION_READ,
            scope: {
              executionHostId: 'local',
              wslDistro: null,
              workspaceId: WORKSPACE,
              workspaceKind: 'git-worktree'
            }
          }
        ]
      })
    }
    for (const reply of replies) {
      expect(
        reply,
        `${build.label}: ${method} was refused: ${JSON.stringify(reply)}`
      ).toMatchObject({ ok: true })
    }
    if (result) {
      // Assert the exact envelope so registered-but-failing handlers cannot pass.
      expect(replies, `${build.label}: ${method} must answer exactly once`).toHaveLength(1)
      expect(replies[0], `${build.label}: ${method} answered off-contract`).toMatchObject({
        ok: true,
        result
      })
    }
  }
}

export function expectItemDetailNegotiation(builds: readonly AgentSessionWireBuild[]): void {
  expect(builds[0]?.capabilities).toContain(AGENT_SESSION_ITEM_DETAIL_RUNTIME_CAPABILITY)
  for (const build of builds) {
    expect(build.capabilities.includes(AGENT_SESSION_ITEM_DETAIL_RUNTIME_CAPABILITY)).toBe(
      build.methodNames.includes('agentSession.item')
    )
  }
}