core/execution/context.py

"""
V11 Execution Context
---------------------

Central runtime object passed into every task.

Responsibilities:
- Hold inputs
- Share memory between tasks
- Store outputs
- Track execution metadata
- Provide filesystem helpers
- Provide logging helpers
"""

from __future__ import annotations

import os
import uuid
import tempfile
from typing import Any, Dict, Optional


# =========================================================
# Context Object
# =========================================================

class ExecutionContext:
    """
    Standard runtime context used by ALL tasks.

    Every task receives:
        async def run(ctx: ExecutionContext)
    """

    # -----------------------------------------------------
    # INIT
    # -----------------------------------------------------
    def __init__(
        self,
        task_name: str,
        inputs: Optional[Dict[str, Any]] = None,
        workspace: Optional[str] = None,
    ):

        self.task_name = task_name
        self.job_id = str(uuid.uuid4())

        self.inputs: Dict[str, Any] = inputs or {}
        self.outputs: Dict[str, Any] = {}
        self.memory: Dict[str, Any] = {}

        self.status: str = "created"
        self.error: Optional[str] = None

        self.workspace = workspace or self._create_workspace()

    # -----------------------------------------------------
    # WORKSPACE
    # -----------------------------------------------------
    def _create_workspace(self) -> str:
        path = tempfile.mkdtemp(prefix="basyx_job_")
        return path

    def path(self, filename: str) -> str:
        """
        Safe workspace path helper
        """
        return os.path.join(self.workspace, filename)

    # -----------------------------------------------------
    # INPUT HELPERS
    # -----------------------------------------------------
    def get(self, key: str, default=None):
        return self.inputs.get(key, default)

    def require(self, key: str):
        if key not in self.inputs:
            raise ValueError(f"Missing required input: {key}")
        return self.inputs[key]

    # -----------------------------------------------------
    # OUTPUT HELPERS
    # -----------------------------------------------------
    def set_output(self, key: str, value: Any):
        self.outputs[key] = value

    def result(self) -> Dict[str, Any]:
        return {
            "job_id": self.job_id,
            "task": self.task_name,
            "status": self.status,
            "outputs": self.outputs,
            "error": self.error,
        }

    # -----------------------------------------------------
    # MEMORY (cross-task sharing)
    # -----------------------------------------------------
    def remember(self, key: str, value: Any):
        """
        Save value for downstream tasks.
        """
        self.memory[key] = value

    def recall(self, key: str, default=None):
        return self.memory.get(key, default)

    # -----------------------------------------------------
    # STATUS MANAGEMENT
    # -----------------------------------------------------
    def mark_running(self):
        self.status = "running"

    def mark_complete(self):
        self.status = "completed"

    def mark_failed(self, error: Exception | str):
        self.status = "failed"
        self.error = str(error)

    # -----------------------------------------------------
    # LOGGING
    # -----------------------------------------------------
    def log(self, message: str):
        print(f"[{self.task_name} | {self.job_id}] {message}")

    # -----------------------------------------------------
    # SERIALIZATION
    # -----------------------------------------------------
    def to_dict(self):
        return {
            "job_id": self.job_id,
            "task_name": self.task_name,
            "inputs": self.inputs,
            "outputs": self.outputs,
            "memory": self.memory,
            "status": self.status,
            "error": self.error,
            "workspace": self.workspace,
        }


# =========================================================
# Context Factory
# =========================================================

def create_context(task_name: str, inputs: Dict[str, Any]) -> ExecutionContext:
    """
    Standardized factory used by executor.
    """
    return ExecutionContext(
        task_name=task_name,
        inputs=inputs,
    )