chore: optimization run

Author: vicary

src/handlers/chat.test.ts

@@ -404,4 +404,28 @@ describe("chat handler", () => {
     assertEquals(graphitiAsync.refreshCalls, []);
     assertEquals(graphitiAsync.drainCalls, []);
   });
+
+  it("skips session resolution and hot-tier work when no text prompt is present", async () => {
+    const sessionManager = new MockSessionManager();
+    const redisEvents = new MockRedisEvents();
+    const graphitiAsync = new MockGraphitiAsync();
+
+    const handler = createChatHandler({
+      sessionManager: sessionManager as never,
+      redisEvents: redisEvents as never,
+      graphitiAsync: graphitiAsync as never,
+      drainTriggerSize: 99,
+    });
+
+    await handler(
+      { sessionID: "session-1" },
+      { parts: [{ type: "file", path: "src/index.ts" }] } as never,
+    );
+
+    assertEquals(sessionManager.activeCalls, []);
+    assertEquals(sessionManager.prepareInjectionCalls, []);
+    assertEquals(redisEvents.calls, []);
+    assertEquals(graphitiAsync.refreshCalls, []);
+    assertEquals(graphitiAsync.drainCalls, []);
+  });
 });

src/handlers/chat.ts

@@ -24,6 +24,12 @@ export function createChatHandler(deps: ChatHandlerDeps): ChatMessageHook {
   return async ({ sessionID }: ChatMessageInput, output: ChatMessageOutput) => {
     try {
       sessionManager.markSessionActive(sessionID);
+
+      const messageText = extractTextFromParts(output.parts);
+      if (!messageText) return;
+      const sanitizedMessageText = sanitizeMemoryInput(messageText);
+      if (!sanitizedMessageText) return;
+
       const { state, resolved, canonicalSessionId } = await sessionManager
         .resolveSessionState(
           sessionID,
@@ -32,11 +38,6 @@ export function createChatHandler(deps: ChatHandlerDeps): ChatMessageHook {
       if (!canonicalSessionId) return;
       sessionManager.markResolvedSessionActive(sessionID, canonicalSessionId);
 
-      const messageText = extractTextFromParts(output.parts);
-      if (!messageText) return;
-      const sanitizedMessageText = sanitizeMemoryInput(messageText);
-      if (!sanitizedMessageText) return;
-
       state.messageCount += 1;
       state.latestUserRequest = sanitizedMessageText;
       state.latestRefreshQuery = sanitizedMessageText;

src/handlers/messages.test.ts

@@ -759,4 +759,26 @@ describe("messages handler", () => {
 
     assertEquals(output.messages[0].parts[0].text, "startup prompt");
   });
+
+  it("skips transform work when the latest user entry has no text part", async () => {
+    const sessionManager = new MockSessionManager();
+    sessionManager.prepareInjectionImpl = () => {
+      throw new Error("prepareInjection should not run");
+    };
+    const handler = createMessagesHandler({
+      sessionManager: sessionManager as never,
+    });
+
+    const output = {
+      messages: [{
+        info: { role: "user", sessionID: "session-1" },
+        parts: [{ type: "file", path: "src/index.ts" }],
+      }],
+    };
+
+    await handler({ message: "should be ignored" } as never, output as never);
+
+    assertEquals(sessionManager.activeCalls, []);
+    assertEquals(sessionManager.state.pendingInjection, undefined);
+  });
 });

src/handlers/messages.ts

@@ -28,17 +28,6 @@ const getTransformMessage = (input: unknown): string | undefined => {
   return typeof message === "string" ? message : undefined;
 };
 
-const getLatestUserText = (
-  output: MessagesTransformOutput,
-): string | undefined => {
-  const lastUserEntry = output.messages
-    .findLast((message) => message.info.role === "user");
-  const textPart = lastUserEntry?.parts.find(isTextPart);
-  return textPart
-    ? sanitizeMemoryInputPreservingMemoryBlocks(textPart.text)
-    : undefined;
-};
-
 const LEADING_INJECTED_MEMORY_BLOCK =
   /^(?:<session_memory\b[\s\S]*?<\/session_memory>|<memory\b[\s\S]*?<\/memory>|<persistent_memory\b[\s\S]*?<\/persistent_memory>)(?:\r?\n){0,2}/;
 
@@ -64,6 +53,10 @@ export function createMessagesHandler(
       .findLast((message) => message.info.role === "user");
     if (!lastUserEntry) return;
 
+    const textPart = lastUserEntry.parts.find(isTextPart);
+    const latestUserText = textPart?.text;
+    if (latestUserText === undefined) return;
+
     const sourceSessionID = lastUserEntry.info.sessionID;
 
     try {
@@ -79,12 +72,9 @@ export function createMessagesHandler(
         canonicalSessionId,
       );
 
-      const textPart = lastUserEntry.parts.find(isTextPart);
-      const latestUserText = textPart?.text;
-
       const recallQuery = sanitizeMemoryInput(
         stripInjectedMemoryBlocks(
-          getTransformMessage(input) ?? getLatestUserText(output) ?? "",
+          getTransformMessage(input) ?? latestUserText,
         ),
       ) || undefined;
       const prepared = state.pendingInjection ??
@@ -93,8 +83,8 @@ export function createMessagesHandler(
           recallQuery,
         );
       if (!prepared) return;
+      if (!textPart) return;
 
-      if (!textPart || latestUserText === undefined) return;
       const scrubbedUserText = scrubPromptMemoryText(latestUserText);
       const effectiveUserText = sanitizeMemoryInputPreservingMemoryBlocks(
         scrubbedUserText,

src/services/batch-drain.test.ts

@@ -303,6 +303,38 @@ describe("batch drain", () => {
     );
   });
 
+  it("avoids an extra ownership refresh before checkpointing skipped entries", async () => {
+    const { events, drain } = await createDeps({
+      drain: { batchSize: 2, claimHeartbeatIntervalMs: null },
+    });
+    const skipped = createSessionEvent("message", "assistant", {
+      summary: "assistant chatter",
+      body: "assistant chatter",
+    });
+    const drained = createSessionEvent("message", "user", {
+      summary: "user message",
+      body: "user message",
+    });
+    await events.recordEvent("session-1", "group-1", skipped);
+    await events.recordEvent("session-1", "group-1", drained);
+
+    const refreshSpy = spy(events, "refreshClaimLease");
+    const added: string[] = [];
+    try {
+      const result = await drain.drainGroup("group-1", {
+        addMemory(input: { name: string }) {
+          added.push(input.name);
+        },
+      } as never);
+
+      assertEquals(result, { status: "success", drained: 1 });
+      assertEquals(added, [`message:${drained.id}`]);
+      assertEquals(refreshSpy.calls.length, 4);
+    } finally {
+      refreshSpy.restore();
+    }
+  });
+
   it("limits batches using serialized Graphiti episode bodies", async () => {
     const first = createSessionEvent("message", "user", {
       summary: "first",

src/services/batch-drain.ts

@@ -33,6 +33,16 @@ class DrainClaimLostError extends Error {
 const makeBatchKey = (entries: DrainQueueEntry[]): string =>
   `${entries[0]?.event.id ?? "empty"}:${entries.at(-1)?.event.id ?? "empty"}`;
 
+const getDrainableEntryIds = (entries: DrainQueueEntry[]): Set<string> => {
+  const drainableEntryIds = new Set<string>();
+  for (const entry of entries) {
+    if (shouldDrainEntry(entry)) {
+      drainableEntryIds.add(entry.event.id);
+    }
+  }
+  return drainableEntryIds;
+};
+
 const shouldDrainEntry = (entry: DrainQueueEntry): boolean => {
   const text = sanitizeMemoryInput(getSessionEventRecallText(entry.event));
   if (!text) return false;
@@ -135,8 +145,9 @@ export class BatchDrainService {
 
     const batch = claimed.entries;
     const batchKey = makeBatchKey(batch);
-    const semanticBatch = batch.filter(shouldDrainEntry);
-    if (semanticBatch.length === 0) {
+    const eventIds = batch.map((entry) => entry.event.id);
+    const drainableEntryIds = getDrainableEntryIds(batch);
+    if (drainableEntryIds.size === 0) {
       await this.events.markBatchSuccess(groupId, claimed.claimToken, batch);
       await this.redis.deleteKey(drainRetryKey(groupId, batchKey));
       return { status: "success", drained: 0 };
@@ -149,19 +160,7 @@ export class BatchDrainService {
     }
 
     let lostClaim = false;
-    const refreshClaimHeartbeat = async (): Promise<void> => {
-      try {
-        const refreshed = await this.events.refreshClaimLease(
-          groupId,
-          claimed.claimToken,
-          claimed.lockTtlSeconds,
-        );
-        if (!refreshed) lostClaim = true;
-      } catch {
-        lostClaim = true;
-      }
-    };
-    const confirmClaimOwnership = async (): Promise<boolean> => {
+    const refreshClaimOwnership = async (): Promise<boolean> => {
       if (lostClaim) return false;
       try {
         const refreshed = await this.events.refreshClaimLease(
@@ -175,6 +174,11 @@ export class BatchDrainService {
       }
       return !lostClaim;
     };
+    const refreshClaimHeartbeat = async (): Promise<void> => {
+      await refreshClaimOwnership();
+    };
+    const confirmClaimOwnership = (): Promise<boolean> =>
+      refreshClaimOwnership();
     const assertClaimOwnership = async (): Promise<void> => {
       if (!await confirmClaimOwnership()) {
         throw new DrainClaimLostError();
@@ -187,8 +191,8 @@ export class BatchDrainService {
 
     try {
       for (const entry of batch) {
-        await assertClaimOwnership();
-        if (shouldDrainEntry(entry)) {
+        if (drainableEntryIds.has(entry.event.id)) {
+          await assertClaimOwnership();
           await graphiti.addMemory({
             name: `${entry.event.category}:${entry.event.id}`,
             episodeBody: entry.episodeBody,
@@ -208,13 +212,13 @@ export class BatchDrainService {
       await assertClaimOwnership();
       await this.events.markBatchSuccess(groupId, claimed.claimToken, batch);
       await this.redis.deleteKey(drainRetryKey(groupId, batchKey));
-      return { status: "success", drained: semanticBatch.length };
+      return { status: "success", drained: drainableEntryIds.size };
     } catch (err) {
       const lostOwnership = err instanceof DrainClaimLostError;
       if (lostOwnership) {
         logger.warn("Drain claim heartbeat lost ownership", {
           groupId,
-          eventIds: batch.map((entry) => entry.event.id),
+          eventIds,
         });
       }
       const attempts = (retryState?.attempts ?? 0) + 1;
@@ -223,7 +227,7 @@ export class BatchDrainService {
         if (!lostOwnership) {
           logger.warn("Drain claim heartbeat lost ownership", {
             groupId,
-            eventIds: batch.map((entry) => entry.event.id),
+            eventIds,
           });
         }
         await this.redis.deleteKey(drainRetryKey(groupId, batchKey));

src/services/connection-manager.ts

@@ -430,18 +430,6 @@ export class GraphitiConnectionManager implements GraphitiToolCaller {
     }
   }
 
-  private async executeConnectedCall(
-    name: string,
-    args: Record<string, unknown>,
-    deadlineMs: number,
-  ): Promise<unknown> {
-    return await this.executeConnectedCallWithinDeadline(
-      name,
-      args,
-      this.now() + deadlineMs,
-    );
-  }
-
   private async executeConnectedCallWithinDeadline(
     name: string,
     args: Record<string, unknown>,
@@ -470,49 +458,58 @@ export class GraphitiConnectionManager implements GraphitiToolCaller {
       }
 
       if (isSessionExpired(err)) {
-        const typedError = new GraphitiSessionExpiredError(
-          getErrorMessage(err) || undefined,
-        );
-
-        if (attempt >= 1) {
-          void this.reconnect();
-          throw typedError;
-        }
-
-        const connected = await this.reconnectWithinDeadline(deadlineAt);
-        if (!connected) throw typedError;
-        return await this.executeConnectedCallWithinDeadline(
+        return await this.retryConnectedCallAfterRecoverableError(
+          new GraphitiSessionExpiredError(
+            getErrorMessage(err) || undefined,
+          ),
           name,
           args,
           deadlineAt,
-          attempt + 1,
+          attempt,
         );
       }
 
       if (isTransportFailure(err)) {
-        const typedError = new GraphitiTransportError(
-          getErrorMessage(err) || undefined,
-        );
-
-        if (attempt >= 1) {
-          void this.reconnect();
-          throw typedError;
-        }
-
-        const connected = await this.reconnectWithinDeadline(deadlineAt);
-        if (!connected) throw typedError;
-        return await this.executeConnectedCallWithinDeadline(
+        return await this.retryConnectedCallAfterRecoverableError(
+          new GraphitiTransportError(
+            getErrorMessage(err) || undefined,
+          ),
           name,
           args,
           deadlineAt,
-          attempt + 1,
+          attempt,
         );
       }
 
       throw err;
     }
   }
 
+  private async retryConnectedCallAfterRecoverableError(
+    typedError: GraphitiSessionExpiredError | GraphitiTransportError,
+    name: string,
+    args: Record<string, unknown>,
+    deadlineAt: number,
+    attempt: number,
+  ): Promise<unknown> {
+    if (attempt >= 1) {
+      void this.reconnect();
+      throw typedError;
+    }
+
+    const connected = await this.reconnectWithinDeadline(deadlineAt);
+    if (!connected) {
+      throw typedError;
+    }
+
+    return await this.executeConnectedCallWithinDeadline(
+      name,
+      args,
+      deadlineAt,
+      attempt + 1,
+    );
+  }
+
   private getRemainingDeadlineMs(deadlineAt: number): number {
     return deadlineAt - this.now();
   }