diff --git a/cli/src/agent/messageConverter.test.ts b/cli/src/agent/messageConverter.test.ts index 4068c0e1..ee794715 100644 --- a/cli/src/agent/messageConverter.test.ts +++ b/cli/src/agent/messageConverter.test.ts @@ -2,6 +2,35 @@ import { describe, expect, it } from 'vitest'; import { convertAgentMessage } from './messageConverter'; describe('convertAgentMessage', () => { + it('preserves a stable text stream id on the message wire payload', () => { + const converted = convertAgentMessage({ + type: 'text', + text: 'partial response', + id: 'text-stream-1', + live: true, + streamSnapshot: true + }); + + expect(converted).toEqual({ + type: 'message', + message: 'partial response', + id: 'text-stream-1', + streamSnapshot: true + }); + }); + + it('keeps legacy text payloads free of a stream id', () => { + const converted = convertAgentMessage({ + type: 'text', + text: 'complete response' + }); + + expect(converted).toEqual({ + type: 'message', + message: 'complete response' + }); + }); + it('keeps tool-call status when converting ACP tool events', () => { const converted = convertAgentMessage({ type: 'tool_call', @@ -37,6 +66,26 @@ describe('convertAgentMessage', () => { }); }); + it('preserves running tool progress without changing the tool input', () => { + const converted = convertAgentMessage({ + type: 'tool_call', + id: 'call-progress', + name: 'Bash', + input: { command: 'bun test' }, + status: 'in_progress', + progress: { stdout: 'running tests...\\n' } + }); + + expect(converted).toEqual({ + type: 'tool-call', + callId: 'call-progress', + name: 'Bash', + input: { command: 'bun test' }, + status: 'in_progress', + progress: { stdout: 'running tests...\\n' } + }); + }); + it('marks failed tool results as error', () => { const converted = convertAgentMessage({ type: 'tool_result', diff --git a/cli/src/agent/messageConverter.ts b/cli/src/agent/messageConverter.ts index 68e5395c..d9867be9 100644 --- a/cli/src/agent/messageConverter.ts +++ b/cli/src/agent/messageConverter.ts @@ -2,7 +2,7 @@ import { randomUUID } from 'node:crypto'; import type { AgentMessage, PlanItem } from './types'; export type CodexMessage = - | { type: 'message'; message: string } + | { type: 'message'; message: string; id?: string; streamSnapshot?: boolean } | { type: 'reasoning'; message: string; id: string } | { type: 'token_count'; @@ -28,6 +28,7 @@ export type CodexMessage = status?: 'pending' | 'in_progress' | 'completed' | 'failed'; nativeTitle?: string; nativeKind?: string; + progress?: unknown; } | { type: 'tool-call-result'; @@ -41,7 +42,12 @@ export type CodexMessage = export function convertAgentMessage(message: AgentMessage, model?: string | null): CodexMessage | null { switch (message.type) { case 'text': - return { type: 'message', message: message.text }; + return { + type: 'message', + message: message.text, + ...(message.id !== undefined ? { id: message.id } : {}), + ...(message.streamSnapshot === true ? { streamSnapshot: true } : {}) + }; case 'reasoning': // AgentMessage uses `text` (consistent with the `text` variant); // the wire-level CodexMessage uses `message` to match the @@ -76,7 +82,8 @@ export function convertAgentMessage(message: AgentMessage, model?: string | null input: message.input, status: message.status, ...(message.title ? { nativeTitle: message.title } : {}), - ...(message.kind ? { nativeKind: message.kind } : {}) + ...(message.kind ? { nativeKind: message.kind } : {}), + ...(message.progress !== undefined ? { progress: message.progress } : {}) }; case 'tool_result': return { diff --git a/cli/src/agent/sessionFactory.test.ts b/cli/src/agent/sessionFactory.test.ts index 77095401..01789e1a 100644 --- a/cli/src/agent/sessionFactory.test.ts +++ b/cli/src/agent/sessionFactory.test.ts @@ -164,6 +164,8 @@ describe('bootstrapExistingSession', () => { }, tools: ['read_file'], slashCommands: ['/compact'], + conversationHistoryPoints: { 'local-user-1': true }, + conversationHistoryEntryIds: { 'local-user-1': 'pi-entry-1' }, capabilities: { terminal: true, conversationHistory: { forkCurrent: true } @@ -208,6 +210,8 @@ describe('bootstrapExistingSession', () => { }, tools: ['read_file'], slashCommands: ['/compact'], + conversationHistoryPoints: { 'local-user-1': true }, + conversationHistoryEntryIds: { 'local-user-1': 'pi-entry-1' }, capabilities: { terminal: true, conversationHistory: { forkCurrent: true } @@ -217,13 +221,15 @@ describe('bootstrapExistingSession', () => { const updateHandler = sessionClient.updateMetadata.mock.calls[0][0] expect(updateHandler(session.metadata)).toEqual(expect.objectContaining({ codexSessionId: 'codex-thread-1', - grokSessionId: 'grok-thread-1' + grokSessionId: 'grok-thread-1', + conversationHistoryEntryIds: { 'local-user-1': 'pi-entry-1' } })) expect(notifyRunnerSessionStartedMock).toHaveBeenCalledWith( 'hapi-session-1', expect.objectContaining({ codexSessionId: 'codex-thread-1', - grokSessionId: 'grok-thread-1' + grokSessionId: 'grok-thread-1', + conversationHistoryEntryIds: { 'local-user-1': 'pi-entry-1' } }) ) }) diff --git a/cli/src/agent/sessionFactory.ts b/cli/src/agent/sessionFactory.ts index d20f89e4..afea22b3 100644 --- a/cli/src/agent/sessionFactory.ts +++ b/cli/src/agent/sessionFactory.ts @@ -128,6 +128,9 @@ function pickExistingSessionMetadata(metadata: Metadata | null | undefined): Par if (metadata.conversationHistoryTurns !== undefined) { preserved.conversationHistoryTurns = metadata.conversationHistoryTurns } + if (metadata.conversationHistoryEntryIds !== undefined) { + preserved.conversationHistoryEntryIds = metadata.conversationHistoryEntryIds + } if (metadata.conversationHistoryDiverged !== undefined) { preserved.conversationHistoryDiverged = metadata.conversationHistoryDiverged } diff --git a/cli/src/agent/types.ts b/cli/src/agent/types.ts index 9c47ab62..65e849d4 100644 --- a/cli/src/agent/types.ts +++ b/cli/src/agent/types.ts @@ -29,7 +29,7 @@ export type PlanItem = { }; export type AgentMessage = - | { type: 'text'; text: string } + | { type: 'text'; text: string; id?: string; live?: boolean; streamSnapshot?: boolean } | { type: 'reasoning'; text: string; id?: string; live?: boolean } | { type: 'tool_call'; @@ -39,6 +39,7 @@ export type AgentMessage = status: 'pending' | 'in_progress' | 'completed' | 'failed'; title?: string; kind?: string; + progress?: unknown; } | { type: 'tool_result'; id: string; output: unknown; status: 'completed' | 'failed' } | { diff --git a/cli/src/api/types.ts b/cli/src/api/types.ts index 9e19cb9c..4a5af91e 100644 --- a/cli/src/api/types.ts +++ b/cli/src/api/types.ts @@ -71,6 +71,9 @@ export type { export const MessageMetaSchema = z.object({ sentFrom: z.string().optional(), + // Queue remains the default for existing clients. Pi-aware callers may + // explicitly request native steering while a turn is streaming. + deliveryMode: z.enum(['queue', 'steer']).optional(), fallbackModel: z.string().nullable().optional(), customSystemPrompt: z.string().nullable().optional(), appendSystemPrompt: z.string().nullable().optional(), diff --git a/cli/src/grok/conversationHistory.test.ts b/cli/src/grok/conversationHistory.test.ts index c04bc7c7..b7bf5d61 100644 --- a/cli/src/grok/conversationHistory.test.ts +++ b/cli/src/grok/conversationHistory.test.ts @@ -64,6 +64,8 @@ describe('GrokConversationHistory', () => { history.rememberPromptIndex('local-y', 1) const result = await history.rewind('local-y') expect(result.success).toBe(true) + expect(result.success).toBe(true) + if (!result.success) throw new Error(result.error) expect(result.truncateFromLocalId).toBe('local-y') }) diff --git a/cli/src/modules/common/attachmentFile.test.ts b/cli/src/modules/common/attachmentFile.test.ts new file mode 100644 index 00000000..8629a59d --- /dev/null +++ b/cli/src/modules/common/attachmentFile.test.ts @@ -0,0 +1,54 @@ +import { mkdtemp, mkdir, rm, truncate, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { afterEach, describe, expect, it } from 'vitest'; +import { MAX_UPLOAD_BYTES } from './attachmentLimits'; +import { readBoundedAttachmentFile } from './attachmentFile'; + +describe('readBoundedAttachmentFile', () => { + const directories: string[] = []; + + afterEach(async () => { + await Promise.all(directories.splice(0).map((directory) => rm(directory, { recursive: true, force: true }))); + }); + + async function createTemporaryDirectory(): Promise { + const directory = await mkdtemp(join(tmpdir(), 'hapi-attachment-file-')); + directories.push(directory); + return directory; + } + + it('reads a small regular file', async () => { + const directory = await createTemporaryDirectory(); + const path = join(directory, 'image.png'); + await writeFile(path, Buffer.from([0x89, 0x50, 0x4e, 0x47])); + + await expect(readBoundedAttachmentFile(path)).resolves.toEqual(Buffer.from([0x89, 0x50, 0x4e, 0x47])); + }); + + it('rejects a directory', async () => { + const directory = await createTemporaryDirectory(); + const path = join(directory, 'not-a-file'); + await mkdir(path); + + await expect(readBoundedAttachmentFile(path)).rejects.toThrow('Attachment must be a regular file'); + }); + + it('rejects an oversized sparse file before loading it into memory', async () => { + const directory = await createTemporaryDirectory(); + const path = join(directory, 'too-large-sparse.png'); + await writeFile(path, ''); + await truncate(path, MAX_UPLOAD_BYTES + 1); + + await expect(readBoundedAttachmentFile(path)).rejects.toThrow('Attachment file too large (max 50MB)'); + }); + + it('enforces a caller-provided remaining aggregate budget before reading', async () => { + const directory = await createTemporaryDirectory(); + const path = join(directory, 'second-image.png'); + await writeFile(path, Buffer.from([1, 2, 3, 4])); + + await expect(readBoundedAttachmentFile(path, 3)).rejects.toThrow('remaining image budget'); + await expect(readBoundedAttachmentFile(path, 0)).rejects.toThrow('remaining image budget'); + }); +}); diff --git a/cli/src/modules/common/attachmentFile.ts b/cli/src/modules/common/attachmentFile.ts new file mode 100644 index 00000000..8ef4e80b --- /dev/null +++ b/cli/src/modules/common/attachmentFile.ts @@ -0,0 +1,63 @@ +import { constants } from 'node:fs'; +import { open } from 'node:fs/promises'; +import { MAX_UPLOAD_BYTES } from './attachmentLimits'; +import type { UploadFileIdentity } from './handlers/uploads'; + +/** + * Reads a user attachment through one file handle with the upload byte limit. + * + * The caller is responsible for authorizing the attachment path for its + * session. This helper intentionally validates only that the opened target is + * a regular file and that its actual bytes fit the shared upload limit; it does + * not infer ownership from an arbitrary filesystem path. + */ +export async function readBoundedAttachmentFile( + path: string, + maxBytes: number = MAX_UPLOAD_BYTES, + authorizeOpenedFile?: (identity: UploadFileIdentity) => boolean, +): Promise { + const byteLimit = Math.min(MAX_UPLOAD_BYTES, Math.max(0, Math.floor(maxBytes))); + if (byteLimit === 0) { + throw new Error('Attachment exceeds the remaining image budget'); + } + // The caller passes the exact registered upload path. O_NOFOLLOW closes the + // final replacement window on POSIX if it is swapped to a symlink + // before open; Windows falls back to its normal read-only flag. + const flags = process.platform === 'win32' + ? 'r' + : constants.O_RDONLY | constants.O_NOFOLLOW; + const handle = await open(path, flags); + try { + // fstat and the subsequent read share this handle so a path replacement + // cannot change the object after its type/size check. + const stats = await handle.stat(); + if (authorizeOpenedFile && !authorizeOpenedFile(stats)) { + throw new Error('invalid upload path'); + } + if (!stats.isFile()) { + throw new Error('Attachment must be a regular file'); + } + if (stats.size > byteLimit) { + throw new Error(byteLimit === MAX_UPLOAD_BYTES + ? 'Attachment file too large (max 50MB)' + : 'Attachment exceeds the remaining image budget'); + } + + // Read into one bounded allocation. The extra sentinel byte detects a + // file that grows after fstat without retaining chunk buffers plus an + // equally large Buffer.concat copy at the 50 MiB boundary. + const buffer = Buffer.allocUnsafe(stats.size + 1); + let bytesRead = 0; + while (bytesRead < buffer.length) { + const result = await handle.read(buffer, bytesRead, buffer.length - bytesRead, bytesRead); + if (result.bytesRead === 0) break; + bytesRead += result.bytesRead; + } + if (bytesRead > stats.size) { + throw new Error('Attachment changed while it was being read'); + } + return buffer.subarray(0, bytesRead); + } finally { + await handle.close(); + } +} diff --git a/cli/src/modules/common/attachmentLimits.ts b/cli/src/modules/common/attachmentLimits.ts new file mode 100644 index 00000000..0be9b65d --- /dev/null +++ b/cli/src/modules/common/attachmentLimits.ts @@ -0,0 +1,2 @@ +/** Maximum decoded bytes accepted for one HAPI upload or one Pi image batch. */ +export const MAX_UPLOAD_BYTES = 50 * 1024 * 1024; diff --git a/cli/src/modules/common/handlers/uploads.test.ts b/cli/src/modules/common/handlers/uploads.test.ts new file mode 100644 index 00000000..68c0a91c --- /dev/null +++ b/cli/src/modules/common/handlers/uploads.test.ts @@ -0,0 +1,139 @@ +import { basename, join } from 'node:path'; +import { realpath, rename, rm, stat, writeFile } from 'node:fs/promises'; +import { describe, expect, it, vi } from 'vitest'; +import { getHapiBlobsDir } from '@/constants/uploadPaths'; +import { RPC_METHODS } from '@hapi/protocol/rpcMethods'; +import { readBoundedAttachmentFile } from '../attachmentFile'; +import { cleanupUploadDir, isAuthorizedUploadFile, isPathWithinUploadDir, registerUploadHandlers } from './uploads'; + +describe('isPathWithinUploadDir', () => { + it('accepts only paths under the matching session upload directory', () => { + const sessionId = 'session-allowed'; + const ownUpload = join(getHapiBlobsDir(), `${sessionId}-random`, 'image.png'); + const otherUpload = join(getHapiBlobsDir(), 'session-other-random', 'image.png'); + + expect(isPathWithinUploadDir(ownUpload, sessionId)).toBe(true); + expect(isPathWithinUploadDir(otherUpload, sessionId)).toBe(false); + expect(isPathWithinUploadDir('/etc/hosts', sessionId)).toBe(false); + }); + + it('binds authorization to the uploaded file identity, not only its path', async () => { + const handlers = new Map Promise>(); + registerUploadHandlers({ + registerHandler: (method: string, handler: (payload: any) => Promise) => handlers.set(method, handler), + } as never); + const sessionId = `session-identity-${Date.now()}`; + const upload = handlers.get(RPC_METHODS.UploadFile)!; + let path: string | null = null; + try { + const result = await upload({ + sessionId, + filename: 'image.png', + mimeType: 'image/png', + content: Buffer.from([1, 2, 3, 4]).toString('base64'), + }); + expect(result).toMatchObject({ success: true }); + path = result.path; + const original = await stat(path!); + expect(isAuthorizedUploadFile(path!, sessionId, original)).toBe(true); + expect(isAuthorizedUploadFile(path!, 'another-session', original)).toBe(false); + // The lexical upload path can canonicalize differently (notably + // /var -> /private/var on Darwin). Authorization and opened-file + // identity deliberately use the exact registered path. + expect(await readBoundedAttachmentFile( + path!, + 50 * 1024 * 1024, + (identity) => isAuthorizedUploadFile(path!, sessionId, identity), + )).toEqual(Buffer.from([1, 2, 3, 4])); + expect(typeof await realpath(path!)).toBe('string'); + + const replacementPath = `${path!}.replacement`; + await writeFile(replacementPath, Buffer.from([5, 6, 7, 8])); + const distinctReplacement = await stat(replacementPath); + expect(`${distinctReplacement.dev}:${distinctReplacement.ino}`).not.toBe(`${original.dev}:${original.ino}`); + await rm(path!, { force: true }); + await rename(replacementPath, path!); + const replacement = await stat(path!); + expect(isAuthorizedUploadFile(path!, sessionId, replacement)).toBe(false); + } finally { + await cleanupUploadDir(sessionId); + } + }); + + it('supports concurrent uploads with the same filename using distinct exclusive paths', async () => { + const handlers = new Map Promise>(); + registerUploadHandlers({ + registerHandler: (method: string, handler: (payload: any) => Promise) => handlers.set(method, handler), + } as never); + const sessionId = `session-concurrent-${Date.now()}`; + const upload = handlers.get(RPC_METHODS.UploadFile)!; + try { + const payload = { + sessionId, + filename: 'screenshot.png', + mimeType: 'image/png', + content: Buffer.from([1, 2, 3, 4]).toString('base64'), + }; + const [first, second] = await Promise.all([upload(payload), upload(payload)]); + + expect(first).toMatchObject({ success: true }); + expect(second).toMatchObject({ success: true }); + expect(first.path).not.toBe(second.path); + expect(isAuthorizedUploadFile(first.path, sessionId, await stat(first.path))).toBe(true); + expect(isAuthorizedUploadFile(second.path, sessionId, await stat(second.path))).toBe(true); + } finally { + await cleanupUploadDir(sessionId); + } + }); + + it('keeps long ASCII and Unicode filename components within 255 UTF-8 bytes', async () => { + const handlers = new Map Promise>(); + registerUploadHandlers({ + registerHandler: (method: string, handler: (payload: any) => Promise) => handlers.set(method, handler), + } as never); + const sessionId = `session-long-name-${Date.now()}`; + const upload = handlers.get(RPC_METHODS.UploadFile)!; + try { + for (const filename of [`${'a'.repeat(300)}.png`, `${'😀'.repeat(300)}.png`]) { + const result = await upload({ + sessionId, + filename, + mimeType: 'image/png', + content: Buffer.from([1]).toString('base64'), + }); + expect(result).toMatchObject({ success: true }); + expect(Buffer.byteLength(basename(result.path))).toBeLessThanOrEqual(255); + expect(basename(result.path)).toMatch(/\.png$/); + } + } finally { + await cleanupUploadDir(sessionId); + } + }); + + it('preserves an extension that exactly fills the remaining component budget', async () => { + vi.useFakeTimers(); + vi.setSystemTime(0); + const handlers = new Map Promise>(); + registerUploadHandlers({ + registerHandler: (method: string, handler: (payload: any) => Promise) => handlers.set(method, handler), + } as never); + const sessionId = 'session-exact-extension'; + const upload = handlers.get(RPC_METHODS.UploadFile)!; + // At time 0 the prefix is `0-<36-byte UUID>-` (39 bytes), leaving 216. + const extension = `.${'e'.repeat(215)}`; + try { + const result = await upload({ + sessionId, + filename: `x${extension}`, + mimeType: 'application/octet-stream', + content: Buffer.from([1]).toString('base64'), + }); + expect(result).toMatchObject({ success: true }); + expect(Buffer.byteLength(basename(result.path))).toBe(255); + expect(basename(result.path).endsWith(extension)).toBe(true); + } finally { + vi.useRealTimers(); + await cleanupUploadDir(sessionId); + } + }); +}); diff --git a/cli/src/modules/common/handlers/uploads.ts b/cli/src/modules/common/handlers/uploads.ts index d74fdc3c..17c3428d 100644 --- a/cli/src/modules/common/handlers/uploads.ts +++ b/cli/src/modules/common/handlers/uploads.ts @@ -1,12 +1,14 @@ import { logger } from '@/ui/logger' -import { mkdir, mkdtemp, rm, writeFile } from 'fs/promises' -import { join, resolve, sep } from 'path' +import { randomUUID } from 'node:crypto' +import { mkdir, mkdtemp, open, rm } from 'fs/promises' +import { extname, join, resolve, sep } from 'path' import { rmSync } from 'node:fs' import type { DeleteUploadResponse, UploadFileResponse } from '@hapi/protocol/apiTypes' import { RPC_METHODS } from '@hapi/protocol/rpcMethods' import type { RpcHandlerManager } from '@/api/rpc/RpcHandlerManager' import { getErrorMessage, rpcError } from '../rpcResponses' import { getHapiBlobsDir } from '@/constants/uploadPaths' +import { MAX_UPLOAD_BYTES } from '../attachmentLimits' interface UploadFileRequest { sessionId?: string @@ -21,23 +23,56 @@ interface DeleteUploadRequest { } const uploadDirs = new Map() +const uploadFileIdentities = new Map>() const uploadDirPromises = new Map>() const uploadDirCleanupRequested = new Set() let cleanupRegistered = false -const MAX_UPLOAD_BYTES = 50 * 1024 * 1024 +const MAX_FILENAME_COMPONENT_BYTES = 255 + +export type UploadFileIdentity = { dev: number; ino: number } + +function uploadIdentityKey(identity: UploadFileIdentity): string { + return `${identity.dev}:${identity.ino}` +} + +export function isAuthorizedUploadFile(path: string, sessionId: string | undefined, identity: UploadFileIdentity): boolean { + return uploadFileIdentities.get(getSessionKey(sessionId))?.get(resolve(path)) === uploadIdentityKey(identity) +} function sanitizeFilename(filename: string): string { - // Remove path separators and limit length + // Remove path separators; byte-safe length fitting happens after the + // unique prefix is known so the original extension can be preserved. const sanitized = filename .replace(/[/\\]/g, '_') .replace(/\.\./g, '_') .replace(/\s+/g, '_') - .slice(0, 255) // If filename is empty after sanitization, use a default return sanitized || 'upload' } +function truncateUtf8(text: string, maxBytes: number): string { + let result = '' + let bytes = 0 + for (const character of text) { + const characterBytes = Buffer.byteLength(character) + if (bytes + characterBytes > maxBytes) break + result += character + bytes += characterBytes + } + return result +} + +function fitFilenameToBytes(filename: string, maxBytes: number): string { + const extension = extname(filename) + const extensionBytes = Buffer.byteLength(extension) + if (extension && extensionBytes <= maxBytes) { + const stem = filename.slice(0, -extension.length) + return `${truncateUtf8(stem, maxBytes - extensionBytes)}${extension}` + } + return truncateUtf8(filename, maxBytes) || 'upload' +} + function getSessionKey(sessionId?: string): string { const trimmed = sessionId?.trim() return trimmed ? trimmed : 'unknown' @@ -62,7 +97,7 @@ async function getOrCreateUploadDir(sessionId?: string): Promise { return await inflight } - const safeKey = sanitizeFilename(sessionKey) + const safeKey = fitFilenameToBytes(sanitizeFilename(sessionKey), 120) const creation = (async () => { try { const blobsDir = getHapiBlobsDir() @@ -102,6 +137,7 @@ export async function cleanupUploadDir(sessionId?: string): Promise { const dir = uploadDirs.get(sessionKey) uploadDirs.delete(sessionKey) + uploadFileIdentities.delete(sessionKey) uploadDirPromises.delete(sessionKey) if (!dir) { @@ -121,6 +157,7 @@ export async function cleanupUploadDir(sessionId?: string): Promise { function cleanupUploadDirsSync(): void { const dirs = Array.from(uploadDirs.values()) uploadDirs.clear() + uploadFileIdentities.clear() uploadDirPromises.clear() uploadDirCleanupRequested.clear() @@ -133,7 +170,7 @@ function cleanupUploadDirsSync(): void { } } -function isPathWithinUploadDir(path: string, sessionId?: string): boolean { +export function isPathWithinUploadDir(path: string, sessionId?: string): boolean { const sessionKey = getSessionKey(sessionId) const resolvedPath = resolve(path) const activeDir = uploadDirs.get(sessionKey) @@ -143,7 +180,7 @@ function isPathWithinUploadDir(path: string, sessionId?: string): boolean { return resolvedPath.startsWith(dirPrefix) } - const safeKey = sanitizeFilename(sessionKey) + const safeKey = fitFilenameToBytes(sanitizeFilename(sessionKey), 120) const resolvedPrefix = resolve(getHapiBlobsDir(), `${safeKey}-`) return resolvedPath.startsWith(resolvedPrefix) } @@ -174,9 +211,15 @@ export function registerUploadHandlers(rpcHandlerManager: RpcHandlerManager): vo const dir = await getOrCreateUploadDir(data.sessionId) const sanitizedFilename = sanitizeFilename(data.filename) - // Add timestamp to avoid collisions + // Combine a readable timestamp with a random suffix so concurrent + // uploads of the same filename remain compatible with exclusive create. const timestamp = Date.now() - const uniqueFilename = `${timestamp}-${sanitizedFilename}` + const prefix = `${timestamp}-${randomUUID()}-` + const boundedFilename = fitFilenameToBytes( + sanitizedFilename, + MAX_FILENAME_COMPONENT_BYTES - Buffer.byteLength(prefix), + ) + const uniqueFilename = `${prefix}${boundedFilename}` const filePath = join(dir, uniqueFilename) // Decode base64 content and write to file @@ -184,7 +227,28 @@ export function registerUploadHandlers(rpcHandlerManager: RpcHandlerManager): vo if (buffer.length > MAX_UPLOAD_BYTES) { return rpcError('File too large (max 50MB)') } - await writeFile(filePath, buffer) + const file = await open(filePath, 'wx', 0o600) + let operationError: unknown + try { + await file.writeFile(buffer) + const info = await file.stat() + const files = uploadFileIdentities.get(getSessionKey(data.sessionId)) ?? new Map() + files.set(resolve(filePath), uploadIdentityKey(info)) + uploadFileIdentities.set(getSessionKey(data.sessionId), files) + } catch (error) { + operationError = error + } + let closeError: unknown + try { + await file.close() + } catch (error) { + closeError = error + } + if (operationError || closeError) { + uploadFileIdentities.get(getSessionKey(data.sessionId))?.delete(resolve(filePath)) + await rm(filePath, { force: true }).catch(() => {}) + throw operationError ?? closeError + } logger.debug('File uploaded successfully:', filePath) return { success: true, path: filePath } @@ -206,6 +270,7 @@ export function registerUploadHandlers(rpcHandlerManager: RpcHandlerManager): vo try { await rm(path, { force: true }) + uploadFileIdentities.get(getSessionKey(data.sessionId))?.delete(resolve(path)) return { success: true } } catch (error) { logger.debug('Failed to delete upload file:', error) diff --git a/cli/src/pi/conversationHistory.test.ts b/cli/src/pi/conversationHistory.test.ts new file mode 100644 index 00000000..c9d1388c --- /dev/null +++ b/cli/src/pi/conversationHistory.test.ts @@ -0,0 +1,694 @@ +import { describe, expect, it, vi } from 'vitest' +import { PI_HISTORY_OPERATION_TIMEOUT_MS, PiConversationHistory, PiHistoryRestoreError } from './conversationHistory' +import { PiSession } from './session' + +function createSession(options?: { nativeReady?: boolean }) { + const metadata: Record = {} + const client = { + keepAlive: vi.fn(), + flushMetadata: vi.fn(async () => true), + updateMetadata: vi.fn((updater: (current: Record) => Record) => { + Object.assign(metadata, updater(metadata)) + }), + sendAgentMessage: vi.fn(), + emitMessagesConsumed: vi.fn(), + sendSessionEvent: vi.fn(), + emitSessionReady: vi.fn(), + } + const session = new PiSession({ + api: {} as never, + client: client as never, + path: '/tmp/project', + logPath: '/tmp/pi.log', + startedBy: 'terminal', + startingMode: 'remote', + }) + if (options?.nativeReady !== false) session.markNativeReady() + return { + metadata, + client, + session, + } +} + +const source = { sessionId: 'source-id', sessionFile: '/tmp/source.jsonl' } +const clone = { sessionId: 'clone-id', sessionFile: '/tmp/clone.jsonl' } + +describe('PiConversationHistory entry mapping', () => { + it('pairs duplicate prompts to native user entries strictly FIFO without reading text', () => { + const { session, metadata } = createSession() + const history = new PiConversationHistory(session, vi.fn()) + history.registerUserEntry('local-1') + history.registerUserEntry('local-2') + + history.observeEntry({ type: 'message', id: 'entry-1', message: { role: 'user', content: 'same' } }) + history.observeEntry({ type: 'message', id: 'assistant-1', message: { role: 'assistant', content: 'same' } }) + // Pi can forward an entry_appended event and later return the same entry + // from get_entries. It must not consume local-2 twice. + history.observeEntry({ type: 'message', id: 'entry-1', message: { role: 'user', content: 'same' } }) + history.observeEntry({ type: 'message', id: 'entry-2', message: { role: 'user', content: 'same' } }) + + expect(history.getEntryIds()).toEqual({ 'local-1': 'entry-1', 'local-2': 'entry-2' }) + expect(metadata).toMatchObject({ + conversationHistoryPoints: { 'local-1': true, 'local-2': true }, + conversationHistoryEntryIds: { 'local-1': 'entry-1', 'local-2': 'entry-2' }, + }) + }) + + it('uses the append cursor for an entry event fallback', async () => { + const { session } = createSession() + const rpc = vi.fn(async (command: Record, _timeoutMs?: number) => { + if (rpc.mock.calls.length === 1) { + expect(command).toEqual({ type: 'get_entries' }) + return { + entries: [{ type: 'message', id: 'native-1', message: { role: 'user' } }], + leafId: 'branch-leaf-that-moved-backward', + } + } + expect(command).toEqual({ type: 'get_entries', since: 'native-1' }) + return { entries: [], leafId: 'older-active-leaf' } + }) + const history = new PiConversationHistory(session, rpc) + history.registerUserEntry('local-1') + await history.syncEntries() + expect(history.getEntryIds()).toEqual({ 'local-1': 'native-1' }) + await history.syncEntries() + }) + + it('serializes concurrent syncs and ignores a duplicate entry_appended/get_entries user entry', async () => { + const { session } = createSession() + let resolveFirstRead!: (value: unknown) => void + let activeReads = 0 + let maxActiveReads = 0 + const rpc = vi.fn((command: Record) => { + activeReads += 1 + maxActiveReads = Math.max(maxActiveReads, activeReads) + if (rpc.mock.calls.length === 1) { + return new Promise((resolve) => { + resolveFirstRead = (value) => { + activeReads -= 1 + resolve(value) + } + }) + } + expect(command).toEqual({ type: 'get_entries', since: 'entry-1' }) + activeReads -= 1 + return Promise.resolve({ entries: [{ type: 'message', id: 'entry-1', message: { role: 'user' } }], leafId: 'entry-1' }) + }) + const history = new PiConversationHistory(session, rpc) + history.registerUserEntry('local-1') + history.registerUserEntry('local-2') + + const first = history.syncEntries() + history.observeEntry({ type: 'message', id: 'entry-1', message: { role: 'user' } }) + const concurrent = history.syncEntries() + resolveFirstRead({ entries: [{ type: 'message', id: 'entry-1', message: { role: 'user' } }], leafId: 'entry-1' }) + await Promise.all([first, concurrent]) + history.observeEntry({ type: 'message', id: 'entry-2', message: { role: 'user' } }) + + expect(maxActiveReads).toBe(1) + expect(history.getEntryIds()).toEqual({ 'local-1': 'entry-1', 'local-2': 'entry-2' }) + }) + + it('disables optional history polling when startup get_entries is unsupported', async () => { + const { session } = createSession() + const rpc = vi.fn(async () => { throw new Error('Unknown command: get_entries') }) + const publishCapabilities = vi.fn(async () => {}) + const history = new PiConversationHistory(session, rpc) + history.setPublishCapabilities(publishCapabilities) + + await history.initialize() + await history.syncEntries() + await history.syncEntries() + history.registerUserEntry('disabled-1') + history.registerUserEntry('disabled-2') + history.observeEntry({ type: 'message', id: 'late-user', message: { role: 'user' } }) + + expect(rpc).toHaveBeenCalledTimes(1) + expect(history.getCapabilitiesForMetadata()).toBeUndefined() + expect(history.getEntryIds()).toEqual({}) + expect(publishCapabilities).toHaveBeenCalledTimes(1) + }) + + it.each([ + ['transient rejection', async () => { throw new Error('temporary read failure') }], + ['malformed response', async () => ({ entries: null, leafId: null })], + ])('disables history when the startup baseline has a %s', async (_label, read) => { + const { session } = createSession() + const rpc = vi.fn(read) + const history = new PiConversationHistory(session, rpc) + + await history.initialize() + history.registerUserEntry('new-local') + history.observeEntry({ type: 'message', id: 'old-native-user', message: { role: 'user' } }) + await history.syncEntries() + + expect(rpc).toHaveBeenCalledTimes(1) + expect(history.getCapabilitiesForMetadata()).toBeUndefined() + expect(history.getEntryIds()).toEqual({}) + await expect(history.fork()).rejects.toThrow('unavailable') + await expect(history.rewind('stale-local')).resolves.toMatchObject({ + success: false, + error: expect.stringContaining('unavailable'), + }) + expect(rpc).toHaveBeenCalledTimes(1) + }) + + it('handles rejection from a detached coalesced sync retry', async () => { + const { session } = createSession() + const rpc = vi.fn(async () => { throw new Error('transient history read failure') }) + const history = new PiConversationHistory(session, rpc) + + const first = history.syncEntries().catch(() => {}) + const concurrent = history.syncEntries().catch(() => {}) + await Promise.all([first, concurrent]) + await vi.waitFor(() => expect(rpc).toHaveBeenCalledTimes(2)) + await Promise.resolve() + }) + + it('removes a failed local FIFO prompt by exact localId without consuming its neighbor', () => { + const { session } = createSession() + const history = new PiConversationHistory(session, vi.fn()) + history.registerUserEntry('prompt-failed') + history.registerUserEntry('prompt-ok') + history.rejectPendingEntry('prompt-failed') + history.observeEntry({ type: 'message', id: 'prompt-entry', message: { role: 'user' } }) + expect(history.getEntryIds()).toEqual({ 'prompt-ok': 'prompt-entry' }) + + history.registerUserEntry('aborted-before-turn') + history.rejectPendingEntry('aborted-before-turn') + history.observeEntry({ type: 'message', id: 'unrelated-user-entry', message: { role: 'user' } }) + expect(history.getEntryIds()).toEqual({ 'prompt-ok': 'prompt-entry' }) + }) +}) + +describe('PiConversationHistory native transactions', () => { + it('rejects before native fork when final source locator metadata does not flush', async () => { + const { session, client } = createSession() + client.flushMetadata.mockResolvedValue(false) + const rpc = vi.fn(async (command: Record, _timeoutMs?: number) => { + if (command.type === 'get_entries') return { entries: [], leafId: null } + throw new Error(`native fork must not run: ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + + await expect(history.fork()).rejects.toThrow('metadata did not persist') + expect(rpc.mock.calls.map(([command]) => command.type)).toEqual(['get_entries']) + }) + + it('forks current by clone then restores the exact source identity', async () => { + const { session } = createSession() + let stateCalls = 0 + const rpc = vi.fn(async (command: Record, _timeoutMs?: number) => { + if (command.type === 'get_entries') return { entries: [], leafId: null } + if (command.type === 'get_state') return [source, clone, clone, source][stateCalls++] + if (command.type === 'clone') return { cancelled: false } + if (command.type === 'switch_session') return { cancelled: false } + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + + await expect(history.fork()).resolves.toEqual({ nativeSessionId: 'clone-id' }) + expect(rpc.mock.calls.map(([command]) => command)).toEqual([ + { type: 'get_entries' }, { type: 'get_state' }, { type: 'clone' }, { type: 'get_state' }, + { type: 'get_state' }, { type: 'switch_session', sessionPath: source.sessionFile }, { type: 'get_state' }, + ]) + const mutationTimeouts = rpc.mock.calls + .filter(([command]) => command.type === 'clone' || command.type === 'switch_session') + .map(([, timeoutMs]) => timeoutMs) + expect(mutationTimeouts).toHaveLength(2) + for (const timeoutMs of mutationTimeouts) { + expect(timeoutMs).toBeGreaterThan(0) + expect(timeoutMs).toBeLessThanOrEqual(PI_HISTORY_OPERATION_TIMEOUT_MS) + } + expect(session.isHistoryTransactionActive).toBe(false) + }) + + it('uses one absolute transaction deadline across clone and source restoration', async () => { + vi.useFakeTimers() + vi.setSystemTime(0) + try { + const { session } = createSession() + let stateCalls = 0 + const mutationTimeouts: number[] = [] + const rpc = vi.fn(async (command: Record, timeoutMs?: number) => { + if (command.type === 'get_entries') return { entries: [], leafId: null } + if (command.type === 'get_state') return [source, clone, clone, source][stateCalls++] + if (command.type === 'clone') { + mutationTimeouts.push(timeoutMs ?? 0) + vi.setSystemTime(60_000) + return { cancelled: false } + } + if (command.type === 'switch_session') { + mutationTimeouts.push(timeoutMs ?? 0) + return { cancelled: false } + } + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + + await expect(history.fork()).resolves.toEqual({ nativeSessionId: 'clone-id' }) + expect(mutationTimeouts).toHaveLength(2) + expect(mutationTimeouts[0]).toBeLessThanOrEqual(PI_HISTORY_OPERATION_TIMEOUT_MS - 10_000) + expect(mutationTimeouts[1]).toBeLessThanOrEqual(PI_HISTORY_OPERATION_TIMEOUT_MS - 60_000) + } finally { + vi.useRealTimers() + } + }) + + it('treats a pre-mutation read timeout as an ordinary rejection and drains queued work', async () => { + const { session } = createSession() + const timeout = Object.assign(new Error('Pi RPC get_state (id=1) timed out after 10000ms'), { + name: 'PiRpcTimeoutError', + }) + const rpc = vi.fn(async (command: Record) => { + if (command.type === 'get_entries') return { entries: [], leafId: null } + if (command.type === 'get_state') throw timeout + throw new Error(`mutation must not run: ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + let delivered = false + + const pending = history.fork() + session.runWhenHistoryIdle(() => { delivered = true }, 'queued-after-read-timeout') + + await expect(pending).rejects.toBe(timeout) + expect(delivered).toBe(true) + expect(session.isHistoryTransactionActive).toBe(false) + expect(rpc.mock.calls.map(([command]) => command.type)).toEqual(['get_entries', 'get_state']) + }) + + it('forks a historical boundary from source and restores source afterward', async () => { + const { session } = createSession() + let stateCalls = 0 + const rpc = vi.fn(async (command: Record, _timeoutMs?: number) => { + if (command.type === 'get_entries') return { entries: [], leafId: null } + if (command.type === 'get_state') return [source, clone, clone, source][stateCalls++] + if (command.type === 'fork' || command.type === 'switch_session') return { cancelled: false } + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + history.restoreEntryIds({ local: 'entry-user' }) + + await expect(history.fork('local')).resolves.toEqual({ nativeSessionId: 'clone-id' }) + expect(rpc.mock.calls.map(([command]) => command)).toEqual([ + { type: 'get_entries' }, { type: 'get_state' }, { type: 'fork', entryId: 'entry-user' }, { type: 'get_state' }, + { type: 'get_state' }, { type: 'switch_session', sessionPath: source.sessionFile }, { type: 'get_state' }, + ]) + const mutationTimeouts = rpc.mock.calls + .filter(([command]) => command.type === 'fork' || command.type === 'switch_session') + .map(([, timeoutMs]) => timeoutMs) + expect(mutationTimeouts).toHaveLength(2) + for (const timeoutMs of mutationTimeouts) { + expect(timeoutMs).toBeGreaterThan(0) + expect(timeoutMs).toBeLessThanOrEqual(PI_HISTORY_OPERATION_TIMEOUT_MS) + } + }) + + it('commits the rewound branch identity, resets its cursor, and maps the next prompt', async () => { + const { session, metadata } = createSession() + const sourceState = { + ...source, + model: { id: 'source-model', provider: 'source-provider' }, + thinkingLevel: 'low', + steeringMode: 'all', + isStreaming: false, + } + const rewound = { + sessionId: 'rewind-id', + sessionFile: '/tmp/rewind.jsonl', + model: { id: 'rewind-model', provider: 'rewind-provider' }, + thinkingLevel: 'high', + steeringMode: 'one-at-a-time', + isStreaming: true, + } + let entriesCalls = 0 + let stateCalls = 0 + const rpc = vi.fn(async (command: Record, _timeoutMs?: number) => { + if (command.type === 'get_entries') { + entriesCalls += 1 + return entriesCalls === 1 + ? { entries: [], leafId: null } + : { entries: [{ id: 'entry-before-user', type: 'message', message: { role: 'assistant' } }], leafId: 'entry-before-user' } + } + if (command.type === 'get_state') return [sourceState, rewound][stateCalls++] + if (command.type === 'get_fork_messages') return { messages: [{ entryId: 'entry-user', text: 'ignored' }] } + if (command.type === 'fork') return { cancelled: false } + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + history.restoreEntryIds({ local: 'entry-user' }) + + await expect(history.rewind('local')).resolves.toEqual({ + success: true, + truncateFromLocalId: 'local', + messages: [], + }) + expect(session.expectedNativeSessionId).toBe('rewind-id') + expect(session.currentNativeSessionFile).toBe('/tmp/rewind.jsonl') + expect(session.currentModel).toBe('rewind-model') + expect(session.currentProvider).toBe('rewind-provider') + expect(session.currentThinkingLevel).toBe('high') + expect(session.currentSteeringMode).toBe('one-at-a-time') + expect(session.piIsStreaming).toBe(true) + expect(metadata).toMatchObject({ + piSessionId: 'rewind-id', + piSelectedModel: { provider: 'rewind-provider', modelId: 'rewind-model' }, + }) + history.registerUserEntry('next-local') + history.observeEntry({ type: 'message', id: 'next-entry', message: { role: 'user' } }) + expect(history.getEntryIds()).toEqual({ 'next-local': 'next-entry' }) + expect(rpc.mock.calls.map(([command]) => command)).toEqual([ + { type: 'get_entries' }, { type: 'get_state' }, { type: 'get_fork_messages' }, { type: 'fork', entryId: 'entry-user' }, + { type: 'get_state' }, { type: 'get_entries' }, + ]) + expect(rpc.mock.calls.find(([command]) => command.type === 'fork')?.[1]) + .toBeLessThanOrEqual(PI_HISTORY_OPERATION_TIMEOUT_MS) + }) + + it('does not combine a changed rewind model with a provider omitted by Pi', async () => { + const { session, metadata } = createSession() + session.currentModel = 'old-model' + session.currentProvider = 'old-provider' + const sourceState = { ...source, model: { id: 'old-model', provider: 'old-provider' }, isStreaming: false } + const rewound = { sessionId: 'rewind-id', sessionFile: '/tmp/rewind.jsonl', model: { id: 'branch-model' }, isStreaming: false } + let stateCalls = 0 + let entriesCalls = 0 + const rpc = vi.fn(async (command: Record) => { + if (command.type === 'get_entries') { + entriesCalls += 1 + return entriesCalls === 1 + ? { entries: [], leafId: null } + : { entries: [{ id: 'prefix', type: 'message', message: { role: 'assistant' } }], leafId: 'prefix' } + } + if (command.type === 'get_state') return [sourceState, rewound][stateCalls++] + if (command.type === 'get_fork_messages') return { messages: [{ entryId: 'entry-user', text: 'user' }] } + if (command.type === 'fork') return { cancelled: false } + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + history.restoreEntryIds({ local: 'entry-user' }) + + await expect(history.rewind('local')).resolves.toMatchObject({ success: true }) + expect(session.currentModel).toBe('branch-model') + expect(session.currentProvider).toBeNull() + expect(metadata.piSelectedModel).toBeUndefined() + }) + + it('restores source and rolls back identity/locators when rewind metadata flush fails', async () => { + const { session, client } = createSession() + client.flushMetadata.mockResolvedValueOnce(true).mockResolvedValueOnce(false).mockResolvedValueOnce(true) + const sourceState = { + ...source, + model: { id: 'source-model', provider: 'source-provider' }, + thinkingLevel: 'medium', + steeringMode: 'all', + isStreaming: false, + } + const rewound = { + sessionId: 'rewind-id', + sessionFile: '/tmp/rewind.jsonl', + model: { id: 'rewind-model', provider: 'rewind-provider' }, + thinkingLevel: 'high', + steeringMode: 'one-at-a-time', + isStreaming: true, + } + let stateCalls = 0 + let entryCalls = 0 + const rpc = vi.fn(async (command: Record) => { + if (command.type === 'get_entries') { + entryCalls += 1 + return entryCalls === 1 + ? { entries: [], leafId: null } + : { entries: [{ id: 'branch-prefix', type: 'message', message: { role: 'assistant' } }], leafId: 'branch-prefix' } + } + if (command.type === 'get_state') return [sourceState, rewound, rewound, sourceState][stateCalls++] + if (command.type === 'get_fork_messages') return { messages: [{ entryId: 'entry-user', text: 'user' }] } + if (command.type === 'fork' || command.type === 'switch_session') return { cancelled: false } + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + history.restoreEntryIds({ local: 'entry-user' }) + + await expect(history.rewind('local')).resolves.toEqual({ + success: false, + error: 'Pi rewind metadata did not persist', + outcome: 'source_restored', + }) + expect(session.expectedNativeSessionId).toBe(source.sessionId) + expect(session.currentModel).toBe('source-model') + expect(session.currentProvider).toBe('source-provider') + expect(session.currentThinkingLevel).toBe('medium') + expect(session.currentSteeringMode).toBe('all') + expect(session.piIsStreaming).toBe(false) + expect(history.getEntryIds()).toEqual({ local: 'entry-user' }) + }) + + it('releases the history gate when rollback persistence reaches the transaction deadline', async () => { + vi.useFakeTimers() + vi.setSystemTime(0) + try { + const { session, client } = createSession() + client.flushMetadata.mockResolvedValueOnce(true).mockResolvedValueOnce(false) + const sourceState = { ...source, model: { id: 'source-model', provider: 'source-provider' }, isStreaming: false } + const rewound = { sessionId: 'rewind-id', sessionFile: '/tmp/rewind.jsonl', model: { id: 'rewind-model', provider: 'rewind-provider' }, isStreaming: false } + let stateCalls = 0 + let entryCalls = 0 + const rpc = vi.fn(async (command: Record) => { + if (command.type === 'get_entries') { + entryCalls += 1 + return entryCalls === 1 + ? { entries: [], leafId: null } + : { entries: [{ id: 'branch-prefix', type: 'message', message: { role: 'assistant' } }], leafId: 'branch-prefix' } + } + if (command.type === 'get_state') { + const state = [sourceState, rewound, rewound, sourceState][stateCalls++] + if (stateCalls === 4) vi.setSystemTime(PI_HISTORY_OPERATION_TIMEOUT_MS + 1) + return state + } + if (command.type === 'get_fork_messages') return { messages: [{ entryId: 'entry-user', text: 'user' }] } + if (command.type === 'fork' || command.type === 'switch_session') return { cancelled: false } + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + history.restoreEntryIds({ local: 'entry-user' }) + let deferredDelivered = false + + const pending = history.rewind('local') + session.runWhenHistoryIdle(() => { deferredDelivered = true }, 'queued-during-failed-rollback') + + await expect(pending).rejects.toBeInstanceOf(PiHistoryRestoreError) + expect(session.isHistoryTransactionActive).toBe(false) + expect(deferredDelivered).toBe(false) + await expect(session.runRuntimeMutation(async () => 'released')).resolves.toBe('released') + } finally { + vi.useRealTimers() + } + }) + + it('waits for a delayed source sync before rewind, then maps the next branch prompt', async () => { + const { session } = createSession() + let resolveStaleRead!: (data: unknown) => void + const rewound = { sessionId: 'rewind-id', sessionFile: '/tmp/rewind.jsonl' } + let getEntriesCalls = 0 + const rpc = vi.fn((command: Record) => { + if (command.type === 'get_entries') { + getEntriesCalls += 1 + if (getEntriesCalls === 1) { + return new Promise((resolve) => { resolveStaleRead = resolve }) + } + return Promise.resolve({ entries: [{ id: 'new-prefix', type: 'message', message: { role: 'assistant' } }], leafId: 'new-prefix' }) + } + if (command.type === 'get_state') { + const states = rpc.mock.calls.filter(([item]) => item.type === 'get_state').length + return Promise.resolve(states === 1 ? source : rewound) + } + if (command.type === 'get_fork_messages') return Promise.resolve({ messages: [{ entryId: 'old-user', text: 'old' }] }) + if (command.type === 'fork') return Promise.resolve({ cancelled: false }) + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + history.restoreEntryIds({ old: 'old-user' }) + const staleSync = history.syncEntries() + const rewindPending = history.rewind('old') + await Promise.resolve() + expect(rpc.mock.calls.some(([command]) => command.type === 'fork')).toBe(false) + resolveStaleRead({ entries: [{ id: 'stale-user', type: 'message', message: { role: 'user' } }], leafId: 'stale-user' }) + await staleSync + const rewind = await rewindPending + expect(rewind.success).toBe(true) + + history.registerUserEntry('next') + history.observeEntry({ type: 'message', id: 'next-user', message: { role: 'user' } }) + expect(history.getEntryIds()).toEqual({ next: 'next-user' }) + }) + + it('refuses history actions before native-ready or while Pi is streaming/prompting', async () => { + const unready = createSession({ nativeReady: false }) + const unreadyHistory = new PiConversationHistory(unready.session, vi.fn()) + await expect(unreadyHistory.fork()).rejects.toThrow('not ready') + + const busy = createSession() + busy.session.piIsStreaming = true + const busyHistory = new PiConversationHistory(busy.session, vi.fn()) + await expect(busyHistory.fork()).rejects.toThrow('busy') + + const preflight = createSession() + const preflightRpc = vi.fn() + const preflightHistory = new PiConversationHistory(preflight.session, preflightRpc) + preflightHistory.registerUserEntry('local-preflight') + preflightHistory.observeEntry({ type: 'message', id: 'native-preflight', message: { role: 'user' } }) + preflight.session.setPromptInFlight(true) + await expect(preflightHistory.fork()).rejects.toThrow('busy') + expect(preflightRpc).not.toHaveBeenCalled() + }) + + it('returns deterministic failure after restoring source instead of throwing/diverging', async () => { + const { session } = createSession() + let stateCalls = 0 + const rpc = vi.fn(async (command: Record) => { + if (command.type === 'get_entries') return { entries: [], leafId: null } + if (command.type === 'get_state') return [source, source][stateCalls++] + if (command.type === 'get_fork_messages') return { messages: [] } + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + history.restoreEntryIds({ local: 'entry-user' }) + + await expect(history.rewind('local')).resolves.toEqual({ + success: false, + error: 'Pi rewind point is no longer available', + outcome: 'source_restored', + }) + expect(rpc.mock.calls.some(([command]) => command.type === 'switch_session')).toBe(false) + }) + + it('returns cancelled without a redundant source switch when Pi fork never leaves source', async () => { + const { session } = createSession() + let stateCalls = 0 + const rpc = vi.fn(async (command: Record) => { + if (command.type === 'get_entries') return { entries: [], leafId: null } + if (command.type === 'get_state') return [source, source][stateCalls++] + if (command.type === 'get_fork_messages') return { messages: [{ entryId: 'entry-user', text: 'user' }] } + if (command.type === 'fork') return { cancelled: true } + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + history.restoreEntryIds({ local: 'entry-user' }) + + await expect(history.rewind('local')).resolves.toEqual({ + success: false, + error: 'Pi rewind was cancelled', + outcome: 'cancelled', + }) + expect(rpc.mock.calls.some(([command]) => command.type === 'switch_session')).toBe(false) + }) + + it('fails closed when source restoration fails', async () => { + const { session } = createSession() + let resolveClone!: (value: unknown) => void + let stateCalls = 0 + const rpc = vi.fn((command: Record) => { + if (command.type === 'get_entries') return Promise.resolve({ entries: [], leafId: null }) + if (command.type === 'get_state') { + return Promise.resolve([source, clone, clone][stateCalls++]) + } + if (command.type === 'clone') { + return new Promise((resolve) => { resolveClone = resolve }) + } + if (command.type === 'switch_session') return Promise.resolve({ cancelled: true }) + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + let delivered = false + const pending = history.fork() + await vi.waitFor(() => expect(resolveClone).toBeTypeOf('function')) + session.runWhenHistoryIdle(() => { delivered = true }, 'queued-during-restore') + resolveClone({ cancelled: false }) + + await expect(pending).rejects.toBeInstanceOf(PiHistoryRestoreError) + expect(delivered).toBe(false) + expect(session.isHistoryTransactionActive).toBe(false) + }) + + it('treats a late clone timeout as indeterminate, discards the history queue, and never classifies state', async () => { + const { session } = createSession() + let cloneTimedOut = false + let rejectClone!: (error: Error) => void + let cloneTimeoutMs = 0 + const rpc = vi.fn(async (command: Record, timeoutMs?: number) => { + if (command.type === 'get_entries') return { entries: [], leafId: null } + if (command.type === 'get_state') { + if (cloneTimedOut) throw new Error('must not read state after an indeterminate clone timeout') + return source + } + if (command.type === 'clone') { + cloneTimeoutMs = timeoutMs ?? 0 + expect(cloneTimeoutMs).toBeGreaterThan(0) + expect(cloneTimeoutMs).toBeLessThan(PI_HISTORY_OPERATION_TIMEOUT_MS) + return new Promise((_, reject) => { + rejectClone = (error) => { + cloneTimedOut = true + reject(error) + } + }) + } + if (command.type === 'switch_session') throw new Error('must not switch after an indeterminate clone timeout') + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + let delivered = false + + const pending = history.fork() + await vi.waitFor(() => expect(rejectClone).toBeTypeOf('function')) + session.runWhenHistoryIdle(() => { delivered = true }, 'queued-after-clone-timeout') + rejectClone(new Error(`Pi RPC clone (id=42) timed out after ${cloneTimeoutMs}ms`)) + + await expect(pending).rejects.toBeInstanceOf(PiHistoryRestoreError) + expect(rpc.mock.calls.map(([command]) => command.type)).toEqual(['get_entries', 'get_state', 'clone']) + expect(delivered).toBe(false) + expect(session.isHistoryTransactionActive).toBe(false) + }) + + it('closes the history gate before waiting for an in-flight config mutation', async () => { + const { session } = createSession() + const releaseConfigMutation = await session.acquireRuntimeMutation() + let stateCalls = 0 + const rpc = vi.fn(async (command: Record) => { + if (command.type === 'get_entries') return { entries: [], leafId: null } + if (command.type === 'get_state') return [source, clone, clone, source][stateCalls++] + if (command.type === 'clone' || command.type === 'switch_session') return { cancelled: false } + throw new Error(`unexpected ${command.type}`) + }) + const history = new PiConversationHistory(session, rpc) + let deferredDelivered = false + + const pending = history.fork() + session.runWhenHistoryIdle(() => { deferredDelivered = true }, 'prompt-after-history-begin') + + expect(session.isHistoryTransactionActive).toBe(true) + expect(rpc).not.toHaveBeenCalled() + releaseConfigMutation() + + await expect(pending).resolves.toEqual({ nativeSessionId: 'clone-id' }) + expect(deferredDelivered).toBe(true) + }) + + it('keeps history operations mutually exclusive and revokes a command that Pi rejects as unknown', async () => { + const { session } = createSession() + let rejectClone!: (reason: Error) => void + const rpc = vi.fn((command: Record) => { + if (command.type === 'get_entries') return Promise.resolve({ entries: [], leafId: null }) + if (command.type === 'get_state') return Promise.resolve(source) + if (command.type === 'clone') return new Promise((_, reject) => { rejectClone = reject }) + return Promise.resolve({ cancelled: false }) + }) + const history = new PiConversationHistory(session, rpc) + const pending = history.fork() + await vi.waitFor(() => expect(rejectClone).toBeTypeOf('function')) + await expect(history.fork()).rejects.toThrow('already in progress') + rejectClone(new Error('Unknown command: clone')) + await expect(pending).rejects.toThrow('Unknown command') + expect(history.getCapabilitiesForMetadata()).toBeUndefined() + }) +}) diff --git a/cli/src/pi/conversationHistory.ts b/cli/src/pi/conversationHistory.ts new file mode 100644 index 00000000..349b7a83 --- /dev/null +++ b/cli/src/pi/conversationHistory.ts @@ -0,0 +1,783 @@ +import type { Metadata } from '@/api/types' +import { + PI_CONVERSATION_HISTORY_INITIAL, + markSupported, + markUnsupported, + toConversationHistoryCapabilities, + type ConversationHistoryCapabilityStates +} from '@hapi/protocol/conversationHistory' +import type { ForkConversationRpcResult, RewindConversationRpcResult } from '@hapi/protocol/apiTypes' +import { PI_THINKING_LEVELS } from '@hapi/protocol' +import type { PiThinkingLevel } from './types' +import type { PiNativeRuntimeState, PiSession } from './session' + +/** Keep the complete native history transaction below the Hub's 120s ceiling. */ +export const PI_HISTORY_OPERATION_TIMEOUT_MS = 110_000 +const PI_HISTORY_RESTORE_RESERVE_MS = 10_000 + +type PiRpc = (command: Record, timeoutMs?: number) => Promise + +type PiIdentity = { + sessionId: string + sessionFile: string +} + +type PiState = PiIdentity & { + runtime: PiNativeRuntimeState +} + +type PiEntry = { + id: string + type: string + message?: { role?: unknown } +} + +type PendingUserEntry = { + localId: string +} + +/** Source identity could not be restored; caller must terminate this Pi wrapper. */ +export class PiHistoryRestoreError extends Error { + constructor(message: string) { + super(message) + this.name = 'PiHistoryRestoreError' + } +} + +class PiHistoryDeadlineError extends Error { + constructor(message: string) { + super(message) + this.name = 'PiHistoryDeadlineError' + } +} + +class PiHistoryIndeterminateMutationError extends Error { + constructor(message: string) { + super(message) + this.name = 'PiHistoryIndeterminateMutationError' + } +} + +function asRecord(value: unknown): Record | null { + return value !== null && typeof value === 'object' && !Array.isArray(value) + ? value as Record + : null +} + +function asString(value: unknown): string | null { + return typeof value === 'string' && value.length > 0 ? value : null +} + +function isUnknownCommand(error: unknown): boolean { + const message = error instanceof Error ? error.message : String(error) + return /unknown command|method not found|-32601/i.test(message) +} + +/** + * PiRpcResolver currently emits this message for elapsed requests. Keep this + * matcher deliberately narrow so ordinary native errors still use the normal + * source-restore path. A typed error name is supported for adapters/tests. + */ +function isPiRpcTimeout(error: unknown): boolean { + if (!(error instanceof Error)) return false + return error.name === 'PiRpcTimeoutError' + || /^Pi RPC \S+ \(id=\d+\) timed out after \d+ms$/.test(error.message) +} + +function wasCancelled(data: unknown): boolean { + return asRecord(data)?.cancelled === true +} + +function readState(data: unknown): PiState { + const state = asRecord(data) + const sessionId = asString(state?.sessionId) + const sessionFile = asString(state?.sessionFile) + if (!sessionId || !sessionFile) { + throw new Error('Pi get_state did not return sessionId and sessionFile') + } + const model = asRecord(state?.model) + const thinkingLevel = asString(state?.thinkingLevel) + const hasModelId = model !== null && ('id' in model || 'modelId' in model) + const hasProvider = model !== null && 'provider' in model + return { + sessionId, + sessionFile, + runtime: { + model: hasModelId ? asString(model.id) ?? asString(model.modelId) ?? null : undefined, + provider: hasProvider ? asString(model.provider) ?? null : undefined, + thinkingLevel: thinkingLevel && PI_THINKING_LEVELS.includes(thinkingLevel as PiThinkingLevel) + ? thinkingLevel as PiThinkingLevel + : undefined, + steeringMode: state?.steeringMode === 'all' || state?.steeringMode === 'one-at-a-time' + ? state.steeringMode + : undefined, + isStreaming: typeof state?.isStreaming === 'boolean' ? state.isStreaming : undefined, + }, + } +} + +function readEntries(data: unknown): { entries: PiEntry[]; leafId: string | null } { + const record = asRecord(data) + const rawEntries = Array.isArray(record?.entries) ? record.entries : null + if (!rawEntries) throw new Error('Pi get_entries returned malformed data') + const entries = rawEntries.flatMap((raw): PiEntry[] => { + const entry = asRecord(raw) + const id = asString(entry?.id) + const type = asString(entry?.type) + if (!id || !type) return [] + const message = asRecord(entry?.message) + return [{ id, type, message: message ? { role: message.role } : undefined }] + }) + return { entries, leafId: asString(record?.leafId) } +} + +function isUserEntry(entry: PiEntry): boolean { + return entry.type === 'message' && entry.message?.role === 'user' +} + +function containsForkEntry(data: unknown, entryId: string): boolean { + const messages = asRecord(data)?.messages + return Array.isArray(messages) && messages.some((message) => asRecord(message)?.entryId === entryId) +} + +/** + * Native Pi history coordinator. Entry association is intentionally FIFO only: + * a HAPI prompt is paired with the next Pi user entry, never with message text. + */ +export class PiConversationHistory { + private states: ConversationHistoryCapabilityStates = { ...PI_CONVERSATION_HISTORY_INITIAL } + private readonly entryIdByLocalId = new Map() + private readonly pendingUserEntries: PendingUserEntry[] = [] + private observedEntryIds = new Set() + private appendCursor: string | null = null + private publishCapabilities: (() => Promise) | null = null + private syncInFlight: Promise | null = null + private syncRequestedWhileInFlight = false + private syncGeneration = 0 + private historySyncDisabled = false + + constructor( + private readonly session: PiSession, + private readonly rpc: PiRpc, + ) {} + + setPublishCapabilities(fn: () => Promise): void { + this.publishCapabilities = fn + } + + getCapabilitiesForMetadata(): Metadata['capabilities'] { + const conversationHistory = toConversationHistoryCapabilities(this.states) + return conversationHistory ? { conversationHistory } : undefined + } + + getHistoryPoints(): Record { + return Object.fromEntries(Array.from(this.entryIdByLocalId.keys(), (localId) => [localId, true])) + } + + getEntryIds(): Record { + return Object.fromEntries(this.entryIdByLocalId.entries()) + } + + restoreEntryIds(entryIds: Record | null | undefined): void { + if (!entryIds) return + for (const [localId, entryId] of Object.entries(entryIds)) { + if (localId && entryId) this.entryIdByLocalId.set(localId, entryId) + } + } + + /** Establish the append-log cursor before any buffered prompt is released. */ + async initializeBaseline(): Promise { + try { + await this.syncEntries() + } catch (error) { + // A trustworthy append-log baseline is mandatory before prompts are + // released. Any startup failure disables history for this wrapper; + // otherwise a later full-log read could pair an old user entry with + // the first new HAPI localId. + this.historySyncDisabled = true + this.pendingUserEntries.length = 0 + this.states = markUnsupported(this.states, 'forkCurrent') + this.states = markUnsupported(this.states, 'forkAtMessage') + this.states = markUnsupported(this.states, 'rewindToMessage') + await this.publishCapabilities?.().catch(() => {}) + return false + } + return true + } + + /** Probe and publish controls only after Pi has a validated native identity. */ + async initialize(): Promise { + if (!await this.initializeBaseline()) return + await this.probeCapabilities().catch(() => {}) + } + + /** + * Register a HAPI user message before its corresponding native command is + * written. Prompts and native steers both append Pi user entries, so their + * associations share one strict FIFO rather than a prompt-only queue. + */ + registerUserEntry(localId: string | undefined): void { + if (this.historySyncDisabled) return + if (localId) this.pendingUserEntries.push({ localId }) + } + + /** Remove a rejected/aborted local FIFO entry by exact localId. */ + rejectPendingEntry(localId: string | undefined): void { + if (!localId) return + const index = this.pendingUserEntries.findIndex((entry) => entry.localId === localId) + if (index !== -1) this.pendingUserEntries.splice(index, 1) + } + + observeEntry(rawEntry: unknown): void { + if (this.historySyncDisabled || this.session.isHistoryTransactionActive) return + const parsed = readEntries({ entries: [rawEntry], leafId: null }) + for (const entry of parsed.entries) this.observeParsedEntry(entry) + } + + async syncEntries(): Promise { + if (this.historySyncDisabled || this.session.isHistoryTransactionActive) return + if (this.syncInFlight) { + // A turn_start can be emitted for tool/retry loops before the prior + // incremental read returns. Coalesce it into one serialized follow-up. + this.syncRequestedWhileInFlight = true + return await this.syncInFlight + } + this.syncInFlight = this.runEntrySync().finally(() => { + // A request can land after runEntrySync observes `false` but before + // finally clears syncInFlight. Preserve that boundary request. + const scheduleFollowUp = this.syncRequestedWhileInFlight && !this.session.isHistoryTransactionActive + this.syncInFlight = null + this.syncRequestedWhileInFlight = false + if (scheduleFollowUp) void this.syncEntries().catch(() => {}) + }) + return await this.syncInFlight + } + + private async runEntrySync(): Promise { + do { + this.syncRequestedWhileInFlight = false + await this.syncEntriesOnce() + } while (this.syncRequestedWhileInFlight && !this.session.isHistoryTransactionActive) + } + + private async syncEntriesOnce(timeoutMs?: number): Promise { + const generation = this.syncGeneration + const data = await this.rpc( + this.appendCursor + ? { type: 'get_entries', since: this.appendCursor } + : { type: 'get_entries' }, + timeoutMs, + ) + if (generation !== this.syncGeneration) return + const result = readEntries(data) + for (const entry of result.entries) this.observeParsedEntry(entry) + // `since` indexes the immutable append log, not the active branch. + // A fork can move leafId backwards; advancing the cursor to it would + // replay entries and break FIFO pairing. Empty increments keep cursor. + if (result.entries.length > 0) this.appendCursor = result.entries[result.entries.length - 1]!.id + } + + async probeCapabilities(): Promise { + if (this.historySyncDisabled) return + if (this.states.forkCurrent !== 'unknown' && this.states.forkAtMessage !== 'unknown' + && this.states.rewindToMessage !== 'unknown') return + try { + // Both reads are side-effect free and exist together with Pi 0.83's + // clone/fork APIs. Do not expose controls before this succeeds. + await this.rpc({ type: 'get_fork_messages' }) + const entries = await this.rpc({ type: 'get_entries', ...(this.appendCursor ? { since: this.appendCursor } : {}) }) + readEntries(entries) + this.states = markSupported(this.states, 'forkCurrent') + this.states = markSupported(this.states, 'forkAtMessage') + this.states = markSupported(this.states, 'rewindToMessage') + } catch (error) { + if (isUnknownCommand(error)) { + this.states = markUnsupported(this.states, 'forkCurrent') + this.states = markUnsupported(this.states, 'forkAtMessage') + this.states = markUnsupported(this.states, 'rewindToMessage') + } + } + await this.publishCapabilities?.() + } + + async fork(messageLocalId?: string): Promise { + this.assertHistoryIdle() + if (messageLocalId) return await this.forkHistorical(messageLocalId) + if (this.states.forkCurrent === 'unsupported') throw new Error('Fork current is not supported') + + return await this.withSourceRestored('forkCurrent', async (source, deadlineAt, markMutationIssued) => { + const clone = await this.cloneAndReadIdentity(source, deadlineAt, markMutationIssued) + return { nativeSessionId: clone.sessionId } + }) + } + + async rewind(messageLocalId: string): Promise { + try { + this.assertHistoryIdle() + } catch (error) { + return { success: false, error: error instanceof Error ? error.message : String(error), outcome: 'rejected' } + } + if (this.states.rewindToMessage === 'unsupported') { + return { success: false, error: 'Rewind is not supported', outcome: 'rejected' } + } + const entryId = this.entryIdByLocalId.get(messageLocalId) + if (!entryId) { + return { success: false, error: `No native history point for message ${messageLocalId}`, outcome: 'rejected' } + } + + const transaction = await this.beginHistoryTransaction() + if (transaction.rejection) { + transaction.release() + return { success: false, error: transaction.rejection, outcome: 'rejected' } + } + const { release, deadlineAt } = transaction + let source: PiState | null = null + let committed = false + let mutationIssued = false + let mutationCompleted = false + let rollbackRewindMetadata = false + let indeterminateTimeout: unknown + const locatorSnapshot = this.captureLocatorState() + let success: Extract | null = null + let failure: { error: string; outcome: 'rejected' | 'cancelled' | 'source_restored' } | null = null + try { + source = await this.getState(deadlineAt) + const forkMessages = await this.rpcWithinDeadline({ type: 'get_fork_messages' }, deadlineAt) + if (!containsForkEntry(forkMessages, entryId)) { + throw new Error('Pi rewind point is no longer available') + } + const result = await this.nativeMutation( + { type: 'fork', entryId }, + deadlineAt, + true, + () => { mutationIssued = true }, + ) + mutationCompleted = true + if (wasCancelled(result)) { + failure = { error: 'Pi rewind was cancelled', outcome: 'cancelled' } + throw new Error(failure.error) + } + const forked = await this.getState(deadlineAt) + this.assertDistinctIdentity(source, forked, 'Pi rewind') + const entries = readEntries(await this.rpcWithinDeadline({ type: 'get_entries' }, deadlineAt)) + this.commitRewindState(forked, entries) + if (!await this.session.flushMetadata(Math.min(5_000, this.remainingMs(deadlineAt)))) { + rollbackRewindMetadata = true + throw new Error('Pi rewind metadata did not persist') + } + committed = true + this.states = markSupported(this.states, 'rewindToMessage') + success = { success: true, truncateFromLocalId: messageLocalId, messages: [] } + } catch (error) { + if (error instanceof PiHistoryIndeterminateMutationError) { + indeterminateTimeout = error + } else if (mutationCompleted && (isPiRpcTimeout(error) || error instanceof PiHistoryDeadlineError)) { + indeterminateTimeout = new PiHistoryIndeterminateMutationError( + `Pi rewind completed but its resulting state is indeterminate: ${error instanceof Error ? error.message : String(error)}`, + ) + } + if (!failure) { + failure = { + error: error instanceof Error ? error.message : String(error), + outcome: 'rejected' + } + } + if (isUnknownCommand(error)) this.states = markUnsupported(this.states, 'rewindToMessage') + } finally { + let restoreError: unknown + let restoredSource: PiState | null = null + // Before commit this is a failed transaction and the old source must + // remain active. After commit, the branched Pi session *is* rewind. + // A timeout is indeterminate: Pi may have committed its native + // mutation after the client gave up. Do not issue a momentary + // get_state/switch classification against an unknown active branch. + if (!indeterminateTimeout && !committed && source && mutationIssued) { + try { + const restored = await this.restoreSource(source, deadlineAt) + restoredSource = restored + this.session.applyNativeRuntimeState(restored.runtime) + } catch (error) { + restoreError = error + } + } + if (!restoreError && rollbackRewindMetadata && restoredSource) { + this.restoreLocatorState(locatorSnapshot) + this.session.commitNativeSessionState(restoredSource, restoredSource.runtime, (metadata) => + this.metadataWithLocators(metadata, locatorSnapshot.entryIds, locatorSnapshot.points) + ) + try { + const timeoutMs = Math.min(5_000, this.remainingMs(deadlineAt)) + if (!await this.session.flushMetadata(timeoutMs)) { + restoreError = new Error('Pi rewind metadata rollback did not persist') + } + } catch (error) { + // Deadline calculation is part of rollback persistence. Keep + // it inside the restoreError path so release() always runs + // and the wrapper fails closed instead of retaining both + // the history gate and runtime-mutation lease. + restoreError = error + } + } + release({ drain: !restoreError && !indeterminateTimeout }) + if (indeterminateTimeout) { + throw new PiHistoryRestoreError(`Pi rewind timed out with indeterminate native state: ${indeterminateTimeout instanceof Error ? indeterminateTimeout.message : String(indeterminateTimeout)}`) + } + if (restoreError) { + await this.publishCapabilities?.().catch(() => {}) + throw new PiHistoryRestoreError(`Pi rewind failed closed: source session restoration failed: ${restoreError instanceof Error ? restoreError.message : String(restoreError)}`) + } + // Capability metadata is advisory; it must not turn a precisely + // restored deterministic rejection into a Hub-diverging exception. + await this.publishCapabilities?.().catch(() => {}) + } + if (success) return success + if (source && failure?.outcome !== 'cancelled') { + return { success: false, ...failure!, outcome: 'source_restored' } + } + return failure + ? { success: false, ...failure } + : { success: false, error: 'Pi rewind did not complete', outcome: 'rejected' } + } + + private async forkHistorical(messageLocalId: string): Promise { + if (this.states.forkAtMessage === 'unsupported') throw new Error('Historical fork is not supported') + const entryId = this.entryIdByLocalId.get(messageLocalId) + if (!entryId) throw new Error(`No native history point for message ${messageLocalId}`) + + return await this.withSourceRestored('forkAtMessage', async (source, deadlineAt, markMutationIssued) => { + let mutationCompleted = false + try { + const result = await this.nativeMutation( + { type: 'fork', entryId }, + deadlineAt, + true, + markMutationIssued, + ) + mutationCompleted = true + if (wasCancelled(result)) throw new Error('Pi historical fork was cancelled') + const afterFork = await this.getState(deadlineAt) + this.assertDistinctIdentity(source, afterFork, 'Pi historical fork') + return { nativeSessionId: afterFork.sessionId } + } catch (error) { + if (mutationCompleted && (isPiRpcTimeout(error) || error instanceof PiHistoryDeadlineError)) { + throw new PiHistoryIndeterminateMutationError( + `Pi historical fork completed but its resulting state is indeterminate: ${error instanceof Error ? error.message : String(error)}`, + ) + } + throw error + } + }) + } + + private async withSourceRestored( + capability: keyof ConversationHistoryCapabilityStates, + work: (source: PiState, deadlineAt: number, markMutationIssued: () => void) => Promise, + ): Promise { + const transaction = await this.beginHistoryTransaction() + if (transaction.rejection) { + transaction.release() + throw new Error(transaction.rejection) + } + const { release, deadlineAt } = transaction + let source: PiState | null = null + let outcome: T | undefined + let operationError: unknown + let indeterminateTimeout: unknown + let mutationIssued = false + try { + source = await this.getState(deadlineAt) + outcome = await work(source, deadlineAt, () => { mutationIssued = true }) + this.states = markSupported(this.states, capability) + } catch (error) { + operationError = error + if (error instanceof PiHistoryIndeterminateMutationError) indeterminateTimeout = error + if (isUnknownCommand(error)) { + this.states = markUnsupported(this.states, capability) + // Both fork flows start with Pi's clone command; a real + // unknown-command response there invalidates both affordances. + if (capability === 'forkCurrent' || capability === 'forkAtMessage') { + this.states = markUnsupported(this.states, 'forkCurrent') + this.states = markUnsupported(this.states, 'forkAtMessage') + } + } + } finally { + let restoreError: unknown + // A timed-out native mutation could have completed late. Any read or + // switch used to classify it is itself a divergent mutation race. + if (!indeterminateTimeout && source && mutationIssued) { + try { + const restored = await this.restoreSource(source, deadlineAt) + this.session.applyNativeRuntimeState(restored.runtime) + } catch (error) { + restoreError = error + } + } + release({ drain: !restoreError && !indeterminateTimeout }) + if (indeterminateTimeout) { + throw new PiHistoryRestoreError(`Pi history operation timed out with indeterminate native state: ${indeterminateTimeout instanceof Error ? indeterminateTimeout.message : String(indeterminateTimeout)}`) + } + if (restoreError) { + await this.publishCapabilities?.().catch(() => {}) + throw new PiHistoryRestoreError(`Pi history operation failed closed: source session restoration failed: ${restoreError instanceof Error ? restoreError.message : String(restoreError)}`) + } + await this.publishCapabilities?.() + } + if (operationError) throw operationError + return outcome as T + } + + private assertHistoryIdle(): void { + if (this.historySyncDisabled) throw new Error('Pi conversation history is unavailable') + if (!this.session.isNativeReady) throw new Error('Pi native session is not ready') + if (this.session.piIsStreaming || this.session.hasPromptInFlight) throw new Error('Pi session is busy') + } + + private remainingMs(deadlineAt: number, reserveMs: number = 0): number { + const remaining = Math.floor(deadlineAt - Date.now() - reserveMs) + if (remaining <= 0) { + throw new PiHistoryDeadlineError('Pi history operation exceeded its transaction deadline') + } + return remaining + } + + private async rpcWithinDeadline( + command: Record, + deadlineAt: number, + reserveMs: number = 0, + ): Promise { + return await this.rpc(command, this.remainingMs(deadlineAt, reserveMs)) + } + + private async nativeMutation( + command: Record, + deadlineAt: number, + reserveRestore: boolean = true, + onIssued?: () => void, + ): Promise { + const timeoutMs = this.remainingMs( + deadlineAt, + reserveRestore ? PI_HISTORY_RESTORE_RESERVE_MS : 0, + ) + onIssued?.() + try { + return await this.rpc(command, timeoutMs) + } catch (error) { + if (isPiRpcTimeout(error)) { + throw new PiHistoryIndeterminateMutationError( + `Pi ${String(command.type)} timed out with indeterminate native state: ${error instanceof Error ? error.message : String(error)}`, + ) + } + throw error + } + } + + private async waitWithinDeadline(promise: Promise, deadlineAt: number, operation: string): Promise { + const timeoutMs = this.remainingMs(deadlineAt) + let timer: ReturnType | null = null + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => reject(new PiHistoryDeadlineError( + `Pi history operation timed out while waiting to ${operation}`, + )), timeoutMs) + timer.unref?.() + }), + ]) + } finally { + if (timer) clearTimeout(timer) + } + } + + /** + * Lock new prompts before waiting for old reads, then take one final source + * snapshot. This closes the mapping race between Pi persistence and a + * history mutation without discarding a pre-existing pending localId. + */ + private async beginHistoryTransaction(): Promise<{ + release: (options?: { drain?: boolean }) => void + deadlineAt: number + rejection?: string + }> { + const deadlineAt = Date.now() + PI_HISTORY_OPERATION_TIMEOUT_MS + // The history gate is intentionally acquired before the first await. + // This queues later runtime mutations behind us while an already-active + // config/abort mutation drains, closing the final-sync/fork race. + const releaseHistoryGate = this.session.beginHistoryTransaction() + let releaseRuntimeMutation: (() => void) | null = null + let released = false + const release = (options?: { drain?: boolean }) => { + if (released) return + released = true + // Drain (or discard) prompt work while the runtime mutex remains + // held, then make the next config/abort mutation eligible. + releaseHistoryGate(options) + releaseRuntimeMutation?.() + } + try { + const acquireRuntimeMutation = this.session.acquireRuntimeMutation() + try { + releaseRuntimeMutation = await this.waitWithinDeadline( + acquireRuntimeMutation, + deadlineAt, + 'acquire the runtime mutation lock', + ) + } catch (error) { + // The FIFO mutex acquisition cannot be cancelled. Release its + // eventual lease immediately so a timed-out history request + // cannot wedge later config/abort operations. + void acquireRuntimeMutation.then((lateRelease) => lateRelease()) + throw error + } + if (this.syncInFlight) { + await this.waitWithinDeadline(this.syncInFlight, deadlineAt, 'finish the previous history sync') + } + await this.syncEntriesOnce(this.remainingMs(deadlineAt)) + if (!await this.session.flushMetadata(Math.min(5_000, this.remainingMs(deadlineAt)))) { + return { release, deadlineAt, rejection: 'Pi history metadata did not persist before native fork' } + } + } catch (error) { + return { + release, + deadlineAt, + rejection: `Pi history synchronization failed: ${error instanceof Error ? error.message : String(error)}` + } + } + if (this.pendingUserEntries.length > 0) { + return { release, deadlineAt, rejection: 'Pi session has pending user entries' } + } + this.invalidatePendingSync() + return { release, deadlineAt } + } + + private async cloneAndReadIdentity( + source: PiState, + deadlineAt: number, + markMutationIssued: () => void, + ): Promise { + let mutationCompleted = false + try { + const cloned = await this.nativeMutation( + { type: 'clone' }, + deadlineAt, + true, + markMutationIssued, + ) + mutationCompleted = true + if (wasCancelled(cloned)) throw new Error('Pi clone was cancelled') + const clone = await this.getState(deadlineAt) + this.assertDistinctIdentity(source, clone, 'Pi clone') + return clone + } catch (error) { + if (mutationCompleted && (isPiRpcTimeout(error) || error instanceof PiHistoryDeadlineError)) { + throw new PiHistoryIndeterminateMutationError( + `Pi clone completed but its resulting state is indeterminate: ${error instanceof Error ? error.message : String(error)}`, + ) + } + throw error + } + } + + private assertDistinctIdentity(source: PiIdentity, next: PiIdentity, operation: string): void { + if (next.sessionId === source.sessionId && next.sessionFile === source.sessionFile) { + throw new Error(`${operation} did not create a distinct native session identity`) + } + } + + /** Reset the append cursor and retain only Pi entry mappings copied into the new branch. */ + private commitRewindState(state: PiState, entries: { entries: PiEntry[]; leafId: string | null }): void { + const validEntryIds = new Set(entries.entries.map((entry) => entry.id)) + for (const [localId, entryId] of this.entryIdByLocalId.entries()) { + if (!validEntryIds.has(entryId)) this.entryIdByLocalId.delete(localId) + } + this.observedEntryIds = validEntryIds + this.appendCursor = entries.entries.length > 0 ? entries.entries[entries.entries.length - 1]!.id : null + this.invalidatePendingSync() + const entryIds = this.getEntryIds() + const points = this.getHistoryPoints() + this.session.commitNativeSessionState(state, state.runtime, (metadata) => this.metadataWithLocators(metadata, entryIds, points)) + } + + private captureLocatorState(): { + entryIds: Record + points: Record + observedEntryIds: Set + appendCursor: string | null + } { + return { + entryIds: this.getEntryIds(), + points: this.getHistoryPoints(), + observedEntryIds: new Set(this.observedEntryIds), + appendCursor: this.appendCursor + } + } + + private restoreLocatorState(snapshot: ReturnType): void { + this.entryIdByLocalId.clear() + for (const [localId, entryId] of Object.entries(snapshot.entryIds)) { + this.entryIdByLocalId.set(localId, entryId) + } + this.observedEntryIds = new Set(snapshot.observedEntryIds) + this.appendCursor = snapshot.appendCursor + } + + private metadataWithLocators( + metadata: Metadata, + entryIds: Record, + points: Record + ): Metadata { + const next: Metadata = { ...metadata, conversationHistoryEntryIds: entryIds, conversationHistoryPoints: points } + if (Object.keys(entryIds).length === 0) delete next.conversationHistoryEntryIds + if (Object.keys(points).length === 0) delete next.conversationHistoryPoints + return next + } + + private async restoreSource(source: PiState, deadlineAt: number): Promise { + const current = await this.getState(deadlineAt) + if (current.sessionId === source.sessionId && current.sessionFile === source.sessionFile) return current + const switched = await this.nativeMutation( + { type: 'switch_session', sessionPath: source.sessionFile }, + deadlineAt, + false, + ) + if (wasCancelled(switched)) throw new Error('Pi source session restoration was cancelled') + const restored = await this.getState(deadlineAt) + if (restored.sessionId !== source.sessionId || restored.sessionFile !== source.sessionFile) { + throw new Error('Pi source session restoration returned a different identity') + } + return restored + } + + private async getState(deadlineAt?: number): Promise { + return readState(await (deadlineAt === undefined + ? this.rpc({ type: 'get_state' }) + : this.rpcWithinDeadline({ type: 'get_state' }, deadlineAt))) + } + + private observeParsedEntry(entry: PiEntry): void { + if (this.observedEntryIds.has(entry.id)) return + this.observedEntryIds.add(entry.id) + this.appendCursor = entry.id + if (!isUserEntry(entry)) return + const pending = this.pendingUserEntries.shift() + if (!pending || this.entryIdByLocalId.has(pending.localId)) return + const localId = pending.localId + this.entryIdByLocalId.set(localId, entry.id) + this.session.updateMetadata((metadata) => ({ + ...metadata, + conversationHistoryPoints: { + ...metadata.conversationHistoryPoints, + [localId]: true as const, + }, + conversationHistoryEntryIds: { + ...metadata.conversationHistoryEntryIds, + [localId]: entry.id, + }, + })) + } + + private invalidatePendingSync(): void { + this.syncGeneration += 1 + this.syncRequestedWhileInFlight = false + } +} diff --git a/cli/src/pi/extensionUiHandler.test.ts b/cli/src/pi/extensionUiHandler.test.ts new file mode 100644 index 00000000..e475347e --- /dev/null +++ b/cli/src/pi/extensionUiHandler.test.ts @@ -0,0 +1,224 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { getExtensionUiCleanupTimeout, PiExtensionUiHandler } from './extensionUiHandler'; + +type PermissionHandler = (response: unknown) => Promise; + +function createHarness() { + let permissionHandler: PermissionHandler | null = null; + let state: Record = { requests: {}, completedRequests: {} }; + const session = { + rpcHandlerManager: { + registerHandler: vi.fn((_method: unknown, handler: PermissionHandler) => { permissionHandler = handler; }), + }, + updateAgentState: vi.fn((updater: (current: never) => unknown) => { state = updater(state as never) as Record; }), + sendAgentMessage: vi.fn(), + sendSessionEvent: vi.fn(), + getMetadata: vi.fn(() => null), + updateMetadata: vi.fn(), + }; + const sendResponse = vi.fn(); + const handler = new PiExtensionUiHandler({ session: session as never, sendResponse }); + return { + handler, + session, + sendResponse, + state: () => state, + respond: async (response: unknown) => permissionHandler?.(response), + }; +} + +describe('PiExtensionUiHandler', () => { + beforeEach(() => vi.useRealTimers()); + + it('maps select into request_user_input and returns its selected option', async () => { + const harness = createHarness(); + harness.handler.handle({ type: 'extension_ui_request', id: 'select-1', method: 'select', title: 'Pick', options: ['one', 'two'] }); + expect(harness.state().requests).toMatchObject({ + 'select-1': { tool: 'request_user_input', arguments: { questions: [{ id: 'select-1', options: [{ label: 'one' }, { label: 'two' }] }] } }, + }); + expect(harness.session.sendAgentMessage).toHaveBeenCalledWith(expect.objectContaining({ + type: 'tool-call', callId: 'select-1', name: 'request_user_input', status: 'in_progress', + })); + + await harness.respond({ + id: 'select-1', + approved: true, + answers: { 'select-1': { answers: ['two', 'user_note: optional note'] } } + }); + expect(harness.sendResponse).toHaveBeenCalledWith({ type: 'extension_ui_response', id: 'select-1', value: 'two' }); + expect(harness.session.sendAgentMessage).toHaveBeenCalledWith({ + type: 'tool-call-result', + callId: 'select-1', + output: { 'select-1': { answers: ['two', 'user_note: optional note'] } }, + is_error: false, + }); + expect(harness.state().completedRequests).toMatchObject({ 'select-1': { status: 'approved' } }); + }); + + it('round-trips select options that use metadata-like prefixes or padded labels', async () => { + const harness = createHarness(); + harness.handler.handle({ + type: 'extension_ui_request', id: 'select-prefix', method: 'select', title: 'Pick', + options: ['user_note: later', ' padded option '], + }); + await harness.respond({ + id: 'select-prefix', approved: true, + answers: { 'select-prefix': { answers: ['user_note: later'] } }, + }); + expect(harness.sendResponse).toHaveBeenLastCalledWith({ + type: 'extension_ui_response', id: 'select-prefix', value: 'user_note: later', + }); + + harness.handler.handle({ + type: 'extension_ui_request', id: 'select-padded', method: 'select', title: 'Pick', + options: [' padded option '], + }); + await harness.respond({ + id: 'select-padded', approved: true, + answers: { 'select-padded': { answers: ['padded option'] } }, + }); + expect(harness.sendResponse).toHaveBeenLastCalledWith({ + type: 'extension_ui_response', id: 'select-padded', value: ' padded option ', + }); + + harness.handler.handle({ + type: 'extension_ui_request', id: 'select-exact-first', method: 'select', title: 'Pick', + options: [' padded option ', 'padded option'], + }); + await harness.respond({ + id: 'select-exact-first', approved: true, + answers: { 'select-exact-first': { answers: ['padded option'] } }, + }); + expect(harness.sendResponse).toHaveBeenLastCalledWith({ + type: 'extension_ui_response', id: 'select-exact-first', value: 'padded option', + }); + + harness.handler.handle({ + type: 'extension_ui_request', id: 'select-ambiguous', method: 'select', title: 'Pick', + options: [' padded option', 'padded option '], + }); + await harness.respond({ + id: 'select-ambiguous', approved: true, + answers: { 'select-ambiguous': { answers: ['padded option'] } }, + }); + expect(harness.sendResponse).toHaveBeenLastCalledWith({ + type: 'extension_ui_response', id: 'select-ambiguous', cancelled: true, + }); + }); + + it('maps confirm to the generic permission card and retains a denied history entry', async () => { + const harness = createHarness(); + harness.handler.handle({ type: 'extension_ui_request', id: 'confirm-1', method: 'confirm', title: 'Proceed?', message: 'Continue this extension?' }); + expect(harness.state().requests).toMatchObject({ 'confirm-1': { tool: 'PiExtensionConfirm' } }); + + // The normal Hub deny route omits an explicit decision. + await harness.respond({ id: 'confirm-1', approved: false }); + expect(harness.sendResponse).toHaveBeenCalledWith({ type: 'extension_ui_response', id: 'confirm-1', confirmed: false }); + expect(harness.state().completedRequests).toMatchObject({ 'confirm-1': { status: 'denied', decision: 'denied' } }); + }); + + it('preserves editor prefill and cancels timeout/session shutdown exactly once', async () => { + vi.useFakeTimers(); + const harness = createHarness(); + harness.handler.handle({ type: 'extension_ui_request', id: 'editor-1', method: 'editor', title: 'Edit', prefill: 'existing text' }); + expect(harness.state().requests).toMatchObject({ + 'editor-1': { arguments: { questions: [{ inputType: 'editor', prefill: 'existing text' }] } }, + }); + await vi.advanceTimersByTimeAsync(120_000); + expect(harness.sendResponse).not.toHaveBeenCalled(); + harness.handler.cancelAll('session shutdown'); + expect(harness.sendResponse).toHaveBeenCalledWith({ type: 'extension_ui_response', id: 'editor-1', cancelled: true }); + expect(harness.state().completedRequests).toMatchObject({ 'editor-1': { status: 'canceled', decision: 'abort' } }); + + harness.handler.handle({ type: 'extension_ui_request', id: 'input-1', method: 'input', title: 'Name', placeholder: 'Ada', timeout: 50 }); + await vi.advanceTimersByTimeAsync(38); + expect(harness.sendResponse).toHaveBeenLastCalledWith({ type: 'extension_ui_response', id: 'input-1', cancelled: true }); + + harness.handler.handle({ type: 'extension_ui_request', id: 'no-timeout', method: 'input', title: 'Name', timeout: 0 }); + await vi.advanceTimersByTimeAsync(1_000); + expect(harness.sendResponse).not.toHaveBeenLastCalledWith({ type: 'extension_ui_response', id: 'no-timeout', cancelled: true }); + }); + + it('returns editor whitespace and an intentionally empty document byte-for-byte', async () => { + const harness = createHarness(); + harness.handler.handle({ type: 'extension_ui_request', id: 'editor-whitespace', method: 'editor', title: 'Edit' }); + await harness.respond({ + id: 'editor-whitespace', approved: true, + answers: { 'editor-whitespace': { answers: ['user_note: code\n'] } }, + }); + expect(harness.sendResponse).toHaveBeenCalledWith({ + type: 'extension_ui_response', id: 'editor-whitespace', value: ' code\n', + }); + + harness.handler.handle({ type: 'extension_ui_request', id: 'editor-empty', method: 'editor', title: 'Edit' }); + await harness.respond({ + id: 'editor-empty', approved: true, + answers: { 'editor-empty': { answers: ['user_note: '] } }, + }); + expect(harness.sendResponse).toHaveBeenCalledWith({ + type: 'extension_ui_response', id: 'editor-empty', value: '', + }); + }); + + it('reserves a capped proportional margin before Pi expires the dialog', async () => { + vi.useFakeTimers(); + expect(getExtensionUiCleanupTimeout(undefined)).toBeUndefined(); + expect(getExtensionUiCleanupTimeout(0)).toBeUndefined(); + expect(getExtensionUiCleanupTimeout(0.5)).toBe(0); + expect(getExtensionUiCleanupTimeout(1)).toBe(0); + expect(getExtensionUiCleanupTimeout(2)).toBe(1); + expect(getExtensionUiCleanupTimeout(50)).toBe(38); + expect(getExtensionUiCleanupTimeout(1_000)).toBe(900); + expect(getExtensionUiCleanupTimeout(100_000)).toBe(95_000); + + const harness = createHarness(); + harness.handler.handle({ type: 'extension_ui_request', id: 'margin-1', method: 'input', title: 'Name', timeout: 1_000 }); + await vi.advanceTimersByTimeAsync(899); + expect(harness.sendResponse).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(1); + expect(harness.sendResponse).toHaveBeenCalledWith({ type: 'extension_ui_response', id: 'margin-1', cancelled: true }); + }); + + it('puts notify on the timeline and ignores unsupported transient UI operations', () => { + const harness = createHarness(); + harness.handler.handle({ type: 'extension_ui_request', id: 'notice', method: 'notify', message: 'Heads up', notifyType: 'warning' }); + harness.handler.handle({ type: 'extension_ui_request', id: 'status', method: 'setStatus', statusKey: 'x', statusText: 'busy' }); + expect(harness.session.sendSessionEvent).toHaveBeenCalledWith({ type: 'message', message: '[Pi warning] Heads up' }); + }); +}); + + +describe('PiExtensionUiHandler duplicate ids', () => { + it('tombstones a reused id so delayed approval cannot bind a replacement dialog', async () => { + const harness = createHarness(); + harness.handler.handle({ type: 'extension_ui_request', id: 'same', method: 'input', title: 'First' }); + harness.handler.handle({ type: 'extension_ui_request', id: 'same', method: 'input', title: 'Second' }); + expect(harness.sendResponse).toHaveBeenCalledWith({ type: 'extension_ui_response', id: 'same', cancelled: true }); + await harness.respond({ id: 'same', approved: true, answers: { same: { answers: ['user_note: late'] } } }); + expect(harness.sendResponse).toHaveBeenCalledTimes(1); + expect(harness.session.sendSessionEvent).toHaveBeenCalledWith(expect.objectContaining({ type: 'message', message: expect.stringContaining('id was reused') })); + }); + + it('retires an id after normal completion and cancels any later reuse', async () => { + const harness = createHarness(); + harness.handler.handle({ type: 'extension_ui_request', id: 'completed', method: 'input', title: 'First' }); + await harness.respond({ id: 'completed', approved: true, answers: { completed: { answers: ['user_note: accepted'] } } }); + harness.handler.handle({ type: 'extension_ui_request', id: 'completed', method: 'input', title: 'Replacement' }); + + expect(harness.sendResponse).toHaveBeenNthCalledWith(1, { type: 'extension_ui_response', id: 'completed', value: 'accepted' }); + expect(harness.sendResponse).toHaveBeenNthCalledWith(2, { type: 'extension_ui_response', id: 'completed', cancelled: true }); + await harness.respond({ id: 'completed', approved: true, answers: { completed: { answers: ['user_note: late'] } } }); + expect(harness.sendResponse).toHaveBeenCalledTimes(2); + }); + + it('retires a timed-out id and immediately cancels a replacement dialog', async () => { + vi.useFakeTimers(); + const harness = createHarness(); + harness.handler.handle({ type: 'extension_ui_request', id: 'timed-out', method: 'input', title: 'First', timeout: 50 }); + await vi.advanceTimersByTimeAsync(38); + harness.handler.handle({ type: 'extension_ui_request', id: 'timed-out', method: 'input', title: 'Replacement' }); + + expect(harness.sendResponse).toHaveBeenNthCalledWith(1, { type: 'extension_ui_response', id: 'timed-out', cancelled: true }); + expect(harness.sendResponse).toHaveBeenNthCalledWith(2, { type: 'extension_ui_response', id: 'timed-out', cancelled: true }); + }); +}); diff --git a/cli/src/pi/extensionUiHandler.ts b/cli/src/pi/extensionUiHandler.ts new file mode 100644 index 00000000..d1461341 --- /dev/null +++ b/cli/src/pi/extensionUiHandler.ts @@ -0,0 +1,310 @@ +import { z } from 'zod'; +import type { ApiSessionClient } from '@/api/apiSession'; +import type { AgentState } from '@/api/types'; +import { createNativeSessionTitleMetadataSync } from '@/agent/nativeSessionTitle'; +import { RPC_METHODS } from '@hapi/protocol/rpcMethods'; +import { logger } from '@/ui/logger'; +import type { PiExtensionUiRequest, PiExtensionUiResponse } from './types'; + +const PiPermissionResponseSchema = z.object({ + id: z.string().min(1), + approved: z.boolean(), + decision: z.enum(['approved', 'approved_for_session', 'denied', 'abort']).optional(), + reason: z.string().optional(), + answers: z.record(z.string(), z.union([ + z.array(z.string()), + z.object({ answers: z.array(z.string()) }), + ])).optional(), +}).passthrough(); + +type PiPermissionResponse = z.infer; + +type PendingExtensionRequest = Extract; + +type PendingEntry = { + request: PendingExtensionRequest; + timer: ReturnType | null; +}; + +type PiExtensionUiHandlerOptions = { + session: Pick; + sendResponse: (response: PiExtensionUiResponse) => void; +}; + +// Pi starts its dialog timeout before the event reaches the HAPI process. Keep +// HAPI's cleanup timer ahead of that deadline so a late browser response cannot +// race Pi's own timeout. The 10% allowance covers transport/UI propagation, +// has a 5s ceiling for long-lived dialogs, and is capped at a quarter of a +// short timeout so valid short dialogs still retain at least 75% of their +// requested duration. +const EXTENSION_UI_TIMEOUT_MARGIN_RATIO = 0.1; +const EXTENSION_UI_TIMEOUT_MARGIN_MIN_MS = 100; +const EXTENSION_UI_TIMEOUT_MARGIN_MAX_MS = 5_000; +const EXTENSION_UI_TIMEOUT_MARGIN_MAX_RATIO = 0.25; + +export function getExtensionUiCleanupTimeout(timeout: number | undefined): number | undefined { + if (timeout === undefined || timeout === 0) return undefined; + + const requestedMargin = Math.min( + EXTENSION_UI_TIMEOUT_MARGIN_MAX_MS, + Math.max(EXTENSION_UI_TIMEOUT_MARGIN_MIN_MS, Math.ceil(timeout * EXTENSION_UI_TIMEOUT_MARGIN_RATIO)), + ); + // Even a very short positive timeout must clean up before Pi's timer. Such + // dialogs are not realistically interactive, so reserve at least 1ms. + const margin = Math.min(timeout, Math.max(1, Math.min( + Math.floor(timeout * EXTENSION_UI_TIMEOUT_MARGIN_MAX_RATIO), + requestedMargin, + ))); + return timeout - margin; +} + +function requestToolName(request: PendingExtensionRequest): string { + return request.method === 'confirm' ? 'PiExtensionConfirm' : 'request_user_input'; +} + +function requestArguments(request: PendingExtensionRequest): Record { + if (request.method === 'confirm') { + return { title: request.title, message: request.message }; + } + + const question: Record = { + id: request.id, + header: request.title, + question: request.title, + required: true, + multiple: false, + options: request.method === 'select' + ? request.options.map((label) => ({ label, description: null })) + : [], + }; + + if (request.method === 'input' && request.placeholder) { + question.placeholder = request.placeholder; + } + if (request.method === 'editor') { + // These fields intentionally travel in the request payload. The web's + // request_user_input renderer owns their presentation; dropping them here + // would make an extension editor lose its initial document on the way to + // the mobile client. + question.inputType = 'editor'; + if (request.prefill !== undefined) question.prefill = request.prefill; + } + + return { questions: [question] }; +} + +function extractAnswer( + response: PiPermissionResponse, + request: Exclude +): string | null { + const raw = response.answers?.[request.id]; + const answers = Array.isArray(raw) ? raw : raw?.answers; + if (!answers || answers.length === 0) return null; + + if (request.method === 'select') { + for (const answer of answers) { + const exact = request.options.find((option) => option === answer); + if (exact !== undefined) return exact; + const trimmedMatches = request.options.filter((option) => option.trim() === answer); + if (trimmedMatches.length === 1) return trimmedMatches[0]!; + if (trimmedMatches.length > 1) return null; + } + return null; + } + const note = answers.find((answer) => answer.startsWith('user_note: ')); + if (note) return note.slice('user_note: '.length); + return answers[0] ?? null; +} + +function normalizeAnswers(answers: PiPermissionResponse['answers']): Record | undefined { + if (!answers) return undefined; + const normalized: Record = {}; + for (const [id, value] of Object.entries(answers)) { + normalized[id] = { answers: Array.isArray(value) ? value : value.answers }; + } + return normalized; +} + +/** + * Bridges Pi RPC extension UI events onto HAPI's existing AgentState / Permission + * RPC contract. Blocking Pi dialogs remain one HAPI request each and are always + * completed with exactly one extension_ui_response, including timeouts and + * process shutdown. + */ +export class PiExtensionUiHandler { + private readonly pending = new Map(); + private readonly tombstonedIds = new Set(); + private readonly syncTitle: (title: unknown) => void; + + constructor(private readonly options: PiExtensionUiHandlerOptions) { + this.syncTitle = createNativeSessionTitleMetadataSync(options.session); + options.session.rpcHandlerManager.registerHandler(RPC_METHODS.Permission, async (rawResponse) => { + const parsed = PiPermissionResponseSchema.safeParse(rawResponse); + if (!parsed.success) { + logger.debug('[pi] Ignoring malformed extension UI permission response'); + return; + } + this.handlePermissionResponse(parsed.data); + }); + } + + handle(request: PiExtensionUiRequest): void { + switch (request.method) { + case 'notify': + this.options.session.sendSessionEvent({ + type: 'message', + message: `[Pi ${request.notifyType ?? 'info'}] ${request.message}`, + }); + return; + case 'setTitle': + this.syncTitle(request.title); + return; + case 'setStatus': + case 'setWidget': + case 'set_editor_text': + logger.debug(`[pi] Extension UI ${request.method} is not exposed by the HAPI transport`); + return; + case 'select': + case 'confirm': + case 'input': + case 'editor': + this.registerPending(request); + return; + } + } + + cancelAll(reason: string, options: { sendResponse?: boolean } = {}): void { + for (const id of Array.from(this.pending.keys())) { + this.cancel(id, reason, options.sendResponse ?? true); + } + } + + private registerPending(request: PendingExtensionRequest): void { + if (this.tombstonedIds.has(request.id)) { + logger.debug(`[pi] Rejecting reused extension request id ${request.id}`); + this.options.session.sendSessionEvent({ type: 'message', message: `Pi extension request ${request.id} was canceled because its id was already retired.` }); + this.options.sendResponse({ type: 'extension_ui_response', id: request.id, cancelled: true }); + return; + } + if (this.pending.has(request.id)) { + // Pi requires unique ids for a pending RPC dialog. Reusing one could + // bind a delayed approval to a different dialog, so cancel the old + // request once and tombstone the id rather than replacing it. + this.cancel(request.id, 'Duplicate extension UI request id'); + this.tombstonedIds.add(request.id); + logger.debug(`[pi] Tombstoned duplicate extension request ${request.id}`); + this.options.session.sendSessionEvent({ type: 'message', message: `Pi extension request ${request.id} was canceled because its id was reused.` }); + return; + } + + const timeout = request.method === 'editor' ? undefined : getExtensionUiCleanupTimeout(request.timeout); + // Start this timer before publishing the HAPI request. Pi's deadline + // began before it emitted the extension_ui_request event. + const timer = timeout === undefined + ? null + : setTimeout(() => this.cancel(request.id, 'Extension UI request timed out'), timeout); + timer?.unref?.(); + this.pending.set(request.id, { request, timer }); + + const tool = requestToolName(request); + const argumentsValue = requestArguments(request); + this.options.session.sendAgentMessage({ + type: 'tool-call', + callId: request.id, + name: tool, + input: argumentsValue, + status: 'in_progress', + }); + this.options.session.updateAgentState((currentState) => ({ + ...currentState, + requests: { + ...currentState.requests, + [request.id]: { + tool, + arguments: argumentsValue, + createdAt: Date.now(), + }, + }, + } satisfies AgentState)); + } + + private handlePermissionResponse(response: PiPermissionResponse): void { + const entry = this.pending.get(response.id); + if (!entry) { + logger.debug(`[pi] Permission response for unknown extension request ${response.id}`); + return; + } + + const { request } = entry; + if (request.method === 'confirm') { + this.complete(request.id, response.approved + ? { type: 'extension_ui_response', id: request.id, confirmed: true } + : { type: 'extension_ui_response', id: request.id, confirmed: false }, response); + return; + } + + const answer = response.approved ? extractAnswer(response, request) : null; + if (answer === null || (request.method === 'select' && !request.options.includes(answer))) { + this.complete(request.id, { type: 'extension_ui_response', id: request.id, cancelled: true }, response); + return; + } + this.complete(request.id, { type: 'extension_ui_response', id: request.id, value: answer }, response); + } + + private cancel(id: string, reason: string, sendResponse = true): void { + const entry = this.pending.get(id); + if (!entry) return; + this.complete(id, { type: 'extension_ui_response', id, cancelled: true }, { + id, + approved: false, + decision: 'abort', + reason, + }, sendResponse); + } + + private complete(id: string, response: PiExtensionUiResponse, permission: PiPermissionResponse, sendResponse = true): void { + const entry = this.pending.get(id); + if (!entry) return; + this.pending.delete(id); + this.tombstonedIds.add(id); + if (entry.timer) clearTimeout(entry.timer); + + const tool = requestToolName(entry.request); + const approved = !('cancelled' in response) && (('confirmed' in response && response.confirmed) || 'value' in response); + // The Hub's deny endpoint normally omits an explicit decision. Match the + // existing permission adapters: approved=false means a user denial unless + // this handler itself marked the completion as an abort/cancellation. + const denied = !approved && permission.decision !== 'abort'; + const answers = normalizeAnswers(permission.answers); + this.options.session.updateAgentState((currentState) => { + const request = currentState.requests?.[id]; + const { [id]: _removed, ...remaining } = currentState.requests ?? {}; + return { + ...currentState, + requests: remaining, + completedRequests: { + ...currentState.completedRequests, + [id]: { + tool, + arguments: request?.arguments ?? requestArguments(entry.request), + createdAt: request?.createdAt ?? Date.now(), + completedAt: Date.now(), + status: approved ? 'approved' : denied ? 'denied' : 'canceled', + reason: permission.reason, + decision: permission.decision ?? (approved ? 'approved' : 'denied'), + ...(answers ? { answers } : {}), + }, + }, + } satisfies AgentState; + }); + this.options.session.sendAgentMessage({ + type: 'tool-call-result', + callId: id, + output: answers ?? response, + is_error: !approved, + }); + if (sendResponse) this.options.sendResponse(response); + } +} diff --git a/cli/src/pi/loop.test.ts b/cli/src/pi/loop.test.ts index 84d7b190..f6d1fe20 100644 --- a/cli/src/pi/loop.test.ts +++ b/cli/src/pi/loop.test.ts @@ -1,9 +1,11 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; -import { parsePiModels, parsePiCommands, parsePiContextUsage, sendPiRpcAndWait, wireTransportEvents } from './loop'; +import { parsePiModels, parsePiCommands, parsePiContextUsage, PiRpcTimeoutError, sendPiRpcAndWait, wireTransportEvents } from './loop'; import type { PiResponseEvent } from './types'; import { PiSession } from './session'; import { PiTransport } from './piTransport'; +import { PiConversationHistory } from './conversationHistory'; import type { PiThinkingLevel } from './types'; +import { PiAgentEventSchema } from './schemas'; // Mock logger vi.mock('@/ui/logger', () => ({ @@ -31,11 +33,12 @@ vi.mock('./piMessageAccumulator', () => { return { PiMessageAccumulator: class { handleEvent = vi.fn(() => []); + flush = vi.fn(() => []); }, }; }); -function createMockSession(): PiSession { +function createMockSession(model?: string): PiSession { return new PiSession({ api: {} as any, client: { @@ -44,12 +47,16 @@ function createMockSession(): PiSession { sendAgentMessage: vi.fn(), emitMessagesConsumed: vi.fn(), sendSessionEvent: vi.fn(), + updateAgentState: vi.fn(), emitSessionReady: vi.fn(), + getMetadata: vi.fn(() => null), + rpcHandlerManager: { registerHandler: vi.fn() }, } as any, path: '/tmp/test', logPath: '/tmp/test.log', startedBy: 'terminal', startingMode: 'local', + model, }); } @@ -217,6 +224,21 @@ describe('parsePiContextUsage', () => { }); }); +describe('Pi lifecycle event normalization', () => { + it('normalizes legacy auto_compaction aliases with lifecycle defaults at the transport boundary', () => { + expect(PiAgentEventSchema.parse({ type: 'auto_compaction_start' })).toMatchObject({ + type: 'compaction_start', + reason: 'threshold', + }); + expect(PiAgentEventSchema.parse({ type: 'auto_compaction_end' })).toMatchObject({ + type: 'compaction_end', + reason: 'threshold', + aborted: false, + willRetry: false, + }); + }); +}); + // --- wireTransportEvents (integration) --- describe('wireTransportEvents', () => { @@ -349,6 +371,8 @@ describe('wireTransportEvents', () => { emitMessagesConsumed: vi.fn(), sendSessionEvent: vi.fn(), emitSessionReady: vi.fn(), + getMetadata: vi.fn(() => null), + rpcHandlerManager: { registerHandler: vi.fn() }, } as any, path: '/tmp/test', logPath: '/tmp/test.log', @@ -384,6 +408,8 @@ describe('wireTransportEvents', () => { emitMessagesConsumed: vi.fn(), sendSessionEvent: vi.fn(), emitSessionReady: vi.fn(), + getMetadata: vi.fn(() => null), + rpcHandlerManager: { registerHandler: vi.fn() }, } as any, path: '/tmp/test', logPath: '/tmp/test.log', @@ -424,6 +450,8 @@ describe('wireTransportEvents', () => { emitMessagesConsumed: vi.fn(), sendSessionEvent: vi.fn(), emitSessionReady: vi.fn(), + getMetadata: vi.fn(() => null), + rpcHandlerManager: { registerHandler: vi.fn() }, } as any, path: '/tmp/test', logPath: '/tmp/test.log', @@ -510,6 +538,19 @@ describe('wireTransportEvents', () => { expect(session.client.emitMessagesConsumed).toHaveBeenCalledWith(['prompt-1'], undefined); }); + it('observes each agent lifecycle start without settling the prompt', () => { + const transport = createMockTransport(); + const onAgentLifecycleStarted = vi.fn(); + const onAgentSettled = vi.fn(); + wireTransportEvents(transport, session, [], { onAgentLifecycleStarted, onAgentSettled }); + + emitEvent({ type: 'agent_start' }); + emitEvent({ type: 'turn_start' }); + + expect(onAgentLifecycleStarted).toHaveBeenCalledTimes(2); + expect(onAgentSettled).not.toHaveBeenCalled(); + }); + it('publishes authoritative context usage after turn_end stats resolve', async () => { const transport = createMockTransport(); wireTransportEvents(transport, session, []); @@ -523,7 +564,7 @@ describe('wireTransportEvents', () => { }, }); - expect(session.piIsStreaming).toBe(false); + expect(session.piIsStreaming).toBe(true); expect(session.client.sendAgentMessage).not.toHaveBeenCalled(); const command = getSentCommand(transport); expect(command).toMatchObject({ type: 'get_session_stats' }); @@ -650,13 +691,15 @@ describe('wireTransportEvents', () => { expect(session.client.sendAgentMessage).toHaveBeenCalledTimes(1); }); - it('handles agent_end — stops streaming', () => { + it('handles agent_settled — stops streaming after an agent_end grace window', () => { const transport = createMockTransport(); wireTransportEvents(transport, session, []); session.piIsStreaming = true; + emitEvent({ type: 'agent_start' }); emitEvent({ type: 'agent_end' }); - + expect(session.piIsStreaming).toBe(true); + emitEvent({ type: 'agent_settled' }); expect(session.piIsStreaming).toBe(false); }); @@ -682,6 +725,38 @@ describe('wireTransportEvents', () => { ]); }); + it('fails closed and poisons the mutation lease when the detached startup model times out', async () => { + vi.useFakeTimers(); + try { + session = createMockSession('startup-model'); + const transport = createMockTransport(); + const onStartupFailure = vi.fn(); + wireTransportEvents(transport, session, [], { onStartupFailure }); + + emitEvent({ + type: 'response', + command: 'get_available_models', + success: true, + data: { models: [{ id: 'startup-model', provider: 'provider' }] }, + }); + await vi.advanceTimersByTimeAsync(0); + expect(transport.send).toHaveBeenCalledWith(expect.objectContaining({ + type: 'set_model', provider: 'provider', modelId: 'startup-model', + })); + + await vi.advanceTimersByTimeAsync(10_000); + await vi.waitFor(() => expect(onStartupFailure).toHaveBeenCalledWith(expect.objectContaining({ + message: expect.stringContaining('startup model outcome is indeterminate'), + }))); + let secondMutationStarted = false; + void session.runRuntimeMutation(async () => { secondMutationStarted = true; }); + await vi.advanceTimersByTimeAsync(0); + expect(secondMutationStarted).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + it('handles get_commands response — caches commands', () => { const transport = createMockTransport(); wireTransportEvents(transport, session, []); @@ -808,6 +883,35 @@ describe('sendPiRpcAndWait', () => { await expect(promise).resolves.toBeUndefined(); }); + it('steer response resolves its matching RPC without changing the main thinking state', async () => { + const handlers = new Map void>(); + const { transport, reply } = recordingTransport(handlers); + const session = createMockSession(); + session.updateThinkingState(true); + wireTransportEvents(transport, session, []); + + const promise = sendPiRpcAndWait(session, transport, { + type: 'steer', message: 'redirect current work', images: [], + }, 10_000); + reply({ command: 'steer', success: true }); + + await expect(promise).resolves.toBeUndefined(); + expect(session.piIsStreaming).toBe(true); + }); + + it('does not double-report a matching steer failure outside its dispatcher', async () => { + const handlers = new Map void>(); + const { transport, reply } = recordingTransport(handlers); + const session = createMockSession(); + wireTransportEvents(transport, session, []); + + const promise = sendPiRpcAndWait(session, transport, { type: 'steer', message: 'reject me' }, 10_000); + reply({ command: 'steer', success: false, error: 'native steer rejected' }); + + await expect(promise).rejects.toThrow('native steer rejected'); + expect(session.client.sendSessionEvent).not.toHaveBeenCalled(); + }); + it('get_available_models response resolves the awaited promise before timeout', async () => { const handlers = new Map void>(); const { transport, reply } = recordingTransport(handlers); @@ -838,14 +942,428 @@ describe('sendPiRpcAndWait', () => { await expect(promise).rejects.toThrow('Unknown provider: bad'); }); - it('rejects with timeout when Pi never responds', async () => { + it('rejects with a typed timeout when Pi never responds', async () => { const handlers = new Map void>(); const { transport } = recordingTransport(handlers); const session = createMockSession(); wireTransportEvents(transport, session, []); // No reply emitted -> must time out (guards against hangs). - await expect(sendPiRpcAndWait(session, transport, { type: 'test' }, 100)) - .rejects.toThrow('timed out'); + const pending = sendPiRpcAndWait(session, transport, { type: 'test' }, 100); + await expect(pending).rejects.toBeInstanceOf(PiRpcTimeoutError); + await expect(pending) + .rejects.toMatchObject({ + name: 'PiRpcTimeoutError', + command: 'test', + requestId: 1, + timeoutMs: 100, + message: 'Pi RPC test (id=1) timed out after 100ms', + }); + }); +}); + +describe('Pi lifecycle timeline', () => { + it('synchronizes authoritative get_state streaming and deduplicates compaction/retry timeline events', () => { + let listener: ((event: Record) => void) | null = null; + const transport = { + onEvent: vi.fn((handler: (event: Record) => void) => { listener = handler; }), + send: vi.fn(), + } as unknown as PiTransport; + const stateSession = createMockSession(); + wireTransportEvents(transport, stateSession, []); + const emit = (event: Record) => listener?.(event); + + emit({ type: 'response', command: 'get_state', success: true, data: { isStreaming: true } }); + expect(stateSession.piIsStreaming).toBe(true); + + emit({ type: 'compaction_start', reason: 'threshold' }); + emit({ type: 'compaction_start', reason: 'threshold' }); + emit({ type: 'compaction_end', reason: 'threshold', aborted: false, willRetry: false }); + emit({ type: 'auto_retry_start', attempt: 1, maxAttempts: 3, delayMs: 10, errorMessage: '429' }); + emit({ type: 'auto_retry_start', attempt: 1, maxAttempts: 3, delayMs: 10, errorMessage: '429' }); + emit({ type: 'auto_retry_end', attempt: 1, success: true }); + // A later compaction/retry episode with the same reason and attempt must + // remain visible; dedupe applies only while its current episode is open. + emit({ type: 'compaction_start', reason: 'threshold' }); + emit({ type: 'compaction_end', reason: 'threshold', aborted: false, willRetry: false }); + emit({ type: 'auto_retry_start', attempt: 1, maxAttempts: 3, delayMs: 10, errorMessage: '429' }); + emit({ type: 'auto_retry_end', attempt: 1, success: true }); + expect(stateSession.client.sendSessionEvent).toHaveBeenCalledWith({ type: 'message', message: '📦 Compaction started' }); + expect(stateSession.client.sendSessionEvent).toHaveBeenCalledWith({ type: 'message', message: '📦 Compaction completed' }); + expect(stateSession.client.sendSessionEvent).toHaveBeenCalledTimes(8); + + // Nested maintenance remains active until each own terminal event. + emit({ type: 'compaction_start', reason: 'manual' }); + emit({ type: 'summarization_retry_scheduled', attempt: 1, maxAttempts: 2, delayMs: 1, errorMessage: 'x' }); + emit({ type: 'summarization_retry_finished' }); + emit({ type: 'agent_end', willRetry: false }); + expect(stateSession.piIsStreaming).toBe(true); + emit({ type: 'compaction_end', reason: 'manual', aborted: false, willRetry: false }); + }); +}); + +describe('Pi prompt-settlement boundaries', () => { + it('does not release the local FIFO between tool-loop turns; only agent_end settles a prompt', () => { + let listener: ((event: Record) => void) | null = null; + const transport = { + onEvent: vi.fn((handler: (event: Record) => void) => { listener = handler; }), + send: vi.fn(), + } as unknown as PiTransport; + const onAgentSettled = vi.fn(); + const stateSession = createMockSession(); + wireTransportEvents(transport, stateSession, [], { onAgentSettled }); + + listener!({ type: 'agent_start' }); + listener!({ type: 'turn_start' }); + listener!({ type: 'turn_end', message: {} }); + listener!({ type: 'turn_start' }); + listener!({ type: 'turn_end', message: {} }); + expect(stateSession.piIsStreaming).toBe(true); + expect(onAgentSettled).not.toHaveBeenCalled(); + + listener!({ type: 'agent_end', willRetry: true }); + expect(stateSession.piIsStreaming).toBe(true); + expect(onAgentSettled).not.toHaveBeenCalled(); + + listener!({ type: 'agent_end', willRetry: false }); + expect(stateSession.piIsStreaming).toBe(true); + expect(onAgentSettled).not.toHaveBeenCalled(); + + listener!({ type: 'compaction_start', reason: 'threshold' }); + listener!({ type: 'agent_settled' }); + expect(stateSession.piIsStreaming).toBe(false); + expect(onAgentSettled).toHaveBeenCalledTimes(1); + }); +}); + + +describe('Pi settlement compatibility fallbacks', () => { + function setup() { + let listener: ((event: Record) => void) | null = null; + const transport = { + onEvent: vi.fn((handler: (event: Record) => void) => { listener = handler; }), + send: vi.fn(), + } as unknown as PiTransport; + const stateSession = createMockSession(); + const onAgentSettled = vi.fn(); + const onPromptLifecycleMissing = vi.fn(); + const onPromptRejected = vi.fn(); + const controller = wireTransportEvents(transport, stateSession, [], { onAgentSettled, onPromptLifecycleMissing, onPromptRejected }); + return { emit: (event: Record) => listener?.(event), stateSession, onAgentSettled, onPromptLifecycleMissing, onPromptRejected, controller, transport }; + } + + it('waits for agent_settled through maintenance and uses grace only for legacy Pi', async () => { + vi.useFakeTimers(); + const h = setup(); + h.stateSession.piIsStreaming = true; + h.emit({ type: 'agent_start' }); + h.emit({ type: 'agent_end', willRetry: false }); + h.emit({ type: 'compaction_start', reason: 'threshold' }); + await vi.advanceTimersByTimeAsync(600); + expect(h.onAgentSettled).not.toHaveBeenCalled(); + h.emit({ type: 'agent_settled' }); + expect(h.onAgentSettled).toHaveBeenCalledTimes(1); + + const legacy = setup(); + legacy.emit({ type: 'agent_end', willRetry: false }); + await vi.advanceTimersByTimeAsync(500); + expect(legacy.onAgentSettled).toHaveBeenCalledTimes(1); + }); + + it('releases command-only prompts only after their lifecycle grace, while agent_start cancels it', async () => { + vi.useFakeTimers(); + const h = setup(); + h.emit({ type: 'response', id: 'command-a', command: 'prompt', success: true }); + await vi.advanceTimersByTimeAsync(1_000); + expect(h.onPromptLifecycleMissing).toHaveBeenCalledTimes(1); + + const normal = setup(); + normal.emit({ type: 'response', command: 'prompt', success: true }); + normal.emit({ type: 'agent_start' }); + await vi.advanceTimersByTimeAsync(1_000); + expect(normal.onPromptLifecycleMissing).not.toHaveBeenCalled(); + }); + + it('syncs conversation history before reporting a missing prompt lifecycle', async () => { + vi.useFakeTimers(); + const order: string[] = []; + let listener: ((event: Record) => void) | null = null; + const transport = { + onEvent: vi.fn((handler: (event: Record) => void) => { listener = handler; }), + send: vi.fn(), + } as unknown as PiTransport; + const stateSession = createMockSession(); + const conversationHistory = { + syncEntries: vi.fn(async () => { order.push('sync'); }), + } as unknown as PiConversationHistory; + const onPromptLifecycleMissing = vi.fn(() => { order.push('missing'); }); + const controller = wireTransportEvents(transport, stateSession, ['command-local'], { + conversationHistory, + onPromptLifecycleMissing, + }); + + controller.beginPromptLifecycle('command-a'); + listener!({ type: 'response', id: 'command-a', command: 'prompt', success: true }); + await vi.advanceTimersByTimeAsync(1_000); + + expect(conversationHistory.syncEntries).toHaveBeenCalledTimes(1); + expect(onPromptLifecycleMissing).toHaveBeenCalledWith('command-local'); + expect(order).toEqual(['sync', 'missing']); + }); + + it('fails closed without retiring the prompt when command-only history sync fails', async () => { + vi.useFakeTimers(); + const pendingLocalIds = ['command-local']; + let listener: ((event: Record) => void) | null = null; + const transport = { + onEvent: vi.fn((handler: (event: Record) => void) => { listener = handler; }), + send: vi.fn(), + } as unknown as PiTransport; + const stateSession = createMockSession(); + const conversationHistory = { + syncEntries: vi.fn(async () => { throw new Error('temporary get_entries failure'); }), + } as unknown as PiConversationHistory; + const onPromptLifecycleMissing = vi.fn(); + const onStartupFailure = vi.fn(); + const controller = wireTransportEvents(transport, stateSession, pendingLocalIds, { + conversationHistory, + onPromptLifecycleMissing, + onStartupFailure, + }); + + controller.beginPromptLifecycle('command-a'); + listener!({ type: 'response', id: 'command-a', command: 'prompt', success: true }); + await vi.advanceTimersByTimeAsync(1_000); + + expect(onStartupFailure).toHaveBeenCalledWith(expect.objectContaining({ + message: 'Pi command-only history sync failed: temporary get_entries failure', + })); + expect(onPromptLifecycleMissing).not.toHaveBeenCalled(); + expect(pendingLocalIds).toEqual(['command-local']); + }); + + it('uses a fresh generation for consecutive command-only prompts', async () => { + vi.useFakeTimers(); + const h = setup(); + h.controller.beginPromptLifecycle('command-a'); + h.emit({ type: 'response', id: 'command-a', command: 'prompt', success: true }); + await vi.advanceTimersByTimeAsync(1_000); + expect(h.onPromptLifecycleMissing).toHaveBeenCalledTimes(1); + + h.controller.beginPromptLifecycle('command-b'); + h.emit({ type: 'response', id: 'command-b', command: 'prompt', success: true }); + await vi.advanceTimersByTimeAsync(1_000); + expect(h.onPromptLifecycleMissing).toHaveBeenCalledTimes(2); + + // A late settled event belongs to no current agent lifecycle and cannot + // cause a third settlement for the command-only generation. + h.emit({ type: 'agent_settled' }); + expect(h.onAgentSettled).not.toHaveBeenCalled(); + }); + + it('rejects a matching prompt after turn_start already consumed its local ID', async () => { + vi.useFakeTimers(); + const pendingLocalIds = ['local-a']; + let listener: ((event: Record) => void) | null = null; + const transport = { + onEvent: vi.fn((handler: (event: Record) => void) => { listener = handler; }), + send: vi.fn(), + } as unknown as PiTransport; + const stateSession = createMockSession(); + const onPromptRejected = vi.fn(); + const onPromptLifecycleMissing = vi.fn(); + const controller = wireTransportEvents(transport, stateSession, pendingLocalIds, { + onPromptRejected, + onPromptLifecycleMissing, + }); + + controller.beginPromptLifecycle('prompt-a'); + listener!({ type: 'agent_start' }); + listener!({ type: 'turn_start' }); + expect(pendingLocalIds).toEqual([]); + + listener!({ type: 'response', id: 'prompt-a', command: 'prompt', success: false, error: 'rejected' }); + expect(onPromptRejected).toHaveBeenCalledTimes(1); + expect(onPromptRejected).toHaveBeenCalledWith('local-a'); + expect(stateSession.piIsStreaming).toBe(false); + + // The failed generation owns no delayed lifecycle work. It must not + // later report lifecycle-missing or reject again. + await vi.advanceTimersByTimeAsync(1_000); + expect(onPromptLifecycleMissing).not.toHaveBeenCalled(); + listener!({ type: 'response', id: 'prompt-a', command: 'prompt', success: false, error: 'duplicate' }); + expect(onPromptRejected).toHaveBeenCalledTimes(1); + + controller.beginPromptLifecycle('prompt-b'); + stateSession.updateThinkingState(true); + listener!({ type: 'response', id: 'prompt-a', command: 'prompt', success: false, error: 'stale' }); + expect(stateSession.piIsStreaming).toBe(true); + expect(onPromptRejected).toHaveBeenCalledTimes(1); + }); + + it('blocks legacy auto compaction until it ends before using the legacy settlement grace', async () => { + vi.useFakeTimers(); + const h = setup(); + h.stateSession.piIsStreaming = true; + h.emit({ type: 'agent_start' }); + h.emit({ type: 'agent_end', willRetry: false }); + h.emit({ type: 'auto_compaction_start', reason: 'threshold' }); + + await vi.advanceTimersByTimeAsync(600); + expect(h.onAgentSettled).not.toHaveBeenCalled(); + expect(h.stateSession.client.sendSessionEvent).toHaveBeenCalledWith({ type: 'message', message: '📦 Compaction started' }); + + h.emit({ type: 'auto_compaction_end', reason: 'threshold', aborted: false, willRetry: false }); + await vi.advanceTimersByTimeAsync(499); + expect(h.onAgentSettled).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(1); + expect(h.onAgentSettled).toHaveBeenCalledTimes(1); + expect(h.stateSession.client.sendSessionEvent).toHaveBeenCalledWith({ type: 'message', message: '📦 Compaction completed' }); + }); + + it('keeps the settlement gate closed across a compaction retry', async () => { + vi.useFakeTimers(); + const h = setup(); + h.stateSession.piIsStreaming = true; + h.emit({ type: 'agent_start' }); + h.emit({ type: 'agent_end', willRetry: false }); + h.emit({ type: 'compaction_start', reason: 'threshold' }); + h.emit({ type: 'compaction_end', reason: 'threshold', aborted: false, willRetry: true }); + + await vi.advanceTimersByTimeAsync(600); + expect(h.onAgentSettled).not.toHaveBeenCalled(); + + h.emit({ type: 'agent_start' }); + h.emit({ type: 'agent_end', willRetry: false }); + await vi.advanceTimersByTimeAsync(499); + expect(h.onAgentSettled).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(1); + expect(h.onAgentSettled).toHaveBeenCalledTimes(1); + }); + + it('uses a bounded fallback when a compaction retry never starts', async () => { + vi.useFakeTimers(); + const h = setup(); + h.stateSession.piIsStreaming = true; + h.emit({ type: 'agent_start' }); + h.emit({ type: 'agent_end', willRetry: false }); + h.emit({ type: 'compaction_start', reason: 'threshold' }); + h.emit({ type: 'compaction_end', reason: 'threshold', aborted: false, willRetry: true }); + + await vi.advanceTimersByTimeAsync(1_499); + expect(h.onAgentSettled).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(1); + expect(h.onAgentSettled).toHaveBeenCalledTimes(1); + }); + + it('rejects all pending RPCs immediately during transport termination', async () => { + const h = setup(); + const pending = sendPiRpcAndWait(h.stateSession, h.transport, { type: 'abort' }, 10_000); + h.controller.terminatePendingRpc(new Error('transport closed')); + h.controller.terminatePendingRpc(new Error('transport closed again')); + await expect(pending).rejects.toThrow('transport closed'); + await expect(sendPiRpcAndWait(h.stateSession, h.transport, { type: 'abort' }, 10_000)).rejects.toThrow('transport closed'); + expect(h.transport.send).toHaveBeenCalledTimes(1); + }); +}); + +describe('Pi abort UI lifecycle', () => { + it('cancels pending extension input and ignores late prompt lifecycle events after abort', async () => { + let listener: ((event: Record) => void) | null = null; + const transport = { + onEvent: vi.fn((handler: (event: Record) => void) => { listener = handler; }), + send: vi.fn(), + } as unknown as PiTransport; + const session = createMockSession(); + const onPromptLifecycleMissing = vi.fn(); + const controller = wireTransportEvents(transport, session, [], { onPromptLifecycleMissing }); + controller.beginPromptLifecycle('prompt-a'); + listener!({ type: 'extension_ui_request', id: 'ui-a', method: 'input', title: 'Need input' }); + controller.cancelPendingExtensionUi('Pi prompt aborted', { sendResponse: true }); + controller.abortPromptLifecycle(); + + expect(transport.send).toHaveBeenCalledWith({ type: 'extension_ui_response', id: 'ui-a', cancelled: true }); + expect(session.client.updateAgentState).toHaveBeenCalled(); + listener!({ type: 'response', id: 'prompt-a', command: 'prompt', success: true }); + listener!({ type: 'agent_settled' }); + await Promise.resolve(); + expect(onPromptLifecycleMissing).not.toHaveBeenCalled(); + }); +}); + +describe('Pi conversation-history transport integration', () => { + function setup(expectedNativeSessionId?: string) { + let listener: ((event: Record) => void) | null = null; + const transport = { + onEvent: vi.fn((handler: (event: Record) => void) => { listener = handler; }), + send: vi.fn(), + } as unknown as PiTransport; + const session = new PiSession({ + api: {} as never, + client: { + keepAlive: vi.fn(), + updateMetadata: vi.fn(), + sendAgentMessage: vi.fn(), + emitMessagesConsumed: vi.fn(), + sendSessionEvent: vi.fn(), + updateAgentState: vi.fn(), + emitSessionReady: vi.fn(), + rpcHandlerManager: { registerHandler: vi.fn() }, + } as never, + path: '/tmp/test', + logPath: '/tmp/test.log', + startedBy: 'terminal', + startingMode: 'remote', + expectedNativeSessionId, + }); + return { + transport, + session, + emit: (event: Record) => listener?.(event), + }; + } + + it('resolves an awaited temporary get_state without publishing clone identity', async () => { + const h = setup('source-id'); + wireTransportEvents(h.transport, h.session, []); + const release = h.session.beginHistoryTransaction(); + const pending = sendPiRpcAndWait(h.session, h.transport, { type: 'get_state' }); + const command = (h.transport.send as ReturnType).mock.calls.at(-1)![0] as { id: string }; + + h.emit({ + type: 'response', + id: command.id, + command: 'get_state', + success: true, + data: { sessionId: 'clone-id', sessionFile: '/tmp/clone.jsonl' }, + }); + + await expect(pending).resolves.toEqual({ sessionId: 'clone-id', sessionFile: '/tmp/clone.jsonl' }); + expect(h.session.client.updateMetadata).not.toHaveBeenCalled(); + expect(h.session.client.emitSessionReady).not.toHaveBeenCalled(); + release(); + }); + + it('maps entry_appended events and completes the final sync before releasing the queue', async () => { + const h = setup(); + const rpc = vi.fn(async () => ({ entries: [], leafId: null })); + const history = new PiConversationHistory(h.session, rpc); + history.registerUserEntry('local-1'); + const onAgentSettled = vi.fn(); + const controller = wireTransportEvents(h.transport, h.session, [], { + conversationHistory: history, + onAgentSettled, + }); + + h.emit({ type: 'entry_appended', entry: { id: 'entry-1', type: 'message', message: { role: 'user' } } }); + expect(history.getEntryIds()).toEqual({ 'local-1': 'entry-1' }); + + controller.beginPromptLifecycle('prompt-1'); + h.emit({ type: 'response', id: 'prompt-1', command: 'prompt', success: true }); + h.emit({ type: 'agent_start' }); + h.emit({ type: 'agent_settled' }); + expect(onAgentSettled).not.toHaveBeenCalled(); + await vi.waitFor(() => expect(onAgentSettled).toHaveBeenCalledTimes(1)); + expect(rpc).toHaveBeenCalledWith({ type: 'get_entries', since: 'entry-1' }, undefined); }); }); diff --git a/cli/src/pi/loop.ts b/cli/src/pi/loop.ts index 82be9348..37fe62ba 100644 --- a/cli/src/pi/loop.ts +++ b/cli/src/pi/loop.ts @@ -3,31 +3,51 @@ import { convertAgentMessage } from '@/agent/messageConverter'; import { PiTransport } from './piTransport'; import { convertPiEvent, convertPiTurnUsage } from './piEventConverter'; import { PiMessageAccumulator } from './piMessageAccumulator'; -import { parsePiModels, parsePiCommands, parsePiContextUsage, PiResponseEventSchema, PiStateDataSchema, PiSetModelDataSchema } from './schemas'; +import { PiExtensionUiHandler } from './extensionUiHandler'; +import { parsePiModels, parsePiCommands, parsePiContextUsage, PiAgentEndEventSchema, PiAgentSettledEventSchema, PiExtensionUiRequestSchema, PiLifecycleEventSchema, PiResponseEventSchema, PiStateDataSchema, PiSetModelDataSchema } from './schemas'; import type { PiContextUsage, PiResponseEvent, PiRpcCommand, PiThinkingLevel, PiTurnEndEvent } from './types'; import type { PiSession } from './session'; +import type { PiConversationHistory } from './conversationHistory'; // --- Response parsers: re-exported from schemas.ts --- export { parsePiModels, parsePiCommands, parsePiContextUsage } from './schemas'; // --- Pending RPC resolver --- // Instance-scoped: created once by wireTransportEvents, stored on PiSession. +export class PiRpcTimeoutError extends Error { + readonly command: string; + readonly requestId: number; + readonly timeoutMs: number; + + constructor(command: string, requestId: number, timeoutMs: number) { + super(`Pi RPC ${command} (id=${requestId}) timed out after ${timeoutMs}ms`); + this.name = 'PiRpcTimeoutError'; + this.command = command; + this.requestId = requestId; + this.timeoutMs = timeoutMs; + } +} + export class PiRpcResolver { private idCounter = 0; + private terminalError: Error | null = null; private readonly pending = new Map void; reject: (error: Error) => void; + timer: ReturnType; }>(); sendAndWait(transport: PiTransport, command: Record, timeoutMs = 10_000): Promise { + if (this.terminalError) return Promise.reject(this.terminalError); const id = ++this.idCounter; return new Promise((resolve, reject) => { const timer = setTimeout(() => { this.pending.delete(id); - reject(new Error(`Pi RPC ${command.type} (id=${id}) timed out after ${timeoutMs}ms`)); + reject(new PiRpcTimeoutError(String(command.type), id, timeoutMs)); }, timeoutMs); this.pending.set(id, { + timer, resolve: (data) => { clearTimeout(timer); this.pending.delete(id); resolve(data); }, reject: (error) => { clearTimeout(timer); this.pending.delete(id); reject(error); }, }); @@ -36,6 +56,15 @@ export class PiRpcResolver { }); } + rejectAll(error: Error): void { + this.terminalError ??= error; + for (const [id, pending] of this.pending) { + clearTimeout(pending.timer); + this.pending.delete(id); + pending.reject(error); + } + } + resolveResponse(raw: unknown): void { const parsed = PiResponseEventSchema.safeParse(raw); if (!parsed.success) return; @@ -90,6 +119,7 @@ function applyGetState( sessionId?: string; thinkingLevel?: string; steeringMode?: 'all' | 'one-at-a-time'; + isStreaming?: boolean; }, session: PiSession, ): void { @@ -131,6 +161,12 @@ function applyGetState( session.currentSteeringMode = data.steeringMode; } + if (data.isStreaming !== undefined) { + // get_state is Pi's authoritative current-session snapshot. Synchronize + // only its explicit boolean, never infer state from unknown event types. + session.updateThinkingState(data.isStreaming); + } + } function handleResponse( @@ -139,7 +175,9 @@ function handleResponse( pendingLocalIds: string[], transport?: PiTransport, onStartupFailure?: (error: Error) => void, -): void { + conversationHistory?: PiConversationHistory, + onReady?: () => void, +): { rejectedPromptLocalId?: string } { const { command, success } = response; const resolver = session.rpcResolver!; @@ -150,12 +188,14 @@ function handleResponse( // get_session_stats is a best-effort compatibility probe. Older Pi // versions may reject it, so fall back silently instead of surfacing an // error event to the user on every completed turn. - if (command !== 'get_session_stats') { + if (command !== 'get_session_stats' && command !== 'steer') { session.sendSessionEvent({ type: 'message', message: error }); } if (command === 'prompt' && pendingLocalIds.length > 0) { const oldestLocalId = pendingLocalIds.shift()!; session.emitMessagesConsumed([oldestLocalId], { clearQueuedThinkingGrace: true }); + conversationHistory?.rejectPendingEntry(oldestLocalId); + return { rejectedPromptLocalId: oldestLocalId }; } // A failed initial get_state means Pi did not load its native session. // Do not leave the HAPI wrapper alive until the hub's ready timeout: the @@ -165,7 +205,7 @@ function handleResponse( if (command === 'get_state' && session.expectedNativeSessionId && !session.isNativeReady) { onStartupFailure?.(new Error(`Pi get_state failed: ${error}`)); } - return; + return {}; } switch (command) { @@ -178,13 +218,14 @@ function handleResponse( // mutating model/metadata state: an invalid resume must not publish a // colliding piSessionId that auto-dedup could merge. if (!parsed.success) { + resolvePendingRpc(resolver, response); if (session.expectedNativeSessionId) { onStartupFailure?.(new Error('Pi get_state returned malformed state data')); } break; } const state = parsed.data; - if (!session.matchesExpectedNativeSessionId(state.sessionId)) { + if (!session.isHistoryTransactionActive && !session.matchesExpectedNativeSessionId(state.sessionId)) { const actual = state.sessionId ? state.sessionId : '(missing)'; const error = `Pi loaded unexpected native session ${actual} instead of ${session.expectedNativeSessionId}`; logger.debug(`[pi] ${error}`); @@ -195,8 +236,15 @@ function handleResponse( // Emit ready before publishing Pi metadata. On native resume, this // ensures the hub can never merge based on a piSessionId before the // get_state identity check has completed. - session.markNativeReady(); - applyGetState(state, session); + // History transactions deliberately switch Pi through temporary + // clone/fork identities. Resolve their awaited get_state request, + // but never publish the temporary identity/model to the source row. + if (!session.isHistoryTransactionActive) { + session.markNativeReady(); + applyGetState(state, session); + onReady?.(); + } + resolvePendingRpc(resolver, response); break; } case 'set_model': { @@ -247,16 +295,22 @@ function handleResponse( if (match) { void (async () => { try { - await sendPiRpcAndWait(session, transport, { - type: 'set_model', - provider: match.provider, - modelId: match.modelId, - }); - session.currentModel = match.modelId; - session.currentProvider = match.provider; - persistSelectedPiModel(session); + await session.runRuntimeMutation(async () => { + await sendPiRpcAndWait(session, transport, { + type: 'set_model', + provider: match.provider, + modelId: match.modelId, + }); + session.currentModel = match.modelId; + session.currentProvider = match.provider; + persistSelectedPiModel(session); + }, { poisonOnError: (error) => error instanceof PiRpcTimeoutError }); logger.debug(`[pi] Startup model applied: ${match.provider}/${match.modelId}`); } catch (error) { + if (error instanceof PiRpcTimeoutError) { + onStartupFailure?.(new Error(`Pi startup model outcome is indeterminate: ${error.message}`)); + return; + } logger.debug(`[pi] Startup model set_model rejected, keeping Pi default: ${error instanceof Error ? error.message : String(error)}`); } })(); @@ -282,15 +336,22 @@ function handleResponse( break; case 'abort': logger.debug('[pi] Abort confirmed'); + resolvePendingRpc(resolver, response); break; case 'prompt': logger.debug('[pi] Prompt accepted'); break; + case 'steer': + logger.debug('[pi] Steer accepted'); + resolvePendingRpc(resolver, response); + break; default: logger.debug(`[pi] Response for ${command}`); resolvePendingRpc(resolver, response); break; } + + return {}; } const PI_CONTEXT_USAGE_RPC_TIMEOUT_MS = 1_000; @@ -332,67 +393,395 @@ async function publishPiTurnUsage( // --- Wire transport events to session --- +export type PiTransportEventController = { + flush: () => void; + cancelPendingExtensionUi: (reason: string, options?: { sendResponse?: boolean }) => void; + terminatePendingRpc: (error: Error) => void; + beginPromptLifecycle: (promptId: string) => void; + abortPromptLifecycle: () => void; +}; + +type PiTransportEventOptions = { + onStartupFailure?: (error: Error) => void; + onReady?: () => void; + /** Observes each Pi agent_start/turn_start without releasing the prompt queue. */ + onAgentLifecycleStarted?: () => void; + onAgentSettled?: () => void; + onPromptRejected?: (localId?: string) => void; + /** Return false to keep this prompt generation open for a late lifecycle. */ + onPromptLifecycleMissing?: (localId?: string) => void | boolean; + conversationHistory?: PiConversationHistory; +}; + +const PI_LEGACY_SETTLE_GRACE_MS = 500; +const PI_PROMPT_LIFECYCLE_GRACE_MS = 1_000; +const PI_COMPACTION_RETRY_START_GRACE_MS = 1_000; + +class PiLifecycleTimeline { + private compacting = false; + private activeRetryKey: string | null = null; + private summaryRetryActive = false; + + emit(event: unknown, session: PiSession): void { + const parsed = PiLifecycleEventSchema.safeParse(event); + if (!parsed.success) return; + const lifecycle = parsed.data; + let message: string | null = null; + + switch (lifecycle.type) { + case 'compaction_start': + if (this.compacting) return; + this.compacting = true; + message = '📦 Compaction started'; + break; + case 'compaction_end': + if (!this.compacting) return; + this.compacting = false; + message = lifecycle.aborted + ? '📦 Compaction canceled' + : lifecycle.errorMessage + ? `📦 Compaction failed: ${lifecycle.errorMessage}` + : lifecycle.willRetry + ? '📦 Compaction will retry' + : '📦 Compaction completed'; + break; + case 'auto_retry_start': { + const key = `${lifecycle.attempt}:${lifecycle.errorMessage}`; + if (this.activeRetryKey === key) return; + this.activeRetryKey = key; + message = `↻ Retrying after error (attempt ${lifecycle.attempt}/${lifecycle.maxAttempts}): ${lifecycle.errorMessage}`; + break; + } + case 'auto_retry_end': + if (this.activeRetryKey === null) return; + this.activeRetryKey = null; + message = lifecycle.success + ? `↻ Retry succeeded (attempt ${lifecycle.attempt})` + : `↻ Retry failed (attempt ${lifecycle.attempt})${lifecycle.finalError ? `: ${lifecycle.finalError}` : ''}`; + break; + case 'summarization_retry_scheduled': + if (this.summaryRetryActive) return; + this.summaryRetryActive = true; + message = `📝 Summary retry scheduled (attempt ${lifecycle.attempt}/${lifecycle.maxAttempts}): ${lifecycle.errorMessage}`; + break; + case 'summarization_retry_attempt_start': + if (!this.summaryRetryActive) return; + message = `📝 Summary retry started (${lifecycle.source})`; + break; + case 'summarization_retry_finished': + if (!this.summaryRetryActive) return; + this.summaryRetryActive = false; + message = '📝 Summary retry completed'; + break; + } + + if (message) session.sendSessionEvent({ type: 'message', message }); + } +} + export function wireTransportEvents( transport: PiTransport, session: PiSession, pendingLocalIds: string[], - options?: { onStartupFailure?: (error: Error) => void }, -): void { + options: PiTransportEventOptions = {}, +): PiTransportEventController { session.rpcResolver = new PiRpcResolver(); const assistantMessageAccumulator = new PiMessageAccumulator(); + const extensionUi = new PiExtensionUiHandler({ + session: session.client, + sendResponse: (response) => transport.send(response), + }); + const lifecycleTimeline = new PiLifecycleTimeline(); let latestContextUsageRequest = 0; + let deliveredSettlement = false; + let legacySettleTimer: ReturnType | null = null; + let promptLifecycleTimer: ReturnType | null = null; + let compactionRetryTimer: ReturnType | null = null; + const maintenanceActive = new Set<'compaction' | 'compactionRetry' | 'autoRetry' | 'summary'>(); + let agentEndObserved = false; + let agentLifecycleSeen = false; + let lifecycleGeneration = 0; + let activePromptId: string | null = null; + let activePromptResponseAccepted = false; + let activeAgentSettledSeen = false; + let promptLifecycleAborted = false; + // turn_start consumes the FIFO entry before the prompt response is known. + // Retain that exact local ID so a later matching response failure can reject + // the history registration even after pendingLocalIds has been drained. + let activePromptLocalId: string | undefined; + + const clearLegacySettleFallback = (): void => { + if (legacySettleTimer) clearTimeout(legacySettleTimer); + legacySettleTimer = null; + }; + const clearPromptLifecycleFallback = (): void => { + if (promptLifecycleTimer) clearTimeout(promptLifecycleTimer); + promptLifecycleTimer = null; + }; + const clearCompactionRetryPending = (): void => { + maintenanceActive.delete('compactionRetry'); + if (compactionRetryTimer) clearTimeout(compactionRetryTimer); + compactionRetryTimer = null; + }; + const beginPromptLifecycle = (promptId: string): void => { + lifecycleGeneration += 1; + activePromptId = promptId; + promptLifecycleAborted = false; + activePromptResponseAccepted = false; + activeAgentSettledSeen = false; + activePromptLocalId = undefined; + deliveredSettlement = false; + agentEndObserved = false; + agentLifecycleSeen = false; + maintenanceActive.clear(); + clearCompactionRetryPending(); + clearLegacySettleFallback(); + clearPromptLifecycleFallback(); + }; + const abortPromptLifecycle = (): void => { + lifecycleGeneration += 1; + activePromptId = null; + promptLifecycleAborted = true; + activePromptResponseAccepted = false; + activeAgentSettledSeen = false; + activePromptLocalId = undefined; + deliveredSettlement = false; + agentEndObserved = false; + agentLifecycleSeen = false; + maintenanceActive.clear(); + clearCompactionRetryPending(); + clearLegacySettleFallback(); + clearPromptLifecycleFallback(); + }; + const rejectPromptLifecycle = (): void => { + // A rejected prompt is terminal for this generation. Invalidate all + // delayed settlement work before notifying runPi so it can immediately + // resume FIFO pumping without a stale grace timer changing its state. + lifecycleGeneration += 1; + activePromptId = null; + promptLifecycleAborted = true; + activePromptResponseAccepted = false; + activeAgentSettledSeen = false; + activePromptLocalId = undefined; + deliveredSettlement = true; + agentEndObserved = false; + agentLifecycleSeen = false; + maintenanceActive.clear(); + clearCompactionRetryPending(); + latestContextUsageRequest += 1; + clearLegacySettleFallback(); + clearPromptLifecycleFallback(); + flushAccumulator(); + session.updateThinkingState(false); + }; + const deliverSettlement = (): void => { + if (deliveredSettlement || (activePromptId !== null && !activePromptResponseAccepted)) return; + deliveredSettlement = true; + clearCompactionRetryPending(); + clearLegacySettleFallback(); + clearPromptLifecycleFallback(); + session.updateThinkingState(false); + if (options.conversationHistory) { + void options.conversationHistory.syncEntries() + .catch(() => {}) + .finally(() => options.onAgentSettled?.()); + } else { + options.onAgentSettled?.(); + } + }; + const scheduleLegacySettleFallback = (): void => { + if ((activePromptId !== null && !activePromptResponseAccepted) || !agentEndObserved || maintenanceActive.size > 0 || deliveredSettlement || legacySettleTimer) return; + legacySettleTimer = setTimeout(() => { + legacySettleTimer = null; + deliverSettlement(); + }, PI_LEGACY_SETTLE_GRACE_MS); + legacySettleTimer.unref?.(); + }; + const schedulePromptLifecycleFallback = (): void => { + const generation = lifecycleGeneration; + clearPromptLifecycleFallback(); + promptLifecycleTimer = setTimeout(() => { + promptLifecycleTimer = null; + void (async () => { + if (generation !== lifecycleGeneration || deliveredSettlement || agentLifecycleSeen) return; + // Some Pi integrations omit entry_appended even for successful + // command-only prompts. Read the append log before runPi retires + // the pending localId, otherwise the next user entry can inherit + // this prompt's native history association. + if (options.conversationHistory) { + try { + await options.conversationHistory.syncEntries(); + } catch (error) { + if (generation !== lifecycleGeneration || deliveredSettlement || agentLifecycleSeen) return; + const detail = error instanceof Error ? error.message : String(error); + // Continuing would discard the only FIFO association for + // an unread native user entry. Fail the wrapper closed + // instead of allowing all later fork/rewind points to + // shift onto the wrong HAPI messages. + options.onStartupFailure?.(new Error(`Pi command-only history sync failed: ${detail}`)); + return; + } + } + if (generation !== lifecycleGeneration || deliveredSettlement || agentLifecycleSeen) return; + const handled = options.onPromptLifecycleMissing?.(pendingLocalIds[0]); + // The callback may synchronously pump the next queued prompt, + // whose beginPromptLifecycle() advances the generation and + // resets state. Never stamp the old timer's settlement onto it. + if (handled === false || generation !== lifecycleGeneration) return; + deliveredSettlement = true; + session.updateThinkingState(false); + })(); + }, PI_PROMPT_LIFECYCLE_GRACE_MS); + promptLifecycleTimer.unref?.(); + }; + + const sendMessages = (messages: ReturnType): void => { + for (const message of messages) { + const converted = convertAgentMessage(message, session.currentModel); + if (converted) session.sendAgentMessage(converted); + } + }; + const flushAccumulator = (): void => sendMessages(assistantMessageAccumulator.flush()); transport.onEvent((event) => { - // Debug: log all event types to diagnose missing Pi output + // Legacy Pi emitted auto_compaction_*; normalize it before every + // lifecycle consumer so both the maintenance gate and timeline see the + // same current event names. PiTransport performs this for subprocess + // traffic too; retaining it here keeps direct/test transports aligned. + const parsedLifecycle = PiLifecycleEventSchema.safeParse(event); + if (parsedLifecycle.success) { + event = parsedLifecycle.data; + } if (event.type !== 'keep_alive') { logger.debug(`[pi][event] ${event.type}`); } if (event.type === 'response') { - handleResponse( - event as unknown as PiResponseEvent, - session, - pendingLocalIds, - transport, - options?.onStartupFailure, - ); + const parsed = PiResponseEventSchema.safeParse(event); + if (parsed.success) { + const isCurrentPrompt = parsed.data.command === 'prompt' + && !deliveredSettlement + && !activePromptResponseAccepted + && (activePromptId === null ? !promptLifecycleAborted : parsed.data.id === activePromptId); + if (parsed.data.command === 'prompt' && !isCurrentPrompt) { + logger.debug(`[pi] Ignoring stale prompt response id=${parsed.data.id ?? 'missing'}`); + return; + } + const responseOutcome = handleResponse( + parsed.data, + session, + pendingLocalIds, + transport, + options.onStartupFailure, + options.conversationHistory, + options.onReady, + ); + if (isCurrentPrompt && !parsed.data.success) { + const rejectedLocalId = responseOutcome.rejectedPromptLocalId ?? activePromptLocalId; + // A Pi 0.83 turn_start can consume the HAPI FIFO before Pi + // replies that prompt failed. handleResponse only has the + // still-pending list, so finish the exact consumed history + // entry here when the FIFO was already shifted. + if (responseOutcome.rejectedPromptLocalId === undefined && rejectedLocalId) { + options.conversationHistory?.rejectPendingEntry(rejectedLocalId); + } + rejectPromptLifecycle(); + options.onPromptRejected?.(rejectedLocalId); + return; + } + if (isCurrentPrompt && parsed.data.success) { + activePromptResponseAccepted = true; + if (activeAgentSettledSeen) { + deliverSettlement(); + } else if (agentLifecycleSeen) { + scheduleLegacySettleFallback(); + } else { + schedulePromptLifecycleFallback(); + } + } + } else { + logger.debug('[pi] Ignoring malformed RPC response'); + } return; } - // Accumulate text/thinking deltas into snapshots, flush on message_end - const accumulated = assistantMessageAccumulator.handleEvent(event); - if (accumulated.length > 0) { - for (const msg of accumulated) { - const converted = convertAgentMessage(msg, session.currentModel); - if (converted) session.sendAgentMessage(converted); - } + if (event.type === 'entry_appended') { + options.conversationHistory?.observeEntry((event as { entry?: unknown }).entry); } - // message_start/update/end handled by accumulator — skip converter + if (event.type === 'extension_ui_request') { + const parsed = PiExtensionUiRequestSchema.safeParse(event); + if (parsed.success) { + extensionUi.handle(parsed.data); + } else { + logger.debug('[pi] Ignoring malformed extension_ui_request'); + } + return; + } + + if (event.type === 'agent_start' || event.type === 'turn_start') { + clearCompactionRetryPending(); + agentLifecycleSeen = true; + clearLegacySettleFallback(); + clearPromptLifecycleFallback(); + options.onAgentLifecycleStarted?.(); + } + if (event.type === 'compaction_start') { + clearCompactionRetryPending(); + maintenanceActive.add('compaction'); + clearLegacySettleFallback(); + } else if (event.type === 'auto_retry_start') { + maintenanceActive.add('autoRetry'); + clearLegacySettleFallback(); + } else if (event.type === 'summarization_retry_scheduled') { + maintenanceActive.add('summary'); + clearLegacySettleFallback(); + } else if (event.type === 'compaction_end') { + maintenanceActive.delete('compaction'); + if ('willRetry' in event && event.willRetry === true) { + maintenanceActive.add('compactionRetry'); + clearLegacySettleFallback(); + if (compactionRetryTimer) clearTimeout(compactionRetryTimer); + compactionRetryTimer = setTimeout(() => { + compactionRetryTimer = null; + if (!maintenanceActive.delete('compactionRetry')) return; + scheduleLegacySettleFallback(); + }, PI_COMPACTION_RETRY_START_GRACE_MS); + compactionRetryTimer.unref?.(); + } + } else if (event.type === 'auto_retry_end') { + maintenanceActive.delete('autoRetry'); + } else if (event.type === 'summarization_retry_finished') { + maintenanceActive.delete('summary'); + } + lifecycleTimeline.emit(event, session); + sendMessages(assistantMessageAccumulator.handleEvent(event)); + if (event.type !== 'message_start' && event.type !== 'message_update' && event.type !== 'message_end') { const messages = convertPiEvent(event); - for (const msg of messages) { - const converted = convertAgentMessage(msg, session.currentModel); + for (const message of messages) { + const converted = convertAgentMessage(message, session.currentModel); if (converted) session.sendAgentMessage(converted); } } - // Keep-alive + streaming state tracking - // - // Pi emits agent_start and turn_start back-to-back for each prompt. - // Only turn_start marks "my prompt was accepted and a turn began", so - // the pending localId is drained there. Draining on both would pop the - // FIFO twice per prompt — once with the real id, then once with - // undefined — and ship a garbage localId to the hub. if (event.type === 'agent_start') { session.updateThinkingState(true); } else if (event.type === 'turn_start') { session.updateThinkingState(true); if (pendingLocalIds.length > 0) { const oldestLocalId = pendingLocalIds.shift()!; + activePromptLocalId = oldestLocalId; session.emitMessagesConsumed([oldestLocalId]); } + // Some Pi integrations omit entry_appended forwarding. Incremental + // get_entries is the durable fallback and still pairs only FIFO. + if (options.conversationHistory) { + void options.conversationHistory.syncEntries().catch(() => {}); + } } else if (event.type === 'turn_end') { - session.updateThinkingState(false); + // Pi emits turn_end for each LLM/tool-loop iteration. The enclosing + // user prompt remains active until agent_end, so keep both streaming + // state and the local FIFO blocked here. const requestVersion = ++latestContextUsageRequest; void publishPiTurnUsage( event as PiTurnEndEvent, @@ -401,7 +790,31 @@ export function wireTransportEvents( () => requestVersion === latestContextUsageRequest, ); } else if (event.type === 'agent_end') { - session.piIsStreaming = false; + const parsed = PiAgentEndEventSchema.safeParse(event); + // Pi can end one attempt before its built-in auto-retry starts. That + // is not a user-prompt settlement, so keep the HAPI FIFO blocked. + if (parsed.success && parsed.data.willRetry === true) return; + agentEndObserved = true; + scheduleLegacySettleFallback(); + } else if (event.type === 'agent_settled') { + // A command-only generation has no Pi agent lifecycle; a delayed + // settled event from an earlier prompt must not settle this one. + if (agentLifecycleSeen && PiAgentSettledEventSchema.safeParse(event).success) { + activeAgentSettledSeen = true; + deliverSettlement(); + } + } + + if (agentEndObserved && maintenanceActive.size === 0 && (event.type === 'compaction_end' || event.type === 'auto_retry_end' || event.type === 'summarization_retry_finished')) { + scheduleLegacySettleFallback(); } }); + + return { + flush: flushAccumulator, + cancelPendingExtensionUi: (reason, options) => extensionUi.cancelAll(reason, options), + terminatePendingRpc: (error) => session.rpcResolver?.rejectAll(error), + beginPromptLifecycle, + abortPromptLifecycle, + }; } diff --git a/cli/src/pi/piEventConverter.test.ts b/cli/src/pi/piEventConverter.test.ts index 17d7a461..78b360b3 100644 --- a/cli/src/pi/piEventConverter.test.ts +++ b/cli/src/pi/piEventConverter.test.ts @@ -73,6 +73,18 @@ describe('convertPiEvent', () => { }]); }); + it('maps tool execution progress onto the running tool call id', () => { + expect(convertPiEvent({ + type: 'tool_execution_update', + toolCallId: 'tc-1', + toolName: 'read_file', + args: { path: '/foo.ts' }, + partialResult: { linesRead: 10 }, + })).toEqual([{ + type: 'tool_call', id: 'tc-1', name: 'read_file', input: { path: '/foo.ts' }, status: 'in_progress', progress: { linesRead: 10 }, + }]); + }); + it('should convert tool_execution_end (success) to tool_result AgentMessage', () => { const result = convertPiEvent({ type: 'tool_execution_end', @@ -105,31 +117,13 @@ describe('convertPiEvent', () => { }]); }); - it('should handle tool_execution_end with missing result', () => { - const result = convertPiEvent({ - type: 'tool_execution_end', - toolCallId: 'tc-1', - toolName: 'read_file', - isError: false - } as any); - expect(result).toEqual([{ - type: 'tool_result', - id: 'tc-1', - output: undefined, - status: 'completed' - }]); - }); - - it('should handle tool_execution_end with missing toolCallId', () => { - const result = convertPiEvent({ - type: 'tool_execution_end', - toolName: 'read_file', - result: 'ok', - isError: false - } as any); - expect(result).toHaveLength(1); - expect(result[0].type).toBe('tool_result'); - expect((result[0] as any).id).toBeUndefined(); + it('drops malformed tool completion events instead of emitting an uncorrelated result', () => { + expect(convertPiEvent({ + type: 'tool_execution_end', toolCallId: 'tc-1', toolName: 'read_file', isError: false, + } as never)).toEqual([]); + expect(convertPiEvent({ + type: 'tool_execution_end', toolName: 'read_file', result: 'ok', isError: false, + } as never)).toEqual([]); }); it('should defer turn usage and convert only turn completion', () => { diff --git a/cli/src/pi/piEventConverter.ts b/cli/src/pi/piEventConverter.ts index f9168dc6..ddc31b2c 100644 --- a/cli/src/pi/piEventConverter.ts +++ b/cli/src/pi/piEventConverter.ts @@ -1,34 +1,23 @@ import { logger } from '@/ui/logger'; import type { AgentMessage } from '@/agent/types'; -import type { - PiAgentEvent, - PiToolExecutionStartEvent, - PiToolExecutionEndEvent, - PiTurnEndEvent, - PiContextUsage, - PiUsage -} from './types'; +import { + PiToolExecutionEndEventSchema, + PiToolExecutionStartEventSchema, + PiToolExecutionUpdateEventSchema, +} from './schemas'; +import type { PiAgentEvent, PiContextUsage, PiTurnEndEvent, PiUsage } from './types'; function hasMeaningfulUsage(usage: PiUsage | undefined): usage is PiUsage { - return usage !== undefined - && Number.isFinite(usage.totalTokens) - && usage.totalTokens > 0; + return usage !== undefined && Number.isFinite(usage.totalTokens) && usage.totalTokens > 0; } -/** - * Builds the turn usage update after Pi's session stats request settles. - * - * undefined stats fall back to the turn's totalTokens for older Pi versions. - * null means Pi explicitly reported an unknown context size, so the previous - * valid HAPI usage state is preserved by not publishing an update. - */ +/** Builds a turn usage update after Pi's session-stats request settles. */ export function convertPiTurnUsage( event: PiTurnEndEvent, contextUsage: PiContextUsage | null | undefined, ): AgentMessage | null { const usage = event.message?.usage; if (!hasMeaningfulUsage(usage) || contextUsage === null) return null; - return { type: 'usage', inputTokens: usage.input ?? 0, @@ -41,67 +30,59 @@ export function convertPiTurnUsage( }; } -/** - * Converts Pi AgentEvent to HAPI AgentMessage array. - * - * Pi events come from `pi --mode rpc` stdout as JSONL. - * Not all Pi events map to HAPI AgentMessages — response/ack events - * are handled directly by the runner, not by this converter. - */ +/** Converts validated Pi lifecycle events to HAPI chat messages. */ export function convertPiEvent(event: PiAgentEvent): AgentMessage[] { - try { - switch (event.type) { - case 'tool_execution_start': { - const e = event as PiToolExecutionStartEvent; - return [{ - type: 'tool_call', - id: e.toolCallId, - name: e.toolName, - input: e.args, - status: 'in_progress' - }]; - } - - case 'tool_execution_end': { - const e = event as PiToolExecutionEndEvent; - return [{ - type: 'tool_result', - id: e.toolCallId, - output: e.result, - status: e.isError ? 'failed' : 'completed' - }]; - } - - case 'turn_end': { - const e = event as PiTurnEndEvent; - return [{ - type: 'turn_complete', - stopReason: e.message?.stopReason ?? 'stop' - }]; - } - - // Lifecycle and other events — not converted to AgentMessage. - // message_start/update/end are handled by PiMessageAccumulator - // in loop.ts before this converter is called — they never reach here, - // but are listed for exhaustive matching. - case 'agent_start': - case 'agent_end': - case 'turn_start': - case 'message_start': - case 'message_update': - case 'message_end': - case 'tool_execution_update': - case 'extension_ui_request': - case 'keep_alive': - case 'response': - return []; - - default: - logger.debug(`[pi] Unknown event type: ${event.type}`); - return []; + switch (event.type) { + case 'tool_execution_start': { + const parsed = PiToolExecutionStartEventSchema.safeParse(event); + if (!parsed.success) return []; + return [{ + type: 'tool_call', + id: parsed.data.toolCallId, + name: parsed.data.toolName, + input: parsed.data.args, + status: 'in_progress', + }]; } - } catch (err) { - logger.debug(`[pi] convertPiEvent failed for type=${event.type}: ${err}`); - return []; + case 'tool_execution_update': { + const parsed = PiToolExecutionUpdateEventSchema.safeParse(event); + if (!parsed.success) return []; + return [{ + type: 'tool_call', + id: parsed.data.toolCallId, + name: parsed.data.toolName, + input: parsed.data.args, + status: 'in_progress', + progress: parsed.data.partialResult, + }]; + } + case 'tool_execution_end': { + const parsed = PiToolExecutionEndEventSchema.safeParse(event); + if (!parsed.success) return []; + return [{ + type: 'tool_result', + id: parsed.data.toolCallId, + output: parsed.data.result, + status: parsed.data.isError ? 'failed' : 'completed', + }]; + } + case 'turn_end': { + const turn = event as PiTurnEndEvent; + return [{ type: 'turn_complete', stopReason: turn.message?.stopReason ?? 'stop' }]; + } + case 'agent_start': + case 'agent_end': + case 'agent_settled': + case 'turn_start': + case 'message_start': + case 'message_update': + case 'message_end': + case 'extension_ui_request': + case 'keep_alive': + case 'response': + return []; + default: + logger.debug(`[pi] Unknown event type: ${event.type}`); + return []; } } diff --git a/cli/src/pi/piMessageAccumulator.test.ts b/cli/src/pi/piMessageAccumulator.test.ts index 1631599a..1f96db25 100644 --- a/cli/src/pi/piMessageAccumulator.test.ts +++ b/cli/src/pi/piMessageAccumulator.test.ts @@ -1,211 +1,99 @@ -import { describe, it, expect } from 'vitest'; +import { describe, expect, it } from 'vitest'; import { PiMessageAccumulator } from './piMessageAccumulator'; +const event = (type: string, extra: Record = {}) => ({ type, ...extra }); + describe('PiMessageAccumulator', () => { - function makeEvent(type: string, extra: Record = {}): any { - return { type, ...extra }; - } + it('streams throttled cumulative text and reasoning snapshots with separate stable ids', () => { + let now = 0; + const accumulator = new PiMessageAccumulator({ now: () => now, streamNonceFactory: () => 'nonce' }); + accumulator.handleEvent(event('turn_start')); + accumulator.handleEvent(event('message_start')); - it('returns empty for events that are not handled', () => { - const acc = new PiMessageAccumulator(); - expect(acc.handleEvent(makeEvent('agent_start'))).toEqual([]); - expect(acc.handleEvent(makeEvent('turn_start'))).toEqual([]); - expect(acc.handleEvent(makeEvent('turn_end'))).toEqual([]); - expect(acc.handleEvent(makeEvent('agent_end'))).toEqual([]); - }); + const first = accumulator.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'thinking_delta', delta: 'think ' }, + })); + expect(first).toEqual([{ + type: 'reasoning', text: 'think ', id: 'pi-nonce-turn-1-message-1-reasoning-0', live: true, + }]); - it('accumulates text deltas and flushes one text message on message_end', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - expect(acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'hello ' } - }))).toEqual([]); - expect(acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'world' } + now = 100; + expect(accumulator.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'text_delta', delta: 'answer' }, }))).toEqual([]); - const flushed = acc.handleEvent(makeEvent('message_end', { message: {} })); - expect(flushed).toEqual([ - { type: 'text', text: 'hello world' } - ]); + now = 250; + expect(accumulator.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'text_delta', delta: '!' }, + }))).toEqual([{ + type: 'text', text: 'answer!', id: 'pi-nonce-turn-1-message-1-text-0', streamSnapshot: true, live: true, + }]); }); - it('accumulates thinking deltas and flushes one reasoning message on message_end', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'thinking_delta', delta: 'let me ' } - })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'thinking_delta', delta: 'think...' } + it('flushes the final changed snapshot once at message_end and close/error boundaries', () => { + let now = 0; + const accumulator = new PiMessageAccumulator({ now: () => now, streamNonceFactory: () => 'nonce' }); + accumulator.handleEvent(event('message_start')); + accumulator.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'text_delta', delta: 'first' }, })); + now = 100; + expect(accumulator.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'text_delta', delta: ' second' }, + }))).toEqual([]); + expect(accumulator.handleEvent(event('message_end'))).toEqual([{ + type: 'text', text: 'first second', id: 'pi-nonce-turn-0-message-1-text-0', streamSnapshot: true, + }]); + expect(accumulator.handleEvent(event('message_end'))).toEqual([]); - const flushed = acc.handleEvent(makeEvent('message_end', { message: {} })); - expect(flushed).toEqual([ - { type: 'reasoning', text: 'let me think...', id: 'pi-stream' } - ]); + accumulator.handleEvent(event('message_start')); + accumulator.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'text_delta', delta: 'partial' }, + })); + now = 200; + accumulator.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'text_delta', delta: ' answer' }, + })); + expect(accumulator.flush()).toEqual([{ + type: 'text', text: 'partial answer', id: 'pi-nonce-turn-0-message-2-text-0', streamSnapshot: true, + }]); + expect(accumulator.flush()).toEqual([]); }); - it('flushes both reasoning and text in order on message_end', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'thinking_delta', delta: 'thinking' } + it('keeps multiple content indexes separate instead of concatenating blocks', () => { + let now = 0; + const accumulator = new PiMessageAccumulator({ now: () => now, streamNonceFactory: () => 'nonce' }); + accumulator.handleEvent(event('message_start')); + const first = accumulator.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'text_delta', contentIndex: 0, delta: 'alpha' }, })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'reply' } + expect(first).toHaveLength(1); + now = 300; + const second = accumulator.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'text_delta', contentIndex: 1, delta: 'beta' }, })); - - const flushed = acc.handleEvent(makeEvent('message_end', { message: {} })); - expect(flushed).toEqual([ - { type: 'reasoning', text: 'thinking', id: 'pi-stream' }, - { type: 'text', text: 'reply' } - ]); + expect(second).toEqual([{ + type: 'text', text: 'beta', id: 'pi-nonce-turn-0-message-1-text-1', streamSnapshot: true, live: true, + }]); }); - it('skips empty content on flush', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'only text' } - })); - - const flushed = acc.handleEvent(makeEvent('message_end', { message: {} })); - expect(flushed).toEqual([ - { type: 'text', text: 'only text' } - ]); - }); - - it('drops empty/missing deltas silently', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta' } - })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: '' } - })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'thinking_delta' } - })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'thinking_delta', delta: ' ' } - })); - const flushed = acc.handleEvent(makeEvent('message_end', { message: {} })); - expect(flushed).toEqual([ - { type: 'reasoning', text: ' ', id: 'pi-stream' } - ]); - }); - - it('uses contentIndex as streamId when provided', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'x', contentIndex: 2 } - })); - const flushed = acc.handleEvent(makeEvent('message_end', { message: {} })); - expect(flushed).toEqual([ - { type: 'text', text: 'x' } - ]); - }); - - it('updates streamId from later deltas', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'a' } - })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'b', contentIndex: 7 } - })); - const flushed = acc.handleEvent(makeEvent('message_end', { message: {} })); - expect(flushed).toEqual([ - { type: 'text', text: 'ab' } - ]); - }); - - it('resets state on the next message_start', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'first' } - })); - acc.handleEvent(makeEvent('message_end', { message: {} })); - - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'second' } - })); - const flushed = acc.handleEvent(makeEvent('message_end', { message: {} })); - expect(flushed).toEqual([ - { type: 'text', text: 'second' } - ]); - }); - - it('flushes on turn_end as a safety net (no message_end received)', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'incomplete' } - })); - // No message_end — older Pi builds, partial streams, etc. - const flushed = acc.handleEvent(makeEvent('turn_end', { - message: { usage: { input: 1, output: 2, cacheRead: 0, cacheWrite: 0, totalTokens: 3 } } - })); - expect(flushed).toEqual([ - { type: 'text', text: 'incomplete' } - ]); - }); - - it('does not flush on turn_end if no message_start was seen', () => { - const acc = new PiMessageAccumulator(); - const flushed = acc.handleEvent(makeEvent('turn_end', { message: {} })); - expect(flushed).toEqual([]); - }); - - it('does not flush twice on message_end', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'once' } - })); - expect(acc.handleEvent(makeEvent('message_end', { message: {} }))).toEqual([ - { type: 'text', text: 'once' } - ]); - // Second message_end with no content buffered — must be empty, - // not a duplicate. - expect(acc.handleEvent(makeEvent('message_end', { message: {} }))).toEqual([]); - }); - - it('ignores text_start / thinking_start / text_end / thinking_end (full snapshots cause duplicates)', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_start' } - })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'thinking_start' } - })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_end' } - })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'thinking_end' } - })); - acc.handleEvent(makeEvent('message_update', { - assistantMessageEvent: { type: 'text_delta', delta: 'real content' } - })); - const flushed = acc.handleEvent(makeEvent('message_end', { message: {} })); - expect(flushed).toEqual([ - { type: 'text', text: 'real content' } - ]); - }); - - it('handles message_update without assistantMessageEvent', () => { - const acc = new PiMessageAccumulator(); - acc.handleEvent(makeEvent('message_start', { message: {} })); - expect(() => acc.handleEvent(makeEvent('message_update'))).not.toThrow(); - const flushed = acc.handleEvent(makeEvent('message_end', { message: {} })); - expect(flushed).toEqual([]); + it('does not collide across accumulator instances after a session resume', () => { + const first = new PiMessageAccumulator({ streamNonceFactory: () => 'before-restart' }); + const second = new PiMessageAccumulator({ streamNonceFactory: () => 'after-restart' }); + for (const accumulator of [first, second]) { + accumulator.handleEvent(event('turn_start')); + accumulator.handleEvent(event('message_start')); + } + const firstMessage = first.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'text_delta', delta: 'one' }, + }))[0]!; + const secondMessage = second.handleEvent(event('message_update', { + assistantMessageEvent: { type: 'text_delta', delta: 'two' }, + }))[0]!; + expect(firstMessage.type).toBe('text'); + expect(secondMessage.type).toBe('text'); + if (firstMessage.type === 'text' && secondMessage.type === 'text') { + expect(firstMessage.id).not.toBe(secondMessage.id); + } }); }); diff --git a/cli/src/pi/piMessageAccumulator.ts b/cli/src/pi/piMessageAccumulator.ts index c8ffc189..9e76f4ae 100644 --- a/cli/src/pi/piMessageAccumulator.ts +++ b/cli/src/pi/piMessageAccumulator.ts @@ -1,96 +1,142 @@ -import type { AgentMessage } from '@/agent/types' -import type { PiAgentEvent, PiAssistantMessageEvent } from './types' -import { PiAssistantMessageEventSchema } from './schemas' +import { randomUUID } from 'node:crypto'; +import type { AgentMessage } from '@/agent/types'; +import { PiAssistantMessageEventSchema } from './schemas'; +import type { PiAgentEvent } from './types'; + +const DEFAULT_SNAPSHOT_INTERVAL_MS = 250; + +type PiMessageAccumulatorOptions = { + snapshotIntervalMs?: number; + now?: () => number; + streamNonceFactory?: () => string; +}; + +type Segment = { + text: string; + lastSnapshot: string; +}; /** - * Accumulates Pi assistant-message text/thinking deltas into a single - * snapshot, flushed on `message_end` (with a `turn_end` safety net). - * - * Without this, every delta would become a separate hub message, and - * the web's reducer would render the last delta as the whole reasoning - * block (the per-message content-array dedup by streamId would only - * see one snapshot) while stacking every text delta as a new agent-text - * block, producing a character-by-character column. - * - * Mirrors codex's `ReasoningProcessor`: accumulate deltas locally, - * emit one reasoning + one text message per assistant message. + * Turns Pi's delta stream into throttled, cumulative snapshots. Pi content + * blocks are independently indexed, so a message with multiple text or thinking + * blocks never concatenates unrelated blocks into one timeline row. Every + * accumulator instance has a nonce as a session resume/restart must not reuse a + * previous timeline stream id. */ export class PiMessageAccumulator { - private active = false - private text = '' - private reasoning = '' - private streamId = 'pi-stream' + private active = false; + private turnSequence = 0; + private messageSequence = 0; + private readonly streamNonce: string; + private readonly textSegments = new Map(); + private readonly reasoningSegments = new Map(); + private lastSnapshotAt: number | null = null; + private readonly snapshotIntervalMs: number; + private readonly now: () => number; + + constructor(options: PiMessageAccumulatorOptions = {}) { + this.snapshotIntervalMs = options.snapshotIntervalMs ?? DEFAULT_SNAPSHOT_INTERVAL_MS; + this.now = options.now ?? Date.now; + this.streamNonce = options.streamNonceFactory?.() ?? randomUUID(); + } - /** - * Apply a Pi event to the accumulator. - * - * @returns AgentMessages to forward to the hub, if this event - * represents a flush point (`message_end` or `turn_end` with - * pending content). Returns an empty array otherwise. - */ handleEvent(event: PiAgentEvent): AgentMessage[] { + if (event.type === 'turn_start') { + this.turnSequence += 1; + return []; + } + if (event.type === 'message_start') { - this.active = true - this.text = '' - this.reasoning = '' - this.streamId = 'pi-stream' - return [] + const pending = this.active ? this.flush() : []; + this.startMessage(); + return pending; } if (event.type === 'message_update') { - const updateEvent = event as { assistantMessageEvent?: PiAssistantMessageEvent } - const rawAme = updateEvent.assistantMessageEvent - if (!rawAme) return [] - const ameResult = PiAssistantMessageEventSchema.safeParse(rawAme) - if (!ameResult.success) return [] - const ame = ameResult.data - const streamId = ame.contentIndex?.toString() ?? 'pi-stream' - this.streamId = streamId - if (ame.type === 'text_delta' && ame.delta) { - this.text += ame.delta - } else if (ame.type === 'thinking_delta' && ame.delta) { - this.reasoning += ame.delta + if (!('assistantMessageEvent' in event)) return []; + const parsed = PiAssistantMessageEventSchema.safeParse(event.assistantMessageEvent); + if (!parsed.success || !this.active || !parsed.data.delta) return []; + + const index = parsed.data.contentIndex ?? 0; + if (parsed.data.type === 'text_delta') { + this.append(this.textSegments, index, parsed.data.delta); + } else if (parsed.data.type === 'thinking_delta') { + this.append(this.reasoningSegments, index, parsed.data.delta); + } else { + return []; } - // Other assistant message events (text_start/thinking_start/ - // text_end/thinking_end) carry the full partial state — we - // already have the deltas, so we ignore them. - return [] + return this.snapshotIfDue(); } - if (event.type === 'message_end') { - if (this.active) return this.flush() - return [] + if (event.type === 'message_end' || event.type === 'turn_end' || event.type === 'agent_end') { + return this.flush(); } - // Safety net: turn_end with pending content means the assistant - // message ended without a clean `message_end` (older Pi builds, - // partial streams, or a stream that crashed mid-flight). - if (event.type === 'turn_end' && this.active) { - return this.flush() - } - - return [] + return []; } - private flush(): AgentMessage[] { - const streamId = this.streamId - const reasoning = this.reasoning - const text = this.text - this.active = false - this.text = '' - this.reasoning = '' - this.streamId = 'pi-stream' + /** Explicit transport-close/error safety net. Idempotent after a flush. */ + flush(): AgentMessage[] { + if (!this.active) return []; + const snapshots = this.createSnapshots(false); + this.resetMessage(); + return snapshots; + } - const out: AgentMessage[] = [] - // Reasoning comes before text in the Pi event sequence, so emit - // in that order. Empty content is dropped so the web doesn't - // render empty bubbles. - if (reasoning) { - out.push({ type: 'reasoning', text: reasoning, id: streamId }) + private startMessage(): void { + this.active = true; + this.messageSequence += 1; + this.textSegments.clear(); + this.reasoningSegments.clear(); + this.lastSnapshotAt = null; + } + + private append(segments: Map, index: number, delta: string): void { + const segment = segments.get(index) ?? { text: '', lastSnapshot: '' }; + segment.text += delta; + segments.set(index, segment); + } + + private snapshotIfDue(): AgentMessage[] { + const now = this.now(); + if (this.lastSnapshotAt !== null && now - this.lastSnapshotAt < this.snapshotIntervalMs) { + return []; } - if (text) { - out.push({ type: 'text', text }) + this.lastSnapshotAt = now; + return this.createSnapshots(true); + } + + private createSnapshots(live: boolean): AgentMessage[] { + const messages: AgentMessage[] = []; + this.addSnapshots(messages, 'reasoning', this.reasoningSegments, live); + this.addSnapshots(messages, 'text', this.textSegments, live); + return messages; + } + + private addSnapshots( + output: AgentMessage[], + kind: 'reasoning' | 'text', + segments: Map, + live: boolean, + ): void { + for (const [index, segment] of Array.from(segments.entries()).sort(([left], [right]) => left - right)) { + if (!segment.text.trim() || segment.text === segment.lastSnapshot) continue; + segment.lastSnapshot = segment.text; + const id = `pi-${this.streamNonce}-turn-${this.turnSequence}-message-${this.messageSequence}-${kind}-${index}`; + output.push({ + type: kind, + text: segment.text, + id, + ...(kind === 'text' ? { streamSnapshot: true } : {}), + ...(live ? { live: true } : {}), + }); } - return out + } + + private resetMessage(): void { + this.active = false; + this.textSegments.clear(); + this.reasoningSegments.clear(); + this.lastSnapshotAt = null; } } diff --git a/cli/src/pi/piTransport.test.ts b/cli/src/pi/piTransport.test.ts index be9c8c0b..ea5772bd 100644 --- a/cli/src/pi/piTransport.test.ts +++ b/cli/src/pi/piTransport.test.ts @@ -185,6 +185,20 @@ describe('PiTransport', () => { }); describe('onClose()', () => { + it('does not report stdout end as a synthetic close before the real exit code', () => { + const transport = new PiTransport({ command: 'pi', args: ['--mode', 'rpc'], cwd: '/work' }); + transport.start(); + const closeHandler = vi.fn(); + transport.onClose(closeHandler); + + mockProcess.stdout.emit('end'); + expect(closeHandler).not.toHaveBeenCalled(); + mockProcess.emit('close', 7, null); + mockProcess.emit('close', 7, null); + expect(closeHandler).toHaveBeenCalledTimes(1); + expect(closeHandler).toHaveBeenCalledWith(7, null); + }); + it('should call handler when subprocess exits', () => { const transport = new PiTransport({ command: 'pi', args: ['--mode', 'rpc'], cwd: '/work' }); transport.start(); diff --git a/cli/src/pi/piTransport.ts b/cli/src/pi/piTransport.ts index 8d4a99ce..85a1fa3a 100644 --- a/cli/src/pi/piTransport.ts +++ b/cli/src/pi/piTransport.ts @@ -8,6 +8,7 @@ export interface PiTransportOptions { command: string; args: string[]; cwd: string; + env?: NodeJS.ProcessEnv; } export class PiTransport extends JsonLineParser { @@ -18,6 +19,7 @@ export class PiTransport extends JsonLineParser { private killed = false; private started = false; private exited = false; + private closeReported = false; private readonly options: PiTransportOptions; constructor(options: PiTransportOptions) { @@ -36,17 +38,17 @@ export class PiTransport extends JsonLineParser { this.process = spawn(this.options.command, this.options.args, { cwd: this.options.cwd, - stdio: ['pipe', 'pipe', 'pipe'] + stdio: ['pipe', 'pipe', 'pipe'], + env: this.options.env }) as ChildProcessWithoutNullStreams; this.process.stdout.setEncoding('utf8'); this.process.stdout.on('data', (chunk: string) => this.feed(chunk)); this.process.stdout.on('end', () => { - if (!this.exited && !this.killed) { - logger.debug('[pi] stdout ended before process close — treating as exit'); - this.exited = true; - this.closeHandler?.(null, null); - } + // stdout can end before ChildProcess emits `close`. Do not publish a + // synthetic close here: doing so loses the real exit code and used to + // invoke lifecycle cleanup twice. `close` below is the one authority. + logger.debug('[pi] stdout ended; awaiting process close for exit status'); }); this.process.stderr.setEncoding('utf8'); @@ -57,6 +59,8 @@ export class PiTransport extends JsonLineParser { this.process.on('close', (code, signal) => { logger.debug(`[pi] Process exited (code=${code}, signal=${signal})`); this.exited = true; + if (this.closeReported) return; + this.closeReported = true; this.closeHandler?.(code, signal); }); @@ -75,7 +79,7 @@ export class PiTransport extends JsonLineParser { } send(message: PiRpcCommand): void { - if (!this.process || this.killed) { + if (!this.process || this.killed || this.exited || this.process.stdin.destroyed || this.process.stdin.writableEnded) { logger.debug('[pi] Dropping message: transport not running'); return; } diff --git a/cli/src/pi/promptQueue.test.ts b/cli/src/pi/promptQueue.test.ts new file mode 100644 index 00000000..adb50a6a --- /dev/null +++ b/cli/src/pi/promptQueue.test.ts @@ -0,0 +1,24 @@ +import { describe, expect, it } from 'vitest'; +import { PiPromptQueue } from './promptQueue'; + +describe('PiPromptQueue', () => { + it('preserves FIFO and permits cancellation before a Pi turn starts', () => { + const queue = new PiPromptQueue(); + queue.enqueue({ message: 'first', images: [], outboundSequence: 1, localId: 'one' }); + queue.enqueue({ message: 'cancel', images: [], outboundSequence: 2, localId: 'two' }); + queue.enqueue({ message: 'third', images: [], outboundSequence: 3, localId: 'three' }); + expect(queue.cancelByLocalId('two')).toBe(true); + expect(queue.dequeue()?.message).toBe('first'); + expect(queue.dequeue()?.message).toBe('third'); + expect(queue.dequeue()).toBeUndefined(); + }); + + it('inserts a delayed steer fallback ahead of a later ordinary prompt', () => { + const queue = new PiPromptQueue(); + queue.enqueue({ message: 'later ordinary', images: [], outboundSequence: 2, localId: 'two' }); + queue.enqueue({ message: 'earlier steer fallback', images: [], outboundSequence: 1, localId: 'one' }); + + expect(queue.dequeue()?.message).toBe('earlier steer fallback'); + expect(queue.dequeue()?.message).toBe('later ordinary'); + }); +}); diff --git a/cli/src/pi/promptQueue.ts b/cli/src/pi/promptQueue.ts new file mode 100644 index 00000000..0addeefe --- /dev/null +++ b/cli/src/pi/promptQueue.ts @@ -0,0 +1,42 @@ +import type { PiImageContent } from './types'; + +export type PiPreparedPrompt = { + message: string; + images: PiImageContent[]; + /** Monotonic arrival reservation assigned before asynchronous preparation. */ + outboundSequence: number; + localId?: string; +}; + +/** + * Small cancellable FIFO: HAPI owns queueing, Pi receives only real turns. + * + * A native steer can asynchronously degrade to a prompt after a later ordinary + * prompt was prepared. Arrival reservations preserve original message order at + * that boundary instead of using completion order. + */ +export class PiPromptQueue { + private readonly entries: PiPreparedPrompt[] = []; + + enqueue(prompt: PiPreparedPrompt): void { + const index = this.entries.findIndex((entry) => entry.outboundSequence > prompt.outboundSequence); + if (index === -1) this.entries.push(prompt); + else this.entries.splice(index, 0, prompt); + } + + dequeue(): PiPreparedPrompt | undefined { + return this.entries.shift(); + } + + cancelByLocalId(localId: string): boolean { + if (!localId) return false; + const index = this.entries.findIndex((entry) => entry.localId === localId); + if (index === -1) return false; + this.entries.splice(index, 1); + return true; + } + + get size(): number { + return this.entries.length; + } +} diff --git a/cli/src/pi/runPi.test.ts b/cli/src/pi/runPi.test.ts index b45b1438..bc01d6dc 100644 --- a/cli/src/pi/runPi.test.ts +++ b/cli/src/pi/runPi.test.ts @@ -8,12 +8,20 @@ const harness = vi.hoisted(() => ({ sent: [] as unknown[], throwOnGetCommands: true, onError: null as ((error: Error) => void) | null, + onEvent: null as ((event: Record) => void) | null, + rpcHandlers: new Map Promise>(), killCount: 0, cleanupCount: 0, session: { + sessionId: 'hapi-session-test', keepAlive: vi.fn(), onUserMessage: vi.fn(), onCancelQueuedMessage: vi.fn(), + emitMessagesConsumed: vi.fn(), + sendSessionEvent: vi.fn(), + updateMetadata: vi.fn(), + getMetadata: vi.fn(() => null), + emitSessionReady: vi.fn(), rpcHandlerManager: { registerHandler: vi.fn() }, }, })); @@ -61,7 +69,9 @@ vi.mock('./piTransport', () => ({ onClose(): void {} - onEvent(): void {} + onEvent(callback: (event: Record) => void): void { + harness.onEvent = callback; + } start(): void {} @@ -78,9 +88,36 @@ vi.mock('./piTransport', () => ({ }, })); -import { buildPiCommandInventory, formatPiUserMessage, rewritePiSkillPrompt, runPi } from './runPi'; +import { buildPiCommandInventory, failPiHistoryOnRestoreError, formatPiUserMessage, rewritePiSkillPrompt, runPi } from './runPi'; import { bootstrapExistingSession } from '@/agent/sessionFactory'; import { PiSession } from './session'; +import { PiHistoryRestoreError } from './conversationHistory'; +import { RPC_METHODS } from '@hapi/protocol/rpcMethods'; + +async function replyToHistoryCommand(type: 'get_entries' | 'get_fork_messages', occurrence: number, data: unknown): Promise { + await vi.waitFor(() => { + expect(harness.sent.filter((item) => (item as { type?: string }).type === type)).toHaveLength(occurrence); + }); + const command = harness.sent.filter((item) => (item as { type?: string }).type === type)[occurrence - 1] as { id: string }; + harness.onEvent!({ type: 'response', id: command.id, command: type, success: true, data }); +} + +async function completeHistoryBaseline( + initialEntries: unknown[] = [], + initialLeafId: string | null = null, +): Promise { + await replyToHistoryCommand('get_entries', 1, { entries: initialEntries, leafId: initialLeafId }); +} + +async function completeHistoryProbe(initialLeafId: string | null = null): Promise { + await replyToHistoryCommand('get_fork_messages', 1, { messages: [] }); + await replyToHistoryCommand('get_entries', 2, { entries: [], leafId: initialLeafId }); +} + +async function completeHistoryInitialization(): Promise { + await completeHistoryBaseline(); + await completeHistoryProbe(); +} describe('Pi command namespaces', () => { const commands = [ @@ -115,7 +152,7 @@ describe('Pi command namespaces', () => { mimeType: 'text/plain', size: 5, path: '/tmp/query.txt', - }], commands)).toBe('/skill:brave-search\n\n@/tmp/query.txt'); + }], commands)).toBe('/skill:brave-search\n\nAttached file: \"/tmp/query.txt\"'); }); }); @@ -125,6 +162,17 @@ describe('runPi startup', () => { harness.sent.length = 0; harness.throwOnGetCommands = true; harness.onError = null; + harness.onEvent = null; + harness.rpcHandlers.clear(); + harness.session.rpcHandlerManager.registerHandler.mockReset(); + harness.session.rpcHandlerManager.registerHandler.mockImplementation((method: string, handler: (payload: unknown) => Promise) => { + harness.rpcHandlers.set(method, handler); + }); + harness.session.onUserMessage.mockReset(); + harness.session.onCancelQueuedMessage.mockReset(); + harness.session.emitMessagesConsumed.mockReset(); + harness.session.sendSessionEvent.mockReset(); + harness.session.updateMetadata.mockReset(); harness.killCount = 0; harness.cleanupCount = 0; vi.useRealTimers(); @@ -133,10 +181,11 @@ describe('runPi startup', () => { it('lets Pi create a fresh session when no resume ID is provided', async () => { await runPi({ workingDirectory: '/work' }); - expect(harness.transportOptions).toEqual({ + expect(harness.transportOptions).toMatchObject({ command: 'pi', args: ['--mode', 'rpc'], cwd: '/work', + env: { PI_RPC_EMIT_TITLE: '1' }, }); expect(harness.sent).toEqual([ { type: 'get_state' }, @@ -151,10 +200,11 @@ describe('runPi startup', () => { resumeSessionId: 'pi-session-123', }); - expect(harness.transportOptions).toEqual({ + expect(harness.transportOptions).toMatchObject({ command: 'pi', args: ['--mode', 'rpc', '--session', 'pi-session-123'], cwd: '/work', + env: { PI_RPC_EMIT_TITLE: '1' }, }); expect(harness.sent).toEqual([ { type: 'get_state' }, @@ -179,21 +229,812 @@ describe('runPi startup', () => { }); }); + it('registers native conversation fork and rewind RPC handlers', async () => { + await runPi({ workingDirectory: '/work' }); + + expect(harness.rpcHandlers.has(RPC_METHODS.ForkConversation)).toBe(true); + expect(harness.rpcHandlers.has(RPC_METHODS.RewindConversation)).toBe(true); + }); + + it('escalates only failed source restoration to lifecycle cleanup', () => { + const failNativeStartup = vi.fn(); + failPiHistoryOnRestoreError(new PiHistoryRestoreError('restore failed'), failNativeStartup); + failPiHistoryOnRestoreError(new Error('ordinary fork failure'), failNativeStartup); + + expect(failNativeStartup).toHaveBeenCalledTimes(1); + expect(failNativeStartup).toHaveBeenCalledWith(expect.any(PiHistoryRestoreError)); + }); + it.each([ - ['fresh', undefined, 1, 0], - ['resume', 'pi-session-1', 0, 1], - ] as const)('applies the startup fallback only to %s sessions', async (_label, resumeSessionId, expectedCalls, expectedKills) => { + ['fresh', undefined], + ['resume', 'pi-session-1'], + ] as const)('applies the startup fallback only to %s sessions', async (_label, resumeSessionId) => { vi.useFakeTimers(); harness.throwOnGetCommands = false; const markReady = vi.spyOn(PiSession.prototype, 'markReady'); const running = runPi({ workingDirectory: '/work', resumeSessionId }); await vi.advanceTimersByTimeAsync(31_000); - expect(markReady).toHaveBeenCalledTimes(expectedCalls); - expect(harness.cleanupCount).toBe(expectedKills); + if (resumeSessionId) { + expect(markReady).not.toHaveBeenCalled(); + expect(harness.cleanupCount).toBe(1); + } else { + expect(markReady).not.toHaveBeenCalled(); + await completeHistoryBaseline(); + await vi.advanceTimersByTimeAsync(0); + expect(markReady).toHaveBeenCalledTimes(1); + expect(harness.cleanupCount).toBe(0); + } harness.onError?.(new Error('stop test transport')); await running; markReady.mockRestore(); }); + + it('establishes the history baseline before a fresh-session fallback drains prompts', async () => { + vi.useFakeTimers(); + harness.throwOnGetCommands = false; + const running = runPi({ workingDirectory: '/work' }); + await Promise.resolve(); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as ( + message: { role: 'user'; content: { type: 'text'; text: string } }, + localId: string + ) => void; + onUserMessage({ role: 'user', content: { type: 'text', text: 'queued before ready' } }, 'fallback-id'); + await Promise.resolve(); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'queued before ready' })); + + await vi.advanceTimersByTimeAsync(31_000); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'queued before ready' })); + await completeHistoryBaseline([ + { id: 'old-native-user', type: 'message', message: { role: 'user' } }, + ], 'old-native-user'); + await vi.advanceTimersByTimeAsync(0); + expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'queued before ready' })); + let preNativeMetadata: Record = {}; + for (const [updater] of harness.session.updateMetadata.mock.calls) { + if (typeof updater === 'function') preNativeMetadata = updater(preNativeMetadata); + } + expect(preNativeMetadata).not.toMatchObject({ capabilities: { conversationHistory: expect.anything() } }); + + const prompt = harness.sent.find((item) => (item as { type?: string }).type === 'prompt') as { id: string }; + harness.onEvent!({ + type: 'response', command: 'get_state', success: true, + data: { sessionId: 'late-session', sessionFile: '/tmp/late-session.jsonl' }, + }); + await completeHistoryProbe('old-native-user'); + harness.onEvent!({ type: 'response', id: prompt.id, command: 'prompt', success: true }); + harness.onEvent!({ type: 'turn_start' }); + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')).toHaveLength(3)); + const incremental = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[2] as { id: string }; + harness.onEvent!({ + type: 'response', id: incremental.id, command: 'get_entries', success: true, + data: { entries: [{ id: 'new-native-user', type: 'message', message: { role: 'user' } }], leafId: 'new-native-user' }, + }); + await vi.advanceTimersByTimeAsync(0); + let metadata: Record = {}; + for (const [updater] of harness.session.updateMetadata.mock.calls) { + if (typeof updater === 'function') metadata = updater(metadata); + } + expect(metadata).toMatchObject({ + capabilities: { conversationHistory: expect.anything() }, + conversationHistoryEntryIds: { 'fallback-id': 'new-native-user' }, + }); + + harness.onError?.(new Error('stop test transport')); + await running; + }); + + it('does not drain a fallback prompt when cleanup races a late native-ready preparation', async () => { + vi.useFakeTimers(); + harness.throwOnGetCommands = false; + const markReady = vi.spyOn(PiSession.prototype, 'markReady'); + const running = runPi({ workingDirectory: '/work' }); + await vi.advanceTimersByTimeAsync(0); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as ( + message: { role: 'user'; content: { type: 'text'; text: string } }, + localId: string + ) => void; + onUserMessage({ role: 'user', content: { type: 'text', text: 'must not drain' } }, 'cleanup-race-id'); + await vi.advanceTimersByTimeAsync(31_000); + expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'get_entries' })); + + harness.onEvent!({ + type: 'response', command: 'get_state', success: true, + data: { sessionId: 'late-session', sessionFile: '/tmp/late-session.jsonl' }, + }); + harness.onError?.(new Error('transport failed during history baseline')); + await Promise.resolve(); + await Promise.resolve(); + await running; + + expect(markReady).not.toHaveBeenCalled(); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'must not drain' })); + markReady.mockRestore(); + }); +}); + + +describe('Pi abort queue boundary', () => { + beforeEach(() => { + vi.useRealTimers(); + harness.sent.length = 0; + harness.throwOnGetCommands = false; + harness.onEvent = null; + harness.rpcHandlers.clear(); + harness.session.onUserMessage.mockReset(); + harness.session.onCancelQueuedMessage.mockReset(); + harness.session.emitMessagesConsumed.mockReset(); + harness.session.sendSessionEvent.mockReset(); + harness.session.updateMetadata.mockReset(); + harness.cleanupCount = 0; + harness.killCount = 0; + harness.session.rpcHandlerManager.registerHandler.mockReset(); + harness.session.rpcHandlerManager.registerHandler.mockImplementation((method: string, handler: (payload: unknown) => Promise) => { + harness.rpcHandlers.set(method, handler); + }); + }); + + it('does not send an empty prompt when every image attachment fails', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: { + role: 'user'; content: { type: 'text'; text: string; attachments: Array<{ id: string; filename: string; mimeType: string; size: number; path: string }> }; + }, localId: string) => void; + onUserMessage({ + role: 'user', + content: { type: 'text', text: '', attachments: [{ id: 'bad', filename: 'missing.png', mimeType: 'image/png', size: 1, path: '/missing/image.png' }] }, + }, 'missing-image-id'); + await vi.waitFor(() => expect(harness.session.sendSessionEvent).toHaveBeenCalledWith(expect.objectContaining({ + type: 'message', message: expect.stringContaining('Could not attach image missing.png'), + }))); + expect(harness.session.emitMessagesConsumed).toHaveBeenCalledWith(['missing-image-id'], { clearQueuedThinkingGrace: true }); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt' })); + harness.onError?.(new Error('finish test')); + await running; + }); + + it('short-circuits a canceled queued preparation before attachment I/O', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + await completeHistoryInitialization(); + let pathReads = 0; + const attachment = { + id: 'cancel-image', filename: 'cancel.png', mimeType: 'image/png', size: 4, + get path() { pathReads += 1; return '/etc/hosts'; }, + }; + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: any, localId: string) => void; + const cancelQueued = harness.session.onCancelQueuedMessage.mock.calls.at(-1)![0] as (localId: string) => boolean; + + onUserMessage({ role: 'user', content: { type: 'text', text: '', attachments: [attachment] } }, 'cancel-id'); + expect(cancelQueued('cancel-id')).toBe(true); + onUserMessage({ role: 'user', content: { type: 'text', text: 'next valid prompt' } }, 'next-id'); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'next valid prompt' }))); + + expect(pathReads).toBe(0); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: '' })); + harness.onError?.(new Error('finish test')); + await running; + }); + + it('does not pump the next prompt when cleanup rejects a pending settlement sync', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + await completeHistoryInitialization(); + + const userMessage = (text: string) => ({ role: 'user', content: { type: 'text', text } }); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: ReturnType, localId: string) => void; + onUserMessage(userMessage('first'), 'first-id'); + onUserMessage(userMessage('second'), 'second-id'); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'first' }))); + const firstPrompt = harness.sent.find((item) => (item as { type?: string; message?: string }).type === 'prompt') as { id: string }; + harness.onEvent!({ type: 'response', id: firstPrompt.id, command: 'prompt', success: true }); + harness.onEvent!({ type: 'agent_start' }); + harness.onEvent!({ type: 'agent_settled' }); + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')).toHaveLength(3)); + + harness.onError?.(new Error('transport failed during settlement sync')); + await Promise.resolve(); + await Promise.resolve(); + await running; + + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'second' })); + }); + + it('compensates a preflight abort when the prompt starts late, then releases the FIFO once settled', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + await completeHistoryInitialization(); + + const userMessage = (text: string) => ({ role: 'user', content: { type: 'text', text } }); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: ReturnType, localId: string) => void; + onUserMessage(userMessage('first'), 'first-id'); + onUserMessage(userMessage('second'), 'second-id'); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'first' }))); + const promptCommand = harness.sent.find((item) => (item as { type?: string; message?: string }).type === 'prompt') as { id: string }; + + const abort = harness.rpcHandlers.get(RPC_METHODS.Abort); + expect(abort).toBeDefined(); + const abortPromise = abort!({}); + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'abort')).toHaveLength(1)); + const firstAbort = harness.sent.filter((item) => (item as { type?: string }).type === 'abort')[0] as { id: string }; + harness.onEvent!({ type: 'response', id: firstAbort.id, command: 'abort', success: true }); + + // Pi 0.83 may acknowledge abort while prompt preflight is still running. + // A later agent_start must issue one compensating abort and keep the next + // FIFO item blocked until both the real settlement and compensation ack. + harness.onEvent!({ type: 'response', id: promptCommand.id, command: 'prompt', success: true }); + let abortResolved = false; + void abortPromise.then(() => { abortResolved = true; }); + await new Promise((resolve) => setTimeout(resolve, 1_100)); + expect(abortResolved).toBe(false); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'second' })); + harness.onEvent!({ type: 'agent_start' }); + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'abort')).toHaveLength(2)); + const compensatingAbort = harness.sent.filter((item) => (item as { type?: string }).type === 'abort')[1] as { id: string }; + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'second' })); + + harness.onEvent!({ type: 'agent_end', messages: [], willRetry: false }); + harness.onEvent!({ type: 'agent_settled' }); + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')).toHaveLength(3)); + const settlementSync = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[2] as { id: string }; + harness.onEvent!({ type: 'response', id: settlementSync.id, command: 'get_entries', success: true, data: { entries: [], leafId: null } }); + // agent_settled requested another read while the command-only fallback + // sync was in flight, so the history layer serializes one follow-up. + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')).toHaveLength(4)); + const followUpSync = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[3] as { id: string }; + harness.onEvent!({ type: 'response', id: followUpSync.id, command: 'get_entries', success: true, data: { entries: [], leafId: null } }); + harness.onEvent!({ type: 'response', id: compensatingAbort.id, command: 'abort', success: true }); + await expect(abortPromise).resolves.toEqual({ success: true }); + + expect(harness.session.emitMessagesConsumed).toHaveBeenCalledWith(['first-id'], { clearQueuedThinkingGrace: true }); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'second' }))); + harness.onError?.(new Error('finish test')); + await running; + }); + + it('keeps the preflight guard when no-active abort rejection arrives before lifecycle fallback', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + await completeHistoryInitialization(); + + const userMessage = (text: string) => ({ role: 'user', content: { type: 'text', text } }); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: ReturnType & { + meta?: { deliveryMode?: 'queue' | 'steer' }; + }, localId: string) => void; + onUserMessage(userMessage('late preflight'), 'late-id'); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'late preflight' }))); + const promptCommand = harness.sent.find((item) => (item as { type?: string; message?: string }).type === 'prompt') as { id: string }; + // Preserve the preflight/no-lifecycle shape while making the native + // streaming generation observable to a steer queued behind abort. + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: { isStreaming: true } }); + + const abort = harness.rpcHandlers.get(RPC_METHODS.Abort)!; + const abortPromise = abort({}); + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'abort')).toHaveLength(1)); + const firstAbort = harness.sent.filter((item) => (item as { type?: string }).type === 'abort')[0] as { id: string }; + onUserMessage({ + ...userMessage('after abort'), + meta: { deliveryMode: 'steer' }, + }, 'post-abort-steer'); + + // Pi rejects before the 1s lifecycle-missing fallback observes that the + // prompt is still in preflight. The guard must remain installed so a + // later agent_start can still trigger the compensating abort. + harness.onEvent!({ type: 'response', id: firstAbort.id, command: 'abort', success: false, error: 'No active agent to abort' }); + let abortResolved = false; + void abortPromise.then(() => { abortResolved = true; }); + await Promise.resolve(); + await Promise.resolve(); + expect(abortResolved).toBe(false); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'after abort' })); + expect(harness.sent.some((item) => (item as { type?: string }).type === 'steer')).toBe(false); + + harness.onEvent!({ type: 'response', id: promptCommand.id, command: 'prompt', success: true }); + harness.onEvent!({ type: 'agent_start' }); + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'abort')).toHaveLength(2)); + const compensatingAbort = harness.sent.filter((item) => (item as { type?: string }).type === 'abort')[1] as { id: string }; + + harness.onEvent!({ type: 'agent_end', messages: [], willRetry: false }); + harness.onEvent!({ type: 'agent_settled' }); + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')).toHaveLength(3)); + const settlementSync = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[2] as { id: string }; + harness.onEvent!({ type: 'response', id: settlementSync.id, command: 'get_entries', success: true, data: { entries: [], leafId: null } }); + harness.onEvent!({ type: 'response', id: compensatingAbort.id, command: 'abort', success: true }); + + await expect(abortPromise).resolves.toEqual({ success: true }); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'after abort' }))); + expect(harness.sent.some((item) => (item as { type?: string }).type === 'steer')).toBe(false); + + harness.onError?.(new Error('finish test')); + await running; + }); + + it('installs the abort barrier before waiting for a config mutation and never aborts the next prompt', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + await completeHistoryInitialization(); + + const userMessage = (text: string) => ({ role: 'user', content: { type: 'text', text } }); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: ReturnType, localId: string) => void; + onUserMessage(userMessage('first'), 'first-id'); + onUserMessage(userMessage('second'), 'second-id'); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'first' }))); + const promptCommand = harness.sent.find((item) => (item as { type?: string; message?: string }).type === 'prompt') as { id: string }; + harness.onEvent!({ type: 'response', id: promptCommand.id, command: 'prompt', success: true }); + harness.onEvent!({ type: 'agent_start' }); + + const setConfig = harness.rpcHandlers.get(RPC_METHODS.SetSessionConfig)!; + const configPromise = setConfig({ model: { provider: 'provider', modelId: 'model' } }); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'set_model' }))); + const setModelCommand = harness.sent.find((item) => (item as { type?: string }).type === 'set_model') as { id: string }; + + const abort = harness.rpcHandlers.get(RPC_METHODS.Abort)!; + const abortPromise = abort({}); + harness.onEvent!({ type: 'agent_end', messages: [], willRetry: false }); + harness.onEvent!({ type: 'agent_settled' }); + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')).toHaveLength(3)); + const settlementSync = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[2] as { id: string }; + harness.onEvent!({ type: 'response', id: settlementSync.id, command: 'get_entries', success: true, data: { entries: [], leafId: null } }); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(harness.sent.filter((item) => (item as { type?: string }).type === 'abort')).toHaveLength(0); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'second' })); + + harness.onEvent!({ + type: 'response', id: setModelCommand.id, command: 'set_model', success: true, + data: { id: 'model', provider: 'provider' }, + }); + await expect(configPromise).resolves.toMatchObject({ applied: { model: { provider: 'provider', modelId: 'model' } } }); + await expect(abortPromise).resolves.toEqual({ success: true }); + expect(harness.sent.filter((item) => (item as { type?: string }).type === 'abort')).toHaveLength(0); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'second' }))); + + harness.onError?.(new Error('finish test')); + await running; + }); + + it('keeps a command-only abort guard after an ordinary abort error until the stability deadline', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + await completeHistoryInitialization(); + + const userMessage = (text: string) => ({ role: 'user', content: { type: 'text', text } }); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: ReturnType, localId: string) => void; + onUserMessage(userMessage('command-only'), 'command-id'); + onUserMessage(userMessage('next'), 'next-id'); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'command-only' }))); + const promptCommand = harness.sent.find((item) => (item as { type?: string; message?: string }).type === 'prompt') as { id: string }; + + vi.useFakeTimers(); + const abort = harness.rpcHandlers.get(RPC_METHODS.Abort)!; + const abortPromise = abort({}); + await Promise.resolve(); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(0); + const abortCommand = harness.sent.find((item) => (item as { type?: string }).type === 'abort') as { id: string }; + expect(abortCommand).toBeDefined(); + harness.onEvent!({ type: 'response', id: promptCommand.id, command: 'prompt', success: true }); + await vi.advanceTimersByTimeAsync(1_000); + const fallbackSync = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[2] as { id: string }; + expect(fallbackSync).toBeDefined(); + harness.onEvent!({ type: 'response', id: fallbackSync.id, command: 'get_entries', success: true, data: { entries: [], leafId: null } }); + await vi.advanceTimersByTimeAsync(0); + harness.onEvent!({ type: 'response', id: abortCommand.id, command: 'abort', success: false, error: 'No active agent to abort' }); + await Promise.resolve(); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'next' })); + + await vi.advanceTimersByTimeAsync(24_000); + await expect(abortPromise).resolves.toEqual({ success: true }); + expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'next' })); + vi.useRealTimers(); + + harness.onError?.(new Error('finish test')); + await running; + }); + + it('still sends native abort when Pi is streaming without a local prompt boundary', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: { isStreaming: true } }); + await completeHistoryInitialization(); + + const abort = harness.rpcHandlers.get(RPC_METHODS.Abort)!; + const abortPromise = abort({}); + await vi.waitFor(() => expect(harness.sent.filter((item) => (item as { type?: string }).type === 'abort')).toHaveLength(1)); + const abortCommand = harness.sent.find((item) => (item as { type?: string }).type === 'abort') as { id: string }; + + // A steer arriving behind the abort mutation targets the generation + // being aborted. Abort success must invalidate it before releasing the + // mutex so the message becomes an ordinary prompt instead of entering + // Pi's now-idle native steer queue. + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: { + role: 'user'; + content: { type: 'text'; text: string }; + meta?: { deliveryMode?: 'queue' | 'steer' }; + }, localId: string) => void; + onUserMessage({ + role: 'user', + content: { type: 'text', text: 'after abort' }, + meta: { deliveryMode: 'steer' }, + }, 'post-abort-steer'); + + harness.onEvent!({ type: 'response', id: abortCommand.id, command: 'abort', success: true }); + + await expect(abortPromise).resolves.toEqual({ success: true }); + expect(harness.session.keepAlive).toHaveBeenLastCalledWith(false, 'remote', undefined); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ + type: 'prompt', message: 'after abort', + }))); + expect(harness.sent.some((item) => (item as { type?: string }).type === 'steer')).toBe(false); + + harness.onError?.(new Error('finish test')); + await running; + }); + + it('fails closed and poisons the mutation lease when configuration times out', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + await completeHistoryInitialization(); + + vi.useFakeTimers(); + const setConfig = harness.rpcHandlers.get(RPC_METHODS.SetSessionConfig)!; + const configPromise = setConfig({ model: { provider: 'provider', modelId: 'model' } }); + const configRejection = expect(configPromise).rejects.toThrow('timed out'); + await vi.advanceTimersByTimeAsync(0); + expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'set_model', provider: 'provider', modelId: 'model' })); + + await vi.advanceTimersByTimeAsync(10_000); + await configRejection; + await Promise.resolve(); + expect(harness.cleanupCount).toBe(1); + vi.useRealTimers(); + + await running; + }); + + it('fails closed when the detached startup effort mutation times out', async () => { + vi.useFakeTimers(); + const running = runPi({ workingDirectory: '/work', effort: 'high' }); + await vi.advanceTimersByTimeAsync(0); + await Promise.resolve(); + expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'set_thinking_level', level: 'high' })); + + await vi.advanceTimersByTimeAsync(10_000); + await Promise.resolve(); + expect(harness.cleanupCount).toBe(1); + const setConfig = harness.rpcHandlers.get(RPC_METHODS.SetSessionConfig)!; + const blockedConfig = setConfig({ model: { provider: 'provider', modelId: 'model' } }); + void blockedConfig.catch(() => {}); + await vi.advanceTimersByTimeAsync(0); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'set_model' })); + vi.useRealTimers(); + + await running; + }); + + it('settles consecutive command-only prompts without stamping the old timer onto the next generation', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + await completeHistoryInitialization(); + + const userMessage = (text: string) => ({ role: 'user', content: { type: 'text', text } }); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: ReturnType, localId: string) => void; + onUserMessage(userMessage('command-a'), 'command-a-id'); + onUserMessage(userMessage('command-b'), 'command-b-id'); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'command-a' }))); + + vi.useFakeTimers(); + const firstPrompt = harness.sent.find((item) => (item as { type?: string; message?: string }).type === 'prompt') as { id: string }; + harness.onEvent!({ type: 'response', id: firstPrompt.id, command: 'prompt', success: true }); + await vi.advanceTimersByTimeAsync(1_000); + const firstFallbackSync = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[2] as { id: string }; + expect(firstFallbackSync).toBeDefined(); + harness.onEvent!({ type: 'response', id: firstFallbackSync.id, command: 'get_entries', success: true, data: { entries: [], leafId: null } }); + await vi.advanceTimersByTimeAsync(0); + const promptsAfterFirst = harness.sent.filter((item) => (item as { type?: string }).type === 'prompt') as Array<{ id: string; message: string }>; + expect(promptsAfterFirst.map((item) => item.message)).toEqual(['command-a', 'command-b']); + + const secondPrompt = promptsAfterFirst[1]!; + harness.onEvent!({ type: 'response', id: secondPrompt.id, command: 'prompt', success: true }); + await vi.advanceTimersByTimeAsync(1_000); + const secondFallbackSync = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[3] as { id: string }; + expect(secondFallbackSync).toBeDefined(); + harness.onEvent!({ type: 'response', id: secondFallbackSync.id, command: 'get_entries', success: true, data: { entries: [], leafId: null } }); + await vi.advanceTimersByTimeAsync(0); + onUserMessage(userMessage('command-c'), 'command-c-id'); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(0); + expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'command-c' })); + vi.useRealTimers(); + + harness.onError?.(new Error('finish test')); + await running; + }); + + it('syncs a command-only append before registering the following prompt history entry', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + await completeHistoryInitialization(); + + const userMessage = (text: string) => ({ role: 'user', content: { type: 'text', text } }); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: ReturnType, localId: string) => void; + onUserMessage(userMessage('command-only'), 'command-local'); + onUserMessage(userMessage('following prompt'), 'following-local'); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'command-only' }))); + + vi.useFakeTimers(); + const commandPrompt = harness.sent.find((item) => (item as { type?: string; message?: string }).type === 'prompt') as { id: string }; + harness.onEvent!({ type: 'response', id: commandPrompt.id, command: 'prompt', success: true }); + await vi.advanceTimersByTimeAsync(1_000); + + // No entry_appended event arrives. The compatibility fallback must + // read Pi's append log and bind this entry before command-local is + // retired and the following prompt is allowed to start. + const commandSync = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[2] as { id: string }; + expect(commandSync).toBeDefined(); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'following prompt' })); + harness.onEvent!({ + type: 'response', id: commandSync.id, command: 'get_entries', success: true, + data: { entries: [{ id: 'native-command', type: 'message', message: { role: 'user' } }], leafId: 'native-command' }, + }); + await vi.advanceTimersByTimeAsync(0); + + const prompts = harness.sent.filter((item) => (item as { type?: string }).type === 'prompt') as Array<{ id: string; message: string }>; + expect(prompts.map((prompt) => prompt.message)).toEqual(['command-only', 'following prompt']); + const followingPrompt = prompts[1]!; + harness.onEvent!({ type: 'response', id: followingPrompt.id, command: 'prompt', success: true }); + harness.onEvent!({ type: 'agent_start' }); + harness.onEvent!({ type: 'turn_start' }); + const followingSync = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[3] as { id: string }; + expect(followingSync).toBeDefined(); + harness.onEvent!({ + type: 'response', id: followingSync.id, command: 'get_entries', success: true, + data: { entries: [{ id: 'native-following', type: 'message', message: { role: 'user' } }], leafId: 'native-following' }, + }); + await vi.advanceTimersByTimeAsync(0); + + let metadata: Record = {}; + for (const [updater] of harness.session.updateMetadata.mock.calls) { + if (typeof updater === 'function') metadata = updater(metadata); + } + expect(metadata).toMatchObject({ + conversationHistoryEntryIds: { + 'command-local': 'native-command', + 'following-local': 'native-following', + }, + }); + vi.useRealTimers(); + + harness.onError?.(new Error('finish test')); + await running; + }); + + it('fails closed without starting the next prompt when command-only history sync fails', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ type: 'response', command: 'get_state', success: true, data: {} }); + await completeHistoryInitialization(); + + const userMessage = (text: string) => ({ role: 'user', content: { type: 'text', text } }); + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: ReturnType, localId: string) => void; + onUserMessage(userMessage('command-only'), 'command-local'); + onUserMessage(userMessage('must remain blocked'), 'following-local'); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'prompt', message: 'command-only' }))); + + vi.useFakeTimers(); + const commandPrompt = harness.sent.find((item) => (item as { type?: string; message?: string }).type === 'prompt') as { id: string }; + harness.onEvent!({ type: 'response', id: commandPrompt.id, command: 'prompt', success: true }); + await vi.advanceTimersByTimeAsync(1_000); + const commandSync = harness.sent.filter((item) => (item as { type?: string }).type === 'get_entries')[2] as { id: string }; + expect(commandSync).toBeDefined(); + harness.onEvent!({ type: 'response', id: commandSync.id, command: 'get_entries', success: false, error: 'temporary read failure' }); + await vi.advanceTimersByTimeAsync(0); + + expect(harness.sent).not.toContainEqual(expect.objectContaining({ type: 'prompt', message: 'must remain blocked' })); + expect(harness.session.emitMessagesConsumed).not.toHaveBeenCalledWith( + ['command-local'], + expect.anything(), + ); + expect(harness.cleanupCount).toBe(1); + vi.useRealTimers(); + await running; + }); +}); + +describe('Pi native steering delivery mode', () => { + beforeEach(() => { + harness.sent.length = 0; + harness.throwOnGetCommands = false; + harness.onEvent = null; + harness.rpcHandlers.clear(); + harness.session.onUserMessage.mockReset(); + harness.session.onCancelQueuedMessage.mockReset(); + harness.session.emitMessagesConsumed.mockReset(); + harness.session.sendSessionEvent.mockReset(); + harness.session.updateMetadata.mockReset(); + harness.cleanupCount = 0; + harness.killCount = 0; + harness.session.rpcHandlerManager.registerHandler.mockReset(); + harness.session.rpcHandlerManager.registerHandler.mockImplementation((method: string, handler: (payload: unknown) => Promise) => { + harness.rpcHandlers.set(method, handler); + }); + }); + + it('routes explicit steer messages natively while streaming and retains explicit queue FIFO', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ + type: 'response', command: 'get_state', success: true, + data: { sessionId: 'pi-steering-session', sessionFile: '/tmp/pi-steering.jsonl', isStreaming: true }, + }); + await completeHistoryInitialization(); + + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: { + role: 'user'; + content: { type: 'text'; text: string }; + meta?: { deliveryMode?: 'queue' | 'steer' }; + }, localId: string) => void; + onUserMessage({ + role: 'user', + content: { type: 'text', text: 'steer the active turn' }, + meta: { deliveryMode: 'steer' }, + }, 'native-steer-id'); + onUserMessage({ + role: 'user', + content: { type: 'text', text: 'keep this in the normal queue' }, + meta: { deliveryMode: 'queue' }, + }, 'queue-id'); + + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ + type: 'steer', message: 'steer the active turn', + }))); + expect(harness.sent).not.toContainEqual(expect.objectContaining({ + type: 'prompt', message: 'keep this in the normal queue', + })); + const steer = harness.sent.find((item) => (item as { type?: string }).type === 'steer') as { id: string }; + harness.onEvent!({ type: 'response', id: steer.id, command: 'steer', success: true }); + await vi.waitFor(() => expect(harness.session.emitMessagesConsumed).toHaveBeenCalledWith(['native-steer-id'], undefined)); + + // The main turn ending releases only the explicit queue mode. The steer + // response itself never changes Pi's main streaming/thinking state. + harness.onEvent!({ + type: 'response', command: 'get_state', success: true, + data: { sessionId: 'pi-steering-session', sessionFile: '/tmp/pi-steering.jsonl', isStreaming: false }, + }); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ + type: 'prompt', message: 'keep this in the normal queue', + }))); + + harness.onError?.(new Error('finish test')); + await running; + }); + + it('does not steer a later streaming generation and preserves fallback arrival order', async () => { + const running = runPi({ workingDirectory: '/work' }); + await vi.waitFor(() => expect(harness.onEvent).not.toBeNull()); + harness.onEvent!({ + type: 'response', command: 'get_state', success: true, + data: { sessionId: 'pi-steering-session', sessionFile: '/tmp/pi-steering.jsonl', isStreaming: true }, + }); + await completeHistoryInitialization(); + + // Hold the shared mutation lock, then simulate turn A ending and turn B + // starting before the queued steer can reach Pi. + const setConfig = harness.rpcHandlers.get(RPC_METHODS.SetSessionConfig)!; + const configRequest = setConfig({ effort: 'low' }); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ type: 'set_thinking_level' }))); + const setThinking = harness.sent.find((item) => (item as { type?: string }).type === 'set_thinking_level') as { id: string }; + + const onUserMessage = harness.session.onUserMessage.mock.calls.at(-1)![0] as (message: { + role: 'user'; + content: { type: 'text'; text: string }; + meta?: { deliveryMode?: 'queue' | 'steer' }; + }, localId: string) => void; + onUserMessage({ + role: 'user', + content: { type: 'text', text: 'earlier steer fallback' }, + meta: { deliveryMode: 'steer' }, + }, 'steer-id'); + onUserMessage({ + role: 'user', + content: { type: 'text', text: 'later ordinary prompt' }, + meta: { deliveryMode: 'queue' }, + }, 'queue-id'); + + harness.onEvent!({ + type: 'response', command: 'get_state', success: true, + data: { sessionId: 'pi-steering-session', sessionFile: '/tmp/pi-steering.jsonl', isStreaming: false }, + }); + harness.onEvent!({ + type: 'response', command: 'get_state', success: true, + data: { sessionId: 'pi-steering-session', sessionFile: '/tmp/pi-steering.jsonl', isStreaming: true }, + }); + harness.onEvent!({ type: 'response', id: setThinking.id, command: 'set_thinking_level', success: true }); + await configRequest; + + // The stale steer is a normal prompt now, but Pi turn B is still + // streaming, so neither fallback nor later queue item may start yet. + await vi.waitFor(() => expect(harness.sent.some((item) => (item as { type?: string }).type === 'steer')).toBe(false)); + expect(harness.sent.some((item) => (item as { type?: string; message?: string }).type === 'prompt' && (item as { message?: string }).message === 'later ordinary prompt')).toBe(false); + + // When B settles, delayed fallback must win its original reservation. + harness.onEvent!({ + type: 'response', command: 'get_state', success: true, + data: { sessionId: 'pi-steering-session', sessionFile: '/tmp/pi-steering.jsonl', isStreaming: false }, + }); + await vi.waitFor(() => expect(harness.sent).toContainEqual(expect.objectContaining({ + type: 'prompt', message: 'earlier steer fallback', + }))); + const prompts = harness.sent.filter((item) => (item as { type?: string }).type === 'prompt') as Array<{ message: string }>; + expect(prompts.map((prompt) => prompt.message)).toEqual(['earlier steer fallback']); + + harness.onError?.(new Error('finish test')); + await running; + }); +}); + +describe('Pi prompt preparation', () => { + it('reads image attachments into Pi RPC image content while retaining safe text references', async () => { + const { mkdtemp, writeFile, rm, symlink } = await import('node:fs/promises'); + const { tmpdir } = await import('node:os'); + const { join, sep } = await import('node:path'); + const imagePath = join(process.env.TMPDIR ?? '/tmp', `pi-image-${Date.now()}.png`); + const uploadDir = await mkdtemp(join(tmpdir(), 'pi-upload-auth-')); + const outsidePath = process.platform === 'win32' ? null : join(tmpdir(), `pi-outside-${Date.now()}.png`); + const symlinkPath = process.platform === 'win32' ? null : join(uploadDir, 'escape.png'); + await writeFile(imagePath, Buffer.from([0x89, 0x50, 0x4e, 0x47])); + if (outsidePath && symlinkPath) { + await writeFile(outsidePath, Buffer.from([1, 2, 3, 4])); + await symlink(outsidePath, symlinkPath); + } + try { + const { preparePiUserMessage } = await import('./runPi'); + const prepared = await preparePiUserMessage('$brave-search explain', [ + { id: 'image', filename: 'plot.png', mimeType: 'image/png', size: 4, path: imagePath }, + { id: 'text', filename: 'notes file.txt', mimeType: 'text/plain', size: 1, path: '/tmp/notes file.txt' }, + ], [{ name: 'skill:brave-search', source: 'skill' }], { + authorizeImagePath: () => true, + authorizeOpenedImage: () => true, + }); + expect(prepared.message).toBe('/skill:brave-search explain\n\nAttached file: \"/tmp/notes file.txt\"'); + expect(prepared.images).toEqual([{ type: 'image', mimeType: 'image/png', data: 'iVBORw==' }]); + expect(prepared.imageReadErrors).toEqual([]); + expect(formatPiUserMessage('', [{ id: 'newline', filename: 'x', mimeType: 'text/plain', size: 1, path: '/tmp/a\nb' }], [])).toBe('Attached file: \"/tmp/a\\nb\"'); + const failed = await preparePiUserMessage('', [{ id: 'missing', filename: 'missing.png', mimeType: 'image/png', size: 1, path: '/missing/image.png' }], [], { + authorizeImagePath: () => true, + authorizeOpenedImage: () => true, + }); + expect(failed).toMatchObject({ message: '', images: [] }); + expect(failed.imageReadErrors[0]).toContain('Could not attach image missing.png'); + const unauthorized = await preparePiUserMessage('', [{ id: 'forged', filename: 'hosts.png', mimeType: 'image/png', size: 1, path: '/etc/hosts' }], [], { + authorizeImagePath: () => false, + authorizeOpenedImage: () => false, + }); + expect(unauthorized).toMatchObject({ message: '', images: [] }); + expect(unauthorized.imageReadErrors).toEqual(['Could not attach image hosts.png: invalid upload path']); + if (symlinkPath) { + const symlinkEscape = await preparePiUserMessage('', [{ id: 'symlink', filename: 'escape.png', mimeType: 'image/png', size: 4, path: symlinkPath }], [], { + authorizeImagePath: (path) => path.startsWith(`${uploadDir}${sep}`), + authorizeOpenedImage: () => false, + }); + expect(symlinkEscape).toMatchObject({ message: '', images: [] }); + expect(symlinkEscape.imageReadErrors[0]).toContain('Could not attach image escape.png'); + } + } finally { + await rm(imagePath, { force: true }); + if (outsidePath) await rm(outsidePath, { force: true }); + await rm(uploadDir, { recursive: true, force: true }); + } + }); }); diff --git a/cli/src/pi/runPi.ts b/cli/src/pi/runPi.ts index d809fae5..f2a3fb3f 100644 --- a/cli/src/pi/runPi.ts +++ b/cli/src/pi/runPi.ts @@ -1,24 +1,35 @@ +import { randomUUID } from 'node:crypto'; import { logger } from '@/ui/logger'; import { bootstrapExistingSession, bootstrapSession } from '@/agent/sessionFactory'; import { registerKillSessionHandler } from '@/claude/registerKillSessionHandler'; import { registerLocalHandoffHandler } from '@/agent/localHandoff'; import { createRunnerLifecycle, createModeChangeHandler, setControlledByUser } from '@/agent/runnerLifecycle'; -import { formatAttachmentsForClaude, formatMessageWithAttachments } from '@/utils/attachmentFormatter'; import { getInvokedCwd } from '@/utils/invokedCwd'; import { PiTransport } from './piTransport'; import { PiSession } from './session'; -import { parsePiModels, parsePiCommands, sendPiRpcAndWait, wireTransportEvents } from './loop'; +import { PiConversationHistory, PiHistoryRestoreError } from './conversationHistory'; +import { parsePiModels, parsePiCommands, PiRpcTimeoutError, sendPiRpcAndWait, wireTransportEvents } from './loop'; import { PiThinkingLevelSchema, SetSessionConfigPayloadSchema } from './schemas'; -import type { PiThinkingLevel } from './types'; +import type { PiImageContent, PiThinkingLevel } from './types'; +import { PiPromptQueue, type PiPreparedPrompt } from './promptQueue'; +import { PiSteerDispatcher } from './steerDispatcher'; import type { ListPiModelsResponse, PiCommandSummary, SlashCommand, SlashCommandsResponse } from '@hapi/protocol/apiTypes'; import { RPC_METHODS } from '@hapi/protocol/rpcMethods'; import type { ListSkillsResponse, SkillSummary } from '@/modules/common/skills'; import type { AttachmentMetadata } from '@/api/types'; +import { readBoundedAttachmentFile } from '@/modules/common/attachmentFile'; +import { MAX_UPLOAD_BYTES } from '@/modules/common/attachmentLimits'; +import { isAuthorizedUploadFile, isPathWithinUploadDir, type UploadFileIdentity } from '@/modules/common/handlers/uploads'; // Grace period before force-draining prompts buffered during Pi startup when no // get_state response arrives. Comfortably above the 10s Pi RPC timeout so a slow // but healthy startup still flips ready via get_state first (issue #1143). const PI_READY_FALLBACK_MS = 30_000; +const PI_ABORT_OPERATION_TIMEOUT_MS = 25_000; + +function isPiNoActiveAbortError(detail: string): boolean { + return /no active|nothing.*abort/i.test(detail); +} function getPiSkillName(commandName: string): string { return commandName.startsWith('skill:') ? commandName.slice('skill:'.length) : commandName; @@ -57,18 +68,80 @@ export function rewritePiSkillPrompt(message: string, commands: readonly PiComma return `${match[1]}/${command?.name ?? `skill:${match[2]}`}${message.slice(match[0].length)}`; } +function formatPiFileNotice(path: string): string { + // Pi 0.83 rejects @file arguments in RPC mode. Keep this as ordinary prompt + // text so the model can use its read tool, with JSON quoting preventing a + // path containing whitespace or a newline from injecting another prompt line. + return `Attached file: ${JSON.stringify(path)}`; +} + +function formatPiTextAttachments(attachments: AttachmentMetadata[] | undefined): string { + if (!attachments) return ''; + return attachments + .filter((attachment) => !attachment.mimeType.toLowerCase().startsWith('image/')) + .map((attachment) => formatPiFileNotice(attachment.path)) + .join('\n'); +} + export function formatPiUserMessage( message: string, attachments: AttachmentMetadata[] | undefined, commands: readonly PiCommandSummary[], ): string { const skillPrompt = rewritePiSkillPrompt(message, commands); - if (skillPrompt === message) return formatMessageWithAttachments(message, attachments); - - const attachmentText = formatAttachmentsForClaude(attachments); + const attachmentText = formatPiTextAttachments(attachments); + if (skillPrompt === message) { + if (!attachmentText) return message; + return message ? `${attachmentText}\n\n${message}` : attachmentText; + } + // Pi parses slash/skill commands only when they are the first line. return attachmentText ? `${skillPrompt}\n\n${attachmentText}` : skillPrompt; } +export type PiPromptPreparation = Omit & { imageReadErrors: string[] }; + +export async function preparePiUserMessage( + message: string, + attachments: AttachmentMetadata[] | undefined, + commands: readonly PiCommandSummary[], + options: { + authorizeImagePath: (path: string) => boolean; + authorizeOpenedImage: (path: string, identity: UploadFileIdentity) => boolean; + }, +): Promise { + const formattedMessage = formatPiUserMessage(message, attachments, commands); + const images: PiImageContent[] = []; + const imageReadErrors: string[] = []; + let totalImageBytes = 0; + + for (const attachment of attachments ?? []) { + if (!attachment.mimeType.toLowerCase().startsWith('image/')) continue; + try { + const uploadPath = attachment.path; + if (!options.authorizeImagePath(uploadPath)) { + throw new Error('invalid upload path'); + } + const data = await readBoundedAttachmentFile( + uploadPath, + MAX_UPLOAD_BYTES - totalImageBytes, + (identity) => options.authorizeOpenedImage(uploadPath, identity), + ); + totalImageBytes += data.length; + images.push({ type: 'image', data: data.toString('base64'), mimeType: attachment.mimeType }); + } catch (error) { + const detail = error instanceof Error ? error.message : String(error); + imageReadErrors.push(`Could not attach image ${attachment.filename}: ${detail}`); + } + } + + return { message: formattedMessage, images, imageReadErrors }; +} + +/** A failed source restore leaves the live wrapper on an unknown native branch. */ +export function failPiHistoryOnRestoreError(error: unknown, failNativeStartup: (error: Error) => void): void { + if (error instanceof PiHistoryRestoreError) failNativeStartup(error); +} + export async function runPi(opts: { startedBy?: 'runner' | 'terminal'; startingMode?: 'local' | 'remote'; @@ -127,7 +200,47 @@ export async function runPi(opts: { if (opts.resumeSessionId) { transportArgs.push('--session', opts.resumeSessionId); } - const transport = new PiTransport({ command: 'pi', args: transportArgs, cwd: workingDirectory }); + const transport = new PiTransport({ + command: 'pi', + args: transportArgs, + cwd: workingDirectory, + env: { ...process.env, PI_RPC_EMIT_TITLE: '1' }, + }); + const conversationHistory = new PiConversationHistory( + piSession, + (command, timeoutMs) => sendPiRpcAndWait(piSession, transport, command, timeoutMs), + ); + let historyBaseline: Promise | null = null; + let historyInitialization: Promise | null = null; + const initializeHistoryBaseline = (): Promise => { + historyBaseline ??= conversationHistory.initializeBaseline(); + return historyBaseline; + }; + const initializeHistory = (): Promise => { + historyInitialization ??= initializeHistoryBaseline().then(async (baselineReady) => { + if (baselineReady) await conversationHistory.probeCapabilities(); + }); + return historyInitialization; + }; + piSession.setNativeReadyPreparation(initializeHistory); + + const publishConversationHistoryCapabilities = async () => { + const conversationHistoryCapabilities = conversationHistory.getCapabilitiesForMetadata()?.conversationHistory; + piSession.updateMetadata((metadata) => { + const capabilities = { ...metadata.capabilities }; + delete capabilities.conversationHistory; + if (conversationHistoryCapabilities) { + capabilities.conversationHistory = conversationHistoryCapabilities; + } + return { ...metadata, capabilities }; + }); + }; + conversationHistory.setPublishCapabilities(publishConversationHistoryCapabilities); + conversationHistory.restoreEntryIds( + typeof apiSession.getMetadata === 'function' + ? apiSession.getMetadata()?.conversationHistoryEntryIds + : undefined, + ); piSession.startKeepAlive(); @@ -148,9 +261,12 @@ export async function runPi(opts: { registerLocalHandoffHandler(apiSession.rpcHandlerManager, lifecycle); let cleanupInitiated = false; + let steerDispatcher: PiSteerDispatcher | null = null; const safeCleanup = async () => { if (cleanupInitiated) return; cleanupInitiated = true; + steerDispatcher?.stop(); + piSession.cancelReadyGate(); await lifecycle.cleanupAndExit(); }; @@ -170,8 +286,14 @@ export async function runPi(opts: { // Pending user-message localIds in FIFO order const pendingLocalIds: string[] = []; + let transportEvents: ReturnType | null = null; + // --- Transport error/close handlers --- transport.onError((error) => { + steerDispatcher?.stop(); + transportEvents?.flush(); + transportEvents?.cancelPendingExtensionUi('Pi transport failed', { sendResponse: false }); + transportEvents?.terminatePendingRpc(error); logger.debug(`[pi] Transport error: ${error.message}`); lifecycle.markCrash(error); lifecycle.setExitCode(1); @@ -181,6 +303,10 @@ export async function runPi(opts: { }); transport.onClose((code, signal) => { + steerDispatcher?.stop(); + transportEvents?.flush(); + transportEvents?.cancelPendingExtensionUi('Pi session ended', { sendResponse: false }); + transportEvents?.terminatePendingRpc(new Error('Pi session ended')); if (killedByCleanup) { logger.debug(`[pi] Pi process closed during lifecycle cleanup (code=${code}, signal=${signal})`); void safeCleanup(); @@ -226,8 +352,217 @@ export async function runPi(opts: { void safeCleanup(); }; - wireTransportEvents(transport, piSession, pendingLocalIds, { + const promptQueue = new PiPromptQueue(); + const preparingLocalIds = new Set(); + const cancelledWhilePreparing = new Set(); + let preparationChain = Promise.resolve(); + let promptCommandInFlight = false; + let abortInFlight = false; + let activePromptLocalId: string | undefined; + let historyPumpDeferred = false; + let agentLifecycleStarted = false; + let nextOutboundSequence = 0; + + const setPromptCommandInFlight = (value: boolean): void => { + promptCommandInFlight = value; + piSession.setPromptInFlight(value); + }; + + type ActiveAbortGuard = { + deadlineAt: number; + lifecycleStartedAtRequest: boolean; + compensationRequired: boolean; + compensationStarted: boolean; + compensationConfirmed: boolean; + lifecycleMissingObserved: boolean; + settled: boolean; + timer: ReturnType; + resolve: () => void; + reject: (error: Error) => void; + }; + let activeAbortGuard: ActiveAbortGuard | null = null; + + const maybeResolveAbortGuard = (): void => { + const guard = activeAbortGuard; + if (!guard?.settled) return; + if (guard.compensationRequired && !guard.compensationConfirmed) return; + clearTimeout(guard.timer); + guard.resolve(); + }; + + const markPromptBoundarySettled = (): void => { + if (activeAbortGuard) { + activeAbortGuard.settled = true; + maybeResolveAbortGuard(); + return; + } + setPromptCommandInFlight(false); + activePromptLocalId = undefined; + agentLifecycleStarted = false; + pumpPromptQueue(); + }; + + const createAbortGuard = (deadlineAt: number): { guard: ActiveAbortGuard; promise: Promise } => { + let resolve!: () => void; + let reject!: (error: Error) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + let guard!: ActiveAbortGuard; + const timer = setTimeout(() => { + // A successful prompt response with no agent lifecycle is how Pi + // reports extension-command-only prompts. During abort we keep the + // guard for the entire stability window so a genuinely late + // agent_start can still trigger compensation. + if (guard.lifecycleMissingObserved && !guard.compensationRequired) { + guard.settled = true; + maybeResolveAbortGuard(); + return; + } + reject(new Error(`Pi prompt did not settle within ${PI_ABORT_OPERATION_TIMEOUT_MS}ms after abort`)); + }, Math.max(1, deadlineAt - Date.now())); + timer.unref?.(); + guard = { + deadlineAt, + lifecycleStartedAtRequest: agentLifecycleStarted, + compensationRequired: false, + compensationStarted: false, + compensationConfirmed: false, + lifecycleMissingObserved: false, + settled: false, + timer, + resolve, + reject, + }; + return { guard, promise }; + }; + + const pumpPromptQueue = (): void => { + if (cleanupInitiated) return; + if (piSession.isHistoryTransactionActive) { + if (!historyPumpDeferred) { + historyPumpDeferred = true; + piSession.runWhenHistoryIdle(() => { + historyPumpDeferred = false; + pumpPromptQueue(); + }); + } + return; + } + if ( + !piSession.isReady + || piSession.piIsStreaming + || promptCommandInFlight + || abortInFlight + // Earlier steers can fall back only after async preparation/runtime + // lock wait. Do not let a later normal prompt overtake that result. + || steerDispatcher?.hasPending + ) return; + const next = promptQueue.dequeue(); + if (!next) return; + setPromptCommandInFlight(true); + activePromptLocalId = next.localId; + agentLifecycleStarted = false; + const promptId = randomUUID(); + transportEvents?.beginPromptLifecycle(promptId); + if (next.localId) { + conversationHistory.registerUserEntry(next.localId); + pendingLocalIds.push(next.localId); + } + transport.send({ id: promptId, type: 'prompt', message: next.message, ...(next.images.length > 0 ? { images: next.images } : {}) }); + }; + + transportEvents = wireTransportEvents(transport, piSession, pendingLocalIds, { onStartupFailure: failNativeStartup, + conversationHistory, + onReady: () => piSession.runWhenReady(pumpPromptQueue), + onAgentLifecycleStarted: () => { + const wasStarted = agentLifecycleStarted; + agentLifecycleStarted = true; + const guard = activeAbortGuard; + if (!guard || guard.lifecycleStartedAtRequest || wasStarted || guard.compensationStarted) return; + + // Pi 0.83 can acknowledge abort while a prompt is still in async + // preflight. If that prompt starts afterwards, compensate only once + // now that agent.abort() has an active run to cancel. + guard.compensationRequired = true; + guard.compensationStarted = true; + const remainingMs = Math.floor(guard.deadlineAt - Date.now()); + if (remainingMs <= 0) { + guard.reject(new Error('Pi compensating abort missed the abort deadline')); + return; + } + void sendPiRpcAndWait(piSession, transport, { type: 'abort' }, Math.min(10_000, remainingMs)) + .then(() => { + if (activeAbortGuard !== guard) return; + guard.compensationConfirmed = true; + maybeResolveAbortGuard(); + }) + .catch((error: unknown) => { + if (activeAbortGuard !== guard) return; + const detail = error instanceof Error ? error.message : String(error); + guard.reject(new Error(`Pi compensating abort failed: ${detail}`)); + }); + }, + onAgentSettled: markPromptBoundarySettled, + onPromptRejected: (localId) => { + const rejectedLocalId = localId ?? activePromptLocalId; + if (rejectedLocalId) conversationHistory.rejectPendingEntry(rejectedLocalId); + markPromptBoundarySettled(); + }, + onPromptLifecycleMissing: (localId) => { + const rejectedLocalId = localId ?? activePromptLocalId; + if (rejectedLocalId) { + conversationHistory.rejectPendingEntry(rejectedLocalId); + if (pendingLocalIds[0] === rejectedLocalId) pendingLocalIds.shift(); + piSession.emitMessagesConsumed([rejectedLocalId], { clearQueuedThinkingGrace: true }); + } + if (activeAbortGuard) { + activeAbortGuard.lifecycleMissingObserved = true; + return false; + } + markPromptBoundarySettled(); + return true; + }, + }); + + steerDispatcher = new PiSteerDispatcher({ + session: piSession, + transport, + conversationHistory, + enqueuePrompt: (entry) => { + if (cleanupInitiated) return; + promptQueue.enqueue(entry); + }, + onIndeterminateTimeout: (error) => { + failNativeStartup(new Error(`Pi steer outcome is indeterminate: ${error.message}`)); + }, + onPendingStateChange: pumpPromptQueue, + }); + + apiSession.rpcHandlerManager.registerHandler(RPC_METHODS.ForkConversation, async (payload: unknown) => { + const messageLocalId = payload && typeof payload === 'object' + && typeof (payload as { messageLocalId?: unknown }).messageLocalId === 'string' + ? (payload as { messageLocalId: string }).messageLocalId + : undefined; + try { + return await conversationHistory.fork(messageLocalId); + } catch (error) { + failPiHistoryOnRestoreError(error, failNativeStartup); + throw error; + } + }); + apiSession.rpcHandlerManager.registerHandler(RPC_METHODS.RewindConversation, async (payload: unknown) => { + if (!payload || typeof payload !== 'object' || typeof (payload as { messageLocalId?: unknown }).messageLocalId !== 'string') { + throw new Error('messageLocalId is required'); + } + try { + return await conversationHistory.rewind((payload as { messageLocalId: string }).messageLocalId); + } catch (error) { + failPiHistoryOnRestoreError(error, failNativeStartup); + throw error; + } }); // --- Session config RPC --- @@ -240,6 +575,7 @@ export async function runPi(opts: { // { provider, modelId } for Pi sessions. apiSession.rpcHandlerManager.registerHandler(RPC_METHODS.SetSessionConfig, async (rawPayload: unknown) => { + piSession.assertNoHistoryTransaction('change session configuration'); const parsed = SetSessionConfigPayloadSchema.safeParse(rawPayload); if (!parsed.success) { throw new Error('Invalid session config payload'); @@ -284,54 +620,63 @@ export async function runPi(opts: { } } - // Forward changes to Pi process — wait for Pi to confirm before - // committing to PiSession or reporting applied, so the hub does not - // persist a model/effort that Pi rejected (e.g. invalid provider/model - // or thinking level) or that the RPC timed out on. - if (requestedModel) { - if (requestedModel.modelId && requestedModel.provider) { - await sendPiRpcAndWait(piSession, transport, { - type: 'set_model', - provider: requestedModel.provider, - modelId: requestedModel.modelId, - }); - piSession.currentModel = requestedModel.modelId; - piSession.currentProvider = requestedModel.provider; - } else if (requestedModel.modelId && !requestedModel.provider) { - // Provider is unknown until get_state/get_available_models resolve. - // Committing now would persist piSelectedModel while Pi never received - // set_model — contradicting the "await Pi confirmation" contract above. - // Throw so the hub returns 409 and the web client can retry once the - // provider is known. - logger.debug('[pi] set_model suppressed: provider unknown until get_state'); - throw new Error('Model cannot be applied yet: provider is not yet known'); - } else if (requestedModel.modelId === null) { - // Clearing the model needs no Pi RPC (nothing to confirm), so commit - // immediately. This path is not reachable from the web Pi picker today. - piSession.currentModel = null; - piSession.currentProvider = null; + try { + return await piSession.runRuntimeMutation(async () => { + // Forward changes to Pi process — wait for Pi to confirm before + // committing to PiSession or reporting applied. The runtime mutation + // lock is shared with clone/fork/switch_session so a slow set_model + // can never finish against a temporary history branch. + if (requestedModel) { + if (requestedModel.modelId && requestedModel.provider) { + await sendPiRpcAndWait(piSession, transport, { + type: 'set_model', + provider: requestedModel.provider, + modelId: requestedModel.modelId, + }); + piSession.currentModel = requestedModel.modelId; + piSession.currentProvider = requestedModel.provider; + } else if (requestedModel.modelId && !requestedModel.provider) { + // Provider is unknown until get_state/get_available_models resolve. + // Committing now would persist piSelectedModel while Pi never received + // set_model — contradicting the "await Pi confirmation" contract above. + // Throw so the hub returns 409 and the web client can retry once the + // provider is known. + logger.debug('[pi] set_model suppressed: provider unknown until get_state'); + throw new Error('Model cannot be applied yet: provider is not yet known'); + } else if (requestedModel.modelId === null) { + // Clearing the model needs no Pi RPC (nothing to confirm), so commit + // immediately. This path is not reachable from the web Pi picker today. + piSession.currentModel = null; + piSession.currentProvider = null; + } + } + if (requestedThinkingLevel !== undefined) { + const level = requestedThinkingLevel ?? 'off'; + await sendPiRpcAndWait(piSession, transport, { type: 'set_thinking_level', level }); + piSession.currentThinkingLevel = requestedThinkingLevel; + } + piSession.pushKeepAlive(); + + // Return provider-qualified model so the hub persists piSelectedModel. + // A bare modelId string would make applySessionConfig clear the + // provider metadata (object check fails), defeating Fix #3. + const appliedModel = piSession.currentModel && piSession.currentProvider + ? { provider: piSession.currentProvider, modelId: piSession.currentModel } + : piSession.currentModel; + + return { + applied: { + model: appliedModel, + effort: piSession.currentThinkingLevel, + }, + }; + }, { poisonOnError: (error) => error instanceof PiRpcTimeoutError }); + } catch (error) { + if (error instanceof PiRpcTimeoutError) { + failNativeStartup(new Error(`Pi configuration outcome is indeterminate: ${error.message}`)); } + throw error; } - if (requestedThinkingLevel !== undefined) { - const level = requestedThinkingLevel ?? 'off'; - await sendPiRpcAndWait(piSession, transport, { type: 'set_thinking_level', level }); - piSession.currentThinkingLevel = requestedThinkingLevel; - } - piSession.pushKeepAlive(); - - // Return provider-qualified model so the hub persists piSelectedModel. - // A bare modelId string would make applySessionConfig clear the - // provider metadata (object check fails), defeating Fix #3. - const appliedModel = piSession.currentModel && piSession.currentProvider - ? { provider: piSession.currentProvider, modelId: piSession.currentModel } - : piSession.currentModel; - - return { - applied: { - model: appliedModel, - effort: piSession.currentThinkingLevel, - }, - }; }); // --- Pi model discovery RPC --- @@ -396,54 +741,233 @@ export async function runPi(opts: { ); // --- User message handler --- + // Preparation reads image files asynchronously. A single promise chain keeps + // attachment completion order identical to user-message arrival order. apiSession.onUserMessage((message, localId) => { - const formattedText = formatPiUserMessage( - message.content.text, - message.content.attachments, - piSession.cachedPiCommands, - ); - // Gate the send behind Pi startup readiness. A prompt POSTed immediately - // after spawn (supported handoff pattern — hapi-ping-peer, intake - // scripts) would otherwise reach Pi before its initial get_state finishes - // and wedge the turn. runWhenReady delivers now if ready, else buffers - // FIFO until the first get_state response drains it (issue #1143). - // piIsStreaming is evaluated at delivery time so a message buffered - // before startup still takes the prompt path. - piSession.runWhenReady(() => { - if (piSession.piIsStreaming) { - // Steer does not start a new turn, so the localId would never be - // drained by turn_start. Mark it consumed immediately so it does - // not poison the FIFO for the next real prompt. - transport.send({ type: 'steer', message: formattedText }); - if (localId) piSession.emitMessagesConsumed([localId]); - } else { - if (localId) pendingLocalIds.push(localId); - transport.send({ type: 'prompt', message: formattedText }); + const deliveryMode = message.meta?.deliveryMode ?? 'queue'; + // Reserve message order and, for native steering, the exact turn before + // attachment preparation or a runtime-lock wait can yield execution. + const outboundSequence = nextOutboundSequence++; + const targetStreamingGeneration = deliveryMode === 'steer' + ? piSession.currentStreamingGeneration + : null; + if (localId) preparingLocalIds.add(localId); + preparationChain = preparationChain.then(async () => { + if (localId && cancelledWhilePreparing.delete(localId)) { + preparingLocalIds.delete(localId); + return; } - }, localId); + const prepared = await preparePiUserMessage( + message.content.text, + message.content.attachments, + piSession.cachedPiCommands, + { + authorizeImagePath: (path) => isPathWithinUploadDir(path, apiSession.sessionId), + authorizeOpenedImage: (path, identity) => isAuthorizedUploadFile(path, apiSession.sessionId, identity), + }, + ); + if (localId) { + preparingLocalIds.delete(localId); + if (cancelledWhilePreparing.delete(localId)) return; + } + for (const imageReadError of prepared.imageReadErrors) { + piSession.sendSessionEvent({ type: 'message', message: imageReadError }); + } + if (!prepared.message.trim() && prepared.images.length === 0) { + // An image-only prompt can lose every image during asynchronous + // preparation. Never issue an empty Pi prompt; resolve the HAPI + // queue row so QueuedMessagesBar cannot remain stuck forever. + if (localId) piSession.emitMessagesConsumed([localId], { clearQueuedThinkingGrace: true }); + return; + } + const entry = { + message: prepared.message, + images: prepared.images, + outboundSequence, + ...(localId ? { localId } : {}), + }; + if (deliveryMode === 'steer') { + steerDispatcher?.enqueue({ ...entry, targetStreamingGeneration }); + } else { + promptQueue.enqueue(entry); + pumpPromptQueue(); + } + }).catch((error: unknown) => { + const wasCancelled = localId ? cancelledWhilePreparing.delete(localId) : false; + if (localId) preparingLocalIds.delete(localId); + const detail = error instanceof Error ? error.message : String(error); + piSession.sendSessionEvent({ type: 'message', message: `Failed to prepare Pi prompt: ${detail}` }); + if (localId && !wasCancelled) { + piSession.emitMessagesConsumed([localId], { clearQueuedThinkingGrace: true }); + } + }); }); // --- Cancel-queued-message handler --- - // A prompt buffered during the startup window (runWhenReady) can be cancelled - // by the hub before it drains. Without this, ApiSessionClient acks - // removed:false, the hub marks the row invoked, yet the closure would still - // fire the cancelled prompt on get_state. Dropping it from the buffer keeps - // the queued-message cancel contract intact (issue #1143 review — MAJOR). - // Once sent to Pi it cannot be recalled — return false (best-effort), which - // matches the other agents' queue.cancelByLocalId semantics. - apiSession.onCancelQueuedMessage((localId) => piSession.cancelBufferedMessage(localId)); + // HAPI owns both asynchronous preparation and the local FIFO. A cancel may + // arrive before image preparation completes or while the ready/settled queue + // holds the item; both cases are removable. Once sent to Pi it is best-effort + // only, matching other harness queues. + apiSession.onCancelQueuedMessage((localId) => { + if (preparingLocalIds.has(localId)) { + cancelledWhilePreparing.add(localId); + return true; + } + return promptQueue.cancelByLocalId(localId) || steerDispatcher?.cancelByLocalId(localId) === true; + }); // --- Abort handler --- // Only cancel the current turn, keep session alive for next prompt. // Pi's `abort` command cancels the active turn but the process stays in RPC mode. + let abortPromise: Promise<{ success: true }> | null = null; apiSession.rpcHandlerManager.registerHandler(RPC_METHODS.Abort, async () => { - transport.send({ type: 'abort' }); - piSession.piIsStreaming = false; - piSession.updateThinkingState(false); - if (pendingLocalIds.length > 0) { - piSession.emitMessagesConsumed([pendingLocalIds.shift()!]); + if (abortPromise) return await abortPromise; + piSession.assertNoHistoryTransaction('abort Pi'); + const deadlineAt = Date.now() + PI_ABORT_OPERATION_TIMEOUT_MS; + abortInFlight = true; + transportEvents?.cancelPendingExtensionUi('Pi prompt aborted', { sendResponse: true }); + const abortedLocalId = activePromptLocalId; + const hadActivePrompt = promptCommandInFlight; + const hadNativeWork = hadActivePrompt || piSession.piIsStreaming; + const abortBoundary = hadActivePrompt ? createAbortGuard(deadlineAt) : null; + if (abortBoundary) { + activeAbortGuard = abortBoundary.guard; + // The runtime lock can be queued behind a config mutation; attach a + // handler immediately so a deadline rejection is never unhandled. + void abortBoundary.promise.catch(() => {}); } - return { success: true }; + // Reserve our mutation slot synchronously, after installing the pump + // barrier. A prompt that settles while an older config mutation owns the + // lock therefore cannot start the next FIFO item or change our target. + const acquireRuntimeMutation = hadNativeWork + ? piSession.acquireRuntimeMutation() + : null; + + let streamingInvalidatedUnderLock = false; + abortPromise = (async (): Promise<{ success: true }> => { + let releaseRuntimeMutation: (() => void) | null = null; + let nativeAbortIssued = false; + let firstAbortConfirmed = false; + try { + if (!acquireRuntimeMutation) { + return { success: true }; + } + + let lockTimer: ReturnType | null = null; + try { + releaseRuntimeMutation = await Promise.race([ + acquireRuntimeMutation, + new Promise((_, reject) => { + lockTimer = setTimeout(() => reject(new PiRpcTimeoutError( + 'abort-lock', + -1, + PI_ABORT_OPERATION_TIMEOUT_MS, + )), Math.max(1, deadlineAt - Date.now())); + lockTimer.unref?.(); + }), + ]); + } finally { + if (lockTimer) clearTimeout(lockTimer); + } + + if (!abortBoundary?.guard.settled) { + const remainingMs = Math.floor(deadlineAt - Date.now()); + if (remainingMs <= 0) { + throw new PiRpcTimeoutError('abort', -1, PI_ABORT_OPERATION_TIMEOUT_MS); + } + nativeAbortIssued = true; + await sendPiRpcAndWait( + piSession, + transport, + { type: 'abort' }, + Math.min(10_000, remainingMs), + ); + firstAbortConfirmed = true; + } + if (abortBoundary) await abortBoundary.promise; + // Invalidate the aborted turn while still holding the runtime + // mutation lease. A queued steer waiter must observe idle (and + // therefore fall back) before it can acquire this same lease. + piSession.updateThinkingState(false); + streamingInvalidatedUnderLock = true; + return { success: true }; + } catch (error) { + const detail = error instanceof Error ? error.message : String(error); + const targetSettled = abortBoundary?.guard.settled === true; + const mayStillStart = Boolean( + !(error instanceof PiRpcTimeoutError) + && !firstAbortConfirmed + && abortBoundary + && !abortBoundary.guard.lifecycleStartedAtRequest + && isPiNoActiveAbortError(detail) + ); + if (mayStillStart) { + try { + // Pi can reject abort before async prompt preflight + // produces agent_start. Keep the guard for the complete + // stability window so a late start is compensated even + // when the 1s lifecycle-missing fallback has not fired. + await abortBoundary!.promise; + // Match the ordinary abort-success path: invalidate the + // target generation before releasing this mutex so a + // queued steer cannot enter the just-aborted turn. + piSession.updateThinkingState(false); + streamingInvalidatedUnderLock = true; + return { success: true }; + } catch (guardError) { + const fatal = new Error(`Pi abort failed closed: ${guardError instanceof Error ? guardError.message : String(guardError)}`); + failNativeStartup(fatal); + throw fatal; + } + } + if (error instanceof PiRpcTimeoutError || firstAbortConfirmed || (!nativeAbortIssued && !targetSettled)) { + const fatal = new Error(`Pi abort failed closed: ${detail}`); + failNativeStartup(fatal); + throw fatal; + } + if (activeAbortGuard === abortBoundary?.guard) { + clearTimeout(abortBoundary.guard.timer); + activeAbortGuard = null; + } + abortInFlight = false; + if (targetSettled) { + setPromptCommandInFlight(false); + activePromptLocalId = undefined; + agentLifecycleStarted = false; + } + pumpPromptQueue(); + piSession.sendSessionEvent({ type: 'message', message: `Pi abort failed: ${detail}` }); + throw new Error(`Pi abort failed: ${detail}`); + } finally { + releaseRuntimeMutation?.(); + if (!releaseRuntimeMutation && acquireRuntimeMutation) { + void acquireRuntimeMutation.then((lateRelease) => lateRelease()); + } + } + })().then((result) => { + if (activeAbortGuard === abortBoundary?.guard) { + clearTimeout(abortBoundary.guard.timer); + activeAbortGuard = null; + } + transportEvents?.abortPromptLifecycle(); + if (abortedLocalId) { + conversationHistory.rejectPendingEntry(abortedLocalId); + if (pendingLocalIds[0] === abortedLocalId) { + pendingLocalIds.shift(); + piSession.emitMessagesConsumed([abortedLocalId], { clearQueuedThinkingGrace: true }); + } + } + activePromptLocalId = undefined; + agentLifecycleStarted = false; + // The native-work path already invalidated the generation under + // the runtime mutex. The no-work fast path reaches only this hook. + if (!streamingInvalidatedUnderLock) piSession.updateThinkingState(false); + setPromptCommandInFlight(false); + abortInFlight = false; + pumpPromptQueue(); + return result; + }).finally(() => { abortPromise = null; }); + return await abortPromise; }); // --- Switch handler --- @@ -469,8 +993,14 @@ export async function runPi(opts: { if (piSession.expectedNativeSessionId) { failNativeStartup(new Error(`Pi native resume did not become ready within ${PI_READY_FALLBACK_MS}ms`)); } else { - logger.debug('[pi] get_state ready signal not seen within grace — draining buffered messages'); - piSession.markReady(); + logger.debug('[pi] get_state ready signal not seen within grace — establishing history baseline before draining buffered messages'); + void initializeHistoryBaseline() + .catch(() => {}) + .finally(() => { + if (cleanupInitiated) return; + piSession.markReady(); + pumpPromptQueue(); + }); } }, PI_READY_FALLBACK_MS); readyFallback.unref?.(); @@ -492,14 +1022,20 @@ export async function runPi(opts: { if (startupThinkingLevel) { void (async () => { try { - await sendPiRpcAndWait(piSession, transport, { - type: 'set_thinking_level', - level: startupThinkingLevel, - }); - piSession.currentThinkingLevel = startupThinkingLevel; - piSession.pushKeepAlive(); + await piSession.runRuntimeMutation(async () => { + await sendPiRpcAndWait(piSession, transport, { + type: 'set_thinking_level', + level: startupThinkingLevel, + }); + piSession.currentThinkingLevel = startupThinkingLevel; + piSession.pushKeepAlive(); + }, { poisonOnError: (error) => error instanceof PiRpcTimeoutError }); logger.debug(`[pi] Startup effort applied: ${startupThinkingLevel}`); } catch (error) { + if (error instanceof PiRpcTimeoutError) { + failNativeStartup(new Error(`Pi startup effort outcome is indeterminate: ${error.message}`)); + return; + } logger.debug(`[pi] Startup effort rejected, keeping Pi default: ${error instanceof Error ? error.message : String(error)}`); } })(); diff --git a/cli/src/pi/schemas.ts b/cli/src/pi/schemas.ts index e588100b..2a4662bf 100644 --- a/cli/src/pi/schemas.ts +++ b/cli/src/pi/schemas.ts @@ -12,7 +12,7 @@ import { z } from 'zod'; import { PI_THINKING_LEVELS } from '@hapi/protocol'; import type { PiModelSummary } from '@hapi/protocol/apiTypes'; -import type { PiContextUsage } from './types'; +import type { PiContextUsage, PiExtensionUiRequest } from './types'; // ============================================================================ // 字段级容错 schema @@ -59,7 +59,163 @@ const asOptThinkingLevelMap = z.unknown().optional().transform((v): Record { + // Legacy Pi used auto_compaction_* for the same maintenance lifecycle that + // current Pi calls compaction_*. Normalize while decoding stdout so every + // downstream event consumer has one canonical protocol vocabulary. + if (event.type === 'auto_compaction_start') { + return { + ...event, + type: 'compaction_start', + reason: event.reason === 'manual' || event.reason === 'threshold' || event.reason === 'overflow' + ? event.reason + : 'threshold', + }; + } + if (event.type === 'auto_compaction_end') { + const { errorMessage: _legacyErrorMessage, ...rest } = event; + return { + ...rest, + type: 'compaction_end', + reason: event.reason === 'manual' || event.reason === 'threshold' || event.reason === 'overflow' + ? event.reason + : 'threshold', + aborted: typeof event.aborted === 'boolean' ? event.aborted : false, + willRetry: typeof event.willRetry === 'boolean' ? event.willRetry : false, + ...(typeof event.errorMessage === 'string' ? { errorMessage: event.errorMessage } : {}), + }; + } + return event; +}); + +// ============================================================================ +// Extension UI requests +// ============================================================================ + +const PiExtensionUiRequestBaseSchema = z.object({ + type: z.literal('extension_ui_request'), + id: z.string().min(1), +}); + +export const PiExtensionUiRequestSchema = z.discriminatedUnion('method', [ + PiExtensionUiRequestBaseSchema.extend({ + method: z.literal('select'), + title: z.string(), + options: z.array(z.string()), + timeout: z.number().finite().nonnegative().optional(), + }), + PiExtensionUiRequestBaseSchema.extend({ + method: z.literal('confirm'), + title: z.string(), + message: z.string(), + timeout: z.number().finite().nonnegative().optional(), + }), + PiExtensionUiRequestBaseSchema.extend({ + method: z.literal('input'), + title: z.string(), + placeholder: z.string().optional(), + timeout: z.number().finite().nonnegative().optional(), + }), + PiExtensionUiRequestBaseSchema.extend({ + method: z.literal('editor'), + title: z.string(), + prefill: z.string().optional(), + }), + PiExtensionUiRequestBaseSchema.extend({ + method: z.literal('notify'), + message: z.string(), + notifyType: z.enum(['info', 'warning', 'error']).optional(), + }), + PiExtensionUiRequestBaseSchema.extend({ + method: z.literal('setStatus'), + statusKey: z.string(), + statusText: z.string().optional(), + }), + PiExtensionUiRequestBaseSchema.extend({ + method: z.literal('setWidget'), + widgetKey: z.string(), + widgetLines: z.array(z.string()).optional(), + widgetPlacement: z.enum(['aboveEditor', 'belowEditor']).optional(), + }), + PiExtensionUiRequestBaseSchema.extend({ + method: z.literal('setTitle'), + title: z.string(), + }), + PiExtensionUiRequestBaseSchema.extend({ + method: z.literal('set_editor_text'), + text: z.string(), + }), +]) satisfies z.ZodType; + +const PiCompactionStartEventSchema = z.object({ + type: z.literal('compaction_start'), + reason: z.enum(['manual', 'threshold', 'overflow']), +}); + +const PiCompactionEndEventSchema = z.object({ + type: z.literal('compaction_end'), + reason: z.enum(['manual', 'threshold', 'overflow']), + aborted: z.boolean(), + willRetry: z.boolean(), + errorMessage: z.string().optional(), +}); + +// Legacy Pi used the auto_compaction_* aliases. Their payloads are equivalent +// in practice, but accept omitted lifecycle detail from older extensions so +// the maintenance gate can still block a legacy agent_end settlement. +const PiLegacyAutoCompactionStartEventSchema = z.object({ + type: z.literal('auto_compaction_start'), + reason: z.enum(['manual', 'threshold', 'overflow']).optional().default('threshold'), +}).passthrough().transform(({ type: _type, ...event }) => ({ + ...event, + type: 'compaction_start' as const, +})); + +const PiLegacyAutoCompactionEndEventSchema = z.object({ + type: z.literal('auto_compaction_end'), + reason: z.enum(['manual', 'threshold', 'overflow']).optional().default('threshold'), + aborted: z.boolean().optional().default(false), + willRetry: z.boolean().optional().default(false), + errorMessage: z.string().optional(), +}).passthrough().transform(({ type: _type, ...event }) => ({ + ...event, + type: 'compaction_end' as const, +})); + +const PiAutoRetryStartEventSchema = z.object({ + type: z.literal('auto_retry_start'), + attempt: z.number().int().positive(), + maxAttempts: z.number().int().positive(), + delayMs: z.number().finite().nonnegative(), + errorMessage: z.string(), +}); + +const PiAutoRetryEndEventSchema = z.object({ + type: z.literal('auto_retry_end'), + success: z.boolean(), + attempt: z.number().int().positive(), + finalError: z.string().optional(), +}); + +const PiSummarizationRetryScheduledEventSchema = z.object({ + type: z.literal('summarization_retry_scheduled'), + attempt: z.number().int().positive(), + maxAttempts: z.number().int().positive(), + delayMs: z.number().finite().nonnegative(), + errorMessage: z.string(), +}); + +export const PiLifecycleEventSchema = z.union([ + PiCompactionStartEventSchema, + PiCompactionEndEventSchema, + PiLegacyAutoCompactionStartEventSchema, + PiLegacyAutoCompactionEndEventSchema, + PiAutoRetryStartEventSchema, + PiAutoRetryEndEventSchema, + PiSummarizationRetryScheduledEventSchema, + z.object({ type: z.literal('summarization_retry_attempt_start'), source: z.enum(['branchSummary', 'compaction']) }).passthrough(), + z.object({ type: z.literal('summarization_retry_finished') }), +]); // ============================================================================ // Pi Response Event (stdout response) @@ -165,8 +321,10 @@ export const PiStateDataSchema = z.object({ provider: z.string().optional(), }).passthrough().optional(), sessionId: z.string().optional(), + sessionFile: z.string().optional(), thinkingLevel: z.string().optional(), steeringMode: z.enum(['all', 'one-at-a-time']).optional(), + isStreaming: z.boolean().optional(), }).passthrough(); // ============================================================================ @@ -209,6 +367,38 @@ export const PiAssistantMessageEventSchema = z.object({ contentIndex: z.number().optional(), }).passthrough(); +export const PiToolExecutionStartEventSchema = z.object({ + type: z.literal('tool_execution_start'), + toolCallId: z.string().min(1), + toolName: z.string().min(1), + args: z.unknown(), +}); + +export const PiToolExecutionUpdateEventSchema = z.object({ + type: z.literal('tool_execution_update'), + toolCallId: z.string().min(1), + toolName: z.string().min(1), + args: z.unknown(), + partialResult: z.unknown(), +}); + +export const PiToolExecutionEndEventSchema = z.object({ + type: z.literal('tool_execution_end'), + toolCallId: z.string().min(1), + toolName: z.string().min(1), + result: z.unknown(), + isError: z.boolean(), +}); + +export const PiAgentEndEventSchema = z.object({ + type: z.literal('agent_end'), + willRetry: z.boolean().optional(), +}).passthrough(); + +export const PiAgentSettledEventSchema = z.object({ + type: z.literal('agent_settled'), +}); + // ============================================================================ // Parse helpers — replace hand-written type guards in loop.ts // ============================================================================ diff --git a/cli/src/pi/session.test.ts b/cli/src/pi/session.test.ts index 7bcca139..d0d6975d 100644 --- a/cli/src/pi/session.test.ts +++ b/cli/src/pi/session.test.ts @@ -95,6 +95,44 @@ describe('PiSession ready gate', () => { expect(nativeSession.client.emitSessionReady).toHaveBeenCalledTimes(1); }); + it('announces native-ready before async history baseline drains prompts', async () => { + const session = createMockSession(); + let finishBaseline!: () => void; + session.setNativeReadyPreparation(() => new Promise((resolve) => { finishBaseline = resolve; })); + const sent = vi.fn(); + session.runWhenReady(sent); + + session.markNativeReady(); + + expect(session.client.emitSessionReady).toHaveBeenCalledTimes(1); + expect(session.isReady).toBe(false); + expect(sent).not.toHaveBeenCalled(); + + finishBaseline(); + await Promise.resolve(); + await Promise.resolve(); + expect(session.isReady).toBe(true); + expect(sent).toHaveBeenCalledTimes(1); + }); + + it('does not drain buffered work when cleanup cancels an in-flight ready preparation', async () => { + const session = createMockSession(); + let finishBaseline!: () => void; + session.setNativeReadyPreparation(() => new Promise((resolve) => { finishBaseline = resolve; })); + const sent = vi.fn(); + session.runWhenReady(sent); + + session.markNativeReady(); + session.cancelReadyGate(); + finishBaseline(); + await Promise.resolve(); + await Promise.resolve(); + session.runWhenReady(sent); + + expect(session.isReady).toBe(false); + expect(sent).not.toHaveBeenCalled(); + }); + it('preserves FIFO across mixed buffered + post-ready enqueues', () => { const session = createMockSession(); const order: string[] = []; @@ -148,3 +186,124 @@ describe('PiSession cancelBufferedMessage', () => { expect(session.cancelBufferedMessage('id-1')).toBe(false); }); }); + +describe('PiSession history transaction gate', () => { + it('defers a prompt during clone/restore and drains it FIFO after source restoration', () => { + const session = createMockSession(); + const sent: string[] = []; + const release = session.beginHistoryTransaction(); + session.runWhenHistoryIdle(() => sent.push('first'), 'first'); + session.runWhenHistoryIdle(() => sent.push('second'), 'second'); + + expect(sent).toEqual([]); + expect(session.cancelBufferedMessage('second')).toBe(true); + release(); + + expect(sent).toEqual(['first']); + expect(session.isHistoryTransactionActive).toBe(false); + }); +}); + +describe('PiSession runtime mutation mutex', () => { + it('serializes config-like mutations FIFO and releases after the active mutation settles', async () => { + const session = createMockSession(); + const order: string[] = []; + let finishFirst!: () => void; + + const first = session.runRuntimeMutation(async () => { + order.push('first-start'); + await new Promise((resolve) => { finishFirst = resolve; }); + order.push('first-end'); + }); + const second = session.runRuntimeMutation(async () => { + order.push('second-start'); + }); + + await vi.waitFor(() => expect(order).toEqual(['first-start'])); + expect(order).not.toContain('second-start'); + + finishFirst(); + await Promise.all([first, second]); + + expect(order).toEqual(['first-start', 'first-end', 'second-start']); + }); + + it('keeps the lease poisoned when a mutation outcome is indeterminate', async () => { + const session = createMockSession(); + const timeout = new Error('timed out'); + let secondStarted = false; + + await expect(session.runRuntimeMutation( + async () => { throw timeout; }, + { poisonOnError: (error) => error === timeout }, + )).rejects.toBe(timeout); + void session.runRuntimeMutation(async () => { secondStarted = true; }); + await Promise.resolve(); + await Promise.resolve(); + + expect(secondStarted).toBe(false); + }); +}); + +describe('PiSession streaming generation', () => { + it('increments only on false-to-true transitions and is hidden while idle', () => { + const session = createMockSession(); + + expect(session.currentStreamingGeneration).toBeNull(); + session.updateThinkingState(true); + const firstGeneration = session.currentStreamingGeneration; + expect(firstGeneration).toBe(1); + + // Repeated lifecycle/get_state confirmations for the same turn do not + // mint a new identity. + session.updateThinkingState(true); + expect(session.currentStreamingGeneration).toBe(firstGeneration); + + session.applyNativeRuntimeState({ isStreaming: false }); + expect(session.currentStreamingGeneration).toBeNull(); + session.applyNativeRuntimeState({ isStreaming: true }); + expect(session.currentStreamingGeneration).toBe(2); + }); + + it('invalidates a streaming generation when rewind commits a new native branch', () => { + const session = createMockSession(); + session.updateThinkingState(true); + const sourceGeneration = session.currentStreamingGeneration; + + session.commitNativeSessionState( + { sessionId: 'rewound-session', sessionFile: '/tmp/rewound-session.jsonl' }, + { isStreaming: true }, + ); + + expect(session.currentStreamingGeneration).not.toBe(sourceGeneration); + expect(session.currentStreamingGeneration).toBe(2); + }); +}); + +describe('PiSession native runtime reconciliation', () => { + it('preserves an omitted provider only when the reported model is unchanged', () => { + const session = createMockSession(); + session.currentModel = 'same-model'; + session.currentProvider = 'known-provider'; + + session.applyNativeRuntimeState({ model: 'same-model' }); + expect(session.currentProvider).toBe('known-provider'); + + session.applyNativeRuntimeState({ model: 'different-model' }); + expect(session.currentProvider).toBeNull(); + }); + + it('infers an omitted provider only from a unique available-model match', () => { + const session = createMockSession(); + session.currentModel = 'old-model'; + session.currentProvider = 'old-provider'; + session.cachedPiModels = [ + { modelId: 'new-model', provider: 'new-provider' }, + { modelId: 'other-model', provider: 'other-provider' }, + ]; + + session.applyNativeRuntimeState({ model: 'new-model' }); + + expect(session.currentProvider).toBe('new-provider'); + }); +}); diff --git a/cli/src/pi/session.ts b/cli/src/pi/session.ts index 7e53b69f..0df63a19 100644 --- a/cli/src/pi/session.ts +++ b/cli/src/pi/session.ts @@ -4,6 +4,19 @@ import type { PiCommandSummary, PiThinkingLevel } from './types'; import type { PiModelSummary } from '@hapi/protocol/apiTypes'; import type { PiRpcResolver } from './loop'; +/** + * The parts of `get_state` which must move atomically with a native session + * identity transition. `undefined` means the Pi version did not report that + * field, so an already-confirmed value must remain intact. + */ +export type PiNativeRuntimeState = { + model?: string | null; + provider?: string | null; + thinkingLevel?: PiThinkingLevel | null; + steeringMode?: 'all' | 'one-at-a-time'; + isStreaming?: boolean; +}; + /** * Pi session state and hub communication wrapper. * @@ -36,10 +49,15 @@ export class PiSession { readonly initialModel: string | null; // A runner/native resume must prove that Pi loaded this exact session with // a non-empty get_state sessionId. Missing or contradictory IDs fail closed. - readonly expectedNativeSessionId: string | null; + expectedNativeSessionId: string | null; + currentNativeSessionFile: string | null = null; - // Streaming state - piIsStreaming = false; + // Streaming state. A generation identifies one concrete Pi turn, rather + // than merely the transient boolean reported by get_state/lifecycle events. + // Native steers capture it at arrival and must not cross into a later turn. + private _piIsStreaming = false; + private streamingGeneration = 0; + private promptInFlight = false; currentSteeringMode: 'all' | 'one-at-a-time' = 'all'; // Cached data from Pi @@ -56,11 +74,21 @@ export class PiSession { // are queued via runWhenReady() and drained FIFO once markReady() fires (on // the first get_state response). private piReady = false; + private readyCancelled = false; private nativeReadyAnnounced = false; + private nativeReadyPreparation: (() => Promise) | null = null; + private nativeReadyPreparationStarted = false; + private historyTransaction: symbol | null = null; + // All commands that mutate the native Pi runtime (history clone/fork/switch, + // model/thinking changes, and abort) share one FIFO lock. A history action + // closes its prompt gate before waiting on this tail, which prevents a later + // runtime mutation from slipping between its final source sync and fork. + private runtimeMutationTail: Promise = Promise.resolve(); // Buffered sends carry their localId so a cancel-queued-message that arrives // while a prompt is still held (before drain) can drop it instead of firing // a cancelled prompt on markReady (issue #1143 review — MAJOR). private readyQueue: Array<{ localId?: string; fn: () => void }> = []; + private historyDeferredQueue: Array<{ localId?: string; fn: () => void }> = []; private keepAliveInterval: NodeJS.Timeout | null = null; @@ -104,11 +132,216 @@ export class PiSession { return this.nativeReadyAnnounced; } + /** + * Establish a stable native baseline before buffered prompts are released. + * Pi's history cursor must be read before the first web prompt is allowed + * to append, otherwise an old duplicate user message could be paired with + * a new HAPI localId. + */ + setNativeReadyPreparation(fn: () => Promise): void { + this.nativeReadyPreparation = fn; + } + + /** Prevent any late startup/history completion from draining outbound work. */ + cancelReadyGate(): void { + this.readyCancelled = true; + this.readyQueue = []; + this.historyDeferredQueue = []; + } + + get isHistoryTransactionActive(): boolean { + return this.historyTransaction !== null; + } + + /** True from prompt dispatch until rejection, abort, or true settlement. */ + get hasPromptInFlight(): boolean { + return this.promptInFlight; + } + + get piIsStreaming(): boolean { + return this._piIsStreaming; + } + + set piIsStreaming(value: boolean) { + // Keep direct state reconciliation (including test adapters) on the + // same transition invariant as updateThinkingState(). + if (value && !this._piIsStreaming) this.streamingGeneration += 1; + this._piIsStreaming = value; + } + + /** Exact identity of the active Pi stream; absent while Pi is idle. */ + get currentStreamingGeneration(): number | null { + return this._piIsStreaming ? this.streamingGeneration : null; + } + + setPromptInFlight(value: boolean): void { + this.promptInFlight = value; + } + + beginHistoryTransaction(): (options?: { drain?: boolean }) => void { + if (this.historyTransaction) throw new Error('Conversation history action already in progress'); + const token = Symbol('pi-history-transaction'); + this.historyTransaction = token; + return (options = {}) => { + if (this.historyTransaction !== token) return; + this.historyTransaction = null; + const deferred = this.historyDeferredQueue; + this.historyDeferredQueue = []; + if (options.drain === false) return; + for (const { fn } of deferred) fn(); + }; + } + + /** + * Acquire exclusive ownership of the native Pi runtime. Callers which need + * a transaction spanning multiple RPCs (conversation history) can retain + * the release callback; single RPC handlers should use runRuntimeMutation. + */ + async acquireRuntimeMutation(): Promise<() => void> { + const previous = this.runtimeMutationTail; + let releaseCurrent!: () => void; + this.runtimeMutationTail = new Promise((resolve) => { + releaseCurrent = resolve; + }); + await previous; + + let released = false; + return () => { + if (released) return; + released = true; + releaseCurrent(); + }; + } + + /** Serialize one Pi runtime mutation behind any active history operation. */ + async runRuntimeMutation( + operation: () => Promise, + options: { poisonOnError?: (error: unknown) => boolean } = {}, + ): Promise { + const release = await this.acquireRuntimeMutation(); + let releaseLease = true; + try { + return await operation(); + } catch (error) { + if (options.poisonOnError?.(error)) releaseLease = false; + throw error; + } finally { + if (releaseLease) release(); + } + } + + assertNoHistoryTransaction(operation: string): void { + if (this.historyTransaction) { + throw new Error(`Cannot ${operation} while a conversation history action is in progress`); + } + } + + /** Queue ordinary outbound work until native source identity is restored. */ + runWhenHistoryIdle(fn: () => void, localId?: string): void { + if (!this.historyTransaction) { + fn(); + return; + } + this.historyDeferredQueue.push({ fn, localId }); + } + matchesExpectedNativeSessionId(actualSessionId: string | undefined): boolean { if (!this.expectedNativeSessionId) return true; return Boolean(actualSessionId) && actualSessionId === this.expectedNativeSessionId; } + /** + * Commit an in-process native session transition (Pi rewind creates a new + * branched session file). Future get_state validation must target the new + * id before its metadata is exposed to the hub. + */ + commitNativeSessionIdentity( + identity: { sessionId: string; sessionFile: string }, + metadataUpdater?: (metadata: Metadata) => Metadata, + ): void { + this.commitNativeSessionState(identity, {}, metadataUpdater); + } + + /** + * Commit a confirmed get_state snapshot after a native session transition. + * This intentionally updates runtime values before metadata is made visible + * so the next keepAlive cannot report the old branch configuration. + */ + commitNativeSessionState( + identity: { sessionId: string; sessionFile: string }, + runtime: PiNativeRuntimeState, + metadataUpdater?: (metadata: Metadata) => Metadata, + ): void { + const identityChanged = this.expectedNativeSessionId !== identity.sessionId + || this.currentNativeSessionFile !== identity.sessionFile; + // A rewind can replace the native branch while both snapshots report + // streaming=true. Invalidate any steer captured for the old branch + // even though the boolean never passed through an observable idle edge. + if (identityChanged && this._piIsStreaming && runtime.isStreaming !== false) { + this.streamingGeneration += 1; + } + this.expectedNativeSessionId = identity.sessionId; + this.currentNativeSessionFile = identity.sessionFile; + this.applyNativeRuntimeState(runtime); + this.updateMetadata((metadata) => { + let next = metadataUpdater?.({ + ...metadata, + piSessionId: identity.sessionId, + }) ?? { + ...metadata, + piSessionId: identity.sessionId, + }; + // A get_state model/provider pair is the authoritative provider- + // qualified selection for the newly active native branch. Preserve + // existing metadata only when an older Pi omitted both fields. + if (runtime.model !== undefined || runtime.provider !== undefined) { + next = { ...next }; + delete next.piSelectedModel; + if (this.currentModel && this.currentProvider) { + next.piSelectedModel = { + provider: this.currentProvider, + modelId: this.currentModel, + }; + } + } + return next; + }); + } + + /** Apply a confirmed state snapshot without changing native identity metadata. */ + applyNativeRuntimeState(runtime: PiNativeRuntimeState): void { + if (runtime.model !== undefined) { + if (runtime.provider === undefined && runtime.model !== this.currentModel) { + if (runtime.model === null) { + this.currentProvider = null; + } else { + const matchingProviders = new Set( + this.cachedPiModels + .filter((model) => model.modelId === runtime.model) + .map((model) => model.provider), + ); + // Never combine a newly reported model id with the previous + // branch's provider. Infer only from an unambiguous catalog. + this.currentProvider = matchingProviders.size === 1 + ? matchingProviders.values().next().value ?? null + : null; + } + } + this.currentModel = runtime.model; + } + if (runtime.provider !== undefined) this.currentProvider = runtime.provider; + if (runtime.thinkingLevel !== undefined) this.currentThinkingLevel = runtime.thinkingLevel; + if (runtime.steeringMode !== undefined) this.currentSteeringMode = runtime.steeringMode; + if (runtime.isStreaming !== undefined) this.piIsStreaming = runtime.isStreaming; + this.pushKeepAlive(); + } + + /** Wait for queued metadata writes before exposing a native history result. */ + async flushMetadata(timeoutMs: number = 5_000): Promise { + const flush = (this.client as Partial).flushMetadata + return flush ? await flush.call(this.client, timeoutMs) : true + } + /** * Run `fn` now if Pi startup is ready, else buffer it FIFO until markReady(). * Used to gate outbound prompt/steer sends so they never reach Pi before its @@ -116,6 +349,7 @@ export class PiSession { * cancel-queued-message can drop it while still buffered. */ runWhenReady(fn: () => void, localId?: string): void { + if (this.readyCancelled) return; if (this.piReady) { fn(); return; @@ -131,8 +365,13 @@ export class PiSession { */ cancelBufferedMessage(localId: string): boolean { const idx = this.readyQueue.findIndex((item) => item.localId === localId); - if (idx === -1) return false; - this.readyQueue.splice(idx, 1); + if (idx !== -1) { + this.readyQueue.splice(idx, 1); + return true; + } + const deferredIdx = this.historyDeferredQueue.findIndex((item) => item.localId === localId); + if (deferredIdx === -1) return false; + this.historyDeferredQueue.splice(deferredIdx, 1); return true; } @@ -142,7 +381,7 @@ export class PiSession { * Pi loaded a requested native session. */ markReady(): boolean { - if (this.piReady) return false; + if (this.piReady || this.readyCancelled) return false; this.piReady = true; const queued = this.readyQueue; this.readyQueue = []; @@ -156,10 +395,26 @@ export class PiSession { * considered successful. */ markNativeReady(): void { - this.markReady(); - if (this.nativeReadyAnnounced) return; + if (this.readyCancelled || this.nativeReadyAnnounced || this.nativeReadyPreparationStarted) return; + // The hub uses session-ready as the validated native identity fence. + // It must precede piSessionId metadata publication, while prompt drain + // may wait for the append-log baseline below. this.nativeReadyAnnounced = true; this.client.emitSessionReady(); + if (this.nativeReadyPreparation) { + this.nativeReadyPreparationStarted = true; + void this.nativeReadyPreparation() + // The history feature is optional. A failed baseline must not + // wedge a normal Pi session; it simply remains unpublished. + .catch(() => {}) + .finally(() => this.finishNativeReady()); + return; + } + this.finishNativeReady(); + } + + private finishNativeReady(): void { + this.markReady(); } startKeepAlive(): void { diff --git a/cli/src/pi/steerDispatcher.test.ts b/cli/src/pi/steerDispatcher.test.ts new file mode 100644 index 00000000..d273e0d0 --- /dev/null +++ b/cli/src/pi/steerDispatcher.test.ts @@ -0,0 +1,249 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { PiConversationHistory } from './conversationHistory'; +import { PiRpcResolver } from './loop'; +import { PiSession } from './session'; +import { PiSteerDispatcher } from './steerDispatcher'; +import type { PiTransport } from './piTransport'; + +function createHarness(options: { streaming?: boolean } = {}) { + const client = { + keepAlive: vi.fn(), + emitMessagesConsumed: vi.fn(), + sendSessionEvent: vi.fn(), + updateMetadata: vi.fn(), + emitSessionReady: vi.fn(), + }; + const session = new PiSession({ + api: {} as never, + client: client as never, + path: '/tmp/project', + logPath: '/tmp/pi.log', + startedBy: 'terminal', + startingMode: 'remote', + }); + session.markNativeReady(); + session.updateThinkingState(options.streaming ?? true); + session.rpcResolver = new PiRpcResolver(); + + const transport = { send: vi.fn() } as unknown as PiTransport; + const history = { + registerUserEntry: vi.fn(), + rejectPendingEntry: vi.fn(), + } as unknown as PiConversationHistory; + const enqueuePrompt = vi.fn(); + const onIndeterminateTimeout = vi.fn(); + const onPendingStateChange = vi.fn(); + const dispatcher = new PiSteerDispatcher({ + session, + transport, + conversationHistory: history, + enqueuePrompt, + onIndeterminateTimeout, + onPendingStateChange, + }); + return { client, dispatcher, enqueuePrompt, history, onIndeterminateTimeout, onPendingStateChange, session, transport }; +} + +function steerCommands(transport: PiTransport): Array<{ id: string; type: 'steer'; message: string; images?: unknown[] }> { + return (transport.send as ReturnType).mock.calls + .map(([command]) => command as { id: string; type: string; message: string; images?: unknown[] }) + .filter((command): command is { id: string; type: 'steer'; message: string; images?: unknown[] } => command.type === 'steer'); +} + +function resolveSteer(session: PiSession, command: { id: string; type: 'steer' }, success = true, error?: string): void { + session.rpcResolver!.resolveResponse({ + type: 'response', + id: command.id, + command: 'steer', + success, + ...(error ? { error } : {}), + }); +} + +afterEach(() => vi.useRealTimers()); + +describe('PiSteerDispatcher', () => { + it('streams an explicit steer, registers history before send, and consumes it without clearing main thinking', async () => { + const h = createHarness(); + h.dispatcher.enqueue({ + localId: 'steer-1', + message: 'change direction', + images: [{ type: 'image', data: 'aW1hZ2U=', mimeType: 'image/png' }], + outboundSequence: 1, + targetStreamingGeneration: h.session.currentStreamingGeneration, + }); + + await vi.waitFor(() => expect(steerCommands(h.transport)).toHaveLength(1)); + const command = steerCommands(h.transport)[0]!; + expect(command).toMatchObject({ type: 'steer', message: 'change direction', images: [{ mimeType: 'image/png' }] }); + expect(h.history.registerUserEntry).toHaveBeenCalledBefore(h.transport.send as ReturnType); + expect(h.client.emitMessagesConsumed).not.toHaveBeenCalled(); + + resolveSteer(h.session, command); + await vi.waitFor(() => expect(h.client.emitMessagesConsumed).toHaveBeenCalledWith(['steer-1'], undefined)); + expect(h.session.piIsStreaming).toBe(true); + }); + + it('sends steers FIFO by matching responses, not by agent settlement or steering mode', async () => { + const h = createHarness(); + h.session.currentSteeringMode = 'one-at-a-time'; + h.dispatcher.enqueue({ localId: 'steer-a', message: 'a', images: [], outboundSequence: 1, targetStreamingGeneration: h.session.currentStreamingGeneration }); + h.dispatcher.enqueue({ localId: 'steer-b', message: 'b', images: [], outboundSequence: 2, targetStreamingGeneration: h.session.currentStreamingGeneration }); + + await vi.waitFor(() => expect(steerCommands(h.transport)).toHaveLength(1)); + expect(steerCommands(h.transport)[0]).toMatchObject({ message: 'a' }); + resolveSteer(h.session, steerCommands(h.transport)[0]!); + await vi.waitFor(() => expect(steerCommands(h.transport)).toHaveLength(2)); + expect(steerCommands(h.transport)[1]).toMatchObject({ message: 'b' }); + expect(h.history.registerUserEntry).toHaveBeenNthCalledWith(1, 'steer-a'); + expect(h.history.registerUserEntry).toHaveBeenNthCalledWith(2, 'steer-b'); + }); + + it('maps native user entries from prompts and steers through one FIFO history registration', async () => { + const h = createHarness(); + const history = new PiConversationHistory(h.session, vi.fn()); + const dispatcher = new PiSteerDispatcher({ + session: h.session, + transport: h.transport, + conversationHistory: history, + enqueuePrompt: h.enqueuePrompt, + onIndeterminateTimeout: h.onIndeterminateTimeout, + }); + dispatcher.enqueue({ localId: 'steer-first', message: 'same text', images: [], outboundSequence: 1, targetStreamingGeneration: h.session.currentStreamingGeneration }); + dispatcher.enqueue({ localId: 'steer-second', message: 'same text', images: [], outboundSequence: 2, targetStreamingGeneration: h.session.currentStreamingGeneration }); + + await vi.waitFor(() => expect(steerCommands(h.transport)).toHaveLength(1)); + resolveSteer(h.session, steerCommands(h.transport)[0]!); + await vi.waitFor(() => expect(steerCommands(h.transport)).toHaveLength(2)); + resolveSteer(h.session, steerCommands(h.transport)[1]!); + await vi.waitFor(() => expect(h.client.emitMessagesConsumed).toHaveBeenCalledTimes(2)); + + history.observeEntry({ id: 'native-first', type: 'message', message: { role: 'user' } }); + history.observeEntry({ id: 'native-second', type: 'message', message: { role: 'user' } }); + expect(history.getEntryIds()).toEqual({ + 'steer-first': 'native-first', + 'steer-second': 'native-second', + }); + }); + + it('rejects an ABA streaming transition behind the shared runtime lock and falls back', async () => { + const h = createHarness(); + const releaseConfig = await h.session.acquireRuntimeMutation(); + const capturedGeneration = h.session.currentStreamingGeneration; + h.dispatcher.enqueue({ + localId: 'locked-steer', + message: 'wait for config', + images: [], + outboundSequence: 1, + targetStreamingGeneration: capturedGeneration, + }); + await Promise.resolve(); + expect(steerCommands(h.transport)).toHaveLength(0); + + h.session.updateThinkingState(false); + h.session.updateThinkingState(true); + expect(h.session.currentStreamingGeneration).not.toBe(capturedGeneration); + releaseConfig(); + await vi.waitFor(() => expect(h.enqueuePrompt).toHaveBeenCalledWith({ + localId: 'locked-steer', message: 'wait for config', images: [], outboundSequence: 1, + })); + expect(steerCommands(h.transport)).toHaveLength(0); + }); + + it('notifies the prompt pump after pending steer work drains', async () => { + const h = createHarness({ streaming: false }); + h.dispatcher.enqueue({ + localId: 'idle-steer', + message: 'ordinary turn now', + images: [], + outboundSequence: 1, + targetStreamingGeneration: null, + }); + h.onPendingStateChange.mockClear(); + + await vi.waitFor(() => expect(h.dispatcher.hasPending).toBe(false)); + + expect(h.enqueuePrompt).toHaveBeenCalledTimes(1); + expect(h.onPendingStateChange).toHaveBeenCalled(); + }); + + it('falls back to the normal prompt queue when a steer reaches dispatch after Pi is idle', async () => { + const h = createHarness({ streaming: false }); + h.dispatcher.enqueue({ localId: 'idle-steer', message: 'ordinary turn now', images: [], outboundSequence: 1, targetStreamingGeneration: h.session.currentStreamingGeneration }); + + await vi.waitFor(() => expect(h.enqueuePrompt).toHaveBeenCalledWith({ + localId: 'idle-steer', message: 'ordinary turn now', images: [], outboundSequence: 1, + })); + expect(steerCommands(h.transport)).toHaveLength(0); + expect(h.history.registerUserEntry).not.toHaveBeenCalled(); + }); + + it('removes failed native steers from history and clears only their queued thinking grace', async () => { + const h = createHarness(); + h.dispatcher.enqueue({ localId: 'failed-steer', message: 'will fail', images: [], outboundSequence: 1, targetStreamingGeneration: h.session.currentStreamingGeneration }); + await vi.waitFor(() => expect(steerCommands(h.transport)).toHaveLength(1)); + + resolveSteer(h.session, steerCommands(h.transport)[0]!, false, 'steer rejected'); + await vi.waitFor(() => expect(h.history.rejectPendingEntry).toHaveBeenCalledWith('failed-steer')); + expect(h.client.emitMessagesConsumed).toHaveBeenCalledWith( + ['failed-steer'], + { clearQueuedThinkingGrace: true }, + ); + expect(h.client.sendSessionEvent).toHaveBeenCalledWith({ + type: 'message', message: 'Pi steer failed: steer rejected', + }); + expect(h.session.piIsStreaming).toBe(true); + }); + + it('allows cancellation before native send but not after the steer has reached stdin', async () => { + const h = createHarness(); + let releaseMutation!: () => void; + const originalRunMutation = h.session.runRuntimeMutation.bind(h.session); + const runMutationSpy = vi.spyOn(h.session, 'runRuntimeMutation').mockImplementation(async (operation, options) => { + await new Promise((resolve) => { releaseMutation = resolve; }); + return await originalRunMutation(operation, options); + }); + h.dispatcher.enqueue({ localId: 'waiting-steer', message: 'waiting', images: [], outboundSequence: 1, targetStreamingGeneration: h.session.currentStreamingGeneration }); + await vi.waitFor(() => expect(h.session.runRuntimeMutation).toHaveBeenCalled()); + expect(h.dispatcher.cancelByLocalId('waiting-steer')).toBe(true); + expect(h.dispatcher.hasPending).toBe(false); + releaseMutation(); + await vi.waitFor(() => expect(h.session.runRuntimeMutation).toHaveBeenCalled()); + expect(steerCommands(h.transport)).toHaveLength(0); + + runMutationSpy.mockRestore(); + h.dispatcher.enqueue({ localId: 'sent-steer', message: 'sent', images: [], outboundSequence: 2, targetStreamingGeneration: h.session.currentStreamingGeneration }); + await vi.waitFor(() => expect(steerCommands(h.transport)).toHaveLength(1)); + expect(h.dispatcher.cancelByLocalId('sent-steer')).toBe(false); + resolveSteer(h.session, steerCommands(h.transport)[0]!); + }); + + it('poisons the shared runtime mutation lease and fails closed when a native steer times out', async () => { + vi.useFakeTimers(); + const h = createHarness(); + h.dispatcher.enqueue({ localId: 'timeout-steer', message: 'timeout', images: [], outboundSequence: 1, targetStreamingGeneration: h.session.currentStreamingGeneration }); + await vi.advanceTimersByTimeAsync(0); + expect(steerCommands(h.transport)).toHaveLength(1); + h.dispatcher.enqueue({ localId: 'after-timeout', message: 'must remain blocked', images: [], outboundSequence: 2, targetStreamingGeneration: h.session.currentStreamingGeneration }); + + await vi.advanceTimersByTimeAsync(10_001); + await vi.waitFor(() => expect(h.onIndeterminateTimeout).toHaveBeenCalledTimes(1)); + expect(h.client.emitMessagesConsumed).toHaveBeenCalledWith( + ['timeout-steer'], + { clearQueuedThinkingGrace: true }, + ); + expect(steerCommands(h.transport)).toHaveLength(1); + }); + + it('stops local steers during transport cleanup without falling back to prompts', async () => { + const h = createHarness(); + h.session.updateThinkingState(false); + h.dispatcher.stop(); + expect(h.onPendingStateChange).not.toHaveBeenCalled(); + h.dispatcher.enqueue({ localId: 'shutdown-steer', message: 'must not queue', images: [], outboundSequence: 1, targetStreamingGeneration: h.session.currentStreamingGeneration }); + await Promise.resolve(); + expect(h.enqueuePrompt).not.toHaveBeenCalled(); + expect(h.onPendingStateChange).not.toHaveBeenCalled(); + expect(steerCommands(h.transport)).toHaveLength(0); + }); +}); diff --git a/cli/src/pi/steerDispatcher.ts b/cli/src/pi/steerDispatcher.ts new file mode 100644 index 00000000..50ea0df0 --- /dev/null +++ b/cli/src/pi/steerDispatcher.ts @@ -0,0 +1,148 @@ +import { PiRpcTimeoutError, sendPiRpcAndWait } from './loop'; +import type { PiConversationHistory } from './conversationHistory'; +import type { PiSession } from './session'; +import type { PiTransport } from './piTransport'; +import type { PiPreparedPrompt } from './promptQueue'; + +export type PiPreparedSteer = PiPreparedPrompt & { + /** Active Pi generation observed when this message arrived at HAPI. */ + targetStreamingGeneration: number | null; +}; + +type ActiveSteer = { + entry: PiPreparedSteer; + cancelled: boolean; + nativeSent: boolean; +}; + +/** + * Serializes explicit native Pi steers by their RPC responses. Pi lifecycle + * events intentionally do not affect this dispatcher: a steer can be accepted + * while the main prompt continues streaming, and Pi's steeringMode is only a + * native policy, not a HAPI ordering signal. + */ +export class PiSteerDispatcher { + private readonly entries: PiPreparedSteer[] = []; + private active: ActiveSteer | null = null; + private stopped = false; + + constructor(private readonly options: { + session: PiSession; + transport: PiTransport; + conversationHistory: PiConversationHistory; + enqueuePrompt: (prompt: PiPreparedPrompt) => void; + onIndeterminateTimeout: (error: PiRpcTimeoutError) => void; + /** Re-attempt normal prompt pumping after local steer work drains. */ + onPendingStateChange?: () => void; + }) {} + + /** A prompt must wait while any earlier steer can still fall back ahead of it. */ + get hasPending(): boolean { + // A locally cancelled active entry cannot reach Pi or fall back, so it + // must release the normal prompt pump even if it is still awaiting the + // runtime lock merely to finish its own cancellation path. + return !this.stopped && ((this.active !== null && !this.active.cancelled) || this.entries.length > 0); + } + + enqueue(entry: PiPreparedSteer): void { + if (this.stopped) return; + this.entries.push(entry); + this.pump(); + this.options.onPendingStateChange?.(); + } + + /** True only while the entry is still local and has not reached Pi stdin. */ + cancelByLocalId(localId: string): boolean { + const index = this.entries.findIndex((entry) => entry.localId === localId); + if (index !== -1) { + this.entries.splice(index, 1); + this.options.onPendingStateChange?.(); + return true; + } + if (this.active?.entry.localId === localId && !this.active.nativeSent) { + this.active.cancelled = true; + this.options.onPendingStateChange?.(); + return true; + } + return false; + } + + /** Stop all local work; never turn a shutdown race into a queued prompt. */ + stop(): void { + this.stopped = true; + this.entries.length = 0; + if (this.active && !this.active.nativeSent) this.active.cancelled = true; + // Shutdown must never wake the ordinary prompt pump: transport error + // handlers call stop() before cleanup marks the runner terminal. + } + + private pump(): void { + if (this.stopped || this.active || this.entries.length === 0) return; + const entry = this.entries.shift()!; + const active: ActiveSteer = { entry, cancelled: false, nativeSent: false }; + this.active = active; + void this.dispatch(active).finally(() => { + if (this.active !== active) return; + this.active = null; + this.pump(); + this.options.onPendingStateChange?.(); + }); + } + + private async dispatch(active: ActiveSteer): Promise { + try { + await this.options.session.runRuntimeMutation(async () => { + if (this.stopped || active.cancelled) return; + + // A steer held behind config/history work may only reach this + // point after the original Pi turn ended. At that boundary, + // preserve normal prompt FIFO semantics instead of issuing an + // invalid/meaningless native steer. + const currentGeneration = this.options.session.currentStreamingGeneration; + if ( + !this.options.session.isReady + || currentGeneration === null + || currentGeneration !== active.entry.targetStreamingGeneration + ) { + this.options.enqueuePrompt({ + message: active.entry.message, + images: active.entry.images, + outboundSequence: active.entry.outboundSequence, + ...(active.entry.localId ? { localId: active.entry.localId } : {}), + }); + return; + } + + // Pi appends a user entry for both prompt and steer. Record it + // immediately before sending so native entry arrival remains + // FIFO-correlated even with identical message text. + this.options.conversationHistory.registerUserEntry(active.entry.localId); + active.nativeSent = true; + await sendPiRpcAndWait(this.options.session, this.options.transport, { + type: 'steer', + message: active.entry.message, + ...(active.entry.images.length > 0 ? { images: active.entry.images } : {}), + }); + + if (this.stopped) return; + if (active.entry.localId) this.options.session.emitMessagesConsumed([active.entry.localId]); + }, { poisonOnError: (error) => error instanceof PiRpcTimeoutError }); + } catch (error) { + // Transport teardown rejects outstanding RPCs. Its cleanup owns the + // session terminal state, so do not emit an error/consume event or + // degrade this steer into the prompt queue afterwards. + if (this.stopped || active.cancelled || !active.nativeSent) return; + + const detail = error instanceof Error ? error.message : String(error); + this.options.conversationHistory.rejectPendingEntry(active.entry.localId); + if (active.entry.localId) { + this.options.session.emitMessagesConsumed( + [active.entry.localId], + { clearQueuedThinkingGrace: true }, + ); + } + this.options.session.sendSessionEvent({ type: 'message', message: `Pi steer failed: ${detail}` }); + if (error instanceof PiRpcTimeoutError) this.options.onIndeterminateTimeout(error); + } + } +} diff --git a/cli/src/pi/types.ts b/cli/src/pi/types.ts index 51107b40..c91c0fed 100644 --- a/cli/src/pi/types.ts +++ b/cli/src/pi/types.ts @@ -44,7 +44,8 @@ export interface PiContextUsage { // Individual event types for proper type narrowing export interface PiAgentStartEvent { type: 'agent_start' } -export interface PiAgentEndEvent { type: 'agent_end'; messages: unknown[] } +export interface PiAgentEndEvent { type: 'agent_end'; messages: unknown[]; willRetry?: boolean } +export interface PiAgentSettledEvent { type: 'agent_settled' } export interface PiTurnStartEvent { type: 'turn_start' } export interface PiTurnEndEvent { type: 'turn_end'; @@ -71,6 +72,84 @@ export interface PiToolExecutionUpdateEvent { args: unknown; partialResult: unknown; } + +export type PiExtensionUiRequest = + | { + type: 'extension_ui_request'; + id: string; + method: 'select'; + title: string; + options: string[]; + timeout?: number; + } + | { + type: 'extension_ui_request'; + id: string; + method: 'confirm'; + title: string; + message: string; + timeout?: number; + } + | { + type: 'extension_ui_request'; + id: string; + method: 'input'; + title: string; + placeholder?: string; + timeout?: number; + } + | { + type: 'extension_ui_request'; + id: string; + method: 'editor'; + title: string; + prefill?: string; + } + | { + type: 'extension_ui_request'; + id: string; + method: 'notify'; + message: string; + notifyType?: 'info' | 'warning' | 'error'; + } + | { + type: 'extension_ui_request'; + id: string; + method: 'setStatus'; + statusKey: string; + statusText?: string; + } + | { + type: 'extension_ui_request'; + id: string; + method: 'setWidget'; + widgetKey: string; + widgetLines?: string[]; + widgetPlacement?: 'aboveEditor' | 'belowEditor'; + } + | { + type: 'extension_ui_request'; + id: string; + method: 'setTitle'; + title: string; + } + | { + type: 'extension_ui_request'; + id: string; + method: 'set_editor_text'; + text: string; + }; + +export type PiExtensionUiResponse = + | { type: 'extension_ui_response'; id: string; value: string } + | { type: 'extension_ui_response'; id: string; confirmed: boolean } + | { type: 'extension_ui_response'; id: string; cancelled: true }; + +export type PiImageContent = { + type: 'image'; + data: string; + mimeType: string; +}; export interface PiToolExecutionEndEvent { type: 'tool_execution_end'; toolCallId: string; @@ -78,10 +157,15 @@ export interface PiToolExecutionEndEvent { result: unknown; isError: boolean; } +export interface PiEntryAppendedEvent { + type: 'entry_appended'; + entry: unknown; +} export type PiAgentEvent = | PiAgentStartEvent | PiAgentEndEvent + | PiAgentSettledEvent | PiTurnStartEvent | PiTurnEndEvent | PiMessageStartEvent @@ -90,6 +174,8 @@ export type PiAgentEvent = | PiToolExecutionStartEvent | PiToolExecutionUpdateEvent | PiToolExecutionEndEvent + | PiExtensionUiRequest + | PiEntryAppendedEvent | { type: string }; // fallback for unknown events // ============================================================================ @@ -101,16 +187,23 @@ import type { PiCommandSummary } from '@hapi/protocol/apiTypes' export type { PiThinkingLevel, PiCommandSummary } export type PiRpcCommand = - | { type: 'prompt'; message: string } - | { type: 'steer'; message: string } - | { type: 'abort' } + | { id?: string; type: 'prompt'; message: string; images?: PiImageContent[]; streamingBehavior?: 'steer' | 'followUp' } + | { id?: string; type: 'steer'; message: string; images?: PiImageContent[] } + | { id?: string; type: 'follow_up'; message: string; images?: PiImageContent[] } + | { id?: string; type: 'abort' } + | PiExtensionUiResponse | { type: 'new_session' } | { type: 'get_state' } | { type: 'set_model'; provider: string; modelId: string } | { type: 'get_available_models' } | { type: 'set_thinking_level'; level: PiThinkingLevel } | { type: 'get_commands' } - | { type: 'get_session_stats' }; + | { id?: string; type: 'get_session_stats' } + | { id?: string; type: 'switch_session'; sessionPath: string } + | { id?: string; type: 'fork'; entryId: string } + | { id?: string; type: 'clone' } + | { id?: string; type: 'get_fork_messages' } + | { id?: string; type: 'get_entries'; since?: string }; // ============================================================================ // Pi RPC Responses (stdout) diff --git a/hub/src/store/index.ts b/hub/src/store/index.ts index 4327cd67..51064abe 100644 --- a/hub/src/store/index.ts +++ b/hub/src/store/index.ts @@ -151,7 +151,7 @@ export class Store { content: unknown, localId?: string, scheduledAt?: number | null - ): { sessionId: string; message: StoredMessage } { + ): { sessionId: string; message: StoredMessage; inserted: boolean } { return this.db.transaction(() => { const row = this.db.prepare('SELECT namespace, metadata FROM sessions WHERE id = ?').get(sessionId) as { namespace: string; metadata: string | null } | undefined if (!row) throw new Error('Message source session not found') @@ -169,7 +169,15 @@ export class Store { .get(targetSessionId, row.namespace) if (!target) throw new Error('OpenCode clear redirect target is unavailable in the source namespace') } - return { sessionId: targetSessionId, message: addMessage(this.db, targetSessionId, content, localId, scheduledAt) } + const alreadyExists = localId + ? Boolean(this.db.prepare('SELECT 1 FROM messages WHERE session_id = ? AND local_id = ? LIMIT 1') + .get(targetSessionId, localId)) + : false + return { + sessionId: targetSessionId, + message: addMessage(this.db, targetSessionId, content, localId, scheduledAt), + inserted: !alreadyExists + } })() } diff --git a/hub/src/sync/conversationHistoryPi.test.ts b/hub/src/sync/conversationHistoryPi.test.ts new file mode 100644 index 00000000..c27aafa2 --- /dev/null +++ b/hub/src/sync/conversationHistoryPi.test.ts @@ -0,0 +1,171 @@ +import { describe, expect, it } from 'bun:test' +import { Store } from '../store' +import { RpcRegistry } from '../socket/rpcRegistry' +import { SyncEngine } from './syncEngine' + +function createEngine() { + const store = new Store(':memory:') + const engine = new SyncEngine(store, {} as never, new RpcRegistry(), { broadcast() {} } as never) + return { store, engine } +} + +describe('Pi conversation-history hub integration', () => { + it('returns 409-style deterministic rewind rejection without marking history diverged', async () => { + const { store, engine } = createEngine() + try { + const session = engine.getOrCreateSession('pi-rewind-rejected', { + path: '/tmp/project', host: 'localhost', flavor: 'pi', + capabilities: { conversationHistory: { rewindToMessage: true } }, + }, null, 'default') + engine.handleSessionAlive({ sid: session.id, time: Date.now(), mode: 'remote' }) + store.messages.addMessage(session.id, { role: 'user', content: 'boundary' }, 'local-boundary') + store.messages.markMessagesInvoked(session.id, ['local-boundary'], Date.now()) + ;(engine as any).rpcGateway.rewindConversation = async () => ({ + success: false, + error: 'Pi rewind was cancelled', + outcome: 'cancelled', + }) + + await expect(engine.rewindConversation(session.id, 'default', 'local-boundary')).resolves.toEqual({ + type: 'error', message: 'Pi rewind was cancelled', + }) + expect(engine.getSession(session.id)?.metadata?.conversationHistoryDiverged).not.toBe(true) + } finally { + engine.stop() + } + }) + + it('shares exact-native bind logic while requiring Pi native-ready but not Grok ready', async () => { + const { engine } = createEngine() + try { + const pi = engine.getOrCreateSession('pi-child', { + path: '/tmp/project', host: 'localhost', flavor: 'pi', piSessionId: 'pi-native', + }, null, 'default') + const grok = engine.getOrCreateSession('grok-child', { + path: '/tmp/project', host: 'localhost', flavor: 'grok', grokSessionId: 'grok-native', + }, null, 'default') + + const piWait = (engine as any).waitForExactNativeForkBound( + pi.id, 'pi-native', 'piSessionId', true + ) + expect(await (engine as any).waitForExactNativeForkBound( + grok.id, 'grok-native', 'grokSessionId', false + )).toBe(true) + + engine.handleSessionReady({ sid: pi.id, time: Date.now() }) + expect(await piWait).toBe(true) + } finally { + engine.stop() + } + }) + + it('copies Pi entry-id locators into a fork child and requires exact native-ready binding', async () => { + const { store, engine } = createEngine() + try { + const source = engine.getOrCreateSession('pi-source', { + path: '/tmp/project', + host: 'localhost', + machineId: 'machine-1', + flavor: 'pi', + piSessionId: 'pi-source-native', + capabilities: { conversationHistory: { forkCurrent: true, forkAtMessage: true, rewindToMessage: true } }, + conversationHistoryPoints: { local1: true, local2: true }, + conversationHistoryEntryIds: { local1: 'entry-1', local2: 'entry-2' }, + }, null, 'default', 'source-model', 'high') + engine.handleSessionAlive({ sid: source.id, time: Date.now(), mode: 'remote' }) + store.messages.addMessage(source.id, { role: 'user', content: 'one' }, 'local1') + store.messages.addMessage(source.id, { role: 'user', content: 'two' }, 'local2') + store.messages.markMessagesInvoked(source.id, ['local1', 'local2'], Date.now()) + // Keep the cache fixture explicit; this is the source used by the + // fork prefix copy path after metadata normalization. + ;(engine.getSession(source.id)!.metadata as any).conversationHistoryEntryIds = { + local1: 'entry-1', local2: 'entry-2', + } + + ;(engine as any).rpcGateway.forkConversation = async () => { + store.messages.addMessage(source.id, { role: 'user', content: 'latest' }, 'local3') + store.messages.markMessagesInvoked(source.id, ['local3'], Date.now()) + const current = store.sessions.getSession(source.id)! + store.sessions.updateSessionMetadata( + source.id, + { + ...(current.metadata as Record), + conversationHistoryPoints: { ...((current.metadata as any).conversationHistoryPoints), local3: true }, + conversationHistoryEntryIds: { ...((current.metadata as any).conversationHistoryEntryIds), local3: 'entry-3' }, + }, + current.metadataVersion, + 'default', + { touchUpdatedAt: false } + ) + return { nativeSessionId: 'pi-clone-native' } + } + let spawnArgs: unknown[] = [] + ;(engine as any).rpcGateway.spawnSession = async (...args: unknown[]) => { + spawnArgs = args + return { type: 'success', sessionId: args[12] } + } + const exactBinds: unknown[][] = [] + ;(engine as any).waitForExactNativeForkBound = async (...args: unknown[]) => { + exactBinds.push(args) + return true + } + let capturedChildMetadata: Record | undefined + const cache = (engine as any).sessionCache + const originalCreate = cache.getOrCreateSession.bind(cache) + cache.getOrCreateSession = (...args: unknown[]) => { + if (typeof args[0] === 'string' && args[0].startsWith('fork:')) { + capturedChildMetadata = args[1] as Record + } + return originalCreate(...args) + } + + const result = await engine.forkConversation(source.id, 'default') + expect(result.type).toBe('success') + if (result.type !== 'success') throw new Error(result.message) + expect(capturedChildMetadata).toMatchObject({ + flavor: 'pi', + piSessionId: 'pi-clone-native', + conversationHistoryEntryIds: { local1: 'entry-1', local2: 'entry-2', local3: 'entry-3' }, + conversationHistoryPoints: { local1: true, local2: true, local3: true }, + }) + expect(store.messages.getAllMessages(result.sessionId).map((message) => message.localId)).toContain('local3') + expect(exactBinds).toEqual([[result.sessionId, 'pi-clone-native', 'piSessionId', true]]) + expect(spawnArgs[3]).toBeUndefined() + expect(spawnArgs[9]).toBeUndefined() + expect(engine.getSession(result.sessionId)?.model).toBeNull() + expect(engine.getSession(result.sessionId)?.effort).toBeNull() + } finally { + engine.stop() + } + }) + + it('scrubs Pi entry-id locators when rewind truncates their HAPI messages', () => { + const { store, engine } = createEngine() + try { + const session = engine.getOrCreateSession('pi-scrub', { + path: '/tmp/project', host: 'localhost', flavor: 'pi', + conversationHistoryEntryIds: { keep: 'entry-keep', remove: 'entry-remove' }, + conversationHistoryPoints: { keep: true, remove: true }, + }, null, 'default') + store.messages.addMessage(session.id, { role: 'user', content: 'keep' }, 'keep') + ;(engine.getSession(session.id)!.metadata as any).conversationHistoryEntryIds = { + keep: 'entry-keep', remove: 'entry-remove', + } + let capturedMetadata: Record | undefined + const originalUpdate = store.sessions.updateSessionMetadata.bind(store.sessions) + ;(store.sessions as any).updateSessionMetadata = (...args: unknown[]) => { + capturedMetadata = args[1] as Record + return { result: 'success' } + } + ;(engine as any).scrubHistoryLocators(session.id, 'default') + ;(store.sessions as any).updateSessionMetadata = originalUpdate + + expect(capturedMetadata).toMatchObject({ + conversationHistoryEntryIds: { keep: 'entry-keep' }, + conversationHistoryPoints: { keep: true }, + }) + } finally { + engine.stop() + } + }) +}) diff --git a/hub/src/sync/messageService.test.ts b/hub/src/sync/messageService.test.ts index 7d1217ab..ba313e90 100644 --- a/hub/src/sync/messageService.test.ts +++ b/hub/src/sync/messageService.test.ts @@ -1106,6 +1106,236 @@ describe('MessageService.sendMessage with scheduledAt', () => { }) }) +describe('MessageService.sendMessage deliveryMode', () => { + function makeTrackingIo(): { io: Server; cliEmitted: unknown[] } { + const cliEmitted: unknown[] = [] + const io = { + of: (ns: string) => ({ + to: (_room: string) => ({ + emit: (_event: string, data: unknown) => { + if (ns === '/cli') cliEmitted.push(data) + }, + timeout: (_ms: number) => ({ emit: () => {} }) + }), + adapter: { rooms: { get: () => undefined } } + }) + } as unknown as Server + return { io, cliEmitted } + } + + it('persists Pi steer provenance but downgrades every deferred CLI delivery to queue', async () => { + const store = makeStore() + const session = store.sessions.getOrCreateSession( + 'delivery-mode-pi', + { path: '/tmp/delivery-mode-pi', host: 'localhost', flavor: 'pi' }, + null, + 'default' + ) + const publisher = makePublisher() + const { io, cliEmitted } = makeTrackingIo() + const service = new MessageService(store, io, publisher as any) + + await service.sendMessage(session.id, { + text: 'steer this Pi turn', + localId: 'pi-steer', + deliveryMode: 'steer' + }) + + const stored = store.messages.getUninvokedLocalMessages(session.id) + expect(stored).toHaveLength(1) + expect(stored[0]?.content).toMatchObject({ + role: 'user', + meta: { sentFrom: 'webapp', deliveryMode: 'steer' } + }) + expect(cliEmitted).toHaveLength(1) + expect(cliEmitted[0]).toMatchObject({ + body: { message: { content: { meta: { deliveryMode: 'steer' } } } } + }) + + expect(service.replayImmediateQueuedMessages(session.id)).toBe(1) + expect(cliEmitted).toHaveLength(2) + expect(cliEmitted[1]).toMatchObject({ + body: { message: { content: { meta: { deliveryMode: 'queue' } } } } + }) + + const backfill = service.getDeliverableMessagesAfter(session.id, { + afterSeq: 0, + limit: 10, + now: Date.now() + }) + expect(backfill).toHaveLength(1) + expect(backfill[0]?.content).toMatchObject({ meta: { deliveryMode: 'queue' } }) + + expect(service.releaseDeliverableQueuedMessages(session.id)).toBe(1) + expect(cliEmitted).toHaveLength(3) + expect(cliEmitted[2]).toMatchObject({ + body: { message: { content: { meta: { deliveryMode: 'queue' } } } } + }) + + // Deferred delivery is a view transformation only. The database keeps + // the original provenance for Web display and diagnostics. + expect(store.messages.getUninvokedLocalMessages(session.id)[0]?.content).toMatchObject({ + role: 'user', + meta: { sentFrom: 'webapp', deliveryMode: 'steer' } + }) + }) + + it('delivers a duplicate-localId retry as queue even when the stored row retains steer', async () => { + const store = makeStore() + const session = store.sessions.getOrCreateSession( + 'delivery-mode-duplicate-pi', + { path: '/tmp/delivery-mode-duplicate-pi', host: 'localhost', flavor: 'pi' }, + null, + 'default' + ) + const { io, cliEmitted } = makeTrackingIo() + const service = new MessageService(store, io, makePublisher() as any) + + await service.sendMessage(session.id, { + text: 'original steer whose response is lost', + localId: 'duplicate-steer', + deliveryMode: 'steer' + }) + await service.sendMessage(session.id, { + text: 'retry requests queue', + localId: 'duplicate-steer', + deliveryMode: 'queue' + }) + + expect(cliEmitted).toHaveLength(2) + expect(cliEmitted[0]).toMatchObject({ + body: { message: { content: { meta: { deliveryMode: 'steer' } } } } + }) + expect(cliEmitted[1]).toMatchObject({ + body: { message: { content: { meta: { deliveryMode: 'queue' } } } } + }) + const rows = store.messages.getUninvokedLocalMessages(session.id) + expect(rows).toHaveLength(1) + expect(rows[0]?.content).toMatchObject({ + role: 'user', + content: { text: 'original steer whose response is lost' }, + meta: { sentFrom: 'webapp', deliveryMode: 'steer' } + }) + }) + + it('downgrades a legacy persisted steer through the mature scheduled scan', () => { + const store = makeStore() + const session = store.sessions.getOrCreateSession( + 'delivery-mode-mature-pi', + { path: '/tmp/delivery-mode-mature-pi', host: 'localhost', flavor: 'pi' }, + null, + 'default' + ) + const publisher = makePublisher() + const { io, cliEmitted } = makeTrackingIo() + const service = new MessageService(store, io, publisher as any) + const scheduledAt = Date.now() - 1_000 + + store.messages.addMessage( + session.id, + { + role: 'user', + content: { type: 'text', text: 'legacy scheduled steer' }, + meta: { sentFrom: 'webapp', deliveryMode: 'steer' } + }, + 'mature-steer', + scheduledAt + ) + + service.releaseMatureScheduledMessages(Date.now()) + + expect(cliEmitted).toHaveLength(1) + expect(cliEmitted[0]).toMatchObject({ + body: { message: { content: { meta: { deliveryMode: 'queue' } } } } + }) + expect(store.messages.getUninvokedLocalMessages(session.id)[0]?.content).toMatchObject({ + role: 'user', + meta: { sentFrom: 'webapp', deliveryMode: 'steer' } + }) + }) + + it('leaves non-user and already-queued content unchanged during CLI backfill', () => { + const store = makeStore() + const session = makeSession(store, 'delivery-mode-backfill-guards') + const service = new MessageService(store, makeTrackingIo().io, makePublisher() as any) + const agentContent = { + role: 'agent', + content: { type: 'text', text: 'agent event' }, + meta: { deliveryMode: 'steer', marker: 'keep-agent' } + } + const userWithoutMeta = { + role: 'user', + content: { type: 'text', text: 'legacy user' } + } + const queuedUser = { + role: 'user', + content: { type: 'text', text: 'already queued' }, + meta: { sentFrom: 'webapp', deliveryMode: 'queue', marker: 'keep-user' } + } + + store.messages.addMessage(session.id, agentContent) + store.messages.addMessage(session.id, userWithoutMeta, 'legacy-user') + store.messages.addMessage(session.id, queuedUser, 'queued-user') + + const backfill = service.getDeliverableMessagesAfter(session.id, { + afterSeq: 0, + limit: 10, + now: Date.now() + }) + + expect(backfill.map((message) => message.content)).toEqual([ + agentContent, + userWithoutMeta, + queuedUser + ]) + }) + + it('downgrades forged steer intent for non-Pi sessions and defaults omitted intent to queue', async () => { + const store = makeStore() + const session = makeSession(store, 'delivery-mode-non-pi') + const service = new MessageService(store, makeTrackingIo().io, makePublisher() as any) + + await service.sendMessage(session.id, { + text: 'forged steer', + localId: 'non-pi-steer', + deliveryMode: 'steer' + }) + await service.sendMessage(session.id, { + text: 'missing mode', + localId: 'missing-mode' + }) + + const rows = store.messages.getUninvokedLocalMessages(session.id) + expect(rows).toHaveLength(2) + expect(rows.map((row) => { + const content = row.content as { meta?: { deliveryMode?: unknown } } + return content.meta?.deliveryMode + })).toEqual(['queue', 'queue']) + }) + + it('normalizes direct scheduled steer requests to queue', async () => { + const store = makeStore() + const session = store.sessions.getOrCreateSession( + 'delivery-mode-pi-scheduled', + { path: '/tmp/delivery-mode-pi-scheduled', host: 'localhost', flavor: 'pi' }, + null, + 'default' + ) + const service = new MessageService(store, makeTrackingIo().io, makePublisher() as any) + + await service.sendMessage(session.id, { + text: 'scheduled steer', + localId: 'pi-scheduled-steer', + scheduledAt: Date.now() + 60_000, + deliveryMode: 'steer' + }) + + expect(store.messages.getUninvokedLocalMessages(session.id)[0]?.content).toMatchObject({ + meta: { deliveryMode: 'queue' } + }) + }) +}) + // --------------------------------------------------------------------------- // releaseMatureScheduledMessages // --------------------------------------------------------------------------- diff --git a/hub/src/sync/messageService.ts b/hub/src/sync/messageService.ts index 5fdd17e2..a3064bc8 100644 --- a/hub/src/sync/messageService.ts +++ b/hub/src/sync/messageService.ts @@ -10,7 +10,7 @@ import { unwrapRoleWrappedRecordEnvelope } from '@hapi/protocol/messages' import { isObject } from '@hapi/protocol' -import type { MessagesResponse, QueuedStateResponse } from '@hapi/protocol/apiTypes' +import type { MessageDeliveryMode, MessagesResponse, QueuedStateResponse } from '@hapi/protocol/apiTypes' import type { Server } from 'socket.io' import { randomUUID } from 'node:crypto' import type { Store, CancelQueuedMessageResult } from '../store' @@ -81,6 +81,39 @@ function isExportVisibleStoredMessage(message: StoredMessageForDelivery): boolea return isClaudeChatVisibleMessage({ type: data.type, subtype: data.subtype }) } +function getNormalizedDeliveryMode( + metadata: unknown, + requestedDeliveryMode: MessageDeliveryMode | undefined, + scheduledAt: number | null | undefined +): MessageDeliveryMode { + if (requestedDeliveryMode !== 'steer' || scheduledAt != null) { + return 'queue' + } + + return isObject(metadata) && metadata.flavor === 'pi' ? 'steer' : 'queue' +} + +/** + * Native steer is scoped to the Pi turn active at the initial live emit. Once + * a durable row is delivered through reconnect, backfill, a clear gate, or a + * scheduled scan, that turn identity is no longer provable. Preserve stored + * provenance for Web diagnostics, but make deferred CLI delivery an ordinary + * queue item so it cannot steer a later generation. + */ +function contentForDeferredDelivery(content: unknown): unknown { + if (!isObject(content) || content.role !== 'user' || !isObject(content.meta)) { + return content + } + if (content.meta.deliveryMode !== 'steer') return content + return { + ...content, + meta: { + ...content.meta, + deliveryMode: 'queue' as const + } + } +} + export class MessageService { /** One scheduled-matured SSE per localId per hub process (cleared on cancel/consume paths here). */ private readonly scheduledMatureNotifiedLocalIds = new Set() @@ -363,7 +396,7 @@ export class MessageService { id: message.id, seq: message.seq, localId: message.localId, - content: message.content, + content: contentForDeferredDelivery(message.content), createdAt: message.createdAt, invokedAt: message.invokedAt, scheduledAt: message.scheduledAt @@ -574,6 +607,7 @@ export class MessageService { attachments?: AttachmentMetadata[] sentFrom?: 'telegram-bot' | 'webapp' scheduledAt?: number | null + deliveryMode?: MessageDeliveryMode } ): Promise { // Defence-in-depth invariant for non-REST callers (Telegram bot, MCP, @@ -589,6 +623,11 @@ export class MessageService { } const sentFrom = payload.sentFrom ?? 'webapp' + const deliveryMode = getNormalizedDeliveryMode( + this.store.sessions.getSession(sessionId)?.metadata, + payload.deliveryMode, + payload.scheduledAt + ) const content = { role: 'user', @@ -598,7 +637,8 @@ export class MessageService { attachments: payload.attachments }, meta: { - sentFrom + sentFrom, + deliveryMode } } @@ -610,6 +650,13 @@ export class MessageService { ) const actualSessionId = inserted.sessionId const msg = inserted.message + // A duplicate localId is an idempotent retry, not proof that the + // original Pi turn still exists. Its stored row may retain steer + // provenance from a POST whose response was lost, so deliver the + // duplicate through the same turn-safe deferred view as reconnect. + const cliContent = inserted.inserted + ? msg.content + : contentForDeferredDelivery(msg.content) this.onSessionActivity?.(actualSessionId, msg.createdAt) // Only emit to CLI if the message is not scheduled for the future. @@ -632,7 +679,7 @@ export class MessageService { seq: msg.seq, createdAt: msg.createdAt, localId: msg.localId, - content: msg.content + content: cliContent } } } @@ -704,7 +751,7 @@ export class MessageService { seq: msg.seq, createdAt: msg.createdAt, localId: msg.localId, - content: msg.content + content: contentForDeferredDelivery(msg.content) } } } @@ -731,7 +778,7 @@ export class MessageService { seq: msg.seq, createdAt: msg.createdAt, localId: msg.localId, - content: msg.content + content: contentForDeferredDelivery(msg.content) } } } @@ -785,7 +832,7 @@ export class MessageService { seq: msg.seq, createdAt: msg.createdAt, localId: msg.localId, - content: msg.content + content: contentForDeferredDelivery(msg.content) } } } diff --git a/hub/src/sync/syncEngine.ts b/hub/src/sync/syncEngine.ts index d2e66713..9fe4c359 100644 --- a/hub/src/sync/syncEngine.ts +++ b/hub/src/sync/syncEngine.ts @@ -8,7 +8,7 @@ */ import { isKnownFlavor, type LocalResumeTarget, type ResumableSession, type SessionEndReason } from '@hapi/protocol' -import type { CursorChatStoreStatus, CursorMigrateOutcome, CursorMigrateToAcpRequest, MessagesResponse, QueuedStateResponse, SlashCommandsResponse } from '@hapi/protocol/apiTypes' +import type { CursorChatStoreStatus, CursorMigrateOutcome, CursorMigrateToAcpRequest, MessageDeliveryMode, MessagesResponse, QueuedStateResponse, SlashCommandsResponse } from '@hapi/protocol/apiTypes' import type { AgentFlavor, CodexCollaborationMode, CopilotAgentMode, DecryptedMessage, PermissionMode, Session, SyncEvent } from '@hapi/protocol/types' import { unwrapRoleWrappedRecordEnvelope } from '@hapi/protocol/messages' import type { Server } from 'socket.io' @@ -918,6 +918,7 @@ async uploadScratchlistAttachment( }> sentFrom?: 'telegram-bot' | 'webapp' scheduledAt?: number | null + deliveryMode?: MessageDeliveryMode } ): Promise { if (this.historyActionsInFlight.has(sessionId)) { @@ -1022,21 +1023,25 @@ async uploadScratchlistAttachment( } /** - * Grok RPC already created `expectedNativeSessionId`. Wait until the child - * binds that exact id — a different id means load failed and fell back. + * A native fork may be created before its runner child has loaded it. Wait + * for the exact persisted native id; Pi additionally requires its + * validated `session-ready` event before the fork is visible to callers. */ - private async waitForGrokForkBound( + private async waitForExactNativeForkBound( childId: string, expectedNativeSessionId: string, + metadataKey: 'grokSessionId' | 'piSessionId', + requireSessionReady: boolean, timeoutMs: number = 60_000 ): Promise { const startedAt = Date.now() while (Date.now() - startedAt < timeoutMs) { this.sessionCache.refreshSession(childId) const child = this.sessionCache.getSession(childId) - const boundId = child?.metadata?.grokSessionId + const boundId = child?.metadata?.[metadataKey] if (typeof boundId === 'string' && boundId.length > 0) { - return boundId === expectedNativeSessionId + if (boundId !== expectedNativeSessionId) return false + if (!requireSessionReady || this.sessionReadyIds.has(childId)) return true } if (child && !child.active && Date.now() - startedAt > 5_000) { return false @@ -1105,6 +1110,21 @@ async uploadScratchlistAttachment( } } + const entryIds = session.metadata.conversationHistoryEntryIds + if (entryIds) { + const nextEntryIds = Object.fromEntries( + Object.entries(entryIds).filter(([localId]) => remainingLocalIds.has(localId)) + ) + if (Object.keys(nextEntryIds).length !== Object.keys(entryIds).length) { + changed = true + if (Object.keys(nextEntryIds).length > 0) { + nextMetadata.conversationHistoryEntryIds = nextEntryIds + } else { + delete nextMetadata.conversationHistoryEntryIds + } + } + } + if (!changed) return const result = this.store.sessions.updateSessionMetadata( @@ -1171,7 +1191,7 @@ async uploadScratchlistAttachment( if (!access.ok) { return { type: 'error', message: access.reason === 'not-found' ? 'Session not found' : 'Access denied' } } - const source = access.session + let source = access.session try { this.assertConversationHistoryIdle(source) } catch (error) { @@ -1212,6 +1232,14 @@ async uploadScratchlistAttachment( return { type: 'error', message: 'Native fork did not return a session id' } } + // Native fork RPC can race CLI metadata/transcript updates. Construct + // the child only from a fresh source snapshot, never the pre-RPC row. + const refreshedSource = this.sessionCache.refreshSession(sessionId) + if (!refreshedSource || refreshedSource.namespace !== namespace) { + return { type: 'error', message: 'Source session disappeared after native fork' } + } + source = refreshedSource + const flavor = this.resolveFlavor(source) const childId = randomUUID() let prefix @@ -1242,17 +1270,29 @@ async uploadScratchlistAttachment( conversationHistoryTurns: Object.fromEntries( Object.entries(source.metadata?.conversationHistoryTurns ?? {}) .filter(([localId]) => copiedLocalIds.has(localId)) + ), + conversationHistoryEntryIds: Object.fromEntries( + Object.entries(source.metadata?.conversationHistoryEntryIds ?? {}) + .filter(([localId]) => copiedLocalIds.has(localId)) ) } if (flavor === 'codex') { childMetadata.codexSessionId = rpcResult.nativeSessionId } else if (flavor === 'grok') { childMetadata.grokSessionId = rpcResult.nativeSessionId + } else if (flavor === 'pi') { + childMetadata.piSessionId = rpcResult.nativeSessionId } else if (flavor === 'claude') { // Child will bind the forked Claude id after --fork-session starts. childMetadata.claudeSessionId = rpcResult.forkSession ? undefined : rpcResult.nativeSessionId } + // A Pi native fork already carries the branch's authoritative model and + // thinking state. Do not replay the source wrapper's current overrides + // onto the child; the resumed child will report its own get_state. + const forkModel = flavor === 'pi' ? undefined : source.model ?? undefined + const forkEffort = flavor === 'pi' ? undefined : source.effort ?? undefined + let childCreated = false let spawnAttempted = false try { @@ -1261,8 +1301,8 @@ async uploadScratchlistAttachment( childMetadata, null, namespace, - source.model ?? undefined, - source.effort ?? undefined, + forkModel, + forkEffort, source.modelReasoningEffort ?? undefined, childId ) @@ -1288,13 +1328,13 @@ async uploadScratchlistAttachment( machineId, directory, flavor, - source.model ?? undefined, + forkModel, source.modelReasoningEffort ?? undefined, undefined, 'simple', undefined, rpcResult.nativeSessionId, - source.effort ?? undefined, + forkEffort, source.permissionMode, source.serviceTier ?? undefined, childId, @@ -1322,12 +1362,23 @@ async uploadScratchlistAttachment( // session if load fails. Do not report success until the child is // bound to the exact forked native id. if (flavor === 'grok') { - const bound = await this.waitForGrokForkBound(childId, rpcResult.nativeSessionId) + const bound = await this.waitForExactNativeForkBound( + childId, rpcResult.nativeSessionId, 'grokSessionId', false + ) if (!bound) { throw new Error('Grok fork could not load the forked native session') } } + if (flavor === 'pi') { + const bound = await this.waitForExactNativeForkBound( + childId, rpcResult.nativeSessionId, 'piSessionId', true + ) + if (!bound) { + throw new Error('Pi fork could not load the exact native session before ready') + } + } + return { type: 'success', sessionId: childId } } catch (error) { if (childCreated) { @@ -1417,7 +1468,7 @@ async uploadScratchlistAttachment( } if (rpcResult?.success !== true) { - return { type: 'error', message: 'Native rewind failed' } + return { type: 'error', message: rpcResult?.error ?? 'Native rewind failed' } } try { diff --git a/hub/src/web/routes/messages.test.ts b/hub/src/web/routes/messages.test.ts index f43b0197..3876d3f8 100644 --- a/hub/src/web/routes/messages.test.ts +++ b/hub/src/web/routes/messages.test.ts @@ -259,6 +259,50 @@ describe('POST /api/sessions/:id/messages — #2 scheduledAt upper bound', () => }) }) +describe('POST /api/sessions/:id/messages — deliveryMode', () => { + it('forwards an immediate steer intent to the hub', async () => { + const { app, sentMessages } = createApp({}) + + const response = await app.request('/api/sessions/session-1/messages', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ text: 'steer now', localId: 'local-steer', deliveryMode: 'steer' }) + }) + + expect(response.status).toBe(200) + expect(sentMessages).toEqual([{ + sessionId: 'session-1', + payload: { + text: 'steer now', + localId: 'local-steer', + attachments: undefined, + sentFrom: 'webapp', + scheduledAt: undefined, + deliveryMode: 'steer' + } + }]) + }) + + it('rejects scheduled steer delivery before calling the hub', async () => { + const { app, sentMessages } = createApp({}) + + const response = await app.request('/api/sessions/session-1/messages', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + text: 'steer later', + localId: 'local-scheduled-steer', + scheduledAt: Date.now() + 60_000, + deliveryMode: 'steer' + }) + }) + + expect(response.status).toBe(400) + expect(JSON.stringify(await response.json())).toContain('deliveryMode') + expect(sentMessages).toHaveLength(0) + }) +}) + // --------------------------------------------------------------------------- // #4 Zod error details in response body // --------------------------------------------------------------------------- diff --git a/hub/src/web/routes/messages.ts b/hub/src/web/routes/messages.ts index 9e3442ac..b4eb79b3 100644 --- a/hub/src/web/routes/messages.ts +++ b/hub/src/web/routes/messages.ts @@ -113,7 +113,8 @@ export function createMessagesRoutes(getSyncEngine: () => SyncEngine | null): Ho localId: parsed.data.localId, attachments: parsed.data.attachments, sentFrom: 'webapp', - scheduledAt: parsed.data.scheduledAt + scheduledAt: parsed.data.scheduledAt, + deliveryMode: parsed.data.deliveryMode }) return c.json({ ok: true }) }) diff --git a/shared/src/apiTypes.test.ts b/shared/src/apiTypes.test.ts index 77cb42ac..8abd5933 100644 --- a/shared/src/apiTypes.test.ts +++ b/shared/src/apiTypes.test.ts @@ -1,5 +1,11 @@ import { describe, expect, it } from 'vitest' -import { ClearOpencodeSessionCallbackRequestSchema, ClearOpencodeSessionResponseSchema, ListCodexSessionsRpcResponseSchema, MessagesQuerySchema } from './apiTypes' +import { + ClearOpencodeSessionCallbackRequestSchema, + ClearOpencodeSessionResponseSchema, + ListCodexSessionsRpcResponseSchema, + MessagesQuerySchema, + SendMessageRequestSchema +} from './apiTypes' describe('ListCodexSessionsRpcResponseSchema', () => { it('preserves Codex session messages when parsing runner RPC responses', () => { @@ -85,3 +91,26 @@ describe('MessagesQuerySchema', () => { expect(MessagesQuerySchema.safeParse({ epoch: 1 }).success).toBe(false) }) }) + +describe('SendMessageRequestSchema deliveryMode', () => { + it('accepts queue and steer delivery modes while leaving the field optional', () => { + expect(SendMessageRequestSchema.parse({ text: 'queue' }).deliveryMode).toBeUndefined() + expect(SendMessageRequestSchema.parse({ text: 'steer', deliveryMode: 'steer' }).deliveryMode).toBe('steer') + expect(SendMessageRequestSchema.parse({ text: 'queue', deliveryMode: 'queue' }).deliveryMode).toBe('queue') + }) + + it('rejects scheduled steer delivery', () => { + const parsed = SendMessageRequestSchema.safeParse({ + text: 'later', + localId: 'scheduled-steer', + scheduledAt: Date.now() + 60_000, + deliveryMode: 'steer' + }) + + expect(parsed.success).toBe(false) + if (!parsed.success) { + expect(parsed.error.issues.some((issue) => issue.path[0] === 'deliveryMode')).toBe(true) + expect(parsed.error.issues.some((issue) => issue.message.includes('cannot use steer'))).toBe(true) + } + }) +}) diff --git a/shared/src/apiTypes.ts b/shared/src/apiTypes.ts index 561b89b9..9f2e42d6 100644 --- a/shared/src/apiTypes.ts +++ b/shared/src/apiTypes.ts @@ -427,11 +427,15 @@ export const MessagesQuerySchema = z.object({ export type MessagesQuery = z.infer +export const MessageDeliveryModeSchema = z.enum(['queue', 'steer']) +export type MessageDeliveryMode = z.infer + export const SendMessageRequestSchema = z.object({ text: z.string(), localId: z.string().min(1).optional(), attachments: z.array(AttachmentMetadataSchema).optional(), - scheduledAt: z.number().int().positive().nullable().optional() + scheduledAt: z.number().int().positive().nullable().optional(), + deliveryMode: MessageDeliveryModeSchema.optional() }).refine( (data) => data.scheduledAt == null || typeof data.localId === 'string', { message: 'scheduledAt requires localId', path: ['localId'] } @@ -441,6 +445,9 @@ export const SendMessageRequestSchema = z.object({ ).refine( (data) => data.scheduledAt == null || !data.attachments?.length, { message: 'scheduled messages with attachments are not supported', path: ['attachments'] } +).refine( + (data) => data.scheduledAt == null || data.deliveryMode !== 'steer', + { message: 'scheduled messages cannot use steer delivery', path: ['deliveryMode'] } ) export type SendMessageRequest = z.infer @@ -482,6 +489,11 @@ export type RewindConversationRpcResult = { createdAt?: number invokedAt?: number | null }> +} | { + success: false + error: string + /** Native state is unchanged, cancelled, or was restored exactly. */ + outcome: 'rejected' | 'cancelled' | 'source_restored' } export const QueuedStateRequestSchema = z.object({ diff --git a/shared/src/conversationHistory.ts b/shared/src/conversationHistory.ts index 0b907c22..f6aedffa 100644 --- a/shared/src/conversationHistory.ts +++ b/shared/src/conversationHistory.ts @@ -62,3 +62,9 @@ export const GROK_CONVERSATION_HISTORY_INITIAL: ConversationHistoryCapabilitySta forkAtMessage: 'unknown', rewindToMessage: 'unknown' } + +export const PI_CONVERSATION_HISTORY_INITIAL: ConversationHistoryCapabilityStates = { + forkCurrent: 'unknown', + forkAtMessage: 'unknown', + rewindToMessage: 'unknown' +} diff --git a/shared/src/schemas.metadata.test.ts b/shared/src/schemas.metadata.test.ts index b29a6706..d3e01894 100644 --- a/shared/src/schemas.metadata.test.ts +++ b/shared/src/schemas.metadata.test.ts @@ -15,4 +15,13 @@ describe('MetadataSchema cursorSessionProtocol', () => { it('rejects unknown protocol values', () => { expect(MetadataSchema.safeParse({ ...base, cursorSessionProtocol: 'websocket' }).success).toBe(false); }); + + it('persists Pi native history entry ids', () => { + const result = MetadataSchema.safeParse({ + ...base, + conversationHistoryEntryIds: { 'local-user-id': 'pi-entry-id' }, + }); + expect(result.success).toBe(true); + expect(result.data?.conversationHistoryEntryIds).toEqual({ 'local-user-id': 'pi-entry-id' }); + }); }); diff --git a/shared/src/schemas.ts b/shared/src/schemas.ts index 1de07670..9e09252d 100644 --- a/shared/src/schemas.ts +++ b/shared/src/schemas.ts @@ -132,6 +132,9 @@ export const MetadataSchema = z.object({ conversationHistoryIndexes: z.record(z.string(), z.number().int().nonnegative()).optional(), // Codex localId → turnId mapping (durable across runner relaunches). conversationHistoryTurns: z.record(z.string(), z.string().min(1)).optional(), + // Pi localId → append-only session entry id mapping. Pi entry ids are the + // only stable native boundary accepted by its fork API. + conversationHistoryEntryIds: z.record(z.string(), z.string().min(1)).optional(), // Set when native rewind succeeded but HAPI truncate/hydrate failed. conversationHistoryDiverged: z.boolean().optional(), worktree: WorktreeMetadataSchema.optional(), diff --git a/web/src/api/client.test.ts b/web/src/api/client.test.ts index 027a654b..35d798aa 100644 --- a/web/src/api/client.test.ts +++ b/web/src/api/client.test.ts @@ -116,6 +116,24 @@ describe('ApiClient error mapping', () => { expect(new Headers(init?.headers).get('content-type')).toBe('application/json') }) + it('forwards the selected delivery mode when sending a message', async () => { + fetchMock.mockResolvedValueOnce(new Response(JSON.stringify({}), { status: 200 })) + + const api = new ApiClient('test-token') + await api.sendMessage('session /?#', 'steer this', 'local-1', undefined, null, 'steer') + + const [url, init] = fetchMock.mock.calls[0] ?? [] + expect(url).toBe('/api/sessions/session%20%2F%3F%23/messages') + expect(init).toMatchObject({ + method: 'POST', + body: JSON.stringify({ + text: 'steer this', + localId: 'local-1', + deliveryMode: 'steer', + }) + }) + }) + it('requests usage buckets in the viewer IANA time zone', async () => { fetchMock.mockResolvedValueOnce(new Response(JSON.stringify({}), { status: 200 })) diff --git a/web/src/api/client.ts b/web/src/api/client.ts index 3c61abe5..ec3e5706 100644 --- a/web/src/api/client.ts +++ b/web/src/api/client.ts @@ -49,7 +49,7 @@ import type { UsageSummaryResponse, UploadFileResponse } from '@hapi/protocol/apiTypes' -import type { AgentFlavor } from '@hapi/protocol' +import type { AgentFlavor, MessageDeliveryMode } from '@hapi/protocol' import type { CancelMessageResponse } from '@hapi/protocol/schemas' import type { TranscriptionMode, TranscriptionProvider, TranscriptionProviderInfo } from '@hapi/protocol/voice' @@ -435,14 +435,22 @@ export class ApiClient { ) } - async sendMessage(sessionId: string, text: string, localId?: string | null, attachments?: AttachmentMetadata[], scheduledAt?: number | null): Promise { + async sendMessage( + sessionId: string, + text: string, + localId?: string | null, + attachments?: AttachmentMetadata[], + scheduledAt?: number | null, + deliveryMode?: MessageDeliveryMode, + ): Promise { await this.request(`/api/sessions/${encodeURIComponent(sessionId)}/messages`, { method: 'POST', body: JSON.stringify({ text, localId: localId ?? undefined, attachments: attachments ?? undefined, - scheduledAt: scheduledAt ?? undefined + scheduledAt: scheduledAt ?? undefined, + deliveryMode: deliveryMode ?? undefined, }) }) } diff --git a/web/src/chat/normalizeAgent.test.ts b/web/src/chat/normalizeAgent.test.ts index 41bdf31f..5bb00d1b 100644 --- a/web/src/chat/normalizeAgent.test.ts +++ b/web/src/chat/normalizeAgent.test.ts @@ -2,6 +2,98 @@ import { describe, expect, it } from 'vitest' import { normalizeAgentRecord } from '@/chat/normalizeAgent' describe('normalizeAgentRecord — agentTimestamp exposure', () => { + it('preserves a wire text message id as its snapshot stream id', () => { + const normalized = normalizeAgentRecord('message-row-1', null, 1, { + type: 'codex', + data: { + type: 'message', + message: 'partial answer', + id: 'text-stream-1' + } + }) + + expect(normalized).toMatchObject({ + role: 'agent', + content: [{ + type: 'text', + text: 'partial answer', + streamId: 'text-stream-1' + }] + }) + }) + + it('keeps legacy wire text messages without a stream id', () => { + const normalized = normalizeAgentRecord('message-row-1', null, 1, { + type: 'codex', + data: { + type: 'message', + message: 'complete answer' + } + }) + + expect(normalized).toMatchObject({ + role: 'agent', + content: [{ + type: 'text', + text: 'complete answer' + }] + }) + expect((normalized as any).content[0].streamId).toBeUndefined() + }) + + it('keeps cumulative review-like snapshots in the same text stream', () => { + const partial = normalizeAgentRecord('review-row-1', null, 1, { + type: 'codex', + data: { type: 'message', id: 'review-stream', streamSnapshot: true, message: '{"overall_correctness":' } + }) + const complete = normalizeAgentRecord('review-row-2', null, 2, { + type: 'codex', + data: { type: 'message', id: 'review-stream', streamSnapshot: true, message: '{"overall_correctness":"patch is correct"}' } + }) + + expect(partial).toMatchObject({ content: [{ type: 'text', streamId: 'review-stream' }] }) + expect(complete).toMatchObject({ + content: [{ + type: 'text', + streamId: 'review-stream', + text: '{"overall_correctness":"patch is correct"}' + }] + }) + }) + + it('keeps legacy Pi snapshot IDs type-stable without the provenance marker', () => { + const streamId = 'pi-legacy-nonce-turn-1-message-1-text-0' + const partial = normalizeAgentRecord('legacy-review-1', null, 1, { + type: 'codex', + data: { type: 'message', id: streamId, message: '{"overall_correctness":' } + }) + const complete = normalizeAgentRecord('legacy-review-2', null, 2, { + type: 'codex', + data: { type: 'message', id: streamId, message: '{"overall_correctness":"patch is correct"}' } + }) + + expect(partial).toMatchObject({ content: [{ type: 'text', streamId }] }) + expect(complete).toMatchObject({ content: [{ type: 'text', streamId }] }) + }) + + it('still parses a standalone review message that has a normal UUID', () => { + const normalized = normalizeAgentRecord('review-row', null, 1, { + type: 'codex', + data: { + type: 'message', + id: '550e8400-e29b-41d4-a716-446655440000', + message: '{"overall_correctness":"patch is correct"}' + } + }) + + expect(normalized).toMatchObject({ + content: [{ + type: 'codex-review', + review: { overallCorrectness: 'patch is correct' } + }] + }) + }) + it('preserves normalized native tool presentation metadata', () => { const normalized = normalizeAgentRecord('msg-native', null, 1, { type: 'codex', @@ -25,6 +117,28 @@ describe('normalizeAgentRecord — agentTimestamp exposure', () => { }) }) + it('preserves tool progress separately from the tool input', () => { + const normalized = normalizeAgentRecord('progress-row-1', null, 1, { + type: 'codex', + data: { + type: 'tool-call', + callId: 'call-progress', + name: 'Bash', + input: { command: 'bun test' }, + progress: { stdout: 'running tests...\\n' } + } + }) + + expect(normalized).toMatchObject({ + role: 'agent', + content: [{ + type: 'tool-call', + input: { command: 'bun test' }, + progress: { stdout: 'running tests...\\n' } + }] + }) + }) + it('parses data.timestamp into agentTimestamp for an assistant tool_use record', () => { const normalized = normalizeAgentRecord('msg-1', null, 1_783_953_478_235, { type: 'output', diff --git a/web/src/chat/normalizeAgent.ts b/web/src/chat/normalizeAgent.ts index 7ec8b88a..37f84398 100644 --- a/web/src/chat/normalizeAgent.ts +++ b/web/src/chat/normalizeAgent.ts @@ -970,7 +970,10 @@ export function normalizeAgentRecord( } if (data.type === 'message' && typeof data.message === 'string') { - const review = parseCodexReviewMessage(data.message) + const streamId = asString(data.id) + const isPiStreamSnapshot = data.streamSnapshot === true + || (streamId !== null && /^pi-.+-turn-\d+-message-\d+-text-\d+$/.test(streamId)) + const review = isPiStreamSnapshot ? null : parseCodexReviewMessage(data.message) if (review) { return { id: messageId, @@ -988,7 +991,13 @@ export function normalizeAgentRecord( createdAt, role: 'agent', isSidechain: false, - content: [{ type: 'text', text: data.message, uuid: messageId, parentUUID: null }], + content: [{ + type: 'text', + text: data.message, + uuid: messageId, + ...(streamId !== null ? { streamId } : {}), + parentUUID: null + }], meta } } @@ -1089,6 +1098,7 @@ export function normalizeAgentRecord( description: asString(data.description), nativeTitle: asString(data.nativeTitle ?? data.title), nativeKind: asString(data.nativeKind ?? data.kind), + ...('progress' in data ? { progress: data.progress } : {}), uuid, parentUUID: null }], diff --git a/web/src/chat/reducerTimeline.test.ts b/web/src/chat/reducerTimeline.test.ts index 6687b60b..3deb1fe0 100644 --- a/web/src/chat/reducerTimeline.test.ts +++ b/web/src/chat/reducerTimeline.test.ts @@ -296,6 +296,143 @@ describe('reduceTimeline', () => { }) }) + it('collapses text snapshots with the same stream id while leaving legacy text separate', () => { + const first = makeAgentMessage('first ', { + id: 'text-row-1', + content: [{ + type: 'text', + text: 'first ', + uuid: 'text-row-1', + streamId: 'text-stream-1', + parentUUID: null + }] + }) + const second = makeAgentMessage('first second', { + id: 'text-row-2', + content: [{ + type: 'text', + text: 'first second', + uuid: 'text-row-2', + streamId: 'text-stream-1', + parentUUID: null + }] + }) + const legacy = makeAgentMessage('separate legacy text', { + id: 'text-row-3', + content: [{ + type: 'text', + text: 'separate legacy text', + uuid: 'text-row-3', + parentUUID: null + }] + }) + + const { blocks } = reduceTimeline([first, second, legacy], makeContext()) + const textBlocks = blocks.filter((block) => block.kind === 'agent-text') + + expect(textBlocks).toHaveLength(2) + expect(textBlocks[0]).toMatchObject({ + id: 'text-row-1:0', + text: 'first second' + }) + expect(textBlocks[1]).toMatchObject({ + id: 'text-row-3:0', + text: 'separate legacy text' + }) + }) + + it('updates one running tool card with progress snapshots before the final result', () => { + const start: TracedMessage = { + id: 'tool-start', + localId: null, + createdAt: 1_700_000_000_000, + role: 'agent', + content: [{ + type: 'tool-call', + id: 'tool-progress', + name: 'Bash', + input: { command: 'bun test' }, + description: null, + uuid: 'tool-start', + parentUUID: null + }], + isSidechain: false + } as TracedMessage + const progressOne: TracedMessage = { + id: 'tool-progress-one', + localId: null, + createdAt: 1_700_000_000_100, + role: 'agent', + content: [{ + type: 'tool-call', + id: 'tool-progress', + name: 'Bash', + input: { command: 'bun test' }, + description: null, + progress: { stdout: 'one\\n' }, + uuid: 'tool-progress-one', + parentUUID: null + }], + isSidechain: false + } as TracedMessage + const progressTwo: TracedMessage = { + id: 'tool-progress-two', + localId: null, + createdAt: 1_700_000_000_200, + role: 'agent', + content: [{ + type: 'tool-call', + id: 'tool-progress', + name: 'Bash', + input: { command: 'bun test' }, + description: null, + progress: { stdout: 'one\\ntwo\\n' }, + uuid: 'tool-progress-two', + parentUUID: null + }], + isSidechain: false + } as TracedMessage + const end: TracedMessage = { + id: 'tool-end', + localId: null, + createdAt: 1_700_000_000_300, + role: 'agent', + content: [{ + type: 'tool-result', + tool_use_id: 'tool-progress', + content: { stdout: 'done\\n', exitCode: 0 }, + is_error: false, + uuid: 'tool-end', + parentUUID: null + }], + isSidechain: false + } as TracedMessage + + const progressive = reduceTimeline([start, progressOne, progressTwo], makeContext()) + const progressiveToolBlocks = progressive.blocks.filter((block) => block.kind === 'tool-call') + expect(progressiveToolBlocks).toHaveLength(1) + expect(progressiveToolBlocks[0]).toMatchObject({ + tool: { + state: 'running', + input: { command: 'bun test' }, + result: { stdout: 'one\\ntwo\\n' } + } + }) + + const { blocks } = reduceTimeline([start, progressOne, progressTwo, end], makeContext()) + const toolBlocks = blocks.filter((block) => block.kind === 'tool-call') + + expect(toolBlocks).toHaveLength(1) + expect(toolBlocks[0]).toMatchObject({ + tool: { + id: 'tool-progress', + state: 'completed', + input: { command: 'bun test' }, + result: { stdout: 'done\\n', exitCode: 0 } + } + }) + }) + it('falls back to the last duration-bearing block when targetMessageId resolves to a non-duration block', () => { // Regression: the matcher used to take the first id-prefix match and // then silently drop the duration when that block was not duration- diff --git a/web/src/chat/reducerTimeline.ts b/web/src/chat/reducerTimeline.ts index bd3c1a31..25816f89 100644 --- a/web/src/chat/reducerTimeline.ts +++ b/web/src/chat/reducerTimeline.ts @@ -297,6 +297,7 @@ export function reduceTimeline( const agentRunCardByAgentId = new Map() const agentRunTraceMessagesByCardId = new Map() const pendingAgentRunCardByFingerprint = new Map() + const textBlocksByStreamId = new Map() const reasoningBlocksByStreamId = new Map() let hasReadyEvent = false @@ -766,7 +767,20 @@ export function reduceTimeline( })) continue } - blocks.push({ + const streamId = asString(c.streamId) + if (streamId) { + const existing = textBlocksByStreamId.get(streamId) + if (existing) { + existing.text = c.text + existing.usage = msg.usage + existing.model = msg.model + existing.meta = msg.meta + existing.invokedAt = msg.invokedAt + continue + } + } + + const block: AgentTextBlock = { kind: 'agent-text', id: `${msg.id}:${idx}`, localId: msg.localId, @@ -776,7 +790,11 @@ export function reduceTimeline( model: msg.model, text: c.text, meta: msg.meta - }) + } + blocks.push(block) + if (streamId) { + textBlocksByStreamId.set(streamId, block) + } continue } @@ -887,6 +905,7 @@ export function reduceTimeline( description: c.description, nativeTitle: c.nativeTitle, nativeKind: c.nativeKind, + progress: c.progress, permission, agentTimestamp: msg.agentTimestamp }) diff --git a/web/src/chat/reducerTools.ts b/web/src/chat/reducerTools.ts index 2e48600a..63b895da 100644 --- a/web/src/chat/reducerTools.ts +++ b/web/src/chat/reducerTools.ts @@ -67,6 +67,7 @@ export function ensureToolBlock( description: string | null nativeTitle?: string | null nativeKind?: string | null + progress?: unknown permission?: ToolPermission /** Claude entry execution-machine timestamp for the tool_use, if known (see `ChatToolCall.execStartedAt`). */ agentTimestamp?: number | null @@ -110,6 +111,9 @@ export function ensureToolBlock( if (seed.nativeKind != null) { existing.tool = { ...existing.tool, nativeKind: seed.nativeKind } } + if (seed.progress !== undefined && existing.tool.state === 'running') { + existing.tool = { ...existing.tool, result: seed.progress } + } // The first call (tool_use) records when the tool was invoked. The // second call (tool_result) carries the result message's invokedAt, // which is when the result was processed — not when the tool was @@ -152,6 +156,7 @@ export function ensureToolBlock( description: seed.description, nativeTitle: seed.nativeTitle ?? null, nativeKind: seed.nativeKind ?? null, + ...(seed.progress !== undefined ? { result: seed.progress } : {}), permission: seed.permission } diff --git a/web/src/chat/types.ts b/web/src/chat/types.ts index c14b84c3..d7ff10b1 100644 --- a/web/src/chat/types.ts +++ b/web/src/chat/types.ts @@ -48,6 +48,7 @@ export type ToolUse = { description: string | null nativeTitle?: string | null nativeKind?: string | null + progress?: unknown uuid: string parentUUID: string | null } @@ -93,6 +94,7 @@ export type NormalizedAgentContent = type: 'text' text: string uuid: string + streamId?: string parentUUID: string | null } | { diff --git a/web/src/components/AssistantChat/ComposerButtons.test.tsx b/web/src/components/AssistantChat/ComposerButtons.test.tsx index b2663df4..8f3f123e 100644 --- a/web/src/components/AssistantChat/ComposerButtons.test.tsx +++ b/web/src/components/AssistantChat/ComposerButtons.test.tsx @@ -3,8 +3,8 @@ import { type ChatModelAdapter, useLocalRuntime, } from '@assistant-ui/react' -import type { ReactElement, ReactNode } from 'react' -import { cleanup, fireEvent, render, screen, within } from '@testing-library/react' +import type { ComponentProps, ReactElement, ReactNode } from 'react' +import { act, cleanup, fireEvent, render, screen, within } from '@testing-library/react' import { afterEach, describe, expect, it, vi } from 'vitest' import { I18nProvider } from '@/lib/i18n-context' import { @@ -107,6 +107,116 @@ describe('UnifiedButton — routesToScratchlist visual state', () => { }) }) +describe('UnifiedButton — touch queue gesture', () => { + afterEach(() => { + cleanup() + vi.useRealTimers() + }) + + function renderSendButton(overrides: Partial> = {}) { + const onSend = vi.fn() + renderInProviders( + {}} + allowQueueGesture + {...overrides} + />, + ) + const label = overrides.voiceStatus === 'connected' + ? 'Stop' + : overrides.routesToScratchlist + ? /scratchlist/i + : 'Send' + return { onSend, button: getButton(label) } + } + + it('keeps normal native click and keyboard activation as the default send intent', () => { + const { onSend, button } = renderSendButton() + + fireEvent.keyDown(button, { key: 'Enter' }) + expect(onSend).not.toHaveBeenCalled() + fireEvent.click(button, { detail: 1 }) + + expect(onSend).toHaveBeenCalledOnce() + expect(onSend).toHaveBeenCalledWith('default') + }) + + it('uses queue only for a mobile touch long-press and suppresses its native click', () => { + vi.useFakeTimers() + const { onSend, button } = renderSendButton() + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + act(() => vi.advanceTimersByTime(500)) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + fireEvent.click(button, { detail: 1 }) + + expect(onSend).toHaveBeenCalledOnce() + expect(onSend).toHaveBeenCalledWith('queue') + + fireEvent.click(button, { detail: 1 }) + expect(onSend).toHaveBeenCalledTimes(2) + expect(onSend).toHaveBeenLastCalledWith('default') + }) + + it('keeps keyboard and assistive send activation after a long touch has no compatibility click', () => { + vi.useFakeTimers() + const { onSend, button } = renderSendButton() + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + act(() => vi.advanceTimersByTime(500)) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + + // The browser does not emit its touch compatibility click. A detail-0 + // click is the native keyboard/assistive activation path. + fireEvent.click(button, { detail: 0 }) + + expect(onSend).toHaveBeenCalledTimes(2) + expect(onSend).toHaveBeenNthCalledWith(1, 'queue') + expect(onSend).toHaveBeenNthCalledWith(2, 'default') + }) + + it('keeps touch tap, desktop mouse hold, and desktop right-click on normal behavior', () => { + vi.useFakeTimers() + const { onSend, button } = renderSendButton() + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + fireEvent.click(button) + + fireEvent.mouseDown(button, { button: 0, clientX: 10, clientY: 10 }) + act(() => vi.advanceTimersByTime(500)) + fireEvent.mouseUp(button, { button: 0, clientX: 10, clientY: 10 }) + fireEvent.click(button) + const contextMenuWasNotPrevented = fireEvent.contextMenu(button, { clientX: 10, clientY: 10 }) + + expect(onSend).toHaveBeenCalledTimes(2) + expect(onSend).toHaveBeenNthCalledWith(1, 'default') + expect(onSend).toHaveBeenNthCalledWith(2, 'default') + expect(contextMenuWasNotPrevented).toBe(true) + }) + + it.each([ + ['voice is active', { voiceStatus: 'connected' as const }], + ['scratchlist route is active', { routesToScratchlist: true }], + ['queue gesture is disabled', { allowQueueGesture: false }], + ])('does not queue on long-press when %s', (_name, overrides) => { + vi.useFakeTimers() + const { onSend, button } = renderSendButton(overrides) + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + act(() => vi.advanceTimersByTime(500)) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + fireEvent.click(button) + + expect(onSend).not.toHaveBeenCalledWith('queue') + }) +}) + describe('DictationButton', () => { afterEach(cleanup) diff --git a/web/src/components/AssistantChat/ComposerButtons.tsx b/web/src/components/AssistantChat/ComposerButtons.tsx index 61fa9213..faa10f04 100644 --- a/web/src/components/AssistantChat/ComposerButtons.tsx +++ b/web/src/components/AssistantChat/ComposerButtons.tsx @@ -8,6 +8,8 @@ import { useFue } from '@/lib/use-fue' import { FueCallout, FueDot } from '@/components/Fue' import { Children, isValidElement, useRef, useState, type ReactElement, type ReactNode } from 'react' import { useComposerToolbarLayout, type ComposerToolbarItemId, type ComposerToolbarLayout } from '@/hooks/useComposerToolbarLayout' +import { useLongPress } from '@/hooks/useLongPress' +import type { ComposerSendIntent } from '@/lib/messageDelivery' function ToolbarItemSlot(props: { item: ComposerToolbarItemId; children: ReactNode }) { return <>{props.children} @@ -473,7 +475,7 @@ export function UnifiedButton(props: { voiceStatus: ConversationStatus voiceEnabled: boolean controlsDisabled: boolean - onSend: () => void + onSend: (intent?: ComposerSendIntent) => void onVoiceToggle: () => void voiceLabel?: string /** @@ -488,6 +490,8 @@ export function UnifiedButton(props: { * would fall back to chat, the button must look like a normal chat send. */ routesToScratchlist?: boolean + /** Pi-only explicit follow-up gesture; never changes the normal click. */ + allowQueueGesture?: boolean }) { const { t } = useTranslation() @@ -501,12 +505,31 @@ export function UnifiedButton(props: { if (isVoiceActive) { props.onVoiceToggle() // Stop voice } else if (hasText) { - props.onSend() // Send message (or scratchlist add — wrapper decides) + props.onSend('default') // Send message (or scratchlist add — wrapper decides) } else if (props.voiceEnabled && !routesToScratchlist) { props.onVoiceToggle() // Start voice (suppressed in scratchlist mode) } } + // This is intentionally narrower than the button's general enabled state: + // a touch hold changes only an active Pi-main-thread chat submission. Voice + // controls, scratchlist routing, scheduled sends, and desktop input retain + // their existing native behavior. + const canQueueGesture = Boolean( + props.allowQueueGesture + && hasText + && !isVoiceActive + && !routesToScratchlist + && !props.controlsDisabled, + ) + const sendButtonHandlers = useLongPress({ + interaction: 'touch-only-native-click', + onClick: handleClick, + onLongPress: () => props.onSend('queue'), + longPressEnabled: canQueueGesture, + disabled: props.controlsDisabled, + }) + let icon: React.ReactNode let className: string let ariaLabel: string @@ -554,7 +577,7 @@ export function UnifiedButton(props: { return ( + ) +} + describe('useLongPress', () => { let now = 10_000 @@ -117,4 +135,192 @@ describe('useLongPress', () => { expect(onClick).toHaveBeenCalledTimes(2) }) + + it('keeps native button click semantics for a touch tap', () => { + const onClick = vi.fn() + const onLongPress = vi.fn() + const { getByTestId } = render() + const button = getByTestId('native-button') + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + fireEvent.click(button) + + expect(onLongPress).not.toHaveBeenCalled() + expect(onClick).toHaveBeenCalledOnce() + }) + + it('fires a touch-only long press once and consumes exactly its following touch-derived click', () => { + const onClick = vi.fn() + const onLongPress = vi.fn() + const { getByTestId } = render() + const button = getByTestId('native-button') + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + act(() => { + now += 500 + vi.advanceTimersByTime(500) + }) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + // A browser compatibility click has click detail 1. Keyboard and + // assistive activation instead reports detail 0. + fireEvent.click(button, { detail: 1 }) + + expect(onLongPress).toHaveBeenCalledOnce() + expect(onClick).not.toHaveBeenCalled() + expect(vi.getTimerCount()).toBe(0) + + // Only the touch compatibility click is suppressed. The next ordinary + // click is the button's normal action. + fireEvent.click(button, { detail: 1 }) + expect(onClick).toHaveBeenCalledOnce() + }) + + it('keeps keyboard and assistive native activation after a long touch with no compatibility click', () => { + const onClick = vi.fn() + const onLongPress = vi.fn() + const { getByTestId } = render() + const button = getByTestId('native-button') + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + act(() => { + now += 500 + vi.advanceTimersByTime(500) + }) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + + // No compatibility click arrives. A detail-zero native click models + // keyboard or assistive-technology activation and must not be lost. + fireEvent.click(button, { detail: 0 }) + + expect(onLongPress).toHaveBeenCalledOnce() + expect(onClick).toHaveBeenCalledOnce() + }) + + it('expires native touch-click suppression so later mouse click and context menu stay native', () => { + const onClick = vi.fn() + const onLongPress = vi.fn() + const { getByTestId } = render() + const button = getByTestId('native-button') + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + act(() => { + now += 500 + vi.advanceTimersByTime(500) + }) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + + // No compatibility click arrives during the bounded suppression window. + act(() => { + now += 700 + vi.advanceTimersByTime(700) + }) + const contextMenuWasNotPrevented = fireEvent.contextMenu(button, { clientX: 10, clientY: 10 }) + fireEvent.click(button, { detail: 1 }) + + expect(onLongPress).toHaveBeenCalledOnce() + expect(contextMenuWasNotPrevented).toBe(true) + expect(onClick).toHaveBeenCalledOnce() + }) + + it('clears pending native click suppression when a new touch is cancelled', () => { + const onClick = vi.fn() + const onLongPress = vi.fn() + const { getByTestId } = render() + const button = getByTestId('native-button') + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + act(() => { + now += 500 + vi.advanceTimersByTime(500) + }) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + + // A new touch clears the old long-touch suppression; cancellation also + // clears the new touch's hold timer before any native click arrives. + fireEvent.touchStart(button, { touches: [{ clientX: 20, clientY: 20 }] }) + fireEvent.touchCancel(button) + fireEvent.click(button, { detail: 1 }) + + expect(onLongPress).toHaveBeenCalledOnce() + expect(onClick).toHaveBeenCalledOnce() + expect(vi.getTimerCount()).toBe(0) + }) + + it('cleans the pending native click suppression timer on unmount', () => { + const onClick = vi.fn() + const onLongPress = vi.fn() + const { getByTestId, unmount } = render() + const button = getByTestId('native-button') + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + act(() => { + now += 500 + vi.advanceTimersByTime(500) + }) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + expect(vi.getTimerCount()).toBe(1) + + unmount() + + expect(vi.getTimerCount()).toBe(0) + }) + + it('does not turn desktop hold or right-click into a touch-only long press', () => { + const onClick = vi.fn() + const onLongPress = vi.fn() + const { getByTestId } = render() + const button = getByTestId('native-button') + + fireEvent.mouseDown(button, { button: 0, clientX: 10, clientY: 10 }) + act(() => { + now += 500 + vi.advanceTimersByTime(500) + }) + fireEvent.mouseUp(button, { button: 0, clientX: 10, clientY: 10 }) + fireEvent.click(button) + const contextMenuWasNotPrevented = fireEvent.contextMenu(button, { clientX: 10, clientY: 10 }) + + expect(onLongPress).not.toHaveBeenCalled() + expect(onClick).toHaveBeenCalledOnce() + expect(contextMenuWasNotPrevented).toBe(true) + }) + + it('can disable only the long-press override without disabling normal click', () => { + const onClick = vi.fn() + const onLongPress = vi.fn() + const { getByTestId } = render( + , + ) + const button = getByTestId('native-button') + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + act(() => { + now += 500 + vi.advanceTimersByTime(500) + }) + fireEvent.touchEnd(button, { changedTouches: [{ clientX: 10, clientY: 10 }] }) + fireEvent.click(button) + + expect(onLongPress).not.toHaveBeenCalled() + expect(onClick).toHaveBeenCalledOnce() + }) + + it('cancels a touch-only long press when the browser cancels the touch', () => { + const onClick = vi.fn() + const onLongPress = vi.fn() + const { getByTestId } = render() + const button = getByTestId('native-button') + + fireEvent.touchStart(button, { touches: [{ clientX: 10, clientY: 10 }] }) + fireEvent.touchCancel(button) + act(() => { + now += 500 + vi.advanceTimersByTime(500) + }) + fireEvent.click(button) + + expect(onLongPress).not.toHaveBeenCalled() + expect(onClick).toHaveBeenCalledOnce() + }) }) diff --git a/web/src/hooks/useLongPress.ts b/web/src/hooks/useLongPress.ts index 4e9d372d..3f8aa3c2 100644 --- a/web/src/hooks/useLongPress.ts +++ b/web/src/hooks/useLongPress.ts @@ -1,17 +1,32 @@ import type React from 'react' -import { useCallback, useRef } from 'react' +import { useCallback, useEffect, useRef } from 'react' type UseLongPressOptions = { onLongPress: (point: { x: number; y: number }) => void onClick?: () => void threshold?: number disabled?: boolean + /** + * `legacy` preserves the original list-row contract: this hook emits + * clicks itself from mouse/touch/key handlers. `touch-only-native-click` + * is for an existing native button: only touch may long-press, while + * click/keyboard accessibility remain browser-native. + */ + interaction?: 'legacy' | 'touch-only-native-click' + /** Disable just the long-press gesture while retaining the normal click. */ + longPressEnabled?: boolean } // How long after a touch interaction to keep ignoring synthesized mouse // events. Android's compatibility mouse events fire ~300ms after touchend; // 700ms covers that with margin without affecting genuine later mouse input. const GHOST_MOUSE_WINDOW_MS = 700 +// Native buttons retain their platform click behavior. A touch long-press has +// already performed its action, so only the compatibility click emitted just +// after that touch should be discarded. Keep this in the same bounded window +// as touch compatibility mouse events; a missing compatibility click must not +// poison a later mouse, keyboard, or assistive activation. +const NATIVE_CLICK_SUPPRESSION_WINDOW_MS = GHOST_MOUSE_WINDOW_MS type UseLongPressHandlers = { onMouseDown: React.MouseEventHandler @@ -20,17 +35,31 @@ type UseLongPressHandlers = { onTouchStart: React.TouchEventHandler onTouchEnd: React.TouchEventHandler onTouchMove: React.TouchEventHandler + onTouchCancel: React.TouchEventHandler onContextMenu: React.MouseEventHandler onKeyDown: React.KeyboardEventHandler + onClick?: React.MouseEventHandler } export function useLongPress(options: UseLongPressOptions): UseLongPressHandlers { - const { onLongPress, onClick, threshold = 500, disabled = false } = options + const { + onLongPress, + onClick, + threshold = 500, + disabled = false, + interaction = 'legacy', + longPressEnabled = true, + } = options const timerRef = useRef | null>(null) const isLongPressRef = useRef(false) const touchMoved = useRef(false) const pressPointRef = useRef<{ x: number; y: number }>({ x: 0, y: 0 }) + // Used only by the native-button mode. A long touch has already sent its + // action when the browser produces the following native click, so consume + // exactly that one click without changing later mouse/keyboard behavior. + const suppressNextNativeClickRef = useRef(false) + const nativeClickSuppressionTimerRef = useRef | null>(null) // Timestamp of the most recent touch interaction. Touch browsers emit // compatibility mouse events (mousedown/mouseup/click) after a tap for any // touch the page did not preventDefault. Since we bind BOTH touch and @@ -46,8 +75,30 @@ export function useLongPress(options: UseLongPressOptions): UseLongPressHandlers } }, []) + const clearNativeClickSuppression = useCallback(() => { + suppressNextNativeClickRef.current = false + if (nativeClickSuppressionTimerRef.current) { + clearTimeout(nativeClickSuppressionTimerRef.current) + nativeClickSuppressionTimerRef.current = null + } + }, []) + + const armNativeClickSuppression = useCallback(() => { + clearNativeClickSuppression() + suppressNextNativeClickRef.current = true + nativeClickSuppressionTimerRef.current = setTimeout(() => { + suppressNextNativeClickRef.current = false + nativeClickSuppressionTimerRef.current = null + }, NATIVE_CLICK_SUPPRESSION_WINDOW_MS) + }, [clearNativeClickSuppression]) + + useEffect(() => () => { + clearTimer() + clearNativeClickSuppression() + }, [clearTimer, clearNativeClickSuppression]) + const startTimer = useCallback((clientX: number, clientY: number) => { - if (disabled) return + if (disabled || !longPressEnabled) return clearTimer() isLongPressRef.current = false @@ -58,7 +109,7 @@ export function useLongPress(options: UseLongPressOptions): UseLongPressHandlers isLongPressRef.current = true onLongPress(pressPointRef.current) }, threshold) - }, [disabled, clearTimer, onLongPress, threshold]) + }, [disabled, longPressEnabled, clearTimer, onLongPress, threshold]) const handleEnd = useCallback((shouldTriggerClick: boolean) => { clearTimer() @@ -130,6 +181,77 @@ export function useLongPress(options: UseLongPressOptions): UseLongPressHandlers } }, [disabled, onClick]) + const onTouchCancel = useCallback(() => { + clearTimer() + isLongPressRef.current = false + touchMoved.current = false + clearNativeClickSuppression() + }, [clearTimer, clearNativeClickSuppression]) + + const onNativeTouchStart = useCallback((e) => { + const touch = e.touches[0] + if (!touch) return + // A new physical touch starts a new native click sequence. If a prior + // long touch did not yield a browser click at all, do not suppress this + // later genuine activation. + clearNativeClickSuppression() + startTimer(touch.clientX, touch.clientY) + }, [clearNativeClickSuppression, startTimer]) + + const onNativeTouchEnd = useCallback(() => { + clearTimer() + if (isLongPressRef.current) { + armNativeClickSuppression() + } + isLongPressRef.current = false + touchMoved.current = false + }, [armNativeClickSuppression, clearTimer]) + + const onNativeTouchMove = useCallback(() => { + touchMoved.current = true + clearTimer() + }, [clearTimer]) + + const onNativeContextMenu = useCallback((e) => { + // A touch long-press may produce a context menu before or after + // touchend. Suppress only that generated menu; desktop right-click + // stays completely native and never becomes a queue gesture. + if (isLongPressRef.current || suppressNextNativeClickRef.current) { + e.preventDefault() + } + }, []) + + const onNativeClick = useCallback((e) => { + if (disabled) return + // Browser-generated keyboard and assistive-technology activation uses + // detail === 0. It is not the touch compatibility click and must retain + // native button semantics even while a stale touch window is pending. + if (suppressNextNativeClickRef.current && e.detail !== 0) { + clearNativeClickSuppression() + e.preventDefault() + e.stopPropagation() + return + } + onClick?.() + }, [clearNativeClickSuppression, disabled, onClick]) + + if (interaction === 'touch-only-native-click') { + + return { + // Existing button semantics own desktop mouse and keyboard clicks. + onMouseDown: () => {}, + onMouseUp: () => {}, + onMouseLeave: () => {}, + onTouchStart: onNativeTouchStart, + onTouchEnd: onNativeTouchEnd, + onTouchMove: onNativeTouchMove, + onTouchCancel, + onContextMenu: onNativeContextMenu, + onKeyDown: () => {}, + onClick: onNativeClick, + } + } + return { onMouseDown, onMouseUp, @@ -137,6 +259,7 @@ export function useLongPress(options: UseLongPressOptions): UseLongPressHandlers onTouchStart, onTouchEnd, onTouchMove, + onTouchCancel, onContextMenu, onKeyDown } diff --git a/web/src/lib/assistant-runtime.ts b/web/src/lib/assistant-runtime.ts index e670c4a2..3105d43b 100644 --- a/web/src/lib/assistant-runtime.ts +++ b/web/src/lib/assistant-runtime.ts @@ -4,6 +4,10 @@ import type { AppendMessage, AttachmentAdapter, ThreadMessageLike } from '@assis import { useExternalMessageConverter, useExternalStoreRuntime } from '@assistant-ui/react' import type { PendingSchedule } from '@/components/AssistantChat/ScheduleTimePicker' import { resolvePendingSchedule } from '@/components/AssistantChat/ScheduleTimePicker' +import { + consumeComposerSendIntent, + type ComposerSendIntent, +} from '@/lib/messageDelivery' import { safeStringify } from '@hapi/protocol' import { renderEventLabel } from '@/chat/presentation' import type { ChatBlock, CliOutputBlock, CodexReview, UsageData } from '@/chat/types' @@ -611,11 +615,22 @@ export function useHappyRuntime(props: { historyVersion: number isSending: boolean isRunning?: boolean - onSendMessage: (text: string, attachments?: AttachmentMetadata[], scheduledAt?: number | null) => void + onSendMessage: ( + text: string, + attachments?: AttachmentMetadata[], + scheduledAt?: number | null, + intent?: ComposerSendIntent, + ) => void onAbort: () => Promise attachmentAdapter?: AttachmentAdapter allowSendWhenInactive?: boolean pendingScheduleRef?: React.RefObject + /** + * Shared one-shot ref with HappyComposer. The composer marks the next + * `api.composer().send()`; this adapter consumes and resets the mark as + * soon as assistant-ui emits the corresponding AppendMessage. + */ + pendingSendIntentRef?: React.MutableRefObject }) { const isRunning = props.isRunning ?? props.session.thinking @@ -681,6 +696,10 @@ export function useHappyRuntime(props: { }) const onNew = useCallback(async (message: AppendMessage) => { + const intent = consumeComposerSendIntent(props.pendingSendIntentRef) + // Reset before any early return so an empty submission, extraction + // failure, or downstream exception cannot leak an explicit queue + // gesture into the next ordinary send. const { text, attachments } = extractMessageContent(message) if (!text && attachments.length === 0) return // Resolve pendingSchedule at send time (Date.now()) so preset-type schedules @@ -688,8 +707,8 @@ export function useHappyRuntime(props: { // moment the user clicked the preset button. const sendNow = Date.now() const scheduledAt = resolvePendingSchedule(props.pendingScheduleRef?.current ?? null, sendNow) - props.onSendMessage(text, attachments.length > 0 ? attachments : undefined, scheduledAt) - }, [props.onSendMessage, props.pendingScheduleRef]) + props.onSendMessage(text, attachments.length > 0 ? attachments : undefined, scheduledAt, intent) + }, [props.onSendMessage, props.pendingScheduleRef, props.pendingSendIntentRef]) const onCancel = useCallback(async () => { await props.onAbort() diff --git a/web/src/lib/messageDelivery.test.ts b/web/src/lib/messageDelivery.test.ts new file mode 100644 index 00000000..646b4984 --- /dev/null +++ b/web/src/lib/messageDelivery.test.ts @@ -0,0 +1,66 @@ +import { describe, expect, it } from 'vitest' +import { + consumeComposerSendIntent, + getRestoredComposerSendIntent, + getRetryDeliveryMode, + resolveMessageDeliveryMode, +} from './messageDelivery' + +describe('consumeComposerSendIntent', () => { + it('consumes an explicit queue request exactly once', () => { + const ref = { current: 'queue' as const } + + expect(consumeComposerSendIntent(ref)).toBe('queue') + expect(ref.current).toBe('default') + expect(consumeComposerSendIntent(ref)).toBe('default') + }) + + it('defaults safely when no composer ref is present', () => { + expect(consumeComposerSendIntent()).toBe('default') + }) +}) + +describe('retry delivery safety', () => { + it('preserves queue and downgrades steer or missing legacy provenance to queue', () => { + expect(getRetryDeliveryMode('queue')).toBe('queue') + expect(getRetryDeliveryMode('steer')).toBe('queue') + expect(getRetryDeliveryMode(undefined)).toBe('queue') + expect(getRestoredComposerSendIntent('queue')).toBe('queue') + expect(getRestoredComposerSendIntent('steer')).toBe('queue') + expect(getRestoredComposerSendIntent(undefined)).toBe('queue') + }) +}) + +describe('resolveMessageDeliveryMode', () => { + const base = { + agentFlavor: 'pi', + isSessionThinking: true, + intent: 'default' as const, + } + + it('steers an immediate fresh Pi send while the main session is thinking', () => { + expect(resolveMessageDeliveryMode(base)).toBe('steer') + }) + + it('keeps an explicit queue gesture queued even while Pi is thinking', () => { + expect(resolveMessageDeliveryMode({ ...base, intent: 'queue' })).toBe('queue') + }) + + it('queues a failed steer retry even when a later Pi generation is active', () => { + const ref = { current: getRestoredComposerSendIntent('steer') } + const retryIntent = consumeComposerSendIntent(ref) + + expect(resolveMessageDeliveryMode({ ...base, intent: retryIntent })).toBe('queue') + expect(ref.current).toBe('default') + expect(resolveMessageDeliveryMode({ ...base, intent: consumeComposerSendIntent(ref) })).toBe('steer') + }) + + it.each([ + { name: 'idle Pi', input: { ...base, isSessionThinking: false } }, + { name: 'non-Pi flavor', input: { ...base, agentFlavor: 'codex' } }, + { name: 'scheduled message', input: { ...base, scheduledAt: Date.now() + 60_000 } }, + { name: 'scratchlist route', input: { ...base, routesToScratchlist: true } }, + ])('queues $name', ({ input }) => { + expect(resolveMessageDeliveryMode(input)).toBe('queue') + }) +}) diff --git a/web/src/lib/messageDelivery.ts b/web/src/lib/messageDelivery.ts new file mode 100644 index 00000000..48242411 --- /dev/null +++ b/web/src/lib/messageDelivery.ts @@ -0,0 +1,63 @@ +import type { MessageDeliveryMode } from '@hapi/protocol' + +/** + * The one-shot UI intent associated with a composer submission. It is not + * the wire delivery mode: `default` is resolved against the current session + * state at the SessionChat boundary, while `queue` is an explicit operator + * request not to steer an in-flight Pi turn. + */ +export type ComposerSendIntent = 'default' | 'queue' + +/** Structural shape shared by React's mutable ref and the runtime adapter. */ +export type ComposerSendIntentRef = { current: ComposerSendIntent } + +/** + * Read exactly one composer intent and immediately return the shared ref to + * the ordinary-send default. This is intentionally independent of React so + * the assistant-ui adapter can consume the value synchronously in `onNew`. + */ +export function consumeComposerSendIntent(ref?: ComposerSendIntentRef): ComposerSendIntent { + const intent = ref?.current ?? 'default' + if (ref) ref.current = 'default' + return intent +} + +/** + * A retry cannot prove that the original Pi turn is still active. Preserve an + * explicit queue, but downgrade turn-scoped steer (and legacy missing mode) to + * the durable HAPI queue instead of binding the retry to a later generation. + */ +export function getRetryDeliveryMode( + deliveryMode: MessageDeliveryMode | undefined, +): 'queue' { + return deliveryMode === 'steer' ? 'queue' : (deliveryMode ?? 'queue') +} + +/** Convert retry delivery into the composer's one-shot intent representation. */ +export function getRestoredComposerSendIntent( + deliveryMode: MessageDeliveryMode | undefined, +): ComposerSendIntent { + return getRetryDeliveryMode(deliveryMode) +} + +/** + * Resolve the web composer intent into the durable message delivery mode. + * + * Fresh steering is deliberately narrow: it only applies to an immediate + * ordinary composer submission while the Pi *main session* reports that it + * is thinking. Scheduled messages, scratchlist additions, and retries never + * steer because none can prove the original turn identity. + */ +export function resolveMessageDeliveryMode(input: { + agentFlavor: string | null | undefined + isSessionThinking: boolean + intent: ComposerSendIntent + scheduledAt?: number | null + routesToScratchlist?: boolean +}): MessageDeliveryMode { + if (input.scheduledAt != null) return 'queue' + if (input.routesToScratchlist === true) return 'queue' + if (input.intent === 'queue') return 'queue' + if (input.agentFlavor !== 'pi') return 'queue' + return input.isSessionThinking ? 'steer' : 'queue' +} diff --git a/web/src/router.tsx b/web/src/router.tsx index 96c865bf..01fd373b 100644 --- a/web/src/router.tsx +++ b/web/src/router.tsx @@ -37,6 +37,7 @@ import type { Suggestion } from '@/hooks/useActiveSuggestions' import { useSendMessage, type SendErrorInfo } from '@/hooks/mutations/useSendMessage' import type { ComposerSendError } from '@/components/AssistantChat/HappyComposer' import { ApiError } from '@/api/client' +import type { MessageDeliveryMode } from '@hapi/protocol' import { queryKeys } from '@/lib/query-keys' import { useToast } from '@/lib/toast-context' import { useTranslation } from '@/lib/use-translation' @@ -381,6 +382,7 @@ function SessionPage() { message: string code: string | null scheduledAt: number | null + deliveryMode: MessageDeliveryMode mutationStarted: boolean restoreSuppressed: boolean } @@ -468,6 +470,7 @@ function SessionPage() { text: rawSendError.text, message: rawSendError.message, scheduledAt: rawSendError.scheduledAt, + deliveryMode: rawSendError.deliveryMode, mutationStarted: rawSendError.mutationStarted, restoreSuppressed: rawSendError.restoreSuppressed, action: rawSendError.code === 'session_inactive' && canOfferInactiveReopen @@ -510,6 +513,7 @@ function SessionPage() { message, code, scheduledAt: info.scheduledAt, + deliveryMode: info.deliveryMode, mutationStarted: info.mutationStarted, restoreSuppressed: false, } diff --git a/web/src/types/api.ts b/web/src/types/api.ts index 37c1d10c..0ecf12a1 100644 --- a/web/src/types/api.ts +++ b/web/src/types/api.ts @@ -100,6 +100,7 @@ export type SessionMetadataSummary = { conversationHistoryPoints?: Record conversationHistoryIndexes?: Record conversationHistoryTurns?: Record + conversationHistoryEntryIds?: Record conversationHistoryDiverged?: boolean worktree?: WorktreeMetadata }