From f801b8eb0d60bcf89da7d9bb0dcd788fe457319c Mon Sep 17 00:00:00 2001
From: shrimbly <falloon.willie@gmail.com>
Date: Wed, 18 Feb 2026 14:28:25 +1300
Subject: [PATCH] feat: connect GenerateAudio output to Output node with audio
 playback
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add audio support to the Output node so GenerateAudio can pipe audio
through for playback and download. Adds audio input handle, <audio>
player, audio→output connection validation, and executeOutput audio
handling.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/components/ConnectionDropMenu.tsx        |  9 ++
 src/components/WorkflowCanvas.tsx            | 16 +++-
 src/components/__tests__/OutputNode.test.tsx |  2 +-
 src/components/nodes/OutputNode.tsx          | 91 +++++++++++++-------
 src/store/execution/simpleNodeExecutors.ts   | 35 +++++++-
 src/types/nodes.ts                           |  3 +-
 6 files changed, 118 insertions(+), 38 deletions(-)

diff --git a/src/components/ConnectionDropMenu.tsx b/src/components/ConnectionDropMenu.tsx
index 876011aa..09bebbc4 100644
--- a/src/components/ConnectionDropMenu.tsx
+++ b/src/components/ConnectionDropMenu.tsx
@@ -290,6 +290,15 @@ const VIDEO_SOURCE_OPTIONS: MenuOption[] = [
 
 // Audio target options (nodes that accept audio input)
 const AUDIO_TARGET_OPTIONS: MenuOption[] = [
+  {
+    type: "output",
+    label: "Output",
+    icon: (
+      <svg className="w-4 h-4" fill="none" viewBox="0 0 24 24" stroke="currentColor" strokeWidth={1.5}>
+        <path strokeLinecap="round" strokeLinejoin="round" d="M3 16.5v2.25A2.25 2.25 0 005.25 21h13.5A2.25 2.25 0 0021 18.75V16.5M16.5 12L12 16.5m0 0L7.5 12m4.5 4.5V3" />
+      </svg>
+    ),
+  },
   {
     type: "videoStitch",
     label: "Video Stitch",
diff --git a/src/components/WorkflowCanvas.tsx b/src/components/WorkflowCanvas.tsx
index 17d3b7fe..f238e7b5 100644
--- a/src/components/WorkflowCanvas.tsx
+++ b/src/components/WorkflowCanvas.tsx
@@ -131,7 +131,7 @@ const getNodeHandles = (nodeType: string): { inputs: string[]; outputs: string[]
     case "splitGrid":
       return { inputs: ["image"], outputs: ["reference"] };
     case "output":
-      return { inputs: ["image", "video"], outputs: [] };
+      return { inputs: ["image", "video", "audio"], outputs: [] };
     case "outputGallery":
       return { inputs: ["image"], outputs: [] };
     case "imageCompare":
@@ -346,8 +346,12 @@ export function WorkflowCanvas() {
         return sourceType === "3d" && targetType === "3d";
       }
 
-      // Audio connections: audio handles can only connect to audio handles
+      // Audio connections: audio handles connect to audio handles, plus output node
       if (sourceType === "audio" || targetType === "audio") {
+        if (sourceType === "audio") {
+          const targetNode = nodes.find((n) => n.id === connection.target);
+          if (targetNode?.type === "output") return true;
+        }
         return sourceType === "audio" && targetType === "audio";
       }
 
@@ -526,6 +530,11 @@ export function WorkflowCanvas() {
           return "image";
         }
 
+        // For audio output connecting to output node, use the "audio" input handle
+        if (handleType === "audio" && needInput && node.type === "output") {
+          return "audio";
+        }
+
         // Then check each handle's type
         for (const h of handleList) {
           if (getHandleType(h) === handleType) return h;
@@ -850,6 +859,9 @@ export function WorkflowCanvas() {
         } else if (nodeType === "videoStitch") {
           // VideoStitch accepts audio
           targetHandleId = "audio";
+        } else if (nodeType === "output") {
+          // Output accepts audio on its audio handle
+          targetHandleId = "audio";
         }
       } else if (handleType === "3d") {
         if (nodeType === "glbViewer") {
diff --git a/src/components/__tests__/OutputNode.test.tsx b/src/components/__tests__/OutputNode.test.tsx
index 9c0e3f07..ed0cf507 100644
--- a/src/components/__tests__/OutputNode.test.tsx
+++ b/src/components/__tests__/OutputNode.test.tsx
@@ -78,7 +78,7 @@ describe("OutputNode", () => {
         </TestWrapper>
       );
 
-      expect(screen.getByText("Waiting for image or video")).toBeInTheDocument();
+      expect(screen.getByText("Waiting for image, video, or audio")).toBeInTheDocument();
     });
 
     it("should render the title 'Output'", () => {
diff --git a/src/components/nodes/OutputNode.tsx b/src/components/nodes/OutputNode.tsx
index b9c9b6c9..40b7d8f8 100644
--- a/src/components/nodes/OutputNode.tsx
+++ b/src/components/nodes/OutputNode.tsx
@@ -15,26 +15,36 @@ export function OutputNode({ id, data, selected }: NodeProps<OutputNodeType>) {
   const updateNodeData = useWorkflowStore((state) => state.updateNodeData);
   const [showLightbox, setShowLightbox] = useState(false);
 
+  // Determine if content is audio
+  const isAudio = useMemo(() => {
+    if (nodeData.audio) return true;
+    if (nodeData.contentType === "audio") return true;
+    if (nodeData.image?.startsWith("data:audio/")) return true;
+    return false;
+  }, [nodeData.audio, nodeData.contentType, nodeData.image]);
+
   // Determine if content is video
   const isVideo = useMemo(() => {
+    if (isAudio) return false;
     if (nodeData.video) return true;
     if (nodeData.contentType === "video") return true;
     if (nodeData.image?.startsWith("data:video/")) return true;
     if (nodeData.image?.includes(".mp4") || nodeData.image?.includes(".webm")) return true;
     return false;
-  }, [nodeData.video, nodeData.contentType, nodeData.image]);
+  }, [isAudio, nodeData.video, nodeData.contentType, nodeData.image]);
 
-  // Get the content source (video or image)
+  // Get the content source (audio, video, or image)
   const contentSrc = useMemo(() => {
+    if (nodeData.audio) return nodeData.audio;
     if (nodeData.video) return nodeData.video;
     return nodeData.image;
-  }, [nodeData.video, nodeData.image]);
+  }, [nodeData.audio, nodeData.video, nodeData.image]);
 
   const handleDownload = useCallback(async () => {
     if (!contentSrc) return;
 
     const timestamp = Date.now();
-    const extension = isVideo ? "mp4" : "png";
+    const extension = isAudio ? "mp3" : isVideo ? "mp4" : "png";
     // Use custom filename if provided, otherwise use timestamp
     const filename = nodeData.outputFilename
       ? `${nodeData.outputFilename}.${extension}`
@@ -67,7 +77,7 @@ export function OutputNode({ id, data, selected }: NodeProps<OutputNodeType>) {
     document.body.appendChild(link);
     link.click();
     document.body.removeChild(link);
-  }, [contentSrc, isVideo, nodeData.outputFilename]);
+  }, [contentSrc, isAudio, isVideo, nodeData.outputFilename]);
 
   return (
     <>
@@ -88,37 +98,54 @@ export function OutputNode({ id, data, selected }: NodeProps<OutputNodeType>) {
           id="image"
           data-handletype="image"
         />
+        <Handle
+          type="target"
+          position={Position.Left}
+          id="audio"
+          data-handletype="audio"
+          style={{ top: "60%", background: "rgb(167, 139, 250)" }}
+        />
 
         {contentSrc ? (
           <div className="flex-1 flex flex-col min-h-0 gap-2">
-            <div
-              className="relative cursor-pointer group flex-1 min-h-0"
-              onClick={() => setShowLightbox(true)}
-            >
-              {isVideo ? (
-                <video
+            {isAudio ? (
+              <div className="flex-1 flex items-center min-h-0 py-2">
+                <audio
                   src={contentSrc}
                   controls
-                  loop
-                  muted
-                  autoPlay
-                  playsInline
-                  className="w-full h-full object-contain rounded"
-                  onClick={(e) => e.stopPropagation()}
-                />
-              ) : (
-                <img
-                  src={contentSrc}
-                  alt="Output"
-                  className="w-full h-full object-contain rounded"
+                  className="w-full rounded"
                 />
-              )}
-              <div className="absolute inset-0 bg-black/0 group-hover:bg-black/10 transition-colors flex items-center justify-center rounded pointer-events-none">
-                <span className="text-[10px] font-medium text-white opacity-0 group-hover:opacity-100 transition-opacity bg-black/50 px-2 py-1 rounded">
-                  View full size
-                </span>
               </div>
-            </div>
+            ) : (
+              <div
+                className="relative cursor-pointer group flex-1 min-h-0"
+                onClick={() => setShowLightbox(true)}
+              >
+                {isVideo ? (
+                  <video
+                    src={contentSrc}
+                    controls
+                    loop
+                    muted
+                    autoPlay
+                    playsInline
+                    className="w-full h-full object-contain rounded"
+                    onClick={(e) => e.stopPropagation()}
+                  />
+                ) : (
+                  <img
+                    src={contentSrc}
+                    alt="Output"
+                    className="w-full h-full object-contain rounded"
+                  />
+                )}
+                <div className="absolute inset-0 bg-black/0 group-hover:bg-black/10 transition-colors flex items-center justify-center rounded pointer-events-none">
+                  <span className="text-[10px] font-medium text-white opacity-0 group-hover:opacity-100 transition-opacity bg-black/50 px-2 py-1 rounded">
+                    View full size
+                  </span>
+                </div>
+              </div>
+            )}
             <button
               onClick={handleDownload}
               className="w-full py-1.5 bg-white hover:bg-neutral-200 text-neutral-900 text-[10px] font-medium rounded transition-colors shrink-0"
@@ -128,7 +155,7 @@ export function OutputNode({ id, data, selected }: NodeProps<OutputNodeType>) {
           </div>
         ) : (
           <div className="w-full flex-1 min-h-[144px] border border-dashed border-neutral-600 rounded flex items-center justify-center">
-            <span className="text-neutral-500 text-[10px]">Waiting for image or video</span>
+            <span className="text-neutral-500 text-[10px]">Waiting for image, video, or audio</span>
           </div>
         )}
 
@@ -144,8 +171,8 @@ export function OutputNode({ id, data, selected }: NodeProps<OutputNodeType>) {
         </div>
       </BaseNode>
 
-      {/* Lightbox Modal */}
-      {showLightbox && contentSrc && (
+      {/* Lightbox Modal (skip for audio) */}
+      {showLightbox && contentSrc && !isAudio && (
         <div
           className="fixed inset-0 bg-black/90 z-[100] flex items-center justify-center p-8"
           onClick={() => setShowLightbox(false)}
diff --git a/src/store/execution/simpleNodeExecutors.ts b/src/store/execution/simpleNodeExecutors.ts
index 0cfcd24d..b451cf7c 100644
--- a/src/store/execution/simpleNodeExecutors.ts
+++ b/src/store/execution/simpleNodeExecutors.ts
@@ -147,9 +147,40 @@ export async function executePromptConstructor(ctx: NodeExecutionContext): Promi
  */
 export async function executeOutput(ctx: NodeExecutionContext): Promise<void> {
   const { node, getConnectedInputs, updateNodeData, saveDirectoryPath } = ctx;
-  const { images, videos } = getConnectedInputs(node.id);
+  const { images, videos, audio } = getConnectedInputs(node.id);
 
-  // Check videos array first (typed data from source)
+  // Check audio array first
+  if (audio.length > 0) {
+    const audioContent = audio[0];
+    updateNodeData(node.id, {
+      audio: audioContent,
+      image: null,
+      video: null,
+      contentType: "audio",
+    });
+
+    // Save to /outputs directory if we have a project path
+    if (saveDirectoryPath) {
+      const outputNodeData = node.data as OutputNodeData;
+      const outputsPath = `${saveDirectoryPath}/outputs`;
+
+      fetch("/api/save-generation", {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({
+          directoryPath: outputsPath,
+          audio: audioContent,
+          customFilename: outputNodeData.outputFilename || undefined,
+          createDirectory: true,
+        }),
+      }).catch((err) => {
+        console.error("Failed to save output:", err);
+      });
+    }
+    return;
+  }
+
+  // Check videos array (typed data from source)
   if (videos.length > 0) {
     const videoContent = videos[0];
     updateNodeData(node.id, {
diff --git a/src/types/nodes.ts b/src/types/nodes.ts
index 80a161c2..19a8fb1e 100644
--- a/src/types/nodes.ts
+++ b/src/types/nodes.ts
@@ -242,7 +242,8 @@ export interface OutputNodeData extends BaseNodeData {
   image: string | null;
   imageRef?: string; // External image reference for storage optimization
   video?: string | null; // Video data URL or HTTP URL
-  contentType?: "image" | "video"; // Explicit content type hint
+  audio?: string | null; // Audio data URL or HTTP URL
+  contentType?: "image" | "video" | "audio"; // Explicit content type hint
   outputFilename?: string; // Custom filename for saved outputs (without extension)
 }