licos-dev-sdk 0.4.6__tar.gz → 0.4.7__tar.gz

{licos_dev_sdk-0.4.6 → licos_dev_sdk-0.4.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: licos-dev-sdk
-Version: 0.4.6
+Version: 0.4.7
 Summary: LICOS Dev SDK - file generation and model capability clients
 Requires-Python: >=3.10
 Requires-Dist: beautifulsoup4>=4.12

{licos_dev_sdk-0.4.6 → licos_dev_sdk-0.4.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "licos-dev-sdk"
-version = "0.4.6"
+version = "0.4.7"
 description = "LICOS Dev SDK - file generation and model capability clients"
 requires-python = ">=3.10"
 dependencies = [

licos_dev_sdk-0.4.7/src/licos_dev_sdk/_usage_identity.py

@@ -0,0 +1,59 @@
+"""Renew existing model billing signatures without changing routing headers."""
+
+from __future__ import annotations
+
+import threading
+import time
+from datetime import datetime
+from typing import Any, Callable
+
+from licos_platform_sdk._runtime import ApiError
+
+IDENTITY_HEADERS = (
+    "X-AIOS-Usage-Context",
+    "X-AIOS-Usage-Context-Signature",
+    "X-AIOS-Provider-Credential-Id",
+    "X-AIOS-Provider-Workspace-Id",
+)
+REFRESH_LEAD_SECS = 30
+
+
+def has_signature(headers: dict[str, str]) -> bool:
+    return any(name.lower() in {value.lower() for value in IDENTITY_HEADERS[:2]} for name in headers)
+
+
+class UsageIdentityLease:
+    def __init__(self, detail: dict[str, Any] | None, parse_headers: Callable[[Any], dict[str, str]]) -> None:
+        self._parse_headers = parse_headers
+        self._headers, self._expires_at = self._parse(detail)
+        self._lock = threading.Lock()
+
+    def _parse(self, detail: dict[str, Any] | None) -> tuple[dict[str, str], float]:
+        if detail is None:
+            raise ApiError("model detail is missing usage identity")
+        raw = self._parse_headers(detail.get("requiredHeaders") or detail.get("required_headers"))
+        normalized = {name.lower(): value for name, value in raw.items()}
+        headers = {}
+        for name in IDENTITY_HEADERS:
+            value = normalized.get(name.lower())
+            if not value:
+                raise ApiError(f"model detail is missing {name}")
+            headers[name] = value
+        expiry = detail.get("usageContextExpiresAt")
+        try:
+            parsed = datetime.fromisoformat(str(expiry).replace("Z", "+00:00"))
+            if parsed.tzinfo is None:
+                raise ValueError("expiry requires timezone")
+            expires_at = parsed.timestamp()
+        except ValueError as exc:
+            raise ApiError("model detail has invalid usageContextExpiresAt") from exc
+        return headers, expires_at
+
+    def current_headers(self, fetch_detail: Callable[[], dict[str, Any] | None]) -> dict[str, str]:
+        with self._lock:
+            if self._expires_at <= time.time() + REFRESH_LEAD_SECS:
+                headers, expires_at = self._parse(fetch_detail())
+                if expires_at <= time.time() + REFRESH_LEAD_SECS:
+                    raise ApiError("renewed model usage identity expires before the request can start")
+                self._headers, self._expires_at = headers, expires_at
+            return dict(self._headers)

{licos_dev_sdk-0.4.6 → licos_dev_sdk-0.4.7}/src/licos_dev_sdk/model.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import atexit
+import asyncio
 import base64
 import json
 import logging
@@ -25,6 +26,7 @@ from licos_platform_sdk._runtime import (
     resolve_user_token,
     should_refresh_user_token,
 )
+from ._usage_identity import UsageIdentityLease, has_signature
 
 
 AVAILABLE_MODELS_PATH_TEMPLATE = "/api/v1/admin/workspaces/{workspace_id}/models/available"
@@ -92,6 +94,7 @@ class ModelEndpoint:
     context_length: int | None = None
     max_input_length: int | None = None
     max_output_length: int | None = None
+    usage_identity: UsageIdentityLease | None = field(default=None, repr=False, compare=False)
 
     @property
     def async_task(self) -> bool:
@@ -575,8 +578,11 @@ def _LangChainUsageHandler(
     endpoint: ModelEndpoint,
 ) -> Any:
     class UsageHandler(base_handler):
+        run_inline = True
+
         def __init__(self) -> None:
             super().__init__()
+            self._routing_headers = _routing_headers(runtime)
             self._request_id: ContextVar[str | None] = ContextVar(
                 "licos_langchain_model_request_id",
                 default=None,
@@ -595,12 +601,37 @@ def _LangChainUsageHandler(
         def on_chat_model_start(self, *_args: Any, **_kwargs: Any) -> None:
             self._current_request_id()
 
+        def _matches_model_request(self, request_value: Any) -> bool:
+            if request_value.method != "POST" or request_value.headers.get("content-type", "").split(";")[0] != "application/json":
+                return False
+            return json.loads(request_value.content).get("model") == endpoint.model and all(
+                request_value.headers.get(name) == value for name, value in self._routing_headers.items()
+            )
+
         def on_http_request(self, request_value: Any) -> None:
+            if not self._matches_model_request(request_value):
+                return
             request_id = self._current_request_id()
             request_value.headers[MODEL_REQUEST_ID_HEADER] = request_id
+            if request_value.extensions.get("licos_usage_identity_managed"):
+                return
+            self._apply_usage_identity(request_value, _current_usage_identity(endpoint, runtime))
+
+        def _apply_usage_identity(self, request_value: Any, headers: dict[str, str]) -> None:
+            for name, value in headers.items():
+                for existing in list(request_value.headers):
+                    if existing.lower() == name.lower():
+                        del request_value.headers[existing]
+                request_value.headers[name] = value
 
         async def on_async_http_request(self, request_value: Any) -> None:
-            self.on_http_request(request_value)
+            if not self._matches_model_request(request_value):
+                return
+            request_value.headers[MODEL_REQUEST_ID_HEADER] = self._current_request_id()
+            if request_value.extensions.get("licos_usage_identity_managed"):
+                return
+            headers = await asyncio.to_thread(_current_usage_identity, endpoint, runtime)
+            self._apply_usage_identity(request_value, headers)
 
         def on_llm_end(self, response: Any, **_kwargs: Any) -> None:
             usage = None
@@ -1640,7 +1671,22 @@ def _enterprise_org_id(value: str | None) -> str | None:
 
 def _request_headers(endpoint: ModelEndpoint, runtime: ModelRuntime) -> dict[str, str]:
     headers = endpoint.request_headers or endpoint.required_headers
-    return _without_authorization(_merge_headers(headers, _routing_headers(runtime)))
+    return _without_authorization(_merge_headers(headers, _routing_headers(runtime), _current_usage_identity(endpoint, runtime)))
+
+
+def _current_usage_identity(endpoint: ModelEndpoint, runtime: ModelRuntime) -> dict[str, str]:
+    if endpoint.usage_identity is None:
+        return {}
+    return endpoint.usage_identity.current_headers(
+        lambda: _fetch_model_detail(runtime, endpoint.model, refresh=True)
+    )
+
+
+def _with_usage_identity(endpoint: ModelEndpoint, detail: dict[str, Any] | None) -> ModelEndpoint:
+    detail_headers = _parse_required_headers((detail or {}).get("requiredHeaders") or (detail or {}).get("required_headers"))
+    if not has_signature(_merge_headers(endpoint.required_headers, endpoint.request_headers, detail_headers)):
+        return endpoint
+    return replace(endpoint, usage_identity=UsageIdentityLease(detail, _parse_required_headers))
 
 
 def _merge_headers(*groups: dict[str, str] | None) -> dict[str, str]:
@@ -2130,7 +2176,7 @@ def _resolve_chat_endpoint(
     selected_model = _selected_model(requested_model, endpoint.model)
     endpoint = replace(endpoint, model=selected_model)
     detail = _fetch_model_detail(runtime, selected_model)
-    return _apply_model_detail(endpoint, detail)
+    return _with_usage_identity(_apply_model_detail(endpoint, detail), detail)
 
 
 def _resolve_endpoint_with_detail(
@@ -2149,7 +2195,7 @@ def _resolve_endpoint_with_detail(
         input_capabilities=input_capabilities,
     )
     detail = _fetch_model_detail(runtime, endpoint.model)
-    return _apply_model_detail(endpoint, detail)
+    return _with_usage_identity(_apply_model_detail(endpoint, detail), detail)
 
 
 def _resolve_endpoint(

licos_dev_sdk-0.4.7/tests/test_usage_identity.py

@@ -0,0 +1,176 @@
+from __future__ import annotations
+
+import asyncio
+import sys
+import unittest
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime, timezone
+from pathlib import Path
+from unittest import mock
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "licos-platform-sdk" / "src"))
+
+import httpx
+
+from licos_dev_sdk import model
+from licos_dev_sdk._usage_identity import IDENTITY_HEADERS, UsageIdentityLease
+from licos_platform_sdk._runtime import ApiError
+
+
+def detail(context: str, expires_at: float) -> dict:
+    return {
+        "usageContextExpiresAt": datetime.fromtimestamp(expires_at, timezone.utc).isoformat(),
+        "requiredHeaders": {
+            IDENTITY_HEADERS[0].lower(): context,
+            IDENTITY_HEADERS[1]: f"{context}-signature",
+            IDENTITY_HEADERS[2]: f"{context}-credential",
+            IDENTITY_HEADERS[3]: f"{context}-workspace",
+            "X-Provider-Unrelated": "must-not-merge",
+            "X-Licos-Owner-User-Id": "must-not-replace",
+        },
+    }
+
+
+def endpoint(expires_at: float) -> model.ModelEndpoint:
+    return model._with_usage_identity(
+        model.ModelEndpoint(
+            provider="bailian", capability="videoGeneration", model="qwen-test",
+            base_url="http://gateway.example/model",
+            required_headers={IDENTITY_HEADERS[0].lower(): "catalog-stale", "X-Provider-Route": "route"},
+            endpoint={"async": True, "taskQueryUrl": "http://gateway.example/tasks/{taskId}"},
+        ),
+        detail("initial", expires_at),
+    )
+
+
+class UsageIdentityTests(unittest.TestCase):
+    def setUp(self) -> None:
+        self.runtime = model.ModelRuntime(
+            base_url="http://platform.example", token="user-token", workspace_id="workspace-1",
+            user_id="owner-1", current_user_id="caller-1", owner_project_id="project-1",
+        )
+
+    def test_valid_lease_reuses_signed_snapshot(self) -> None:
+        with mock.patch("time.time", return_value=100):
+            lease = UsageIdentityLease(detail("initial", 200), model._parse_required_headers)
+            fetch = mock.Mock()
+            snapshot = lease.current_headers(fetch)
+            snapshot[IDENTITY_HEADERS[0]] = "mutated"
+            self.assertEqual(lease.current_headers(fetch)[IDENTITY_HEADERS[0]], "initial")
+            fetch.assert_not_called()
+
+    def test_concurrent_requests_share_one_near_expiry_refresh(self) -> None:
+        with mock.patch("time.time", return_value=100):
+            lease = UsageIdentityLease(detail("initial", 120), model._parse_required_headers)
+            fetch = mock.Mock(return_value=detail("renewed", 200))
+            with ThreadPoolExecutor(max_workers=8) as executor:
+                snapshots = list(executor.map(lambda _: lease.current_headers(fetch), range(16)))
+            fetch.assert_called_once()
+            self.assertTrue(all(value[IDENTITY_HEADERS[0]] == "renewed" for value in snapshots))
+
+    def test_invalid_signed_detail_is_rejected(self) -> None:
+        for missing in (*IDENTITY_HEADERS[1:], "usageContextExpiresAt"):
+            invalid = detail("initial", 200)
+            if missing == "usageContextExpiresAt":
+                del invalid[missing]
+            else:
+                del invalid["requiredHeaders"][missing]
+            with self.subTest(missing=missing), self.assertRaises(ApiError):
+                UsageIdentityLease(invalid, model._parse_required_headers)
+
+    def test_expired_renewal_never_sends_model_request(self) -> None:
+        with mock.patch("time.time", return_value=100):
+            selected = endpoint(90)
+            with mock.patch.object(model, "_fetch_model_detail", return_value=detail("still-expired", 110)), \
+                    mock.patch.object(model, "_request_json") as request:
+                with self.assertRaisesRegex(ApiError, "expires before"):
+                    model._post_model_json(selected, self.runtime, {"model": "qwen-test"})
+                request.assert_not_called()
+
+    def test_refresh_failure_never_sends_model_request(self) -> None:
+        with mock.patch("time.time", return_value=100):
+            selected = endpoint(90)
+            with mock.patch.object(model, "_fetch_model_detail", side_effect=ApiError("catalog unavailable")), \
+                    mock.patch.object(model, "_request_json") as request:
+                with self.assertRaisesRegex(ApiError, "catalog unavailable"):
+                    model._post_model_json(selected, self.runtime, {"model": "qwen-test"})
+                request.assert_not_called()
+
+    def test_submit_replaces_only_four_headers_and_retains_owner_and_caller(self) -> None:
+        with mock.patch("time.time", return_value=100):
+            selected = endpoint(120)
+            with mock.patch.object(model, "_fetch_model_detail", return_value=detail("renewed", 200)) as fetch, \
+                    mock.patch.object(model, "_request_json", return_value={}) as request:
+                for _ in range(2):
+                    model._post_model_json(selected, self.runtime, {"model": "qwen-test"})
+                fetch.assert_called_once_with(self.runtime, "qwen-test", refresh=True)
+                headers = request.call_args.kwargs["headers"]
+                self.assertEqual(headers[IDENTITY_HEADERS[0]], "renewed")
+                self.assertNotIn(IDENTITY_HEADERS[0].lower(), headers)
+                self.assertEqual(headers["X-Licos-Owner-User-Id"], "owner-1")
+                self.assertEqual(headers["X-Licos-Current-User-Id"], "caller-1")
+                self.assertEqual(headers["X-Provider-Route"], "route")
+                self.assertEqual(headers["X-DashScope-Async"], "enable")
+                self.assertNotIn("X-Provider-Unrelated", headers)
+
+    def test_media_poll_refreshes_after_submit_identity_nears_expiry(self) -> None:
+        clock = [100]
+        with mock.patch("time.time", side_effect=lambda: clock[0]):
+            selected = endpoint(200)
+            captured = []
+
+            def request(method, url, **kwargs):
+                captured.append(kwargs["headers"])
+                return {"output": {"task_status": "RUNNING" if len(captured) == 1 else "SUCCEEDED"}}
+
+            with mock.patch.object(model, "_request_json", side_effect=request), \
+                    mock.patch.object(model, "_fetch_model_detail", return_value=detail("renewed", 500)) as fetch, \
+                    mock.patch.object(model, "_enqueue_model_usage"), \
+                    mock.patch("time.sleep", side_effect=lambda _: clock.__setitem__(0, 180)):
+                model._await_async_model_result(selected, self.runtime, {"output": {"task_id": "task-1"}})
+            fetch.assert_called_once_with(self.runtime, "qwen-test", refresh=True)
+            self.assertEqual([headers[IDENTITY_HEADERS[0]] for headers in captured], ["initial", "renewed"])
+            self.assertTrue(all(headers["X-Licos-Current-User-Id"] == "caller-1" for headers in captured))
+
+    def test_langchain_sync_hook_renews_static_default_headers(self) -> None:
+        with mock.patch("time.time", return_value=100):
+            handler = model._LangChainUsageHandler(object, self.runtime, endpoint(90))
+            req = httpx.Request("POST", "http://gateway.example/model", headers={
+                IDENTITY_HEADERS[0].lower(): "stale", **model._routing_headers(self.runtime),
+            }, json={"model": "qwen-test"})
+            with mock.patch.object(model, "_fetch_model_detail", return_value=detail("renewed", 200)) as fetch:
+                handler.on_http_request(req)
+                handler.on_http_request(req)
+            fetch.assert_called_once()
+            self.assertEqual(req.headers.get_list(IDENTITY_HEADERS[0]), ["renewed"])
+            self.assertEqual(req.headers["X-Licos-Owner-User-Id"], "owner-1")
+
+    def test_langchain_async_hooks_share_renewal(self) -> None:
+        with mock.patch("time.time", return_value=100):
+            handler = model._LangChainUsageHandler(object, self.runtime, endpoint(90))
+            requests = [httpx.Request("POST", "http://gateway.example/model", json={"model": "qwen-test"},
+                                      headers=model._routing_headers(self.runtime)) for _ in range(8)]
+
+            async def run():
+                await asyncio.gather(*(handler.on_async_http_request(req) for req in requests))
+
+            with mock.patch.object(model, "_fetch_model_detail", return_value=detail("renewed", 200)) as fetch:
+                asyncio.run(run())
+            fetch.assert_called_once()
+            self.assertTrue(all(req.headers[IDENTITY_HEADERS[0]] == "renewed" for req in requests))
+
+    def test_shared_http_client_hooks_do_not_replace_other_user_or_model_signatures(self) -> None:
+        with mock.patch("time.time", return_value=100):
+            handler = model._LangChainUsageHandler(object, self.runtime, endpoint(90))
+            for body, headers in [
+                ({"model": "another-model"}, model._routing_headers(self.runtime)),
+                ({"model": "qwen-test"}, {**model._routing_headers(self.runtime), "X-Licos-Current-User-Id": "another-user"}),
+            ]:
+                with self.subTest(body=body, headers=headers):
+                    req = httpx.Request("POST", "http://gateway.example/model", json=body,
+                                        headers={**headers, IDENTITY_HEADERS[0]: "other-context"})
+                    with mock.patch.object(model, "_fetch_model_detail") as fetch:
+                        handler.on_http_request(req)
+                    fetch.assert_not_called()
+                    self.assertEqual(req.headers[IDENTITY_HEADERS[0]], "other-context")