Skip to content

Commit ba84e08

Browse files
steipetetiffanychum
steipete
and
tiffanychum
committed
fix(ui): bind stale run state to run identity
Co-authored-by: Tiffany Chum <[email protected]>
1 parent 60cf7ea commit ba84e08

18 files changed

Lines changed: 523 additions & 102 deletions

‎src/gateway/server-chat.agent-events.test.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ describe("agent event handler", () => {
102102
trackTrackedRunTerminalPersistence?: AgentEventHandlerOptions["trackTrackedRunTerminalPersistence"];
103103
resolveActiveLifecycleGenerationForRun?: (runId: string) => string | undefined;
104104
updateRunToolErrorSummary?: AgentEventHandlerOptions["updateRunToolErrorSummary"];
105+
resolveSessionActiveRunState?: AgentEventHandlerOptions["resolveSessionActiveRunState"];
105106
}) {
106107
const nowSpy =
107108
params?.now === undefined ? undefined : vi.spyOn(Date, "now").mockReturnValue(params.now);
@@ -136,6 +137,7 @@ describe("agent event handler", () => {
136137
trackTrackedRunTerminalPersistence: params?.trackTrackedRunTerminalPersistence,
137138
resolveActiveLifecycleGenerationForRun: params?.resolveActiveLifecycleGenerationForRun,
138139
updateRunToolErrorSummary: params?.updateRunToolErrorSummary,
140+
resolveSessionActiveRunState: params?.resolveSessionActiveRunState,
139141
});
140142

141143
return {
@@ -2063,8 +2065,13 @@ describe("agent event handler", () => {
20632065
runtimeMs: 800,
20642066
abortedLastRun: false,
20652067
});
2068+
const resolveSessionActiveRunState = vi
2069+
.fn<NonNullable<AgentEventHandlerOptions["resolveSessionActiveRunState"]>>()
2070+
.mockReturnValueOnce({ active: true, runIds: ["run-finished"] })
2071+
.mockReturnValue({ active: false, runIds: [] });
20662072
const { broadcastToConnIds, sessionEventSubscribers, handler } = createHarness({
20672073
resolveSessionKeyForRun: () => "session-finished",
2074+
resolveSessionActiveRunState,
20682075
});
20692076

20702077
sessionEventSubscribers.subscribe("conn-session");
@@ -2104,16 +2111,28 @@ describe("agent event handler", () => {
21042111
([event]) => event === "sessions.changed",
21052112
);
21062113
expect(sessionsChangedCalls).toHaveLength(2);
2114+
expectPayloadFields(sessionsChangedCalls[0]?.[1], {
2115+
sessionKey: "session-finished",
2116+
phase: "start",
2117+
hasActiveRun: true,
2118+
activeRunIds: ["run-finished"],
2119+
});
21072120
expectPayloadFields(sessionsChangedCalls[1]?.[1], {
21082121
sessionKey: "session-finished",
21092122
phase: "end",
2123+
hasActiveRun: false,
2124+
activeRunIds: [],
21102125
status: "done",
21112126
startedAt: 900,
21122127
endedAt: 1_700,
21132128
runtimeMs: 800,
21142129
updatedAt: 1_700,
21152130
abortedLastRun: false,
21162131
});
2132+
expect(resolveSessionActiveRunState).toHaveBeenCalledWith({
2133+
requestedKey: "session-finished",
2134+
canonicalKey: "session-finished",
2135+
});
21172136
const persistParams = requireRecord(
21182137
persistGatewaySessionLifecycleEventMock.mock.calls
21192138
.map((call) => call[0])

‎src/gateway/server-chat.ts‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,12 @@ export type AgentEventHandlerOptions = {
299299
clientRunId: string;
300300
summary: string | undefined;
301301
}) => void;
302+
resolveSessionActiveRunState?: (params: {
303+
requestedKey: string;
304+
canonicalKey: string;
305+
sessionId?: string;
306+
agentId?: string;
307+
}) => { active: boolean; runIds: string[] };
302308
};
303309

304310
function roundedChatSendTimingMs(value: number): number {
@@ -324,6 +330,7 @@ export function createAgentEventHandler({
324330
trackTrackedRunTerminalPersistence,
325331
resolveActiveLifecycleGenerationForRun = () => undefined,
326332
updateRunToolErrorSummary,
333+
resolveSessionActiveRunState,
327334
}: AgentEventHandlerOptions) {
328335
type TerminalLifecycleOptions = {
329336
skipChatErrorFinal?: boolean;
@@ -414,6 +421,7 @@ export function createAgentEventHandler({
414421
sessionKey: string,
415422
evt?: AgentEventPayload,
416423
agentId?: string,
424+
includeActiveRunState = false,
417425
) => {
418426
const row = loadGatewaySessionRowForSnapshot(sessionKey, agentId ? { agentId } : undefined);
419427
const omitUnscopedGlobalGoal = sessionKey === "global" && !agentId;
@@ -437,7 +445,20 @@ export function createAgentEventHandler({
437445
event: evt,
438446
})
439447
: {};
440-
const session = row ? { ...row, ...lifecyclePatch } : undefined;
448+
const activeRunState = includeActiveRunState
449+
? resolveSessionActiveRunState?.({
450+
requestedKey: sessionKey,
451+
canonicalKey: row?.key ?? sessionKey,
452+
...(row?.sessionId ? { sessionId: row.sessionId } : {}),
453+
...(agentId ? { agentId } : {}),
454+
})
455+
: undefined;
456+
// Agent lifecycle broadcasts merge into cached session rows in the UI.
457+
// Always replace run identity so a newer start cannot inherit a completed run.
458+
const activeRunFields = activeRunState
459+
? { hasActiveRun: activeRunState.active, activeRunIds: activeRunState.runIds }
460+
: {};
461+
const session = row ? { ...row, ...lifecyclePatch, ...activeRunFields } : undefined;
441462
if (session && omitUnscopedGlobalGoal) {
442463
delete session.goal;
443464
}
@@ -490,6 +511,7 @@ export function createAgentEventHandler({
490511
effectiveResponseUsage: row?.effectiveResponseUsage,
491512
modelProvider: row?.modelProvider,
492513
model: row?.model,
514+
...activeRunFields,
493515
status: snapshotSource.status,
494516
startedAt: snapshotSource.startedAt,
495517
endedAt: snapshotSource.endedAt,
@@ -715,7 +737,7 @@ export function createAgentEventHandler({
715737
runId: evt.runId,
716738
...(eventRunId !== evt.runId ? { clientRunId: eventRunId } : {}),
717739
ts: evt.ts,
718-
...buildSessionEventSnapshot(sessionKey, snapshotEvent, sessionAgentId),
740+
...buildSessionEventSnapshot(sessionKey, snapshotEvent, sessionAgentId, true),
719741
},
720742
sessionEventConnIds,
721743
{ dropIfSlow: true },
@@ -1494,7 +1516,7 @@ export function createAgentEventHandler({
14941516
runId: evt.runId,
14951517
...(eventRunId !== evt.runId ? { clientRunId: eventRunId } : {}),
14961518
ts: evt.ts,
1497-
...buildSessionEventSnapshot(sessionKey, evt, sessionAgentId),
1519+
...buildSessionEventSnapshot(sessionKey, evt, sessionAgentId, true),
14981520
},
14991521
sessionEventConnIds,
15001522
{ dropIfSlow: true },

‎src/gateway/server-methods/chat.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,10 @@ import {
227227
loadOptionalServerMethodModelCatalog,
228228
startOptionalServerMethodModelCatalogLoad,
229229
} from "./optional-model-catalog.js";
230-
import { hasTrackedActiveSessionRun, hasVisibleActiveSessionRun } from "./session-active-runs.js";
230+
import {
231+
hasTrackedActiveSessionRun,
232+
resolveVisibleActiveSessionRunState,
233+
} from "./session-active-runs.js";
231234
import { emitSessionsChanged } from "./session-change-event.js";
232235
import type {
233236
GatewayClient,
@@ -3207,14 +3210,16 @@ async function handleChatHistoryRequest({
32073210
});
32083211
const activeRunAgentId =
32093212
canonicalKey === "global" ? (selectedAgent.agentId ?? defaultAgentId) : selectedAgent.agentId;
3210-
sessionInfo.hasActiveRun = hasVisibleActiveSessionRun({
3213+
const activeRunState = resolveVisibleActiveSessionRunState({
32113214
context,
32123215
requestedKey: sessionKey,
32133216
canonicalKey,
32143217
sessionId: entry?.sessionId,
32153218
...(activeRunAgentId ? { agentId: activeRunAgentId } : {}),
32163219
defaultAgentId,
32173220
});
3221+
sessionInfo.hasActiveRun = activeRunState.active;
3222+
sessionInfo.activeRunIds = activeRunState.runIds;
32183223
const defaults = getSessionDefaults(cfg, modelCatalog, { allowPluginNormalization: false });
32193224
const thinkingLevel = sessionInfo.thinkingLevel ?? sessionInfo.thinkingDefault;
32203225
const verboseLevel = entry?.verboseLevel ?? cfg.agents?.defaults?.verboseDefault;

‎src/gateway/server-methods/session-active-runs.test.ts‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
// Tests gateway active-run matching by logical session key and backing id.
22
import { expect, it } from "vitest";
3-
import { hasVisibleActiveSessionRun } from "./session-active-runs.js";
3+
import {
4+
hasVisibleActiveSessionRun,
5+
resolveVisibleActiveSessionRunState,
6+
} from "./session-active-runs.js";
47

58
it("matches session-id-only gateway runs during archive admission", () => {
69
const context = {
@@ -25,3 +28,22 @@ it("matches session-id-only gateway runs during archive admission", () => {
2528
}),
2629
).toBe(true);
2730
});
31+
32+
it("returns deterministic visible run ids for the selected session", () => {
33+
const context = {
34+
chatAbortControllers: new Map([
35+
["run-z", { sessionKey: "main" }],
36+
["run-hidden", { sessionKey: "main", controlUiVisible: false }],
37+
["run-other", { sessionKey: "other" }],
38+
["run-a", { sessionKey: "main" }],
39+
]),
40+
} as never;
41+
42+
expect(
43+
resolveVisibleActiveSessionRunState({
44+
context,
45+
requestedKey: "main",
46+
canonicalKey: "main",
47+
}),
48+
).toEqual({ active: true, runIds: ["run-a", "run-z"] });
49+
});

‎src/gateway/server-methods/session-active-runs.ts‎

Lines changed: 38 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import type { GatewayRequestContext } from "./types.js";
1111
* do not make a session look busy to user-facing session operations.
1212
*/
1313
type TrackedActiveSessionRun = {
14+
runId: string;
1415
sessionKey?: string;
1516
sessionId?: string;
1617
agentId?: string;
@@ -23,14 +24,15 @@ function collectTrackedActiveSessionRuns(
2324
if (!(context.chatAbortControllers instanceof Map)) {
2425
return runs;
2526
}
26-
for (const active of context.chatAbortControllers.values()) {
27+
for (const [runId, active] of context.chatAbortControllers) {
2728
if (active.projectSessionActive !== false && active.controlUiVisible !== false) {
2829
const sessionKey = active.sessionKey?.trim();
2930
const sessionId = active.sessionId?.trim();
3031
if (!sessionKey && !sessionId) {
3132
continue;
3233
}
3334
runs.push({
35+
runId,
3436
...(sessionKey ? { sessionKey } : {}),
3537
...(sessionId ? { sessionId } : {}),
3638
agentId: typeof active.agentId === "string" ? normalizeAgentId(active.agentId) : undefined,
@@ -88,6 +90,40 @@ export function hasTrackedActiveSessionRun(params: {
8890
);
8991
}
9092

93+
export function resolveVisibleActiveSessionRunState(params: {
94+
context: Partial<Pick<GatewayRequestContext, "chatAbortControllers">>;
95+
requestedKey: string;
96+
canonicalKey: string;
97+
sessionId?: string;
98+
agentId?: string;
99+
defaultAgentId?: string;
100+
}): { active: boolean; runIds: string[] } {
101+
const sessionId = params.sessionId?.trim();
102+
const runIds = collectTrackedActiveSessionRuns(params.context)
103+
.filter(
104+
(active) =>
105+
isTrackedActiveSessionRunForKey(
106+
active,
107+
params.canonicalKey,
108+
params.agentId,
109+
params.defaultAgentId,
110+
) ||
111+
isTrackedActiveSessionRunForKey(
112+
active,
113+
params.requestedKey,
114+
params.agentId,
115+
params.defaultAgentId,
116+
) ||
117+
(sessionId !== undefined && active.sessionId === sessionId),
118+
)
119+
.map((active) => active.runId)
120+
.toSorted();
121+
return {
122+
active: runIds.length > 0 || (sessionId !== undefined && isEmbeddedAgentRunActive(sessionId)),
123+
runIds,
124+
};
125+
}
126+
91127
export function hasVisibleActiveSessionRun(params: {
92128
context: Partial<Pick<GatewayRequestContext, "chatAbortControllers">>;
93129
requestedKey: string;
@@ -96,17 +132,5 @@ export function hasVisibleActiveSessionRun(params: {
96132
agentId?: string;
97133
defaultAgentId?: string;
98134
}): boolean {
99-
if (hasTrackedActiveSessionRun(params)) {
100-
return true;
101-
}
102-
const sessionId = params.sessionId?.trim();
103-
if (!sessionId) {
104-
return false;
105-
}
106-
if (
107-
collectTrackedActiveSessionRuns(params.context).some((active) => active.sessionId === sessionId)
108-
) {
109-
return true;
110-
}
111-
return isEmbeddedAgentRunActive(sessionId);
135+
return resolveVisibleActiveSessionRunState(params).active;
112136
}

‎src/gateway/server-methods/session-change-event.ts‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
import { resolveDefaultAgentId } from "../../agents/agent-scope.js";
33
import { buildGatewaySessionEventFields } from "../session-event-payload.js";
44
import { loadGatewaySessionRow } from "../session-utils.js";
5-
import { hasVisibleActiveSessionRun } from "./session-active-runs.js";
5+
import { resolveVisibleActiveSessionRunState } from "./session-active-runs.js";
66
import type { GatewayRequestContext } from "./types.js";
77

88
export type SessionChangedPayload = {
@@ -35,6 +35,16 @@ export function emitSessionsChanged(
3535
)
3636
: null;
3737
const defaultAgentId = resolveDefaultAgentId(context.getRuntimeConfig());
38+
const activeRunState = sessionRow
39+
? resolveVisibleActiveSessionRunState({
40+
context,
41+
requestedKey: payload.sessionKey ?? sessionRow.key,
42+
canonicalKey: sessionRow.key,
43+
sessionId: sessionRow.sessionId,
44+
agentId: sessionRow.key === "global" ? payload.agentId : undefined,
45+
defaultAgentId,
46+
})
47+
: null;
3848
context.broadcastToConnIds(
3949
"sessions.changed",
4050
{
@@ -45,14 +55,8 @@ export function emitSessionsChanged(
4555
...buildGatewaySessionEventFields({
4656
sessionRow,
4757
agentId: payload.agentId,
48-
hasActiveRun: hasVisibleActiveSessionRun({
49-
context,
50-
requestedKey: payload.sessionKey ?? sessionRow.key,
51-
canonicalKey: sessionRow.key,
52-
sessionId: sessionRow.sessionId,
53-
agentId: sessionRow.key === "global" ? payload.agentId : undefined,
54-
defaultAgentId,
55-
}),
58+
hasActiveRun: activeRunState?.active,
59+
activeRunIds: activeRunState?.runIds,
5660
}),
5761
effectiveFastMode: sessionRow.effectiveFastMode,
5862
effectiveFastModeSource: sessionRow.effectiveFastModeSource,

‎src/gateway/server-methods/sessions.ts‎

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,11 @@ import { resolveSessionKeyFromResolveParams } from "../sessions-resolve.js";
124124
import { setGatewayDedupeEntry } from "./agent-wait-dedupe.js";
125125
import { chatHandlers } from "./chat.js";
126126
import { loadOptionalServerMethodModelCatalog } from "./optional-model-catalog.js";
127-
import { hasTrackedActiveSessionRun, hasVisibleActiveSessionRun } from "./session-active-runs.js";
127+
import {
128+
hasTrackedActiveSessionRun,
129+
hasVisibleActiveSessionRun,
130+
resolveVisibleActiveSessionRunState,
131+
} from "./session-active-runs.js";
128132
import { emitSessionsChanged } from "./session-change-event.js";
129133
import type {
130134
GatewayClient,
@@ -815,18 +819,22 @@ export const sessionsHandlers: GatewayRequestHandlers = {
815819
const sessions = measureDiagnosticsTimelineSpanSync(
816820
"gateway.sessions.list.active_run_flags",
817821
() => {
818-
return result.sessions.map((session) =>
819-
Object.assign({}, session, {
820-
hasActiveRun: hasVisibleActiveSessionRun({
821-
context,
822-
requestedKey: session.key,
823-
canonicalKey: session.key,
824-
sessionId: session.sessionId,
825-
...(session.key === "global" && p.agentId ? { agentId: p.agentId } : {}),
826-
defaultAgentId: resolveDefaultAgentId(cfg),
827-
}),
828-
}),
829-
);
822+
return result.sessions.map((session) => {
823+
const activeRunState = resolveVisibleActiveSessionRunState({
824+
context,
825+
requestedKey: session.key,
826+
canonicalKey: session.key,
827+
sessionId: session.sessionId,
828+
...(session.key === "global" && p.agentId ? { agentId: p.agentId } : {}),
829+
defaultAgentId: resolveDefaultAgentId(cfg),
830+
});
831+
return Object.assign({}, session, {
832+
hasActiveRun: activeRunState.active,
833+
...(activeRunState.runIds.length > 0
834+
? { activeRunIds: activeRunState.runIds }
835+
: {}),
836+
});
837+
});
830838
},
831839
{
832840
config: cfg,

0 commit comments

Comments
 (0)