fix(agents-api): Fix task execution logical errors

agents-api/agents_api/activities/task_steps/transition_step.py

@@ -4,11 +4,10 @@
 from temporalio import activity
 
 from ...autogen.openapi_model import CreateTransitionRequest, Transition
+from ...clients.cozo import get_cozo_client
 from ...common.protocol.tasks import StepContext
 from ...env import testing
-from ...models.execution.create_execution_transition import (
-    create_execution_transition as create_execution_transition_query,
-)
+from ...models.execution.create_execution_transition import create_execution_transition
 
 
 @beartype
@@ -24,23 +23,26 @@ async def transition_step(
         transition_info.task_token = task_token
 
     # Create transition
-    transition = create_execution_transition_query(
-        developer_id=context.execution_input.developer_id,
-        execution_id=context.execution_input.execution.id,
-        task_id=context.execution_input.task.id,
-        data=transition_info,
-        update_execution_status=True,
-    )
+    try:
+        transition = create_execution_transition(
+            developer_id=context.execution_input.developer_id,
+            execution_id=context.execution_input.execution.id,
+            task_id=context.execution_input.task.id,
+            data=transition_info,
+            update_execution_status=True,
+        )
+
+    except Exception as exc:
+        client = get_cozo_client()
+        print(client)
+        # breakpoint()
+
+        raise exc
 
     return transition
 
-async def mock_transition_step(
-    context: StepContext,
-    transition_info: CreateTransitionRequest,
-) -> None:
-    # Does nothing
-    return None
 
+mock_transition_step = transition_step
 
 transition_step = activity.defn(name="transition_step")(
     transition_step if not testing else mock_transition_step

agents-api/agents_api/clients/cozo.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict
+from typing import Dict
 
 from pycozo.client import Client
 
@@ -10,7 +10,7 @@
     options.update({"auth": cozo_auth})
 
 
-def get_cozo_client() -> Any:
+def get_cozo_client() -> Client:
     client = getattr(app.state, "cozo_client", Client("http", options=options))
     if not hasattr(app.state, "cozo_client"):
         app.state.cozo_client = client

agents-api/agents_api/common/protocol/tasks.py

@@ -26,17 +26,51 @@
     WorkflowStep,
 )
 
-### NOTE: Here, "init" is NOT a real state, but a placeholder for the start state of the state machine
+# State Machine
+#
+# init -> wait | error | step | cancelled | init_branch | finish
+# init_branch -> wait | error | step | cancelled | finish_branch
+# wait -> resume | error | cancelled
+# resume -> wait | error | cancelled | step | finish | finish_branch | init_branch
+# step -> wait | error | cancelled | step | finish | finish_branch | init_branch
+# finish_branch -> wait | error | cancelled | step | finish | init_branch
+# error ->
+
+## Mermaid Diagram
+# ```mermaid
+# ---
+# title: Execution state machine
+# ---
+# stateDiagram-v2
+#     [*] --> queued
+#     queued --> starting
+#     queued --> cancelled
+#     starting --> cancelled
+#     starting --> failed
+#     starting --> running
+#     running --> running
+#     running --> awaiting_input
+#     running --> cancelled
+#     running --> failed
+#     running --> succeeded
+#     awaiting_input --> running
+#     awaiting_input --> cancelled
+#     cancelled --> [*]
+#     succeeded --> [*]
+#     failed --> [*]
+
+# ```
+# TODO: figure out how to type this
 valid_transitions: dict[TransitionType, list[TransitionType]] = {
     # Start state
-    "init": ["wait", "error", "step", "cancelled", "init_branch"],
-    "init_branch": ["wait", "error", "step", "cancelled"],
+    "init": ["wait", "error", "step", "cancelled", "init_branch", "finish"],
+    "init_branch": ["wait", "error", "step", "cancelled", "finish_branch"],
     # End states
     "finish": [],
     "error": [],
     "cancelled": [],
     # Intermediate states
-    "wait": ["resume", "error", "cancelled"],
+    "wait": ["resume", "cancelled"],
     "resume": [
         "wait",
         "error",
@@ -59,8 +93,13 @@
 }  # type: ignore
 
 valid_previous_statuses: dict[ExecutionStatus, list[ExecutionStatus]] = {
-    "running": ["queued", "starting", "awaiting_input"],
+    "running": ["starting", "awaiting_input", "running"],
+    "starting": ["queued"],
+    "queued": [],
+    "awaiting_input": ["starting", "running"],
     "cancelled": ["queued", "starting", "awaiting_input", "running"],
+    "succeeded": ["starting", "running"],
+    "failed": ["starting", "running"],
 }  # type: ignore
 
 transition_to_execution_status: dict[TransitionType | None, ExecutionStatus] = {
@@ -130,6 +169,11 @@ def is_last_step(self) -> Annotated[bool, Field(exclude=True)]:
     def is_first_step(self) -> Annotated[bool, Field(exclude=True)]:
         return self.cursor.step == 0
 
+    @computed_field
+    @property
+    def is_main(self) -> Annotated[bool, Field(exclude=True)]:
+        return self.cursor.workflow == "main"
+
     def model_dump(self, *args, **kwargs) -> dict[str, Any]:
         dump = super().model_dump(*args, **kwargs)
         dump["_"] = self.current_input

agents-api/agents_api/models/execution/create_execution_transition.py

@@ -134,7 +134,8 @@ def create_execution_transition(
     ?[valid] :=
         matched[prev_transitions],
         found = length(prev_transitions),
-        valid = assert(found > 0, "Invalid transition"),
+        valid = if($next_type == "init", found == 0, found > 0),
+        assert(valid, "Invalid transition"),
     """
 
     # Prepare the insert query