400 lines
13 KiB
TypeScript
400 lines
13 KiB
TypeScript
import type {
|
|
CodingRuntimeErrorCode,
|
|
ConversationContentBlock,
|
|
ConversationMessageNode,
|
|
ConversationNode,
|
|
ConversationSnapshot,
|
|
ConversationToolNode,
|
|
PublicUsage,
|
|
} from '../contracts';
|
|
import type {
|
|
PiEventProjectorOptions,
|
|
PiImageProjectionInput,
|
|
PiProjectedAttachment,
|
|
} from './event-projector';
|
|
import { isProductToolName, productToolDetailsOfResult } from '../product-tool-protocol';
|
|
import { subagentDetailsOfResult } from '../subagent-protocol';
|
|
|
|
export interface PiSessionSnapshotInput {
|
|
snapshot: ConversationSnapshot;
|
|
workerGeneration: number;
|
|
state: unknown;
|
|
entries: unknown;
|
|
stats?: unknown;
|
|
projectImage?(image: PiImageProjectionInput): Promise<PiProjectedAttachment>;
|
|
}
|
|
|
|
interface SessionEntryRecord extends Record<string, unknown> {
|
|
id: string;
|
|
parentId: string | null;
|
|
type: string;
|
|
}
|
|
|
|
export class PiSessionProjectionError extends Error {
|
|
readonly code: CodingRuntimeErrorCode = 'CODING_SESSION_UNREADABLE';
|
|
readonly recoverable = true;
|
|
|
|
constructor() {
|
|
super('Pi session data is unreadable; the source session was left unchanged.');
|
|
this.name = 'PiSessionProjectionError';
|
|
}
|
|
}
|
|
|
|
function unreadable(): never {
|
|
throw new PiSessionProjectionError();
|
|
}
|
|
|
|
function asRecord(value: unknown): Record<string, unknown> | null {
|
|
return value !== null && typeof value === 'object' && !Array.isArray(value)
|
|
? value as Record<string, unknown>
|
|
: null;
|
|
}
|
|
|
|
function sessionEntry(value: unknown): SessionEntryRecord | null {
|
|
const entry = asRecord(value);
|
|
if (!entry
|
|
|| typeof entry.id !== 'string'
|
|
|| (entry.parentId !== null && typeof entry.parentId !== 'string')
|
|
|| typeof entry.type !== 'string') return null;
|
|
return entry as SessionEntryRecord;
|
|
}
|
|
|
|
function usageOf(value: unknown): PublicUsage | undefined {
|
|
const usage = asRecord(value);
|
|
if (!usage
|
|
|| typeof usage.input !== 'number'
|
|
|| typeof usage.output !== 'number') return undefined;
|
|
return {
|
|
inputTokens: usage.input,
|
|
outputTokens: usage.output,
|
|
...(typeof usage.cacheRead === 'number' ? { cacheReadTokens: usage.cacheRead } : {}),
|
|
...(typeof usage.cacheWrite === 'number' ? { cacheWriteTokens: usage.cacheWrite } : {}),
|
|
};
|
|
}
|
|
|
|
function stopReasonOf(value: unknown): ConversationMessageNode['stopReason'] | undefined {
|
|
if (value === 'toolUse') return 'tool-use';
|
|
if (value === 'stop' || value === 'length' || value === 'error' || value === 'aborted') {
|
|
return value;
|
|
}
|
|
return undefined;
|
|
}
|
|
|
|
function activePath(entriesValue: unknown, leafId: unknown): SessionEntryRecord[] {
|
|
if (!Array.isArray(entriesValue)) unreadable();
|
|
if (leafId !== null && typeof leafId !== 'string') unreadable();
|
|
const index = new Map<string, SessionEntryRecord>();
|
|
for (const value of entriesValue) {
|
|
const entry = sessionEntry(value);
|
|
if (!entry || index.has(entry.id)) unreadable();
|
|
index.set(entry.id, entry);
|
|
}
|
|
if (leafId === null) return [];
|
|
const path: SessionEntryRecord[] = [];
|
|
const visited = new Set<string>();
|
|
let current = index.get(leafId);
|
|
if (!current) unreadable();
|
|
while (current) {
|
|
if (visited.has(current.id)) unreadable();
|
|
visited.add(current.id);
|
|
path.push(current);
|
|
if (current.parentId === null) break;
|
|
const parent = index.get(current.parentId);
|
|
if (!parent) unreadable();
|
|
current = parent;
|
|
}
|
|
return path.reverse();
|
|
}
|
|
|
|
function retainedTail(path: SessionEntryRecord[]): SessionEntryRecord[] {
|
|
const compactionIndex = path.findLastIndex((entry) => entry.type === 'compaction');
|
|
if (compactionIndex < 0) return path;
|
|
const compaction = path[compactionIndex];
|
|
const firstKeptEntryId = compaction.firstKeptEntryId;
|
|
if (typeof firstKeptEntryId !== 'string') unreadable();
|
|
const firstKeptIndex = path.findIndex(
|
|
(entry, index) => index < compactionIndex && entry.id === firstKeptEntryId,
|
|
);
|
|
if (firstKeptIndex < 0) unreadable();
|
|
return [
|
|
compaction,
|
|
...path.slice(firstKeptIndex, compactionIndex),
|
|
...path.slice(compactionIndex + 1),
|
|
];
|
|
}
|
|
|
|
async function contentBlocks(
|
|
messageId: string,
|
|
content: unknown,
|
|
conversationId: string,
|
|
projectImage: PiEventProjectorOptions['projectImage'],
|
|
idPart: 'content' | 'output',
|
|
): Promise<ConversationContentBlock[]> {
|
|
const values = typeof content === 'string'
|
|
? [{ type: 'text', text: content }]
|
|
: Array.isArray(content) ? content : [];
|
|
const blocks: ConversationContentBlock[] = [];
|
|
for (let contentIndex = 0; contentIndex < values.length; contentIndex += 1) {
|
|
const block = asRecord(values[contentIndex]);
|
|
if (block?.type === 'thinking' && typeof block.thinking === 'string') {
|
|
blocks.push({
|
|
kind: 'thinking',
|
|
id: `${messageId}:${idPart}:${contentIndex}`,
|
|
text: block.thinking,
|
|
status: 'complete',
|
|
});
|
|
} else if (block?.type === 'text' && typeof block.text === 'string') {
|
|
blocks.push({
|
|
kind: 'text',
|
|
id: `${messageId}:${idPart}:${contentIndex}`,
|
|
text: block.text,
|
|
status: 'complete',
|
|
});
|
|
} else if (block?.type === 'image'
|
|
&& typeof block.data === 'string'
|
|
&& typeof block.mimeType === 'string'
|
|
&& projectImage) {
|
|
const attachment = await projectImage({
|
|
conversationId,
|
|
data: block.data,
|
|
mime: block.mimeType,
|
|
source: 'session',
|
|
});
|
|
blocks.push({
|
|
kind: 'image',
|
|
id: `${messageId}:${idPart}:${contentIndex}`,
|
|
attachmentId: attachment.attachmentId,
|
|
mime: attachment.mime,
|
|
});
|
|
}
|
|
}
|
|
return blocks;
|
|
}
|
|
|
|
function inputTextOf(value: unknown): string {
|
|
if (typeof value === 'string') return value;
|
|
try {
|
|
return JSON.stringify(value ?? {}) ?? '{}';
|
|
} catch {
|
|
return '{}';
|
|
}
|
|
}
|
|
|
|
async function projectEntries(
|
|
path: SessionEntryRecord[],
|
|
input: PiSessionSnapshotInput,
|
|
): Promise<ConversationNode[]> {
|
|
const nodes: ConversationNode[] = [];
|
|
const tools = new Map<string, ConversationToolNode>();
|
|
const subagentIds = new Set<string>();
|
|
for (const entry of path) {
|
|
if (entry.type === 'compaction') {
|
|
nodes.push({
|
|
kind: 'compaction',
|
|
id: `entry:${entry.id}`,
|
|
runId: `session:${entry.id}`,
|
|
source: 'automatic',
|
|
status: 'complete',
|
|
willRetry: false,
|
|
});
|
|
continue;
|
|
}
|
|
if (entry.type !== 'message') continue;
|
|
const message = asRecord(entry.message);
|
|
if (!message) unreadable();
|
|
if (message.role === 'toolResult') {
|
|
if (typeof message.toolCallId !== 'string') unreadable();
|
|
const tool = tools.get(message.toolCallId);
|
|
if (!tool) continue;
|
|
tool.output = await contentBlocks(
|
|
tool.id,
|
|
message.content,
|
|
input.snapshot.conversation.id,
|
|
input.projectImage,
|
|
'output',
|
|
);
|
|
tool.status = message.isError === true ? 'error' : 'complete';
|
|
const subagentDetails = subagentDetailsOfResult(message);
|
|
const details = subagentDetails ?? productToolDetailsOfResult(message);
|
|
if (details) tool.details = details;
|
|
if (subagentDetails) {
|
|
const id = `subagent:${subagentDetails.dispatchId}`;
|
|
if (!subagentIds.has(id)) {
|
|
subagentIds.add(id);
|
|
nodes.push({
|
|
kind: 'subagent',
|
|
id,
|
|
runId: input.snapshot.run.runId ?? `session:${entry.id}`,
|
|
details: subagentDetails,
|
|
});
|
|
}
|
|
} else if ((tool.toolName === 'subagent' || isProductToolName(tool.toolName))
|
|
&& message.details !== undefined
|
|
&& !details) {
|
|
tool.output = [{
|
|
kind: 'text',
|
|
id: `${tool.id}:output:unavailable`,
|
|
text: tool.toolName === 'subagent'
|
|
? 'Subagent details are unavailable for this version.'
|
|
: 'Tool details are unavailable for this version.',
|
|
status: 'complete',
|
|
}];
|
|
}
|
|
continue;
|
|
}
|
|
if (message.role !== 'user' && message.role !== 'assistant') continue;
|
|
const messageId = `entry:${entry.id}`;
|
|
const usage = usageOf(message.usage);
|
|
const stopReason = stopReasonOf(message.stopReason);
|
|
const node: ConversationMessageNode = {
|
|
kind: 'message',
|
|
id: messageId,
|
|
sourceEntryId: entry.id,
|
|
role: message.role,
|
|
status: message.stopReason === 'aborted'
|
|
? 'aborted'
|
|
: message.stopReason === 'error' ? 'error' : 'complete',
|
|
blocks: await contentBlocks(
|
|
messageId,
|
|
message.content,
|
|
input.snapshot.conversation.id,
|
|
input.projectImage,
|
|
'content',
|
|
),
|
|
...(usage ? { usage } : {}),
|
|
...(stopReason ? { stopReason } : {}),
|
|
};
|
|
nodes.push(node);
|
|
if (!Array.isArray(message.content)) continue;
|
|
for (const value of message.content) {
|
|
const block = asRecord(value);
|
|
if (block?.type !== 'toolCall'
|
|
|| typeof block.id !== 'string'
|
|
|| typeof block.name !== 'string') continue;
|
|
const tool: ConversationToolNode = {
|
|
kind: 'tool',
|
|
id: `tool:${block.id}`,
|
|
toolCallId: block.id,
|
|
toolName: block.name,
|
|
title: block.name,
|
|
inputText: inputTextOf(block.arguments),
|
|
status: 'declared',
|
|
output: [],
|
|
};
|
|
nodes.push(tool);
|
|
tools.set(block.id, tool);
|
|
}
|
|
}
|
|
return nodes;
|
|
}
|
|
|
|
function blockSignature(block: ConversationContentBlock): string {
|
|
if (block.kind === 'image') return `image:${block.mime}`;
|
|
return `${block.kind}:${block.text}`;
|
|
}
|
|
|
|
function messageSignature(node: ConversationMessageNode): string {
|
|
return `${node.role}|${node.blocks.map(blockSignature).join('|')}`;
|
|
}
|
|
|
|
function rebaseBlocks(
|
|
blocks: ConversationContentBlock[],
|
|
nodeId: string,
|
|
idPart: 'content' | 'output',
|
|
): ConversationContentBlock[] {
|
|
return blocks.map((block, index) => ({ ...block, id: `${nodeId}:${idPart}:${index}` }));
|
|
}
|
|
|
|
function reconcileLiveIds(
|
|
durableNodes: ConversationNode[],
|
|
liveNodes: ConversationNode[],
|
|
): ConversationNode[] {
|
|
const used = new Set<string>();
|
|
return durableNodes.map((node) => {
|
|
if (node.kind === 'message') {
|
|
const live = liveNodes.find((candidate) => candidate.kind === 'message'
|
|
&& !used.has(candidate.id)
|
|
&& (candidate.sourceEntryId === node.sourceEntryId
|
|
|| messageSignature(candidate) === messageSignature(node)));
|
|
if (!live || live.kind !== 'message') return node;
|
|
used.add(live.id);
|
|
return {
|
|
...node,
|
|
id: live.id,
|
|
blocks: rebaseBlocks(node.blocks, live.id, 'content'),
|
|
...(live.clientRequestId ? { clientRequestId: live.clientRequestId } : {}),
|
|
};
|
|
}
|
|
if (node.kind === 'tool') {
|
|
const live = liveNodes.find((candidate) => candidate.kind === 'tool'
|
|
&& !used.has(candidate.id)
|
|
&& candidate.toolCallId === node.toolCallId);
|
|
if (!live || live.kind !== 'tool') return node;
|
|
used.add(live.id);
|
|
return {
|
|
...node,
|
|
id: live.id,
|
|
output: rebaseBlocks(node.output, live.id, 'output'),
|
|
};
|
|
}
|
|
if (node.kind === 'compaction') {
|
|
const live = liveNodes.findLast(
|
|
(candidate) => candidate.kind === 'compaction' && !used.has(candidate.id),
|
|
);
|
|
if (!live || live.kind !== 'compaction') return node;
|
|
used.add(live.id);
|
|
return { ...node, id: live.id, runId: live.runId };
|
|
}
|
|
return node;
|
|
});
|
|
}
|
|
|
|
function projectedContext(input: PiSessionSnapshotInput, state: Record<string, unknown>) {
|
|
const stats = asRecord(input.stats);
|
|
const contextUsage = asRecord(stats?.contextUsage);
|
|
if (!contextUsage) {
|
|
return {
|
|
...structuredClone(input.snapshot.context),
|
|
compaction: state.isCompacting === true ? 'running' as const : 'idle' as const,
|
|
};
|
|
}
|
|
const tokens = contextUsage.tokens;
|
|
const recalculating = tokens === null;
|
|
return {
|
|
usedTokens: typeof tokens === 'number' ? tokens : 0,
|
|
contextWindow: typeof contextUsage.contextWindow === 'number' ? contextUsage.contextWindow : 0,
|
|
compaction: state.isCompacting === true ? 'running' as const : 'idle' as const,
|
|
...(recalculating ? { recalculating: true } : {}),
|
|
};
|
|
}
|
|
|
|
export async function projectPiSessionSnapshot(
|
|
input: PiSessionSnapshotInput,
|
|
): Promise<ConversationSnapshot> {
|
|
if (!Number.isSafeInteger(input.workerGeneration) || input.workerGeneration <= 0) unreadable();
|
|
const response = asRecord(input.entries);
|
|
const state = asRecord(input.state);
|
|
if (!response || !state) unreadable();
|
|
const path = retainedTail(activePath(response.entries, response.leafId));
|
|
const durableNodes = await projectEntries(path, input);
|
|
const nodes = reconcileLiveIds(durableNodes, input.snapshot.nodes);
|
|
return {
|
|
...structuredClone(input.snapshot),
|
|
nodes,
|
|
run: state.isStreaming === true
|
|
? { ...structuredClone(input.snapshot.run), status: 'running' }
|
|
: { status: 'idle' },
|
|
queue: { items: [] },
|
|
context: projectedContext(input, state),
|
|
pendingInteractions: [],
|
|
worker: { status: 'ready', generation: input.workerGeneration },
|
|
cursor: {
|
|
workerGeneration: input.workerGeneration,
|
|
seq: input.snapshot.cursor.workerGeneration === input.workerGeneration
|
|
? input.snapshot.cursor.seq
|
|
: 0,
|
|
...(typeof response.leafId === 'string' ? { leafEntryId: response.leafId } : {}),
|
|
},
|
|
};
|
|
}
|