fix(backend): Include webhook block executions in GraphExecution queries

autogpt_platform/backend/backend/data/block.py

@@ -1,3 +1,4 @@
+import functools
 import inspect
 from abc import ABC, abstractmethod
 from enum import Enum
@@ -8,6 +9,7 @@
     Generator,
     Generic,
     Optional,
+    Sequence,
     Type,
     TypeVar,
     cast,
@@ -523,3 +525,21 @@ async def initialize_blocks() -> None:
 def get_block(block_id: str) -> Block[BlockSchema, BlockSchema] | None:
     cls = get_blocks().get(block_id)
     return cls() if cls else None
+
+
+@functools.cache
+def get_webhook_block_ids() -> Sequence[str]:
+    return [
+        id
+        for id, B in get_blocks().items()
+        if B().block_type in (BlockType.WEBHOOK, BlockType.WEBHOOK_MANUAL)
+    ]
+
+
+@functools.cache
+def get_io_block_ids() -> Sequence[str]:
+    return [
+        id
+        for id, B in get_blocks().items()
+        if B().block_type in (BlockType.INPUT, BlockType.OUTPUT)
+    ]

autogpt_platform/backend/backend/data/execution.py

@@ -38,12 +38,19 @@
 from backend.util import type as type_utils
 from backend.util.settings import Config
 
-from .block import BlockInput, BlockType, CompletedBlockOutput, get_block
+from .block import (
+    BlockInput,
+    BlockType,
+    CompletedBlockOutput,
+    get_block,
+    get_io_block_ids,
+    get_webhook_block_ids,
+)
 from .db import BaseDbModel
 from .includes import (
     EXECUTION_RESULT_INCLUDE,
-    GRAPH_EXECUTION_INCLUDE,
     GRAPH_EXECUTION_INCLUDE_WITH_NODES,
+    graph_execution_include,
 )
 from .model import CredentialsMetaInput, GraphExecutionStats, NodeExecutionStats
 from .queue import AsyncRedisEventBus, RedisEventBus
@@ -411,7 +418,9 @@ async def get_graph_execution(
         include=(
             GRAPH_EXECUTION_INCLUDE_WITH_NODES
             if include_node_executions
-            else GRAPH_EXECUTION_INCLUDE
+            else graph_execution_include(
+                [*get_io_block_ids(), *get_webhook_block_ids()]
+            )
         ),
     )
     if not execution:
@@ -563,7 +572,9 @@ async def update_graph_execution_start_time(
             "executionStatus": ExecutionStatus.RUNNING,
             "startedAt": datetime.now(tz=timezone.utc),
         },
-        include=GRAPH_EXECUTION_INCLUDE,
+        include=graph_execution_include(
+            [*get_io_block_ids(), *get_webhook_block_ids()]
+        ),
     )
     return GraphExecution.from_db(res) if res else None
 
@@ -598,7 +609,9 @@ async def update_graph_execution_stats(
 
     graph_exec = await AgentGraphExecution.prisma().find_unique_or_raise(
         where={"id": graph_exec_id},
-        include=GRAPH_EXECUTION_INCLUDE,
+        include=graph_execution_include(
+            [*get_io_block_ids(), *get_webhook_block_ids()]
+        ),
     )
     return GraphExecution.from_db(graph_exec)
 

autogpt_platform/backend/backend/data/includes.py

@@ -1,10 +1,8 @@
-from typing import cast
+from typing import Sequence, cast
 
 import prisma.enums
 import prisma.types
 
-from backend.blocks.io import IO_BLOCK_IDs
-
 AGENT_NODE_INCLUDE: prisma.types.AgentNodeInclude = {
     "Input": True,
     "Output": True,
@@ -42,18 +40,26 @@
     }
 }
 
-GRAPH_EXECUTION_INCLUDE: prisma.types.AgentGraphExecutionInclude = {
-    "NodeExecutions": {
-        **cast(
-            prisma.types.FindManyAgentNodeExecutionArgsFromAgentGraphExecution,
-            GRAPH_EXECUTION_INCLUDE_WITH_NODES["NodeExecutions"],
-        ),
-        "where": {
-            "Node": {"is": {"AgentBlock": {"is": {"id": {"in": IO_BLOCK_IDs}}}}},
-            "NOT": [{"executionStatus": prisma.enums.AgentExecutionStatus.INCOMPLETE}],
-        },
+
+def graph_execution_include(
+    include_block_ids: Sequence[str],
+) -> prisma.types.AgentGraphExecutionInclude:
+    return {
+        "NodeExecutions": {
+            **cast(
+                prisma.types.FindManyAgentNodeExecutionArgsFromAgentGraphExecution,
+                GRAPH_EXECUTION_INCLUDE_WITH_NODES["NodeExecutions"],  # type: ignore
+            ),
+            "where": {
+                "Node": {
+                    "is": {"AgentBlock": {"is": {"id": {"in": include_block_ids}}}}
+                },
+                "NOT": [
+                    {"executionStatus": prisma.enums.AgentExecutionStatus.INCOMPLETE}
+                ],
+            },
+        }
     }
-}
 
 
 INTEGRATION_WEBHOOK_INCLUDE: prisma.types.IntegrationWebhookInclude = {