From 36ca8304814f5ad16b5d36f33eb3957d73945594 Mon Sep 17 00:00:00 2001
From: chihlasm <michael@chihlas.com>
Date: Tue, 24 Mar 2026 05:02:34 +0000
Subject: [PATCH] fix: wire image uploads into correct chat endpoint
 (unified_chat_service)

The frontend calls /ai-sessions/{id}/chat (unified_chat_service), not
/assistant/chats/{id}/messages (assistant_chat_service). The previous
commit wired images into the wrong backend. This fixes it:

- ai_session.py schema: add upload_ids to ChatMessageRequest
- ai_sessions.py endpoint: fetch images via _fetch_upload_images
- unified_chat_service: accept and forward images to _call_ai

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 backend/app/api/endpoints/ai_sessions.py     | 7 +++++++
 backend/app/schemas/ai_session.py            | 1 +
 backend/app/services/unified_chat_service.py | 6 ++++++
 3 files changed, 14 insertions(+)

diff --git a/backend/app/api/endpoints/ai_sessions.py b/backend/app/api/endpoints/ai_sessions.py
index df0be33b..0be32bbc 100644
--- a/backend/app/api/endpoints/ai_sessions.py
+++ b/backend/app/api/endpoints/ai_sessions.py
@@ -280,6 +280,12 @@ async def send_chat_message(
     user_id = current_user.id
     account_id = current_user.account_id
 
+    # Fetch attached images from S3 (if any)
+    images = None
+    if data.upload_ids:
+        from app.api.endpoints.assistant_chat import _fetch_upload_images
+        images = await _fetch_upload_images(data.upload_ids, account_id, db) or None
+
     try:
         ai_content, suggested_flows, session = await unified_chat_service.send_chat_message(
             session_id=session_id,
@@ -287,6 +293,7 @@ async def send_chat_message(
             account_id=account_id,
             message=data.message,
             db=db,
+            images=images,
         )
     except ValueError as e:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
diff --git a/backend/app/schemas/ai_session.py b/backend/app/schemas/ai_session.py
index 3def117a..973de98c 100644
--- a/backend/app/schemas/ai_session.py
+++ b/backend/app/schemas/ai_session.py
@@ -245,6 +245,7 @@ class ChatSessionCreateResponse(BaseModel):
 class ChatMessageRequest(BaseModel):
     """Send a message in a chat session."""
     message: str = Field(..., min_length=1, max_length=8000)
+    upload_ids: list[UUID] = Field(default_factory=list, max_length=10)
 
 
 class ChatMessageResponse(BaseModel):
diff --git a/backend/app/services/unified_chat_service.py b/backend/app/services/unified_chat_service.py
index 83aca3e5..e04054fc 100644
--- a/backend/app/services/unified_chat_service.py
+++ b/backend/app/services/unified_chat_service.py
@@ -57,9 +57,14 @@ async def send_chat_message(
     account_id: UUID,
     message: str,
     db: AsyncSession,
+    images: list[dict[str, Any]] | None = None,
 ) -> tuple[str, list[dict[str, Any]], AISession]:
     """Send a message in a chat session and get AI response.
 
+    Args:
+        images: Optional list of {"media_type": str, "data": str (base64)}
+                for vision content attached to this message.
+
     Returns (ai_content, suggested_flows, session).
     """
     result = await db.execute(
@@ -105,6 +110,7 @@ async def send_chat_message(
         rag_context=rag_context,
         history=ai_messages,
         new_message=message,
+        images=images,
     )
 
     # Append messages to conversation_messages