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..2e000043 100644 --- a/packages/client/README.md +++ b/packages/client/README.md @@ -379,6 +379,29 @@ elif result and task.evaluation_metric_key: get_client().track(task.evaluation_metric_key, context, result.track_data, result.score) ``` +#### Shadow requests + +A shadow request runs another variation of an AI Config on the same input and reports it on the normal AI metrics with `isShadow: true`. You attach shadow requests to 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 in the background, each with a 30 second timeout that you can change with `shadow_timeout_s`. A shadow makes no tool calls and runs no judges. Graphs and judges never run shadows. + +With `skip_shadows=True`, `invoke()` returns `shadow_tasks` instead of running them, and `stream()` puts them on the `done` event. A task can be pickled, so you can run it later on your own queue or worker with `run_shadow(task, handlers)`, the same way as `run_judge`: + +```python +from launchdarkly_ai_server import config, run_shadow + +result = await config(key="my-ai-config", handler=handlers, skip_shadows=True).invoke( + user_input, context +) +for task in result.shadow_tasks or []: + queue.put_nowait(task) + +# in the worker +outcome = await run_shadow(task, handlers) +``` + +`run_shadow` never raises. A failure or timeout is in `outcome.error`, and the shadow's events are already tracked. + ### `graph(key, **options)` Runs a multi-agent workflow defined in a LaunchDarkly agent graph flag. The SDK uses a **model-driven router**: it starts at the root node, presents outgoing edges as handoff choices to the model, and follows whichever edge the model selects. The loop terminates when the model produces a final answer, a leaf is reached, a cycle is detected, or the step cap is hit. @@ -1032,7 +1055,7 @@ All types are exported from this package. Handler packages import them from here | `AiConfigRep` | The AI configuration object fetched from a LaunchDarkly flag variation | | `Tool` | A tool definition (name, description, JSON Schema parameters) | | `ProviderHandler` | The callable type that all handler packages produce | -| `ProviderResponse` | The value returned to callers: `response`, `usage`, `track_data`, `judge_results?`, `judge_tasks?`. `judge_results` is populated when `skip_judges=False`; `judge_tasks` (a `list[JudgeTask]`) is populated when `skip_judges=True`. | +| `ProviderResponse` | The value returned to callers: `response`, `usage`, `track_data`, `judge_results?`, `judge_tasks?`, `shadow_tasks?`. `judge_results` is populated when `skip_judges=False`; `judge_tasks` (a `list[JudgeTask]`) is populated when `skip_judges=True`; `shadow_tasks` (a `list[ShadowTask]`) is populated when `skip_shadows=True`. | | `ConfigArgs` | Arguments accepted by `config()` (key, handler, tool_handlers, registry) | | `NativeTool` | Marker class for provider built-in tools | | `LDContext` | Standard LaunchDarkly context dict. Import from `launchdarkly_ai_server`. | diff --git a/packages/client/api-report/launchdarkly_ai_server.txt b/packages/client/api-report/launchdarkly_ai_server.txt index 4ac0e30a..39ad14cc 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 = 30.0) -> 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 = 30.0) 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) -> 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, track_data: TrackData | None = None, shadow_tasks: list[ShadowTask] | 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: UsageDict, track_data: TrackData, error: str | None = None) + error: str | None = None + track_data: TrackData + usage: UsageDict +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]) + config_key: str + history: list[dict[str, Any]] | None + shadow_config: AiConfigRep + shadow_meta: VariationMeta + shadow_request_id: str + span_attributes: dict[str, Any] + 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..345ec602 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", @@ -276,6 +281,8 @@ "build_judge_tasks", "run_judge", "run_judges", + # shadows + "run_shadow", # client "config", "ConfigInstance", diff --git a/packages/client/src/launchdarkly_ai_server/client.py b/packages/client/src/launchdarkly_ai_server/client.py index 48d1109c..2ea3c8db 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 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, ) @@ -55,12 +59,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 = 30.0, ) -> 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: @@ -69,6 +79,33 @@ def _normalize_handlers(self) -> list[ProviderHandler] | None: return self._handler return [self._handler] + def _shadows( + self, + variation: dict[str, Any], + context: LDContext, + user_input: str | None, + history: list[dict[str, Any]] | None, + variables: dict[str, Any] | None, + handlers: list[ProviderHandler], + run_id: str, + ) -> list[ShadowTask] | None: + """Starts the sampled shadows, or returns them when ``skip_shadows`` is set.""" + tasks = build_shadow_tasks( + config_key=self._key, + variation=variation, + context=context, + user_input=user_input, + history=history, + variables=variables, + handlers=handlers, + run_id=run_id, + timeout_s=self._shadow_timeout_s, + ) + if self._skip_shadows: + return tasks + start_shadows(tasks, handlers) + return None + async def invoke( self, user_input: str | None, @@ -91,13 +128,23 @@ 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 = self._shadows( + variation, + context, + user_input, + history, + variables, + resolved_handler_list, + run_id, + ) result = await execute_and_track( config_key=self._key, @@ -109,6 +156,7 @@ async def _invoke( tool_handlers=resolved_tools, variables=variables, history=history, + run_id=run_id, ) raw_response = result["response"] @@ -140,6 +188,7 @@ async def _invoke( response=parsed_response, usage=usage_obj, judge_tasks=judge_tasks, + shadow_tasks=shadow_tasks, track_data=track_data, ) @@ -158,6 +207,7 @@ async def _invoke( response=parsed_response, usage=usage_obj, judge_results=judge_results if judge_results else None, + shadow_tasks=shadow_tasks, track_data=track_data, ) @@ -210,7 +260,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 +278,16 @@ async def _execute_stream( sink["tools"] = resolved_tools sink["provider"] = provider sink["model"] = model + run_id = str(uuid.uuid4()) + sink["shadow_tasks"] = self._shadows( + variation, + context, + user_input, + history, + variables, + resolved_handler_list, + run_id, + ) async with aclosing( execute_and_stream( @@ -240,6 +300,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 +327,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 +371,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 = 30.0, ) -> ConfigInstance: """ Creates a ``ConfigInstance`` bound to *key*. Accepts a single handler or a @@ -319,6 +385,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 on the matched rule or fallthrough run in the background, + each limited to *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 +397,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 +409,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 = 30.0, ) -> ConfigInstance: """Non-reporting :func:`config` for SDK wrappers.""" return _InternalConfigInstance( @@ -344,4 +419,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..a7d87e29 100644 --- a/packages/client/src/launchdarkly_ai_server/lifecycle.py +++ b/packages/client/src/launchdarkly_ai_server/lifecycle.py @@ -586,12 +586,18 @@ 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. Returns ``{"config": AiConfigRep, "meta": VariationMeta}``. Raises on disabled or invalid variations. + Also returns ``reason`` and ``shadow_configuration``, which ``config()`` + uses to pick shadow requests. ``reason`` is ``None`` unless *with_reason* is + set and the client has ``variation_detail``. + Lazily initializes the LD client when ``LD_SDK_KEY`` is set, matching the TypeScript SDK's behaviour where any AI call auto-initializes. """ @@ -603,13 +609,25 @@ 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 + variation_detail = ( + getattr(client, "variation_detail", None) if with_reason else None ) + reason: Any = 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 = detail.value + reason = detail.reason + 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 +649,8 @@ async def extract_variation( **model_stamps_from_meta(ld_meta), } - # Strip _ldMeta 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 +661,9 @@ 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, + "reason": reason, + "shadow_configuration": raw.get("shadowConfiguration"), + } 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..89d15bed --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/shadows.py @@ -0,0 +1,252 @@ +""" +Shadow requests: run another variation of the AI Config 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 sends shadow requests inside the served variation as +``shadowConfiguration``. ``config().invoke()`` and ``.stream()`` call +:func:`build_shadow_tasks`, then either :func:`start_shadows` or, with +``skip_shadows=True``, hand the tasks back for :func:`run_shadow`. +""" + +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, + select_handler, + to_usage_dict, +) + +logger = logging.getLogger("launchdarkly_ai_server.shadows") + +_STRIPPED_KEYS = ("_ldMeta", "tools", "judgeConfiguration", "shadowConfiguration") +_BUCKET_SCALE = float(0xFFFFFFFFFFFFFFF) + +_running: set[asyncio.Task[ShadowResult]] = set() + + +def build_shadow_tasks( + *, + config_key: str, + variation: dict[str, Any], + 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 *variation*, skipping entries that have ended, are not sampled for + this context, or have no matching handler. Never raises. + """ + configuration = variation.get("shadow_configuration") + slot = _slot(variation.get("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) + if not isinstance(entries, list): + return [] + + primary_variation_key = str(variation["meta"].get("variationKey", "")) + now_ms = time.time() * 1000 + tasks: list[ShadowTask] = [] + for entry in entries: + if not _valid_entry(entry) or entry["endsAt"] <= now_ms: + continue + if not _in_sample(entry["id"], context, 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.""" + for task in tasks: + running = asyncio.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 _slot(reason: Any) -> str | None: + if not isinstance(reason, dict): + return None + if reason.get("kind") == "FALLTHROUGH": + return "fallthrough" + if reason.get("kind") == "RULE_MATCH" and isinstance(reason.get("ruleId"), str): + return str(reason["ruleId"]) + return None + + +def _valid_entry(entry: Any) -> bool: + if not isinstance(entry, dict): + return False + for key in ("id", "variationKey"): + if not isinstance(entry.get(key), str) or not entry[key]: + return False + for key in ("samplingRate", "endsAt"): + if not isinstance(entry.get(key), (int, float)) or isinstance(entry[key], bool): + return False + return True + + +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) + return int(digest.hexdigest()[:15], 16) / _BUCKET_SCALE < sampling_rate + + +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: + run_id = task.track_data.get("runId") + extra = {k: v for k, v in task.track_data.items() if k != "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) + 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, + 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(to_usage_dict({}), track_data, error="timeout") + except Exception as err: + return ShadowResult(to_usage_dict({}), track_data, error=str(err)) + return ShadowResult(to_usage_dict(result["usage"]), 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 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..9d1c6aa9 100644 --- a/packages/client/src/launchdarkly_ai_server/tracking.py +++ b/packages/client/src/launchdarkly_ai_server/tracking.py @@ -20,7 +20,12 @@ TrackData, VariationMeta, ) -from .utils import model_stamps_from_meta, parse_usage, to_ld_context +from .utils import ( + SPAN_ATTRIBUTES_KEY, + model_stamps_from_meta, + parse_usage, + to_ld_context, +) @asynccontextmanager @@ -52,6 +57,55 @@ def _try_get_environment_id() -> str | None: return os.environ.get("LD_ENVIRONMENT_ID") or None +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.""" + 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 = 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,31 +221,27 @@ 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 ``{usage, response, track_data}``. + + *extra_track_data* is added to every event, and *span_attributes* to the + handler's root span. """ 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 +257,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 +314,7 @@ 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, ) -> AsyncGenerator[ExecuteStreamEvent, None]: """ Streaming counterpart to ``execute_and_track``. Yields ``chunk`` events, @@ -278,24 +327,13 @@ 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, + ) client = get_client() ld_ctx = to_ld_context(client, user_context) @@ -307,11 +345,7 @@ 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) 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..1d897f24 100644 --- a/packages/client/src/launchdarkly_ai_server/types.py +++ b/packages/client/src/launchdarkly_ai_server/types.py @@ -294,6 +294,11 @@ class ProviderResponse(Generic[T]): so that background judge results are attributed to the originating request (run ID, config key, graph key, etc.). """ + shadow_tasks: list[ShadowTask] | None = None + """ + Sampled shadow requests, produced when ``skip_shadows=True``. Run each with + ``run_shadow(task, handlers)``. ``None`` when shadows ran in the background. + """ TrackData = dict[str, Any] @@ -357,6 +362,37 @@ 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. Plain data, so it can be pickled and queued. + """ + + config_key: str + shadow_request_id: str + shadow_config: AiConfigRep + """The shadow variation, without tools, judges or shadow configuration.""" + 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 served run's ``runId``.""" + span_attributes: dict[str, Any] + + +@dataclass(frozen=True) +class ShadowResult: + """The outcome of one shadow run. The shadow's response is discarded.""" + + usage: UsageDict + 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..ed225ebe 100644 --- a/packages/client/src/launchdarkly_ai_server/utils.py +++ b/packages/client/src/launchdarkly_ai_server/utils.py @@ -768,6 +768,10 @@ def _context_identity_from_ld_context( return None +SPAN_ATTRIBUTES_KEY = "__ld_span" +"""Variables key for extra root-span attributes. Only the SDK sets it.""" + + def set_ld_span_attributes(span: Any, variables: dict[str, Any] | None) -> None: """ Sets LaunchDarkly config-identifying attributes on an OTel span and emits @@ -806,6 +810,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(SPAN_ATTRIBUTES_KEY) + 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..5037685f 100644 --- a/packages/client/tests/test_public_surface.py +++ b/packages/client/tests/test_public_surface.py @@ -48,6 +48,8 @@ "RunUsage", "SDK_INFO_CONTEXT", "SDK_INFO_EVENT", + "ShadowResult", + "ShadowTask", "SpanMessage", "SpanMessagePart", "SpanUsage", @@ -102,6 +104,7 @@ "resolve_tools", "run_judge", "run_judges", + "run_shadow", "set_conversation_id_if_absent", "set_input_content_attributes", "set_ld_span_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..0e029ab4 --- /dev/null +++ b/packages/client/tests/test_shadows.py @@ -0,0 +1,639 @@ +"""Tests for shadow requests.""" + +import asyncio +import dataclasses +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, + extract_variation, + 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 +FALLTHROUGH = {"kind": "FALLTHROUGH"} +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(enabled: bool = True) -> dict[str, Any]: + return { + "_ldMeta": { + "enabled": enabled, + "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 _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=SimpleNamespace(value=value, reason=reason) + ) + return c + + +class _Recorder: + def __init__( + self, + provider: str, + output: str = "ok", + *, + sleep_s: float = 0, + error: Exception | None = None, + ) -> None: + self.calls: list[dict[str, Any]] = [] + self.streamed = 0 + self._output = output + self._sleep_s = sleep_s + self._error = error + self.handler = ProviderHandler( + fn=self._fn, # type: ignore[arg-type] + provides_for=(provider, "messages"), + stream_fn=self._stream, # 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 + return { + "output": self._output, + "usage": {"input_tokens": 3, "output_tokens": 2}, + } + + async def _stream(self, *args: Any) -> Any: + self.streamed += 1 + yield {"type": "chunk", "text": self._output} + yield {"type": "done", "usage": {"input_tokens": 3, "output_tokens": 2}} + + +def _events(client: MagicMock) -> list[tuple[str, dict[str, Any]]]: + return [ + (c[0][0], c[0][2]) + 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]]]: + return [e for e in _events(client) if e[1].get("isShadow")] + + +def _real_events(client: MagicMock) -> list[tuple[str, dict[str, Any]]]: + return [e for e in _events(client) if not e[1].get("isShadow")] + + +def _usage_helpers(client: MagicMock) -> list[str]: + return [ + c[0][2].get("helper") + for c in client.track.call_args_list + if c[0][0] == "$ld:ai:sdk:usage" + ] + + +@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, + b: _Recorder | None = None, + **config_args: Any, +) -> tuple[Any, _Recorder, _Recorder]: + _use(client) + a = _Recorder("TestProvider", "real answer") + b = b or _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 TestPayload: + async def test_shadow_configuration_never_reaches_the_handler(self) -> None: + _, a, b = await _invoke(_make_client(_variation_a())) + assert "shadowConfiguration" not in a.calls[0]["config"] + assert "shadowConfiguration" not in b.calls[0]["config"] + + @pytest.mark.parametrize( + "shadow_configuration", + [ + "nope", + {"shadows": "nope", "variations": {"b": _variation_b()}}, + {"shadows": {"fallthrough": "nope"}, "variations": {"b": _variation_b()}}, + {"shadows": {"fallthrough": ["nope"]}, "variations": {"b": _variation_b()}}, + {"shadows": {"fallthrough": [{"id": "x"}]}, "variations": {"b": {}}}, + { + "shadows": {"fallthrough": [_entry(id=None)]}, + "variations": {"b": _variation_b()}, + }, + { + "shadows": {"fallthrough": [_entry(variationKey=None)]}, + "variations": {"b": _variation_b()}, + }, + { + "shadows": {"fallthrough": [_entry(samplingRate="high")]}, + "variations": {"b": _variation_b()}, + }, + { + "shadows": {"fallthrough": [_entry(endsAt=None)]}, + "variations": {"b": _variation_b()}, + }, + ], + ) + async def test_malformed_configuration_is_ignored( + self, shadow_configuration: Any + ) -> None: + value = _variation_a() + value["shadowConfiguration"] = shadow_configuration + response, a, b = await _invoke(_make_client(value)) + assert response.response == "real answer" + assert len(a.calls) == 1 + assert b.calls == [] + + +class TestSelection: + async def test_evaluates_once_with_variation_detail(self) -> None: + client = _make_client(_variation_a()) + await _invoke(client) + client.variation_detail.assert_called_once() + client.variation.assert_not_called() + + async def test_other_callers_keep_using_variation(self) -> None: + client = _use(_make_client(_variation_a())) + variation = await extract_variation("flag", CONTEXT) + client.variation.assert_called_once() + client.variation_detail.assert_not_called() + assert "shadowConfiguration" not in variation["config"] + + 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 == [] + + 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"} + ) + await _invoke(client) + assert {e[1]["shadowRequestId"] for e in _shadow_events(client)} == { + "rule-shadow" + } + + async def test_fallthrough_uses_the_fallthrough_slot(self) -> None: + shadows = {"r1": [_entry(id="rule-shadow")], "fallthrough": [_entry(id="ft")]} + client = _make_client(_variation_a(shadows)) + await _invoke(client) + assert {e[1]["shadowRequestId"] for e in _shadow_events(client)} == {"ft"} + + @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) + assert {e[1]["shadowRequestId"] for e in _shadow_events(client)} == {"live"} + + async def test_missing_variation_is_skipped(self) -> None: + shadows = {"fallthrough": [_entry(variationKey="gone")]} + _, _, b = await _invoke(_make_client(_variation_a(shadows))) + assert b.calls == [] + + async def test_disabled_variation_is_skipped(self) -> None: + value = _variation_a(variations={"b": _variation_b(enabled=False)}) + _, _, b = await _invoke(_make_client(value)) + 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: + answers = {sh._in_sample("s1", CONTEXT, 0.5) for _ in range(20)} + assert len(answers) == 1 + + 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_bucket_value(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) + + async def test_stream_runs_the_shadow_as_a_plain_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 a.streamed == 1 + assert len(b.calls) == 1 + assert b.streamed == 0 + + +class TestRunning: + async def test_same_input_on_the_stripped_config(self) -> None: + history = [{"role": "user", "content": "earlier"}] + _, _, 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 call["config"]["instructions"] == "Shadow." + assert "tools" not in call["config"] + assert "judgeConfiguration" 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_background_shadow_does_not_delay_the_caller(self) -> None: + _use(_make_client(_variation_a())) + a = _Recorder("TestProvider", "real answer") + b = _Recorder("ShadowProvider", sleep_s=5) + response = await asyncio.wait_for( + config(key="flag", handler=[a.handler, b.handler]).invoke("hi", CONTEXT), + timeout=1, + ) + assert response.response == "real answer" + assert sh._running + await shutdown() + + async def test_failing_shadow_never_reaches_the_caller(self) -> None: + baseline = _make_client(_variation_a(shadows={})) + expected, _, _ = await _invoke(baseline) + lifecycle_module._reset_for_testing() + client = _make_client(_variation_a()) + response, _, _ = await _invoke( + client, b=_Recorder("ShadowProvider", error=RuntimeError("boom")) + ) + assert response.response == expected.response + assert response.usage == expected.usage + assert [e[0] for e in _real_events(client)] == [ + e[0] for e in _real_events(baseline) + ] + assert "$ld:ai:generation:error" in [e[0] for e in _shadow_events(client)] + + async def test_timed_out_shadow_never_reaches_the_caller(self) -> None: + client = _make_client(_variation_a()) + response, _, _ = await _invoke( + client, b=_Recorder("ShadowProvider", sleep_s=1), shadow_timeout_s=0.05 + ) + assert response.response == "real answer" + assert "$ld:ai:generation:success" in [e[0] for e in _real_events(client)] + + @pytest.mark.parametrize("timeout_s", [0, -1, float("inf"), float("nan")]) + def test_timeout_must_be_finite_and_positive(self, timeout_s: float) -> None: + with pytest.raises(ValueError): + config(key="flag", shadow_timeout_s=timeout_s) + + async def test_timeout_is_tracked_once(self) -> None: + client = _make_client(_variation_a()) + task = await _deferred_task(client) + task = dataclasses.replace(task, timeout_s=0.05) + result = await run_shadow( + task, [_Recorder("ShadowProvider", sleep_s=1).handler] + ) + await asyncio.sleep(0) + names = [e[0] for e in _shadow_events(client)] + assert result.error == "timeout" + assert names.count("$ld:ai:generation:error") == 1 + assert names.count("$ld:ai:duration:total") == 1 + assert "$ld:ai:generation:success" not in names + + async def test_handler_timeout_error_is_an_ordinary_failure(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_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) + running = list(sh._running) + assert running + await shutdown() + assert not sh._running + await asyncio.gather(*running, return_exceptions=True) + assert all(t.cancelled() for t in running) + + async def test_background_shadows_report_only_the_config_call(self) -> None: + client = _make_client(_variation_a()) + await _invoke(client) + assert _usage_helpers(client) == ["client.config.invoke"] + + +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_skip_shadows_returns_an_empty_list_when_none_sampled(self) -> None: + response, _, _ = await _invoke( + _make_client(_variation_a(shadows={})), skip_shadows=True + ) + assert response.shadow_tasks == [] + + async def test_background_mode_returns_no_tasks(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") + b = _Recorder("ShadowProvider") + events = [ + e + async for e in config( + key="flag", handler=[a.handler, b.handler], skip_shadows=True + ).stream("hi", CONTEXT) + ] + done = events[-1] + assert done["type"] == "done" + assert [t.shadow_request_id for t in done["shadow_tasks"]] == ["shadow-1"] + assert b.calls == [] + + async def test_task_is_self_contained_and_survives_pickle(self) -> None: + history = [{"role": "user", "content": "earlier"}] + response, _, _ = await _invoke( + _make_client(_variation_a()), + user_input="question", + history=history, + skip_shadows=True, + shadow_timeout_s=12.5, + ) + task = response.shadow_tasks[0] + restored = pickle.loads(pickle.dumps(task)) + assert restored == task + assert restored.config_key == "flag" + assert restored.shadow_request_id == "shadow-1" + assert restored.shadow_meta["variationKey"] == "b" + assert "tools" not in restored.shadow_config + assert restored.user_input == "question" + assert restored.history == history + assert restored.variables == {"topic": "x"} + assert restored.user_context == CONTEXT + assert restored.timeout_s == 12.5 + assert restored.track_data["runId"] == response.track_data["runId"] + assert restored.span_attributes["launchdarkly.is_shadow"] is True + + 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.usage.total == 5 + assert result.track_data["isShadow"] is True + deferred = [ + (name, {k: v for k, v in data.items() if k != "runId"}) + for name, data in _shadow_events(deferred_client) + ] + + lifecycle_module._reset_for_testing() + background_client = _make_client(_variation_a()) + await _invoke(background_client) + background = [ + (name, {k: v for k, v in data.items() if k != "runId"}) + for name, data 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())) + failed = await run_shadow( + task, [_Recorder("ShadowProvider", error=RuntimeError("boom")).handler] + ) + unrouted = await run_shadow(task, []) + assert failed.error == "boom" + assert unrouted.error + assert unrouted.track_data["isShadow"] is True + + 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]) + assert "client.runShadow" in _usage_helpers(client) + + +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)) + + async def test_shadow_span_attributes_reach_the_root_span( + self, mock_span: MagicMock + ) -> None: + _, _, b = await _invoke(_make_client(_variation_a())) + set_ld_span_attributes(mock_span, b.calls[0]["variables"]) + mock_span.set_attribute.assert_any_call("launchdarkly.is_shadow", True) + mock_span.set_attribute.assert_any_call( + "launchdarkly.shadow.request.id", "shadow-1" + ) + mock_span.set_attribute.assert_any_call( + "launchdarkly.shadow.primary.variation.key", "a" + ) + + async def test_template_variables_cannot_set_span_attributes( + self, mock_span: MagicMock + ) -> None: + _use(_make_client(_variation_a(shadows={}))) + a = _Recorder("TestProvider", "real answer") + await config(key="flag", handler=a.handler).invoke( + "hi", CONTEXT, {"__ld_span": {"launchdarkly.is_shadow": True}} + ) + set_ld_span_attributes(mock_span, a.calls[0]["variables"]) + names = [c[0][0] for c in mock_span.set_attribute.call_args_list] + assert "launchdarkly.is_shadow" not in names 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()