fix(chat): 解决token传输、本地文件存储顺序、读取的问题

This commit is contained in:
2026-06-04 18:19:29 +08:00
parent 10c11a5254
commit fc0e76439d
3 changed files with 76 additions and 16 deletions
+34 -12
View File
@@ -823,12 +823,30 @@ export const buildChatRouter = (
};
activeRuns.set(clientSessionId, activeRun);
lastRunStatuses.set(clientSessionId, "running");
await sessionUiStateStore.write(toSessionUiStateContext(activeSessionRecord), {
const sessionUiStateContext = toSessionUiStateContext(activeSessionRecord);
let persistQueue = sessionUiStateStore.write(sessionUiStateContext, {
sessionId: activeSessionRecord.sessionId,
isTitleManuallyEdited: initialSessionState?.isTitleManuallyEdited ?? false,
messages: initialMessages,
branchGroups,
});
const queueSessionUiStatePersist = () => {
const snapshot = {
sessionId: activeSessionRecord.sessionId,
isTitleManuallyEdited: initialSessionState?.isTitleManuallyEdited ?? false,
messages: activeRun.messages,
branchGroups,
};
persistQueue = persistQueue
.catch((error) => {
logger.warn(
{ err: error, sessionId: clientSessionId },
"failed to persist previous chat stream state",
);
})
.then(() => sessionUiStateStore.write(sessionUiStateContext, snapshot));
return persistQueue;
};
const primarySubscriber: StreamSubscriber = {
write: (event, data) => {
if (!streamClosed && !res.writableEnded && !res.destroyed) {
@@ -850,7 +868,7 @@ export const buildChatRouter = (
req.on("close", handleClientClose);
res.on("close", handleClientClose);
const publish = async (event: string, data: Record<string, unknown>) => {
const publish = (event: string, data: Record<string, unknown>) => {
if (event === "token") {
activeRun.messages = updateLastAssistantMessage(activeRun.messages, (message) => ({
...message,
@@ -887,15 +905,12 @@ export const buildChatRouter = (
}));
}
await sessionUiStateStore.write(toSessionUiStateContext(activeSessionRecord), {
sessionId: activeSessionRecord.sessionId,
isTitleManuallyEdited: initialSessionState?.isTitleManuallyEdited ?? false,
messages: activeRun.messages,
branchGroups,
});
for (const subscriber of activeRun.subscribers) {
subscriber.write(event, data);
}
void queueSessionUiStatePersist().catch((error) => {
logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
});
};
try {
@@ -920,11 +935,12 @@ export const buildChatRouter = (
projectId: requestContext.projectId,
signal: abortController.signal,
write: (event, data) => {
void publish(event, data).catch((error) => {
logger.warn({ err: error, sessionId: clientSessionId }, "failed to publish chat stream event");
});
publish(event, data);
},
});
await persistQueue.catch((error) => {
logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
});
if (!streamResult.aborted && !streamResult.failed) {
const messages = await runtime.messages(binding.sessionId, 60);
@@ -965,13 +981,19 @@ export const buildChatRouter = (
sessionTitle &&
sessionTitle !== existingSessionTitle
) {
await publish("session_title", {
publish("session_title", {
session_id: clientSessionId,
title: sessionTitle,
});
await persistQueue.catch((error) => {
logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
});
}
}
} finally {
await persistQueue.catch((error) => {
logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
});
sessionBridge.finalizeRequest(clientSessionId);
activeRun.status = abortController.signal.aborted
? activeRun.status === "aborted"
+9 -4
View File
@@ -1,4 +1,4 @@
import { createHash } from "node:crypto";
import { createHash, randomUUID } from "node:crypto";
import { mkdir, readFile, readdir, rename, rm, stat, writeFile } from "node:fs/promises";
import { basename, dirname, join, relative } from "node:path";
@@ -13,9 +13,14 @@ export const ensureDirectory = async (path: string) => {
export const atomicWriteFile = async (path: string, content: string) => {
await ensureDirectory(dirname(path));
const tempPath = `${path}.${process.pid}.${Date.now().toString(36)}.tmp`;
await writeFile(tempPath, content, "utf8");
await rename(tempPath, path);
const tempPath = `${path}.${process.pid}.${Date.now().toString(36)}.${randomUUID()}.tmp`;
try {
await writeFile(tempPath, content, "utf8");
await rename(tempPath, path);
} catch (error) {
await removeFileIfExists(tempPath);
throw error;
}
};
type HistoricalWriteOptions = {
+33
View File
@@ -0,0 +1,33 @@
import { afterEach, beforeEach, describe, expect, it } from "bun:test";
import { mkdtemp, readdir, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { atomicWriteFile, readTextFile } from "../../src/utils/fileStore.js";
describe("fileStore", () => {
const originalDateNow = Date.now;
let tempDir: string;
beforeEach(async () => {
tempDir = await mkdtemp(join(tmpdir(), "tjwater-file-store-"));
});
afterEach(async () => {
Date.now = originalDateNow;
await rm(tempDir, { force: true, recursive: true });
});
it("uses unique temp paths for concurrent writes in the same millisecond", async () => {
Date.now = () => 1_801_578_600_000;
const path = join(tempDir, "state.json");
const values = Array.from({ length: 24 }, (_, index) => `value-${index}`);
await Promise.all(values.map((value) => atomicWriteFile(path, value)));
const written = await readTextFile(path);
expect(written).not.toBeNull();
expect(values).toContain(written as string);
expect((await readdir(tempDir)).filter((name) => name.endsWith(".tmp"))).toEqual([]);
});
});