Skip to content

Commit 22b6034

Browse files
Oribotocarlos-sarmiento
Oriboto
authored and
carlos-sarmiento
committed
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 <think> tag reasoning flows (not just native events) - Hoist tag transition tracking outside if(next) gate for pure-thinking chunks
1 parent 6ab6e7a commit 22b6034

17 files changed

Lines changed: 665 additions & 197 deletions

‎docs/.generated/plugin-sdk-api-baseline.json‎

Lines changed: 86 additions & 86 deletions
Large diffs are not rendered by default.

‎docs/.generated/plugin-sdk-api-baseline.jsonl‎

Lines changed: 86 additions & 86 deletions
Large diffs are not rendered by default.

‎extensions/lobster/src/lobster-tool.test.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ function fakeApi(overrides: Partial<OpenClawPluginApi> = {}): OpenClawPluginApi
5656
registerContextEngine() {},
5757
registerMemoryPromptSection() {},
5858
on() {},
59+
onAgentEvent() {
60+
return () => {};
61+
},
5962
resolvePath: (p) => p,
6063
...overrides,
6164
};

‎src/agents/pi-embedded-runner/run/attempt.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2591,6 +2591,7 @@ export async function runEmbeddedAttempt(
25912591
const {
25922592
assistantTexts,
25932593
toolMetas,
2594+
getTotalToolCallCount,
25942595
unsubscribe,
25952596
waitForCompactionRetry,
25962597
isCompactionInFlight,
@@ -3072,13 +3073,24 @@ export async function runEmbeddedAttempt(
30723073
// This is fire-and-forget, so we don't await
30733074
// Run even on compaction timeout so plugins can log/cleanup
30743075
if (hookRunner?.hasHooks("agent_end")) {
3076+
const usage = getUsageTotals?.();
30753077
hookRunner
30763078
.runAgentEnd(
30773079
{
30783080
messages: messagesSnapshot,
30793081
success: !aborted && !promptError,
30803082
error: promptError ? describeUnknownError(promptError) : undefined,
30813083
durationMs: Date.now() - promptStartedAt,
3084+
tokenUsage: usage
3085+
? {
3086+
input: usage.input ?? 0,
3087+
output: usage.output ?? 0,
3088+
cacheRead: usage.cacheRead ?? 0,
3089+
cacheWrite: usage.cacheWrite ?? 0,
3090+
total: usage.total ?? 0,
3091+
}
3092+
: undefined,
3093+
toolCallCount: getTotalToolCallCount(),
30823094
},
30833095
{
30843096
agentId: hookAgentId,

‎src/agents/pi-embedded-subscribe.handlers.messages.ts‎

Lines changed: 152 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { parseReplyDirectives } from "../auto-reply/reply/reply-directives.js";
44
import { SILENT_REPLY_TOKEN } from "../auto-reply/tokens.js";
55
import { emitAgentEvent } from "../infra/agent-events.js";
66
import { createInlineCodeState } from "../markdown/code-spans.js";
7+
import { getGlobalHookRunner } from "../plugins/hook-runner-global.js";
78
import {
89
isMessagingToolDuplicateNormalized,
910
normalizeTextForComparison,
@@ -16,6 +17,7 @@ import type {
1617
import { appendRawStream } from "./pi-embedded-subscribe.raw-stream.js";
1718
import {
1819
extractAssistantText,
20+
extractAssistantTextWithThinkingTags,
1921
extractAssistantThinking,
2022
extractThinkingFromTaggedStream,
2123
extractThinkingFromTaggedText,
@@ -52,6 +54,8 @@ function emitReasoningEnd(ctx: EmbeddedPiSubscribeContext) {
5254
return;
5355
}
5456
ctx.state.reasoningStreamOpen = false;
57+
ctx.state.partialBlockState.thinking = false;
58+
(ctx.state.partialBlockState as { thinkingTransitioned?: boolean }).thinkingTransitioned = false;
5559
void ctx.params.onReasoningEnd?.();
5660
}
5761

@@ -179,8 +183,25 @@ export function handleMessageUpdate(
179183
const evtType = typeof assistantRecord?.type === "string" ? assistantRecord.type : "";
180184

181185
if (evtType === "thinking_start" || evtType === "thinking_delta" || evtType === "thinking_end") {
186+
// Mark that native thinking events are driving the lifecycle, so tagged
187+
// <think> transitions in text deltas don't duplicate hooks.
188+
ctx.state.nativeThinkingActive = true;
182189
if (evtType === "thinking_start" || evtType === "thinking_delta") {
183-
ctx.state.reasoningStreamOpen = true;
190+
if (!ctx.state.reasoningStreamOpen) {
191+
ctx.state.reasoningStreamOpen = true;
192+
ctx.state.thinkingStartedAt = Date.now();
193+
// Emit thinking_start hook on first thinking event
194+
const hookRunner = ctx.hookRunner ?? getGlobalHookRunner();
195+
if (hookRunner?.hasHooks("thinking_start")) {
196+
void hookRunner.runThinkingStart(
197+
{ runId: ctx.params.runId },
198+
{
199+
agentId: ctx.params.agentId,
200+
sessionKey: ctx.params.sessionKey,
201+
},
202+
);
203+
}
204+
}
184205
}
185206
const thinkingDelta = typeof assistantRecord?.delta === "string" ? assistantRecord.delta : "";
186207
const thinkingContent =
@@ -194,16 +215,46 @@ export function handleMessageUpdate(
194215
delta: thinkingDelta,
195216
content: thinkingContent,
196217
});
197-
if (ctx.state.streamReasoning) {
218+
if (ctx.state.streamReasoning || ctx.params.onAgentEvent) {
198219
// Prefer full partial-message thinking when available; fall back to event payloads.
199220
const partialThinking = extractAssistantThinking(msg);
200221
ctx.emitReasoningStream(partialThinking || thinkingContent || thinkingDelta);
201222
}
202223
if (evtType === "thinking_end") {
224+
// Ensure plugins always get a matching thinking_start before thinking_end
203225
if (!ctx.state.reasoningStreamOpen) {
204226
ctx.state.reasoningStreamOpen = true;
227+
ctx.state.thinkingStartedAt = Date.now();
228+
const hookRunner = ctx.hookRunner ?? getGlobalHookRunner();
229+
if (hookRunner?.hasHooks("thinking_start")) {
230+
void hookRunner.runThinkingStart(
231+
{ runId: ctx.params.runId },
232+
{
233+
agentId: ctx.params.agentId,
234+
sessionKey: ctx.params.sessionKey,
235+
},
236+
);
237+
}
205238
}
206239
emitReasoningEnd(ctx);
240+
// Emit thinking_end hook
241+
const hookRunner = ctx.hookRunner ?? getGlobalHookRunner();
242+
if (hookRunner?.hasHooks("thinking_end")) {
243+
const fullThinking = extractAssistantThinking(msg);
244+
void hookRunner.runThinkingEnd(
245+
{
246+
runId: ctx.params.runId,
247+
text: fullThinking || thinkingContent || undefined,
248+
durationMs: ctx.state.thinkingStartedAt
249+
? Date.now() - ctx.state.thinkingStartedAt
250+
: undefined,
251+
},
252+
{
253+
agentId: ctx.params.agentId,
254+
sessionKey: ctx.params.sessionKey,
255+
},
256+
);
257+
}
207258
}
208259
return;
209260
}
@@ -253,7 +304,7 @@ export function handleMessageUpdate(
253304
}
254305
}
255306

256-
if (ctx.state.streamReasoning) {
307+
if (ctx.state.streamReasoning || ctx.params.onAgentEvent) {
257308
// Handle partial <think> tags: stream whatever reasoning is visible so far.
258309
ctx.emitReasoningStream(extractThinkingFromTaggedStream(ctx.state.deltaBuffer));
259310
}
@@ -265,16 +316,78 @@ export function handleMessageUpdate(
265316
inlineCode: createInlineCodeState(),
266317
})
267318
.trim();
268-
if (next) {
269-
const wasThinking = ctx.state.partialBlockState.thinking;
270-
const visibleDelta = chunk ? ctx.stripBlockTags(chunk, ctx.state.partialBlockState) : "";
271-
if (!wasThinking && ctx.state.partialBlockState.thinking) {
319+
320+
// Track <think> tag transitions independently of visible text.
321+
// stripBlockTags updates partialBlockState as a side effect, so we must
322+
// call it even when `next` is empty (pure-thinking chunks).
323+
const wasThinking = ctx.state.partialBlockState.thinking;
324+
if (!chunk) {
325+
(ctx.state.partialBlockState as { thinkingTransitioned?: boolean }).thinkingTransitioned =
326+
false;
327+
}
328+
const visibleDelta = chunk ? ctx.stripBlockTags(chunk, ctx.state.partialBlockState) : "";
329+
330+
// Handle single-chunk transitions: if the chunk contained both <think> and </think>,
331+
// partialBlockState.thinking ends as false. Detect this by checking whether
332+
// stripBlockTags consumed any thinking content even though final state is not-thinking.
333+
const nowThinking = ctx.state.partialBlockState.thinking;
334+
const enteredThinking = !wasThinking && nowThinking;
335+
const exitedThinking = wasThinking && !nowThinking;
336+
// Single-chunk: was not thinking, is not thinking, but stripBlockTags found and consumed
337+
// a complete thinking block (open+close in one delta). Uses the thinkingTransitioned flag
338+
// set by stripBlockTags, which correctly respects code-span filtering.
339+
const singleChunkThinking =
340+
!wasThinking &&
341+
!nowThinking &&
342+
!!(ctx.state.partialBlockState as { thinkingTransitioned?: boolean }).thinkingTransitioned;
343+
344+
// Skip tagged transitions entirely when native thinking events already manage the lifecycle.
345+
// This prevents duplicated thinking_start/thinking_end when providers mirror reasoning in
346+
// both native events and <think> text tags simultaneously.
347+
const nativeThinkingActive = ctx.state.nativeThinkingActive ?? false;
348+
349+
if ((enteredThinking || singleChunkThinking) && !nativeThinkingActive) {
350+
// Only fire tagged thinking_start if native thinking hasn't already opened it
351+
if (!ctx.state.reasoningStreamOpen) {
272352
ctx.state.reasoningStreamOpen = true;
353+
ctx.state.thinkingStartedAt = Date.now();
354+
// Fire thinking_start hook for <think> tag flows
355+
const hookRunner = ctx.hookRunner ?? getGlobalHookRunner();
356+
if (hookRunner?.hasHooks("thinking_start")) {
357+
void hookRunner.runThinkingStart(
358+
{ runId: ctx.params.runId },
359+
{
360+
agentId: ctx.params.agentId,
361+
sessionKey: ctx.params.sessionKey,
362+
},
363+
);
364+
}
273365
}
274-
// Detect when thinking block ends (</think> tag processed)
275-
if (wasThinking && !ctx.state.partialBlockState.thinking) {
276-
emitReasoningEnd(ctx);
366+
}
367+
// Detect when thinking block ends (</think> tag processed)
368+
if ((exitedThinking || singleChunkThinking) && !nativeThinkingActive) {
369+
emitReasoningEnd(ctx);
370+
// Fire thinking_end hook for <think> tag flows
371+
const hookRunner = ctx.hookRunner ?? getGlobalHookRunner();
372+
if (hookRunner?.hasHooks("thinking_end")) {
373+
const fullThinking = extractThinkingFromTaggedText(ctx.state.deltaBuffer);
374+
void hookRunner.runThinkingEnd(
375+
{
376+
runId: ctx.params.runId,
377+
text: fullThinking || undefined,
378+
durationMs: ctx.state.thinkingStartedAt
379+
? Date.now() - ctx.state.thinkingStartedAt
380+
: undefined,
381+
},
382+
{
383+
agentId: ctx.params.agentId,
384+
sessionKey: ctx.params.sessionKey,
385+
},
386+
);
277387
}
388+
}
389+
390+
if (next) {
278391
const parsedDelta = visibleDelta ? ctx.consumePartialReplyDirectives(visibleDelta) : null;
279392
const parsedFull = parseReplyDirectives(stripTrailingDirective(next));
280393
const cleanedText = parsedFull.text;
@@ -339,6 +452,33 @@ export function handleMessageEnd(
339452
const assistantMessage = msg;
340453
ctx.noteLastAssistant(assistantMessage);
341454
ctx.recordAssistantUsage((assistantMessage as { usage?: unknown }).usage);
455+
456+
// Safety net: if thinking_start was fired but no thinking_end arrived (provider
457+
// skipped the explicit end event), close the thinking lifecycle now so plugins
458+
// don't get stuck with an unclosed thinking indicator.
459+
if (ctx.state.reasoningStreamOpen) {
460+
emitReasoningEnd(ctx);
461+
const hookRunner = ctx.hookRunner ?? getGlobalHookRunner();
462+
if (hookRunner?.hasHooks("thinking_end")) {
463+
const fullThinking =
464+
extractAssistantThinking(msg) ||
465+
extractThinkingFromTaggedText(extractAssistantTextWithThinkingTags(msg));
466+
void hookRunner.runThinkingEnd(
467+
{
468+
runId: ctx.params.runId,
469+
text: fullThinking || undefined,
470+
durationMs: ctx.state.thinkingStartedAt
471+
? Date.now() - ctx.state.thinkingStartedAt
472+
: undefined,
473+
},
474+
{
475+
agentId: ctx.params.agentId,
476+
sessionKey: ctx.params.sessionKey,
477+
},
478+
);
479+
}
480+
}
481+
342482
if (ctx.state.deterministicApprovalPromptSent) {
343483
return;
344484
}
@@ -359,7 +499,7 @@ export function handleMessageEnd(
359499
messagingToolSentTexts: ctx.state.messagingToolSentTexts,
360500
});
361501
const rawThinking =
362-
ctx.state.includeReasoning || ctx.state.streamReasoning
502+
ctx.state.includeReasoning || ctx.state.streamReasoning || Boolean(ctx.params.onAgentEvent)
363503
? extractAssistantThinking(assistantMessage) || extractThinkingFromTaggedText(rawText)
364504
: "";
365505
const formattedReasoning = rawThinking ? formatReasoningMessage(rawThinking) : "";
@@ -490,7 +630,7 @@ export function handleMessageEnd(
490630
if (!shouldEmitReasoningBeforeAnswer) {
491631
maybeEmitReasoning();
492632
}
493-
if (ctx.state.streamReasoning && rawThinking) {
633+
if ((ctx.state.streamReasoning || ctx.params.onAgentEvent) && rawThinking) {
494634
ctx.emitReasoningStream(rawThinking);
495635
}
496636

‎src/agents/pi-embedded-subscribe.handlers.tools.test.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ function createTestContext(): {
4848
messagingToolSentTargets: [],
4949
successfulCronAdds: 0,
5050
deterministicApprovalPromptSent: false,
51+
totalToolCallCount: 0,
5152
},
5253
shouldEmitToolResult: () => false,
5354
shouldEmitToolOutput: () => false,

‎src/agents/pi-embedded-subscribe.handlers.tools.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -472,6 +472,7 @@ export async function handleToolExecutionEnd(
472472
const callSummary = ctx.state.toolMetaById.get(toolCallId);
473473
const meta = callSummary?.meta;
474474
ctx.state.toolMetas.push({ toolName, meta });
475+
ctx.state.totalToolCallCount++;
475476
ctx.state.toolMetaById.delete(toolCallId);
476477
ctx.state.toolSummaryById.delete(toolCallId);
477478
if (isToolError) {

‎src/agents/pi-embedded-subscribe.handlers.types.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ export type ToolCallSummary = {
3434
export type EmbeddedPiSubscribeState = {
3535
assistantTexts: string[];
3636
toolMetas: Array<{ toolName?: string; meta?: string }>;
37+
totalToolCallCount: number;
3738
toolMetaById: Map<string, ToolCallSummary>;
3839
toolSummaryById: Set<string>;
3940
lastToolError?: ToolErrorSummary;
@@ -54,6 +55,10 @@ export type EmbeddedPiSubscribeState = {
5455
lastStreamedReasoning?: string;
5556
lastBlockReplyText?: string;
5657
reasoningStreamOpen: boolean;
58+
/** True once native thinking_start/thinking_delta/thinking_end events arrive,
59+
* used to suppress duplicate hooks from mirrored <think> tags in text deltas. */
60+
nativeThinkingActive?: boolean;
61+
thinkingStartedAt?: number;
5762
assistantMessageIndex: number;
5863
lastAssistantTextMessageIndex: number;
5964
lastAssistantTextNormalized?: string;
@@ -163,6 +168,7 @@ export type ToolHandlerState = Pick<
163168
| "messagingToolSentTargets"
164169
| "successfulCronAdds"
165170
| "deterministicApprovalPromptSent"
171+
| "totalToolCallCount"
166172
>;
167173

168174
export type ToolHandlerContext = {

0 commit comments

Comments
 (0)