@dyyz1993/pi-coding-agent 0.69.18 → 0.69.24

package/dist/core/agent-session.js

@@ -12,10 +12,13 @@
  *
  * Modes use this class and add their own I/O layer on top.
  */
+import { randomUUID } from "node:crypto";
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { basename, dirname, join, resolve } from "node:path";
 import { Agent, } from "@dyyz1993/pi-agent-core";
 import { complete, isContextOverflow, modelsAreEqual, resetApiProviders, supportsXhigh } from "@dyyz1993/pi-ai";
+import { Compile } from "typebox/compile";
+import { Value } from "typebox/value";
 import { getDocsPath } from "../config.js";
 import { theme } from "../modes/interactive/theme/theme.js";
 import { stripFrontmatter } from "../utils/frontmatter.js";
@@ -33,6 +36,7 @@ import { createSyntheticSourceInfo } from "./source-info.js";
 import { buildSystemPrompt } from "./system-prompt.js";
 import { createLocalBashOperations } from "./tools/bash.js";
 import { createAllToolDefinitions, createTool } from "./tools/index.js";
+import { stripMarkdownCodeBlock } from "./tools/strip-markdown.js";
 import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.js";
 /**
  * Parse a skill block from message text.
@@ -91,6 +95,7 @@ export class AgentSession {
     // Extension system
     _extensionRunner;
     _turnIndex = 0;
+    _toolCallVariables;
     _resourceLoader;
     _customTools;
     _baseToolDefinitions = new Map();
@@ -106,6 +111,8 @@ export class AgentSession {
     _extensionErrorListener;
     _extensionErrorUnsubscriber;
     _registerChannel;
+    _sessionAbortController = new AbortController();
+    _backgroundTasks = new Set();
     // Model registry for API key resolution
     _modelRegistry;
     // Tool registry for extension getTools/setTools
@@ -181,6 +188,7 @@ export class AgentSession {
                     toolName: toolCall.name,
                     toolCallId: toolCall.id,
                     input: args,
+                    variables: this._toolCallVariables,
                 });
             }
             catch (err) {
@@ -494,6 +502,12 @@ export class AgentSession {
      */
     dispose() {
         this._extensionRunner.invalidate("This extension instance is stale after session replacement or reload. Use the provided replacement-session context instead.");
+        this._sessionAbortController.abort();
+        for (const task of this._backgroundTasks) {
+            task.cancel();
+        }
+        void Promise.allSettled([...this._backgroundTasks].map((t) => t.promise));
+        this._backgroundTasks.clear();
         this._disconnectFromAgent();
         this._eventListeners = [];
     }
@@ -512,10 +526,20 @@ export class AgentSession {
     get thinkingLevel() {
         return this.agent.state.thinkingLevel;
     }
+    set toolCallVariables(vars) {
+        this._toolCallVariables = vars;
+    }
+    get toolCallVariables() {
+        return this._toolCallVariables;
+    }
     /** Whether agent is currently streaming a response */
     get isStreaming() {
         return this.agent.state.isStreaming;
     }
+    /** Signal that aborts on session shutdown */
+    get sessionSignal() {
+        return this._sessionAbortController.signal;
+    }
     /** Current effective system prompt (includes any per-turn extension modifications) */
     get systemPrompt() {
         return this.agent.state.systemPrompt;
@@ -1774,10 +1798,14 @@ export class AgentSession {
                     throw new Error(`registerChannel("${name}") is only available in RPC mode`);
                 }),
             callLLM: (options) => this.callLLM(options),
+            callLLMStructured: (opts) => this.callLLMStructured(opts),
+            forkAgent: (prompt, options) => this.forkAgent(prompt, options),
+            background: (fn) => this.background(fn),
         }, {
             getModel: () => this.model,
             isIdle: () => !this.isStreaming,
             getSignal: () => this.agent.signal,
+            getSessionSignal: () => this._sessionAbortController.signal,
             abort: () => this.abort(),
             hasPendingMessages: () => this.pendingMessageCount > 0,
             shutdown: () => {
@@ -1912,6 +1940,186 @@ export class AgentSession {
         }
         return resultText;
     }
+    async callLLMStructured(options) {
+        const maxRetries = options.maxRetries ?? 0;
+        let lastError;
+        const schemaJson = JSON.stringify(options.schema);
+        const structuredSystemPrompt = (options.systemPrompt ?? "") +
+            `\n\nRespond with valid JSON matching this schema:\n${schemaJson}\n\nRespond with JSON only, no markdown.`;
+        for (let attempt = 0; attempt <= maxRetries; attempt++) {
+            const messages = attempt === 0
+                ? options.messages
+                : [
+                    ...options.messages,
+                    { role: "assistant", content: lastError?.raw ?? "" },
+                    {
+                        role: "user",
+                        content: `Your previous response was invalid: ${lastError?.message}. Please respond with valid JSON matching the schema.`,
+                    },
+                ];
+            const raw = await this.callLLM({
+                ...options,
+                systemPrompt: structuredSystemPrompt,
+                messages,
+            });
+            try {
+                const cleaned = stripMarkdownCodeBlock(raw);
+                const parsed = JSON.parse(cleaned);
+                const check = Compile(options.schema);
+                const coerced = Value.Convert(options.schema, parsed);
+                if (!check.Check(coerced)) {
+                    const errors = check
+                        .Errors(coerced)
+                        .map((e) => `${e.instancePath}: ${e.message}`)
+                        .join("; ");
+                    const err = new Error(`Schema validation failed: ${errors}`);
+                    err.raw = raw;
+                    err.reason = "schema_validation";
+                    lastError = err;
+                    if (attempt >= maxRetries)
+                        throw err;
+                    continue;
+                }
+                return coerced;
+            }
+            catch (e) {
+                if (e instanceof SyntaxError) {
+                    const err = new Error(`JSON parse failed: ${e.message}`);
+                    err.raw = raw;
+                    err.reason = "json_parse";
+                    lastError = err;
+                    if (attempt >= maxRetries)
+                        throw err;
+                    continue;
+                }
+                if (e.reason) {
+                    lastError = e;
+                    if (attempt >= maxRetries)
+                        throw lastError;
+                    continue;
+                }
+                throw e;
+            }
+        }
+        throw lastError ?? new Error("callLLMStructured failed");
+    }
+    async forkAgent(promptText, options) {
+        const model = this.model;
+        if (!model)
+            throw new Error("No model selected");
+        const auth = await this._modelRegistry.getApiKeyAndHeaders(model);
+        if (!auth?.ok) {
+            throw new Error("error" in auth ? auth.error : `No API key configured for ${model.provider}`);
+        }
+        if (!auth.apiKey) {
+            throw new Error(`No API key configured for ${model.provider}`);
+        }
+        if (options?.signal?.aborted)
+            throw new Error("Aborted");
+        const opts = options ?? {};
+        let toolNames = opts.tools ?? ["read", "grep", "find", "ls"];
+        if (opts.bash === "deny") {
+            toolNames = toolNames.filter((t) => t !== "bash");
+        }
+        const toolInstances = toolNames
+            .map((name) => {
+            try {
+                const registered = this._toolRegistry.get(name);
+                if (registered)
+                    return registered;
+                return createTool(name, this._cwd);
+            }
+            catch {
+                return undefined;
+            }
+        })
+            .filter((t) => t !== undefined);
+        const effectiveSystemPrompt = opts.inheritSystemPrompt
+            ? (this.agent.state.systemPrompt ?? opts.systemPrompt ?? "")
+            : (opts.systemPrompt ?? "");
+        const messages = opts.shareContext ? [...this.agent.state.messages] : [];
+        const maxTurns = opts.maxTurns ?? 5;
+        let turnCount = 0;
+        const forkedAgent = new Agent({
+            getApiKey: () => auth.apiKey,
+            initialState: {
+                systemPrompt: effectiveSystemPrompt,
+                model,
+                thinkingLevel: "off",
+                tools: toolInstances,
+                messages,
+            },
+            sessionId: opts.shareContext ? this.agent.sessionId : undefined,
+        });
+        let resultText = "";
+        const usage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0 };
+        const abortHandler = opts.signal ? () => forkedAgent.abort() : undefined;
+        if (abortHandler && opts.signal) {
+            opts.signal.addEventListener("abort", abortHandler, { once: true });
+        }
+        const unsub = forkedAgent.subscribe((event) => {
+            if (event.type === "turn_end") {
+                turnCount++;
+                if (turnCount >= maxTurns) {
+                    forkedAgent.abort();
+                }
+            }
+            if (event.type === "message_end") {
+                const msg = event.message;
+                if (msg.role === "assistant") {
+                    const asst = msg;
+                    const content = asst.content;
+                    if (Array.isArray(content)) {
+                        resultText = content
+                            .filter((c) => c.type === "text")
+                            .map((c) => c.text)
+                            .join("\n");
+                    }
+                    if (asst.usage) {
+                        usage.input = asst.usage.input ?? 0;
+                        usage.output = asst.usage.output ?? 0;
+                        usage.cacheRead = asst.usage.cacheRead ?? 0;
+                        usage.cacheWrite = asst.usage.cacheWrite ?? 0;
+                        if (model.cost) {
+                            usage.cost =
+                                (usage.input * model.cost.input +
+                                    usage.output * model.cost.output +
+                                    usage.cacheRead * model.cost.cacheRead +
+                                    usage.cacheWrite * model.cost.cacheWrite) /
+                                    1_000_000;
+                        }
+                    }
+                }
+            }
+        });
+        try {
+            await forkedAgent.prompt({
+                role: "user",
+                content: [{ type: "text", text: promptText }],
+                timestamp: Date.now(),
+            });
+        }
+        finally {
+            unsub();
+            if (abortHandler && opts.signal) {
+                opts.signal.removeEventListener("abort", abortHandler);
+            }
+        }
+        return { text: resultText, usage };
+    }
+    background(fn) {
+        const controller = new AbortController();
+        const promise = fn(controller.signal);
+        const task = {
+            id: randomUUID(),
+            signal: controller.signal,
+            promise: promise,
+            cancel: () => controller.abort(),
+        };
+        this._backgroundTasks.add(task);
+        promise.finally(() => this._backgroundTasks.delete(task));
+        return task;
+    }
     _refreshToolRegistry(options) {
         const previousRegistryNames = new Set(this._toolRegistry.keys());
         const previousActiveToolNames = this.getActiveToolNames();
@@ -2043,6 +2251,14 @@ export class AgentSession {
             await this.extendResourcesFromExtensions("reload");
         }
     }
+    async setCwd(newCwd) {
+        this._cwd = newCwd;
+        this._buildRuntime({
+            activeToolNames: this.getActiveToolNames(),
+            flagValues: this._extensionRunner.getFlagValues(),
+            includeAllExtensionTools: true,
+        });
+    }
     // =========================================================================
     // Auto-Retry
     // =========================================================================