adk-telemetry: address PR #237 review comments (round 2)

Reviewer feedback and design nits from daeunJe0ng's copilot-review:

- Fix A (should-fix #1: auth-path latency): the /organization and /me
  silent lookups on the interactive auth hot path now use a bounded 3s
  socket timeout via plain requests.get (not the module-level _SESSION,
  whose 3x 429/5xx retry with backoff_factor=1 would add ~7s per lookup
  on a throttled tenant). Worst-case authenticate() overhead is now
  bounded to ~6s instead of ~15-40s in the admin-consent-gap case.
  tenant_class is driven by the tid claim already in the token, so a
  timeout here only degrades tenant_name to blank (recoverable on the
  next FlightCheck run) — never a misclassification.

- Fix B (inline nit on common_dimensions): sanitize the explicit
  tenant_id kwarg too. Every tenant_id source (explicit, in-memory
  identity, disk cache) now flows through _sanitize_tenant_id at the
  single choke point in common_dimensions, closing the last back-door
  where a caller-supplied non-GUID could leak into the customer bucket.

- Fix C (inline nit on cache_tenant_name): atomic write via
  tempfile+os.replace, matching cache_tenant_id. A concurrent reader
  (subprocess spawned from a SKILL.md step) can no longer see a
  truncated file and fall back to a blank tenant_name for the whole
  session.

- Fix D (inline nit on get_cached_tenant_id): the legacy raw-string
  compatibility shim now validates against _GUID_RE (case-insensitive)
  before returning, matching the docstring's guarantee. Any garbage /
  torn / hand-edited pre-versioned file yields "" instead of leaking a
  bad value onto real events.

- New regression tests: kwarg-sanitization in common_dimensions, legacy
  raw-string rejection in get_cached_tenant_id.

- Test-side: tests/flightcheck/test_graph_client.py now patches
  requests.get instead of graph_client._SESSION.get to match the
  no-retry hot-path pattern.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 70ddf006-7f8d-48e7-9afa-3fbae73b3864

Author: Avery Milandin

solutions/ess-maker-skills/scripts/adk_telemetry.py

@@ -437,17 +437,19 @@ def common_dimensions(
     tenant_name: str | None = None,
 ) -> dict[str, Any]:
     """Build the dimensions present on every event (spec Common Dimensions)."""
-    tid = _IDENTITY["tenant_id"] if tenant_id is None else tenant_id
-    # Fall back to the disk-persisted tenant_id when this process never
-    # called set_identity (subprocess launched by emit_capability.py from a
-    # SKILL.md step). Without this, capability events emitted by the shim
-    # carry an empty tenant_id, classify as "unknown" via classify_tenant(""),
-    # and never appear on the customer-filtered External dashboard even
-    # though the maker's earlier auth flow knew the tenant. Re-run the cache
-    # value through the same GUID sanitizer used in ``set_identity`` so a
-    # torn / legacy / hand-edited file can never bypass the check.
-    if not tid and tenant_id is None:
-        tid = _sanitize_tenant_id(_fc.get_cached_tenant_id())
+    # Single sanitization choke point: EVERY tenant_id source (explicit kwarg,
+    # in-memory identity, disk cache) flows through ``_sanitize_tenant_id``
+    # here, so a non-GUID value from *any* source normalizes to "" and the
+    # event lands in the "unknown" bucket instead of leaking into "customer".
+    # This matches the guarantee ``set_identity`` gives on ingress and closes
+    # the last back-door where an explicit ``common_dimensions(tenant_id=...)``
+    # kwarg (used by the ``emit_*`` helpers) could bypass the check.
+    if tenant_id is None:
+        tid = _IDENTITY["tenant_id"] or _sanitize_tenant_id(
+            _fc.get_cached_tenant_id()
+        )
+    else:
+        tid = _sanitize_tenant_id(tenant_id)
     tname = _IDENTITY["tenant_name"] if tenant_name is None else tenant_name
     # Fall back to the org display name a prior Graph-capable run (FlightCheck)
     # cached for THIS tenant, so pure-ADK events (session/build/deploy/

solutions/ess-maker-skills/scripts/flightcheck/graph_client.py

@@ -56,6 +56,16 @@
     "https://graph.microsoft.com/ExternalConnection.Read.All",
 ]
 
+# Wall-clock ceiling for the two silent lookups called from the interactive
+# auth hot path (``resolve_tenant_display_name_silent`` → /organization then
+# /me). Kept small — the sign-in critical path must not stall on a
+# best-effort telemetry-only label. ``tenant_class`` is driven by the ``tid``
+# claim already extracted from the token, so a timeout here only degrades
+# ``tenant_name`` to blank (recoverable on the next FlightCheck run), never
+# to a misclassification. Total worst case per authenticate() is
+# 2 × _SILENT_LOOKUP_TIMEOUT.
+_SILENT_LOOKUP_TIMEOUT = 3.0
+
 # Module-level requests Session with bounded retry-with-backoff for 429/5xx.
 # Mirrors the auth.py pattern - Graph throttles on /users and /servicePrincipals
 # in larger tenants, and one transient 503 mid-FlightCheck would otherwise blow
@@ -197,6 +207,13 @@ def _try_organization_lookup(app, account) -> tuple[str, bool]:
     Graph just refused this specific query — so the caller can meaningfully
     fall through to ``/me?$select=companyName``. When the transport is broken
     the ``/me`` call would fail identically, so we tell the caller to skip it.
+
+    Called on the interactive auth hot path so the socket timeout is short
+    (``_SILENT_LOOKUP_TIMEOUT``) and a plain ``requests.get`` is used instead
+    of the module-level ``_SESSION`` — the session's 3x 429/5xx retry with
+    ``backoff_factor=1`` would add up to ~7s of extra wait per lookup on a
+    throttled tenant, which is unacceptable on the sign-in critical path
+    for a best-effort telemetry label.
     """
     try:
         result = app.acquire_token_silent(_ORG_READ_SCOPE, account=account)
@@ -207,13 +224,13 @@ def _try_organization_lookup(app, account) -> tuple[str, bool]:
     except Exception:  # noqa: BLE001
         return "", True
     try:
-        resp = _SESSION.get(
+        resp = requests.get(
             f"{GRAPH_BASE}/organization",
             headers={
                 "Authorization": f"Bearer {result['access_token']}",
                 "Accept": "application/json",
             },
-            timeout=15,
+            timeout=_SILENT_LOOKUP_TIMEOUT,
         )
     except (requests.ConnectionError, requests.Timeout):
         return "", False
@@ -237,18 +254,26 @@ def _try_me_company_name_lookup(app, account) -> str:
     the same classification we already have for ``tenant_name`` — and never
     EUPI: we only read a single organizational attribute, not the user's
     name / UPN / id.
+
+    Runs on the interactive auth hot path in the common admin-consent-gap
+    case where ``/organization`` returned nothing, so the same bounded
+    ``_SILENT_LOOKUP_TIMEOUT`` + no-retry policy as the organization lookup
+    applies here. ``tenant_class`` is driven by the ``tid`` claim already in
+    the token, not by this call, so failing fast never hurts dashboard
+    correctness — it only degrades ``tenant_name`` to blank until the next
+    (interactive) FlightCheck run resolves and caches it.
     """
     try:
         result = app.acquire_token_silent(_USER_READ_SCOPE, account=account)
         if not result or "access_token" not in result:
             return ""
-        resp = _SESSION.get(
+        resp = requests.get(
             f"{GRAPH_BASE}/me?$select=companyName",
             headers={
                 "Authorization": f"Bearer {result['access_token']}",
                 "Accept": "application/json",
             },
-            timeout=15,
+            timeout=_SILENT_LOOKUP_TIMEOUT,
         )
         if resp.status_code != 200:
             return ""

solutions/ess-maker-skills/scripts/flightcheck/telemetry.py

@@ -66,6 +66,7 @@
 import json
 import os
 import platform
+import re
 import tempfile
 import uuid
 from datetime import datetime, timezone
@@ -122,6 +123,17 @@
 TENANT_CLASS_CUSTOMER = "customer"
 TENANT_CLASS_UNKNOWN = "unknown"
 
+# Canonical Entra tenant GUID (8-4-4-4-12 lowercase hex). Used by
+# ``get_cached_tenant_id`` to validate the legacy raw-string cache format
+# so a torn / hand-edited / garbage ``.tenant_id`` file can never leak onto
+# real events. Duplicated from ``adk_telemetry._GUID_RE`` (this module
+# cannot import from adk_telemetry — the dependency direction goes the
+# other way); the two regexes are kept in lock-step by a regression test
+# in ``test_adk_telemetry.py``.
+_GUID_RE = re.compile(
+    r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
+)
+
 
 def _internal_tenant_ids() -> frozenset[str]:
     """Allow-list of tenant GUIDs treated as internal Microsoft tenancy.
@@ -299,15 +311,26 @@ def cache_tenant_name(
         pass
     try:
         os.makedirs(local_dir, exist_ok=True)
-        with open(path, "w", encoding="utf-8") as f:
-            json.dump(
-                {
-                    "tenant_id": tenant_id,
-                    "tenant_name": tenant_name,
-                    "source": source,
-                },
-                f,
-            )
+        # Atomic write via tempfile + os.replace, matching cache_tenant_id.
+        # A concurrent reader (a sibling emit_capability.py subprocess spawned
+        # from a SKILL.md step) must never see a truncated file — otherwise
+        # the reader would fall back to a blank tenant_name for the whole
+        # session, undoing the very cache we're populating.
+        payload = {
+            "tenant_id": tenant_id,
+            "tenant_name": tenant_name,
+            "source": source,
+        }
+        fd, tmp = tempfile.mkstemp(prefix=".tenant_name.", dir=local_dir)
+        try:
+            with os.fdopen(fd, "w", encoding="utf-8") as f:
+                json.dump(payload, f)
+            os.replace(tmp, path)
+        except OSError:
+            try:
+                os.remove(tmp)
+            except OSError:
+                pass
     except OSError:
         pass
 
@@ -371,9 +394,11 @@ def get_cached_tenant_id(local_dir: str = ".local") -> str:
     real events.
 
     A very small compatibility shim recognizes the pre-versioning raw-string
-    format (a single line whose contents look like a GUID) so a maker who
-    upgrades mid-session doesn't lose their cached tenant. Any other content
-    (truncated GUID, JSON at a future schema version, garbage) is discarded.
+    format (a single line whose contents look like a canonical GUID) so a
+    maker who upgrades mid-session doesn't lose their cached tenant. The
+    raw string is validated against ``_GUID_RE`` (case-insensitive) before
+    being returned; any other content (truncated GUID, non-hex, garbage,
+    JSON at a future schema version) is discarded.
     """
     path = os.path.join(local_dir, _TENANT_ID_FILE)
     try:
@@ -391,8 +416,13 @@ def get_cached_tenant_id(local_dir: str = ".local") -> str:
         if not isinstance(obj, dict) or obj.get("version") != 1:
             return ""
         return str(obj.get("tenant_id", "")).strip()
-    # Legacy raw-string format from pre-versioned ADK builds.
-    return raw
+    # Legacy raw-string format from pre-versioned ADK builds: only trust the
+    # value when it matches the canonical GUID shape. Otherwise return ""
+    # and let the caller's fallback path (or classify_tenant) treat it as
+    # unknown — never leak a torn / hand-edited / garbage cache onto real
+    # events. Matches the docstring guarantee above.
+    v = raw.lower()
+    return v if _GUID_RE.match(v) else ""
 
 
 def get_cached_tenant_name(tenant_id: str, local_dir: str = ".local") -> str:

tests/flightcheck/test_graph_client.py

@@ -71,7 +71,7 @@ def test_success_returns_display_name():
     )
     resp = _resp(200, {"value": [{"displayName": "Contoso Ltd"}]})
     with patch.object(graph_client.msal, "PublicClientApplication", return_value=app), \
-         patch.object(graph_client._SESSION, "get", return_value=resp) as mock_get:
+         patch.object(graph_client.requests, "get", return_value=resp) as mock_get:
         assert (
             graph_client.resolve_tenant_display_name_silent("tenant-Z") == "Contoso Ltd"
         )
@@ -88,7 +88,7 @@ def test_non_200_response_returns_empty():
         accounts=[SimpleNamespace()], silent_result={"access_token": "tok"}
     )
     with patch.object(graph_client.msal, "PublicClientApplication", return_value=app), \
-         patch.object(graph_client._SESSION, "get", return_value=_resp(403)):
+         patch.object(graph_client.requests, "get", return_value=_resp(403)):
         assert graph_client.resolve_tenant_display_name_silent("tenant-Z") == ""
 
 
@@ -97,7 +97,7 @@ def test_empty_org_list_returns_empty():
         accounts=[SimpleNamespace()], silent_result={"access_token": "tok"}
     )
     with patch.object(graph_client.msal, "PublicClientApplication", return_value=app), \
-         patch.object(graph_client._SESSION, "get", return_value=_resp(200, {"value": []})):
+         patch.object(graph_client.requests, "get", return_value=_resp(200, {"value": []})):
         assert graph_client.resolve_tenant_display_name_silent("tenant-Z") == ""
 
 
@@ -138,7 +138,7 @@ def fake_get(url, headers, timeout):
         return _resp(200, {"companyName": "  Fabrikam Corp  "})
 
     with patch.object(graph_client.msal, "PublicClientApplication", return_value=app), \
-         patch.object(graph_client._SESSION, "get", side_effect=fake_get):
+         patch.object(graph_client.requests, "get", side_effect=fake_get):
         # Whitespace on companyName is stripped so downstream label matches
         # what /organization would return.
         assert (
@@ -184,5 +184,5 @@ def silent(scopes, account):
 
     app.acquire_token_silent.side_effect = silent
     with patch.object(graph_client.msal, "PublicClientApplication", return_value=app), \
-         patch.object(graph_client._SESSION, "get", return_value=_resp(200, {})):
+         patch.object(graph_client.requests, "get", return_value=_resp(200, {})):
         assert graph_client.resolve_tenant_display_name_silent("tenant-Z") == ""

tests/test_adk_telemetry.py

@@ -296,6 +296,39 @@ def test_classify_tenant_env_allowlist_extends(monkeypatch):
     assert _fc.classify_tenant("11111111-1111-1111-1111-111111111111") == "customer"
 
 
+def test_get_cached_tenant_id_rejects_non_guid_legacy_string(tmp_path):
+    # The legacy raw-string compatibility shim in get_cached_tenant_id must
+    # validate against _GUID_RE before returning — otherwise a torn /
+    # hand-edited / garbage .tenant_id file would ride onto real events.
+    d = tmp_path / "state"
+    d.mkdir()
+    (d / _fc._TENANT_ID_FILE).write_text("tenant-id", encoding="utf-8")
+    assert _fc.get_cached_tenant_id(local_dir=str(d)) == ""
+    # A canonical raw-string GUID (older ADK builds) is still honored.
+    (d / _fc._TENANT_ID_FILE).write_text(
+        "ABCDEF01-2345-6789-abcd-ef0123456789", encoding="utf-8"
+    )
+    assert (
+        _fc.get_cached_tenant_id(local_dir=str(d))
+        == "abcdef01-2345-6789-abcd-ef0123456789"
+    )
+
+
+def test_common_dimensions_sanitizes_explicit_tenant_id_kwarg(monkeypatch):
+    # Single choke point: EVERY tenant_id source (explicit kwarg included)
+    # is sanitized in common_dimensions, so an emit_* helper that forwards
+    # a caller-supplied non-GUID never lands in the customer bucket.
+    monkeypatch.setattr(_fc, "get_instance_id", lambda: "install-guid-1")
+    dims = adk.common_dimensions(adk.SURFACE_CLI, tenant_id="tenant-id")
+    assert dims["tenant_id"] == ""
+    assert dims["tenant_class"] == _fc.TENANT_CLASS_UNKNOWN
+    # A well-formed GUID is preserved (and lowercased).
+    dims = adk.common_dimensions(
+        adk.SURFACE_CLI, tenant_id="ABCDEF01-2345-6789-abcd-ef0123456789"
+    )
+    assert dims["tenant_id"] == "abcdef01-2345-6789-abcd-ef0123456789"
+
+
 def test_tenant_class_flows_into_dimensions(monkeypatch):
     monkeypatch.setattr(_fc, "get_instance_id", lambda: "install-guid-1")
     adk.set_identity(tenant_id=_fc.MICROSOFT_CORP_TENANT_ID, instance_id="inst-9")