@tduyng/pi-subagent 0.0.0-stage → 0.0.1

package/src/subagent/recover.ts

@@ -0,0 +1,268 @@
+import { existsSync } from 'node:fs'
+
+import type { ExtensionAPI, ExtensionContext, AgentSession, ModelRuntime } from '@earendil-works/pi-coding-agent'
+
+import { createChildModelRuntime } from '../lib/child-runtime.js'
+import { loadAgents } from './agent.js'
+import { drivePrompt } from './drive.js'
+import {
+    claimRun,
+    CONTINUE_PROMPT,
+    deliveredKeys,
+    detailsFromRecord,
+    listRuns,
+    pendingTasks,
+    releaseClaim,
+    resolveSessionFile,
+    runUpdatedAt,
+    saveRun,
+    syncRecord,
+    type DurableRun,
+    type DurableTask
+} from './durable.js'
+import { findLastAssistant, resultText } from './format.js'
+import {
+    acquirePermit,
+    failureReason,
+    releasePermit,
+    resolveMaxTurns,
+    resolveTimeoutMs,
+    salvageFromMessages,
+    salvageOutput
+} from './helpers.js'
+import { createChildSession } from './session.js'
+import { liveSessions, registry } from './state.js'
+import { THINKING_LEVELS, type RunDetails, type TaskDetails, type ThinkingLevel } from './types.js'
+
+const RECOVERY_WINDOW_MS = 10 * 60_000
+const INTERRUPTED_ERROR = 'Pi exited before this child returned. Resume it with subagent({ resume: "<id>" }).'
+
+export interface ResumeTaskInput {
+    pi: ExtensionAPI
+    ctx: ExtensionContext
+    run: DurableRun
+    record: DurableTask
+    task: TaskDetails
+    signal?: AbortSignal
+    followUp?: string
+    maxTurns?: number
+    timeoutMs?: number
+    modelRuntime?: ModelRuntime
+    onEmit: (force?: boolean) => void
+}
+
+function thinkingLevel(value: string): ThinkingLevel {
+    return THINKING_LEVELS.find((level) => level === value) ?? 'medium'
+}
+
+async function openSession(input: ResumeTaskInput): Promise<AgentSession> {
+    const live = liveSessions.get(input.task.id)
+    if (live) return live
+    const { record, ctx } = input
+    if (record.worktree && !existsSync(record.worktree)) {
+        throw new Error(`The worktree for "${record.label}" is gone; start the task again.`)
+    }
+    const sessionFile = resolveSessionFile(record)
+    if (!sessionFile) throw new Error(`No saved child session for "${record.label}"; start the task again.`)
+    const model = ctx.modelRegistry.find(record.provider, record.model) ?? ctx.model
+    if (!model) throw new Error(`Model ${record.provider}/${record.model} is no longer available.`)
+    const modelRuntime = input.modelRuntime ?? (await createChildModelRuntime(ctx))
+    const agent = record.agent ? loadAgents(ctx.cwd).get(record.agent) : undefined
+    return createChildSession({
+        cwd: record.cwd,
+        sessionDir: record.sessionDir,
+        model,
+        modelRuntime,
+        thinking: thinkingLevel(record.thinking),
+        parentSystemPrompt: ctx.getSystemPrompt(),
+        allTools: input.pi.getAllTools().map((tool) => tool.name),
+        agent,
+        schema: record.schema,
+        sessionFile,
+        onStructured: (value) => {
+            input.task.structured = value
+        }
+    })
+}
+
+export async function resumeTask(input: ResumeTaskInput): Promise<void> {
+    const { record, task, run } = input
+    let acquired = false
+    try {
+        await acquirePermit(input.signal)
+        acquired = true
+        const session = await openSession(input)
+        liveSessions.set(task.id, session)
+        task.status = 'running'
+        task.startedAt ??= Date.now()
+        task.error = undefined
+        task.partial = undefined
+        record.status = 'running'
+        input.onEmit(true)
+        const maxTurns = resolveMaxTurns(input.maxTurns, record.maxTurns)
+        const timeoutMs = resolveTimeoutMs(input.timeoutMs) ?? record.timeoutMs
+        const resumed = await drivePrompt({
+            session,
+            prompt: input.followUp ?? CONTINUE_PROMPT,
+            maxTurns,
+            schema: record.schema,
+            cwd: record.cwd,
+            timeoutMs,
+            abortSignal: input.signal,
+            onEmit: input.onEmit,
+            detail: task
+        })
+        const assistant = resumed.assistant ?? findLastAssistant(session.messages)
+        if (resumed.output) task.output = resumed.output
+        if (resumed.touched.length > 0) {
+            task.touched = [...new Set([...(task.touched ?? []), ...resumed.touched])].slice(0, 12)
+        }
+        task.sessionFile = session.sessionFile ?? task.sessionFile
+        task.finishedAt = Date.now()
+        if (!task.output || resumed.timedOut) {
+            const salvaged = salvageOutput(assistant) ?? salvageFromMessages(session.messages)
+            if (salvaged && !task.output) {
+                task.output = salvaged
+                task.partial = true
+            }
+        }
+        if (resumed.schemaFallback) task.hasSchema = false
+        if (task.hasSchema && task.structured === undefined) task.partial = true
+        const reason = failureReason({
+            hasAssistant: assistant !== undefined,
+            timedOut: resumed.timedOut,
+            duration: undefined,
+            hitTurnLimit: resumed.turnCount >= maxTurns,
+            maxTurns,
+            aborted: input.signal?.aborted === true,
+            hasSchema: task.hasSchema ?? false,
+            hasStructured: task.structured !== undefined,
+            output: task.output ?? '',
+            stopReason: assistant?.stopReason,
+            errorMessage: assistant?.errorMessage
+        })
+        if (reason) {
+            task.status = 'failed'
+            task.error = reason
+        } else {
+            task.status = 'completed'
+        }
+    } catch (error) {
+        task.status = 'failed'
+        task.error = error instanceof Error ? error.message : String(error)
+        task.finishedAt ??= Date.now()
+    } finally {
+        if (task.status === 'running') {
+            task.status = 'failed'
+            task.error ??= 'Subagent aborted.'
+            task.finishedAt ??= Date.now()
+        }
+        syncRecord(record, task)
+        saveRun(run)
+        if (acquired) releasePermit()
+        input.onEmit(true)
+    }
+}
+
+export interface RecoverRunsInput {
+    pi: ExtensionAPI
+    ctx: ExtensionContext
+    onEmit: (force?: boolean) => void
+}
+
+export async function recoverRuns(input: RecoverRunsInput): Promise<void> {
+    const sessionId = input.ctx.sessionManager.getSessionId()
+    const delivered = deliveredKeys(input.ctx.sessionManager)
+    for (const run of listRuns(sessionId)) {
+        if (delivered.has(run.runId)) continue
+        if (run.toolCallId !== undefined && delivered.has(run.toolCallId)) continue
+        const pending = pendingTasks(run)
+        if (pending.length === 0) {
+            run.finishedAt ??= Date.now()
+            saveRun(run)
+            await postResult(input.pi, run, 'This run finished after the session had already closed.')
+            continue
+        }
+        if (!claimRun(sessionId, run.runId)) continue
+        try {
+            const tasks = run.tasks.map(detailsFromRecord)
+            for (const task of tasks) {
+                if (!registry.some((known) => known.id === task.id)) registry.push(task)
+            }
+            input.onEmit(true)
+            const updatedAt = runUpdatedAt(sessionId, run.runId) ?? run.startedAt
+            const resumable =
+                Date.now() - updatedAt <= RECOVERY_WINDOW_MS &&
+                pending.some((record) => resolveSessionFile(record) !== undefined)
+            if (!resumable) {
+                markInterrupted(run, tasks)
+                await postResult(
+                    input.pi,
+                    run,
+                    'Pi exited before this run returned. Failed children keep their saved sessions and can be resumed by id.'
+                )
+                continue
+            }
+            const modelRuntime = await createChildModelRuntime(input.ctx)
+            await Promise.allSettled(
+                run.tasks.map((record, index) => {
+                    const task = tasks[index]
+                    if (!task || (record.status !== 'queued' && record.status !== 'running')) return undefined
+                    return resumeTask({
+                        pi: input.pi,
+                        ctx: input.ctx,
+                        run,
+                        record,
+                        task,
+                        modelRuntime,
+                        onEmit: input.onEmit
+                    })
+                })
+            )
+            run.finishedAt = Date.now()
+            saveRun(run)
+            await postResult(
+                input.pi,
+                run,
+                'Recovered after Pi exited: interrupted children resumed from their saved sessions.'
+            )
+        } finally {
+            releaseClaim(sessionId, run.runId)
+        }
+    }
+}
+
+function markInterrupted(run: DurableRun, tasks: TaskDetails[]): void {
+    for (const task of tasks) {
+        if (task.status !== 'queued' && task.status !== 'running') continue
+        task.status = 'failed'
+        task.error = INTERRUPTED_ERROR
+        task.finishedAt = Date.now()
+    }
+    for (const [index, record] of run.tasks.entries()) {
+        const task = tasks[index]
+        if (task) syncRecord(record, task)
+    }
+    run.finishedAt = Date.now()
+    saveRun(run)
+}
+
+async function postResult(pi: ExtensionAPI, run: DurableRun, note: string): Promise<void> {
+    const tasks = run.tasks.map(detailsFromRecord)
+    const details = {
+        status: tasks.every((task) => task.status === 'completed') ? 'completed' : 'failed',
+        runId: run.runId,
+        startedAt: run.startedAt,
+        finishedAt: run.finishedAt,
+        tasks
+    } satisfies RunDetails
+    await pi.sendMessage(
+        {
+            customType: 'subagent-durable',
+            content: [{ type: 'text', text: `${note}\n\n${resultText(details)}` }],
+            display: true,
+            details: { runId: run.runId }
+        },
+        { triggerTurn: false }
+    )
+}

package/src/subagent/session.ts

@@ -0,0 +1,87 @@
+import type { Model } from '@earendil-works/pi-ai'
+import {
+    createAgentSession,
+    DefaultResourceLoader,
+    getAgentDir,
+    ModelRuntime,
+    SessionManager,
+    type AgentSession
+} from '@earendil-works/pi-coding-agent'
+
+import cloak from '../../../pi-guard/src/cloak.js'
+import gitInterceptor from '../../../pi-guard/src/git-interceptor.js'
+import { childPromptExtras } from './helpers.js'
+import { makeStructuredOutputTool, structuredOutputInstruction, type JsonPayload } from './structured.js'
+import type { AgentProfile, ThinkingLevel } from './types.js'
+
+export interface ChildSessionInput {
+    cwd: string
+    sessionDir: string
+    model: Model<any>
+    modelRuntime: ModelRuntime
+    thinking: ThinkingLevel
+    parentSystemPrompt: string
+    allTools: string[]
+    agent?: AgentProfile
+    schema?: JsonPayload
+    sessionFile?: string
+    onStructured?: (value: JsonPayload) => void
+}
+
+export async function createChildSession(input: ChildSessionInput): Promise<AgentSession> {
+    const { cwd, sessionDir, model, modelRuntime, thinking, parentSystemPrompt, allTools, agent, schema, sessionFile } =
+        input
+    const loaderOptions: ConstructorParameters<typeof DefaultResourceLoader>[0] = {
+        cwd,
+        agentDir: getAgentDir(),
+        noExtensions: true,
+        extensionFactories: [cloak, gitInterceptor],
+        noThemes: true,
+        noPromptTemplates: true,
+        noSkills: Boolean(agent && !agent.inheritSkills),
+        noContextFiles: Boolean(agent && !agent.inheritProjectContext)
+    }
+    const extras = childPromptExtras(model, parentSystemPrompt, structuredOutputInstruction(schema))
+    if (agent?.systemPrompt && agent.systemPromptMode === 'replace') {
+        loaderOptions['systemPromptOverride'] = () => agent.systemPrompt
+        loaderOptions['appendSystemPromptOverride'] = () => extras
+    } else {
+        loaderOptions['appendSystemPromptOverride'] = (base: string[]) => [
+            ...base,
+            ...(agent?.systemPrompt ? [agent.systemPrompt] : []),
+            ...extras
+        ]
+    }
+    const loader = new DefaultResourceLoader(loaderOptions)
+    await loader.reload()
+
+    const childTools = agent?.tools.length
+        ? agent.disallowedTools.length
+            ? agent.tools.filter((tool) => !agent.disallowedTools.includes(tool))
+            : agent.tools
+        : agent?.disallowedTools.length
+          ? allTools.filter((tool) => !agent.disallowedTools.includes(tool))
+          : undefined
+    const childToolList = schema !== undefined && childTools ? [...childTools, 'structured_output'] : childTools
+
+    const sessionOptions: Parameters<typeof createAgentSession>[0] = {
+        cwd,
+        modelRuntime,
+        model,
+        thinkingLevel: thinking,
+        resourceLoader: loader,
+        sessionManager: sessionFile
+            ? SessionManager.open(sessionFile, sessionDir, cwd)
+            : SessionManager.create(cwd, sessionDir),
+        tools: childToolList
+    }
+    if (schema !== undefined) {
+        sessionOptions.customTools = [
+            makeStructuredOutputTool(schema, (value) => {
+                input.onStructured?.(value)
+            })
+        ]
+    }
+    const { session } = await createAgentSession(sessionOptions)
+    return session
+}

package/src/subagent/state.ts

@@ -0,0 +1,59 @@
+import type { AgentSession } from '@earendil-works/pi-coding-agent'
+
+import { FAILED_RETENTION_MS, FINISHED_RETENTION_MS, MAX_REGISTRY, type TaskDetails } from './types.js'
+
+export const registry: TaskDetails[] = []
+
+export const liveSessions = new Map<string, AgentSession>()
+
+function evict(index: number): void {
+    const task = registry[index]
+    if (!task) return
+    liveSessions.get(task.id)?.dispose()
+    liveSessions.delete(task.id)
+    registry.splice(index, 1)
+}
+
+function retentionFor(task: TaskDetails): number {
+    return task.status === 'failed' ? FAILED_RETENTION_MS : FINISHED_RETENTION_MS
+}
+
+export function sweepRegistry(): void {
+    const now = Date.now()
+    for (let i = registry.length - 1; i >= 0; i--) {
+        const task = registry[i]
+        if (!task) continue
+        const finished = task.status === 'completed' || task.status === 'failed'
+        if (finished && task.finishedAt !== undefined && now - task.finishedAt > retentionFor(task)) evict(i)
+    }
+    while (registry.length > MAX_REGISTRY) evict(0)
+}
+
+export function nextExpiryAt(): number | undefined {
+    let earliest: number | undefined
+    for (const task of registry) {
+        const finished = task.status === 'completed' || task.status === 'failed'
+        if (!finished || task.finishedAt === undefined) continue
+        const due = task.finishedAt + retentionFor(task)
+        if (earliest === undefined || due < earliest) earliest = due
+    }
+    return earliest
+}
+
+export function dismissTask(id: string): boolean {
+    const index = registry.findIndex((task) => task.id === id)
+    const task = registry[index]
+    if (!task || task.status === 'queued' || task.status === 'running') return false
+    evict(index)
+    return true
+}
+
+export function roster(): TaskDetails[] {
+    sweepRegistry()
+    return registry.toSorted((a, b) => {
+        const aActive = a.status === 'queued' || a.status === 'running'
+        const bActive = b.status === 'queued' || b.status === 'running'
+        if (aActive !== bActive) return aActive ? -1 : 1
+        return (b.startedAt ?? 0) - (a.startedAt ?? 0)
+    })
+}

package/src/subagent/structured.ts

@@ -0,0 +1,202 @@
+import { defineTool, type ToolDefinition } from '@earendil-works/pi-coding-agent'
+import { Validator, type JSONSchema } from 'ata-validator'
+import { t } from 'ata-validator/t'
+
+const isStringArray = new Validator(t.array(t.string())).isValidObject
+const isStringOrBoolean = new Validator(t.union([t.string(), t.boolean()])).isValidObject
+const isNumber = new Validator(t.number()).isValidObject
+const isString = new Validator(t.string()).isValidObject
+const isBoolean = new Validator(t.boolean()).isValidObject
+
+const MAX_SCHEMA_CHARS = 100_000
+const SCHEMA_MAP_KEYWORDS = new Set(['$defs', 'definitions', 'dependentSchemas', 'patternProperties', 'properties'])
+const SCHEMA_VALUE_KEYWORDS = new Set([
+    'additionalItems',
+    'additionalProperties',
+    'contains',
+    'else',
+    'if',
+    'items',
+    'not',
+    'then',
+    'unevaluatedItems',
+    'unevaluatedProperties'
+])
+
+const STRUCTURED_OUTPUT_SYSTEM_INSTRUCTION =
+    'When your task is complete, call the `structured_output` tool exactly once as your final action, with fields matching the required schema. Do not write any other text after it.'
+
+const STRUCTURED_OUTPUT_TOOL_DESCRIPTION =
+    'Return your final result as structured data matching the required schema. Call this exactly once, as your last action; do not write any other text after it.'
+
+type JsonValue = null | boolean | number | string | JsonValue[] | JsonRecord
+
+export type JsonPayload = JsonValue | JSONSchema
+
+interface JsonRecord {
+    [key: string]: JsonValue
+}
+
+const ObjectRecordSchema = t.object({}, { additionalProperties: true })
+const isObjectRecordSchema = new Validator(ObjectRecordSchema).isValidObject
+
+function isRecord(value: JsonPayload | undefined): value is JsonRecord {
+    if (!isObjectRecordSchema(value)) return false
+    const prototype = Object.getPrototypeOf(value)
+    return prototype === Object.prototype || prototype === null
+}
+
+export function isJsonSchema(value: unknown): value is JSONSchema {
+    if (!isObjectRecordSchema(value)) return false
+    const prototype = Object.getPrototypeOf(value)
+    if (prototype !== Object.prototype && prototype !== null) return false
+    const ancestors = new Set<object>()
+    let nodes = 0
+    const validate = (current: JsonValue, depth: number): boolean => {
+        if (++nodes > 10_000 || depth > 24) return false
+        if (current === null || isStringOrBoolean(current)) return true
+        if (isNumber(current)) return Number.isFinite(current)
+        if (Array.isArray(current)) return current.every((item) => validate(item, depth + 1))
+        if (!isRecord(current) || ancestors.has(current)) return false
+
+        ancestors.add(current)
+        const valid = Object.entries(current).every(([key, child]) => {
+            if (key === '__proto__' || key === 'constructor' || key === 'prototype') return false
+            if (key === 'required' && !isStringArray(child)) return false
+            if (key === 'type' && !isString(child) && !isStringArray(child)) return false
+            if (key === 'additionalProperties' && !isBoolean(child) && !isRecord(child)) return false
+            if (
+                SCHEMA_MAP_KEYWORDS.has(key) &&
+                (!isRecord(child) || !Object.values(child).every((item) => validate(item, depth + 1)))
+            ) {
+                return false
+            }
+            if (SCHEMA_VALUE_KEYWORDS.has(key) && child !== undefined && !validate(child, depth + 1)) return false
+            return validate(child, depth + 1)
+        })
+        ancestors.delete(current)
+        return valid
+    }
+    return validate(value, 0)
+}
+
+function normalizeRequired(value: JsonValue | undefined): string[] | undefined {
+    if (isString(value)) return [value]
+    if (isRecord(value)) {
+        const item = value['item']
+        const nested = normalizeRequired(item)
+        if (nested) return nested
+        return Object.keys(value).length > 0 &&
+            Object.values(value).every((entry) => isRecord(entry) || entry === undefined)
+            ? Object.keys(value)
+            : undefined
+    }
+    if (!Array.isArray(value)) return undefined
+    const names: string[] = []
+    for (const item of value) {
+        if (isString(item)) {
+            names.push(item)
+            continue
+        }
+        if (!isRecord(item)) return undefined
+        const nested = item['item']
+        if (isString(nested)) names.push(nested)
+        else if (isStringArray(nested)) names.push(...nested)
+        else if (isRecord(nested) && isStringArray(nested['item'])) names.push(...nested['item'])
+        else return undefined
+    }
+    return names
+}
+
+function normalizeSchemaNode(value: JsonValue, stack: Set<object>): JsonValue {
+    if (Array.isArray(value)) return value.map((item) => normalizeSchemaNode(item, stack))
+    if (!isRecord(value)) return value
+    if (stack.has(value)) return value
+
+    stack.add(value)
+    const normalized: JsonRecord = { ...value }
+    const required = normalizeRequired(normalized['required'])
+    if (required) normalized['required'] = required
+    if (normalized['additionalProperties'] === 'false') normalized['additionalProperties'] = false
+    if (Array.isArray(normalized['type'])) {
+        const types = [...new Set(normalized['type'])]
+        if (types.length === 1 && types[0] !== undefined) normalized['type'] = types[0]
+        else if (isRecord(normalized['properties']) && types.includes('object') && normalized['items'] === undefined) {
+            normalized['type'] = 'object'
+        } else if (
+            normalized['items'] !== undefined &&
+            types.includes('array') &&
+            !isRecord(normalized['properties'])
+        ) {
+            normalized['type'] = 'array'
+        }
+    }
+    if (normalized['type'] === undefined) {
+        if (isRecord(normalized['properties'])) normalized['type'] = 'object'
+        else if (normalized['items'] !== undefined) normalized['type'] = 'array'
+    }
+
+    for (const key of SCHEMA_MAP_KEYWORDS) {
+        const child = normalized[key]
+        if (isRecord(child)) {
+            normalized[key] = Object.fromEntries(
+                Object.entries(child).map(([name, schema]) => [name, normalizeSchemaNode(schema, stack)])
+            )
+        }
+    }
+    for (const key of SCHEMA_VALUE_KEYWORDS) {
+        if (normalized[key] !== undefined) normalized[key] = normalizeSchemaNode(normalized[key], stack)
+    }
+    stack.delete(value)
+    return normalized
+}
+
+export function normalizeJsonSchema(value: JsonPayload | undefined): JSONSchema | undefined {
+    if (isString(value)) {
+        if (value.length > MAX_SCHEMA_CHARS) return undefined
+        try {
+            value = JSON.parse(value)
+        } catch {
+            return undefined
+        }
+    }
+    if (value === undefined) return undefined
+    if (!isRecord(value)) return undefined
+    const normalized: unknown = normalizeSchemaNode(value, new Set())
+    return isJsonSchema(normalized) ? normalized : undefined
+}
+
+function structuredParameters(schema: JsonPayload | undefined): JSONSchema {
+    const normalized = normalizeJsonSchema(schema)
+    if (!normalized) {
+        throw new Error('Structured output schema must be a bounded plain JSON object.')
+    }
+    return normalized
+}
+
+export function makeStructuredOutputTool(
+    schema: JsonPayload | undefined,
+    capture: (value: JsonPayload) => void
+): ToolDefinition {
+    return defineTool({
+        name: 'structured_output',
+        label: 'Structured Output',
+        description: STRUCTURED_OUTPUT_TOOL_DESCRIPTION,
+        parameters: structuredParameters(schema),
+        constrainedSampling: { type: 'json_schema', strict: 'prefer' },
+        async execute(_toolCallId, params: JsonPayload) {
+            capture(params)
+            return {
+                content: [{ type: 'text', text: 'Recorded structured result.' }],
+                details: params,
+                terminate: true
+            }
+        }
+    })
+}
+
+export function structuredOutputInstruction(schema: JsonPayload | undefined): string | undefined {
+    const normalized = normalizeJsonSchema(schema)
+    if (!normalized) return undefined
+    return `${STRUCTURED_OUTPUT_SYSTEM_INSTRUCTION}\n\nRequired schema:\n${JSON.stringify(normalized, null, 2)}`
+}