chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,116 @@
|
||||
"""Pytest configuration for strands showcase unit tests.
|
||||
|
||||
Ensures ``src/`` (so ``agents.agent`` imports) is on ``sys.path``.
|
||||
Shared tools are accessible via the ``tools`` symlink at the project root.
|
||||
|
||||
Also installs minimal stubs for ``ag_ui_strands`` and ``strands`` so the
|
||||
unit tests can run in environments where those heavy runtime deps aren't
|
||||
installed. Tests that need real behavior monkey-patch the specific
|
||||
symbols they touch.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
|
||||
_HERE = os.path.dirname(__file__)
|
||||
_PKG_ROOT = os.path.abspath(os.path.join(_HERE, "..", ".."))
|
||||
|
||||
# src/ holds agent_server.py and agents/
|
||||
sys.path.insert(0, os.path.join(_PKG_ROOT, "src"))
|
||||
# tools/ symlink at project root points to shared/python/tools
|
||||
sys.path.insert(0, _PKG_ROOT)
|
||||
|
||||
|
||||
class _Permissive:
|
||||
"""Base stub that accepts any ``__init__`` args and exposes a writable
|
||||
``_agents_by_thread`` attribute so ``build_showcase_agent`` can swap
|
||||
the per-thread dict in place."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
self._agents_by_thread: dict = {}
|
||||
|
||||
|
||||
def _install_stub_modules() -> None:
|
||||
"""Install minimal stub modules so ``agents.agent`` can be imported.
|
||||
|
||||
Unit tests only exercise pure-Python logic (cap hook counter, dict
|
||||
injection). We stub out strands / ag_ui_strands symbols with plain
|
||||
placeholder classes; tests that care about behavior monkey-patch the
|
||||
relevant attributes.
|
||||
"""
|
||||
if "ag_ui_strands" not in sys.modules:
|
||||
m = types.ModuleType("ag_ui_strands")
|
||||
m.StrandsAgent = _Permissive # type: ignore[attr-defined]
|
||||
m.StrandsAgentConfig = _Permissive # type: ignore[attr-defined]
|
||||
m.ToolBehavior = _Permissive # type: ignore[attr-defined]
|
||||
|
||||
class _FakeFastAPI:
|
||||
"""Accepts the decorators agent_server applies (``@app.get`` etc.)."""
|
||||
|
||||
def _decorator(self, *a, **k):
|
||||
def _wrap(fn):
|
||||
return fn
|
||||
|
||||
return _wrap
|
||||
|
||||
get = post = put = delete = patch = _decorator
|
||||
|
||||
# agent_server.py also calls ``app.add_middleware(...)`` to
|
||||
# install HealthMiddleware; accept that on the stub too.
|
||||
def add_middleware(self, *a, **k):
|
||||
return None
|
||||
|
||||
m.create_strands_app = lambda *a, **k: _FakeFastAPI() # type: ignore[attr-defined]
|
||||
sys.modules["ag_ui_strands"] = m
|
||||
|
||||
if "strands" not in sys.modules:
|
||||
m = types.ModuleType("strands")
|
||||
m.Agent = _Permissive # type: ignore[attr-defined]
|
||||
|
||||
def _tool_decorator(func=None, **_kwargs):
|
||||
if callable(func):
|
||||
return func
|
||||
|
||||
def _wrap(f):
|
||||
return f
|
||||
|
||||
return _wrap
|
||||
|
||||
m.tool = _tool_decorator # type: ignore[attr-defined]
|
||||
sys.modules["strands"] = m
|
||||
|
||||
if "strands.hooks" not in sys.modules:
|
||||
m = types.ModuleType("strands.hooks")
|
||||
for name in (
|
||||
"AfterToolCallEvent",
|
||||
"BeforeInvocationEvent",
|
||||
"BeforeToolCallEvent",
|
||||
"HookProvider",
|
||||
"HookRegistry",
|
||||
):
|
||||
setattr(m, name, type(name, (), {}))
|
||||
sys.modules["strands.hooks"] = m
|
||||
|
||||
if "strands.models" not in sys.modules:
|
||||
sys.modules["strands.models"] = types.ModuleType("strands.models")
|
||||
|
||||
if "strands.models.openai" not in sys.modules:
|
||||
m = types.ModuleType("strands.models.openai")
|
||||
m.OpenAIModel = _Permissive # type: ignore[attr-defined]
|
||||
sys.modules["strands.models.openai"] = m
|
||||
|
||||
if "uvicorn" not in sys.modules:
|
||||
m = types.ModuleType("uvicorn")
|
||||
m.run = lambda *a, **k: None # type: ignore[attr-defined]
|
||||
sys.modules["uvicorn"] = m
|
||||
|
||||
if "dotenv" not in sys.modules:
|
||||
m = types.ModuleType("dotenv")
|
||||
m.load_dotenv = lambda *a, **k: None # type: ignore[attr-defined]
|
||||
sys.modules["dotenv"] = m
|
||||
|
||||
|
||||
_install_stub_modules()
|
||||
@@ -0,0 +1,312 @@
|
||||
"""Red→green tests for the strands backend CVDIAG boundary instrumentation.
|
||||
|
||||
Exercises the REAL emit surface — every assertion reads the actual
|
||||
``CVDIAG {<json>}`` lines that ``_shared.cvdiag_bootstrap.emit_cvdiag`` writes
|
||||
to stdout (captured via ``capsys``), driven through the real
|
||||
``CvdiagBackendMiddleware`` and the real ``LlmCallScope`` / agent helpers. No
|
||||
mocks of the emit path.
|
||||
|
||||
What's covered (spec §3 / §5 / §6):
|
||||
* All 11 backend boundaries emit to stdout across the three request shapes
|
||||
that collectively exercise them (happy streaming, aborted stream, raised
|
||||
exception) for synthetic requests with ``CVDIAG_BACKEND_EMITTER=1`` (run at
|
||||
DEBUG tier so the verbose+debug boundaries are permitted).
|
||||
* PII scrub: a synthetic ``sk-test-12345`` in an exception message never
|
||||
appears in the emitted ``backend.error.caught`` JSON.
|
||||
* Heartbeat fires within ~12s of a slow-LLM simulation.
|
||||
* Default-OFF: with the flag unset, NO CVDIAG backend line is emitted.
|
||||
|
||||
RED before instrumentation: ``agents._cvdiag_backend`` does not exist →
|
||||
ImportError; the 11-boundary / heartbeat / scrub assertions cannot pass.
|
||||
GREEN after: every boundary, the scrub, and the heartbeat assert true.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Dict, List
|
||||
|
||||
import pytest
|
||||
from starlette.applications import Starlette
|
||||
from starlette.responses import StreamingResponse
|
||||
from starlette.routing import Route
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from agents._cvdiag_backend import (
|
||||
CvdiagBackendMiddleware,
|
||||
LlmCallScope,
|
||||
_RequestCtx,
|
||||
emit_agent_enter,
|
||||
emit_agent_exit,
|
||||
scrub,
|
||||
)
|
||||
|
||||
# The 11 backend boundaries (spec §5).
|
||||
ALL_BACKEND_BOUNDARIES = {
|
||||
"backend.request.ingress",
|
||||
"backend.agent.enter",
|
||||
"backend.llm.call.start",
|
||||
"backend.llm.call.heartbeat",
|
||||
"backend.llm.call.response",
|
||||
"backend.sse.first_byte",
|
||||
"backend.sse.event",
|
||||
"backend.sse.aborted",
|
||||
"backend.agent.exit",
|
||||
"backend.response.complete",
|
||||
"backend.error.caught",
|
||||
}
|
||||
|
||||
VALID_TEST_ID = "0190a9c0-1a2b-7c3d-8e4f-5a6b7c8d9e0f"
|
||||
|
||||
|
||||
def _parse_cvdiag_lines(captured: str) -> List[Dict]:
|
||||
"""Extract every ``CVDIAG {<json>}`` envelope line from captured stdout."""
|
||||
out: List[Dict] = []
|
||||
for line in captured.splitlines():
|
||||
if line.startswith("CVDIAG {"):
|
||||
out.append(json.loads(line[len("CVDIAG ") :]))
|
||||
return out
|
||||
|
||||
|
||||
def _boundaries(envelopes: List[Dict]) -> set:
|
||||
return {e["boundary"] for e in envelopes}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _debug_tier(monkeypatch):
|
||||
"""Run each test at DEBUG tier so verbose+debug boundaries are permitted.
|
||||
|
||||
``current_tier()`` is resolved once at bootstrap import; re-resolve it under
|
||||
a non-production env with ``CVDIAG_DEBUG=1`` so the §6 matrix lets
|
||||
``backend.sse.event`` (debug) and the verbose LLM boundaries through.
|
||||
"""
|
||||
import _shared.cvdiag_bootstrap as bootstrap
|
||||
|
||||
monkeypatch.setenv("SHOWCASE_ENV", "test")
|
||||
monkeypatch.setenv("CVDIAG_DEBUG", "1")
|
||||
bootstrap.setup({"SHOWCASE_ENV": "test", "CVDIAG_DEBUG": "1"})
|
||||
yield
|
||||
bootstrap.setup({"SHOWCASE_ENV": "test"})
|
||||
|
||||
|
||||
def _make_client(*, raise_server_exceptions: bool = True) -> TestClient:
|
||||
"""An app exposing three routes — happy stream, aborted stream, raise — each
|
||||
wrapped by the CVDIAG middleware. The endpoints emit the agent/LLM
|
||||
boundaries the middleware cannot observe, all keyed on the per-request ctx.
|
||||
"""
|
||||
|
||||
async def happy_stream(request):
|
||||
ctx = getattr(request.state, "cvdiag", None)
|
||||
if ctx is not None:
|
||||
emit_agent_enter(ctx, agent_name="showcase", model_id="gpt-4o-mini")
|
||||
|
||||
async def gen():
|
||||
if ctx is not None:
|
||||
async with LlmCallScope(
|
||||
ctx, provider="openai", model="gpt-4o-mini", interval_s=0.02
|
||||
):
|
||||
await asyncio.sleep(0.05) # let the heartbeat tick once
|
||||
yield b"data: hello\n\n"
|
||||
yield b"data: world\n\n"
|
||||
emit_agent_exit(ctx, terminal_outcome="ok", total_duration_ms=1)
|
||||
else:
|
||||
yield b"data: hello\n\n"
|
||||
|
||||
return StreamingResponse(gen(), media_type="text/event-stream")
|
||||
|
||||
async def raises(request):
|
||||
raise RuntimeError("upstream rejected key sk-test-12345 Bearer abc.def.ghi")
|
||||
|
||||
app = Starlette(
|
||||
routes=[
|
||||
Route("/", happy_stream, methods=["POST"]),
|
||||
Route("/boom", raises, methods=["POST"]),
|
||||
]
|
||||
)
|
||||
app.add_middleware(CvdiagBackendMiddleware)
|
||||
return TestClient(app, raise_server_exceptions=raise_server_exceptions)
|
||||
|
||||
|
||||
async def _drive_abort() -> None:
|
||||
"""Drive the CVDIAG middleware over an unbounded stream and disconnect.
|
||||
|
||||
Builds the middleware around an unbounded inner stream, calls ``dispatch``
|
||||
to get the wrapped ``body_iterator``, reads one chunk, then ``aclose()``s it
|
||||
— the deterministic equivalent of a client disconnecting mid-stream. This
|
||||
raises ``GeneratorExit`` into the wrapper → ``backend.sse.aborted``.
|
||||
"""
|
||||
from starlette.requests import Request
|
||||
|
||||
async def unbounded():
|
||||
i = 0
|
||||
while True:
|
||||
yield f"data: chunk-{i}\n\n".encode()
|
||||
i += 1
|
||||
|
||||
inner_response = StreamingResponse(unbounded(), media_type="text/event-stream")
|
||||
|
||||
async def call_next(_request):
|
||||
return inner_response
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/",
|
||||
"headers": [(b"x-aimock-context", b"strands")],
|
||||
"query_string": b"",
|
||||
}
|
||||
|
||||
async def receive():
|
||||
return {"type": "http.request", "body": b""}
|
||||
|
||||
mw = CvdiagBackendMiddleware(app=lambda *a: None)
|
||||
request = Request(scope, receive)
|
||||
wrapped = await mw.dispatch(request, call_next)
|
||||
|
||||
body = wrapped.body_iterator
|
||||
await body.__anext__() # first chunk
|
||||
await body.aclose() # client disconnect mid-stream
|
||||
|
||||
|
||||
def test_all_eleven_backend_boundaries_emit(monkeypatch, capsys):
|
||||
"""All 11 backend boundaries emit across the three request shapes.
|
||||
|
||||
The happy stream yields ingress / agent.enter / llm.* / sse.first_byte /
|
||||
sse.event / agent.exit / response.complete; a disconnected stream yields
|
||||
sse.aborted; the raising route yields error.caught. Their union is the full
|
||||
eleven.
|
||||
"""
|
||||
monkeypatch.setenv("CVDIAG_BACKEND_EMITTER", "1")
|
||||
client = _make_client(raise_server_exceptions=False)
|
||||
|
||||
headers = {"x-test-id": VALID_TEST_ID, "x-aimock-context": "strands"}
|
||||
resp = client.post("/", headers=headers)
|
||||
assert resp.status_code == 200
|
||||
|
||||
# Client-disconnect abort surface (→ backend.sse.aborted), driven directly
|
||||
# because Starlette's sync TestClient cannot reliably tear a stream down
|
||||
# mid-flight.
|
||||
asyncio.run(_drive_abort())
|
||||
|
||||
client.post("/boom", headers=headers)
|
||||
|
||||
envelopes = _parse_cvdiag_lines(capsys.readouterr().out)
|
||||
seen = _boundaries(envelopes)
|
||||
|
||||
missing = ALL_BACKEND_BOUNDARIES - seen
|
||||
assert not missing, (
|
||||
f"missing backend boundaries: {sorted(missing)}; saw {sorted(seen)}"
|
||||
)
|
||||
|
||||
# Correlation: every backend envelope carries the slug. The header-bearing
|
||||
# HTTP requests forward x-test-id verbatim; the directly driven abort
|
||||
# request mints its own UUIDv7 (no inbound header). Assert the forwarded
|
||||
# test_id appears on the header-bearing envelopes, and every minted id is a
|
||||
# well-formed UUIDv7.
|
||||
backend = [e for e in envelopes if e["layer"] == "backend"]
|
||||
assert backend, "no backend-layer envelopes emitted"
|
||||
assert all(e["slug"] == "strands" for e in backend)
|
||||
forwarded = [e for e in backend if e["test_id"] == VALID_TEST_ID]
|
||||
assert forwarded, "forwarded x-test-id never appeared on any backend envelope"
|
||||
uuid7_re = __import__("re").compile(
|
||||
r"^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$"
|
||||
)
|
||||
assert all(uuid7_re.match(e["test_id"]) for e in backend)
|
||||
# Closed 9-key edge-header bag always present on a header-bearing ingress.
|
||||
ingress = next(
|
||||
e
|
||||
for e in backend
|
||||
if e["boundary"] == "backend.request.ingress" and e["test_id"] == VALID_TEST_ID
|
||||
)
|
||||
assert set(ingress["edge_headers"].keys()) == {
|
||||
"cf-ray",
|
||||
"cf-mitigated",
|
||||
"cf-cache-status",
|
||||
"x-railway-edge",
|
||||
"x-railway-request-id",
|
||||
"x-hikari-trace",
|
||||
"retry-after",
|
||||
"via",
|
||||
"server",
|
||||
}
|
||||
|
||||
|
||||
def test_error_caught_scrubs_secret(monkeypatch, capsys):
|
||||
"""A synthetic ``sk-test-12345`` in an exception never reaches the emitted
|
||||
``backend.error.caught`` envelope."""
|
||||
monkeypatch.setenv("CVDIAG_BACKEND_EMITTER", "1")
|
||||
client = _make_client(raise_server_exceptions=False)
|
||||
|
||||
client.post("/boom", headers={"x-aimock-context": "strands"})
|
||||
|
||||
out = capsys.readouterr().out
|
||||
envelopes = _parse_cvdiag_lines(out)
|
||||
errs = [e for e in envelopes if e["boundary"] == "backend.error.caught"]
|
||||
assert errs, "backend.error.caught not emitted"
|
||||
err = errs[0]
|
||||
assert err["metadata"]["exception_type"] == "RuntimeError"
|
||||
blob = json.dumps(err)
|
||||
assert "sk-test-12345" not in blob, "raw secret leaked into error envelope"
|
||||
assert "Bearer abc" not in blob, "raw bearer token leaked into error envelope"
|
||||
assert "[REDACTED]" in err["metadata"]["message_scrubbed"]
|
||||
|
||||
|
||||
def test_scrub_helper_redacts_known_secret_shapes():
|
||||
"""Unit-level: the scrub helper redacts bearer/sk-/pk-/userinfo shapes."""
|
||||
assert "sk-test-12345" not in scrub("key sk-test-12345 here")
|
||||
assert "sk-abcdefghijklmnopqrstuvwx" not in scrub("sk-abcdefghijklmnopqrstuvwx")
|
||||
assert "Bearer secrettoken" not in scrub("auth Bearer secrettoken")
|
||||
assert "pw" not in scrub("https://user:pw@host/path")
|
||||
assert scrub(None) == ""
|
||||
|
||||
|
||||
def test_heartbeat_fires_within_window(monkeypatch, capsys):
|
||||
"""``backend.llm.call.heartbeat`` fires while a slow LLM call is outstanding.
|
||||
|
||||
Uses a short interval so the test is fast; the production interval is ~10s
|
||||
and the spec requires a heartbeat within ~12s of a slow-LLM simulation —
|
||||
proven here by the same code path firing within its interval.
|
||||
"""
|
||||
monkeypatch.setenv("CVDIAG_BACKEND_EMITTER", "1")
|
||||
|
||||
async def run():
|
||||
ctx = _RequestCtx(test_id=VALID_TEST_ID, slug="strands", demo="default")
|
||||
async with LlmCallScope(ctx, provider="openai", model="m", interval_s=0.05):
|
||||
await asyncio.sleep(0.18) # ~3 heartbeat intervals
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
envelopes = _parse_cvdiag_lines(capsys.readouterr().out)
|
||||
hb = [e for e in envelopes if e["boundary"] == "backend.llm.call.heartbeat"]
|
||||
assert hb, "no heartbeat emitted during a slow LLM call"
|
||||
assert all("elapsed_ms_since_start" in e["metadata"] for e in hb)
|
||||
|
||||
|
||||
def test_sse_aborted_on_client_disconnect(monkeypatch, capsys):
|
||||
"""Tearing the response stream down mid-flight emits ``backend.sse.aborted``
|
||||
with a ``termination_kind`` and the bytes streamed before the abort."""
|
||||
monkeypatch.setenv("CVDIAG_BACKEND_EMITTER", "1")
|
||||
|
||||
asyncio.run(_drive_abort())
|
||||
|
||||
envelopes = _parse_cvdiag_lines(capsys.readouterr().out)
|
||||
aborts = [e for e in envelopes if e["boundary"] == "backend.sse.aborted"]
|
||||
assert aborts, "backend.sse.aborted not emitted on client disconnect"
|
||||
meta = aborts[0]["metadata"]
|
||||
assert meta["termination_kind"] in {"rst", "timeout", "chunk_error"}
|
||||
assert meta["bytes_before_abort"] > 0
|
||||
# A disconnected stream must NOT also report a clean response.complete.
|
||||
completes = [e for e in envelopes if e["boundary"] == "backend.response.complete"]
|
||||
assert not completes, "clean response.complete emitted for an aborted stream"
|
||||
|
||||
|
||||
def test_disabled_by_default_emits_nothing(monkeypatch, capsys):
|
||||
"""With ``CVDIAG_BACKEND_EMITTER`` unset, NO backend CVDIAG line is emitted."""
|
||||
monkeypatch.delenv("CVDIAG_BACKEND_EMITTER", raising=False)
|
||||
client = _make_client()
|
||||
client.post("/", headers={"x-aimock-context": "strands"})
|
||||
|
||||
envelopes = _parse_cvdiag_lines(capsys.readouterr().out)
|
||||
backend = [e for e in envelopes if e["layer"] == "backend"]
|
||||
assert backend == [], f"emitter fired while disabled: {backend}"
|
||||
@@ -0,0 +1,351 @@
|
||||
"""Tests for the hardened ``generate_a2ui`` error-handling surface.
|
||||
|
||||
Mirrors the google-adk sibling agent's hardening pattern: every failure
|
||||
branch returns a structured ``{error, message, remediation}`` dict
|
||||
(JSON-serialized, since the strands tool returns a string) instead of
|
||||
letting raw OpenAI exceptions bubble up through the strands tool
|
||||
machinery.
|
||||
|
||||
Covers:
|
||||
* OpenAI APIError / RateLimitError / APIConnectionError / AuthenticationError
|
||||
* Empty ``response.choices``
|
||||
* Empty / missing ``tool_calls[0]``
|
||||
* Malformed ``json.loads(tool_call.function.arguments)``
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
# ---- Fake ``openai`` module ---------------------------------------------
|
||||
#
|
||||
# The real ``openai`` package may not be installed in the test venv. We
|
||||
# install a stub that provides the exception classes used in the except
|
||||
# branches and an ``OpenAI`` class whose ``chat.completions.create`` we
|
||||
# patch per-test via monkeypatch.
|
||||
|
||||
|
||||
def _install_openai_stub():
|
||||
if "openai" in sys.modules:
|
||||
return
|
||||
m = types.ModuleType("openai")
|
||||
|
||||
class OpenAIError(Exception):
|
||||
"""Base class for all openai-SDK errors. Matches the real SDK's hierarchy."""
|
||||
|
||||
pass
|
||||
|
||||
class APIError(OpenAIError):
|
||||
pass
|
||||
|
||||
class RateLimitError(APIError):
|
||||
pass
|
||||
|
||||
class APIConnectionError(APIError):
|
||||
pass
|
||||
|
||||
class AuthenticationError(APIError):
|
||||
pass
|
||||
|
||||
# Placeholder OpenAI client; tests replace ``OpenAI`` on the module
|
||||
# with a class that returns whatever ``chat.completions.create`` we want.
|
||||
class OpenAI:
|
||||
def __init__(self, *a, **kw):
|
||||
raise AssertionError("tests must override openai.OpenAI")
|
||||
|
||||
m.OpenAIError = OpenAIError
|
||||
m.APIError = APIError
|
||||
m.RateLimitError = RateLimitError
|
||||
m.APIConnectionError = APIConnectionError
|
||||
m.AuthenticationError = AuthenticationError
|
||||
m.OpenAI = OpenAI
|
||||
sys.modules["openai"] = m
|
||||
|
||||
|
||||
def _install_httpx_stub():
|
||||
if "httpx" in sys.modules:
|
||||
return
|
||||
m = types.ModuleType("httpx")
|
||||
|
||||
class HTTPError(Exception):
|
||||
"""Base class for all httpx transport errors."""
|
||||
|
||||
pass
|
||||
|
||||
class ConnectError(HTTPError):
|
||||
pass
|
||||
|
||||
class ReadTimeout(HTTPError):
|
||||
pass
|
||||
|
||||
m.HTTPError = HTTPError
|
||||
m.ConnectError = ConnectError
|
||||
m.ReadTimeout = ReadTimeout
|
||||
sys.modules["httpx"] = m
|
||||
|
||||
|
||||
_install_openai_stub()
|
||||
_install_httpx_stub()
|
||||
|
||||
|
||||
# ---- Helpers ------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_fake_openai_client(*, create_behavior):
|
||||
"""Build a fake ``OpenAI`` class whose ``chat.completions.create``
|
||||
invokes ``create_behavior(**kwargs)``.
|
||||
|
||||
``create_behavior`` may either raise (to simulate an API failure) or
|
||||
return a ``SimpleNamespace`` standing in for the OpenAI response.
|
||||
"""
|
||||
|
||||
class _FakeCompletions:
|
||||
def create(self, **kwargs):
|
||||
return create_behavior(**kwargs)
|
||||
|
||||
class _FakeChat:
|
||||
def __init__(self):
|
||||
self.completions = _FakeCompletions()
|
||||
|
||||
class _FakeOpenAI:
|
||||
def __init__(self, *a, **kw):
|
||||
self.chat = _FakeChat()
|
||||
|
||||
return _FakeOpenAI
|
||||
|
||||
|
||||
def _response_with_tool_args(args_json: str):
|
||||
"""Build a fake OpenAI response whose first tool call has the given JSON args."""
|
||||
tool_call = SimpleNamespace(function=SimpleNamespace(arguments=args_json))
|
||||
message = SimpleNamespace(tool_calls=[tool_call])
|
||||
choice = SimpleNamespace(message=message)
|
||||
return SimpleNamespace(choices=[choice])
|
||||
|
||||
|
||||
def _response_with_no_tool_calls():
|
||||
message = SimpleNamespace(tool_calls=[])
|
||||
choice = SimpleNamespace(message=message)
|
||||
return SimpleNamespace(choices=[choice])
|
||||
|
||||
|
||||
def _response_with_no_choices():
|
||||
return SimpleNamespace(choices=[])
|
||||
|
||||
|
||||
def _invoke_generate_a2ui(context: str = "test context"):
|
||||
"""Call ``generate_a2ui`` on the module and return the parsed result.
|
||||
|
||||
The ``@tool`` decorator in the conftest stub is a no-op, so the
|
||||
underlying function is directly callable.
|
||||
"""
|
||||
from agents.agent import generate_a2ui
|
||||
|
||||
raw = generate_a2ui(context)
|
||||
return json.loads(raw)
|
||||
|
||||
|
||||
# ---- Tests --------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"exc_name",
|
||||
["APIError", "RateLimitError", "APIConnectionError", "AuthenticationError"],
|
||||
)
|
||||
def test_openai_exceptions_return_structured_error(monkeypatch, exc_name):
|
||||
"""OpenAI exception subclasses must be caught and returned as a
|
||||
structured ``{error, message, remediation}`` payload, not raised.
|
||||
|
||||
We construct a subclass of the real openai exception class that
|
||||
accepts a bare message — the real classes have varied/restrictive
|
||||
constructors (e.g. ``AuthenticationError`` requires ``response`` +
|
||||
``body`` kwargs). The subclass keeps ``isinstance(exc, openai.APIError)``
|
||||
true so the except branch in ``generate_a2ui`` catches it the same
|
||||
way, while letting the test instantiate without provider SDK internals.
|
||||
"""
|
||||
import openai
|
||||
|
||||
base_cls = getattr(openai, exc_name)
|
||||
|
||||
# Build a lightweight subclass that carries the bare string message.
|
||||
TestExc = type(
|
||||
f"Test{exc_name}",
|
||||
(base_cls,),
|
||||
{
|
||||
"__init__": lambda self, msg: Exception.__init__(self, msg),
|
||||
},
|
||||
)
|
||||
|
||||
def _raise(**_kwargs):
|
||||
raise TestExc(f"simulated {exc_name}")
|
||||
|
||||
FakeOpenAI = _make_fake_openai_client(create_behavior=_raise)
|
||||
monkeypatch.setattr(openai, "OpenAI", FakeOpenAI)
|
||||
|
||||
result = _invoke_generate_a2ui()
|
||||
|
||||
assert result["error"] == "a2ui_llm_error"
|
||||
# The message carries the subclass name, which includes the parent
|
||||
# exception name as a substring.
|
||||
assert exc_name in result["message"]
|
||||
assert "remediation" in result
|
||||
assert "OPENAI_API_KEY" in result["remediation"]
|
||||
|
||||
|
||||
def test_empty_choices_returns_structured_error(monkeypatch):
|
||||
import openai
|
||||
|
||||
FakeOpenAI = _make_fake_openai_client(
|
||||
create_behavior=lambda **_: _response_with_no_choices()
|
||||
)
|
||||
monkeypatch.setattr(openai, "OpenAI", FakeOpenAI)
|
||||
|
||||
result = _invoke_generate_a2ui()
|
||||
|
||||
assert result["error"] == "a2ui_empty_response"
|
||||
assert "no choices" in result["message"].lower()
|
||||
assert result["remediation"]
|
||||
|
||||
|
||||
def test_missing_tool_calls_returns_structured_error(monkeypatch):
|
||||
import openai
|
||||
|
||||
FakeOpenAI = _make_fake_openai_client(
|
||||
create_behavior=lambda **_: _response_with_no_tool_calls()
|
||||
)
|
||||
monkeypatch.setattr(openai, "OpenAI", FakeOpenAI)
|
||||
|
||||
result = _invoke_generate_a2ui()
|
||||
|
||||
assert result["error"] == "a2ui_no_tool_call"
|
||||
assert "render_a2ui" in result["message"]
|
||||
assert result["remediation"]
|
||||
|
||||
|
||||
def test_malformed_tool_args_returns_structured_error(monkeypatch):
|
||||
import openai
|
||||
|
||||
FakeOpenAI = _make_fake_openai_client(
|
||||
create_behavior=lambda **_: _response_with_tool_args("{not valid json"),
|
||||
)
|
||||
monkeypatch.setattr(openai, "OpenAI", FakeOpenAI)
|
||||
|
||||
result = _invoke_generate_a2ui()
|
||||
|
||||
assert result["error"] == "a2ui_invalid_arguments"
|
||||
assert (
|
||||
"parse" in result["message"].lower() or "arguments" in result["message"].lower()
|
||||
)
|
||||
assert result["remediation"]
|
||||
|
||||
|
||||
def test_openai_base_error_returns_structured(monkeypatch):
|
||||
"""``OpenAIError`` is the base class the SDK raises from the
|
||||
``OpenAI()`` constructor when the API key is missing or malformed —
|
||||
it is NOT a subclass of ``APIError``. The except clause must catch
|
||||
``OpenAIError`` (or broader) so config-time failures become a
|
||||
structured tool result, not an uncaught exception.
|
||||
"""
|
||||
import openai
|
||||
|
||||
class _ConfigError(openai.OpenAIError):
|
||||
def __init__(self, msg):
|
||||
Exception.__init__(self, msg)
|
||||
|
||||
def _raise(**_kwargs):
|
||||
raise _ConfigError("simulated missing/malformed API key")
|
||||
|
||||
FakeOpenAI = _make_fake_openai_client(create_behavior=_raise)
|
||||
monkeypatch.setattr(openai, "OpenAI", FakeOpenAI)
|
||||
|
||||
result = _invoke_generate_a2ui()
|
||||
|
||||
assert result["error"] == "a2ui_llm_error"
|
||||
assert "ConfigError" in result["message"] or "OpenAIError" in result["message"]
|
||||
assert "OPENAI_API_KEY" in result["remediation"]
|
||||
|
||||
|
||||
def test_openai_constructor_openai_error_caught(monkeypatch):
|
||||
"""Analog of the above but the ``OpenAIError`` is raised from the
|
||||
``OpenAI()`` constructor itself (missing env var path). The client
|
||||
construction must sit inside the try block; otherwise the error
|
||||
bypasses the except clause and escapes.
|
||||
"""
|
||||
import openai
|
||||
|
||||
class _ConstructorError(openai.OpenAIError):
|
||||
def __init__(self, msg):
|
||||
Exception.__init__(self, msg)
|
||||
|
||||
class _FailingOpenAI:
|
||||
def __init__(self, *a, **kw):
|
||||
raise _ConstructorError("OPENAI_API_KEY must be set")
|
||||
|
||||
monkeypatch.setattr(openai, "OpenAI", _FailingOpenAI)
|
||||
|
||||
result = _invoke_generate_a2ui()
|
||||
|
||||
assert result["error"] == "a2ui_llm_error"
|
||||
assert "ConstructorError" in result["message"] or "OpenAIError" in result["message"]
|
||||
|
||||
|
||||
def test_httpx_transport_error_returns_structured(monkeypatch):
|
||||
"""``httpx.HTTPError`` (and subclasses like ``ConnectError`` /
|
||||
``ReadTimeout``) can escape below the OpenAI SDK's wrap layer in some
|
||||
failure modes. The except clause must catch them so transport
|
||||
failures surface as a structured tool result.
|
||||
"""
|
||||
import httpx
|
||||
import openai
|
||||
|
||||
def _raise(**_kwargs):
|
||||
raise httpx.ConnectError("simulated DNS failure")
|
||||
|
||||
FakeOpenAI = _make_fake_openai_client(create_behavior=_raise)
|
||||
monkeypatch.setattr(openai, "OpenAI", FakeOpenAI)
|
||||
|
||||
result = _invoke_generate_a2ui()
|
||||
|
||||
assert result["error"] == "a2ui_llm_error"
|
||||
assert "ConnectError" in result["message"]
|
||||
assert result["remediation"]
|
||||
|
||||
|
||||
def test_happy_path_returns_a2ui_operations(monkeypatch):
|
||||
"""Sanity check: a well-formed response goes through
|
||||
``build_a2ui_operations_from_tool_call`` and returns a non-error payload."""
|
||||
import openai
|
||||
|
||||
valid_args = json.dumps(
|
||||
{
|
||||
"surfaceId": "test-surface",
|
||||
"catalogId": "test-catalog",
|
||||
"components": [],
|
||||
"data": {},
|
||||
}
|
||||
)
|
||||
FakeOpenAI = _make_fake_openai_client(
|
||||
create_behavior=lambda **_: _response_with_tool_args(valid_args),
|
||||
)
|
||||
monkeypatch.setattr(openai, "OpenAI", FakeOpenAI)
|
||||
|
||||
# Stub build_a2ui_operations_from_tool_call to return a marker payload
|
||||
# so we don't depend on the shared tool's real implementation shape.
|
||||
import agents.agent as agent_mod
|
||||
|
||||
monkeypatch.setattr(
|
||||
agent_mod,
|
||||
"build_a2ui_operations_from_tool_call",
|
||||
lambda args: {"a2ui_marker": True, "surfaceId": args.get("surfaceId")},
|
||||
)
|
||||
|
||||
result = _invoke_generate_a2ui()
|
||||
|
||||
assert "error" not in result
|
||||
assert result.get("a2ui_marker") is True
|
||||
assert result.get("surfaceId") == "test-surface"
|
||||
@@ -0,0 +1,298 @@
|
||||
"""Tests for _HookInjectingAgentDict in src/agents/agent.py.
|
||||
|
||||
Verifies:
|
||||
* hook is injected when an Agent is inserted via ``__setitem__``,
|
||||
``update()``, ``setdefault()``, and ``|=`` (``__ior__``),
|
||||
* existing entries are preserved when the factory swaps in the dict,
|
||||
* no double-injection on re-insert of the same thread_id.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class _FakeHookRegistry:
|
||||
"""Minimal stand-in for strands' HookRegistry exposing what the cap hook uses."""
|
||||
|
||||
def __init__(self):
|
||||
self._hook_providers = []
|
||||
self._callbacks = []
|
||||
|
||||
def add_hook(self, provider):
|
||||
self._hook_providers.append(provider)
|
||||
provider.register_hooks(self)
|
||||
|
||||
def add_callback(self, event_cls, cb):
|
||||
self._callbacks.append((event_cls, cb))
|
||||
|
||||
|
||||
class _FakeAgent:
|
||||
"""Duck-typed stand-in for strands.Agent — must pass isinstance(Agent) check.
|
||||
|
||||
We monkey-patch ``agents.agent.Agent`` in each test to our fake class so
|
||||
``isinstance(value, Agent)`` inside the dict routes correctly.
|
||||
"""
|
||||
|
||||
def __init__(self, label: str = "", **kwargs):
|
||||
self.label = label
|
||||
self.hooks = _FakeHookRegistry()
|
||||
# Accept (and stash) whatever kwargs the real ``strands.Agent``
|
||||
# accepts (``model``, ``system_prompt``, ``tools``, ...). Tests
|
||||
# don't inspect these — the point is to let factory code that
|
||||
# calls ``Agent(model=..., tools=[...])`` construct this fake
|
||||
# without a TypeError.
|
||||
self.kwargs = kwargs
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patched_agent(monkeypatch):
|
||||
"""Swap ``agents.agent.Agent`` for ``_FakeAgent`` for the duration of the test."""
|
||||
import agents.agent as agent_mod
|
||||
|
||||
monkeypatch.setattr(agent_mod, "Agent", _FakeAgent)
|
||||
return agent_mod
|
||||
|
||||
|
||||
def _count_cap_hooks(agent, cap_hook_cls) -> int:
|
||||
return sum(1 for p in agent.hooks._hook_providers if isinstance(p, cap_hook_cls))
|
||||
|
||||
|
||||
def test_setitem_injects_hook(patched_agent):
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
a = _FakeAgent("t1")
|
||||
|
||||
d["thread-1"] = a
|
||||
|
||||
assert _count_cap_hooks(a, patched_agent._ToolCallCapHook) == 1
|
||||
|
||||
|
||||
def test_update_injects_hook(patched_agent):
|
||||
"""``dict.update`` bypasses ``__setitem__`` in CPython's bulk path;
|
||||
the override must still run injection."""
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
a, b = _FakeAgent("a"), _FakeAgent("b")
|
||||
|
||||
d.update({"thread-a": a, "thread-b": b})
|
||||
|
||||
assert _count_cap_hooks(a, patched_agent._ToolCallCapHook) == 1
|
||||
assert _count_cap_hooks(b, patched_agent._ToolCallCapHook) == 1
|
||||
|
||||
|
||||
def test_update_with_kwargs_injects_hook(patched_agent):
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
a = _FakeAgent("kw")
|
||||
|
||||
d.update(threadk=a)
|
||||
|
||||
assert _count_cap_hooks(a, patched_agent._ToolCallCapHook) == 1
|
||||
|
||||
|
||||
def test_update_with_iterable_of_pairs_injects_hook(patched_agent):
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
a = _FakeAgent("p")
|
||||
|
||||
d.update([("thread-p", a)])
|
||||
|
||||
assert _count_cap_hooks(a, patched_agent._ToolCallCapHook) == 1
|
||||
|
||||
|
||||
def test_setdefault_injects_hook(patched_agent):
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
a = _FakeAgent("sd")
|
||||
|
||||
d.setdefault("thread-sd", a)
|
||||
|
||||
assert _count_cap_hooks(a, patched_agent._ToolCallCapHook) == 1
|
||||
|
||||
|
||||
def test_setdefault_existing_key_skips_default(patched_agent):
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
first = _FakeAgent("first")
|
||||
second = _FakeAgent("second")
|
||||
|
||||
d["x"] = first
|
||||
result = d.setdefault("x", second)
|
||||
|
||||
# setdefault returns the existing value and never inserts second.
|
||||
assert result is first
|
||||
assert _count_cap_hooks(second, patched_agent._ToolCallCapHook) == 0
|
||||
|
||||
|
||||
def test_ior_injects_hook(patched_agent):
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
a = _FakeAgent("ior")
|
||||
|
||||
d |= {"thread-ior": a}
|
||||
|
||||
assert _count_cap_hooks(a, patched_agent._ToolCallCapHook) == 1
|
||||
|
||||
|
||||
def test_existing_entries_preserved_on_wrap(patched_agent):
|
||||
"""When ``build_showcase_agent`` copies the original dict into the
|
||||
injecting dict, pre-existing entries must survive (and gain the hook)."""
|
||||
original = {"preexisting-thread": _FakeAgent("pre")}
|
||||
hook_dict = patched_agent._HookInjectingAgentDict()
|
||||
hook_dict.update(original)
|
||||
|
||||
assert "preexisting-thread" in hook_dict
|
||||
assert hook_dict["preexisting-thread"].label == "pre"
|
||||
assert (
|
||||
_count_cap_hooks(
|
||||
hook_dict["preexisting-thread"], patched_agent._ToolCallCapHook
|
||||
)
|
||||
== 1
|
||||
)
|
||||
|
||||
|
||||
def test_no_double_injection_on_reinsert(patched_agent):
|
||||
"""Re-inserting the same agent for the same thread_id must NOT add a
|
||||
second cap hook (otherwise the effective cap would be halved)."""
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
a = _FakeAgent("re")
|
||||
|
||||
d["thread-re"] = a
|
||||
d["thread-re"] = a # re-insert same agent
|
||||
|
||||
assert _count_cap_hooks(a, patched_agent._ToolCallCapHook) == 1
|
||||
|
||||
|
||||
def test_no_double_injection_on_bulk_reinsert(patched_agent):
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
a = _FakeAgent("bulk")
|
||||
d["t"] = a
|
||||
d.update({"t": a})
|
||||
d.setdefault("t", a)
|
||||
assert _count_cap_hooks(a, patched_agent._ToolCallCapHook) == 1
|
||||
|
||||
|
||||
def test_update_with_dict_items_view(patched_agent):
|
||||
"""``dict.items()`` is a ``Mapping``-like view, but iterating it yields
|
||||
``(k, v)`` pairs (not keys). The ``update`` override must handle this
|
||||
input shape — otherwise ``.items()`` would fall through to the
|
||||
pair-iterable branch and work, but we want an explicit assertion.
|
||||
|
||||
Concretely: strands / ag_ui_strands can legitimately pass a
|
||||
``dict_items`` view (e.g. filtering a source dict). Injection must
|
||||
still fire for each contained Agent.
|
||||
"""
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
a, b = _FakeAgent("iv-a"), _FakeAgent("iv-b")
|
||||
source = {"thread-iv-a": a, "thread-iv-b": b}
|
||||
|
||||
d.update(source.items())
|
||||
|
||||
assert _count_cap_hooks(a, patched_agent._ToolCallCapHook) == 1
|
||||
assert _count_cap_hooks(b, patched_agent._ToolCallCapHook) == 1
|
||||
assert d["thread-iv-a"] is a
|
||||
assert d["thread-iv-b"] is b
|
||||
|
||||
|
||||
def test_update_with_mapping_subtype(patched_agent):
|
||||
"""``collections.ChainMap`` is a ``collections.abc.Mapping`` subtype.
|
||||
The ``update`` override must correctly route it through the Mapping
|
||||
branch so every contained Agent gets a cap hook attached.
|
||||
|
||||
The assertions pin correctness only: every value in the chain lands
|
||||
in the injecting dict with exactly one cap hook.
|
||||
"""
|
||||
from collections import ChainMap
|
||||
|
||||
d = patched_agent._HookInjectingAgentDict()
|
||||
a1, a2 = _FakeAgent("m-a1"), _FakeAgent("m-a2")
|
||||
primary = {"thread-a1": a1}
|
||||
fallback = {"thread-a2": a2}
|
||||
cm = ChainMap(primary, fallback)
|
||||
|
||||
d.update(cm)
|
||||
|
||||
assert "thread-a1" in d
|
||||
assert "thread-a2" in d
|
||||
assert d["thread-a1"] is a1
|
||||
assert d["thread-a2"] is a2
|
||||
assert _count_cap_hooks(a1, patched_agent._ToolCallCapHook) == 1
|
||||
assert _count_cap_hooks(a2, patched_agent._ToolCallCapHook) == 1
|
||||
|
||||
|
||||
def test_build_showcase_agent_swaps_hook_dict(monkeypatch, patched_agent):
|
||||
"""Factory integration: ``build_showcase_agent()`` must replace the
|
||||
``StrandsAgent._agents_by_thread`` dict with ``_HookInjectingAgentDict``,
|
||||
preserve any pre-existing entries, and ensure every entry has a cap
|
||||
hook attached.
|
||||
|
||||
The conftest stubs out ``StrandsAgent`` / ``StrandsAgentConfig`` /
|
||||
``ToolBehavior`` as permissive classes. We patch ``StrandsAgent`` to
|
||||
seed one pre-existing entry in ``_agents_by_thread`` during
|
||||
construction, so the factory's copy-and-wrap logic is actually
|
||||
exercised.
|
||||
"""
|
||||
import agents.agent as agent_mod
|
||||
|
||||
# Pre-existing Agent (with a FakeAgent stand-in that matches the
|
||||
# isinstance check in ``_HookInjectingAgentDict.__setitem__``).
|
||||
preexisting_agent = _FakeAgent("pre")
|
||||
|
||||
class _SeededStrandsAgent:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
# Emulate ag_ui_strands seeding the dict in ``__init__``.
|
||||
self._agents_by_thread = {"preexisting-thread": preexisting_agent}
|
||||
|
||||
# Patch the ``StrandsAgent`` reference bound in the ``agents.agent``
|
||||
# module (not the source in ``ag_ui_strands``). The module already
|
||||
# captured the original class at import time — patching the source
|
||||
# module would have no effect on the factory's call site.
|
||||
monkeypatch.setattr(agent_mod, "StrandsAgent", _SeededStrandsAgent)
|
||||
|
||||
# The factory calls ``_build_model`` which requires OPENAI_API_KEY.
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "test-key-for-factory")
|
||||
|
||||
# Ensure Agent isinstance checks inside the dict succeed for our fake.
|
||||
# ``patched_agent`` already swapped ``agents.agent.Agent`` → _FakeAgent.
|
||||
|
||||
from agents.agent import (
|
||||
_HookInjectingAgentDict,
|
||||
_ToolCallCapHook,
|
||||
build_showcase_agent,
|
||||
)
|
||||
|
||||
agui_agent = build_showcase_agent()
|
||||
|
||||
# 1. The per-thread dict is the hook-injecting variant.
|
||||
assert isinstance(agui_agent._agents_by_thread, _HookInjectingAgentDict)
|
||||
|
||||
# 2. Pre-existing entries survived the swap.
|
||||
assert "preexisting-thread" in agui_agent._agents_by_thread
|
||||
assert agui_agent._agents_by_thread["preexisting-thread"] is preexisting_agent
|
||||
|
||||
# 3. Every surviving entry has a cap hook attached.
|
||||
for agent in agui_agent._agents_by_thread.values():
|
||||
assert _count_cap_hooks(agent, _ToolCallCapHook) == 1
|
||||
|
||||
|
||||
def test_agent_has_cap_hook_uses_sentinel_not_private_attrs(patched_agent):
|
||||
"""``_agent_has_cap_hook`` must check a sentinel attribute we own,
|
||||
NOT spelunk HookRegistry privates. If an upstream ``HookRegistry``
|
||||
rename drops ``_hook_providers`` / ``hook_providers``, double-injection
|
||||
would silently return — which halves the effective cap.
|
||||
|
||||
We simulate the rename by constructing a registry WITHOUT those
|
||||
attributes but WITH the sentinel, and assert ``_agent_has_cap_hook``
|
||||
still returns True.
|
||||
"""
|
||||
from agents.agent import _agent_has_cap_hook, _CAP_HOOK_SENTINEL_ATTR
|
||||
|
||||
class _RegistryWithoutPrivates:
|
||||
# Deliberately missing _hook_providers AND hook_providers.
|
||||
pass
|
||||
|
||||
agent = _FakeAgent("sentinel")
|
||||
agent.hooks = _RegistryWithoutPrivates()
|
||||
# Without the sentinel, no cap hook is known.
|
||||
assert not _agent_has_cap_hook(agent)
|
||||
|
||||
# With the sentinel, the check must return True regardless of what
|
||||
# HookRegistry looks like internally.
|
||||
setattr(agent, _CAP_HOOK_SENTINEL_ATTR, True)
|
||||
assert _agent_has_cap_hook(agent)
|
||||
@@ -0,0 +1,398 @@
|
||||
"""Tests for the ThreadingInstrumentor patch installed by agent_server.py.
|
||||
|
||||
The patch exists because strands-agents calls
|
||||
``ThreadingInstrumentor().instrument()`` at Tracer construction time, which
|
||||
recursively wraps ThreadPoolExecutor.submit and triggers RecursionError
|
||||
during tool-rendering requests.
|
||||
|
||||
We verify:
|
||||
* ``ThreadingInstrumentor.instrument`` has been replaced,
|
||||
* the replacement returns ``self`` (not ``None``) so fluent callers
|
||||
don't AttributeError,
|
||||
* calling ``.instrument()`` is a no-op (doesn't actually wrap anything),
|
||||
* the import-order guard is implemented as ``if/raise`` (NOT ``assert``)
|
||||
so it survives ``python -O``,
|
||||
* pre-imported strands is detected and raises.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _install_patch_without_imports():
|
||||
"""Apply the same patch agent_server.py applies, without importing
|
||||
ag_ui_strands / strands (which aren't available in the local venv).
|
||||
|
||||
Mirrors the logic in ``agent_server.py`` verbatim so this test catches
|
||||
regressions where the two diverge.
|
||||
"""
|
||||
from opentelemetry.instrumentation.threading import ThreadingInstrumentor
|
||||
|
||||
def _disabled_instrument(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
ThreadingInstrumentor.instrument = _disabled_instrument # type: ignore[method-assign]
|
||||
return ThreadingInstrumentor, _disabled_instrument
|
||||
|
||||
|
||||
def test_instrument_returns_self_not_none():
|
||||
ThreadingInstrumentor, _ = _install_patch_without_imports()
|
||||
|
||||
instance = ThreadingInstrumentor()
|
||||
result = instance.instrument()
|
||||
|
||||
# Must return ``self`` so fluent chains don't blow up with
|
||||
# AttributeError: 'NoneType' object has no attribute ...
|
||||
assert result is instance
|
||||
assert result is not None
|
||||
|
||||
|
||||
def test_instrument_accepts_args_and_kwargs():
|
||||
"""Upstream signature may evolve; the patch must accept arbitrary args."""
|
||||
ThreadingInstrumentor, _ = _install_patch_without_imports()
|
||||
|
||||
instance = ThreadingInstrumentor()
|
||||
# Any combination of args/kwargs must be accepted without error.
|
||||
assert instance.instrument() is instance
|
||||
assert instance.instrument("x") is instance
|
||||
assert instance.instrument(foo="bar") is instance
|
||||
assert instance.instrument("x", y=1) is instance
|
||||
|
||||
|
||||
def test_patch_replaces_original_method():
|
||||
ThreadingInstrumentor, sentinel = _install_patch_without_imports()
|
||||
|
||||
assert ThreadingInstrumentor.instrument is sentinel
|
||||
|
||||
|
||||
def test_agent_server_guards_use_if_raise_not_assert():
|
||||
"""The import-order guards in agent_server.py must be implemented as
|
||||
``if not ...: raise RuntimeError(...)``, NOT as ``assert`` statements.
|
||||
|
||||
Why this matters: ``assert`` is stripped when Python runs with ``-O``
|
||||
(some Docker base images, optimized CPython builds). If the guards
|
||||
use ``assert``, they silently vanish under ``-O`` and the documented
|
||||
RecursionError returns with no signal.
|
||||
|
||||
This test reads the source of ``agent_server.py`` and fails if any
|
||||
guard is written as ``assert``.
|
||||
"""
|
||||
import importlib.util
|
||||
|
||||
spec = importlib.util.find_spec("agent_server")
|
||||
if spec is None or spec.origin is None:
|
||||
pytest.skip("agent_server module not locatable on sys.path")
|
||||
|
||||
with open(spec.origin, "r", encoding="utf-8") as fh:
|
||||
source = fh.read()
|
||||
|
||||
# The guards we care about are named functions. The invariant we
|
||||
# enforce: their bodies must raise, not assert. We grep for the
|
||||
# former and reject the latter.
|
||||
assert "_assert_strands_not_preimported" in source, (
|
||||
"agent_server.py must expose _assert_strands_not_preimported() "
|
||||
"so tests can monkey-patch the guard cleanly"
|
||||
)
|
||||
assert "_assert_instrumentor_patched" in source, (
|
||||
"agent_server.py must expose _assert_instrumentor_patched() for the same reason"
|
||||
)
|
||||
|
||||
# Check that neither of the guard strings appears behind an ``assert``.
|
||||
# We match ``assert "strands" not in sys.modules`` and ``assert
|
||||
# _ThreadingInstrumentor.instrument is``. Both are the shapes that
|
||||
# existed before the -O fix. If either regex matches, we've regressed.
|
||||
import re
|
||||
|
||||
forbidden_patterns = [
|
||||
r'^\s*assert\s+"strands"\s+not\s+in\s+sys\.modules\b',
|
||||
r"^\s*assert\s+_ThreadingInstrumentor\.instrument\s+is\b",
|
||||
]
|
||||
for pat in forbidden_patterns:
|
||||
assert not re.search(pat, source, flags=re.MULTILINE), (
|
||||
f"agent_server.py contains a raw ``assert`` guard matching {pat!r}; "
|
||||
"these are stripped under ``python -O``. Use ``if/raise RuntimeError`` instead."
|
||||
)
|
||||
|
||||
# Positive assertion: the named functions must raise RuntimeError
|
||||
# (not AssertionError) so the guard survives -O and also so callers
|
||||
# don't need ``__debug__`` to be true.
|
||||
assert "raise RuntimeError(" in source, (
|
||||
"agent_server.py guards must use ``raise RuntimeError(...)`` — "
|
||||
"RuntimeError (not AssertionError) so the semantics are identical "
|
||||
"under ``python -O``"
|
||||
)
|
||||
|
||||
|
||||
def test_agent_server_module_installs_patch():
|
||||
"""Importing ``agent_server`` must leave the instrumentor patched.
|
||||
|
||||
We stub out the strands imports that ``agent_server`` would otherwise
|
||||
require, so the test can run in environments where strands isn't
|
||||
installed.
|
||||
"""
|
||||
import sys
|
||||
import types
|
||||
|
||||
# Stub modules that agent_server transitively imports, so the module
|
||||
# loads without needing the real strands stack installed.
|
||||
class _AcceptsAnything:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
self._agents_by_thread: dict = {}
|
||||
|
||||
class _FakeFastAPI:
|
||||
"""Accepts the decorator calls agent_server applies to ``app``."""
|
||||
|
||||
def _decorator(self, *a, **k):
|
||||
def _wrap(fn):
|
||||
return fn
|
||||
|
||||
return _wrap
|
||||
|
||||
get = post = put = delete = patch = _decorator
|
||||
|
||||
# agent_server.py installs a ``HealthMiddleware`` via
|
||||
# ``app.add_middleware(...)`` after creating the Strands app — the
|
||||
# stub must accept that call so module execution completes.
|
||||
def add_middleware(self, *a, **k):
|
||||
return None
|
||||
|
||||
# agent_server.py mounts the voice sub-app via ``app.mount(...)``
|
||||
def mount(self, *a, **k):
|
||||
return None
|
||||
|
||||
fake_ag_ui_strands = types.ModuleType("ag_ui_strands")
|
||||
fake_ag_ui_strands.create_strands_app = lambda *a, **k: _FakeFastAPI() # type: ignore[attr-defined]
|
||||
fake_ag_ui_strands.StrandsAgent = _AcceptsAnything # type: ignore[attr-defined]
|
||||
fake_ag_ui_strands.StrandsAgentConfig = _AcceptsAnything # type: ignore[attr-defined]
|
||||
fake_ag_ui_strands.ToolBehavior = _AcceptsAnything # type: ignore[attr-defined]
|
||||
|
||||
fake_strands = types.ModuleType("strands")
|
||||
fake_strands.Agent = _AcceptsAnything # type: ignore[attr-defined]
|
||||
fake_strands.tool = lambda f=None, **_: f if callable(f) else (lambda g: g) # type: ignore[attr-defined]
|
||||
|
||||
fake_hooks = types.ModuleType("strands.hooks")
|
||||
for name in (
|
||||
"AfterToolCallEvent",
|
||||
"BeforeInvocationEvent",
|
||||
"BeforeToolCallEvent",
|
||||
"HookProvider",
|
||||
"HookRegistry",
|
||||
):
|
||||
setattr(fake_hooks, name, type(name, (), {}))
|
||||
|
||||
fake_openai_mod = types.ModuleType("strands.models.openai")
|
||||
|
||||
class _FakeOpenAIModel:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
fake_openai_mod.OpenAIModel = _FakeOpenAIModel # type: ignore[attr-defined]
|
||||
fake_models = types.ModuleType("strands.models")
|
||||
|
||||
# uvicorn is imported at module level but only invoked from ``main()``.
|
||||
if "uvicorn" not in sys.modules:
|
||||
fake_uvicorn = types.ModuleType("uvicorn")
|
||||
fake_uvicorn.run = lambda *a, **k: None # type: ignore[attr-defined]
|
||||
sys.modules["uvicorn"] = fake_uvicorn
|
||||
|
||||
# dotenv similarly only invoked at load_dotenv() call time.
|
||||
if "dotenv" not in sys.modules:
|
||||
fake_dotenv = types.ModuleType("dotenv")
|
||||
fake_dotenv.load_dotenv = lambda *a, **k: None # type: ignore[attr-defined]
|
||||
sys.modules["dotenv"] = fake_dotenv
|
||||
|
||||
# ``build_showcase_agent`` fails fast if OPENAI_API_KEY is missing;
|
||||
# set a dummy value so the module-level call succeeds under the stubs.
|
||||
import os as _os
|
||||
|
||||
_os.environ.setdefault("OPENAI_API_KEY", "test-key-for-instrumentor-patch")
|
||||
|
||||
# Drop any pre-existing strands / agent_server entries so the
|
||||
# import-order guard passes cleanly. After the guard runs, install
|
||||
# the stubs via a meta-path finder so the post-patch
|
||||
# ``from ag_ui_strands import ...`` lines resolve.
|
||||
sys.modules.pop("agent_server", None)
|
||||
for _stale in (
|
||||
"strands",
|
||||
"strands.hooks",
|
||||
"strands.models",
|
||||
"strands.models.openai",
|
||||
"ag_ui_strands",
|
||||
):
|
||||
sys.modules.pop(_stale, None)
|
||||
|
||||
_STUB_MAP = {
|
||||
"strands": fake_strands,
|
||||
"strands.hooks": fake_hooks,
|
||||
"strands.models": fake_models,
|
||||
"strands.models.openai": fake_openai_mod,
|
||||
"ag_ui_strands": fake_ag_ui_strands,
|
||||
}
|
||||
|
||||
class _StubLoader:
|
||||
"""Minimal loader that returns a pre-built module object."""
|
||||
|
||||
def __init__(self, module):
|
||||
self._module = module
|
||||
|
||||
def create_module(self, spec):
|
||||
return self._module
|
||||
|
||||
def exec_module(self, module):
|
||||
# Module body already populated; nothing to execute.
|
||||
return None
|
||||
|
||||
class _LazyStubFinder:
|
||||
"""Serves stub modules for strands / ag_ui_strands names on demand.
|
||||
|
||||
agent_server's ``_assert_strands_not_preimported()`` runs BEFORE
|
||||
any of these modules are referenced, so sys.modules is clean at
|
||||
guard time. The first post-patch ``from ag_ui_strands import ...``
|
||||
triggers this finder, which returns a proper ModuleSpec whose
|
||||
loader yields our pre-built stub.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def find_spec(cls, name, path, target=None):
|
||||
mod = _STUB_MAP.get(name)
|
||||
if mod is None:
|
||||
return None
|
||||
import importlib.util as _u
|
||||
|
||||
return _u.spec_from_loader(name, _StubLoader(mod))
|
||||
|
||||
import importlib.util as _util
|
||||
|
||||
spec = _util.find_spec("agent_server")
|
||||
if spec is None or spec.origin is None:
|
||||
pytest.skip("agent_server module not locatable on sys.path")
|
||||
|
||||
sys.meta_path.insert(0, _LazyStubFinder)
|
||||
try:
|
||||
with open(spec.origin, "r", encoding="utf-8") as fh:
|
||||
source = fh.read()
|
||||
module_ns: dict = {
|
||||
"__name__": "agent_server_under_test",
|
||||
"__file__": spec.origin,
|
||||
}
|
||||
# Execute agent_server's source verbatim — no regex surgery. The
|
||||
# guard runs against an empty-of-strands ``sys.modules`` and passes;
|
||||
# downstream imports hit the lazy stub finder.
|
||||
exec(compile(source, spec.origin, "exec"), module_ns)
|
||||
finally:
|
||||
sys.meta_path.remove(_LazyStubFinder)
|
||||
|
||||
from opentelemetry.instrumentation.threading import ThreadingInstrumentor
|
||||
|
||||
# The replacement must return self.
|
||||
instance = ThreadingInstrumentor()
|
||||
assert instance.instrument() is instance
|
||||
|
||||
|
||||
def test_import_order_guard_catches_preimported_strands():
|
||||
"""agent_server.py contains a ``_assert_strands_not_preimported()``
|
||||
guard BEFORE the ThreadingInstrumentor patch. If strands was already
|
||||
imported (directly or transitively) above that line, the OTel patch
|
||||
would be too late — strands' Tracer may have already been constructed.
|
||||
|
||||
Simulate the failure mode: pre-seed ``sys.modules['strands']``, then
|
||||
try to import agent_server, and verify the guard raises ``RuntimeError``
|
||||
(NOT ``AssertionError`` — the latter would silently disappear under
|
||||
``python -O``).
|
||||
"""
|
||||
import sys
|
||||
import types
|
||||
|
||||
# Install a fake 'strands' before agent_server imports. In the real
|
||||
# failure mode this would be the genuine package that already ran
|
||||
# ThreadingInstrumentor().instrument() — here the presence alone is
|
||||
# what the guard checks.
|
||||
preexisting_strands = types.ModuleType("strands")
|
||||
sys.modules["strands"] = preexisting_strands
|
||||
|
||||
# Ensure agent_server is re-imported fresh so the module-level
|
||||
# guard executes.
|
||||
sys.modules.pop("agent_server", None)
|
||||
|
||||
try:
|
||||
# RuntimeError, NOT AssertionError — the guard is an explicit
|
||||
# ``raise`` so ``python -O`` doesn't strip it.
|
||||
with pytest.raises(RuntimeError, match="strands imported before"):
|
||||
import agent_server # noqa: F401
|
||||
finally:
|
||||
# Cleanup: remove the fake strands so subsequent tests can
|
||||
# install their own stubs.
|
||||
sys.modules.pop("strands", None)
|
||||
sys.modules.pop("agent_server", None)
|
||||
|
||||
|
||||
def test_real_strands_agent_signature_integration():
|
||||
"""Integration-style test: when the real ``strands-agents`` package
|
||||
IS installed, construct a real ``strands.Agent`` via the shapes the
|
||||
conftest stubs cover. This catches drift between our stubs and the
|
||||
real signature — if strands renames an Agent kwarg, the stub still
|
||||
passes the unit tests but this integration test fails.
|
||||
|
||||
Skipped gracefully when ``strands-agents`` isn't installed (the
|
||||
default in the unit-test venv; it's available in the Docker image
|
||||
and CI integration environments).
|
||||
"""
|
||||
import sys
|
||||
import types
|
||||
|
||||
# Don't touch the stub modules installed by conftest unless we're
|
||||
# actually going to use the real package. Probe without importing.
|
||||
try:
|
||||
import importlib.util
|
||||
|
||||
spec = importlib.util.find_spec("strands")
|
||||
except Exception:
|
||||
spec = None
|
||||
|
||||
if spec is None:
|
||||
pytest.skip("strands-agents not installed; integration test skipped")
|
||||
|
||||
# If the stub is already in sys.modules, drop it so the real package
|
||||
# gets imported fresh.
|
||||
for mod_name in [
|
||||
"strands",
|
||||
"strands.hooks",
|
||||
"strands.models",
|
||||
"strands.models.openai",
|
||||
]:
|
||||
mod = sys.modules.get(mod_name)
|
||||
if (
|
||||
mod is not None
|
||||
and isinstance(mod, types.ModuleType)
|
||||
and not getattr(mod, "__file__", None)
|
||||
):
|
||||
sys.modules.pop(mod_name, None)
|
||||
|
||||
try:
|
||||
from strands import Agent # type: ignore[import-not-found]
|
||||
except Exception as exc: # pragma: no cover - defensive
|
||||
pytest.skip(f"strands package present but not importable cleanly: {exc}")
|
||||
return
|
||||
|
||||
# The minimal constructor should at least accept a ``tools`` kwarg
|
||||
# (which our factory relies on). We don't need a real model here --
|
||||
# just verify the signature accepts the shape we use.
|
||||
try:
|
||||
agent = Agent(tools=[]) # type: ignore[call-arg]
|
||||
except TypeError as exc:
|
||||
# TypeError here == real strands Agent signature drifted from
|
||||
# what the factory passes in build_showcase_agent.
|
||||
pytest.fail(
|
||||
f"real strands.Agent no longer accepts the kwargs build_showcase_agent "
|
||||
f"relies on: {exc}"
|
||||
)
|
||||
except Exception:
|
||||
# Any other exception (e.g. model required) is fine — the point
|
||||
# is that the Agent class accepted the kwargs.
|
||||
pass
|
||||
else:
|
||||
assert agent is not None
|
||||
@@ -0,0 +1,196 @@
|
||||
"""Tests for ``sales_state_from_args`` in src/agents/agent.py.
|
||||
|
||||
``sales_state_from_args`` runs as a tool-result callback for
|
||||
``manage_sales_todos`` and emits a ``{"todos": [...]}`` snapshot for the
|
||||
UI. The function uses shape-directed dispatch (``isinstance`` checks on
|
||||
``tool_input``), not try/except-AttributeError, to support three
|
||||
realistic input shapes:
|
||||
|
||||
* ``tool_input`` as a ``dict`` with a ``"todos"`` key (the usual path
|
||||
when the LLM calls the tool through the strands machinery),
|
||||
* ``tool_input`` as a JSON ``str`` (parsed via ``json.loads``; the only
|
||||
exception-guarded branch, catching ``json.JSONDecodeError``),
|
||||
* ``tool_input`` as the bare list (the LLM occasionally inlines the
|
||||
list without the ``"todos"`` wrapper — the ``isinstance(_, list)``
|
||||
branch treats it as the todos payload directly).
|
||||
|
||||
Unsupported types (int, None, missing attribute) short-circuit with a
|
||||
warning log and a ``None`` return — silently dropping the state snapshot
|
||||
is preferable to blowing up the tool response pipeline over malformed
|
||||
input.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _run(coro):
|
||||
"""Run an async coroutine synchronously for tests.
|
||||
|
||||
``sales_state_from_args`` is ``async def`` (strands invokes it as a
|
||||
tool-result callback in an event loop). Using ``asyncio.run`` keeps
|
||||
each test self-contained and avoids event-loop leakage between tests.
|
||||
"""
|
||||
return asyncio.run(coro)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _patched_impl(monkeypatch):
|
||||
"""Patch ``manage_sales_todos_impl`` with a deterministic pass-through
|
||||
so the test doesn't depend on the shared tool's real filtering logic.
|
||||
|
||||
The real impl is covered by tests in the shared-python tools package.
|
||||
"""
|
||||
import agents.agent as agent_mod
|
||||
|
||||
def _identity(todos):
|
||||
# Return a list of dicts (the real impl returns the same shape).
|
||||
return [dict(t) for t in todos]
|
||||
|
||||
monkeypatch.setattr(agent_mod, "manage_sales_todos_impl", _identity)
|
||||
return agent_mod
|
||||
|
||||
|
||||
def test_dict_tool_input_with_todos_key(_patched_impl):
|
||||
"""Happy path: ``tool_input`` is a dict with ``"todos"`` → returns
|
||||
``{"todos": [...]}``."""
|
||||
from agents.agent import sales_state_from_args
|
||||
|
||||
todos = [{"id": "t1", "title": "Call Acme"}, {"id": "t2", "title": "Follow up"}]
|
||||
ctx = SimpleNamespace(tool_input={"todos": todos})
|
||||
|
||||
result = _run(sales_state_from_args(ctx))
|
||||
|
||||
assert result == {"todos": todos}
|
||||
|
||||
|
||||
def test_json_string_tool_input(_patched_impl):
|
||||
"""``tool_input`` as a JSON string is parsed and processed identically."""
|
||||
from agents.agent import sales_state_from_args
|
||||
|
||||
todos = [{"id": "s1", "title": "Prospect"}]
|
||||
ctx = SimpleNamespace(tool_input=json.dumps({"todos": todos}))
|
||||
|
||||
result = _run(sales_state_from_args(ctx))
|
||||
|
||||
assert result == {"todos": todos}
|
||||
|
||||
|
||||
def test_bare_list_tool_input(_patched_impl):
|
||||
"""When the LLM inlines the list directly (no ``"todos"`` wrapper),
|
||||
the function emits the state snapshot. The shape-directed dispatch
|
||||
uses ``isinstance(tool_input, list)`` to treat the bare list as the
|
||||
todos payload — no ``.get``, no exception handling for this branch.
|
||||
"""
|
||||
from agents.agent import sales_state_from_args
|
||||
|
||||
todos = [{"id": "b1", "title": "Bare list"}]
|
||||
ctx = SimpleNamespace(tool_input=todos)
|
||||
|
||||
result = _run(sales_state_from_args(ctx))
|
||||
|
||||
# isinstance(tool_input, list) path: bare list IS the todos payload.
|
||||
assert result == {"todos": todos}
|
||||
|
||||
|
||||
def test_malformed_json_string_returns_none(_patched_impl):
|
||||
"""Invalid JSON in ``tool_input`` must NOT propagate — returns None."""
|
||||
from agents.agent import sales_state_from_args
|
||||
|
||||
ctx = SimpleNamespace(tool_input="{not valid json")
|
||||
|
||||
result = _run(sales_state_from_args(ctx))
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_missing_tool_input_attribute_returns_none(_patched_impl):
|
||||
"""If the context lacks ``tool_input`` entirely, we fall back to None.
|
||||
|
||||
The function uses ``getattr(context, "tool_input", None)`` for the
|
||||
initial fetch and short-circuits on ``None`` via an explicit check —
|
||||
no exception is raised or caught for this branch.
|
||||
"""
|
||||
from agents.agent import sales_state_from_args
|
||||
|
||||
# An object without ``tool_input`` at all. getattr returns None, which
|
||||
# the function handles via an explicit None check (no AttributeError).
|
||||
class _Ctx:
|
||||
pass
|
||||
|
||||
result = _run(sales_state_from_args(_Ctx()))
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_non_dict_non_string_tool_input_returns_none(_patched_impl):
|
||||
"""Edge case: ``tool_input`` is an int (LLM misfire). Must not crash."""
|
||||
from agents.agent import sales_state_from_args
|
||||
|
||||
ctx = SimpleNamespace(tool_input=42)
|
||||
|
||||
result = _run(sales_state_from_args(ctx))
|
||||
|
||||
# Shape-directed dispatch: int is neither dict, list, nor str; the
|
||||
# unsupported-type branch logs a warning and returns None.
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_empty_todos_list_returns_empty_snapshot(_patched_impl):
|
||||
"""Empty todos list is valid input and produces ``{"todos": []}``."""
|
||||
from agents.agent import sales_state_from_args
|
||||
|
||||
ctx = SimpleNamespace(tool_input={"todos": []})
|
||||
|
||||
result = _run(sales_state_from_args(ctx))
|
||||
|
||||
assert result == {"todos": []}
|
||||
|
||||
|
||||
def test_processed_todos_pass_through_shared_impl(monkeypatch):
|
||||
"""Confirms the function actually routes through
|
||||
``manage_sales_todos_impl`` (not just returning the raw input)."""
|
||||
import agents.agent as agent_mod
|
||||
from agents.agent import sales_state_from_args
|
||||
|
||||
def _transformer(todos):
|
||||
# Simulate the real impl adding an index or filtering.
|
||||
return [{"id": t.get("id"), "normalized": True} for t in todos]
|
||||
|
||||
monkeypatch.setattr(agent_mod, "manage_sales_todos_impl", _transformer)
|
||||
|
||||
ctx = SimpleNamespace(tool_input={"todos": [{"id": "x"}, {"id": "y"}]})
|
||||
|
||||
result = _run(sales_state_from_args(ctx))
|
||||
|
||||
assert result == {
|
||||
"todos": [
|
||||
{"id": "x", "normalized": True},
|
||||
{"id": "y", "normalized": True},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def test_logger_warning_on_malformed_input(_patched_impl, caplog):
|
||||
"""Malformed input should emit a WARNING log with a truncated
|
||||
excerpt so on-call can debug without a full traceback."""
|
||||
import logging
|
||||
|
||||
from agents.agent import sales_state_from_args
|
||||
|
||||
ctx = SimpleNamespace(tool_input="{not json")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="agents.agent"):
|
||||
result = _run(sales_state_from_args(ctx))
|
||||
|
||||
assert result is None
|
||||
# At least one warning record must mention the parse failure.
|
||||
warning_records = [r for r in caplog.records if r.levelname == "WARNING"]
|
||||
assert warning_records, "expected a WARNING log on malformed tool input"
|
||||
combined = " ".join(r.getMessage() for r in warning_records)
|
||||
assert "sales_state_from_args" in combined
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Tests for _ToolCallCapHook in src/agents/agent.py.
|
||||
|
||||
Exercises the cap behavior by firing synthetic BeforeInvocationEvent /
|
||||
BeforeToolCallEvent / AfterToolCallEvent instances at the hook and
|
||||
asserting:
|
||||
|
||||
* the cap fires at exactly ``_max_calls + 1`` (i.e. the (N+1)-th call is
|
||||
cancelled, not the N-th),
|
||||
* ``BeforeInvocationEvent`` resets the counter between invocations,
|
||||
* ``AfterToolCallEvent`` sets the ``stop_event_loop`` sentinel on the
|
||||
invocation state once the cap is hit.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hook_cls():
|
||||
from agents.agent import _ToolCallCapHook
|
||||
|
||||
return _ToolCallCapHook
|
||||
|
||||
|
||||
def _make_before_event():
|
||||
# ``BeforeToolCallEvent`` exposes a mutable ``cancel_tool`` attribute.
|
||||
# We fake the event with a SimpleNamespace that accepts the assignment.
|
||||
return SimpleNamespace(cancel_tool=None)
|
||||
|
||||
|
||||
def _make_after_event(invocation_state=None):
|
||||
return SimpleNamespace(
|
||||
invocation_state=invocation_state if invocation_state is not None else {}
|
||||
)
|
||||
|
||||
|
||||
def test_cap_fires_on_call_n_plus_one(hook_cls):
|
||||
hook = hook_cls(max_calls=3)
|
||||
|
||||
# Calls 1..3 should pass through; call 4 (N+1) should cancel.
|
||||
for i in range(1, 4):
|
||||
ev = _make_before_event()
|
||||
hook._on_before_tool(ev)
|
||||
assert ev.cancel_tool is None, f"call {i} should not be cancelled"
|
||||
|
||||
trip_event = _make_before_event()
|
||||
hook._on_before_tool(trip_event)
|
||||
assert trip_event.cancel_tool is not None
|
||||
assert "3" in trip_event.cancel_tool # max_calls surfaced in message
|
||||
|
||||
|
||||
def test_before_invocation_resets_counter(hook_cls):
|
||||
hook = hook_cls(max_calls=2)
|
||||
|
||||
# Exhaust the cap.
|
||||
hook._on_before_tool(_make_before_event())
|
||||
hook._on_before_tool(_make_before_event())
|
||||
trip = _make_before_event()
|
||||
hook._on_before_tool(trip)
|
||||
assert trip.cancel_tool is not None
|
||||
|
||||
# Reset via BeforeInvocationEvent.
|
||||
hook._on_invocation_start(SimpleNamespace())
|
||||
|
||||
# The counter should be back to zero; the next 2 calls must pass.
|
||||
next_ev = _make_before_event()
|
||||
hook._on_before_tool(next_ev)
|
||||
assert next_ev.cancel_tool is None
|
||||
|
||||
second = _make_before_event()
|
||||
hook._on_before_tool(second)
|
||||
assert second.cancel_tool is None
|
||||
|
||||
|
||||
def test_after_tool_sets_stop_event_loop_sentinel(hook_cls):
|
||||
"""Once the counter reaches ``max_calls``, ``_on_after_tool`` must set
|
||||
the ``stop_event_loop`` sentinel on the invocation state so strands halts
|
||||
the event loop at the end of the current cycle.
|
||||
|
||||
Note on sentinel timing: the sentinel fires at ``_count >= _max_calls``
|
||||
(one call earlier than the cancellation, which fires at
|
||||
``_count > _max_calls``). The sentinel and the cancellation are
|
||||
orthogonal mechanisms: the sentinel halts the event loop before a
|
||||
potential (N+1)-th call is ever attempted, and the cancellation is a
|
||||
belt-and-suspenders guard for the case where strands dispatches the
|
||||
(N+1)-th call anyway (e.g. because the sentinel was set too late in
|
||||
the cycle, or the tool dispatch was already in flight).
|
||||
"""
|
||||
hook = hook_cls(max_calls=3)
|
||||
|
||||
# Calls under the cap must not set the sentinel.
|
||||
for _ in range(2):
|
||||
hook._on_before_tool(_make_before_event())
|
||||
state = {}
|
||||
hook._on_after_tool(_make_after_event(state))
|
||||
assert not state.get("request_state", {}).get("stop_event_loop")
|
||||
|
||||
# Reaching the cap (count == max) sets the sentinel.
|
||||
hook._on_before_tool(_make_before_event()) # count now == 3
|
||||
at_cap_state = {}
|
||||
hook._on_after_tool(_make_after_event(at_cap_state))
|
||||
assert at_cap_state["request_state"]["stop_event_loop"] is True
|
||||
|
||||
# Over-cap call is cancelled AND sets the sentinel.
|
||||
tripping = _make_before_event()
|
||||
hook._on_before_tool(tripping) # count now == 4
|
||||
assert tripping.cancel_tool is not None
|
||||
|
||||
over_state = {}
|
||||
hook._on_after_tool(_make_after_event(over_state))
|
||||
assert over_state["request_state"]["stop_event_loop"] is True
|
||||
|
||||
|
||||
def test_default_cap_matches_module_constant(hook_cls):
|
||||
from agents.agent import _MAX_TOOL_CALLS_PER_INVOCATION
|
||||
|
||||
hook = hook_cls()
|
||||
assert hook._max_calls == _MAX_TOOL_CALLS_PER_INVOCATION
|
||||
|
||||
|
||||
def test_concurrent_before_tool_calls_respect_cap(hook_cls):
|
||||
"""Fire 100 concurrent ``_on_before_tool`` calls against a cap of 50
|
||||
and assert the cap holds: exactly 50 calls pass through and 50 are
|
||||
cancelled.
|
||||
|
||||
The hook's ``_lock`` guards ``_count`` mutation so that under
|
||||
concurrent invocation (e.g. strands dispatching tools on a
|
||||
ThreadPoolExecutor, or misuse via two concurrent requests on the same
|
||||
thread_id) we degrade gracefully rather than race silently. Without
|
||||
the lock, the classic read-modify-write race would allow more than 50
|
||||
calls to pass the ``current > max_calls`` gate.
|
||||
"""
|
||||
import threading
|
||||
|
||||
max_calls = 50
|
||||
total = 100
|
||||
hook = hook_cls(max_calls=max_calls)
|
||||
|
||||
events = [_make_before_event() for _ in range(total)]
|
||||
barrier = threading.Barrier(total)
|
||||
|
||||
def _fire(ev):
|
||||
barrier.wait()
|
||||
hook._on_before_tool(ev)
|
||||
|
||||
threads = [threading.Thread(target=_fire, args=(ev,)) for ev in events]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
passed = sum(1 for ev in events if ev.cancel_tool is None)
|
||||
cancelled = sum(1 for ev in events if ev.cancel_tool is not None)
|
||||
|
||||
assert passed == max_calls, f"expected exactly {max_calls} passes, got {passed}"
|
||||
assert cancelled == total - max_calls, (
|
||||
f"expected exactly {total - max_calls} cancellations, got {cancelled}"
|
||||
)
|
||||
# And the internal counter should land at ``total`` (every call was counted).
|
||||
assert hook._count == total
|
||||
|
||||
|
||||
def test_tool_call_cap_validates_max_calls(hook_cls):
|
||||
"""``max_calls < 1`` silently cancels every tool call because the
|
||||
first ``_on_before_tool`` increment-then-compare ends up with
|
||||
``1 > 0`` -> cancel. Constructor must reject this up front."""
|
||||
with pytest.raises(ValueError, match="max_calls must be >= 1"):
|
||||
hook_cls(max_calls=0)
|
||||
|
||||
with pytest.raises(ValueError, match="max_calls must be >= 1"):
|
||||
hook_cls(max_calls=-1)
|
||||
|
||||
# Boundary: 1 is valid. The very next call would cancel, but the
|
||||
# hook itself must construct without error.
|
||||
hook = hook_cls(max_calls=1)
|
||||
assert hook._max_calls == 1
|
||||
Reference in New Issue
Block a user