fix(AIC-3570): re-resolve graph flags on every invoke and stream

packages/client/src/launchdarkly_ai_server/graph.py

@@ -2,7 +2,6 @@
 
 import asyncio
 import inspect as _inspect
-import json
 import logging
 import re as _re
 import time
@@ -39,11 +38,10 @@
 logger = logging.getLogger(__name__)
 
 MAX_TRAVERSAL_DEPTH = 100
-MAX_GRAPH_CACHE_SIZE = 512
 
-# graph() caches a built graph per context, track data included. Each run passes its own copy
+# The built graph's track data is minted before the run's runId. Each run passes its own copy
 # with a fresh runId through route/stream_route opts under this key, so the handoff events
-# emitted inside the cached build carry the run's id instead of the build's.
+# emitted inside the build carry the run's id instead of the build's.
 _RUN_TRACK_DATA = "_run_track_data"
 
 
@@ -878,9 +876,6 @@ def __init__(
     ) -> None:
         self._key = key
         self._options = options
-        self._cache: dict[
-            str, tuple[GraphDefinition, TrackData, Callable[..., Any]]
-        ] = {}
 
     async def invoke(
         self,
@@ -927,19 +922,7 @@ async def _invoke(
                 "use resolve_graph() with a framework-native runner."
             )
 
-        try:
-            cache_key: str | None = json.dumps(context, sort_keys=True)
-        except (TypeError, ValueError):
-            cache_key = None
-        if cache_key is not None and cache_key in self._cache:
-            built = self._cache[cache_key]
-        else:
-            built = await _build_graph(self._key, context, resolved_options)
-            if cache_key is not None:
-                if len(self._cache) >= MAX_GRAPH_CACHE_SIZE:
-                    # Evict an arbitrary entry to keep the cache bounded.
-                    self._cache.pop(next(iter(self._cache)))
-                self._cache[cache_key] = built
+        built = await _build_graph(self._key, context, resolved_options)
         graph_def, built_track_data, _ = built
 
         if not graph_def.enabled:
@@ -1165,18 +1148,7 @@ async def _stream_events(
                 "use resolve_graph() with a framework-native runner."
             )
 
-        try:
-            cache_key: str | None = json.dumps(context, sort_keys=True)
-        except (TypeError, ValueError):
-            cache_key = None
-        if cache_key is not None and cache_key in self._cache:
-            built = self._cache[cache_key]
-        else:
-            built = await _build_graph(self._key, context, resolved_options)
-            if cache_key is not None:
-                if len(self._cache) >= MAX_GRAPH_CACHE_SIZE:
-                    self._cache.pop(next(iter(self._cache)))
-                self._cache[cache_key] = built
+        built = await _build_graph(self._key, context, resolved_options)
         graph_def, built_track_data, stream_route = built
 
         if not graph_def.enabled:

packages/client/tests/test_graph.py

@@ -324,14 +324,80 @@ async def test_graph_node_tracked_per_visited_node(
         assert node_events[0][2]["graphKey"] == "graph-key"
         assert "path" not in node_events[0][2]
 
-    async def test_node_variations_resolved_only_once(
+    async def test_node_variations_resolved_on_every_call(
         self, mock_ld_client: MagicMock
     ) -> None:
+        """Each invoke() re-resolves the graph and node flags. See TESTING.md §3.11."""
         g = graph("graph-key", handlers=[_make_handler()])
         await g.invoke("hi", CONTEXT)
+        after_first = mock_ld_client.variation.call_count
         await g.invoke("hi again", CONTEXT)
-        # Second call should use the cache — variation calls should not double
-        assert mock_ld_client.variation.call_count <= 4  # graph + 2 nodes × 1 cache hit
+        assert mock_ld_client.variation.call_count == after_first * 2
+
+    async def test_flag_change_reaches_reused_graph(
+        self, mock_ld_client: MagicMock
+    ) -> None:
+        """A node flag change between calls is seen on the next call. See TESTING.md §3.11."""
+        state = {"instructions": "v1", "graph_enabled": True}
+
+        async def variation(key: str, ctx: dict, default: Any) -> Any:
+            if key == "graph-key":
+                return (
+                    {"root": "only-node", "edges": {}} if state["graph_enabled"] else {}
+                )
+            return {
+                "model": {"name": "gpt-4"},
+                "provider": {"name": "TestProvider"},
+                "instructions": state["instructions"],
+                "_ldMeta": {"enabled": True, "variationKey": "v1", "version": 1},
+            }
+
+        mock_ld_client.variation = AsyncMock(side_effect=variation)
+        seen: list[str] = []
+
+        async def fn(
+            config, user_input, tool_handlers, variables, history=None
+        ) -> dict:  # type: ignore[no-untyped-def]
+            seen.append(config["instructions"])
+            return {"output": "ok", "usage": {"input_tokens": 1, "output_tokens": 1}}
+
+        g = graph(
+            "graph-key",
+            handlers=[
+                ProviderHandler(fn=fn, provides_for=("TestProvider", "messages"))
+            ],  # type: ignore[arg-type]
+        )
+        await g.invoke("first", CONTEXT)
+        state["instructions"] = "v2"
+        await g.invoke("second", CONTEXT)
+        assert seen == ["v1", "v2"]
+
+        state["graph_enabled"] = False
+        with pytest.raises(ValueError, match="disabled"):
+            await g.invoke("third", CONTEXT)
+        with pytest.raises(ValueError, match="disabled"):
+            async for _ in g.stream("fourth", CONTEXT):
+                pass
+
+    async def test_failed_resolution_is_not_reused(
+        self, mock_ld_client: MagicMock
+    ) -> None:
+        """A transient variation failure on one call does not stick to the next."""
+        original = mock_ld_client.variation.side_effect
+        failures = {"left": 1}
+
+        async def flaky(key: str, ctx: dict, default: Any) -> Any:
+            if key != "graph-key" and failures["left"]:
+                failures["left"] -= 1
+                raise RuntimeError("transient")
+            return await original(key, ctx, default)
+
+        mock_ld_client.variation = AsyncMock(side_effect=flaky)
+        g = graph("graph-key", handlers=[_make_handler()])
+        with pytest.raises(ValueError, match="disabled"):
+            await g.invoke("first", CONTEXT)
+        result = await g.invoke("second", CONTEXT)
+        assert result.response == "ok"
 
     async def test_handoff_success_tracked_on_each_handoff(
         self, mock_ld_client: MagicMock
@@ -464,35 +530,6 @@ async def test_disabled_graph_reverse_traverse_key_is_snake_case(
         assert hasattr(gd, "reverse_traverse")
         assert not hasattr(gd, "reverseTraverse")
 
-    async def test_cache_is_bounded(self, mock_ld_client: MagicMock) -> None:
-        """GraphInstance._cache must not grow beyond MAX_GRAPH_CACHE_SIZE.
-        See TESTING.md §3.11."""
-        from launchdarkly_ai_server.graph import MAX_GRAPH_CACHE_SIZE
-
-        g = graph("graph-key", handlers=[_make_handler()])
-        # Invoke with more distinct contexts than the max
-        limit = MAX_GRAPH_CACHE_SIZE + 10
-        for i in range(limit):
-            await g.invoke("hi", {"kind": "user", "key": f"u{i}"})
-        assert len(g._cache) <= MAX_GRAPH_CACHE_SIZE
-
-    async def test_equal_content_contexts_share_cache_entry(
-        self, mock_ld_client: MagicMock
-    ) -> None:
-        """Two distinct dict objects with identical JSON content must share one
-        cache entry. The cache key must be json.dumps(context, sort_keys=True),
-        not id(context). See TESTING.md §3.11."""
-        g = graph("graph-key", handlers=[_make_handler()])
-        ctx_a = {"kind": "user", "key": "u1"}
-        ctx_b = {"kind": "user", "key": "u1"}  # distinct object, same content
-        assert ctx_a is not ctx_b, "test requires two different dict objects"
-        await g.invoke("first call", ctx_a)
-        await g.invoke("second call", ctx_b)
-        assert len(g._cache) == 1, (
-            "Equal-content contexts must share a single cache entry — "
-            "the cache key must be json.dumps(context), not id(context)."
-        )
-
     async def test_error_logging_when_node_variation_fails(
         self, mock_ld_client: MagicMock
     ) -> None: