Fix public ExecutionContext attribute types

.agents/setup

@@ -1,6 +1,16 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
+if ! git lfs version >/dev/null 2>&1; then
+  echo "Installing Git LFS"
+  sudo apt-get update
+  sudo apt-get install -y git-lfs
+fi
+
+echo "Fetching Git LFS fixtures"
+git lfs install --local
+git lfs pull
+
 if ! command -v uv >/dev/null 2>&1; then
   python -m pip install --user uv
   export PATH="$HOME/.local/bin:$PATH"

tilebox-workflows/tests/runner/test_runner.py

@@ -5,18 +5,37 @@
 import sys
 from pathlib import Path
 from unittest.mock import MagicMock, patch
+from uuid import uuid4
 
 import pytest
 
 from _tilebox.grpc.replay import open_recording_channel, open_replay_channel
 from tilebox.workflows import ExecutionContext, Runner, Task
 from tilebox.workflows.cache import InMemoryCache, JobCache
 from tilebox.workflows.client import Client
-from tilebox.workflows.data import JobState, ProgressIndicator, RunnerContext, TaskState
+from tilebox.workflows.data import JobState, ProgressIndicator, RunnerContext, TaskIdentifier, TaskState
+from tilebox.workflows.data import Task as TaskData
 from tilebox.workflows.observability.tracing import NoopWorkflowTracer
+from tilebox.workflows.runner.executor import ExecutionContext as RunnerExecutionContext
 from tilebox.workflows.runner.task_runner import TaskRunner
 
 
+def test_public_execution_context_attributes() -> None:
+    class RenameTask(Task):
+        def execute(self, context: ExecutionContext) -> None:
+            context.current_task.display = "Processing"
+            context.job_cache["result"] = b"processed"
+
+    task = TaskData(uuid4(), TaskIdentifier("RenameTask", "v0.0"), display="Queued")
+    cache = InMemoryCache()
+    context = RunnerExecutionContext(MagicMock(), task, cache)
+
+    RenameTask().execute(context)
+
+    assert task.display == "Processing"
+    assert cache["result"] == b"processed"
+
+
 def test_task_authoring_imports_are_lazy() -> None:
     code = (
         "import sys\n"
@@ -45,7 +64,7 @@ class FibonacciTask(Task):
 
     async def execute(self, context: ExecutionContext) -> None:
         await asyncio.sleep(0)
-        cache: JobCache = context.job_cache  # ty: ignore[unresolved-attribute]
+        cache: JobCache = context.job_cache
         key = f"fib_{self.n}"
         if f"fib_{self.n}" in cache:
             # If the result is already in the cache, we can skip the calculation
@@ -66,7 +85,7 @@ class SumResultTask(Task):
     n: int
 
     def execute(self, context: ExecutionContext) -> None:
-        cache: JobCache = context.job_cache  # ty: ignore[unresolved-attribute]
+        cache: JobCache = context.job_cache
         fib_n_1 = bytes_to_int(cache[f"fib_{self.n - 1}"])
         fib_n_2 = bytes_to_int(cache[f"fib_{self.n - 2}"])
 
@@ -96,7 +115,7 @@ async def test_runner_with_fibonacci_workflow() -> None:
 class FlakyTask(Task):
     async def execute(self, context: ExecutionContext) -> None:
         await asyncio.sleep(0)
-        cache: JobCache = context.job_cache  # ty: ignore[unresolved-attribute]
+        cache: JobCache = context.job_cache
         if "succeed" in cache:
             return  # finally succeed
 
@@ -288,14 +307,14 @@ def execute(self, context: ExecutionContext) -> None:
 
 class FailingTask(Task):
     def execute(self, context: ExecutionContext) -> None:
-        cache = context.job_cache  # ty: ignore[unresolved-attribute]
+        cache = context.job_cache
         cache["failing_task"] = b"1"  # to make sure it actually ran
         raise ValueError("This task always fails")
 
 
 class SucceedingTask(Task):
     def execute(self, context: ExecutionContext) -> None:
-        cache = context.job_cache  # ty: ignore[unresolved-attribute]
+        cache = context.job_cache
         cache["succeeding_task"] = b"1"  # to make sure it actually ran
 
 

tilebox-workflows/tests/runner/test_worker_concurrency.py

@@ -47,7 +47,7 @@ async def execute(self, context: ExecutionContext) -> None:
                     (self.label, id(self), id(context), id(context.runner_context), id(asyncio.get_running_loop()))
                 )
 
-            cache: JobCache = context.job_cache  # ty: ignore[unresolved-attribute]
+            cache: JobCache = context.job_cache
             cache[self.label] = self.label.encode()
             context.logger.info("Concurrent task executing", label=self.label)
             context.progress(self.label).add(1)

tilebox-workflows/tilebox/workflows/runner/executor.py

@@ -9,7 +9,7 @@
 from contextlib import AbstractContextManager, contextmanager
 from contextvars import copy_context
 from threading import RLock
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 from uuid import UUID
 from warnings import warn
 
@@ -26,8 +26,14 @@
 from tilebox.workflows.observability.logging import StructuredLogger
 from tilebox.workflows.observability.tracing import NoopWorkflowTracer, WorkflowTracer, start_job_span
 from tilebox.workflows.runner.runner import Runner
+from tilebox.workflows.task import (
+    CurrentTask,
+    FutureTask,
+    ProgressUpdate,
+    RunnerContext,
+    merge_future_tasks_to_submissions,
+)
 from tilebox.workflows.task import ExecutionContext as ExecutionContextBase
-from tilebox.workflows.task import FutureTask, ProgressUpdate, RunnerContext, merge_future_tasks_to_submissions
 from tilebox.workflows.task import Task as TaskInstance
 
 if TYPE_CHECKING:
@@ -133,7 +139,8 @@ def execute_task(
 class ExecutionContext(ExecutionContextBase):
     def __init__(self, executor: TaskExecutor, task: Task, job_cache: JobCache) -> None:
         self._executor = executor
-        self.current_task = task
+        # Executing tasks have a job, unlike general Task data objects.
+        self.current_task = cast(CurrentTask, task)
         self.job_cache = job_cache
         self._sub_tasks: list[FutureTask] = []
         self._progress_indicators: dict[str | None, ProgressUpdate] = {}

tilebox-workflows/tilebox/workflows/task.py

@@ -6,7 +6,8 @@
 from contextlib import suppress
 from dataclasses import dataclass, fields, is_dataclass
 from types import NoneType, UnionType
-from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast, get_args, get_origin
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast, get_args, get_origin
+from uuid import UUID
 
 # from python 3.11 onwards this is available as typing.dataclass_transform:
 from typing_extensions import dataclass_transform
@@ -15,6 +16,7 @@
 from tilebox.workflows.data import RunnerContext, TaskIdentifier, TaskSubmissionGroup, TaskSubmissions
 
 if TYPE_CHECKING:
+    from tilebox.workflows.cache import JobCache
     from tilebox.workflows.observability.logging import StructuredLogger
     from tilebox.workflows.observability.tracing import WorkflowTracer
 else:
@@ -375,9 +377,31 @@ def done(self, count: int) -> None:
         self._done += count
 
 
+class CurrentJob(Protocol):
+    """Read-only job information available during task execution."""
+
+    @property
+    def id(self) -> UUID: ...
+
+    @property
+    def name(self) -> str: ...
+
+
+class CurrentTask(Protocol):
+    """Task information available during execution, with an editable display label."""
+
+    display: str | None
+
+    @property
+    def job(self) -> CurrentJob: ...
+
+
 class ExecutionContext(ABC):
     """The execution context for a task."""
 
+    current_task: CurrentTask
+    job_cache: "JobCache"
+
     @abstractmethod
     def submit_subtask(
         self,