@tduyng/pi-claude 0.0.0-stage → 0.0.2

package/src/turn.ts

@@ -0,0 +1,657 @@
+import {
+    calculateCost,
+    createAssistantMessageEventStream,
+    type Api,
+    type AssistantMessage,
+    type AssistantMessageEventStream,
+    type Model,
+    type SimpleStreamOptions,
+    type ToolResultMessage,
+    type Usage
+} from '@earendil-works/pi-ai'
+import { Validator } from 'ata-validator'
+import { t } from 'ata-validator/t'
+
+
+import { jsonArray, jsonFlag, jsonObject, jsonStrings, jsonText, sha, type JsonObject, type JsonValue } from './json.js'
+import type { McpServerHandle } from './mcp.js'
+import { ClaudeProcess, claudeBin, type ClaudeSpec } from './protocol.js'
+import {
+    allSessions,
+    attachSession,
+    createSessionStore,
+    deliverToolResult,
+    dropSession,
+    failMissingResults,
+    newSessionToken,
+    parkToolCall,
+    sessionFor,
+    updateChain,
+    type ClaudeSession,
+    type ClaudeSessionStore,
+    type SessionKey
+} from './session.js'
+import {
+    mcpToolDefs,
+    readTranscript,
+    replayContent,
+    systemPreamble,
+    toolNameFromMcp,
+    userContent,
+    type ClaudeContext,
+    type Transcript
+} from './transcript.js'
+
+export type ClaudeRuntime = {
+    sessions: ClaudeSessionStore
+    ensureMcp(): Promise<McpServerHandle>
+}
+
+export function createRuntime(ensureMcp: () => Promise<McpServerHandle>): ClaudeRuntime {
+    return { sessions: createSessionStore(), ensureMcp }
+}
+
+const INTERRUPT_MS = 5000
+
+type TurnReason = 'stop' | 'length' | 'toolUse'
+type StepResult = { ok: true; reason: TurnReason } | { ok: false; error: string } | 'aborted'
+
+const UsageSchema = t.object(
+    {
+        input_tokens: t.optional(t.number()),
+        output_tokens: t.optional(t.number()),
+        cache_read_input_tokens: t.optional(t.number()),
+        cache_creation_input_tokens: t.optional(t.number()),
+        cache_creation: t.optional(
+            t.object({ ephemeral_1h_input_tokens: t.optional(t.number()) }, { additionalProperties: true })
+        )
+    },
+    { additionalProperties: true }
+)
+const isUsageSchema = new Validator(UsageSchema).isValidObject
+
+const TextBlockSchema = t.object({ type: t.literal('text') }, { additionalProperties: true })
+const ThinkingBlockSchema = t.object({ type: t.literal('thinking') }, { additionalProperties: true })
+const RedactedBlockSchema = t.object(
+    { type: t.literal('redacted_thinking'), data: t.string() },
+    { additionalProperties: true }
+)
+const ToolUseSchema = t.object(
+    { type: t.literal('tool_use'), id: t.string(), name: t.string() },
+    { additionalProperties: true }
+)
+const ContentBlockSchema = t.union([TextBlockSchema, ThinkingBlockSchema, RedactedBlockSchema, ToolUseSchema])
+
+const TextDeltaSchema = t.object(
+    { type: t.literal('text_delta'), text: t.string() },
+    { additionalProperties: true }
+)
+const ThinkingDeltaSchema = t.object(
+    { type: t.literal('thinking_delta'), thinking: t.string() },
+    { additionalProperties: true }
+)
+const SignatureDeltaSchema = t.object(
+    { type: t.literal('signature_delta'), signature: t.string() },
+    { additionalProperties: true }
+)
+const JsonDeltaSchema = t.object(
+    { type: t.literal('input_json_delta'), partial_json: t.string() },
+    { additionalProperties: true }
+)
+const DeltaSchema = t.union([TextDeltaSchema, ThinkingDeltaSchema, SignatureDeltaSchema, JsonDeltaSchema])
+
+const StreamEventSchema = t.union([
+    t.object(
+        {
+            type: t.literal('message_start'),
+            message: t.object(
+                {
+                    id: t.optional(t.string()),
+                    model: t.optional(t.string()),
+                    usage: t.optional(UsageSchema)
+                },
+                { additionalProperties: true }
+            )
+        },
+        { additionalProperties: true }
+    ),
+    t.object(
+        {
+            type: t.literal('content_block_start'),
+            index: t.number(),
+            content_block: ContentBlockSchema
+        },
+        { additionalProperties: true }
+    ),
+    t.object(
+        {
+            type: t.literal('content_block_delta'),
+            index: t.number(),
+            delta: DeltaSchema
+        },
+        { additionalProperties: true }
+    ),
+    t.object({ type: t.literal('content_block_stop'), index: t.number() }, { additionalProperties: true }),
+    t.object(
+        {
+            type: t.literal('message_delta'),
+            delta: t.optional(
+                t.object({ stop_reason: t.optional(t.string()) }, { additionalProperties: true })
+            ),
+            usage: t.optional(UsageSchema)
+        },
+        { additionalProperties: true }
+    ),
+    t.object({ type: t.literal('message_stop') }, { additionalProperties: true })
+])
+const isStreamEventSchema = new Validator(StreamEventSchema).isValidObject
+
+type Totals = { input: number; output: number; cacheRead: number; cacheWrite: number; cacheWrite1h: number }
+
+function emptyTotals(): Totals {
+    return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cacheWrite1h: 0 }
+}
+
+function emptyUsage(): Usage {
+    return {
+        input: 0,
+        output: 0,
+        cacheRead: 0,
+        cacheWrite: 0,
+        totalTokens: 0,
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }
+    }
+}
+
+function effortFor(model: Model<Api>, level: SimpleStreamOptions['reasoning']): string | undefined {
+    if (!level) return undefined
+    const mapped = model.thinkingLevelMap?.[level]
+    if (mapped === null) return undefined
+    if (mapped === 'low' || mapped === 'medium' || mapped === 'high' || mapped === 'xhigh' || mapped === 'max')
+        return mapped
+    if (level === 'minimal') return 'low'
+    return level
+}
+
+function mapStop(reason: string | undefined): TurnReason {
+    if (reason === 'tool_use') return 'toolUse'
+    if (reason === 'max_tokens' || reason === 'model_context_window_exceeded') return 'length'
+    return 'stop'
+}
+
+export function streamClaude(
+    runtime: ClaudeRuntime,
+    model: Model<Api>,
+    context: ClaudeContext,
+    options?: SimpleStreamOptions
+): AssistantMessageEventStream {
+    const stream = createAssistantMessageEventStream()
+    void runTurn(runtime, stream, model, context, options)
+    return stream
+}
+
+async function runTurn(
+    runtime: ClaudeRuntime,
+    stream: AssistantMessageEventStream,
+    model: Model<Api>,
+    context: ClaudeContext,
+    options?: SimpleStreamOptions
+): Promise<void> {
+    const output: AssistantMessage = {
+        role: 'assistant',
+        content: [],
+        api: model.api,
+        provider: model.provider,
+        model: model.id,
+        usage: emptyUsage(),
+        stopReason: 'stop',
+        timestamp: Date.now()
+    }
+    let session: ClaudeSession | undefined
+
+    const fail = (reason: 'error' | 'aborted', message: string): void => {
+        output.stopReason = reason
+        output.errorMessage = message
+        stream.push({ type: 'error', reason, error: output })
+        stream.end()
+    }
+
+    try {
+        if (options?.signal?.aborted) {
+            fail('aborted', 'Request aborted before dispatch')
+            return
+        }
+        if (options?.maxTokens === 1) {
+            output.stopReason = 'length'
+            stream.push({ type: 'start', partial: output })
+            stream.push({ type: 'done', reason: 'length', message: output })
+            stream.end()
+            return
+        }
+
+        const transcript = readTranscript(context)
+        const last = transcript.messages[transcript.messages.length - 1]
+        if (!last || last.role === 'assistant') {
+            throw new Error('the transcript must end with a user message or tool results')
+        }
+
+        const mcp = await runtime.ensureMcp()
+        const cwd = process.cwd()
+        const systemPrompt = [transcript.systemPrompt.trim(), systemPreamble(cwd)]
+            .filter((part) => part.length > 0)
+            .join('\n\n')
+        const effort = effortFor(model, options?.reasoning)
+        const maxTokens = options?.maxTokens ?? model.maxTokens
+        const key: SessionKey = {
+            model: model.id,
+            systemPrompt,
+            toolsHash: transcript.toolsHash,
+            effort: effort ?? '',
+            cwd,
+            contextWindow: model.contextWindow,
+            maxTokens
+        }
+
+        const sessionId = options?.sessionId ?? sha(key.systemPrompt)
+        await Promise.all(allSessions(runtime.sessions).map((candidate) => candidate.settling))
+        let attached = attachSession(runtime.sessions, transcript, key)
+        if (attached && attached.session.spec.mcpUrl !== mcp.urlFor(attached.session.token)) {
+            dropSession(runtime.sessions, attached.session)
+            attached = undefined
+        }
+        if (attached?.session.proc?.isDead()) {
+            dropSession(runtime.sessions, attached.session)
+            attached = undefined
+        }
+
+        if (attached) {
+            const active = attached.session
+            session = active
+            active.busy = true
+            if (!active.proc) throw new Error('Claude Code process is missing')
+            mcp.register(active.token, {
+                listTools: () => mcpToolDefs(active.tools),
+                callTool: (name, args, meta) => parkToolCall(active, name, args, meta)
+            })
+
+            const tail = transcript.messages.slice(attached.prefix)
+            const results = tail.filter((message): message is ToolResultMessage => message.role === 'toolResult')
+            const rest = tail.filter((message) => message.role !== 'toolResult')
+            const content = userContent(rest)
+
+            if (active.awaiting.size > 0) {
+                if (content.length > 0) active.proc.sendUser(content)
+                const seen = new Set<string>()
+                for (const result of results) {
+                    if (!active.awaiting.has(result.toolCallId)) continue
+                    deliverToolResult(active, result)
+                    seen.add(result.toolCallId)
+                }
+                failMissingResults(active, seen)
+                active.awaiting.clear()
+            } else {
+                if (content.length === 0) throw new Error('the transcript has no new user message')
+                active.pendingTools.length = 0
+                active.proc.sendUser(content)
+            }
+        } else {
+            const previous = runtime.sessions.sessions.get(sessionId)
+            if (previous?.busy) throw new Error('this Pi session already has a request in flight')
+            if (previous) dropSession(runtime.sessions, previous)
+
+            const live = await runtime.ensureMcp()
+            const token = newSessionToken()
+            const spec: ClaudeSpec = {
+                bin: claudeBin(),
+                model: model.id,
+                systemPrompt,
+                effort,
+                cwd,
+                mcpUrl: live.urlFor(token),
+                contextWindow: model.contextWindow,
+                maxTokens
+            }
+            const created = sessionFor(runtime.sessions, sessionId, token, key, spec, transcript.tools, transcript)
+            session = created
+            created.busy = true
+            live.register(created.token, {
+                listTools: () => mcpToolDefs(created.tools),
+                callTool: (name, args, meta) => parkToolCall(created, name, args, meta)
+            })
+            created.proc = new ClaudeProcess(spec)
+            await created.proc.start(replayContent(transcript.messages))
+        }
+
+        const active = session
+        const step = await readStep(active, stream, output, model, options)
+
+        if (step === 'aborted') {
+            active.settling = settleAbort(runtime.sessions, active, transcript)
+            fail('aborted', 'Request aborted')
+            return
+        }
+        if (!step.ok) {
+            dropSession(runtime.sessions, active)
+            fail('error', step.error)
+            return
+        }
+
+        updateChain(active, transcript, output)
+        const callIds = new Set<string>()
+        for (const block of output.content) if (block.type === 'toolCall') callIds.add(block.id)
+        active.awaiting = step.reason === 'toolUse' ? callIds : new Set()
+
+        stream.push({ type: 'done', reason: step.reason, message: output })
+        stream.end()
+    } catch (error) {
+        if (session) dropSession(runtime.sessions, session)
+        fail(options?.signal?.aborted ? 'aborted' : 'error', error instanceof Error ? error.message : String(error))
+    } finally {
+        if (session && !session.settling) {
+            session.busy = false
+            session.lastUsed = Date.now()
+        }
+    }
+}
+
+async function settleAbort(store: ClaudeSessionStore, session: ClaudeSession, transcript: Transcript): Promise<void> {
+    const proc = session.proc
+    proc?.interrupt()
+    const deadline = AbortSignal.timeout(INTERRUPT_MS)
+    for (;;) {
+        const event = await proc?.next(deadline).catch(() => undefined)
+        if (!event) {
+            dropSession(store, session)
+            break
+        }
+        if (jsonText(event, 'type') !== 'result') continue
+        session.chain = transcript.chain[transcript.messages.length] ?? session.chain
+        session.count = transcript.messages.length
+        session.awaiting = new Set()
+        session.pendingTools.length = 0
+        break
+    }
+    session.settling = undefined
+    session.busy = false
+    session.lastUsed = Date.now()
+}
+
+async function readStep(
+    session: ClaudeSession,
+    stream: AssistantMessageEventStream,
+    output: AssistantMessage,
+    model: Model<Api>,
+    options?: SimpleStreamOptions
+): Promise<StepResult> {
+    const proc = session.proc
+    if (!proc) return { ok: false, error: 'Claude Code process is missing' }
+
+    const blocks = output.content
+    const positions = new Map<number, number>()
+    const partialJson = new Map<number, string>()
+    const finished = emptyTotals()
+    let current = emptyTotals()
+    let responseStart = 0
+    let inResponse = false
+    let started = false
+    let stopReason: string | undefined
+    let turnStop: TurnReason | undefined
+    let lastTotal = 0
+
+    const start = (): void => {
+        if (started) return
+        started = true
+        stream.push({ type: 'start', partial: output })
+    }
+
+    const setUsage = (raw: JsonObject | undefined): void => {
+        if (!raw || !isUsageSchema(raw)) return
+        const usage = raw
+        if (usage.input_tokens !== undefined) current.input = usage.input_tokens
+        if (usage.output_tokens !== undefined) current.output = usage.output_tokens
+        if (usage.cache_read_input_tokens !== undefined) current.cacheRead = usage.cache_read_input_tokens
+        if (usage.cache_creation_input_tokens !== undefined) current.cacheWrite = usage.cache_creation_input_tokens
+        const oneHour = usage.cache_creation?.ephemeral_1h_input_tokens
+        if (oneHour !== undefined && oneHour > 0) current.cacheWrite1h = oneHour
+        output.usage.input = finished.input + current.input
+        output.usage.output = finished.output + current.output
+        output.usage.cacheRead = finished.cacheRead + current.cacheRead
+        output.usage.cacheWrite = finished.cacheWrite + current.cacheWrite
+        const oneHourTotal = finished.cacheWrite1h + current.cacheWrite1h
+        if (oneHourTotal > 0) output.usage.cacheWrite1h = oneHourTotal
+        output.usage.totalTokens = current.input + current.cacheRead + current.cacheWrite + current.output
+        output.usage.cost = calculateCost(model, output.usage)
+    }
+
+    const foldUsage = (): void => {
+        lastTotal = current.input + current.cacheRead + current.cacheWrite + current.output
+        finished.input += current.input
+        finished.output += current.output
+        finished.cacheRead += current.cacheRead
+        finished.cacheWrite += current.cacheWrite
+        finished.cacheWrite1h += current.cacheWrite1h
+        current = emptyTotals()
+        output.usage.totalTokens = lastTotal
+        output.usage.cost = calculateCost(model, output.usage)
+    }
+
+    for (;;) {
+        let event: JsonObject | undefined
+        try {
+            event = await proc.next(options?.signal)
+        } catch {
+            return 'aborted'
+        }
+        if (!event)
+            return { ok: false, error: `Claude Code exited${proc.stderrTail() ? `\n${proc.stderrTail()}` : ''}` }
+        session.lastUsed = Date.now()
+
+        const type = jsonText(event, 'type')
+        if (type === 'stream_event') {
+            const parent = event['parent_tool_use_id']
+            if (parent !== undefined && parent !== null) continue
+            const raw = jsonObject(event['event'])
+            if (!raw || !isStreamEventSchema(raw)) continue
+            const streamEvent = raw
+            await options?.onProviderStreamEvent?.(streamEvent, model)
+
+            if (streamEvent.type === 'message_start') {
+                if (inResponse) {
+                    blocks.length = responseStart
+                    positions.clear()
+                    partialJson.clear()
+                }
+                responseStart = blocks.length
+                inResponse = true
+                current = emptyTotals()
+                stopReason = undefined
+                if (streamEvent.message.id) output.responseId = streamEvent.message.id
+                if (streamEvent.message.model) output.responseModel = streamEvent.message.model
+                setUsage(streamEvent.message.usage)
+                start()
+                continue
+            }
+
+            if (streamEvent.type === 'content_block_start') {
+                start()
+                const block = streamEvent.content_block
+                const index = blocks.length
+                positions.set(streamEvent.index, index)
+                if (block.type === 'text') {
+                    blocks.push({ type: 'text', text: '' })
+                    stream.push({ type: 'text_start', contentIndex: index, partial: output })
+                } else if (block.type === 'thinking') {
+                    blocks.push({ type: 'thinking', thinking: '', thinkingSignature: '' })
+                    stream.push({ type: 'thinking_start', contentIndex: index, partial: output })
+                } else if (block.type === 'redacted_thinking') {
+                    blocks.push({
+                        type: 'thinking',
+                        thinking: '[Reasoning redacted]',
+                        thinkingSignature: block.data,
+                        redacted: true
+                    })
+                    stream.push({ type: 'thinking_start', contentIndex: index, partial: output })
+                } else {
+                    blocks.push({
+                        type: 'toolCall',
+                        id: block.id,
+                        name: toolNameFromMcp(block.name),
+                        arguments: {}
+                    })
+                    partialJson.set(streamEvent.index, '')
+                    stream.push({ type: 'toolcall_start', contentIndex: index, partial: output })
+                }
+                continue
+            }
+
+            if (streamEvent.type === 'content_block_delta') {
+                const index = positions.get(streamEvent.index)
+                const block = index === undefined ? undefined : blocks[index]
+                if (index === undefined || !block) continue
+                const delta = streamEvent.delta
+                if (delta.type === 'text_delta' && block.type === 'text') {
+                    block.text += delta.text
+                    stream.push({ type: 'text_delta', contentIndex: index, delta: delta.text, partial: output })
+                } else if (delta.type === 'thinking_delta' && block.type === 'thinking') {
+                    block.thinking += delta.thinking
+                    stream.push({
+                        type: 'thinking_delta',
+                        contentIndex: index,
+                        delta: delta.thinking,
+                        partial: output
+                    })
+                } else if (delta.type === 'signature_delta' && block.type === 'thinking') {
+                    block.thinkingSignature = (block.thinkingSignature ?? '') + delta.signature
+                } else if (delta.type === 'input_json_delta' && block.type === 'toolCall') {
+                    partialJson.set(streamEvent.index, (partialJson.get(streamEvent.index) ?? '') + delta.partial_json)
+                    stream.push({
+                        type: 'toolcall_delta',
+                        contentIndex: index,
+                        delta: delta.partial_json,
+                        partial: output
+                    })
+                }
+                continue
+            }
+
+            if (streamEvent.type === 'content_block_stop') {
+                const index = positions.get(streamEvent.index)
+                const block = index === undefined ? undefined : blocks[index]
+                if (index === undefined || !block) continue
+                positions.delete(streamEvent.index)
+                if (block.type === 'text') {
+                    stream.push({ type: 'text_end', contentIndex: index, content: block.text, partial: output })
+                } else if (block.type === 'thinking') {
+                    stream.push({
+                        type: 'thinking_end',
+                        contentIndex: index,
+                        content: block.thinking,
+                        partial: output
+                    })
+                } else {
+                    const rawJson = partialJson.get(streamEvent.index) ?? ''
+                    partialJson.delete(streamEvent.index)
+                    const parsed = rawJson.trim().length > 0 ? safeObject(rawJson) : undefined
+                    if (parsed) block.arguments = parsed
+                    else if (rawJson.trim().length > 0) block.argumentsError = `Invalid JSON arguments: ${rawJson}`
+                    const args: JsonObject = JSON.parse(JSON.stringify(block.arguments))
+                    session.pendingTools.push({ id: block.id, name: block.name, args })
+                    stream.push({ type: 'toolcall_end', contentIndex: index, toolCall: block, partial: output })
+                }
+                continue
+            }
+
+            if (streamEvent.type === 'message_delta') {
+                const reason = streamEvent.delta?.stop_reason
+                if (reason) stopReason = reason
+                setUsage(streamEvent.usage)
+                continue
+            }
+
+            if (streamEvent.type === 'message_stop') {
+                foldUsage()
+                inResponse = false
+                turnStop = mapStop(stopReason)
+                output.stopReason = turnStop
+                if (turnStop === 'toolUse') return { ok: true, reason: turnStop }
+            }
+            continue
+        }
+
+        if (type === 'assistant') {
+            const error = jsonText(event, 'error')
+            if (!error) continue
+            const message = jsonObject(event['message'])
+            const content = message ? textFromContent(message['content']) : ''
+            return { ok: false, error: content || error }
+        }
+
+        if (type === 'user') {
+            markServed(session, event)
+            continue
+        }
+
+        if (type === 'rate_limit_event' || jsonText(event, 'subtype') === 'rate_limit_event') {
+            if (jsonFlag(event, 'isUsingOverage')) {
+                return {
+                    ok: false,
+                    error: 'Claude Code charged this turn to usage credits instead of the plan limits'
+                }
+            }
+            continue
+        }
+
+        if (type === 'result') {
+            const subtype = jsonText(event, 'subtype') ?? 'success'
+            if (jsonFlag(event, 'is_error') || subtype !== 'success') {
+                const errors = jsonStrings(event['errors']) ?? []
+                const text = errors.join('\n') || jsonText(event, 'result') || `Claude Code ${subtype}`
+                return { ok: false, error: text }
+            }
+            start()
+            output.stopReason = turnStop ?? 'stop'
+            if (output.usage.totalTokens === 0) {
+                output.usage.totalTokens = finished.input + finished.cacheRead + finished.cacheWrite + finished.output
+            }
+            output.usage.cost = calculateCost(model, output.usage)
+            return { ok: true, reason: turnStop ?? 'stop' }
+        }
+    }
+}
+
+function safeObject(raw: string): JsonObject | undefined {
+    try {
+        const parsed = JSON.parse(raw)
+        return jsonObject(parsed)
+    } catch {
+        return undefined
+    }
+}
+
+function textFromContent(value: JsonValue | undefined): string {
+    const blocks = jsonArray(value)
+    if (!blocks) return ''
+    const parts: string[] = []
+    for (const item of blocks) {
+        const block = jsonObject(item)
+        if (!block) continue
+        const text = jsonText(block, 'text')
+        if (text) parts.push(text)
+    }
+    return parts.join('')
+}
+
+function markServed(session: ClaudeSession, event: JsonObject): void {
+    const message = jsonObject(event['message'])
+    if (!message) return
+    const content = jsonArray(message['content'])
+    if (!content) return
+    for (const item of content) {
+        const block = jsonObject(item)
+        if (!block || jsonText(block, 'type') !== 'tool_result') continue
+        const id = jsonText(block, 'tool_use_id')
+        if (!id) continue
+        if (session.parked.has(id) || session.ready.has(id) || session.awaiting.has(id)) continue
+        session.servedByClaude.add(id)
+    }
+}