@tduyng/pi-cursor 0.0.0-stage → 0.0.2

package/src/session.ts

@@ -0,0 +1,461 @@
+import type {
+    AgentOptions,
+    Run,
+    RunResult,
+    SDKAgent,
+    SDKCustomTool,
+    SDKCustomToolResult,
+    TokenUsage
+} from '@cursor/sdk'
+import type { AssistantMessage, Tool, ToolResultMessage } from '@earendil-works/pi-ai'
+import type { AgentRecord } from './resume.js'
+
+import { sha, type JsonObject } from './json.js'
+import { noSystemPrompt, patchState } from './state.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
+    reported?: 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,
+        promptFallback: noSystemPrompt(),
+        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()
+}
+
+function agentOptions(session: CursorSession, model: string, systemPrompt: string | undefined): AgentOptions {
+    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 options
+}
+
+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) => Agent.create(agentOptions(session, model, systemPrompt))
+    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)
+            patchState({ noSystemPrompt: true })
+            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 })
+    }
+}
+
+export async function resumeSession(
+    store: CursorSessionStore,
+    id: string,
+    key: SessionKey,
+    systemPrompt: string,
+    tools: Tool[],
+    transcript: Transcript,
+    model: string,
+    record: AgentRecord
+): Promise<CursorSession | undefined> {
+    const session = createSession(store, id, key, systemPrompt, tools, transcript)
+    session.promptFallback = noSystemPrompt()
+    session.promptSent = true
+    session.count = record.count
+    session.chain = record.chain
+    try {
+        const { Agent } = await cursorSdk()
+        const prompt = session.promptFallback ? undefined : systemPrompt
+        session.agent = await Agent.resume(record.agentId, agentOptions(session, model, prompt))
+        return session
+    } catch {
+        dropSession(store, session)
+        return undefined
+    }
+}
+
+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
+    session.reported = 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) {
+                    session.usage = addUsage(session.usage, normalizeUsage(update.usage))
+                }
+            }
+        })
+    } catch (error) {
+        const message = error instanceof Error ? error.message : String(error)
+        if (!session.promptFallback && /system.?prompt/i.test(message)) {
+            await restartWithoutPrompt(session, input, selection)
+            return
+        }
+        throw error
+    }
+    session.run = run
+    void run
+        .wait()
+        .then(async (result) => {
+            if (result.usage) session.usage = largerUsage(session.usage, normalizeUsage(result.usage))
+            const outcome = outcomeOf(result, session.promptFallback === true)
+            if (outcome.kind === 'retry') {
+                await restartWithoutPrompt(session, input, selection)
+                return
+            }
+            if (outcome.kind === 'error') {
+                session.queue.push({ kind: 'error', message: outcome.message })
+                return
+            }
+            session.queue.push({ kind: 'done' })
+        })
+        .catch((error: Error) => {
+            session.queue.push({ kind: 'error', message: error.message || String(error) })
+        })
+}
+
+type ReportedUsage = Omit<TokenUsage, 'totalTokens'>
+
+function normalizeUsage(usage: ReportedUsage): TokenUsage {
+    return {
+        ...usage,
+        totalTokens: usage.inputTokens + usage.outputTokens + usage.cacheReadTokens + usage.cacheWriteTokens
+    }
+}
+
+export function addUsage(current: TokenUsage | undefined, next: TokenUsage): TokenUsage {
+    if (!current) return { ...next }
+    const reasoning = current.reasoningTokens
+    const extra = next.reasoningTokens
+    const total: TokenUsage = {
+        inputTokens: current.inputTokens + next.inputTokens,
+        outputTokens: current.outputTokens + next.outputTokens,
+        cacheReadTokens: current.cacheReadTokens + next.cacheReadTokens,
+        cacheWriteTokens: current.cacheWriteTokens + next.cacheWriteTokens,
+        totalTokens: current.totalTokens + next.totalTokens
+    }
+    if (reasoning !== undefined || extra !== undefined) total.reasoningTokens = (reasoning ?? 0) + (extra ?? 0)
+    return total
+}
+
+export function largerUsage(current: TokenUsage | undefined, next: TokenUsage): TokenUsage {
+    if (!current || next.totalTokens > current.totalTokens) return { ...next }
+    return current
+}
+
+export function usageDelta(total: TokenUsage | undefined, reported: TokenUsage | undefined): TokenUsage | undefined {
+    if (!total) return undefined
+    if (!reported) return total
+    const delta: TokenUsage = {
+        inputTokens: Math.max(total.inputTokens - reported.inputTokens, 0),
+        outputTokens: Math.max(total.outputTokens - reported.outputTokens, 0),
+        cacheReadTokens: Math.max(total.cacheReadTokens - reported.cacheReadTokens, 0),
+        cacheWriteTokens: Math.max(total.cacheWriteTokens - reported.cacheWriteTokens, 0),
+        totalTokens: Math.max(total.totalTokens - reported.totalTokens, 0)
+    }
+    if (total.reasoningTokens !== undefined) {
+        delta.reasoningTokens = Math.max(total.reasoningTokens - (reported.reasoningTokens ?? 0), 0)
+    }
+    return delta
+}
+
+export type RunOutcome = { kind: 'done' } | { kind: 'retry' } | { kind: 'error'; message: string }
+
+export function outcomeOf(result: RunResult, promptFallback: boolean): RunOutcome {
+    if (result.status !== 'error') return { kind: 'done' }
+    const message = result.error?.message ?? 'the Cursor run failed'
+    if (!promptFallback && /system.?prompt/i.test(message)) return { kind: 'retry' }
+    return { kind: 'error', message }
+}
+
+async function restartWithoutPrompt(
+    session: CursorSession,
+    input: CursorInput,
+    selection: ModelSelection
+): Promise<void> {
+    session.promptFallback = true
+    session.agent?.close()
+    session.agent = undefined
+    await startRun(session, input, selection)
+    patchState({ noSystemPrompt: true })
+}
+
+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/state.ts

@@ -0,0 +1,43 @@
+import { readFileSync, renameSync, writeFileSync } from 'node:fs'
+import { join } from 'node:path'
+
+import { getAgentDir } from '@earendil-works/pi-coding-agent'
+import { Validator, type Infer } from 'ata-validator'
+import { t } from 'ata-validator/t'
+
+const CursorPreferences = t.object({
+    noSystemPrompt: t.optional(t.boolean()),
+    fast: t.optional(t.boolean())
+})
+
+const isPreferences = new Validator(CursorPreferences).isValidObject
+
+export type CursorState = Infer<typeof CursorPreferences>
+
+function statePath(): string {
+    return join(getAgentDir(), 'pi-cursor.json')
+}
+
+export function readState(): CursorState {
+    try {
+        const parsed: unknown = JSON.parse(readFileSync(statePath(), 'utf8'))
+        return isPreferences(parsed) ? parsed : {}
+    } catch {
+        return {}
+    }
+}
+
+export function patchState(patch: CursorState): void {
+    const path = statePath()
+    const temp = `${path}.${process.pid}.tmp`
+    writeFileSync(temp, JSON.stringify({ ...readState(), ...patch }), { mode: 0o600 })
+    renameSync(temp, path)
+}
+
+export function noSystemPrompt(): boolean {
+    return readState().noSystemPrompt === true
+}
+
+export function fastMode(): boolean {
+    return readState().fast === true
+}