fix(gateway): cache spawnedBy lineage lookups

Signed-off-by: samzong <samzong.lu@gmail.com>
This commit is contained in:
samzong
2026-04-29 14:29:03 +08:00
committed by Frank Yang
parent c3d49d70f9
commit 849307d3a4
2 changed files with 110 additions and 3 deletions

View File

@@ -1866,5 +1866,101 @@ describe("agent event handler", () => {
resetAgentRunContextForTest();
});
it("caches spawnedBy lookup so repeated events for the same subagent session only load the row once", () => {
vi.mocked(loadGatewaySessionRow).mockClear();
vi.mocked(loadGatewaySessionRow).mockReturnValue({
key: "agent:coder:subagent:cache-test",
kind: "direct",
updatedAt: null,
spawnedBy: "agent:conductor:task:parent-cache",
});
const { broadcast, handler, chatRunState } = createHarness({
resolveSessionKeyForRun: () => "agent:coder:subagent:cache-test",
});
chatRunState.registry.add("run-cache", {
sessionKey: "agent:coder:subagent:cache-test",
clientRunId: "client-cache",
});
// Fire multiple events for the same session
handler({
runId: "run-cache",
seq: 1,
stream: "assistant",
ts: Date.now(),
data: { text: "chunk 1" },
});
handler({
runId: "run-cache",
seq: 2,
stream: "assistant",
ts: Date.now(),
data: { text: "chunk 2" },
});
handler({
runId: "run-cache",
seq: 3,
stream: "lifecycle",
ts: Date.now(),
data: { phase: "end" },
});
// Key assertion: loadGatewaySessionRow called exactly once despite 3 events
expect(loadGatewaySessionRow).toHaveBeenCalledTimes(1);
expect(loadGatewaySessionRow).toHaveBeenCalledWith("agent:coder:subagent:cache-test");
// All broadcasts still have correct spawnedBy
const chatCalls = chatBroadcastCalls(broadcast);
for (const [, payload] of chatCalls) {
expect(payload).toMatchObject({
spawnedBy: "agent:conductor:task:parent-cache",
});
}
});
it("caches null spawnedBy for eligible subagent sessions that lack a spawnedBy value", () => {
vi.mocked(loadGatewaySessionRow).mockClear();
vi.mocked(loadGatewaySessionRow).mockReturnValue({
key: "agent:coder:subagent:no-lineage",
kind: "direct",
updatedAt: null,
// no spawnedBy field
});
const { broadcast, handler, chatRunState } = createHarness({
resolveSessionKeyForRun: () => "agent:coder:subagent:no-lineage",
});
chatRunState.registry.add("run-null", {
sessionKey: "agent:coder:subagent:no-lineage",
clientRunId: "client-null",
});
handler({
runId: "run-null",
seq: 1,
stream: "assistant",
ts: Date.now(),
data: { text: "chunk 1" },
});
handler({
runId: "run-null",
seq: 2,
stream: "assistant",
ts: Date.now(),
data: { text: "chunk 2" },
});
// null result is cached — only one DB call despite two events
expect(loadGatewaySessionRow).toHaveBeenCalledTimes(1);
const chatCalls = chatBroadcastCalls(broadcast);
for (const [, payload] of chatCalls) {
expect(payload).not.toHaveProperty("spawnedBy");
}
});
});
});

View File

@@ -188,16 +188,27 @@ export function createAgentEventHandler({
// Only subagent/acp keys can carry spawnedBy (mirrors supportsSpawnLineage in
// sessions-patch.ts). Short-circuit everyone else so high-volume chat streams
// do not touch the session store.
// do not touch the session store. Results are cached per sessionKey because
// spawnedBy is immutable once set and resolveSpawnedBy sits on the hot event
// path (delta, flush, final, agent, seq-gap).
const spawnedByCache = new Map<string, string | null>();
const resolveSpawnedBy = (sessionKey: string): string | null => {
if (spawnedByCache.has(sessionKey)) {
return spawnedByCache.get(sessionKey)!;
}
// Non-lineage keys return null without polluting the cache; only
// subagent/ACP results (positive or null) are worth memoising.
if (!isSubagentSessionKey(sessionKey) && !isAcpSessionKey(sessionKey)) {
return null;
}
let result: string | null = null;
try {
return loadGatewaySessionRow(sessionKey)?.spawnedBy ?? null;
result = loadGatewaySessionRow(sessionKey)?.spawnedBy ?? null;
} catch {
return null;
// result stays null
}
spawnedByCache.set(sessionKey, result);
return result;
};
const buildSessionEventSnapshot = (sessionKey: string, evt?: AgentEventPayload) => {