Skip to content

Commit 3bbb763

Browse files
chriswritescode-devchriswritescode-dev
committed
feat: set optimistic busy status on send/command/shell mutations
1 parent 0d384d6 commit 3bbb763

6 files changed

Lines changed: 214 additions & 10 deletions

File tree

‎frontend/src/hooks/__tests__/useLoadSkill.test.tsx‎

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { useLoadSkill } from '../useOpenCode'
55
import type { MessageWithParts } from '../../api/types'
66

77
import { showToast } from '../../lib/toast'
8+
import { useSessionStatus } from '../../stores/sessionStatusStore'
89

910
const mocks = vi.hoisted(() => ({
1011
sendCommand: vi.fn(),
@@ -38,6 +39,7 @@ const createWrapper = () => {
3839
describe('useLoadSkill', () => {
3940
beforeEach(() => {
4041
vi.clearAllMocks()
42+
useSessionStatus.getState().replaceStatuses({})
4143
})
4244

4345
it('calls sendCommand with correct parameters when mutate is called', async () => {
@@ -161,13 +163,29 @@ describe('useLoadSkill', () => {
161163
})
162164
})
163165

164-
it('sets session status to busy on mutate and idle on error', async () => {
166+
it('sets session status to busy on mutate', async () => {
167+
mocks.sendCommand.mockReturnValue(new Promise(() => {}))
168+
169+
const { result } = renderHook(
170+
() => useLoadSkill('http://localhost:5551', 'test-session-id', '/test/dir'),
171+
{ wrapper: createWrapper() },
172+
)
173+
174+
result.current.mutate({ skillName: 'my-skill' })
175+
176+
await waitFor(() => {
177+
expect(useSessionStatus.getState().getStatus('test-session-id').type).toBe('busy')
178+
})
179+
})
180+
181+
it('rolls back to idle on mutation error', async () => {
165182
const queryClient = new QueryClient({
166183
defaultOptions: {
167184
queries: { retry: false },
168185
mutations: { retry: false },
169186
},
170187
})
188+
171189
const testError = new Error('Command failed')
172190
mocks.sendCommand.mockRejectedValue(testError)
173191

@@ -182,6 +200,33 @@ describe('useLoadSkill', () => {
182200
expect(result.current.isError).toBe(true)
183201
})
184202

203+
expect(useSessionStatus.getState().getStatus('test-session-id').type).toBe('idle')
204+
})
205+
206+
it('does not set session status when sessionID is undefined', async () => {
207+
const queryClient = new QueryClient({
208+
defaultOptions: {
209+
queries: { retry: false },
210+
mutations: { retry: false },
211+
},
212+
})
213+
214+
mocks.sendCommand.mockReturnValue(new Promise(() => {}))
215+
216+
const { result } = renderHook(
217+
() => useLoadSkill('http://localhost:5551', undefined, '/test/dir'),
218+
{ wrapper: ({ children }) => <QueryClientProvider client={queryClient}>{children}</QueryClientProvider> },
219+
)
220+
221+
result.current.mutate({ skillName: 'my-skill' })
222+
223+
await waitFor(() => {
224+
expect(result.current.isError).toBe(true)
225+
})
226+
227+
expect(result.current.error).toBeInstanceOf(Error)
228+
expect((result.current.error as Error).message).toBe('No active session')
229+
expect(useSessionStatus.getState().getStatus('test-session-id').type).toBe('idle')
185230
})
186231

187232
it('removes optimistic message from cache on error', async () => {

‎frontend/src/hooks/useOpenCode.ts‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import type { paths, components } from "../api/opencode-types";
1212
import { parseNetworkError } from "../lib/opencode-errors";
1313
import { showToast } from "../lib/toast";
1414
import { useSendErrorStore } from "../stores/sendErrorStore";
15+
import { beginOptimisticBusy, rollbackOptimisticBusy } from "../stores/sessionStatusStore";
16+
import type { SessionStatusType } from "../stores/sessionStatusStore";
1517
import { invalidateSessionListCaches, messagesQueryKey } from "../lib/queryInvalidation";
1618

1719
type AssistantMessage = components["schemas"]["AssistantMessage"];
@@ -22,6 +24,8 @@ type SendPromptRequest = NonNullable<
2224

2325
type SendCommandResponse = paths["/session/{sessionID}/command"]["post"]["responses"]["200"]["content"]["application/json"];
2426

27+
type OptimisticBusyContext = { previousBusyStatus: SessionStatusType };
28+
2529
const parseModelString = (model: string) => {
2630
const [providerID, ...rest] = model.split("/");
2731
const modelID = rest.join("/");
@@ -477,8 +481,13 @@ export const useSendPrompt = (opcodeUrl: string | null | undefined, directory?:
477481

478482
return { optimisticUserID, response, queued: false };
479483
},
480-
onError: (error, variables) => {
484+
onMutate: ({ sessionID }): OptimisticBusyContext => ({
485+
previousBusyStatus: beginOptimisticBusy(sessionID),
486+
}),
487+
onError: (error, variables, context: OptimisticBusyContext | undefined) => {
481488
const { sessionID } = variables;
489+
rollbackOptimisticBusy(sessionID, context?.previousBusyStatus ?? { type: 'idle' });
490+
482491
const queryKey = messagesQueryKey(opcodeUrl, sessionID, directory);
483492

484493
queryClient.setQueryData<MessageWithParts[]>(
@@ -726,8 +735,12 @@ export const useSendShell = (opcodeUrl: string | null | undefined, directory?: s
726735

727736
return { optimisticUserID, response };
728737
},
729-
onError: (_, variables) => {
738+
onMutate: ({ sessionID }): OptimisticBusyContext => ({
739+
previousBusyStatus: beginOptimisticBusy(sessionID),
740+
}),
741+
onError: (_error, variables, context: OptimisticBusyContext | undefined) => {
730742
const { sessionID } = variables;
743+
rollbackOptimisticBusy(sessionID, context?.previousBusyStatus ?? { type: 'idle' });
731744
queryClient.setQueryData<MessageWithParts[]>(
732745
messagesQueryKey(opcodeUrl, sessionID, directory),
733746
(old) => {
@@ -788,7 +801,8 @@ export const useLoadSkill = (
788801
return useMutation<
789802
{ optimisticUserID: string; response: SendCommandResponse },
790803
Error,
791-
{ skillName: string }
804+
{ skillName: string },
805+
OptimisticBusyContext | undefined
792806
>({
793807
mutationFn: async ({ skillName }: { skillName: string }) => {
794808
if (!client) throw new Error("No OpenCode client available");
@@ -817,7 +831,14 @@ export const useLoadSkill = (
817831
const response = await client.sendCommand(sessionID, { command: skillName, arguments: "" });
818832
return { optimisticUserID, response };
819833
},
820-
onError: (error) => {
834+
onMutate: () => {
835+
if (!sessionID) return undefined;
836+
return { previousBusyStatus: beginOptimisticBusy(sessionID) };
837+
},
838+
onError: (error, _variables, context: OptimisticBusyContext | undefined) => {
839+
if (sessionID && context?.previousBusyStatus) {
840+
rollbackOptimisticBusy(sessionID, context.previousBusyStatus);
841+
}
821842
if (sessionID) {
822843
const queryKey = messagesQueryKey(opcodeUrl, sessionID, directory);
823844
queryClient.setQueryData<MessageWithParts[]>(

‎frontend/src/hooks/useSendPrompt.test.ts‎

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
import { describe, it, expect, vi, beforeEach } from 'vitest'
2-
import { renderHook } from '@testing-library/react'
2+
import { renderHook, waitFor } from '@testing-library/react'
33
import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
44
import { createElement } from 'react'
55
import { useSendPrompt } from './useOpenCode'
66
import { FetchError } from '../api/fetchWrapper'
7+
import { useSessionStatus } from '../stores/sessionStatusStore'
78

89
const mockSendPrompt = vi.fn()
910
const mockSendPromptAsync = vi.fn()
@@ -19,10 +20,6 @@ vi.mock('../api/opencode', async () => {
1920
}
2021
})
2122

22-
vi.mock('@/stores/sessionStatusStore', () => ({
23-
useSessionStatus: vi.fn(() => vi.fn()),
24-
}))
25-
2623
vi.mock('../lib/toast', () => ({
2724
showToast: { error: vi.fn() },
2825
}))
@@ -62,6 +59,7 @@ describe('useSendPrompt', () => {
6259

6360
beforeEach(() => {
6461
vi.clearAllMocks()
62+
useSessionStatus.getState().replaceStatuses({})
6563
queryClient = createTestQueryClient()
6664
mockSendPrompt.mockResolvedValue({
6765
info: { id: 'test-response' },
@@ -194,4 +192,25 @@ describe('useSendPrompt', () => {
194192

195193
expect(mockClearError).toHaveBeenCalledWith('session-2')
196194
})
195+
196+
it('sets session status to busy immediately on send', async () => {
197+
mockSendPrompt.mockReturnValue(new Promise(() => {}))
198+
199+
const { result } = renderHookWithProviders()
200+
result.current.mutate({ sessionID: 'busy-session', prompt: 'Hi' })
201+
202+
await waitFor(() => {
203+
expect(useSessionStatus.getState().getStatus('busy-session').type).toBe('busy')
204+
})
205+
})
206+
207+
it('rolls back to idle when send fails', async () => {
208+
mockSendPrompt.mockRejectedValue(new Error('boom'))
209+
210+
const { result } = renderHookWithProviders()
211+
result.current.mutate({ sessionID: 'busy-session', prompt: 'Hi' })
212+
213+
await waitFor(() => expect(result.current.isError).toBe(true))
214+
expect(useSessionStatus.getState().getStatus('busy-session').type).toBe('idle')
215+
})
197216
})

‎frontend/src/hooks/useSendShell.test.ts‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
import { describe, it, expect, vi, beforeEach } from 'vitest'
2+
import { renderHook, waitFor } from '@testing-library/react'
3+
import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
4+
import { createElement } from 'react'
5+
import { useSendShell } from './useOpenCode'
6+
import { useSessionStatus } from '../stores/sessionStatusStore'
7+
8+
const mockSendShell = vi.fn()
9+
10+
vi.mock('../api/opencode', async () => {
11+
const actual = await vi.importActual('../api/opencode')
12+
return {
13+
...actual,
14+
OpenCodeClient: vi.fn().mockImplementation(() => ({
15+
sendShell: mockSendShell,
16+
})),
17+
}
18+
})
19+
20+
vi.mock('../lib/toast', () => ({
21+
showToast: { error: vi.fn() },
22+
}))
23+
24+
const createTestQueryClient = () =>
25+
new QueryClient({
26+
defaultOptions: {
27+
queries: {
28+
retry: false,
29+
},
30+
},
31+
})
32+
33+
describe('useSendShell', () => {
34+
let queryClient: QueryClient
35+
36+
beforeEach(() => {
37+
vi.clearAllMocks()
38+
useSessionStatus.getState().replaceStatuses({})
39+
queryClient = createTestQueryClient()
40+
})
41+
42+
const renderHookWithProviders = () =>
43+
renderHook(
44+
() => useSendShell('http://localhost:5551', '/test'),
45+
{
46+
wrapper: ({ children }) =>
47+
createElement(QueryClientProvider, { client: queryClient }, children),
48+
}
49+
)
50+
51+
it('sets session status to busy on shell send', async () => {
52+
mockSendShell.mockReturnValue(new Promise(() => {}))
53+
54+
const { result } = renderHookWithProviders()
55+
result.current.mutate({ sessionID: 'shell-session', command: 'ls', agent: 'general' })
56+
57+
await waitFor(() => {
58+
expect(useSessionStatus.getState().getStatus('shell-session').type).toBe('busy')
59+
})
60+
})
61+
62+
it('rolls back to idle when shell send fails', async () => {
63+
mockSendShell.mockRejectedValue(new Error('boom'))
64+
65+
const { result } = renderHookWithProviders()
66+
result.current.mutate({ sessionID: 'shell-session', command: 'ls', agent: 'general' })
67+
68+
await waitFor(() => expect(result.current.isError).toBe(true))
69+
expect(useSessionStatus.getState().getStatus('shell-session').type).toBe('idle')
70+
})
71+
})

‎frontend/src/stores/sessionStatusStore.test.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
import { describe, it, expect, beforeEach } from 'vitest'
2+
import { useSessionStatus, beginOptimisticBusy, rollbackOptimisticBusy } from './sessionStatusStore'
3+
4+
describe('optimistic busy helpers', () => {
5+
beforeEach(() => {
6+
useSessionStatus.getState().replaceStatuses({})
7+
})
8+
9+
it('beginOptimisticBusy on idle session sets busy and returns idle', () => {
10+
const previous = beginOptimisticBusy('session-1')
11+
expect(previous).toEqual({ type: 'idle' })
12+
expect(useSessionStatus.getState().getStatus('session-1')).toEqual({ type: 'busy' })
13+
})
14+
15+
it('rollbackOptimisticBusy returns session to idle', () => {
16+
beginOptimisticBusy('session-1')
17+
rollbackOptimisticBusy('session-1', { type: 'idle' })
18+
expect(useSessionStatus.getState().getStatus('session-1')).toEqual({ type: 'idle' })
19+
})
20+
21+
it('beginOptimisticBusy on already-busy session returns busy and keeps busy', () => {
22+
useSessionStatus.getState().setStatus('session-1', { type: 'busy' })
23+
const previous = beginOptimisticBusy('session-1')
24+
expect(previous).toEqual({ type: 'busy' })
25+
expect(useSessionStatus.getState().getStatus('session-1')).toEqual({ type: 'busy' })
26+
rollbackOptimisticBusy('session-1', { type: 'busy' })
27+
expect(useSessionStatus.getState().getStatus('session-1')).toEqual({ type: 'busy' })
28+
})
29+
30+
it('busy state is scoped per sessionID', () => {
31+
beginOptimisticBusy('a')
32+
expect(useSessionStatus.getState().getStatus('b')).toEqual({ type: 'idle' })
33+
})
34+
})

‎frontend/src/stores/sessionStatusStore.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,3 +104,17 @@ export const useSessionStatusForSession = (sessionID: string | undefined): Sessi
104104
sessionID ? (state.statuses.get(sessionID) ?? DEFAULT_STATUS) : DEFAULT_STATUS
105105
)
106106
}
107+
108+
export function beginOptimisticBusy(sessionID: string): SessionStatusType {
109+
const { getStatus, setStatus } = useSessionStatus.getState()
110+
const previous = getStatus(sessionID)
111+
setStatus(sessionID, { type: 'busy' })
112+
return previous
113+
}
114+
115+
export function rollbackOptimisticBusy(
116+
sessionID: string,
117+
previous: SessionStatusType,
118+
): void {
119+
useSessionStatus.getState().setStatus(sessionID, previous)
120+
}

0 commit comments

Comments
 (0)