diff --git a/.env.example b/.env.example index 96f1206..f147a6c 100644 --- a/.env.example +++ b/.env.example @@ -18,20 +18,25 @@ LLM_MODEL=gemini-3.1-flash-lite GEMINI_API_KEY=your_gemini_api_key_here OPENAI_API_KEY=your_openai_api_key_here +# Retry policy for transient LLM failures (429 / 5xx / network). Non-transient +# errors such as 401 or 400 are never retried. +LLM_MAX_RETRIES=3 +LLM_RETRY_BASE_DELAY_S=0.5 +LLM_RETRY_MAX_DELAY_S=10 + # External Tools & Sandboxing SEARCH_API_KEY=your_search_api_key_here FILE_TOOL_BASE_DIR=/tmp/agentkit_sandbox # Security & API Authentication -# Comma-separated list of authorized client API keys. -# docker-compose.yml ships no default: leave this unset and the api service -# exits on settings validation instead of starting with a published key. -API_KEYS=ak_test_key_12345,ak_admin_key_67890 +# Comma-separated list of authorized client API keys. Required: AgentKit refuses +# to start when this is empty or unset, so no deployment can come up with a key +# published in this repository. Replace the values below with your own. +API_KEYS=replace_me_key_1,replace_me_key_2 RATE_LIMIT_PER_MIN=30 # Agent Execution Limits & Guardrails MAX_STEPS=10 -MAX_RETRIES=3 RUN_TIMEOUT_S=60 TOOL_TIMEOUT_S=15 SESSION_TTL_S=3600 diff --git a/README.md b/README.md index ff46277..047bcdf 100644 --- a/README.md +++ b/README.md @@ -110,16 +110,22 @@ OPENAI_API_KEY=your_openai_api_key_here DATABASE_URL=postgresql+asyncpg://agentkit:agentkit_password@localhost:5432/agentkit REDIS_URL=redis://localhost:6379/0 -# Security (comma-separated authorized keys) -API_KEYS=ak_test_key_12345,ak_admin_key_67890 +# Security (API_KEYS is required; the app refuses to start without it) +API_KEYS=dev_key_1,dev_key_2 + +# Optional: retry policy for transient LLM failures (429/5xx/network) +LLM_MAX_RETRIES=3 +LLM_RETRY_BASE_DELAY_S=0.5 +LLM_RETRY_MAX_DELAY_S=10 ``` ### 4. Running with Docker Compose (Recommended) Launch the API, PostgreSQL, and Redis together with automatic migrations. -`docker-compose.yml` ships no default `API_KEYS`. Leave it unset and the api service -fails immediately on settings validation rather than starting with a published key; -`down`, `ps` and `logs` still work either way. +`API_KEYS` is required in every environment, not just under compose. `Settings` rejects an +empty value, so a deployment that forgets to configure its keys fails at startup rather than +serving unauthenticated traffic with a key published in the repository. `docker-compose.yml` +ships no default either; `down`, `ps` and `logs` still work without it. ```bash export API_KEYS=ak_local_dev_key # or place API_KEYS in your .env diff --git a/agentkit/api/deps.py b/agentkit/api/deps.py index cdf6287..121b2ee 100644 --- a/agentkit/api/deps.py +++ b/agentkit/api/deps.py @@ -138,16 +138,26 @@ def get_llm_client( request: Request, settings: Settings = Depends(get_current_settings), ) -> LLMClient: - """Resolve configured LLM client. + """Resolve the configured LLM client. - Instantiates the configured provider from settings. Explicit test overrides - can be supplied via request.app.state.llm_client. + Instantiates the configured provider from settings, which also applies the + configured transient-error retry policy. Explicit test overrides can be supplied + via request.app.state.llm_client. A provider construction failure is surfaced as a + 503: falling back to a scripted client would answer a misconfigured deployment with + fake text and HTTP 200 instead of surfacing the misconfiguration. """ - if hasattr(request.app.state, "llm_client") and request.app.state.llm_client is not None: - client: LLMClient = request.app.state.llm_client - return client - - return create_llm_client_from_settings(settings) + override: LLMClient | None = getattr(request.app.state, "llm_client", None) + if override is not None: + return override + + try: + return create_llm_client_from_settings(settings) + except Exception as exc: + logger.error("Failed to initialize LLM provider %s: %s", settings.LLM_PROVIDER, exc) + raise ServiceUnavailableError( + "llm", + f"LLM provider '{settings.LLM_PROVIDER}' could not be initialized.", + ) from exc def get_agent( @@ -158,10 +168,15 @@ def get_agent( memory: MemoryStore = Depends(get_memory_store), trace: TraceSink = Depends(get_trace_sink), ) -> Agent: - """Assemble configured Agent instance with all runtime dependencies.""" - if hasattr(request.app.state, "agent") and request.app.state.agent is not None: - agent: Agent = request.app.state.agent - return agent + """Assemble a configured Agent from the currently resolved dependencies. + + The agent is rebuilt per request rather than memoized on ``app.state``: caching the first + assembled instance would pin every later request to that request's LLM client, registry, + memory store and trace sink, silently ignoring any later override or settings change. + """ + override: Agent | None = getattr(request.app.state, "agent_override", None) + if override is not None: + return override config = AgentConfig( max_steps=settings.MAX_STEPS, diff --git a/agentkit/api/ratelimit.py b/agentkit/api/ratelimit.py index 610662b..755377d 100644 --- a/agentkit/api/ratelimit.py +++ b/agentkit/api/ratelimit.py @@ -10,7 +10,7 @@ from fastapi import Depends from redis.asyncio import Redis -from redis.exceptions import RedisError, ResponseError +from redis.exceptions import RedisError, ResponseError, WatchError from agentkit.api.auth import get_current_settings, verify_api_key from agentkit.api.deps import get_redis_client @@ -57,6 +57,10 @@ return {1, 0} """ +#: Retries for the WATCH/MULTI fallback when concurrent writers keep invalidating the watch. +#: Only reached on Redis deployments without scripting support. +_WATCH_MAX_RETRIES = 5 + class RateLimiter: """Sliding window rate limiter enforcing request volume limits per client key. @@ -203,23 +207,48 @@ async def _check_window_unatomic( member: str, key_ttl: int, ) -> tuple[bool, int]: - """Best-effort sliding-window check for servers without scripting support.""" - pipe = redis_client.pipeline(transaction=True) - pipe.zremrangebyscore(redis_key, "-inf", now - window) - pipe.zcard(redis_key) - results = await pipe.execute() - if int(results[1]) >= limit: - oldest = await redis_client.zrange(redis_key, 0, 0, withscores=True) - retry_after = ( - max(1, math.ceil(float(oldest[0][1]) + window - now)) if oldest else int(window) - ) - return False, retry_after + """Best-effort sliding-window check for servers without scripting support. + + An earlier version trimmed and counted in one transaction, then inserted in a second. + That second round trip leaves a gap in which concurrent requests all read the same + pre-insert cardinality and are admitted together, so a parallel burst could exceed the + quota. WATCH/MULTI/EXEC closes the gap: the whole read-decide-write cycle aborts and + retries whenever another client modifies the key in between. + """ + for _ in range(_WATCH_MAX_RETRIES): + async with redis_client.pipeline(transaction=True) as pipe: + try: + await pipe.watch(redis_key) + await pipe.zremrangebyscore(redis_key, "-inf", now - window) + count = await pipe.zcard(redis_key) - pipe = redis_client.pipeline(transaction=True) - pipe.zadd(redis_key, {member: now}) - pipe.expire(redis_key, key_ttl) - await pipe.execute() - return True, 0 + if int(count) >= limit: + oldest = await pipe.zrange(redis_key, 0, 0, withscores=True) + await pipe.reset() + retry_after = ( + max(1, math.ceil(float(oldest[0][1]) + window - now)) + if oldest + else int(window) + ) + return False, retry_after + + pipe.multi() + pipe.zadd(redis_key, {member: now}) + pipe.expire(redis_key, key_ttl) + await pipe.execute() + return True, 0 + except WatchError: + # Another client changed the window between WATCH and EXEC; re-read and retry. + continue + + # Contention never settled. Denying is the safe direction for a limiter: a false + # reject costs one request, a false admit lets a client exceed its quota. + logger.warning( + "Rate limit check for key %s hit %d watch retries under contention; denying request.", + redis_key[:8] + "...", + _WATCH_MAX_RETRIES, + ) + return False, max(1, math.ceil(window)) async def __call__( self, diff --git a/agentkit/api/routes.py b/agentkit/api/routes.py index 18dce63..6f75248 100644 --- a/agentkit/api/routes.py +++ b/agentkit/api/routes.py @@ -31,6 +31,11 @@ logger = logging.getLogger(__name__) +#: Fixed client-facing strings for failed dependency checks. Raw driver errors can contain +#: connection URLs, credentials or internal hostnames, and /health requires no API key. +DATABASE_UNHEALTHY = "Database connectivity check failed." +REDIS_UNHEALTHY = "Redis connectivity check failed." + runs_router = APIRouter( prefix="/runs", tags=["runs"], @@ -47,7 +52,10 @@ async def create_run( memory_store: PostgresMemoryStore | MemoryStore = Depends(get_memory_store), ) -> RunResponse: """Execute an agent goal and return execution summary.""" - logger.info("Executing run for goal: %s (session_id: %s)", payload.goal, payload.session_id) + # The goal is caller-supplied prompt data, so only its length is logged at INFO. + logger.info( + "Executing run (goal_chars=%d, session_id: %s)", len(payload.goal), payload.session_id + ) result = await agent.run(goal=payload.goal, session_id=payload.session_id) if hasattr(memory_store, "get_run"): @@ -153,7 +161,8 @@ async def health_check(request: Request) -> JSONResponse: redis_status = "healthy" details: dict[str, str] = {} - # Check Database Engine + # /health is unauthenticated, so connection error strings are logged rather than echoed: + # driver messages routinely embed host, port and database name. engine = getattr(request.app.state, "db_engine", None) if engine is not None: try: @@ -161,23 +170,24 @@ async def health_check(request: Request) -> JSONResponse: await conn.execute(text("SELECT 1")) except Exception as exc: db_status = "unhealthy" - details["database"] = str(exc) + details["database"] = DATABASE_UNHEALTHY + logger.warning("Database health check failed: %s", exc) status_code = 503 else: db_status = "disabled" - # Check Redis Client redis_client = getattr(request.app.state, "redis_client", None) if redis_client is not None: try: ping_res = await redis_client.ping() if not ping_res: redis_status = "unhealthy" - details["redis"] = "Ping returned False" + details["redis"] = REDIS_UNHEALTHY status_code = 503 except Exception as exc: redis_status = "unhealthy" - details["redis"] = str(exc) + details["redis"] = REDIS_UNHEALTHY + logger.warning("Redis health check failed: %s", exc) status_code = 503 else: redis_status = "disabled" diff --git a/agentkit/config.py b/agentkit/config.py index e0aa074..8963f8f 100644 --- a/agentkit/config.py +++ b/agentkit/config.py @@ -49,6 +49,22 @@ class Settings(BaseSettings): default=None, description="External search API key (Tavily/SerpAPI).", ) + LLM_MAX_RETRIES: int = Field( + default=3, + ge=0, + le=10, + description="Retry attempts for transient LLM provider failures (429/5xx/network).", + ) + LLM_RETRY_BASE_DELAY_S: float = Field( + default=0.5, + gt=0, + description="Base delay in seconds for exponential LLM retry backoff.", + ) + LLM_RETRY_MAX_DELAY_S: float = Field( + default=10.0, + gt=0, + description="Maximum delay in seconds between LLM retry attempts.", + ) FILE_TOOL_BASE_DIR: str = Field( default="/tmp/agentkit_sandbox", description="Base directory for sandboxed file reading operations.", @@ -56,8 +72,8 @@ class Settings(BaseSettings): # Security & API Authentication API_KEYS: str = Field( - default="ak_test_key_12345", - description="Comma-separated list of valid client API keys.", + default="", + description="Comma-separated list of valid client API keys. Must be set explicitly.", ) RATE_LIMIT_PER_MIN: int = Field( default=30, @@ -71,11 +87,6 @@ class Settings(BaseSettings): gt=0, description="Maximum execution steps allowed per agent run.", ) - MAX_RETRIES: int = Field( - default=3, - ge=0, - description="Maximum retries on transient LLM errors.", - ) RUN_TIMEOUT_S: int = Field( default=60, gt=0, @@ -100,10 +111,17 @@ class Settings(BaseSettings): @field_validator("API_KEYS") @classmethod def validate_api_keys(cls, v: str) -> str: - """Ensure API keys are not empty.""" + """Reject blank or placeholder keys so a deployment cannot start unauthenticated. + + The default is intentionally empty: starting with a key published in the repository + would expose any deployment that forgets to configure its own. + """ cleaned = v.strip() if not cleaned: - raise ValueError("API_KEYS must not be empty") + raise ValueError( + "API_KEYS must be set to a comma-separated list of client API keys. " + "AgentKit refuses to start without explicit authentication credentials." + ) return cleaned @property diff --git a/agentkit/core/agent.py b/agentkit/core/agent.py index 94b67f1..2c95115 100644 --- a/agentkit/core/agent.py +++ b/agentkit/core/agent.py @@ -142,7 +142,10 @@ async def run( run_id = str(uuid.uuid4()) with log_context(run_id=run_id): - logger.info("Starting agent run for goal: %s", goal) + # The goal is caller-supplied prompt data, so it stays out of INFO per the + # logging policy; only its length is recorded there. + logger.info("Starting agent run (goal_chars=%d, session_id=%s)", len(goal), session_id) + logger.debug("Run goal content: %s", goal) if self.memory is not None and hasattr(self.memory, "create_run"): with contextlib.suppress(Exception): await self.memory.create_run( diff --git a/agentkit/llm/factory.py b/agentkit/llm/factory.py index bd1bc41..f31e2e1 100644 --- a/agentkit/llm/factory.py +++ b/agentkit/llm/factory.py @@ -76,10 +76,12 @@ def create_llm_client_from_settings(settings: Settings) -> LLMClient: model=settings.LLM_MODEL, ) - if provider in ("gemini", "openai") and settings.MAX_RETRIES > 0: + if provider in ("gemini", "openai") and settings.LLM_MAX_RETRIES > 0: return RetryingLLMClient( client=client, - max_retries=settings.MAX_RETRIES, + max_retries=settings.LLM_MAX_RETRIES, + base_delay=settings.LLM_RETRY_BASE_DELAY_S, + max_delay=settings.LLM_RETRY_MAX_DELAY_S, ) return client diff --git a/agentkit/llm/gemini.py b/agentkit/llm/gemini.py index 1857dc8..b527e85 100644 --- a/agentkit/llm/gemini.py +++ b/agentkit/llm/gemini.py @@ -152,48 +152,65 @@ async def chat( status_code=int(status_code) if isinstance(status_code, int) else None, ) from exc - # Parse candidate parts - text_parts: list[str] = [] - tool_calls: list[ToolCall] = [] - - if raw_response.candidates: - candidate = raw_response.candidates[0] - if candidate.content and candidate.content.parts: - for part in candidate.content.parts: - if getattr(part, "text", None): - text_parts.append(part.text) - if getattr(part, "function_call", None): - fc = part.function_call - raw_sig = getattr(part, "thought_signature", None) - sig = raw_sig if isinstance(raw_sig, bytes) else None - tool_calls.append( - ToolCall( - id=str(uuid.uuid4())[:8], - name=fc.name, - arguments=dict(fc.args or {}), - thought_signature=sig, + # Response assembly stays inside the error boundary: reading ``raw_response.text`` raises on + # blocked prompts, an empty ``candidates`` list would IndexError, and a mock or future + # SDK shape can drop attributes. All of those are provider faults, so they have to + # surface as LLMProviderError for callers to map, not as raw SDK exceptions. + try: + text_parts: list[str] = [] + tool_calls: list[ToolCall] = [] + + candidates = getattr(raw_response, "candidates", None) + if candidates: + candidate = candidates[0] + if candidate.content and candidate.content.parts: + for part in candidate.content.parts: + if getattr(part, "text", None): + text_parts.append(part.text) + if getattr(part, "function_call", None): + fc = part.function_call + raw_sig = getattr(part, "thought_signature", None) + sig = raw_sig if isinstance(raw_sig, bytes) else None + tool_calls.append( + ToolCall( + # Full UUID: tool_call_id is the join key between an + # assistant tool call and its tool result, so a truncated + # 32-bit id risks collisions within a long session. + id=str(uuid.uuid4()), + name=fc.name, + arguments=dict(fc.args or {}), + thought_signature=sig, + ) ) - ) - - # Extract usage metrics - input_tokens = 0 - output_tokens = 0 - if getattr(raw_response, "usage_metadata", None): - input_tokens = getattr(raw_response.usage_metadata, "prompt_token_count", 0) or 0 - output_tokens = getattr(raw_response.usage_metadata, "candidates_token_count", 0) or 0 - text = "".join(text_parts).strip() - if not text and getattr(raw_response, "text", None): - text = raw_response.text + # Extract usage metrics + input_tokens = 0 + output_tokens = 0 + if getattr(raw_response, "usage_metadata", None): + input_tokens = getattr(raw_response.usage_metadata, "prompt_token_count", 0) or 0 + output_tokens = ( + getattr(raw_response.usage_metadata, "candidates_token_count", 0) or 0 + ) - metadata: dict[str, Any] = {} - if hasattr(raw_response, "model_dump"): - try: - dumped = raw_response.model_dump() - if isinstance(dumped, dict): - metadata = dumped - except Exception: - metadata = {} + text = "".join(text_parts).strip() + if not text and getattr(raw_response, "text", None): + text = raw_response.text + + metadata: dict[str, Any] = {} + if hasattr(raw_response, "model_dump"): + try: + dumped = raw_response.model_dump() + if isinstance(dumped, dict): + metadata = dumped + except Exception: + metadata = {} + except (LLMProviderError, AuthenticationError): + raise + except Exception as exc: + raise LLMProviderError( + provider="gemini", + message=f"Malformed Gemini response: {exc}", + ) from exc return LLMResponse( text=text, diff --git a/agentkit/llm/openai.py b/agentkit/llm/openai.py index 5ec2d6e..6acaabf 100644 --- a/agentkit/llm/openai.py +++ b/agentkit/llm/openai.py @@ -128,44 +128,80 @@ async def chat( except Exception as exc: raise LLMProviderError(provider="openai", message=str(exc)) from exc - choice = raw_response.choices[0] - choice_message = choice.message - text = choice_message.content or "" - tool_calls: list[ToolCall] = [] - - if getattr(choice_message, "tool_calls", None): - for tc in choice_message.tool_calls: - call_id = tc.id or str(uuid.uuid4())[:8] - func_name = tc.function.name - raw_args = tc.function.arguments or "{}" - if isinstance(raw_args, str): - try: - parsed_args = json.loads(raw_args) - except Exception: - parsed_args = {"raw": raw_args} - else: - parsed_args = dict(raw_args) - - tool_calls.append( - ToolCall( - id=call_id, - name=func_name, - arguments=parsed_args, - ) + # Response assembly stays inside the error boundary: an empty ``choices`` list would + # IndexError and unexpected SDK shapes would raise attribute errors. Those are + # provider faults and must surface as LLMProviderError so callers can map them. + try: + if not raw_response.choices: + raise LLMProviderError( + provider="openai", + message="Response contained no choices.", ) - input_tokens = 0 - output_tokens = 0 - if getattr(raw_response, "usage", None): - input_tokens = getattr(raw_response.usage, "prompt_tokens", 0) or 0 - output_tokens = getattr(raw_response.usage, "completion_tokens", 0) or 0 - - metadata: dict[str, Any] = {} - if hasattr(raw_response, "model_dump"): - try: - metadata = raw_response.model_dump() - except Exception: - metadata = {} + choice = raw_response.choices[0] + choice_message = choice.message + text = choice_message.content or "" + tool_calls: list[ToolCall] = [] + + if getattr(choice_message, "tool_calls", None): + for tc in choice_message.tool_calls: + # Full UUID: tool_call_id is the join key between an assistant tool + # call and its tool result, so a truncated 32-bit id risks collisions. + call_id = tc.id or str(uuid.uuid4()) + func_name = tc.function.name + raw_args = tc.function.arguments or "{}" + if isinstance(raw_args, str): + try: + parsed_args = json.loads(raw_args) + except json.JSONDecodeError as exc: + # Substituting {"raw": ...} here would surface downstream as a + # misleading "missing required field" validation error instead of + # naming the real fault: the model emitted malformed JSON. + raise LLMProviderError( + provider="openai", + message=( + f"Tool call '{func_name}' carried malformed JSON arguments: {exc}" + ), + ) from exc + else: + parsed_args = dict(raw_args) + + if not isinstance(parsed_args, dict): + raise LLMProviderError( + provider="openai", + message=( + f"Tool call '{func_name}' arguments must decode to an object, " + f"got {type(parsed_args).__name__}." + ), + ) + + tool_calls.append( + ToolCall( + id=call_id, + name=func_name, + arguments=parsed_args, + ) + ) + + input_tokens = 0 + output_tokens = 0 + if getattr(raw_response, "usage", None): + input_tokens = getattr(raw_response.usage, "prompt_tokens", 0) or 0 + output_tokens = getattr(raw_response.usage, "completion_tokens", 0) or 0 + + metadata: dict[str, Any] = {} + if hasattr(raw_response, "model_dump"): + try: + metadata = raw_response.model_dump() + except Exception: + metadata = {} + except (LLMProviderError, AuthenticationError, RateLimitExceededError): + raise + except Exception as exc: + raise LLMProviderError( + provider="openai", + message=f"Malformed OpenAI response: {exc}", + ) from exc return LLMResponse( text=text, diff --git a/agentkit/llm/retry.py b/agentkit/llm/retry.py index 6ae5f89..4790b04 100644 --- a/agentkit/llm/retry.py +++ b/agentkit/llm/retry.py @@ -167,5 +167,21 @@ async def chat( return await self._retrying_chat(messages=messages, tools=tools) def __getattr__(self, name: str) -> Any: - """Delegate attribute lookups to the underlying client.""" - return getattr(self._client, name) + """Delegate attribute lookups to the underlying client. + + ``_client`` is resolved via ``__dict__`` rather than ``getattr`` on purpose: it is + absent when ``__init__`` raises part-way, and ``getattr(self, "_client")`` would + re-enter this method and recurse until the stack blew up. + """ + if name == "_client": + raise AttributeError( + f"{type(self).__name__!r} object has no attribute '_client'; " + "the wrapped client was never assigned." + ) + client = self.__dict__.get("_client") + if client is None: + raise AttributeError( + f"{type(self).__name__!r} object has no attribute {name!r}; " + "the wrapped client was never assigned." + ) + return getattr(client, name) diff --git a/agentkit/memory/redis_store.py b/agentkit/memory/redis_store.py index 39bb896..dc818d7 100644 --- a/agentkit/memory/redis_store.py +++ b/agentkit/memory/redis_store.py @@ -43,8 +43,11 @@ def truncate_messages(messages: list[Message], max_messages: int) -> list[Messag if messages and messages[0].role == Role.SYSTEM: system_msg = messages[0] non_system = messages[1:] - truncated = non_system[-max_messages:] - return [system_msg, *truncated] + # The system prompt occupies one of the max_messages slots, so only + # max_messages - 1 conversational turns fit. Taking a full tail would silently + # return max_messages + 1 entries and grow the history without bound. + budget = max_messages - 1 + return [system_msg, *non_system[-budget:]] if budget > 0 else [system_msg] return messages[-max_messages:] diff --git a/agentkit/memory/tiered.py b/agentkit/memory/tiered.py index 5d2b2c6..8718e9d 100644 --- a/agentkit/memory/tiered.py +++ b/agentkit/memory/tiered.py @@ -5,6 +5,7 @@ import logging from typing import Any +from agentkit.core.trace import StepTrace from agentkit.core.types import RunStatus from agentkit.llm.base import Message from agentkit.memory.base import MemoryStore, RunRecord, SessionStore @@ -74,3 +75,26 @@ async def get_run(self, run_id: str) -> RunRecord | None: if isinstance(record, RunRecord): return record return None + + async def get_run_trace(self, run_id: str) -> list[StepTrace]: + """Retrieve execution steps for a run from the run store. + + Forwarded so callers relying on duck-typed ``get_run_trace`` detection (such as the + trace endpoint) read the durable ``steps`` table instead of silently falling back to a + trace sink that may hold no rows. + """ + getter = getattr(self.run_store, "get_run_trace", None) + if getter is None: + return [] + return list(await getter(run_id)) + + async def list_runs( + self, + session_id: str | None = None, + limit: int = 50, + ) -> list[RunRecord]: + """List historical runs from the run store, optionally scoped to a session.""" + lister = getattr(self.run_store, "list_runs", None) + if lister is None: + return [] + return list(await lister(session_id=session_id, limit=limit)) diff --git a/agentkit/tools/builtin/calculator.py b/agentkit/tools/builtin/calculator.py index 6ea6ff4..828094b 100644 --- a/agentkit/tools/builtin/calculator.py +++ b/agentkit/tools/builtin/calculator.py @@ -8,6 +8,11 @@ MAX_AST_NODES = 100 MAX_EXPONENT = 1000 +#: Rough ceiling on the bit length of an exponentiation result. 10**1000 is ~3322 bits and +#: still cheap once, but the ReAct loop can repeat the call, so the result size itself is +#: capped rather than only the exponent. +MAX_RESULT_BITS = 4096 + def _evaluate_node(node: ast.AST) -> int | float: """Recursively evaluate an AST node against a strict whitelist of arithmetic operations.""" @@ -53,6 +58,20 @@ def _evaluate_node(node: ast.AST) -> int | float: if isinstance(node.op, ast.Pow): if abs(right) > MAX_EXPONENT or (abs(left) > 1000 and right > 100): raise ValueError(f"Exponent too large (maximum allowed is {MAX_EXPONENT})") + + # The exponent bound alone still permits enormous big integers (10**1000 is a + # 1001-digit result), and the agent loop can repeat the call. Bound the size of + # the result itself so a single expression cannot consume noticeable CPU. + # Only integer bases can grow; float exponentiation always yields a float. + if ( + isinstance(left, int) + and isinstance(right, int) + and abs(left) > 1 + and abs(right) * left.bit_length() > MAX_RESULT_BITS + ): + raise ValueError( + f"Result too large to compute safely (exceeds {MAX_RESULT_BITS} bits)." + ) return left**right raise ValueError(f"Unsupported binary operator: {type(node.op).__name__}") diff --git a/agentkit/tools/builtin/http_fetch.py b/agentkit/tools/builtin/http_fetch.py index 204f831..b4c5bfa 100644 --- a/agentkit/tools/builtin/http_fetch.py +++ b/agentkit/tools/builtin/http_fetch.py @@ -1,6 +1,10 @@ """SSRF-protected HTTP fetch tool blocking private, loopback, and metadata endpoints.""" +from __future__ import annotations + +import asyncio import ipaddress +import logging import socket from urllib.parse import urljoin, urlparse @@ -8,13 +12,24 @@ from agentkit.tools.registry import tool +logger = logging.getLogger(__name__) + DEFAULT_TIMEOUT_S = 10.0 DEFAULT_MAX_CHARS = 100_000 MAX_REDIRECTS = 3 +#: Hard ceiling on bytes read from a response body. ``max_chars`` bounds what is *returned*; +#: this bounds what is *buffered*, so an oversized or malicious payload cannot exhaust memory +#: before truncation has a chance to apply. +MAX_RESPONSE_BYTES = 8 * 1024 * 1024 + def resolve_host_ips(hostname: str) -> list[str]: - """Resolve a hostname to a list of IPv4/IPv6 address strings.""" + """Resolve a hostname to a list of IPv4/IPv6 address strings. + + Blocking by design (``socket.getaddrinfo`` performs synchronous network I/O); async callers + must go through :func:`resolve_host_ips_async` so the event loop is not stalled. + """ try: # Check if already a valid literal IP ipaddress.ip_address(hostname) @@ -29,11 +44,43 @@ def resolve_host_ips(hostname: str) -> list[str]: raise ValueError(f"Failed to resolve host '{hostname}': {exc}") from exc -def validate_url_ssrf(url: str) -> None: +async def resolve_host_ips_async(hostname: str) -> list[str]: + """Resolve a hostname off the event loop. + + DNS resolution is blocking network I/O; running it inline would stall every concurrent + request for the duration of the resolver timeout. + """ + return await asyncio.to_thread(resolve_host_ips, hostname) + + +def is_blocked_ip(ip_str: str) -> bool: + """Return True when an IP address must not be reachable through this tool.""" + try: + ip = ipaddress.ip_address(ip_str) + except ValueError: + # An unparseable address cannot be vetted, so treat it as blocked. + return True + + return bool( + ip.is_loopback + or ip.is_private + or ip.is_link_local + or ip.is_multicast + or ip.is_reserved + or ip.is_unspecified + ) + + +def validate_url_ssrf(url: str, resolved_ips: list[str] | None = None) -> list[str]: """Validate a URL against Server-Side Request Forgery (SSRF) guardrails. Args: url: The candidate URL to inspect. + resolved_ips: Optional pre-resolved addresses, letting a caller resolve once and reuse + the result instead of paying for a second lookup. + + Returns: + The list of validated public IP addresses the host resolves to. Raises: ValueError: If scheme is non-HTTP/HTTPS, or host resolves to loopback, @@ -53,28 +100,73 @@ def validate_url_ssrf(url: str) -> None: raise ValueError(f"SSRF protection blocked access to restricted host '{clean_host}'.") # Resolve all candidate IPs (both IPv4 and IPv6) - resolved_ips = resolve_host_ips(clean_host) - if not resolved_ips: + ips = resolved_ips if resolved_ips is not None else resolve_host_ips(clean_host) + if not ips: raise ValueError(f"No IP addresses could be resolved for host '{clean_host}'.") - for ip_str in resolved_ips: - try: - ip = ipaddress.ip_address(ip_str) - except ValueError as exc: - raise ValueError(f"Invalid resolved IP '{ip_str}' for host '{clean_host}'.") from exc - - if ( - ip.is_loopback - or ip.is_private - or ip.is_link_local - or ip.is_multicast - or ip.is_reserved - or ip.is_unspecified - ): + for ip_str in ips: + if is_blocked_ip(ip_str): raise ValueError( f"SSRF protection blocked access to restricted IP address: {ip_str} for host '{clean_host}'." ) + return ips + + +async def validate_url_ssrf_async(url: str) -> list[str]: + """Async SSRF validation returning the resolved, vetted public IPs for the host.""" + parsed = urlparse(url) + hostname = parsed.hostname + scheme = (parsed.scheme or "").lower() + if scheme not in ("http", "https"): + raise ValueError(f"Disallowed URL scheme: '{scheme}'. Only http and https are allowed.") + if not hostname: + raise ValueError("URL must contain a valid hostname.") + + ips = await resolve_host_ips_async(hostname.strip().lower()) + return validate_url_ssrf(url, resolved_ips=ips) + + +def verify_peer_address(response: httpx.Response, allowed_ips: list[str]) -> None: + """Confirm the socket actually connected to an IP that passed validation. + + Validating a hostname and then letting the HTTP client resolve it again leaves a + time-of-check/time-of-use gap: a short-TTL DNS answer can hand back a public address + during validation and a loopback or metadata address during connection. Checking the + real peer address of the established socket narrows that gap: a rebound connection is + refused before any response byte is read, followed, or returned. It cannot unsend the + outbound request itself, so this complements rather than replaces pre-flight validation. + + Raises: + ValueError: If the connection landed on an address outside ``allowed_ips``. + """ + stream = response.extensions.get("network_stream") + if stream is None: + # No transport introspection available (some transports and test doubles); the + # pre-flight validation above is the only guard in that case. + logger.debug("No network stream available to verify peer address; skipping peer check.") + return + + server_addr = stream.get_extra_info("server_addr") + if not isinstance(server_addr, (tuple, list)) or not server_addr: + # No usable peer address from this transport; the pre-flight validation above + # remains the only guard, same as when no network stream is exposed at all. + logger.debug("No verifiable peer address; skipping peer check.") + return + + peer_ip = server_addr[0] + if not isinstance(peer_ip, str) or is_blocked_ip(peer_ip): + raise ValueError( + f"SSRF protection blocked connection to '{peer_ip}': the peer address is restricted." + ) + normalized_allowed = {str(ipaddress.ip_address(ip)) for ip in allowed_ips} + normalized_peer = str(ipaddress.ip_address(peer_ip)) + if normalized_peer not in normalized_allowed: + raise ValueError( + f"SSRF protection blocked connection to '{peer_ip}': it was not among the validated " + f"addresses for the requested host ({', '.join(sorted(normalized_allowed))})." + ) + @tool async def http_fetch( @@ -85,25 +177,41 @@ async def http_fetch( """Fetch the text content of a public web page via HTTP GET. Protected against SSRF: private IPs, localhost, and cloud metadata endpoints are blocked. - Redirects are validated individually against SSRF restrictions. + Redirects are validated individually against SSRF restrictions, and the connected peer + address is re-checked against the validated addresses. """ current_url = url + read_budget = min(MAX_RESPONSE_BYTES, max(max_chars, 1)) async with httpx.AsyncClient(timeout=httpx.Timeout(timeout_s)) as client: for _ in range(MAX_REDIRECTS + 1): - validate_url_ssrf(current_url) - - response = await client.get(current_url, follow_redirects=False) - - if response.status_code in (301, 302, 303, 307, 308) and "Location" in response.headers: - next_url = urljoin(current_url, response.headers["Location"]) - current_url = next_url - continue - - response.raise_for_status() - text = response.text - if len(text) > max_chars: + allowed_ips = await validate_url_ssrf_async(current_url) + + async with client.stream("GET", current_url, follow_redirects=False) as response: + verify_peer_address(response, allowed_ips) + + if ( + response.status_code in (301, 302, 303, 307, 308) + and "Location" in response.headers + ): + next_url = urljoin(current_url, response.headers["Location"]) + current_url = next_url + continue + + response.raise_for_status() + + chunks: list[bytes] = [] + total = 0 + async for chunk in response.aiter_bytes(): + remaining = read_budget - total + if remaining <= 0: + break + if len(chunk) > remaining: + chunk = chunk[:remaining] + chunks.append(chunk) + total += len(chunk) + + text = b"".join(chunks).decode(response.encoding or "utf-8", errors="replace") return text[:max_chars] - return text raise ValueError(f"Exceeded maximum allowed redirects ({MAX_REDIRECTS}).") diff --git a/agentkit/tools/builtin/sql_readonly.py b/agentkit/tools/builtin/sql_readonly.py index 88f04c7..fa3e905 100644 --- a/agentkit/tools/builtin/sql_readonly.py +++ b/agentkit/tools/builtin/sql_readonly.py @@ -1,16 +1,19 @@ """Read-only SQL tool with strict single-SELECT validation and DDL/DML rejection.""" import asyncio +import logging import re from collections.abc import Sequence from typing import Any from sqlalchemy import text -from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine +from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine from agentkit.config import get_settings from agentkit.tools.registry import tool +logger = logging.getLogger(__name__) + DEFAULT_ROW_LIMIT = 50 DEFAULT_STATEMENT_TIMEOUT_S = 10.0 @@ -45,18 +48,44 @@ re.IGNORECASE, ) +#: A '#' that is not the start of a JSON path operator (#>, #>>, #-), i.e. a comment marker. +_BARE_HASH_REGEX = re.compile(r"#(?![>~-])") + +#: A row-locking clause. LIMIT cannot simply be appended after one of these, so such +#: statements take the subquery-wrap fallback instead. +_LOCKING_CLAUSE_REGEX = re.compile( + r"\bfor\s+(update|share|no\s+key\s+update|key\s+share)\b", + re.IGNORECASE, +) + +#: A clause that already bounds the statement's own output. Such statements are executed +#: untouched so a caller-supplied bound is never widened. +_EXISTING_BOUND_REGEX = re.compile( + r"(\blimit\b|\bfetch\s+first\b)", + re.IGNORECASE, +) + +#: Process-wide engine, created on first use. Connection pools are expensive to rebuild and +#: are safe to share, but the engine is still injectable per call so tests and alternative +#: deployments can supply their own without mutating module state. _readonly_engine: AsyncEngine | None = None def get_readonly_engine() -> AsyncEngine: - """Return singleton read-only SQLAlchemy AsyncEngine.""" + """Return the shared read-only SQLAlchemy AsyncEngine, creating it on first use.""" global _readonly_engine if _readonly_engine is None: settings = get_settings() - _readonly_engine = create_async_engine(settings.DATABASE_URL) + _readonly_engine = create_async_engine(settings.DATABASE_URL, pool_pre_ping=True) return _readonly_engine +def set_readonly_engine(engine: AsyncEngine | None) -> None: + """Override or clear the shared read-only engine (dependency injection seam for tests).""" + global _readonly_engine + _readonly_engine = engine + + def validate_sql_query(query: str) -> str: """Validate a SQL query string to ensure it is a safe, single read-only SELECT statement. @@ -74,15 +103,30 @@ def validate_sql_query(query: str) -> str: if not stripped: raise ValueError("Empty or whitespace-only SQL query.") - # Reject comments that could hide payload structure - if "--" in stripped or "/*" in stripped or "*/" in stripped or "#" in stripped: + # Reject comments that could hide payload structure. The scan runs on the query with + # string literals and quoted identifiers blanked out, so comment markers that are merely + # data (WHERE note = 'x -- y', or a '#' inside a literal) do not cause false rejections. + # A marker can only be hidden inside quotes when it really is inside a literal, in which + # case the database parses it as data too. A bare '#' is allowed when it begins a + # PostgreSQL JSON path operator (#>, #>>, #-); those are legitimate in a read-only + # projection. + masked_structure = _mask_literals(stripped) + if ( + "--" in masked_structure + or "/*" in masked_structure + or "*/" in masked_structure + or _BARE_HASH_REGEX.search(masked_structure) + ): raise ValueError( "Comments are not permitted in SQL queries to prevent security circumvention." ) - # Strip single trailing semicolon if present - cleaned = stripped.rstrip(";").strip() - if ";" in cleaned: + # Strip a single trailing semicolon if present. rstrip(";") would also swallow the + # repeated separators that signal a stacked-statement payload. The multi-statement scan + # likewise runs on the masked structure so a ';' inside a string literal is data, not a + # statement separator; a real separator outside any literal is still caught. + cleaned = stripped[:-1].strip() if stripped.endswith(";") else stripped + if ";" in _mask_literals(cleaned): raise ValueError("Multiple SQL statements separated by semicolons are strictly prohibited.") # Validate statement header: must begin with SELECT or WITH @@ -90,8 +134,12 @@ def validate_sql_query(query: str) -> str: if not (lower.startswith("select") or lower.startswith("with")): raise ValueError("Only SELECT statements are permitted.") - # Check for forbidden DDL, DML, or administrative keywords - matches = _KEYWORD_REGEX.findall(cleaned) + # Check for forbidden DDL, DML, or administrative keywords. Scanning the query with its + # string literals and quoted identifiers blanked out avoids rejecting harmless queries + # like SELECT * FROM events WHERE action = 'insert', where the keyword only appears as + # data. The literals are still covered by the single-statement check above, since an + # embedded ';' inside a literal cannot be resolved without full SQL parsing. + matches = _KEYWORD_REGEX.findall(_mask_literals(cleaned)) if matches: prohibited = sorted({m.upper() for m in matches}) raise ValueError(f"Query contains prohibited SQL keyword(s): {', '.join(prohibited)}") @@ -99,6 +147,52 @@ def validate_sql_query(query: str) -> str: return cleaned +def _mask_literals(sql: str) -> str: + """Replace string literals and quoted identifiers with placeholders. + + Keeps the keyword scan focused on SQL structure instead of matching words that merely + appear inside a quoted value or a quoted column name. + """ + masked: list[str] = [] + index = 0 + length = len(sql) + + while index < length: + char = sql[index] + + if char == "'": + masked.append(" ") + index += 1 + while index < length: + if sql[index] == "'": + # Doubled quote is an escaped quote inside the literal, not its end. + if index + 1 < length and sql[index + 1] == "'": + index += 2 + continue + index += 1 + break + index += 1 + continue + + if char == '"': + masked.append(" ") + index += 1 + while index < length: + if sql[index] == '"': + if index + 1 < length and sql[index + 1] == '"': + index += 2 + continue + index += 1 + break + index += 1 + continue + + masked.append(char) + index += 1 + + return "".join(masked) + + def format_rows_as_markdown(columns: Sequence[str], rows: Sequence[Any]) -> str: """Format tabular SQL results into clean Markdown table string.""" if not rows: @@ -121,6 +215,11 @@ async def execute_sql_query( ) -> str: """Validate and execute a read-only SQL query, returning a formatted markdown table. + The row limit is applied in SQL, not just client-side. ``fetchmany(row_limit)`` alone still + makes the database plan and execute the full scan, so a query over a large table burns + server resources even though only a handful of rows are returned; the client-side + ``asyncio.timeout`` cancels late and leaves the work already done. + Args: query: SQL SELECT statement. engine: Optional custom AsyncEngine override (e.g. for testing). @@ -135,12 +234,77 @@ async def execute_sql_query( async with asyncio.timeout(timeout_s): async with target_engine.connect() as conn: - cursor = await conn.execute(text(clean_sql)) + await _apply_statement_timeout(conn, timeout_s) + cursor = await conn.execute(text(_apply_row_limit(clean_sql, row_limit))) columns = list(cursor.keys()) rows = cursor.fetchmany(row_limit) return format_rows_as_markdown(columns, rows) +async def _apply_statement_timeout(conn: AsyncConnection, timeout_s: float) -> None: + """Push the timeout down to the database where the dialect supports it. + + The client-side ``asyncio.timeout`` only cancels the wait; the server keeps running the + statement. ``statement_timeout`` aborts it at the source. SQLite has no equivalent, so the + client-side timeout remains the only guard there. + """ + dialect = conn.dialect.name + if dialect not in ("postgresql", "mysql", "mariadb"): + return + + try: + if dialect == "postgresql": + await conn.execute(text(f"SET LOCAL statement_timeout = {int(timeout_s * 1000)}")) + else: + # MAX_EXECUTION_TIME is in milliseconds and applies session-wide to subsequent + # top-level SELECT statements, unlike SET STATEMENT ... FOR which would only + # cover the single statement it prefixes. + await conn.execute(text(f"SET SESSION MAX_EXECUTION_TIME = {int(timeout_s * 1000)}")) + except Exception as exc: + # A restricted role may lack permission to set this; the client-side timeout and the + # in-SQL row limit still bound the work, so this must not fail the query. + logger.warning("Could not apply server-side statement timeout: %s", exc) + + +def _top_level_tail(masked_sql: str) -> str: + """Return the trailing top-level clause of a masked statement. + + Clause keywords are only meaningful at parenthesis depth zero: a LIMIT inside a subquery + does not bound the outer statement. Parenthesis counting on the literal-masked text is + exact for real queries, since quoted parentheses are already blanked out. + """ + depth = 0 + tail_start = 0 + for index, char in enumerate(masked_sql): + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth <= 0: + depth = 0 + tail_start = index + 1 + return masked_sql[tail_start:] + + +def _apply_row_limit(sql: str, row_limit: int) -> str: + """Bound the statement's output in SQL, preferring an appended LIMIT. + + Appending ``LIMIT`` leaves an already-valid statement untouched in shape, so constructs + like duplicate output column names in a JOIN keep working exactly as written. Only two + cases need different handling, both detected on the trailing top-level clause: + + - the statement already bounds itself (LIMIT / FETCH FIRST): executed untouched; + - the statement ends in a row-locking clause (FOR UPDATE / FOR SHARE), after which a + LIMIT cannot legally appear: wrapped in a limited subquery instead. + """ + tail = _top_level_tail(_mask_literals(sql)) + if _EXISTING_BOUND_REGEX.search(tail): + return sql + if _LOCKING_CLAUSE_REGEX.search(tail): + return f"SELECT * FROM ({sql}) AS agentkit_limited_query LIMIT {int(row_limit)}" + return f"{sql} LIMIT {int(row_limit)}" + + @tool async def sql_readonly(query: str) -> str: """Execute a read-only SQL query against the database and return results as text. diff --git a/docker-compose.yml b/docker-compose.yml index 0a7ed87..c880a60 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -49,6 +49,10 @@ services: condition: service_healthy environment: - DATABASE_URL=${DATABASE_URL:-postgresql+asyncpg://postgres:postgres@postgres:5432/agentkit} + # Settings validation requires API_KEYS, and both alembic and the seed script + # construct Settings. Passed through so the compose stack fails closed here, + # with the same validation error, rather than in the api service. + - API_KEYS=${API_KEYS:-} networks: - agentkit-net diff --git a/docs/DECISIONS.md b/docs/DECISIONS.md index 3aa7471..5fc1675 100644 --- a/docs/DECISIONS.md +++ b/docs/DECISIONS.md @@ -118,10 +118,63 @@ This log records significant architectural, technical, and design decisions made - **Decision**: - Update `GeminiClient.chat()` to wrap function responses with `role="user"` using `types.Part.from_function_response(name=tool_name, response={"result": msg.content})`. - Standardize the default active Gemini model to `gemini-3.1-flash-lite` across `Settings`, `GeminiClient`, `create_llm_client`, `.env.example`, `docker-compose.yml`, and documentation. - - Wire `RetryingLLMClient` with configurable `MAX_RETRIES` (default 3) and exponential backoff into `create_llm_client_from_settings()`. + - Wire `RetryingLLMClient` with configurable `LLM_MAX_RETRIES` (default 3) and exponential backoff into `create_llm_client_from_settings()`. - **Consequences**: - Eliminates 400 Invalid Argument errors on multi-turn ReAct tool execution with Gemini. - Consistent model behavior across local, containerized, and CI environments. - Automatic resilience against transient 429 quota exhaustion and 5xx network spikes. +--- + +## ADR-009: Fail-Closed Configuration and Dependency Resolution + +- **Date**: 2026-10-01 +- **Status**: Accepted +- **Context**: A whole-project audit found several paths where a misconfiguration or a transient fault was silently converted into a working-looking success. `Settings.API_KEYS` defaulted to `ak_test_key_12345`, a key published in the repository, so any deployment without a `.env` came up openly accessible. `get_llm_client` caught every exception from the provider factory and returned `FakeLLMClient`, so a missing `GEMINI_API_KEY` produced HTTP 200 with scripted text instead of a 503. The `app.state.agent` override was indistinguishable from the memoized first assembly that AUD-01 removed. `RetryingLLMClient` was wired by AUD-04 but with only a retry count: the backoff delays were fixed defaults rather than configurable settings. `TieredMemoryStore` forwarded only the run methods, so duck-typed `get_run_trace` detection in the trace endpoint fell through to a trace sink that may hold no rows. +- **Decision**: + - Make `API_KEYS` required with no default; `Settings` validation rejects an empty or whitespace-only value. + - Remove the `FakeLLMClient` fallback from the request path. Provider construction failure raises `ServiceUnavailableError` (503). `LLM_PROVIDER=fake` remains an explicit opt-in resolved by the factory. + - Keep retry wrapping in `create_llm_client_from_settings`, now configurable through `LLM_MAX_RETRIES`, `LLM_RETRY_BASE_DELAY_S` and `LLM_RETRY_MAX_DELAY_S`. An explicit `app.state.llm_client` override still wins and is not wrapped, keeping tests deterministic. + - Assemble the `Agent` per request from resolved dependencies; an explicit `app.state.agent_override` remains the only override mechanism. + - Forward `get_run_trace` and `list_runs` through `TieredMemoryStore` to the underlying run store, returning empty lists when the run store does not support them. +- **Consequences**: + - A misconfigured deployment fails loudly at startup or on the first request instead of serving fake or unauthenticated traffic. + - No request is served by a dependency graph assembled for an earlier request. + - Transient provider failures (429/5xx/network) are retried with configurable exponential backoff and jitter; authentication and other 4xx failures are not. + - The test suite seeds a fixed `API_KEYS` before importing `agentkit.api.main`, whose module-level `create_app()` would otherwise fail validation. + +--- + +## ADR-010: Post-Connect SSRF Verification and Bounded Resource Use in Tools + +- **Date**: 2026-10-01 +- **Status**: Accepted +- **Context**: `http_fetch` resolved the hostname, checked the resolved addresses, then let `httpx` resolve the hostname again when connecting. That time-of-check/time-of-use gap lets a short-TTL DNS answer hand back a public address during validation and a loopback or `169.254.169.254` address during connection. DNS resolution was synchronous `socket.getaddrinfo` called directly from async code, stalling the event loop, and the full response body was buffered before being truncated to `max_chars`, so a large payload could exhaust memory. Separately, `sql_readonly` applied `fetchmany(row_limit)` client-side only, which still makes the database plan and execute the full scan, and its keyword scan rejected any query containing a write keyword even inside a string literal, while its bare `#` check rejected the legitimate PostgreSQL JSON operators `#>` and `#-`. +- **Decision**: + - Resolve DNS via `asyncio.to_thread` and reuse the resolved addresses for both validation and the connection decision. + - After connecting, read `server_addr` from the response's network stream and reject the request if the peer is not among the validated addresses. + - Stream the body with an explicit byte budget instead of buffering then slicing. + - Apply the row limit in SQL by wrapping an unbounded statement, leaving a caller-supplied top-level `LIMIT` untouched, and push `statement_timeout` to PostgreSQL where the dialect supports it. + - Mask string literals and quoted identifiers before the keyword scan, and treat `#` as a comment marker only when it does not begin a JSON path operator. +- **Consequences**: + - DNS rebinding to a private or metadata address is refused after the connection is established, not merely before it. + - No blocking DNS on the event loop, and response memory is bounded regardless of `max_chars`. + - Legitimate read-only queries that merely mention a write keyword as data, or use JSON path operators, are accepted; write statements remain blocked. + +--- +## ADR-011: Atomic Fallbacks and Domain-Error Boundaries + +- **Date**: 2026-10-01 +- **Status**: Accepted +- **Context**: Two fallbacks reintroduced the exact races their primary paths had been written to eliminate, and several error boundaries leaked raw SDK exceptions. The rate limiter's non-scripting fallback trimmed and counted in one pipeline then inserted in a second, so concurrent requests between the two observed the same pre-insert cardinality and were admitted together. In both LLM adapters the `try` block wrapped only the SDK call, so response parsing escaped as raw provider exceptions: reading `raw_response.text` raises on a blocked prompt, an empty `choices`/`candidates` list raised `IndexError`, and malformed tool-call JSON was rewritten to `{"raw": ...}`, which later failed argument validation with a misleading "missing required field" message. `RetryingLLMClient.__getattr__` recursed infinitely when `_client` was absent, and `/health` echoed `str(exc)` from SQLAlchemy and Redis on an unauthenticated endpoint. +- **Decision**: + - Implement the non-scripting rate limit fallback with `WATCH`/`MULTI`/`EXEC` optimistic locking, denying the request if contention never settles within the retry budget. + - Wrap response parsing in both adapters inside the error boundary, re-raising domain errors unchanged and mapping anything else to `LLMProviderError`. Report malformed tool-call JSON as such instead of substituting a placeholder object. + - Use full UUIDs for generated `tool_call_id` values; a truncated 32-bit id risks collisions within a long session. + - Guard `RetryingLLMClient.__getattr__` against re-entering itself when `_client` was never assigned. + - Return fixed client-facing strings from `/health` and log the underlying driver error instead. +- **Consequences**: + - The fallback enforces the documented limit under concurrency, erring toward denial where it cannot decide. + - Every provider fault is matchable on `LLMProviderError`, so HTTP status mapping stays correct for malformed as well as failed responses. + - The unauthenticated health endpoint no longer discloses internal hostnames, ports or database names. diff --git a/docs/DEPLOY.md b/docs/DEPLOY.md index fba2b0e..15d3be9 100644 --- a/docs/DEPLOY.md +++ b/docs/DEPLOY.md @@ -204,7 +204,7 @@ GEMINI_API_KEY=${GEMINI_KEY} API_KEYS=${API_KEYS} FILE_TOOL_BASE_DIR=/tmp/agentkit_sandbox RUN_TIMEOUT_S=60 -MAX_RETRIES=3 +LLM_MAX_RETRIES=3 ENVFILE chmod 600 /opt/agentkit/.env diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 2af443f..2c7b115 100644 --- a/docs/PROGRESS.md +++ b/docs/PROGRESS.md @@ -60,7 +60,7 @@ This document tracks execution progress across all backlog tasks, including pull | AUD-08 | AUD | Expose Configurable Timeout in Docker Compose | `fix/docker-compose-timeout` | #52 | 2 | Done | Exposed RUN_TIMEOUT_S in docker-compose.yml | | AUD-09 | AUD | Live End-to-End Infrastructure Verification & Documentation Proof | `docs/e2e-verification-report` | #53 | 3 | Done | Standardized model, verified live stack, and documented in E2E_REPORT.md | | AUD-10 | AUD | Rate Limiter Outage Fallback & Provider-Valid Session History | `fix/ratelimit-and-session-repair` | #56 | 3 | Done | In-memory fallback on Redis connection failures; synthetic tool results repair aborted sessions | - +| M13-T06 | M13 | Whole-Project Audit Remediation | `fix/audit-hardening` | #47 | 3 | Done | Fail-closed config and LLM resolution, SSRF post-connect check, atomic fallback | --- @@ -401,3 +401,9 @@ This document tracks execution progress across all backlog tasks, including pull - **Commits**: 3 - **Notes**: A full-project review (issue #55) surfaced two medium-severity bugs. `RateLimiter.check` only caught `ResponseError`, so a real Redis outage (`ConnectionError` / `TimeoutError`) escaped the dependency and returned HTTP 500; it now catches `RedisError`, logs a warning, and degrades to the in-memory sliding window, extracted into `_check_window_in_memory`. `Agent._finalize_and_return` could persist histories whose trailing assistant turn requested tool calls with no matching tool messages when a run timed out mid-tool or the duplicate-loop guard aborted mid-batch, which OpenAI and Gemini reject on the next run in that session; `_repair_trailing_tool_calls` now appends synthetic error results for unanswered call ids before saving, on every terminal path. Added regression tests in `tests/unit/test_api_ratelimit.py`, `tests/unit/test_agent_guards.py`, and `tests/unit/test_agent_loop_detector.py`, each confirmed to fail pre-fix. + +### M13-T06: Whole-Project Audit Remediation +- **Status**: Done +- **PR**: #47 +- **Commits**: 3 +- **Notes**: Remediated 17 findings from a full-codebase audit, rebased onto the AUD-01/AUD-04/AUD-05 follow-ups. Critical: removed the `ak_test_key_12345` default from `Settings.API_KEYS` and made the field required; replaced the `FakeLLMClient` fallback in `get_llm_client` with a `ServiceUnavailableError` (503) boundary around provider construction; kept the AUD-04 factory-level `RetryingLLMClient` wrapping and made it configurable through new `LLM_MAX_RETRIES`/`LLM_RETRY_BASE_DELAY_S`/`LLM_RETRY_MAX_DELAY_S` settings; made the per-request `Agent` assembly's override explicit as `app.state.agent_override`; forwarded `get_run_trace`/`list_runs` through `TieredMemoryStore`; moved `http_fetch` DNS off the event loop with `asyncio.to_thread`, added a post-connect `server_addr` check against the validated addresses to close the DNS-rebinding TOCTOU window, and replaced full-body buffering with a byte-budgeted stream; replaced the rate limiter's non-scripting fallback with `WATCH`/`MULTI`/`EXEC`, leaving the AUD-10 in-memory fallback for Redis outages intact. Moderate: corrected the `RedisMemoryStore.truncate_messages` off-by-one that returned `max_messages + 1` entries when a system prompt was present; switched generated `tool_call_id`s to full UUIDs; wrapped response parsing in both LLM adapters inside the error boundary so malformed responses and malformed tool-argument JSON raise `LLMProviderError`; guarded `RetryingLLMClient.__getattr__` against recursion on a partially constructed instance; sanitized `/health` details so the unauthenticated endpoint no longer echoes driver error strings; made `sql_readonly` mask string literals and quoted identifiers before its keyword scan, allow `#>`/`#>>`/`#-`, apply the row limit in SQL and push `statement_timeout` to PostgreSQL. Minor: stopped logging raw run goals at INFO (length only, content at DEBUG); replaced `split(";")` in `scripts/seed.py` with a statement-aware splitter handling literals, comments and dollar-quotes; bounded calculator exponentiation by result bit length. Added `tests/unit/test_audit_regressions.py` and a real-Redis concurrency regression test pinning the fallback fix. Recorded in ADR-009, ADR-010, ADR-011. Suite: 298 passed, 25 skipped; coverage 92.86% on core/tools/llm; `ruff check`, `ruff format --check` and `mypy agentkit tests` all clean. diff --git a/examples/demo.py b/examples/demo.py index 3d9e68b..d909487 100644 --- a/examples/demo.py +++ b/examples/demo.py @@ -144,7 +144,12 @@ def build_fake_llm() -> LLMClient: def build_llm(offline: bool, provider: str) -> LLMClient: - """Resolve the LLM client, requiring valid API keys unless running in offline/fake mode.""" + """Resolve the LLM client for the demo. + + Offline execution has to be requested with ``--offline``. A missing API key is otherwise + an error rather than a silent downgrade to the scripted client, which would otherwise + print a canned answer that looks like a real run. + """ if offline or provider == "fake": return build_fake_llm() @@ -152,8 +157,8 @@ def build_llm(offline: bool, provider: str) -> LLMClient: key_present = settings.GEMINI_API_KEY if provider == "gemini" else settings.OPENAI_API_KEY if not key_present: key_name = "GEMINI_API_KEY" if provider == "gemini" else "OPENAI_API_KEY" - raise ValueError( - f"API key for provider '{provider}' is missing ({key_name} is not set). " + raise SystemExit( + f"No API key configured for provider '{provider}' ({key_name} is not set). " f"To run offline without real API credentials, pass --offline or --provider fake." ) @@ -180,8 +185,9 @@ async def run_demo( Returns: The agent's final answer string. """ - settings = get_settings() - target_db_url = IN_MEMORY_DB_URL if offline else (db_url or settings.DATABASE_URL) + # Settings are only read on the live path. Offline mode deliberately avoids + # get_settings() so the demo runs with no API keys and no services at all. + target_db_url = IN_MEMORY_DB_URL if offline else (db_url or get_settings().DATABASE_URL) query_engine = create_async_engine(target_db_url, pool_pre_ping=True) await seed_demo_database(query_engine) diff --git a/scripts/seed.py b/scripts/seed.py index bf697b8..0714872 100644 --- a/scripts/seed.py +++ b/scripts/seed.py @@ -15,6 +15,97 @@ logger = logging.getLogger("agentkit.seed") +def split_sql_statements(sql_content: str) -> list[str]: + """Split a SQL script into individual statements. + + ``str.split(";")`` is wrong for anything but the simplest script: a semicolon inside a + string literal, a dollar-quoted body or a line comment would split a statement in half + and leave a fragment that fails to parse. + + Args: + sql_content: Full text of a SQL script. + + Returns: + Non-empty, stripped statement strings in source order. + """ + statements: list[str] = [] + current: list[str] = [] + index = 0 + length = len(sql_content) + + while index < length: + char = sql_content[index] + + if char == "-" and sql_content.startswith("--", index): + end = sql_content.find("\n", index) + index = length if end == -1 else end + continue + + if char == "/" and sql_content.startswith("/*", index): + end = sql_content.find("*/", index + 2) + index = length if end == -1 else end + 2 + continue + + if char == "'": + closing = _scan_quoted(sql_content, index, "'") + current.append(sql_content[index:closing]) + index = closing + continue + + if char == '"': + closing = _scan_quoted(sql_content, index, '"') + current.append(sql_content[index:closing]) + index = closing + continue + + if char == "$": + tag_end = sql_content.find("$", index + 1) + if tag_end != -1 and _is_dollar_tag(sql_content[index + 1 : tag_end]): + tag = sql_content[index : tag_end + 1] + closing = sql_content.find(tag, tag_end + 1) + end = length if closing == -1 else closing + len(tag) + current.append(sql_content[index:end]) + index = end + continue + + if char == ";": + statements.append("".join(current).strip()) + current = [] + index += 1 + continue + + current.append(char) + index += 1 + + statements.append("".join(current).strip()) + return [stmt for stmt in statements if stmt] + + +def _is_dollar_tag(candidate: str) -> bool: + """Return True when a run of characters is a valid dollar-quote tag body. + + An empty body is valid and means the plain ``$$`` delimiter form. + """ + return candidate == "" or candidate.replace("_", "").isalnum() + + +def _scan_quoted(sql_content: str, start: int, quote: str) -> int: + """Return the index just past a quoted run starting at ``start``. + + A doubled quote character is an escaped quote inside the literal, not its terminator. + """ + index = start + 1 + length = len(sql_content) + while index < length: + if sql_content[index] == quote: + if index + 1 < length and sql_content[index + 1] == quote: + index += 2 + continue + return index + 1 + index += 1 + return length + + async def seed_database(engine: AsyncEngine, script_path: Path | None = None) -> None: """Execute demo SQL seed script against the provided database engine. @@ -29,7 +120,7 @@ async def seed_database(engine: AsyncEngine, script_path: Path | None = None) -> raise FileNotFoundError(f"Seed script not found at {script_path}") sql_content = script_path.read_text(encoding="utf-8") - statements = [stmt.strip() for stmt in sql_content.split(";") if stmt.strip()] + statements = split_sql_statements(sql_content) logger.info("Applying %d SQL statements from %s", len(statements), script_path.name) async with engine.begin() as conn: diff --git a/tests/conftest.py b/tests/conftest.py index dafca8a..0e6c617 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -34,6 +34,13 @@ REAL_INFRA_ENV_FLAG = "AGENTKIT_REAL_INFRA" TRUTHY_VALUES = frozenset({"1", "true", "yes", "on"}) +# Settings now refuses to validate without an explicit API_KEYS, which is the intended +# fail-closed behaviour for a deployment but would break any test that reaches the module +# level ``app = create_app()`` in agentkit.api.main. A fixed test key is seeded before that +# import happens, and an operator-supplied value always wins so the suite stays hermetic +# regardless of the developer's local .env. +os.environ.setdefault("API_KEYS", "test_api_key_env_default") + #: Only AgentKit-owned key prefixes are removed between tests so pointing the #: suite at a shared development Redis cannot destroy unrelated data. MANAGED_REDIS_PATTERNS = ("session:*", "ratelimit:*") diff --git a/tests/integration/test_api_deps_strict.py b/tests/integration/test_api_deps_strict.py index 188eb0f..2f84dca 100644 --- a/tests/integration/test_api_deps_strict.py +++ b/tests/integration/test_api_deps_strict.py @@ -13,7 +13,7 @@ ) from agentkit.api.main import create_app from agentkit.config import Settings -from agentkit.core.errors import AuthenticationError, ServiceUnavailableError +from agentkit.core.errors import ServiceUnavailableError from agentkit.llm.fake import FakeLLMClient @@ -63,14 +63,15 @@ def test_llm_client_does_not_silently_fallback_to_fake() -> None: GEMINI_API_KEY="", ) req = MockRequest() - # If GEMINI_API_KEY is empty/invalid and ADC fails, it must raise AuthenticationError, not return FakeLLMClient + # An empty key must surface as ServiceUnavailableError (503), never a silent FakeLLMClient. try: client = get_llm_client(req, settings=settings) # type: ignore[arg-type] + except ServiceUnavailableError as exc: + assert exc.service == "llm" + else: assert not isinstance(client, FakeLLMClient), ( "get_llm_client must not return FakeLLMClient when LLM_PROVIDER is 'gemini'" ) - except AuthenticationError: - pass # Expected authentic error, not swallowed into FakeLLMClient def test_llm_client_returns_fake_when_explicitly_configured() -> None: diff --git a/tests/integration/test_demo.py b/tests/integration/test_demo.py index e75b8a3..df250c3 100644 --- a/tests/integration/test_demo.py +++ b/tests/integration/test_demo.py @@ -150,11 +150,11 @@ def test_build_llm_strict_live_mode(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("OPENAI_API_KEY", "") get_settings.cache_clear() - # In live mode without key, it must raise ValueError - with pytest.raises(ValueError, match="API key for provider 'gemini' is missing"): + # In live mode without key, it must exit with a clear CLI error rather than fake an answer. + with pytest.raises(SystemExit, match="No API key configured for provider 'gemini'"): build_llm(offline=False, provider="gemini") - with pytest.raises(ValueError, match="API key for provider 'openai' is missing"): + with pytest.raises(SystemExit, match="No API key configured for provider 'openai'"): build_llm(offline=False, provider="openai") # In offline or fake mode, it returns FakeLLMClient diff --git a/tests/real/test_redis_real.py b/tests/real/test_redis_real.py index 1075755..d116ed8 100644 --- a/tests/real/test_redis_real.py +++ b/tests/real/test_redis_real.py @@ -10,6 +10,7 @@ import pytest from redis.asyncio import Redis +from redis.exceptions import ResponseError from agentkit.api.ratelimit import RateLimiter from agentkit.core.errors import RateLimitExceededError @@ -81,7 +82,9 @@ async def test_sliding_window_truncation_is_persisted(real_redis: Redis[str]) -> ) loaded = await store.get_messages("sess-truncate") - assert len(loaded) == 4 + # max_messages is a total budget and the system prompt spends one of those slots, + # so 3 means the system prompt plus the 2 most recent messages. + assert len(loaded) == 3 assert loaded[0].content == "System" assert loaded[-1].content == "five" @@ -163,3 +166,38 @@ async def test_concurrent_checks_respect_limit(real_redis: Redis[str]) -> None: assert admitted == 5 assert rejected == 15 assert await real_redis.zcard(f"ratelimit:{key}") == 5 + + +async def test_non_scripting_fallback_does_not_overshoot_limit( + real_redis: Redis[str], monkeypatch: pytest.MonkeyPatch +) -> None: + """The scripting-free fallback must enforce the limit under concurrency too. + + The previous fallback trimmed and counted in one transaction, then inserted in a second. + Against live Redis that admitted 39 of 50 simultaneous requests against a limit of 5, + because the intervening round trip let every concurrent caller observe the same + pre-insert cardinality. The WATCH/MULTI/EXEC fallback closes that window. + """ + limiter = RateLimiter(requests_per_window=5, window_seconds=60) + + async def _no_scripting(*_args: object, **_kwargs: object) -> None: + raise ResponseError("NOSCRIPTING this Redis has scripting disabled") + + monkeypatch.setattr(Redis, "eval", _no_scripting) + + outcomes = await asyncio.gather( + *( + limiter._check_window_unatomic( + real_redis, "ratelimit:ak_fallback_concurrent", 1000.0, 60.0, 5, f"m{i}", 65 + ) + for i in range(50) + ), + return_exceptions=True, + ) + + assert not [o for o in outcomes if isinstance(o, BaseException)], outcomes + admitted = sum(1 for o in outcomes if o == (True, 0)) + denied = sum(1 for o in outcomes if isinstance(o, tuple) and o[0] is False) + + assert admitted == 5, f"expected 5 admissions, got {admitted}" + assert denied == 45 diff --git a/tests/unit/test_audit_regressions.py b/tests/unit/test_audit_regressions.py new file mode 100644 index 0000000..e7c0d9d --- /dev/null +++ b/tests/unit/test_audit_regressions.py @@ -0,0 +1,592 @@ +"""Regression tests for the audit hardening pass. + +Each test pins behaviour that was previously wrong, so a later refactor cannot silently +reintroduce the bug. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, patch + +import fakeredis.aioredis +import pytest +from fastapi.testclient import TestClient +from pydantic import ValidationError +from redis.asyncio import Redis +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from agentkit.api.deps import get_llm_client +from agentkit.api.main import create_app +from agentkit.api.ratelimit import RateLimiter +from agentkit.config import Settings +from agentkit.core.errors import AuthenticationError, ServiceUnavailableError +from agentkit.core.trace import InMemoryTraceSink, StepTrace +from agentkit.llm.base import LLMClient, LLMResponse, Message +from agentkit.llm.retry import RetryingLLMClient +from agentkit.memory.base import InMemoryMemoryStore, RunRecord +from agentkit.memory.tiered import TieredMemoryStore +from agentkit.tools.builtin.calculator import MAX_RESULT_BITS, safe_calculate +from agentkit.tools.builtin.sql_readonly import validate_sql_query + + +class _State: + """Stand-in for ``app.state`` exposing only the attributes ``get_llm_client`` reads.""" + + def __init__(self, **overrides: Any) -> None: + self.llm_client: Any = None + self.settings = Settings(API_KEYS="k", **overrides) + + +class _Request: + """Minimal Request stand-in exposing only the ``app.state`` surface deps.py reads. + + Typed as ``Any`` at the call site because ``get_llm_client`` declares a Starlette + ``Request``, and constructing a real one just to read ``app.state`` adds nothing. + """ + + def __init__(self, state: Any) -> None: + self.app: Any = type("App", (), {"state": state})() + + +def _request(state: _State) -> Any: + """Build the Request stand-in, typed loosely to match the FastAPI dependency signature.""" + return _Request(state) + + +# --- config: fail closed on API_KEYS ------------------------------------------------- + + +def test_api_keys_must_be_explicit() -> None: + """A deployment cannot start with a key published in the repository.""" + with pytest.raises(ValidationError, match="API_KEYS must be set"): + Settings(API_KEYS="") + + +def test_api_keys_rejects_whitespace_only() -> None: + with pytest.raises(ValidationError, match="API_KEYS must be set"): + Settings(API_KEYS=" ") + + +# --- deps: LLM client fail-closed and retry-wired ------------------------------------- + + +def test_get_llm_client_delegates_to_factory() -> None: + """Deps must return the factory client unchanged. + + Retry wiring lives in ``create_llm_client_from_settings``; wrapping again here would + double the retry budget and skip the app.state override contract. + """ + settings = Settings(API_KEYS="k", LLM_PROVIDER="openai", OPENAI_API_KEY="sk-test") + produced = AsyncMock(spec=LLMClient) + + with patch( + "agentkit.api.deps.create_llm_client_from_settings", + return_value=produced, + ) as mock_factory: + client = get_llm_client(_request(_State()), settings=settings) + + mock_factory.assert_called_once_with(settings) + assert client is produced + + +def test_get_llm_client_raises_when_provider_unavailable() -> None: + """Misconfiguration must surface as an error, not as a scripted 200 response.""" + settings = Settings(API_KEYS="k", LLM_PROVIDER="gemini") + + with ( + patch( + "agentkit.api.deps.create_llm_client_from_settings", + side_effect=AuthenticationError("no key"), + ), + pytest.raises(ServiceUnavailableError) as excinfo, + ): + get_llm_client(_request(_State()), settings=settings) + + assert excinfo.value.service == "llm" + + +def test_get_llm_client_honours_explicit_override() -> None: + """An app.state override wins and is not wrapped, so tests stay deterministic.""" + settings = Settings(API_KEYS="k", LLM_PROVIDER="gemini") + state = _State() + sentinel = AsyncMock(spec=LLMClient) + state.llm_client = sentinel + + assert get_llm_client(_request(state), settings=settings) is sentinel + + +# --- deps: the agent must not be pinned to the first request's dependencies ----------- + + +class _ProbeLLM(LLMClient): + """LLM stub that identifies itself, so a stale-agent bug is visible in the answer.""" + + def __init__(self, tag: str) -> None: + self.tag = tag + + async def chat( + self, + messages: list[Message], # noqa: ARG002 - LLMClient signature + tools: list[Any] | None = None, # noqa: ARG002 - LLMClient signature + ) -> LLMResponse: + return LLMResponse(text=f"answer from {self.tag}") + + +async def test_agent_reflects_dependency_changes_between_requests() -> None: + """Rebuilding per request means a later LLM override is not served by a stale agent. + + The old implementation memoized the first assembled Agent on app.state, so the second + request reused the first request's LLM client regardless of any override. + """ + from agentkit.db.models import Base + + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + factory = async_sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + app = create_app(settings=Settings(API_KEYS="secret", LLM_PROVIDER="fake")) + app.state.session_factory = factory + # In-memory overrides keep the request path hermetic: strict dependency resolution + # (AUD-01) otherwise requires live Redis and PostgreSQL for memory and traces. + app.state.memory_store = InMemoryMemoryStore() + app.state.trace_sink = InMemoryTraceSink() + app.state.llm_client = _ProbeLLM("first") + # No context manager: that would run the lifespan, which opens a real Redis pool. + client = TestClient(app, raise_server_exceptions=False) + headers = {"X-API-Key": "secret"} + + resp = client.post("/runs", json={"goal": "hi"}, headers=headers) + assert resp.status_code == 200, resp.text + assert resp.json()["answer"] == "answer from first" + + app.state.llm_client = _ProbeLLM("second") + resp = client.post("/runs", json={"goal": "hi"}, headers=headers) + assert resp.status_code == 200, resp.text + assert resp.json()["answer"] == "answer from second" + + await engine.dispose() + + +# --- tiered: trace forwarding ---------------------------------------------------------- + + +async def test_tiered_forwards_get_run_trace() -> None: + """TieredMemoryStore must expose the durable steps, not silently return empty.""" + rows = [ + StepTrace( + run_id="r1", + step_no=1, + tool_name="calculator", + args={"expression": "2+2"}, + result={"output": "4"}, + ) + ] + + class _RunStore: + def __init__(self) -> None: + self.calls: list[str] = [] + + async def get_run_trace(self, run_id: str) -> list[StepTrace]: + self.calls.append(run_id) + return rows + + run_store = _RunStore() + tiered = TieredMemoryStore(session_store=InMemoryMemoryStore(), run_store=run_store) + + traces = await tiered.get_run_trace("r1") + + assert traces == rows + assert run_store.calls == ["r1"] + + +async def test_tiered_forwards_list_runs() -> None: + class _RunStore: + def __init__(self) -> None: + self.calls: list[tuple[str | None, int]] = [] + + async def list_runs(self, session_id: str | None = None, limit: int = 50) -> list[Any]: + self.calls.append((session_id, limit)) + return [RunRecord(id="r1", goal="g")] + + run_store = _RunStore() + tiered = TieredMemoryStore(session_store=InMemoryMemoryStore(), run_store=run_store) + + runs = await tiered.list_runs(session_id="s1", limit=5) + + assert len(runs) == 1 + assert run_store.calls == [("s1", 5)] + + +async def test_tiered_trace_helpers_tolerate_bare_run_store() -> None: + """A run store without trace support yields empty lists rather than raising.""" + + class _Bare: + pass + + tiered = TieredMemoryStore(session_store=InMemoryMemoryStore(), run_store=_Bare()) + + assert await tiered.get_run_trace("r1") == [] + assert await tiered.list_runs() == [] + + +# --- retry: no __getattr__ recursion ---------------------------------------------------- + + +def test_retrying_client_getattr_on_uninitialized_instance() -> None: + """A partially constructed wrapper must raise AttributeError, not recurse.""" + wrapper = RetryingLLMClient.__new__(RetryingLLMClient) + + with pytest.raises(AttributeError): + _ = wrapper.anything + + +def test_retrying_client_delegates_attributes() -> None: + inner = AsyncMock(spec=LLMClient) + inner.some_marker = "visible" + wrapper = RetryingLLMClient(inner) + + assert wrapper.some_marker == "visible" + + +# --- health endpoint: no infrastructure detail leakage ---------------------------------- + + +def test_health_does_not_leak_connection_error_details() -> None: + """str(exc) from a driver can embed host, port and database name.""" + + class _Engine: + @asynccontextmanager + async def connect(self) -> AsyncIterator[Any]: + raise RuntimeError( + 'connection to server at "db.internal.corp:5432", database agentkit failed' + ) + yield # pragma: no cover - unreachable, keeps this an async context manager + + app = create_app(settings=Settings(API_KEYS="k")) + app.state.db_engine = _Engine() + + client = TestClient(app, raise_server_exceptions=False) + resp = client.get("/health") + + assert resp.status_code == 503 + body = resp.text + assert "db.internal.corp" not in body + assert "5432" not in body + assert "Database connectivity check failed." in body + + +# --- sql_readonly: literal-aware keyword scan -------------------------------------------- + + +def test_keyword_inside_string_literal_is_allowed() -> None: + query = "SELECT * FROM audit_events WHERE action = 'insert'" + assert validate_sql_query(query) == query + + +def test_comment_markers_inside_literals_are_data_not_comments() -> None: + """Masking applies to the comment and separator scans, not just the keyword scan.""" + assert ( + validate_sql_query("SELECT * FROM t WHERE note = 'x -- y'") + == "SELECT * FROM t WHERE note = 'x -- y'" + ) + assert validate_sql_query("SELECT * FROM t WHERE x = '#5'") == "SELECT * FROM t WHERE x = '#5'" + assert ( + validate_sql_query("SELECT 'it''s; still one' FROM t") == "SELECT 'it''s; still one' FROM t" + ) + + +def test_real_comment_after_literal_still_rejected() -> None: + """An empty literal must not let a following comment hide behind the mask.""" + with pytest.raises(ValueError, match="Comments are not permitted"): + validate_sql_query("SELECT * FROM t WHERE a = '' -- '") + with pytest.raises(ValueError, match="Comments are not permitted"): + validate_sql_query("SELECT 1 -- comment") + + +def test_keyword_outside_literal_is_still_blocked() -> None: + with pytest.raises(ValueError, match="prohibited SQL keyword"): + validate_sql_query("SELECT 'drop me' FROM logs UNION SELECT grant FROM users") + + +def test_keyword_inside_identifier_not_split_on_underscore() -> None: + """A column literally named delete_flag is data, not a write statement.""" + assert validate_sql_query("SELECT delete_flag FROM events") == "SELECT delete_flag FROM events" + + +def test_write_keyword_beside_literal_still_blocked() -> None: + """Masking literals must not let a real keyword in the surrounding SQL slip through.""" + with pytest.raises(ValueError, match="prohibited SQL keyword"): + validate_sql_query("SELECT 'safe text' FROM t WHERE grant = 1") + + +def test_quoted_identifier_containing_keyword_allowed() -> None: + query = 'SELECT "update_count" FROM metrics' + assert validate_sql_query(query) == query + + +def test_hash_json_operator_is_not_treated_as_comment() -> None: + """'#>' is a legitimate PostgreSQL JSON path operator, not a comment.""" + query = "SELECT payload #>> '{user,name}' FROM events" + assert validate_sql_query(query) == query + + +def test_bare_hash_comment_still_rejected() -> None: + with pytest.raises(ValueError, match="Comments are not permitted"): + validate_sql_query("SELECT 1 # sneaky comment") + + +def test_trailing_semicolon_stripped_once() -> None: + assert validate_sql_query("SELECT 1;") == "SELECT 1" + with pytest.raises(ValueError, match="Multiple SQL statements"): + validate_sql_query("SELECT 1;;") + + +# --- llm adapters: parsing failures must use the domain error taxonomy ------------------- + + +class _Msg: + def __init__(self, content: str | None = None, tool_calls: Any = None) -> None: + self.content = content + self.tool_calls = tool_calls + + +class _Choice: + def __init__(self, message: _Msg) -> None: + self.message = message + + +class _Usage: + prompt_tokens = 3 + completion_tokens = 4 + + +class _RawResponse: + def __init__(self, choices: Any, usage: Any = None) -> None: + self.choices = choices + self.usage = usage + + +def _openai_client(choices: Any) -> Any: + from agentkit.llm.openai import OpenAIClient + + sdk_client = AsyncMock() + sdk_client.chat.completions.create = AsyncMock(return_value=_RawResponse(choices, _Usage())) + return OpenAIClient(api_key="sk-test", client=sdk_client) + + +def _tool_call(name: str, arguments: str, call_id: str = "call_1") -> Any: + func = AsyncMock() + func.name = name + func.arguments = arguments + tc = AsyncMock() + tc.id = call_id + tc.function = func + return tc + + +async def test_openai_malformed_tool_arguments_raise_provider_error() -> None: + """Malformed JSON must name the real fault, not fail later as a missing field.""" + from agentkit.core.errors import LLMProviderError as _ProviderError + + client = _openai_client([_Choice(_Msg(tool_calls=[_tool_call("calc", "{not json")]))]) + + with pytest.raises(_ProviderError, match="malformed JSON arguments"): + await client.chat(messages=[]) + + +async def test_openai_non_object_tool_arguments_raise_provider_error() -> None: + from agentkit.core.errors import LLMProviderError as _ProviderError + + client = _openai_client([_Choice(_Msg(tool_calls=[_tool_call("calc", "[1, 2, 3]")]))]) + + with pytest.raises(_ProviderError, match="must decode to an object"): + await client.chat(messages=[]) + + +async def test_openai_empty_choices_raise_provider_error() -> None: + """An empty choices list must surface as LLMProviderError, not an IndexError.""" + from agentkit.core.errors import LLMProviderError as _ProviderError + + client = _openai_client([]) + + with pytest.raises(_ProviderError, match="no choices"): + await client.chat(messages=[]) + + +async def test_openai_tool_call_ids_are_full_uuids() -> None: + """A truncated 32-bit tool_call_id risks collisions within a long session.""" + client = _openai_client( + [_Choice(_Msg(tool_calls=[_tool_call("calc", "{}", call_id="")]))], + ) + + response = await client.chat(messages=[]) + + call_id = response.tool_calls[0].id + assert len(call_id) == 36, f"expected a full UUID, got {call_id!r}" + + +async def test_gemini_tool_call_ids_are_full_uuids() -> None: + from agentkit.llm.gemini import GeminiClient + + function_call = SimpleNamespace(name="calc", args={"expression": "2+2"}) + part = SimpleNamespace(text=None, function_call=function_call, thought_signature=None) + candidate = SimpleNamespace(content=SimpleNamespace(parts=[part])) + raw = SimpleNamespace(candidates=[candidate], usage_metadata=None, text="") + + sdk_client = AsyncMock() + sdk_client.aio.models.generate_content = AsyncMock(return_value=raw) + + client = GeminiClient(api_key="key", client=sdk_client) + response = await client.chat(messages=[]) + + call_id = response.tool_calls[0].id + assert len(call_id) == 36, f"expected a full UUID, got {call_id!r}" + + +async def test_gemini_malformed_response_raises_provider_error() -> None: + """Blocked prompts raise on attribute access; that must map to LLMProviderError.""" + from agentkit.core.errors import LLMProviderError as _ProviderError + from agentkit.llm.gemini import GeminiClient + + class _Exploding: + @property + def candidates(self) -> Any: + raise ValueError("prompt was blocked") + + raw = _Exploding() + sdk_client = AsyncMock() + sdk_client.aio.models.generate_content = AsyncMock(return_value=raw) + + client = GeminiClient(api_key="key", client=sdk_client) + + with pytest.raises(_ProviderError, match="Malformed Gemini response"): + await client.chat(messages=[]) + + +# --- sql_readonly: server-side row limit --------------------------------------------------- + + +async def test_sql_query_is_limited_in_sql_not_just_client_side() -> None: + """fetchmany alone still makes the database run the full scan.""" + from sqlalchemy import text + + from agentkit.tools.builtin.sql_readonly import execute_sql_query + + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as conn: + await conn.execute(text("CREATE TABLE t (id INTEGER PRIMARY KEY, v TEXT)")) + await conn.execute( + text("INSERT INTO t (id, v) VALUES (1, 'a'), (2, 'b'), (3, 'c')"), + ) + + output = await execute_sql_query("SELECT v FROM t", engine=engine, row_limit=2) + + lines = [line for line in output.strip().split("\n") if line.strip()] + assert lines == ["| v |", "| --- |", "| a |", "| b |"] + await engine.dispose() + + +def test_row_limit_rewrite_prefers_appending_over_wrapping() -> None: + """Appending LIMIT keeps duplicate-column JOINs working; wrapping is the fallback. + + Verified against PostgreSQL 16, which accepts both forms: the append form is preferred + because it cannot change the shape of an already-valid statement. + """ + from agentkit.tools.builtin.sql_readonly import _apply_row_limit + + assert _apply_row_limit("SELECT o.id, p.id FROM o JOIN p ON true", 50).endswith("LIMIT 50") + assert "agentkit_limited_query" not in _apply_row_limit("SELECT 1", 50) + assert _apply_row_limit("SELECT * FROM t LIMIT 5", 50) == "SELECT * FROM t LIMIT 5" + wrapped = _apply_row_limit("SELECT * FROM t FOR UPDATE", 50) + assert "agentkit_limited_query" in wrapped and wrapped.endswith("LIMIT 50") + + +# --- seed script: statement-aware splitting ------------------------------------------------- + + +def test_seed_splitter_respects_literals_and_dollar_quotes() -> None: + """A naive split(';') would tear these statements in half.""" + import sys + + sys.path.insert(0, ".") + from scripts.seed import split_sql_statements + + sql = ( + "INSERT INTO t (v) VALUES ('semi; colon');\n" + "-- a comment; with a semicolon\n" + "SELECT * FROM t WHERE v = 'it''s fine';\n" + "CREATE FUNCTION f() RETURNS int AS $$ BEGIN RETURN 1; END; $$ LANGUAGE plpgsql;\n" + ) + + statements = split_sql_statements(sql) + + assert len(statements) == 3, statements + assert statements[0] == "INSERT INTO t (v) VALUES ('semi; colon')" + assert statements[1] == "SELECT * FROM t WHERE v = 'it''s fine'" + assert "BEGIN RETURN 1; END;" in statements[2] + + +def test_seed_splitter_ignores_trailing_content() -> None: + import sys + + sys.path.insert(0, ".") + from scripts.seed import split_sql_statements + + assert split_sql_statements("SELECT 1;\n\n") == ["SELECT 1"] + assert split_sql_statements("") == [] + + +# --- calculator: bounded result size ------------------------------------------------------ + + +def test_exponentiation_result_size_is_bounded() -> None: + """The per-exponent guard cannot catch a nested form that balloons the result. + + ``(999**100)**5`` keeps every individual exponent inside MAX_EXPONENT and under the + existing large-base guard, but the inner power is already ~997 bits and raising it to the + fifth produces roughly 5000. The result size check is what rejects it. + """ + two_power = safe_calculate("2**1000") + assert isinstance(two_power, int) + assert two_power.bit_length() <= MAX_RESULT_BITS + + assert safe_calculate("10**1000") == 10**1000 + + with pytest.raises(ValueError, match="too large to compute safely"): + safe_calculate("(999**100)**5") + + +# --- ratelimit: the non-scripting fallback stays atomic ------------------------------------- + + +@pytest.mark.asyncio +async def test_unatomic_fallback_does_not_over_admit_under_concurrency() -> None: + """WATCH/MULTI must keep a parallel burst from exceeding the quota. + + The previous two-phase trim/count/insert let every concurrent request observe the same + pre-insert cardinality and be admitted together. + """ + client: Redis[Any] = fakeredis.aioredis.FakeRedis(decode_responses=True) + limiter = RateLimiter(requests_per_window=5, window_seconds=60) + + results = await asyncio.gather( + *( + limiter._check_window_unatomic(client, "ratelimit:k", 1000.0, 60.0, 5, f"m{i}", 65) + for i in range(20) + ), + return_exceptions=True, + ) + + admitted = [r for r in results if r == (True, 0)] + denied = [r for r in results if isinstance(r, tuple) and r[0] is False] + errors = [r for r in results if isinstance(r, Exception)] + + assert not errors, f"unexpected exceptions: {errors}" + assert len(admitted) == 5, f"expected exactly 5 admissions, got {len(admitted)}" + assert len(denied) == 15 diff --git a/tests/unit/test_llm_factory.py b/tests/unit/test_llm_factory.py index fe08a7e..c775f46 100644 --- a/tests/unit/test_llm_factory.py +++ b/tests/unit/test_llm_factory.py @@ -52,13 +52,17 @@ def test_create_from_settings_gemini() -> None: LLM_PROVIDER="gemini", LLM_MODEL="gemini-2.0-flash", GEMINI_API_KEY="test-gemini-key", - MAX_RETRIES=3, + LLM_MAX_RETRIES=3, + LLM_RETRY_BASE_DELAY_S=0.25, + LLM_RETRY_MAX_DELAY_S=2.5, ) client = create_llm_client_from_settings(settings) assert isinstance(client, RetryingLLMClient) assert isinstance(client._client, GeminiClient) assert client.model == "gemini-2.0-flash" assert client._max_retries == 3 + assert client._base_delay == 0.25 + assert client._max_delay == 2.5 def test_create_from_settings_openai() -> None: @@ -66,7 +70,7 @@ def test_create_from_settings_openai() -> None: LLM_PROVIDER="openai", LLM_MODEL="gpt-4o", OPENAI_API_KEY="test-openai-key", - MAX_RETRIES=2, + LLM_MAX_RETRIES=2, ) client = create_llm_client_from_settings(settings) assert isinstance(client, RetryingLLMClient) @@ -80,7 +84,7 @@ def test_create_from_settings_zero_retries() -> None: LLM_PROVIDER="gemini", LLM_MODEL="gemini-2.0-flash", GEMINI_API_KEY="test-gemini-key", - MAX_RETRIES=0, + LLM_MAX_RETRIES=0, ) client = create_llm_client_from_settings(settings) assert isinstance(client, GeminiClient) diff --git a/tests/unit/test_redis_store.py b/tests/unit/test_redis_store.py index 8a7fde2..7b7511a 100644 --- a/tests/unit/test_redis_store.py +++ b/tests/unit/test_redis_store.py @@ -107,13 +107,13 @@ async def test_sliding_window_truncation_with_system_prompt(fake_redis: Redis[An await store.save_messages(session_id, messages) loaded = await store.get_messages(session_id) - # Total messages should be: 1 system prompt + 3 most recent messages = 4 - assert len(loaded) == 4 + # max_messages is a total budget, and the system prompt spends one of those slots, + # so the result is the system prompt plus the 2 most recent conversational messages. + assert len(loaded) == 3 assert loaded[0].role == Role.SYSTEM assert loaded[0].content == "System prompt instructions" - assert loaded[1].content == "Message 2" - assert loaded[2].content == "Response 2" - assert loaded[3].content == "Message 3" + assert loaded[1].content == "Response 2" + assert loaded[2].content == "Message 3" @pytest.mark.asyncio diff --git a/tests/unit/test_tool_http_fetch.py b/tests/unit/test_tool_http_fetch.py index 27f4688..935c87c 100644 --- a/tests/unit/test_tool_http_fetch.py +++ b/tests/unit/test_tool_http_fetch.py @@ -1,14 +1,51 @@ """Unit and security tests for SSRF-protected HTTP fetch tool.""" -from unittest.mock import AsyncMock, patch +from __future__ import annotations + +from collections.abc import AsyncIterator +from contextlib import AbstractContextManager, asynccontextmanager +from typing import Any +from unittest.mock import patch import httpx import pytest -from agentkit.tools.builtin.http_fetch import http_fetch, validate_url_ssrf +from agentkit.tools.builtin.http_fetch import ( + http_fetch, + is_blocked_ip, + validate_url_ssrf, + verify_peer_address, +) from agentkit.tools.models import ToolCall from agentkit.tools.registry import ToolRegistry +PUBLIC_IP = "93.184.216.34" + + +class _FakeNetworkStream: + """Minimal stand-in for httpcore's AnyIOStream exposing ``server_addr``.""" + + def __init__(self, server_addr: object) -> None: + self._server_addr = server_addr + + def get_extra_info(self, key: str) -> object: + return self._server_addr if key == "server_addr" else None + + +def _patch_stream(response: httpx.Response) -> AbstractContextManager[Any]: + """Patch ``AsyncClient.stream`` to yield ``response`` from an async context manager. + + ``http_fetch`` now streams the body rather than buffering it, so the tests drive + ``AsyncClient.stream`` instead of ``AsyncClient.get``. A response built with ``text=`` or + ``content=`` already streams correctly, so no body stubbing is needed. + """ + + @asynccontextmanager + async def fake_stream(*_args: object, **_kwargs: object) -> AsyncIterator[httpx.Response]: + yield response + + return patch.object(httpx.AsyncClient, "stream", fake_stream) + @pytest.mark.parametrize( "invalid_scheme", @@ -71,16 +108,15 @@ def test_dns_rebinding_or_private_resolution_blocked() -> None: async def test_http_fetch_success() -> None: with patch( "agentkit.tools.builtin.http_fetch.resolve_host_ips", - return_value=["93.184.216.34"], + return_value=[PUBLIC_IP], ): mock_response = httpx.Response( status_code=200, text="Hello World", request=httpx.Request("GET", "https://example.com"), ) - with patch.object(httpx.AsyncClient, "get", new_callable=AsyncMock) as mock_get: - mock_get.return_value = mock_response + with _patch_stream(mock_response): content = await http_fetch("https://example.com") assert "Hello World" in content @@ -90,7 +126,7 @@ async def test_http_fetch_blocks_private_redirect() -> None: with patch( "agentkit.tools.builtin.http_fetch.resolve_host_ips", side_effect=[ - ["93.184.216.34"], # initial public url + [PUBLIC_IP], # initial public url ["169.254.169.254"], # redirect to metadata ], ): @@ -99,18 +135,18 @@ async def test_http_fetch_blocks_private_redirect() -> None: headers={"Location": "http://169.254.169.254/latest/meta-data/"}, request=httpx.Request("GET", "https://example.com"), ) - with patch.object(httpx.AsyncClient, "get", new_callable=AsyncMock) as mock_get: - mock_get.return_value = mock_redirect - - with pytest.raises(ValueError, match="SSRF protection blocked"): - await http_fetch("https://example.com") + with ( + _patch_stream(mock_redirect), + pytest.raises(ValueError, match="SSRF protection blocked"), + ): + await http_fetch("https://example.com") @pytest.mark.asyncio async def test_http_fetch_size_limit() -> None: with patch( "agentkit.tools.builtin.http_fetch.resolve_host_ips", - return_value=["93.184.216.34"], + return_value=[PUBLIC_IP], ): large_body = "x" * 2000 mock_response = httpx.Response( @@ -118,13 +154,110 @@ async def test_http_fetch_size_limit() -> None: text=large_body, request=httpx.Request("GET", "https://example.com"), ) - with patch.object(httpx.AsyncClient, "get", new_callable=AsyncMock) as mock_get: - mock_get.return_value = mock_response + with _patch_stream(mock_response): content = await http_fetch("https://example.com", max_chars=500) assert len(content) == 500 +@pytest.mark.asyncio +async def test_http_fetch_rejects_rebind_to_unvalidated_peer() -> None: + """A connection that lands off the validated addresses must be refused. + + Validation resolves the host to a public IP, but the client resolves again when it + connects. This asserts the post-connect peer check closes that gap. + """ + with patch( + "agentkit.tools.builtin.http_fetch.resolve_host_ips", + return_value=[PUBLIC_IP], + ): + mock_response = httpx.Response( + status_code=200, + text="internal secret", + request=httpx.Request("GET", "https://rebind.example.com"), + ) + mock_response.extensions["network_stream"] = _FakeNetworkStream(("127.0.0.1", 80)) + + with ( + _patch_stream(mock_response), + pytest.raises(ValueError, match="SSRF protection blocked"), + ): + await http_fetch("https://rebind.example.com") + + +@pytest.mark.asyncio +async def test_http_fetch_rejects_peer_outside_validated_set() -> None: + """A connection that lands on a public IP outside the validated set is refused. + + Distinct from the loopback case above: the peer is a routable address, just not one of + the addresses validation approved for this host. + """ + with patch( + "agentkit.tools.builtin.http_fetch.resolve_host_ips", + return_value=[PUBLIC_IP], + ): + mock_response = httpx.Response( + status_code=200, + text="internal secret", + request=httpx.Request("GET", "https://rebind.example.com"), + ) + # 1.1.1.1 is globally routable, so it passes the restricted-address screen and + # reaches the validated-set comparison, which must refuse it. + mock_response.extensions["network_stream"] = _FakeNetworkStream(("1.1.1.1", 443)) + + with ( + _patch_stream(mock_response), + pytest.raises(ValueError, match="was not among the validated addresses"), + ): + await http_fetch("https://rebind.example.com") + + +def test_verify_peer_address_allows_validated_peer() -> None: + response = httpx.Response(status_code=200, request=httpx.Request("GET", "https://example.com")) + response.extensions["network_stream"] = _FakeNetworkStream((PUBLIC_IP, 443)) + + verify_peer_address(response, [PUBLIC_IP]) + + +def test_verify_peer_address_no_stream_is_permissive() -> None: + """Without transport introspection the pre-flight validation remains the only guard.""" + response = httpx.Response(status_code=200, request=httpx.Request("GET", "https://example.com")) + + verify_peer_address(response, [PUBLIC_IP]) + + +@pytest.mark.parametrize("addr", [None, 12345, (), "not-a-tuple"]) +def test_verify_peer_address_malformed_server_addr_skips_cleanly(addr: object) -> None: + """A transport that exposes no usable peer address must not crash the check.""" + response = httpx.Response(status_code=200, request=httpx.Request("GET", "https://example.com")) + response.extensions["network_stream"] = _FakeNetworkStream(addr) + + verify_peer_address(response, [PUBLIC_IP]) + + +def test_verify_peer_address_empty_host_is_blocked() -> None: + """An empty peer host fails closed instead of escaping as a parsing error.""" + response = httpx.Response(status_code=200, request=httpx.Request("GET", "https://example.com")) + response.extensions["network_stream"] = _FakeNetworkStream(("", 80)) + + with pytest.raises(ValueError, match="SSRF protection blocked"): + verify_peer_address(response, [PUBLIC_IP]) + + +@pytest.mark.parametrize( + ("ip", "blocked"), + [ + ("127.0.0.1", True), + ("169.254.169.254", True), + ("10.1.2.3", True), + ("not-an-ip", True), + (PUBLIC_IP, False), + ], +) +def test_is_blocked_ip(ip: str, blocked: bool) -> None: + assert is_blocked_ip(ip) is blocked + + @pytest.mark.asyncio async def test_tool_registry_integration() -> None: registry = ToolRegistry() @@ -132,16 +265,15 @@ async def test_tool_registry_integration() -> None: with patch( "agentkit.tools.builtin.http_fetch.resolve_host_ips", - return_value=["93.184.216.34"], + return_value=[PUBLIC_IP], ): mock_response = httpx.Response( status_code=200, text="API Docs", request=httpx.Request("GET", "https://api.example.com"), ) - with patch.object(httpx.AsyncClient, "get", new_callable=AsyncMock) as mock_get: - mock_get.return_value = mock_response + with _patch_stream(mock_response): call = ToolCall( id="c1", name="http_fetch", arguments={"url": "https://api.example.com"} )