From 10f7436a68e6ebb283f69a68c5c041fa98679647 Mon Sep 17 00:00:00 2001 From: Oriboto Date: Mon, 16 Mar 2026 16:16:54 -0500 Subject: [PATCH] feat: agent event hooks, thinking lifecycle, and token usage for plugins - Add thinking_start/thinking_end plugin hooks with durationMs tracking - Add token usage (input/output/cache) and toolCallCount to agent_end hook - Expose onAgentEvent() in plugin API with optional sessionKey filtering - Gate onAgentEvent behind registrationMode='full' to prevent listener leaks - Emit synthetic thinking_start before orphaned thinking_end events - Track toolCallCount independently of compaction retry resets - Fire thinking hooks for tag reasoning flows (not just native events) - Hoist tag transition tracking outside if(next) gate for pure-thinking chunks --- src/agents/pi-embedded-runner/run/attempt.ts | 12 +++ ...pi-embedded-subscribe.handlers.messages.ts | 100 ++++++++++++++++-- .../pi-embedded-subscribe.handlers.tools.ts | 1 + .../pi-embedded-subscribe.handlers.types.ts | 2 + src/agents/pi-embedded-subscribe.ts | 2 + src/infra/agent-events.ts | 8 +- src/plugins/hooks.ts | 31 ++++++ src/plugins/registry.ts | 14 +++ src/plugins/types.ts | 55 ++++++++++ 9 files changed, 214 insertions(+), 11 deletions(-) diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index b02e8a59fb8..1c8aa3f84d2 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -2240,6 +2240,7 @@ export async function runEmbeddedAttempt( const { assistantTexts, toolMetas, + getTotalToolCallCount, unsubscribe, waitForCompactionRetry, isCompactionInFlight, @@ -2735,6 +2736,7 @@ export async function runEmbeddedAttempt( // This is fire-and-forget, so we don't await // Run even on compaction timeout so plugins can log/cleanup if (hookRunner?.hasHooks("agent_end")) { + const usage = getUsageTotals?.(); hookRunner .runAgentEnd( { @@ -2742,6 +2744,16 @@ export async function runEmbeddedAttempt( success: !aborted && !promptError, error: promptError ? describeUnknownError(promptError) : undefined, durationMs: Date.now() - promptStartedAt, + tokenUsage: usage + ? { + input: usage.input, + output: usage.output, + cacheRead: usage.cacheRead, + cacheWrite: usage.cacheWrite, + total: usage.total, + } + : undefined, + toolCallCount: getTotalToolCallCount(), }, { agentId: hookAgentId, diff --git a/src/agents/pi-embedded-subscribe.handlers.messages.ts b/src/agents/pi-embedded-subscribe.handlers.messages.ts index 04f47e67cde..3c377c7962c 100644 --- a/src/agents/pi-embedded-subscribe.handlers.messages.ts +++ b/src/agents/pi-embedded-subscribe.handlers.messages.ts @@ -2,6 +2,7 @@ import type { AgentEvent, AgentMessage } from "@mariozechner/pi-agent-core"; import { parseReplyDirectives } from "../auto-reply/reply/reply-directives.js"; import { SILENT_REPLY_TOKEN } from "../auto-reply/tokens.js"; import { emitAgentEvent } from "../infra/agent-events.js"; +import { getGlobalHookRunner } from "../plugins/hook-runner-global.js"; import { createInlineCodeState } from "../markdown/code-spans.js"; import { isMessagingToolDuplicateNormalized, @@ -98,7 +99,21 @@ export function handleMessageUpdate( if (evtType === "thinking_start" || evtType === "thinking_delta" || evtType === "thinking_end") { if (evtType === "thinking_start" || evtType === "thinking_delta") { - ctx.state.reasoningStreamOpen = true; + if (!ctx.state.reasoningStreamOpen) { + ctx.state.reasoningStreamOpen = true; + ctx.state.thinkingStartedAt = Date.now(); + // Emit thinking_start hook on first thinking event + const hookRunner = ctx.hookRunner ?? getGlobalHookRunner(); + if (hookRunner?.hasHooks("thinking_start")) { + void hookRunner.runThinkingStart( + { runId: ctx.params.runId }, + { + agentId: ctx.params.agentId, + sessionKey: ctx.params.sessionKey, + }, + ); + } + } } const thinkingDelta = typeof assistantRecord?.delta === "string" ? assistantRecord.delta : ""; const thinkingContent = @@ -118,10 +133,40 @@ export function handleMessageUpdate( ctx.emitReasoningStream(partialThinking || thinkingContent || thinkingDelta); } if (evtType === "thinking_end") { + // Ensure plugins always get a matching thinking_start before thinking_end if (!ctx.state.reasoningStreamOpen) { ctx.state.reasoningStreamOpen = true; + ctx.state.thinkingStartedAt = Date.now(); + const hookRunner = ctx.hookRunner ?? getGlobalHookRunner(); + if (hookRunner?.hasHooks("thinking_start")) { + void hookRunner.runThinkingStart( + { runId: ctx.params.runId }, + { + agentId: ctx.params.agentId, + sessionKey: ctx.params.sessionKey, + }, + ); + } } emitReasoningEnd(ctx); + // Emit thinking_end hook + const hookRunner = ctx.hookRunner ?? getGlobalHookRunner(); + if (hookRunner?.hasHooks("thinking_end")) { + const fullThinking = extractAssistantThinking(msg); + void hookRunner.runThinkingEnd( + { + runId: ctx.params.runId, + text: fullThinking || thinkingContent || undefined, + durationMs: ctx.state.thinkingStartedAt + ? Date.now() - ctx.state.thinkingStartedAt + : undefined, + }, + { + agentId: ctx.params.agentId, + sessionKey: ctx.params.sessionKey, + }, + ); + } } return; } @@ -183,16 +228,51 @@ export function handleMessageUpdate( inlineCode: createInlineCodeState(), }) .trim(); + + // Track tag transitions independently of visible text. + // stripBlockTags updates partialBlockState as a side effect, so we must + // call it even when `next` is empty (pure-thinking chunks). + const wasThinking = ctx.state.partialBlockState.thinking; + const visibleDelta = chunk ? ctx.stripBlockTags(chunk, ctx.state.partialBlockState) : ""; + if (!wasThinking && ctx.state.partialBlockState.thinking) { + ctx.state.reasoningStreamOpen = true; + ctx.state.thinkingStartedAt = Date.now(); + // Fire thinking_start hook for tag flows + const hookRunner = ctx.hookRunner ?? getGlobalHookRunner(); + if (hookRunner?.hasHooks("thinking_start")) { + void hookRunner.runThinkingStart( + { runId: ctx.params.runId }, + { + agentId: ctx.params.agentId, + sessionKey: ctx.params.sessionKey, + }, + ); + } + } + // Detect when thinking block ends ( tag processed) + if (wasThinking && !ctx.state.partialBlockState.thinking) { + emitReasoningEnd(ctx); + // Fire thinking_end hook for tag flows + const hookRunner = ctx.hookRunner ?? getGlobalHookRunner(); + if (hookRunner?.hasHooks("thinking_end")) { + const fullThinking = extractThinkingFromTaggedText(ctx.state.deltaBuffer); + void hookRunner.runThinkingEnd( + { + runId: ctx.params.runId, + text: fullThinking || undefined, + durationMs: ctx.state.thinkingStartedAt + ? Date.now() - ctx.state.thinkingStartedAt + : undefined, + }, + { + agentId: ctx.params.agentId, + sessionKey: ctx.params.sessionKey, + }, + ); + } + } + if (next) { - const wasThinking = ctx.state.partialBlockState.thinking; - const visibleDelta = chunk ? ctx.stripBlockTags(chunk, ctx.state.partialBlockState) : ""; - if (!wasThinking && ctx.state.partialBlockState.thinking) { - ctx.state.reasoningStreamOpen = true; - } - // Detect when thinking block ends ( tag processed) - if (wasThinking && !ctx.state.partialBlockState.thinking) { - emitReasoningEnd(ctx); - } const parsedDelta = visibleDelta ? ctx.consumePartialReplyDirectives(visibleDelta) : null; const parsedFull = parseReplyDirectives(stripTrailingDirective(next)); const cleanedText = parsedFull.text; diff --git a/src/agents/pi-embedded-subscribe.handlers.tools.ts b/src/agents/pi-embedded-subscribe.handlers.tools.ts index 70f6b54639c..074235ac78d 100644 --- a/src/agents/pi-embedded-subscribe.handlers.tools.ts +++ b/src/agents/pi-embedded-subscribe.handlers.tools.ts @@ -441,6 +441,7 @@ export async function handleToolExecutionEnd( const callSummary = ctx.state.toolMetaById.get(toolCallId); const meta = callSummary?.meta; ctx.state.toolMetas.push({ toolName, meta }); + ctx.state.totalToolCallCount++; ctx.state.toolMetaById.delete(toolCallId); ctx.state.toolSummaryById.delete(toolCallId); if (isToolError) { diff --git a/src/agents/pi-embedded-subscribe.handlers.types.ts b/src/agents/pi-embedded-subscribe.handlers.types.ts index 4436e6f6aa3..52050766a8c 100644 --- a/src/agents/pi-embedded-subscribe.handlers.types.ts +++ b/src/agents/pi-embedded-subscribe.handlers.types.ts @@ -33,6 +33,7 @@ export type ToolCallSummary = { export type EmbeddedPiSubscribeState = { assistantTexts: string[]; toolMetas: Array<{ toolName?: string; meta?: string }>; + totalToolCallCount: number; toolMetaById: Map; toolSummaryById: Set; lastToolError?: ToolErrorSummary; @@ -53,6 +54,7 @@ export type EmbeddedPiSubscribeState = { lastStreamedReasoning?: string; lastBlockReplyText?: string; reasoningStreamOpen: boolean; + thinkingStartedAt?: number; assistantMessageIndex: number; lastAssistantTextMessageIndex: number; lastAssistantTextNormalized?: string; diff --git a/src/agents/pi-embedded-subscribe.ts b/src/agents/pi-embedded-subscribe.ts index 83592372e80..5eb54b60103 100644 --- a/src/agents/pi-embedded-subscribe.ts +++ b/src/agents/pi-embedded-subscribe.ts @@ -38,6 +38,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar const state: EmbeddedPiSubscribeState = { assistantTexts: [], toolMetas: [], + totalToolCallCount: 0, toolMetaById: new Map(), toolSummaryById: new Set(), lastToolError: undefined, @@ -679,6 +680,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar return { assistantTexts, toolMetas, + getTotalToolCallCount: () => state.totalToolCallCount, unsubscribe, isCompacting: () => state.compactionInFlight || state.pendingCompactionRetry > 0, isCompactionInFlight: () => state.compactionInFlight, diff --git a/src/infra/agent-events.ts b/src/infra/agent-events.ts index 3b2e219574b..3a6f076bbb2 100644 --- a/src/infra/agent-events.ts +++ b/src/infra/agent-events.ts @@ -1,6 +1,12 @@ import type { VerboseLevel } from "../auto-reply/thinking.js"; -export type AgentEventStream = "lifecycle" | "tool" | "assistant" | "error" | (string & {}); +export type AgentEventStream = + | "lifecycle" + | "tool" + | "thinking" + | "assistant" + | "error" + | (string & {}); export type AgentEventPayload = { runId: string; diff --git a/src/plugins/hooks.ts b/src/plugins/hooks.ts index cffafd6645d..70499654cee 100644 --- a/src/plugins/hooks.ts +++ b/src/plugins/hooks.ts @@ -30,6 +30,8 @@ import type { PluginHookGatewayContext, PluginHookGatewayStartEvent, PluginHookGatewayStopEvent, + PluginHookThinkingStartEvent, + PluginHookThinkingEndEvent, PluginHookMessageContext, PluginHookMessageReceivedEvent, PluginHookMessageSendingEvent, @@ -100,6 +102,8 @@ export type { PluginHookGatewayContext, PluginHookGatewayStartEvent, PluginHookGatewayStopEvent, + PluginHookThinkingStartEvent, + PluginHookThinkingEndEvent, }; export type HookRunnerLogger = { @@ -901,6 +905,30 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp return runVoidHook("gateway_stop", event, ctx); } + /** + * Run thinking_start hook. + * Fires when the agent begins a thinking/reasoning phase. + * Runs in parallel (fire-and-forget). + */ + async function runThinkingStart( + event: PluginHookThinkingStartEvent, + ctx: PluginHookAgentContext, + ): Promise { + return runVoidHook("thinking_start", event, ctx); + } + + /** + * Run thinking_end hook. + * Fires when the agent completes a thinking/reasoning phase. + * Runs in parallel (fire-and-forget). + */ + async function runThinkingEnd( + event: PluginHookThinkingEndEvent, + ctx: PluginHookAgentContext, + ): Promise { + return runVoidHook("thinking_end", event, ctx); + } + // ========================================================================= // Utility // ========================================================================= @@ -953,6 +981,9 @@ export function createHookRunner(registry: PluginRegistry, options: HookRunnerOp // Gateway hooks runGatewayStart, runGatewayStop, + // Thinking hooks + runThinkingStart, + runThinkingEnd, // Utility hasHooks, getHookCount, diff --git a/src/plugins/registry.ts b/src/plugins/registry.ts index 9b450af26e7..d57f16f4b3c 100644 --- a/src/plugins/registry.ts +++ b/src/plugins/registry.ts @@ -3,6 +3,7 @@ import type { AnyAgentTool } from "../agents/tools/common.js"; import type { ChannelDock } from "../channels/dock.js"; import type { ChannelPlugin } from "../channels/plugins/types.js"; import { registerContextEngineForOwner } from "../context-engine/registry.js"; +import { onAgentEvent, type AgentEventPayload } from "../infra/agent-events.js"; import type { GatewayRequestHandler, GatewayRequestHandlers, @@ -851,6 +852,19 @@ export function createPluginRegistry(registryParams: PluginRegistryParams) { } }, resolvePath: (input: string) => resolveUserPath(input), + onAgentEvent: registrationMode === "full" + ? (listener: (evt: AgentEventPayload) => void, filter?: { sessionKey?: string }) => { + if (!filter?.sessionKey) { + return onAgentEvent(listener); + } + const sk = filter.sessionKey; + return onAgentEvent((evt) => { + if (evt.sessionKey === sk) { + listener(evt); + } + }); + } + : () => () => {}, on: (hookName, handler, opts) => registrationMode === "full" ? registerTypedHook(record, hookName, handler, opts, params.hookPolicy) diff --git a/src/plugins/types.ts b/src/plugins/types.ts index 6dc5788b9eb..c8912c69ea2 100644 --- a/src/plugins/types.ts +++ b/src/plugins/types.ts @@ -3,6 +3,7 @@ import type { TopLevelComponents } from "@buape/carbon"; import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { StreamFn } from "@mariozechner/pi-agent-core"; import type { Api, Model } from "@mariozechner/pi-ai"; +import type { AgentEventPayload } from "../infra/agent-events.js"; import type { ModelRegistry } from "@mariozechner/pi-coding-agent"; import type { Command } from "commander"; import type { @@ -1216,6 +1217,22 @@ export type OpenClawPluginApi = { factory: import("../context-engine/registry.js").ContextEngineFactory, ) => void; resolvePath: (input: string) => string; + /** + * Subscribe to the real-time agent event stream. + * Receives all agent events (thinking deltas, tool start/result, assistant text, + * lifecycle phases) as they happen. Returns an unsubscribe function. + * + * For high-frequency events (thinking deltas, assistant text), prefer this over + * hooks. For lifecycle moments (tool complete, agent end), prefer `on()` hooks. + * + * @param listener - Callback invoked for each event. + * @param filter - Optional filter. When `sessionKey` is provided, only events + * matching that session are delivered. Omit for the global firehose. + */ + onAgentEvent: ( + listener: (evt: AgentEventPayload) => void, + filter?: { sessionKey?: string }, + ) => () => void; /** Register a lifecycle hook handler */ on: ( hookName: K, @@ -1265,6 +1282,8 @@ export type PluginHookName = | "subagent_delivery_target" | "subagent_spawned" | "subagent_ended" + | "thinking_start" + | "thinking_end" | "gateway_start" | "gateway_stop"; @@ -1292,6 +1311,8 @@ export const PLUGIN_HOOK_NAMES = [ "subagent_delivery_target", "subagent_spawned", "subagent_ended", + "thinking_start", + "thinking_end", "gateway_start", "gateway_stop", ] as const satisfies readonly PluginHookName[]; @@ -1447,6 +1468,16 @@ export type PluginHookAgentEndEvent = { success: boolean; error?: string; durationMs?: number; + /** Token usage totals for this agent run. */ + tokenUsage?: { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + total: number; + }; + /** Number of tool calls made during this agent run. */ + toolCallCount?: number; }; // Compaction hooks @@ -1739,6 +1770,22 @@ export type PluginHookGatewayStopEvent = { reason?: string; }; +// thinking_start hook +export type PluginHookThinkingStartEvent = { + /** Stable run identifier for this agent invocation. */ + runId?: string; +}; + +// thinking_end hook +export type PluginHookThinkingEndEvent = { + /** Stable run identifier for this agent invocation. */ + runId?: string; + /** Full thinking/reasoning text. */ + text?: string; + /** Duration of the thinking phase in milliseconds. */ + durationMs?: number; +}; + // Hook handler types mapped by hook name export type PluginHookHandlerMap = { before_model_resolve: ( @@ -1841,6 +1888,14 @@ export type PluginHookHandlerMap = { event: PluginHookGatewayStopEvent, ctx: PluginHookGatewayContext, ) => Promise | void; + thinking_start: ( + event: PluginHookThinkingStartEvent, + ctx: PluginHookAgentContext, + ) => Promise | void; + thinking_end: ( + event: PluginHookThinkingEndEvent, + ctx: PluginHookAgentContext, + ) => Promise | void; }; export type PluginHookRegistration = {