diff --git a/backend/windmill-worker/src/ai/types.rs b/backend/windmill-worker/src/ai/types.rs
index 4a2de30926..457eb3b215 100644
--- a/backend/windmill-worker/src/ai/types.rs
+++ b/backend/windmill-worker/src/ai/types.rs
@@ -1,8 +1,8 @@
 use crate::ai::providers::openai::OpenAIToolCall;
-use windmill_common::mcp_client::McpToolSource;
 use serde::{Deserialize, Serialize};
 use serde_json::value::RawValue;
 use std::collections::HashMap;
+use windmill_common::mcp_client::McpToolSource;
 use windmill_common::{
     ai_providers::AIProvider, db::DB, error::Error, flow_status::AgentAction, flows::FlowModule,
     s3_helpers::S3Object,
@@ -386,3 +386,11 @@ impl OpenAPISchema {
         self
     }
 }
+
+/// Wrapper for S3Object with type discriminator for conversation storage
+#[derive(Serialize)]
+pub struct S3ObjectWithType {
+    #[serde(flatten)]
+    pub s3_object: S3Object,
+    pub r#type: String,
+}
diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs
index 9032a0ad07..cb696d8a6c 100644
--- a/backend/windmill-worker/src/ai_executor.rs
+++ b/backend/windmill-worker/src/ai_executor.rs
@@ -698,9 +698,64 @@ pub async fn run_agent(
                         used_structured_output_tool = tool_used_structured_output;
                     }
                     ParsedResponse::Image { base64_data } => {
-                        // For image output with tools, we got an image response
+                        // For image output, upload to S3 and track in conversation
                         let s3_object = upload_image_to_s3(&base64_data, job, client).await?;
-                        return Ok(to_raw_value(&s3_object));
+
+                        let content = to_raw_value(&s3_object);
+
+                        // Add assistant message to conversation if chat_input_enabled
+                        let chat_enabled = flow_context
+                            .flow_status
+                            .as_ref()
+                            .and_then(|fs| fs.chat_input_enabled)
+                            .unwrap_or(false);
+                        if chat_enabled {
+                            if let Some(memory_id) = flow_context
+                                .flow_status
+                                .as_ref()
+                                .and_then(|fs| fs.memory_id)
+                            {
+                                let agent_job_id = job.id;
+                                let db_clone = db.clone();
+                                let flow_step_id_owned = job.flow_step_id.clone();
+                                let summary_owned = summary.map(|s| s.to_string());
+
+                                // Create extended version with type discriminator for conversation storage
+                                // This avoids conflicts with outputs that are of the same format as S3 objects
+                                let s3_with_type = S3ObjectWithType {
+                                    s3_object: s3_object.clone(),
+                                    r#type: "windmill_s3_object".to_string(),
+                                };
+
+                                let message_content = serde_json::to_string(&s3_with_type)
+                                    .unwrap_or_else(|_| content.get().to_string());
+
+                                // Spawn task because we do not need to wait for the result
+                                tokio::spawn(async move {
+                                    let step_name = get_step_name_from_flow(
+                                        summary_owned.as_deref(),
+                                        flow_step_id_owned.as_deref(),
+                                    );
+
+                                    if let Err(e) = add_message_to_conversation(
+                                        &db_clone,
+                                        &memory_id,
+                                        Some(agent_job_id),
+                                        &message_content,
+                                        MessageType::Assistant,
+                                        &step_name,
+                                        true,
+                                    )
+                                    .await
+                                    {
+                                        tracing::warn!("Failed to add assistant message to conversation {}: {}", memory_id, e);
+                                    }
+                                });
+                            }
+                        }
+
+                        // Return early since image generation is complete
+                        return Ok(content);
                     }
                 }
             }
diff --git a/frontend/src/lib/components/flows/conversations/FlowChatMessage.svelte b/frontend/src/lib/components/flows/conversations/FlowChatMessage.svelte
index ecb614321a..803137677d 100644
--- a/frontend/src/lib/components/flows/conversations/FlowChatMessage.svelte
+++ b/frontend/src/lib/components/flows/conversations/FlowChatMessage.svelte
@@ -4,7 +4,9 @@
 	import { Loader2, CheckCircle2, AlertTriangle } from 'lucide-svelte'
 	import CodeDisplay from '$lib/components/copilot/chat/script/CodeDisplay.svelte'
 	import LinkRenderer from '$lib/components/copilot/chat/LinkRenderer.svelte'
+	import DisplayResult from '$lib/components/DisplayResult.svelte'
 	import { type ChatMessage } from './FlowChatManager.svelte'
+	import { workspaceStore } from '$lib/stores'
 
 	interface Props {
 		message: ChatMessage
@@ -12,6 +14,22 @@
 
 	let { message }: Props = $props()
 
+	// Parse content to detect S3 objects
+	const s3Object: any | undefined = $derived.by(() => {
+		if (message.message_type === 'assistant' && message.content) {
+			try {
+				const parsed = JSON.parse(message.content)
+				// Check if it's a Windmill S3 object with type discriminator
+				if (parsed?.type === 'windmill_s3_object' && parsed?.s3 && typeof parsed.s3 === 'string') {
+					return parsed
+				}
+			} catch (e) {
+				// Not JSON, treat as regular text
+			}
+		}
+		return undefined
+	})
+
 	const messageClass = $derived.by(() => {
 		const base = 'max-w-[90%] min-w-0 rounded-lg w-fit'
 		if (message.message_type === 'user') {
@@ -36,33 +54,39 @@
 			<span>Processing...</span>
 		</div>
 	{:else if message.content}
-		<div
-			class="flex flex-row items-center gap-2 px-3 pb-3 text-sm {!message.step_name
-				? 'pt-3'
-				: ''} overflow-x-auto"
-		>
-			{#if message.message_type === 'tool'}
-				{#if message.success !== false}
-					<CheckCircle2 class="w-3.5 h-3.5 text-green-500" />
-				{:else}
-					<AlertTriangle class="w-3.5 h-3.5 text-red-500" />
-				{/if}
-			{/if}
-			<div class="prose prose-sm dark:prose-invert break-words prose-headings:!text-base">
-				<Markdown
-					md={message.content}
-					plugins={[
-						gfmPlugin(),
-						{
-							renderer: {
-								pre: CodeDisplay,
-								a: LinkRenderer
-							}
-						}
-					]}
-				/>
+		{#if s3Object}
+			<div class="px-3 pb-3 {!message.step_name ? 'pt-3' : ''}">
+				<DisplayResult result={s3Object} workspaceId={$workspaceStore} noControls={true} />
 			</div>
-		</div>
+		{:else}
+			<div
+				class="flex flex-row items-center gap-2 px-3 pb-3 text-sm {!message.step_name
+					? 'pt-3'
+					: ''} overflow-x-auto"
+			>
+				{#if message.message_type === 'tool'}
+					{#if message.success !== false}
+						<CheckCircle2 class="w-3.5 h-3.5 text-green-500" />
+					{:else}
+						<AlertTriangle class="w-3.5 h-3.5 text-red-500" />
+					{/if}
+				{/if}
+				<div class="prose prose-sm dark:prose-invert break-words prose-headings:!text-base">
+					<Markdown
+						md={message.content}
+						plugins={[
+							gfmPlugin(),
+							{
+								renderer: {
+									pre: CodeDisplay,
+									a: LinkRenderer
+								}
+							}
+						]}
+					/>
+				</div>
+			</div>
+		{/if}
 	{:else}
 		<p class="text-tertiary text-sm">No result</p>
 	{/if}