ai: forward X-LC-UID so User API Keys can run org AI sessions

limacharlie/sdk/ai.py

@@ -40,6 +40,21 @@ def _resolve_map_secrets(self, m: dict[str, str] | None) -> dict[str, str] | Non
             return m
         return {k: self._resolve_secret(v) for k, v in m.items()}
 
+    def _org_auth_headers(self) -> dict[str, str]:
+        """Identity headers for the org-scoped ai-sessions endpoints.
+
+        Always carries ``X-LC-OID``.  A User API Key is scoped to a
+        user rather than an org, so jwt.limacharlie.io can only resolve
+        it when the UID accompanies the secret -- forward it via
+        ``X-LC-UID``.  ai-sessions' OrgDualAuthMiddleware reads
+        ``X-LC-UID`` only for the raw-API-key path and ignores it for
+        JWT auth, so sending it whenever a uid is known is always safe.
+        """
+        headers = {"X-LC-OID": self.client._oid}
+        if self.client._uid:
+            headers["X-LC-UID"] = self.client._uid
+        return headers
+
     # Fields copied verbatim from an ai_agent hive record into the
     # request's ``profile`` section.  Every entry in this tuple maps
     # one-to-one to a field on the server's ``ProfileContent`` type
@@ -210,7 +225,7 @@ def start_session(self, definition_name: str, prompt: str | None = None,
         if profile:
             request_body["profile"] = profile
 
-        extra = {"X-LC-OID": self.client._oid}
+        extra = self._org_auth_headers()
 
         # Use the raw API key when available (works with current and future
         # ai-sessions deployments).  Fall back to JWT auth for OAuth users
@@ -331,7 +346,7 @@ def _org_request(self, verb: str, path: str,
         OrgDualAuthMiddleware. We send the raw API key when available
         (same pattern as start_session) for maximum compatibility.
         """
-        extra: dict[str, str] = {"X-LC-OID": self.client._oid}
+        extra: dict[str, str] = self._org_auth_headers()
         if self.client._api_key is not None:
             extra["Authorization"] = f"Bearer {self.client._api_key}"
             return self.client.request(

tests/unit/test_sdk_ai_sessions.py

@@ -19,6 +19,10 @@ def mock_org():
     org.client = MagicMock()
     org.client._api_key = "test-api-key"
     org.client._oid = "test-oid"
+    # Org-scoped credentials by default: no user id.  (Without this the
+    # MagicMock would auto-vivify a truthy _uid and leak an X-LC-UID
+    # header into every request.)
+    org.client._uid = None
     return org
 
 
@@ -85,6 +89,52 @@ def test_sends_correct_request(self, ai, mock_org):
         assert result == {"session_id": "sess-123", "status": "pending"}
 
 
+class TestOrgRequestUidHeader:
+    """X-LC-UID is forwarded for User API Keys so jwt.limacharlie.io
+    can resolve a user-scoped key against the requested org."""
+
+    def _start(self, ai, mock_org):
+        with patch("limacharlie.sdk.hive.Hive") as MockHive:
+            hive_instance = MagicMock()
+            MockHive.return_value = hive_instance
+            hive_instance.get.return_value = _make_hive_record(
+                {"prompt": "p", "anthropic_secret": "sk", "lc_api_key_secret": "k"}
+            )
+            mock_org.client.request.return_value = {"session_id": "s1"}
+            ai.start_session("agent")
+        return mock_org.client.request.call_args[1]["extra_headers"]
+
+    def test_org_scoped_key_omits_uid_header(self, ai, mock_org):
+        # mock_org defaults to _uid = None (org-scoped).
+        headers = self._start(ai, mock_org)
+        assert headers["X-LC-OID"] == "test-oid"
+        assert "X-LC-UID" not in headers
+
+    def test_user_scoped_key_sends_uid_header(self, ai, mock_org):
+        mock_org.client._uid = "user-123"
+        headers = self._start(ai, mock_org)
+        assert headers["X-LC-OID"] == "test-oid"
+        assert headers["X-LC-UID"] == "user-123"
+        assert headers["Authorization"] == "Bearer test-api-key"
+
+    def test_org_request_path_forwards_uid(self, ai, mock_org):
+        # The same identity headers must reach the org session
+        # management endpoints (list/get/terminate), not just
+        # start-session.
+        mock_org.client._uid = "user-123"
+        mock_org.client.request.return_value = {"session": {}}
+        ai.get_session("sess-1")
+        headers = mock_org.client.request.call_args[1]["extra_headers"]
+        assert headers["X-LC-OID"] == "test-oid"
+        assert headers["X-LC-UID"] == "user-123"
+
+    def test_org_request_omits_uid_when_org_scoped(self, ai, mock_org):
+        mock_org.client.request.return_value = {"session": {}}
+        ai.get_session("sess-1")
+        headers = mock_org.client.request.call_args[1]["extra_headers"]
+        assert "X-LC-UID" not in headers
+
+
 class TestStartSessionSecretResolution:
     """Credentials that use hive://secret/ references get resolved."""