diff --git a/AUDIT.md b/AUDIT.md index 71c6783..3c50900 100644 --- a/AUDIT.md +++ b/AUDIT.md @@ -205,3 +205,22 @@ Added `StreamReservation` and `AsyncStreamReservation` context managers that aut - **Error handling:** `RESERVATION_FINALIZED`, `RESERVATION_EXPIRED`, and `IDEMPOTENCY_MISMATCH` do not trigger release; other 4xx client errors do trigger release — matches lifecycle.py behavior exactly Protocol conformance: No new endpoints or protocol changes. All reservation, commit, release, and extend calls use the same client methods and body formats as the decorator path. Verified by 64 unit tests covering success, deny, error, retry, heartbeat, cost resolution, context propagation, spec validation, and all commit error-code branches. + +--- + +## Dynamic Subject & Action Fields on `@cycles` (added 2026-04-27) + +**Issue:** [#45](https://github.com/runcycles/cycles-client-python/issues/45) +**Files:** `runcycles/lifecycle.py`, `runcycles/decorator.py` +**Test files:** `tests/test_lifecycle.py`, `tests/test_decorator.py` +**Version:** 0.4.0 + +Widened the `@cycles` decorator to accept callables — in addition to constants — for every field that previously had to be static at decoration time. Mirrors the existing `estimate` / `actual` callable contract and re-aligns the Python client with the Java client's `@Cycles(workspace = "#workspaceId")` SpEL behavior shipped in `cycles-spring-boot-starter` 0.2.1 ([java#50](https://github.com/runcycles/cycles-spring-boot-starter/pull/50)). + +- **Newly callable fields:** `tenant`, `workspace`, `app`, `workflow`, `agent`, `toolset`, `action_kind`, `action_name`, `action_tags`, `dimensions`. Each accepts `T | Callable[..., T | None] | None`. +- **Resolution:** new `_resolve_value(val, args, kwargs)` helper in `lifecycle.py` invokes the callable with the decorated function's `*args, **kwargs` at reservation time; constants pass through untouched. +- **Fallback semantics preserved:** subject callables returning `None` fall through to `default_subject_fields` (client config); `action_kind` / `action_name` returning `None` fall through to `"unknown"`; `action_tags` / `dimensions` returning `None` are omitted. Constants behave identically to today (regression-tested). +- **Fail-fast:** exceptions raised inside a user callable propagate to the decorator caller without creating a reservation. +- **Signature change:** `_build_reservation_body` now takes `args` and `kwargs` parameters; both `CyclesLifecycle.execute` and `AsyncCyclesLifecycle.execute` thread them through. + +Protocol conformance: No protocol or wire-format changes. The reservation request body shape is unchanged — only the source of each field's value is widened. Verified by new unit tests in `TestCallableSubjectFields`, `TestCallableActionFields`, `TestCallableDimensions` plus an end-to-end decorator test asserting the captured request body. diff --git a/CHANGELOG.md b/CHANGELOG.md index bb8e8c1..4299eb6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,18 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog 1.1.0](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.4.0] - 2026-04-27 + +Dynamic subject and action fields on the `@cycles` decorator. + +### Added + +- Subject fields (`tenant`, `workspace`, `app`, `workflow`, `agent`, `toolset`), action fields (`action_kind`, `action_name`, `action_tags`), and `dimensions` now accept callables in addition to constants. Callables are invoked with the decorated function's `*args, **kwargs` at reservation time, enabling per-call budget routing and dynamic action labeling. Mirrors the Java client's SpEL behavior. (#45) + +### Changed + +- `_build_reservation_body` signature widened to thread `args` / `kwargs` through to the new `_resolve_value` helper. Internal API only; no protocol or wire-format changes. + ## [0.3.0] - 2026-04-08 Add streaming support. @@ -108,6 +120,7 @@ Initial public release. - Comprehensive error handling and improved API model validation (#1) +[0.4.0]: https://github.com/runcycles/cycles-client-python/compare/v0.3.0...v0.4.0 [0.3.0]: https://github.com/runcycles/cycles-client-python/compare/v0.2.0...v0.3.0 [0.2.0]: https://github.com/runcycles/cycles-client-python/compare/v0.1.3...v0.2.0 [0.1.3]: https://github.com/runcycles/cycles-client-python/compare/v0.1.2...v0.1.3 diff --git a/README.md b/README.md index f4cd8b6..162f65d 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,30 @@ result = call_llm("Hello", tokens=100) > ``` > The key (e.g. `cyc_live_abc123...`) is shown only once — save it immediately. For key rotation and lifecycle details, see [API Key Management](https://runcycles.io/how-to/api-key-management-in-cycles). +### Dynamic subject and action fields + +Subject fields (`tenant`, `workspace`, `app`, `workflow`, `agent`, `toolset`), action fields (`action_kind`, `action_name`, `action_tags`), and `dimensions` all accept either a constant or a callable. When given a callable, it is invoked with the decorated function's `*args, **kwargs` at reservation time — useful for routing per-call to different budget scopes or labeling actions dynamically: + +```python +@cycles( + estimate=lambda req, workspace_id: req.tokens * 10, + workspace=lambda req, workspace_id: workspace_id, # per-call budget routing + action_kind=lambda req, *_: f"llm.{req.provider}", # dynamic action label + action_name=lambda req, *_: req.model, + dimensions=lambda req, *_: {"region": req.region}, + client=client, +) +def run_request(req: ResponseRequest, workspace_id: str) -> Response: + ... +``` + +Fallback semantics mirror the constant case: + +- Subject callables returning `None` fall through to the client-config default (`CyclesConfig(workspace=...)`). +- `action_kind` / `action_name` returning `None` fall through to `"unknown"`. +- `action_tags` / `dimensions` returning `None` are omitted from the request. +- A callable that raises propagates the exception — fail-fast — without creating a reservation. + ### Budget lifecycle The `@cycles` decorator wraps your function in a reserve → execute → commit/release lifecycle: diff --git a/examples/decorator_usage.py b/examples/decorator_usage.py index 9bb5d36..567fdb8 100644 --- a/examples/decorator_usage.py +++ b/examples/decorator_usage.py @@ -49,6 +49,19 @@ def call_llm(prompt: str, tokens: int) -> str: return "Generated response for: " + prompt +# Per-call subject / action routing via callables — resolved at reservation time +# against the wrapped function's *args, **kwargs +@cycles( + estimate=1000, + workspace=lambda req, workspace_id: workspace_id, + action_kind=lambda req, *_: f"llm.{req['provider']}", + action_name=lambda req, *_: req["model"], + client=client, +) +def run_request(req: dict[str, str], workspace_id: str) -> str: + return f"Routed {req['model']} to {workspace_id}" + + def main() -> None: print("Simple call:") result1 = simple_call() @@ -58,6 +71,10 @@ def main() -> None: result2 = call_llm("Tell me a joke", tokens=200) print(f" Result: {result2}") + print("\nPer-call subject/action routing:") + result3 = run_request({"provider": "openai", "model": "gpt-4"}, workspace_id="ws-42") + print(f" Result: {result3}") + if __name__ == "__main__": main() diff --git a/pyproject.toml b/pyproject.toml index c51a36f..675c41c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "runcycles" -version = "0.3.0" +version = "0.4.0" description = "Python client for the Cycles budget-management protocol" readme = "README.md" license = "Apache-2.0" diff --git a/runcycles/decorator.py b/runcycles/decorator.py index b4751bb..abf3ac4 100644 --- a/runcycles/decorator.py +++ b/runcycles/decorator.py @@ -56,46 +56,57 @@ def cycles( estimate: int | Callable[..., int], *, actual: int | Callable[..., int] | None = None, - action_kind: str | None = None, - action_name: str | None = None, - action_tags: list[str] | None = None, + action_kind: str | Callable[..., str | None] | None = None, + action_name: str | Callable[..., str | None] | None = None, + action_tags: list[str] | Callable[..., list[str] | None] | None = None, unit: Unit | str = Unit.USD_MICROCENTS, ttl_ms: int = 60_000, grace_period_ms: int | None = None, overage_policy: str = "ALLOW_IF_AVAILABLE", dry_run: bool = False, - tenant: str | None = None, - workspace: str | None = None, - app: str | None = None, - workflow: str | None = None, - agent: str | None = None, - toolset: str | None = None, - dimensions: dict[str, str] | None = None, + tenant: str | Callable[..., str | None] | None = None, + workspace: str | Callable[..., str | None] | None = None, + app: str | Callable[..., str | None] | None = None, + workflow: str | Callable[..., str | None] | None = None, + agent: str | Callable[..., str | None] | None = None, + toolset: str | Callable[..., str | None] | None = None, + dimensions: dict[str, str] | Callable[..., dict[str, str] | None] | None = None, client: CyclesClient | AsyncCyclesClient | None = None, use_estimate_if_actual_not_provided: bool = True, ) -> Callable[[F], F]: """Decorator that wraps a function with the Cycles reserve/execute/commit lifecycle. + Subject and action fields accept either a constant or a callable. When given a + callable, it is invoked with the decorated function's ``*args, **kwargs`` at + reservation time. Subject callables returning ``None`` fall through to the + client-config default; ``action_kind`` / ``action_name`` returning ``None`` fall + through to ``"unknown"``; ``action_tags`` / ``dimensions`` returning ``None`` are + omitted. + Args: estimate: Estimated cost. Either an int constant or a callable that receives the decorated function's ``*args, **kwargs`` and returns an int. actual: Actual cost. Either an int constant or a callable that receives the function's return value and returns an int. Defaults to the estimate. - action_kind: Action category (e.g. "llm.completion"). - action_name: Action identifier (e.g. "gpt-4"). - action_tags: Optional tags for filtering/reporting. + action_kind: Action category (e.g. "llm.completion"). Constant or callable + receiving ``*args, **kwargs``. + action_name: Action identifier (e.g. "gpt-4"). Constant or callable + receiving ``*args, **kwargs``. + action_tags: Optional tags for filtering/reporting. Constant list or callable + receiving ``*args, **kwargs`` returning a list. unit: Cost unit. Default: USD_MICROCENTS. ttl_ms: Reservation TTL in milliseconds. Default: 60000. grace_period_ms: Grace period after TTL expiry in milliseconds. overage_policy: REJECT, ALLOW_IF_AVAILABLE (default), or ALLOW_WITH_OVERDRAFT. dry_run: If True, evaluate without persisting (method won't execute). - tenant: Subject tenant override. - workspace: Subject workspace override. - app: Subject app override. - workflow: Subject workflow override. - agent: Subject agent override. - toolset: Subject toolset override. - dimensions: Custom dimensions for the subject. + tenant: Subject tenant override. Constant or callable receiving ``*args, **kwargs``. + workspace: Subject workspace override. Constant or callable receiving ``*args, **kwargs``. + app: Subject app override. Constant or callable receiving ``*args, **kwargs``. + workflow: Subject workflow override. Constant or callable receiving ``*args, **kwargs``. + agent: Subject agent override. Constant or callable receiving ``*args, **kwargs``. + toolset: Subject toolset override. Constant or callable receiving ``*args, **kwargs``. + dimensions: Custom dimensions for the subject. Constant dict or callable + receiving ``*args, **kwargs`` returning a dict. client: Explicit Cycles client to use. Falls back to module-level default. use_estimate_if_actual_not_provided: If True and actual is None, use estimate as actual. @@ -116,6 +127,17 @@ def call_llm(prompt: str) -> str: ) def call_llm(prompt: str, tokens: int) -> str: return openai.complete(prompt, max_tokens=tokens) + + # Per-call subject/action routing via callables + @cycles( + estimate=1000, + workspace=lambda req, workspace_id: workspace_id, + action_kind=lambda req, *_: f"llm.{req.provider}", + action_name=lambda req, *_: req.model, + client=my_client, + ) + def run_request(req: Request, workspace_id: str) -> Response: + ... """ unit_str = unit.value if isinstance(unit, Unit) else str(unit) diff --git a/runcycles/lifecycle.py b/runcycles/lifecycle.py index 26b0f98..0053e69 100644 --- a/runcycles/lifecycle.py +++ b/runcycles/lifecycle.py @@ -47,21 +47,21 @@ class DecoratorConfig: estimate: int | Callable[..., int] actual: int | Callable[..., int] | None = None - action_kind: str | None = None - action_name: str | None = None - action_tags: list[str] | None = None + action_kind: str | Callable[..., str | None] | None = None + action_name: str | Callable[..., str | None] | None = None + action_tags: list[str] | Callable[..., list[str] | None] | None = None unit: str = "USD_MICROCENTS" ttl_ms: int = 60_000 grace_period_ms: int | None = None overage_policy: str = "ALLOW_IF_AVAILABLE" dry_run: bool = False - tenant: str | None = None - workspace: str | None = None - app: str | None = None - workflow: str | None = None - agent: str | None = None - toolset: str | None = None - dimensions: dict[str, str] | None = None + tenant: str | Callable[..., str | None] | None = None + workspace: str | Callable[..., str | None] | None = None + app: str | Callable[..., str | None] | None = None + workflow: str | Callable[..., str | None] | None = None + agent: str | Callable[..., str | None] | None = None + toolset: str | Callable[..., str | None] | None = None + dimensions: dict[str, str] | Callable[..., dict[str, str] | None] | None = None use_estimate_if_actual_not_provided: bool = True @@ -72,6 +72,13 @@ def _evaluate_amount(expr: int | Callable[..., int], args: tuple[Any, ...], kwar return int(expr) +def _resolve_value(val: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + """Resolve a decorator value: invoke if callable, else return as-is.""" + if callable(val): + return val(*args, **kwargs) + return val + + def _evaluate_actual( expr: int | Callable[..., int] | None, result: Any, @@ -89,7 +96,11 @@ def _evaluate_actual( def _build_reservation_body( - cfg: DecoratorConfig, estimate: int, default_subject_fields: dict[str, str | None], + cfg: DecoratorConfig, + estimate: int, + default_subject_fields: dict[str, str | None], + args: tuple[Any, ...], + kwargs: dict[str, Any], ) -> dict[str, Any]: """Build the reservation create request body.""" validate_non_negative(estimate, "estimate") @@ -97,21 +108,27 @@ def _build_reservation_body( subject: dict[str, Any] = {} for field_name in ("tenant", "workspace", "app", "workflow", "agent", "toolset"): - val = getattr(cfg, field_name, None) or default_subject_fields.get(field_name) + val = _resolve_value(getattr(cfg, field_name, None), args, kwargs) + if not val: + val = default_subject_fields.get(field_name) if val: subject[field_name] = val - if cfg.dimensions: - subject["dimensions"] = cfg.dimensions + dims = _resolve_value(cfg.dimensions, args, kwargs) + if dims: + subject["dimensions"] = dims subject_model = Subject(**subject) validate_subject(subject_model) + kind = _resolve_value(cfg.action_kind, args, kwargs) + name = _resolve_value(cfg.action_name, args, kwargs) + tags = _resolve_value(cfg.action_tags, args, kwargs) action: dict[str, Any] = { - "kind": cfg.action_kind or "unknown", - "name": cfg.action_name or "unknown", + "kind": kind or "unknown", + "name": name or "unknown", } - if cfg.action_tags: - action["tags"] = cfg.action_tags + if tags: + action["tags"] = tags body: dict[str, Any] = { "idempotency_key": str(uuid.uuid4()), @@ -234,7 +251,7 @@ def execute( logger.debug("Estimated usage: estimate=%d", estimate) # Create reservation - create_body = _build_reservation_body(cfg, estimate, self._default_subject) + create_body = _build_reservation_body(cfg, estimate, self._default_subject, args, kwargs) logger.debug("Creating reservation: body=%s", create_body) res_t1 = time.monotonic() @@ -416,7 +433,7 @@ async def execute( estimate = _evaluate_amount(cfg.estimate, args, kwargs) logger.debug("Estimated usage: estimate=%d", estimate) - create_body = _build_reservation_body(cfg, estimate, self._default_subject) + create_body = _build_reservation_body(cfg, estimate, self._default_subject, args, kwargs) res_response = await self._client.create_reservation(create_body) if not res_response.is_success: diff --git a/tests/test_decorator.py b/tests/test_decorator.py index 62daa65..5e48e75 100644 --- a/tests/test_decorator.py +++ b/tests/test_decorator.py @@ -82,6 +82,49 @@ def compute(x: int) -> str: assert result == "hello" client.close() + def test_callable_subject_and_action_fields(self, config: CyclesConfig, httpx_mock) -> None: # type: ignore[no-untyped-def] + """Per-call subject/action callables resolve against function args at reservation time.""" + import json + + httpx_mock.add_response( + method="POST", + url="http://localhost:7878/v1/reservations", + json={ + "decision": "ALLOW", "reservation_id": "res_dec_dyn", + "expires_at_ms": 9999999999, "affected_scopes": ["tenant:acme"], + }, + status_code=200, + ) + httpx_mock.add_response( + method="POST", + url="http://localhost:7878/v1/reservations/res_dec_dyn/commit", + json={"status": "COMMITTED", "charged": {"unit": "USD_MICROCENTS", "amount": 1000}}, + status_code=200, + ) + + client = CyclesClient(config) + + @cycles( + estimate=1000, + workspace=lambda req, workspace_id: workspace_id, + action_kind=lambda req, workspace_id: f"llm.{req['provider']}", + action_name=lambda req, workspace_id: req["model"], + client=client, + ) + def run_request(req: dict[str, str], workspace_id: str) -> str: + return "ok" + + result = run_request({"provider": "openai", "model": "gpt-4"}, workspace_id="ws-42") + assert result == "ok" + + sent = httpx_mock.get_request(method="POST", url="http://localhost:7878/v1/reservations") + assert sent is not None + body = json.loads(sent.content) + assert body["subject"]["workspace"] == "ws-42" + assert body["action"]["kind"] == "llm.openai" + assert body["action"]["name"] == "gpt-4" + client.close() + def test_denied_raises(self, config: CyclesConfig, httpx_mock) -> None: # type: ignore[no-untyped-def] httpx_mock.add_response( method="POST", diff --git a/tests/test_lifecycle.py b/tests/test_lifecycle.py index a839cf9..f70201f 100644 --- a/tests/test_lifecycle.py +++ b/tests/test_lifecycle.py @@ -37,7 +37,7 @@ class TestBuildReservationBody: def test_action_defaults_to_unknown(self) -> None: """When action_kind and action_name are not provided, they default to 'unknown'.""" cfg = DecoratorConfig(estimate=1000, tenant="acme") - body = _build_reservation_body(cfg, 1000, {}) + body = _build_reservation_body(cfg, 1000, {}, (), {}) assert body["action"]["kind"] == "unknown" assert body["action"]["name"] == "unknown" @@ -49,7 +49,7 @@ def test_action_fields_set(self) -> None: action_tags=["prod"], tenant="acme", ) - body = _build_reservation_body(cfg, 1000, {}) + body = _build_reservation_body(cfg, 1000, {}, (), {}) assert body["action"]["kind"] == "llm.completion" assert body["action"]["name"] == "gpt-4" assert body["action"]["tags"] == ["prod"] @@ -58,37 +58,37 @@ def test_validates_estimate_non_negative(self) -> None: """Spec: Amount.amount has minimum: 0, so 0 is valid but negative is not.""" cfg = DecoratorConfig(estimate=0, tenant="acme") # 0 should be valid per spec - body = _build_reservation_body(cfg, 0, {}) + body = _build_reservation_body(cfg, 0, {}, (), {}) assert body["estimate"]["amount"] == 0 cfg_neg = DecoratorConfig(estimate=-1, tenant="acme") with pytest.raises(ValueError, match="estimate"): - _build_reservation_body(cfg_neg, -1, {}) + _build_reservation_body(cfg_neg, -1, {}, (), {}) def test_validates_ttl_range(self) -> None: cfg = DecoratorConfig(estimate=1000, ttl_ms=500, tenant="acme") with pytest.raises(ValueError, match="ttl_ms"): - _build_reservation_body(cfg, 1000, {}) + _build_reservation_body(cfg, 1000, {}, (), {}) def test_validates_subject_has_standard_field(self) -> None: cfg = DecoratorConfig(estimate=1000, dimensions={"custom": "val"}) with pytest.raises(ValueError, match="at least one standard field"): - _build_reservation_body(cfg, 1000, {}) + _build_reservation_body(cfg, 1000, {}, (), {}) def test_validates_grace_period_ms_range(self) -> None: cfg = DecoratorConfig(estimate=1000, grace_period_ms=60001, tenant="acme") with pytest.raises(ValueError, match="grace_period_ms"): - _build_reservation_body(cfg, 1000, {}) + _build_reservation_body(cfg, 1000, {}, (), {}) def test_validates_grace_period_ms_negative(self) -> None: cfg = DecoratorConfig(estimate=1000, grace_period_ms=-1, tenant="acme") with pytest.raises(ValueError, match="grace_period_ms"): - _build_reservation_body(cfg, 1000, {}) + _build_reservation_body(cfg, 1000, {}, (), {}) def test_merges_default_subject_fields(self) -> None: cfg = DecoratorConfig(estimate=1000, workflow="task-1") defaults = {"tenant": "acme", "workspace": "prod"} - body = _build_reservation_body(cfg, 1000, defaults) + body = _build_reservation_body(cfg, 1000, defaults, (), {}) assert body["subject"]["tenant"] == "acme" assert body["subject"]["workspace"] == "prod" assert body["subject"]["workflow"] == "task-1" @@ -96,20 +96,128 @@ def test_merges_default_subject_fields(self) -> None: def test_decorator_subject_overrides_defaults(self) -> None: cfg = DecoratorConfig(estimate=1000, tenant="override") defaults = {"tenant": "default-tenant"} - body = _build_reservation_body(cfg, 1000, defaults) + body = _build_reservation_body(cfg, 1000, defaults, (), {}) assert body["subject"]["tenant"] == "override" def test_dry_run_flag(self) -> None: cfg = DecoratorConfig(estimate=1000, dry_run=True, tenant="acme") - body = _build_reservation_body(cfg, 1000, {}) + body = _build_reservation_body(cfg, 1000, {}, (), {}) assert body["dry_run"] is True def test_grace_period_included(self) -> None: cfg = DecoratorConfig(estimate=1000, grace_period_ms=10000, tenant="acme") - body = _build_reservation_body(cfg, 1000, {}) + body = _build_reservation_body(cfg, 1000, {}, (), {}) assert body["grace_period_ms"] == 10000 +class TestCallableSubjectFields: + def test_callable_resolves_against_args_kwargs(self) -> None: + cfg = DecoratorConfig( + estimate=1000, + workspace=lambda req, workspace_id: workspace_id, + tenant="acme", + ) + body = _build_reservation_body(cfg, 1000, {}, ("req-payload",), {"workspace_id": "ws-42"}) + assert body["subject"]["workspace"] == "ws-42" + + def test_callable_returning_none_falls_through_to_default(self) -> None: + cfg = DecoratorConfig( + estimate=1000, + workspace=lambda *_, **__: None, + tenant="acme", + ) + body = _build_reservation_body(cfg, 1000, {"workspace": "fallback"}, (), {}) + assert body["subject"]["workspace"] == "fallback" + + def test_constant_subject_field_regression(self) -> None: + cfg = DecoratorConfig(estimate=1000, workspace="prod", tenant="acme") + body = _build_reservation_body(cfg, 1000, {}, (), {}) + assert body["subject"]["workspace"] == "prod" + + def test_callable_exception_propagates(self) -> None: + def boom(*_: object, **__: object) -> str: + raise RuntimeError("boom") + + cfg = DecoratorConfig(estimate=1000, workspace=boom, tenant="acme") + with pytest.raises(RuntimeError, match="boom"): + _build_reservation_body(cfg, 1000, {}, (), {}) + + @pytest.mark.parametrize("field_name", ["tenant", "app", "workflow", "agent", "toolset"]) + def test_callable_resolves_for_all_subject_fields(self, field_name: str) -> None: + cfg = DecoratorConfig(estimate=1000, workspace="ws") + setattr(cfg, field_name, lambda *_, **__: f"resolved-{field_name}") + body = _build_reservation_body(cfg, 1000, {}, (), {}) + assert body["subject"][field_name] == f"resolved-{field_name}" + + +class TestCallableActionFields: + def test_callable_action_kind_resolves(self) -> None: + cfg = DecoratorConfig( + estimate=1000, + action_kind=lambda req: f"llm.{req['provider']}", + tenant="acme", + ) + body = _build_reservation_body(cfg, 1000, {}, ({"provider": "openai"},), {}) + assert body["action"]["kind"] == "llm.openai" + + def test_callable_action_name_resolves(self) -> None: + cfg = DecoratorConfig( + estimate=1000, + action_name=lambda req: req["model"], + tenant="acme", + ) + body = _build_reservation_body(cfg, 1000, {}, ({"model": "gpt-4"},), {}) + assert body["action"]["name"] == "gpt-4" + + def test_callable_action_kind_returning_none_falls_through_to_unknown(self) -> None: + cfg = DecoratorConfig( + estimate=1000, + action_kind=lambda *_, **__: None, + tenant="acme", + ) + body = _build_reservation_body(cfg, 1000, {}, (), {}) + assert body["action"]["kind"] == "unknown" + + def test_callable_action_tags_resolves(self) -> None: + cfg = DecoratorConfig( + estimate=1000, + action_tags=lambda req: [f"env:{req['env']}"], + tenant="acme", + ) + body = _build_reservation_body(cfg, 1000, {}, ({"env": "prod"},), {}) + assert body["action"]["tags"] == ["env:prod"] + + def test_constant_action_kind_regression(self) -> None: + cfg = DecoratorConfig(estimate=1000, action_kind="llm.completion", tenant="acme") + body = _build_reservation_body(cfg, 1000, {}, (), {}) + assert body["action"]["kind"] == "llm.completion" + + +class TestCallableDimensions: + def test_callable_dimensions_resolves(self) -> None: + cfg = DecoratorConfig( + estimate=1000, + dimensions=lambda req: {"region": req["region"]}, + tenant="acme", + ) + body = _build_reservation_body(cfg, 1000, {}, ({"region": "us-east-1"},), {}) + assert body["subject"]["dimensions"] == {"region": "us-east-1"} + + def test_callable_dimensions_returning_none_omits_key(self) -> None: + cfg = DecoratorConfig( + estimate=1000, + dimensions=lambda *_, **__: None, + tenant="acme", + ) + body = _build_reservation_body(cfg, 1000, {}, (), {}) + assert "dimensions" not in body["subject"] + + def test_constant_dimensions_regression(self) -> None: + cfg = DecoratorConfig(estimate=1000, dimensions={"k": "v"}, tenant="acme") + body = _build_reservation_body(cfg, 1000, {}, (), {}) + assert body["subject"]["dimensions"] == {"k": "v"} + + class TestBuildCommitBody: def test_basic(self) -> None: body = _build_commit_body(500, "USD_MICROCENTS", None, None) diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 04625de..88d7578 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -671,6 +671,57 @@ def test_heartbeat_starts_and_stops(self) -> None: mock.extend_reservation.assert_called() + def test_heartbeat_skipped_when_ttl_zero(self) -> None: + mock = _make_mock_client() + sr = StreamReservation( + mock, + subject=_default_subject(), + action=_default_action(), + estimate=_default_estimate(), + ttl_ms=0, + ) + assert sr._start_heartbeat() is None + + def test_heartbeat_extend_failure_logged(self) -> None: + mock = _make_mock_client() + mock.create_reservation.return_value = _allow_response() + mock.commit_reservation.return_value = _commit_success() + mock.extend_reservation.return_value = CyclesResponse.http_error( + 500, "Server error", body={"error": "INTERNAL"}, + ) + + sr = StreamReservation( + mock, + subject=_default_subject(), + action=_default_action(), + estimate=_default_estimate(), + ttl_ms=2000, + ) + + with sr: + time.sleep(1.2) + + mock.extend_reservation.assert_called() + + def test_heartbeat_extend_exception_logged(self) -> None: + mock = _make_mock_client() + mock.create_reservation.return_value = _allow_response() + mock.commit_reservation.return_value = _commit_success() + mock.extend_reservation.side_effect = RuntimeError("boom") + + sr = StreamReservation( + mock, + subject=_default_subject(), + action=_default_action(), + estimate=_default_estimate(), + ttl_ms=2000, + ) + + with sr: + time.sleep(1.2) + + mock.extend_reservation.assert_called() + # --------------------------------------------------------------------------- # AsyncStreamReservation tests @@ -1281,6 +1332,77 @@ async def test_caps_propagated(self) -> None: assert reservation.caps is not None assert reservation.caps.max_tokens == 256 + @pytest.mark.asyncio + async def test_decision_property(self) -> None: + mock = _make_async_mock_client() + mock.create_reservation.return_value = _allow_response() + mock.commit_reservation.return_value = _commit_success() + + asr = AsyncStreamReservation( + mock, + subject=_default_subject(), + action=_default_action(), + estimate=_default_estimate(), + ttl_ms=1000, + ) + + async with asr as reservation: + assert reservation.decision is not None + + @pytest.mark.asyncio + async def test_heartbeat_skipped_when_ttl_zero(self) -> None: + mock = _make_async_mock_client() + asr = AsyncStreamReservation( + mock, + subject=_default_subject(), + action=_default_action(), + estimate=_default_estimate(), + ttl_ms=0, + ) + assert asr._start_heartbeat() is None + + @pytest.mark.asyncio + async def test_heartbeat_extend_failure_logged(self) -> None: + mock = _make_async_mock_client() + mock.create_reservation.return_value = _allow_response() + mock.commit_reservation.return_value = _commit_success() + mock.extend_reservation.return_value = CyclesResponse.http_error( + 500, "Server error", body={"error": "INTERNAL"}, + ) + + asr = AsyncStreamReservation( + mock, + subject=_default_subject(), + action=_default_action(), + estimate=_default_estimate(), + ttl_ms=2000, + ) + + async with asr: + await asyncio.sleep(1.2) + + mock.extend_reservation.assert_called() + + @pytest.mark.asyncio + async def test_heartbeat_extend_exception_logged(self) -> None: + mock = _make_async_mock_client() + mock.create_reservation.return_value = _allow_response() + mock.commit_reservation.return_value = _commit_success() + mock.extend_reservation.side_effect = RuntimeError("boom") + + asr = AsyncStreamReservation( + mock, + subject=_default_subject(), + action=_default_action(), + estimate=_default_estimate(), + ttl_ms=2000, + ) + + async with asr: + await asyncio.sleep(1.2) + + mock.extend_reservation.assert_called() + class TestBudgetExceeded: def test_budget_exceeded_raises_typed_error(self) -> None: