diff --git a/packages/claude-agents/tests/test_skills.py b/packages/claude-agents/tests/test_skills.py index 5fdb836d..b2a8a2bd 100644 --- a/packages/claude-agents/tests/test_skills.py +++ b/packages/claude-agents/tests/test_skills.py @@ -672,6 +672,7 @@ async def during(kwargs: dict[str, Any]) -> None: async def ld_client() -> Any: await shutdown() client = MagicMock() + del client.variation_detail client.track = MagicMock() client.flush = MagicMock() client.close = MagicMock() diff --git a/packages/client/README.md b/packages/client/README.md index 040c7b63..46dfafad 100644 --- a/packages/client/README.md +++ b/packages/client/README.md @@ -305,6 +305,29 @@ result = await evals.run( **The list is checked before any network I/O.** A blank key, an uppercase key, a `schema` that is not a JSON object, a schema that is not JSON-serializable (including a `NaN` or `Infinity` value), and a non-callable implementation each fail with zero requests issued. A repeated key fails too, and keys are compared without case. A `NativeTool` is valid only for a library tool, because the provider supplies its schema. +### Shadow requests + +A shadow request runs another variation of an AI Config in the background on the same input, and reports it on the normal AI metrics with `isShadow: true`. You set shadow requests up on a targeting rule or the fallthrough in LaunchDarkly. The shadow's response is discarded, and your response, latency and events are unchanged. Shadow runs use your provider account, so they are billed like real calls. + +`config().invoke()` and `config().stream()` run them for you in the background, with a 30 second timeout. Shadows don't run for graphs or judges, and a shadow makes no tool calls and runs no judges. + +```python +caller = config(key="my-ai-config", handler=handlers, shadow_timeout_s=10) +``` + +To run shadows on your own queue or worker instead, pass `skip_shadows=True`. Each response then carries `shadow_tasks`, which you can pickle and run later with `run_shadow`, the same way `skip_judges` and `run_judge` work: + +```python +from launchdarkly_ai_server import config, run_shadow + +result = await config(key="my-ai-config", handler=handlers, skip_shadows=True).invoke(...) +for task in result.shadow_tasks or []: + queue.put_nowait(task) + +# in the worker +outcome = await run_shadow(task, handlers) +``` + ### `config(**args)` The primary entry point for AI config invocations. Accepts either a single handler or a list of handlers and routes to the correct one at invoke-time based on the flag variation's provider and mode. diff --git a/packages/client/api-report/launchdarkly_ai_server.txt b/packages/client/api-report/launchdarkly_ai_server.txt index 4ac0e30a..0dd5940a 100644 --- a/packages/client/api-report/launchdarkly_ai_server.txt +++ b/packages/client/api-report/launchdarkly_ai_server.txt @@ -10,10 +10,10 @@ any_multimodal(turns: list[dict[str, Any]]) -> bool build_judge_tasks(*, config: AiConfigRep, user_context: LDContext, handler: ProviderHandler, handlers: list[ProviderHandler] | None = None, llm_response: str, base_track_data: TrackData, user_input: str | None = None, trajectory: str = '') -> list[JudgeTask] [async] compose(a: Registry, b: Registry) -> Registry compose_history(*, history: list[dict[str, Any]], user_input: str | None = None, config_messages: list[dict[str, Any]] | None = None) -> list[dict[str, Any]] -config(*, key: str, handler: ProviderHandler | list[ProviderHandler] | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, registry: Any = None, skip_judges: bool = False) -> ConfigInstance +config(*, key: str, handler: ProviderHandler | list[ProviderHandler] | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, registry: Any = None, skip_judges: bool = False, skip_shadows: bool = False, shadow_timeout_s: float = DEFAULT_SHADOW_TIMEOUT_S) -> ConfigInstance config_skills(key: str, context: LDContext) -> SkillSource ConfigInstance [class] - __init__(key: str, handler: ProviderHandler | list[ProviderHandler] | None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None, registry: Any, skip_judges: bool = False) + __init__(key: str, handler: ProviderHandler | list[ProviderHandler] | None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None, registry: Any, skip_judges: bool = False, skip_shadows: bool = False, shadow_timeout_s: float = DEFAULT_SHADOW_TIMEOUT_S) invoke(user_input: str | None, context: LDContext, variables: dict[str, Any] | None = None, history: list[dict[str, Any]] | None = None) -> ProviderResponse[Any] [async] stream(user_input: str | None, context: LDContext, variables: dict[str, Any] | None = None, history: list[dict[str, Any]] | None = None) -> AsyncGenerator[StreamEvent, None] configure(*, options: SDKOptions) -> None @@ -61,11 +61,11 @@ EvaluationsModule [class] sdk_key: str | None [property] tools: ToolsClient [property] ui_base_uri: str [property] -execute_and_stream(*, config_key: str, config: AiConfigRep, meta: VariationMeta, user_context: LDContext, handler: ProviderHandler, user_input: str | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None) -> AsyncGenerator[ExecuteStreamEvent, None] [async] -execute_and_track(*, config_key: str, config: AiConfigRep, meta: VariationMeta, user_context: LDContext, handler: ProviderHandler, user_input: str | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None) -> dict[str, Any] [async] +execute_and_stream(*, config_key: str, config: AiConfigRep, meta: VariationMeta, user_context: LDContext, handler: ProviderHandler, user_input: str | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None, run_id: str | None = None, extra_track_data: dict[str, Any] | None = None, span_attributes: dict[str, Any] | None = None) -> AsyncGenerator[ExecuteStreamEvent, None] [async] +execute_and_track(*, config_key: str, config: AiConfigRep, meta: VariationMeta, user_context: LDContext, handler: ProviderHandler, user_input: str | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None, run_id: str | None = None, extra_track_data: dict[str, Any] | None = None, span_attributes: dict[str, Any] | None = None) -> dict[str, Any] [async] ExecuteStreamDoneEvent = dict[str, Any] ExecuteStreamEvent = dict[str, Any] -extract_variation(config_key: str, context: dict[str, Any]) -> dict[str, Any] [async] +extract_variation(config_key: str, context: dict[str, Any], *, with_reason: bool = False) -> dict[str, Any] [async] FDv2SkillStore = LaunchDarklySkillStore GenerationConfig(TypedDict, total=False) [class] instructions: str @@ -255,10 +255,11 @@ ProviderHandler [class] provides_for: tuple[str, Literal['agent', 'messages']] | None = provides_for stream(config: AiConfigRep, user_input: str | None = None, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, variables: dict[str, Any] | None = None, history: list[dict[str, Any]] | None = None) -> AsyncGenerator[dict[str, Any], None] [async] ProviderResponse(Generic[T]) [class, @dataclass] - __init__(response: T, usage: UsageDict, judge_results: dict[str, JudgeResult] | None = None, judge_tasks: list[JudgeTask] | None = None, track_data: TrackData | None = None) + __init__(response: T, usage: UsageDict, judge_results: dict[str, JudgeResult] | None = None, judge_tasks: list[JudgeTask] | None = None, shadow_tasks: list[ShadowTask] | None = None, track_data: TrackData | None = None) judge_results: dict[str, JudgeResult] | None = None judge_tasks: list[JudgeTask] | None = None response: T + shadow_tasks: list[ShadowTask] | None = None track_data: TrackData | None = None usage: UsageDict ReconcileAction [class, @dataclass(frozen=True)] @@ -286,6 +287,7 @@ resolve_handlers(registry: Registry | None, local_handlers: list[ProviderHandler resolve_tools(registry: Registry | None, local_tools: Mapping[str, Callable[..., Any] | NativeTool] | None) -> Mapping[str, Callable[..., Any] | NativeTool] | None run_judge(task: JudgeTask, handlers: list[ProviderHandler]) -> JudgeRunResult | None [async] run_judges(*, config: AiConfigRep, user_context: LDContext, handler: ProviderHandler, handlers: list[ProviderHandler] | None = None, user_input: str | None, llm_response: str, base_track_data: TrackData, tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, graph_key: str | None = None, trajectory: str = '') -> dict[str, JudgeResult] [async] +run_shadow(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult [async] RunSummary [class, @dataclass] __init__(total_rows: int = 0, passed_rows: int = 0, failed_rows: int = 0, error_rows: int = 0, pending_rows: int = 0) error_rows: int = 0 @@ -324,6 +326,24 @@ set_skill_store(store: SkillStore | None) -> None set_tool_call_content_attributes(span: Any, capture: bool, *, arguments: Any = None, result: Any = None) -> None set_tool_definition_attributes(span: Any, capture: bool, tools: list[ToolDefinitionInput]) -> None set_usage_span_attributes(span: Any, usage: SpanUsage) -> None +ShadowResult [class, @dataclass(frozen=True)] + __init__(usage: dict[str, int], track_data: TrackData, error: str | None = None) + error: str | None = None + track_data: TrackData + usage: dict[str, int] +ShadowTask [class, @dataclass(frozen=True)] + __init__(config_key: str, shadow_request_id: str, shadow_config: AiConfigRep, shadow_meta: VariationMeta, user_input: str | None, history: list[dict[str, Any]] | None, variables: dict[str, Any] | None, user_context: LDContext, timeout_s: float, track_data: TrackData, span_attributes: dict[str, Any] = dict()) + config_key: str + history: list[dict[str, Any]] | None + shadow_config: AiConfigRep + shadow_meta: VariationMeta + shadow_request_id: str + span_attributes: dict[str, Any] = field(default_factory=dict) + timeout_s: float + track_data: TrackData + user_context: LDContext + user_input: str | None + variables: dict[str, Any] | None shutdown() -> None [async] Skill [class, @dataclass(frozen=True)] __init__(key: str, version: int, content: bytes, content_hash: str, name: str | None = None, description: str | None = None) diff --git a/packages/client/src/launchdarkly_ai_server/__init__.py b/packages/client/src/launchdarkly_ai_server/__init__.py index 48355f8c..84c70c96 100644 --- a/packages/client/src/launchdarkly_ai_server/__init__.py +++ b/packages/client/src/launchdarkly_ai_server/__init__.py @@ -71,6 +71,7 @@ ) from .sdk_info import SDK_INFO_CONTEXT, SDK_INFO_EVENT, register_ai_sdk_package from .sdk_usage import report_usage +from .shadows import run_shadow from .skills import ( InMemorySkillStore, all_skills, @@ -128,6 +129,8 @@ ReconcileActionKind, ReconcileReport, SDKOptions, + ShadowResult, + ShadowTask, Skill, SkillOutcome, SkillOutcomeReason, @@ -200,6 +203,8 @@ "ProviderGraphResponse", "ProviderHandler", "ProviderResponse", + "ShadowResult", + "ShadowTask", "StreamChunkEvent", "StreamDoneEvent", "StreamEvent", @@ -275,6 +280,7 @@ # judges "build_judge_tasks", "run_judge", + "run_shadow", "run_judges", # client "config", diff --git a/packages/client/src/launchdarkly_ai_server/client.py b/packages/client/src/launchdarkly_ai_server/client.py index 48d1109c..618a5f17 100644 --- a/packages/client/src/launchdarkly_ai_server/client.py +++ b/packages/client/src/launchdarkly_ai_server/client.py @@ -1,5 +1,7 @@ from __future__ import annotations +import math +import uuid from collections.abc import AsyncGenerator, Callable, Mapping from contextlib import aclosing from typing import Any @@ -11,6 +13,7 @@ from .lifecycle import extract_variation from .registry import resolve_handlers, resolve_tools from .sdk_usage import report_usage +from .shadows import DEFAULT_SHADOW_TIMEOUT_S, build_shadow_tasks, start_shadows from .tracking import execute_and_stream, execute_and_track from .types import ( AiConfigRep, @@ -18,6 +21,7 @@ NativeTool, ProviderHandler, ProviderResponse, + ShadowTask, StreamEvent, VariationMeta, ) @@ -45,6 +49,34 @@ def _resolve_output_format_response( return parsed if parsed is not None else raw_response +def _shadow_tasks( + *, + config_key: str, + variation: dict[str, Any], + meta: VariationMeta, + context: LDContext, + user_input: str | None, + history: list[dict[str, Any]] | None, + variables: dict[str, Any] | None, + handlers: list[ProviderHandler], + run_id: str, + timeout_s: float, +) -> list[ShadowTask]: + return build_shadow_tasks( + config_key=config_key, + raw_variation=variation.get("raw"), + reason=variation.get("reason"), + primary_variation_key=str(meta.get("variationKey", "")), + context=context, + user_input=user_input, + history=history, + variables=variables, + handlers=handlers, + run_id=run_id, + timeout_s=timeout_s, + ) + + class ConfigInstance: """Return type of ``config()``.""" @@ -55,12 +87,18 @@ def __init__( tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None, registry: Any, # Registry | None skip_judges: bool = False, + skip_shadows: bool = False, + shadow_timeout_s: float = DEFAULT_SHADOW_TIMEOUT_S, ) -> None: + if not (math.isfinite(shadow_timeout_s) and shadow_timeout_s > 0): + raise ValueError("shadow_timeout_s must be a finite number greater than 0") self._key = key self._handler = handler self._tool_handlers = tool_handlers self._registry = registry self._skip_judges = skip_judges + self._skip_shadows = skip_shadows + self._shadow_timeout_s = shadow_timeout_s def _normalize_handlers(self) -> list[ProviderHandler] | None: if self._handler is None: @@ -91,13 +129,28 @@ async def _invoke( ) resolved_tools = resolve_tools(self._registry, self._tool_handlers) - variation = await extract_variation(self._key, context) + variation = await extract_variation(self._key, context, with_reason=True) config: AiConfigRep = variation["config"] meta: VariationMeta = variation["meta"] if not resolved_handler_list: raise ValueError("No handlers provided to config()") handler = select_handler(config, meta, resolved_handler_list) + run_id = str(uuid.uuid4()) + shadow_tasks = _shadow_tasks( + config_key=self._key, + variation=variation, + meta=meta, + context=context, + user_input=user_input, + history=history, + variables=variables, + handlers=resolved_handler_list, + run_id=run_id, + timeout_s=self._shadow_timeout_s, + ) + if not self._skip_shadows: + start_shadows(shadow_tasks, resolved_handler_list) result = await execute_and_track( config_key=self._key, @@ -109,6 +162,7 @@ async def _invoke( tool_handlers=resolved_tools, variables=variables, history=history, + run_id=run_id, ) raw_response = result["response"] @@ -140,6 +194,7 @@ async def _invoke( response=parsed_response, usage=usage_obj, judge_tasks=judge_tasks, + shadow_tasks=shadow_tasks if self._skip_shadows else None, track_data=track_data, ) @@ -158,6 +213,7 @@ async def _invoke( response=parsed_response, usage=usage_obj, judge_results=judge_results if judge_results else None, + shadow_tasks=shadow_tasks if self._skip_shadows else None, track_data=track_data, ) @@ -210,7 +266,7 @@ async def _execute_stream( ) resolved_tools = resolve_tools(self._registry, self._tool_handlers) - variation = await extract_variation(self._key, context) + variation = await extract_variation(self._key, context, with_reason=True) config: AiConfigRep = variation["config"] meta: VariationMeta = variation["meta"] @@ -228,6 +284,22 @@ async def _execute_stream( sink["tools"] = resolved_tools sink["provider"] = provider sink["model"] = model + run_id = str(uuid.uuid4()) + shadow_tasks = _shadow_tasks( + config_key=self._key, + variation=variation, + meta=meta, + context=context, + user_input=user_input, + history=history, + variables=variables, + handlers=resolved_handler_list, + run_id=run_id, + timeout_s=self._shadow_timeout_s, + ) + if not self._skip_shadows: + start_shadows(shadow_tasks, resolved_handler_list) + sink["shadow_tasks"] = shadow_tasks if self._skip_shadows else None async with aclosing( execute_and_stream( @@ -240,6 +312,7 @@ async def _execute_stream( tool_handlers=resolved_tools, variables=variables, history=history, + run_id=run_id, ) ) as source: async for event in source: @@ -266,12 +339,15 @@ async def _stream_events( judge_results = await judge_completed_stream( self, sink, done_event, user_input, context ) - yield { + done: StreamEvent = { "type": "done", "response": done_event.get("response", ""), "usage": done_event.get("usage"), "judge_results": judge_results if judge_results else None, } + if sink.get("shadow_tasks") is not None: + done["shadow_tasks"] = sink["shadow_tasks"] + yield done class _InternalConfigInstance(ConfigInstance): @@ -307,6 +383,8 @@ def config( tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, registry: Any = None, skip_judges: bool = False, + skip_shadows: bool = False, + shadow_timeout_s: float = DEFAULT_SHADOW_TIMEOUT_S, ) -> ConfigInstance: """ Creates a ``ConfigInstance`` bound to *key*. Accepts a single handler or a @@ -319,6 +397,11 @@ def config( ``.invoke()`` / ``.stream()``. When set, ``invoke()`` returns ``judge_tasks: list[JudgeTask]`` — pre-packaged tasks ready for a background thread calling ``run_judge(task, handlers)``. + + Shadow requests attached to the matched rule or fallthrough run in the + background with *shadow_timeout_s*. Pass ``skip_shadows=True`` to get them + back as ``shadow_tasks`` instead, and run each with + ``run_shadow(task, handlers)``. """ return ConfigInstance( key=key, @@ -326,6 +409,8 @@ def config( tool_handlers=tool_handlers, registry=registry, skip_judges=skip_judges, + skip_shadows=skip_shadows, + shadow_timeout_s=shadow_timeout_s, ) @@ -336,6 +421,8 @@ def _config( tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None = None, registry: Any = None, skip_judges: bool = False, + skip_shadows: bool = False, + shadow_timeout_s: float = DEFAULT_SHADOW_TIMEOUT_S, ) -> ConfigInstance: """Non-reporting :func:`config` for SDK wrappers.""" return _InternalConfigInstance( @@ -344,4 +431,6 @@ def _config( tool_handlers=tool_handlers, registry=registry, skip_judges=skip_judges, + skip_shadows=skip_shadows, + shadow_timeout_s=shadow_timeout_s, ) diff --git a/packages/client/src/launchdarkly_ai_server/lifecycle.py b/packages/client/src/launchdarkly_ai_server/lifecycle.py index 335d9917..0501424a 100644 --- a/packages/client/src/launchdarkly_ai_server/lifecycle.py +++ b/packages/client/src/launchdarkly_ai_server/lifecycle.py @@ -586,6 +586,8 @@ async def inspect_config( async def extract_variation( config_key: str, context: dict[str, Any], + *, + with_reason: bool = False, ) -> dict[str, Any]: """ Fetches an AI config variation from the LD client and parses it. @@ -603,13 +605,26 @@ async def extract_variation( # The real ldclient.variation() expects an ldclient.Context, not a plain dict. # Convert when running against the real SDK; test mocks accept dicts directly. ld_context = to_ld_context(client, context) - # variation() is synchronous in the real Python LD SDK; AsyncMock in tests. - variation_result = client.variation(config_key, ld_context, None) - raw = ( - await variation_result - if inspect.isawaitable(variation_result) - else variation_result + # The evaluation reason picks the shadow requests for the matched rule. + reason: Any = None + variation_detail = ( + getattr(client, "variation_detail", None) if with_reason else None ) + if variation_detail is not None: + detail_result = variation_detail(config_key, ld_context, None) + detail = ( + await detail_result if inspect.isawaitable(detail_result) else detail_result + ) + raw = getattr(detail, "value", None) + reason = getattr(detail, "reason", None) + else: + # variation() is synchronous in the real Python LD SDK; AsyncMock in tests. + variation_result = client.variation(config_key, ld_context, None) + raw = ( + await variation_result + if inspect.isawaitable(variation_result) + else variation_result + ) if raw is None: raise RuntimeError( @@ -631,9 +646,9 @@ async def extract_variation( **model_stamps_from_meta(ld_meta), } - # Strip _ldMeta for config parsing + # Strip _ldMeta and shadowConfiguration for config parsing config_raw = ( - {k: v for k, v in raw.items() if k != "_ldMeta"} + {k: v for k, v in raw.items() if k not in ("_ldMeta", "shadowConfiguration")} if isinstance(raw, dict) else raw ) @@ -644,4 +659,4 @@ async def extract_variation( f"Invalid AI config variation for '{config_key}': {result.error['message']}" ) - return {"config": result.data, "meta": meta} + return {"config": result.data, "meta": meta, "raw": raw, "reason": reason} diff --git a/packages/client/src/launchdarkly_ai_server/shadows.py b/packages/client/src/launchdarkly_ai_server/shadows.py new file mode 100644 index 00000000..00abfaa9 --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/shadows.py @@ -0,0 +1,283 @@ +""" +Shadow requests: run another variation of the AI Config in the background on +the same input, and track it on the normal ``$ld:ai:*`` events with +``isShadow``. The shadow's response is discarded, and the caller's response, +latency and events are unchanged. + +LaunchDarkly attaches shadow requests to a targeting rule or the fallthrough +and sends them in the served variation as ``shadowConfiguration``. +``config().invoke()`` and ``.stream()`` call :func:`build_shadow_tasks`, then +either :func:`start_shadows` or hand the tasks back when ``skip_shadows=True``. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import time +from typing import Any + +from .sdk_usage import report_usage +from .shutdown_hooks import register_shutdown_hook +from .tracking import build_track_data, execute_and_track +from .types import ( + AiConfigRep, + LDContext, + ProviderHandler, + ShadowResult, + ShadowTask, + VariationMeta, +) +from .types_validation import parse_ai_config +from .utils import ( + _context_identity_from_ld_context, + model_stamps_from_meta, + parse_usage, + select_handler, +) + +logger = logging.getLogger("launchdarkly_ai_server.shadows") + +DEFAULT_SHADOW_TIMEOUT_S = 30.0 +_FALLTHROUGH = "fallthrough" +_STRIPPED_KEYS = ("_ldMeta", "tools", "judgeConfiguration", "shadowConfiguration") +_BUCKET_SCALE = float(0xFFFFFFFFFFFFFFF) + +_running: set[asyncio.Task[ShadowResult]] = set() + + +def build_shadow_tasks( + *, + config_key: str, + raw_variation: Any, + reason: Any, + primary_variation_key: str, + context: LDContext, + user_input: str | None, + history: list[dict[str, Any]] | None, + variables: dict[str, Any] | None, + handlers: list[ProviderHandler], + run_id: str, + timeout_s: float, +) -> list[ShadowTask]: + """ + Returns a task for each shadow request on the rule or fallthrough that + served this call, skipping entries that have ended, are not sampled for this + context, or have no matching handler. Never raises. + """ + try: + entries, variations = _entries_for(raw_variation, reason) + except Exception: + logger.debug("Ignoring a malformed shadowConfiguration", exc_info=True) + return [] + + tasks: list[ShadowTask] = [] + now_ms = time.time() * 1000 + for entry in entries: + if not _valid_entry(entry) or entry["endsAt"] <= now_ms: + continue + if not _in_sample(entry["id"], context, float(entry["samplingRate"])): + continue + parsed = _parse_variation(variations.get(entry["variationKey"])) + if parsed is None: + continue + shadow_config, shadow_meta = parsed + try: + select_handler(shadow_config, shadow_meta, handlers) + except Exception: + logger.debug("No handler for shadow request %s", entry["id"]) + continue + tasks.append( + ShadowTask( + config_key=config_key, + shadow_request_id=entry["id"], + shadow_config=shadow_config, + shadow_meta=shadow_meta, + user_input=user_input, + history=list(history) if history else None, + variables=dict(variables) if variables else None, + user_context=dict(context), + timeout_s=timeout_s, + track_data={ + "runId": run_id, + "isShadow": True, + "primaryVariationKey": primary_variation_key, + "shadowRequestId": entry["id"], + "ownerKind": entry.get("ownerKind", ""), + "ownerId": entry.get("ownerId", ""), + }, + span_attributes={ + "launchdarkly.is_shadow": True, + "launchdarkly.shadow.request.id": entry["id"], + "launchdarkly.shadow.primary.variation.key": primary_variation_key, + }, + ) + ) + return tasks + + +def start_shadows(tasks: list[ShadowTask], handlers: list[ProviderHandler]) -> None: + """Runs each task in the background. The caller does not wait for them.""" + if not tasks: + return + loop = asyncio.get_running_loop() + for task in tasks: + running = loop.create_task(_run(task, handlers)) + _running.add(running) + running.add_done_callback(_running.discard) + + +async def run_shadow(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult: + """Runs one :class:`ShadowTask`. Never raises; a failure is in ``error``.""" + report_usage("client.runShadow") + return await _run(task, handlers) + + +def _cancel_running() -> None: + for running in list(_running): + running.cancel() + _running.clear() + + +register_shutdown_hook("Shadow requests", _cancel_running) + + +async def _wait_for_shadows_for_testing() -> None: + while _running: + await asyncio.gather(*list(_running), return_exceptions=True) + + +def _entries_for(raw_variation: Any, reason: Any) -> tuple[list[Any], dict[str, Any]]: + configuration = ( + raw_variation.get("shadowConfiguration") + if isinstance(raw_variation, dict) + else None + ) + slot = _slot(reason) + if slot is None or not isinstance(configuration, dict): + return [], {} + shadows = configuration.get("shadows") + variations = configuration.get("variations") + if not isinstance(shadows, dict) or not isinstance(variations, dict): + return [], {} + entries = shadows.get(slot) + return (entries if isinstance(entries, list) else []), variations + + +def _slot(reason: Any) -> str | None: + if not isinstance(reason, dict): + return None + kind = reason.get("kind") + if kind == "FALLTHROUGH": + return _FALLTHROUGH + if kind == "RULE_MATCH" and isinstance(reason.get("ruleId"), str): + return str(reason["ruleId"]) + return None + + +def _in_sample(shadow_id: str, context: LDContext, sampling_rate: float) -> bool: + identity = _context_identity_from_ld_context(context) + if identity is None: + return False + canonical, _ = identity + digest = hashlib.sha1(f"{shadow_id}.{canonical}".encode(), usedforsecurity=False) + bucket = int(digest.hexdigest()[:15], 16) / _BUCKET_SCALE + return bucket < sampling_rate + + +def _valid_entry(entry: Any) -> bool: + if not isinstance(entry, dict): + return False + numbers_ok = all( + isinstance(entry.get(k), (int, float)) and not isinstance(entry.get(k), bool) + for k in ("samplingRate", "endsAt") + ) + strings_ok = all( + isinstance(entry.get(k), str) and entry.get(k) for k in ("id", "variationKey") + ) + return numbers_ok and strings_ok + + +def _parse_variation(raw: Any) -> tuple[AiConfigRep, VariationMeta] | None: + if not isinstance(raw, dict): + return None + ld_meta = raw.get("_ldMeta") + if not isinstance(ld_meta, dict) or not ld_meta.get("enabled"): + return None + parsed = parse_ai_config({k: v for k, v in raw.items() if k not in _STRIPPED_KEYS}) + if not parsed.success: + return None + meta: VariationMeta = { + "enabled": True, + "variationKey": ld_meta.get("variationKey", ""), + "version": ld_meta.get("version", 1), + "mode": ld_meta.get("mode"), + **model_stamps_from_meta(ld_meta), + } + return parsed.data, meta + + +async def _run(task: ShadowTask, handlers: list[ProviderHandler]) -> ShadowResult: + extra = {k: v for k, v in task.track_data.items() if k != "runId"} + run_id = task.track_data.get("runId") + track_data = build_track_data( + config_key=task.config_key, + config=task.shadow_config, + meta=task.shadow_meta, + run_id=run_id, + extra_track_data=extra, + ) + try: + handler = select_handler(task.shadow_config, task.shadow_meta, handlers) + except Exception as err: + return ShadowResult( + usage=parse_usage({}), track_data=track_data, error=str(err) + ) + + try: + result = await execute_and_track( + config_key=task.config_key, + config=task.shadow_config, + meta=task.shadow_meta, + user_context=task.user_context, + handler=_with_deadline(handler, task.timeout_s), + user_input=task.user_input, + tool_handlers=None, + variables=task.variables, + history=list(task.history) if task.history else None, + run_id=run_id, + extra_track_data=extra, + span_attributes=task.span_attributes, + ) + except _ShadowTimeout: + return ShadowResult( + usage=parse_usage({}), track_data=track_data, error="timeout" + ) + except Exception as err: + return ShadowResult( + usage=parse_usage({}), track_data=track_data, error=str(err) + ) + return ShadowResult(usage=result["usage"], track_data=result["track_data"]) + + +class _ShadowTimeout(Exception): + pass + + +def _with_deadline(handler: ProviderHandler, timeout_s: float) -> ProviderHandler: + async def timed(*args: Any) -> Any: + deadline = asyncio.timeout(timeout_s) + try: + async with deadline: + return await handler(*args) + except TimeoutError: + if deadline.expired(): + raise _ShadowTimeout("timeout") from None + raise + + return ProviderHandler( + fn=timed, + provides_for=handler.provides_for, + capture_content=handler.capture_content, + ) diff --git a/packages/client/src/launchdarkly_ai_server/tracking.py b/packages/client/src/launchdarkly_ai_server/tracking.py index e93f6918..3c0a740f 100644 --- a/packages/client/src/launchdarkly_ai_server/tracking.py +++ b/packages/client/src/launchdarkly_ai_server/tracking.py @@ -52,6 +52,59 @@ def _try_get_environment_id() -> str | None: return os.environ.get("LD_ENVIRONMENT_ID") or None +SPAN_ATTRIBUTES_KEY = "__ld_span" +"""Variables key for extra root-span attributes. Set only by the SDK, never from template variables.""" + + +def build_track_data( + *, + config_key: str, + config: AiConfigRep, + meta: VariationMeta, + graph_key: str | None = None, + run_id: str | None = None, + extra_track_data: dict[str, Any] | None = None, +) -> TrackData: + """The tracking payload for one invocation. *extra_track_data* is added last.""" + environment_id = _try_get_environment_id() + track_data: TrackData = { + "runId": run_id or str(uuid.uuid4()), + "configKey": config_key, + "variationKey": meta.get("variationKey", "") if isinstance(meta, dict) else "", + "version": meta.get("version", 1) if isinstance(meta, dict) else 1, + "modelName": config.get("model", {}).get("name", "") + if isinstance(config, dict) + else "", + "providerName": config.get("provider", {}).get("name", "") + if isinstance(config, dict) + else "", + **model_stamps_from_meta(meta), + } + if graph_key: + track_data["graphKey"] = graph_key + if environment_id: + track_data["environmentId"] = environment_id + if extra_track_data: + track_data.update(extra_track_data) + return track_data + + +def _merge_variables( + variables: dict[str, Any] | None, + user_context: LDContext, + track_data: TrackData, + span_attributes: dict[str, Any] | None, +) -> dict[str, Any]: + merged: dict[str, Any] = { + **{k: v for k, v in (variables or {}).items() if k != SPAN_ATTRIBUTES_KEY}, + "ldContext": {**user_context}, + "__ld": track_data, + } + if span_attributes: + merged[SPAN_ATTRIBUTES_KEY] = dict(span_attributes) + return merged + + def wrap_tool_handlers( tool_handlers: Mapping[str, Callable[..., Any] | NativeTool] | None, user_context: LDContext, @@ -167,6 +220,9 @@ async def execute_and_track( variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None, + run_id: str | None = None, + extra_track_data: dict[str, Any] | None = None, + span_attributes: dict[str, Any] | None = None, ) -> dict[str, Any]: """ Calls *handler*, emits LD telemetry events, and returns @@ -174,24 +230,14 @@ async def execute_and_track( """ from .lifecycle import get_client - environment_id = _try_get_environment_id() - track_data: TrackData = { - "runId": str(uuid.uuid4()), - "configKey": config_key, - "variationKey": meta.get("variationKey", "") if isinstance(meta, dict) else "", - "version": meta.get("version", 1) if isinstance(meta, dict) else 1, - "modelName": config.get("model", {}).get("name", "") - if isinstance(config, dict) - else "", - "providerName": config.get("provider", {}).get("name", "") - if isinstance(config, dict) - else "", - **model_stamps_from_meta(meta), - } - if graph_key: - track_data["graphKey"] = graph_key - if environment_id: - track_data["environmentId"] = environment_id + track_data = build_track_data( + config_key=config_key, + config=config, + meta=meta, + graph_key=graph_key, + run_id=run_id, + extra_track_data=extra_track_data, + ) client = get_client() ld_ctx = to_ld_context(client, user_context) @@ -207,11 +253,9 @@ async def execute_and_track( ld_ctx, track_data, ) - merged_variables: dict[str, Any] = { - **(variables or {}), - "ldContext": {**user_context}, - "__ld": track_data, - } + merged_variables = _merge_variables( + variables, user_context, track_data, span_attributes + ) start_time = time.monotonic() @@ -266,6 +310,9 @@ async def execute_and_stream( variables: dict[str, Any] | None = None, graph_key: str | None = None, history: list[dict[str, Any]] | None = None, + run_id: str | None = None, + extra_track_data: dict[str, Any] | None = None, + span_attributes: dict[str, Any] | None = None, ) -> AsyncGenerator[ExecuteStreamEvent, None]: """ Streaming counterpart to ``execute_and_track``. Yields ``chunk`` events, @@ -278,24 +325,14 @@ async def execute_and_stream( """ from .lifecycle import get_client - environment_id = _try_get_environment_id() - track_data: TrackData = { - "runId": str(uuid.uuid4()), - "configKey": config_key, - "variationKey": meta.get("variationKey", "") if isinstance(meta, dict) else "", - "version": meta.get("version", 1) if isinstance(meta, dict) else 1, - "modelName": config.get("model", {}).get("name", "") - if isinstance(config, dict) - else "", - "providerName": config.get("provider", {}).get("name", "") - if isinstance(config, dict) - else "", - **model_stamps_from_meta(meta), - } - if graph_key: - track_data["graphKey"] = graph_key - if environment_id: - track_data["environmentId"] = environment_id + track_data = build_track_data( + config_key=config_key, + config=config, + meta=meta, + graph_key=graph_key, + run_id=run_id, + extra_track_data=extra_track_data, + ) client = get_client() ld_ctx = to_ld_context(client, user_context) @@ -307,11 +344,9 @@ async def execute_and_stream( ld_ctx, track_data, ) - merged_variables: dict[str, Any] = { - **(variables or {}), - "ldContext": {**user_context}, - "__ld": track_data, - } + merged_variables = _merge_variables( + variables, user_context, track_data, span_attributes + ) start_time = time.monotonic() full_text = "" diff --git a/packages/client/src/launchdarkly_ai_server/types.py b/packages/client/src/launchdarkly_ai_server/types.py index 079f9d2e..99161147 100644 --- a/packages/client/src/launchdarkly_ai_server/types.py +++ b/packages/client/src/launchdarkly_ai_server/types.py @@ -288,6 +288,12 @@ class ProviderResponse(Generic[T]): ``None`` when ``skip_judges=False`` (judges ran inline). """ + shadow_tasks: list[ShadowTask] | None = None + """ + Shadow requests sampled for this call, produced when ``skip_shadows=True``. + Run each with ``run_shadow(task, handlers)``. ``None`` when shadows ran in + the background. + """ track_data: TrackData | None = None """ Tracking payload from this invocation. Carried inside each :class:`JudgeTask` @@ -357,6 +363,39 @@ class JudgeTask: """LD metric key to track the score against.""" +@dataclass(frozen=True) +class ShadowTask: + """ + Everything needed to run one shadow request later, without evaluating the + flag again. Returned in ``ProviderResponse.shadow_tasks`` when + ``skip_shadows=True``. Plain data, so it can be pickled and queued. Run it + with ``run_shadow(task, handlers)``. + """ + + config_key: str + shadow_request_id: str + shadow_config: AiConfigRep + """The shadow variation, with tools, judges and shadow configuration removed.""" + shadow_meta: VariationMeta + user_input: str | None + history: list[dict[str, Any]] | None + variables: dict[str, Any] | None + user_context: LDContext + timeout_s: float + track_data: TrackData + """Fields added to every event of the shadow run, including the real run's ``runId``.""" + span_attributes: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ShadowResult: + """The outcome of one shadow run. The shadow's response is discarded.""" + + usage: dict[str, int] + track_data: TrackData + error: str | None = None + + @dataclass class JudgeRunResult: """ diff --git a/packages/client/src/launchdarkly_ai_server/utils.py b/packages/client/src/launchdarkly_ai_server/utils.py index 5e1fabba..8fa5c1c6 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -806,6 +806,10 @@ def set_ld_span_attributes(span: Any, variables: dict[str, Any] | None) -> None: span.set_attribute("launchdarkly.run.id", ld.get("runId", "")) if ld.get("graphKey"): span.set_attribute("launchdarkly.graph.key", ld["graphKey"]) + extra = variables.get("__ld_span") + if isinstance(extra, dict): + for name, value in extra.items(): + span.set_attribute(name, value) feature_flag_attrs: dict[str, str] = { "feature_flag.key": ld.get("configKey", ""), diff --git a/packages/client/tests/conftest.py b/packages/client/tests/conftest.py index de69af0c..5c40d4c7 100644 --- a/packages/client/tests/conftest.py +++ b/packages/client/tests/conftest.py @@ -14,6 +14,7 @@ def mock_ld_client() -> MagicMock: """Stub LaunchDarkly client with variation/track/flush/close.""" client = MagicMock() + del client.variation_detail client.variation = AsyncMock(return_value=None) client.track = MagicMock() client.flush = AsyncMock() diff --git a/packages/client/tests/test_agui.py b/packages/client/tests/test_agui.py index 68613a08..16c7f3d1 100644 --- a/packages/client/tests/test_agui.py +++ b/packages/client/tests/test_agui.py @@ -32,6 +32,7 @@ def _client() -> MagicMock: client = MagicMock() + del client.variation_detail client.track = MagicMock() client.flush = AsyncMock() client.close = AsyncMock() diff --git a/packages/client/tests/test_client.py b/packages/client/tests/test_client.py index 051f8d62..231d4344 100644 --- a/packages/client/tests/test_client.py +++ b/packages/client/tests/test_client.py @@ -34,6 +34,7 @@ def _generation_calls(client: MagicMock) -> list[Any]: def _make_client(response: Any = "Hello", usage: dict | None = None) -> MagicMock: c = MagicMock() + del c.variation_detail c.track = MagicMock() c.flush = AsyncMock() c.close = AsyncMock() diff --git a/packages/client/tests/test_public_surface.py b/packages/client/tests/test_public_surface.py index 29177a3a..a22e7070 100644 --- a/packages/client/tests/test_public_surface.py +++ b/packages/client/tests/test_public_surface.py @@ -44,6 +44,8 @@ "ProviderGraphResponse", "ProviderHandler", "ProviderResponse", + "ShadowResult", + "ShadowTask", "Registry", "RunUsage", "SDK_INFO_CONTEXT", @@ -101,6 +103,7 @@ "resolve_handlers", "resolve_tools", "run_judge", + "run_shadow", "run_judges", "set_conversation_id_if_absent", "set_input_content_attributes", diff --git a/packages/client/tests/test_sdk_usage.py b/packages/client/tests/test_sdk_usage.py index 9f2a8d8e..0a649706 100644 --- a/packages/client/tests/test_sdk_usage.py +++ b/packages/client/tests/test_sdk_usage.py @@ -43,6 +43,7 @@ def _client(variation: Any = None) -> MagicMock: client = MagicMock() + del client.variation_detail client.track = MagicMock() client.variation = AsyncMock(return_value=variation) client.flush = AsyncMock() diff --git a/packages/client/tests/test_shadows.py b/packages/client/tests/test_shadows.py new file mode 100644 index 00000000..d318113b --- /dev/null +++ b/packages/client/tests/test_shadows.py @@ -0,0 +1,550 @@ +"""Tests for shadow requests.""" + +import asyncio +import pickle +import time +from collections.abc import Iterator +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import launchdarkly_ai_server.lifecycle as lifecycle_module +from launchdarkly_ai_server import ( + ProviderHandler, + ShadowTask, + config, + run_shadow, + shutdown, +) +from launchdarkly_ai_server import shadows as sh +from launchdarkly_ai_server.sdk_usage import reset_sdk_usage +from launchdarkly_ai_server.utils import set_ld_span_attributes + +CONTEXT = {"kind": "user", "key": "u1"} +FUTURE = int(time.time() * 1000) + 3_600_000 +PAST = int(time.time() * 1000) - 1 +_SHADOW_FIELDS = { + "isShadow", + "primaryVariationKey", + "shadowRequestId", + "ownerKind", + "ownerId", +} +_BOOKKEEPING = {"$ld:ai:sdk:usage", "$ld:ai:sdk:info"} + + +def _entry(**overrides: Any) -> dict[str, Any]: + entry = { + "id": "shadow-1", + "variationKey": "b", + "samplingRate": 1.0, + "endsAt": FUTURE, + "ownerKind": "member", + "ownerId": "member-1", + } + entry.update(overrides) + return entry + + +def _variation_b() -> dict[str, Any]: + return { + "_ldMeta": { + "enabled": True, + "variationKey": "b", + "version": 7, + "mode": "messages", + "modelKey": "claude", + "modelVersion": 2, + }, + "model": {"name": "claude-x"}, + "provider": {"name": "ShadowProvider"}, + "instructions": "Shadow.", + "tools": {"search": {"name": "search", "type": "function", "parameters": {}}}, + "judgeConfiguration": {"judges": [{"key": "j", "samplingRate": 1}]}, + } + + +def _variation_a( + shadows: dict[str, Any] | None = None, variations: dict[str, Any] | None = None +) -> dict[str, Any]: + return { + "_ldMeta": { + "enabled": True, + "variationKey": "a", + "version": 1, + "mode": "messages", + }, + "model": {"name": "gpt-4"}, + "provider": {"name": "TestProvider"}, + "instructions": "Be helpful.", + "shadowConfiguration": { + "shadows": shadows if shadows is not None else {"fallthrough": [_entry()]}, + "variations": variations + if variations is not None + else {"b": _variation_b()}, + }, + } + + +def _detail(value: dict[str, Any], reason: dict[str, Any] | None) -> SimpleNamespace: + return SimpleNamespace(value=value, reason=reason) + + +FALLTHROUGH = {"kind": "FALLTHROUGH"} + + +def _make_client( + value: dict[str, Any], reason: dict[str, Any] | None = FALLTHROUGH +) -> MagicMock: + c = MagicMock() + c.track = MagicMock() + c.flush = AsyncMock() + c.close = AsyncMock() + c.variation = AsyncMock(return_value=value) + c.variation_detail = AsyncMock(return_value=_detail(value, reason)) + return c + + +class _Recorder: + def __init__( + self, + provider: str, + output: str = "ok", + *, + raises: bool = False, + sleep_s: float = 0, + error: Exception | None = None, + ) -> None: + self.calls: list[dict[str, Any]] = [] + self._output = output + self._raises = raises + self._sleep_s = sleep_s + self._error = error + self.handler = ProviderHandler(fn=self._fn, provides_for=(provider, "messages")) # type: ignore[arg-type] + + async def _fn( + self, + cfg: Any, + user_input: Any, + tool_handlers: Any, + variables: Any, + history: Any = None, + ) -> dict[str, Any]: + self.calls.append( + { + "config": cfg, + "input": user_input, + "tools": tool_handlers, + "variables": variables, + "history": history, + } + ) + if self._sleep_s: + await asyncio.sleep(self._sleep_s) + if self._error is not None: + raise self._error + if self._raises: + raise RuntimeError("shadow boom") + return { + "output": self._output, + "usage": {"input_tokens": 3, "output_tokens": 2}, + } + + +def _events(client: MagicMock) -> list[tuple[str, dict[str, Any], Any]]: + return [ + (c[0][0], c[0][2], c[0][3]) + for c in client.track.call_args_list + if c[0][0] not in _BOOKKEEPING + ] + + +def _shadow_events(client: MagicMock) -> list[tuple[str, dict[str, Any], Any]]: + return [e for e in _events(client) if e[1].get("isShadow")] + + +def _real_events(client: MagicMock) -> list[tuple[str, dict[str, Any], Any]]: + return [e for e in _events(client) if not e[1].get("isShadow")] + + +@pytest.fixture(autouse=True) +def _reset() -> Iterator[None]: + lifecycle_module._reset_for_testing() + reset_sdk_usage() + yield + lifecycle_module._reset_for_testing() + + +def _use(client: MagicMock) -> MagicMock: + lifecycle_module._set_client_for_testing(client) + return client + + +async def _invoke( + client: MagicMock, + *, + user_input: str = "hi", + history: list[dict[str, Any]] | None = None, + **config_args: Any, +) -> tuple[Any, _Recorder, _Recorder]: + _use(client) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", "shadow answer") + response = await config( + key="flag", handler=[a.handler, b.handler], **config_args + ).invoke(user_input, CONTEXT, {"topic": "x"}, history) + await sh._wait_for_shadows_for_testing() + return response, a, b + + +async def _deferred_task(client: MagicMock) -> ShadowTask: + response, _, _ = await _invoke(client, skip_shadows=True) + assert response.shadow_tasks is not None + return response.shadow_tasks[0] + + +class TestRunning: + async def test_runs_the_fallthrough_shadow(self) -> None: + client = _make_client(_variation_a()) + _, a, b = await _invoke(client) + assert len(a.calls) == 1 + assert len(b.calls) == 1 + client.variation_detail.assert_called_once() + client.variation.assert_not_called() + + async def test_client_without_variation_detail_runs_no_shadows(self) -> None: + client = _make_client(_variation_a()) + del client.variation_detail + response, a, b = await _invoke(client) + assert response.response == "real answer" + assert len(a.calls) == 1 + assert b.calls == [] + + @pytest.mark.parametrize("timeout_s", [0, -1, float("inf"), float("nan")]) + def test_timeout_must_be_positive_and_finite(self, timeout_s: float) -> None: + with pytest.raises(ValueError): + config(key="flag", shadow_timeout_s=timeout_s) + + async def test_shutdown_cancels_running_shadows(self) -> None: + _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", sleep_s=5) + await config(key="flag", handler=[a.handler, b.handler]).invoke("hi", CONTEXT) + assert sh._running + await shutdown() + assert not sh._running + + async def test_same_input_and_stripped_config(self) -> None: + history = [{"role": "user", "content": "earlier"}] + _, a, b = await _invoke( + _make_client(_variation_a()), user_input="question", history=history + ) + call = b.calls[0] + assert call["input"] == "question" + assert call["history"] == history + assert call["variables"]["topic"] == "x" + assert "tools" not in call["config"] + assert "judgeConfiguration" not in call["config"] + assert "shadowConfiguration" not in a.calls[0]["config"] + assert "shadowConfiguration" not in call["config"] + assert not call["tools"] + + async def test_shadow_answer_never_joins_history(self) -> None: + history = [{"role": "user", "content": "earlier"}] + await _invoke(_make_client(_variation_a()), history=history) + assert history == [{"role": "user", "content": "earlier"}] + + async def test_failing_shadow_never_reaches_the_caller(self) -> None: + client = _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", raises=True) + response = await config(key="flag", handler=[a.handler, b.handler]).invoke( + "hi", CONTEXT + ) + await sh._wait_for_shadows_for_testing() + assert response.response == "real answer" + assert "$ld:ai:generation:success" in [e[0] for e in _real_events(client)] + assert "$ld:ai:generation:error" in [e[0] for e in _shadow_events(client)] + + async def test_timeout_is_tracked_once(self) -> None: + client = _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", sleep_s=1) + response = await config( + key="flag", handler=[a.handler, b.handler], shadow_timeout_s=0.05 + ).invoke("hi", CONTEXT) + await sh._wait_for_shadows_for_testing() + assert response.response == "real answer" + names = [e[0] for e in _shadow_events(client)] + assert names.count("$ld:ai:generation:error") == 1 + assert names.count("$ld:ai:duration:total") == 1 + + async def test_handler_timeout_error_is_tracked_once(self) -> None: + client = _make_client(_variation_a()) + task = await _deferred_task(client) + b = _Recorder("ShadowProvider", error=TimeoutError("provider timeout")) + result = await run_shadow(task, [b.handler]) + names = [e[0] for e in _shadow_events(client)] + assert result.error == "provider timeout" + assert names.count("$ld:ai:generation:error") == 1 + + async def test_real_run_is_unchanged(self) -> None: + baseline_client = _make_client(_variation_a(shadows={})) + baseline, _, _ = await _invoke(baseline_client) + lifecycle_module._reset_for_testing() + client = _make_client(_variation_a()) + response, _, _ = await _invoke(client) + assert response.response == baseline.response + assert response.usage == baseline.usage + assert [e[0] for e in _real_events(client)] == [ + e[0] for e in _real_events(baseline_client) + ] + + async def test_stream_runs_the_shadow_as_invoke(self) -> None: + _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", "shadow answer") + events = [ + e + async for e in config(key="flag", handler=[a.handler, b.handler]).stream( + "hi", CONTEXT + ) + ] + await sh._wait_for_shadows_for_testing() + assert events[-1]["type"] == "done" + assert "shadow_tasks" not in events[-1] + assert len(b.calls) == 1 + + +class TestPayload: + @pytest.mark.parametrize( + "shadow_configuration", + [ + "nope", + {"shadows": {"fallthrough": "nope"}, "variations": {"b": _variation_b()}}, + { + "shadows": {"fallthrough": [{"id": "x"}]}, + "variations": {"b": _variation_b()}, + }, + { + "shadows": {"fallthrough": [_entry(samplingRate="high")]}, + "variations": {"b": _variation_b()}, + }, + ], + ) + async def test_malformed_configuration_is_ignored( + self, shadow_configuration: Any + ) -> None: + value = _variation_a() + value["shadowConfiguration"] = shadow_configuration + response, _, b = await _invoke(_make_client(value)) + assert response.response == "real answer" + assert b.calls == [] + + +class TestSelection: + async def test_rule_match_uses_only_that_rule(self) -> None: + shadows = { + "r1": [_entry(id="rule-shadow")], + "fallthrough": [_entry(id="ft-shadow")], + } + client = _make_client( + _variation_a(shadows), + {"kind": "RULE_MATCH", "ruleId": "r1", "ruleIndex": 0}, + ) + await _invoke(client) + ids = {e[1]["shadowRequestId"] for e in _shadow_events(client)} + assert ids == {"rule-shadow"} + + @pytest.mark.parametrize( + "reason", + [ + {"kind": "TARGET_MATCH"}, + {"kind": "OFF"}, + {"kind": "PREREQUISITE_FAILED"}, + {"kind": "ERROR"}, + None, + ], + ) + async def test_other_reasons_run_nothing( + self, reason: dict[str, Any] | None + ) -> None: + _, _, b = await _invoke(_make_client(_variation_a(), reason)) + assert b.calls == [] + + async def test_entry_past_its_end_is_skipped(self) -> None: + shadows = {"fallthrough": [_entry(id="ended", endsAt=PAST), _entry(id="live")]} + client = _make_client(_variation_a(shadows)) + await _invoke(client) + ids = {e[1]["shadowRequestId"] for e in _shadow_events(client)} + assert ids == {"live"} + + async def test_missing_variation_is_skipped(self) -> None: + _, _, b = await _invoke( + _make_client(_variation_a({"fallthrough": [_entry(variationKey="gone")]})) + ) + assert b.calls == [] + + async def test_no_matching_handler_is_skipped(self) -> None: + client = _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + response = await config(key="flag", handler=a.handler).invoke("hi", CONTEXT) + await sh._wait_for_shadows_for_testing() + assert response.response == "real answer" + assert _shadow_events(client) == [] + + def test_sampling_is_sticky(self) -> None: + assert sh._in_sample("s1", CONTEXT, 0.5) == sh._in_sample("s1", CONTEXT, 0.5) + + def test_sampling_honors_the_rate(self) -> None: + hits = sum( + sh._in_sample("s1", {"kind": "user", "key": f"k{i}"}, 0.2) + for i in range(10_000) + ) + assert 1800 <= hits <= 2200 + + def test_sampling_edges(self) -> None: + contexts = [{"kind": "user", "key": f"k{i}"} for i in range(200)] + assert all(sh._in_sample("s1", c, 1.0) for c in contexts) + assert not any(sh._in_sample("s1", c, 0.0) for c in contexts) + + def test_sampling_matches_typescript(self) -> None: + assert sh._in_sample("s1", CONTEXT, 0.732) + assert not sh._in_sample("s1", CONTEXT, 0.731) + + def test_context_without_key_is_not_sampled(self) -> None: + assert not sh._in_sample("s1", {"kind": "user"}, 1.0) + + +class TestDeferred: + async def test_skip_shadows_returns_tasks_and_runs_nothing(self) -> None: + client = _make_client(_variation_a()) + response, a, b = await _invoke(client, skip_shadows=True) + assert len(a.calls) == 1 + assert b.calls == [] + assert response.shadow_tasks is not None + assert len(response.shadow_tasks) == 1 + assert _shadow_events(client) == [] + + async def test_tasks_are_none_when_shadows_run_in_the_background(self) -> None: + response, _, _ = await _invoke(_make_client(_variation_a())) + assert response.shadow_tasks is None + + async def test_stream_done_event_carries_tasks(self) -> None: + _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + events = [ + e + async for e in config( + key="flag", handler=a.handler, skip_shadows=True + ).stream("hi", CONTEXT) + ] + assert events[-1]["type"] == "done" + assert events[-1]["shadow_tasks"] == [] + + async def test_task_survives_pickle(self) -> None: + task = await _deferred_task(_make_client(_variation_a())) + restored = pickle.loads(pickle.dumps(task)) + assert restored == task + assert restored.shadow_request_id == "shadow-1" + + async def test_run_shadow_matches_background_mode(self) -> None: + deferred_client = _make_client(_variation_a()) + task = await _deferred_task(deferred_client) + result = await run_shadow(task, [_Recorder("ShadowProvider").handler]) + assert result.error is None + assert result.track_data["isShadow"] is True + deferred = [ + (e[0], {k: v for k, v in e[1].items() if k != "runId"}) + for e in _shadow_events(deferred_client) + ] + + lifecycle_module._reset_for_testing() + background_client = _make_client(_variation_a()) + await _invoke(background_client) + background = [ + (e[0], {k: v for k, v in e[1].items() if k != "runId"}) + for e in _shadow_events(background_client) + ] + assert deferred == background + + async def test_run_shadow_never_raises(self) -> None: + task = await _deferred_task(_make_client(_variation_a())) + result = await run_shadow( + task, [_Recorder("ShadowProvider", raises=True).handler] + ) + assert result.error + + async def test_run_shadow_reports_usage(self) -> None: + client = _make_client(_variation_a()) + task = await _deferred_task(client) + await run_shadow(task, [_Recorder("ShadowProvider").handler]) + helpers = [ + c[0][2].get("helper") + for c in client.track.call_args_list + if c[0][0] == "$ld:ai:sdk:usage" + ] + assert "client.runShadow" in helpers + + +class TestEvents: + async def test_shadow_events_describe_the_shadow_variation(self) -> None: + client = _make_client(_variation_a()) + response, _, _ = await _invoke(client) + events = _shadow_events(client) + assert {e[0] for e in events} >= { + "$ld:ai:duration:total", + "$ld:ai:generation:success", + "$ld:ai:tokens:total", + } + data = events[0][1] + assert data["configKey"] == "flag" + assert data["variationKey"] == "b" + assert data["version"] == 7 + assert data["modelName"] == "claude-x" + assert data["providerName"] == "ShadowProvider" + assert data["modelKey"] == "claude" + assert data["modelVersion"] == 2 + assert data["runId"] == response.track_data["runId"] + assert data["isShadow"] is True + assert data["primaryVariationKey"] == "a" + assert data["shadowRequestId"] == "shadow-1" + assert data["ownerKind"] == "member" + assert data["ownerId"] == "member-1" + + async def test_nothing_is_added_to_real_events(self) -> None: + client = _make_client(_variation_a()) + response, _, _ = await _invoke(client) + assert not (_SHADOW_FIELDS & set(response.track_data)) + for _, data, _ in _real_events(client): + assert not (_SHADOW_FIELDS & set(data)) + + def test_span_attributes_are_set_on_the_root_span( + self, mock_span: MagicMock + ) -> None: + set_ld_span_attributes( + mock_span, + { + "__ld": {"configKey": "flag", "variationKey": "b", "runId": "r"}, + "__ld_span": {"launchdarkly.is_shadow": True}, + }, + ) + mock_span.set_attribute.assert_any_call("launchdarkly.is_shadow", True) + + async def test_shadow_span_attributes_reach_the_handler(self) -> None: + _, _, b = await _invoke(_make_client(_variation_a())) + assert b.calls[0]["variables"]["__ld_span"] == { + "launchdarkly.is_shadow": True, + "launchdarkly.shadow.request.id": "shadow-1", + "launchdarkly.shadow.primary.variation.key": "a", + } + + async def test_template_variables_cannot_set_span_attributes(self) -> None: + _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + await config(key="flag", handler=a.handler).invoke( + "hi", CONTEXT, {"__ld_span": {"x": 1}} + ) + assert "__ld_span" not in a.calls[0]["variables"] diff --git a/packages/client/tests/test_stream_conversation.py b/packages/client/tests/test_stream_conversation.py index 198f993c..a94dc57c 100644 --- a/packages/client/tests/test_stream_conversation.py +++ b/packages/client/tests/test_stream_conversation.py @@ -47,6 +47,7 @@ def _reset_exporter() -> Iterator[None]: @pytest.fixture def mock_ld_client() -> Iterator[MagicMock]: c = MagicMock() + del c.variation_detail c.track = MagicMock() c.flush = AsyncMock() c.close = AsyncMock() diff --git a/packages/opencode-agents/tests/test_handler.py b/packages/opencode-agents/tests/test_handler.py index 6adc2f20..d321f9b7 100644 --- a/packages/opencode-agents/tests/test_handler.py +++ b/packages/opencode-agents/tests/test_handler.py @@ -277,6 +277,7 @@ def _reset_sdk() -> Iterator[None]: @pytest.fixture def mock_ld_client() -> Iterator[MagicMock]: client = MagicMock() + del client.variation_detail client.track = MagicMock() client.flush = AsyncMock() client.close = AsyncMock()