Skip to content

Commit 1c65500

Browse files
obviyusobviyus
committed
fix(hooks): tighten reply usage state correlation
1 parent 618d781 commit 1c65500

10 files changed

Lines changed: 214 additions & 400 deletions

‎docs/plugins/hooks.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -425,6 +425,10 @@ even when the channel payload has no visible text/caption. Rewriting that
425425
`content` updates the hook-visible transcript only; it is not rendered as a
426426
media caption.
427427

428+
`reply_payload_sending` events may include `usageState`, a best-effort live
429+
per-turn model/usage/context snapshot. Durable delivery, recovered replay, and
430+
replies without exact run correlation omit it.
431+
428432
Message hook contexts expose stable correlation fields when available:
429433
`ctx.sessionKey`, `ctx.runId`, `ctx.messageId`, `ctx.senderId`, `ctx.trace`,
430434
`ctx.traceId`, `ctx.spanId`, `ctx.parentSpanId`, and `ctx.callDepth`. Inbound

‎src/auto-reply/dispatch.test.ts‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,8 @@ const {
7373
dispatchInboundMessageWithBufferedDispatcher,
7474
withReplyDispatcher,
7575
} = await import("./dispatch.js");
76+
const { clearReplyUsageStateForTest, recordReplyUsageState } =
77+
await import("./reply/reply-usage-state.js");
7678

7779
function createDispatcher(record: string[]): ReplyDispatcher {
7880
return {
@@ -110,6 +112,7 @@ function requireReplyDispatcherOptions(index = 0): Parameters<CreateReplyDispatc
110112
describe("withReplyDispatcher", () => {
111113
beforeEach(() => {
112114
vi.clearAllMocks();
115+
clearReplyUsageStateForTest();
113116
hoisted.finalizeInboundContextMock.mockImplementation((ctx: unknown) => ctx);
114117
hoisted.deriveInboundMessageHookContextMock.mockReturnValue({
115118
channelId: "threads",
@@ -424,6 +427,57 @@ describe("withReplyDispatcher", () => {
424427
);
425428
});
426429

430+
it("correlates reply_payload_sending usageState with the generated run id", async () => {
431+
const usageState = { provider: "openai", model: "gpt-5.5" };
432+
const runReplyPayloadSending = vi.fn(async ({ payload }: { payload: { text?: string } }) => ({
433+
payload,
434+
}));
435+
hoisted.getGlobalHookRunnerMock.mockReturnValue({
436+
hasHooks: vi.fn((hookName?: string) => hookName === "reply_payload_sending"),
437+
runMessageSending: vi.fn(async () => undefined),
438+
runReplyPayloadSending,
439+
});
440+
hoisted.createReplyDispatcherMock.mockReturnValueOnce(createDispatcher([]));
441+
hoisted.dispatchReplyFromConfigMock.mockImplementationOnce(async ({ replyOptions }) => {
442+
replyOptions?.onAgentRunStart?.("generated-run");
443+
recordReplyUsageState("generated-run", usageState);
444+
return { text: "ok" };
445+
});
446+
447+
await dispatchInboundMessageWithDispatcher({
448+
ctx: buildTestCtx({ Surface: "telegram", SessionKey: "agent:test:session" }),
449+
cfg: {} as OpenClawConfig,
450+
dispatcherOptions: {
451+
deliver: async () => undefined,
452+
},
453+
replyResolver: async () => ({ text: "ok" }),
454+
});
455+
456+
const dispatcherOptions = requireReplyDispatcherOptions();
457+
if (!dispatcherOptions?.beforeDeliver) {
458+
throw new Error("expected beforeDeliver hook");
459+
}
460+
461+
await dispatcherOptions.beforeDeliver({ text: "original reply" }, { kind: "final" });
462+
463+
expect(runReplyPayloadSending).toHaveBeenCalledWith(
464+
{
465+
payload: { text: "original reply" },
466+
kind: "final",
467+
channel: "telegram",
468+
sessionKey: "agent:test:session",
469+
runId: "generated-run",
470+
usageState,
471+
},
472+
{
473+
accountId: "acct-1",
474+
channelId: "threads",
475+
conversationId: "conv-1",
476+
runId: "generated-run",
477+
},
478+
);
479+
});
480+
427481
it("runs message_sending after reply_payload_sending for inbound dispatcher delivery", async () => {
428482
const runReplyPayloadSending = vi.fn(async ({ payload }: { payload: { text?: string } }) => ({
429483
payload: {

‎src/auto-reply/dispatch.ts‎

Lines changed: 46 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@ type ForegroundReplyFenceSnapshot = {
5252
generation: number;
5353
};
5454

55+
type ReplyPayloadRunState = {
56+
runId?: string;
57+
};
58+
5559
const foregroundReplyFenceByKey = new Map<string, ForegroundReplyFenceState>();
5660
const replyPayloadSendingDispatchers = new WeakSet<ReplyDispatcher>();
5761

@@ -349,37 +353,52 @@ function buildMessageSendingBeforeDeliver(
349353

350354
function buildReplyPayloadSendingBeforeDeliver(
351355
ctx: MsgContext | FinalizedMsgContext,
352-
opts?: { runId?: string },
356+
runState: ReplyPayloadRunState,
353357
): ReplyDispatchBeforeDeliver {
354358
const finalized = finalizeInboundContext(ctx);
355359
const hookCtx = deriveInboundMessageHookContext(finalized);
356360

357361
return async (payload: ReplyPayload, info): Promise<ReplyPayload | null> => {
362+
const runId = runState.runId;
358363
const hookedPayload = await runReplyPayloadSendingHook({
359364
payload,
360365
kind: info.kind,
361366
channel: finalized.Surface ?? finalized.Provider,
362367
sessionKey: finalized.SessionKey,
363-
runId: opts?.runId,
364-
usageState: consumeReplyUsageState(opts?.runId, finalized.SessionKey),
368+
runId,
369+
usageState: consumeReplyUsageState(runId),
365370
context: {
366371
...toPluginMessageContext(hookCtx),
367-
runId: opts?.runId,
372+
runId,
368373
},
369374
});
370375
return hookedPayload && hasOutboundReplyContent(hookedPayload) ? hookedPayload : null;
371376
};
372377
}
373378

379+
function bindReplyPayloadRunState(
380+
replyOptions: Omit<GetReplyOptions, "onBlockReply"> | undefined,
381+
runState: ReplyPayloadRunState,
382+
): Omit<GetReplyOptions, "onBlockReply"> {
383+
const onAgentRunStart = replyOptions?.onAgentRunStart;
384+
return {
385+
...replyOptions,
386+
onAgentRunStart: (runId) => {
387+
runState.runId = runId;
388+
onAgentRunStart?.(runId);
389+
},
390+
};
391+
}
392+
374393
function installReplyPayloadSendingBeforeDeliver(
375394
dispatcher: ReplyDispatcher,
376395
ctx: MsgContext | FinalizedMsgContext,
377-
opts?: { runId?: string },
396+
runState: ReplyPayloadRunState,
378397
): void {
379398
if (replyPayloadSendingDispatchers.has(dispatcher)) {
380399
return;
381400
}
382-
const beforeDeliver = buildReplyPayloadSendingBeforeDeliver(ctx, opts);
401+
const beforeDeliver = buildReplyPayloadSendingBeforeDeliver(ctx, runState);
383402
if (!beforeDeliver || !dispatcher.appendBeforeDeliver) {
384403
return;
385404
}
@@ -483,8 +502,13 @@ export async function dispatchInboundMessage(params: {
483502
replyOptions?: Omit<GetReplyOptions, "onBlockReply">;
484503
replyResolver?: GetReplyFromConfig;
485504
onSessionMetadataChanges?: (changes: CommandSessionMetadataChange[]) => void;
505+
replyPayloadRunState?: ReplyPayloadRunState;
486506
}): Promise<DispatchInboundResult> {
487507
const replyOptions = applyRuntimeToolsAllow(params.replyOptions, params.toolsAllow);
508+
const replyPayloadRunState = params.replyPayloadRunState ?? {
509+
runId: replyOptions?.runId,
510+
};
511+
const replyOptionsWithRunState = bindReplyPayloadRunState(replyOptions, replyPayloadRunState);
488512
const finalized = measureDiagnosticsTimelineSpanSync(
489513
"auto_reply.finalize_context",
490514
() => finalizeInboundContext(params.ctx),
@@ -503,9 +527,7 @@ export async function dispatchInboundMessage(params: {
503527
source: "dispatchInboundMessage",
504528
});
505529
}
506-
installReplyPayloadSendingBeforeDeliver(params.dispatcher, finalized, {
507-
runId: replyOptions?.runId,
508-
});
530+
installReplyPayloadSendingBeforeDeliver(params.dispatcher, finalized, replyPayloadRunState);
509531
const result = await withReplyDispatcher({
510532
dispatcher: params.dispatcher,
511533
run: () =>
@@ -516,7 +538,7 @@ export async function dispatchInboundMessage(params: {
516538
ctx: finalized,
517539
cfg: params.cfg,
518540
dispatcher: params.dispatcher,
519-
replyOptions,
541+
replyOptions: replyOptionsWithRunState,
520542
replyResolver: params.replyResolver,
521543
onSessionMetadataChanges: params.onSessionMetadataChanges,
522544
}),
@@ -543,9 +565,13 @@ export async function dispatchInboundMessageWithBufferedDispatcher(params: {
543565
const finalized = finalizeInboundContext(params.ctx);
544566
const foregroundReplyFence = beginForegroundReplyFence(finalized);
545567
const silentReplyContext = resolveDispatcherSilentReplyContext(finalized, params.cfg);
546-
const replyPayloadBeforeDeliver = buildReplyPayloadSendingBeforeDeliver(finalized, {
568+
const replyPayloadRunState = {
547569
runId: params.replyOptions?.runId,
548-
});
570+
};
571+
const replyPayloadBeforeDeliver = buildReplyPayloadSendingBeforeDeliver(
572+
finalized,
573+
replyPayloadRunState,
574+
);
549575
const globalBeforeDeliver = combineBeforeDeliverHooks(
550576
replyPayloadBeforeDeliver,
551577
buildMessageSendingBeforeDeliver(finalized),
@@ -605,6 +631,7 @@ export async function dispatchInboundMessageWithBufferedDispatcher(params: {
605631
...params.replyOptions,
606632
...replyOptions,
607633
},
634+
replyPayloadRunState,
608635
onSessionMetadataChanges: params.onSessionMetadataChanges,
609636
});
610637
} finally {
@@ -637,9 +664,13 @@ export async function dispatchInboundMessageWithDispatcher(params: {
637664
replyResolver?: GetReplyFromConfig;
638665
}): Promise<DispatchInboundResult> {
639666
const silentReplyContext = resolveDispatcherSilentReplyContext(params.ctx, params.cfg);
640-
const replyPayloadBeforeDeliver = buildReplyPayloadSendingBeforeDeliver(params.ctx, {
667+
const replyPayloadRunState = {
641668
runId: params.replyOptions?.runId,
642-
});
669+
};
670+
const replyPayloadBeforeDeliver = buildReplyPayloadSendingBeforeDeliver(
671+
params.ctx,
672+
replyPayloadRunState,
673+
);
643674
const globalBeforeDeliver = combineBeforeDeliverHooks(
644675
replyPayloadBeforeDeliver,
645676
buildMessageSendingBeforeDeliver(params.ctx),
@@ -660,5 +691,6 @@ export async function dispatchInboundMessageWithDispatcher(params: {
660691
toolsAllow: params.toolsAllow,
661692
replyResolver: params.replyResolver,
662693
replyOptions: params.replyOptions,
694+
replyPayloadRunState,
663695
});
664696
}

0 commit comments

Comments
 (0)