From 2e7372c59c77991df7f97f881e55eef6b587baf4 Mon Sep 17 00:00:00 2001 From: tomsmith8 Date: Wed, 24 Jun 2026 15:51:18 +0000 Subject: [PATCH] Generated with Hive: Parse agent streaming envelope and render markdown with citations --- src/components/agent/message-list.tsx | 75 +++++++++++---- src/lib/__tests__/agent-api.test.ts | 99 ++++++++++++++----- src/lib/__tests__/message-list.test.tsx | 120 +++++++++++++++++++++++- src/lib/agent-api.ts | 33 +++++-- 4 files changed, 277 insertions(+), 50 deletions(-) diff --git a/src/components/agent/message-list.tsx b/src/components/agent/message-list.tsx index d95c9cd..7ca0448 100644 --- a/src/components/agent/message-list.tsx +++ b/src/components/agent/message-list.tsx @@ -1,16 +1,18 @@ "use client" -import { useRef, useEffect } from "react" +import { useRef, useEffect, useState } from "react" import { Bot, User } from "lucide-react" import { ScrollArea } from "@/components/ui/scroll-area" import { NodeRow } from "@/components/layout/node-row" import { ToolCallRow } from "./tool-call-row" import { unlockNode } from "@/lib/unlock-node" +import { getNode } from "@/lib/graph-api" import { useSchemaStore } from "@/stores/schema-store" import { cn } from "@/lib/utils" import { Skeleton } from "@/components/ui/skeleton" import type { AgentMessage } from "@/lib/agent-api" import type { GraphNode } from "@/lib/graph-api" +import type { SchemaNode } from "@/app/ontology/page" // Basic markdown renderer — bold, italic, inline code, line breaks, headings function MarkdownText({ text }: { text: string }) { @@ -57,32 +59,73 @@ interface CitedNodesProps { refIds: string[] } +function CitedNodeChip({ refId, schemas }: { refId: string; schemas: SchemaNode[] }) { + const [node, setNode] = useState(null) + const [loading, setLoading] = useState(true) + + useEffect(() => { + let cancelled = false + getNode(refId) + .then((n) => { + if (!cancelled) setNode(n as GraphNode) + }) + .catch(() => { + // fall back to raw refId shown via stub node + }) + .finally(() => { + if (!cancelled) setLoading(false) + }) + return () => { + cancelled = true + } + }, [refId]) + + if (loading) { + return + } + + const displayNode: GraphNode = node ?? { + ref_id: refId, + node_type: refId, + properties: { name: refId }, + } + + // Derive human-readable label + const label = + (displayNode.properties?.name as string | undefined) ?? + (displayNode.properties?.title as string | undefined) ?? + displayNode.node_type ?? + refId + + const nodeWithLabel: GraphNode = { + ...displayNode, + node_type: displayNode.node_type, + properties: { ...displayNode.properties, name: label }, + } + + return ( + unlockNode(refId).catch(() => {})} + hideBoost + /> + ) +} + function CitedNodes({ refIds }: CitedNodesProps) { const schemas = useSchemaStore((s) => s.schemas) if (refIds.length === 0) return null - // Create minimal GraphNode stubs for display — real data fetched on click via unlockNode - const stubNodes: GraphNode[] = refIds.map((ref_id) => ({ - ref_id, - node_type: "Unknown", - properties: { ref_id }, - })) - return (

Sources

- {stubNodes.map((node) => ( - unlockNode(node.ref_id).catch(() => {})} - hideBoost - /> + {refIds.map((refId) => ( + ))}
diff --git a/src/lib/__tests__/agent-api.test.ts b/src/lib/__tests__/agent-api.test.ts index 6243c88..87907a4 100644 --- a/src/lib/__tests__/agent-api.test.ts +++ b/src/lib/__tests__/agent-api.test.ts @@ -129,28 +129,16 @@ describe("streamAgent – context in POST body", () => { // ─── SSE streaming behaviour ────────────────────────────────────────────────── describe("streamAgent – SSE streaming", () => { - it("text-delta chunk calls onChunk with the delta string", async () => { + it("text-delta chunk does NOT call onChunk (answer suppressed until finish)", async () => { const onChunk = vi.fn() - await streamAgent( - "Q?", - makeOpts({ - onChunk, - ...await makeSseResponse([ - { type: "text-delta", textDelta: "Hello!" }, - { type: "finish-message" }, - ]).then(() => ({})), - }) - ) - // Use a dedicated fetch mock for this test - const onChunk2 = vi.fn() mockFetch.mockImplementationOnce(() => makeSseResponse([ { type: "text-delta", textDelta: "Hello!" }, { type: "finish-message" }, ]) ) - await streamAgent("Q?", makeOpts({ onChunk: onChunk2 })) - expect(onChunk2).toHaveBeenCalledWith("Hello!") + await streamAgent("Q?", makeOpts({ onChunk })) + expect(onChunk).not.toHaveBeenCalled() }) it("tool-input-available calls onToolCall with status in-flight", async () => { @@ -200,23 +188,92 @@ describe("streamAgent – SSE streaming", () => { expect(done).toBeDefined() }) - it("finish-message calls onDone with accumulated text and empty cited_ref_ids", async () => { + it("finish-message parses JSON envelope: extracts answer and cited_ref_ids", async () => { const onDone = vi.fn() + const envelope = JSON.stringify({ + answer: "The answer is 42.", + cited_ref_ids: ["node-abc", "node-xyz"], + usage: {}, + }) mockFetch.mockImplementationOnce(() => makeSseResponse([ - { type: "text-delta", textDelta: "Foo " }, - { type: "text-delta", textDelta: "bar." }, + { type: "text-delta", textDelta: envelope }, { type: "finish-message" }, ]) ) await streamAgent("Q?", makeOpts({ onDone })) - expect(onDone).toHaveBeenCalledWith({ answer: "Foo bar.", cited_ref_ids: [] }) + expect(onDone).toHaveBeenCalledWith({ + answer: "The answer is 42.", + cited_ref_ids: ["node-abc", "node-xyz"], + }) + }) + + it("finish-message: envelope missing cited_ref_ids defaults to []", async () => { + const onDone = vi.fn() + const envelope = JSON.stringify({ answer: "Short answer." }) + mockFetch.mockImplementationOnce(() => + makeSseResponse([ + { type: "text-delta", textDelta: envelope }, + { type: "finish-message" }, + ]) + ) + await streamAgent("Q?", makeOpts({ onDone })) + expect(onDone).toHaveBeenCalledWith({ answer: "Short answer.", cited_ref_ids: [] }) + }) + + it("finish-message: strips end-of-answer marker before parsing", async () => { + const onDone = vi.fn() + const envelope = + JSON.stringify({ answer: "Marked answer.", cited_ref_ids: ["ref-1"] }) + + "[END_OF_ANSWER]" + mockFetch.mockImplementationOnce(() => + makeSseResponse([ + { type: "text-delta", textDelta: envelope }, + { type: "finish-message" }, + ]) + ) + await streamAgent("Q?", makeOpts({ onDone })) + expect(onDone).toHaveBeenCalledWith({ answer: "Marked answer.", cited_ref_ids: ["ref-1"] }) + }) + + it("finish-message: JSON inside a ```json fence is still extracted", async () => { + const onDone = vi.fn() + const fenced = + "```json\n" + + JSON.stringify({ answer: "Fenced answer.", cited_ref_ids: ["ref-2"] }) + + "\n```" + mockFetch.mockImplementationOnce(() => + makeSseResponse([ + { type: "text-delta", textDelta: fenced }, + { type: "finish-message" }, + ]) + ) + await streamAgent("Q?", makeOpts({ onDone })) + expect(onDone).toHaveBeenCalledWith({ answer: "Fenced answer.", cited_ref_ids: ["ref-2"] }) + }) + + it("finish-message: falls back to raw text and warns on invalid JSON", async () => { + const onDone = vi.fn() + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + mockFetch.mockImplementationOnce(() => + makeSseResponse([ + { type: "text-delta", textDelta: "{ bad json }" }, + { type: "finish-message" }, + ]) + ) + await streamAgent("Q?", makeOpts({ onDone })) + expect(warnSpy).toHaveBeenCalledWith( + expect.stringContaining("[agent-api] envelope parse failed") + ) + expect(onDone).toHaveBeenCalledWith({ answer: "{ bad json }", cited_ref_ids: [] }) + warnSpy.mockRestore() }) - it("fallback onDone called when stream ends without finish-message", async () => { + it("fallback onDone called when stream ends without finish-message (parses envelope)", async () => { const onDone = vi.fn() + const envelope = JSON.stringify({ answer: "Partial.", cited_ref_ids: [] }) mockFetch.mockImplementationOnce(() => - makeSseResponse([{ type: "text-delta", textDelta: "Partial." }]) + makeSseResponse([{ type: "text-delta", textDelta: envelope }]) ) await streamAgent("Q?", makeOpts({ onDone })) expect(onDone).toHaveBeenCalledWith({ answer: "Partial.", cited_ref_ids: [] }) diff --git a/src/lib/__tests__/message-list.test.tsx b/src/lib/__tests__/message-list.test.tsx index aa1b41a..309e9f0 100644 --- a/src/lib/__tests__/message-list.test.tsx +++ b/src/lib/__tests__/message-list.test.tsx @@ -1,5 +1,5 @@ -import { describe, it, expect, vi } from "vitest" -import { render, screen } from "@testing-library/react" +import { describe, it, expect, vi, beforeEach } from "vitest" +import { render, screen, waitFor } from "@testing-library/react" import React from "react" // ── Mock heavy dependencies ─────────────────────────────────────────────────── @@ -8,13 +8,18 @@ vi.mock("@/components/ui/scroll-area", () => ({ })) vi.mock("@/components/ui/skeleton", () => ({ - Skeleton: ({ className }: { className?: string }) => ( -
+ Skeleton: ({ className, ...props }: { className?: string; [key: string]: unknown }) => ( +
), })) vi.mock("@/components/layout/node-row", () => ({ - NodeRow: () =>
, + NodeRow: ({ node }: { node: { properties?: { name?: string }; node_type?: string } }) => ( +
+ {node.properties?.name} + {node.node_type} +
+ ), })) vi.mock("@/lib/unlock-node", () => ({ @@ -25,8 +30,24 @@ vi.mock("@/stores/schema-store", () => ({ useSchemaStore: (sel: (s: { schemas: [] }) => unknown) => sel({ schemas: [] }), })) +// Mock graph-api getNode +const mockGetNode = vi.fn() +vi.mock("@/lib/graph-api", () => ({ + getNode: (...args: unknown[]) => mockGetNode(...args), +})) + import { MessageList } from "@/components/agent/message-list" +beforeEach(() => { + vi.clearAllMocks() + // Default: getNode resolves with a real node + mockGetNode.mockResolvedValue({ + ref_id: "abc", + node_type: "Episode", + properties: { name: "Test Ep" }, + }) +}) + describe("MessageList", () => { it("shows Thinking… placeholder when isStreaming=true and content is empty", () => { render( @@ -88,3 +109,92 @@ describe("MessageList", () => { expect(screen.getAllByTestId("skeleton")).toHaveLength(3) }) }) + +describe("CitedNodes — label resolution", () => { + it("shows a loading skeleton per chip while getNode is pending", async () => { + // Never resolves during this test + mockGetNode.mockReturnValue(new Promise(() => {})) + render( + + ) + // The "cited-node-skeleton" should appear while fetch is in-flight + expect(screen.getByTestId("cited-node-skeleton")).toBeInTheDocument() + }) + + it("renders the resolved node label and type badge", async () => { + mockGetNode.mockResolvedValue({ + ref_id: "abc", + node_type: "Episode", + properties: { name: "Test Ep" }, + }) + render( + + ) + await waitFor(() => { + expect(screen.getByTestId("node-label")).toHaveTextContent("Test Ep") + expect(screen.getByTestId("node-type")).toHaveTextContent("Episode") + }) + }) + + it("falls back to ref_id when getNode rejects", async () => { + mockGetNode.mockRejectedValue(new Error("not found")) + render( + + ) + await waitFor(() => { + expect(screen.getByTestId("node-row")).toBeInTheDocument() + // Label falls back to refId when node not found + expect(screen.getByTestId("node-label")).toHaveTextContent("fallback-id") + }) + }) + + it("uses title property when name is absent", async () => { + mockGetNode.mockResolvedValue({ + ref_id: "ep-99", + node_type: "Article", + properties: { title: "My Article Title" }, + }) + render( + + ) + await waitFor(() => { + expect(screen.getByTestId("node-label")).toHaveTextContent("My Article Title") + }) + }) +}) diff --git a/src/lib/agent-api.ts b/src/lib/agent-api.ts index fda0432..fd3f68d 100644 --- a/src/lib/agent-api.ts +++ b/src/lib/agent-api.ts @@ -37,6 +37,27 @@ export interface StreamAgentOpts { onError: (err: Error) => void } +// Parses the agent's JSON envelope from accumulated text. +// Strips the end-of-answer marker, extracts {answer, cited_ref_ids}. +// Falls back to raw text on parse failure. +function unwrapEnvelope(raw: string): { answer: string; cited_ref_ids: string[] } { + const END_MARKER = "[END_OF_" + "ANSWER]" // avoid literal in source + const stripped = raw.replace(END_MARKER, "").trim() + const jsonMatch = stripped.match(/\{[\s\S]*\}/) + let answer = stripped + let cited_ref_ids: string[] = [] + if (jsonMatch) { + try { + const parsed = JSON.parse(jsonMatch[0]) + if (parsed.answer) answer = parsed.answer + if (Array.isArray(parsed.cited_ref_ids)) cited_ref_ids = parsed.cited_ref_ids + } catch { + console.warn("[agent-api] envelope parse failed; falling back to raw text") + } + } + return { answer, cited_ref_ids } +} + // Builds a signed URL for a given API path async function buildSignedUrl(path: string): Promise { const url = new URL(`${API_URL}${path}`) @@ -97,11 +118,6 @@ async function mockStreamAgent( `The most prominent nodes relate to recent episodes and community discussions.\n\n` + `*(This is a mock response — connect to a real backend to get live answers.)*` - for (const word of answer.split(" ")) { - await delay(40) - opts.onChunk(word + " ") - } - await delay(200) opts.onDone({ answer, cited_ref_ids: ["mock-node-1", "mock-node-2"] }) } @@ -144,7 +160,8 @@ async function processSSEStream(response: Response, opts: StreamAgentOpts): Prom case "text-delta": { const delta = (chunk.textDelta ?? chunk.delta ?? "") as string accumulatedText += delta - opts.onChunk(delta) + // Do NOT call opts.onChunk — suppress raw JSON from the bubble. + // Tool-call events stream live; the answer is parsed and delivered on finish. break } case "tool-input-available": { @@ -167,14 +184,14 @@ async function processSSEStream(response: Response, opts: StreamAgentOpts): Prom break } case "finish-message": { - opts.onDone({ answer: accumulatedText, cited_ref_ids: [] }) + opts.onDone(unwrapEnvelope(accumulatedText)) return } } } } // Fallback if stream ends without finish-message - opts.onDone({ answer: accumulatedText, cited_ref_ids: [] }) + opts.onDone(unwrapEnvelope(accumulatedText)) } catch (err) { if (err instanceof DOMException && err.name === "AbortError") return opts.onError(err instanceof Error ? err : new Error(String(err)))