Skip to content

Commit c9050c9

Browse files
vincentkocvincentkoc
authored
perf(control-ui): trace first assistant event
Add Control UI first-assistant chat send timing instrumentation and tests.
1 parent 0933726 commit c9050c9

7 files changed

Lines changed: 377 additions & 9 deletions

File tree

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,19 @@
22
// Tracks active runs, delta buffers, tool recipients, and session subscribers.
33
import type { AgentEventPayload } from "../infra/agent-events.js";
44

5+
export type ChatRunTiming = {
6+
ackedAtMs: number;
7+
connId: string;
8+
dispatchStartedAtMs?: number;
9+
firstAssistantEventSent?: boolean;
10+
receivedAtMs: number;
11+
};
12+
513
export type ChatRunEntry = {
614
sessionKey: string;
715
agentId?: string;
816
clientRunId: string;
17+
chatSendTiming?: ChatRunTiming;
918
};
1019

1120
export type BufferedAgentEvent = {

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

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -413,6 +413,89 @@ describe("agent event handler", () => {
413413
nowSpy?.mockRestore();
414414
});
415415

416+
it("emits the first assistant chat.send timing event to the originating Control UI", () => {
417+
const { broadcastToConnIds, chatRunState, handler, nowSpy } = createHarness({ now: 1_000 });
418+
chatRunState.registry.add("run-1", {
419+
sessionKey: "session-1",
420+
clientRunId: "client-1",
421+
chatSendTiming: {
422+
ackedAtMs: 0,
423+
connId: "conn-control-ui",
424+
dispatchStartedAtMs: 0,
425+
receivedAtMs: 0,
426+
},
427+
});
428+
429+
handler({
430+
runId: "run-1",
431+
seq: 1,
432+
stream: "assistant",
433+
ts: Date.now(),
434+
data: { text: "Hello world" },
435+
});
436+
handler({
437+
runId: "run-1",
438+
seq: 2,
439+
stream: "assistant",
440+
ts: Date.now(),
441+
data: { text: "Hello world again" },
442+
});
443+
444+
expect(broadcastToConnIds).toHaveBeenCalledTimes(1);
445+
expect(broadcastToConnIds).toHaveBeenCalledWith(
446+
"chat.send_timing",
447+
expect.objectContaining({
448+
phase: "first-assistant-event",
449+
runId: "client-1",
450+
sessionKey: "session-1",
451+
ackToPhaseMs: expect.any(Number),
452+
dispatchStartedToPhaseMs: expect.any(Number),
453+
receivedToPhaseMs: expect.any(Number),
454+
}),
455+
new Set(["conn-control-ui"]),
456+
{ dropIfSlow: true },
457+
);
458+
nowSpy?.mockRestore();
459+
});
460+
461+
it("emits first assistant chat.send timing when text first flushes on final", () => {
462+
const { broadcastToConnIds, chatRunState, handler, nowSpy } = createHarness({ now: 1_000 });
463+
chatRunState.registry.add("run-final-only", {
464+
sessionKey: "session-1",
465+
clientRunId: "client-final",
466+
chatSendTiming: {
467+
ackedAtMs: 0,
468+
connId: "conn-control-ui",
469+
dispatchStartedAtMs: 0,
470+
receivedAtMs: 0,
471+
},
472+
});
473+
chatRunState.buffers.set("client-final", "Final only reply");
474+
475+
handler({
476+
runId: "run-final-only",
477+
seq: 1,
478+
stream: "lifecycle",
479+
ts: Date.now(),
480+
data: { phase: "end" },
481+
});
482+
483+
expect(broadcastToConnIds).toHaveBeenCalledWith(
484+
"chat.send_timing",
485+
expect.objectContaining({
486+
phase: "first-assistant-event",
487+
runId: "client-final",
488+
sessionKey: "session-1",
489+
ackToPhaseMs: expect.any(Number),
490+
dispatchStartedToPhaseMs: expect.any(Number),
491+
receivedToPhaseMs: expect.any(Number),
492+
}),
493+
new Set(["conn-control-ui"]),
494+
{ dropIfSlow: true },
495+
);
496+
nowSpy?.mockRestore();
497+
});
498+
416499
it("coalesces assistant agent events under the chat delta throttle", () => {
417500
let now = 10_000;
418501
const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now);

‎src/gateway/server-chat.ts‎

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// Gateway chat runtime projects agent events into chat/session subscriber
22
// streams, lifecycle persistence, heartbeat visibility, and live UI updates.
3+
import { performance } from "node:perf_hooks";
34
import { resolveDefaultAgentId } from "../agents/agent-scope.js";
45
import { resolveToolSearchCodeDisplayTarget } from "../agents/tool-display-common.js";
56
import { DEFAULT_HEARTBEAT_ACK_MAX_CHARS, stripHeartbeatToken } from "../auto-reply/heartbeat.js";
@@ -18,6 +19,7 @@ import {
1819
} from "./live-chat-projector.js";
1920
import type {
2021
BufferedAgentEvent,
22+
ChatRunEntry,
2123
ChatRunState,
2224
SessionEventSubscriberRegistry,
2325
SessionMessageSubscriberRegistry,
@@ -257,6 +259,10 @@ export type AgentEventHandlerOptions = {
257259
}) => void;
258260
};
259261

262+
function roundedChatSendTimingMs(value: number): number {
263+
return Math.max(0, Math.round(value * 1000) / 1000);
264+
}
265+
260266
export function createAgentEventHandler({
261267
broadcast,
262268
broadcastToConnIds,
@@ -446,6 +452,35 @@ export function createAgentEventHandler({
446452
}
447453
};
448454

455+
const emitFirstAssistantChatSendTiming = (chatLink: ChatRunEntry | undefined) => {
456+
const timing = chatLink?.chatSendTiming;
457+
if (!timing || timing.firstAssistantEventSent) {
458+
return;
459+
}
460+
timing.firstAssistantEventSent = true;
461+
const nowMs = performance.now();
462+
broadcastToConnIds(
463+
"chat.send_timing",
464+
{
465+
phase: "first-assistant-event",
466+
runId: chatLink.clientRunId,
467+
sessionKey: chatLink.sessionKey,
468+
...(chatLink.agentId ? { agentId: chatLink.agentId } : {}),
469+
ackToPhaseMs: roundedChatSendTimingMs(nowMs - timing.ackedAtMs),
470+
receivedToPhaseMs: roundedChatSendTimingMs(nowMs - timing.receivedAtMs),
471+
...(timing.dispatchStartedAtMs !== undefined
472+
? {
473+
dispatchStartedToPhaseMs: roundedChatSendTimingMs(
474+
nowMs - timing.dispatchStartedAtMs,
475+
),
476+
}
477+
: {}),
478+
},
479+
new Set([timing.connId]),
480+
{ dropIfSlow: true },
481+
);
482+
};
483+
449484
const finalizeLifecycleEvent = (evt: AgentEventPayload, opts?: TerminalLifecycleOptions) => {
450485
const lifecyclePhase =
451486
evt.stream === "lifecycle" && typeof evt.data?.phase === "string" ? evt.data.phase : null;
@@ -496,7 +531,11 @@ export function createAgentEventHandler({
496531
evt.data?.error,
497532
evtStopReason,
498533
evtErrorKind,
499-
{ agentId: finished.agentId, controlUiVisible: isControlUiVisible },
534+
{
535+
agentId: finished.agentId,
536+
controlUiVisible: isControlUiVisible,
537+
firstAssistantTimingEntry: finished,
538+
},
500539
);
501540
}
502541
} else if (!(opts?.skipChatErrorFinal && lifecyclePhase === "error")) {
@@ -632,6 +671,7 @@ export function createAgentEventHandler({
632671
timestamp: now,
633672
},
634673
};
674+
emitFirstAssistantChatSendTiming(chatRunState.registry.peek(sourceRunId));
635675
sendChatPayload(sessionKey, payload, {
636676
agentId,
637677
controlUiVisible: opts?.controlUiVisible ?? true,
@@ -667,7 +707,7 @@ export function createAgentEventHandler({
667707
clientRunId: string,
668708
sourceRunId: string,
669709
seq: number,
670-
opts?: { controlUiVisible?: boolean },
710+
opts?: { controlUiVisible?: boolean; firstAssistantTimingEntry?: ChatRunEntry },
671711
) => {
672712
const { text, shouldSuppressSilent } = resolveBufferedChatTextState(clientRunId, sourceRunId, {
673713
suppressLeadFragments: true,
@@ -704,6 +744,9 @@ export function createAgentEventHandler({
704744
timestamp: now,
705745
},
706746
};
747+
emitFirstAssistantChatSendTiming(
748+
opts?.firstAssistantTimingEntry ?? chatRunState.registry.peek(sourceRunId),
749+
);
707750
sendChatPayload(sessionKey, flushPayload, {
708751
agentId,
709752
controlUiVisible: opts?.controlUiVisible ?? true,
@@ -740,7 +783,11 @@ export function createAgentEventHandler({
740783
error?: unknown,
741784
stopReason?: string,
742785
errorKind?: ErrorKind,
743-
opts?: { agentId?: string; controlUiVisible?: boolean },
786+
opts?: {
787+
agentId?: string;
788+
controlUiVisible?: boolean;
789+
firstAssistantTimingEntry?: ChatRunEntry;
790+
},
744791
) => {
745792
const { text, shouldSuppressSilent } = resolveBufferedChatTextState(clientRunId, sourceRunId, {
746793
suppressLeadFragments: false,

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

Lines changed: 57 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ import {
138138
createManagedOutgoingImageBlocks,
139139
} from "../managed-image-attachments.js";
140140
import { ADMIN_SCOPE } from "../method-scopes.js";
141+
import type { ChatRunTiming } from "../server-chat-state.js";
141142
import { getMaxChatHistoryMessagesBytes, MAX_PAYLOAD_BYTES } from "../server-constants.js";
142143
import { resolveSessionHistoryTailReadOptions } from "../session-history-state.js";
143144
import { readSessionTranscriptIndex } from "../session-transcript-index.fs.js";
@@ -237,6 +238,7 @@ type ChatSendServerTimingPhase =
237238
| "dispatch-started"
238239
| "model-selected"
239240
| "agent-run-started"
241+
| "first-assistant-event"
240242
| "dispatch-completed"
241243
| "post-dispatch-completed";
242244

@@ -969,6 +971,26 @@ function hasAssistantDisplayMediaContent(
969971
return Boolean(content?.some((block) => block?.type !== "text"));
970972
}
971973

974+
function hasVisibleAssistantFinalMessage(message: Record<string, unknown> | undefined): boolean {
975+
if (!message) {
976+
return false;
977+
}
978+
if (typeof message.text === "string" && message.text.trim()) {
979+
return true;
980+
}
981+
const content = Array.isArray(message.content) ? message.content : [];
982+
return content.some((block) => {
983+
if (!block || typeof block !== "object") {
984+
return false;
985+
}
986+
const record = block as Record<string, unknown>;
987+
if (record.type === "text") {
988+
return typeof record.text === "string" && record.text.trim().length > 0;
989+
}
990+
return true;
991+
});
992+
}
993+
972994
function hasManagedOutgoingAssistantContent(
973995
content: readonly AssistantDisplayContentBlock[] | undefined,
974996
): boolean {
@@ -3352,18 +3374,27 @@ export const chatHandlers: GatewayRequestHandlers = {
33523374
}
33533375

33543376
try {
3355-
context.addChatRun(clientRunId, {
3356-
sessionKey,
3357-
agentId: selectedAgent.agentId,
3358-
clientRunId,
3359-
});
33603377
const serverTiming = shouldIncludeChatSendAckServerTiming(clientInfo)
33613378
? {
33623379
receivedToAckMs: roundedChatSendTimingMs(performance.now() - chatSendReceivedAtMs),
33633380
loadSessionMs: sessionLoadMs,
33643381
...(prepareAttachmentsMs !== undefined ? { prepareAttachmentsMs } : {}),
33653382
}
33663383
: undefined;
3384+
const chatSendTiming: ChatRunTiming | undefined =
3385+
serverTiming && typeof client?.connId === "string" && client.connId.trim()
3386+
? {
3387+
ackedAtMs: performance.now(),
3388+
connId: client.connId.trim(),
3389+
receivedAtMs: chatSendReceivedAtMs,
3390+
}
3391+
: undefined;
3392+
context.addChatRun(clientRunId, {
3393+
sessionKey,
3394+
agentId: selectedAgent.agentId,
3395+
clientRunId,
3396+
...(chatSendTiming ? { chatSendTiming } : {}),
3397+
});
33673398
const ackPayload = {
33683399
runId: clientRunId,
33693400
status: "started" as const,
@@ -3383,7 +3414,7 @@ export const chatHandlers: GatewayRequestHandlers = {
33833414
{ config: cfg },
33843415
);
33853416
respond(true, ackPayload, undefined, { runId: clientRunId });
3386-
const chatSendAckedAtMs = performance.now();
3417+
const chatSendAckedAtMs = chatSendTiming?.ackedAtMs ?? performance.now();
33873418
const persistedImagesPromise = persistChatSendImages({
33883419
images: parsedImages,
33893420
imageOrder,
@@ -3707,7 +3738,21 @@ export const chatHandlers: GatewayRequestHandlers = {
37073738
});
37083739
};
37093740
const dispatchStartedAtMs = performance.now();
3741+
if (chatSendTiming) {
3742+
chatSendTiming.dispatchStartedAtMs = dispatchStartedAtMs;
3743+
}
37103744
emitServerTiming("dispatch-started");
3745+
let firstAssistantServerTimingEmitted = false;
3746+
const emitFirstAssistantServerTiming = () => {
3747+
if (firstAssistantServerTimingEmitted || chatSendTiming?.firstAssistantEventSent) {
3748+
return;
3749+
}
3750+
firstAssistantServerTimingEmitted = true;
3751+
if (chatSendTiming) {
3752+
chatSendTiming.firstAssistantEventSent = true;
3753+
}
3754+
emitServerTiming("first-assistant-event", undefined, dispatchStartedAtMs);
3755+
};
37113756
void measureDiagnosticsTimelineSpan(
37123757
"gateway.chat_send.dispatch_inbound",
37133758
async () => {
@@ -4038,6 +4083,9 @@ export const chatHandlers: GatewayRequestHandlers = {
40384083
};
40394084
}
40404085
}
4086+
if (hasVisibleAssistantFinalMessage(message)) {
4087+
emitFirstAssistantServerTiming();
4088+
}
40414089
broadcastChatFinal({
40424090
context,
40434091
runId: clientRunId,
@@ -4360,6 +4408,9 @@ export const chatHandlers: GatewayRequestHandlers = {
43604408
stopReason: "stop",
43614409
usage: { input: 0, output: 0, totalTokens: 0 },
43624410
};
4411+
if (hasVisibleAssistantFinalMessage(message)) {
4412+
emitFirstAssistantServerTiming();
4413+
}
43634414
broadcastChatFinal({
43644415
context,
43654416
runId: clientRunId,

0 commit comments

Comments
 (0)