feat(workbench): redesign investigation reports and diagrams

src/features/workbench/hooks/agent-session-message-state.ts

@@ -2,6 +2,7 @@ import type { UIMessage } from "ai";
 import type { Dispatch, SetStateAction } from "react";
 import {
   applyPermissionResponse,
+  agentTaskContextReferenceSchema,
   applyQuestionResponse,
   toTodoUpdate,
   upsertPermission,
@@ -16,6 +17,7 @@ import {
 } from "@/features/agent";
 
 type AgentUiDataParts = {
+  task_context: Record<string, unknown>;
   execution_plan: Record<string, unknown>;
   progress: Record<string, unknown>;
   todo_update: Record<string, unknown>;
@@ -123,7 +125,10 @@ export function toAgentChatMessages(
       if (!isAgentDataPart(part)) {
         continue;
       }
-      if (part.type === "data-progress") {
+      if (part.type === "data-task_context") {
+        const parsed = agentTaskContextReferenceSchema.safeParse(part.data);
+        if (parsed.success) next = { ...next, taskContext: parsed.data };
+      } else if (part.type === "data-progress") {
         next = {
           ...next,
           progress: upsertProgress(next.progress, part.data)