@tduyng/pi-cursor 0.0.0-stage → 0.0.1

package/src/turn.ts

@@ -0,0 +1,336 @@
+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 { selectionFor } from './models.js'
+import {
+    attachSession,
+    createSession,
+    createSessionStore,
+    deliverToolResult,
+    dropSession,
+    failMissingResults,
+    startRun,
+    updateChain,
+    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 }
+
+export function createRuntime(): CursorRuntime {
+    return { sessions: createSessionStore() }
+}
+
+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)
+        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()
+            } 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 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
+        }
+
+        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 usage = session.usage
+    if (!usage) return
+    output.usage.input = usage.inputTokens
+    output.usage.output = usage.outputTokens
+    output.usage.cacheRead = usage.cacheReadTokens
+    output.usage.cacheWrite = usage.cacheWriteTokens
+    output.usage.totalTokens =
+        usage.totalTokens || usage.inputTokens + usage.outputTokens + usage.cacheReadTokens + usage.cacheWriteTokens
+    if (usage.reasoningTokens !== undefined) output.usage.reasoning = usage.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)
+    }
+}