voodoo-framework 2.6.0__tar.gz → 2.6.1__tar.gz

{voodoo_framework-2.6.0/src/voodoo_framework.egg-info → voodoo_framework-2.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: voodoo-framework
-Version: 2.6.0
+Version: 2.6.1
 Summary: The programmable runtime for adaptive applications and operational systems — web, APIs, agents, workers, human workflows, distributed systems, and physical systems in one runtime.
 Author-email: Helder Perez <contact@helderperez.com>
 Maintainer-email: Helder Perez <contact@helderperez.com>

{voodoo_framework-2.6.0 → voodoo_framework-2.6.1}/src/voodoo/__init__.py

@@ -102,7 +102,7 @@ from .ui.styles import (
 from .ui.styles.theme import Theme, ThemeColors, create_theme
 from .workers import task
 
-__version__ = "2.6.0"
+__version__ = "2.6.1"
 
 __all__ = [
     # Core runtime

{voodoo_framework-2.6.0 → voodoo_framework-2.6.1}/src/voodoo/config.py

@@ -284,6 +284,11 @@ class EdgeConfig(BaseModel):
     max_state_size: int = 65536  # 64 KiB
     max_pending_effects: int = 1000
     heartbeat_interval: int = 30
+    # Enrollment security (Sprint 23.1).
+    enrollment_auth_required: bool = True
+    enrollment_admin_token: str = ""
+    # Effect delivery (Sprint 23.1).
+    max_effect_retries: int = 3
 
 
 class VoodooConfig(BaseModel):

{voodoo_framework-2.6.0 → voodoo_framework-2.6.1}/src/voodoo/edge/__init__.py

@@ -27,6 +27,8 @@ from voodoo.edge.auth import (
 from voodoo.edge.errors import (
     AuthenticationFailedError,
     AuthorizationFailedError,
+    DeliveryFailedError,
+    DeviceIdMismatchError,
     DeviceNotFoundError,
     DeviceRevokedError,
     DuplicateMessageError,
@@ -37,6 +39,9 @@ from voodoo.edge.errors import (
     InvalidMessageError,
     InvalidProtocolVersionError,
     InvalidStateVersionError,
+    PayloadTooLargeError,
+    SessionInvalidError,
+    SessionRequiredError,
     TransportError,
     error_response,
 )
@@ -109,7 +114,9 @@ __all__ = [
     "error_response",
     "AuthenticationFailedError",
     "AuthorizationFailedError",
+    "DeliveryFailedError",
     "DeviceNotFoundError",
+    "DeviceIdMismatchError",
     "DeviceRevokedError",
     "DuplicateMessageError",
     "EffectExpiredError",
@@ -118,5 +125,8 @@ __all__ = [
     "InvalidMessageError",
     "InvalidProtocolVersionError",
     "InvalidStateVersionError",
+    "PayloadTooLargeError",
+    "SessionInvalidError",
+    "SessionRequiredError",
     "TransportError",
 ]

{voodoo_framework-2.6.0 → voodoo_framework-2.6.1}/src/voodoo/edge/auth.py

@@ -223,6 +223,13 @@ async def authenticate_device(
 
     await store.mark_credential_used(record.credential_id)
 
+    # Reconnect: invalidate any existing sessions for this device before
+    # creating a new one.  The device transitions RECONNECTING → CONNECTED
+    # so downstream observers see a clean lifecycle signal (EDGE §30).
+    old_session_count = await store.delete_device_sessions(device.device_id)
+    if old_session_count > 0:
+        await store.update_device_status(device.device_id, DeviceStatus.RECONNECTING)
+
     session = DeviceSession(
         device_id=device.device_id,
         transport=transport,

{voodoo_framework-2.6.0 → voodoo_framework-2.6.1}/src/voodoo/edge/errors.py

@@ -17,6 +17,7 @@ __all__ = [
     "AuthenticationFailedError",
     "AuthorizationFailedError",
     "DeviceNotFoundError",
+    "DeviceIdMismatchError",
     "DeviceRevokedError",
     "InvalidMessageError",
     "InvalidProtocolVersionError",
@@ -25,6 +26,9 @@ __all__ = [
     "DuplicateMessageError",
     "EffectNotFoundError",
     "EffectExpiredError",
+    "SessionRequiredError",
+    "SessionInvalidError",
+    "DeliveryFailedError",
     "TransportError",
     "HTTP_STATUS",
     "error_response",
@@ -50,6 +54,9 @@ class EdgeErrorCode:
     DUPLICATE_MESSAGE = "DUPLICATE_MESSAGE"
     EFFECT_NOT_FOUND = "EFFECT_NOT_FOUND"
     EFFECT_EXPIRED = "EFFECT_EXPIRED"
+    SESSION_REQUIRED = "SESSION_REQUIRED"
+    SESSION_INVALID = "SESSION_INVALID"
+    DELIVERY_FAILED = "DELIVERY_FAILED"
     TRANSPORT_ERROR = "TRANSPORT_ERROR"
 
 
@@ -66,6 +73,9 @@ HTTP_STATUS: dict[str, int] = {
     EdgeErrorCode.DUPLICATE_MESSAGE: 200,  # idempotent success, not an error for the client
     EdgeErrorCode.EFFECT_NOT_FOUND: 404,
     EdgeErrorCode.EFFECT_EXPIRED: 410,
+    EdgeErrorCode.SESSION_REQUIRED: 401,
+    EdgeErrorCode.SESSION_INVALID: 401,
+    EdgeErrorCode.DELIVERY_FAILED: 502,
     EdgeErrorCode.TRANSPORT_ERROR: 502,
 }
 
@@ -163,6 +173,41 @@ class TransportError(EdgeError):
     status = 502
 
 
+class SessionRequiredError(AuthenticationFailedError):
+    """Request requires an authenticated session."""
+
+    code = EdgeErrorCode.SESSION_REQUIRED
+    status = 401
+
+
+class SessionInvalidError(AuthenticationFailedError):
+    """Session ID is invalid or expired."""
+
+    code = EdgeErrorCode.SESSION_INVALID
+    status = 401
+
+
+class DeliveryFailedError(TransportError):
+    """Effect delivery failed after max retries."""
+
+    code = EdgeErrorCode.DELIVERY_FAILED
+    status = 502
+
+
+class DeviceIdMismatchError(AuthorizationFailedError):
+    """Message device_id does not match the authenticated device."""
+
+    code = EdgeErrorCode.AUTHORIZATION_FAILED
+
+
+class PayloadTooLargeError(InvalidMessageError):
+    """Message or state payload exceeds configured size limit."""
+
+    code = EdgeErrorCode.INVALID_MESSAGE
+    status = 413
+    status = 403
+
+
 # ---------------------------------------------------------------------------
 # Serialization
 # ---------------------------------------------------------------------------

{voodoo_framework-2.6.0 → voodoo_framework-2.6.1}/src/voodoo/edge/gateway.py

@@ -22,10 +22,12 @@ from voodoo.edge.auth import authenticate_device, consume_enrollment
 from voodoo.edge.errors import (
     AuthenticationFailedError,
     AuthorizationFailedError,
+    DeviceIdMismatchError,
     DeviceNotFoundError,
     EffectNotFoundError,
     InvalidMessageError,
     InvalidStateVersionError,
+    PayloadTooLargeError,
 )
 from voodoo.edge.models import (
     AuthenticatedDeviceContext,
@@ -50,6 +52,17 @@ logger = logging.getLogger("voodoo.edge")
 __all__ = ["DeviceGateway"]
 
 
+def _redact(value: str, *, show: int = 4) -> str:
+    """Redact a sensitive value, showing only the last *show* characters.
+
+    Used for debug logging of device IDs and session tokens — credentials
+    are never logged (EDGE §9).
+    """
+    if len(value) <= show:
+        return "***"
+    return f"{'*' * (len(value) - show)}{value[-show:]}"
+
+
 # Events that are pure telemetry — never become Executions (EDGE §21).
 _TELEMETRY_EVENTS = frozenset({"heartbeat", "device.heartbeat"})
 
@@ -68,10 +81,15 @@ class DeviceGateway:
     """
 
     def __init__(
-        self, store: DeviceStoreProtocol, engine: ExecutionEngine | None = None
+        self,
+        store: DeviceStoreProtocol,
+        engine: ExecutionEngine | None = None,
+        *,
+        config: Any | None = None,
     ) -> None:
         self._store = store
         self._engine = engine
+        self._config = config
         # Authenticated contexts by session_id — reconnects re-attach.
         self._contexts: dict[str, AuthenticatedDeviceContext] = {}
 
@@ -117,7 +135,8 @@ class DeviceGateway:
         This is the canonical way transports authenticate — direct calls
         to ``authenticate_device`` bypass session tracking. Returns
         ``(context, session)`` with the context registered for subsequent
-        messages.
+        messages.  On reconnect, stale contexts for the same device are
+        evicted first.
         """
 
         ctx, session = await authenticate_device(
@@ -126,6 +145,12 @@ class DeviceGateway:
             claimed_device_id=claimed_device_id,
             transport=transport,
         )
+        # Evict any stale contexts for this device (reconnect cleanup).
+        stale_sids = [
+            sid for sid, c in self._contexts.items() if c.device_id == ctx.device_id
+        ]
+        for sid in stale_sids:
+            self._contexts.pop(sid, None)
         self._contexts[session.session_id] = ctx
         await self._emit(
             "device.connected",
@@ -138,6 +163,31 @@ class DeviceGateway:
         )
         return ctx, session
 
+    # ------------------------------------------------------------------
+    # Stateless authentication (HTTP — no session created)
+    # ------------------------------------------------------------------
+
+    async def authenticate_request(
+        self,
+        credential: str,
+        *,
+        transport: TransportKind = TransportKind.HTTP,
+    ) -> AuthenticatedDeviceContext:
+        """Authenticate a device credential without creating a session.
+
+        For stateless transports (HTTP) where each request carries its
+        own credential. Validates the credential, checks the device is
+        not revoked, and updates ``last_seen_at``. Returns the
+        authenticated context only — no session is persisted.
+        """
+        ctx, _session = await authenticate_device(
+            self._store,
+            credential,
+            transport=transport,
+        )
+        await self._store.update_last_seen(ctx.device_id)
+        return ctx
+
     # ------------------------------------------------------------------
     # Message entry point (EDGE §72 — thin route → gateway)
     # ------------------------------------------------------------------
@@ -158,11 +208,34 @@ class DeviceGateway:
         """
         message = decode_message(raw)
 
+        # Resource limit: message size (EDGE §76).
+        if self._config is not None:
+            edge_cfg = getattr(self._config, "edge", None)
+            if edge_cfg is not None:
+                max_size = getattr(edge_cfg, "max_message_size", 0)
+                if max_size > 0:
+                    raw_size = (
+                        len(raw) if isinstance(raw, (str, bytes)) else len(str(raw))
+                    )
+                    if raw_size > max_size:
+                        raise PayloadTooLargeError(
+                            f"Message size {raw_size} exceeds limit {max_size}"
+                        )
+
         # Protocol-version gate happens inside decode; re-check payload types.
         if message.type is EdgeMessageType.AUTH:
             return await self.handle_auth(message, transport=transport)
 
         ctx = context or await self._require_context(message)
+
+        # Identity validation: message device_id must match the
+        # authenticated context (EDGE §77).
+        if message.device_id and message.device_id != ctx.device_id:
+            raise DeviceIdMismatchError(
+                f"Message device_id '{message.device_id}' does not match "
+                f"authenticated device '{ctx.device_id}'"
+            )
+
         handler = {
             EdgeMessageType.HELLO: self.handle_hello,
             EdgeMessageType.STATE_SYNC: self.handle_state_sync,
@@ -236,13 +309,18 @@ class DeviceGateway:
         """Ingest a device event (EDGE §19–§21).
 
         Duplicate message_ids (protocol retries) are idempotent — the
-        first delivery wins and retries receive the same response.
-        Telemetry events update state without creating Executions;
-        semantic events become Intents executed by the engine.
+        first delivery wins and retries receive the **same** response
+        (replayed from the message log).  Telemetry events update state
+        without creating Executions; semantic events become Intents
+        executed by the engine.
         """
         payload = message.typed_payload()  # EventPayload
         if await self._store.seen_message(message.message_id):
-            # Idempotent duplicate: acknowledge without re-processing.
+            # Idempotent duplicate: replay the original response.
+            stored = await self._store.get_stored_response(message.message_id)
+            if stored is not None:
+                return EdgeMessage.model_validate_json(stored)
+            # Fallback for messages logged before response storage existed.
             return make_message(
                 EdgeMessageType.EVENT,
                 device_id=ctx.device_id,
@@ -269,7 +347,7 @@ class DeviceGateway:
 
         if event_name in _TELEMETRY_EVENTS:
             await self._store.update_last_seen(ctx.device_id)
-            return make_message(
+            response = make_message(
                 EdgeMessageType.EVENT,
                 device_id=ctx.device_id,
                 payload={
@@ -280,6 +358,10 @@ class DeviceGateway:
                 correlation_id=message.message_id,
                 trace_id=message.trace_id,
             )
+            await self._store.store_response(
+                message.message_id, response.model_dump_json()
+            )
+            return response
 
         execution_id: str | None = None
         if self._engine is not None:
@@ -302,7 +384,7 @@ class DeviceGateway:
             )
             execution_id = execution.id
 
-        return make_message(
+        response = make_message(
             EdgeMessageType.EVENT,
             device_id=ctx.device_id,
             payload={
@@ -313,6 +395,8 @@ class DeviceGateway:
             correlation_id=message.message_id,
             trace_id=message.trace_id,
         )
+        await self._store.store_response(message.message_id, response.model_dump_json())
+        return response
 
     async def handle_state_sync(
         self, message: EdgeMessage, ctx: AuthenticatedDeviceContext
@@ -321,6 +405,20 @@ class DeviceGateway:
         payload = message.typed_payload()  # StateSyncPayload
         device = await self._require_device(ctx.device_id)
 
+        # Resource limit: state size (EDGE §76).
+        if self._config is not None:
+            edge_cfg = getattr(self._config, "edge", None)
+            if edge_cfg is not None:
+                max_size = getattr(edge_cfg, "max_state_size", 0)
+                if max_size > 0:
+                    import json
+
+                    state_size = len(json.dumps(payload.state).encode())
+                    if state_size > max_size:
+                        raise PayloadTooLargeError(
+                            f"State size {state_size} exceeds limit {max_size}"
+                        )
+
         applied = await self._store.update_device_state(
             ctx.device_id, payload.state, payload.state_version
         )
@@ -415,7 +513,8 @@ class DeviceGateway:
 
         Enforces capability authorization at the runtime boundary
         (EDGE §50) — a device without the required capability never
-        receives the effect.
+        receives the effect.  The retry limit is read from
+        ``config.edge.max_effect_retries`` (default 3).
         """
         device = await self._require_device(device_id)
         if capability not in device.capabilities:
@@ -426,12 +525,31 @@ class DeviceGateway:
                     "device_capabilities": device.capabilities,
                 },
             )
+        # Resource limit: pending effects per device (EDGE §76).
+        if self._config is not None:
+            edge_cfg = getattr(self._config, "edge", None)
+            if edge_cfg is not None:
+                max_pending = getattr(edge_cfg, "max_pending_effects", 0)
+                if max_pending > 0:
+                    pending = await self._store.pending_effects(device_id)
+                    if len(pending) >= max_pending:
+                        raise PayloadTooLargeError(
+                            f"Device '{device_id}' has {len(pending)} pending effects "
+                            f"(limit {max_pending})"
+                        )
+        max_retries = 3
+        if self._config is not None:
+            try:
+                max_retries = int(self._config.edge.max_effect_retries)
+            except (AttributeError, TypeError, ValueError):
+                pass
         delivery = EffectDelivery(
             effect_id=effect_id,
             execution_id=execution_id,
             device_id=device_id,
             capability=capability,
             payload=payload,
+            max_retries=max_retries,
         )
         # Idempotent on effect_id — duplicate submissions are no-ops.
         await self._store.add_effect_delivery(delivery)
@@ -449,18 +567,39 @@ class DeviceGateway:
     async def pending_effects(
         self, device_id: str, ctx: AuthenticatedDeviceContext
     ) -> list[EffectDelivery]:
-        """Retrieve pending effects for a device (device isolation enforced)."""
+        """Retrieve and claim pending effects for a device.
+
+        Device isolation is enforced — a device can only list its own
+        effects.  Each effect is atomically claimed (``pending →
+        delivering``) so concurrent HTTP polls never receive the same
+        effect twice.
+        """
         if device_id != ctx.device_id:
             raise AuthorizationFailedError("Devices may only list their own effects")
         deliveries = await self._store.pending_effects(device_id)
+        claimed: list[EffectDelivery] = []
         for d in deliveries:
-            await self._store.mark_effect_delivered(d.effect_id)
+            if await self._store.claim_effect(device_id, d.effect_id):
+                claimed.append(d)
         await self._emit(
             "device.effect.delivered",
-            {"device_id": device_id, "count": len(deliveries)},
+            {"device_id": device_id, "count": len(claimed)},
             trace_id=None,
         )
-        return deliveries
+        return claimed
+
+    async def retry_effects(self, device_id: str) -> int:
+        """Reset stale ``delivering`` effects back to ``pending`` for retry.
+
+        Returns the number of effects that were reset or marked as
+        ``delivery_failed`` (exceeded max retries).
+        """
+        stale = await self._store.stale_deliveries(device_id)
+        retried = 0
+        for d in stale:
+            if await self._store.retry_effect(d.effect_id):
+                retried += 1
+        return retried
 
     # ------------------------------------------------------------------
     # Sessions & revocation

{voodoo_framework-2.6.0 → voodoo_framework-2.6.1}/src/voodoo/edge/http.py

@@ -74,9 +74,28 @@ def _envelope_response(message: Any) -> JSONResponse:
 
 
 def _make_enrollments_endpoint(gateway: DeviceGateway):
-    """POST /v1/edge/enrollments — issue a single-use enrollment key."""
+    """POST /v1/edge/enrollments — issue a single-use enrollment key.
+
+    Protected by ``X-Admin-Token`` when
+    ``config.edge.enrollment_auth_required`` is set (Sprint 23.1).
+    """
 
     async def enrollments(request: Request) -> Response:
+        # Enrollment auth guard (Sprint 23.1).
+        config = getattr(gateway, "_config", None)
+        edge_cfg = getattr(config, "edge", None) if config else None
+        if edge_cfg is not None and getattr(
+            edge_cfg, "enrollment_auth_required", False
+        ):
+            admin_token = getattr(edge_cfg, "enrollment_admin_token", "")
+            if admin_token:
+                provided = request.headers.get("x-admin-token", "")
+                if provided != admin_token:
+                    return _error_to_response(
+                        AuthenticationFailedError(
+                            "Missing or invalid admin token for enrollment"
+                        )
+                    )
         body = await _json_body(request)
         raw_key = await create_enrollment(
             gateway.store,
@@ -151,7 +170,12 @@ def _make_auth_endpoint(gateway: DeviceGateway):
 
 
 def _make_device_endpoint(gateway: DeviceGateway, message_type: EdgeMessageType):
-    """Shared factory for authenticated device → runtime endpoints."""
+    """Shared factory for stateless authenticated device → runtime endpoints.
+
+    Each HTTP request authenticates independently via
+    ``gateway.authenticate_request()`` — no persistent session is
+    created. This matches HTTP's stateless nature.
+    """
 
     async def handler(request: Request) -> Response:
         credential = _credential_from_request(request)
@@ -161,7 +185,7 @@ def _make_device_endpoint(gateway: DeviceGateway, message_type: EdgeMessageType)
             )
         body = await _json_body(request)
         try:
-            ctx, _session = await gateway.connect(
+            ctx = await gateway.authenticate_request(
                 credential, transport=TransportKind.HTTP
             )
         except EdgeError as e:
@@ -171,6 +195,7 @@ def _make_device_endpoint(gateway: DeviceGateway, message_type: EdgeMessageType)
             message_type,
             device_id=ctx.device_id,
             payload=body,
+            message_id=body.pop("message_id", None),
             trace_id=request.headers.get("x-trace-id"),
         )
         try:
@@ -196,7 +221,7 @@ def _make_effects_endpoint(gateway: DeviceGateway):
                 AuthenticationFailedError("Missing device credential")
             )
         try:
-            ctx, _session = await gateway.connect(
+            ctx = await gateway.authenticate_request(
                 credential, transport=TransportKind.HTTP
             )
         except EdgeError as e:
@@ -238,7 +263,7 @@ def _make_effect_ack_endpoint(gateway: DeviceGateway):
         effect_id = str(request.path_params.get("effect_id", ""))
         body.setdefault("effect_id", effect_id)
         try:
-            ctx, _session = await gateway.connect(
+            ctx = await gateway.authenticate_request(
                 credential, transport=TransportKind.HTTP
             )
         except EdgeError as e:

{voodoo_framework-2.6.0 → voodoo_framework-2.6.1}/src/voodoo/edge/mqtt.py

@@ -9,6 +9,10 @@ QoS 1 (at-least-once) is deliberate: the stable ``message_id`` /
 ``effect_id`` provide application-level idempotency, which remains
 required regardless of QoS (EDGE §37).
 
+MQTT is session-oriented: AUTH establishes a persistent context keyed
+by ``device_id``. Subsequent messages reuse the stored context. Only
+AUTH is allowed before a session is established (Sprint 23.1).
+
 The paho-mqtt dependency is optional (``voodoo[edge]`` extra); importing
 this module without it raises a clear ImportError.
 """
@@ -22,7 +26,7 @@ from typing import TYPE_CHECKING, Any
 
 from voodoo.edge.errors import TransportError
 from voodoo.edge.gateway import DeviceGateway
-from voodoo.edge.models import TransportKind
+from voodoo.edge.models import AuthenticatedDeviceContext, TransportKind
 from voodoo.edge.protocol import EdgeMessageType, make_message
 
 if TYPE_CHECKING:
@@ -130,6 +134,8 @@ class EdgeMQTTTransport:
         self._connected = asyncio.Event()
         # Per-device subscriptions for outbound effects.
         self._device_queues: dict[str, asyncio.Queue[Any]] = {}
+        # MQTT session registry — AUTH establishes persistent context.
+        self._sessions: dict[str, AuthenticatedDeviceContext] = {}
 
     # -- lifecycle ----------------------------------------------------------
 
@@ -168,6 +174,14 @@ class EdgeMQTTTransport:
     def connected(self) -> bool:
         return self._client is not None and self._connected.is_set()
 
+    def session_for(self, device_id: str) -> AuthenticatedDeviceContext | None:
+        """Return the authenticated session for a device, if any."""
+        return self._sessions.get(device_id)
+
+    def disconnect_device(self, device_id: str) -> None:
+        """Remove a device's MQTT session (e.g. on revoke or disconnect)."""
+        self._sessions.pop(device_id, None)
+
     # -- outbound: effects to devices ---------------------------------------
 
     async def publish_effect(self, device_id: str, effect: dict[str, Any]) -> None:
@@ -215,12 +229,19 @@ class EdgeMQTTTransport:
         )
 
     async def _handle_inbox(self, topic: str, payload: str) -> None:
-        """Route a decoded inbox message through the shared gateway."""
+        """Route a decoded inbox message through the shared gateway.
+
+        Sprint 23.1 hardening:
+        - Topic identity binding: topic device_id must match message device_id.
+        - Pre-AUTH rejection: only AUTH allowed before session is established.
+        - Session registry: AUTH stores context; subsequent messages reuse it.
+        - Topic authorization: device can only publish to its own namespace.
+        """
         parsed = parse_topic(topic)
         if parsed is None:
             logger.warning("Ignoring malformed topic %s", topic)
             return
-        device_id, kind = parsed
+        topic_device_id, kind = parsed
         expected_type = _INBOX_TOPICS.get(kind)
         if expected_type is None:
             logger.warning("Ignoring unknown topic kind %s", kind)
@@ -230,10 +251,73 @@ class EdgeMQTTTransport:
         except json.JSONDecodeError:
             logger.warning("Malformed JSON on %s", topic)
             return
+
+        if not self._validate_message(raw, topic_device_id, kind, expected_type):
+            return
+
         # Normalize into the canonical envelope typed for the topic kind.
         raw.setdefault("type", expected_type.value)
-        raw.setdefault("device_id", device_id)
+        raw.setdefault("device_id", topic_device_id)
         try:
-            await self._gateway.handle_message(raw, transport=TransportKind.MQTT)
+            response = await self._gateway.handle_message(
+                raw, transport=TransportKind.MQTT
+            )
         except Exception as e:  # noqa: BLE001 — broker handler must never crash
             logger.warning("Gateway rejected message on %s: %s", topic, e)
+            return
+
+        # AUTH success → store session context for subsequent messages.
+        if expected_type is EdgeMessageType.AUTH:
+            self._register_session(response, topic_device_id)
+
+    def _validate_message(
+        self,
+        raw: dict[str, Any],
+        topic_device_id: str,
+        kind: str,
+        expected_type: EdgeMessageType,
+    ) -> bool:
+        """Validate topic/message identity and pre-AUTH rules. Returns True if OK."""
+        message_device_id = raw.get("device_id", "")
+
+        # Topic identity binding: topic device_id must match message device_id.
+        if message_device_id and message_device_id != topic_device_id:
+            logger.warning(
+                "Rejecting %s: message device_id '%s' != topic device_id '%s'",
+                kind,
+                message_device_id,
+                topic_device_id,
+            )
+            return False
+
+        # Pre-AUTH rejection: only AUTH messages allowed before session.
+        if expected_type is not EdgeMessageType.AUTH:
+            if topic_device_id not in self._sessions:
+                logger.warning(
+                    "Rejecting %s from %s: no authenticated session",
+                    kind,
+                    topic_device_id,
+                )
+                return False
+            # Topic authorization: device can only publish to own namespace.
+            session = self._sessions[topic_device_id]
+            if message_device_id and message_device_id != session.device_id:
+                logger.warning(
+                    "Rejecting %s: message device_id '%s' != session device '%s'",
+                    kind,
+                    message_device_id,
+                    session.device_id,
+                )
+                return False
+
+        return True
+
+    def _register_session(self, response: Any, topic_device_id: str) -> None:
+        """Store the MQTT session context after successful AUTH."""
+        session_id = response.payload.get("session_id", "")
+        device_id = response.payload.get("device_id", topic_device_id)
+        if session_id:
+            ctx = self._gateway.context_for_session(session_id)
+            if ctx is not None:
+                self._sessions[device_id] = ctx
+                logger.info("MQTT session established for %s", device_id)