Merge pull request #352 from stakwork/bugfix/CMQS8FYG5-fix-agent-streaming-envelope-markdown-citations-1782316262

Generated with Hive: Parse agent streaming envelope and render markdown with citations

Author: tomsmith8

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<GraphNode | null>(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 <Skeleton data-testid="cited-node-skeleton" className="h-8 w-full rounded-md" />
+  }
+
+  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 (
+    <NodeRow
+      node={nodeWithLabel}
+      schemas={schemas}
+      onClick={() => 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 (
     <div className="mt-3 pt-3 border-t border-border/40">
       <p className="text-[10px] font-mono uppercase tracking-widest text-muted-foreground mb-2">
         Sources
       </p>
       <div className="flex flex-col gap-0.5">
-        {stubNodes.map((node) => (
-          <NodeRow
-            key={node.ref_id}
-            node={node}
-            schemas={schemas}
-            onClick={() => unlockNode(node.ref_id).catch(() => {})}
-            hideBoost
-          />
+        {refIds.map((refId) => (
+          <CitedNodeChip key={refId} refId={refId} schemas={schemas} />
         ))}
       </div>
     </div>

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: [] })

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 }) => (
-    <div data-testid="skeleton" className={className} />
+  Skeleton: ({ className, ...props }: { className?: string; [key: string]: unknown }) => (
+    <div data-testid="skeleton" className={className} {...props} />
   ),
 }))
 
 vi.mock("@/components/layout/node-row", () => ({
-  NodeRow: () => <div data-testid="node-row" />,
+  NodeRow: ({ node }: { node: { properties?: { name?: string }; node_type?: string } }) => (
+    <div data-testid="node-row">
+      <span data-testid="node-label">{node.properties?.name}</span>
+      <span data-testid="node-type">{node.node_type}</span>
+    </div>
+  ),
 }))
 
 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(
+      <MessageList
+        messages={[
+          {
+            role: "agent",
+            content: "Hello",
+            isStreaming: false,
+            citedRefIds: ["abc"],
+          },
+        ]}
+      />
+    )
+    // 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(
+      <MessageList
+        messages={[
+          {
+            role: "agent",
+            content: "Hello",
+            isStreaming: false,
+            citedRefIds: ["abc"],
+          },
+        ]}
+      />
+    )
+    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(
+      <MessageList
+        messages={[
+          {
+            role: "agent",
+            content: "Hello",
+            isStreaming: false,
+            citedRefIds: ["fallback-id"],
+          },
+        ]}
+      />
+    )
+    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(
+      <MessageList
+        messages={[
+          {
+            role: "agent",
+            content: "Answer",
+            isStreaming: false,
+            citedRefIds: ["ep-99"],
+          },
+        ]}
+      />
+    )
+    await waitFor(() => {
+      expect(screen.getByTestId("node-label")).toHaveTextContent("My Article Title")
+    })
+  })
+})