Skip to content

Commit 5fb6c2a

Browse files
mbenhamdmbenhamd
authored
feat(core): add coalesced harness display-state subscriptions (#15974)
Adds Harness.subscribeDisplayState() for UI, SSE, TUI, and bridge consumers that render from HarnessDisplayState. Raw subscribe() still delivers every event for logs, debugging, analytics, and replay. Example usage: ```ts render(harness.getDisplayState()); const unsubscribe = harness.subscribeDisplayState(render, { windowMs: 250, maxWaitMs: 500, }); ``` The scheduler coalesces high-frequency message and tool updates, emits fresh display-state snapshots, and flushes critical lifecycle updates immediately. Fixes #15973. Validated with `pnpm --filter ./packages/core test -- display-state`, `pnpm --filter ./packages/core check`, and `pnpm build:core`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## ELI5 This PR adds a smarter way for UI components to listen to harness updates. Instead of getting bombarded with every single internal event, they can now use `subscribeDisplayState()` to receive batched updates that combine rapid changes into fewer snapshots, while keeping the old raw `subscribe()` method unchanged for debugging and logging. ## Overview This PR introduces `Harness.subscribeDisplayState()`, a new UI-facing subscription API that delivers coalesced `HarnessDisplayState` snapshots. It implements intelligent batching of high-frequency internal events (messages, tool updates) into consolidated display-state snapshots, while preserving the raw event semantics of the existing `subscribe()` API. Critical lifecycle events (agent start/end, errors, approvals, suspensions, questions, plans) bypass coalescing and flush immediately. ## Key Changes ### New Display State Scheduler Module - **`display-state-scheduler.ts`**: Introduces `DisplayStateScheduler` class that buffers non-critical display state updates and dispatches them using two configurable timers: - `windowMs` (default 250ms): throttle window that resets on each non-critical update - `maxWaitMs` (default 500ms): maximum wait time to prevent starvation - Critical updates bypass buffering and flush immediately - Deep-clones entire `HarnessDisplayState` before invoking listeners to ensure isolation - Catches and logs listener errors via `console.error` without disrupting other listeners - Exports constants: `DEFAULT_DISPLAY_STATE_SUBSCRIPTION_OPTIONS` and `CRITICAL_DISPLAY_STATE_EVENT_TYPES` ### API Additions to Harness - **`getDisplayState()`**: Returns a snapshot of current `HarnessDisplayState` for initial rendering - **`subscribeDisplayState(listener, options?)`**: Registers a coalesced display-state listener that returns an unsubscribe function - Accepts optional `windowMs` and `maxWaitMs` tuning parameters - Listener is not invoked immediately (caller must use `getDisplayState()` for initial render) - All subscriptions are cleaned up during `harness.destroy()` ### Integration Updates - Modified `Harness.emit()` to notify schedulers after dispatching events, passing `isCritical` flag based on event type - Event type `display_state_changed` is skipped to avoid double-notification - Harness destroy method extended to dispose and clear all registered schedulers ### Type System Enhancements - **New types in `types.ts`**: - `HarnessDisplayStateListener`: callback type receiving `HarnessDisplayState`, may be sync or async - `HarnessDisplayStateSubscriptionOptions`: configuration interface for `windowMs` and `maxWaitMs` - **Index re-exports**: Both new types exported from harness module ### Documentation - Updated harness-class.mdx with comprehensive documentation of both new APIs - Clarified that `subscribe()` is intended for raw event logs while `subscribeDisplayState()` is preferred for UI rendering - Added changeset documenting the new API with usage examples ## Testing Comprehensive test suite in `display-state.test.ts` covering: - Coalescing behavior driven by `windowMs` with fake timers - Forced flushes at `maxWaitMs` to prevent starvation - Critical event types trigger immediate (non-coalesced) emission - Emitted snapshots are fresh deep-clones isolated from harness internal state - Unsubscribe and `harness.destroy()` properly cancel pending timer flushes - Custom coalescing options change flush timing as configured - Raw subscriptions continue to receive every event while display subscribers remain coalesced - Listener errors are caught and logged without preventing other listeners - Validation of key snapshot fields like `isRunning`, `pendingApproval`, `pendingSuspension`, etc. ## Validation Changes have been validated with: - `pnpm --filter ./packages/core test -- display-state` - `pnpm --filter ./packages/core check` - `pnpm build:core` Fixes issue #15973. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent 2c83efc commit 5fb6c2a

7 files changed

Lines changed: 602 additions & 1 deletion

File tree

‎.changeset/five-hands-punch.md‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
---
2+
'@mastra/core': patch
3+
---
4+
5+
Added a coalesced display state subscription API for Harness.
6+
7+
This helps UI clients render fewer updates while still receiving the latest state. The example below renders the initial state, then subscribes to coalesced updates with the default `windowMs` and `maxWaitMs` timing options.
8+
9+
```ts
10+
render(harness.getDisplayState());
11+
12+
const unsubscribe = harness.subscribeDisplayState(render, {
13+
windowMs: 250,
14+
maxWaitMs: 500,
15+
});
16+
```

‎docs/src/content/en/reference/harness/harness-class.mdx‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,14 @@ Return a read-only snapshot of the current harness state.
347347
const state = harness.getState()
348348
```
349349

350+
#### `getDisplayState()`
351+
352+
Return the current `HarnessDisplayState` snapshot for UI rendering.
353+
354+
```typescript
355+
const displayState = harness.getDisplayState()
356+
```
357+
350358
#### `setState(updates)`
351359

352360
Update the harness state. Validates against `stateSchema` if provided, and emits a `state_changed` event with the new state and changed keys.
@@ -917,10 +925,35 @@ Forked mode trades isolation for context inheritance. If the subagent should run
917925

918926
### Events
919927

928+
#### `subscribeDisplayState(listener, options?)`
929+
930+
Register a listener for coalesced display state snapshots. Use this method for UI, Server-Sent Events (SSE), terminal UI (TUI), and bridge rendering paths that only need the latest `HarnessDisplayState`. Use [`subscribe`](#subscribelistener) when you need the raw event log.
931+
932+
```typescript
933+
const unsubscribe = harness.subscribeDisplayState(displayState => {
934+
render(displayState)
935+
})
936+
937+
// Optional tuning:
938+
harness.subscribeDisplayState(render, {
939+
windowMs: 250,
940+
maxWaitMs: 500,
941+
})
942+
943+
// Later:
944+
unsubscribe()
945+
```
946+
947+
Returns: `() => void`
948+
949+
`subscribeDisplayState()` does not call the listener immediately. Call [`getDisplayState`](#getdisplaystate) first if the UI needs an initial render before the next harness event.
950+
920951
#### `subscribe(listener)`
921952

922953
Register an event listener. Returns an unsubscribe function.
923954

955+
Use this method for audit logs, debugging, analytics, deterministic replay, or consumers that need every raw event. For display rendering, prefer [`subscribeDisplayState`](#subscribedisplaystatelistener-options) so high-frequency events such as `message_update`, `tool_update`, and `tool_input_delta` are coalesced into the latest display state snapshot.
956+
924957
```typescript
925958
const unsubscribe = harness.subscribe(event => {
926959
switch (event.type) {

‎packages/core/src/harness/display-state-scheduler.ts‎

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
import type { HarnessDisplayState, HarnessDisplayStateListener, HarnessEvent } from './types';
2+
3+
export const DEFAULT_DISPLAY_STATE_SUBSCRIPTION_OPTIONS = {
4+
windowMs: 250,
5+
maxWaitMs: 500,
6+
} as const;
7+
8+
export const CRITICAL_DISPLAY_STATE_EVENT_TYPES: ReadonlySet<HarnessEvent['type']> = new Set([
9+
'agent_start',
10+
'agent_end',
11+
'error',
12+
'tool_approval_required',
13+
'tool_suspended',
14+
'ask_question',
15+
'plan_approval_required',
16+
'plan_approved',
17+
'thread_changed',
18+
'thread_created',
19+
'thread_deleted',
20+
'mode_changed',
21+
'model_changed',
22+
'subagent_model_changed',
23+
'state_changed',
24+
'tool_input_end',
25+
'tool_end',
26+
'subagent_end',
27+
]);
28+
29+
function cloneValue(value: unknown, seen = new WeakMap<object, unknown>()): unknown {
30+
if (value === null || typeof value !== 'object') {
31+
return value;
32+
}
33+
34+
if (value instanceof Date) {
35+
return new Date(value.getTime());
36+
}
37+
38+
if (seen.has(value)) {
39+
return seen.get(value);
40+
}
41+
42+
if (Array.isArray(value)) {
43+
const cloned: unknown[] = [];
44+
seen.set(value, cloned);
45+
for (const item of value) {
46+
cloned.push(cloneValue(item, seen));
47+
}
48+
return cloned;
49+
}
50+
51+
if (value instanceof Map) {
52+
const cloned = new Map<unknown, unknown>();
53+
seen.set(value, cloned);
54+
for (const [key, mapValue] of value) {
55+
cloned.set(cloneValue(key, seen), cloneValue(mapValue, seen));
56+
}
57+
return cloned;
58+
}
59+
60+
if (value instanceof Set) {
61+
const cloned = new Set<unknown>();
62+
seen.set(value, cloned);
63+
for (const item of value) {
64+
cloned.add(cloneValue(item, seen));
65+
}
66+
return cloned;
67+
}
68+
69+
const cloned: Record<PropertyKey, unknown> = {};
70+
seen.set(value, cloned);
71+
for (const key of Reflect.ownKeys(value)) {
72+
cloned[key] = cloneValue((value as Record<PropertyKey, unknown>)[key], seen);
73+
}
74+
return cloned;
75+
}
76+
77+
function cloneUnknown<T>(value: T): T {
78+
return cloneValue(value) as T;
79+
}
80+
81+
function cloneDisplayState(state: HarnessDisplayState): HarnessDisplayState {
82+
return {
83+
...state,
84+
currentMessage: state.currentMessage
85+
? {
86+
...state.currentMessage,
87+
createdAt: new Date(state.currentMessage.createdAt.getTime()),
88+
content: state.currentMessage.content.map(part => cloneUnknown(part)),
89+
}
90+
: null,
91+
tokenUsage: { ...state.tokenUsage },
92+
activeTools: new Map(
93+
Array.from(state.activeTools, ([id, tool]) => [
94+
id,
95+
{
96+
...tool,
97+
args: cloneUnknown(tool.args),
98+
result: cloneUnknown(tool.result),
99+
},
100+
]),
101+
),
102+
toolInputBuffers: new Map(Array.from(state.toolInputBuffers, ([id, buffer]) => [id, { ...buffer }])),
103+
pendingApproval: state.pendingApproval
104+
? { ...state.pendingApproval, args: cloneUnknown(state.pendingApproval.args) }
105+
: null,
106+
pendingSuspension: state.pendingSuspension
107+
? {
108+
...state.pendingSuspension,
109+
args: cloneUnknown(state.pendingSuspension.args),
110+
suspendPayload: cloneUnknown(state.pendingSuspension.suspendPayload),
111+
}
112+
: null,
113+
pendingQuestion: state.pendingQuestion
114+
? {
115+
...state.pendingQuestion,
116+
options: state.pendingQuestion.options?.map(option => cloneUnknown(option)),
117+
}
118+
: null,
119+
pendingPlanApproval: state.pendingPlanApproval ? { ...state.pendingPlanApproval } : null,
120+
activeSubagents: new Map(
121+
Array.from(state.activeSubagents, ([id, subagent]) => [
122+
id,
123+
{
124+
...subagent,
125+
toolCalls: subagent.toolCalls.map(toolCall => cloneUnknown(toolCall)),
126+
},
127+
]),
128+
),
129+
omProgress: {
130+
...state.omProgress,
131+
buffered: {
132+
observations: { ...state.omProgress.buffered.observations },
133+
reflection: { ...state.omProgress.buffered.reflection },
134+
},
135+
},
136+
modifiedFiles: new Map(
137+
Array.from(state.modifiedFiles, ([path, modifiedFile]) => [
138+
path,
139+
{
140+
...modifiedFile,
141+
firstModified: new Date(modifiedFile.firstModified.getTime()),
142+
operations: [...modifiedFile.operations],
143+
},
144+
]),
145+
),
146+
tasks: state.tasks.map(task => cloneUnknown(task)),
147+
previousTasks: state.previousTasks.map(task => cloneUnknown(task)),
148+
};
149+
}
150+
151+
export class DisplayStateScheduler {
152+
private disposed = false;
153+
private pendingState: HarnessDisplayState | null = null;
154+
private windowTimer: ReturnType<typeof setTimeout> | null = null;
155+
private maxWaitTimer: ReturnType<typeof setTimeout> | null = null;
156+
157+
constructor(
158+
private readonly listener: HarnessDisplayStateListener,
159+
private readonly windowMs: number,
160+
private readonly maxWaitMs: number,
161+
) {}
162+
163+
notify(state: HarnessDisplayState, isCritical: boolean): void {
164+
if (this.disposed) return;
165+
166+
if (isCritical) {
167+
this.flush(state);
168+
return;
169+
}
170+
171+
this.pendingState = state;
172+
173+
if (this.windowTimer) {
174+
clearTimeout(this.windowTimer);
175+
}
176+
this.windowTimer = setTimeout(() => this.flushPending(), this.windowMs);
177+
178+
if (!this.maxWaitTimer) {
179+
this.maxWaitTimer = setTimeout(() => this.flushPending(), this.maxWaitMs);
180+
}
181+
}
182+
183+
dispose(): void {
184+
this.disposed = true;
185+
this.pendingState = null;
186+
this.clearTimers();
187+
}
188+
189+
private flushPending(): void {
190+
if (!this.pendingState) return;
191+
this.flush(this.pendingState);
192+
}
193+
194+
private flush(state: HarnessDisplayState): void {
195+
if (this.disposed) return;
196+
197+
this.pendingState = null;
198+
this.clearTimers();
199+
200+
try {
201+
const result = this.listener(cloneDisplayState(state));
202+
if (result && typeof result === 'object' && 'catch' in result && typeof result.catch === 'function') {
203+
(result as Promise<void>).catch(err => console.error('Error in harness display state listener:', err));
204+
}
205+
} catch (err) {
206+
console.error('Error in harness display state listener:', err);
207+
}
208+
}
209+
210+
private clearTimers(): void {
211+
if (this.windowTimer) {
212+
clearTimeout(this.windowTimer);
213+
}
214+
if (this.maxWaitTimer) {
215+
clearTimeout(this.maxWaitTimer);
216+
}
217+
this.windowTimer = null;
218+
this.maxWaitTimer = null;
219+
}
220+
}

0 commit comments

Comments
 (0)