iac-code 0.1.0__py3-none-any.whl

iac_code/services/telemetry/fallback.py

@@ -0,0 +1,59 @@
+"""FallbackStore — on-disk persistence of failed event batches."""
+
+from __future__ import annotations
+
+import json
+import uuid
+from collections.abc import Iterable, Iterator
+from pathlib import Path
+
+_FAILED_PREFIX = "failed_events."
+_FAILED_SUFFIX = ".jsonl"
+
+
+class FallbackStore:
+    """JSONL-backed store for event batches that failed default-backend export."""
+
+    def __init__(self, base_dir: Path) -> None:
+        self._base_dir = base_dir
+
+    def _ensure_dir(self) -> Path:
+        self._base_dir.mkdir(parents=True, exist_ok=True)
+        return self._base_dir
+
+    def write(self, session_id: str, events: Iterable[dict]) -> Path:
+        """Write one JSONL file per batch. One line per event."""
+        batch_uuid = uuid.uuid4().hex[:12]
+        path = self._ensure_dir() / f"{_FAILED_PREFIX}{session_id}.{batch_uuid}{_FAILED_SUFFIX}"
+        with path.open("w", encoding="utf-8") as fh:
+            for event in events:
+                fh.write(json.dumps(event, ensure_ascii=False) + "\n")
+        return path
+
+    def list_pending(self) -> Iterator[Path]:
+        """Yield every failed-batch file currently on disk."""
+        if not self._base_dir.exists():
+            return
+        for p in self._base_dir.iterdir():
+            if p.is_file() and p.name.startswith(_FAILED_PREFIX) and p.suffix == _FAILED_SUFFIX:
+                yield p
+
+    def remove(self, path: Path) -> None:
+        """Delete a batch file after successful re-export. Silent on missing."""
+        try:
+            path.unlink()
+        except FileNotFoundError:
+            pass
+
+    def read(self, path: Path) -> list[dict]:
+        """Parse a JSONL batch file. Unparseable lines skipped."""
+        out: list[dict] = []
+        for line in path.read_text(encoding="utf-8").splitlines():
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                out.append(json.loads(line))
+            except json.JSONDecodeError:
+                continue
+        return out

iac_code/services/telemetry/identity.py

@@ -0,0 +1,73 @@
+"""Identity generation for telemetry.
+
+user.id    = iac_user_<uuid4>, persisted to a settings.yml path
+session.id = iac_sess_<uuid4>, per Identity instance (per process)
+tenant.id  = iac_tenant_<user-defined>, from IAC_CODE_TENANT_ID
+"""
+
+from __future__ import annotations
+
+import os
+import uuid
+from pathlib import Path
+
+from iac_code.config import _load_yaml, _save_yaml
+
+USER_ID_PREFIX = "iac_user_"
+SESSION_ID_PREFIX = "iac_sess_"
+TENANT_ID_PREFIX = "iac_tenant_"
+
+_USER_ID_KEY = "userID"
+_TENANT_ENV_VAR = "IAC_CODE_TENANT_ID"
+
+
+class Identity:
+    """Owns user.id / session.id / tenant.id.
+
+    `settings_path` is injected so tests can pass a tmp path instead of the
+    real ~/.iac-code/settings.yml.
+    """
+
+    def __init__(self, settings_path: Path, session_id: str | None = None) -> None:
+        self._settings_path = settings_path
+        self._user_id: str | None = None
+        self._session_id: str | None = f"{SESSION_ID_PREFIX}{session_id}" if session_id else None
+        self._was_first_run = False
+
+    def get_user_id(self) -> str:
+        """Return the persistent user.id; generate + persist on first miss."""
+        if self._user_id is not None:
+            return self._user_id
+        settings = _load_yaml(self._settings_path)
+        existing = settings.get(_USER_ID_KEY)
+        if isinstance(existing, str) and existing.startswith(USER_ID_PREFIX):
+            self._user_id = existing
+            return existing
+        new_id = f"{USER_ID_PREFIX}{uuid.uuid4()}"
+        settings[_USER_ID_KEY] = new_id
+        _save_yaml(self._settings_path, settings)
+        self._user_id = new_id
+        self._was_first_run = True
+        return new_id
+
+    def get_session_id(self) -> str:
+        """Return per-instance session.id; generate on first call."""
+        if self._session_id is None:
+            self._session_id = f"{SESSION_ID_PREFIX}{uuid.uuid4()}"
+        return self._session_id
+
+    def get_tenant_id(self) -> str | None:
+        """Return tenant.id if IAC_CODE_TENANT_ID is set, else None.
+
+        Read fresh each call so monkeypatching in tests works reliably.
+        """
+        raw = os.environ.get(_TENANT_ENV_VAR, "").strip()
+        if not raw:
+            return None
+        if raw.startswith(TENANT_ID_PREFIX):
+            return raw
+        return f"{TENANT_ID_PREFIX}{raw}"
+
+    def was_first_run(self) -> bool:
+        """True iff get_user_id() minted a new id on this instance."""
+        return self._was_first_run

iac_code/services/telemetry/metrics.py

@@ -0,0 +1,62 @@
+"""MetricsRegistry — wraps OTel Counter/Histogram instruments."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger
+from opentelemetry.metrics import Meter
+
+from iac_code.services.telemetry.names import Metrics as M  # noqa: N817
+
+METRIC_NAMES: tuple[str, ...] = (
+    M.SESSION_COUNT,
+    M.ACTIVE_TIME_TOTAL,
+    M.TOKEN_USAGE,
+    M.API_REQUEST_COUNT,
+    M.API_REQUEST_DURATION,
+    M.TOOL_USE_COUNT,
+    M.TEMPLATE_GENERATED_COUNT,
+    M.TEMPLATE_VALIDATED_COUNT,
+    M.DEPLOYMENT_COUNT,
+    M.DEPLOYMENT_DURATION,
+    M.RESOURCE_TYPE_OBSERVED_COUNT,
+    M.ALIYUN_API_CALLED_COUNT,
+    M.ALIYUN_API_CALLED_DURATION,
+    M.TERRAFORM_PROVIDER_OBSERVED_COUNT,
+)
+
+_HISTOGRAM_NAMES: frozenset[str] = frozenset(
+    {
+        M.API_REQUEST_DURATION,
+        M.DEPLOYMENT_DURATION,
+        M.ALIYUN_API_CALLED_DURATION,
+    }
+)
+
+
+class MetricsRegistry:
+    """Holds instrument objects and dispatches add/record calls."""
+
+    def __init__(self, instruments: dict[str, Any] | None = None) -> None:
+        self._instruments: dict[str, Any] = dict(instruments) if instruments else {}
+
+    def register_all(self, meter: Meter) -> None:
+        """Create a Counter or Histogram per known metric name."""
+        self._instruments.clear()
+        for name in METRIC_NAMES:
+            if name in _HISTOGRAM_NAMES:
+                self._instruments[name] = meter.create_histogram(name=name, description=name)
+            else:
+                self._instruments[name] = meter.create_counter(name=name, description=name)
+
+    def add(self, name: str, value: int | float, attributes: dict[str, Any]) -> None:
+        """Route to the correct instrument method; silently drop unknown names."""
+        logger.info("[metric] {} value={} {}", name, value, attributes)
+        inst = self._instruments.get(name)
+        if inst is None:
+            return
+        if name in _HISTOGRAM_NAMES:
+            inst.record(value, attributes)
+        else:
+            inst.add(value, attributes)

iac_code/services/telemetry/names.py

@@ -0,0 +1,199 @@
+"""Centralized constants for all telemetry signal names."""
+
+from __future__ import annotations
+
+# =====================================================================
+# ARMS LLM semantic conventions (gen_ai.*)
+# https://help.aliyun.com/zh/arms/application-monitoring/developer-reference/llm-trace-field-definition-description
+# =====================================================================
+
+
+class GenAiSpanKind:
+    """gen_ai.span.kind enumeration values."""
+
+    ENTRY = "ENTRY"
+    LLM = "LLM"
+    TOOL = "TOOL"
+    STEP = "STEP"
+    AGENT = "AGENT"
+    CHAIN = "CHAIN"
+    TASK = "TASK"
+    RETRIEVER = "RETRIEVER"
+    EMBEDDING = "EMBEDDING"
+    RERANKER = "RERANKER"
+
+
+class GenAiOperationName:
+    """gen_ai.operation.name enumeration values."""
+
+    ENTER = "enter"
+    CHAT = "chat"
+    TEXT_COMPLETION = "text_completion"
+    GENERATE_CONTENT = "generate_content"
+    EXECUTE_TOOL = "execute_tool"
+    INVOKE_AGENT = "invoke_agent"
+    CREATE_AGENT = "create_agent"
+    REACT = "react"
+    RETRIEVAL = "retrieval"
+    EMBEDDINGS = "embeddings"
+
+
+class GenAiAttr:
+    """gen_ai.* span attribute key constants."""
+
+    # --- Common (all spans) ---
+    SPAN_KIND = "gen_ai.span.kind"
+    OPERATION_NAME = "gen_ai.operation.name"
+    SESSION_ID = "gen_ai.session.id"
+    USER_ID = "gen_ai.user.id"
+    FRAMEWORK = "gen_ai.framework"
+
+    # --- Provider & Model ---
+    PROVIDER_NAME = "gen_ai.provider.name"
+    REQUEST_MODEL = "gen_ai.request.model"
+    RESPONSE_MODEL = "gen_ai.response.model"
+    RESPONSE_ID = "gen_ai.response.id"
+    CONVERSATION_ID = "gen_ai.conversation.id"
+
+    # --- Request parameters ---
+    REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens"
+    REQUEST_TEMPERATURE = "gen_ai.request.temperature"
+    REQUEST_TOP_P = "gen_ai.request.top_p"
+    REQUEST_TOP_K = "gen_ai.request.top_k"
+    REQUEST_FREQUENCY_PENALTY = "gen_ai.request.frequency_penalty"
+    REQUEST_PRESENCE_PENALTY = "gen_ai.request.presence_penalty"
+    REQUEST_STOP_SEQUENCES = "gen_ai.request.stop_sequences"
+    REQUEST_SEED = "gen_ai.request.seed"
+    REQUEST_CHOICE_COUNT = "gen_ai.request.choice.count"
+
+    # --- Response ---
+    RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons"
+    RESPONSE_TIME_TO_FIRST_TOKEN = "gen_ai.response.time_to_first_token"
+    USER_TIME_TO_FIRST_TOKEN = "gen_ai.user.time_to_first_token"
+    RESPONSE_REASONING_TIME = "gen_ai.response.reasoning_time"
+    OUTPUT_TYPE = "gen_ai.output.type"
+
+    # --- Usage ---
+    USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens"
+    USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
+    USAGE_TOTAL_TOKENS = "gen_ai.usage.total_tokens"
+    USAGE_CACHE_CREATION_INPUT_TOKENS = "gen_ai.usage.cache_creation.input_tokens"
+    USAGE_CACHE_READ_INPUT_TOKENS = "gen_ai.usage.cache_read.input_tokens"
+
+    # --- Content (debug mode only) ---
+    INPUT_MESSAGES = "gen_ai.input.messages"
+    OUTPUT_MESSAGES = "gen_ai.output.messages"
+    SYSTEM_INSTRUCTIONS = "gen_ai.system_instructions"
+    TOOL_DEFINITIONS = "gen_ai.tool.definitions"
+
+    # --- Tool ---
+    TOOL_NAME = "gen_ai.tool.name"
+    TOOL_TYPE = "gen_ai.tool.type"
+    TOOL_CALL_ID = "gen_ai.tool.call.id"
+    TOOL_DESCRIPTION = "gen_ai.tool.description"
+    TOOL_CALL_ARGUMENTS = "gen_ai.tool.call.arguments"
+    TOOL_CALL_RESULT = "gen_ai.tool.call.result"
+
+    # --- Agent ---
+    AGENT_NAME = "gen_ai.agent.name"
+    AGENT_ID = "gen_ai.agent.id"
+    AGENT_DESCRIPTION = "gen_ai.agent.description"
+    DATA_SOURCE_ID = "gen_ai.data_source.id"
+
+    # --- ReAct Step ---
+    REACT_FINISH_REASON = "gen_ai.react.finish_reason"
+    REACT_ROUND = "gen_ai.react.round"
+
+
+class ArmsResourceAttr:
+    """ARMS-specific resource attribute keys."""
+
+    SERVICE_FEATURE = "acs.arms.service.feature"
+    CMS_WORKSPACE = "acs.cms.workspace"
+    SERVICE_ID = "acs.arms.service.id"
+
+
+ARMS_FEATURE_GENAI_APP = "genai_app"
+FRAMEWORK_IAC_CODE = "iac-code"
+
+
+# =====================================================================
+# iac-code application signals (iac.*)
+# =====================================================================
+
+
+class Events:
+    """All iac.* event names (OTel Logs signal)."""
+
+    # --- Lifecycle (5) ---
+    INIT = "iac.init"
+    SESSION_STARTED = "iac.session.started"
+    SESSION_EXITED = "iac.session.exited"
+    SESSION_CANCELLED = "iac.session.cancelled"
+    AUTH_CONFIGURED = "iac.auth.configured"
+
+    # --- API / LLM (5) ---
+    API_REQUEST_STARTED = "iac.api.request.started"
+    API_REQUEST_SUCCEEDED = "iac.api.request.succeeded"
+    API_REQUEST_FAILED = "iac.api.request.failed"
+    API_REQUEST_RETRIED = "iac.api.request.retried"
+    MODEL_FALLBACK_TRIGGERED = "iac.model.fallback.triggered"
+
+    # --- Tool (4) ---
+    TOOL_USE_SUCCEEDED = "iac.tool.use.succeeded"
+    TOOL_USE_FAILED = "iac.tool.use.failed"
+    TOOL_USE_GRANTED_IN_PROMPT = "iac.tool.use.granted_in_prompt"
+    TOOL_USE_REJECTED_IN_PROMPT = "iac.tool.use.rejected_in_prompt"
+
+    # --- IaC core (9) ---
+    TEMPLATE_GENERATED = "iac.template.generated"
+    TEMPLATE_VALIDATED = "iac.template.validated"
+    DEPLOYMENT_STARTED = "iac.deployment.started"
+    DEPLOYMENT_SUCCEEDED = "iac.deployment.succeeded"
+    DEPLOYMENT_FAILED = "iac.deployment.failed"
+    DEPLOYMENT_CANCELLED = "iac.deployment.cancelled"
+    DOC_SEARCHED = "iac.doc.searched"
+    SKILL_INVOKED = "iac.skill.invoked"
+    SKILL_COMPLETED = "iac.skill.completed"
+
+    # --- Aliyun API (1) ---
+    ALIYUN_API_CALLED = "iac.aliyun.api.called"
+
+    # --- Memory (2) ---
+    MEMORY_COMPACT_SUCCEEDED = "iac.memory.compact.succeeded"
+    MEMORY_COMPACT_FAILED = "iac.memory.compact.failed"
+
+    # --- Crash / error (3) ---
+    EXCEPTION_UNCAUGHT = "iac.exception.uncaught"
+    EXCEPTION_UNHANDLED = "iac.exception.unhandled"
+    QUERY_FAILED = "iac.query.failed"
+
+
+class Metrics:
+    """All iac.* metric names."""
+
+    SESSION_COUNT = "iac.session.count"
+    ACTIVE_TIME_TOTAL = "iac.active_time.total"
+    TOKEN_USAGE = "iac.token.usage"
+    API_REQUEST_COUNT = "iac.api.request.count"
+    API_REQUEST_DURATION = "iac.api.request.duration"
+    TOOL_USE_COUNT = "iac.tool.use.count"
+    TEMPLATE_GENERATED_COUNT = "iac.template.generated.count"
+    TEMPLATE_VALIDATED_COUNT = "iac.template.validated.count"
+    DEPLOYMENT_COUNT = "iac.deployment.count"
+    DEPLOYMENT_DURATION = "iac.deployment.duration"
+    RESOURCE_TYPE_OBSERVED_COUNT = "iac.resource_type.observed.count"
+    ALIYUN_API_CALLED_COUNT = "iac.aliyun.api.called.count"
+    ALIYUN_API_CALLED_DURATION = "iac.aliyun.api.called.duration"
+    TERRAFORM_PROVIDER_OBSERVED_COUNT = "iac.terraform.provider.observed.count"
+
+
+class Spans:
+    """Span name constants (ARMS LLM convention: '{operation} {identifier}')."""
+
+    ENTRY = "enter_ai_application_system"
+    LLM_CHAT = "chat"
+    TOOL_EXECUTE = "execute_tool"
+    REACT_STEP = "react step"
+    SKILL_EXECUTE = "iac.skill.execute"
+    TEMPLATE_VALIDATE = "iac.template.validate"

iac_code/services/telemetry/sanitize.py

@@ -0,0 +1,88 @@
+"""Sanitization helpers for telemetry text fields."""
+
+from __future__ import annotations
+
+import re
+from typing import Literal
+
+from iac_code.services.telemetry.config import is_essential_traffic_only
+from iac_code.services.telemetry.constants import (
+    BUNDLED_SKILLS,
+    CUSTOM_ROS_RESOURCE_PLACEHOLDER,
+    CUSTOM_SKILL_PLACEHOLDER,
+    CUSTOM_TF_PROVIDER_PLACEHOLDER,
+    CUSTOM_TF_RESOURCE_PLACEHOLDER,
+    KNOWN_MODELS,
+    MCP_TOOL_PLACEHOLDER,
+    OTHER_MODEL_PLACEHOLDER,
+    ROS_ALLOWED_PREFIXES,
+    TERRAFORM_OFFICIAL_PROVIDERS,
+)
+
+_CONTROL_CHARS_RE = re.compile(r"[\n\r\t\x00-\x1f]+")
+_MAX_ERROR_MSG_BYTES = 512
+_TRUNCATION_MARKER = "... (truncated)"
+_DEV_VERSION_SUFFIX_RE = re.compile(r"-\d{8}$")
+
+
+def sanitize_error_message(raw: str | None) -> str | None:
+    """Clean and truncate an error message. None under essential-traffic mode."""
+    if raw is None:
+        return None
+    if is_essential_traffic_only():
+        return None
+    cleaned = _CONTROL_CHARS_RE.sub(" ", raw).strip()
+    encoded = cleaned.encode("utf-8")
+    if len(encoded) > _MAX_ERROR_MSG_BYTES:
+        keep = _MAX_ERROR_MSG_BYTES - len(_TRUNCATION_MARKER.encode("utf-8"))
+        cleaned = encoded[:keep].decode("utf-8", errors="ignore") + _TRUNCATION_MARKER
+    return cleaned
+
+
+def sanitize_skill_name(raw: str | None) -> str | None:
+    if raw is None:
+        return None
+    return raw if raw in BUNDLED_SKILLS else CUSTOM_SKILL_PLACEHOLDER
+
+
+def sanitize_resource_type(raw: str, kind: Literal["ros", "terraform"]) -> str:
+    """Keep official resource types; replace custom/unknown with placeholder."""
+    if kind == "ros":
+        if raw.startswith(ROS_ALLOWED_PREFIXES):
+            return raw
+        return CUSTOM_ROS_RESOURCE_PLACEHOLDER
+    # Terraform: resource types are of the form `<provider>_<resource>`.
+    if "_" not in raw:
+        return CUSTOM_TF_RESOURCE_PLACEHOLDER
+    provider = raw.split("_", 1)[0]
+    if provider in TERRAFORM_OFFICIAL_PROVIDERS:
+        return raw
+    return CUSTOM_TF_RESOURCE_PLACEHOLDER
+
+
+def sanitize_terraform_provider(raw: str) -> str:
+    return raw if raw in TERRAFORM_OFFICIAL_PROVIDERS else CUSTOM_TF_PROVIDER_PLACEHOLDER
+
+
+def sanitize_model_name(raw: str) -> str:
+    """Trim dev suffix (-YYYYMMDD) and map to known set or 'other'."""
+    base = _DEV_VERSION_SUFFIX_RE.sub("", raw)
+    return base if base in KNOWN_MODELS else OTHER_MODEL_PLACEHOLDER
+
+
+def sanitize_tool_name(raw: str) -> str:
+    """MCP tools are collapsed to a single placeholder."""
+    if raw.startswith("mcp__"):
+        return MCP_TOOL_PLACEHOLDER
+    return raw
+
+
+def bucket_resource_count(n: int) -> str:
+    """bucket for iac.deployment.duration histogram."""
+    if n <= 5:
+        return "1-5"
+    if n <= 20:
+        return "6-20"
+    if n <= 50:
+        return "21-50"
+    return "50+"

iac_code/services/telemetry/sink.py

@@ -0,0 +1,67 @@
+"""AnalyticsSink — routes events through the privacy gate to the Events pipeline.
+
+Design:
+  - Before `activate()`, events are queued in-memory (bounded by maxlen=10k).
+  - After `activate()`, events go directly to the EventEmitter.
+  - `drain_sync()` / `drain_soon()` flushes the pre-queue once activated.
+  - The privacy gate blocks emission when no-telemetry / essential-traffic is on.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections import deque
+from threading import Lock
+from typing import Any
+
+from loguru import logger
+
+from iac_code.services.telemetry.config import is_telemetry_disabled
+from iac_code.services.telemetry.events import EventEmitter
+
+
+class AnalyticsSink:
+    """Privacy-gated event router with pre-activation queue."""
+
+    def __init__(self, emitter: EventEmitter, queue_max: int = 10_000) -> None:
+        self._emitter = emitter
+        self._queue: deque[tuple[str, dict[str, Any]]] = deque(maxlen=queue_max)
+        self._lock = Lock()
+        self._active = False
+
+    def log_event(self, event_name: str, metadata: dict[str, Any]) -> None:
+        """Queue the event if not yet active, else gate+emit directly."""
+        with self._lock:
+            if not self._active:
+                self._queue.append((event_name, metadata))
+                return
+        self._dispatch(event_name, metadata)
+
+    def activate(self) -> None:
+        """Mark the sink active. Idempotent. Does NOT drain — call drain_*()."""
+        with self._lock:
+            self._active = True
+
+    def drain_sync(self) -> None:
+        """Synchronously flush the pre-activation queue."""
+        with self._lock:
+            queued = list(self._queue)
+            self._queue.clear()
+        for name, meta in queued:
+            self._dispatch(name, meta)
+
+    def drain_soon(self) -> None:
+        """Schedule an async drain on the running loop, if any. Else sync."""
+        try:
+            loop = asyncio.get_running_loop()
+            loop.call_soon(self.drain_sync)
+            return
+        except RuntimeError:
+            pass
+        self.drain_sync()
+
+    def _dispatch(self, event_name: str, metadata: dict[str, Any]) -> None:
+        logger.info("[event] {} {}", event_name, metadata)
+        if is_telemetry_disabled():
+            return
+        self._emitter.emit(event_name, metadata)

iac_code/services/telemetry/tracing.py

@@ -0,0 +1,38 @@
+"""SpanFactory — wraps the OTel Tracer."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Any, Iterator
+
+from loguru import logger
+from opentelemetry import trace
+from opentelemetry.trace import Tracer
+
+_TRACER_NAME = "com.iac-code.tracing"
+_TRACER_VERSION = "1.0.0"
+
+
+class SpanFactory:
+    """Context-manager span producer."""
+
+    def __init__(self) -> None:
+        self._tracer: Tracer | None = None
+
+    def attach(self, tracer: Tracer) -> None:
+        self._tracer = tracer
+
+    def _get_tracer(self) -> Tracer:
+        if self._tracer is None:
+            return trace.get_tracer(_TRACER_NAME, _TRACER_VERSION)
+        return self._tracer
+
+    @contextmanager
+    def start(self, name: str, attributes: dict[str, Any] | None = None) -> Iterator[Any]:
+        """Start a span as a context manager."""
+        tracer = self._get_tracer()
+        with tracer.start_as_current_span(name, attributes=attributes or {}) as span:
+            yield span
+            attrs = getattr(span, "attributes", None)
+            if attrs:
+                logger.debug("[span] {} {}", name, {k: v for k, v in attrs.items()})

iac_code/services/telemetry/types.py

@@ -0,0 +1,13 @@
+"""Type definitions for telemetry payloads."""
+
+from __future__ import annotations
+
+from typing import Union
+
+# Primitive values allowed directly in event/metric metadata.
+# Strings are NOT in this list by design — callers must sanitize text explicitly.
+AllowedMetadataValue = Union[bool, int, float, None, list, str]
+# Note: `str` is included because sanitize_*() returns str. Callers must only pass
+# strings that came from a sanitize_*() function or from a well-known enum constant.
+
+EventMetadata = dict[str, AllowedMetadataValue]

iac_code/services/token_budget.py

@@ -0,0 +1,54 @@
+"""Token budget management for controlling LLM usage."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+
+@dataclass
+class TokenBudget:
+    """Tracks token consumption against a total budget."""
+
+    total: int | None = None  # None means unlimited
+    used: int = field(default=0, init=False)
+
+    @property
+    def remaining(self) -> int | float:
+        if self.total is None:
+            return float("inf")
+        return max(0, self.total - self.used)
+
+    @property
+    def is_exhausted(self) -> bool:
+        if self.total is None:
+            return False
+        return self.used >= self.total
+
+    @property
+    def usage_percent(self) -> float:
+        if self.total is None or self.total == 0:
+            return 0.0
+        return (self.used / self.total) * 100.0
+
+    def consume(self, tokens: int) -> None:
+        self.used += tokens
+
+    @staticmethod
+    def parse_shorthand(text: str) -> int:
+        cleaned = text.strip().lstrip("+")
+        match = re.match(r"^(\d+(?:\.\d+)?)\s*([kmKM])?$", cleaned)
+        if not match:
+            raise ValueError(f"Invalid token shorthand: '{text}'")
+        value = float(match.group(1))
+        suffix = (match.group(2) or "").lower()
+        multipliers = {"k": 1_000, "m": 1_000_000}
+        return int(value * multipliers.get(suffix, 1))
+
+    @classmethod
+    def unlimited(cls) -> TokenBudget:
+        return cls(total=None)
+
+    @classmethod
+    def from_shorthand(cls, text: str) -> TokenBudget:
+        return cls(total=cls.parse_shorthand(text))