kanna-code 0.26.5 → 0.26.7

package/src/server/ws-router.ts

@@ -15,9 +15,32 @@ import { deriveChatSnapshot, deriveLocalProjectsSnapshot, deriveSidebarData } fr
 
 const DEFAULT_CHAT_RECENT_LIMIT = 200
 
+function isSendToStartingProfilingEnabled() {
+  return process.env.KANNA_PROFILE_SEND_TO_STARTING === "1"
+}
+
+function logSendToStartingProfile(
+  traceId: string | null | undefined,
+  startedAt: number | null | undefined,
+  stage: string,
+  details?: Record<string, unknown>
+) {
+  if (!traceId || startedAt === undefined || startedAt === null || !isSendToStartingProfilingEnabled()) {
+    return
+  }
+
+  console.log("[kanna/send->starting][server]", JSON.stringify({
+    traceId,
+    stage,
+    elapsedMs: Number((performance.now() - startedAt).toFixed(1)),
+    ...details,
+  }))
+}
+
 export interface ClientState {
   subscriptions: Map<string, SubscriptionTopic>
   snapshotSignatures: Map<string, string>
+  protectedDraftChatIds?: Set<string>
 }
 
 interface CreateWsRouterArgs {
@@ -83,9 +106,28 @@ export function createWsRouter({
     ])
   }
 
-  async function maybePruneStaleEmptyChats() {
+  function getProtectedDraftChatIds(extraSockets?: Iterable<ServerWebSocket<ClientState>>) {
+    const protectedChatIds = new Set<string>()
+
+    for (const socket of sockets) {
+      for (const chatId of socket.data.protectedDraftChatIds ?? []) {
+        protectedChatIds.add(chatId)
+      }
+    }
+
+    for (const socket of extraSockets ?? []) {
+      for (const chatId of socket.data.protectedDraftChatIds ?? []) {
+        protectedChatIds.add(chatId)
+      }
+    }
+
+    return protectedChatIds
+  }
+
+  async function maybePruneStaleEmptyChats(extraSockets?: Iterable<ServerWebSocket<ClientState>>) {
     await store.pruneStaleEmptyChats?.({
       activeChatIds: getProtectedChatIds(),
+      protectedChatIds: getProtectedDraftChatIds(extraSockets),
     })
   }
 
@@ -194,7 +236,7 @@ export function createWsRouter({
 
   async function pushSnapshots(ws: ServerWebSocket<ClientState>, options?: { skipPrune?: boolean }) {
     if (!options?.skipPrune) {
-      await maybePruneStaleEmptyChats()
+      await maybePruneStaleEmptyChats([ws])
     }
     const snapshotSignatures = ensureSnapshotSignatures(ws)
     for (const [id, topic] of ws.data.subscriptions.entries()) {
@@ -205,6 +247,14 @@ export function createWsRouter({
         continue
       }
       snapshotSignatures.set(id, signature)
+      if (topic.type === "chat" && envelope.snapshot.type === "chat" && envelope.snapshot.data?.runtime.status === "starting") {
+        const profile = agent.getActiveTurnProfile(topic.chatId)
+        logSendToStartingProfile(profile?.traceId, profile?.startedAt, "ws.snapshot_sent", {
+          chatId: topic.chatId,
+          status: envelope.snapshot.data.runtime.status,
+          messageCount: envelope.snapshot.data.messages.length,
+        })
+      }
       send(ws, envelope)
     }
   }
@@ -410,8 +460,19 @@ export function createWsRouter({
           send(ws, { v: PROTOCOL_VERSION, type: "ack", id })
           break
         }
+        case "chat.setDraftProtection": {
+          ws.data.protectedDraftChatIds = new Set(command.chatIds)
+          send(ws, { v: PROTOCOL_VERSION, type: "ack", id })
+          break
+        }
         case "chat.send": {
           const result = await agent.send(command)
+          const profile = command.clientTraceId && result.chatId
+            ? agent.getActiveTurnProfile(result.chatId)
+            : null
+          logSendToStartingProfile(profile?.traceId ?? command.clientTraceId, profile?.startedAt, "ws.chat_send_ack", {
+            chatId: result.chatId ?? null,
+          })
           send(ws, { v: PROTOCOL_VERSION, type: "ack", id, result })
           break
         }

package/src/shared/protocol.ts

@@ -67,11 +67,13 @@ export type ClientCommand =
   | { type: "chat.create"; projectId: string }
   | { type: "chat.rename"; chatId: string; title: string }
   | { type: "chat.delete"; chatId: string }
+  | { type: "chat.setDraftProtection"; chatIds: string[] }
   | { type: "chat.markRead"; chatId: string }
   | {
       type: "chat.send"
       chatId?: string
       projectId?: string
+      clientTraceId?: string
       provider?: AgentProvider
       content: string
       attachments?: ChatAttachment[]