diff --git a/server/src/ws.ts b/server/src/ws.ts index a9b90b1..bb573f4 100644 --- a/server/src/ws.ts +++ b/server/src/ws.ts @@ -15,6 +15,7 @@ import { userFacingError, type ChatMessage, } from "./services/llm" +import { classifyMessage, inBackground } from "./services/classifier" type Kind = "prompt" | "guidance" @@ -116,11 +117,12 @@ async function handlePrompt(ws: WS, data: any) { const model = String(data.model ?? "") if (!prompt) return - ws.data.pendingUserMessageId = insertMessage({ + const userMessageId = insertMessage({ conversation_id: ws.data.conversationId, role: "user", content: prompt, }) + ws.data.pendingUserMessageId = userMessageId try { const history: ChatMessage[] = conversationHistory(ws.data.conversationId).map((m) => ({ role: m.role as ChatMessage["role"], @@ -146,6 +148,8 @@ async function handlePrompt(ws: WS, data: any) { code_css: code.css, code_js: code.js, }) + // 有了回复,这条提问才留下来;在后台打 L1–L6 层级 + inBackground(classifyMessage(userMessageId)) ws.data.pendingUserMessageId = null send(ws, { type: "complete", code, message_id: messageId }) } finally { @@ -157,12 +161,13 @@ async function handleGuidance(ws: WS, data: any) { const prompt = String(data.content ?? "").trim() if (!prompt) return - ws.data.pendingUserMessageId = insertMessage({ + const userMessageId = insertMessage({ conversation_id: ws.data.conversationId, role: "user", content: prompt, source: "guidance", }) + ws.data.pendingUserMessageId = userMessageId ws.data.guidanceHistory.push({ role: "user", content: prompt }) try { let full: string @@ -185,6 +190,7 @@ async function handleGuidance(ws: WS, data: any) { content: clean, source: "guidance", }) + inBackground(classifyMessage(userMessageId)) ws.data.pendingUserMessageId = null send(ws, { type: "complete", message_id: messageId, is_ready: isReady }) } finally { diff --git a/server/tests/ws.test.ts b/server/tests/ws.test.ts index 94f0b41..6cc17db 100644 --- a/server/tests/ws.test.ts +++ b/server/tests/ws.test.ts @@ -8,12 +8,18 @@ import { client, freshDb, makeTask, makeUser } from "./helpers" // 假的模型接口:每段之间停一下,方便测试中途断开 const PIECES = ["好的", "```html\n
hi
\n```"] -// 最近一次发给模型的请求 +// 最近一次发给模型的请求;分类(非流式)单独记,固定回“3” let lastRequest: any = null +let lastClassify: any = null const llm = Bun.serve({ port: 0, async fetch(req) { - lastRequest = await req.json() + const body: any = await req.json() + if (!body.stream) { + lastClassify = body + return Response.json({ choices: [{ message: { content: "3" } }] }) + } + lastRequest = body const enc = new TextEncoder() const stream = new ReadableStream({ async start(ctrl) { @@ -106,6 +112,19 @@ describe("WebSocket 代码生成", () => { c.ws.close() }) + test("回复完后在后台给提问打层级", async () => { + const { task, token } = await setup() + const c = (await connect(`/ws/prompt/${task}/?token=${token}`))! + await c.next("init") + c.ws.send(JSON.stringify({ type: "message", content: "做个页面" })) + await c.next("complete") + const level = () => all(`SELECT prompt_level FROM prompt_message WHERE role = 'user'`)[0]?.prompt_level + for (let i = 0; i < 50 && level() == null; i++) await Bun.sleep(20) + expect(level()).toBe(3) + expect(lastClassify.thinking).toEqual({ type: "disabled" }) + c.ws.close() + }) + test("AI 回复中再发消息会被拒绝", async () => { const { task, token } = await setup() const c = (await connect(`/ws/prompt/${task}/?token=${token}`))!