mirror of
https://github.com/NanmiCoder/cc-haha
synced 2026-07-16 13:03:31 +08:00
fix: isolate desktop session streaming state
Desktop chat sessions can stream and reconnect concurrently, so shared in-process buffers must not be keyed only by process lifetime. This change scopes streaming delta buffers, task-tool refresh bookkeeping, and CLI task mutations to the session that produced the event. It also ignores stale client socket closes after a newer socket has replaced the active connection for the same session. Constraint: Desktop users can keep multiple sessions and tabs active at the same time. Rejected: Serialize all desktop streaming through one active session | would hide the race instead of preserving multi-session behavior. Confidence: high Scope-risk: moderate Directive: Do not reintroduce process-global chat/task pending state without session keys and multi-session regression tests. Tested: cd desktop && bun run test -- src/stores/chatStore.test.ts src/stores/cliTaskStore.test.ts Tested: bun test src/server/__tests__/websocket-handler.test.ts Tested: cd desktop && bun run lint Tested: bun run check:desktop Tested: bun run check:server Tested: bun run check:coverage | changed-lines 100%, fails only existing agent-utils baseline Tested: bun run verify | 8 passed, 1 failed on existing agent-utils coverage baseline, 2 skipped Not-tested: Real desktop dual-window manual smoke. Related: https://github.com/NanmiCoder/cc-haha/issues/302 Related: https://github.com/NanmiCoder/cc-haha/issues/303
This commit is contained in:
程序员阿江(Relakkes)
parent
165d914d5a
commit
fd5dc097a7
@ -122,11 +122,39 @@ vi.mock('./cliTaskStore', () => ({
|
||||
},
|
||||
}))
|
||||
|
||||
import { mapHistoryMessagesToUiMessages, reconstructAgentNotifications, useChatStore } from './chatStore'
|
||||
import { sessionsApi } from '../api/sessions'
|
||||
import {
|
||||
mapHistoryMessagesToUiMessages,
|
||||
reconstructAgentNotifications,
|
||||
type PerSessionState,
|
||||
useChatStore,
|
||||
} from './chatStore'
|
||||
|
||||
const TEST_SESSION_ID = 'test-session-1'
|
||||
const initialState = useChatStore.getState()
|
||||
|
||||
function makeSession(overrides: Partial<PerSessionState> = {}): PerSessionState {
|
||||
return {
|
||||
messages: [],
|
||||
chatState: 'streaming',
|
||||
connectionState: 'connected',
|
||||
streamingText: '',
|
||||
streamingToolInput: '',
|
||||
activeToolUseId: null,
|
||||
activeToolName: null,
|
||||
activeThinkingId: null,
|
||||
pendingPermission: null,
|
||||
pendingComputerUsePermission: null,
|
||||
tokenUsage: { input_tokens: 0, output_tokens: 0 },
|
||||
elapsedSeconds: 0,
|
||||
statusVerb: '',
|
||||
slashCommands: [],
|
||||
agentTaskNotifications: {},
|
||||
elapsedTimer: null,
|
||||
...overrides,
|
||||
}
|
||||
}
|
||||
|
||||
describe('chatStore history mapping', () => {
|
||||
beforeEach(() => {
|
||||
sendMock.mockReset()
|
||||
@ -527,6 +555,124 @@ describe('chatStore history mapping', () => {
|
||||
])
|
||||
})
|
||||
|
||||
it('hydrates TodoWrite history into the currently tracked task store only', async () => {
|
||||
const todos = [{ content: 'Session task', status: 'in_progress' }]
|
||||
vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({
|
||||
messages: [
|
||||
{
|
||||
id: 'assistant-todo',
|
||||
type: 'assistant',
|
||||
timestamp: '2026-04-06T00:00:00.000Z',
|
||||
content: [
|
||||
{ type: 'tool_use', name: 'TodoWrite', id: 'todo-1', input: { todos } },
|
||||
],
|
||||
},
|
||||
],
|
||||
})
|
||||
cliTaskStoreSnapshot.sessionId = TEST_SESSION_ID
|
||||
cliTaskStoreSnapshot.tasks = []
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
[TEST_SESSION_ID]: makeSession({ messages: [] }),
|
||||
},
|
||||
})
|
||||
|
||||
await useChatStore.getState().loadHistory(TEST_SESSION_ID)
|
||||
|
||||
expect(setTasksFromTodosMock).toHaveBeenCalledWith(todos, TEST_SESSION_ID)
|
||||
})
|
||||
|
||||
it('marks history task completion dismissed when the user already continued', async () => {
|
||||
vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({
|
||||
messages: [
|
||||
{
|
||||
id: 'assistant-task',
|
||||
type: 'assistant',
|
||||
timestamp: '2026-04-06T00:00:00.000Z',
|
||||
content: [
|
||||
{ type: 'tool_use', name: 'TaskCreate', id: 'task-1', input: { subject: 'Done' } },
|
||||
],
|
||||
},
|
||||
{
|
||||
id: 'user-next',
|
||||
type: 'user',
|
||||
timestamp: '2026-04-06T00:00:01.000Z',
|
||||
content: '继续下一步',
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
[TEST_SESSION_ID]: makeSession({ messages: [] }),
|
||||
},
|
||||
})
|
||||
|
||||
await useChatStore.getState().loadHistory(TEST_SESSION_ID)
|
||||
|
||||
expect(setTasksFromTodosMock).toHaveBeenCalledWith([], TEST_SESSION_ID)
|
||||
expect(markCompletedAndDismissedMock).toHaveBeenCalledWith(TEST_SESSION_ID)
|
||||
})
|
||||
|
||||
it('reloads history task state for the requested session', async () => {
|
||||
const todos = [{ content: 'Reloaded task', status: 'pending' }]
|
||||
vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({
|
||||
messages: [
|
||||
{
|
||||
id: 'assistant-todo',
|
||||
type: 'assistant',
|
||||
timestamp: '2026-04-06T00:00:00.000Z',
|
||||
content: [
|
||||
{ type: 'tool_use', name: 'TodoWrite', id: 'todo-1', input: { todos } },
|
||||
],
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
[TEST_SESSION_ID]: makeSession({ messages: [{ id: 'old', type: 'assistant_text', content: 'old', timestamp: 1 }] }),
|
||||
},
|
||||
})
|
||||
|
||||
await useChatStore.getState().reloadHistory(TEST_SESSION_ID)
|
||||
|
||||
expect(setTasksFromTodosMock).toHaveBeenCalledWith(todos, TEST_SESSION_ID)
|
||||
})
|
||||
|
||||
it('clears reloaded task state after completed history is followed by a user turn', async () => {
|
||||
vi.mocked(sessionsApi.getMessages).mockResolvedValueOnce({
|
||||
messages: [
|
||||
{
|
||||
id: 'assistant-task',
|
||||
type: 'assistant',
|
||||
timestamp: '2026-04-06T00:00:00.000Z',
|
||||
content: [
|
||||
{ type: 'tool_use', name: 'TaskUpdate', id: 'task-1', input: { subject: 'Done' } },
|
||||
],
|
||||
},
|
||||
{
|
||||
id: 'user-next',
|
||||
type: 'user',
|
||||
timestamp: '2026-04-06T00:00:01.000Z',
|
||||
content: '新的问题',
|
||||
},
|
||||
],
|
||||
})
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
[TEST_SESSION_ID]: makeSession({ messages: [{ id: 'old', type: 'assistant_text', content: 'old', timestamp: 1 }] }),
|
||||
},
|
||||
})
|
||||
|
||||
await useChatStore.getState().reloadHistory(TEST_SESSION_ID)
|
||||
|
||||
expect(setTasksFromTodosMock).toHaveBeenCalledWith([], TEST_SESSION_ID)
|
||||
expect(markCompletedAndDismissedMock).toHaveBeenCalledWith(TEST_SESSION_ID)
|
||||
})
|
||||
|
||||
it('keeps parent tool linkage for live tool events', () => {
|
||||
// Initialize the session first
|
||||
useChatStore.setState({
|
||||
@ -582,6 +728,24 @@ describe('chatStore history mapping', () => {
|
||||
])
|
||||
})
|
||||
|
||||
it('syncs live TodoWrite tool input into the task store for that session', () => {
|
||||
const todos = [{ content: 'Live todo', status: 'in_progress' }]
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
[TEST_SESSION_ID]: makeSession({ chatState: 'tool_executing' }),
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
|
||||
type: 'tool_use_complete',
|
||||
toolName: 'TodoWrite',
|
||||
toolUseId: 'todo-live',
|
||||
input: { todos },
|
||||
})
|
||||
|
||||
expect(setTasksFromTodosMock).toHaveBeenCalledWith(todos, TEST_SESSION_ID)
|
||||
})
|
||||
|
||||
it('replays saved runtime selection when reconnecting a session', () => {
|
||||
useSessionRuntimeStore.getState().setSelection(TEST_SESSION_ID, {
|
||||
providerId: 'provider-1',
|
||||
@ -811,6 +975,8 @@ describe('chatStore history mapping', () => {
|
||||
})
|
||||
|
||||
it('clears local desktop chat state when the server confirms /clear', () => {
|
||||
vi.useFakeTimers()
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
[TEST_SESSION_ID]: {
|
||||
@ -837,6 +1003,10 @@ describe('chatStore history mapping', () => {
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
|
||||
type: 'content_delta',
|
||||
text: 'stale throttled delta',
|
||||
})
|
||||
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
|
||||
type: 'system_notification',
|
||||
subtype: 'session_cleared',
|
||||
@ -849,7 +1019,34 @@ describe('chatStore history mapping', () => {
|
||||
expect(session?.chatState).toBe('idle')
|
||||
expect(session?.tokenUsage).toEqual({ input_tokens: 0, output_tokens: 0 })
|
||||
expect(session?.slashCommands).toEqual([])
|
||||
expect(clearTasksMock).toHaveBeenCalled()
|
||||
expect(clearTasksMock).toHaveBeenCalledWith(TEST_SESSION_ID)
|
||||
|
||||
vi.advanceTimersByTime(60)
|
||||
expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.streamingText).toBe('')
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
it('clears local message state for only the requested session', () => {
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
'session-a': makeSession({
|
||||
messages: [{ id: 'a1', type: 'assistant_text', content: 'A old', timestamp: 1 }],
|
||||
streamingText: 'A pending',
|
||||
}),
|
||||
'session-b': makeSession({
|
||||
messages: [{ id: 'b1', type: 'assistant_text', content: 'B old', timestamp: 1 }],
|
||||
streamingText: 'B pending',
|
||||
}),
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().clearMessages('session-a')
|
||||
|
||||
expect(useChatStore.getState().sessions['session-a']?.messages).toEqual([])
|
||||
expect(useChatStore.getState().sessions['session-a']?.streamingText).toBe('')
|
||||
expect(useChatStore.getState().sessions['session-b']?.messages).toMatchObject([
|
||||
{ content: 'B old' },
|
||||
])
|
||||
})
|
||||
|
||||
it('renders compact boundary notifications as system messages', () => {
|
||||
@ -958,7 +1155,7 @@ describe('chatStore history mapping', () => {
|
||||
|
||||
useChatStore.getState().sendMessage(TEST_SESSION_ID, '继续下一轮')
|
||||
|
||||
expect(resetCompletedTasksMock).toHaveBeenCalledTimes(1)
|
||||
expect(resetCompletedTasksMock).toHaveBeenCalledWith(TEST_SESSION_ID)
|
||||
expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toMatchObject([
|
||||
{
|
||||
type: 'task_summary',
|
||||
@ -974,6 +1171,71 @@ describe('chatStore history mapping', () => {
|
||||
])
|
||||
})
|
||||
|
||||
it('does not attach completed tasks from another tracked session to a new user turn', () => {
|
||||
cliTaskStoreSnapshot.sessionId = 'session-b'
|
||||
cliTaskStoreSnapshot.tasks = [
|
||||
{ id: '1', subject: 'Session B completed task', status: 'completed' },
|
||||
]
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
'session-a': makeSession({ chatState: 'idle' }),
|
||||
'session-b': makeSession({ chatState: 'idle' }),
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().sendMessage('session-a', '继续 A 会话')
|
||||
|
||||
expect(resetCompletedTasksMock).not.toHaveBeenCalled()
|
||||
expect(useChatStore.getState().sessions['session-a']?.messages).toMatchObject([
|
||||
{
|
||||
type: 'user_text',
|
||||
content: '继续 A 会话',
|
||||
},
|
||||
])
|
||||
})
|
||||
|
||||
it('tracks task tool results independently per session even when tool IDs collide', () => {
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
'session-a': makeSession({
|
||||
activeToolUseId: 'tool-same',
|
||||
activeToolName: 'TaskCreate',
|
||||
}),
|
||||
'session-b': makeSession({
|
||||
activeToolUseId: 'tool-same',
|
||||
activeToolName: 'TaskCreate',
|
||||
}),
|
||||
},
|
||||
})
|
||||
|
||||
for (const sessionId of ['session-a', 'session-b']) {
|
||||
useChatStore.getState().handleServerMessage(sessionId, {
|
||||
type: 'tool_use_complete',
|
||||
toolName: 'TaskCreate',
|
||||
toolUseId: 'tool-same',
|
||||
input: { subject: sessionId },
|
||||
})
|
||||
}
|
||||
|
||||
useChatStore.getState().handleServerMessage('session-a', {
|
||||
type: 'tool_result',
|
||||
toolUseId: 'tool-same',
|
||||
content: 'created A',
|
||||
isError: false,
|
||||
})
|
||||
useChatStore.getState().handleServerMessage('session-b', {
|
||||
type: 'tool_result',
|
||||
toolUseId: 'tool-same',
|
||||
content: 'created B',
|
||||
isError: false,
|
||||
})
|
||||
|
||||
expect(refreshTasksMock).toHaveBeenCalledTimes(2)
|
||||
expect(refreshTasksMock).toHaveBeenNthCalledWith(1, 'session-a')
|
||||
expect(refreshTasksMock).toHaveBeenNthCalledWith(2, 'session-b')
|
||||
})
|
||||
|
||||
it('tracks Computer Use approval requests separately from generic tool permissions', () => {
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
@ -1104,6 +1366,218 @@ describe('chatStore history mapping', () => {
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
it('keeps throttled streaming deltas isolated per session', () => {
|
||||
vi.useFakeTimers()
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
'session-a': makeSession(),
|
||||
'session-b': makeSession(),
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().handleServerMessage('session-a', {
|
||||
type: 'content_start',
|
||||
blockType: 'text',
|
||||
})
|
||||
useChatStore.getState().handleServerMessage('session-a', {
|
||||
type: 'content_delta',
|
||||
text: 'A-only response',
|
||||
})
|
||||
useChatStore.getState().handleServerMessage('session-b', {
|
||||
type: 'content_start',
|
||||
blockType: 'text',
|
||||
})
|
||||
useChatStore.getState().handleServerMessage('session-b', {
|
||||
type: 'content_delta',
|
||||
text: 'B-only response',
|
||||
})
|
||||
|
||||
useChatStore.getState().handleServerMessage('session-a', {
|
||||
type: 'message_complete',
|
||||
usage: { input_tokens: 1, output_tokens: 1 },
|
||||
})
|
||||
useChatStore.getState().handleServerMessage('session-b', {
|
||||
type: 'message_complete',
|
||||
usage: { input_tokens: 1, output_tokens: 1 },
|
||||
})
|
||||
|
||||
expect(useChatStore.getState().sessions['session-a']?.messages).toMatchObject([
|
||||
{ type: 'assistant_text', content: 'A-only response' },
|
||||
])
|
||||
expect(useChatStore.getState().sessions['session-b']?.messages).toMatchObject([
|
||||
{ type: 'assistant_text', content: 'B-only response' },
|
||||
])
|
||||
|
||||
vi.runOnlyPendingTimers()
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
it('flushes pending text before appending a thinking block', () => {
|
||||
vi.useFakeTimers()
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
[TEST_SESSION_ID]: makeSession({ chatState: 'streaming' }),
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
|
||||
type: 'content_delta',
|
||||
text: 'visible answer before thinking',
|
||||
})
|
||||
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
|
||||
type: 'thinking',
|
||||
text: 'internal note',
|
||||
})
|
||||
|
||||
expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toMatchObject([
|
||||
{ type: 'assistant_text', content: 'visible answer before thinking' },
|
||||
{ type: 'thinking', content: 'internal note' },
|
||||
])
|
||||
|
||||
vi.runOnlyPendingTimers()
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
it('flushes pending text before appending an error message', () => {
|
||||
vi.useFakeTimers()
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
[TEST_SESSION_ID]: makeSession({ chatState: 'streaming' }),
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
|
||||
type: 'content_delta',
|
||||
text: 'partial answer before error',
|
||||
})
|
||||
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
|
||||
type: 'error',
|
||||
message: 'provider failed',
|
||||
code: 'provider_error',
|
||||
})
|
||||
|
||||
expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.messages).toMatchObject([
|
||||
{ type: 'assistant_text', content: 'partial answer before error' },
|
||||
{ type: 'error', message: 'provider failed' },
|
||||
])
|
||||
|
||||
vi.runOnlyPendingTimers()
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
it('flushes throttled deltas only for the stopped session', () => {
|
||||
vi.useFakeTimers()
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
'session-a': makeSession(),
|
||||
'session-b': makeSession(),
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().handleServerMessage('session-a', {
|
||||
type: 'content_delta',
|
||||
text: 'A-only response',
|
||||
})
|
||||
useChatStore.getState().handleServerMessage('session-b', {
|
||||
type: 'content_delta',
|
||||
text: 'B-only response',
|
||||
})
|
||||
|
||||
useChatStore.getState().stopGeneration('session-a')
|
||||
|
||||
expect(useChatStore.getState().sessions['session-a']?.streamingText).toBe('A-only response')
|
||||
expect(useChatStore.getState().sessions['session-b']?.streamingText).toBe('')
|
||||
|
||||
useChatStore.getState().handleServerMessage('session-b', {
|
||||
type: 'message_complete',
|
||||
usage: { input_tokens: 1, output_tokens: 1 },
|
||||
})
|
||||
|
||||
expect(useChatStore.getState().sessions['session-b']?.messages).toMatchObject([
|
||||
{ type: 'assistant_text', content: 'B-only response' },
|
||||
])
|
||||
|
||||
vi.runOnlyPendingTimers()
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
it('does not flush one session throttled delta into another disconnected session', () => {
|
||||
vi.useFakeTimers()
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
'session-a': makeSession(),
|
||||
'session-b': makeSession(),
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().handleServerMessage('session-a', {
|
||||
type: 'content_delta',
|
||||
text: 'A-only response',
|
||||
})
|
||||
useChatStore.getState().handleServerMessage('session-b', {
|
||||
type: 'content_delta',
|
||||
text: 'B-only response',
|
||||
})
|
||||
|
||||
useChatStore.getState().disconnectSession('session-a')
|
||||
|
||||
expect(useChatStore.getState().sessions['session-a']).toBeUndefined()
|
||||
expect(useChatStore.getState().sessions['session-b']?.streamingText).toBe('')
|
||||
|
||||
useChatStore.getState().handleServerMessage('session-b', {
|
||||
type: 'message_complete',
|
||||
usage: { input_tokens: 1, output_tokens: 1 },
|
||||
})
|
||||
|
||||
expect(useChatStore.getState().sessions['session-b']?.messages).toMatchObject([
|
||||
{ type: 'assistant_text', content: 'B-only response' },
|
||||
])
|
||||
|
||||
vi.runOnlyPendingTimers()
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
it('ignores late throttled deltas after a session has disconnected', () => {
|
||||
vi.useFakeTimers()
|
||||
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
'session-a': makeSession(),
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().handleServerMessage('session-a', {
|
||||
type: 'content_delta',
|
||||
text: 'before disconnect',
|
||||
})
|
||||
useChatStore.getState().disconnectSession('session-a')
|
||||
|
||||
useChatStore.getState().handleServerMessage('session-a', {
|
||||
type: 'content_delta',
|
||||
text: 'late stale delta',
|
||||
})
|
||||
useChatStore.setState({
|
||||
sessions: {
|
||||
'session-a': makeSession({ chatState: 'idle' }),
|
||||
},
|
||||
})
|
||||
|
||||
useChatStore.getState().sendMessage('session-a', 'fresh turn')
|
||||
|
||||
expect(useChatStore.getState().sessions['session-a']?.messages).toMatchObject([
|
||||
{ type: 'user_text', content: 'fresh turn' },
|
||||
])
|
||||
|
||||
vi.runOnlyPendingTimers()
|
||||
expect(useChatStore.getState().sessions['session-a']?.streamingText).toBe('')
|
||||
vi.useRealTimers()
|
||||
})
|
||||
|
||||
it('does not split one streamed markdown reply when task progress arrives mid-stream', () => {
|
||||
vi.useFakeTimers()
|
||||
|
||||
|
||||
@ -124,26 +124,62 @@ type ChatStore = {
|
||||
}
|
||||
|
||||
const TASK_TOOL_NAMES = new Set(['TaskCreate', 'TaskUpdate', 'TaskGet', 'TaskList', 'TodoWrite'])
|
||||
const pendingTaskToolUseIds = new Set<string>()
|
||||
const pendingTaskToolUseIdsBySession = new Map<string, Set<string>>()
|
||||
|
||||
function addPendingTaskToolUseId(sessionId: string, toolUseId: string): void {
|
||||
const ids = pendingTaskToolUseIdsBySession.get(sessionId) ?? new Set<string>()
|
||||
ids.add(toolUseId)
|
||||
pendingTaskToolUseIdsBySession.set(sessionId, ids)
|
||||
}
|
||||
|
||||
function consumePendingTaskToolUseId(sessionId: string, toolUseId: string): boolean {
|
||||
const ids = pendingTaskToolUseIdsBySession.get(sessionId)
|
||||
if (!ids?.has(toolUseId)) return false
|
||||
ids.delete(toolUseId)
|
||||
if (ids.size === 0) pendingTaskToolUseIdsBySession.delete(sessionId)
|
||||
return true
|
||||
}
|
||||
|
||||
function clearPendingTaskToolUseIds(sessionId: string): void {
|
||||
pendingTaskToolUseIdsBySession.delete(sessionId)
|
||||
}
|
||||
const AGENT_COMPLETION_NOTIFICATION_PREVIEW_CHARS = 160
|
||||
|
||||
let msgCounter = 0
|
||||
const nextId = () => `msg-${++msgCounter}-${Date.now()}`
|
||||
|
||||
// Streaming throttle for content_delta
|
||||
let pendingDelta = ''
|
||||
let flushTimer: ReturnType<typeof setTimeout> | null = null
|
||||
// Streaming throttle for content_delta. Buffers must be per-session because
|
||||
// multiple desktop tabs can stream at the same time.
|
||||
const pendingDeltaBySession = new Map<string, string>()
|
||||
const flushTimerBySession = new Map<string, ReturnType<typeof setTimeout>>()
|
||||
|
||||
function consumePendingDelta(): string {
|
||||
function consumePendingDelta(sessionId: string): string {
|
||||
const flushTimer = flushTimerBySession.get(sessionId)
|
||||
if (flushTimer) {
|
||||
clearTimeout(flushTimer)
|
||||
flushTimer = null
|
||||
flushTimerBySession.delete(sessionId)
|
||||
}
|
||||
const text = pendingDelta
|
||||
pendingDelta = ''
|
||||
const text = pendingDeltaBySession.get(sessionId) ?? ''
|
||||
pendingDeltaBySession.delete(sessionId)
|
||||
return text
|
||||
}
|
||||
|
||||
function appendPendingDelta(sessionId: string, text: string): void {
|
||||
pendingDeltaBySession.set(
|
||||
sessionId,
|
||||
`${pendingDeltaBySession.get(sessionId) ?? ''}${text}`,
|
||||
)
|
||||
}
|
||||
|
||||
function clearPendingDelta(sessionId: string): void {
|
||||
const flushTimer = flushTimerBySession.get(sessionId)
|
||||
if (flushTimer) {
|
||||
clearTimeout(flushTimer)
|
||||
flushTimerBySession.delete(sessionId)
|
||||
}
|
||||
pendingDeltaBySession.delete(sessionId)
|
||||
}
|
||||
|
||||
function appendAssistantTextMessage(
|
||||
messages: UIMessage[],
|
||||
content: string,
|
||||
@ -283,11 +319,11 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
disconnectSession: (sessionId) => {
|
||||
const session = get().sessions[sessionId]
|
||||
if (session?.elapsedTimer) clearInterval(session.elapsedTimer)
|
||||
if (flushTimer) { clearTimeout(flushTimer); flushTimer = null }
|
||||
if (pendingDelta) {
|
||||
const text = consumePendingDelta()
|
||||
if (pendingDeltaBySession.has(sessionId)) {
|
||||
const text = consumePendingDelta(sessionId)
|
||||
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) }))
|
||||
}
|
||||
clearPendingTaskToolUseIds(sessionId)
|
||||
wsManager.disconnect(sessionId)
|
||||
set((s) => {
|
||||
const { [sessionId]: _, ...rest } = s.sessions
|
||||
@ -317,13 +353,14 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
: undefined
|
||||
|
||||
const taskStore = useCLITaskStore.getState()
|
||||
const allTasksDone = taskStore.tasks.length > 0 && taskStore.tasks.every((t) => t.status === 'completed')
|
||||
const sessionTasks = taskStore.sessionId === sessionId ? taskStore.tasks : []
|
||||
const allTasksDone = sessionTasks.length > 0 && sessionTasks.every((t) => t.status === 'completed')
|
||||
const completedTaskSummary = allTasksDone
|
||||
? taskStore.tasks.map((t) => ({ id: t.id, subject: t.subject, status: t.status, activeForm: t.activeForm }))
|
||||
? sessionTasks.map((t) => ({ id: t.id, subject: t.subject, status: t.status, activeForm: t.activeForm }))
|
||||
: []
|
||||
|
||||
if (!isMemberSession && allTasksDone) {
|
||||
void taskStore.resetCompletedTasks()
|
||||
void taskStore.resetCompletedTasks(sessionId)
|
||||
}
|
||||
|
||||
if (!isMemberSession) {
|
||||
@ -332,11 +369,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
|
||||
set((s) => {
|
||||
const session = s.sessions[sessionId] ?? createDefaultSessionState()
|
||||
if (flushTimer) {
|
||||
clearTimeout(flushTimer)
|
||||
flushTimer = null
|
||||
}
|
||||
const bufferedDelta = consumePendingDelta()
|
||||
const bufferedDelta = consumePendingDelta(sessionId)
|
||||
const pendingAssistantText = `${session.streamingText}${bufferedDelta}`
|
||||
|
||||
const newMessages = pendingAssistantText.trim()
|
||||
@ -449,9 +482,8 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
|
||||
stopGeneration: (sessionId) => {
|
||||
wsManager.send(sessionId, { type: 'stop_generation' })
|
||||
if (flushTimer) { clearTimeout(flushTimer); flushTimer = null }
|
||||
if (pendingDelta) {
|
||||
const text = consumePendingDelta()
|
||||
if (pendingDeltaBySession.has(sessionId)) {
|
||||
const text = consumePendingDelta(sessionId)
|
||||
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, (sess) => ({ streamingText: sess.streamingText + text })) }))
|
||||
}
|
||||
set((s) => {
|
||||
@ -491,12 +523,12 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
})
|
||||
if (lastTodos && lastTodos.length > 0) {
|
||||
const taskStore = useCLITaskStore.getState()
|
||||
if (taskStore.tasks.length === 0) taskStore.setTasksFromTodos(lastTodos)
|
||||
if (taskStore.sessionId === sessionId && taskStore.tasks.length === 0) taskStore.setTasksFromTodos(lastTodos, sessionId)
|
||||
} else {
|
||||
useCLITaskStore.getState().setTasksFromTodos([])
|
||||
useCLITaskStore.getState().setTasksFromTodos([], sessionId)
|
||||
}
|
||||
if (hasMessagesAfterTaskCompletion) {
|
||||
useCLITaskStore.getState().markCompletedAndDismissed()
|
||||
useCLITaskStore.getState().markCompletedAndDismissed(sessionId)
|
||||
}
|
||||
} catch {
|
||||
// Session may not have messages yet
|
||||
@ -535,12 +567,12 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
})
|
||||
|
||||
if (lastTodos && lastTodos.length > 0) {
|
||||
useCLITaskStore.getState().setTasksFromTodos(lastTodos)
|
||||
useCLITaskStore.getState().setTasksFromTodos(lastTodos, sessionId)
|
||||
} else {
|
||||
useCLITaskStore.getState().setTasksFromTodos([])
|
||||
useCLITaskStore.getState().setTasksFromTodos([], sessionId)
|
||||
}
|
||||
if (hasMessagesAfterTaskCompletion) {
|
||||
useCLITaskStore.getState().markCompletedAndDismissed()
|
||||
useCLITaskStore.getState().markCompletedAndDismissed(sessionId)
|
||||
}
|
||||
} catch {
|
||||
// Session may not have messages yet
|
||||
@ -560,6 +592,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
},
|
||||
|
||||
clearMessages: (sessionId) => {
|
||||
clearPendingTaskToolUseIds(sessionId)
|
||||
set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ messages: [], streamingText: '', chatState: 'idle' })) }))
|
||||
},
|
||||
|
||||
@ -574,7 +607,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
|
||||
case 'status':
|
||||
update((session) => {
|
||||
const pendingText = `${session.streamingText}${consumePendingDelta()}`
|
||||
const pendingText = `${session.streamingText}${consumePendingDelta(sessionId)}`
|
||||
const hasPendingStreamText =
|
||||
session.chatState === 'streaming' && pendingText.trim().length > 0
|
||||
// Background task progress can arrive while the assistant is still
|
||||
@ -608,7 +641,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
case 'content_start': {
|
||||
const session = get().sessions[sessionId]
|
||||
if (!session) break
|
||||
const pendingText = `${session.streamingText}${consumePendingDelta()}`
|
||||
const pendingText = `${session.streamingText}${consumePendingDelta(sessionId)}`
|
||||
if (msg.blockType !== 'text' && pendingText.trim()) {
|
||||
update((s) => ({
|
||||
messages: appendAssistantTextMessage(s.messages, pendingText, Date.now()),
|
||||
@ -635,14 +668,16 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
|
||||
case 'content_delta':
|
||||
if (msg.text !== undefined) {
|
||||
pendingDelta += msg.text
|
||||
if (!flushTimer) {
|
||||
flushTimer = setTimeout(() => {
|
||||
const text = pendingDelta
|
||||
pendingDelta = ''
|
||||
flushTimer = null
|
||||
if (!get().sessions[sessionId]) break
|
||||
appendPendingDelta(sessionId, msg.text)
|
||||
if (!flushTimerBySession.has(sessionId)) {
|
||||
const timer = setTimeout(() => {
|
||||
const text = pendingDeltaBySession.get(sessionId) ?? ''
|
||||
pendingDeltaBySession.delete(sessionId)
|
||||
flushTimerBySession.delete(sessionId)
|
||||
update((s) => ({ streamingText: s.streamingText + text }))
|
||||
}, 50)
|
||||
flushTimerBySession.set(sessionId, timer)
|
||||
}
|
||||
}
|
||||
if (msg.toolInput !== undefined) update((s) => ({ streamingToolInput: s.streamingToolInput + msg.toolInput }))
|
||||
@ -650,7 +685,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
|
||||
case 'thinking':
|
||||
update((s) => {
|
||||
const pendingText = `${s.streamingText}${consumePendingDelta()}`
|
||||
const pendingText = `${s.streamingText}${consumePendingDelta(sessionId)}`
|
||||
const base = pendingText.trim()
|
||||
? appendAssistantTextMessage(s.messages, pendingText, Date.now())
|
||||
: s.messages
|
||||
@ -682,10 +717,10 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
activeToolUseId: null, activeToolName: null, activeThinkingId: null, streamingToolInput: '',
|
||||
}))
|
||||
if (toolName === 'TodoWrite' && Array.isArray((msg.input as any)?.todos)) {
|
||||
useCLITaskStore.getState().setTasksFromTodos((msg.input as any).todos)
|
||||
useCLITaskStore.getState().setTasksFromTodos((msg.input as any).todos, sessionId)
|
||||
} else if (TASK_TOOL_NAMES.has(toolName)) {
|
||||
const useId = msg.toolUseId || session?.activeToolUseId
|
||||
if (useId) pendingTaskToolUseIds.add(useId)
|
||||
if (useId) addPendingTaskToolUseId(sessionId, useId)
|
||||
}
|
||||
break
|
||||
}
|
||||
@ -698,9 +733,8 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
}],
|
||||
chatState: 'thinking', activeThinkingId: null,
|
||||
}))
|
||||
if (pendingTaskToolUseIds.has(msg.toolUseId)) {
|
||||
pendingTaskToolUseIds.delete(msg.toolUseId)
|
||||
useCLITaskStore.getState().refreshTasks()
|
||||
if (consumePendingTaskToolUseId(sessionId, msg.toolUseId)) {
|
||||
useCLITaskStore.getState().refreshTasks(sessionId)
|
||||
}
|
||||
break
|
||||
|
||||
@ -766,7 +800,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
const session = get().sessions[sessionId]
|
||||
if (!session) break
|
||||
const wasAgentRunning = session.chatState !== 'idle'
|
||||
const text = `${session.streamingText}${consumePendingDelta()}`
|
||||
const text = `${session.streamingText}${consumePendingDelta(sessionId)}`
|
||||
let completionMessages = session.messages
|
||||
if (text.trim()) {
|
||||
completionMessages = appendAssistantTextMessage(session.messages, text, Date.now())
|
||||
@ -803,7 +837,7 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
|
||||
case 'error':
|
||||
update((s) => {
|
||||
const pendingText = `${s.streamingText}${consumePendingDelta()}`
|
||||
const pendingText = `${s.streamingText}${consumePendingDelta(sessionId)}`
|
||||
let newMessages = s.messages
|
||||
if (pendingText.trim()) {
|
||||
newMessages = appendAssistantTextMessage(newMessages, pendingText, Date.now())
|
||||
@ -866,7 +900,9 @@ export const useChatStore = create<ChatStore>((set, get) => ({
|
||||
tokenUsage: { input_tokens: 0, output_tokens: 0 },
|
||||
slashCommands: [],
|
||||
}))
|
||||
useCLITaskStore.getState().clearTasks()
|
||||
clearPendingDelta(sessionId)
|
||||
clearPendingTaskToolUseIds(sessionId)
|
||||
useCLITaskStore.getState().clearTasks(sessionId)
|
||||
useSessionStore.getState().updateSessionTitle(sessionId, 'New Session')
|
||||
useTabStore.getState().updateTabTitle(sessionId, 'New Session')
|
||||
useTabStore.getState().updateTabStatus(sessionId, 'idle')
|
||||
|
||||
@ -106,4 +106,78 @@ describe('cliTaskStore', () => {
|
||||
|
||||
expect(useCLITaskStore.getState().resetting).toBe(false)
|
||||
})
|
||||
|
||||
it('refreshes tasks for the currently tracked session by default', async () => {
|
||||
vi.mocked(cliTasksApi.getTasksForList).mockResolvedValue({
|
||||
tasks: [makeTask('session-1', 'in_progress')],
|
||||
})
|
||||
|
||||
useCLITaskStore.setState({
|
||||
sessionId: 'session-1',
|
||||
tasks: [],
|
||||
expanded: false,
|
||||
completedAndDismissed: false,
|
||||
dismissedCompletionKey: null,
|
||||
})
|
||||
|
||||
await useCLITaskStore.getState().refreshTasks()
|
||||
|
||||
expect(cliTasksApi.getTasksForList).toHaveBeenCalledWith('session-1')
|
||||
expect(useCLITaskStore.getState().tasks).toMatchObject([
|
||||
{ taskListId: 'session-1', status: 'in_progress' },
|
||||
])
|
||||
})
|
||||
|
||||
it('marks completed tasks dismissed for the currently tracked session by default', () => {
|
||||
useCLITaskStore.setState({
|
||||
sessionId: 'session-1',
|
||||
tasks: [makeTask('session-1', 'completed')],
|
||||
expanded: true,
|
||||
completedAndDismissed: false,
|
||||
dismissedCompletionKey: null,
|
||||
})
|
||||
|
||||
useCLITaskStore.getState().markCompletedAndDismissed()
|
||||
|
||||
expect(useCLITaskStore.getState()).toMatchObject({
|
||||
completedAndDismissed: true,
|
||||
dismissedCompletionKey: 'session-1::1::Keep current session isolated::completed::::',
|
||||
expanded: false,
|
||||
})
|
||||
})
|
||||
|
||||
it('ignores TodoWrite updates for a session that is not currently tracked', () => {
|
||||
useCLITaskStore.setState({
|
||||
sessionId: 'session-1',
|
||||
tasks: [makeTask('session-1', 'in_progress')],
|
||||
expanded: true,
|
||||
completedAndDismissed: false,
|
||||
dismissedCompletionKey: null,
|
||||
})
|
||||
|
||||
useCLITaskStore.getState().setTasksFromTodos([
|
||||
{ content: 'Session 2 task', status: 'completed' },
|
||||
], 'session-2')
|
||||
|
||||
expect(useCLITaskStore.getState().tasks).toMatchObject([
|
||||
{ taskListId: 'session-1', subject: 'Keep current session isolated' },
|
||||
])
|
||||
})
|
||||
|
||||
it('does not reset completed tasks for a different session', async () => {
|
||||
useCLITaskStore.setState({
|
||||
sessionId: 'session-1',
|
||||
tasks: [makeTask('session-1', 'completed')],
|
||||
expanded: true,
|
||||
completedAndDismissed: false,
|
||||
dismissedCompletionKey: null,
|
||||
})
|
||||
|
||||
await useCLITaskStore.getState().resetCompletedTasks('session-2')
|
||||
|
||||
expect(vi.mocked(cliTasksApi.resetTaskList)).not.toHaveBeenCalled()
|
||||
expect(useCLITaskStore.getState().tasks).toMatchObject([
|
||||
{ taskListId: 'session-1', status: 'completed' },
|
||||
])
|
||||
})
|
||||
})
|
||||
|
||||
@ -25,16 +25,16 @@ type CLITaskStore = {
|
||||
|
||||
/** Fetch tasks for a given session (uses sessionId as taskListId) */
|
||||
fetchSessionTasks: (sessionId: string) => Promise<void>
|
||||
/** Refresh tasks for the currently tracked session */
|
||||
refreshTasks: () => Promise<void>
|
||||
/** Refresh tasks for the currently tracked session, or a specific session if provided */
|
||||
refreshTasks: (sessionId?: string) => Promise<void>
|
||||
/** Update tasks from TodoWrite V1 tool input (in-memory, no disk read needed) */
|
||||
setTasksFromTodos: (todos: TodoItem[]) => void
|
||||
setTasksFromTodos: (todos: TodoItem[], sessionId?: string) => void
|
||||
/** Mark that completed tasks were already dismissed (conversation continued) */
|
||||
markCompletedAndDismissed: () => void
|
||||
markCompletedAndDismissed: (sessionId?: string) => void
|
||||
/** Clear a completed task list locally and remotely so the next cycle starts clean */
|
||||
resetCompletedTasks: () => Promise<void>
|
||||
resetCompletedTasks: (sessionId?: string) => Promise<void>
|
||||
/** Clear task tracking state */
|
||||
clearTasks: () => void
|
||||
clearTasks: (sessionId?: string) => void
|
||||
/** Toggle expanded state */
|
||||
toggleExpanded: () => void
|
||||
}
|
||||
@ -116,8 +116,8 @@ export const useCLITaskStore = create<CLITaskStore>((set, get) => ({
|
||||
}
|
||||
},
|
||||
|
||||
refreshTasks: async () => {
|
||||
const { sessionId } = get()
|
||||
refreshTasks: async (targetSessionId) => {
|
||||
const sessionId = targetSessionId ?? get().sessionId
|
||||
if (!sessionId) return
|
||||
try {
|
||||
const { tasks } = await cliTasksApi.getTasksForList(sessionId)
|
||||
@ -132,15 +132,19 @@ export const useCLITaskStore = create<CLITaskStore>((set, get) => ({
|
||||
}
|
||||
},
|
||||
|
||||
setTasksFromTodos: (todos) => {
|
||||
const tasks = mapTodosToTasks(todos, get().sessionId)
|
||||
setTasksFromTodos: (todos, targetSessionId) => {
|
||||
const sessionId = targetSessionId ?? get().sessionId
|
||||
if (!sessionId || get().sessionId !== sessionId) return
|
||||
const tasks = mapTodosToTasks(todos, sessionId)
|
||||
set((state) => ({
|
||||
tasks,
|
||||
...resolveDismissState(tasks, state.dismissedCompletionKey),
|
||||
}))
|
||||
},
|
||||
|
||||
markCompletedAndDismissed: () => {
|
||||
markCompletedAndDismissed: (targetSessionId) => {
|
||||
const sessionId = targetSessionId ?? get().sessionId
|
||||
if (!sessionId || get().sessionId !== sessionId) return
|
||||
const completionKey = buildCompletedTaskKey(get().tasks)
|
||||
if (!completionKey) return
|
||||
|
||||
@ -151,10 +155,12 @@ export const useCLITaskStore = create<CLITaskStore>((set, get) => ({
|
||||
})
|
||||
},
|
||||
|
||||
resetCompletedTasks: async () => {
|
||||
const { sessionId, tasks } = get()
|
||||
resetCompletedTasks: async (targetSessionId) => {
|
||||
const sessionId = targetSessionId ?? get().sessionId
|
||||
if (!sessionId || get().sessionId !== sessionId) return
|
||||
const { tasks } = get()
|
||||
const completionKey = buildCompletedTaskKey(tasks)
|
||||
if (!sessionId || !completionKey) return
|
||||
if (!completionKey) return
|
||||
|
||||
set({
|
||||
tasks: [],
|
||||
@ -173,7 +179,8 @@ export const useCLITaskStore = create<CLITaskStore>((set, get) => ({
|
||||
}
|
||||
},
|
||||
|
||||
clearTasks: () => {
|
||||
clearTasks: (targetSessionId) => {
|
||||
if (targetSessionId && get().sessionId !== targetSessionId) return
|
||||
set({
|
||||
sessionId: null,
|
||||
tasks: [],
|
||||
|
||||
55
src/server/__tests__/websocket-handler.test.ts
Normal file
55
src/server/__tests__/websocket-handler.test.ts
Normal file
@ -0,0 +1,55 @@
|
||||
import { afterEach, describe, expect, it, mock, spyOn } from 'bun:test'
|
||||
import type { ServerWebSocket } from 'bun'
|
||||
import {
|
||||
__resetWebSocketHandlerStateForTests,
|
||||
getActiveSessionIds,
|
||||
handleWebSocket,
|
||||
type WebSocketData,
|
||||
} from '../ws/handler.js'
|
||||
import { conversationService } from '../services/conversationService.js'
|
||||
import { computerUseApprovalService } from '../services/computerUseApprovalService.js'
|
||||
|
||||
function makeClientSocket(sessionId: string) {
|
||||
const sent: string[] = []
|
||||
return {
|
||||
data: {
|
||||
sessionId,
|
||||
connectedAt: Date.now(),
|
||||
channel: 'client',
|
||||
sdkToken: null,
|
||||
serverPort: 0,
|
||||
serverHost: '127.0.0.1',
|
||||
},
|
||||
send: mock((payload: string) => {
|
||||
sent.push(payload)
|
||||
}),
|
||||
close: mock(() => {}),
|
||||
sent,
|
||||
} as unknown as ServerWebSocket<WebSocketData> & { sent: string[] }
|
||||
}
|
||||
|
||||
describe('WebSocket handler session isolation', () => {
|
||||
afterEach(() => {
|
||||
__resetWebSocketHandlerStateForTests()
|
||||
mock.restore()
|
||||
})
|
||||
|
||||
it('ignores stale disconnects from an older socket for the same session', () => {
|
||||
const sessionId = `duplicate-${crypto.randomUUID()}`
|
||||
const first = makeClientSocket(sessionId)
|
||||
const second = makeClientSocket(sessionId)
|
||||
const clearCallbacks = spyOn(conversationService, 'clearOutputCallbacks')
|
||||
const cancelComputerUse = spyOn(computerUseApprovalService, 'cancelSession')
|
||||
|
||||
handleWebSocket.open(first)
|
||||
handleWebSocket.open(second)
|
||||
clearCallbacks.mockClear()
|
||||
cancelComputerUse.mockClear()
|
||||
|
||||
handleWebSocket.close(first, 1000, 'stale tab closed')
|
||||
|
||||
expect(getActiveSessionIds()).toContain(sessionId)
|
||||
expect(clearCallbacks).not.toHaveBeenCalled()
|
||||
expect(cancelComputerUse).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
@ -193,6 +193,10 @@ export const handleWebSocket = {
|
||||
}
|
||||
|
||||
console.log(`[WS] Client disconnected from session: ${sessionId} (${code}: ${reason})`)
|
||||
if (activeSessions.get(sessionId) !== ws) {
|
||||
console.log(`[WS] Ignoring stale client disconnect for session: ${sessionId}`)
|
||||
return
|
||||
}
|
||||
computerUseApprovalService.cancelSession(sessionId)
|
||||
activeSessions.delete(sessionId)
|
||||
conversationService.clearOutputCallbacks(sessionId)
|
||||
@ -1498,3 +1502,11 @@ export function sendToSession(sessionId: string, message: ServerMessage): boolea
|
||||
export function getActiveSessionIds(): string[] {
|
||||
return Array.from(activeSessions.keys())
|
||||
}
|
||||
|
||||
export function __resetWebSocketHandlerStateForTests(): void {
|
||||
for (const timer of sessionCleanupTimers.values()) clearTimeout(timer)
|
||||
for (const timer of prewarmIdleTimers.values()) clearTimeout(timer)
|
||||
activeSessions.clear()
|
||||
sessionCleanupTimers.clear()
|
||||
prewarmIdleTimers.clear()
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user