@draht/coding-agent 2026.3.14 → 2026.3.25-1

package/dist/core/agent-session.js

@@ -12,8 +12,8 @@
  *
  * Modes use this class and add their own I/O layer on top.
  */
-import { readFileSync } from "node:fs";
-import { basename, dirname, join } from "node:path";
+import { copyFileSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
+import { basename, dirname, join, resolve } from "node:path";
 import { isContextOverflow, modelsAreEqual, resetApiProviders, supportsXhigh } from "@draht/ai";
 import { getDocsPath } from "../config.js";
 import { theme } from "../modes/interactive/theme/theme.js";
@@ -24,12 +24,13 @@ import { calculateContextTokens, collectEntriesForBranchSummary, compact, estima
 import { DEFAULT_THINKING_LEVEL } from "./defaults.js";
 import { exportSessionToHtml } from "./export-html/index.js";
 import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
-import { ExtensionRunner, wrapRegisteredTools, wrapToolsWithExtensions, } from "./extensions/index.js";
+import { ExtensionRunner, wrapRegisteredTools, } from "./extensions/index.js";
 import { expandPromptTemplate } from "./prompt-templates.js";
-import { getLatestCompactionEntry } from "./session-manager.js";
-import { BUILTIN_SLASH_COMMANDS } from "./slash-commands.js";
+import { CURRENT_SESSION_VERSION, getLatestCompactionEntry } from "./session-manager.js";
+import { createSyntheticSourceInfo } from "./source-info.js";
 import { buildSystemPrompt } from "./system-prompt.js";
-import { createAllTools } from "./tools/index.js";
+import { createAllToolDefinitions } from "./tools/index.js";
+import { createToolDefinitionFromAgentTool, wrapToolDefinition } from "./tools/tool-definition-wrapper.js";
 /**
  * Parse a skill block from message text.
  * Returns null if the text doesn't contain a skill block.
@@ -89,7 +90,7 @@ export class AgentSession {
     _turnIndex = 0;
     _resourceLoader;
     _customTools;
-    _baseToolRegistry = new Map();
+    _baseToolDefinitions = new Map();
     _cwd;
     _extensionRunnerRef;
     _initialActiveToolNames;
@@ -103,6 +104,7 @@ export class AgentSession {
     _modelRegistry;
     // Tool registry for extension getTools/setTools
     _toolRegistry = new Map();
+    _toolDefinitions = new Map();
     _toolPromptSnippets = new Map();
     _toolPromptGuidelines = new Map();
     // Base system prompt (without extension appends) - used to apply fresh appends each turn
@@ -122,6 +124,7 @@ export class AgentSession {
         // Always subscribe to agent events for internal handling
         // (session persistence, extensions, auto-compaction, retry logic)
         this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
+        this._installAgentToolHooks();
         this._buildRuntime({
             activeToolNames: this._initialActiveToolNames,
             includeAllExtensionTools: true,
@@ -131,6 +134,59 @@ export class AgentSession {
     get modelRegistry() {
         return this._modelRegistry;
     }
+    /**
+     * Install tool hooks once on the Agent instance.
+     *
+     * The callbacks read `this._extensionRunner` at execution time, so extension reload swaps in the
+     * new runner without reinstalling hooks. Extension-specific tool wrappers are still used to adapt
+     * registered tool execution to the extension context. Tool call and tool result interception now
+     * happens here instead of in wrappers.
+     */
+    _installAgentToolHooks() {
+        this.agent.setBeforeToolCall(async ({ toolCall, args }) => {
+            const runner = this._extensionRunner;
+            if (!runner?.hasHandlers("tool_call")) {
+                return undefined;
+            }
+            await this._agentEventQueue;
+            try {
+                return await runner.emitToolCall({
+                    type: "tool_call",
+                    toolName: toolCall.name,
+                    toolCallId: toolCall.id,
+                    input: args,
+                });
+            }
+            catch (err) {
+                if (err instanceof Error) {
+                    throw err;
+                }
+                throw new Error(`Extension failed, blocking execution: ${String(err)}`);
+            }
+        });
+        this.agent.setAfterToolCall(async ({ toolCall, args, result, isError }) => {
+            const runner = this._extensionRunner;
+            if (!runner?.hasHandlers("tool_result")) {
+                return undefined;
+            }
+            const hookResult = await runner.emitToolResult({
+                type: "tool_result",
+                toolName: toolCall.name,
+                toolCallId: toolCall.id,
+                input: args,
+                content: result.content,
+                details: isError ? undefined : result.details,
+                isError,
+            });
+            if (!hookResult || isError) {
+                return undefined;
+            }
+            return {
+                content: hookResult.content,
+                details: hookResult.details,
+            };
+        });
+    }
     // =========================================================================
     // Event Subscription
     // =========================================================================
@@ -234,7 +290,6 @@ export class AgentSession {
                         attempt: this._retryAttempt,
                     });
                     this._retryAttempt = 0;
-                    this._resolveRetry();
                 }
             }
         }
@@ -248,6 +303,7 @@ export class AgentSession {
                 if (didRetry)
                     return; // Retry was initiated, don't proceed to compaction
             }
+            this._resolveRetry();
             await this._checkCompaction(msg);
         }
     }
@@ -439,15 +495,19 @@ export class AgentSession {
         return this.agent.state.tools.map((t) => t.name);
     }
     /**
-     * Get all configured tools with name, description, and parameter schema.
+     * Get all configured tools with name, description, parameter schema, and source metadata.
      */
     getAllTools() {
-        return Array.from(this._toolRegistry.values()).map((t) => ({
-            name: t.name,
-            description: t.description,
-            parameters: t.parameters,
+        return Array.from(this._toolDefinitions.values()).map(({ definition, sourceInfo }) => ({
+            name: definition.name,
+            description: definition.description,
+            parameters: definition.parameters,
+            sourceInfo,
         }));
     }
+    getToolDefinition(name) {
+        return this._toolDefinitions.get(name)?.definition;
+    }
     /**
      * Set active tools by name.
      * Only tools in the registry can be enabled. Unknown tool names are ignored.
@@ -748,8 +808,9 @@ export class AgentSession {
         }
     }
     /**
-     * Queue a steering message to interrupt the agent mid-run.
-     * Delivered after current tool execution, skips remaining tools.
+     * Queue a steering message while the agent is running.
+     * Delivered after the current assistant turn finishes executing its tool calls,
+     * before the next LLM call.
      * Expands skill commands and prompt templates. Errors on extension commands.
      * @param images Optional image attachments to include with the message
      * @throws Error if text is an extension command
@@ -1371,19 +1432,9 @@ export class AgentSession {
         // This ensures sessions that hit persistent API errors (e.g. 529) can still compact.
         let contextTokens;
         if (assistantMessage.stopReason === "error") {
-            const messages = this.agent.state.messages;
-            const estimate = estimateContextTokens(messages);
+            const estimate = estimateContextTokens(this.agent.state.messages);
             if (estimate.lastUsageIndex === null)
                 return; // No usage data at all
-            // Verify the usage source is post-compaction. Kept pre-compaction messages
-            // have stale usage reflecting the old (larger) context and would falsely
-            // trigger compaction right after one just finished.
-            const usageMsg = messages[estimate.lastUsageIndex];
-            if (compactionEntry &&
-                usageMsg.role === "assistant" &&
-                usageMsg.timestamp <= new Date(compactionEntry.timestamp).getTime()) {
-                return;
-            }
             contextTokens = estimate.tokens;
         }
         else {
@@ -1589,37 +1640,36 @@ export class AgentSession {
             ? runner.onError(this._extensionErrorListener)
             : undefined;
     }
+    _refreshCurrentModelFromRegistry() {
+        const currentModel = this.model;
+        if (!currentModel) {
+            return;
+        }
+        const refreshedModel = this._modelRegistry.find(currentModel.provider, currentModel.id);
+        if (!refreshedModel || refreshedModel === currentModel) {
+            return;
+        }
+        this.agent.setModel(refreshedModel);
+    }
     _bindExtensionCore(runner) {
-        const normalizeLocation = (source) => {
-            if (source === "user" || source === "project" || source === "path") {
-                return source;
-            }
-            return undefined;
-        };
-        const reservedBuiltins = new Set(BUILTIN_SLASH_COMMANDS.map((command) => command.name));
         const getCommands = () => {
-            const extensionCommands = runner
-                .getRegisteredCommandsWithPaths()
-                .filter(({ command }) => !reservedBuiltins.has(command.name))
-                .map(({ command, extensionPath }) => ({
-                name: command.name,
+            const extensionCommands = runner.getRegisteredCommands().map((command) => ({
+                name: command.invocationName,
                 description: command.description,
                 source: "extension",
-                path: extensionPath,
+                sourceInfo: command.sourceInfo,
             }));
             const templates = this.promptTemplates.map((template) => ({
                 name: template.name,
                 description: template.description,
                 source: "prompt",
-                location: normalizeLocation(template.source),
-                path: template.filePath,
+                sourceInfo: template.sourceInfo,
             }));
             const skills = this._resourceLoader.getSkills().skills.map((skill) => ({
                 name: `skill:${skill.name}`,
                 description: skill.description,
                 source: "skill",
-                location: normalizeLocation(skill.source),
-                path: skill.filePath,
+                sourceInfo: skill.sourceInfo,
             }));
             return [...extensionCommands, ...templates, ...skills];
         };
@@ -1690,6 +1740,15 @@ export class AgentSession {
                 })();
             },
             getSystemPrompt: () => this.systemPrompt,
+        }, {
+            registerProvider: (name, config) => {
+                this._modelRegistry.registerProvider(name, config);
+                this._refreshCurrentModelFromRegistry();
+            },
+            unregisterProvider: (name) => {
+                this._modelRegistry.unregisterProvider(name);
+                this._refreshCurrentModelFromRegistry();
+            },
         });
     }
     _refreshToolRegistry(options) {
@@ -1698,34 +1757,48 @@ export class AgentSession {
         const registeredTools = this._extensionRunner?.getAllRegisteredTools() ?? [];
         const allCustomTools = [
             ...registeredTools,
-            ...this._customTools.map((def) => ({ definition: def, extensionPath: "<sdk>" })),
+            ...this._customTools.map((definition) => ({
+                definition,
+                sourceInfo: createSyntheticSourceInfo(`<sdk:${definition.name}>`, { source: "sdk" }),
+            })),
         ];
-        this._toolPromptSnippets = new Map(allCustomTools
-            .map((registeredTool) => {
-            const snippet = this._normalizePromptSnippet(registeredTool.definition.promptSnippet ?? registeredTool.definition.description);
-            return snippet ? [registeredTool.definition.name, snippet] : undefined;
+        const definitionRegistry = new Map(Array.from(this._baseToolDefinitions.entries()).map(([name, definition]) => [
+            name,
+            {
+                definition,
+                sourceInfo: createSyntheticSourceInfo(`<builtin:${name}>`, { source: "builtin" }),
+            },
+        ]));
+        for (const tool of allCustomTools) {
+            definitionRegistry.set(tool.definition.name, {
+                definition: tool.definition,
+                sourceInfo: tool.sourceInfo,
+            });
+        }
+        this._toolDefinitions = definitionRegistry;
+        this._toolPromptSnippets = new Map(Array.from(definitionRegistry.values())
+            .map(({ definition }) => {
+            const snippet = this._normalizePromptSnippet(definition.promptSnippet);
+            return snippet ? [definition.name, snippet] : undefined;
         })
             .filter((entry) => entry !== undefined));
-        this._toolPromptGuidelines = new Map(allCustomTools
-            .map((registeredTool) => {
-            const guidelines = this._normalizePromptGuidelines(registeredTool.definition.promptGuidelines);
-            return guidelines.length > 0 ? [registeredTool.definition.name, guidelines] : undefined;
+        this._toolPromptGuidelines = new Map(Array.from(definitionRegistry.values())
+            .map(({ definition }) => {
+            const guidelines = this._normalizePromptGuidelines(definition.promptGuidelines);
+            return guidelines.length > 0 ? [definition.name, guidelines] : undefined;
         })
             .filter((entry) => entry !== undefined));
         const wrappedExtensionTools = this._extensionRunner
             ? wrapRegisteredTools(allCustomTools, this._extensionRunner)
             : [];
-        const toolRegistry = new Map(this._baseToolRegistry);
+        const toolRegistry = new Map(Array.from(this._baseToolDefinitions.values()).map((definition) => [
+            definition.name,
+            wrapToolDefinition(definition),
+        ]));
         for (const tool of wrappedExtensionTools) {
             toolRegistry.set(tool.name, tool);
         }
-        if (this._extensionRunner) {
-            const wrappedAllTools = wrapToolsWithExtensions(Array.from(toolRegistry.values()), this._extensionRunner);
-            this._toolRegistry = new Map(wrappedAllTools.map((tool) => [tool.name, tool]));
-        }
-        else {
-            this._toolRegistry = toolRegistry;
-        }
+        this._toolRegistry = toolRegistry;
         const nextActiveToolNames = options?.activeToolNames
             ? [...options.activeToolNames]
             : [...previousActiveToolNames];
@@ -1746,13 +1819,16 @@ export class AgentSession {
     _buildRuntime(options) {
         const autoResizeImages = this.settingsManager.getImageAutoResize();
         const shellCommandPrefix = this.settingsManager.getShellCommandPrefix();
-        const baseTools = this._baseToolsOverride
-            ? this._baseToolsOverride
-            : createAllTools(this._cwd, {
+        const baseToolDefinitions = this._baseToolsOverride
+            ? Object.fromEntries(Object.entries(this._baseToolsOverride).map(([name, tool]) => [
+                name,
+                createToolDefinitionFromAgentTool(tool),
+            ]))
+            : createAllToolDefinitions(this._cwd, {
                 read: { autoResizeImages },
                 bash: { commandPrefix: shellCommandPrefix },
             });
-        this._baseToolRegistry = new Map(Object.entries(baseTools).map(([name, tool]) => [name, tool]));
+        this._baseToolDefinitions = new Map(Object.entries(baseToolDefinitions).map(([name, tool]) => [name, tool]));
         const extensionsResult = this._resourceLoader.getExtensions();
         if (options.flagValues) {
             for (const [name, value] of options.flagValues) {
@@ -1816,8 +1892,8 @@ export class AgentSession {
         if (isContextOverflow(message, contextWindow))
             return false;
         const err = message.errorMessage;
-        // Match: overloaded_error, rate limit, 429, 500, 502, 503, 504, service unavailable, connection errors, fetch failed, terminated, retry delay exceeded
-        return /overloaded|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server error|internal error|connection.?error|connection.?refused|other side closed|fetch failed|upstream.?connect|reset before headers|terminated|retry delay/i.test(err);
+        // Match: overloaded_error, provider returned error, rate limit, 429, 500, 502, 503, 504, service unavailable, network/connection errors, fetch failed, terminated, retry delay exceeded
+        return /overloaded|provider.?returned.?error|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server.?error|internal.?error|network.?error|connection.?error|connection.?refused|other side closed|fetch failed|upstream.?connect|reset before headers|socket hang up|timed? out|timeout|terminated|retry delay/i.test(err);
     }
     /**
      * Handle retryable errors with exponential backoff.
@@ -2380,6 +2456,7 @@ export class AgentSession {
                 total: totalInput + totalOutput + totalCacheRead + totalCacheWrite,
             },
             cost: totalCost,
+            contextUsage: this.getContextUsage(),
         };
     }
     getContextUsage() {
@@ -2431,19 +2508,70 @@ export class AgentSession {
     async exportToHtml(outputPath) {
         const themeName = this.settingsManager.getTheme();
         // Create tool renderer if we have an extension runner (for custom tool HTML rendering)
-        let toolRenderer;
-        if (this._extensionRunner) {
-            toolRenderer = createToolHtmlRenderer({
-                getToolDefinition: (name) => this._extensionRunner.getToolDefinition(name),
-                theme,
-            });
-        }
+        const toolRenderer = createToolHtmlRenderer({
+            getToolDefinition: (name) => this.getToolDefinition(name),
+            theme,
+        });
         return await exportSessionToHtml(this.sessionManager, this.state, {
             outputPath,
             themeName,
             toolRenderer,
         });
     }
+    /**
+     * Export the current session branch to a JSONL file.
+     * Writes the session header followed by all entries on the current branch path.
+     * @param outputPath Target file path. If omitted, generates a timestamped file in cwd.
+     * @returns The resolved output file path.
+     */
+    exportToJsonl(outputPath) {
+        const filePath = resolve(outputPath ?? `session-${new Date().toISOString().replace(/[:.]/g, "-")}.jsonl`);
+        const dir = dirname(filePath);
+        if (!existsSync(dir)) {
+            mkdirSync(dir, { recursive: true });
+        }
+        const header = {
+            type: "session",
+            version: CURRENT_SESSION_VERSION,
+            id: this.sessionManager.getSessionId(),
+            timestamp: new Date().toISOString(),
+            cwd: this.sessionManager.getCwd(),
+        };
+        const branchEntries = this.sessionManager.getBranch();
+        const lines = [JSON.stringify(header)];
+        // Re-chain parentIds to form a linear sequence
+        let prevId = null;
+        for (const entry of branchEntries) {
+            const linear = { ...entry, parentId: prevId };
+            lines.push(JSON.stringify(linear));
+            prevId = entry.id;
+        }
+        writeFileSync(filePath, `${lines.join("\n")}\n`);
+        return filePath;
+    }
+    /**
+     * Import a JSONL session file.
+     * Copies the file into the session directory and switches to it (like /resume).
+     * @param inputPath Path to the JSONL file to import.
+     * @returns true if the session was switched successfully.
+     */
+    async importFromJsonl(inputPath) {
+        const resolved = resolve(inputPath);
+        if (!existsSync(resolved)) {
+            throw new Error(`File not found: ${resolved}`);
+        }
+        // Copy into the session directory so we don't modify the original
+        const sessionDir = this.sessionManager.getSessionDir();
+        if (!existsSync(sessionDir)) {
+            mkdirSync(sessionDir, { recursive: true });
+        }
+        const destPath = join(sessionDir, basename(resolved));
+        // Avoid overwriting if source and destination are the same file
+        if (resolve(destPath) !== resolved) {
+            copyFileSync(resolved, destPath);
+        }
+        return this.switchSession(destPath);
+    }
     // =========================================================================
     // Utilities
     // =========================================================================