fix(tarko): fallback to beforeActionImage in afterAction strategy to prevent flickering (#1352)

Author: ulivz

multimodal/tarko/agent-web-ui/src/standalone/workspace/renderers/browser-control/useScreenshots.ts

@@ -65,35 +65,32 @@ export const useScreenshots = ({
 
     let foundBeforeImage = false;
     let foundAfterImage = false;
+    let beforeImageData: { url: string; pageUrl: string | null } | null = null;
+    let afterImageData: { url: string; pageUrl: string | null } | null = null;
 
     // Search for screenshots BEFORE the current tool call
-    if (currentStrategy === 'beforeAction' || currentStrategy === 'both') {
-      for (let i = currentToolCallIndex - 1; i >= 0; i--) {
-        const msg = sessionMessages[i];
-        if (msg.role === 'environment' && Array.isArray(msg.content)) {
-          const imgContent = msg.content.find(
-            (c) => typeof c === 'object' && 'type' in c && c.type === 'image_url',
-          );
+    // Always search for before action image as it may be used as fallback
+    for (let i = currentToolCallIndex - 1; i >= 0; i--) {
+      const msg = sessionMessages[i];
+      if (msg.role === 'environment' && Array.isArray(msg.content)) {
+        const imgContent = msg.content.find(
+          (c) => typeof c === 'object' && 'type' in c && c.type === 'image_url',
+        );
 
-          if (imgContent && 'image_url' in imgContent && imgContent.image_url.url) {
-            const url =
-              msg.metadata?.type === 'screenshot' && 'url' in msg.metadata
-                ? msg.metadata.url
-                : null;
-            setBeforeActionImage(imgContent.image_url.url);
-            setBeforeActionImageUrl(url || null);
-            if (currentStrategy === 'beforeAction') {
-              setRelatedImage(imgContent.image_url.url);
-              setRelatedImageUrl(url || null);
-            }
-            foundBeforeImage = true;
-            break;
-          }
+        if (imgContent && 'image_url' in imgContent && imgContent.image_url.url) {
+          const url =
+            msg.metadata?.type === 'screenshot' && 'url' in msg.metadata ? msg.metadata.url : null;
+          beforeImageData = { url: imgContent.image_url.url, pageUrl: url };
+          setBeforeActionImage(imgContent.image_url.url);
+          setBeforeActionImageUrl(url || null);
+          foundBeforeImage = true;
+          break;
         }
       }
     }
 
     // Search for screenshots AFTER the current tool call
+    // Always search for after action image as it may be needed for afterAction or both strategies
     if (currentStrategy === 'afterAction' || currentStrategy === 'both') {
       for (let i = currentToolCallIndex + 1; i < sessionMessages.length; i++) {
         const msg = sessionMessages[i];
@@ -107,41 +104,56 @@ export const useScreenshots = ({
               msg.metadata?.type === 'screenshot' && 'url' in msg.metadata
                 ? msg.metadata.url
                 : null;
+            afterImageData = { url: imgContent.image_url.url, pageUrl: url };
             setAfterActionImage(imgContent.image_url.url);
             setAfterActionImageUrl(url || null);
-            if (currentStrategy === 'afterAction') {
-              setRelatedImage(imgContent.image_url.url);
-              setRelatedImageUrl(url || null);
-            }
             foundAfterImage = true;
             break;
           }
         }
       }
     }
 
-    // Handle strategy-specific warnings and fallbacks
+    // Handle strategy-specific logic and fallbacks
     if (!environmentImage) {
-      if (currentStrategy === 'beforeAction' && !foundBeforeImage) {
-        console.warn(
-          `[BrowserControlRenderer] No valid screenshot found before toolCallId: ${toolCallId}. Clearing screenshot display.`,
-        );
-        setRelatedImage(null);
-        setRelatedImageUrl(null);
-      } else if (currentStrategy === 'afterAction' && !foundAfterImage) {
-        console.warn(
-          `[BrowserControlRenderer] No valid screenshot found after toolCallId: ${toolCallId}. Clearing screenshot display.`,
-        );
-        setRelatedImage(null);
-        setRelatedImageUrl(null);
+      if (currentStrategy === 'beforeAction') {
+        if (beforeImageData) {
+          setRelatedImage(beforeImageData.url);
+          setRelatedImageUrl(beforeImageData.pageUrl);
+        } else {
+          console.warn(
+            `[BrowserControlRenderer] No valid screenshot found before toolCallId: ${toolCallId}. Clearing screenshot display.`,
+          );
+          setRelatedImage(null);
+          setRelatedImageUrl(null);
+        }
+      } else if (currentStrategy === 'afterAction') {
+        if (afterImageData) {
+          // Use after action image when available
+          setRelatedImage(afterImageData.url);
+          setRelatedImageUrl(afterImageData.pageUrl);
+        } else if (beforeImageData) {
+          // Fallback to before action image to prevent flickering
+          console.warn(
+            `[BrowserControlRenderer] No valid screenshot found after toolCallId: ${toolCallId}. Falling back to before action image.`,
+          );
+          setRelatedImage(beforeImageData.url);
+          setRelatedImageUrl(beforeImageData.pageUrl);
+        } else {
+          console.warn(
+            `[BrowserControlRenderer] No valid screenshots found for toolCallId: ${toolCallId}. Clearing screenshot display.`,
+          );
+          setRelatedImage(null);
+          setRelatedImageUrl(null);
+        }
       } else if (currentStrategy === 'both') {
         // For 'both' strategy, use the after action image as primary if available
-        if (foundAfterImage) {
-          setRelatedImage(afterActionImage);
-          setRelatedImageUrl(afterActionImageUrl);
-        } else if (foundBeforeImage) {
-          setRelatedImage(beforeActionImage);
-          setRelatedImageUrl(beforeActionImageUrl);
+        if (afterImageData) {
+          setRelatedImage(afterImageData.url);
+          setRelatedImageUrl(afterImageData.pageUrl);
+        } else if (beforeImageData) {
+          setRelatedImage(beforeImageData.url);
+          setRelatedImageUrl(beforeImageData.pageUrl);
         } else {
           console.warn(
             `[BrowserControlRenderer] No valid screenshots found for toolCallId: ${toolCallId}. Clearing screenshot display.`,