Merge pull request #1063 from Dhruvkumar1-Microsoft/psl-WafImageIssue

fix: Resolve the fix for intermittent issue image not showing in final answer and also fixed the Image is not getting rendered in WAF deployment

Author: Roopan-Microsoft

src/App/src/api/config.tsx

@@ -85,6 +85,28 @@ export function getApiUrl() {
 
     return API_URL;
 }
+
+export function resolveApiAssetUrl(url: string): string {
+   if (!url) {
+       return url;
+   }
+   const marker = "/api/v4/images/";
+   const idx = url.indexOf(marker);
+   if (idx === -1) {
+       return url;
+   }
+   const path = url.slice(idx); // drop any scheme+host before the image path
+   const apiUrl = getApiUrl();
+   if (apiUrl && /^https?:\/\//i.test(apiUrl)) {
+       try {
+           return new URL(apiUrl).origin + path;
+       } catch {
+           return path;
+       }
+   }
+   return path;
+}
+
 export function getUserInfoGlobal() {
     if (!USER_INFO) {
         // Check if window.userInfo exists

src/App/src/commonComponents/modules/Chat.tsx

@@ -9,6 +9,7 @@ import {
   ToolbarDivider,
 } from "@fluentui/react-components";
 import { Copy, Send } from "../imports/bundleicons";
+import { resolveApiAssetUrl } from "@/api/config";
 import { ChatDismiss20Regular, HeartRegular } from "@fluentui/react-icons";
 import ChatInput from "./ChatInput";
 import "./Chat.css";
@@ -182,7 +183,7 @@ const Chat: React.FC<ChatProps> = ({
           <div key={index} className={`message ${msg.role}`}>
             <Body1>
               <div style={{ display: "flex", flexDirection: "column", whiteSpace: "pre-wrap", width: "100%" }}>
-                <ReactMarkdown remarkPlugins={[remarkGfm]} rehypePlugins={[rehypePrism]}>
+                <ReactMarkdown remarkPlugins={[remarkGfm]} rehypePlugins={[rehypePrism]} urlTransform={resolveApiAssetUrl}>
                   {msg.content}
                 </ReactMarkdown>
                 {msg.role === "assistant" && (

src/App/src/components/content/streaming/StreamingAgentMessage.tsx

@@ -8,6 +8,7 @@ import { TaskService } from "@/store";
 import { PersonRegular, ArrowDownloadRegular } from "@fluentui/react-icons";
 import { getAgentIcon, getAgentDisplayName } from '@/utils/agentIconUtils';
 import { formatJsonInText } from '@/utils/jsonFormatter';
+import { resolveApiAssetUrl } from "@/api/config";
 
 interface StreamingAgentMessageProps {
   agentMessages: AgentMessageData[];
@@ -205,7 +206,7 @@ const renderAgentMessages = (
                 <ReactMarkdown
                   remarkPlugins={[remarkGfm]}
                   rehypePlugins={[rehypePrism]}
-                  urlTransform={(url: string) => url}
+                  urlTransform={resolveApiAssetUrl}
                   components={{
                       a: ({ node: _node, ...props }) => (
                         <a

src/App/src/components/content/streaming/StreamingBufferMessage.tsx

@@ -7,6 +7,7 @@ import ReactMarkdown from "react-markdown";
 import remarkGfm from "remark-gfm";
 import rehypePrism from "rehype-prism";
 import { formatJsonInText } from "@/utils/jsonFormatter";
+import { resolveApiAssetUrl } from "@/api/config";
 
 interface StreamingBufferMessageProps {
     streamingMessageBuffer: string;
@@ -164,6 +165,7 @@ const StreamingBufferMessage: React.FC<StreamingBufferMessageProps> = ({
                                 <ReactMarkdown
                                     remarkPlugins={[remarkGfm]}
                                     rehypePlugins={[rehypePrism]}
+                                    urlTransform={resolveApiAssetUrl}
                                     components={{
                                         a: ({ node, ...props }) => (
                                             <a
@@ -217,6 +219,7 @@ const StreamingBufferMessage: React.FC<StreamingBufferMessageProps> = ({
                         <ReactMarkdown
                             remarkPlugins={[remarkGfm]}
                             rehypePlugins={[rehypePrism]}
+                            urlTransform={resolveApiAssetUrl}
                             components={{
                                         a: ({ node, ...props }) => (
                                             <a

src/backend/agents/image_agent.py

@@ -11,7 +11,6 @@
 
 import base64
 import logging
-import os
 import uuid
 from typing import Any, AsyncIterable, Awaitable
 
@@ -178,12 +177,11 @@ async def _invoke_stream(
             blob_name = await _upload_image_to_blob(png_bytes, image_id)
 
             if blob_name:
-                # Build the image URL pointing at the backend proxy endpoint
-                backend_base = (config.AZURE_AI_AGENT_ENDPOINT or "").rstrip("/")
-                backend_origin = os.environ.get("BACKEND_URL", "").rstrip("/")
-                if not backend_origin:
-                    backend_origin = backend_base
-                image_src = f"{backend_origin}/api/v4/images/{blob_name}"
+                # Relative backend proxy path. The browser resolves it against the
+                # SPA's configured API origin, so it works in both the standard
+                # deployment (public backend) and the WAF/private deployment (the
+                # frontend's same-origin reverse proxy forwards to the internal backend).
+                image_src = f"/api/v4/images/{blob_name}"
                 image_content = f"![Generated Marketing Image]({image_src})"
             else:
                 # Fallback: embed base64 directly

src/backend/orchestration/orchestration_manager.py

@@ -35,12 +35,18 @@
 apply_tool_history_leak_patch()
 
 _BARE_IMAGE_URL_RE = re.compile(
-    r"(?<![\(\]])"
-    r"(?<!\]\()"
-    r"(https?://[^\s)]+?"
-    r"(?:/api/v4/images/[^\s)]+?|[^\s)]+?\.(?:png|jpe?g|gif|webp)))"
-    r"(?=[\s)\]]|$)",
-    re.IGNORECASE,
+   r"(?<![\(\]])"
+   r"(?<!\]\()"
+   r"("
+   # Absolute image URL (any host, or a backend /api/v4/images path)
+   r"https?://[^\s)]+?(?:/api/v4/images/[^\s)]+?|[^\s)]+?\.(?:png|jpe?g|gif|webp))"
+   # Bare relative backend image path (emitted by the MCP/backend image tools).
+   # The (?<![^\s]) guard requires the path to start at whitespace/string-start so
+   # it never matches the same substring inside an absolute URL.
+   r"|(?<![^\s])/api/v4/images/[^\s)]+?\.(?:png|jpe?g|gif|webp)"
+   r")"
+   r"(?=[\s)\]]|$)",
+   re.IGNORECASE,
 )
 
 

src/backend/orchestration/plan_review_helpers.py

@@ -114,6 +114,9 @@ def get_magentic_prompt_kwargs(*, has_user_responses: bool = False) -> dict:
 - Compile ONLY from messages agents actually produced. Quote verbatim where appropriate.
 - Do NOT fabricate URLs, results, or content that no agent produced.
 - If a required agent step did not run, state it plainly — do not pretend it did.
+- If an agent produced an image (a markdown image ![alt](url) or an image URL such as one
+  under /api/v4/images/), embed it in your answer using markdown image syntax
+  ![description](url). NEVER present an image as a bare URL or a plain link.
 - Do NOT offer further help. Provide the answer and end with a polite closing.
 """
 

src/mcp_server/services/image_service.py

@@ -9,15 +9,13 @@
 import base64
 import logging
 import uuid
-from datetime import datetime, timedelta, timezone
+
 
 import httpx
 from azure.identity import DefaultAzureCredential, ManagedIdentityCredential, get_bearer_token_provider
 from azure.storage.blob import (
-    BlobSasPermissions,
-    BlobServiceClient,
-    ContentSettings,
-    generate_blob_sas,
+   BlobServiceClient,
+   ContentSettings,
 )
 
 from config.settings import config
@@ -26,7 +24,7 @@
 logger = logging.getLogger(__name__)
 
 _IMAGE_API_VERSION = "2025-04-01-preview"
-_SAS_VALIDITY_DAYS = 7
+
 
 
 def _get_credential():
@@ -48,11 +46,14 @@ def _ensure_container(blob_service: BlobServiceClient, container_name: str) -> N
 
 
 def _upload_png_and_get_url(png_bytes: bytes) -> str:
-    """Upload PNG bytes to blob storage, return an accessible URL.
+    """Upload PNG bytes to blob storage, return a relative backend-proxy path.
+
 
-    Prefers the backend image-proxy URL (BACKEND_URL/api/v4/images/{blob}) so
-    the browser never needs direct blob access.  Falls back to a user-delegation
-    SAS URL, or bare blob URL as a last resort.
+    Returns "/api/v4/images/{blob}" — a relative path the browser resolves
+    against the SPA's configured API origin. The backend proxy reads the blob
+    with its own credential, so the browser never needs direct blob access (no
+    SAS or public blob access required), and it works in both the standard and
+    WAF/private-networking deployments.
     """
     if not config.azure_storage_blob_url:
         raise RuntimeError("AZURE_STORAGE_BLOB_URL is not configured on the MCP server")
@@ -72,33 +73,16 @@ def _upload_png_and_get_url(png_bytes: bytes) -> str:
         content_settings=ContentSettings(content_type="image/png"),
     )
 
-    # Prefer backend proxy URL — the browser fetches via the backend which has
-    # its own credential for blob access; no SAS or public access needed.
-    if config.backend_url:
-        backend_origin = config.backend_url.rstrip("/")
-        return f"{backend_origin}/api/v4/images/{blob_name}"
 
-    # Fallback: generate a user-delegation SAS URL for direct blob access.
-    blob_url = f"{account_url}/{container_name}/{blob_name}"
-    try:
-        now = datetime.now(timezone.utc)
-        delegation_key = blob_service.get_user_delegation_key(
-            key_start_time=now - timedelta(minutes=5),
-            key_expiry_time=now + timedelta(days=_SAS_VALIDITY_DAYS),
-        )
-        sas = generate_blob_sas(
-            account_name=blob_service.account_name,
-            container_name=container_name,
-            blob_name=blob_name,
-            user_delegation_key=delegation_key,
-            permission=BlobSasPermissions(read=True),
-            expiry=now + timedelta(days=_SAS_VALIDITY_DAYS),
-            start=now - timedelta(minutes=5),
-        )
-        return f"{blob_url}?{sas}"
-    except Exception as exc:  # noqa: BLE001
-        logger.warning("Failed to generate user-delegation SAS, returning bare URL: %s", exc)
-        return blob_url
+    # Return a relative backend proxy path. The browser resolves it against the
+    # SPA's configured API origin, so it works in both the standard deployment
+    # (resolves to the public backend) and the WAF/private deployment (resolves
+    # to the frontend's same-origin reverse proxy, which forwards to the internal
+    # backend over the VNet). The backend proxy reads the blob with its own
+    # credential, so no SAS or public blob access is needed.
+    return f"/api/v4/images/{blob_name}"
+
+
 
 
 class ImageService(MCPToolBase):