CM-73403: Address review: typed models and one session cache check

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: Altruistus

cycode/cli/apps/ai_guardrails/scan/mcp_server_status.py

@@ -5,27 +5,23 @@
 
 import json
 import time
-from dataclasses import InitVar, dataclass, field
-from enum import Enum
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional
+from typing import Any, Optional
+
+from marshmallow import EXCLUDE, Schema, fields, post_load
 
 from cycode.cli.apps.ai_guardrails.scan.guardrail_config import DEFAULT_TTL_SECONDS
 from cycode.cli.consts import CYCODE_CONFIGURATION_DIRECTORY
 from cycode.cli.utils.path_utils import atomic_write_text, quarantine_corrupt_file
+from cycode.cyclient.models import McpServerAuthorizationStatus, McpServerStatus, McpServerStatusSchema
 from cycode.logger import get_logger
 
 logger = get_logger('AI Guardrails')
 
 MCP_SERVER_STATUSES_FILE_NAME = 'ai-guardrails-mcp-servers.json'
 
 
-class McpServerAuthorizationStatus(str, Enum):
-    AUTHORIZED = 'Authorized'
-    UNREVIEWED = 'Unreviewed'
-    UNAUTHORIZED = 'Unauthorized'
-
-
 # An alias shared by several servers (e.g. configured differently per device) gets the most restrictive status.
 _RESTRICTIVENESS = {
     McpServerAuthorizationStatus.AUTHORIZED: 0,
@@ -38,61 +34,57 @@ def get_mcp_server_statuses_cache_path() -> Path:
     return Path.home() / CYCODE_CONFIGURATION_DIRECTORY / MCP_SERVER_STATUSES_FILE_NAME
 
 
-def parse_status(raw_status: object) -> McpServerAuthorizationStatus:
-    """An unknown status reads as Unreviewed: no decision was made on the server."""
-    lowered = str(raw_status or '').lower()
-    for status in McpServerAuthorizationStatus:
-        if status.value.lower() == lowered:
-            return status
-    return McpServerAuthorizationStatus.UNREVIEWED
-
-
 def is_enforced(status: Optional[McpServerAuthorizationStatus]) -> bool:
     """``status`` is None when the platform never saw the server."""
     return status == McpServerAuthorizationStatus.UNAUTHORIZED
 
 
 @dataclass
 class McpServerStatuses:
-    servers: InitVar[list]
+    servers: list[McpServerStatus]
     fetched_at: float
     tenant_id: Optional[str] = None
     ttl_seconds: float = DEFAULT_TTL_SECONDS
     _by_alias: dict = field(init=False, repr=False)
 
-    def __post_init__(self, servers: list) -> None:
+    def __post_init__(self) -> None:
         # The platform matches aliases case-insensitively, so the CLI does too.
         self._by_alias = {}
-        for server in servers:
-            if not isinstance(server, dict) or not server.get('alias'):
+        for server in self.servers:
+            if not server.alias:
                 continue
-            key = str(server['alias']).lower()
-            status = parse_status(server.get('status'))
+            key = server.alias.lower()
             known = self._by_alias.get(key)
-            if known is None or _RESTRICTIVENESS[status] > _RESTRICTIVENESS[known]:
-                self._by_alias[key] = status
+            if known is None or _RESTRICTIVENESS[server.status] > _RESTRICTIVENESS[known]:
+                self._by_alias[key] = server.status
 
     def status_of(self, alias: str) -> Optional[McpServerAuthorizationStatus]:
         return self._by_alias.get(alias.lower())
 
-    def is_expired(self) -> bool:
-        return time.time() - self.fetched_at > self.ttl_seconds
 
-    def needs_refresh(self, tenant_id: Optional[str]) -> bool:
-        return self.is_expired() or self.tenant_id != tenant_id
+class McpServerStatusesSchema(Schema):
+    class Meta:
+        unknown = EXCLUDE
+
+    servers = fields.List(fields.Nested(McpServerStatusSchema), required=True)
+    fetched_at = fields.Float(required=True)
+    tenant_id = fields.String(allow_none=True, load_default=None)
+    ttl_seconds = fields.Float(allow_none=True, load_default=None)
+
+    @post_load
+    def build_dto(self, data: dict[str, Any], **_) -> McpServerStatuses:
+        data['ttl_seconds'] = data['ttl_seconds'] or DEFAULT_TTL_SECONDS
+        return McpServerStatuses(**data)
 
 
-def save_mcp_server_statuses(servers: list, tenant_id: Optional[str], ttl_seconds: float = DEFAULT_TTL_SECONDS) -> None:
+def save_mcp_server_statuses(
+    servers: list[McpServerStatus], tenant_id: Optional[str], ttl_seconds: float = DEFAULT_TTL_SECONDS
+) -> None:
     path = get_mcp_server_statuses_cache_path()
-    content = {
-        'fetched_at': time.time(),
-        'tenant_id': tenant_id,
-        'ttl_seconds': ttl_seconds,
-        'servers': servers,
-    }
+    statuses = McpServerStatuses(servers=servers, fetched_at=time.time(), tenant_id=tenant_id, ttl_seconds=ttl_seconds)
     try:
         path.parent.mkdir(parents=True, exist_ok=True)
-        atomic_write_text(str(path), json.dumps(content))
+        atomic_write_text(str(path), json.dumps(McpServerStatusesSchema().dump(statuses)))
     except Exception as e:
         logger.debug('Failed to save MCP server statuses cache', exc_info=e)
 
@@ -104,16 +96,7 @@ def load_mcp_server_statuses() -> Optional[McpServerStatuses]:
 
     try:
         with open(path, encoding='UTF-8') as file:
-            content = json.load(file)
-        servers = content['servers']
-        if not isinstance(servers, list):
-            raise ValueError('servers is not a list')
-        return McpServerStatuses(
-            servers=servers,
-            fetched_at=float(content['fetched_at']),
-            tenant_id=content.get('tenant_id'),
-            ttl_seconds=float(content.get('ttl_seconds') or DEFAULT_TTL_SECONDS),
-        )
+            return McpServerStatusesSchema().load(json.load(file))
     except Exception as e:
         logger.warning('MCP server statuses cache is corrupt and will be moved aside', exc_info=e)
         quarantine_corrupt_file(str(path))

cycode/cli/apps/ai_guardrails/session_start_command.py

@@ -21,10 +21,7 @@
     load_guardrail_config,
     save_guardrail_config,
 )
-from cycode.cli.apps.ai_guardrails.scan.mcp_server_status import (
-    load_mcp_server_statuses,
-    save_mcp_server_statuses,
-)
+from cycode.cli.apps.ai_guardrails.scan.mcp_server_status import save_mcp_server_statuses
 from cycode.cli.apps.ai_guardrails.scan.types import BlockReason
 from cycode.cli.apps.ai_guardrails.scan.utils import read_stdin_text, safe_json_parse
 from cycode.cli.apps.auth.auth_common import get_authorization_info
@@ -59,14 +56,21 @@ def _session_context_digest(report: dict) -> str:
     return hashlib.sha256(canonical.encode('utf-8')).hexdigest()
 
 
-def _should_skip_report(digest: str, tenant_id: Optional[str]) -> bool:
+def _report_ttl_seconds(config: Optional[GuardrailConfig]) -> float:
+    """The response carries the MCP server statuses, so the guardrail's TTL applies while it is enabled."""
+    if config is None or config.is_off_for_every_agent(BlockReason.UNAUTHORIZED_MCP_SERVER):
+        return _SESSION_CONTEXT_TTL_SECONDS
+    return config.ttl_seconds
+
+
+def _should_skip_report(digest: str, tenant_id: Optional[str], ttl_seconds: float) -> bool:
     """Skip when the same payload was already sent for this tenant and the TTL hasn't expired."""
     try:
         cache = json.loads(_session_context_cache_path().read_text(encoding='utf-8'))
         return (
             cache.get('hash') == digest
             and cache.get('tenant_id') == tenant_id
-            and time.time() - float(cache.get('sent_at', 0)) < _SESSION_CONTEXT_TTL_SECONDS
+            and time.time() - float(cache.get('sent_at', 0)) < ttl_seconds
         )
     except Exception:
         # Missing/corrupt cache reads as a miss - over-sending is harmless
@@ -84,13 +88,6 @@ def _save_report_cache(digest: str, tenant_id: Optional[str]) -> None:
         logger.debug('Failed to write session context cache', exc_info=e)
 
 
-def _needs_mcp_server_statuses(config: Optional[GuardrailConfig], tenant_id: Optional[str]) -> bool:
-    if config is None or config.is_off_for_every_agent(BlockReason.UNAUTHORIZED_MCP_SERVER):
-        return False
-    cached = load_mcp_server_statuses()
-    return cached is None or cached.needs_refresh(tenant_id)
-
-
 def _report_session_context(
     ai_client: 'AISecurityManagerClient',
     user_email: Optional[str],
@@ -101,7 +98,7 @@ def _report_session_context(
 
     The device context is always reported. MCP configs and skills are collected from every
     registered IDE, not just the triggering one. Unchanged payloads are skipped via a hash cache
-    until the TTL expires, unless the MCP server statuses the response carries are needed.
+    until the TTL expires.
     """
     try:
         config_files, enabled_plugins = collect_all_session_contexts()
@@ -121,7 +118,7 @@ def _report_session_context(
         }
 
         digest = _session_context_digest(report)
-        if _should_skip_report(digest, tenant_id) and not _needs_mcp_server_statuses(config, tenant_id):
+        if _should_skip_report(digest, tenant_id, _report_ttl_seconds(config)):
             logger.debug('Session context unchanged; skipping report')
             return
 

cycode/cyclient/ai_security_manager_client.py

@@ -1,21 +1,18 @@
 """Client for AI Security Manager service."""
 
-from typing import TYPE_CHECKING, NamedTuple, Optional
+from typing import TYPE_CHECKING, Optional
 
 from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError
 from cycode.cyclient.cycode_client_base import CycodeClientBase
 from cycode.cyclient.logger import logger
+from cycode.cyclient.models import SessionContextResponse, SessionContextResponseSchema
 
 if TYPE_CHECKING:
     from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload
     from cycode.cli.apps.ai_guardrails.scan.types import AiHookEventType, AIHookOutcome, BlockReason
     from cycode.cyclient.ai_security_manager_service_config import AISecurityManagerServiceConfigBase
 
 
-class SessionContextResponse(NamedTuple):
-    mcp_servers: Optional[list]  # [{alias, normalized_id, status}]; None when missing or malformed
-
-
 class AISecurityManagerClient:
     """Client for interacting with AI Security Manager service."""
 
@@ -140,10 +137,7 @@ def report_session_context(
             return None
 
         try:
-            mcp_servers = response.json().get('mcp_servers')
-            if not isinstance(mcp_servers, list):
-                raise ValueError('mcp_servers is not a list')
+            return SessionContextResponseSchema().load(response.json())
         except Exception as e:
             logger.debug('Failed to parse the session context response', exc_info=e)
-            mcp_servers = None
-        return SessionContextResponse(mcp_servers=mcp_servers)
+            return SessionContextResponse()

cycode/cyclient/models.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass
+from enum import Enum
 from typing import Any, Optional
 
-from marshmallow import EXCLUDE, Schema, fields, post_load
+from marshmallow import EXCLUDE, Schema, fields, missing, post_load
 
 
 class Detection(Schema):
@@ -578,3 +579,60 @@ class Meta:
     @post_load
     def build_dto(self, data: dict[str, Any], **_) -> 'ScanConfiguration':
         return ScanConfiguration(**data)
+
+
+class McpServerAuthorizationStatus(str, Enum):
+    AUTHORIZED = 'Authorized'
+    UNREVIEWED = 'Unreviewed'
+    UNAUTHORIZED = 'Unauthorized'
+
+    @classmethod
+    def parse(cls, raw_status: object) -> 'McpServerAuthorizationStatus':
+        """An unknown status reads as Unreviewed: no decision was made on the server."""
+        lowered = str(raw_status or '').lower()
+        return next((status for status in cls if status.value.lower() == lowered), cls.UNREVIEWED)
+
+
+class McpServerAuthorizationStatusField(fields.Field):
+    def _serialize(self, value: Optional[McpServerAuthorizationStatus], *_, **__) -> Optional[str]:
+        return value.value if value is not None else None
+
+    def deserialize(self, value: object, *_, **__) -> McpServerAuthorizationStatus:
+        return McpServerAuthorizationStatus.parse(None if value is missing else value)
+
+
+@dataclass
+class McpServerStatus:
+    alias: Optional[str] = None
+    normalized_id: Optional[str] = None
+    status: McpServerAuthorizationStatus = McpServerAuthorizationStatus.UNREVIEWED
+
+
+class McpServerStatusSchema(Schema):
+    class Meta:
+        unknown = EXCLUDE
+
+    alias = fields.String(allow_none=True, load_default=None)
+    normalized_id = fields.String(allow_none=True, load_default=None)
+    status = McpServerAuthorizationStatusField()
+
+    @post_load
+    def build_dto(self, data: dict[str, Any], **_) -> McpServerStatus:
+        return McpServerStatus(**data)
+
+
+@dataclass
+class SessionContextResponse:
+    # None when the response carries no (or malformed) MCP server statuses.
+    mcp_servers: Optional[list[McpServerStatus]] = None
+
+
+class SessionContextResponseSchema(Schema):
+    class Meta:
+        unknown = EXCLUDE
+
+    mcp_servers = fields.List(fields.Nested(McpServerStatusSchema), allow_none=True, load_default=None)
+
+    @post_load
+    def build_dto(self, data: dict[str, Any], **_) -> SessionContextResponse:
+        return SessionContextResponse(**data)

tests/cli/commands/ai_guardrails/scan/test_handlers.py

@@ -28,7 +28,7 @@
 from cycode.cli.apps.ai_guardrails.scan.types import AiHookEventType, AIHookOutcome, BlockReason
 from cycode.cli.apps.ai_guardrails.scan.utils import MAX_VIOLATION_DETAIL_LINES, build_violation_summary
 from cycode.cli.models import Document, DocumentDetections, LocalScanResult
-from cycode.cyclient.models import Detection
+from cycode.cyclient.models import Detection, McpServerAuthorizationStatus, McpServerStatus
 
 
 @pytest.fixture
@@ -604,7 +604,8 @@ def _mcp_payload(server: Optional[str] = 'github') -> AIHookPayload:
 
 
 def _cached_statuses(*rows: tuple[str, str]) -> McpServerStatuses:
-    return McpServerStatuses(servers=[{'alias': a, 'status': s} for a, s in rows], fetched_at=time.time())
+    servers = [McpServerStatus(alias=a, status=McpServerAuthorizationStatus.parse(s)) for a, s in rows]
+    return McpServerStatuses(servers=servers, fetched_at=time.time())
 
 
 def _reported_event(mock_ctx: MagicMock) -> tuple[AIHookOutcome, Optional[BlockReason]]: