feat: project Pi conversation events
This commit is contained in:
368
electron/coding-runtime/pi/session-projector.ts
Normal file
368
electron/coding-runtime/pi/session-projector.ts
Normal file
@@ -0,0 +1,368 @@
|
||||
import type {
|
||||
CodingRuntimeErrorCode,
|
||||
ConversationContentBlock,
|
||||
ConversationMessageNode,
|
||||
ConversationNode,
|
||||
ConversationSnapshot,
|
||||
ConversationToolNode,
|
||||
PublicUsage,
|
||||
} from '../contracts';
|
||||
import type {
|
||||
PiEventProjectorOptions,
|
||||
PiImageProjectionInput,
|
||||
PiProjectedAttachment,
|
||||
} from './event-projector';
|
||||
|
||||
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>();
|
||||
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';
|
||||
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: 0,
|
||||
...(typeof response.leafId === 'string' ? { leafEntryId: response.leafId } : {}),
|
||||
},
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user