@tduyng/pi-cursor 0.0.0-stage → 0.0.2

package/src/transcript.ts

@@ -0,0 +1,210 @@
+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 type { SDKCustomToolContent, SDKCustomToolResult, SDKJsonValue, SDKUserMessage } from '@cursor/sdk'
+
+import { sha, stableStringify, type JsonValue } from './json.js'
+
+export type CursorContext = { messages: Message[] }
+export type TurnMessage = UserMessage | AssistantMessage | ToolResultMessage | SystemMessage
+export type CursorImage = { data: string; mimeType: string }
+export type CursorInput = { text: string; images: CursorImage[] }
+
+export type Transcript = {
+    systemPrompt: string
+    tools: Tool[]
+    toolsHash: string
+    messages: TurnMessage[]
+    chain: string[]
+}
+
+export function readTranscript(context: CursorContext): 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 toolSchema(tool: Tool): Record<string, SDKJsonValue> {
+    const schema: Record<string, SDKJsonValue> = JSON.parse(JSON.stringify(tool.parameters))
+    if (schema['type'] === undefined) schema['type'] = 'object'
+    if (schema['required'] === undefined) schema['required'] = []
+    return schema
+}
+
+export function toolResultCustom(message: ToolResultMessage): SDKCustomToolResult {
+    const content: SDKCustomToolContent[] = message.content.map((block) =>
+        block.type === 'text'
+            ? { type: 'text', text: block.text }
+            : { type: 'image', data: block.data, mimeType: block.mimeType }
+    )
+    return {
+        content: content.length > 0 ? content : [{ type: 'text', text: '(no output)' }],
+        isError: message.isError
+    }
+}
+
+function image(block: ImageContent): CursorImage {
+    return { data: block.data, mimeType: block.mimeType }
+}
+
+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 userInput(messages: TurnMessage[]): CursorInput {
+    const parts: string[] = []
+    const images: CursorImage[] = []
+    for (const message of messages) {
+        if (message.role === 'user') {
+            if (!Array.isArray(message.content)) {
+                if (message.content.length > 0) parts.push(message.content)
+                continue
+            }
+            for (const block of message.content) {
+                if (block.type === 'text') {
+                    if (block.text.length > 0) parts.push(block.text)
+                } else {
+                    images.push(image(block))
+                }
+            }
+            continue
+        }
+        if (message.role === 'system') {
+            const text = renderSystemMessageUpdate(message).trim()
+            if (text) parts.push(`<system-reminder>\n${text}\n</system-reminder>`)
+        }
+    }
+    return { text: parts.join('\n\n'), images }
+}
+
+export function replayInput(messages: TurnMessage[]): CursorInput {
+    if (messages.length === 1 && messages[0]?.role === 'user') return userInput(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 pending = messages.findLast((message) => message.role === 'user')
+    const images: CursorImage[] = []
+    if (pending && Array.isArray(pending.content)) {
+        for (const block of pending.content) {
+            if (block.type === 'image') images.push(image(block))
+        }
+    }
+    return { text: parts.join('\n\n'), images }
+}
+
+export function withPreamble(input: CursorInput, prompt: string): CursorInput {
+    if (prompt.trim().length === 0) return input
+    return { text: `<system>\n${prompt.trim()}\n</system>\n\n${input.text}`, images: input.images }
+}
+
+export function toSdkMessage(input: CursorInput): string | SDKUserMessage {
+    if (input.images.length === 0) return input.text
+    return { text: input.text, images: input.images }
+}
+
+export function systemPreamble(cwd: string): string {
+    return [
+        'Your tools are Pi tools, exposed to you as custom tools. Call them by their declared names. Cursor built-in tools are disabled, and no shell, read, edit, write, grep, glob, ls or task tool exists.',
+        'Pi owns the conversation, approvals, file edits, shell, and compaction. Every tool call runs in Pi. Do not claim work you have not done with a Pi tool.',
+        `Working directory: ${cwd}`
+    ].join('\n')
+}

package/src/turn.ts

@@ -0,0 +1,389 @@
+import {
+    createAssistantMessageEventStream,
+    type Api,
+    type AssistantMessage,
+    type AssistantMessageEventStream,
+    type Model,
+    type SimpleStreamOptions,
+    type ToolCall,
+    type ToolResultMessage,
+    type Usage
+} from '@earendil-works/pi-ai'
+
+import { sha } from './json.js'
+import { fastMode } from './state.js'
+import { dropRecord, readRecord, resumeTail, writeRecord } from './resume.js'
+import { selectionFor } from './models.js'
+import {
+    attachSession,
+    createSession,
+    createSessionStore,
+    deliverToolResult,
+    dropSession,
+    failMissingResults,
+    resumeSession,
+    startRun,
+    updateChain,
+    usageDelta,
+    type Attached,
+    type CursorSession,
+    type CursorEvent,
+    type CursorSessionStore,
+    type SessionKey
+} from './session.js'
+import { readTranscript, replayInput, systemPreamble, userInput, type CursorContext } from './transcript.js'
+
+export type CursorRuntime = { sessions: CursorSessionStore; fast: boolean }
+
+export function createRuntime(): CursorRuntime {
+    return { sessions: createSessionStore(), fast: fastMode() }
+}
+
+const TOOL_GRACE_MS = 250
+
+type TurnReason = 'stop' | 'length' | 'toolUse'
+type StepResult = { ok: true; reason: TurnReason } | { ok: false; error: string } | 'aborted'
+
+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 }
+    }
+}
+
+export function streamCursor(
+    runtime: CursorRuntime,
+    model: Model<Api>,
+    context: CursorContext,
+    options?: SimpleStreamOptions
+): AssistantMessageEventStream {
+    const stream = createAssistantMessageEventStream()
+    void runTurn(runtime, stream, model, context, options)
+    return stream
+}
+
+async function runTurn(
+    runtime: CursorRuntime,
+    stream: AssistantMessageEventStream,
+    model: Model<Api>,
+    context: CursorContext,
+    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: CursorSession | 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 cwd = process.cwd()
+        const systemPrompt = [transcript.systemPrompt.trim(), systemPreamble(cwd)]
+            .filter((part) => part.length > 0)
+            .join('\n\n')
+        const key: SessionKey = { systemPrompt, toolsHash: transcript.toolsHash, cwd }
+        const sessionId = options?.sessionId ?? sha(systemPrompt)
+        const selection = selectionFor(model.id, options?.reasoning, runtime.fast)
+        let attached: Attached | undefined = attachSession(runtime.sessions, transcript, key)
+        if (attached?.session.dead) {
+            dropSession(runtime.sessions, attached.session)
+            attached = undefined
+        }
+
+        stream.push({ type: 'start', partial: output })
+
+        if (attached) {
+            const active = attached.session
+            session = active
+            active.abortCleanup?.()
+            active.abortCleanup = undefined
+            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')
+
+            if (active.awaiting.size > 0) {
+                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()
+                const steering = userInput(rest)
+                if (steering.text.length > 0 && active.run?.steer) {
+                    const note =
+                        steering.images.length > 0
+                            ? `${steering.text}\n[${steering.images.length} image(s) could not be injected mid-turn]`
+                            : steering.text
+                    await active.run.steer(note)
+                }
+            } else {
+                const input = userInput(rest)
+                if (input.text.length === 0 && input.images.length === 0) {
+                    throw new Error('the transcript has no new user message')
+                }
+                await startRun(active, input, selection)
+            }
+        } 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 record = options?.sessionId ? readRecord(options.sessionId) : undefined
+            const tail = record ? resumeTail(record, transcript, key) : undefined
+            const resumed =
+                record && tail
+                    ? await resumeSession(
+                          runtime.sessions,
+                          sessionId,
+                          key,
+                          systemPrompt,
+                          transcript.tools,
+                          transcript,
+                          model.id,
+                          record
+                      )
+                    : undefined
+            if (resumed && tail) {
+                session = resumed
+                await startRun(resumed, userInput(tail), selection)
+            } else {
+                if (record && options?.sessionId) dropRecord(options.sessionId)
+                const created = createSession(
+                    runtime.sessions,
+                    sessionId,
+                    key,
+                    systemPrompt,
+                    transcript.tools,
+                    transcript
+                )
+                session = created
+                await startRun(created, replayInput(transcript.messages), selection)
+            }
+        }
+
+        const active = session
+        const step = await readStep(active, stream, output, options)
+        if (step === 'aborted') {
+            dropSession(runtime.sessions, active)
+            fail('aborted', 'Request aborted')
+            return
+        }
+        if (!step.ok) {
+            dropSession(runtime.sessions, active)
+            fail('error', step.error)
+            return
+        }
+
+        applyUsage(output, active)
+        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()
+        if (step.reason === 'toolUse') {
+            watchAbort(runtime.sessions, active, options?.signal)
+        } else {
+            active.abortCleanup?.()
+            active.abortCleanup = undefined
+            active.run = undefined
+            const agentId = active.agent?.agentId
+            if (options?.sessionId && agentId) {
+                writeRecord({
+                    sessionId: options.sessionId,
+                    agentId,
+                    cwd: key.cwd,
+                    toolsHash: key.toolsHash,
+                    count: active.count,
+                    chain: active.chain,
+                    at: Date.now()
+                })
+            }
+        }
+
+        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.busy = false
+            session.lastUsed = Date.now()
+        }
+    }
+}
+
+function applyUsage(output: AssistantMessage, session: CursorSession): void {
+    const delta = usageDelta(session.usage, session.reported)
+    if (!delta) return
+    session.reported = session.usage ? { ...session.usage } : undefined
+    output.usage.input = delta.inputTokens
+    output.usage.output = delta.outputTokens
+    output.usage.cacheRead = delta.cacheReadTokens
+    output.usage.cacheWrite = delta.cacheWriteTokens
+    output.usage.totalTokens =
+        delta.inputTokens + delta.outputTokens + delta.cacheReadTokens + delta.cacheWriteTokens
+    if (delta.reasoningTokens !== undefined) output.usage.reasoning = delta.reasoningTokens
+    output.usage.cost = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }
+}
+
+function watchAbort(store: CursorSessionStore, session: CursorSession, signal: AbortSignal | undefined): void {
+    if (!signal) return
+    const onAbort = () => dropSession(store, session)
+    if (signal.aborted) {
+        onAbort()
+        return
+    }
+    signal.addEventListener('abort', onAbort, { once: true })
+    session.abortCleanup = () => signal.removeEventListener('abort', onAbort)
+}
+
+async function readStep(
+    session: CursorSession,
+    stream: AssistantMessageEventStream,
+    output: AssistantMessage,
+    options?: SimpleStreamOptions
+): Promise<StepResult> {
+    const blocks = output.content
+    let textIndex: number | undefined
+    let thinkingIndex: number | undefined
+    const pending: ToolCall[] = []
+    let grace: ReturnType<typeof setTimeout> | undefined
+
+    const closeText = (): void => {
+        if (textIndex === undefined) return
+        const block = blocks[textIndex]
+        if (block?.type === 'text') {
+            stream.push({ type: 'text_end', contentIndex: textIndex, content: block.text, partial: output })
+        }
+        textIndex = undefined
+    }
+
+    const closeThinking = (): void => {
+        if (thinkingIndex === undefined) return
+        const block = blocks[thinkingIndex]
+        if (block?.type === 'thinking') {
+            stream.push({
+                type: 'thinking_end',
+                contentIndex: thinkingIndex,
+                content: block.thinking,
+                partial: output
+            })
+        }
+        thinkingIndex = undefined
+    }
+
+    const closeBlocks = (): void => {
+        closeText()
+        closeThinking()
+    }
+
+    try {
+        for (;;) {
+            let event: CursorEvent | undefined
+            try {
+                event = await session.queue.next(options?.signal)
+            } catch {
+                return 'aborted'
+            }
+            if (!event) return { ok: false, error: 'the Cursor run ended without a result' }
+
+            if (event.kind === 'text') {
+                closeThinking()
+                if (textIndex === undefined) {
+                    blocks.push({ type: 'text', text: '' })
+                    textIndex = blocks.length - 1
+                    stream.push({ type: 'text_start', contentIndex: textIndex, partial: output })
+                }
+                const block = blocks[textIndex]
+                if (block?.type !== 'text') continue
+                block.text += event.text
+                stream.push({ type: 'text_delta', contentIndex: textIndex, delta: event.text, partial: output })
+                continue
+            }
+
+            if (event.kind === 'thinking') {
+                closeText()
+                if (thinkingIndex === undefined) {
+                    blocks.push({ type: 'thinking', thinking: '', thinkingSignature: '' })
+                    thinkingIndex = blocks.length - 1
+                    stream.push({ type: 'thinking_start', contentIndex: thinkingIndex, partial: output })
+                }
+                const block = blocks[thinkingIndex]
+                if (block?.type !== 'thinking') continue
+                block.thinking += event.text
+                stream.push({
+                    type: 'thinking_delta',
+                    contentIndex: thinkingIndex,
+                    delta: event.text,
+                    partial: output
+                })
+                continue
+            }
+
+            if (event.kind === 'tool') {
+                closeBlocks()
+                blocks.push({ type: 'toolCall', id: event.id, name: event.name, arguments: event.args })
+                const index = blocks.length - 1
+                const toolCall = blocks[index]
+                if (toolCall?.type !== 'toolCall') continue
+                pending.push(toolCall)
+                stream.push({ type: 'toolcall_start', contentIndex: index, partial: output })
+                stream.push({ type: 'toolcall_end', contentIndex: index, toolCall, partial: output })
+                if (grace) clearTimeout(grace)
+                grace = setTimeout(() => session.queue.push({ kind: 'grace' }), TOOL_GRACE_MS)
+                continue
+            }
+
+            if (event.kind === 'grace') {
+                if (pending.length > 0) return { ok: true, reason: 'toolUse' }
+                continue
+            }
+
+            if (event.kind === 'error') {
+                closeBlocks()
+                return event.message === 'aborted' ? 'aborted' : { ok: false, error: event.message }
+            }
+
+            closeBlocks()
+            if (pending.length > 0) return { ok: true, reason: 'toolUse' }
+            return { ok: true, reason: 'stop' }
+        }
+    } finally {
+        if (grace) clearTimeout(grace)
+    }
+}