From 6239efb34a67cf8bb454b9437524b04a166c45f5 Mon Sep 17 00:00:00 2001 From: Kelvin Amoaba <97001695+AmoabaKelvin@users.noreply.github.com> Date: Fri, 6 Nov 2026 18:12:19 +0000 Subject: [PATCH] perf(agent-session): stream chat text as appends (#39051) * perf(agent-session): stream chat text as appends * perf(mobile): read text-append support from the connection's status answer * perf(mobile): open the chat stream once the host's append support is known * refactor(mobile): seed text-append support from the shell's status * fix(agent-session): stop asking for appends after one is refused --- .../use-mobile-structured-agent-state.ts | 26 ++++ .../structured-agent-session.ts | 81 +---------- ...ructured-agent-session-subscribe-method.ts | 65 +++++++++ .../agent-session-runtime-capabilities.ts | 3 + .../agent-session-text-append.test.ts | 72 +++++++++ .../agent-session-text-append.ts | 137 ++++++++++++++++++ ...structured-agent-session-live-read.test.ts | 71 +++++++++ .../structured-agent-session-live-read.ts | 51 ++++++- .../structured-agent-session-params.ts | 6 +- 9 files changed, 433 insertions(+), 79 deletions(-) create mode 100644 src/main/runtime/rpc/methods/structured-sessions/subscriptions/structured-agent-session-subscribe-method.ts create mode 100644 src/shared/agent-session/agent-session-text-append.test.ts create mode 100644 src/shared/agent-session/agent-session-text-append.ts create mode 100644 src/shared/agent-session/structured-agent-session-live-read.test.ts diff --git a/mobile/src/session/use-mobile-structured-agent-state.ts b/mobile/src/session/use-mobile-structured-agent-state.ts index 1342defdebe..05626a309b3 100644 --- a/mobile/src/session/use-mobile-structured-agent-state.ts +++ b/mobile/src/session/use-mobile-structured-agent-state.ts @@ -19,6 +19,10 @@ import { type StructuredAgentSessionState } from '../../../src/shared/agent-session/structured-agent-session-reducer' import { subscribeConnectionRevival } from '../transport/connection-revival-signal' +import { readRuntimeStatusProbe, recordRuntimeStatusProbe } from '../transport/runtime-status-probe' +import { hostStatusProbe, readProbedHostStatus } from '../transport/host-status-probe-operations' +import { hostStatusSchema } from '../transport/host-status-reply-schema' +import { AGENT_SESSION_TEXT_APPEND_RUNTIME_CAPABILITY } from '../../../src/shared/agent-session/agent-session-runtime-capabilities' import type { RpcClient } from '../transport/rpc-client' import { agentSessionReadFailureRefusal, @@ -30,6 +34,27 @@ import { useMobileOlderPageLoader, type OlderPagePhase } from './use-mobile-olde const MAX_RETAINED_SESSION_STATES = 32 +// Read alongside the hold, so the stream's first open knows whether to ask for appends: this +// connection's probe, else the status the shell proved (direct hosts only), else one read. +async function hostStreamsTextAppends(client: RpcClient): Promise { + let status = readRuntimeStatusProbe(client) + const shell = + status === undefined ? hostStatusSchema.safeParse(client.getShellHostStatus?.()) : null + if (shell?.success) { + status = shell.data + } else if (status === undefined) { + const reply = await hostStatusProbe + .request(client, undefined, { timeoutMs: 15_000, failWhenDisconnected: true }) + .catch(() => null) + const landed = reply ? readProbedHostStatus(reply) : null + if (landed) { + recordRuntimeStatusProbe(client, landed.status) + } + status = landed?.status + } + return status?.capabilities?.includes(AGENT_SESSION_TEXT_APPEND_RUNTIME_CAPABILITY) === true +} + export function useMobileStructuredAgentState(args: { onCommandsChanged?: () => void client: RpcClient | null @@ -132,6 +157,7 @@ export function useMobileStructuredAgentState(args: { callAgentSession(client, 'agentSession.currentRequests', request), sessionId, subscribe: subscribeTranscriptAfterHold(client, held, () => ended), + textAppend: hostStreamsTextAppends(client), subscribeRetryNow: subscribeConnectionRevival }) return () => { diff --git a/src/main/runtime/rpc/methods/structured-sessions/structured-agent-session.ts b/src/main/runtime/rpc/methods/structured-sessions/structured-agent-session.ts index 763daa6ffc4..576daa69bde 100644 --- a/src/main/runtime/rpc/methods/structured-sessions/structured-agent-session.ts +++ b/src/main/runtime/rpc/methods/structured-sessions/structured-agent-session.ts @@ -1,21 +1,15 @@ import { STRUCTURED_AGENT_SESSION_ACCOUNT_USAGE_METHODS } from './catalog/structured-agent-session-account-usage' import { STRUCTURED_AGENT_SESSION_USAGE_LIMIT_CONTINUE_METHODS } from './catalog/structured-agent-session-usage-limit-continue' import { STRUCTURED_AGENT_SESSION_CLOSE_METHOD } from './conversation/structured-agent-session-close-method' -import { clientRequestsCommands } from './structured-agent-session-policy' +import { STRUCTURED_AGENT_SESSION_SUBSCRIBE_METHOD } from './subscriptions/structured-agent-session-subscribe-method' import { clientReadsCurrentRequests, STRUCTURED_AGENT_SESSION_CURRENT_REQUESTS_METHODS } from './conversation/structured-agent-session-current-requests' import { clientSupportsSnapshotCatchUp } from './history/structured-agent-session-history-capability' -import { - projectBackgroundTaskEvent, - projectBackgroundTaskHistory -} from './history/structured-agent-session-background-task-capability' -import { - projectTurnItemEvent, - projectTurnItemHistory -} from './history/structured-agent-session-turn-item-capability' -import { defineMethod, defineStreamingMethod } from '../../core' +import { projectBackgroundTaskHistory } from './history/structured-agent-session-background-task-capability' +import { projectTurnItemHistory } from './history/structured-agent-session-turn-item-capability' +import { defineMethod } from '../../core' import { ensureStructuredHostInstalled as ensureHostInstalled, requireInstalledStructuredHost as requireInstalledHost, @@ -36,14 +30,8 @@ import { STRUCTURED_AGENT_SESSION_REGENERATE_NAME_METHODS } from './conversation import { STRUCTURED_AGENT_SESSION_RENAME_METHODS } from './conversation/structured-agent-session-rename' import { STRUCTURED_AGENT_SESSION_QUEUED_METHODS } from './conversation/structured-agent-session-queued-methods' import { STRUCTURED_AGENT_SESSION_RESTART_RESUME_METHODS } from './conversation/structured-agent-session-restart-resume' -import { - bindStructuredAgentSessionStream, - STRUCTURED_AGENT_SESSION_STATUS_METHODS -} from './subscriptions/structured-agent-session-status-stream' -import { - structuredAgentSessionSubscriptionBase as subscriptionBaseFor, - structuredAgentSessionSubscriptionId as subscriptionIdFor -} from './subscriptions/structured-agent-session-subscription-id' +import { STRUCTURED_AGENT_SESSION_STATUS_METHODS } from './subscriptions/structured-agent-session-status-stream' +import { structuredAgentSessionSubscriptionBase as subscriptionBaseFor } from './subscriptions/structured-agent-session-subscription-id' import { STRUCTURED_AGENT_SESSION_TURN_COMPLETION_METHODS } from './subscriptions/structured-agent-session-turn-completion-stream' import { STRUCTURED_AGENT_SESSION_THREAD_GOAL_METHODS } from './conversation/structured-agent-session-thread-goal' import { STRUCTURED_AGENT_SESSION_CONVERSATION_OUTLINE_METHODS } from './history/structured-agent-session-conversation-outline' @@ -59,24 +47,19 @@ import { RewindParams, SendParams, SleepWorkspaceParams, - SubscribeParams, UnsubscribeParams } from '../../../../../shared/rpc-contract/structured-agent-session-params' import '../../../../../shared/rpc-contract/structured-agent-session-continue-params' import '../../../../../shared/rpc-contract/structured-agent-session-rename-params' import '../../../../../shared/rpc-contract/structured-agent-session-queued-edit-params' import { sendStructuredAgentSessionForClient } from './conversation/structured-agent-session-send-compatibility' -import { recordMobileViewOpened } from '../../../../telemetry/mobile-client-telemetry' import { runConversationCommandForClient } from './conversation/structured-agent-session-conversation-command' import { getStructuredAgentSessionHost } from '../../../../native-chat/agent-session-wire/host/structured-agent-session-registry' import { LOCAL_EXECUTION_HOST_ID } from '../../../../../shared/execution-host' import { ORCA_LAYOUT_OWNER_IS_MAIN } from '../../../../../shared/layout-owner-switch' import { refusingStructuredAgentSessionCreatesClosedFirst } from './creation/structured-agent-session-close-before-create' import * as queuedImages from './attachments/structured-agent-session-queued-image-capability' -import { - projectReceiptDeviceEvent, - projectReceiptDeviceHistory -} from './subscriptions/structured-agent-session-receipt-devices' +import { projectReceiptDeviceHistory } from './subscriptions/structured-agent-session-receipt-devices' export const STRUCTURED_AGENT_SESSION_METHODS = [ defineMethod({ @@ -207,55 +190,7 @@ export const STRUCTURED_AGENT_SESSION_METHODS = [ ) } }), - defineStreamingMethod({ - name: 'agentSession.subscribe', - permission: 'workspace', - params: SubscribeParams, - handler: async (params, ctx, emit) => { - const host = await requireInstalledHost(ctx, params.sessionId) - const subscriptionId = subscriptionIdFor(ctx, params.sessionId) - // A stream reads; it never keeps an agent alive or starts one. - let dispose = (): void => {} - const stream = bindStructuredAgentSessionStream(ctx, subscriptionId, () => dispose()) - if (stream.isClosed()) { - return - } - // Resolve after opening snapshot/missed batch; a racing close disposes the bound stream. - dispose = await host.subscribe({ - commandsOnRequest: clientRequestsCommands(ctx), - readsCurrentRequests: clientReadsCurrentRequests(ctx), - id: subscriptionId, - sessionId: params.sessionId, - ...queuedImages.queuedImageSubscriber(ctx), - projectQueuePublication: (publication) => - queuedImages.projectQueuedImageCardPublication(publication, ctx), - emit: (event) => - emit( - projectReceiptDeviceEvent( - projectTurnItemEvent( - projectBackgroundTaskEvent(event, ctx), - ctx, - host.sessionAgent(params.sessionId) - ), - ctx - ) - ), - ...(params.cursor ? { cursor: params.cursor } : {}), - snapshotCatchUp: clientSupportsSnapshotCatchUp(ctx), - // Every phone's app, whatever its version, ends a stream over a frame its bridge cannot carry. - ...(ctx.clientKind === 'mobile' ? { boundedFrames: true } : {}) - }) - if (stream.isClosed()) { - dispose() - return - } - recordMobileViewOpened(ctx, { - surface: 'chat', - viewKey: `agentSession:${params.sessionId}`, - agent: () => host.sessionAgent(params.sessionId) - }) - } - }), + STRUCTURED_AGENT_SESSION_SUBSCRIBE_METHOD, defineMethod({ name: 'agentSession.unsubscribe', permission: 'workspace', diff --git a/src/main/runtime/rpc/methods/structured-sessions/subscriptions/structured-agent-session-subscribe-method.ts b/src/main/runtime/rpc/methods/structured-sessions/subscriptions/structured-agent-session-subscribe-method.ts new file mode 100644 index 00000000000..685ac7f0385 --- /dev/null +++ b/src/main/runtime/rpc/methods/structured-sessions/subscriptions/structured-agent-session-subscribe-method.ts @@ -0,0 +1,65 @@ +import { defineStreamingMethod } from '../../../core' +import { sendingTextAppends } from '../../../../../../shared/agent-session/agent-session-text-append' +import { recordMobileViewOpened } from '../../../../../telemetry/mobile-client-telemetry' +import { clientRequestsCommands } from '../structured-agent-session-policy' +import { clientReadsCurrentRequests } from '../conversation/structured-agent-session-current-requests' +import { clientSupportsSnapshotCatchUp } from '../history/structured-agent-session-history-capability' +import { projectBackgroundTaskEvent } from '../history/structured-agent-session-background-task-capability' +import { projectTurnItemEvent } from '../history/structured-agent-session-turn-item-capability' +import { requireInstalledStructuredHost as requireInstalledHost } from '../creation/structured-agent-session-gate' +import { bindStructuredAgentSessionStream } from './structured-agent-session-status-stream' +import { structuredAgentSessionSubscriptionId as subscriptionIdFor } from './structured-agent-session-subscription-id' +import { SubscribeParams } from '../../../../../../shared/rpc-contract/structured-agent-session-params' +import * as queuedImages from '../attachments/structured-agent-session-queued-image-capability' +import { projectReceiptDeviceEvent } from './structured-agent-session-receipt-devices' + +export const STRUCTURED_AGENT_SESSION_SUBSCRIBE_METHOD = defineStreamingMethod({ + name: 'agentSession.subscribe', + permission: 'workspace', + params: SubscribeParams, + handler: async (params, ctx, emitFrame) => { + const emit = params.capabilities?.textAppend === 1 ? sendingTextAppends(emitFrame) : emitFrame + const host = await requireInstalledHost(ctx, params.sessionId) + const subscriptionId = subscriptionIdFor(ctx, params.sessionId) + // A stream reads; it never keeps an agent alive or starts one. + let dispose = (): void => {} + const stream = bindStructuredAgentSessionStream(ctx, subscriptionId, () => dispose()) + if (stream.isClosed()) { + return + } + // Resolve after opening snapshot/missed batch; a racing close disposes the bound stream. + dispose = await host.subscribe({ + commandsOnRequest: clientRequestsCommands(ctx), + readsCurrentRequests: clientReadsCurrentRequests(ctx), + id: subscriptionId, + sessionId: params.sessionId, + ...queuedImages.queuedImageSubscriber(ctx), + projectQueuePublication: (publication) => + queuedImages.projectQueuedImageCardPublication(publication, ctx), + emit: (event) => + emit( + projectReceiptDeviceEvent( + projectTurnItemEvent( + projectBackgroundTaskEvent(event, ctx), + ctx, + host.sessionAgent(params.sessionId) + ), + ctx + ) + ), + ...(params.cursor ? { cursor: params.cursor } : {}), + snapshotCatchUp: clientSupportsSnapshotCatchUp(ctx), + // Every phone's app, whatever its version, ends a stream over a frame its bridge cannot carry. + ...(ctx.clientKind === 'mobile' ? { boundedFrames: true } : {}) + }) + if (stream.isClosed()) { + dispose() + return + } + recordMobileViewOpened(ctx, { + surface: 'chat', + viewKey: `agentSession:${params.sessionId}`, + agent: () => host.sessionAgent(params.sessionId) + }) + } +}) diff --git a/src/shared/agent-session/agent-session-runtime-capabilities.ts b/src/shared/agent-session/agent-session-runtime-capabilities.ts index 5ad2b316fd1..f2e24aa61b0 100644 --- a/src/shared/agent-session/agent-session-runtime-capabilities.ts +++ b/src/shared/agent-session/agent-session-runtime-capabilities.ts @@ -113,6 +113,8 @@ export const AGENT_SESSION_TEXT_ATTACHMENT_READ_CAPABILITY = // Readers must explicitly honor stoppable:false on live child rows. The Stop capability alone predates this field and cannot prove safe rendering. export const AGENT_SESSION_BACKGROUND_TASK_ROW_STOP_CAPABILITY = 'agent-session.background-task-row-stop.v1' as const +// Strict subscribe params refuse `capabilities` on older hosts; a client sends it only when advertised. +export const AGENT_SESSION_TEXT_APPEND_RUNTIME_CAPABILITY = 'agent-session.text-append.v1' as const // Individual client flags must not retain the unused host capability aggregate. export const AGENT_SESSION_RUNTIME_CAPABILITIES = /* @__PURE__ */ (() => @@ -160,5 +162,6 @@ export const AGENT_SESSION_RUNTIME_CAPABILITIES = /* @__PURE__ */ (() => AGENT_SESSION_TURN_ITEM_CAPABILITY, AGENT_SESSION_BACKGROUND_TASK_ROW_STOP_CAPABILITY, AGENT_SESSION_BACKGROUND_TASK_CHILD_VIEWS_CAPABILITY, + AGENT_SESSION_TEXT_APPEND_RUNTIME_CAPABILITY, ...AGENT_SESSION_RESUME_RUNTIME_CAPABILITIES ] as const)() diff --git a/src/shared/agent-session/agent-session-text-append.test.ts b/src/shared/agent-session/agent-session-text-append.test.ts new file mode 100644 index 00000000000..2d756d5d9b7 --- /dev/null +++ b/src/shared/agent-session/agent-session-text-append.test.ts @@ -0,0 +1,72 @@ +import { describe, expect, it } from 'vitest' +import type { AgentJournalRenderItem } from './journal/agent-session-journal-types' +import type { AgentSessionSubscribeEvent } from './agent-session-wire' +import { + type AgentSessionTextAppendItem, + createAgentSessionTextAppendReader, + sendingTextAppends +} from './agent-session-text-append' + +function reply(text: string, revision: number): AgentJournalRenderItem { + return { + itemId: 'reply-1', + revision, + sequence: 4, + observedAt: 1000, + body: { kind: 'message', role: 'assistant', blocks: [{ type: 'text', text }] } + } +} + +function appended(revision: number, from: number, text: string): AgentSessionTextAppendItem { + return { ...reply('', revision), textAppend: { from, text } } +} + +function batch(sequence: number, items: AgentJournalRenderItem[]): AgentSessionSubscribeEvent { + return { + type: 'batch', + sessionId: 'session-1', + batch: { cursor: { epoch: 'epoch-1', sequence }, items, removedItemIds: [], submissions: [] } + } +} + +const snapshot: AgentSessionSubscribeEvent = { + type: 'snapshot', + sessionId: 'session-1', + fence: 1, + page: { + sessionId: 'session-1', + epoch: 'epoch-1', + direction: 'before', + items: [], + removedItemIds: [], + submissions: [], + window: { oldest: null, newest: null, nextCursor: { epoch: 'epoch-1', sequence: 3 } }, + hasOlder: false, + hasNewer: false + } +} + +function sentItems(event: AgentSessionSubscribeEvent): unknown[] { + return event.type === 'batch' ? event.batch.items : [] +} + +describe('agent session text appends', () => { + it('sends a growing reply whole once, then only its new text until a snapshot', () => { + const send = sendingTextAppends((event: AgentSessionSubscribeEvent) => event) + send(snapshot) + expect(sentItems(send(batch(4, [reply('Hel', 1)])))).toEqual([reply('Hel', 1)]) + expect(sentItems(send(batch(5, [reply('Hello wor', 2)])))).toEqual([appended(2, 3, 'lo wor')]) + // A catch-up snapshot replaces what the reader holds, so the reply goes whole again. + send(snapshot) + expect(sentItems(send(batch(6, [reply('Hello world', 3)])))).toEqual([reply('Hello world', 3)]) + }) + + it('restores the whole reply from an append', () => { + const read = createAgentSessionTextAppendReader() + read(snapshot) + read(batch(4, [reply('Hel', 1)])) + expect(sentItems(read(batch(5, [appended(2, 3, 'lo wor')])) ?? batch(5, []))).toEqual([ + reply('Hello wor', 2) + ]) + }) +}) diff --git a/src/shared/agent-session/agent-session-text-append.ts b/src/shared/agent-session/agent-session-text-append.ts new file mode 100644 index 00000000000..8a33fc26890 --- /dev/null +++ b/src/shared/agent-session/agent-session-text-append.ts @@ -0,0 +1,137 @@ +import type { AgentJournalRenderItem } from './journal/agent-session-journal-types' +import type { AgentSessionSubscribeEvent } from './agent-session-wire' + +/** A growing reply sent as only its new text: its text block is empty on the wire, and the reader + * restores it as the text it holds for the item (`from` long) followed by `text`. */ +export type AgentSessionTextAppendItem = AgentJournalRenderItem & { + textAppend: { from: number; text: string } +} + +// Bases are only unfinished replies, so a few cover every stream a turn grows at once. +const MAX_TEXT_APPEND_BASES = 8 + +function growingText(item: AgentJournalRenderItem): string | null { + const { body } = item + if ( + item.shortened || + body.kind !== 'message' || + (body.role !== 'assistant' && body.role !== 'reasoning') || + body.blocks.length !== 1 || + body.blocks[0]?.type !== 'text' + ) { + return null + } + return body.blocks[0].text +} + +function isTextAppendItem(item: AgentJournalRenderItem): item is AgentSessionTextAppendItem { + if (!('textAppend' in item)) { + return false + } + const append: unknown = item.textAppend + return ( + typeof append === 'object' && + append !== null && + 'from' in append && + typeof append.from === 'number' && + 'text' in append && + typeof append.text === 'string' + ) +} + +function withText(item: AgentJournalRenderItem, text: string): AgentJournalRenderItem { + if (item.body.kind !== 'message' || item.body.blocks[0]?.type !== 'text') { + return item + } + return { ...item, body: { ...item.body, blocks: [{ ...item.body.blocks[0], text }] } } +} + +/** One stream's sender or reader. Both walk every frame of the stream in order with the same + * rules, so they hold the same base text for each item without exchanging anything else. */ +function textAppendStream( + mapItem: (base: string | undefined, item: AgentJournalRenderItem) => AgentJournalRenderItem | null +): (event: AgentSessionSubscribeEvent) => AgentSessionSubscribeEvent | null { + const bases = new Map() + let epoch: string | null = null + const note = (item: AgentJournalRenderItem): void => { + bases.delete(item.itemId) + const text = growingText(item) + if (text === null || (item.body.kind === 'message' && item.body.completedAt !== undefined)) { + return + } + bases.set(item.itemId, text) + const oldest = bases.size > MAX_TEXT_APPEND_BASES ? bases.keys().next().value : undefined + if (oldest !== undefined) { + bases.delete(oldest) + } + } + return (event) => { + if (event.type === 'end') { + return event + } + if (event.type === 'snapshot' || event.type === 'reset') { + bases.clear() + epoch = event.page.epoch + event.page.items.forEach(note) + return event + } + // A new journal epoch renumbers rows, so nothing from the old one is a base. + if (event.batch.cursor.epoch !== epoch) { + bases.clear() + epoch = event.batch.cursor.epoch + } + for (const itemId of event.batch.removedItemIds) { + bases.delete(itemId) + } + let changed = false + const items: AgentJournalRenderItem[] = [] + for (const item of event.batch.items) { + const mapped = mapItem(bases.get(item.itemId), item) + if (mapped === null) { + return null + } + changed ||= mapped !== item + items.push(mapped) + note('textAppend' in mapped ? item : mapped) + } + return changed ? { ...event, batch: { ...event.batch, items } } : event + } +} + +/** Host side, wrapping a stream's emit after every other projection: a reply that extends what + * this stream last sent for it goes out as its new text. */ +export function sendingTextAppends( + emit: (event: AgentSessionSubscribeEvent) => R +): (event: AgentSessionSubscribeEvent) => R { + const send = textAppendStream((base, item) => { + const text = growingText(item) + if (base === undefined || text === null || !text.startsWith(base)) { + return item + } + return { + ...withText(item, ''), + textAppend: { from: base.length, text: text.slice(base.length) } + } satisfies AgentSessionTextAppendItem + }) + return (event) => emit(send(event) ?? event) +} + +/** Client side: whole items back from appends; null when an append does not start where this + * reader's text for the item ends, and only a new stream (which sends whole items) can recover. */ +export function createAgentSessionTextAppendReader(): ( + event: AgentSessionSubscribeEvent +) => AgentSessionSubscribeEvent | null { + return textAppendStream((base, item) => { + if (!('textAppend' in item)) { + return item + } + if (!isTextAppendItem(item)) { + return null + } + const { textAppend, ...rest } = item + if (base === undefined || base.length !== textAppend.from || growingText(rest) !== '') { + return null + } + return withText(rest, base + textAppend.text) + }) +} diff --git a/src/shared/agent-session/structured-agent-session-live-read.test.ts b/src/shared/agent-session/structured-agent-session-live-read.test.ts new file mode 100644 index 00000000000..3039063c31f --- /dev/null +++ b/src/shared/agent-session/structured-agent-session-live-read.test.ts @@ -0,0 +1,71 @@ +import { describe, expect, it, vi } from 'vitest' +import type { AgentJournalRenderItem } from './journal/agent-session-journal-types' +import type { AgentSessionSubscribeEvent } from './agent-session-wire' +import type { AgentSessionTextAppendItem } from './agent-session-text-append' +import { + startStructuredAgentSessionLiveRead, + type StructuredAgentSessionLiveSubscribe +} from './structured-agent-session-live-read' + +function reply(text: string, revision: number): AgentJournalRenderItem { + return { + itemId: 'reply-1', + revision, + sequence: 4, + observedAt: 1000, + body: { kind: 'message', role: 'assistant', blocks: [{ type: 'text', text }] } + } +} + +function appended(revision: number, from: number, text: string): AgentSessionTextAppendItem { + return { ...reply('', revision), textAppend: { from, text } } +} + +function batch(sequence: number, items: AgentJournalRenderItem[]): AgentSessionSubscribeEvent { + return { + type: 'batch', + sessionId: 'session-1', + batch: { cursor: { epoch: 'epoch-1', sequence }, items, removedItemIds: [], submissions: [] } + } +} + +describe('structured agent session live read with text appends', () => { + it('drops a stream whose append does not match and resumes without appends', async () => { + const opened: { params: unknown; emit: (event: AgentSessionSubscribeEvent) => void }[] = [] + const unsubscribed: number[] = [] + const subscribe: StructuredAgentSessionLiveSubscribe = async (params, onEvent) => { + const index = opened.push({ params, emit: onEvent }) - 1 + return { unsubscribe: () => unsubscribed.push(index) } + } + const applied: AgentSessionSubscribeEvent[] = [] + const read = startStructuredAgentSessionLiveRead({ + applyEvent: (event) => applied.push(event), + applyError: () => {}, + getCursor: () => ({ epoch: 'epoch-1', sequence: 4 }), + onContactLost: () => {}, + onHistoryReadInvalidated: () => {}, + sessionId: 'session-1', + subscribe, + textAppend: Promise.resolve(true), + readCurrentRequests: vi.fn(), + subscribeRetryNow: () => () => {} + }) + await vi.waitFor(() => expect(opened).toHaveLength(1)) + expect(opened[0]?.params).toEqual({ + sessionId: 'session-1', + cursor: { epoch: 'epoch-1', sequence: 4 }, + capabilities: { textAppend: 1 } + }) + opened[0]?.emit(batch(4, [reply('Hel', 1)])) + opened[0]?.emit(batch(5, [appended(2, 7, 'lo')])) + + await vi.waitFor(() => expect(opened).toHaveLength(2)) + expect(opened[1]?.params).toEqual({ + sessionId: 'session-1', + cursor: { epoch: 'epoch-1', sequence: 4 } + }) + expect(unsubscribed).toEqual([0]) + expect(applied).toEqual([batch(4, [reply('Hel', 1)])]) + read.dispose() + }) +}) diff --git a/src/shared/agent-session/structured-agent-session-live-read.ts b/src/shared/agent-session/structured-agent-session-live-read.ts index de4c73a9d24..a66fe54fbfa 100644 --- a/src/shared/agent-session/structured-agent-session-live-read.ts +++ b/src/shared/agent-session/structured-agent-session-live-read.ts @@ -1,6 +1,7 @@ import type { AgentJournalCursor } from './journal/agent-session-journal-identity' import type { AgentSessionSubscribeEvent } from './agent-session-wire' import { createStructuredAgentSessionEventCoalescer } from './structured-agent-session-coalescer' +import { createAgentSessionTextAppendReader } from './agent-session-text-append' import { AGENT_SESSION_UNATTACHED_READ_GRACE_MS, isFinalAgentSessionReadRefusal, @@ -18,7 +19,7 @@ import { createReconnectScheduler } from './structured-agent-session-reconnect' /** One `agentSession.subscribe` stream, opened by whichever transport the client has. */ export type StructuredAgentSessionLiveSubscribe = ( - params: { sessionId: string; cursor?: AgentJournalCursor }, + params: { sessionId: string; cursor?: AgentJournalCursor; capabilities?: { textAppend: 1 } }, onEvent: (event: AgentSessionSubscribeEvent) => void, onError: (error: unknown) => void, onClose: () => void @@ -56,6 +57,8 @@ export function startStructuredAgentSessionLiveRead(args: { hydrate?: (shouldStop: () => boolean) => Promise sessionId: string subscribe: StructuredAgentSessionLiveSubscribe + /** Whether the host advertises AGENT_SESSION_TEXT_APPEND_RUNTIME_CAPABILITY; opens wait for it. */ + textAppend?: Promise readCurrentRequests: ( request: AgentSessionCurrentRequestsRead ) => Promise @@ -76,6 +79,11 @@ export function startStructuredAgentSessionLiveRead(args: { let openingFrameGeneration = 0 let stateGeneration = 0 let unsubscribe = (): void => {} + let readTextAppends: ReturnType | null = null + /** An open whose stream must not be applied further; the next open replaces it. */ + let replacedOpenGeneration = 0 + // A refused append means the two sides disagree; later streams take whole items so it cannot loop. + let appendsRefused = false let shouldStopCoalescedEvent = (): boolean => true const coalescer = createStructuredAgentSessionEventCoalescer((event) => { if (!shouldStopCoalescedEvent()) { @@ -136,8 +144,30 @@ export function startStructuredAgentSessionLiveRead(args: { replaceMarker: (marker) => args.applyRequests?.({ type: 'requests-marker', marker }), error: reportReadFailure }) - const handleEvent = (event: AgentSessionSubscribeEvent, eventOpenGeneration: number): void => { - if (!isCurrentOpenGeneration(eventOpenGeneration)) { + /** Resume from the last applied cursor on a new stream, which starts with whole items. */ + const replaceStream = (eventOpenGeneration: number): void => { + replacedOpenGeneration = eventOpenGeneration + appendsRefused = true + if (!connected) { + // Still opening: retry later; the open drops its handle when it lands. + loseContact() + return + } + unsubscribe() + unsubscribe = (): void => {} + connected = false + void open() + } + const handleEvent = (received: AgentSessionSubscribeEvent, eventOpenGeneration: number): void => { + if ( + !isCurrentOpenGeneration(eventOpenGeneration) || + eventOpenGeneration === replacedOpenGeneration + ) { + return + } + const event = readTextAppends ? readTextAppends(received) : received + if (!event) { + replaceStream(eventOpenGeneration) return } clearUnattachedReadGrace() @@ -199,13 +229,20 @@ export function startStructuredAgentSessionLiveRead(args: { unsubscribe() unsubscribe = (): void => {} try { + // Desktop passes none, so its subscribe still goes out in the same turn. + const textAppend = args.textAppend && !appendsRefused ? await args.textAppend : false if (!isCurrentOpenGeneration(currentOpenGeneration)) { return } let closedDuringOpen = false const cursor = args.getCursor() + readTextAppends = textAppend ? createAgentSessionTextAppendReader() : null const handle = await args.subscribe( - { sessionId: args.sessionId, ...(cursor ? { cursor } : {}) }, + { + sessionId: args.sessionId, + ...(cursor ? { cursor } : {}), + ...(textAppend ? { capabilities: { textAppend: 1 as const } } : {}) + }, (event) => handleEvent(event, currentOpenGeneration), (error) => { if (!isCurrentOpenGeneration(currentOpenGeneration)) { @@ -223,7 +260,11 @@ export function startStructuredAgentSessionLiveRead(args: { loseContact() } ) - if (!isCurrentOpenGeneration(currentOpenGeneration) || closedDuringOpen) { + if ( + !isCurrentOpenGeneration(currentOpenGeneration) || + closedDuringOpen || + replacedOpenGeneration === currentOpenGeneration + ) { handle.unsubscribe() if (isCurrentOpenGeneration(currentOpenGeneration)) { reconnectScheduler.schedule() diff --git a/src/shared/rpc-contract/structured-agent-session-params.ts b/src/shared/rpc-contract/structured-agent-session-params.ts index fe6bc8334db..cc38e5354eb 100644 --- a/src/shared/rpc-contract/structured-agent-session-params.ts +++ b/src/shared/rpc-contract/structured-agent-session-params.ts @@ -335,7 +335,11 @@ export const HistoryParams = z }) .strict() -export const SubscribeParams = OptionsParams.extend({ cursor: JournalCursor.optional() }).strict() +export const SubscribeParams = OptionsParams.extend({ + cursor: JournalCursor.optional(), + // Not strict: this host drops keys a newer client adds; a new value for a known key still needs a capability. + capabilities: z.object({ textAppend: z.literal(1).optional() }).optional() +}).strict() // Not strict: an older host ignores these params entirely, and this host must ignore a newer client's. export const SubscribeTurnCompletionsParams = z.object({ includePrompts: z.boolean().optional() })