fix: stop typing after dispatcher idle

This commit is contained in:
Peter Steinberger
2026-01-06 03:05:11 +00:00
parent cbc39bd005
commit d5f088978a
13 changed files with 238 additions and 138 deletions

View File

@@ -233,6 +233,7 @@ export async function getReplyFromConfig(
silentToken: SILENT_REPLY_TOKEN,
log: defaultRuntime.log,
});
opts?.onTypingController?.(typing);
let transcribedText: string | undefined;
if (cfg.routing?.transcribeAudio && isAudio(ctx.MediaType)) {

View File

@@ -50,6 +50,8 @@ function createTyping(): TypingController {
startTypingLoop: vi.fn(async () => {}),
startTypingOnText: vi.fn(async () => {}),
refreshTypingTtl: vi.fn(),
markRunComplete: vi.fn(),
markDispatchIdle: vi.fn(),
cleanup: vi.fn(),
};
}

View File

@@ -514,6 +514,6 @@ export async function runReplyAgent(params: {
finalPayloads.length === 1 ? finalPayloads[0] : finalPayloads,
);
} finally {
typing.cleanup();
typing.markRunComplete();
}
}

View File

@@ -37,6 +37,8 @@ function createTyping(): TypingController {
startTypingLoop: vi.fn(async () => {}),
startTypingOnText: vi.fn(async () => {}),
refreshTypingTtl: vi.fn(),
markRunComplete: vi.fn(),
markDispatchIdle: vi.fn(),
cleanup: vi.fn(),
};
}

View File

@@ -58,153 +58,157 @@ export function createFollowupRunner(params: {
};
return async (queued: FollowupRun) => {
const runId = crypto.randomUUID();
if (queued.run.sessionKey) {
registerAgentRunContext(runId, { sessionKey: queued.run.sessionKey });
}
let autoCompactionCompleted = false;
let runResult: Awaited<ReturnType<typeof runEmbeddedPiAgent>>;
let fallbackProvider = queued.run.provider;
let fallbackModel = queued.run.model;
try {
const fallbackResult = await runWithModelFallback({
cfg: queued.run.config,
provider: queued.run.provider,
model: queued.run.model,
run: (provider, model) =>
runEmbeddedPiAgent({
sessionId: queued.run.sessionId,
sessionKey: queued.run.sessionKey,
surface: queued.run.surface,
sessionFile: queued.run.sessionFile,
workspaceDir: queued.run.workspaceDir,
config: queued.run.config,
skillsSnapshot: queued.run.skillsSnapshot,
prompt: queued.prompt,
extraSystemPrompt: queued.run.extraSystemPrompt,
ownerNumbers: queued.run.ownerNumbers,
enforceFinalTag: queued.run.enforceFinalTag,
provider,
model,
authProfileId: queued.run.authProfileId,
thinkLevel: queued.run.thinkLevel,
verboseLevel: queued.run.verboseLevel,
bashElevated: queued.run.bashElevated,
timeoutMs: queued.run.timeoutMs,
runId,
blockReplyBreak: queued.run.blockReplyBreak,
onAgentEvent: (evt) => {
if (evt.stream !== "compaction") return;
const phase = String(evt.data.phase ?? "");
const willRetry = Boolean(evt.data.willRetry);
if (phase === "end" && !willRetry) {
autoCompactionCompleted = true;
}
},
}),
});
runResult = fallbackResult.result;
fallbackProvider = fallbackResult.provider;
fallbackModel = fallbackResult.model;
} catch (err) {
const message = err instanceof Error ? err.message : String(err);
defaultRuntime.error?.(`Followup agent failed before reply: ${message}`);
return;
}
const payloadArray = runResult.payloads ?? [];
if (payloadArray.length === 0) return;
const sanitizedPayloads = payloadArray.flatMap((payload) => {
const text = payload.text;
if (!text || !text.includes("HEARTBEAT_OK")) return [payload];
const stripped = stripHeartbeatToken(text, { mode: "message" });
const hasMedia =
Boolean(payload.mediaUrl) || (payload.mediaUrls?.length ?? 0) > 0;
if (stripped.shouldSkip && !hasMedia) return [];
return [{ ...payload, text: stripped.text }];
});
const replyTaggedPayloads: ReplyPayload[] = sanitizedPayloads
.map((payload) => {
const { cleaned, replyToId } = extractReplyToTag(payload.text);
return {
...payload,
text: cleaned ? cleaned : undefined,
replyToId: replyToId ?? payload.replyToId,
};
})
.filter(
(payload) =>
payload.text ||
payload.mediaUrl ||
(payload.mediaUrls && payload.mediaUrls.length > 0),
);
if (replyTaggedPayloads.length === 0) return;
if (autoCompactionCompleted) {
const count = await incrementCompactionCount({
sessionEntry,
sessionStore,
sessionKey,
storePath,
});
if (queued.run.verboseLevel === "on") {
const suffix = typeof count === "number" ? ` (count ${count})` : "";
replyTaggedPayloads.unshift({
text: `🧹 Auto-compaction complete${suffix}.`,
const runId = crypto.randomUUID();
if (queued.run.sessionKey) {
registerAgentRunContext(runId, { sessionKey: queued.run.sessionKey });
}
let autoCompactionCompleted = false;
let runResult: Awaited<ReturnType<typeof runEmbeddedPiAgent>>;
let fallbackProvider = queued.run.provider;
let fallbackModel = queued.run.model;
try {
const fallbackResult = await runWithModelFallback({
cfg: queued.run.config,
provider: queued.run.provider,
model: queued.run.model,
run: (provider, model) =>
runEmbeddedPiAgent({
sessionId: queued.run.sessionId,
sessionKey: queued.run.sessionKey,
surface: queued.run.surface,
sessionFile: queued.run.sessionFile,
workspaceDir: queued.run.workspaceDir,
config: queued.run.config,
skillsSnapshot: queued.run.skillsSnapshot,
prompt: queued.prompt,
extraSystemPrompt: queued.run.extraSystemPrompt,
ownerNumbers: queued.run.ownerNumbers,
enforceFinalTag: queued.run.enforceFinalTag,
provider,
model,
authProfileId: queued.run.authProfileId,
thinkLevel: queued.run.thinkLevel,
verboseLevel: queued.run.verboseLevel,
bashElevated: queued.run.bashElevated,
timeoutMs: queued.run.timeoutMs,
runId,
blockReplyBreak: queued.run.blockReplyBreak,
onAgentEvent: (evt) => {
if (evt.stream !== "compaction") return;
const phase = String(evt.data.phase ?? "");
const willRetry = Boolean(evt.data.willRetry);
if (phase === "end" && !willRetry) {
autoCompactionCompleted = true;
}
},
}),
});
runResult = fallbackResult.result;
fallbackProvider = fallbackResult.provider;
fallbackModel = fallbackResult.model;
} catch (err) {
const message = err instanceof Error ? err.message : String(err);
defaultRuntime.error?.(`Followup agent failed before reply: ${message}`);
return;
}
}
if (sessionStore && sessionKey) {
const usage = runResult.meta.agentMeta?.usage;
const modelUsed =
runResult.meta.agentMeta?.model ?? fallbackModel ?? defaultModel;
const contextTokensUsed =
agentCfgContextTokens ??
lookupContextTokens(modelUsed) ??
sessionEntry?.contextTokens ??
DEFAULT_CONTEXT_TOKENS;
const payloadArray = runResult.payloads ?? [];
if (payloadArray.length === 0) return;
const sanitizedPayloads = payloadArray.flatMap((payload) => {
const text = payload.text;
if (!text || !text.includes("HEARTBEAT_OK")) return [payload];
const stripped = stripHeartbeatToken(text, { mode: "message" });
const hasMedia =
Boolean(payload.mediaUrl) || (payload.mediaUrls?.length ?? 0) > 0;
if (stripped.shouldSkip && !hasMedia) return [];
return [{ ...payload, text: stripped.text }];
});
if (usage) {
const entry = sessionStore[sessionKey];
if (entry) {
const input = usage.input ?? 0;
const output = usage.output ?? 0;
const promptTokens =
input + (usage.cacheRead ?? 0) + (usage.cacheWrite ?? 0);
sessionStore[sessionKey] = {
...entry,
inputTokens: input,
outputTokens: output,
totalTokens:
promptTokens > 0 ? promptTokens : (usage.total ?? input),
modelProvider: fallbackProvider ?? entry.modelProvider,
model: modelUsed,
contextTokens: contextTokensUsed ?? entry.contextTokens,
updatedAt: Date.now(),
const replyTaggedPayloads: ReplyPayload[] = sanitizedPayloads
.map((payload) => {
const { cleaned, replyToId } = extractReplyToTag(payload.text);
return {
...payload,
text: cleaned ? cleaned : undefined,
replyToId: replyToId ?? payload.replyToId,
};
if (storePath) {
await saveSessionStore(storePath, sessionStore);
}
})
.filter(
(payload) =>
payload.text ||
payload.mediaUrl ||
(payload.mediaUrls && payload.mediaUrls.length > 0),
);
if (replyTaggedPayloads.length === 0) return;
if (autoCompactionCompleted) {
const count = await incrementCompactionCount({
sessionEntry,
sessionStore,
sessionKey,
storePath,
});
if (queued.run.verboseLevel === "on") {
const suffix = typeof count === "number" ? ` (count ${count})` : "";
replyTaggedPayloads.unshift({
text: `🧹 Auto-compaction complete${suffix}.`,
});
}
} else if (modelUsed || contextTokensUsed) {
const entry = sessionStore[sessionKey];
if (entry) {
sessionStore[sessionKey] = {
...entry,
modelProvider: fallbackProvider ?? entry.modelProvider,
model: modelUsed ?? entry.model,
contextTokens: contextTokensUsed ?? entry.contextTokens,
};
if (storePath) {
await saveSessionStore(storePath, sessionStore);
}
if (sessionStore && sessionKey) {
const usage = runResult.meta.agentMeta?.usage;
const modelUsed =
runResult.meta.agentMeta?.model ?? fallbackModel ?? defaultModel;
const contextTokensUsed =
agentCfgContextTokens ??
lookupContextTokens(modelUsed) ??
sessionEntry?.contextTokens ??
DEFAULT_CONTEXT_TOKENS;
if (usage) {
const entry = sessionStore[sessionKey];
if (entry) {
const input = usage.input ?? 0;
const output = usage.output ?? 0;
const promptTokens =
input + (usage.cacheRead ?? 0) + (usage.cacheWrite ?? 0);
sessionStore[sessionKey] = {
...entry,
inputTokens: input,
outputTokens: output,
totalTokens:
promptTokens > 0 ? promptTokens : (usage.total ?? input),
modelProvider: fallbackProvider ?? entry.modelProvider,
model: modelUsed,
contextTokens: contextTokensUsed ?? entry.contextTokens,
updatedAt: Date.now(),
};
if (storePath) {
await saveSessionStore(storePath, sessionStore);
}
}
} else if (modelUsed || contextTokensUsed) {
const entry = sessionStore[sessionKey];
if (entry) {
sessionStore[sessionKey] = {
...entry,
modelProvider: fallbackProvider ?? entry.modelProvider,
model: modelUsed ?? entry.model,
contextTokens: contextTokensUsed ?? entry.contextTokens,
};
if (storePath) {
await saveSessionStore(storePath, sessionStore);
}
}
}
}
}
await sendFollowupPayloads(replyTaggedPayloads);
await sendFollowupPayloads(replyTaggedPayloads);
} finally {
typing.markRunComplete();
}
};
}

View File

@@ -79,4 +79,18 @@ describe("createReplyDispatcher", () => {
await dispatcher.waitForIdle();
expect(delivered).toEqual(["tool", "block", "final"]);
});
it("fires onIdle when the queue drains", async () => {
const deliver = vi.fn(
async () => await new Promise((resolve) => setTimeout(resolve, 5)),
);
const onIdle = vi.fn();
const dispatcher = createReplyDispatcher({ deliver, onIdle });
dispatcher.sendToolResult({ text: "one" });
dispatcher.sendFinalReply({ text: "two" });
await dispatcher.waitForIdle();
expect(onIdle).toHaveBeenCalledTimes(1);
});
});

View File

@@ -18,6 +18,7 @@ export type ReplyDispatcherOptions = {
deliver: ReplyDispatchDeliverer;
responsePrefix?: string;
onHeartbeatStrip?: () => void;
onIdle?: () => void;
onError?: ReplyDispatchErrorHandler;
};
@@ -70,6 +71,8 @@ export function createReplyDispatcher(
options: ReplyDispatcherOptions,
): ReplyDispatcher {
let sendChain: Promise<void> = Promise.resolve();
// Track in-flight deliveries so we can emit a reliable "idle" signal.
let pending = 0;
// Serialize outbound replies to preserve tool/block/final order.
const queuedCounts: Record<ReplyDispatchKind, number> = {
tool: 0,
@@ -81,10 +84,17 @@ export function createReplyDispatcher(
const normalized = normalizeReplyPayload(payload, options);
if (!normalized) return false;
queuedCounts[kind] += 1;
pending += 1;
sendChain = sendChain
.then(() => options.deliver(normalized, { kind }))
.catch((err) => {
options.onError?.(err, { kind });
})
.finally(() => {
pending -= 1;
if (pending === 0) {
options.onIdle?.();
}
});
return true;
};

View File

@@ -3,6 +3,8 @@ export type TypingController = {
startTypingLoop: () => Promise<void>;
startTypingOnText: (text?: string) => Promise<void>;
refreshTypingTtl: () => void;
markRunComplete: () => void;
markDispatchIdle: () => void;
cleanup: () => void;
};
@@ -21,6 +23,9 @@ export function createTypingController(params: {
log,
} = params;
let started = false;
let active = false;
let runComplete = false;
let dispatchIdle = false;
let typingTimer: NodeJS.Timeout | undefined;
let typingTtlTimer: NodeJS.Timeout | undefined;
const typingIntervalMs = typingIntervalSeconds * 1000;
@@ -30,6 +35,13 @@ export function createTypingController(params: {
return `${Math.round(ms / 1000)}s`;
};
const resetCycle = () => {
started = false;
active = false;
runComplete = false;
dispatchIdle = false;
};
const cleanup = () => {
if (typingTtlTimer) {
clearTimeout(typingTtlTimer);
@@ -39,6 +51,7 @@ export function createTypingController(params: {
clearInterval(typingTimer);
typingTimer = undefined;
}
resetCycle();
};
const refreshTypingTtl = () => {
@@ -61,11 +74,22 @@ export function createTypingController(params: {
};
const ensureStart = async () => {
if (!active) {
active = true;
runComplete = false;
dispatchIdle = false;
}
if (started) return;
started = true;
await triggerTyping();
};
const maybeStopOnIdle = () => {
if (!active) return;
// Stop only when the model run is done and the dispatcher queue is empty.
if (runComplete && dispatchIdle) cleanup();
};
const startTypingLoop = async () => {
if (!onReplyStart) return;
if (typingIntervalMs <= 0) return;
@@ -85,11 +109,23 @@ export function createTypingController(params: {
await startTypingLoop();
};
const markRunComplete = () => {
runComplete = true;
maybeStopOnIdle();
};
const markDispatchIdle = () => {
dispatchIdle = true;
maybeStopOnIdle();
};
return {
onReplyStart: ensureStart,
startTypingLoop,
startTypingOnText,
refreshTypingTtl,
markRunComplete,
markDispatchIdle,
cleanup,
};
}

View File

@@ -1,5 +1,8 @@
import type { TypingController } from "./reply/typing.js";
export type GetReplyOptions = {
onReplyStart?: () => Promise<void> | void;
onTypingController?: (typing: TypingController) => void;
isHeartbeat?: boolean;
onPartialReply?: (payload: ReplyPayload) => Promise<void> | void;
onBlockReply?: (payload: ReplyPayload) => Promise<void> | void;