@tduyng/pi-claude 0.0.0-stage → 0.0.2

package/src/session.ts

@@ -0,0 +1,237 @@
+import { randomBytes } from 'node:crypto'
+
+import type { AssistantMessage, Tool, ToolResultMessage } from '@earendil-works/pi-ai'
+
+import { jsonText, sha, stableStringify, type JsonObject } from './json.js'
+import type { McpResult } from './mcp.js'
+import { ClaudeProcess, type ClaudeSpec } from './protocol.js'
+import { fingerprint, toolNameFromMcp, toolResultMcp, type Transcript } from './transcript.js'
+
+const MAX_SESSIONS = 8
+
+export type SessionKey = {
+    model: string
+    systemPrompt: string
+    toolsHash: string
+    effort: string
+    cwd: string
+    contextWindow: number
+    maxTokens: number
+}
+
+export type ClaudeSession = {
+    token: string
+    key: SessionKey
+    spec: ClaudeSpec
+    tools: Tool[]
+    proc?: ClaudeProcess
+    chain: string
+    count: number
+    awaiting: Set<string>
+    parked: Map<string, (result: McpResult) => void>
+    ready: Map<string, McpResult>
+    servedByClaude: Set<string>
+    pendingTools: { id: string; name: string; args: JsonObject }[]
+    busy: boolean
+    dead: boolean
+    lastUsed: number
+    settling?: Promise<void>
+}
+
+export type ClaudeSessionStore = { sessions: Map<string, ClaudeSession> }
+
+export type Attached = { session: ClaudeSession; prefix: number }
+
+export function createSessionStore(): ClaudeSessionStore {
+    return { sessions: new Map() }
+}
+
+export function allSessions(store: ClaudeSessionStore): ClaudeSession[] {
+    return [...store.sessions.values()]
+}
+
+export function newSessionToken(): string {
+    return randomBytes(12).toString('hex')
+}
+
+export function sessionFor(
+    store: ClaudeSessionStore,
+    id: string,
+    token: string,
+    key: SessionKey,
+    spec: ClaudeSpec,
+    tools: Tool[],
+    transcript: Transcript
+): ClaudeSession {
+    const session: ClaudeSession = {
+        token,
+        key,
+        spec,
+        tools,
+        chain: transcript.chain[transcript.messages.length] ?? transcript.chain[0] ?? sha('prompt'),
+        count: transcript.messages.length,
+        awaiting: new Set(),
+        parked: new Map(),
+        ready: new Map(),
+        servedByClaude: new Set(),
+        pendingTools: [],
+        busy: true,
+        dead: false,
+        lastUsed: Date.now()
+    }
+    store.sessions.set(id, session)
+    evict(store)
+    return session
+}
+
+export function attachSession(
+    store: ClaudeSessionStore,
+    transcript: Transcript,
+    key: SessionKey
+): Attached | undefined {
+    for (const session of store.sessions.values()) {
+        if (session.dead || session.busy) continue
+        if (!sameKey(session.key, key)) continue
+        if (session.count > transcript.messages.length) continue
+        if (transcript.chain[session.count] !== session.chain) continue
+
+        const tail = transcript.messages.slice(session.count)
+        if (tail.some((message) => message.role === 'assistant')) continue
+
+        const resultIds = tail.flatMap((message) => (message.role === 'toolResult' ? [message.toolCallId] : []))
+        if (!resultIds.every((id) => session.awaiting.has(id))) continue
+
+        session.busy = true
+        session.lastUsed = Date.now()
+        return { session, prefix: session.count }
+    }
+    return undefined
+}
+
+function sameKey(left: SessionKey, right: SessionKey): boolean {
+    return (
+        left.model === right.model &&
+        left.systemPrompt === right.systemPrompt &&
+        left.toolsHash === right.toolsHash &&
+        left.effort === right.effort &&
+        left.cwd === right.cwd &&
+        left.contextWindow === right.contextWindow &&
+        left.maxTokens === right.maxTokens
+    )
+}
+
+export function updateChain(session: ClaudeSession, transcript: Transcript, output: AssistantMessage): void {
+    const base = transcript.chain[transcript.messages.length] ?? session.chain
+    session.chain = sha(`${base}\0${fingerprint(output)}`)
+    session.count = transcript.messages.length + 1
+    session.lastUsed = Date.now()
+}
+
+export function parkToolCall(
+    session: ClaudeSession,
+    name: string,
+    args: JsonObject,
+    meta: JsonObject
+): Promise<McpResult> {
+    const toolName = toolNameFromMcp(name)
+    if (!session.tools.some((tool) => tool.name === toolName)) {
+        return Promise.resolve({
+            content: [{ type: 'text', text: `Pi bridge: unknown tool ${toolName}` }],
+            isError: true
+        })
+    }
+    let id = jsonText(meta, 'claudecode/toolUseId')
+    if (id) {
+        const index = session.pendingTools.findIndex((call) => call.id === id)
+        if (index >= 0) session.pendingTools.splice(index, 1)
+    } else {
+        const index = session.pendingTools.findIndex(
+            (call) => call.name === toolName && stableStringify(call.args) === stableStringify(args)
+        )
+        const call = index >= 0 ? session.pendingTools.splice(index, 1)[0] : undefined
+        id = call?.id
+    }
+    if (!id) {
+        return Promise.resolve({
+            content: [{ type: 'text', text: 'Pi bridge: tool call without a tool_use id' }],
+            isError: true
+        })
+    }
+    const ready = session.ready.get(id)
+    if (ready) {
+        session.ready.delete(id)
+        return Promise.resolve(ready)
+    }
+    if (session.parked.has(id)) {
+        return Promise.resolve({ content: [{ type: 'text', text: 'Pi bridge: duplicate tool call' }], isError: true })
+    }
+    return new Promise((resolve) => {
+        session.parked.set(id, resolve)
+    })
+}
+
+export function deliverToolResult(session: ClaudeSession, message: ToolResultMessage): void {
+    const result = toolResultMcp(message)
+    const resolve = session.parked.get(message.toolCallId)
+    if (resolve) {
+        session.parked.delete(message.toolCallId)
+        resolve(result)
+        return
+    }
+    session.ready.set(message.toolCallId, result)
+}
+
+export function failMissingResults(session: ClaudeSession, seen: Set<string>): void {
+    for (const id of session.awaiting) {
+        if (seen.has(id)) continue
+        const resolve = session.parked.get(id)
+        const result: McpResult = {
+            content: [{ type: 'text', text: 'Pi did not return a result for this tool call.' }],
+            isError: true
+        }
+        if (resolve) {
+            session.parked.delete(id)
+            resolve(result)
+        } else {
+            session.ready.set(id, result)
+        }
+    }
+}
+
+export function failParked(session: ClaudeSession, reason: string): void {
+    for (const resolve of session.parked.values()) {
+        resolve({ content: [{ type: 'text', text: reason }], isError: true })
+    }
+    session.parked.clear()
+    session.ready.clear()
+    session.pendingTools.length = 0
+}
+
+export function dropSession(store: ClaudeSessionStore, session: ClaudeSession): void {
+    session.dead = true
+    for (const [id, candidate] of store.sessions) if (candidate === session) store.sessions.delete(id)
+    failParked(session, 'The Pi session closed this Claude Code process.')
+    void session.proc?.close()
+}
+
+export function dropSessions(store: ClaudeSessionStore): void {
+    for (const session of store.sessions.values()) dropSession(store, session)
+}
+
+export function sweep(store: ClaudeSessionStore, idleMs: number): void {
+    const now = Date.now()
+    for (const session of store.sessions.values()) {
+        if (!session.busy && now - session.lastUsed > idleMs) dropSession(store, session)
+    }
+}
+
+function evict(store: ClaudeSessionStore): void {
+    if (store.sessions.size <= MAX_SESSIONS) return
+    const idle = [...store.sessions.values()]
+        .filter((session) => !session.busy && session.awaiting.size === 0 && session.parked.size === 0)
+        .toSorted((left, right) => left.lastUsed - right.lastUsed)
+    while (store.sessions.size > MAX_SESSIONS && idle.length > 0) {
+        const session = idle.shift()
+        if (session) dropSession(store, session)
+    }
+}

package/src/transcript.ts

@@ -0,0 +1,213 @@
+import {
+    getCurrentTools,
+    getInitialSystemMessage,
+    getSystemMessageText,
+    renderSystemMessageUpdate,
+    toToolDeclaration,
+    type AssistantMessage,
+    type ImageContent,
+    type Message,
+    type SystemMessage,
+    type TextContent,
+    type Tool,
+    type ToolResultMessage,
+    type UserMessage
+} from '@earendil-works/pi-ai'
+
+import { jsonObject, sha, stableStringify, type JsonValue } from './json.js'
+import { mcpInputSchema, type McpResult, type McpToolDef } from './mcp.js'
+
+export type ClaudeContext = { messages: Message[] }
+export type TurnMessage = UserMessage | AssistantMessage | ToolResultMessage | SystemMessage
+
+export type Transcript = {
+    systemPrompt: string
+    tools: Tool[]
+    toolsHash: string
+    messages: TurnMessage[]
+    chain: string[]
+}
+
+export type CliTextBlock = { type: 'text'; text: string }
+export type CliImageBlock = { type: 'image'; source: { type: 'base64'; media_type: string; data: string } }
+export type CliBlock = CliTextBlock | CliImageBlock
+
+const MCP_PREFIX = 'mcp__pi__'
+
+export function readTranscript(context: ClaudeContext): Transcript {
+    const leading = getInitialSystemMessage(context.messages)
+    const systemPrompt = leading ? getSystemMessageText(leading) : ''
+    const tools = getCurrentTools(context.messages).map(toToolDeclaration)
+    const messages: TurnMessage[] = []
+    const dropped = new Set<string>()
+
+    for (const message of context.messages) {
+        if (message === leading) continue
+        if (message.role === 'assistant' && (message.stopReason === 'error' || message.stopReason === 'aborted')) {
+            for (const block of message.content) if (block.type === 'toolCall') dropped.add(block.id)
+            continue
+        }
+        if (message.role === 'toolResult' && dropped.has(message.toolCallId)) continue
+        messages.push(message)
+    }
+
+    const chain = [sha(`prompt\0${systemPrompt}`)]
+    for (const message of messages) chain.push(sha(`${chain[chain.length - 1]}\0${fingerprint(message)}`))
+    return { systemPrompt, tools, toolsHash: toolsHash(tools), messages, chain }
+}
+
+export function fingerprint(message: TurnMessage): string {
+    switch (message.role) {
+        case 'user':
+            return stableStringify(['u', normalContent(message.content)])
+        case 'toolResult':
+            return stableStringify(['r', message.toolCallId, normalContent(message.content), message.isError])
+        case 'system':
+            return stableStringify(['s', renderSystemMessageUpdate(message)])
+        case 'assistant': {
+            const blocks: JsonValue[] = []
+            for (const block of message.content) {
+                if (block.type === 'text') {
+                    if (block.text.length > 0) blocks.push(['t', block.text])
+                } else if (block.type === 'thinking') {
+                    blocks.push(['k', block.thinking, block.thinkingSignature ?? ''])
+                } else {
+                    const args: JsonValue = JSON.parse(JSON.stringify(block.arguments))
+                    blocks.push(['c', block.id, block.name, args])
+                }
+            }
+            return stableStringify(['a', blocks])
+        }
+    }
+}
+
+function normalContent(content: string | (TextContent | ImageContent)[]): JsonValue {
+    if (!Array.isArray(content)) return [['t', content]]
+    return content.map((block) => (block.type === 'text' ? ['t', block.text] : ['i', block.mimeType, sha(block.data)]))
+}
+
+export function toolsHash(tools: Tool[]): string {
+    return sha(
+        JSON.stringify(
+            tools.map((tool) => ({ name: tool.name, description: tool.description, parameters: tool.parameters }))
+        )
+    )
+}
+
+export function toolNameFromMcp(name: string): string {
+    return name.startsWith(MCP_PREFIX) ? name.slice(MCP_PREFIX.length) : name
+}
+
+export function mcpToolDefs(tools: Tool[]): McpToolDef[] {
+    return tools.map((tool) => {
+        const parameters = jsonObject(JSON.parse(JSON.stringify(tool.parameters)))
+        return {
+            name: tool.name,
+            description: tool.description,
+            inputSchema: mcpInputSchema(parameters)
+        }
+    })
+}
+
+export function toolResultMcp(message: ToolResultMessage): McpResult {
+    const content = message.content.map((block) =>
+        block.type === 'text'
+            ? { type: 'text' as const, text: block.text }
+            : { type: 'image' as const, data: block.data, mimeType: block.mimeType }
+    )
+    return {
+        content: content.length > 0 ? content : [{ type: 'text', text: '(no output)' }],
+        isError: message.isError
+    }
+}
+
+function imageBlock(block: ImageContent): CliBlock {
+    return { type: 'image', source: { type: 'base64', media_type: block.mimeType, data: block.data } }
+}
+
+function messageText(content: string | (TextContent | ImageContent)[]): string {
+    if (!Array.isArray(content)) return content
+    return content
+        .map((block) => (block.type === 'text' ? block.text : `[image: ${block.mimeType}]`))
+        .join('\n')
+        .trim()
+}
+
+export function userContent(messages: TurnMessage[]): CliBlock[] {
+    const content: CliBlock[] = []
+    for (const message of messages) {
+        if (message.role === 'user') {
+            if (!Array.isArray(message.content)) {
+                if (message.content.length > 0) content.push({ type: 'text', text: message.content })
+                continue
+            }
+            for (const block of message.content) {
+                if (block.type === 'text') {
+                    if (block.text.length > 0) content.push({ type: 'text', text: block.text })
+                } else {
+                    content.push(imageBlock(block))
+                }
+            }
+            continue
+        }
+        if (message.role === 'system') {
+            const text = renderSystemMessageUpdate(message).trim()
+            if (text) content.push({ type: 'text', text: `<system-reminder>\n${text}\n</system-reminder>` })
+        }
+    }
+    return content
+}
+
+export function replayContent(messages: TurnMessage[]): CliBlock[] {
+    if (messages.length === 1 && messages[0]?.role === 'user') return userContent(messages)
+
+    const parts: string[] = [
+        'This is the existing Pi conversation, replayed because the transport restarted. Treat it as prior turns, not as a new request. Continue after the final entry.'
+    ]
+    for (const message of messages) {
+        if (message.role === 'user') {
+            const text = messageText(message.content)
+            if (text) parts.push(`USER:\n${text}`)
+            continue
+        }
+        if (message.role === 'toolResult') {
+            const text = messageText(message.content)
+            const status = message.isError ? 'failed' : 'ok'
+            parts.push(`TOOL RESULT ${message.toolName} (${status}):\n${text}`)
+            continue
+        }
+        if (message.role === 'system') {
+            const text = renderSystemMessageUpdate(message).trim()
+            if (text) parts.push(`SYSTEM UPDATE:\n${text}`)
+            continue
+        }
+        const rendered = message.content
+            .map((block) => {
+                if (block.type === 'text') return block.text
+                if (block.type === 'thinking') return ''
+                return `CALLED ${block.name}(${JSON.stringify(block.arguments)})`
+            })
+            .filter((line) => line.length > 0)
+            .join('\n')
+        if (rendered) parts.push(`ASSISTANT:\n${rendered}`)
+    }
+
+    const content: CliBlock[] = [{ type: 'text', text: parts.join('\n\n') }]
+    for (const message of messages) {
+        if (message.role !== 'user' || !Array.isArray(message.content)) continue
+        for (const block of message.content) {
+            if (block.type !== 'image') continue
+            content.push({ type: 'text', text: 'Image from an earlier user message:' })
+            content.push(imageBlock(block))
+        }
+    }
+    return content
+}
+
+export function systemPreamble(cwd: string): string {
+    return [
+        'Your tools are MCP tools served by Pi. Call them by their declared names, which Claude Code shows as mcp__pi__<name>. No built-in tools are enabled.',
+        'Pi owns the conversation, approvals, file edits, shell, and compaction. Do not claim work you have not done with a Pi tool.',
+        `Working directory: ${cwd}`
+    ].join('\n')
+}