From 7f34098cf57fb2b4f28221278fd7226cc9978246 Mon Sep 17 00:00:00 2001 From: Huarch Date: Tue, 7 Jul 2026 12:05:46 +0800 Subject: [PATCH] refactor(chat): split stream helpers --- src/routes/chat.ts | 31 +----- src/routes/chatSession.ts | 30 ++++++ src/routes/chatStream.ts | 178 ++----------------------------- src/routes/chatTokenSmoother.ts | 172 +++++++++++++++++++++++++++++ tests/routes/chatSession.test.ts | 4 +- 5 files changed, 211 insertions(+), 204 deletions(-) create mode 100644 src/routes/chatTokenSmoother.ts diff --git a/src/routes/chat.ts b/src/routes/chat.ts index c34f262..08e333f 100644 --- a/src/routes/chat.ts +++ b/src/routes/chat.ts @@ -35,9 +35,11 @@ import { isPressureTrendDemoPrompt, } from "./pressureTrendDemo.js"; import { + buildForkedSessionUiState, buildPromptWithLearningContext, extractLatestFrontendTurn, generateSessionTitle, + resolveRequestedStreamSessionId, shouldGenerateSessionTitle, shouldRestoreConversationForRuntime, } from "./chatSession.js"; @@ -98,20 +100,6 @@ const toSessionUiStateContext = (sessionRecord: SessionRecord) => ({ sessionId: sessionRecord.sessionId, }); -export const buildForkedSessionUiState = ( - sourceState: { messages?: unknown[] } | null | undefined, - input: { - keepMessageCount: number; - targetSessionId: string; - }, -) => ({ - sessionId: input.targetSessionId, - isTitleManuallyEdited: false, - messages: Array.isArray(sourceState?.messages) - ? sourceState.messages.slice(0, input.keepMessageCount) - : [], -}); - const getSessionRunStatus = (sessionId: string) => activeRuns.get(sessionId)?.status ?? lastRunStatuses.get(sessionId); @@ -193,21 +181,6 @@ const pipeUiChunk = ( } }; -export const resolveRequestedStreamSessionId = (input: { - sessionId?: string; - promptText: string; - createClientSessionId: () => string; -}) => { - const sessionId = input.sessionId?.trim(); - if (sessionId) { - return sessionId; - } - if (isPressureTrendDemoPrompt(input.promptText)) { - return input.createClientSessionId(); - } - return undefined; -}; - export const buildChatRouter = ( sessionBridge: ChatSessionBridge, runtime: OpencodeRuntimeAdapter, diff --git a/src/routes/chatSession.ts b/src/routes/chatSession.ts index d6f0aa6..834e504 100644 --- a/src/routes/chatSession.ts +++ b/src/routes/chatSession.ts @@ -4,6 +4,7 @@ import { MemoryStore } from "../memory/store.js"; import { type OpencodeRuntimeAdapter } from "../runtime/opencode.js"; import { collectTextContent } from "./chatStream.js"; +import { isPressureTrendDemoPrompt } from "./pressureTrendDemo.js"; const TITLE_PROMPT_TIMEOUT_MS = 5000; const TITLE_CONTEXT_MESSAGE_LIMIT = 40; @@ -217,6 +218,35 @@ export const shouldRestoreConversationForRuntime = (options: { runtimeHasConversation: boolean; }) => !options.hadExistingSessionRecord || !options.runtimeHasConversation; +export const buildForkedSessionUiState = ( + sourceState: { messages?: unknown[] } | null | undefined, + input: { + keepMessageCount: number; + targetSessionId: string; + }, +) => ({ + sessionId: input.targetSessionId, + isTitleManuallyEdited: false, + messages: Array.isArray(sourceState?.messages) + ? sourceState.messages.slice(0, input.keepMessageCount) + : [], +}); + +export const resolveRequestedStreamSessionId = (input: { + sessionId?: string; + promptText: string; + createClientSessionId: () => string; +}) => { + const sessionId = input.sessionId?.trim(); + if (sessionId) { + return sessionId; + } + if (isPressureTrendDemoPrompt(input.promptText)) { + return input.createClientSessionId(); + } + return undefined; +}; + const buildRestoredConversationContext = (recentTurns: SessionTurnRecord[]) => { const formattedTurns = recentTurns .slice(-RESTORE_TURN_LIMIT) diff --git a/src/routes/chatStream.ts b/src/routes/chatStream.ts index d3e7833..650a1f4 100644 --- a/src/routes/chatStream.ts +++ b/src/routes/chatStream.ts @@ -1,5 +1,4 @@ import type { Event as OpencodeEvent, Part } from "@opencode-ai/sdk/v2"; -import { smoothStream, type TextStreamPart } from "ai"; import { writeLlmRequestAuditLog } from "../audit/llmRequestAudit.js"; import { type SupportedModel } from "../chat/models.js"; @@ -50,6 +49,12 @@ import { type TodoItemPayload, type TodoUpdatePayload, } from "./chatStreamEvents.js"; +import { createTokenSmoother } from "./chatTokenSmoother.js"; + +export { + createTokenSmoother, + detectTokenSmoothingChunk, +} from "./chatTokenSmoother.js"; export { collectTextContent, @@ -84,177 +89,6 @@ type ProgressPayload = { detail?: string; }; -type TokenSmootherOptions = { - enabled: boolean; - delayMs: number; - locale: string; - sessionId: string; - write: (event: string, data: Record) => void; -}; - -type SmoothTextStreamPart = TextStreamPart>; -type SmoothTextDeltaPart = Extract; -type SmoothTextEndPart = Extract; - -const segmenters = new Map(); -const CJK_SCRIPT_PATTERN = /[\p{Script=Han}\p{Script=Hiragana}\p{Script=Katakana}\p{Script=Hangul}]/u; -const CJK_PUNCTUATION_PATTERN = /[。！？!?，、；：,;:]/u; -const CJK_FALLBACK_TRIGGER_CHARS = 24; -const CJK_FALLBACK_MIN_CHARS = 12; -const CJK_FALLBACK_TARGET_CHARS = 18; - -function getTokenSmoothingSegmenter(locale: string) { - if (segmenters.has(locale)) { - return segmenters.get(locale) ?? null; - } - - const segmenter = - typeof Intl !== "undefined" && "Segmenter" in Intl - ? new Intl.Segmenter(locale, { granularity: "word" }) - : null; - segmenters.set(locale, segmenter); - return segmenter; -} - -function sliceCodePoints(value: string, count: number) { - return Array.from(value).slice(0, count).join(""); -} - -function codePointLength(value: string) { - return Array.from(value).length; -} - -function isCjkSmoothingContent(content: string, locale: string) { - return ( - CJK_SCRIPT_PATTERN.test(content) || - (/^(zh|ja|ko)(?:-|$)/i.test(locale) && CJK_PUNCTUATION_PATTERN.test(content)) - ); -} - -function detectCjkFallbackChunk(content: string, locale: string) { - if (codePointLength(content) < CJK_FALLBACK_TRIGGER_CHARS) { - return null; - } - - const segmenter = getTokenSmoothingSegmenter(locale); - if (segmenter) { - let candidate: string | null = null; - for (const segment of segmenter.segment(content)) { - const nextCandidate = content.slice(0, segment.index + segment.segment.length); - const nextLength = codePointLength(nextCandidate); - if (nextLength > CJK_FALLBACK_TARGET_CHARS) { - break; - } - if (nextLength >= CJK_FALLBACK_MIN_CHARS) { - candidate = nextCandidate; - } - } - if (candidate) { - return candidate; - } - } - - return sliceCodePoints(content, CJK_FALLBACK_TARGET_CHARS); -} - -export function detectTokenSmoothingChunk(content: string, locale = "zh") { - if (!content) { - return null; - } - - if (isCjkSmoothingContent(content, locale)) { - return detectCjkFallbackChunk(content, locale); - } - - const wordChunk = content.match(/^\S+\s*/u)?.[0]; - if (wordChunk) { - return wordChunk; - } - - return sliceCodePoints(content, 1); -} - -export function createTokenSmoother({ - enabled, - delayMs, - locale, - sessionId, - write, -}: TokenSmootherOptions) { - const emitChunk = (content: string) => { - write("token", { - session_id: sessionId, - content, - }); - }; - - if (!enabled) { - return { - writeToken(content: string) { - if (content) { - emitChunk(content); - } - }, - async flush() { - return; - }, - }; - } - - const textPartId = `smooth-text-${sessionId}`; - const stream = smoothStream>({ - delayInMs: delayMs, - chunking: (buffer) => detectTokenSmoothingChunk(buffer, locale), - })({ tools: {} }); - const writer = stream.writable.getWriter(); - const reader = stream.readable.getReader(); - let closed = false; - let writeChain = Promise.resolve(); - const readPromise = (async () => { - while (true) { - const result = await reader.read(); - if (result.done) { - return; - } - if (result.value.type === "text-delta") { - emitChunk(result.value.text); - } - } - })(); - - return { - writeToken(content: string) { - if (!content) { - return; - } - if (closed) { - return; - } - writeChain = writeChain.then(() => - writer.write({ - type: "text-delta", - id: textPartId, - text: content, - } satisfies SmoothTextDeltaPart), - ); - }, - async flush() { - if (closed) { - await readPromise; - return; - } - closed = true; - await writeChain; - await writer.write({ - type: "text-end", - id: textPartId, - } satisfies SmoothTextEndPart); - await writer.close(); - await readPromise; - }, - }; -} - const getPermissionTarget = (metadata: unknown) => { if (!isObjectRecord(metadata)) { return undefined; diff --git a/src/routes/chatTokenSmoother.ts b/src/routes/chatTokenSmoother.ts new file mode 100644 index 0000000..783be53 --- /dev/null +++ b/src/routes/chatTokenSmoother.ts @@ -0,0 +1,172 @@ +import { smoothStream, type TextStreamPart } from "ai"; + +type TokenSmootherOptions = { + enabled: boolean; + delayMs: number; + locale: string; + sessionId: string; + write: (event: string, data: Record) => void; +}; + +type SmoothTextStreamPart = TextStreamPart>; +type SmoothTextDeltaPart = Extract; +type SmoothTextEndPart = Extract; + +const segmenters = new Map(); +const CJK_SCRIPT_PATTERN = /[\p{Script=Han}\p{Script=Hiragana}\p{Script=Katakana}\p{Script=Hangul}]/u; +const CJK_PUNCTUATION_PATTERN = /[。！？!?，、；：,;:]/u; +const CJK_FALLBACK_TRIGGER_CHARS = 24; +const CJK_FALLBACK_MIN_CHARS = 12; +const CJK_FALLBACK_TARGET_CHARS = 18; + +function getTokenSmoothingSegmenter(locale: string) { + if (segmenters.has(locale)) { + return segmenters.get(locale) ?? null; + } + + const segmenter = + typeof Intl !== "undefined" && "Segmenter" in Intl + ? new Intl.Segmenter(locale, { granularity: "word" }) + : null; + segmenters.set(locale, segmenter); + return segmenter; +} + +function sliceCodePoints(value: string, count: number) { + return Array.from(value).slice(0, count).join(""); +} + +function codePointLength(value: string) { + return Array.from(value).length; +} + +function isCjkSmoothingContent(content: string, locale: string) { + return ( + CJK_SCRIPT_PATTERN.test(content) || + (/^(zh|ja|ko)(?:-|$)/i.test(locale) && CJK_PUNCTUATION_PATTERN.test(content)) + ); +} + +function detectCjkFallbackChunk(content: string, locale: string) { + if (codePointLength(content) < CJK_FALLBACK_TRIGGER_CHARS) { + return null; + } + + const segmenter = getTokenSmoothingSegmenter(locale); + if (segmenter) { + let candidate: string | null = null; + for (const segment of segmenter.segment(content)) { + const nextCandidate = content.slice(0, segment.index + segment.segment.length); + const nextLength = codePointLength(nextCandidate); + if (nextLength > CJK_FALLBACK_TARGET_CHARS) { + break; + } + if (nextLength >= CJK_FALLBACK_MIN_CHARS) { + candidate = nextCandidate; + } + } + if (candidate) { + return candidate; + } + } + + return sliceCodePoints(content, CJK_FALLBACK_TARGET_CHARS); +} + +export function detectTokenSmoothingChunk(content: string, locale = "zh") { + if (!content) { + return null; + } + + if (isCjkSmoothingContent(content, locale)) { + return detectCjkFallbackChunk(content, locale); + } + + const wordChunk = content.match(/^\S+\s*/u)?.[0]; + if (wordChunk) { + return wordChunk; + } + + return sliceCodePoints(content, 1); +} + +export function createTokenSmoother({ + enabled, + delayMs, + locale, + sessionId, + write, +}: TokenSmootherOptions) { + const emitChunk = (content: string) => { + write("token", { + session_id: sessionId, + content, + }); + }; + + if (!enabled) { + return { + writeToken(content: string) { + if (content) { + emitChunk(content); + } + }, + async flush() { + return; + }, + }; + } + + const textPartId = `smooth-text-${sessionId}`; + const stream = smoothStream>({ + delayInMs: delayMs, + chunking: (buffer) => detectTokenSmoothingChunk(buffer, locale), + })({ tools: {} }); + const writer = stream.writable.getWriter(); + const reader = stream.readable.getReader(); + let closed = false; + let writeChain = Promise.resolve(); + const readPromise = (async () => { + while (true) { + const result = await reader.read(); + if (result.done) { + return; + } + if (result.value.type === "text-delta") { + emitChunk(result.value.text); + } + } + })(); + + return { + writeToken(content: string) { + if (!content) { + return; + } + if (closed) { + return; + } + writeChain = writeChain.then(() => + writer.write({ + type: "text-delta", + id: textPartId, + text: content, + } satisfies SmoothTextDeltaPart), + ); + }, + async flush() { + if (closed) { + await readPromise; + return; + } + closed = true; + await writeChain; + await writer.write({ + type: "text-end", + id: textPartId, + } satisfies SmoothTextEndPart); + await writer.close(); + await readPromise; + }, + }; +} diff --git a/tests/routes/chatSession.test.ts b/tests/routes/chatSession.test.ts index 426068a..72dd707 100644 --- a/tests/routes/chatSession.test.ts +++ b/tests/routes/chatSession.test.ts @@ -2,12 +2,10 @@ import { describe, expect, it } from "bun:test"; import { buildForkedSessionUiState, - resolveRequestedStreamSessionId, -} from "../../src/routes/chat.js"; -import { buildPromptWithLearningContext, extractLatestFrontendTurn, generateSessionTitle, + resolveRequestedStreamSessionId, shouldRestoreConversationForRuntime, shouldGenerateSessionTitle, } from "../../src/routes/chatSession.js";