From 1fc32ff0351271ab0275a6d52f331583198d3416 Mon Sep 17 00:00:00 2001 From: Hamza Imran Date: Wed, 21 Oct 2026 02:56:18 +0500 Subject: [PATCH] fix(tools): preserve in generated tool schemas (fixes #70) --- agentkit/tools/schema.py | 4 ++ docs/PROGRESS.md | 8 +++ tests/unit/test_tool_schema.py | 107 +++++++++++++++++++++++++++++++++ 3 files changed, 119 insertions(+) diff --git a/agentkit/tools/schema.py b/agentkit/tools/schema.py index 33f9a31..70fc101 100644 --- a/agentkit/tools/schema.py +++ b/agentkit/tools/schema.py @@ -70,6 +70,10 @@ def create_tool_schema( } if required: parameters["required"] = required + if "$defs" in raw_json_schema: + parameters["$defs"] = raw_json_schema["$defs"] + if "additionalProperties" in raw_json_schema: + parameters["additionalProperties"] = raw_json_schema["additionalProperties"] schema = ToolSchema( name=tool_name, diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 2c7b115..8712803 100644 --- a/docs/PROGRESS.md +++ b/docs/PROGRESS.md @@ -61,6 +61,7 @@ This document tracks execution progress across all backlog tasks, including pull | 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 | +| AUDIT-13 | AUD | Preserve $defs and definitions in Generated Tool Schemas | `fix/tool-schema-defs` | #81 | 1 | Done | Preserve $defs and additionalProperties in tool parameters; eliminate unresolvable $ref pointers | --- @@ -407,3 +408,10 @@ This document tracks execution progress across all backlog tasks, including pull - **PR**: #47 - **Commits**: 3 - **Notes**: Remediated 17 findings from a full-codebase audit, rebased onto the AUD-01/AUD-04/AUD-05 follow-ups. Critical: removed the `ak_test_key_12345` default from `Settings.API_KEYS` and made the field required; replaced the `FakeLLMClient` fallback in `get_llm_client` with a `ServiceUnavailableError` (503) boundary around provider construction; kept the AUD-04 factory-level `RetryingLLMClient` wrapping and made it configurable through new `LLM_MAX_RETRIES`/`LLM_RETRY_BASE_DELAY_S`/`LLM_RETRY_MAX_DELAY_S` settings; made the per-request `Agent` assembly's override explicit as `app.state.agent_override`; forwarded `get_run_trace`/`list_runs` through `TieredMemoryStore`; moved `http_fetch` DNS off the event loop with `asyncio.to_thread`, added a post-connect `server_addr` check against the validated addresses to close the DNS-rebinding TOCTOU window, and replaced full-body buffering with a byte-budgeted stream; replaced the rate limiter's non-scripting fallback with `WATCH`/`MULTI`/`EXEC`, leaving the AUD-10 in-memory fallback for Redis outages intact. Moderate: corrected the `RedisMemoryStore.truncate_messages` off-by-one that returned `max_messages + 1` entries when a system prompt was present; switched generated `tool_call_id`s to full UUIDs; wrapped response parsing in both LLM adapters inside the error boundary so malformed responses and malformed tool-argument JSON raise `LLMProviderError`; guarded `RetryingLLMClient.__getattr__` against recursion on a partially constructed instance; sanitized `/health` details so the unauthenticated endpoint no longer echoes driver error strings; made `sql_readonly` mask string literals and quoted identifiers before its keyword scan, allow `#>`/`#>>`/`#-`, apply the row limit in SQL and push `statement_timeout` to PostgreSQL. Minor: stopped logging raw run goals at INFO (length only, content at DEBUG); replaced `split(";")` in `scripts/seed.py` with a statement-aware splitter handling literals, comments and dollar-quotes; bounded calculator exponentiation by result bit length. Added `tests/unit/test_audit_regressions.py` and a real-Redis concurrency regression test pinning the fallback fix. Recorded in ADR-009, ADR-010, ADR-011. Suite: 298 passed, 25 skipped; coverage 92.86% on core/tools/llm; `ruff check`, `ruff format --check` and `mypy agentkit tests` all clean. + + +### AUDIT-13: Preserve $defs in Generated Tool Schemas +- **Status**: Done +- **PR**: #81 +- **Commits**: 1 +- **Notes**: In `agentkit/tools/schema.py`, updated `create_tool_schema` to retain `$defs` and `additionalProperties` emitted by Pydantic's `model_json_schema()` in the final `parameters` dictionary. Tools with `Enum` or nested `BaseModel` parameters now emit valid JSON Schema where every `$ref` successfully resolves within the document root, while existing simple built-in tools remain byte-identical without superfluous keys. Added unit and pointer-resolution tests in `tests/unit/test_tool_schema.py`. diff --git a/tests/unit/test_tool_schema.py b/tests/unit/test_tool_schema.py index 00a24aa..04d8878 100644 --- a/tests/unit/test_tool_schema.py +++ b/tests/unit/test_tool_schema.py @@ -1,5 +1,16 @@ """Unit tests for tool schema generation from callable inspection.""" +from collections.abc import Callable +from enum import StrEnum +from typing import Any + +from pydantic import BaseModel, Field + +from agentkit.tools.builtin.calculator import calculator +from agentkit.tools.builtin.http_fetch import http_fetch +from agentkit.tools.builtin.read_file import read_file +from agentkit.tools.builtin.sql_readonly import sql_readonly +from agentkit.tools.builtin.web_search import web_search from agentkit.tools.schema import create_tool_schema @@ -67,3 +78,99 @@ def test_create_tool_schema_custom_name_and_description() -> None: ) assert schema.name == "custom_calc" assert schema.description == "Custom description" + + +class QueryMode(StrEnum): + QUICK = "quick" + EXHAUSTIVE = "exhaustive" + + +class FilterConfig(BaseModel): + max_results: int = Field(default=10, description="Maximum number of results to fetch.") + include_raw: bool = False + + +def tool_with_enum_and_nested_model( + query: str, + mode: QueryMode = QueryMode.QUICK, + filter_config: FilterConfig | None = None, +) -> str: + """Tool that utilizes Enum and nested BaseModel schemas.""" + return f"{query}:{mode}:{filter_config}" + + +def _assert_all_refs_resolve(schema: dict[str, Any]) -> None: + """Validate that every JSON Pointer in a $ref successfully resolves inside the schema.""" + refs: list[str] = [] + + def _collect_refs(node: Any) -> None: + if isinstance(node, dict): + for k, v in node.items(): + if k == "$ref" and isinstance(v, str): + refs.append(v) + else: + _collect_refs(v) + elif isinstance(node, list): + for item in node: + _collect_refs(item) + + _collect_refs(schema) + assert refs, "Expected schema to contain $ref references to test" + + for ref in refs: + assert ref.startswith("#/"), f"Expected URI fragment reference, got {ref}" + parts = ref[2:].split("/") + target: Any = schema + for part in parts: + assert isinstance(target, dict), f"Cannot navigate into non-dict at {part}" + assert part in target, ( + f"Dangling reference: key '{part}' not found in {list(target.keys())}" + ) + target = target[part] + + +def test_create_tool_schema_preserves_defs_for_enums_and_models() -> None: + """AUDIT-13: Preserves $defs and ensures all $ref pointers resolve.""" + schema, args_model = create_tool_schema(tool_with_enum_and_nested_model) + params = schema.parameters + + assert "$defs" in params + assert "QueryMode" in params["$defs"] + assert "FilterConfig" in params["$defs"] + + # Verify Enum definition + mode_def = params["$defs"]["QueryMode"] + assert mode_def["enum"] == ["quick", "exhaustive"] + assert mode_def["type"] == "string" + + # Verify nested BaseModel definition + config_def = params["$defs"]["FilterConfig"] + assert config_def["type"] == "object" + assert "max_results" in config_def["properties"] + + # Verify all references resolve + _assert_all_refs_resolve(params) + + # Verify model validation still functions + val = args_model(query="search query", mode=QueryMode.EXHAUSTIVE) + assert val.model_dump() == { + "query": "search query", + "mode": QueryMode.EXHAUSTIVE, + "filter_config": None, + } + + +def test_builtin_tools_schemas_unaffected() -> None: + """AUDIT-13: Simple builtin tools must not emit unnecessary $defs or extra keys.""" + builtin_tools: list[Callable[..., Any]] = [ + calculator, + read_file, + http_fetch, + sql_readonly, + web_search, + ] + for fn in builtin_tools: + schema, _ = create_tool_schema(fn) + assert "$defs" not in schema.parameters + assert "additionalProperties" not in schema.parameters + assert set(schema.parameters.keys()) <= {"type", "properties", "required"}