From 74584625d26060fc764095dc3533637d4cf5e9c9 Mon Sep 17 00:00:00 2001 From: Amar K P Date: Thu, 1 Oct 2026 12:47:01 +0530 Subject: [PATCH 1/3] fix(audit): remediate whole-project audit findings Addresses 17 findings from a full-codebase audit. Several paths converted a misconfiguration or a transient fault into a working-looking success. Fail-closed configuration and dependency resolution: - Settings.API_KEYS no longer defaults to a key published in the repository; validation rejects an empty or whitespace-only value. - get_llm_client no longer falls back to FakeLLMClient. Provider construction failure raises ServiceUnavailableError (503). LLM_PROVIDER=fake remains an explicit opt-in that short-circuits before the factory. - get_agent builds the Agent per request instead of memoizing the first assembly on app.state, which pinned every later request to the first request's LLM client, registry, memory store and trace sink. - Wired RetryingLLMClient into the request path (previously dead code) with new LLM_MAX_RETRIES / LLM_RETRY_BASE_DELAY_S / LLM_RETRY_MAX_DELAY_S. - TieredMemoryStore forwards get_run_trace and list_runs to the run store, so the trace endpoint no longer falls through to an empty trace sink. - examples/demo.py requires --offline for scripted execution rather than silently downgrading on a missing key. Post-connect SSRF verification and bounded resource use: - http_fetch resolves DNS via asyncio.to_thread and reuses those addresses. - Verifies the connected socket's server_addr against the validated addresses, closing the DNS-rebinding window between validation and connection. - Streams the body under an explicit byte budget instead of buffering then slicing. - sql_readonly applies the row limit in SQL, pushes statement_timeout to PostgreSQL, masks string literals and quoted identifiers before its keyword scan, and allows the #> / #>> / #- JSON path operators. Atomic fallbacks and domain-error boundaries: - The rate limiter's non-scripting fallback uses WATCH/MULTI/EXEC optimistic locking instead of a trim/count-then-insert split. Against live Redis the old version admitted 39 of 50 simultaneous requests against a limit of 5; the new one admits exactly 5. - Both LLM adapters wrap response parsing in the error boundary, so blocked prompts, empty choices and malformed tool-argument JSON raise LLMProviderError rather than escaping as raw SDK exceptions. - Generated tool_call_ids use full UUIDs instead of 32-bit truncations. - RetryingLLMClient.__getattr__ no longer recurses on a partial instance. - /health returns fixed client-facing strings and logs the driver error, so the unauthenticated endpoint stops disclosing hostnames, ports and database names. Smaller fixes: fixed the RedisMemoryStore.truncate_messages off-by-one that returned max_messages + 1 entries with a system prompt; stopped logging raw run goals at INFO; replaced split(";") in scripts/seed.py with a statement-aware splitter handling literals, comments and dollar-quotes; bounded calculator exponentiation by result bit length. Suite: 277 hermetic passed, 25 real-infra passed against live PostgreSQL and Redis. Coverage 92% overall. ruff check, ruff format --check and mypy agentkit tests are clean. --- .env.example | 15 +- README.md | 16 +- agentkit/api/deps.py | 39 +- agentkit/api/ratelimit.py | 63 ++- agentkit/api/routes.py | 22 +- agentkit/config.py | 36 +- agentkit/core/agent.py | 5 +- agentkit/llm/factory.py | 6 +- agentkit/llm/gemini.py | 95 +++-- agentkit/llm/openai.py | 108 +++-- agentkit/llm/retry.py | 20 +- agentkit/memory/redis_store.py | 7 +- agentkit/memory/tiered.py | 24 ++ agentkit/tools/builtin/calculator.py | 19 + agentkit/tools/builtin/http_fetch.py | 163 ++++++-- agentkit/tools/builtin/sql_readonly.py | 139 ++++++- docker-compose.yml | 4 + docs/DECISIONS.md | 55 ++- docs/DEPLOY.md | 2 +- docs/PROGRESS.md | 8 +- examples/demo.py | 16 +- scripts/seed.py | 93 ++++- tests/conftest.py | 7 + tests/real/test_redis_real.py | 40 +- tests/unit/test_audit_regressions.py | 553 +++++++++++++++++++++++++ tests/unit/test_llm_factory.py | 10 +- tests/unit/test_redis_store.py | 10 +- tests/unit/test_tool_http_fetch.py | 121 +++++- 28 files changed, 1483 insertions(+), 213 deletions(-) create mode 100644 tests/unit/test_audit_regressions.py 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..1b4b60b 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,64 @@ 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 closes that gap. + + 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 server_addr: + return + + peer_ip = server_addr[0] + 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 +168,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..e00203c 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,34 @@ 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 LIMIT or OFFSET clause at the end of the top-level statement. Statements with their own +#: bound are left alone so wrapping never widens a caller-supplied limit. +_TOP_LEVEL_LIMIT_REGEX = re.compile(r"\blimit\s+\d+\s*(?:;)?\s*$", 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,14 +93,22 @@ 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. A bare '#' is allowed when it + # begins a PostgreSQL JSON path operator (#>, #>>, #-); those are legitimate in a + # read-only projection, whereas a '#' followed by whitespace opens a MySQL-style comment. + if ( + "--" in stripped + or "/*" in stripped + or "*/" in stripped + or _BARE_HASH_REGEX.search(stripped) + ): raise ValueError( "Comments are not permitted in SQL queries to prevent security circumvention." ) - # Strip single trailing semicolon if present - cleaned = stripped.rstrip(";").strip() + # Strip a single trailing semicolon if present. rstrip(";") would also swallow the + # repeated separators that signal a stacked-statement payload. + cleaned = stripped[:-1].strip() if stripped.endswith(";") else stripped if ";" in cleaned: raise ValueError("Multiple SQL statements separated by semicolons are strictly prohibited.") @@ -90,8 +117,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 +130,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 +198,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 +217,49 @@ 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: + await conn.execute( + text(f"SET STATEMENT max_execution_time={int(timeout_s * 1000)} FOR SELECT 1") + ) + 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 _apply_row_limit(sql: str, row_limit: int) -> str: + """Wrap the statement in a limited subquery when it carries no top-level LIMIT. + + An existing LIMIT or OFFSET at the end of the statement is left untouched so a caller + supplied bound is never widened by the wrapper. + """ + bounded = f"SELECT * FROM ({sql}) AS agentkit_limited_query LIMIT {int(row_limit)}" + if _TOP_LEVEL_LIMIT_REGEX.search(sql): + return sql + return bounded + + @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..cba14d8 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: 330 passed, 24 skipped; coverage 92% 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/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..e37cfe4 --- /dev/null +++ b/tests/unit/test_audit_regressions.py @@ -0,0 +1,553 @@ +"""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 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 + 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_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() + + +# --- 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..d059e4d 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: tuple[str, int]) -> None: + self._server_addr = server_addr + + def get_extra_info(self, key: str) -> tuple[str, int] | None: + 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,65 @@ 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="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( + ("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 +220,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"} ) From 4cafbd8722d20ae0c9bc7669b9b5ff741a572fa2 Mon Sep 17 00:00:00 2001 From: Amar K P Date: Thu, 1 Oct 2026 13:08:12 +0530 Subject: [PATCH 2/3] fix(audit): harden peer check, masked scans, row-limit rewrite, mysql timeout - verify_peer_address tolerates malformed server_addr without crashing and fails closed on restricted or unparseable peer hosts. - Run the comment and multi-statement scans on the literal-masked query so markers inside string data no longer cause false rejections. - Prefer appending LIMIT over subquery wrapping; wrap only for row-locking tail clauses. Add FETCH FIRST to the existing-bound detection. - MySQL timeout uses SET SESSION MAX_EXECUTION_TIME, which actually covers subsequent statements, instead of a no-op SET STATEMENT ... FOR SELECT 1. --- agentkit/tools/builtin/http_fetch.py | 15 ++++- agentkit/tools/builtin/sql_readonly.py | 87 +++++++++++++++++++------- tests/unit/test_audit_regressions.py | 35 +++++++++++ tests/unit/test_tool_http_fetch.py | 49 ++++++++++++++- 4 files changed, 160 insertions(+), 26 deletions(-) diff --git a/agentkit/tools/builtin/http_fetch.py b/agentkit/tools/builtin/http_fetch.py index 1b4b60b..b4c5bfa 100644 --- a/agentkit/tools/builtin/http_fetch.py +++ b/agentkit/tools/builtin/http_fetch.py @@ -132,8 +132,10 @@ def verify_peer_address(response: httpx.Response, allowed_ips: list[str]) -> Non 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 closes that gap. + 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``. @@ -146,10 +148,17 @@ def verify_peer_address(response: httpx.Response, allowed_ips: list[str]) -> Non return server_addr = stream.get_extra_info("server_addr") - if not 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: diff --git a/agentkit/tools/builtin/sql_readonly.py b/agentkit/tools/builtin/sql_readonly.py index e00203c..fa3e905 100644 --- a/agentkit/tools/builtin/sql_readonly.py +++ b/agentkit/tools/builtin/sql_readonly.py @@ -51,9 +51,19 @@ #: A '#' that is not the start of a JSON path operator (#>, #>>, #-), i.e. a comment marker. _BARE_HASH_REGEX = re.compile(r"#(?![>~-])") -#: A LIMIT or OFFSET clause at the end of the top-level statement. Statements with their own -#: bound are left alone so wrapping never widens a caller-supplied limit. -_TOP_LEVEL_LIMIT_REGEX = re.compile(r"\blimit\s+\d+\s*(?:;)?\s*$", re.IGNORECASE) +#: 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 @@ -93,23 +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. A bare '#' is allowed when it - # begins a PostgreSQL JSON path operator (#>, #>>, #-); those are legitimate in a - # read-only projection, whereas a '#' followed by whitespace opens a MySQL-style comment. + # 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 stripped - or "/*" in stripped - or "*/" in stripped - or _BARE_HASH_REGEX.search(stripped) + "--" 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 a single trailing semicolon if present. rstrip(";") would also swallow the - # repeated separators that signal a stacked-statement payload. + # 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 cleaned: + 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 @@ -239,25 +256,53 @@ async def _apply_statement_timeout(conn: AsyncConnection, timeout_s: float) -> N if dialect == "postgresql": await conn.execute(text(f"SET LOCAL statement_timeout = {int(timeout_s * 1000)}")) else: - await conn.execute( - text(f"SET STATEMENT max_execution_time={int(timeout_s * 1000)} FOR SELECT 1") - ) + # 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: - """Wrap the statement in a limited subquery when it carries no top-level LIMIT. + """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: - An existing LIMIT or OFFSET at the end of the statement is left untouched so a caller - supplied bound is never widened by the wrapper. + - 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. """ - bounded = f"SELECT * FROM ({sql}) AS agentkit_limited_query LIMIT {int(row_limit)}" - if _TOP_LEVEL_LIMIT_REGEX.search(sql): + tail = _top_level_tail(_mask_literals(sql)) + if _EXISTING_BOUND_REGEX.search(tail): return sql - return bounded + 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 diff --git a/tests/unit/test_audit_regressions.py b/tests/unit/test_audit_regressions.py index e37cfe4..d79b256 100644 --- a/tests/unit/test_audit_regressions.py +++ b/tests/unit/test_audit_regressions.py @@ -285,6 +285,26 @@ def test_keyword_inside_string_literal_is_allowed() -> None: 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") @@ -468,6 +488,21 @@ async def test_sql_query_is_limited_in_sql_not_just_client_side() -> None: 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 ------------------------------------------------- diff --git a/tests/unit/test_tool_http_fetch.py b/tests/unit/test_tool_http_fetch.py index d059e4d..935c87c 100644 --- a/tests/unit/test_tool_http_fetch.py +++ b/tests/unit/test_tool_http_fetch.py @@ -25,10 +25,10 @@ class _FakeNetworkStream: """Minimal stand-in for httpcore's AnyIOStream exposing ``server_addr``.""" - def __init__(self, server_addr: tuple[str, int]) -> None: + def __init__(self, server_addr: object) -> None: self._server_addr = server_addr - def get_extra_info(self, key: str) -> tuple[str, int] | None: + def get_extra_info(self, key: str) -> object: return self._server_addr if key == "server_addr" else None @@ -178,6 +178,33 @@ async def test_http_fetch_rejects_rebind_to_unvalidated_peer() -> None: ) 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"), @@ -199,6 +226,24 @@ def test_verify_peer_address_no_stream_is_permissive() -> None: 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"), [ From e6ead1fb7ebabe5e01c95ea497106112134a1784 Mon Sep 17 00:00:00 2001 From: Amar K P Date: Sun, 18 Oct 2026 22:37:16 +0530 Subject: [PATCH 3/3] test(audit): reconcile regression tests with strict dependency resolution Adapt the rebased audit tests to the AUD-01/AUD-04 behaviour now on main: deps delegate to the factory instead of wrapping retry itself, build_llm exits via SystemExit on a missing key, and the per-request probe test uses in-memory memory/trace overrides so it stays hermetic without Redis. --- docs/PROGRESS.md | 2 +- tests/integration/test_api_deps_strict.py | 9 +++++---- tests/integration/test_demo.py | 6 +++--- tests/unit/test_audit_regressions.py | 6 +++++- 4 files changed, 14 insertions(+), 9 deletions(-) diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index cba14d8..2c7b115 100644 --- a/docs/PROGRESS.md +++ b/docs/PROGRESS.md @@ -406,4 +406,4 @@ This document tracks execution progress across all backlog tasks, including pull - **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: 330 passed, 24 skipped; coverage 92% on core/tools/llm; `ruff check`, `ruff format --check` and `mypy agentkit tests` all clean. +- **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/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/unit/test_audit_regressions.py b/tests/unit/test_audit_regressions.py index d79b256..e7c0d9d 100644 --- a/tests/unit/test_audit_regressions.py +++ b/tests/unit/test_audit_regressions.py @@ -25,7 +25,7 @@ from agentkit.api.ratelimit import RateLimiter from agentkit.config import Settings from agentkit.core.errors import AuthenticationError, ServiceUnavailableError -from agentkit.core.trace import StepTrace +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 @@ -152,6 +152,10 @@ async def test_agent_reflects_dependency_changes_between_requests() -> None: 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)