fix: artifact_bucket

schedule/pyproject.toml

@@ -31,7 +31,12 @@ build-backend = "hatchling.build"
 [dependency-groups]
 dev = [
     "pytest>=9.1.1",
+    "pytest-asyncio>=1.4.0",
 ]
 
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests"]
+
 [tool.hatch.build.targets.wheel]
 packages = ["src/schedule"]

schedule/src/schedule/service.py

@@ -186,16 +186,25 @@ class SchedulerService:
         return await self.orchestrator.dispatch_run(run_id)
 
 
-def build_object_store() -> Any:
+def build_object_store(bucket_name: str | None = None) -> Any:
     """Construct the AsyncStorageBackend the worker reads version artifacts from.
 
     Respects ``settings.storage_backend``: in ``local`` mode points at the
     local ``${local_storage_base_dir}/version`` directory (the same place
-    ``backend.storage_api`` writes to), in ``s3`` mode points at
+    ``backend.storage_api`` writes to), in ``s3`` mode points at the
-    ``settings.s3_version_bucket``. The worker only reads version artifacts
+    bucket named by ``bucket_name`` (default ``settings.s3_version_bucket``).
-    via this store (run logs / results still go through the backend's
+
-    HTTP storage API), so pointing at the version bucket/directory is the
+    The bucket override exists because some workspaces configure a custom
-    correct resolution regardless of the artifact's workspace.
+    ``Workspaces.artifact_bucket``; their version artifacts are uploaded
+    into that bucket by the backend, so the worker must read from the
+    same place — not the global version bucket. Run logs / results still
+    go through the backend's HTTP storage API (which honors the override
+    on its own), so this factory only matters for the artifact-download
+    path.
+
+    For local mode the override is ignored: custom bucket subdirectories
+    only make sense in S3 mode, and the version base_dir is the right
+    destination either way.
     """
     if settings.storage_backend == "local":
         return create_storage(
@@ -211,7 +220,7 @@ def build_object_store() -> Any:
         {
             "type": "s3",
             "mode": "async",
-            "bucket": settings.s3_version_bucket,
+            "bucket": bucket_name or settings.s3_version_bucket,
             "endpoint_url": settings.s3_endpoint,
             "aws_access_key_id": settings.s3_access_key,
             "aws_secret_access_key": settings.s3_secret_key,

schedule/src/schedule/worker.py

@@ -58,6 +58,30 @@ class NodeExecutor:
         self.session_factory = session_factory
         self.object_store = object_store
         self.storage_client = storage_client
+        # Per-bucket object store cache. The injected ``object_store`` is
+        # the default version-bucket store; workspaces that override
+        # ``Workspaces.artifact_bucket`` need a store bound to that custom
+        # bucket (P0-3 fix). Build lazily so the common (no-override) path
+        # incurs no extra cost.
+        self._bucket_stores: dict[str, Any] = {
+            settings.s3_version_bucket: object_store,
+        }
+
+    def _store_for(self, bucket_name: str) -> Any:
+        """Return the AsyncStorageBackend bound to ``bucket_name``.
+
+        Caches per-bucket stores on first use; the default version bucket
+        always reuses the injected ``object_store`` so the common path
+        stays zero-allocation.
+        """
+        store = self._bucket_stores.get(bucket_name)
+        if store is not None:
+            return store
+        from schedule.service import build_object_store
+
+        store = build_object_store(bucket_name=bucket_name)
+        self._bucket_stores[bucket_name] = store
+        return store
 
     async def handle_node_execute(
         self,
@@ -380,7 +404,11 @@ class NodeExecutor:
         object_key: str,
         content_hash: str,
     ) -> bytes:
-        content = await self.object_store.get(object_key)
+        # Honor the artifact's actual bucket (P0-3 fix): the artifact may
+        # live in ``Workspaces.artifact_bucket`` rather than the global
+        # version bucket the default ``object_store`` is bound to.
+        store = self._store_for(bucket_name)
+        content = await store.get(object_key)
         if hashlib.sha256(content).hexdigest() != content_hash:
             raise ValueError("stable version artifact hash mismatch")
         logger.debug(

schedule/tests/test_worker.py

@@ -0,0 +1,154 @@
+"""Tests for NodeExecutor bucket-aware artifact download (P0-3).
+
+Pre-fix the worker's ``object_store`` was bound to the global version
+bucket at startup, so a workspace whose ``Workspaces.artifact_bucket``
+points at a custom S3 bucket would always 404 on download. The fix
+introduces a per-bucket store cache so the worker reads from whatever
+bucket the artifact actually lives in.
+
+These tests cover only the routing logic — no live S3 / MySQL.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+
+def _make_executor() -> tuple[SimpleNamespace, MagicMock, MagicMock]:
+    """Construct a NodeExecutor with the bucket-store cache pre-seeded.
+
+    Returns ``(executor, default_store, storage_client)``. The default
+    store is whatever the worker would receive via ``object_store=`` at
+    construction time; it must be returned untouched by ``_store_for`` for
+    the global version bucket.
+    """
+    from schedule.worker import NodeExecutor
+
+    default_store = MagicMock(name="default_store")
+    storage_client = MagicMock(name="storage_client")
+    executor = NodeExecutor(
+        session_factory=MagicMock(),
+        object_store=default_store,
+        storage_client=storage_client,
+    )
+    return executor, default_store, storage_client
+
+
+def test_store_for_default_bucket_returns_injected_store() -> None:
+    """The default version bucket must reuse the injected ``object_store``.
+
+    Otherwise the common (no-override) path would pay for an extra
+    ``create_storage`` call on every run — wasted work.
+    """
+    from common.config import settings
+
+    executor, default_store, _ = _make_executor()
+    store = executor._store_for(settings.s3_version_bucket)
+    assert store is default_store
+
+
+def test_store_for_custom_bucket_uses_build_factory(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A custom bucket_name must produce a store via ``build_object_store``."""
+    custom_store = MagicMock(name="custom_store")
+    with patch(
+        "schedule.service.build_object_store",
+        return_value=custom_store,
+    ) as mock_build:
+        executor, _, _ = _make_executor()
+        store = executor._store_for("my-workspace-artifacts")
+
+    assert store is custom_store
+    mock_build.assert_called_once_with(bucket_name="my-workspace-artifacts")
+
+
+def test_store_for_custom_bucket_is_cached(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Repeated lookups for the same custom bucket must hit the cache."""
+    custom_store = MagicMock(name="custom_store")
+    with patch(
+        "schedule.service.build_object_store",
+        return_value=custom_store,
+    ) as mock_build:
+        executor, _, _ = _make_executor()
+        executor._store_for("my-workspace-artifacts")
+        executor._store_for("my-workspace-artifacts")
+        executor._store_for("my-workspace-artifacts")
+
+    mock_build.assert_called_once_with(bucket_name="my-workspace-artifacts")
+
+
+@pytest.mark.asyncio
+async def test_download_artifact_uses_per_bucket_store() -> None:
+    """_download_artifact must call ``store.get(object_key)`` on the
+    bucket-bound store, NOT the default version-bucket store.
+
+    Pre-fix this would route every workspace's reads through the global
+    version bucket, silently 404'ing for any workspace that overrode
+    ``artifact_bucket``.
+    """
+    custom_store = MagicMock(name="custom_store")
+    custom_store.get = AsyncMock(return_value=b"hello world")
+    executor, default_store, _ = _make_executor()
+    # Pretend the cache already has the custom bucket wired up.
+    executor._bucket_stores["my-workspace-artifacts"] = custom_store
+
+    payload = b"hello world"
+    expected_hash = hashlib.sha256(payload).hexdigest()
+    result = await executor._download_artifact(
+        bucket_name="my-workspace-artifacts",
+        object_key="ws-1/user-1/script.py",
+        content_hash=expected_hash,
+    )
+
+    assert result == payload
+    custom_store.get.assert_awaited_once_with("ws-1/user-1/script.py")
+    # The default store must NOT have been touched — that was the bug.
+    default_store.get.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_download_artifact_uses_default_store_when_bucket_matches() -> None:
+    """Sanity: default bucket reads still flow through the injected store."""
+    from common.config import settings
+
+    default_store = MagicMock(name="default_store")
+    default_store.get = AsyncMock(return_value=b"payload")
+    executor = type(_make_executor()[0])(
+        session_factory=MagicMock(),
+        object_store=default_store,
+        storage_client=MagicMock(),
+    )
+    payload = b"payload"
+    result = await executor._download_artifact(
+        bucket_name=settings.s3_version_bucket,
+        object_key="ws-1/user-1/script.py",
+        content_hash=hashlib.sha256(payload).hexdigest(),
+    )
+
+    assert result == payload
+    default_store.get.assert_awaited_once_with("ws-1/user-1/script.py")
+
+
+@pytest.mark.asyncio
+async def test_download_artifact_rejects_hash_mismatch() -> None:
+    """Hash mismatch must raise regardless of which store was used —
+    protects against corrupted bytes landing in any bucket.
+    """
+    custom_store = MagicMock(name="custom_store")
+    custom_store.get = AsyncMock(return_value=b"corrupted")
+    executor, default_store, _ = _make_executor()
+    executor._bucket_stores["my-workspace-artifacts"] = custom_store
+
+    with pytest.raises(ValueError, match="hash mismatch"):
+        await executor._download_artifact(
+            bucket_name="my-workspace-artifacts",
+            object_key="ws-1/user-1/script.py",
+            content_hash="0" * 64,  # cannot match corrupted bytes
+        )
+
+    # Reads still went through the bucket-bound store, not the default.
+    custom_store.get.assert_awaited_once()
+    default_store.get.assert_not_called()

uv.lock

@@ -2585,6 +2585,7 @@ dependencies = [
 [package.dev-dependencies]
 dev = [
     { name = "pytest" },
+    { name = "pytest-asyncio" },
 ]
 
 [package.metadata]
@@ -2603,7 +2604,10 @@ requires-dist = [
 ]
 
 [package.metadata.requires-dev]
-dev = [{ name = "pytest", specifier = ">=9.1.1" }]
+dev = [
+    { name = "pytest", specifier = ">=9.1.1" },
+    { name = "pytest-asyncio", specifier = ">=1.4.0" },
+]
 
 [[package]]
 name = "send2trash"