@orion-agents/orion-code 0.1.1 → 0.1.2

package/dist/runtime/chat-controller.js

@@ -40,6 +40,7 @@ exports.createToolEventPresenter = createToolEventPresenter;
 exports.loadSessionIntoRuntime = loadSessionIntoRuntime;
 exports.closeSession = closeSession;
 const commands_1 = require("../commands");
+const async_hooks_1 = require("async_hooks");
 const parser_1 = require("../commands/parser");
 const path = __importStar(require("path"));
 const session_storage_1 = require("../services/session-storage");
@@ -397,7 +398,14 @@ function appendAssistantNotice(messages, notice) {
 function recordTraceEvent(events, sessionId, event) {
     if (!sessionId)
         return null;
-    const traceEvent = (0, session_storage_1.appendSessionTraceEvent)(sessionId, event);
+    const goal = goalTraceContext.getStore();
+    const traceEvent = (0, session_storage_1.appendSessionTraceEvent)(sessionId, {
+        ...event,
+        goalId: goal?.goalId,
+        goalRevision: goal?.goalRevision,
+        goalInputKind: goal?.goalInputKind,
+        goalStopReason: goal?.getStopReason(),
+    });
     if (traceEvent) {
         events.traceEventRecorded?.(traceEvent);
     }
@@ -485,9 +493,12 @@ function structuredToolFinishActivity(event, seq, options = {}) {
     const outputBytes = event.outputBytes ?? Buffer.byteLength(durableOutput, 'utf8');
     const aggregatePresentation = (0, aggregate_tool_presenter_1.presentAggregateToolResult)(event.name, durableOutput, outputBytes);
     const aggregate = aggregatePresentation?.view;
-    const storedArtifact = event.artifactRef ?? (options.projectPath && durableOutput && (outputBytes > tool_output_presentation_1.DEFAULT_TOOL_OUTPUT_POLICY.inlineMaxBytes || aggregate)
-        ? (0, tool_artifacts_1.storeArtifact)(options.projectPath, event.name, durableOutput, outputBytes) ?? undefined
-        : undefined);
+    const storedArtifact = event.artifactRef ??
+        (options.projectPath &&
+            durableOutput &&
+            (outputBytes > tool_output_presentation_1.DEFAULT_TOOL_OUTPUT_POLICY.inlineMaxBytes || aggregate)
+            ? ((0, tool_artifacts_1.storeArtifact)(options.projectPath, event.name, durableOutput, outputBytes) ?? undefined)
+            : undefined);
     const artifactRef = storedArtifact
         ? { id: storedArtifact.id, outputBytes: storedArtifact.outputBytes }
         : undefined;
@@ -645,13 +656,13 @@ function sessionMessagesToTranscriptEntries(sessionId, options = {}) {
                 : durable;
             const durableOutput = typeof durable.output === 'string' ? durable.output : '';
             const displayOutput = typeof modelVisible.output === 'string' ? modelVisible.output : durableOutput;
-            const outputBytes = trace?.outputBytes
-                ?? (durable.schemaVersion === 1 ? durable.outputBytes : undefined)
-                ?? Buffer.byteLength(durableOutput, 'utf8');
+            const outputBytes = trace?.outputBytes ??
+                (durable.schemaVersion === 1 ? durable.outputBytes : undefined) ??
+                Buffer.byteLength(durableOutput, 'utf8');
             const artifactId = durable.artifactRef?.id ?? trace?.artifactId;
             fallbackToolSequence += 1;
             const sequence = message.toolCallId
-                ? sequenceByCallId.get(message.toolCallId) ?? fallbackToolSequence
+                ? (sequenceByCallId.get(message.toolCallId) ?? fallbackToolSequence)
                 : fallbackToolSequence;
             const outputView = options.includeToolOutputViews && call && message.toolCallId
                 ? (0, tool_output_presentation_1.createToolOutputView)({
@@ -764,6 +775,7 @@ function createAssistantStreamPresenter(events, abortSignal) {
         },
     };
 }
+const goalTraceContext = new async_hooks_1.AsyncLocalStorage();
 class LocalFastPathBlockedError extends Error {
     constructor(message) {
         super(message);
@@ -902,6 +914,7 @@ function createToolEventPresenter(events, options = {}) {
                         outputBytes: toolActivity.outputView.detailRef.outputBytes,
                     }
                     : event.artifactRef,
+                externalAssertion: event.externalAssertion,
                 sequence: seq,
                 batchCount: event.batchCount,
                 batchIndex: event.batchIndex,
@@ -979,6 +992,14 @@ class AgentChatController {
         const text = input.trim();
         if (!text)
             return;
+        // Every root input owns a fresh accounting snapshot. Commands and local
+        // fast paths are known zero-provider turns; runChat changes tokenUsage to
+        // null before networking so missing provider usage remains fail-closed.
+        const storeWithState = this.runtime.store;
+        storeWithState.setState?.({
+            tokenUsage: { promptTokens: 0, completionTokens: 0 },
+            lastLoopStats: undefined,
+        });
         const parsed = (0, parser_1.parseInput)(text);
         if (!parsed.isCommand) {
             const localFastPath = parseLocalFastPath(text);
@@ -1022,7 +1043,67 @@ class AgentChatController {
             return;
         }
         const ctx = this.createCommandContext(options.abortSignal, options.turnId);
-        const { result, output } = await captureConsoleOutput(() => command.execute(ctx, parsed.args));
+        let commandProviderAttempts = 0;
+        const commandProviderPreflight = async (context) => {
+            commandProviderAttempts += 1;
+            return this.controllerOptions.beforeProviderRequest
+                ? this.controllerOptions.beforeProviderRequest(context)
+                : { available: true };
+        };
+        const restoreProviderPreflight = typeof this.runtime.llm?.setProviderRequestPreflight === 'function'
+            ? this.runtime.llm.setProviderRequestPreflight(commandProviderPreflight)
+            : undefined;
+        const commandUsage = { promptTokens: 0, completionTokens: 0 };
+        let commandUsageEvents = 0;
+        const unsubscribeCommandUsage = typeof this.runtime.llm?.subscribeUsage === 'function'
+            ? this.runtime.llm.subscribeUsage(event => {
+                commandUsageEvents += 1;
+                commandUsage.promptTokens += event.usage.promptTokens;
+                commandUsage.completionTokens += event.usage.completionTokens;
+            })
+            : undefined;
+        let result;
+        let output;
+        try {
+            ({ result, output } = await captureConsoleOutput(() => command.execute(ctx, parsed.args)));
+        }
+        finally {
+            unsubscribeCommandUsage?.();
+            restoreProviderPreflight?.();
+            if (commandUsage.promptTokens > 0 || commandUsage.completionTokens > 0) {
+                this.runtime.store.setTokenUsage(commandUsage);
+            }
+            const providerRequests = Math.max(commandProviderAttempts, commandUsageEvents);
+            if (providerRequests > 0) {
+                const diagnostics = getLastRequestDiagnostics(this.runtime.llm);
+                const retryCount = diagnostics?.retryCount ?? 0;
+                const fallbackCount = diagnostics?.fallbackTriggered ? 1 : 0;
+                this.setLoopStats({
+                    turnsStarted: 1,
+                    llmRequests: providerRequests,
+                    toolCalls: 0,
+                    readOnlyToolCalls: 0,
+                    unsafeToolCalls: 0,
+                    toolResultBytes: 0,
+                    modelVisibleToolBytes: 0,
+                    summarizedBytes: 0,
+                    finishReason: commandUsageEvents > 0 ? 'completed' : 'failed',
+                    providerRetryCount: retryCount,
+                    providerRetryDelayMs: diagnostics?.retryDelayMs,
+                    providerRetryErrorTypes: diagnostics?.retryErrorTypes,
+                    providerLastRetryErrorType: diagnostics?.lastRetryErrorType,
+                    providerLastRetryStatus: diagnostics?.lastRetryStatus,
+                    providerFallbackCount: fallbackCount,
+                    providerFallbackFromModel: diagnostics?.fallbackFromModel,
+                    providerFallbackToModel: diagnostics?.fallbackToModel,
+                    providerFinalModel: diagnostics?.finalModel,
+                    providerUsingFallback: diagnostics?.usingFallback,
+                    singleReadOnlyStreak: 0,
+                    batchReadSuggestionCount: 0,
+                    localFastPathUsed: false,
+                });
+            }
+        }
         if (output) {
             this.events.append({
                 role: result.success ? 'system' : 'error',
@@ -1057,6 +1138,25 @@ class AgentChatController {
             await this.runChat(result.chatInput ?? parsed.args, options);
         }
     }
+    async runRequest(request, options = {}) {
+        const internalInput = request.text?.trim() ||
+            'Continue pursuing the active goal from its persisted plan and evidence.';
+        const run = () => this.runInput(internalInput, {
+            ...options,
+            persistAsUserMessage: request.persistAsUserMessage,
+            inputKind: request.inputKind,
+        });
+        if (!request.goal) {
+            await run();
+            return;
+        }
+        await goalTraceContext.run({
+            goalId: request.goal.goalId,
+            goalRevision: request.goal.revision,
+            goalInputKind: request.inputKind,
+            getStopReason: () => this.goalCoordinator?.goal?.stopReason?.message,
+        }, run);
+    }
     async runLocalFastPath(input, action, options = {}) {
         const activeSession = this.runtime.getSession() ??
             this.runtime.ensureSession() ??
@@ -1088,7 +1188,9 @@ class AgentChatController {
                 note: compactMiddle(action.label, 160),
             });
         }
-        this.runtime.store.addMessage({ role: 'user', content: input });
+        if (options.persistAsUserMessage !== false) {
+            this.runtime.store.addMessage({ role: 'user', content: input });
+        }
         this.events.setStatus(`Running local ${action.label}...`);
         try {
             const tool = (0, tools_1.getRuntimeTools)().find(candidate => candidate.name === action.tool);
@@ -1253,9 +1355,7 @@ class AgentChatController {
             ensureSession: this.runtime.ensureSession,
             setSession: session => {
                 this.runtime.setSession(session);
-                const renderer = this.controllerOptions.uiRenderer
-                    ?? this.runtime.config.ui?.renderer
-                    ?? 'terminal';
+                const renderer = this.controllerOptions.uiRenderer ?? this.runtime.config.ui?.renderer ?? 'terminal';
                 this.events.replaceTranscript(sessionMessagesToTranscriptEntries(session.id, {
                     includeToolOutputViews: renderer === 'tui',
                 }));
@@ -1349,7 +1449,7 @@ class AgentChatController {
         if (!bundle)
             return stats;
         const usage = bundle.getAggregateUsage();
-        if (usage.modelRequests === 0 && usage.toolCalls === 0)
+        if (usage.modelRequests === 0 && usage.toolCalls === 0 && usage.usageComplete === true)
             return stats;
         const subtaskCount = bundle.getSubtaskCount();
         const subagentSuffix = `subagents: ${usage.modelRequests} req/${usage.toolCalls} calls across ${subtaskCount} task(s)`;
@@ -1358,6 +1458,11 @@ class AgentChatController {
             llmRequests: (stats.llmRequests ?? 0) + usage.modelRequests,
             toolCalls: (stats.toolCalls ?? 0) + usage.toolCalls,
             readOnlyToolCalls: (stats.readOnlyToolCalls ?? 0) + usage.toolCalls,
+            subagentPromptTokens: usage.promptTokens,
+            subagentCompletionTokens: usage.completionTokens,
+            subagentTotalTokens: usage.promptTokens + usage.completionTokens,
+            usageAccountingComplete: stats.usageAccountingComplete !== false &&
+                usage.usageComplete === true,
             continuationHint: stats.continuationHint
                 ? `${stats.continuationHint} (${subagentSuffix})`
                 : subagentSuffix,
@@ -1380,6 +1485,9 @@ class AgentChatController {
             });
             return;
         }
+        // Usage belongs to a single root turn. Clear the previous snapshot before
+        // any provider request so failures without usage cannot inherit stale data.
+        this.runtime.store.setState({ tokenUsage: null });
         const abortSignal = options.abortSignal;
         const turnId = traceTurnId(options.turnId);
         const activeSession = this.runtime.getSession() ??
@@ -1396,7 +1504,7 @@ class AgentChatController {
             sessionId,
         });
         const appliedSkillNames = skillResolution.skills.map(skill => skill.name);
-        if (sessionId) {
+        if (sessionId && options.persistAsUserMessage !== false) {
             (0, session_storage_1.appendSessionMessage)(sessionId, {
                 role: 'user',
                 content: input,
@@ -1411,7 +1519,9 @@ class AgentChatController {
             });
             appendWorkspaceSnapshotTrace(this.events, sessionId, turnId, 'pre_turn', preWorkspace);
         }
-        this.runtime.store.addMessage({ role: 'user', content: input });
+        if (options.persistAsUserMessage !== false) {
+            this.runtime.store.addMessage({ role: 'user', content: input });
+        }
         (0, prompt_context_1.refreshProjectInstructions)(this.runtime.store, this.runtime.cwd);
         const snapshot = this.runtime.store.getSnapshot();
         const harness = (0, harness_1.createContextHarness)({
@@ -1424,7 +1534,9 @@ class AgentChatController {
                 completionGate: true,
             },
         });
-        const intent = harness.updateContractFromUserInput(input);
+        const intent = harness.updateContractFromUserInput(options.inputKind === 'goal_continuation'
+            ? (this.goalCoordinator?.goal?.objective ?? input)
+            : input);
         harness.recordAppliedSkills(skillResolution.skills);
         // Reconcile diagnostic: when harness state is present but objective may be incomplete
         if (snapshot.harnessState &&
@@ -1480,6 +1592,7 @@ class AgentChatController {
                 // R6: wire child usage callback so CostTracker records subagent
                 // token consumption. Injected by AgentRuntimeController via chatOptions.
                 onChildUsage: this.controllerOptions.onChildUsage,
+                beforeProviderRequest: this.controllerOptions.beforeProviderRequest,
             })
             : null;
         const subtaskTool = subagentBundle?.tool ?? null;
@@ -1503,6 +1616,9 @@ class AgentChatController {
             { role: 'system', content: systemPrompt.static, cacheControl: { type: 'ephemeral' } },
             ...(systemPrompt.dynamic ? [{ role: 'system', content: systemPrompt.dynamic }] : []),
             ...snapshot.conversationHistory,
+            ...(options.persistAsUserMessage === false
+                ? [{ role: 'user', content: input }]
+                : []),
         ];
         let finalContent = '';
         let finalUsage;
@@ -1563,6 +1679,26 @@ class AgentChatController {
         let observedUnsafeToolCalls = 0;
         let observedToolResultBytes = 0;
         let observedModelVisibleToolBytes = 0;
+        const restoreProviderPreflight = typeof this.runtime.llm.setProviderRequestPreflight === 'function'
+            ? this.runtime.llm.setProviderRequestPreflight(this.controllerOptions.beforeProviderRequest)
+            : undefined;
+        const persistInterruptedAccounting = (rootUsage, stats, usageAccountingComplete) => {
+            const accountingStats = {
+                ...stats,
+                finishReason: 'cancelled',
+                usageAccountingComplete,
+            };
+            const interruptedStats = this.foldSubagentUsage(accountingStats, subagentBundle);
+            this.setLoopStats(interruptedStats);
+            const subagentUsage = subagentBundle?.getAggregateUsage();
+            const knownZeroRootUsage = usageAccountingComplete && interruptedStats.llmRequests === 0;
+            if (rootUsage || knownZeroRootUsage) {
+                this.runtime.store.setTokenUsage({
+                    promptTokens: (rootUsage?.promptTokens ?? 0) + (subagentUsage?.promptTokens ?? 0),
+                    completionTokens: (rootUsage?.completionTokens ?? 0) + (subagentUsage?.completionTokens ?? 0),
+                });
+            }
+        };
         try {
             for await (const event of (0, framework_1.query)({
                 messages,
@@ -1861,6 +1997,19 @@ class AgentChatController {
             }
             const wasAborted = abortSignal?.aborted === true;
             if (wasAborted) {
+                const stats = pendingCompleteStats ??
+                    (0, framework_1.createFailedLoopStats)({
+                        loopBudget,
+                        diagnostics: observedLlmRequests > 0 ? getLastRequestDiagnostics(this.runtime.llm) : undefined,
+                        turnsStarted: observedTurnsStarted,
+                        llmRequests: observedLlmRequests,
+                        toolCalls: observedToolCalls,
+                        readOnlyToolCalls: observedReadOnlyToolCalls,
+                        unsafeToolCalls: observedUnsafeToolCalls,
+                        toolResultBytes: observedToolResultBytes,
+                        modelVisibleToolBytes: observedModelVisibleToolBytes,
+                    });
+                persistInterruptedAccounting(finalUsage, stats, finalUsage !== undefined || stats.llmRequests === 0);
                 assistantStream.discardSegment();
                 this.events.setStatus('Interrupted.');
                 removeTrailingUserMessage(this.runtime);
@@ -2087,10 +2236,30 @@ class AgentChatController {
                         gc.deferContinuation();
                     }
                 }
-                catch { /* best effort */ }
+                catch {
+                    /* best effort */
+                }
                 return;
             }
             if (isAbortError(error, abortSignal)) {
+                const interruptedStats = error instanceof framework_1.QueryLoopError
+                    ? error.stats
+                    : (0, framework_1.createFailedLoopStats)({
+                        loopBudget,
+                        diagnostics: observedLlmRequests > 0 ? getLastRequestDiagnostics(this.runtime.llm) : undefined,
+                        turnsStarted: observedTurnsStarted,
+                        llmRequests: observedLlmRequests,
+                        toolCalls: observedToolCalls,
+                        readOnlyToolCalls: observedReadOnlyToolCalls,
+                        unsafeToolCalls: observedUnsafeToolCalls,
+                        toolResultBytes: observedToolResultBytes,
+                        modelVisibleToolBytes: observedModelVisibleToolBytes,
+                    });
+                const interruptedRootUsage = error instanceof framework_1.QueryLoopError ? error.aggregateUsage : undefined;
+                // A thrown abort after a provider request started can only prove a
+                // lower bound: the in-flight request may be billable even when its
+                // final usage chunk never arrived.
+                persistInterruptedAccounting(interruptedRootUsage, interruptedStats, interruptedStats.llmRequests === 0);
                 assistantStream.discardSegment();
                 this.events.setStatus('Interrupted.');
                 removeTrailingUserMessage(this.runtime);
@@ -2123,7 +2292,7 @@ class AgentChatController {
                 errorLayer: errorLayerForChatError(error),
             });
             this.events.setStatus('Turn failed. Ready for the next input.');
-            const failedStats = error instanceof framework_1.QueryLoopError
+            const failedStats = this.foldSubagentUsage(error instanceof framework_1.QueryLoopError
                 ? error.stats
                 : (0, framework_1.createFailedLoopStats)({
                     loopBudget,
@@ -2135,7 +2304,16 @@ class AgentChatController {
                     unsafeToolCalls: observedUnsafeToolCalls,
                     toolResultBytes: observedToolResultBytes,
                     modelVisibleToolBytes: observedModelVisibleToolBytes,
+                }), subagentBundle);
+            const rootUsage = error instanceof framework_1.QueryLoopError ? error.aggregateUsage : undefined;
+            const subagentUsage = subagentBundle?.getAggregateUsage();
+            const hasSubagentUsage = (subagentUsage?.promptTokens ?? 0) > 0 || (subagentUsage?.completionTokens ?? 0) > 0;
+            if (rootUsage || hasSubagentUsage) {
+                this.runtime.store.setTokenUsage({
+                    promptTokens: (rootUsage?.promptTokens ?? 0) + (subagentUsage?.promptTokens ?? 0),
+                    completionTokens: (rootUsage?.completionTokens ?? 0) + (subagentUsage?.completionTokens ?? 0),
                 });
+            }
             this.setLoopStats(failedStats);
             if (sessionId) {
                 recordProviderTraceEvents(this.events, sessionId, turnId, failedStats);
@@ -2183,6 +2361,9 @@ class AgentChatController {
                 this.runtime.store.setState({ conversationHistory: history.slice(0, -1) });
             }
         }
+        finally {
+            restoreProviderPreflight?.();
+        }
     }
 }
 exports.AgentChatController = AgentChatController;