@tduyng/pi-cursor 0.0.0-stage → 0.0.1

package/src/session.ts

@@ -0,0 +1,352 @@
+import type { AgentOptions, Run, SDKAgent, SDKCustomTool, SDKCustomToolResult, TokenUsage } from '@cursor/sdk'
+import type { AssistantMessage, Tool, ToolResultMessage } from '@earendil-works/pi-ai'
+
+import { sha, type JsonObject } from './json.js'
+import { selectionFor, type ModelSelection } from './models.js'
+import { cursorApiKey, cursorSdk } from './sdk.js'
+import {
+    fingerprint,
+    toSdkMessage,
+    withPreamble,
+    toolResultCustom,
+    toolSchema,
+    type CursorInput,
+    type Transcript
+} from './transcript.js'
+
+const MAX_SESSIONS = 6
+
+export type SessionKey = {
+    systemPrompt: string
+    toolsHash: string
+    cwd: string
+}
+
+export type CursorEvent =
+    | { kind: 'text'; text: string }
+    | { kind: 'thinking'; text: string }
+    | { kind: 'tool'; id: string; name: string; args: JsonObject }
+    | { kind: 'grace' }
+    | { kind: 'done' }
+    | { kind: 'error'; message: string }
+
+export type CursorSession = {
+    id: string
+    key: SessionKey
+    systemPrompt: string
+    tools: Tool[]
+    agent?: SDKAgent
+    run?: Run
+    queue: CursorEventQueue
+    parked: Map<string, (result: SDKCustomToolResult) => void>
+    ready: Map<string, SDKCustomToolResult>
+    awaiting: Set<string>
+    chain: string
+    count: number
+    usage?: TokenUsage
+    toolCounter: number
+    promptFallback?: boolean
+    promptSent?: boolean
+    busy: boolean
+    dead: boolean
+    lastUsed: number
+    abortCleanup?: () => void
+}
+
+export type CursorSessionStore = { sessions: Map<string, CursorSession> }
+export type Attached = { session: CursorSession; prefix: number }
+
+export class CursorEventQueue {
+    private readonly items: CursorEvent[] = []
+    private readonly waiters: ((event: CursorEvent | undefined) => void)[] = []
+    private closed = false
+
+    push(event: CursorEvent): void {
+        const waiter = this.waiters.shift()
+        if (waiter) waiter(event)
+        else this.items.push(event)
+    }
+
+    close(): void {
+        if (this.closed) return
+        this.closed = true
+        for (const waiter of this.waiters.splice(0)) waiter(undefined)
+    }
+
+    next(signal?: AbortSignal): Promise<CursorEvent | undefined> {
+        const queued = this.items.shift()
+        if (queued) return Promise.resolve(queued)
+        if (this.closed) return Promise.resolve(undefined)
+        return new Promise((resolve, reject) => {
+            let settled = false
+            const waiter = (event: CursorEvent | undefined) => {
+                if (settled) return
+                settled = true
+                signal?.removeEventListener('abort', onAbort)
+                resolve(event)
+            }
+            const onAbort = () => {
+                if (settled) return
+                settled = true
+                this.waiters.splice(this.waiters.indexOf(waiter), 1)
+                reject(new Error('aborted'))
+            }
+            this.waiters.push(waiter)
+            if (signal?.aborted) {
+                onAbort()
+                return
+            }
+            signal?.addEventListener('abort', onAbort, { once: true })
+        })
+    }
+}
+
+export function createSessionStore(): CursorSessionStore {
+    return { sessions: new Map() }
+}
+
+export function allSessions(store: CursorSessionStore): CursorSession[] {
+    return [...store.sessions.values()]
+}
+
+export function createSession(
+    store: CursorSessionStore,
+    id: string,
+    key: SessionKey,
+    systemPrompt: string,
+    tools: Tool[],
+    transcript: Transcript
+): CursorSession {
+    const session: CursorSession = {
+        id,
+        key,
+        systemPrompt,
+        tools,
+        queue: new CursorEventQueue(),
+        parked: new Map(),
+        ready: new Map(),
+        awaiting: new Set(),
+        chain: transcript.chain[transcript.messages.length] ?? transcript.chain[0] ?? sha('prompt'),
+        count: transcript.messages.length,
+        toolCounter: 0,
+        busy: true,
+        dead: false,
+        lastUsed: Date.now()
+    }
+    store.sessions.set(id, session)
+    evict(store)
+    return session
+}
+
+export function attachSession(
+    store: CursorSessionStore,
+    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 results = tail.filter((message) => message.role === 'toolResult')
+        const resultIds = new Set(results.map((message) => message.toolCallId))
+        if (session.awaiting.size > 0) {
+            if (![...session.awaiting].every((id) => resultIds.has(id))) continue
+        } else if (resultIds.size > 0) {
+            continue
+        }
+
+        session.busy = true
+        session.lastUsed = Date.now()
+        return { session, prefix: session.count }
+    }
+    return undefined
+}
+
+function sameKey(left: SessionKey, right: SessionKey): boolean {
+    return left.systemPrompt === right.systemPrompt && left.toolsHash === right.toolsHash && left.cwd === right.cwd
+}
+
+export function updateChain(session: CursorSession, 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 async function ensureAgent(session: CursorSession, model: string): Promise<SDKAgent> {
+    if (session.agent) return session.agent
+    const { Agent } = await cursorSdk()
+    const create = (systemPrompt: string | undefined) => {
+        const options: AgentOptions = {
+            apiKey: cursorApiKey(),
+            model: selectionFor(model, undefined),
+            tools: ['mcp'],
+            local: {
+                cwd: session.key.cwd,
+                settingSources: [],
+                customTools: customTools(session)
+            }
+        }
+        if (systemPrompt !== undefined) options.systemPrompt = systemPrompt
+        return Agent.create(options)
+    }
+    try {
+        session.agent = await create(session.promptFallback ? undefined : session.systemPrompt)
+        return session.agent
+    } catch (error) {
+        const message = error instanceof Error ? error.message : String(error)
+        if (!session.promptFallback && /system.?prompt/i.test(message)) {
+            session.promptFallback = true
+            session.agent = await create(undefined)
+            return session.agent
+        }
+        if (/auth|api.?key|login|credential/i.test(message)) {
+            throw new Error(`Cursor is not logged in. Run /cursor login, then retry.\n${message}`, { cause: error })
+        }
+        throw new Error(`Cursor agent could not start: ${message}`, { cause: error })
+    }
+}
+
+function customTools(session: CursorSession) {
+    const tools: Record<string, SDKCustomTool> = {}
+    for (const tool of session.tools) {
+        tools[tool.name] = {
+            description: tool.description,
+            inputSchema: toolSchema(tool),
+            execute: (args, context) => {
+                const id = context.toolCallId ?? `pi_${++session.toolCounter}`
+                const ready = session.ready.get(id)
+                if (ready) {
+                    session.ready.delete(id)
+                    return ready
+                }
+                const value: JsonObject = JSON.parse(JSON.stringify(args))
+                return new Promise<SDKCustomToolResult>((resolve) => {
+                    session.parked.set(id, resolve)
+                    session.queue.push({ kind: 'tool', id, name: tool.name, args: value })
+                })
+            }
+        }
+    }
+    return tools
+}
+
+export async function startRun(session: CursorSession, input: CursorInput, selection: ModelSelection): Promise<void> {
+    const agent = await ensureAgent(session, selection.id)
+    const payload = session.promptFallback && !session.promptSent ? withPreamble(input, session.systemPrompt) : input
+    if (session.promptFallback) session.promptSent = true
+    session.usage = undefined
+    let run: Run
+    try {
+        run = await agent.send(toSdkMessage(payload), {
+            model: selection,
+            onDelta: ({ update }) => {
+                if (update.type === 'text-delta') session.queue.push({ kind: 'text', text: update.text })
+                else if (update.type === 'thinking-delta') session.queue.push({ kind: 'thinking', text: update.text })
+                else if (update.type === 'turn-ended' && update.usage) {
+                    const usage = update.usage
+                    session.usage = {
+                        ...usage,
+                        totalTokens:
+                            usage.inputTokens + usage.outputTokens + usage.cacheReadTokens + usage.cacheWriteTokens
+                    }
+                }
+            }
+        })
+    } catch (error) {
+        const message = error instanceof Error ? error.message : String(error)
+        if (!session.promptFallback && /system.?prompt/i.test(message)) {
+            session.promptFallback = true
+            session.agent?.close()
+            session.agent = undefined
+            await startRun(session, input, selection)
+            return
+        }
+        throw error
+    }
+    session.run = run
+    void run
+        .wait()
+        .then((result) => {
+            if (!session.usage && result.usage) session.usage = result.usage
+            session.queue.push({ kind: 'done' })
+        })
+        .catch((error: Error) => {
+            session.queue.push({ kind: 'error', message: error.message || String(error) })
+        })
+}
+
+export function deliverToolResult(session: CursorSession, message: ToolResultMessage): void {
+    const result = toolResultCustom(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: CursorSession, seen: Set<string>): void {
+    for (const id of session.awaiting) {
+        if (seen.has(id)) continue
+        const result: SDKCustomToolResult = {
+            content: [{ type: 'text', text: 'Pi did not return a result for this tool call.' }],
+            isError: true
+        }
+        const resolve = session.parked.get(id)
+        if (resolve) {
+            session.parked.delete(id)
+            resolve(result)
+        } else {
+            session.ready.set(id, result)
+        }
+    }
+}
+
+export function failParked(session: CursorSession, reason: string): void {
+    for (const resolve of session.parked.values()) {
+        resolve({ content: [{ type: 'text', text: reason }], isError: true })
+    }
+    session.parked.clear()
+    session.ready.clear()
+}
+
+export function dropSession(store: CursorSessionStore, session: CursorSession): void {
+    session.dead = true
+    for (const [id, candidate] of store.sessions) if (candidate === session) store.sessions.delete(id)
+    session.abortCleanup?.()
+    session.abortCleanup = undefined
+    failParked(session, 'Pi closed this Cursor session.')
+    session.queue.close()
+    void session.run?.cancel().catch(() => undefined)
+    session.agent?.close()
+}
+
+export function dropSessions(store: CursorSessionStore): void {
+    for (const session of store.sessions.values()) dropSession(store, session)
+}
+
+export function sweep(store: CursorSessionStore, idleMs: number): void {
+    const now = Date.now()
+    for (const session of store.sessions.values()) {
+        if (session.busy || session.awaiting.size > 0) continue
+        if (now - session.lastUsed > idleMs) dropSession(store, session)
+    }
+}
+
+function evict(store: CursorSessionStore): 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,212 @@
+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 images: CursorImage[] = []
+    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
+            images.push(image(block))
+        }
+    }
+    if (images.length > 0) parts.push(`The replay includes ${images.length} image(s) from earlier user messages.`)
+    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')
+}