Skip to content

Commit 9b81ba1

Browse files
jeanduplessisjeanduplessis
authored
fix(security-agent): harden async command settlement (#3871)
* fix(security-agent): harden async command settlement * fix(security-agent): address command review feedback
1 parent 32f26ca commit 9b81ba1

12 files changed

Lines changed: 845 additions & 84 deletions

File tree

‎apps/web/src/components/security-agent/SecurityAgentContext.tsx‎

Lines changed: 152 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
'use client';
22

3-
import { createContext, use, useState, useCallback, useEffect, useMemo, useRef } from 'react';
3+
import { createContext, use, useCallback, useEffect, useMemo, useReducer, useRef } from 'react';
44
import { useTRPC } from '@/lib/trpc/utils';
55
import { useQuery, useMutation, useQueries, useQueryClient } from '@tanstack/react-query';
66
import { toast } from 'sonner';
@@ -126,6 +126,83 @@ type SecurityAgentCommand = {
126126
lastErrorRedacted: string | null;
127127
};
128128

129+
type SecurityAgentProviderState = {
130+
optimisticStartingAnalysisIds: Set<string>;
131+
trackedCommandIds: Set<string>;
132+
processedTerminalCommandIds: Set<string>;
133+
gitHubError: string | null;
134+
};
135+
136+
type SecurityAgentProviderAction =
137+
| { type: 'track-command'; commandId: string }
138+
| { type: 'add-optimistic-analysis'; findingId: string }
139+
| { type: 'remove-optimistic-analysis'; findingId: string }
140+
| { type: 'settle-commands'; commands: SecurityAgentCommand[]; gitHubError?: string }
141+
| { type: 'prune-processed-commands'; polledCommandIds: Set<string> }
142+
| { type: 'set-github-error'; error: string | null };
143+
144+
function createSecurityAgentProviderState(): SecurityAgentProviderState {
145+
return {
146+
optimisticStartingAnalysisIds: new Set(),
147+
trackedCommandIds: new Set(),
148+
processedTerminalCommandIds: new Set(),
149+
gitHubError: null,
150+
};
151+
}
152+
153+
function securityAgentProviderReducer(
154+
state: SecurityAgentProviderState,
155+
action: SecurityAgentProviderAction
156+
): SecurityAgentProviderState {
157+
switch (action.type) {
158+
case 'track-command':
159+
return {
160+
...state,
161+
trackedCommandIds: new Set(state.trackedCommandIds).add(action.commandId),
162+
};
163+
case 'add-optimistic-analysis':
164+
return {
165+
...state,
166+
optimisticStartingAnalysisIds: new Set(state.optimisticStartingAnalysisIds).add(
167+
action.findingId
168+
),
169+
};
170+
case 'remove-optimistic-analysis': {
171+
const optimisticStartingAnalysisIds = new Set(state.optimisticStartingAnalysisIds);
172+
optimisticStartingAnalysisIds.delete(action.findingId);
173+
return { ...state, optimisticStartingAnalysisIds };
174+
}
175+
case 'settle-commands': {
176+
const optimisticStartingAnalysisIds = new Set(state.optimisticStartingAnalysisIds);
177+
const trackedCommandIds = new Set(state.trackedCommandIds);
178+
const processedTerminalCommandIds = new Set(state.processedTerminalCommandIds);
179+
for (const command of action.commands) {
180+
if (command.findingId) optimisticStartingAnalysisIds.delete(command.findingId);
181+
trackedCommandIds.delete(command.id);
182+
processedTerminalCommandIds.add(command.id);
183+
}
184+
return {
185+
optimisticStartingAnalysisIds,
186+
trackedCommandIds,
187+
processedTerminalCommandIds,
188+
gitHubError: action.gitHubError ?? state.gitHubError,
189+
};
190+
}
191+
case 'prune-processed-commands': {
192+
const processedTerminalCommandIds = new Set(
193+
[...state.processedTerminalCommandIds].filter(commandId =>
194+
action.polledCommandIds.has(commandId)
195+
)
196+
);
197+
return processedTerminalCommandIds.size === state.processedTerminalCommandIds.size
198+
? state
199+
: { ...state, processedTerminalCommandIds };
200+
}
201+
case 'set-github-error':
202+
return { ...state, gitHubError: action.error };
203+
}
204+
}
205+
129206
function commandFailureDescription(command: SecurityAgentCommand): string {
130207
switch (command.resultCode) {
131208
case 'OWNER_CAP_REACHED':
@@ -156,19 +233,22 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
156233
const queryClient = useQueryClient();
157234
const isOrg = !!organizationId;
158235

159-
const [optimisticStartingAnalysisIds, setOptimisticStartingAnalysisIds] = useState<Set<string>>(
160-
new Set()
236+
const [providerState, dispatchProviderState] = useReducer(
237+
securityAgentProviderReducer,
238+
undefined,
239+
createSecurityAgentProviderState
161240
);
162-
const [trackedCommandIds, setTrackedCommandIds] = useState<Set<string>>(new Set());
163-
const [gitHubError, setGitHubError] = useState<string | null>(null);
164241
const toggleEnabledInFlightRef = useRef(false);
165-
const processedTerminalCommandIdsRef = useRef(new Set<string>());
166-
const recoveredCommandIdsRef = useRef(new Set<string>());
167-
const commandSuccessCallbacksRef = useRef(new Map<string, () => void>());
242+
const commandSuccessCallbacksRef = useRef<Map<string, () => void>>(null);
168243

169244
const trackCommand = useCallback((commandId: string, onSuccess?: () => void) => {
170-
if (onSuccess) commandSuccessCallbacksRef.current.set(commandId, onSuccess);
171-
setTrackedCommandIds(previous => new Set(previous).add(commandId));
245+
if (onSuccess) {
246+
if (commandSuccessCallbacksRef.current === null) {
247+
commandSuccessCallbacksRef.current = new Map();
248+
}
249+
commandSuccessCallbacksRef.current.set(commandId, onSuccess);
250+
}
251+
dispatchProviderState({ type: 'track-command', commandId });
172252
}, []);
173253

174254
const invalidateAcceptedQueueQueries = useCallback(() => {
@@ -269,12 +349,24 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
269349
query.state.data && query.state.data.length > 0 ? COMMAND_POLL_INTERVAL_MS : false,
270350
});
271351

272-
useEffect(() => {
273-
for (const command of activeCommandsData ?? []) recoveredCommandIdsRef.current.add(command.id);
274-
}, [activeCommandsData]);
352+
const commandIdsToPoll = useMemo(() => {
353+
const commandIds = new Set(providerState.trackedCommandIds);
354+
for (const command of activeCommandsData ?? []) commandIds.add(command.id);
355+
return commandIds;
356+
}, [activeCommandsData, providerState.trackedCommandIds]);
275357

276-
const commandIdsToPoll = new Set([...trackedCommandIds, ...recoveredCommandIdsRef.current]);
277-
for (const command of activeCommandsData ?? []) commandIdsToPoll.add(command.id);
358+
useEffect(() => {
359+
if (
360+
[...providerState.processedTerminalCommandIds].some(
361+
commandId => !commandIdsToPoll.has(commandId)
362+
)
363+
) {
364+
dispatchProviderState({
365+
type: 'prune-processed-commands',
366+
polledCommandIds: commandIdsToPoll,
367+
});
368+
}
369+
}, [commandIdsToPoll, providerState.processedTerminalCommandIds]);
278370

279371
const commandStatusQueries = useQueries({
280372
queries: [...commandIdsToPoll].map(commandId => ({
@@ -304,32 +396,32 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
304396
command => command.commandType === 'dismiss_finding'
305397
);
306398
const startingAnalysisIds = useMemo(() => {
307-
const ids = new Set(optimisticStartingAnalysisIds);
399+
const ids = new Set(providerState.optimisticStartingAnalysisIds);
308400
for (const command of activeCommands) {
309401
if (command.commandType === 'start_analysis' && command.findingId) {
310402
ids.add(command.findingId);
311403
}
312404
}
313405
return ids;
314-
}, [activeCommands, optimisticStartingAnalysisIds]);
406+
}, [activeCommands, providerState.optimisticStartingAnalysisIds]);
315407

316408
useEffect(() => {
317-
for (const query of commandStatusQueries) {
409+
const terminalCommands = commandStatusQueries.flatMap(query => {
318410
const command = query.data;
319-
if (!command || command.status === 'accepted' || command.status === 'running') continue;
320-
if (processedTerminalCommandIdsRef.current.has(command.id)) continue;
321-
processedTerminalCommandIdsRef.current.add(command.id);
322-
recoveredCommandIdsRef.current.delete(command.id);
411+
return command && command.status !== 'accepted' && command.status !== 'running'
412+
? [command]
413+
: [];
414+
});
415+
const unprocessedTerminalCommands = terminalCommands.filter(
416+
command => !providerState.processedTerminalCommandIds.has(command.id)
417+
);
418+
if (unprocessedTerminalCommands.length === 0) return;
419+
420+
let terminalGitHubError: string | undefined;
421+
for (const command of unprocessedTerminalCommands) {
323422
invalidateAcceptedQueueQueries();
324-
if (command.findingId) {
325-
setOptimisticStartingAnalysisIds(previous => {
326-
const next = new Set(previous);
327-
next.delete(command.findingId ?? '');
328-
return next;
329-
});
330-
}
331-
const successCallback = commandSuccessCallbacksRef.current.get(command.id);
332-
commandSuccessCallbacksRef.current.delete(command.id);
423+
const successCallback = commandSuccessCallbacksRef.current?.get(command.id);
424+
commandSuccessCallbacksRef.current?.delete(command.id);
333425
if (command.status === 'failed') {
334426
const title =
335427
command.commandType === 'sync'
@@ -338,7 +430,7 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
338430
? 'Failed to dismiss finding'
339431
: 'Failed to start analysis';
340432
if (command.resultCode === 'GITHUB_AUTH_INVALID') {
341-
setGitHubError(commandFailureDescription(command));
433+
terminalGitHubError = commandFailureDescription(command);
342434
}
343435
toast.error(title, { description: commandFailureDescription(command), duration: 8000 });
344436
} else {
@@ -349,26 +441,31 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
349441
);
350442
}
351443
}
352-
setTrackedCommandIds(previous => {
353-
const next = new Set(previous);
354-
next.delete(command.id);
355-
return next;
356-
});
357444
}
358-
}, [commandStatusQueries, invalidateAcceptedQueueQueries]);
445+
446+
dispatchProviderState({
447+
type: 'settle-commands',
448+
commands: unprocessedTerminalCommands,
449+
gitHubError: terminalGitHubError,
450+
});
451+
}, [
452+
commandStatusQueries,
453+
providerState.processedTerminalCommandIds,
454+
invalidateAcceptedQueueQueries,
455+
]);
359456

360457
// ---- Mutations (org) ----
361458
const { mutate: orgSyncMutate, isPending: isOrgSyncPending } = useMutation(
362459
trpc.organizations.securityAgent.triggerSync.mutationOptions({
363460
onSuccess: data => {
364-
setGitHubError(null);
461+
dispatchProviderState({ type: 'set-github-error', error: null });
365462
toast.success('Sync queued');
366463
trackCommand(data.commandId);
367464
},
368465
onError: error => {
369466
const message = error instanceof Error ? error.message : String(error);
370467
if (isGitHubIntegrationError(error)) {
371-
setGitHubError(message);
468+
dispatchProviderState({ type: 'set-github-error', error: message });
372469
toast.error('GitHub Integration Error', {
373470
description:
374471
'The GitHub App may have been uninstalled. Please check your integrations.',
@@ -440,14 +537,14 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
440537
const { mutate: orgStartAnalysisMutate } = useMutation(
441538
trpc.organizations.securityAgent.startAnalysis.mutationOptions({
442539
onSuccess: async data => {
443-
setGitHubError(null);
540+
dispatchProviderState({ type: 'set-github-error', error: null });
444541
toast.success(manualAnalysisAdmissionCopy.successTitle);
445542
trackCommand(data.commandId);
446543
},
447544
onError: (error, variables) => {
448545
const message = error instanceof Error ? error.message : String(error);
449546
if (isGitHubIntegrationError(error)) {
450-
setGitHubError(message);
547+
dispatchProviderState({ type: 'set-github-error', error: message });
451548
toast.error('GitHub Integration Error', {
452549
description:
453550
'The GitHub App may have been uninstalled. Please check your integrations.',
@@ -459,10 +556,9 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
459556
});
460557
}
461558
void queryClient.invalidateQueries();
462-
setOptimisticStartingAnalysisIds(prev => {
463-
const next = new Set(prev);
464-
next.delete(variables.findingId);
465-
return next;
559+
dispatchProviderState({
560+
type: 'remove-optimistic-analysis',
561+
findingId: variables.findingId,
466562
});
467563
},
468564
})
@@ -486,14 +582,14 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
486582
const { mutate: personalSyncMutate, isPending: isPersonalSyncPending } = useMutation(
487583
trpc.securityAgent.triggerSync.mutationOptions({
488584
onSuccess: data => {
489-
setGitHubError(null);
585+
dispatchProviderState({ type: 'set-github-error', error: null });
490586
toast.success('Sync queued');
491587
trackCommand(data.commandId);
492588
},
493589
onError: error => {
494590
const message = error instanceof Error ? error.message : String(error);
495591
if (isGitHubIntegrationError(error)) {
496-
setGitHubError(message);
592+
dispatchProviderState({ type: 'set-github-error', error: message });
497593
toast.error('GitHub Integration Error', {
498594
description:
499595
'The GitHub App may have been uninstalled. Please check your integrations.',
@@ -565,14 +661,14 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
565661
const { mutate: personalStartAnalysisMutate } = useMutation(
566662
trpc.securityAgent.startAnalysis.mutationOptions({
567663
onSuccess: async data => {
568-
setGitHubError(null);
664+
dispatchProviderState({ type: 'set-github-error', error: null });
569665
toast.success(manualAnalysisAdmissionCopy.successTitle);
570666
trackCommand(data.commandId);
571667
},
572668
onError: (error, variables) => {
573669
const message = error instanceof Error ? error.message : String(error);
574670
if (isGitHubIntegrationError(error)) {
575-
setGitHubError(message);
671+
dispatchProviderState({ type: 'set-github-error', error: message });
576672
toast.error('GitHub Integration Error', {
577673
description:
578674
'The GitHub App may have been uninstalled. Please check your integrations.',
@@ -584,10 +680,9 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
584680
});
585681
}
586682
void queryClient.invalidateQueries();
587-
setOptimisticStartingAnalysisIds(prev => {
588-
const next = new Set(prev);
589-
next.delete(variables.findingId);
590-
return next;
683+
dispatchProviderState({
684+
type: 'remove-optimistic-analysis',
685+
findingId: variables.findingId,
591686
});
592687
},
593688
})
@@ -721,7 +816,7 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
721816

722817
const handleStartAnalysis = useCallback(
723818
(findingId: string, { retrySandboxOnly }: { retrySandboxOnly?: boolean } = {}) => {
724-
setOptimisticStartingAnalysisIds(prev => new Set(prev).add(findingId));
819+
dispatchProviderState({ type: 'add-optimistic-analysis', findingId });
725820
if (isOrg && organizationId) {
726821
orgStartAnalysisMutate({ organizationId, findingId, retrySandboxOnly });
727822
} else {
@@ -802,7 +897,7 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
802897
isTogglingEnabled: isOrg ? isOrgSetEnabledPending : isPersonalSetEnabledPending,
803898
isDeletingFindings: isOrg ? isOrgDeleteFindingsPending : isPersonalDeleteFindingsPending,
804899
startingAnalysisIds,
805-
gitHubError,
900+
gitHubError: providerState.gitHubError,
806901
orphanedRepositories: orphanedReposData ?? [],
807902
}),
808903
[
@@ -837,7 +932,7 @@ export function SecurityAgentProvider({ organizationId, children }: SecurityAgen
837932
isOrgDeleteFindingsPending,
838933
isPersonalDeleteFindingsPending,
839934
startingAnalysisIds,
840-
gitHubError,
935+
providerState.gitHubError,
841936
orphanedReposData,
842937
triageModelSlug,
843938
analysisModelSlug,

0 commit comments

Comments
 (0)