gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/tracing/identity.py

@@ -0,0 +1,288 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Agent identity for GatewaySDK tracing.
+
+The agent is the one identity a developer declares — everything else
+(user, session, context items) flows inherently:
+
+- **Agent**: declared once via ``tracing.init(agent=...)`` or the
+  ``GATEWAY_AGENT_ID`` / ``GATEWAY_AGENT_NAME`` / ``GATEWAY_AGENT_VERSION``
+  environment variables. Stamped on every span and registered with the
+  platform control plane at init.
+- **User**: resolved from ``GATEWAY_USER_ID`` (or set per-request by
+  middleware / explicitly per span). Never hand-stamped in app code.
+
+The identity is process-wide: one process is one agent. It is stored at
+module level (not a context variable) and read by the
+``TracingContextSpanProcessor`` on every span start.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+import threading
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+logger = logging.getLogger(__name__)
+
+_AGENT_ID_RE = re.compile(r"[^a-z0-9._-]+")
+
+
+@dataclass(frozen=True)
+class AgentIdentity:
+    """The declared identity of the agent running in this process."""
+
+    id: str
+    name: Optional[str] = None
+    version: Optional[str] = None
+
+
+_current_agent: Optional[AgentIdentity] = None
+
+
+def resolve_agent_identity(
+    agent: Union[str, dict, AgentIdentity, None] = None,
+) -> Optional[AgentIdentity]:
+    """Resolve the agent identity: explicit argument first, env second.
+
+    Environment variables: ``GATEWAY_AGENT_ID``, ``GATEWAY_AGENT_NAME``,
+    ``GATEWAY_AGENT_VERSION``. Returns None when nothing declares an agent —
+    tracing still works, traces are just not agent-attributed.
+    """
+    if isinstance(agent, AgentIdentity):
+        return agent
+    if isinstance(agent, dict):
+        agent_id = agent.get("id") or agent.get("agent_id")
+        if not agent_id:
+            raise ValueError("agent dict requires an 'id' key")
+        return AgentIdentity(
+            id=str(agent_id),
+            name=agent.get("name"),
+            version=agent.get("version"),
+        )
+    if isinstance(agent, str) and agent.strip():
+        return AgentIdentity(
+            id=agent.strip(),
+            name=os.getenv("GATEWAY_AGENT_NAME"),
+            version=os.getenv("GATEWAY_AGENT_VERSION"),
+        )
+
+    env_id = os.getenv("GATEWAY_AGENT_ID")
+    if env_id and env_id.strip():
+        return AgentIdentity(
+            id=env_id.strip(),
+            name=os.getenv("GATEWAY_AGENT_NAME"),
+            version=os.getenv("GATEWAY_AGENT_VERSION"),
+        )
+    return None
+
+
+def set_agent_identity(identity: Optional[AgentIdentity]) -> None:
+    """Set the process-wide agent identity (called by ``tracing.init``)."""
+    global _current_agent
+    _current_agent = identity
+
+
+def get_agent_identity() -> Optional[AgentIdentity]:
+    """The process-wide agent identity, or None when undeclared."""
+    return _current_agent
+
+
+_current_session_id: Optional[str] = None
+
+
+def initialize_ambient_session() -> str:
+    """Establish the process-wide ambient session id (called by init).
+
+    Sessions auto-populate: every trace belongs to a session without any
+    caller code. Resolution: GATEWAY_SESSION_ID env (explicit pin, e.g. a
+    stable name across restarts) → auto-generated per process. Finer-grained
+    boundaries (request middleware, run()/rollout()) override this ambient
+    default per span.
+    """
+    global _current_session_id
+    env_session = os.getenv("GATEWAY_SESSION_ID")
+    if env_session and env_session.strip():
+        _current_session_id = env_session.strip()
+    elif _current_session_id is None:
+        from uuid import uuid4
+
+        agent = get_agent_identity()
+        prefix = f"{agent.id}-" if agent else ""
+        _current_session_id = f"{prefix}{uuid4().hex[:12]}"
+    return _current_session_id
+
+
+def get_ambient_session_id() -> Optional[str]:
+    """The process-wide session id, or None before ``tracing.init``."""
+    return _current_session_id
+
+
+def reset_ambient_session() -> None:
+    """Testing hook: clear the ambient session."""
+    global _current_session_id
+    _current_session_id = None
+
+
+# Per-context user identity set via identify() — ContextVar so concurrent
+# requests serving different users never bleed into each other.
+from contextvars import ContextVar
+
+_current_user_id: ContextVar[Optional[str]] = ContextVar(
+    "gateway_current_user_id", default=None
+)
+
+
+def resolve_user_id() -> Optional[str]:
+    """The active end-user id: identify() context > GATEWAY_USER_ID env.
+
+    Per-request user attribution (e.g. from auth middleware) can also
+    override per span by setting the ``gateway.user.id`` attribute directly —
+    explicit set wins over the ambient default stamped at span start.
+    """
+    contextual = _current_user_id.get()
+    if contextual:
+        return contextual
+    user_id = os.getenv("GATEWAY_USER_ID")
+    return user_id.strip() if user_id and user_id.strip() else None
+
+
+def identify(
+    user_id: str,
+    *,
+    attributes: Optional[dict] = None,
+    display_name: Optional[str] = None,
+    email: Optional[str] = None,
+) -> None:
+    """Attribute everything that follows (and the current trace) to a user.
+
+    - Sets the user for all spans started in this context, and stamps the
+      currently open span so the in-flight trace is attributed immediately.
+    - When ``attributes``/``display_name``/``email`` are given, the user's
+      platform profile is updated in the background (attributes merge with
+      what's stored — ideal for incremental ROI facts like plan or
+      account_value). Requires GATEWAY_HOST/GATEWAY_PUBLIC_KEY/
+      GATEWAY_SECRET_KEY for the profile update; tracing attribution works
+      without them.
+
+    Example::
+
+        gatewaysdk.identify(
+            "user-123",
+            attributes={"plan": "enterprise", "account_value": "48000"},
+        )
+    """
+    cleaned = user_id.strip()
+    if not cleaned:
+        return
+    _current_user_id.set(cleaned)
+
+    # Stamp the open span so the in-flight trace is upserted with the user —
+    # ingestion treats user.id on ANY span as a trace-level update.
+    try:
+        from opentelemetry import trace as _otel_trace
+
+        span = _otel_trace.get_current_span()
+        if span is not None and span.get_span_context().is_valid:
+            span.set_attribute("gateway.user.id", cleaned)
+    except Exception:  # pragma: no cover - tracing not initialized
+        logger.debug("identify(): no active span to stamp", exc_info=True)
+
+    if attributes or display_name or email:
+        def _push_profile() -> None:
+            try:
+                from ..users import UsersClient
+
+                UsersClient.from_env().update(
+                    cleaned,
+                    display_name=display_name,
+                    email=email,
+                    attributes={
+                        key: str(value) for key, value in (attributes or {}).items()
+                    },
+                )
+            except Exception:
+                logger.debug(
+                    "identify(): profile update failed", exc_info=True
+                )
+
+        threading.Thread(
+            target=_push_profile, name="gateway-identify", daemon=True
+        ).start()
+
+
+def register_agent_with_platform(
+    identity: AgentIdentity,
+    *,
+    host: Optional[str] = None,
+    public_key: Optional[str] = None,
+    secret_key: Optional[str] = None,
+) -> None:
+    """Best-effort, non-blocking registration with the platform control plane.
+
+    Upserts the agent (idempotent by project + name) via
+    ``POST /api/public/agents/register`` so it appears on the Agents page
+    before its first trace, with version history in the manifest. Failures
+    are logged at debug level and never break tracing.
+    """
+    resolved_host = host or os.getenv("GATEWAY_HOST")
+    resolved_public = public_key or os.getenv("GATEWAY_PUBLIC_KEY")
+    resolved_secret = secret_key or os.getenv("GATEWAY_SECRET_KEY")
+    if not (resolved_host and resolved_public and resolved_secret):
+        logger.debug(
+            "Agent registration skipped: missing GATEWAY_HOST/PUBLIC_KEY/SECRET_KEY"
+        )
+        return
+
+    def _register() -> None:
+        import base64
+        from urllib import request as urllib_request
+
+        image_name = _AGENT_ID_RE.sub("-", identity.id.lower()).strip("-") or "agent"
+        payload: dict[str, Any] = {
+            "name": identity.id,
+            # The control plane requires a container image reference; SDK
+            # declared agents are not containerized, so a stable synthetic
+            # ref keeps the upsert idempotent.
+            "image": f"gateway/{image_name}:declared",
+            "capabilities": ["gateway_tracing"],
+            "manifest": {
+                "source": "gatewaysdk.tracing.init",
+                "agent": {
+                    "id": identity.id,
+                    "name": identity.name,
+                    "version": identity.version,
+                },
+            },
+            "resources": {"runtime": "sdk"},
+        }
+        if identity.name:
+            payload["description"] = identity.name
+
+        token = base64.b64encode(
+            f"{resolved_public}:{resolved_secret}".encode("utf-8")
+        ).decode("utf-8")
+        req = urllib_request.Request(
+            f"{resolved_host.rstrip('/')}/api/public/agents/register",
+            data=json.dumps(payload).encode("utf-8"),
+            headers={
+                "Authorization": f"Basic {token}",
+                "Content-Type": "application/json",
+            },
+            method="POST",
+        )
+        try:
+            with urllib_request.urlopen(req, timeout=10) as response:
+                response.read()
+            logger.debug("Registered agent %s with platform", identity.id)
+        except Exception as exc:  # noqa: BLE001 — registration must never break apps
+            logger.debug("Agent registration failed (non-fatal): %s", exc)
+
+    thread = threading.Thread(
+        target=_register, name="gatewaysdk-agent-register", daemon=True
+    )
+    thread.start()