@@ -0,0 +1,295 @@
1+export type ChatRunEntry = {
2+sessionKey: string;
3+clientRunId: string;
4+};
5+6+export type ChatRunRegistry = {
7+add: (sessionId: string, entry: ChatRunEntry) => void;
8+peek: (sessionId: string) => ChatRunEntry | undefined;
9+shift: (sessionId: string) => ChatRunEntry | undefined;
10+remove: (sessionId: string, clientRunId: string, sessionKey?: string) => ChatRunEntry | undefined;
11+clear: () => void;
12+};
13+14+export function createChatRunRegistry(): ChatRunRegistry {
15+const chatRunSessions = new Map<string, ChatRunEntry[]>();
16+17+const add = (sessionId: string, entry: ChatRunEntry) => {
18+const queue = chatRunSessions.get(sessionId);
19+if (queue) {
20+queue.push(entry);
21+} else {
22+chatRunSessions.set(sessionId, [entry]);
23+}
24+};
25+26+const peek = (sessionId: string) => chatRunSessions.get(sessionId)?.[0];
27+28+const shift = (sessionId: string) => {
29+const queue = chatRunSessions.get(sessionId);
30+if (!queue || queue.length === 0) {
31+return undefined;
32+}
33+const entry = queue.shift();
34+if (!queue.length) {
35+chatRunSessions.delete(sessionId);
36+}
37+return entry;
38+};
39+40+const remove = (sessionId: string, clientRunId: string, sessionKey?: string) => {
41+const queue = chatRunSessions.get(sessionId);
42+if (!queue || queue.length === 0) {
43+return undefined;
44+}
45+const idx = queue.findIndex(
46+(entry) =>
47+entry.clientRunId === clientRunId && (sessionKey ? entry.sessionKey === sessionKey : true),
48+);
49+if (idx < 0) {
50+return undefined;
51+}
52+const [entry] = queue.splice(idx, 1);
53+if (!queue.length) {
54+chatRunSessions.delete(sessionId);
55+}
56+return entry;
57+};
58+59+const clear = () => {
60+chatRunSessions.clear();
61+};
62+63+return { add, peek, shift, remove, clear };
64+}
65+66+export type ChatRunState = {
67+registry: ChatRunRegistry;
68+rawBuffers: Map<string, string>;
69+buffers: Map<string, string>;
70+deltaSentAt: Map<string, number>;
71+/** Length of text at the time of the last broadcast, used to avoid duplicate flushes. */
72+deltaLastBroadcastLen: Map<string, number>;
73+abortedRuns: Map<string, number>;
74+clear: () => void;
75+};
76+77+export function createChatRunState(): ChatRunState {
78+const registry = createChatRunRegistry();
79+const rawBuffers = new Map<string, string>();
80+const buffers = new Map<string, string>();
81+const deltaSentAt = new Map<string, number>();
82+const deltaLastBroadcastLen = new Map<string, number>();
83+const abortedRuns = new Map<string, number>();
84+85+const clear = () => {
86+registry.clear();
87+rawBuffers.clear();
88+buffers.clear();
89+deltaSentAt.clear();
90+deltaLastBroadcastLen.clear();
91+abortedRuns.clear();
92+};
93+94+return {
95+ registry,
96+ rawBuffers,
97+ buffers,
98+ deltaSentAt,
99+ deltaLastBroadcastLen,
100+ abortedRuns,
101+ clear,
102+};
103+}
104+105+export type ToolEventRecipientRegistry = {
106+add: (runId: string, connId: string) => void;
107+get: (runId: string) => ReadonlySet<string> | undefined;
108+markFinal: (runId: string) => void;
109+};
110+111+export type SessionEventSubscriberRegistry = {
112+subscribe: (connId: string) => void;
113+unsubscribe: (connId: string) => void;
114+getAll: () => ReadonlySet<string>;
115+clear: () => void;
116+};
117+118+export type SessionMessageSubscriberRegistry = {
119+subscribe: (connId: string, sessionKey: string) => void;
120+unsubscribe: (connId: string, sessionKey: string) => void;
121+unsubscribeAll: (connId: string) => void;
122+get: (sessionKey: string) => ReadonlySet<string>;
123+clear: () => void;
124+};
125+126+type ToolRecipientEntry = {
127+connIds: Set<string>;
128+updatedAt: number;
129+finalizedAt?: number;
130+};
131+132+const TOOL_EVENT_RECIPIENT_TTL_MS = 10 * 60 * 1000;
133+const TOOL_EVENT_RECIPIENT_FINAL_GRACE_MS = 30 * 1000;
134+135+export function createSessionEventSubscriberRegistry(): SessionEventSubscriberRegistry {
136+const connIds = new Set<string>();
137+const empty = new Set<string>();
138+139+return {
140+subscribe: (connId: string) => {
141+const normalized = connId.trim();
142+if (!normalized) {
143+return;
144+}
145+connIds.add(normalized);
146+},
147+unsubscribe: (connId: string) => {
148+const normalized = connId.trim();
149+if (!normalized) {
150+return;
151+}
152+connIds.delete(normalized);
153+},
154+getAll: () => (connIds.size > 0 ? connIds : empty),
155+clear: () => {
156+connIds.clear();
157+},
158+};
159+}
160+161+export function createSessionMessageSubscriberRegistry(): SessionMessageSubscriberRegistry {
162+const sessionToConnIds = new Map<string, Set<string>>();
163+const connToSessionKeys = new Map<string, Set<string>>();
164+const empty = new Set<string>();
165+166+const normalize = (value: string): string => value.trim();
167+168+return {
169+subscribe: (connId: string, sessionKey: string) => {
170+const normalizedConnId = normalize(connId);
171+const normalizedSessionKey = normalize(sessionKey);
172+if (!normalizedConnId || !normalizedSessionKey) {
173+return;
174+}
175+const connIds = sessionToConnIds.get(normalizedSessionKey) ?? new Set<string>();
176+connIds.add(normalizedConnId);
177+sessionToConnIds.set(normalizedSessionKey, connIds);
178+179+const sessionKeys = connToSessionKeys.get(normalizedConnId) ?? new Set<string>();
180+sessionKeys.add(normalizedSessionKey);
181+connToSessionKeys.set(normalizedConnId, sessionKeys);
182+},
183+unsubscribe: (connId: string, sessionKey: string) => {
184+const normalizedConnId = normalize(connId);
185+const normalizedSessionKey = normalize(sessionKey);
186+if (!normalizedConnId || !normalizedSessionKey) {
187+return;
188+}
189+const connIds = sessionToConnIds.get(normalizedSessionKey);
190+if (connIds) {
191+connIds.delete(normalizedConnId);
192+if (connIds.size === 0) {
193+sessionToConnIds.delete(normalizedSessionKey);
194+}
195+}
196+const sessionKeys = connToSessionKeys.get(normalizedConnId);
197+if (sessionKeys) {
198+sessionKeys.delete(normalizedSessionKey);
199+if (sessionKeys.size === 0) {
200+connToSessionKeys.delete(normalizedConnId);
201+}
202+}
203+},
204+unsubscribeAll: (connId: string) => {
205+const normalizedConnId = normalize(connId);
206+if (!normalizedConnId) {
207+return;
208+}
209+const sessionKeys = connToSessionKeys.get(normalizedConnId);
210+if (!sessionKeys) {
211+return;
212+}
213+for (const sessionKey of sessionKeys) {
214+const connIds = sessionToConnIds.get(sessionKey);
215+if (!connIds) {
216+continue;
217+}
218+connIds.delete(normalizedConnId);
219+if (connIds.size === 0) {
220+sessionToConnIds.delete(sessionKey);
221+}
222+}
223+connToSessionKeys.delete(normalizedConnId);
224+},
225+get: (sessionKey: string) => {
226+const normalizedSessionKey = normalize(sessionKey);
227+if (!normalizedSessionKey) {
228+return empty;
229+}
230+return sessionToConnIds.get(normalizedSessionKey) ?? empty;
231+},
232+clear: () => {
233+sessionToConnIds.clear();
234+connToSessionKeys.clear();
235+},
236+};
237+}
238+239+export function createToolEventRecipientRegistry(): ToolEventRecipientRegistry {
240+const recipients = new Map<string, ToolRecipientEntry>();
241+242+const prune = () => {
243+if (recipients.size === 0) {
244+return;
245+}
246+const now = Date.now();
247+for (const [runId, entry] of recipients) {
248+const cutoff = entry.finalizedAt
249+ ? entry.finalizedAt + TOOL_EVENT_RECIPIENT_FINAL_GRACE_MS
250+ : entry.updatedAt + TOOL_EVENT_RECIPIENT_TTL_MS;
251+if (now >= cutoff) {
252+recipients.delete(runId);
253+}
254+}
255+};
256+257+const add = (runId: string, connId: string) => {
258+if (!runId || !connId) {
259+return;
260+}
261+const now = Date.now();
262+const existing = recipients.get(runId);
263+if (existing) {
264+existing.connIds.add(connId);
265+existing.updatedAt = now;
266+} else {
267+recipients.set(runId, {
268+connIds: new Set([connId]),
269+updatedAt: now,
270+});
271+}
272+prune();
273+};
274+275+const get = (runId: string) => {
276+const entry = recipients.get(runId);
277+if (!entry) {
278+return undefined;
279+}
280+entry.updatedAt = Date.now();
281+prune();
282+return entry.connIds;
283+};
284+285+const markFinal = (runId: string) => {
286+const entry = recipients.get(runId);
287+if (!entry) {
288+return;
289+}
290+entry.finalizedAt = Date.now();
291+prune();
292+};
293+294+return { add, get, markFinal };
295+}
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。