From 28afc7e67db0e874bcc4ea6d2a26f53ff0c568bc Mon Sep 17 00:00:00 2001
From: balibabu <cike8899@users.noreply.github.com>
Date: Mon, 22 Sep 2025 18:08:04 +0800
Subject: [PATCH] Feat: Exporting the results of data flow tests #9869 (#10209)

### What problem does this PR solve?

Feat: Exporting the results of data flow tests #9869

### Type of change


- [x] New Feature (non-breaking change which adds functionality)
---
 .../data-flow/hooks/use-download-output.ts    | 35 ++++++++++++++++++
 .../data-flow/log-sheet/dataflow-timeline.tsx | 17 +++------
 web/src/pages/data-flow/log-sheet/index.tsx   | 36 ++++++++++++++++---
 3 files changed, 71 insertions(+), 17 deletions(-)
 create mode 100644 web/src/pages/data-flow/hooks/use-download-output.ts

diff --git a/web/src/pages/data-flow/hooks/use-download-output.ts b/web/src/pages/data-flow/hooks/use-download-output.ts
new file mode 100644
index 000000000..b2bd8c5e7
--- /dev/null
+++ b/web/src/pages/data-flow/hooks/use-download-output.ts
@@ -0,0 +1,35 @@
+import { ITraceData } from '@/interfaces/database/agent';
+import { downloadJsonFile } from '@/utils/file-util';
+import { get, isEmpty } from 'lodash';
+import { useCallback } from 'react';
+
+export function findEndOutput(list?: ITraceData[]) {
+  if (Array.isArray(list)) {
+    const trace = list.find((x) => x.component_id === 'END')?.trace;
+
+    const str = get(trace, '0.message');
+
+    try {
+      if (!isEmpty(str)) {
+        const json = JSON.parse(str);
+        return json;
+      }
+    } catch (error) {}
+  }
+}
+
+export function isEndOutputEmpty(list?: ITraceData[]) {
+  return isEmpty(findEndOutput(list));
+}
+export function useDownloadOutput(data?: ITraceData[]) {
+  const handleDownloadJson = useCallback(() => {
+    const output = findEndOutput(data);
+    if (!isEndOutputEmpty(data)) {
+      downloadJsonFile(output, `output.json`);
+    }
+  }, [data]);
+
+  return {
+    handleDownloadJson,
+  };
+}
diff --git a/web/src/pages/data-flow/log-sheet/dataflow-timeline.tsx b/web/src/pages/data-flow/log-sheet/dataflow-timeline.tsx
index b2ada516f..5e0410058 100644
--- a/web/src/pages/data-flow/log-sheet/dataflow-timeline.tsx
+++ b/web/src/pages/data-flow/log-sheet/dataflow-timeline.tsx
@@ -8,9 +8,8 @@ import {
   TimelineSeparator,
   TimelineTitle,
 } from '@/components/originui/timeline';
-import { useFetchMessageTrace } from '@/hooks/use-agent-request';
+import { ITraceData } from '@/interfaces/database/agent';
 import { Aperture } from 'lucide-react';
-import { useEffect } from 'react';
 
 const items = [
   {
@@ -50,7 +49,9 @@ const items = [
   },
 ];
 
-export type DataflowTimelineProps = { messageId: string };
+export type DataflowTimelineProps = {
+  traceList?: ITraceData[];
+};
 
 interface DataflowTrace {
   datetime: string;
@@ -59,15 +60,7 @@ interface DataflowTrace {
   progress: number;
   timestamp: number;
 }
-export function DataflowTimeline({ messageId }: DataflowTimelineProps) {
-  const { setMessageId, data } = useFetchMessageTrace(false);
-
-  useEffect(() => {
-    if (messageId) {
-      setMessageId(messageId);
-    }
-  }, [messageId, setMessageId]);
-
+export function DataflowTimeline({ traceList }: DataflowTimelineProps) {
   return (
     <Timeline>
       {items.map((item) => (
diff --git a/web/src/pages/data-flow/log-sheet/index.tsx b/web/src/pages/data-flow/log-sheet/index.tsx
index b4aafa837..e2912e3cb 100644
--- a/web/src/pages/data-flow/log-sheet/index.tsx
+++ b/web/src/pages/data-flow/log-sheet/index.tsx
@@ -1,28 +1,54 @@
+import { Button } from '@/components/ui/button';
 import {
   Sheet,
   SheetContent,
   SheetHeader,
   SheetTitle,
 } from '@/components/ui/sheet';
+import { useFetchMessageTrace } from '@/hooks/use-agent-request';
 import { IModalProps } from '@/interfaces/common';
 import { cn } from '@/lib/utils';
-import { NotebookText } from 'lucide-react';
+import { NotebookText, SquareArrowOutUpRight } from 'lucide-react';
+import { useEffect } from 'react';
 import 'react18-json-view/src/style.css';
-import { DataflowTimeline, DataflowTimelineProps } from './dataflow-timeline';
+import {
+  isEndOutputEmpty,
+  useDownloadOutput,
+} from '../hooks/use-download-output';
+import { DataflowTimeline } from './dataflow-timeline';
 
-type LogSheetProps = IModalProps<any> & DataflowTimelineProps;
+type LogSheetProps = IModalProps<any> & { messageId?: string };
 
 export function LogSheet({ hideModal, messageId }: LogSheetProps) {
+  const { setMessageId, data } = useFetchMessageTrace(false);
+
+  const { handleDownloadJson } = useDownloadOutput(data);
+
+  useEffect(() => {
+    if (messageId) {
+      setMessageId(messageId);
+    }
+  }, [messageId, setMessageId]);
+
   return (
     <Sheet open onOpenChange={hideModal} modal={false}>
       <SheetContent className={cn('top-20 right-[620px]')}>
         <SheetHeader>
           <SheetTitle className="flex items-center gap-1">
             <NotebookText className="size-4" />
-            <DataflowTimeline messageId={messageId}></DataflowTimeline>
           </SheetTitle>
         </SheetHeader>
-        <section className="max-h-[82vh] overflow-auto mt-6"></section>
+        <section className="max-h-[82vh] overflow-auto mt-6">
+          <DataflowTimeline traceList={data}></DataflowTimeline>
+          <Button
+            onClick={handleDownloadJson}
+            disabled={isEndOutputEmpty(data)}
+            className="w-full mt-8"
+          >
+            <SquareArrowOutUpRight />
+            Export JSON
+          </Button>
+        </section>
       </SheetContent>
     </Sheet>
   );