feat: extend capability gate to all task types

#74 gated dispatch on a single SSH capability. Generalize it: each
Executor declares the task types it serves, the worker advertises the
union over the executors it initialized as WorkerCapabilities.supported_task_types,
and the dispatcher routes a task only to workers advertising its type.
A worker missing an executor (e.g. no Docker for SSH, or absent GPU/training
deps) is no longer handed a task it would fail inside the runner.

Signed-off-by: Noppanat Wadlom <noppanat.wad@gmail.com>

Author: kaiitunnz

docs/ARCHITECTURE.md

@@ -127,12 +127,12 @@ scripts/dev/            compile_protos, sync_requirements, check_env_examples
   `WorkerHardware`. The dispatcher's `_cached_worker_candidates` filters
   to workers whose cache covers the task's references; entries older
   than `WORKER_CACHE_TTL_SEC` are ignored.
-- **Worker capabilities.** Beyond hardware fit, each worker advertises a
-  `WorkerCapabilities` describing the task kinds it can service. The dispatcher
-  routes a task that requires a capability only to workers advertising it, and
-  excludes a worker that advertises none from those tasks rather than failing
-  them on it. SSH is one such capability — a worker advertises it when it can
-  launch session containers.
+- **Worker capabilities.** Beyond hardware fit, each worker advertises the set
+  of task types it can service, and the dispatcher routes a task only to workers
+  that advertise its type. A worker advertises a type only when its executor came
+  up — e.g. SSH requires a reachable Docker daemon, and training or omni types
+  require their (often GPU-only) dependencies — so a worker missing that executor
+  isn't a candidate, rather than being handed a task it would fail.
 - **Cursor pagination.** List endpoints accept `limit` and `before` /
   `after` cursors. The cursor is an opaque base64 of `(timestamp, id)`;
   do not parse client-side.

sdk/src/flowmesh/models/workers.py

@@ -4,6 +4,8 @@
 
 from pydantic import BaseModel, Field, field_validator
 
+from .common import TaskType
+
 
 class CPUInfo(BaseModel):
     logical_cores: int | None = None
@@ -80,7 +82,7 @@ class SSHLimits(BaseModel):
 
 
 class WorkerCapabilities(BaseModel):
-    ssh: bool = False
+    supported_task_types: frozenset[TaskType] = Field(default_factory=frozenset)
 
 
 class Worker(BaseModel):

src/server/registries/worker.py

@@ -550,9 +550,7 @@ def hw_satisfies(worker: Worker, task: TaskEnvelope) -> bool:
 
 
 def capability_satisfies(worker: Worker, task: TaskEnvelope) -> bool:
-    if isinstance(task.spec, (SSHSpecStrict, SSHSpecTemplate)):
-        return worker.capabilities.ssh
-    return True
+    return task.spec.taskType in worker.capabilities.supported_task_types
 
 
 def _gpu_meets_requirements(hw: WorkerHardware, gpu_req: GPURequirements) -> bool:

src/shared/schemas/worker.py

@@ -2,6 +2,8 @@
 
 from pydantic import BaseModel, Field
 
+from shared.tasks.task_type import TaskType
+
 
 class WorkerStatus(StrEnum):
     UNKNOWN = "UNKNOWN"
@@ -33,8 +35,9 @@ class SSHLimits(BaseModel):
 class WorkerCapabilities(BaseModel):
     """Task capabilities a worker advertises to the dispatcher."""
 
-    ssh: bool = Field(
-        default=False, description="Whether the worker can run SSH session tasks."
+    supported_task_types: frozenset[TaskType] = Field(
+        default_factory=frozenset,
+        description="Types of tasks this worker can service.",
     )
 
 

src/worker/README.md

@@ -119,10 +119,11 @@ The SSH session container is labelled with `flowmesh.ssh.task_id` and
 `flowmesh.ssh.worker_id` so the server can clean it up if the worker
 container is stopped externally.
 
-At startup the worker initializes the SSH executor only if the Docker daemon is
-reachable, and advertises its `ssh` capability accordingly; the dispatcher then
-selects it for SSH tasks only when it can actually run them. (The supervisor
-mounts the Docker socket only when SSH is enabled for the node; see
+At startup the worker advertises the task types it can service, and the
+dispatcher only routes a task to a worker that advertises its type. A worker
+advertises a type only when its executor came up: SSH requires a reachable
+Docker daemon, so `ssh` is offered only on Docker-capable workers. (The
+supervisor mounts the Docker socket only when SSH is enabled for the node; see
 `ENABLE_SSH_BY_DEFAULT` below.)
 
 Relevant env vars for SSH tasks:

src/worker/executors/agent_executor.py

@@ -19,6 +19,7 @@
 from shared.schemas.artifact import ArtifactRef
 from shared.schemas.result import BaseExecutorResult
 from shared.tasks.specs import AgentSpecStrict
+from shared.tasks.task_type import TaskType
 
 from .base_executor import ExecutionError, Executor, ExecutorTask
 from .utils.checkpoints import maybe_upload_artifacts, write_executor_result
@@ -66,6 +67,7 @@ class AgentExecutor(Executor):
     """Agent executor using youtu-agent (utu) framework"""
 
     name = "agent"
+    supported_task_types = frozenset({TaskType.AGENT})
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)

src/worker/executors/api_executor.py

@@ -9,6 +9,7 @@
 
 from shared.schemas.result import BaseExecutorResult
 from shared.tasks.specs import ApiSpecStrict
+from shared.tasks.task_type import TaskType
 
 from .base_executor import ExecutionError, Executor, ExecutorTask
 
@@ -40,6 +41,7 @@ class APIExecutor(Executor):
     """
 
     name = "api"
+    supported_task_types = frozenset({TaskType.API})
 
     # ---- Class-level connection pool (shared across all instances) ----
     _clients: ClassVar[dict[_ClientKey, httpx.Client]] = {}

src/worker/executors/base_executor.py

@@ -29,11 +29,12 @@ def run(self, task: ExecutorTask, out_dir: Path) -> MyResult:
 import json
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import Any, ClassVar, TypeVar
 
 from shared.schemas.result import BaseExecutorResult
 from shared.tasks import MergedChildTaskStrict
 from shared.tasks.specs import TaskSpecStrictBase
+from shared.tasks.task_type import TaskType
 from shared.tasks.worker_message import WorkerHardware, WorkerTaskMessage
 from worker.config import WorkerConfig
 from worker.lifecycle import Lifecycle
@@ -67,8 +68,10 @@ class Executor(ABC):
     Subclasses must implement `run` and may override `prepare` and `teardown`.
     """
 
-    #: Human-readable identifier for logging/telemetry
     name: str = "executor"
+    """Human-readable identifier for logging/telemetry"""
+    supported_task_types: ClassVar[frozenset[TaskType]] = frozenset()
+    """Types of tasks this executor can service"""
 
     def __init__(
         self,
@@ -172,6 +175,7 @@ class EchoResult(BaseExecutorResult):
 
 class EchoExecutor(Executor):
     name = "echo"
+    supported_task_types = frozenset({TaskType.ECHO})
 
     def run(self, task: ExecutorTask, out_dir: Path) -> EchoResult:
         return EchoResult(

src/worker/executors/data_profiling_executor.py

@@ -10,6 +10,7 @@
 
 from shared.schemas.result import BaseExecutorResult
 from shared.tasks.specs import DataProfilingSpecStrict
+from shared.tasks.task_type import TaskType
 from shared.utils.json import to_json_serializable, validate_keys
 
 from ..connectors import get_connector_from_spec
@@ -31,6 +32,7 @@ class DataProfilingExecutor(DataMixin, Executor):
     """Executor that estimates SQL query costs by sampling SQL template params."""
 
     name = "data_profiling"
+    supported_task_types = frozenset({TaskType.DATA_PROFILING})
 
     def run(self, task: ExecutorTask, out_dir: Path) -> DataProfilingResult:
         spec = self.require_spec(task, DataProfilingSpecStrict)

src/worker/executors/data_retrieval_executor.py

@@ -19,6 +19,7 @@
 from shared.schemas.artifact import ArtifactRef
 from shared.schemas.result import BaseExecutorResult
 from shared.tasks.specs import DataRetrievalSpecStrict
+from shared.tasks.task_type import TaskType
 from shared.utils.json import validate_keys
 
 from ..connectors import LumidDataConnector, PostgreSQLConnector, S3Connector
@@ -40,6 +41,7 @@ class DataRetrievalResult(BaseExecutorResult):
 
 class DataRetrievalExecutor(DataMixin, Executor):
     name = "data_retrieval"
+    supported_task_types = frozenset({TaskType.DATA_RETRIEVAL})
 
     def run(self, task: ExecutorTask, out_dir: Path) -> DataRetrievalResult:
         spec = self.require_spec(task, DataRetrievalSpecStrict)