From 999eda9b5d750af6989ccbc9b4ef0a73f7ee8582 Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Tue, 2 Apr 2024 15:01:46 -0700
Subject: [PATCH] Return action screenshots from get task endpoint as well
 (#149)

---
 skyvern/forge/agent.py                     |  2 +-
 skyvern/forge/sdk/routes/agent_protocol.py | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py
index 8dcb5cdc0..7ed79e0bf 100644
--- a/skyvern/forge/agent.py
+++ b/skyvern/forge/agent.py
@@ -787,7 +787,7 @@ async def execute_task_webhook(self, task: Task, last_step: Step, api_key: str |
         if recording_artifact:
             recording_url = await app.ARTIFACT_MANAGER.get_share_link(recording_artifact)
 
-        # get the artifact of the last  screenshot and get the screenshot_url
+        # get the artifact of the last TASK_RESPONSE_ACTION_SCREENSHOT_COUNT screenshots and get the screenshot_url
         latest_action_screenshot_artifacts = await app.DATABASE.get_latest_n_artifacts(
             task_id=task.task_id,
             organization_id=task.organization_id,
diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py
index 3c34bfc70..d84346238 100644
--- a/skyvern/forge/sdk/routes/agent_protocol.py
+++ b/skyvern/forge/sdk/routes/agent_protocol.py
@@ -218,6 +218,27 @@ async def get_task(
     if recording_artifact:
         recording_url = await app.ARTIFACT_MANAGER.get_share_link(recording_artifact)
 
+    # get the artifact of the last  screenshot and get the screenshot_url
+    latest_action_screenshot_artifacts = await app.DATABASE.get_latest_n_artifacts(
+        task_id=task_obj.task_id,
+        organization_id=task_obj.organization_id,
+        artifact_types=[ArtifactType.SCREENSHOT_ACTION],
+        n=SettingsManager.get_settings().TASK_RESPONSE_ACTION_SCREENSHOT_COUNT,
+    )
+    latest_action_screenshot_urls = []
+    if latest_action_screenshot_artifacts:
+        for artifact in latest_action_screenshot_artifacts:
+            screenshot_url = await app.ARTIFACT_MANAGER.get_share_link(artifact)
+            if screenshot_url:
+                latest_action_screenshot_urls.append(screenshot_url)
+            else:
+                LOG.error(
+                    "Failed to get share link for action screenshot",
+                    artifact_id=artifact.artifact_id,
+                )
+    else:
+        LOG.error("Failed to get latest action screenshots")
+
     failure_reason = None
     if task_obj.status == TaskStatus.failed and (latest_step.output or task_obj.failure_reason):
         failure_reason = ""
@@ -231,6 +252,7 @@ async def get_task(
             )
 
     return task_obj.to_task_response(
+        action_screenshot_urls=latest_action_screenshot_urls,
         screenshot_url=screenshot_url,
         recording_url=recording_url,
         failure_reason=failure_reason,