@ponythewhite/base-context 1.0.14 → 1.0.15

package/dist/modes/daemon/daemon-mode.js

@@ -6383,6 +6383,10 @@ Use prime_context read/search with selected lines or a query.`;
             if (!state || !state.clients.has(client)) {
                 continue;
             }
+            const streamsSnapshot = client.transport === "private-framed" &&
+                daemonClientCapabilitiesForSession(client, activeSessionId).has("chunked_snapshot");
+            const snapshotSignal = streamsSnapshot ? markClientSnapshotStreaming(client, activeSessionId) : undefined;
+            let ownsSnapshot = streamsSnapshot;
             try {
                 const result = await this.createAttachResult(client, state, {
                     type: "attach",
@@ -6391,8 +6395,7 @@ Use prime_context read/search with selected lines or a query.`;
                 if (this.sessions.get(activeSessionId) !== state || !state.clients.has(client)) {
                     continue;
                 }
-                if (client.transport === "private-framed" &&
-                    daemonClientCapabilitiesForSession(client, activeSessionId).has("chunked_snapshot")) {
+                if (snapshotSignal) {
                     if (purpose === "replacement") {
                         this.write(client, {
                             type: "session_replaced",
@@ -6404,21 +6407,14 @@ Use prime_context read/search with selected lines or a query.`;
                         });
                     }
                     const snapshotId = `${result.activeSessionId}-${result.lastEventCursor?.generation}-${result.lastEventSequence}`;
-                    const snapshotSignal = markClientSnapshotStreaming(client, activeSessionId);
-                    let transcript;
-                    try {
-                        transcript = createSnapshotTranscriptChunks({
-                            activeSessionId,
-                            snapshotId,
-                            messages: result.snapshot.messages,
-                            targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES,
-                            signal: snapshotSignal,
-                        });
-                    }
-                    catch (error) {
-                        finishClientSnapshotStreaming(client, activeSessionId);
-                        throw error;
-                    }
+                    const transcript = createSnapshotTranscriptChunks({
+                        activeSessionId,
+                        snapshotId,
+                        messages: result.snapshot.messages,
+                        targetChunkBytes: SNAPSHOT_TARGET_CHUNK_BYTES,
+                        signal: snapshotSignal,
+                    });
+                    ownsSnapshot = false;
                     await this.streamWorkerSnapshot(client, {
                         ...result,
                         messages: result.messages ? [] : undefined,
@@ -6456,6 +6452,10 @@ Use prime_context read/search with selected lines or a query.`;
                 this.scheduleClientCatchupRetry(client);
                 return "retry-later";
             }
+            finally {
+                if (ownsSnapshot)
+                    finishClientSnapshotStreaming(client, activeSessionId);
+            }
         }
         return "drained";
     }