Integrate Cognee service updates

ai/pyproject.toml

@@ -16,7 +16,8 @@ dependencies = [
     "fastmcp",
     "mcp",
     "typing-extensions",
-    "cognee>=0.3.0",
+    "cognee @ file://../cognee/cognee",
+    "python-dotenv",
 ]
 
 [project.optional-dependencies]

ai/src/fuzzforge_ai/cognee_api_client.py

@@ -0,0 +1,316 @@
+"""Async HTTP client for interacting with a remote Cognee service."""
+
+from __future__ import annotations
+
+import json
+import logging
+import mimetypes
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Sequence
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+
+class CogneeAPIError(RuntimeError):
+    """Raised when the Cognee service returns an unexpected response."""
+
+
+class CogneeAPIClient:
+    """Thin async wrapper around Cognee's HTTP API."""
+
+    def __init__(self, base_url: str, timeout: float = 180.0, verify_ssl: Optional[bool] = None):
+        if not base_url:
+            raise ValueError("Cognee service URL must be provided")
+        self.base_url = base_url.rstrip("/")
+        kwargs: Dict[str, Any] = {"timeout": timeout, "base_url": self.base_url, "follow_redirects": True}
+        if verify_ssl is not None:
+            kwargs["verify"] = verify_ssl
+        self._client = httpx.AsyncClient(**kwargs)
+        self._token: Optional[str] = None
+
+    async def close(self) -> None:
+        await self._client.aclose()
+
+    # ------------------------------------------------------------------
+    # Authentication helpers
+    # ------------------------------------------------------------------
+    async def login(self, email: str, password: str) -> bool:
+        """Attempt to authenticate the user and store the JWT token."""
+        form = {"username": email, "password": password}
+        try:
+            response = await self._client.post("/api/v1/auth/login", data=form)
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee login request failed: {exc}") from exc
+
+        if response.status_code == 200:
+            data = response.json()
+            self._token = data.get("access_token")
+            if self._token:
+                # Update client headers with Bearer token
+                self._client.headers.update({"Authorization": f"Bearer {self._token}"})
+            return True
+        if response.status_code in {400, 401}:  # Invalid credentials / user not found
+            return False
+        raise CogneeAPIError(
+            f"Unexpected login status {response.status_code}: {response.text.strip()}"
+        )
+
+    async def register(self, email: str, password: str) -> bool:
+        """Register a new user; returns True when created, False if it already exists."""
+        payload = {"email": email, "password": password}
+        try:
+            response = await self._client.post("/api/v1/auth/register", json=payload)
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee register request failed: {exc}") from exc
+
+        if response.status_code in {200, 201}:
+            return True
+        if response.status_code in {400, 409}:  # already exists / validation error
+            logger.debug("Cognee register response %s: %s", response.status_code, response.text)
+            return False
+        raise CogneeAPIError(
+            f"Unexpected register status {response.status_code}: {response.text.strip()}"
+        )
+
+    async def ensure_user(self, email: str, password: str, tenant_name: str) -> Dict[str, Any]:
+        """Ensure a user exists, authenticate, and provision tenant if needed."""
+        logged_in = await self.login(email, password)
+        if not logged_in:
+            created = await self.register(email, password)
+            if not created:
+                # User may already exist with different password; try login again to surface error
+                logged_in = await self.login(email, password)
+                if not logged_in:
+                    raise CogneeAPIError("Unable to authenticate with Cognee service")
+            else:
+                # Newly created -> login to capture cookies
+                if not await self.login(email, password):
+                    raise CogneeAPIError("Login failed after registering Cognee service user")
+
+        # Fetch user profile - if this fails with 401, auth may not be enforced
+        try:
+            profile = await self._get_me()
+        except CogneeAPIError as exc:
+            if "401" in str(exc):
+                # Auth endpoints exist but may not be enforced - create minimal profile
+                logger.debug("Auth endpoints returned 401, using minimal profile")
+                profile = {"email": email, "id": None}
+            else:
+                raise
+
+        # Ensure tenant exists and is associated with the user
+        if tenant_name:
+            await self._create_tenant_if_needed(tenant_name)
+            profile["tenant_name"] = tenant_name
+
+        return profile
+
+    async def _get_me(self) -> Dict[str, Any]:
+        try:
+            response = await self._client.get("/api/v1/auth/me")
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee profile request failed: {exc}") from exc
+
+        if response.status_code == 200:
+            return response.json()
+        raise CogneeAPIError(
+            f"Failed to fetch Cognee profile ({response.status_code}): {response.text.strip()}"
+        )
+
+    async def _create_tenant_if_needed(self, tenant_name: str) -> None:
+        try:
+            response = await self._client.post(
+                "/api/v1/permissions/tenants", params={"tenant_name": tenant_name}
+            )
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee tenant request failed: {exc}") from exc
+
+        if response.status_code in {200, 409}:  # created or already exists
+            return
+        # FastAPI may surface validation errors as 400/422; treat as already exists
+        if response.status_code in {400, 422}:
+            logger.debug("Tenant creation returned %s: %s", response.status_code, response.text)
+            return
+        raise CogneeAPIError(
+            f"Failed to ensure Cognee tenant ({response.status_code}): {response.text.strip()}"
+        )
+
+    # ------------------------------------------------------------------
+    # Dataset helpers
+    # ------------------------------------------------------------------
+    async def ensure_dataset(self, dataset_name: str) -> Dict[str, Any]:
+        payload = {"name": dataset_name}
+        try:
+            response = await self._client.post("/api/v1/datasets", json=payload)
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee dataset request failed: {exc}") from exc
+
+        if response.status_code in {200, 201}:
+            return response.json()
+        raise CogneeAPIError(
+            f"Failed to ensure dataset {dataset_name} ({response.status_code}): {response.text.strip()}"
+        )
+
+    async def delete_dataset(self, dataset_id: str) -> None:
+        try:
+            response = await self._client.delete(f"/api/v1/datasets/{dataset_id}")
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee dataset delete failed: {exc}") from exc
+
+        if response.status_code in {200, 204}:
+            return
+        if response.status_code == 404:
+            return
+        raise CogneeAPIError(
+            f"Failed to delete dataset ({response.status_code}): {response.text.strip()}"
+        )
+
+    async def list_datasets(self) -> List[Dict[str, Any]]:
+        try:
+            response = await self._client.get("/api/v1/datasets")
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee datasets list failed: {exc}") from exc
+
+        if response.status_code == 200:
+            return response.json()
+        raise CogneeAPIError(
+            f"Failed to list datasets ({response.status_code}): {response.text.strip()}"
+        )
+
+    # ------------------------------------------------------------------
+    # Ingestion and processing
+    # ------------------------------------------------------------------
+    async def ingest_files(self, dataset_name: str, files: Sequence[Path]) -> Dict[str, Any]:
+        form_data = {"datasetName": dataset_name}
+        multipart_files: List[tuple[str, tuple[str, Any, Optional[str]]]] = []
+        for path in files:
+            mime_type, _ = mimetypes.guess_type(path.name)
+            try:
+                content = path.read_bytes()
+            except OSError as exc:
+                raise CogneeAPIError(f"Failed to read file for Cognee upload: {path}: {exc}") from exc
+            multipart_files.append(
+                (
+                    "data",
+                    (path.name, content, mime_type or "application/octet-stream"),
+                )
+            )
+
+        try:
+            response = await self._client.post(
+                "/api/v1/add", data=form_data, files=multipart_files
+            )
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee ingestion request failed: {exc}") from exc
+
+        if response.status_code in {200, 201}:
+            result = response.json()
+            result.setdefault("status", response.status_code)
+            return result
+
+        raise CogneeAPIError(
+            f"Cognee ingestion failed ({response.status_code}): {response.text.strip()}"
+        )
+
+    async def cognify(self, dataset_name: str) -> Dict[str, Any]:
+        payload = {"datasets": [dataset_name]}
+        try:
+            response = await self._client.post("/api/v1/cognify", json=payload)
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee cognify request failed: {exc}") from exc
+
+        if response.status_code in {200, 201}:
+            return response.json()
+        raise CogneeAPIError(
+            f"Cognee cognify failed ({response.status_code}): {response.text.strip()}"
+        )
+
+    async def ingest_text(self, dataset_name: str, content: str) -> Dict[str, Any]:
+        data_path = Path("virtual.txt")
+        multipart_files = [
+            (
+                "data",
+                (data_path.name, content.encode("utf-8"), "text/plain"),
+            )
+        ]
+        form_data = {"datasetName": dataset_name}
+        try:
+            response = await self._client.post(
+                "/api/v1/add", data=form_data, files=multipart_files
+            )
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee text ingestion request failed: {exc}") from exc
+
+        if response.status_code in {200, 201}:
+            return response.json()
+        raise CogneeAPIError(
+            f"Cognee text ingestion failed ({response.status_code}): {response.text.strip()}"
+        )
+
+    # ------------------------------------------------------------------
+    # Search / insights
+    # ------------------------------------------------------------------
+    async def search(
+        self,
+        query: str,
+        datasets: Optional[Iterable[str]] = None,
+        search_type: str = "INSIGHTS",
+        top_k: Optional[int] = None,
+    ) -> Any:
+        payload: Dict[str, Any] = {
+            "query": query,
+            "search_type": search_type,
+        }
+        if datasets:
+            payload["datasets"] = list(datasets)
+        if top_k is not None:
+            payload["top_k"] = top_k
+
+        try:
+            response = await self._client.post("/api/v1/search", json=payload)
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee search request failed: {exc}") from exc
+
+        if response.status_code in {200, 201}:
+            # search may already be JSON serialisable but ensure data type is python object
+            return response.json()
+        raise CogneeAPIError(
+            f"Cognee search failed ({response.status_code}): {response.text.strip()}"
+        )
+
+    async def status(self) -> Dict[str, Any]:
+        """Return health information along with dataset list."""
+        info: Dict[str, Any] = {}
+        try:
+            response = await self._client.get("/health")
+            if response.status_code == 200:
+                info["health"] = response.json()
+            else:
+                info["health_error"] = {
+                    "status_code": response.status_code,
+                    "body": response.text.strip(),
+                }
+        except httpx.RequestError as exc:
+            info["health_error"] = {"message": str(exc)}
+
+        try:
+            datasets = await self.list_datasets()
+            info["datasets"] = datasets
+        except CogneeAPIError as exc:
+            info.setdefault("dataset_error", str(exc))
+
+        return info
+
+    # ------------------------------------------------------------------
+    async def raw_request(self, method: str, url: str, **kwargs: Any) -> httpx.Response:
+        """Expose a raw request helper for advanced scenarios."""
+        try:
+            response = await self._client.request(method, url, **kwargs)
+        except httpx.RequestError as exc:
+            raise CogneeAPIError(f"Cognee request failed: {exc}") from exc
+        return response
+
+
+__all__ = ["CogneeAPIClient", "CogneeAPIError"]

ai/src/fuzzforge_ai/cognee_service.py

@@ -15,14 +15,16 @@ Provides integrated Cognee functionality for codebase analysis and knowledge gra
 
 
 import os
-import asyncio
+import json
 import logging
 from pathlib import Path
-from typing import Dict, List, Any, Optional
+from typing import Dict, List, Any, Optional, Tuple
 from datetime import datetime
 
 logger = logging.getLogger(__name__)
 
+from .cognee_api_client import CogneeAPIClient, CogneeAPIError
+
 
 class CogneeService:
     """
@@ -33,17 +35,108 @@ class CogneeService:
     def __init__(self, config):
         """Initialize with FuzzForge config"""
         self.config = config
-        self.cognee_config = config.get_cognee_config()
-        self.project_context = config.get_project_context()
         self._cognee = None
         self._user = None
         self._initialized = False
+        self._s3_client = None
+        self._api_client: Optional[CogneeAPIClient] = None
+        self._dataset_id: Optional[str] = None
+        self._sync_config_state()
+
+    def _sync_config_state(self) -> None:
+        """Refresh cached configuration from the project manager."""
+        self.cognee_config = self.config.get_cognee_config()
+        self.project_context = self.config.get_project_context()
+        self.project_dir = Path(self.project_context.get("project_dir", Path.cwd()))
+        self.storage_backend = self.cognee_config.get("storage_backend", "filesystem")
+        mode_env = os.getenv("COGNEE_STORAGE_MODE")
+        if mode_env:
+            self.service_mode = mode_env
+        else:
+            self.service_mode = self.cognee_config.get("mode", "embedded")
+
+        bucket_env = os.getenv("COGNEE_S3_BUCKET") or os.getenv("S3_BUCKET")
+        prefix_env = os.getenv("COGNEE_S3_PREFIX")
+        self.s3_bucket = bucket_env or self.cognee_config.get("s3_bucket")
+        self.s3_prefix = prefix_env or self.cognee_config.get("s3_prefix")
+
+        email_env = os.getenv("COGNEE_SERVICE_USER_EMAIL") or os.getenv("DEFAULT_USER_EMAIL")
+        self.service_user_email = email_env or self.cognee_config.get("service_user_email")
+        self.service_env_dir = self.cognee_config.get("service_env_dir")
+
+        url_env = os.getenv("COGNEE_SERVICE_URL") or os.getenv("COGNEE_API_URL")
+        self.service_url = url_env or self.cognee_config.get("service_url")
+
+    def _using_service(self) -> bool:
+        return (self.service_mode or "").lower() == "service" and bool(self.service_url)
+
+    def _get_service_credentials(self) -> Tuple[str, str]:
+        email = (
+            self.service_user_email
+            or os.getenv("COGNEE_SERVICE_USER_EMAIL")
+            or os.getenv("DEFAULT_USER_EMAIL")
+            or f"project_{self.project_context['project_id']}@cognee.dev"
+        )
+        password = (
+            os.getenv("COGNEE_SERVICE_USER_PASSWORD")
+            or os.getenv("DEFAULT_USER_PASSWORD")
+            or f"{self.project_context['project_id'][:8]}_C0gn33!"
+        )
+        os.environ.setdefault("COGNEE_SERVICE_USER_EMAIL", email)
+        os.environ.setdefault("DEFAULT_USER_EMAIL", email)
+        os.environ.setdefault("COGNEE_SERVICE_USER_PASSWORD", password)
+        os.environ.setdefault("DEFAULT_USER_PASSWORD", password)
+        return email, password
+
+    async def _initialize_remote(self) -> None:
+        if not self.service_url:
+            raise ValueError(
+                "COGNEE_SERVICE_URL must be configured when Cognee mode is set to 'service'"
+            )
+
+        # Ensure env variables are populated so downstream components (LLM config etc.) stay aligned
+        self.config.setup_cognee_environment()
+        self._sync_config_state()
+
+        email, password = self._get_service_credentials()
+        tenant_name = self.project_context.get("tenant_id")
+
+        if self._api_client:
+            await self._api_client.close()
+
+        self._api_client = CogneeAPIClient(self.service_url)
+
+        try:
+            profile = await self._api_client.ensure_user(email, password, tenant_name)
+            self._user = profile
+            dataset_name = self.get_project_dataset_name()
+            dataset_info = await self._api_client.ensure_dataset(dataset_name)
+            self._dataset_id = dataset_info.get("id")
+            logger.info(
+                "Connected to Cognee service",
+                extra={
+                    "service_url": self.service_url,
+                    "project": self.project_context.get("project_name"),
+                    "dataset": dataset_name,
+                },
+            )
+        except Exception:
+            if self._api_client:
+                await self._api_client.close()
+            self._api_client = None
+            raise
     
     async def initialize(self):
         """Initialize Cognee with project-specific configuration"""
         try:
-            # Ensure environment variables for Cognee are set before import
+            self.config.refresh()
             self.config.setup_cognee_environment()
+            self._sync_config_state()
+            if self._using_service():
+                await self._initialize_remote()
+                self._initialized = True
+                return
+
             logger.debug(
                 "Cognee environment configured",
                 extra={
@@ -119,8 +212,14 @@ class CogneeService:
             await self._setup_user_context()
             
             self._initialized = True
-            logger.info(f"Cognee initialized for project {self.project_context['project_name']} "
-                       f"with Kuzu at {system_dir}")
+            logger.info(
+                "Cognee initialized",
+                extra={
+                    "project": self.project_context.get("project_name"),
+                    "storage": self.storage_backend,
+                    "data_root": data_dir,
+                },
+            )
             
         except ImportError:
             logger.error("Cognee not installed. Install with: pip install cognee")
@@ -137,7 +236,7 @@ class CogneeService:
         try:
             # Dataset creation is handled automatically by Cognee when adding files
             # We just ensure we have the right context set up
-            dataset_name = f"{self.project_context['project_name']}_codebase"
+            dataset_name = self.get_project_dataset_name()
             logger.info(f"Dataset {dataset_name} ready for project {self.project_context['project_name']}")
             return dataset_name
         except Exception as e:
@@ -146,47 +245,187 @@ class CogneeService:
     
     async def _setup_user_context(self):
         """Setup user context for multi-tenant isolation"""
+        if self._using_service():
+            return
         try:
             from cognee.modules.users.methods import create_user, get_user
-            
-            # Always try fallback email first to avoid validation issues
-            fallback_email = f"project_{self.project_context['project_id']}@fuzzforge.example"
-            user_tenant = self.project_context['tenant_id']
-            
-            # Try to get existing fallback user first
+            from cognee.modules.users.methods.get_user_by_email import get_user_by_email
+            from cognee.modules.users.tenants.methods import create_tenant
+            from cognee.infrastructure.databases.exceptions import EntityAlreadyExistsError
+
+            fallback_email = self.service_user_email or f"project_{self.project_context['project_id']}@cognee.dev"
+            fallback_password = (
+                os.getenv("COGNEE_SERVICE_USER_PASSWORD")
+                or os.getenv("DEFAULT_USER_PASSWORD")
+                or f"{self.project_context['project_id'][:8]}_C0gn33!"
+            )
+
+            os.environ.setdefault("DEFAULT_USER_EMAIL", fallback_email)
+            os.environ.setdefault("DEFAULT_USER_PASSWORD", fallback_password)
+
             try:
-                self._user = await get_user(fallback_email)
-                logger.info(f"Using existing user: {fallback_email}")
-                return
-            except:
-                # User doesn't exist, try to create fallback
-                pass
-            
-            # Create fallback user
-            try:
-                self._user = await create_user(fallback_email, user_tenant)
-                logger.info(f"Created fallback user: {fallback_email} for tenant: {user_tenant}")
-                return
-            except Exception as fallback_error:
-                logger.warning(f"Fallback user creation failed: {fallback_error}")
-                self._user = None
-                return
-            
+                user = await get_user_by_email(fallback_email)
+                logger.info("Using existing Cognee service user", extra={"email": fallback_email})
+            except Exception:
+                user = None
+
+            if user is None:
+                try:
+                    user = await create_user(
+                        email=fallback_email,
+                        password=fallback_password,
+                        is_superuser=True,
+                        is_active=True,
+                        is_verified=True,
+                        auto_login=True,
+                    )
+                    logger.info(
+                        "Created Cognee service user",
+                        extra={"email": fallback_email, "tenant": self.project_context.get("tenant_id")},
+                    )
+                except Exception as exc:
+                    logger.warning("Failed to create Cognee service user", exc_info=True)
+                    self._user = None
+                    return
+
+            tenant_name = self.project_context.get("tenant_id") or f"fuzzforge_project_{self.project_context['project_id']}"
+            if not getattr(user, "tenant_id", None):
+                try:
+                    await create_tenant(tenant_name, user.id)
+                except EntityAlreadyExistsError:
+                    logger.debug("Tenant already exists for project", extra={"tenant": tenant_name})
+                except Exception as exc:
+                    logger.warning("Failed to ensure Cognee tenant", exc_info=True)
+
+                try:
+                    user = await get_user(user.id)
+                except Exception:
+                    pass
+
+            self._user = user
+            self.service_user_email = fallback_email
+
         except Exception as e:
             logger.warning(f"Could not setup multi-tenant user context: {e}")
             logger.info("Proceeding with default context")
             self._user = None
+
+    def _uses_s3_storage(self) -> bool:
+        if (self.storage_backend or "").lower() == "s3":
+            return True
+        if os.getenv("STORAGE_BACKEND", "").lower() == "s3":
+            return True
+
+        data_dir = (self.cognee_config.get("data_directory") or "").lower()
+        system_dir = (self.cognee_config.get("system_directory") or "").lower()
+        return data_dir.startswith("s3://") or system_dir.startswith("s3://")
+
+    def _relative_ingest_path(self, file_path: Path) -> str:
+        try:
+            relative = file_path.relative_to(self.project_dir)
+        except ValueError:
+            relative = file_path.name
+        return str(relative).replace(os.sep, "/")
+
+    def _resolve_s3_paths(self) -> Tuple[str, str]:
+        bucket = (
+            self.s3_bucket
+            or os.getenv("STORAGE_BUCKET_NAME")
+            or os.getenv("COGNEE_S3_BUCKET")
+            or os.getenv("S3_BUCKET")
+        )
+        if not bucket:
+            raise ValueError("S3 bucket not configured for Cognee service storage")
+
+        prefix = (
+            self.s3_prefix
+            or os.getenv("COGNEE_S3_PREFIX")
+            or f"cognee/projects/{self.project_context['project_id']}"
+        )
+        return bucket, prefix.rstrip("/")
+
+    def _get_s3_client(self):
+        if self._s3_client is not None:
+            return self._s3_client
+
+        try:
+            import boto3
+        except ImportError as exc:
+            raise RuntimeError("boto3 is required for Cognee S3 ingestion") from exc
+
+        endpoint = os.getenv("AWS_ENDPOINT_URL") or os.getenv("S3_ENDPOINT")
+        access_key = os.getenv("AWS_ACCESS_KEY_ID") or os.getenv("S3_ACCESS_KEY")
+        secret_key = os.getenv("AWS_SECRET_ACCESS_KEY") or os.getenv("S3_SECRET_KEY")
+        region = (
+            os.getenv("AWS_REGION")
+            or os.getenv("AWS_DEFAULT_REGION")
+            or os.getenv("S3_REGION")
+            or "us-east-1"
+        )
+        use_ssl = os.getenv("S3_USE_SSL", "false").lower() == "true"
+
+        if not (access_key and secret_key):
+            raise ValueError("AWS/S3 credentials missing for Cognee service ingestion")
+
+        self._s3_client = boto3.client(
+            "s3",
+            endpoint_url=endpoint,
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key,
+            region_name=region,
+            use_ssl=use_ssl,
+        )
+        return self._s3_client
+
+    def _upload_files_to_s3(self, file_paths: List[Path]) -> Tuple[List[str], Dict[str, str]]:
+        bucket, prefix = self._resolve_s3_paths()
+        client = self._get_s3_client()
+
+        timestamp = datetime.utcnow().strftime("%Y%m%dT%H%M%S")
+        try:
+            from uuid import uuid4
+        except ImportError:  # pragma: no cover - standard lib always available
+            uuid4 = None  # type: ignore
+
+        batch_suffix = uuid4().hex[:8] if uuid4 else "batch"
+        base_prefix = f"{prefix}/uploads/{self.project_context['project_id']}/{timestamp}_{batch_suffix}"
+
+        uploaded: Dict[str, str] = {}
+        uris: List[str] = []
+        for file_path in file_paths:
+            key = f"{base_prefix}/{self._relative_ingest_path(file_path)}"
+            with open(file_path, "rb") as handle:
+                client.upload_fileobj(handle, bucket, key)
+            uri = f"s3://{bucket}/{key}"
+            uploaded[str(file_path)] = uri
+            uris.append(uri)
+
+        return uris, uploaded
+
+    def _prepare_ingest_payload(self, file_paths: List[Path]) -> Tuple[List[str], Dict[str, str]]:
+        if self._uses_s3_storage():
+            return self._upload_files_to_s3(file_paths)
+        return [str(path) for path in file_paths], {}
     
     def get_project_dataset_name(self, dataset_suffix: str = "codebase") -> str:
         """Get project-specific dataset name"""
         return f"{self.project_context['project_name']}_{dataset_suffix}"
     
-    async def ingest_text(self, content: str, dataset: str = "fuzzforge") -> bool:
+    async def ingest_text(self, content: str, dataset: str = None) -> bool:
         """Ingest text content into knowledge graph"""
         if not self._initialized:
             await self.initialize()
+        dataset = dataset or self.get_project_dataset_name()
         
         try:
+            if self._using_service():
+                if not self._api_client:
+                    raise RuntimeError("Cognee API client not initialized")
+                await self._api_client.ensure_dataset(dataset)
+                await self._api_client.ingest_text(dataset, content)
+                await self._api_client.cognify(dataset)
+                return True
+
             await self._cognee.add([content], dataset)
             await self._cognee.cognify([dataset])
             return True
@@ -194,45 +433,76 @@ class CogneeService:
             logger.error(f"Failed to ingest text: {e}")
             return False
     
-    async def ingest_files(self, file_paths: List[Path], dataset: str = "fuzzforge") -> Dict[str, Any]:
+    async def ingest_files(self, file_paths: List[Path], dataset: str = None) -> Dict[str, Any]:
         """Ingest multiple files into knowledge graph"""
         if not self._initialized:
             await self.initialize()
+        self._sync_config_state()
+        dataset = dataset or self.get_project_dataset_name()
         
         results = {
             "success": 0,
             "failed": 0,
-            "errors": []
+            "errors": [],
+            "uploaded": {},
         }
-        
-        try:
-            ingest_paths: List[str] = []
-            for file_path in file_paths:
-                try:
-                    with open(file_path, 'r', encoding='utf-8'):
-                        ingest_paths.append(str(file_path))
-                    results["success"] += 1
-                except (UnicodeDecodeError, PermissionError) as exc:
-                    results["failed"] += 1
-                    results["errors"].append(f"{file_path}: {exc}")
-                    logger.warning("Skipping %s: %s", file_path, exc)
+        readable_files: List[Path] = []
 
-            if ingest_paths:
-                await self._cognee.add(ingest_paths, dataset_name=dataset)
+        for file_path in file_paths:
+            try:
+                with open(file_path, "r", encoding="utf-8"):
+                    pass
+                readable_files.append(Path(file_path))
+            except (UnicodeDecodeError, PermissionError) as exc:
+                results["failed"] += 1
+                results["errors"].append(f"{file_path}: {exc}")
+                logger.warning("Skipping %s: %s", file_path, exc)
+
+        if not readable_files:
+            return results
+
+        try:
+            if self._using_service():
+                if not self._api_client:
+                    raise RuntimeError("Cognee API client not initialized")
+                await self._api_client.ensure_dataset(dataset)
+                service_result = await self._api_client.ingest_files(dataset, readable_files)
+                await self._api_client.cognify(dataset)
+                results["success"] = len(readable_files)
+                results["service"] = service_result
+            else:
+                ingest_payload, uploaded = self._prepare_ingest_payload(readable_files)
+                await self._cognee.add(ingest_payload, dataset_name=dataset)
                 await self._cognee.cognify([dataset])
-            
+                results["success"] = len(readable_files)
+                if uploaded:
+                    results["uploaded"] = uploaded
         except Exception as e:
-            logger.error(f"Failed to ingest files: {e}")
+            logger.error("Failed to ingest files", exc_info=True)
             results["errors"].append(f"Cognify error: {str(e)}")
-        
+            results["failed"] += len(readable_files)
+            results["success"] = 0
+            return results
+
         return results
     
     async def search_insights(self, query: str, dataset: str = None) -> List[str]:
         """Search for insights in the knowledge graph"""
         if not self._initialized:
             await self.initialize()
-        
+
         try:
+            if self._using_service():
+                if not self._api_client:
+                    raise RuntimeError("Cognee API client not initialized")
+                dataset_name = dataset or self.get_project_dataset_name()
+                results = await self._api_client.search(
+                    query,
+                    datasets=[dataset_name] if dataset_name else None,
+                    search_type="INSIGHTS",
+                )
+                return self._normalise_service_results(results)
+
             from cognee.modules.search.types import SearchType
             
             kwargs = {
@@ -256,6 +526,17 @@ class CogneeService:
             await self.initialize()
         
         try:
+            if self._using_service():
+                if not self._api_client:
+                    raise RuntimeError("Cognee API client not initialized")
+                dataset_name = dataset or self.get_project_dataset_name()
+                results = await self._api_client.search(
+                    query,
+                    datasets=[dataset_name] if dataset_name else None,
+                    search_type="CHUNKS",
+                )
+                return self._normalise_service_results(results)
+
             from cognee.modules.search.types import SearchType
             
             kwargs = {
@@ -279,6 +560,16 @@ class CogneeService:
             await self.initialize()
         
         try:
+            if self._using_service():
+                if not self._api_client:
+                    raise RuntimeError("Cognee API client not initialized")
+                results = await self._api_client.search(
+                    query,
+                    datasets=[self.get_project_dataset_name()],
+                    search_type="GRAPH_COMPLETION",
+                )
+                return self._normalise_service_results(results)
+
             from cognee.modules.search.types import SearchType
             
             results = await self._cognee.search(
@@ -290,34 +581,77 @@ class CogneeService:
         except Exception as e:
             logger.error(f"Failed to search graph completion: {e}")
             return []
-    
+
     async def get_status(self) -> Dict[str, Any]:
         """Get service status and statistics"""
+        self._sync_config_state()
         status = {
             "initialized": self._initialized,
             "enabled": self.cognee_config.get("enabled", True),
             "provider": self.cognee_config.get("graph_database_provider", "kuzu"),
             "data_directory": self.cognee_config.get("data_directory"),
             "system_directory": self.cognee_config.get("system_directory"),
+            "storage_backend": self.storage_backend,
+            "service_mode": self.service_mode,
+            "s3_bucket": self.s3_bucket,
+            "s3_prefix": self.s3_prefix,
+            "dataset_name": self.get_project_dataset_name(),
         }
-        
-        if self._initialized:
+        if self._dataset_id:
+            status["dataset_id"] = self._dataset_id
+
+        if self._initialized and self._using_service() and self._api_client:
+            try:
+                status.update(await self._api_client.status())
+            except CogneeAPIError as exc:
+                status["service_error"] = str(exc)
+
+        if self._initialized and not self._uses_s3_storage() and not self._using_service():
             try:
-                # Check if directories exist and get sizes
                 data_dir = Path(status["data_directory"])
                 system_dir = Path(status["system_directory"])
-                
-                status.update({
-                    "data_dir_exists": data_dir.exists(),
-                    "system_dir_exists": system_dir.exists(),
-                    "kuzu_db_exists": (system_dir / "kuzu_db").exists(),
-                    "lancedb_exists": (system_dir / "lancedb").exists(),
-                })
-                
+
+                status.update(
+                    {
+                        "data_dir_exists": data_dir.exists(),
+                        "system_dir_exists": system_dir.exists(),
+                        "kuzu_db_exists": (system_dir / "kuzu_db").exists(),
+                        "lancedb_exists": (system_dir / "lancedb").exists(),
+                    }
+                )
+
             except Exception as e:
                 status["status_error"] = str(e)
-        
+        else:
+            status["s3_storage_active"] = self._uses_s3_storage()
+
         return status
+
+    def _normalise_service_results(self, payload: Any) -> List[str]:
+        if payload is None:
+            return []
+        if isinstance(payload, list):
+            return [self._stringify_service_item(item) for item in payload]
+        if isinstance(payload, dict):
+            if "results" in payload and isinstance(payload["results"], list):
+                return [self._stringify_service_item(item) for item in payload["results"]]
+            if "data" in payload and isinstance(payload["data"], list):
+                return [self._stringify_service_item(item) for item in payload["data"]]
+        return [self._stringify_service_item(payload)]
+
+    @staticmethod
+    def _stringify_service_item(item: Any) -> str:
+        if isinstance(item, str):
+            return item
+        try:
+            return json.dumps(item, ensure_ascii=False)
+        except Exception:
+            return str(item)
+
+    async def close(self) -> None:
+        if self._api_client:
+            await self._api_client.close()
+            self._api_client = None
     
     async def clear_data(self, confirm: bool = False):
         """Clear all ingested data (dangerous!)"""
@@ -327,6 +661,25 @@ class CogneeService:
         if not self._initialized:
             await self.initialize()
         
+        if self._using_service():
+            if not self._api_client:
+                raise RuntimeError("Cognee API client not initialized")
+            dataset_name = self.get_project_dataset_name()
+            try:
+                datasets = await self._api_client.list_datasets()
+                target = next((d for d in datasets if d.get("name") == dataset_name), None)
+                if target and target.get("id"):
+                    await self._api_client.delete_dataset(target["id"])
+                    logger.info("Deleted Cognee dataset via service", extra={"dataset": dataset_name})
+                else:
+                    logger.info(
+                        "No matching Cognee dataset found to delete", extra={"dataset": dataset_name}
+                    )
+            except Exception as exc:
+                logger.error("Failed to clear Cognee dataset via service", exc_info=True)
+                raise
+            return
+
         try:
             await self._cognee.prune.prune_data()
             await self._cognee.prune.prune_system(metadata=True)

ai/src/fuzzforge_ai/mcp/__init__.py

@@ -0,0 +1 @@
+"""MCP server utilities for FuzzForge."""

ai/src/fuzzforge_ai/mcp/cognee_mcp_server.py

@@ -0,0 +1,259 @@
+"""Multi-tenant aware MCP server backed by the Cognee service API."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import logging
+import os
+from pathlib import Path
+from typing import List, Optional
+
+from dotenv import load_dotenv
+from mcp.server import FastMCP
+import mcp.types as types
+
+from fuzzforge_ai.cognee_api_client import CogneeAPIClient, CogneeAPIError
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=os.getenv("MCP_LOG_LEVEL", "INFO"))
+
+mcp = FastMCP("FuzzForge-Cognee")
+
+
+def _load_project_env() -> None:
+    """Load environment variables from `.fuzzforge/.env` if present."""
+
+    env_path = Path.cwd() / ".fuzzforge" / ".env"
+    if env_path.exists():
+        load_dotenv(env_path, override=False)
+
+
+class CogneeMCPContext:
+    """Holds shared state for MCP tools and manages the Cognee API client."""
+
+    def __init__(self) -> None:
+        _load_project_env()
+        self.service_url = os.getenv("COGNEE_SERVICE_URL", "").rstrip("/")
+        self.verify_ssl = os.getenv("COGNEE_VERIFY_SSL", "true").lower() != "false"
+        self.email = os.getenv("COGNEE_SERVICE_USER_EMAIL") or os.getenv("DEFAULT_USER_EMAIL")
+        self.password = os.getenv("COGNEE_SERVICE_USER_PASSWORD") or os.getenv("DEFAULT_USER_PASSWORD")
+        if not (self.service_url and self.email and self.password):
+            raise RuntimeError(
+                "COGNEE_SERVICE_URL, COGNEE_SERVICE_USER_EMAIL and"
+                " COGNEE_SERVICE_USER_PASSWORD must be set in the environment"
+            )
+        self._client: Optional[CogneeAPIClient] = None
+        self.active_dataset_name: Optional[str] = os.getenv("COGNEE_DATASET_NAME") or self._default_dataset_name()
+        self.active_dataset_id: Optional[str] = None
+
+    async def start(self) -> None:
+        self._client = CogneeAPIClient(self.service_url, verify_ssl=self.verify_ssl)
+        await self._client.ensure_user(self.email, self.password, tenant_name=self._tenant_name())
+        if self.active_dataset_name:
+            dataset_info = await self._client.ensure_dataset(self.active_dataset_name)
+            self.active_dataset_id = dataset_info.get("id")
+        logger.info(
+            "Connected to Cognee service",
+            extra={
+                "service_url": self.service_url,
+                "dataset": self.active_dataset_name,
+                "dataset_id": self.active_dataset_id,
+            },
+        )
+
+    async def close(self) -> None:
+        if self._client:
+            await self._client.close()
+
+    @property
+    def client(self) -> CogneeAPIClient:
+        if not self._client:
+            raise RuntimeError("Cognee API client is not initialised yet")
+        return self._client
+
+    def _tenant_name(self) -> str:
+        return os.getenv("COGNEE_TENANT_NAME") or os.getenv("COGNEE_PROJECT_ID") or "default-tenant"
+
+    def _default_dataset_name(self) -> str:
+        project_name = os.getenv("FUZZFORGE_PROJECT_NAME") or "project"
+        return f"{project_name}_codebase"
+
+    async def ensure_dataset(self, dataset_name: str) -> str:
+        """Ensure the dataset exists and cache its id."""
+
+        info = await self.client.ensure_dataset(dataset_name)
+        dataset_id = info.get("id")
+        if not dataset_id:
+            raise RuntimeError(f"Failed to resolve dataset id for '{dataset_name}'")
+        self.active_dataset_name = dataset_name
+        self.active_dataset_id = dataset_id
+        return dataset_id
+
+    async def resolve_dataset(self, dataset: Optional[str]) -> tuple[str, str]:
+        name = dataset or self.active_dataset_name
+        if not name:
+            raise ValueError("Dataset is not specified; set COGNEE_DATASET_NAME or call set_dataset first")
+        # if the requested dataset differs from the cached one, ensure it
+        if dataset and dataset != self.active_dataset_name:
+            dataset_id = await self.ensure_dataset(dataset)
+        else:
+            dataset_id = self.active_dataset_id
+            if not dataset_id:
+                dataset_id = await self.ensure_dataset(name)
+        return name, dataset_id
+
+
+CONTEXT = CogneeMCPContext()
+
+
+def _text(message: str) -> List[types.TextContent]:
+    return [types.TextContent(type="text", text=message)]
+
+
+@mcp.tool()
+async def set_dataset(dataset: str) -> List[types.TextContent]:
+    """Switch the active dataset for subsequent operations."""
+
+    dataset_id = await CONTEXT.ensure_dataset(dataset)
+    return _text(f"Active dataset set to '{dataset}' (id: {dataset_id}).")
+
+
+@mcp.tool()
+async def ingest_file(path: str, dataset: Optional[str] = None) -> List[types.TextContent]:
+    """Upload a local file into the Cognee dataset."""
+
+    file_path = Path(path).expanduser().resolve()
+    if not file_path.is_file():
+        raise ValueError(f"File not found: {file_path}")
+
+    dataset_name, _ = await CONTEXT.resolve_dataset(dataset)
+    result = await CONTEXT.client.ingest_files(dataset_name, [file_path])
+    success = result.get("success", 0)
+    failed = result.get("failed", 0)
+    errors = result.get("errors", [])
+
+    message = [
+        f"Uploaded {success} file(s) to dataset '{dataset_name}'.",
+        f"Failed: {failed}",
+    ]
+    if errors:
+        message.append("Errors:\n" + "\n".join(errors))
+
+    return _text("\n".join(message))
+
+
+@mcp.tool()
+async def cognify(run_in_background: bool = False, dataset: Optional[str] = None) -> List[types.TextContent]:
+    """Trigger the Cognify pipeline for the dataset."""
+
+    dataset_name, _ = await CONTEXT.resolve_dataset(dataset)
+    response = await CONTEXT.client.cognify(dataset_name, run_in_background=run_in_background)
+    return _text(json.dumps(response, indent=2))
+
+
+@mcp.tool()
+async def cognify_status(dataset: Optional[str] = None) -> List[types.TextContent]:
+    """Return the status of recent Cognify runs."""
+
+    _, dataset_id = await CONTEXT.resolve_dataset(dataset)
+    status = await CONTEXT.client.status()
+    status.setdefault("dataset_ids", []).append(dataset_id)
+    return _text(json.dumps(status, indent=2))
+
+
+@mcp.tool()
+async def search_chunks(query: str, top_k: int = 5, dataset: Optional[str] = None) -> List[types.TextContent]:
+    """Search dataset for relevant document chunks."""
+
+    dataset_name, _ = await CONTEXT.resolve_dataset(dataset)
+    result = await CONTEXT.client.search(
+        query,
+        datasets=[dataset_name],
+        search_type="CHUNKS",
+        top_k=top_k,
+    )
+    return _text(json.dumps(result, indent=2))
+
+
+@mcp.tool()
+async def search_insights(query: str, top_k: int = 5, dataset: Optional[str] = None) -> List[types.TextContent]:
+    """Search dataset for graph insights."""
+
+    dataset_name, _ = await CONTEXT.resolve_dataset(dataset)
+    result = await CONTEXT.client.search(
+        query,
+        datasets=[dataset_name],
+        search_type="INSIGHTS",
+        top_k=top_k,
+    )
+    return _text(json.dumps(result, indent=2))
+
+
+@mcp.tool()
+async def list_datasets() -> List[types.TextContent]:
+    """List datasets accessible to the current service user."""
+
+    datasets = await CONTEXT.client.list_datasets()
+    return _text(json.dumps(datasets, indent=2))
+
+
+@mcp.tool()
+async def prune_dataset(dataset: Optional[str] = None) -> List[types.TextContent]:
+    """Delete the dataset and recreate it empty."""
+
+    dataset_name, dataset_id = await CONTEXT.resolve_dataset(dataset)
+    await CONTEXT.client.delete_dataset(dataset_id)
+    dataset_info = await CONTEXT.client.ensure_dataset(dataset_name)
+    CONTEXT.active_dataset_name = dataset_name
+    CONTEXT.active_dataset_id = dataset_info.get("id")
+    return _text(f"Dataset '{dataset_name}' has been pruned and recreated.")
+
+
+async def _async_main(args: argparse.Namespace) -> None:
+    try:
+        await CONTEXT.start()
+    except CogneeAPIError as exc:
+        raise SystemExit(f"Failed to initialise Cognee MCP: {exc}") from exc
+
+    transport = args.transport
+    host = args.host
+    port = args.port
+    path = args.path
+
+    logger.info(
+        "Starting Cognee MCP server",
+        extra={
+            "transport": transport,
+            "host": host,
+            "port": port,
+            "path": path,
+            "dataset": CONTEXT.dataset_name,
+        },
+    )
+
+    try:
+        if transport == "http":
+            await mcp.run_http(host=host, port=port, path=path)
+        elif transport == "sse":
+            await mcp.run_sse(host=host, port=port, path=path)
+        else:
+            await mcp.run_stdio()
+    finally:
+        await CONTEXT.close()
+
+
+def main(argv: Optional[List[str]] = None) -> None:
+    parser = argparse.ArgumentParser(description="Cognee MCP server (service-backed)")
+    parser.add_argument("--transport", choices=["stdio", "sse", "http"], default=os.getenv("TRANSPORT_MODE", "stdio"))
+    parser.add_argument("--host", default=os.getenv("MCP_HOST", "127.0.0.1"))
+    parser.add_argument("--port", type=int, default=int(os.getenv("MCP_PORT", "8000")))
+    parser.add_argument("--path", default=os.getenv("MCP_HTTP_PATH", "/mcp"))
+    args = parser.parse_args(argv)
+
+    asyncio.run(_async_main(args))
+
+
+if __name__ == "__main__":
+    main()