refactor(auto-reply): split reply flow

src/auto-reply/reply/followup-runner.ts

@@ -0,0 +1,168 @@
+import crypto from "node:crypto";
+import { lookupContextTokens } from "../../agents/context.js";
+import { DEFAULT_CONTEXT_TOKENS } from "../../agents/defaults.js";
+import { runEmbeddedPiAgent } from "../../agents/pi-embedded.js";
+import { type SessionEntry, saveSessionStore } from "../../config/sessions.js";
+import { logVerbose } from "../../globals.js";
+import { registerAgentRunContext } from "../../infra/agent-events.js";
+import { defaultRuntime } from "../../runtime.js";
+import { stripHeartbeatToken } from "../heartbeat.js";
+import { SILENT_REPLY_TOKEN } from "../tokens.js";
+import type { GetReplyOptions, ReplyPayload } from "../types.js";
+import type { FollowupRun } from "./queue.js";
+import { extractReplyToTag } from "./reply-tags.js";
+import type { TypingController } from "./typing.js";
+
+export function createFollowupRunner(params: {
+  opts?: GetReplyOptions;
+  typing: TypingController;
+  sessionEntry?: SessionEntry;
+  sessionStore?: Record<string, SessionEntry>;
+  sessionKey?: string;
+  storePath?: string;
+  defaultModel: string;
+  agentCfgContextTokens?: number;
+}): (queued: FollowupRun) => Promise<void> {
+  const {
+    opts,
+    typing,
+    sessionEntry,
+    sessionStore,
+    sessionKey,
+    storePath,
+    defaultModel,
+    agentCfgContextTokens,
+  } = params;
+
+  const sendFollowupPayloads = async (payloads: ReplyPayload[]) => {
+    if (!opts?.onBlockReply) {
+      logVerbose("followup queue: no onBlockReply handler; dropping payloads");
+      return;
+    }
+    for (const payload of payloads) {
+      if (!payload?.text && !payload?.mediaUrl && !payload?.mediaUrls?.length) {
+        continue;
+      }
+      if (
+        payload.text?.trim() === SILENT_REPLY_TOKEN &&
+        !payload.mediaUrl &&
+        !payload.mediaUrls?.length
+      ) {
+        continue;
+      }
+      await typing.startTypingOnText(payload.text);
+      await opts.onBlockReply(payload);
+    }
+  };
+
+  return async (queued: FollowupRun) => {
+    const runId = crypto.randomUUID();
+    if (queued.run.sessionKey) {
+      registerAgentRunContext(runId, { sessionKey: queued.run.sessionKey });
+    }
+    let runResult: Awaited<ReturnType<typeof runEmbeddedPiAgent>>;
+    try {
+      runResult = await 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: queued.run.provider,
+        model: queued.run.model,
+        thinkLevel: queued.run.thinkLevel,
+        verboseLevel: queued.run.verboseLevel,
+        timeoutMs: queued.run.timeoutMs,
+        runId,
+        blockReplyBreak: queued.run.blockReplyBreak,
+      });
+    } 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 (sessionStore && sessionKey) {
+      const usage = runResult.meta.agentMeta?.usage;
+      const modelUsed = runResult.meta.agentMeta?.model ?? 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),
+            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,
+            model: modelUsed ?? entry.model,
+            contextTokens: contextTokensUsed ?? entry.contextTokens,
+          };
+          if (storePath) {
+            await saveSessionStore(storePath, sessionStore);
+          }
+        }
+      }
+    }
+
+    await sendFollowupPayloads(replyTaggedPayloads);
+  };
+}