diff --git a/autogpt_platform/backend/backend/blocks/agent.py b/autogpt_platform/backend/backend/blocks/agent.py index afbe410e4291..427de51106d8 100644 --- a/autogpt_platform/backend/backend/blocks/agent.py +++ b/autogpt_platform/backend/backend/blocks/agent.py @@ -76,7 +76,11 @@ def run(self, input_data: Input, **kwargs) -> BlockOutput: ) if not event.node_id: - if event.status in [ExecutionStatus.COMPLETED, ExecutionStatus.FAILED]: + if event.status in [ + ExecutionStatus.COMPLETED, + ExecutionStatus.TERMINATED, + ExecutionStatus.FAILED, + ]: logger.info(f"Execution {log_id} ended with status {event.status}") break else: diff --git a/autogpt_platform/backend/backend/data/execution.py b/autogpt_platform/backend/backend/data/execution.py index 4775d94553df..230c8e497034 100644 --- a/autogpt_platform/backend/backend/data/execution.py +++ b/autogpt_platform/backend/backend/data/execution.py @@ -270,9 +270,9 @@ async def update_graph_execution_start_time(graph_exec_id: str): async def update_graph_execution_stats( graph_exec_id: str, + status: ExecutionStatus, stats: dict[str, Any], ) -> ExecutionResult: - status = ExecutionStatus.FAILED if stats.get("error") else ExecutionStatus.COMPLETED res = await AgentGraphExecution.prisma().update( where={"id": graph_exec_id}, data={ diff --git a/autogpt_platform/backend/backend/executor/manager.py b/autogpt_platform/backend/backend/executor/manager.py index bae06ee70107..046da905a11c 100644 --- a/autogpt_platform/backend/backend/executor/manager.py +++ b/autogpt_platform/backend/backend/executor/manager.py @@ -597,7 +597,7 @@ def on_graph_execution( node_eid="*", block_name="-", ) - timing_info, (exec_stats, error) = cls._on_graph_execution( + timing_info, (exec_stats, status, error) = cls._on_graph_execution( graph_exec, cancel, log_metadata ) exec_stats["walltime"] = timing_info.wall_time @@ -605,6 +605,7 @@ def on_graph_execution( exec_stats["error"] = str(error) if error else None result = cls.db_client.update_graph_execution_stats( graph_exec_id=graph_exec.graph_exec_id, + status=status, stats=exec_stats, ) cls.db_client.send_execution_update(result) @@ -616,11 +617,12 @@ def _on_graph_execution( graph_exec: GraphExecutionEntry, cancel: threading.Event, log_metadata: LogMetadata, - ) -> tuple[dict[str, Any], Exception | None]: + ) -> tuple[dict[str, Any], ExecutionStatus, Exception | None]: """ Returns: - The execution statistics of the graph execution. - The error that occurred during the execution. + dict: The execution statistics of the graph execution. + ExecutionStatus: The final status of the graph execution. + Exception | None: The error that occurred during the execution, if any. """ log_metadata.info(f"Start graph execution {graph_exec.graph_exec_id}") exec_stats = { @@ -665,8 +667,7 @@ def callback(result: object): while not queue.empty(): if cancel.is_set(): - error = RuntimeError("Execution is cancelled") - return exec_stats, error + return exec_stats, ExecutionStatus.TERMINATED, error exec_data = queue.get() @@ -696,8 +697,7 @@ def callback(result: object): ) for node_id, execution in list(running_executions.items()): if cancel.is_set(): - error = RuntimeError("Execution is cancelled") - return exec_stats, error + return exec_stats, ExecutionStatus.TERMINATED, error if not queue.empty(): break # yield to parent loop to execute new queue items @@ -716,7 +716,12 @@ def callback(result: object): finished = True cancel.set() cancel_thread.join() - return exec_stats, error + + return ( + exec_stats, + ExecutionStatus.FAILED if error else ExecutionStatus.COMPLETED, + error, + ) class ExecutionManager(AppService): @@ -882,11 +887,8 @@ def cancel_execution(self, graph_exec_id: str) -> None: ExecutionStatus.COMPLETED, ExecutionStatus.FAILED, ): - self.db_client.upsert_execution_output( - node_exec.node_exec_id, "error", "TERMINATED" - ) exec_update = self.db_client.update_execution_status( - node_exec.node_exec_id, ExecutionStatus.FAILED + node_exec.node_exec_id, ExecutionStatus.TERMINATED ) self.db_client.send_execution_update(exec_update) diff --git a/autogpt_platform/backend/backend/util/test.py b/autogpt_platform/backend/backend/util/test.py index a471ce68183d..09a26b9f80c4 100644 --- a/autogpt_platform/backend/backend/util/test.py +++ b/autogpt_platform/backend/backend/util/test.py @@ -65,6 +65,9 @@ async def is_execution_completed(): if status == ExecutionStatus.FAILED: log.info("Execution failed") raise Exception("Execution failed") + if status == ExecutionStatus.TERMINATED: + log.info("Execution terminated") + raise Exception("Execution terminated") return status == ExecutionStatus.COMPLETED # Wait for the executions to complete diff --git a/autogpt_platform/backend/migrations/20250103143207_add_terminated_execution_status/migration.sql b/autogpt_platform/backend/migrations/20250103143207_add_terminated_execution_status/migration.sql new file mode 100644 index 000000000000..0b3a2a27b9be --- /dev/null +++ b/autogpt_platform/backend/migrations/20250103143207_add_terminated_execution_status/migration.sql @@ -0,0 +1,2 @@ +-- Add "TERMINATED" to execution status enum type +ALTER TYPE "AgentExecutionStatus" ADD VALUE 'TERMINATED'; diff --git a/autogpt_platform/backend/schema.prisma b/autogpt_platform/backend/schema.prisma index 7cd0ed36f4b5..ea1865791be0 100644 --- a/autogpt_platform/backend/schema.prisma +++ b/autogpt_platform/backend/schema.prisma @@ -216,6 +216,7 @@ enum AgentExecutionStatus { QUEUED RUNNING COMPLETED + TERMINATED FAILED } @@ -638,4 +639,4 @@ enum APIKeyStatus { ACTIVE REVOKED SUSPENDED -} \ No newline at end of file +} diff --git a/autogpt_platform/frontend/src/components/CustomNode.tsx b/autogpt_platform/frontend/src/components/CustomNode.tsx index 2ff382879938..410093204826 100644 --- a/autogpt_platform/frontend/src/components/CustomNode.tsx +++ b/autogpt_platform/frontend/src/components/CustomNode.tsx @@ -848,8 +848,10 @@ export function CustomNode({ data.status === "COMPLETED", "border-yellow-600 bg-yellow-600 text-white": data.status === "RUNNING", - "border-red-600 bg-red-600 text-white": - data.status === "FAILED", + "border-red-600 bg-red-600 text-white": [ + "FAILED", + "TERMINATED", + ].includes(data.status || ""), "border-blue-600 bg-blue-600 text-white": data.status === "QUEUED", "border-gray-600 bg-gray-600 font-black": diff --git a/autogpt_platform/frontend/src/hooks/useAgentGraph.ts b/autogpt_platform/frontend/src/hooks/useAgentGraph.ts index cd1fe4747444..e275be505bca 100644 --- a/autogpt_platform/frontend/src/hooks/useAgentGraph.ts +++ b/autogpt_platform/frontend/src/hooks/useAgentGraph.ts @@ -558,8 +558,9 @@ export default function useAgentGraph( return; } if ( - nodeResult.status != "COMPLETED" && - nodeResult.status != "FAILED" + !["COMPLETED", "TERMINATED", "FAILED"].includes( + nodeResult.status, + ) ) { pendingNodeExecutions.add(nodeResult.node_exec_id); } else { diff --git a/autogpt_platform/frontend/src/lib/autogpt-server-api/types.ts b/autogpt_platform/frontend/src/lib/autogpt-server-api/types.ts index 9872ace436dc..c21a93b55a69 100644 --- a/autogpt_platform/frontend/src/lib/autogpt-server-api/types.ts +++ b/autogpt_platform/frontend/src/lib/autogpt-server-api/types.ts @@ -196,7 +196,7 @@ export type GraphExecution = { ended_at: number; duration: number; total_run_time: number; - status: "INCOMPLETE" | "QUEUED" | "RUNNING" | "COMPLETED" | "FAILED"; + status: "QUEUED" | "RUNNING" | "COMPLETED" | "TERMINATED" | "FAILED"; graph_id: string; graph_version: number; }; @@ -246,7 +246,13 @@ export type NodeExecutionResult = { node_exec_id: string; node_id: string; block_id: string; - status: "INCOMPLETE" | "QUEUED" | "RUNNING" | "COMPLETED" | "FAILED"; + status: + | "INCOMPLETE" + | "QUEUED" + | "RUNNING" + | "COMPLETED" + | "TERMINATED" + | "FAILED"; input_data: { [key: string]: any }; output_data: { [key: string]: Array }; add_time: Date;