Merge pull request #5268 from Agenta-AI/feat/agent-mounts-v2

feat(platform): add durable files shared across agent sessions

Author: mmabrouk

api/oss/src/apis/fastapi/mounts/models.py

@@ -31,6 +31,11 @@ class MountQueryRequest(BaseModel):
     windowing: Optional[Windowing] = None
 
 
+class AgentMountQueryRequest(BaseModel):
+    artifact_id: str
+    name: str = "default"
+
+
 # ---------------------------------------------------------------------------
 # Response models
 # ---------------------------------------------------------------------------

api/oss/src/apis/fastapi/mounts/router.py

@@ -12,6 +12,7 @@
 
 from oss.src.core.mounts.service import MountsService
 from oss.src.core.mounts.types import (
+    MountArtifactIdInvalid,
     MountDataInvalid,
     MountFileNotFound,
     MountImmutableField,
@@ -24,6 +25,7 @@
 )
 
 from oss.src.apis.fastapi.mounts.models import (
+    AgentMountQueryRequest,
     MountCreateRequest,
     MountCredentialsResponse,
     MountEditRequest,
@@ -65,6 +67,11 @@ async def wrapper(*args, **kwargs):
                     status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                     detail=e.message,
                 ) from e
+            except MountArtifactIdInvalid as e:
+                raise HTTPException(
+                    status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                    detail=e.message,
+                ) from e
             except MountSlugConflict as e:
                 raise HTTPException(
                     status_code=status.HTTP_409_CONFLICT,
@@ -130,6 +137,25 @@ def __init__(
             response_model_exclude_none=True,
             status_code=status.HTTP_200_OK,
         )
+        # Fixed agent sub-paths must be registered before "/{mount_id}" so they win.
+        self.router.add_api_route(
+            "/agents/sign",
+            self.sign_agent_mount_credentials,
+            methods=["POST"],
+            operation_id="sign_agent_mount_credentials",
+            response_model=MountCredentialsResponse,
+            response_model_exclude_none=True,
+            status_code=status.HTTP_200_OK,
+        )
+        self.router.add_api_route(
+            "/agents/query",
+            self.query_agent_mount,
+            methods=["POST"],
+            operation_id="query_agent_mount",
+            response_model=MountsResponse,
+            response_model_exclude_none=True,
+            status_code=status.HTTP_200_OK,
+        )
         self.router.add_api_route(
             "/{mount_id}",
             self.fetch_mount,
@@ -284,6 +310,48 @@ async def query_mounts(
 
         return MountsResponse(count=len(mounts), mounts=mounts)
 
+    @intercept_exceptions()
+    @handle_mount_exceptions()
+    async def sign_agent_mount_credentials(
+        self,
+        request: Request,
+        *,
+        artifact_id: str = Query(...),
+        name: str = Query(default="default"),
+    ) -> MountCredentialsResponse:
+        await self._check(request, Permission.RUN_SESSIONS)
+
+        mount = await self.mounts_service.get_or_create_agent_mount(
+            project_id=UUID(request.state.project_id),
+            user_id=UUID(str(request.state.user_id)),
+            artifact_id=artifact_id,
+            name=name,
+        )
+        credentials = await sign_mount_credentials(
+            mounts_service=self.mounts_service,
+            project_id=UUID(request.state.project_id),
+            mount_id=mount.id,
+        )
+        return MountCredentialsResponse(count=1, mount=mount, credentials=credentials)
+
+    @intercept_exceptions()
+    @handle_mount_exceptions()
+    async def query_agent_mount(
+        self,
+        request: Request,
+        *,
+        body: AgentMountQueryRequest,
+    ) -> MountsResponse:
+        await self._check(request, Permission.VIEW_SESSIONS)
+
+        mount = await self.mounts_service.fetch_agent_mount(
+            project_id=UUID(request.state.project_id),
+            artifact_id=body.artifact_id,
+            name=body.name,
+        )
+        mounts = [mount] if mount else []
+        return MountsResponse(count=len(mounts), mounts=mounts)
+
     @intercept_exceptions()
     async def fetch_mount(
         self,

api/oss/src/core/mounts/interfaces.py

@@ -36,6 +36,15 @@ async def fetch_mount(
         mount_id: UUID,
     ) -> Optional[Mount]: ...
 
+    @abstractmethod
+    async def fetch_mount_by_slug(
+        self,
+        *,
+        project_id: UUID,
+        #
+        slug: str,
+    ) -> Optional[Mount]: ...
+
     @abstractmethod
     async def edit_mount(
         self,

api/oss/src/core/mounts/service.py

@@ -18,6 +18,7 @@
 from oss.src.core.mounts.interfaces import MountsDAOInterface
 from oss.src.core.store.storage import ObjectStore
 from oss.src.core.mounts.types import (
+    MountArtifactIdInvalid,
     MountFileNotFound,
     MountNameInvalid,
     MountNotFound,
@@ -73,6 +74,21 @@ def mint_session_slug(*, session_id: str, name: str) -> str:
     return f"{_RESERVED_SLUG_PREFIX}{uuid5(_MOUNTS_NAMESPACE, session_id)}__{slugify_mount_name(name)}"
 
 
+def mint_agent_slug(*, artifact_id: str, name: str) -> str:
+    """Mint the deterministic reserved slug for an artifact mount.
+
+    Artifact IDs are UUID-parsed and rendered lowercase. Sign and query must use
+    this same derivation byte-identically so they address the same mount.
+    """
+    try:
+        canonical_artifact_id = UUID(str(artifact_id))
+    except (ValueError, TypeError, AttributeError) as e:
+        raise MountArtifactIdInvalid(str(artifact_id)) from e
+
+    slug_name = slugify_mount_name(name)
+    return f"{_RESERVED_SLUG_PREFIX}agent__{canonical_artifact_id}__{slug_name}"
+
+
 def reject_reserved_slug(slug: str) -> None:
     """A caller may not author a slug in the reserved namespace (the service mints those)."""
     if slug.startswith(_RESERVED_SLUG_PREFIX):
@@ -188,6 +204,26 @@ async def get_or_create_session_mount(
             mount_create=mount_create,
         )
 
+    async def get_or_create_agent_mount(
+        self,
+        *,
+        project_id: UUID,
+        user_id: UUID,
+        artifact_id: str,
+        name: str = "default",
+    ) -> Mount:
+        """Bind idempotently one durable mount for an artifact, keyed by name."""
+        slug_name = slugify_mount_name(name)
+        mount_create = MountCreate(
+            slug=mint_agent_slug(artifact_id=artifact_id, name=name),
+            name=slug_name,
+        )
+        return await self.mounts_dao.upsert_mount(
+            project_id=project_id,
+            user_id=user_id,
+            mount_create=mount_create,
+        )
+
     async def get_or_create_session_cwd(
         self,
         *,
@@ -216,6 +252,20 @@ async def fetch_mount(
             mount_id=mount_id,
         )
 
+    async def fetch_agent_mount(
+        self,
+        *,
+        project_id: UUID,
+        artifact_id: str,
+        name: str = "default",
+    ) -> Optional[Mount]:
+        """Fetch the active artifact mount keyed by name without creating it."""
+        slug = mint_agent_slug(artifact_id=artifact_id, name=name)
+        return await self.mounts_dao.fetch_mount_by_slug(
+            project_id=project_id,
+            slug=slug,
+        )
+
     async def edit_mount(
         self,
         *,

api/oss/src/core/mounts/types.py

@@ -33,6 +33,12 @@ def __init__(self, name: str = "name"):
         self.name = name
 
 
+class MountArtifactIdInvalid(MountError):
+    def __init__(self, artifact_id: str = "artifact_id"):
+        super().__init__(f"Artifact id '{artifact_id}' must be a valid UUID.")
+        self.artifact_id = artifact_id
+
+
 class MountImmutableField(MountError):
     def __init__(self, field: str = "field"):
         super().__init__(f"Mount field '{field}' is immutable after creation.")

api/oss/src/dbs/postgres/mounts/dao.py

@@ -115,6 +115,29 @@ async def fetch_mount(
 
             return map_mount_dbe_to_dto(mount_dbe=mount_dbe)
 
+    async def fetch_mount_by_slug(
+        self,
+        *,
+        project_id: UUID,
+        #
+        slug: str,
+    ) -> Optional[Mount]:
+        """Fetch by slug, excluding archived rows for agent-mount reads."""
+        async with self.engine.session() as session:
+            stmt = select(MountDBE).where(
+                MountDBE.project_id == project_id,
+                MountDBE.slug == slug,
+                MountDBE.deleted_at.is_(None),
+            )
+
+            result = await session.execute(stmt)
+            mount_dbe = result.scalar_one_or_none()
+
+            if not mount_dbe:
+                return None
+
+            return map_mount_dbe_to_dto(mount_dbe=mount_dbe)
+
     async def edit_mount(
         self,
         *,

api/oss/tests/pytest/acceptance/mounts/test_agent_mounts_basics.py

@@ -0,0 +1,69 @@
+"""Acceptance tests for artifact-scoped agent mounts.
+
+Query-only tests run without an object store. Credential-signing tests skip when
+the running API has no mount storage backend configured.
+"""
+
+from uuid import uuid4
+
+import pytest
+
+
+def _query_agent_mount(authed_api, artifact_id, *, name="default"):
+    return authed_api(
+        "POST",
+        "/mounts/agents/query",
+        json={"artifact_id": artifact_id, "name": name},
+    )
+
+
+def _sign_agent_mount(authed_api, artifact_id, *, name="default"):
+    response = authed_api(
+        "POST",
+        "/mounts/agents/sign",
+        params={"artifact_id": artifact_id, "name": name},
+    )
+    if response.status_code == 503:
+        pytest.skip("Mount storage backend not configured in this environment")
+    return response
+
+
+class TestAgentMountReads:
+    def test_query_rejects_non_uuid_artifact_id(self, authed_api):
+        response = _query_agent_mount(authed_api, "not-a-uuid")
+        assert response.status_code == 422, response.text
+
+    def test_query_unknown_uuid_stays_empty(self, authed_api):
+        artifact_id = str(uuid4())
+
+        first = _query_agent_mount(authed_api, artifact_id)
+        assert first.status_code == 200, first.text
+        assert first.json()["count"] == 0
+        assert first.json()["mounts"] == []
+
+        second = _query_agent_mount(authed_api, artifact_id)
+        assert second.status_code == 200, second.text
+        assert second.json()["count"] == 0
+        assert second.json()["mounts"] == []
+
+
+class TestAgentMountSign:
+    def test_sign_then_query_returns_same_mount(self, authed_api):
+        artifact_id = str(uuid4())
+        signed = _sign_agent_mount(authed_api, artifact_id)
+        assert signed.status_code == 200, signed.text
+        mount_id = signed.json()["mount"]["id"]
+
+        queried = _query_agent_mount(authed_api, artifact_id)
+        assert queried.status_code == 200, queried.text
+        assert queried.json()["count"] == 1
+        assert queried.json()["mounts"][0]["id"] == mount_id
+
+    def test_sign_twice_returns_same_mount(self, authed_api):
+        artifact_id = str(uuid4())
+        first = _sign_agent_mount(authed_api, artifact_id)
+        assert first.status_code == 200, first.text
+
+        second = _sign_agent_mount(authed_api, artifact_id)
+        assert second.status_code == 200, second.text
+        assert second.json()["mount"]["id"] == first.json()["mount"]["id"]