Skip to content

Commit 8c2fca4

Browse files
author
NIO
committed
fix(session-fork): add 2s timeout to parent-fork token count resolution
Wrap the transcript read + byte-stat path in a withTimeout so a slow or very large parent transcript cannot stall session creation indefinitely. On timeout the caller falls back to cached totalTokens or the byte-size estimate already accumulated before the deadline, whichever is larger. Fixes #101718
1 parent c69724e commit 8c2fca4

3 files changed

Lines changed: 171 additions & 37 deletions

File tree

‎src/auto-reply/reply/session-fork.runtime.test.ts‎

Lines changed: 96 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,9 @@
22
import fs from "node:fs/promises";
33
import os from "node:os";
44
import path from "node:path";
5-
import { afterEach, describe, expect, it } from "vitest";
5+
import { afterEach, describe, expect, it, vi } from "vitest";
66
import type { SessionEntry } from "../../config/sessions/types.js";
7+
import * as sessionUtilsFs from "../../gateway/session-utils.fs.js";
78
import { resolveParentForkTokenCountRuntime } from "./session-fork.runtime.js";
89

910
const roots: string[] = [];
@@ -15,6 +16,8 @@ async function makeRoot(prefix: string): Promise<string> {
1516
}
1617

1718
afterEach(async () => {
19+
vi.useRealTimers();
20+
vi.restoreAllMocks();
1821
await Promise.all(roots.splice(0).map((root) => fs.rm(root, { recursive: true, force: true })));
1922
});
2023

@@ -363,4 +366,96 @@ describe("resolveParentForkTokenCountRuntime", () => {
363366
expect(tokens).toBeGreaterThan(100_000);
364367
expect(tokens).toBeLessThan(110_000);
365368
});
369+
370+
it("falls back to cached parent tokens when transcript usage hangs", async () => {
371+
vi.useFakeTimers();
372+
const root = await makeRoot("openclaw-parent-fork-timeout-cached-");
373+
const sessionsDir = path.join(root, "sessions");
374+
await fs.mkdir(sessionsDir);
375+
376+
const sessionId = "parent-timeout-cached";
377+
const sessionFile = path.join(sessionsDir, "parent.jsonl");
378+
await fs.writeFile(
379+
sessionFile,
380+
[
381+
JSON.stringify({
382+
type: "session",
383+
version: 3,
384+
id: sessionId,
385+
timestamp: new Date().toISOString(),
386+
cwd: process.cwd(),
387+
}),
388+
].join("\n"),
389+
"utf-8",
390+
);
391+
vi.spyOn(sessionUtilsFs, "readLatestRecentSessionUsageFromTranscriptAsync").mockImplementation(
392+
() => new Promise(() => {}),
393+
);
394+
395+
const tokensPromise = resolveParentForkTokenCountRuntime({
396+
parentEntry: {
397+
sessionId,
398+
sessionFile,
399+
updatedAt: Date.now(),
400+
totalTokens: 4_567,
401+
totalTokensFresh: false,
402+
},
403+
storePath: path.join(root, "sessions.json"),
404+
});
405+
406+
await vi.advanceTimersByTimeAsync(2_000);
407+
408+
await expect(tokensPromise).resolves.toBe(4_567);
409+
});
410+
411+
it("keeps the byte estimate when transcript usage times out after stat succeeds", async () => {
412+
vi.useFakeTimers();
413+
const root = await makeRoot("openclaw-parent-fork-timeout-bytes-");
414+
const sessionsDir = path.join(root, "sessions");
415+
await fs.mkdir(sessionsDir);
416+
417+
const sessionId = "parent-timeout-bytes";
418+
const sessionFile = path.join(sessionsDir, "parent.jsonl");
419+
const lines = [
420+
JSON.stringify({
421+
type: "session",
422+
version: 3,
423+
id: sessionId,
424+
timestamp: new Date().toISOString(),
425+
cwd: process.cwd(),
426+
}),
427+
];
428+
for (let index = 0; index < 24; index += 1) {
429+
lines.push(
430+
JSON.stringify({
431+
type: "message",
432+
id: `u${index}`,
433+
parentId: index === 0 ? null : `a${index - 1}`,
434+
timestamp: new Date().toISOString(),
435+
message: { role: "user", content: `turn-${index} ${"x".repeat(24_000)}` },
436+
}),
437+
);
438+
}
439+
await fs.writeFile(sessionFile, `${lines.join("\n")}\n`, "utf-8");
440+
vi.spyOn(sessionUtilsFs, "readLatestRecentSessionUsageFromTranscriptAsync").mockImplementation(
441+
() => new Promise(() => {}),
442+
);
443+
const stat = await fs.stat(sessionFile);
444+
vi.spyOn(fs, "stat").mockResolvedValue(stat);
445+
const expected = Math.ceil(stat.size / 4);
446+
447+
const tokensPromise = resolveParentForkTokenCountRuntime({
448+
parentEntry: {
449+
sessionId,
450+
sessionFile,
451+
updatedAt: Date.now(),
452+
totalTokensFresh: false,
453+
},
454+
storePath: path.join(root, "sessions.json"),
455+
});
456+
457+
await vi.advanceTimersByTimeAsync(2_000);
458+
459+
await expect(tokensPromise).resolves.toBe(expected);
460+
});
366461
});

‎src/auto-reply/reply/session-fork.runtime.ts‎

Lines changed: 53 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,10 @@ import {
1414
type SessionEntry as StoreSessionEntry,
1515
} from "../../config/sessions/types.js";
1616
import { readLatestRecentSessionUsageFromTranscriptAsync } from "../../gateway/session-utils.fs.js";
17+
import { withTimeout } from "../../utils/with-timeout.js";
1718

1819
const FALLBACK_TRANSCRIPT_BYTES_PER_TOKEN = 4;
20+
const PARENT_FORK_TOKEN_COUNT_TIMEOUT_MS = 2_000;
1921

2022
function resolvePositiveTokenCount(value: number | undefined): number | undefined {
2123
return typeof value === "number" && Number.isFinite(value) && value > 0
@@ -61,44 +63,60 @@ export async function resolveParentForkTokenCountRuntime(params: {
6163
return freshPersistedTokens;
6264
}
6365

64-
const cachedTokens = resolvePositiveTokenCount(params.parentEntry.totalTokens);
65-
const byteEstimateTokens = await estimateParentTranscriptTokensFromBytes(params);
66+
let bestEffortTokens = resolvePositiveTokenCount(params.parentEntry.totalTokens);
6667
try {
67-
const usage = await readLatestRecentSessionUsageFromTranscriptAsync(
68-
params.parentEntry.sessionId,
69-
params.storePath,
70-
params.parentEntry.sessionFile,
71-
undefined,
72-
1024 * 1024,
68+
return await withTimeout(
69+
(async () => {
70+
const byteEstimateTokens = await estimateParentTranscriptTokensFromBytes(params);
71+
bestEffortTokens = maxPositiveTokenCount(bestEffortTokens, byteEstimateTokens);
72+
const usage = await readLatestRecentSessionUsageFromTranscriptAsync(
73+
params.parentEntry.sessionId,
74+
params.storePath,
75+
params.parentEntry.sessionFile,
76+
undefined,
77+
1024 * 1024,
78+
);
79+
let transcriptTokens: number | undefined;
80+
if (usage?.contextUsage?.state === "available") {
81+
const trailingTokens = Math.ceil(
82+
(usage.trailingBytes ?? 0) / FALLBACK_TRANSCRIPT_BYTES_PER_TOKEN,
83+
);
84+
transcriptTokens = resolvePositiveTokenCount(
85+
usage.contextUsage.totalTokens + trailingTokens,
86+
);
87+
if (typeof transcriptTokens === "number") {
88+
return transcriptTokens;
89+
}
90+
} else if (usage?.contextUsage?.state !== "unavailable") {
91+
const promptTokens = resolvePositiveTokenCount(
92+
derivePromptTokens({
93+
input: usage?.inputTokens,
94+
cacheRead: usage?.cacheRead,
95+
cacheWrite: usage?.cacheWrite,
96+
}),
97+
);
98+
const outputTokens = resolvePositiveTokenCount(usage?.outputTokens);
99+
if (typeof promptTokens === "number") {
100+
transcriptTokens = promptTokens + (outputTokens ?? 0);
101+
}
102+
}
103+
if (typeof transcriptTokens === "number") {
104+
return maxPositiveTokenCount(transcriptTokens, bestEffortTokens);
105+
}
106+
return bestEffortTokens;
107+
})(),
108+
PARENT_FORK_TOKEN_COUNT_TIMEOUT_MS,
109+
{
110+
createError: () =>
111+
new Error(
112+
`parent fork token count timed out after ${PARENT_FORK_TOKEN_COUNT_TIMEOUT_MS}ms`,
113+
),
114+
},
73115
);
74-
let transcriptTokens: number | undefined;
75-
if (usage?.contextUsage?.state === "available") {
76-
const trailingTokens = Math.ceil(
77-
(usage.trailingBytes ?? 0) / FALLBACK_TRANSCRIPT_BYTES_PER_TOKEN,
78-
);
79-
transcriptTokens = resolvePositiveTokenCount(usage.contextUsage.totalTokens + trailingTokens);
80-
if (typeof transcriptTokens === "number") {
81-
return transcriptTokens;
82-
}
83-
} else if (usage?.contextUsage?.state !== "unavailable") {
84-
const promptTokens = resolvePositiveTokenCount(
85-
derivePromptTokens({
86-
input: usage?.inputTokens,
87-
cacheRead: usage?.cacheRead,
88-
cacheWrite: usage?.cacheWrite,
89-
}),
90-
);
91-
const outputTokens = resolvePositiveTokenCount(usage?.outputTokens);
92-
if (typeof promptTokens === "number") {
93-
transcriptTokens = promptTokens + (outputTokens ?? 0);
94-
}
95-
}
96-
if (typeof transcriptTokens === "number") {
97-
return maxPositiveTokenCount(transcriptTokens, cachedTokens, byteEstimateTokens);
98-
}
99116
} catch {
100-
// Fall back to cached totals when recent transcript usage cannot be read.
117+
// Keep parent-fork session creation moving even when transcript metadata
118+
// reads stall; the caller can still use cached or size-based estimates.
101119
}
102120

103-
return maxPositiveTokenCount(cachedTokens, byteEstimateTokens);
121+
return bestEffortTokens;
104122
}

‎src/auto-reply/reply/session-fork.test.ts‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import os from "node:os";
44
import path from "node:path";
55
import { afterEach, describe, expect, it, vi } from "vitest";
66
import type { OpenClawConfig } from "../../config/types.openclaw.js";
7-
import { forkSessionEntryFromParent } from "./session-fork.js";
7+
import { forkSessionEntryFromParent, resolveParentForkDecision } from "./session-fork.js";
88

99
const runtimeMocks = vi.hoisted(() => ({
1010
resolveParentForkTokenCountRuntime: vi.fn(),
@@ -101,3 +101,24 @@ describe("forkSessionEntryFromParent", () => {
101101
expect(stored[sessionKey]?.sessionFile).toBe(result.fork.sessionFile);
102102
});
103103
});
104+
105+
describe("resolveParentForkDecision", () => {
106+
it("keeps forking when the runtime cannot produce a bounded parent token count", async () => {
107+
runtimeMocks.resolveParentForkTokenCountRuntime.mockResolvedValue(undefined);
108+
109+
await expect(
110+
resolveParentForkDecision({
111+
parentEntry: {
112+
sessionId: "parent-session",
113+
sessionFile: "/tmp/parent.jsonl",
114+
updatedAt: Date.now(),
115+
totalTokensFresh: false,
116+
},
117+
storePath: "/tmp/sessions.json",
118+
}),
119+
).resolves.toEqual({
120+
status: "fork",
121+
maxTokens: 100_000,
122+
});
123+
});
124+
});

0 commit comments

Comments
 (0)