fix(chat): 解决token传输、本地文件存储顺序、读取的问题
This commit is contained in:
+34
-12
@@ -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"
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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([]);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user