chuk-artifacts 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

chuk_artifacts/store.py

@@ -1,12 +1,18 @@
 # -*- coding: utf-8 -*-
-# chuk_artifacts/store.py (ENHANCED)
+# chuk_artifacts/store.py
 """
-Asynchronous, object-store-backed artefact manager with MCP server support.
+Clean ArtifactStore with mandatory sessions and grid architecture.
+
+Grid Architecture:
+- Mandatory session allocation (no anonymous artifacts)
+- Grid paths: grid/{sandbox_id}/{session_id}/{artifact_id}
+- Clean, focused implementation
 """
 
 from __future__ import annotations
 
-import os, logging
+import os, logging, uuid
+from datetime import datetime
 from typing import Any, Dict, List, Callable, AsyncContextManager, Optional, Union
 
 try:
@@ -25,7 +31,8 @@ except ImportError:
     logger.debug("python-dotenv not available, skipping .env file loading")
 
 # Import exceptions
-from .exceptions import ArtifactStoreError
+from .exceptions import ArtifactStoreError, ProviderError
+from .session.session_manager import SessionManager
 
 # Configure structured logging
 logger = logging.getLogger(__name__)
@@ -51,104 +58,73 @@ def _default_session_factory() -> Callable[[], AsyncContextManager]:
 # ─────────────────────────────────────────────────────────────────────
 class ArtifactStore:
     """
-    Asynchronous artifact storage with MCP server support.
+    Clean ArtifactStore with grid architecture and mandatory sessions.
     
-    Enhanced with MCP-specific operations for file management within sessions.
+    Simple rules:
+    - Always allocate a session (no anonymous artifacts)
+    - Grid paths only: grid/{sandbox_id}/{session_id}/{artifact_id}
+    - Clean, focused implementation
     """
 
     def __init__(
         self,
         *,
         bucket: Optional[str] = None,
-        s3_factory: Optional[Callable[[], AsyncContextManager]] = None,
         storage_provider: Optional[str] = None,
-        session_factory: Optional[Callable[[], AsyncContextManager]] = None,
         session_provider: Optional[str] = None,
+        sandbox_id: Optional[str] = None,
+        session_ttl_hours: int = 24,
         max_retries: int = 3,
-        # Backward compatibility - deprecated but still supported
-        redis_url: Optional[str] = None,
-        provider: Optional[str] = None,
     ):
-        # Read from environment variables with memory as defaults
-        bucket = bucket or os.getenv("ARTIFACT_BUCKET", "mcp-bucket")
+        # Configuration
+        self.bucket = bucket or os.getenv("ARTIFACT_BUCKET", "artifacts")
+        self.sandbox_id = sandbox_id or self._detect_sandbox_id()
+        self.session_ttl_hours = session_ttl_hours
+        self.max_retries = max_retries
+        self._closed = False
+        
+        # Storage provider
         storage_provider = storage_provider or os.getenv("ARTIFACT_PROVIDER", "memory")
+        self._s3_factory = self._load_storage_provider(storage_provider)
+        self._storage_provider_name = storage_provider
+        
+        # Session provider
         session_provider = session_provider or os.getenv("SESSION_PROVIDER", "memory")
+        self._session_factory = self._load_session_provider(session_provider)
+        self._session_provider_name = session_provider
         
-        # Handle backward compatibility
-        if redis_url is not None:
-            import warnings
-            warnings.warn(
-                "redis_url parameter is deprecated. Use session_provider='redis' "
-                "and set SESSION_REDIS_URL environment variable instead.",
-                DeprecationWarning,
-                stacklevel=2
-            )
-            os.environ["SESSION_REDIS_URL"] = redis_url
-            session_provider = "redis"
-            
-        if provider is not None:
-            import warnings
-            warnings.warn(
-                "provider parameter is deprecated. Use storage_provider instead.",
-                DeprecationWarning,
-                stacklevel=2
-            )
-            storage_provider = provider
-
-        # Validate factory/provider combinations
-        if s3_factory and storage_provider:
-            raise ValueError("Specify either s3_factory or storage_provider—not both")
-        if session_factory and session_provider:
-            raise ValueError("Specify either session_factory or session_provider—not both")
-
-        # Initialize storage factory
-        if s3_factory:
-            self._s3_factory = s3_factory
-        elif storage_provider:
-            self._s3_factory = self._load_storage_provider(storage_provider)
-        else:
-            self._s3_factory = _default_storage_factory()
-
-        # Initialize session factory
-        if session_factory:
-            self._session_factory = session_factory
-        elif session_provider:
-            self._session_factory = self._load_session_provider(session_provider)
-        else:
-            self._session_factory = _default_session_factory()
-
-        self.bucket = bucket
-        self.max_retries = max_retries
-        self._storage_provider_name = storage_provider or "memory"
-        self._session_provider_name = session_provider or "memory"
-        self._closed = False
-
-        # Initialize operation modules
-        from .core import CoreStorageOperations
-        from .presigned import PresignedURLOperations
-        from .metadata import MetadataOperations
-        from .batch import BatchOperations
-        from .admin import AdminOperations
-        from .session_operations import SessionOperations 
+        # Session manager (always enabled)
+        self._session_manager = SessionManager(
+            sandbox_id=self.sandbox_id,
+            session_factory=self._session_factory,
+            default_ttl_hours=session_ttl_hours,
+        )
+        
+        # Operation modules
+        from .core import CoreStorageOperations as CoreOps
+        from .metadata import MetadataOperations as MetaOps
+        from .presigned import PresignedURLOperations as PresignedOps
+        from .batch import BatchOperations as BatchOps
+        from .admin import AdminOperations as AdminOps
+        
+        self._core = CoreOps(self)
+        self._metadata = MetaOps(self)
+        self._presigned = PresignedOps(self)
+        self._batch = BatchOps(self)
+        self._admin = AdminOps(self)
         
-        self._core = CoreStorageOperations(self)
-        self._presigned = PresignedURLOperations(self)
-        self._metadata = MetadataOperations(self)
-        self._batch = BatchOperations(self)
-        self._admin = AdminOperations(self)
-        self._session = SessionOperations(self)
-
         logger.info(
-            "ArtifactStore initialized with session operations support",
+            "ArtifactStore initialized",
             extra={
-                "bucket": bucket,
-                "storage_provider": self._storage_provider_name,
-                "session_provider": self._session_provider_name,
+                "bucket": self.bucket,
+                "sandbox_id": self.sandbox_id,
+                "storage_provider": storage_provider,
+                "session_provider": session_provider,
             }
         )
 
     # ─────────────────────────────────────────────────────────────────
-    # Core storage operations (delegated to CoreStorageOperations)
+    # Core operations
     # ─────────────────────────────────────────────────────────────────
 
     async def store(
@@ -160,11 +136,22 @@ class ArtifactStore:
         meta: Dict[str, Any] | None = None,
         filename: str | None = None,
         session_id: str | None = None,
+        user_id: str | None = None,
         ttl: int = _DEFAULT_TTL,
     ) -> str:
-        """Store artifact data with metadata."""
-        return await self._core.store(
-            data,
+        """Store artifact with mandatory session allocation."""
+        # Always allocate/validate session
+        session_id = await self._session_manager.allocate_session(
+            session_id=session_id,
+            user_id=user_id,
+        )
+        
+        # Work around the naming conflict in CoreStorageOperations
+        # where self.store is the artifact store instance, not the method
+        core_store_method = getattr(self._core.__class__, 'store')
+        return await core_store_method(
+            self._core,
+            data=data,
             mime=mime,
             summary=summary,
             meta=meta,
@@ -174,11 +161,226 @@ class ArtifactStore:
         )
 
     async def retrieve(self, artifact_id: str) -> bytes:
-        """Retrieve artifact data directly."""
+        """Retrieve artifact data."""
         return await self._core.retrieve(artifact_id)
 
+    async def metadata(self, artifact_id: str) -> Dict[str, Any]:
+        """Get artifact metadata."""
+        return await self._metadata.get_metadata(artifact_id)
+
+    async def exists(self, artifact_id: str) -> bool:
+        """Check if artifact exists."""
+        return await self._metadata.exists(artifact_id)
+
+    async def delete(self, artifact_id: str) -> bool:
+        """Delete artifact."""
+        return await self._metadata.delete(artifact_id)
+
+    async def list_by_session(self, session_id: str, limit: int = 100) -> List[Dict[str, Any]]:
+        """List artifacts in session."""
+        return await self._metadata.list_by_session(session_id, limit)
+
+    # ─────────────────────────────────────────────────────────────────
+    # Session operations
+    # ─────────────────────────────────────────────────────────────────
+
+    async def create_session(
+        self,
+        user_id: Optional[str] = None,
+        ttl_hours: Optional[int] = None,
+    ) -> str:
+        """Create a new session."""
+        return await self._session_manager.allocate_session(
+            user_id=user_id,
+            ttl_hours=ttl_hours,
+        )
+
+    async def validate_session(self, session_id: str) -> bool:
+        """Validate session."""
+        return await self._session_manager.validate_session(session_id)
+
+    async def get_session_info(self, session_id: str) -> Optional[Dict[str, Any]]:
+        """Get session information."""
+        return await self._session_manager.get_session_info(session_id)
+
+    # ─────────────────────────────────────────────────────────────────
+    # Grid operations
+    # ─────────────────────────────────────────────────────────────────
+
+    def get_canonical_prefix(self, session_id: str) -> str:
+        """Get grid path prefix for session."""
+        return self._session_manager.get_canonical_prefix(session_id)
+
+    def generate_artifact_key(self, session_id: str, artifact_id: str) -> str:
+        """Generate grid artifact key."""
+        return self._session_manager.generate_artifact_key(session_id, artifact_id)
+
+    # ─────────────────────────────────────────────────────────────────
+    # File operations
+    # ─────────────────────────────────────────────────────────────────
+
+    async def write_file(
+        self,
+        content: Union[str, bytes],
+        *,
+        filename: str,
+        mime: str = "text/plain",
+        summary: str = "",
+        session_id: str = None,
+        user_id: str = None,
+        meta: Dict[str, Any] = None,
+        encoding: str = "utf-8",
+    ) -> str:
+        """Write content to file."""
+        if isinstance(content, str):
+            data = content.encode(encoding)
+        else:
+            data = content
+        
+        return await self.store(
+            data=data,
+            mime=mime,
+            summary=summary or f"File: {filename}",
+            filename=filename,
+            session_id=session_id,
+            user_id=user_id,
+            meta=meta,
+        )
+
+    async def read_file(
+        self,
+        artifact_id: str,
+        *,
+        encoding: str = "utf-8",
+        as_text: bool = True
+    ) -> Union[str, bytes]:
+        """Read file content."""
+        data = await self.retrieve(artifact_id)
+        
+        if as_text:
+            return data.decode(encoding)
+        return data
+
+    async def list_files(
+        self,
+        session_id: str,
+        prefix: str = "",
+        limit: int = 100
+    ) -> List[Dict[str, Any]]:
+        """List files in session with optional prefix filter."""
+        return await self._metadata.list_by_prefix(session_id, prefix, limit)
+
+    async def get_directory_contents(
+        self,
+        session_id: str,
+        directory_prefix: str = "",
+        limit: int = 100
+    ) -> List[Dict[str, Any]]:
+        """
+        List files in a directory-like structure within a session.
+        """
+        try:
+            return await self._metadata.list_by_prefix(session_id, directory_prefix, limit)
+        except Exception as e:
+            logger.error(
+                "Directory listing failed for session %s: %s",
+                session_id,
+                str(e),
+                extra={
+                    "session_id": session_id,
+                    "directory_prefix": directory_prefix,
+                    "operation": "get_directory_contents"
+                }
+            )
+            raise ProviderError(f"Directory listing failed: {e}") from e
+
+    async def copy_file(
+        self,
+        artifact_id: str,
+        *,
+        new_filename: str = None,
+        target_session_id: str = None,
+        new_meta: Dict[str, Any] = None,
+        summary: str = None
+    ) -> str:
+        """Copy a file WITHIN THE SAME SESSION only (security enforced)."""
+        # Get original metadata to check session
+        original_meta = await self.metadata(artifact_id)
+        original_session = original_meta.get("session_id")
+        
+        # STRICT SECURITY: Block ALL cross-session copies
+        if target_session_id and target_session_id != original_session:
+            raise ArtifactStoreError(
+                f"Cross-session copies are not permitted for security reasons. "
+                f"Artifact {artifact_id} belongs to session '{original_session}', "
+                f"cannot copy to session '{target_session_id}'. Files can only be "
+                f"copied within the same session."
+            )
+        
+        # Get original data
+        original_data = await self.retrieve(artifact_id)
+        
+        # Prepare copy metadata
+        copy_filename = new_filename or (
+            (original_meta.get("filename", "file") or "file") + "_copy"
+        )
+        copy_summary = summary or f"Copy of {original_meta.get('summary', 'artifact')}"
+        
+        # Merge metadata
+        copy_meta = {**original_meta.get("meta", {})}
+        if new_meta:
+            copy_meta.update(new_meta)
+        
+        # Add copy tracking
+        copy_meta["copied_from"] = artifact_id
+        copy_meta["copy_timestamp"] = datetime.utcnow().isoformat() + "Z"
+        
+        # Store the copy in the same session
+        return await self.store(
+            data=original_data,
+            mime=original_meta["mime"],
+            summary=copy_summary,
+            filename=copy_filename,
+            session_id=original_session,  # Always same session
+            meta=copy_meta
+        )
+
+    async def move_file(
+        self,
+        artifact_id: str,
+        *,
+        new_filename: str = None,
+        new_session_id: str = None,
+        new_meta: Dict[str, Any] = None
+    ) -> Dict[str, Any]:
+        """Move/rename a file WITHIN THE SAME SESSION only (security enforced)."""
+        # Get current metadata
+        record = await self.metadata(artifact_id)
+        current_session = record.get("session_id")
+        
+        # STRICT SECURITY: Block ALL cross-session moves
+        if new_session_id and new_session_id != current_session:
+            raise ArtifactStoreError(
+                f"Cross-session moves are not permitted for security reasons. "
+                f"Artifact {artifact_id} belongs to session '{current_session}', "
+                f"cannot move to session '{new_session_id}'. Use copy operations within "
+                f"the same session only."
+            )
+        
+        # For now, just simulate a move by updating metadata
+        # A full implementation would update the metadata record
+        if new_filename:
+            # This is a simplified move - just return updated record
+            record["filename"] = new_filename
+        if new_meta:
+            existing_meta = record.get("meta", {})
+            existing_meta.update(new_meta)
+            record["meta"] = existing_meta
+        
+        return record
+
     # ─────────────────────────────────────────────────────────────────
-    # Presigned URL operations (delegated to PresignedURLOperations)
+    # Presigned URL operations
     # ─────────────────────────────────────────────────────────────────
 
     async def presign(self, artifact_id: str, expires: int = _DEFAULT_PRESIGN_EXPIRES) -> str:
@@ -205,12 +407,7 @@ class ArtifactStore:
         expires: int = _DEFAULT_PRESIGN_EXPIRES
     ) -> tuple[str, str]:
         """Generate a presigned URL for uploading a new artifact."""
-        return await self._presigned.presign_upload(
-            session_id=session_id,
-            filename=filename,
-            mime_type=mime_type,
-            expires=expires
-        )
+        return await self._presigned.presign_upload(session_id, filename, mime_type, expires)
 
     async def register_uploaded_artifact(
         self,
@@ -257,76 +454,50 @@ class ArtifactStore:
         )
 
     # ─────────────────────────────────────────────────────────────────
-    # Metadata operations (delegated to MetadataOperations)
+    # Batch operations
     # ─────────────────────────────────────────────────────────────────
 
-    async def metadata(self, artifact_id: str) -> Dict[str, Any]:
-        """Get artifact metadata."""
-        return await self._metadata.metadata(artifact_id)
-
-    async def exists(self, artifact_id: str) -> bool:
-        """Check if artifact exists and hasn't expired."""
-        return await self._metadata.exists(artifact_id)
+    async def store_batch(
+        self,
+        items: List[Dict[str, Any]],
+        session_id: str | None = None,
+        ttl: int = _DEFAULT_TTL,
+    ) -> List[str]:
+        """Store multiple artifacts in a batch operation."""
+        return await self._batch.store_batch(items, session_id, ttl)
 
-    async def delete(self, artifact_id: str) -> bool:
-        """Delete artifact and its metadata."""
-        return await self._metadata.delete(artifact_id)
+    # ─────────────────────────────────────────────────────────────────
+    # Metadata operations
+    # ─────────────────────────────────────────────────────────────────
 
     async def update_metadata(
-        self, 
-        artifact_id: str, 
+        self,
+        artifact_id: str,
         *,
         summary: str = None,
         meta: Dict[str, Any] = None,
-        filename: str = None,
-        ttl: int = None,
-        # NEW: MCP-specific parameters
-        new_meta: Dict[str, Any] = None,
-        merge: bool = True
+        merge: bool = True,
+        **kwargs
     ) -> Dict[str, Any]:
-        """Update artifact metadata with MCP server compatibility."""
+        """Update artifact metadata."""
         return await self._metadata.update_metadata(
             artifact_id,
             summary=summary,
             meta=meta,
-            filename=filename,
-            ttl=ttl,
-            new_meta=new_meta,
-            merge=merge
+            merge=merge,
+            **kwargs
         )
 
-    async def extend_ttl(self, artifact_id: str, additional_seconds: int) -> Dict[str, Any]:
-        """Extend the TTL of an artifact's metadata."""
-        return await self._metadata.extend_ttl(artifact_id, additional_seconds)
-
-    async def list_by_session(self, session_id: str, limit: int = 100) -> List[Dict[str, Any]]:
-        """List artifacts for a specific session."""
-        return await self._metadata.list_by_session(session_id, limit)
-
-    async def list_by_prefix(
-        self, 
-        session_id: str, 
-        prefix: str = "", 
-        limit: int = 100
-    ) -> List[Dict[str, Any]]:
-        """List artifacts in a session with filename prefix filtering."""
-        return await self._metadata.list_by_prefix(session_id, prefix, limit)
-
-    # ─────────────────────────────────────────────────────────────────
-    # Batch operations (delegated to BatchOperations)
-    # ─────────────────────────────────────────────────────────────────
-
-    async def store_batch(
+    async def extend_ttl(
         self,
-        items: List[Dict[str, Any]],
-        session_id: str | None = None,
-        ttl: int = _DEFAULT_TTL,
-    ) -> List[str]:
-        """Store multiple artifacts in a batch operation."""
-        return await self._batch.store_batch(items, session_id, ttl)
+        artifact_id: str,
+        additional_seconds: int
+    ) -> Dict[str, Any]:
+        """Extend artifact TTL."""
+        return await self._metadata.extend_ttl(artifact_id, additional_seconds)
 
     # ─────────────────────────────────────────────────────────────────
-    # Administrative operations (delegated to AdminOperations)
+    # Administrative operations
     # ─────────────────────────────────────────────────────────────────
 
     async def validate_configuration(self) -> Dict[str, Any]:
@@ -338,100 +509,53 @@ class ArtifactStore:
         return await self._admin.get_stats()
 
     # ─────────────────────────────────────────────────────────────────
-    # Session-based file operations (delegated to SessionOperations)
+    # Helpers
     # ─────────────────────────────────────────────────────────────────
 
-    async def move_file(
-        self,
-        artifact_id: str,
-        *,
-        new_filename: str = None,
-        new_session_id: str = None,
-        new_meta: Dict[str, Any] = None
-    ) -> Dict[str, Any]:
-        """Move a file within sessions or rename it."""
-        return await self._session.move_file(
-            artifact_id,
-            new_filename=new_filename,
-            new_session_id=new_session_id,
-            new_meta=new_meta
-        )
-
-    async def copy_file(
-        self,
-        artifact_id: str,
-        *,
-        new_filename: str = None,
-        target_session_id: str = None,
-        new_meta: Dict[str, Any] = None,
-        summary: str = None
-    ) -> str:
-        """Copy a file within or across sessions."""
-        return await self._session.copy_file(
-            artifact_id,
-            new_filename=new_filename,
-            target_session_id=target_session_id,
-            new_meta=new_meta,
-            summary=summary
-        )
-
-    async def read_file(
-        self,
-        artifact_id: str,
-        *,
-        encoding: str = "utf-8",
-        as_text: bool = True
-    ) -> Union[str, bytes]:
-        """Read file content directly."""
-        return await self._session.read_file(
-            artifact_id,
-            encoding=encoding,
-            as_text=as_text
-        )
+    def _detect_sandbox_id(self) -> str:
+        """Auto-detect sandbox ID."""
+        candidates = [
+            os.getenv("ARTIFACT_SANDBOX_ID"),
+            os.getenv("SANDBOX_ID"),
+            os.getenv("HOSTNAME"),
+        ]
+        
+        for candidate in candidates:
+            if candidate:
+                clean_id = "".join(c for c in candidate if c.isalnum() or c in "-_")[:32]
+                if clean_id:
+                    return clean_id
+        
+        # Generate fallback
+        return f"sandbox-{uuid.uuid4().hex[:8]}"
 
-    async def write_file(
-        self,
-        content: Union[str, bytes],
-        *,
-        filename: str,
-        mime: str = "text/plain",
-        summary: str = "",
-        session_id: str = None,
-        meta: Dict[str, Any] = None,
-        encoding: str = "utf-8",
-        overwrite_artifact_id: str = None
-    ) -> str:
-        """Write content to a new file or overwrite existing."""
-        return await self._session.write_file(
-            content,
-            filename=filename,
-            mime=mime,
-            summary=summary,
-            session_id=session_id,
-            meta=meta,
-            encoding=encoding,
-            overwrite_artifact_id=overwrite_artifact_id
-        )
+    def _load_storage_provider(self, name: str) -> Callable[[], AsyncContextManager]:
+        """Load storage provider."""
+        from importlib import import_module
+        
+        try:
+            mod = import_module(f"chuk_artifacts.providers.{name}")
+            return mod.factory()
+        except ModuleNotFoundError as exc:
+            available = ["memory", "filesystem", "s3", "ibm_cos"]
+            raise ValueError(f"Unknown storage provider '{name}'. Available: {', '.join(available)}") from exc
 
-    async def get_directory_contents(
-        self,
-        session_id: str,
-        directory_prefix: str = "",
-        limit: int = 100
-    ) -> List[Dict[str, Any]]:
-        """List files in a directory-like structure within a session."""
-        return await self._session.get_directory_contents(
-            session_id,
-            directory_prefix,
-            limit
-        )
+    def _load_session_provider(self, name: str) -> Callable[[], AsyncContextManager]:
+        """Load session provider."""
+        from importlib import import_module
+        
+        try:
+            mod = import_module(f"chuk_sessions.providers.{name}")
+            return mod.factory()
+        except ModuleNotFoundError as exc:
+            raise ValueError(f"Unknown session provider '{name}'") from exc
 
     # ─────────────────────────────────────────────────────────────────
     # Resource management
     # ─────────────────────────────────────────────────────────────────
 
     async def close(self):
-        """Mark store as closed."""
+        """Close the store."""
         if not self._closed:
             self._closed = True
             logger.info("ArtifactStore closed")
@@ -440,42 +564,4 @@ class ArtifactStore:
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.close()
-
-    # ─────────────────────────────────────────────────────────────────
-    # Helper functions (still needed for provider loading)
-    # ─────────────────────────────────────────────────────────────────
-
-    def _load_storage_provider(self, name: str) -> Callable[[], AsyncContextManager]:
-        """Load storage provider by name."""
-        from importlib import import_module
-
-        try:
-            mod = import_module(f"chuk_artifacts.providers.{name}")
-        except ModuleNotFoundError as exc:
-            available = ["memory", "filesystem", "s3", "ibm_cos", "ibm_cos_iam"]
-            raise ValueError(
-                f"Unknown storage provider '{name}'. "
-                f"Available providers: {', '.join(available)}"
-            ) from exc
-
-        if not hasattr(mod, "factory"):
-            raise AttributeError(f"Storage provider '{name}' lacks factory()")
-        
-        logger.info(f"Loaded storage provider: {name}")
-        return mod.factory()
-
-    def _load_session_provider(self, name: str) -> Callable[[], AsyncContextManager]:
-        """Load session provider by name."""
-        from importlib import import_module
-
-        try:
-            mod = import_module(f"chuk_sessions.providers.{name}")
-        except ModuleNotFoundError as exc:
-            raise ValueError(f"Unknown session provider '{name}'") from exc
-
-        if not hasattr(mod, "factory"):
-            raise AttributeError(f"Session provider '{name}' lacks factory()")
-        
-        logger.info(f"Loaded session provider: {name}")
-        return mod.factory()
+        await self.close()