Skip to content

Commit b67371a

Browse files
MarioCadenasMarioCadenas
committed
refactor(playground): migrate useAgentStream to wrap useAgentChat
Reviewer agentic finding #21: dev-playground reimplemented the same SSE parsing the template did, creating two divergent implementations of one concern. The new `useAgentChat` hook in @databricks/appkit-ui already owns the fetch + frame parsing + state plumbing; the playground hook is now a thin wrapper that adds the two playground- specific concerns and re-exports the legacy API surface. The wrapper keeps two layers on top of the shared hook: 1. `contextPrefix` on send() — the Smart Dashboard injects active filter / highlight state into the user message so the agent sees the UI state on every turn. Composed in this wrapper rather than polluting the shared hook surface. 2. Stream-inspector wiring — every send opens a `StreamRecord` via `beginStreamRun` and forwards every event to `recordStreamEvent` so the inspector drawer can render the timeline. The inspector is a dev-playground feature, not a shared concern. API stays drop-in compatible: `agentName` (not `agent`), `isLoading` (not `isStreaming`), `send(message, { contextPrefix })`. `SSEEvent` is re-exported as a type alias for `AgentChatEvent` so `use-stream-inspector.ts` and other consumers don't need to be touched. `routes/smart-dashboard.route.tsx`, `agent-sidebar.tsx`, and `quick-actions-bar.tsx` work unchanged. Error handling preserves the prior UX: fetch-level failures now project from `useAgentChat.error` into the streamed `content` as `"Error: ..."`, matching the dashboard's assistant-message rendering of the previous hook's `setContent('Error: ...')` path. Net: 159 -> ~100 lines on the playground hook, zero duplicated SSE parsing across consumers, single source of truth in `@databricks/appkit-ui/react`. Workspace tests: 2310 pass unchanged. Signed-off-by: MarioCadenas <MarioCadenas@users.noreply.github.com>
1 parent 17117dc commit b67371a

1 file changed

Lines changed: 76 additions & 111 deletions

File tree

‎apps/dev-playground/client/src/features/smart-dashboard/hooks/use-agent-stream.ts‎

Lines changed: 76 additions & 111 deletions
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,15 @@
1-
import { useCallback, useRef, useState } from "react";
1+
import { type AgentChatEvent, useAgentChat } from "@databricks/appkit-ui/react";
2+
import { useCallback, useMemo, useRef } from "react";
23
import { beginStreamRun, recordStreamEvent } from "./use-stream-inspector";
34

4-
export interface SSEEvent {
5-
type: string;
6-
delta?: string;
7-
item_id?: string;
8-
item?: {
9-
type?: string;
10-
id?: string;
11-
call_id?: string;
12-
name?: string;
13-
arguments?: string;
14-
output?: string;
15-
status?: string;
16-
};
17-
content?: string;
18-
data?: Record<string, unknown>;
19-
error?: string;
20-
sequence_number?: number;
21-
output_index?: number;
22-
// appkit.approval_pending payload
23-
approval_id?: string;
24-
stream_id?: string;
25-
tool_name?: string;
26-
args?: unknown;
27-
annotations?: {
28-
readOnly?: boolean;
29-
destructive?: boolean;
30-
idempotent?: boolean;
31-
};
32-
}
5+
/**
6+
* Backwards-compatible alias for the SSE event shape that the rest of
7+
* the smart-dashboard code (stream inspector, chat section, action
8+
* dispatcher) already knows about. Identical to {@link AgentChatEvent}
9+
* from `@databricks/appkit-ui/react` — keeping the name in this module
10+
* means downstream callers don't need to be touched.
11+
*/
12+
export type SSEEvent = AgentChatEvent;
3313

3414
interface UseAgentStreamOptions {
3515
agentName: string;
@@ -54,105 +34,90 @@ interface UseAgentStreamReturn {
5434
reset: () => void;
5535
}
5636

37+
/**
38+
* Smart-Dashboard wrapper around `useAgentChat` from
39+
* `@databricks/appkit-ui/react`. The shared hook owns the fetch + SSE
40+
* parsing + state plumbing; this wrapper adds two playground-specific
41+
* concerns:
42+
*
43+
* 1. **`contextPrefix`** on `send()` — the dashboard injects active
44+
* filters / highlights into the user message so the agent always
45+
* sees the UI state. The shared hook stays narrow and lets us
46+
* compose the message here.
47+
* 2. **Stream inspector wiring** — every send opens a `StreamRecord`
48+
* via {@link beginStreamRun} and forwards every event to
49+
* {@link recordStreamEvent} so the inspector drawer can render a
50+
* human-legible timeline. None of that belongs in the shared hook.
51+
*
52+
* Aside from those two layers this hook is a re-export: the SSE parsing
53+
* code that used to live here moved into `useAgentChat`, and the API
54+
* surface is preserved so existing callers (`smart-dashboard.route`,
55+
* `agent-sidebar`) keep working.
56+
*/
5757
export function useAgentStream({
5858
agentName,
5959
onEvent,
6060
}: UseAgentStreamOptions): UseAgentStreamReturn {
61-
const [content, setContent] = useState("");
62-
const [events, setEvents] = useState<SSEEvent[]>([]);
63-
const [isLoading, setIsLoading] = useState(false);
64-
const [threadId, setThreadId] = useState<string | null>(null);
65-
const contentRef = useRef("");
61+
// `runId` is captured at `send()` time so every event of the same run
62+
// lands in the same StreamRecord. Stored as a ref to avoid re-mounting
63+
// the chat hook every time the inspector dispatches.
64+
const runIdRef = useRef<string | null>(null);
6665
const onEventRef = useRef(onEvent);
6766
onEventRef.current = onEvent;
6867

69-
const reset = useCallback(() => {
70-
setContent("");
71-
setEvents([]);
72-
contentRef.current = "";
68+
const handleEvent = useCallback((event: AgentChatEvent) => {
69+
if (runIdRef.current) {
70+
recordStreamEvent(runIdRef.current, event);
71+
}
72+
onEventRef.current?.(event);
7373
}, []);
7474

75+
const {
76+
content: chatContent,
77+
events,
78+
isStreaming,
79+
threadId,
80+
error,
81+
send: chatSend,
82+
reset,
83+
} = useAgentChat({ agent: agentName, onEvent: handleEvent });
84+
7585
const send = useCallback(
7686
async (message: string, opts?: SendOptions) => {
77-
setIsLoading(true);
78-
setContent("");
79-
setEvents([]);
80-
contentRef.current = "";
81-
87+
runIdRef.current = beginStreamRun(
88+
`${agentName}: ${message.slice(0, 80)}`,
89+
);
8290
const wire = opts?.contextPrefix
8391
? `${opts.contextPrefix}${message}`
8492
: message;
85-
86-
const runId = beginStreamRun(`${agentName}: ${message.slice(0, 80)}`);
87-
8893
try {
89-
const res = await fetch("/api/agents/chat", {
90-
method: "POST",
91-
headers: { "Content-Type": "application/json" },
92-
body: JSON.stringify({
93-
message: wire,
94-
agent: agentName,
95-
...(threadId && { threadId }),
96-
}),
97-
});
98-
99-
if (!res.ok) {
100-
const errText = await res.text();
101-
try {
102-
const err = JSON.parse(errText);
103-
setContent(`Error: ${err.error}`);
104-
} catch {
105-
setContent(`Error: Server returned ${res.status}`);
106-
}
107-
return;
108-
}
109-
110-
const reader = res.body?.getReader();
111-
if (!reader) return;
112-
113-
const decoder = new TextDecoder();
114-
let buffer = "";
115-
116-
while (true) {
117-
const { done, value } = await reader.read();
118-
if (done) break;
119-
buffer += decoder.decode(value, { stream: true });
120-
const lines = buffer.split("\n");
121-
buffer = lines.pop() ?? "";
122-
123-
for (const line of lines) {
124-
if (!line.startsWith("data: ")) continue;
125-
const data = line.slice(6).trim();
126-
if (!data || data === "[DONE]") continue;
127-
try {
128-
const event: SSEEvent = JSON.parse(data);
129-
if (!event.type) continue;
130-
setEvents((prev) => [...prev, event]);
131-
recordStreamEvent(runId, event);
132-
onEventRef.current?.(event);
133-
134-
if (event.type === "appkit.metadata" && event.data?.threadId) {
135-
setThreadId(event.data.threadId as string);
136-
}
137-
if (event.type === "response.output_text.delta" && event.delta) {
138-
contentRef.current += event.delta;
139-
setContent(contentRef.current);
140-
}
141-
} catch {
142-
/* skip malformed events */
143-
}
144-
}
145-
}
146-
} catch (err) {
147-
setContent(
148-
`Error: ${err instanceof Error ? err.message : "Unknown error"}`,
149-
);
94+
await chatSend(wire);
15095
} finally {
151-
setIsLoading(false);
96+
runIdRef.current = null;
15297
}
15398
},
154-
[agentName, threadId],
99+
[agentName, chatSend],
155100
);
156101

157-
return { content, events, isLoading, threadId, send, reset };
102+
// Surface fetch-level failures in the displayed content so the
103+
// dashboard's assistant message turns into a visible error row,
104+
// mirroring the prior hook's UX (it wrote "Error: ..." into the
105+
// streamed content on `!res.ok`). `useAgentChat` exposes the error
106+
// via a dedicated `error` field; we project it into `content` only
107+
// when the stream actually failed, otherwise pass `chat.content`
108+
// through verbatim.
109+
const content = useMemo(() => {
110+
if (error) return `Error: ${error}`;
111+
return chatContent;
112+
}, [error, chatContent]);
113+
114+
return {
115+
content,
116+
events,
117+
// `isLoading` is the legacy name; the shared hook uses `isStreaming`.
118+
isLoading: isStreaming,
119+
threadId,
120+
send,
121+
reset,
122+
};
158123
}

0 commit comments

Comments
 (0)