agent-platform/services/runtime-service/app/schemas/run.py

from datetime import datetime
from typing import TYPE_CHECKING

from core_domain import (
    InitialNodeContract,
    NodeExecutionRequestContract,
    NodeRunContract,
    NodeRunStatusUpdateContract,
    RunBootstrapContract,
    RunCreateContract,
    RunExecutionRequestContract,
    WorkflowRunContract,
    WorkflowRunStatusUpdateContract,
)
from core_shared import JSONValue
from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from app.db.models import ExecutionLog, NodeArtifact, NodeRun, TraceSpan, WorkflowRun


class InitialNodeCreateRequest(InitialNodeContract):
    pass


class RunCreateRequest(RunCreateContract):
    initial_node: InitialNodeCreateRequest | None = None


class WorkflowRunListRequest(BaseModel):
    session_id: str | None = None
    limit: int = Field(default=50, ge=1, le=500)


class NodeRunListRequest(BaseModel):
    run_id: str


class ExecutionLogListRequest(BaseModel):
    run_id: str | None = None
    node_run_id: str | None = None


class NodeArtifactListRequest(BaseModel):
    run_id: str | None = None
    node_run_id: str | None = None
    artifact_type: str | None = None


class TraceSpanListRequest(BaseModel):
    run_id: str | None = None
    node_run_id: str | None = None
    span_type: str | None = None


class RuntimeDebugSnapshotRequest(BaseModel):
    run_id: str


class WorkflowRunResponse(WorkflowRunContract):

    @classmethod
    def from_entity(cls, entity: "WorkflowRun") -> "WorkflowRunResponse":
        return cls.model_validate(entity, from_attributes=True)


class NodeRunResponse(NodeRunContract):

    @classmethod
    def from_entity(cls, entity: "NodeRun") -> "NodeRunResponse":
        return cls.model_validate(entity, from_attributes=True)


class RunBootstrapResponse(RunBootstrapContract):
    run: WorkflowRunResponse
    initial_node: NodeRunResponse | None = None


class WorkflowRunStatusUpdateRequest(WorkflowRunStatusUpdateContract):
    pass


class NodeRunStatusUpdateRequest(NodeRunStatusUpdateContract):
    pass


class NodeRunExecuteRequest(NodeExecutionRequestContract):
    pass


class NodeRunExecuteResponse(BaseModel):
    run: WorkflowRunResponse
    node_run: NodeRunResponse
    executor_name: str


class RunExecuteRequest(RunExecutionRequestContract):
    pass


class RunExecuteResponse(BaseModel):
    run: WorkflowRunResponse
    node_runs: list[NodeRunResponse]
    executor_names: list[str]


class RuntimeDebugContinueRequest(BaseModel):
    worker_key: str | None = None
    max_steps: int = Field(default=32, ge=1, le=500)
    breakpoint_node_ids: list[str] = Field(default_factory=list)


class WorkerExecuteNextRequest(BaseModel):
    worker_key: str
    lease_seconds: int | None = Field(default=None, gt=0)


class WorkerExecuteNextResponse(BaseModel):
    run: WorkflowRunResponse
    node_run: NodeRunResponse
    executor_name: str
    released_lease_count: int = 0


class HumanNodeResumeRequest(BaseModel):
    human_task_id: str
    worker_key: str | None = None


class ExecutionLogResponse(BaseModel):
    id: str
    run_id: str
    node_run_id: str | None = None
    event_type: str
    level: str
    message: str
    detail_json: dict[str, JSONValue] | None = None
    created_time: datetime

    @classmethod
    def from_entity(cls, entity: "ExecutionLog") -> "ExecutionLogResponse":
        return cls.model_validate(entity, from_attributes=True)


class NodeArtifactResponse(BaseModel):
    id: str
    run_id: str
    node_run_id: str
    node_id: str
    artifact_type: str
    name: str
    mime_type: str | None = None
    content_text: str | None = None
    content_json: dict[str, JSONValue] | None = None
    storage_uri: str | None = None
    size_bytes: int | None = None
    created_time: datetime

    @classmethod
    def from_entity(cls, entity: "NodeArtifact") -> "NodeArtifactResponse":
        return cls.model_validate(entity, from_attributes=True)


class TraceSpanResponse(BaseModel):
    id: str
    run_id: str
    node_run_id: str | None = None
    parent_span_id: str | None = None
    span_type: str
    name: str
    status: str
    started_time: datetime
    ended_time: datetime | None = None
    duration_ms: int | None = None
    attributes_json: dict[str, JSONValue] | None = None
    error_code: str | None = None
    error_message: str | None = None
    created_time: datetime

    @classmethod
    def from_entity(cls, entity: "TraceSpan") -> "TraceSpanResponse":
        return cls.model_validate(entity, from_attributes=True)


class RuntimeDebugSnapshotResponse(BaseModel):
    run: WorkflowRunResponse
    node_runs: list[NodeRunResponse]
    run_state_json: dict[str, JSONValue]
    node_output_json_by_node_id: dict[str, dict[str, JSONValue]]
    node_output_text_by_node_id: dict[str, str]
    queued_node_ids: list[str]
    running_node_ids: list[str]
    completed_node_ids: list[str]
    failed_node_ids: list[str]
    execution_logs: list[ExecutionLogResponse]
    node_artifacts: list[NodeArtifactResponse]
    trace_spans: list[TraceSpanResponse]


class RuntimeDebugStepResponse(BaseModel):
    snapshot: RuntimeDebugSnapshotResponse
    executed_node_runs: list[NodeRunResponse]
    executor_names: list[str]
    paused_before_node_id: str | None = None
    reason: str