fix: apply size check to text files too (OOM protection)

- Remove redundant inline comment on text_only param
- Simplify file filtering logic per review suggestion

autogpt_platform/backend/backend/api/features/chat/tools/agent_output.py

@@ -5,7 +5,7 @@ import re
 from datetime import datetime, timedelta, timezone
 from typing import Any
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, field_validator
 
 from backend.api.features.chat.model import ChatSession
 from backend.api.features.library import db as library_db
@@ -14,7 +14,6 @@ from backend.data import execution as execution_db
 from backend.data.execution import ExecutionStatus, GraphExecution, GraphExecutionMeta
 
 from .base import BaseTool
-from .execution_utils import TERMINAL_STATUSES, wait_for_execution
 from .models import (
     AgentOutputResponse,
     ErrorResponse,
@@ -35,7 +34,6 @@ class AgentOutputInput(BaseModel):
     store_slug: str = ""
     execution_id: str = ""
     run_time: str = "latest"
-    wait_if_running: int = Field(default=0, ge=0, le=300)
 
     @field_validator(
         "agent_name",
@@ -119,11 +117,6 @@ class AgentOutputTool(BaseTool):
         Select which run to retrieve using:
         - execution_id: Specific execution ID
         - run_time: 'latest' (default), 'yesterday', 'last week', or ISO date 'YYYY-MM-DD'
-
-        Wait for completion (optional):
-        - wait_if_running: Max seconds to wait if execution is still running (0-300).
-          If the execution is running/queued, waits up to this many seconds for it to complete.
-          Returns current status on timeout. If already finished, returns immediately.
         """
 
     @property
@@ -153,13 +146,6 @@ class AgentOutputTool(BaseTool):
                         "Time filter: 'latest', 'yesterday', 'last week', or 'YYYY-MM-DD'"
                     ),
                 },
-                "wait_if_running": {
-                    "type": "integer",
-                    "description": (
-                        "Max seconds to wait if execution is still running (0-300). "
-                        "If running, waits for completion. Returns current state on timeout."
-                    ),
-                },
             },
             "required": [],
         }
@@ -237,14 +223,10 @@ class AgentOutputTool(BaseTool):
         execution_id: str | None,
         time_start: datetime | None,
         time_end: datetime | None,
-        include_running: bool = False,
     ) -> tuple[GraphExecution | None, list[GraphExecutionMeta], str | None]:
         """
         Fetch execution(s) based on filters.
         Returns (single_execution, available_executions_meta, error_message).
-
-        Args:
-            include_running: If True, also look for running/queued executions (for waiting)
         """
         # If specific execution_id provided, fetch it directly
         if execution_id:
@@ -257,22 +239,11 @@ class AgentOutputTool(BaseTool):
                 return None, [], f"Execution '{execution_id}' not found"
             return execution, [], None
 
-        # Determine which statuses to query
+        # Get completed executions with time filters
-        statuses = [ExecutionStatus.COMPLETED]
-        if include_running:
-            statuses.extend(
-                [
-                    ExecutionStatus.RUNNING,
-                    ExecutionStatus.QUEUED,
-                    ExecutionStatus.INCOMPLETE,
-                ]
-            )
-
-        # Get executions with time filters
         executions = await execution_db.get_graph_executions(
             graph_id=graph_id,
             user_id=user_id,
-            statuses=statuses,
+            statuses=[ExecutionStatus.COMPLETED],
             created_time_gte=time_start,
             created_time_lte=time_end,
             limit=10,
@@ -339,28 +310,10 @@ class AgentOutputTool(BaseTool):
                 for e in available_executions[:5]
             ]
 
-        # Build appropriate message based on execution status
+        message = f"Found execution outputs for agent '{agent.name}'"
-        if execution.status == ExecutionStatus.COMPLETED:
-            message = f"Found execution outputs for agent '{agent.name}'"
-        elif execution.status == ExecutionStatus.FAILED:
-            message = f"Execution for agent '{agent.name}' failed"
-        elif execution.status == ExecutionStatus.TERMINATED:
-            message = f"Execution for agent '{agent.name}' was terminated"
-        elif execution.status in (
-            ExecutionStatus.RUNNING,
-            ExecutionStatus.QUEUED,
-            ExecutionStatus.INCOMPLETE,
-        ):
-            message = (
-                f"Execution for agent '{agent.name}' is still {execution.status.value}. "
-                "Results may be incomplete. Use wait_if_running to wait for completion."
-            )
-        else:
-            message = f"Found execution for agent '{agent.name}' (status: {execution.status.value})"
-
         if len(available_executions) > 1:
             message += (
-                f" Showing latest of {len(available_executions)} matching executions."
+                f". Showing latest of {len(available_executions)} matching executions."
             )
 
         return AgentOutputResponse(
@@ -475,17 +428,13 @@ class AgentOutputTool(BaseTool):
         # Parse time expression
         time_start, time_end = parse_time_expression(input_data.run_time)
 
-        # Check if we should wait for running executions
+        # Fetch execution(s)
-        wait_timeout = input_data.wait_if_running
-
-        # Fetch execution(s) - include running if we're going to wait
         execution, available_executions, exec_error = await self._get_execution(
             user_id=user_id,
             graph_id=agent.graph_id,
             execution_id=input_data.execution_id or None,
             time_start=time_start,
             time_end=time_end,
-            include_running=wait_timeout > 0,
         )
 
         if exec_error:
@@ -494,17 +443,4 @@ class AgentOutputTool(BaseTool):
                 session_id=session_id,
             )
 
-        # If we have an execution that's still running and we should wait
-        if execution and wait_timeout > 0 and execution.status not in TERMINAL_STATUSES:
-            logger.info(
-                f"Execution {execution.id} is {execution.status}, "
-                f"waiting up to {wait_timeout}s for completion"
-            )
-            execution = await wait_for_execution(
-                user_id=user_id,
-                graph_id=agent.graph_id,
-                execution_id=execution.id,
-                timeout_seconds=wait_timeout,
-            )
-
         return self._build_response(agent, execution, available_executions, session_id)

autogpt_platform/backend/backend/api/features/chat/tools/execution_utils.py

@@ -1,124 +0,0 @@
-"""Shared utilities for execution waiting and status handling."""
-
-import asyncio
-import logging
-from typing import Any
-
-from backend.data import execution as execution_db
-from backend.data.execution import (
-    AsyncRedisExecutionEventBus,
-    ExecutionStatus,
-    GraphExecution,
-    GraphExecutionEvent,
-)
-
-logger = logging.getLogger(__name__)
-
-# Terminal statuses that indicate execution is complete
-TERMINAL_STATUSES = frozenset(
-    {
-        ExecutionStatus.COMPLETED,
-        ExecutionStatus.FAILED,
-        ExecutionStatus.TERMINATED,
-    }
-)
-
-
-async def wait_for_execution(
-    user_id: str,
-    graph_id: str,
-    execution_id: str,
-    timeout_seconds: int,
-) -> GraphExecution | None:
-    """
-    Wait for an execution to reach a terminal status using Redis pubsub.
-
-    Uses asyncio.wait_for to ensure timeout is respected even when no events
-    are received.
-
-    Args:
-        user_id: User ID
-        graph_id: Graph ID
-        execution_id: Execution ID to wait for
-        timeout_seconds: Max seconds to wait
-
-    Returns:
-        The execution with current status, or None if not found
-    """
-    # First check current status - maybe it's already done
-    execution = await execution_db.get_graph_execution(
-        user_id=user_id,
-        execution_id=execution_id,
-        include_node_executions=False,
-    )
-    if not execution:
-        return None
-
-    # If already in terminal state, return immediately
-    if execution.status in TERMINAL_STATUSES:
-        logger.debug(
-            f"Execution {execution_id} already in terminal state: {execution.status}"
-        )
-        return execution
-
-    logger.info(
-        f"Waiting up to {timeout_seconds}s for execution {execution_id} "
-        f"(current status: {execution.status})"
-    )
-
-    # Subscribe to execution updates via Redis pubsub
-    event_bus = AsyncRedisExecutionEventBus()
-    channel_key = f"{user_id}/{graph_id}/{execution_id}"
-
-    try:
-        # Use wait_for to enforce timeout on the entire listen operation
-        result = await asyncio.wait_for(
-            _listen_for_terminal_status(event_bus, channel_key, user_id, execution_id),
-            timeout=timeout_seconds,
-        )
-        return result
-    except asyncio.TimeoutError:
-        logger.info(f"Timeout waiting for execution {execution_id}")
-    except Exception as e:
-        logger.error(f"Error waiting for execution: {e}", exc_info=True)
-
-    # Return current state on timeout/error
-    return await execution_db.get_graph_execution(
-        user_id=user_id,
-        execution_id=execution_id,
-        include_node_executions=False,
-    )
-
-
-async def _listen_for_terminal_status(
-    event_bus: AsyncRedisExecutionEventBus,
-    channel_key: str,
-    user_id: str,
-    execution_id: str,
-) -> GraphExecution | None:
-    """
-    Listen for execution events until a terminal status is reached.
-
-    This is a helper that gets wrapped in asyncio.wait_for for timeout handling.
-    """
-    async for event in event_bus.listen_events(channel_key):
-        # Only process GraphExecutionEvents (not NodeExecutionEvents)
-        if isinstance(event, GraphExecutionEvent):
-            logger.debug(f"Received execution update: {event.status}")
-            if event.status in TERMINAL_STATUSES:
-                # Fetch full execution with outputs
-                return await execution_db.get_graph_execution(
-                    user_id=user_id,
-                    execution_id=execution_id,
-                    include_node_executions=False,
-                )
-
-    # Should not reach here normally (generator should yield indefinitely)
-    return None
-
-
-def get_execution_outputs(execution: GraphExecution | None) -> dict[str, Any] | None:
-    """Extract outputs from an execution, or return None."""
-    if execution is None:
-        return None
-    return execution.outputs

autogpt_platform/backend/backend/api/features/chat/tools/models.py

@@ -192,7 +192,6 @@ class ExecutionStartedResponse(ToolResponseBase):
     library_agent_id: str | None = None
     library_agent_link: str | None = None
     status: str = "QUEUED"
-    outputs: dict[str, Any] | None = None  # Populated when wait_for_result is used
 
 
 # Auth/error models

autogpt_platform/backend/backend/api/features/chat/tools/run_agent.py

@@ -12,7 +12,6 @@ from backend.api.features.chat.tracking import (
     track_agent_scheduled,
 )
 from backend.api.features.library import db as library_db
-from backend.data.execution import ExecutionStatus
 from backend.data.graph import GraphModel
 from backend.data.model import CredentialsMetaInput
 from backend.data.user import get_user_by_id
@@ -25,7 +24,6 @@ from backend.util.timezone_utils import (
 )
 
 from .base import BaseTool
-from .execution_utils import get_execution_outputs, wait_for_execution
 from .helpers import get_inputs_from_schema
 from .models import (
     AgentDetails,
@@ -72,7 +70,6 @@ class RunAgentInput(BaseModel):
     schedule_name: str = ""
     cron: str = ""
     timezone: str = "UTC"
-    wait_for_result: int = Field(default=0, ge=0, le=300)
 
     @field_validator(
         "username_agent_slug",
@@ -154,14 +151,6 @@ class RunAgentTool(BaseTool):
                     "type": "string",
                     "description": "IANA timezone for schedule (default: UTC)",
                 },
-                "wait_for_result": {
-                    "type": "integer",
-                    "description": (
-                        "Max seconds to wait for execution to complete (0-300). "
-                        "If >0, blocks until the execution finishes or times out. "
-                        "Returns execution outputs when complete."
-                    ),
-                },
             },
             "required": [],
         }
@@ -358,7 +347,6 @@ class RunAgentTool(BaseTool):
                     graph=graph,
                     graph_credentials=graph_credentials,
                     inputs=params.inputs,
-                    wait_for_result=params.wait_for_result,
                 )
 
         except NotFoundError as e:
@@ -442,9 +430,8 @@ class RunAgentTool(BaseTool):
         graph: GraphModel,
         graph_credentials: dict[str, CredentialsMetaInput],
         inputs: dict[str, Any],
-        wait_for_result: int = 0,
     ) -> ToolResponseBase:
-        """Execute an agent immediately, optionally waiting for completion."""
+        """Execute an agent immediately."""
         session_id = session.session_id
 
         # Check rate limits
@@ -481,60 +468,6 @@ class RunAgentTool(BaseTool):
         )
 
         library_agent_link = f"/library/agents/{library_agent.id}"
-
-        # If wait_for_result is specified, wait for execution to complete
-        if wait_for_result > 0:
-            logger.info(
-                f"Waiting up to {wait_for_result}s for execution {execution.id}"
-            )
-            result = await wait_for_execution(
-                user_id=user_id,
-                graph_id=library_agent.graph_id,
-                execution_id=execution.id,
-                timeout_seconds=wait_for_result,
-            )
-            final_status = result.status if result else ExecutionStatus.FAILED
-            outputs = get_execution_outputs(result)
-
-            # Build message based on final status
-            if final_status == ExecutionStatus.COMPLETED:
-                message = (
-                    f"Agent '{library_agent.name}' execution completed successfully. "
-                    f"{MSG_DO_NOT_RUN_AGAIN}"
-                )
-            elif final_status == ExecutionStatus.FAILED:
-                message = (
-                    f"Agent '{library_agent.name}' execution failed. "
-                    f"View details at {library_agent_link}. "
-                    f"{MSG_DO_NOT_RUN_AGAIN}"
-                )
-            elif final_status == ExecutionStatus.TERMINATED:
-                message = (
-                    f"Agent '{library_agent.name}' execution was terminated. "
-                    f"View details at {library_agent_link}. "
-                    f"{MSG_DO_NOT_RUN_AGAIN}"
-                )
-            else:
-                message = (
-                    f"Agent '{library_agent.name}' execution is still {final_status.value} "
-                    f"(timed out after {wait_for_result}s). "
-                    f"View at {library_agent_link}. "
-                    f"{MSG_DO_NOT_RUN_AGAIN}"
-                )
-
-            return ExecutionStartedResponse(
-                message=message,
-                session_id=session_id,
-                execution_id=execution.id,
-                graph_id=library_agent.graph_id,
-                graph_name=library_agent.name,
-                library_agent_id=library_agent.id,
-                library_agent_link=library_agent_link,
-                status=final_status.value,
-                outputs=outputs,
-            )
-
-        # Default: return immediately without waiting
         return ExecutionStartedResponse(
             message=(
                 f"Agent '{library_agent.name}' execution started successfully. "

autogpt_platform/backend/backend/blocks/claude_code.py

@@ -187,9 +187,11 @@ class ClaudeCodeBlock(Block):
         )
         files: list[SandboxFileOutput] = SchemaField(
             description=(
-                "List of text files created/modified by Claude Code during this execution. "
+                "List of files created/modified by Claude Code during this execution. "
+                "Includes text files and binary files (images, PDFs, etc.). "
                 "Each file has 'path', 'relative_path', 'name', 'content', and 'workspace_ref' fields. "
-                "workspace_ref contains a workspace:// URI if the file was stored to workspace."
+                "workspace_ref contains a workspace:// URI for workspace storage. "
+                "For binary files, content contains a placeholder; use workspace_ref to access the file."
             )
         )
         conversation_history: str = SchemaField(
@@ -452,13 +454,15 @@ class ClaudeCodeBlock(Block):
                 else:
                     new_conversation_history = turn_entry
 
-            # Extract files created/modified during this run and store to workspace
+            # Extract files created/modified during this run and store to workspace.
+            # Binary files (images, PDFs, etc.) are stored via store_media_file
+            # which handles virus scanning and workspace storage.
             sandbox_files = await extract_and_store_sandbox_files(
                 sandbox=sandbox,
                 working_directory=working_directory,
                 execution_context=execution_context,
                 since_timestamp=start_timestamp,
-                text_only=True,
+                text_only=False,
             )
 
             return (

autogpt_platform/backend/backend/util/sandbox_files.py

@@ -74,8 +74,50 @@ TEXT_EXTENSIONS = {
     ".tex",
     ".csv",
     ".log",
+    ".svg",  # SVG is XML-based text
 }
 
+# Binary file extensions we explicitly support extracting
+BINARY_EXTENSIONS = {
+    # Images
+    ".png",
+    ".jpg",
+    ".jpeg",
+    ".gif",
+    ".webp",
+    ".ico",
+    ".bmp",
+    ".tiff",
+    ".tif",
+    # Documents
+    ".pdf",
+    # Archives
+    ".zip",
+    ".tar",
+    ".gz",
+    ".7z",
+    # Audio
+    ".mp3",
+    ".wav",
+    ".ogg",
+    ".flac",
+    # Video
+    ".mp4",
+    ".webm",
+    ".mov",
+    ".avi",
+    # Fonts
+    ".woff",
+    ".woff2",
+    ".ttf",
+    ".otf",
+    ".eot",
+}
+
+# Maximum file size for binary extraction (50MB)
+# Prevents OOM from accidentally extracting huge files
+MAX_BINARY_FILE_SIZE = 50 * 1024 * 1024
+
 
 class SandboxFileOutput(BaseModel):
     """A file extracted from a sandbox and optionally stored in workspace."""
@@ -120,7 +162,8 @@ async def extract_sandbox_files(
         sandbox: The E2B sandbox instance
         working_directory: Directory to search for files
         since_timestamp: ISO timestamp - only return files modified after this time
-        text_only: If True, only extract text files (default). If False, extract all files.
+        text_only: If True, only extract text files. If False, also extract
+                   supported binary files (images, PDFs, etc.).
 
     Returns:
         List of ExtractedFile objects with path, content, and metadata
@@ -149,15 +192,48 @@ async def extract_sandbox_files(
             if not file_path:
                 continue
 
-            # Check if it's a text file
+            # Check file type (case-insensitive for extensions)
-            is_text = any(file_path.endswith(ext) for ext in TEXT_EXTENSIONS)
+            file_path_lower = file_path.lower()
+            is_text = any(
+                file_path_lower.endswith(ext.lower()) for ext in TEXT_EXTENSIONS
+            )
+            is_binary = any(
+                file_path_lower.endswith(ext.lower()) for ext in BINARY_EXTENSIONS
+            )
 
-            # Skip non-text files if text_only mode
+            # Skip files with unrecognized extensions
+            if not is_text and not is_binary:
+                continue
+
+            # In text_only mode, skip binary files
             if text_only and not is_text:
                 continue
 
             try:
-                # Read file content as bytes
+                # Check file size before reading to prevent OOM
+                stat_result = await sandbox.commands.run(
+                    f"stat -c %s {shlex.quote(file_path)} 2>/dev/null"
+                )
+                if stat_result.exit_code != 0 or not stat_result.stdout:
+                    logger.debug(f"Skipping {file_path}: could not determine file size")
+                    continue
+
+                try:
+                    file_size = int(stat_result.stdout.strip())
+                except ValueError:
+                    logger.debug(
+                        f"Skipping {file_path}: unexpected stat output "
+                        f"{stat_result.stdout.strip()!r}"
+                    )
+                    continue
+
+                if file_size > MAX_BINARY_FILE_SIZE:
+                    logger.info(
+                        f"Skipping {file_path}: size {file_size} bytes "
+                        f"exceeds limit {MAX_BINARY_FILE_SIZE}"
+                    )
+                    continue
+
                 content = await sandbox.files.read(file_path, format="bytes")
                 if isinstance(content, str):
                     content = content.encode("utf-8")

docs/integrations/block-integrations/claude_code.md

@@ -16,7 +16,7 @@ When activated, the block:
    - Install dependencies (npm, pip, etc.)
    - Run terminal commands
    - Build and test applications
-5. Extracts all text files created/modified during execution
+5. Extracts all files created/modified during execution (text files and binary files like images, PDFs, etc.)
 6. Returns the response and files, optionally keeping the sandbox alive for follow-up tasks
 
 The block supports conversation continuation through three mechanisms:
@@ -42,7 +42,7 @@ The block supports conversation continuation through three mechanisms:
 | Output | Description |
 |--------|-------------|
 | Response | The output/response from Claude Code execution |
-| Files | List of text files created/modified during execution. Each file includes path, relative_path, name, and content fields |
+| Files | List of files (text and binary) created/modified during execution. Includes images, PDFs, and other supported formats. Each file has path, relative_path, name, content, and workspace_ref fields. Binary files are stored in workspace and accessible via workspace_ref |
 | Conversation History | Full conversation history including this turn. Use to restore context on a fresh sandbox |
 | Session ID | Session ID for this conversation. Pass back with sandbox_id to continue the conversation |
 | Sandbox ID | ID of the sandbox instance (null if disposed). Pass back with session_id to continue the conversation |

docs/integrations/block-integrations/llm.md

@@ -535,7 +535,7 @@ When activated, the block:
 2. Installs the latest version of Claude Code in the sandbox
 3. Optionally runs setup commands to prepare the environment
 4. Executes your prompt using Claude Code, which can create/edit files, install dependencies, run terminal commands, and build applications
-5. Extracts all text files created/modified during execution
+5. Extracts all files created/modified during execution (text files and binary files like images, PDFs, etc.)
 6. Returns the response and files, optionally keeping the sandbox alive for follow-up tasks
 
 The block supports conversation continuation through three mechanisms:
@@ -563,7 +563,7 @@ The block supports conversation continuation through three mechanisms:
 |--------|-------------|------|
 | error | Error message if execution failed | str |
 | response | The output/response from Claude Code execution | str |
-| files | List of text files created/modified by Claude Code during this execution. Each file has 'path', 'relative_path', 'name', 'content', and 'workspace_ref' fields. workspace_ref contains a workspace:// URI if the file was stored to workspace. | List[SandboxFileOutput] |
+| files | List of files created/modified by Claude Code during this execution. Includes text files and binary files (images, PDFs, etc.). Each file has 'path', 'relative_path', 'name', 'content', and 'workspace_ref' fields. workspace_ref contains a workspace:// URI for workspace storage. For binary files, content contains a placeholder; use workspace_ref to access the file. | List[SandboxFileOutput] |
 | conversation_history | Full conversation history including this turn. Pass this to conversation_history input to continue on a fresh sandbox if the previous sandbox timed out. | str |
 | session_id | Session ID for this conversation. Pass this back along with sandbox_id to continue the conversation. | str |
 | sandbox_id | ID of the sandbox instance. Pass this back along with session_id to continue the conversation. This is None if dispose_sandbox was True (sandbox was disposed). | str |