Skip to content

Commit d5e25e0

Browse files
steipetesteipete
committed
refactor: centralize dispatcher lifecycle ownership
1 parent 5caf829 commit d5e25e0

6 files changed

Lines changed: 107 additions & 88 deletions

File tree

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

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import { describe, expect, it, vi } from "vitest";
2+
import type { OpenClawConfig } from "../config/config.js";
23
import type { ReplyDispatcher } from "./reply/reply-dispatcher.js";
3-
import { withReplyDispatcher } from "./dispatch.js";
4+
import { dispatchInboundMessage, withReplyDispatcher } from "./dispatch.js";
5+
import { buildTestCtx } from "./reply/test-ctx.js";
46

57
function createDispatcher(record: string[]): ReplyDispatcher {
68
return {
@@ -58,4 +60,32 @@ describe("withReplyDispatcher", () => {
5860
expect(onSettled).toHaveBeenCalledTimes(1);
5961
expect(order).toEqual(["run", "markComplete", "waitForIdle", "onSettled"]);
6062
});
63+
64+
it("dispatchInboundMessage owns dispatcher lifecycle", async () => {
65+
const order: string[] = [];
66+
const dispatcher = {
67+
sendToolResult: () => true,
68+
sendBlockReply: () => true,
69+
sendFinalReply: () => {
70+
order.push("sendFinalReply");
71+
return true;
72+
},
73+
getQueuedCounts: () => ({ tool: 0, block: 0, final: 0 }),
74+
markComplete: () => {
75+
order.push("markComplete");
76+
},
77+
waitForIdle: async () => {
78+
order.push("waitForIdle");
79+
},
80+
} satisfies ReplyDispatcher;
81+
82+
await dispatchInboundMessage({
83+
ctx: buildTestCtx(),
84+
cfg: {} as OpenClawConfig,
85+
dispatcher,
86+
replyResolver: async () => ({ text: "ok" }),
87+
});
88+
89+
expect(order).toEqual(["sendFinalReply", "markComplete", "waitForIdle"]);
90+
});
6191
});

‎src/auto-reply/dispatch.ts‎

Lines changed: 28 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -40,12 +40,16 @@ export async function dispatchInboundMessage(params: {
4040
replyResolver?: typeof import("./reply.js").getReplyFromConfig;
4141
}): Promise<DispatchInboundResult> {
4242
const finalized = finalizeInboundContext(params.ctx);
43-
return await dispatchReplyFromConfig({
44-
ctx: finalized,
45-
cfg: params.cfg,
43+
return await withReplyDispatcher({
4644
dispatcher: params.dispatcher,
47-
replyOptions: params.replyOptions,
48-
replyResolver: params.replyResolver,
45+
run: () =>
46+
dispatchReplyFromConfig({
47+
ctx: finalized,
48+
cfg: params.cfg,
49+
dispatcher: params.dispatcher,
50+
replyOptions: params.replyOptions,
51+
replyResolver: params.replyResolver,
52+
}),
4953
});
5054
}
5155

@@ -59,23 +63,20 @@ export async function dispatchInboundMessageWithBufferedDispatcher(params: {
5963
const { dispatcher, replyOptions, markDispatchIdle } = createReplyDispatcherWithTyping(
6064
params.dispatcherOptions,
6165
);
62-
return await withReplyDispatcher({
63-
dispatcher,
64-
run: async () =>
65-
dispatchInboundMessage({
66-
ctx: params.ctx,
67-
cfg: params.cfg,
68-
dispatcher,
69-
replyResolver: params.replyResolver,
70-
replyOptions: {
71-
...params.replyOptions,
72-
...replyOptions,
73-
},
74-
}),
75-
onSettled: () => {
76-
markDispatchIdle();
77-
},
78-
});
66+
try {
67+
return await dispatchInboundMessage({
68+
ctx: params.ctx,
69+
cfg: params.cfg,
70+
dispatcher,
71+
replyResolver: params.replyResolver,
72+
replyOptions: {
73+
...params.replyOptions,
74+
...replyOptions,
75+
},
76+
});
77+
} finally {
78+
markDispatchIdle();
79+
}
7980
}
8081

8182
export async function dispatchInboundMessageWithDispatcher(params: {
@@ -86,15 +87,11 @@ export async function dispatchInboundMessageWithDispatcher(params: {
8687
replyResolver?: typeof import("./reply.js").getReplyFromConfig;
8788
}): Promise<DispatchInboundResult> {
8889
const dispatcher = createReplyDispatcher(params.dispatcherOptions);
89-
return await withReplyDispatcher({
90+
return await dispatchInboundMessage({
91+
ctx: params.ctx,
92+
cfg: params.cfg,
9093
dispatcher,
91-
run: async () =>
92-
dispatchInboundMessage({
93-
ctx: params.ctx,
94-
cfg: params.cfg,
95-
dispatcher,
96-
replyResolver: params.replyResolver,
97-
replyOptions: params.replyOptions,
98-
}),
94+
replyResolver: params.replyResolver,
95+
replyOptions: params.replyOptions,
9996
});
10097
}

‎src/auto-reply/reply/dispatch-from-config.ts‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,6 @@ export async function dispatchReplyFromConfig(params: {
278278
} else {
279279
queuedFinal = dispatcher.sendFinalReply(payload);
280280
}
281-
await dispatcher.waitForIdle();
282281
const counts = dispatcher.getQueuedCounts();
283282
counts.final += routedFinalCount;
284283
recordProcessed("completed", { reason: "fast_abort" });
@@ -443,8 +442,6 @@ export async function dispatchReplyFromConfig(params: {
443442
}
444443
}
445444

446-
await dispatcher.waitForIdle();
447-
448445
const counts = dispatcher.getQueuedCounts();
449446
counts.final += routedFinalCount;
450447
recordProcessed("completed");
@@ -454,9 +451,5 @@ export async function dispatchReplyFromConfig(params: {
454451
recordProcessed("error", { error: String(err) });
455452
markIdle("message_error");
456453
throw err;
457-
} finally {
458-
// Always clear the dispatcher reservation so a leaked pending count
459-
// can never permanently block gateway restarts.
460-
dispatcher.markComplete();
461454
}
462455
}

‎src/discord/monitor/message-handler.process.test.ts‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,14 @@ vi.mock("../../auto-reply/reply/dispatch-from-config.js", () => ({
2020

2121
vi.mock("../../auto-reply/reply/reply-dispatcher.js", () => ({
2222
createReplyDispatcherWithTyping: vi.fn(() => ({
23-
dispatcher: {},
23+
dispatcher: {
24+
sendToolResult: vi.fn(() => true),
25+
sendBlockReply: vi.fn(() => true),
26+
sendFinalReply: vi.fn(() => true),
27+
waitForIdle: vi.fn(async () => {}),
28+
getQueuedCounts: vi.fn(() => ({ tool: 0, block: 0, final: 0 })),
29+
markComplete: vi.fn(),
30+
},
2431
replyOptions: {},
2532
markDispatchIdle: vi.fn(),
2633
})),

‎src/gateway/server-methods/chat.ts‎

Lines changed: 29 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import type { GatewayRequestContext, GatewayRequestHandlers } from "./types.js";
66
import { resolveSessionAgentId } from "../../agents/agent-scope.js";
77
import { resolveThinkingDefault } from "../../agents/model-selection.js";
88
import { resolveAgentTimeoutMs } from "../../agents/timeout.js";
9-
import { dispatchInboundMessage, withReplyDispatcher } from "../../auto-reply/dispatch.js";
9+
import { dispatchInboundMessage } from "../../auto-reply/dispatch.js";
1010
import { createReplyDispatcher } from "../../auto-reply/reply/reply-dispatcher.js";
1111
import { createReplyPrefixOptions } from "../../channels/reply-prefix.js";
1212
import { resolveSessionFilePath } from "../../config/sessions.js";
@@ -524,40 +524,36 @@ export const chatHandlers: GatewayRequestHandlers = {
524524
});
525525

526526
let agentRunStarted = false;
527-
void withReplyDispatcher({
527+
void dispatchInboundMessage({
528+
ctx,
529+
cfg,
528530
dispatcher,
529-
run: () =>
530-
dispatchInboundMessage({
531-
ctx,
532-
cfg,
533-
dispatcher,
534-
replyOptions: {
535-
runId: clientRunId,
536-
abortSignal: abortController.signal,
537-
images: parsedImages.length > 0 ? parsedImages : undefined,
538-
disableBlockStreaming: true,
539-
onAgentRunStart: (runId) => {
540-
agentRunStarted = true;
541-
const connId = typeof client?.connId === "string" ? client.connId : undefined;
542-
const wantsToolEvents = hasGatewayClientCap(
543-
client?.connect?.caps,
544-
GATEWAY_CLIENT_CAPS.TOOL_EVENTS,
545-
);
546-
if (connId && wantsToolEvents) {
547-
context.registerToolEventRecipient(runId, connId);
548-
// Register for any other active runs *in the same session* so
549-
// late-joining clients (e.g. page refresh mid-response) receive
550-
// in-progress tool events without leaking cross-session data.
551-
for (const [activeRunId, active] of context.chatAbortControllers) {
552-
if (activeRunId !== runId && active.sessionKey === p.sessionKey) {
553-
context.registerToolEventRecipient(activeRunId, connId);
554-
}
555-
}
531+
replyOptions: {
532+
runId: clientRunId,
533+
abortSignal: abortController.signal,
534+
images: parsedImages.length > 0 ? parsedImages : undefined,
535+
disableBlockStreaming: true,
536+
onAgentRunStart: (runId) => {
537+
agentRunStarted = true;
538+
const connId = typeof client?.connId === "string" ? client.connId : undefined;
539+
const wantsToolEvents = hasGatewayClientCap(
540+
client?.connect?.caps,
541+
GATEWAY_CLIENT_CAPS.TOOL_EVENTS,
542+
);
543+
if (connId && wantsToolEvents) {
544+
context.registerToolEventRecipient(runId, connId);
545+
// Register for any other active runs *in the same session* so
546+
// late-joining clients (e.g. page refresh mid-response) receive
547+
// in-progress tool events without leaking cross-session data.
548+
for (const [activeRunId, active] of context.chatAbortControllers) {
549+
if (activeRunId !== runId && active.sessionKey === p.sessionKey) {
550+
context.registerToolEventRecipient(activeRunId, connId);
556551
}
557-
},
558-
onModelSelected,
559-
},
560-
}),
552+
}
553+
}
554+
},
555+
onModelSelected,
556+
},
561557
})
562558
.then(() => {
563559
if (!agentRunStarted) {

‎src/imessage/monitor/monitor-provider.ts‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import type { IMessagePayload, MonitorIMessageOpts } from "./types.js";
33
import { resolveHumanDelayConfig } from "../../agents/identity.js";
44
import { resolveTextChunkLimit } from "../../auto-reply/chunk.js";
55
import { hasControlCommand } from "../../auto-reply/command-detection.js";
6-
import { dispatchInboundMessage, withReplyDispatcher } from "../../auto-reply/dispatch.js";
6+
import { dispatchInboundMessage } from "../../auto-reply/dispatch.js";
77
import {
88
formatInboundEnvelope,
99
formatInboundFromLabel,
@@ -647,21 +647,17 @@ export async function monitorIMessageProvider(opts: MonitorIMessageOpts = {}): P
647647
},
648648
});
649649

650-
const { queuedFinal } = await withReplyDispatcher({
650+
const { queuedFinal } = await dispatchInboundMessage({
651+
ctx: ctxPayload,
652+
cfg,
651653
dispatcher,
652-
run: () =>
653-
dispatchInboundMessage({
654-
ctx: ctxPayload,
655-
cfg,
656-
dispatcher,
657-
replyOptions: {
658-
disableBlockStreaming:
659-
typeof accountInfo.config.blockStreaming === "boolean"
660-
? !accountInfo.config.blockStreaming
661-
: undefined,
662-
onModelSelected,
663-
},
664-
}),
654+
replyOptions: {
655+
disableBlockStreaming:
656+
typeof accountInfo.config.blockStreaming === "boolean"
657+
? !accountInfo.config.blockStreaming
658+
: undefined,
659+
onModelSelected,
660+
},
665661
});
666662

667663
if (!queuedFinal) {

0 commit comments

Comments
 (0)