Skip to content

Commit 6eb6fab

Browse files
mcintyre94claude
andcommitted
Fix VM leak on chat close/delete, remove stale connecting guard
- Call chatSessionManager.remove() on both closeChat and deleteChat in SpriteNavigationPanel and ChatSwitcherSheet, so cached VMs are detached and evicted immediately rather than leaking with an active WebSocket stream - Remove .disabled(isConnecting) on the new chat button — switching chats no longer cancels streams so the guard is unnecessary - Add ChatSessionManagerTests covering cache hit, eviction, detachAll, and isStreaming behaviour - Add runReconnectLoop status transition test (.reconnecting → .idle) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 577d4cd commit 6eb6fab

File tree

6 files changed

+151
-1
lines changed

6 files changed

+151
-1
lines changed

Wisp.xcodeproj/project.pbxproj

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
AA3400BB00CC00DD00EE0001 /* QuickChatView.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA3400BB00CC00DD00EE0002 /* QuickChatView.swift */; };
4747
AA3500BB00CC00DD00EE0001 /* BashQuickView.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA3500BB00CC00DD00EE0002 /* BashQuickView.swift */; };
4848
AA3700BB00CC00DD00EE0001 /* BashQuickViewModelTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA3700BB00CC00DD00EE0002 /* BashQuickViewModelTests.swift */; };
49+
AA3800BB00CC00DD00EE0001 /* ChatSessionManagerTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA3800BB00CC00DD00EE0002 /* ChatSessionManagerTests.swift */; };
4950
FE9876543210ABCDEF987601 /* QuickChatViewModelTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = FE9876543210ABCDEF987602 /* QuickChatViewModelTests.swift */; };
5051
AA00BB11CC22DD3300000001 /* WispAskCard.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA00BB11CC22DD3300000002 /* WispAskCard.swift */; };
5152
AA11BB22CC33DD44EE55FF02 /* GitHubDeviceFlowClient.swift in Sources */ = {isa = PBXBuildFile; fileRef = AA11BB22CC33DD44EE55FF01 /* GitHubDeviceFlowClient.swift */; };
@@ -251,6 +252,7 @@
251252
AA3400BB00CC00DD00EE0002 /* QuickChatView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = QuickChatView.swift; sourceTree = "<group>"; };
252253
AA3500BB00CC00DD00EE0002 /* BashQuickView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = BashQuickView.swift; sourceTree = "<group>"; };
253254
AA3700BB00CC00DD00EE0002 /* BashQuickViewModelTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = BashQuickViewModelTests.swift; sourceTree = "<group>"; };
255+
AA3800BB00CC00DD00EE0002 /* ChatSessionManagerTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ChatSessionManagerTests.swift; sourceTree = "<group>"; };
254256
FE9876543210ABCDEF987602 /* QuickChatViewModelTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = QuickChatViewModelTests.swift; sourceTree = "<group>"; };
255257
/* End PBXFileReference section */
256258

@@ -474,6 +476,7 @@
474476
EE22FF334455AA66BB770001 /* ExecSessionURLTests.swift */,
475477
CC11DD22EE33FF4400110005 /* FileEntryTests.swift */,
476478
AA3700BB00CC00DD00EE0002 /* BashQuickViewModelTests.swift */,
479+
AA3800BB00CC00DD00EE0002 /* ChatSessionManagerTests.swift */,
477480
FE9876543210ABCDEF987602 /* QuickChatViewModelTests.swift */,
478481
);
479482
path = WispTests;
@@ -817,6 +820,7 @@
817820
0DFF41BB14AC3922E95DF781 /* WorktreeTests.swift in Sources */,
818821
CC11DD22EE33FF4400110006 /* FileEntryTests.swift in Sources */,
819822
AA3700BB00CC00DD00EE0001 /* BashQuickViewModelTests.swift in Sources */,
823+
AA3800BB00CC00DD00EE0001 /* ChatSessionManagerTests.swift in Sources */,
820824
FE9876543210ABCDEF987601 /* QuickChatViewModelTests.swift in Sources */,
821825
);
822826
runOnlyForDeploymentPostprocessing = 0;

Wisp/Views/SpriteDetail/Chat/ChatSwitcherSheet.swift

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import SwiftData
44
struct ChatSwitcherSheet: View {
55
@Bindable var viewModel: SpriteChatListViewModel
66
@Environment(SpritesAPIClient.self) private var apiClient
7+
@Environment(ChatSessionManager.self) private var chatSessionManager
78
@Environment(\.modelContext) private var modelContext
89
@Environment(\.dismiss) private var dismiss
910
@State private var chatToDelete: SpriteChat?
@@ -39,6 +40,7 @@ struct ChatSwitcherSheet: View {
3940
}
4041
if !chat.isClosed {
4142
Button {
43+
chatSessionManager.remove(chatId: chat.id, modelContext: modelContext)
4244
viewModel.closeChat(chat, apiClient: apiClient, modelContext: modelContext)
4345
} label: {
4446
Label("Close", systemImage: "xmark.circle")
@@ -71,6 +73,7 @@ struct ChatSwitcherSheet: View {
7173

7274
if !chat.isClosed {
7375
Button {
76+
chatSessionManager.remove(chatId: chat.id, modelContext: modelContext)
7477
viewModel.closeChat(chat, apiClient: apiClient, modelContext: modelContext)
7578
} label: {
7679
Label("Close", systemImage: "xmark.circle")
@@ -87,6 +90,7 @@ struct ChatSwitcherSheet: View {
8790
titleVisibility: .visible
8891
) {
8992
Button("Delete", role: .destructive) {
93+
chatSessionManager.remove(chatId: chat.id, modelContext: modelContext)
9094
viewModel.deleteChat(chat, apiClient: apiClient, modelContext: modelContext)
9195
chatToDelete = nil
9296
}

Wisp/Views/SpriteDetail/SpriteDetailView.swift

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,6 @@ struct SpriteDetailView: View {
185185
} label: {
186186
Image(systemName: "square.and.pencil")
187187
}
188-
.disabled(chatViewModel?.status.isConnecting == true)
189188
}
190189
}
191190
} else if selectedTab == .overview {

Wisp/Views/SpriteDetail/SpriteNavigationPanel.swift

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ struct SpriteNavigationPanel: View {
1313
let chatListViewModel: SpriteChatListViewModel
1414
let onCreateChat: () -> Void
1515
@Environment(SpritesAPIClient.self) private var apiClient
16+
@Environment(ChatSessionManager.self) private var chatSessionManager
1617
@Environment(\.modelContext) private var modelContext
1718
@State private var chatToRename: SpriteChat?
1819
@State private var renameText = ""
@@ -97,6 +98,7 @@ struct SpriteNavigationPanel: View {
9798
}
9899
if !chat.isClosed {
99100
Button {
101+
chatSessionManager.remove(chatId: chat.id, modelContext: modelContext)
100102
chatListViewModel.closeChat(chat, apiClient: apiClient, modelContext: modelContext)
101103
} label: {
102104
Label("Close", systemImage: "xmark.circle")
@@ -117,6 +119,7 @@ struct SpriteNavigationPanel: View {
117119
}
118120
if !chat.isClosed {
119121
Button {
122+
chatSessionManager.remove(chatId: chat.id, modelContext: modelContext)
120123
chatListViewModel.closeChat(chat, apiClient: apiClient, modelContext: modelContext)
121124
} label: {
122125
Label("Close", systemImage: "xmark.circle")
@@ -133,6 +136,7 @@ struct SpriteNavigationPanel: View {
133136
titleVisibility: .visible
134137
) {
135138
Button("Delete", role: .destructive) {
139+
chatSessionManager.remove(chatId: chat.id, modelContext: modelContext)
136140
chatListViewModel.deleteChat(chat, apiClient: apiClient, modelContext: modelContext)
137141
chatToDelete = nil
138142
}
Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
import Testing
2+
import Foundation
3+
import SwiftData
4+
@testable import Wisp
5+
6+
@MainActor
7+
@Suite("ChatSessionManager")
8+
struct ChatSessionManagerTests {
9+
10+
private func makeModelContext() throws -> ModelContext {
11+
let config = ModelConfiguration(isStoredInMemoryOnly: true)
12+
let container = try ModelContainer(for: SpriteChat.self, SpriteSession.self, configurations: config)
13+
return ModelContext(container)
14+
}
15+
16+
private func makeChat(number: Int = 1, modelContext: ModelContext) -> SpriteChat {
17+
let chat = SpriteChat(spriteName: "test", chatNumber: number)
18+
modelContext.insert(chat)
19+
try? modelContext.save()
20+
return chat
21+
}
22+
23+
// MARK: - viewModel(for:)
24+
25+
@Test func viewModel_returnsSameInstanceOnSecondCall() throws {
26+
let ctx = try makeModelContext()
27+
let manager = ChatSessionManager()
28+
let chat = makeChat(modelContext: ctx)
29+
let apiClient = SpritesAPIClient()
30+
31+
let vm1 = manager.viewModel(for: chat, spriteName: "test", apiClient: apiClient, modelContext: ctx)
32+
let vm2 = manager.viewModel(for: chat, spriteName: "test", apiClient: apiClient, modelContext: ctx)
33+
34+
#expect(vm1 === vm2)
35+
}
36+
37+
@Test func viewModel_createsDistinctInstancesForDifferentChats() throws {
38+
let ctx = try makeModelContext()
39+
let manager = ChatSessionManager()
40+
let chat1 = makeChat(number: 1, modelContext: ctx)
41+
let chat2 = makeChat(number: 2, modelContext: ctx)
42+
let apiClient = SpritesAPIClient()
43+
44+
let vm1 = manager.viewModel(for: chat1, spriteName: "test", apiClient: apiClient, modelContext: ctx)
45+
let vm2 = manager.viewModel(for: chat2, spriteName: "test", apiClient: apiClient, modelContext: ctx)
46+
47+
#expect(vm1 !== vm2)
48+
}
49+
50+
// MARK: - isStreaming
51+
52+
@Test func isStreaming_returnsFalseForUncachedId() {
53+
let manager = ChatSessionManager()
54+
#expect(manager.isStreaming(chatId: UUID()) == false)
55+
}
56+
57+
@Test func isStreaming_returnsFalseForCachedIdleVM() throws {
58+
let ctx = try makeModelContext()
59+
let manager = ChatSessionManager()
60+
let chat = makeChat(modelContext: ctx)
61+
62+
let vm = manager.viewModel(for: chat, spriteName: "test", apiClient: SpritesAPIClient(), modelContext: ctx)
63+
vm.status = .idle
64+
65+
#expect(manager.isStreaming(chatId: chat.id) == false)
66+
}
67+
68+
// MARK: - remove(chatId:)
69+
70+
@Test func remove_evictsFromCacheSoNextCallCreatesNewVM() throws {
71+
let ctx = try makeModelContext()
72+
let manager = ChatSessionManager()
73+
let chat = makeChat(modelContext: ctx)
74+
let apiClient = SpritesAPIClient()
75+
76+
let vm1 = manager.viewModel(for: chat, spriteName: "test", apiClient: apiClient, modelContext: ctx)
77+
manager.remove(chatId: chat.id, modelContext: ctx)
78+
let vm2 = manager.viewModel(for: chat, spriteName: "test", apiClient: apiClient, modelContext: ctx)
79+
80+
#expect(vm1 !== vm2)
81+
}
82+
83+
@Test func remove_isNoopForUncachedId() throws {
84+
let ctx = try makeModelContext()
85+
let manager = ChatSessionManager()
86+
// Should not crash when removing an ID that was never cached
87+
manager.remove(chatId: UUID(), modelContext: ctx)
88+
}
89+
90+
// MARK: - detachAll()
91+
92+
@Test func detachAll_clearsAllCachedVMs() throws {
93+
let ctx = try makeModelContext()
94+
let manager = ChatSessionManager()
95+
let chat1 = makeChat(number: 1, modelContext: ctx)
96+
let chat2 = makeChat(number: 2, modelContext: ctx)
97+
let apiClient = SpritesAPIClient()
98+
99+
let vm1 = manager.viewModel(for: chat1, spriteName: "test", apiClient: apiClient, modelContext: ctx)
100+
let vm2 = manager.viewModel(for: chat2, spriteName: "test", apiClient: apiClient, modelContext: ctx)
101+
102+
manager.detachAll(modelContext: ctx)
103+
104+
let vm1After = manager.viewModel(for: chat1, spriteName: "test", apiClient: apiClient, modelContext: ctx)
105+
let vm2After = manager.viewModel(for: chat2, spriteName: "test", apiClient: apiClient, modelContext: ctx)
106+
107+
#expect(vm1 !== vm1After)
108+
#expect(vm2 !== vm2After)
109+
}
110+
}

WispTests/ChatViewModelTests.swift

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1117,6 +1117,35 @@ struct ChatViewModelTests {
11171117
}
11181118
}
11191119

1120+
// MARK: - reconnecting → streaming status transition
1121+
1122+
@Test func runReconnectLoop_transitionsReconnectingToStreamingOnStdoutData() async throws {
1123+
// Verifies the status trajectory: idle → reconnecting → streaming → idle
1124+
// When stdout data arrives in .reconnecting state, status must move to .streaming (line 935-936).
1125+
// When the result event arrives and the loop exits cleanly, status returns to .idle.
1126+
let ctx = try makeModelContext()
1127+
let (vm, _) = makeChatViewModel(modelContext: ctx)
1128+
1129+
let assistantMsg = ChatMessage(role: .assistant)
1130+
vm.messages.append(assistantMsg)
1131+
vm.setCurrentAssistantMessage(assistantMsg)
1132+
1133+
let textLine = #"{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Hello"}]}}"# + "\n"
1134+
let resultLine = #"{"type":"result","session_id":"s1","subtype":"success"}"# + "\n"
1135+
1136+
let stream = AsyncThrowingStream<ServiceLogEvent, Error> { continuation in
1137+
continuation.yield(ServiceLogEvent(type: .stdout, data: textLine, exitCode: nil, timestamp: nil, logFiles: nil))
1138+
continuation.yield(ServiceLogEvent(type: .stdout, data: resultLine, exitCode: nil, timestamp: nil, logFiles: nil))
1139+
continuation.finish()
1140+
}
1141+
let mock = MockServiceLogsProvider(streams: [stream], statuses: ["stopped"])
1142+
1143+
#expect(vm.status == .idle)
1144+
await vm.runReconnectLoop(apiClient: mock, modelContext: ctx, serviceAlreadyStopped: true)
1145+
// Loop completed cleanly: reconnecting → streaming → idle
1146+
#expect(vm.status == .idle)
1147+
}
1148+
11201149
// MARK: - addAttachedFile
11211150

11221151
@Test func addAttachedFile_appendsWithLastPathComponent() throws {

0 commit comments

Comments
 (0)