Skip to content

Commit c60e473

Browse files
vincentkocvincentkoc
committed
fix(tui): consume early accepted run finals
1 parent d4bedce commit c60e473

4 files changed

Lines changed: 67 additions & 3 deletions

File tree

‎src/tui/tui-command-handlers.test.ts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ type SelectableOverlay = {
1515
};
1616
type SetActivityStatusMock = ReturnType<typeof vi.fn> & ((text: string) => void);
1717
type SetSessionMock = ReturnType<typeof vi.fn> & ((key: string) => Promise<void>);
18+
type ConsumeCompletedRunMock = ReturnType<typeof vi.fn> & ((runId: string) => boolean);
19+
type FlushPendingHistoryRefreshMock = ReturnType<typeof vi.fn> & (() => void);
1820

1921
async function flushAsyncSelect() {
2022
await new Promise<void>((resolve) => setImmediate(resolve));
@@ -81,6 +83,8 @@ function createHarness(params?: {
8183
currentAgentId?: string;
8284
currentSessionKey?: string;
8385
abortActive?: AbortActiveMock;
86+
consumeCompletedRunForPendingSend?: ConsumeCompletedRunMock;
87+
flushPendingHistoryRefreshIfIdle?: FlushPendingHistoryRefreshMock;
8488
}) {
8589
const sendChat = params?.sendChat ?? vi.fn().mockResolvedValue({ runId: "r1" });
8690
const getGatewayStatus = params?.getGatewayStatus ?? vi.fn().mockResolvedValue({});
@@ -153,6 +157,8 @@ function createHarness(params?: {
153157
noteLocalBtwRunId,
154158
forgetLocalRunId,
155159
forgetLocalBtwRunId: vi.fn(),
160+
consumeCompletedRunForPendingSend: params?.consumeCompletedRunForPendingSend,
161+
flushPendingHistoryRefreshIfIdle: params?.flushPendingHistoryRefreshIfIdle,
156162
runAuthFlow,
157163
requestExit,
158164
});
@@ -562,6 +568,29 @@ describe("tui command handlers", () => {
562568
expect(noteLocalRunId).toHaveBeenCalledWith("run-accepted");
563569
});
564570

571+
it("does not reintroduce a backend-accepted runId after an early terminal event", async () => {
572+
const sendChat = vi.fn().mockResolvedValue({ runId: "run-accepted" });
573+
const consumeCompletedRunForPendingSend = vi.fn((runId: string) => runId === "run-accepted");
574+
const flushPendingHistoryRefreshIfIdle = vi.fn();
575+
const { handleCommand, state, noteLocalRunId, forgetLocalRunId, setActivityStatus } =
576+
createHarness({
577+
sendChat,
578+
consumeCompletedRunForPendingSend,
579+
flushPendingHistoryRefreshIfIdle,
580+
});
581+
582+
await handleCommand("hello");
583+
584+
const sentRunId = (firstMockArg(sendChat, "sendChat") as { runId: string }).runId;
585+
expect(consumeCompletedRunForPendingSend).toHaveBeenCalledWith("run-accepted");
586+
expect(forgetLocalRunId).toHaveBeenCalledWith(sentRunId);
587+
expect(noteLocalRunId).not.toHaveBeenCalledWith("run-accepted");
588+
expect(state.pendingChatRunId).toBeNull();
589+
expect(state.pendingOptimisticUserMessage).toBe(false);
590+
expect(setActivityStatus).toHaveBeenCalledWith("idle");
591+
expect(flushPendingHistoryRefreshIfIdle).toHaveBeenCalledTimes(1);
592+
});
593+
565594
it("clears the pending runId if sendChat fails", async () => {
566595
const sendChat = vi.fn().mockRejectedValue(new Error("boom"));
567596
const { handleCommand, state } = createHarness({ sendChat });

‎src/tui/tui-command-handlers.ts‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,8 @@ type CommandHandlerContext = {
5959
noteLocalBtwRunId?: (runId: string) => void;
6060
forgetLocalRunId?: (runId: string) => void;
6161
forgetLocalBtwRunId?: (runId: string) => void;
62+
consumeCompletedRunForPendingSend?: (runId: string) => boolean;
63+
flushPendingHistoryRefreshIfIdle?: () => void;
6264
runAuthFlow?: (params: {
6365
provider?: string;
6466
}) => Promise<{ exitCode: number | null; signal: NodeJS.Signals | null }>;
@@ -111,6 +113,8 @@ export function createCommandHandlers(context: CommandHandlerContext) {
111113
noteLocalBtwRunId,
112114
forgetLocalRunId,
113115
forgetLocalBtwRunId,
116+
consumeCompletedRunForPendingSend,
117+
flushPendingHistoryRefreshIfIdle,
114118
runAuthFlow,
115119
requestExit,
116120
} = context;
@@ -750,13 +754,24 @@ export function createCommandHandlers(context: CommandHandlerContext) {
750754
});
751755
if (!isBtw) {
752756
const acceptedRunId = sendResult.runId || runId;
757+
const acceptedRunAlreadyCompleted =
758+
acceptedRunId !== runId && (consumeCompletedRunForPendingSend?.(acceptedRunId) ?? false);
753759
if (acceptedRunId !== runId) {
754760
forgetLocalRunId?.(runId);
755-
noteLocalRunId?.(acceptedRunId);
761+
if (!acceptedRunAlreadyCompleted) {
762+
noteLocalRunId?.(acceptedRunId);
763+
}
756764
}
757765
if (state.pendingOptimisticUserMessage) {
758-
state.pendingChatRunId = acceptedRunId;
759-
setActivityStatus("waiting");
766+
if (acceptedRunAlreadyCompleted) {
767+
state.pendingOptimisticUserMessage = false;
768+
state.pendingChatRunId = null;
769+
setActivityStatus("idle");
770+
flushPendingHistoryRefreshIfIdle?.();
771+
} else {
772+
state.pendingChatRunId = acceptedRunId;
773+
setActivityStatus("waiting");
774+
}
760775
tui.requestRender();
761776
}
762777
}

‎src/tui/tui-event-handlers.ts‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ export function createEventHandlers(context: EventHandlerContext) {
7575
const sessionRuns = new Map<string, number>();
7676
const finalizedRuns = new Map<string, number>();
7777
const finalizedRunsWithDisplay = new Map<string, number>();
78+
const completedRuns = new Map<string, number>();
7879
const postFinalizingRuns = new Map<string, number>();
7980
let streamAssembler = new TuiStreamAssembler();
8081
let lastSessionKey = state.currentSessionKey;
@@ -181,6 +182,7 @@ export function createEventHandlers(context: EventHandlerContext) {
181182
}
182183
finalizedRuns.clear();
183184
finalizedRunsWithDisplay.clear();
185+
completedRuns.clear();
184186
sessionRuns.clear();
185187
postFinalizingRuns.clear();
186188
streamAssembler = new TuiStreamAssembler();
@@ -248,13 +250,15 @@ export function createEventHandlers(context: EventHandlerContext) {
248250

249251
const noteFinalizedRun = (runId: string, opts?: { displayedFinal?: boolean }) => {
250252
finalizedRuns.set(runId, Date.now());
253+
completedRuns.set(runId, Date.now());
251254
if (opts?.displayedFinal === true) {
252255
finalizedRunsWithDisplay.set(runId, Date.now());
253256
}
254257
sessionRuns.delete(runId);
255258
streamAssembler.drop(runId);
256259
pruneRunMap(finalizedRuns);
257260
pruneRunMap(finalizedRunsWithDisplay);
261+
pruneRunMap(completedRuns);
258262
};
259263

260264
const notePostFinalizingRun = (runId: string) => {
@@ -329,6 +333,8 @@ export function createEventHandlers(context: EventHandlerContext) {
329333
wasActiveRun: boolean;
330334
status: "aborted" | "error";
331335
}) => {
336+
completedRuns.set(params.runId, Date.now());
337+
pruneRunMap(completedRuns);
332338
streamAssembler.drop(params.runId);
333339
sessionRuns.delete(params.runId);
334340
clearActiveRunIfMatch(params.runId);
@@ -742,12 +748,22 @@ export function createEventHandlers(context: EventHandlerContext) {
742748
clearStreamingWatchdog();
743749
};
744750

751+
const consumeCompletedRunForPendingSend = (runId: string) => {
752+
if (!completedRuns.has(runId)) {
753+
return false;
754+
}
755+
completedRuns.delete(runId);
756+
return true;
757+
};
758+
745759
return {
746760
handleChatEvent,
747761
handleAgentEvent,
748762
handleBtwEvent,
749763
pauseStreamingWatchdog,
750764
reconnectStreamingWatchdog,
765+
consumeCompletedRunForPendingSend,
766+
flushPendingHistoryRefreshIfIdle,
751767
dispose,
752768
};
753769
}

‎src/tui/tui.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1225,6 +1225,8 @@ export async function runTui(opts: RunTuiOptions): Promise<TuiResult> {
12251225
handleBtwEvent,
12261226
pauseStreamingWatchdog,
12271227
reconnectStreamingWatchdog,
1228+
consumeCompletedRunForPendingSend,
1229+
flushPendingHistoryRefreshIfIdle,
12281230
} = createEventHandlers({
12291231
chatLog,
12301232
btw,
@@ -1314,6 +1316,8 @@ export async function runTui(opts: RunTuiOptions): Promise<TuiResult> {
13141316
noteLocalBtwRunId,
13151317
forgetLocalRunId,
13161318
forgetLocalBtwRunId,
1319+
consumeCompletedRunForPendingSend,
1320+
flushPendingHistoryRefreshIfIdle,
13171321
runAuthFlow,
13181322
requestExit,
13191323
});

0 commit comments

Comments
 (0)