Normalize extract schema mappings safely

Co-authored-by: Shri Sukhani <shrisukhani@users.noreply.github.com>

Author: cursoragent

hyperbrowser/tools/__init__.py

@@ -52,6 +52,33 @@ def _format_tool_param_key_for_error(key: str) -> str:
     return f"{normalized_key[:available_length]}{_TRUNCATED_KEY_DISPLAY_SUFFIX}"
 
 
+def _normalize_extract_schema_mapping(schema_value: MappingABC[object, Any]) -> Dict[str, Any]:
+    try:
+        schema_keys = list(schema_value.keys())
+    except HyperbrowserError:
+        raise
+    except Exception as exc:
+        raise HyperbrowserError(
+            "Failed to read extract tool `schema` object keys",
+            original_error=exc,
+        ) from exc
+    normalized_schema: Dict[str, Any] = {}
+    for key in schema_keys:
+        if not isinstance(key, str):
+            raise HyperbrowserError("Extract tool `schema` object keys must be strings")
+        try:
+            normalized_schema[key] = schema_value[key]
+        except HyperbrowserError:
+            raise
+        except Exception as exc:
+            key_display = _format_tool_param_key_for_error(key)
+            raise HyperbrowserError(
+                f"Failed to read extract tool `schema` value for key '{key_display}'",
+                original_error=exc,
+            ) from exc
+    return normalized_schema
+
+
 def _prepare_extract_tool_params(params: Mapping[str, Any]) -> Dict[str, Any]:
     normalized_params = _to_param_dict(params)
     schema_value = normalized_params.get("schema")
@@ -69,11 +96,13 @@ def _prepare_extract_tool_params(params: Mapping[str, Any]) -> Dict[str, Any]:
                 "Invalid JSON string provided for `schema` in extract tool params",
                 original_error=exc,
             ) from exc
-        if parsed_schema is not None and not isinstance(parsed_schema, MappingABC):
+        if not isinstance(parsed_schema, MappingABC):
             raise HyperbrowserError(
                 "Extract tool `schema` must decode to a JSON object"
             )
-        normalized_params["schema"] = parsed_schema
+        normalized_params["schema"] = _normalize_extract_schema_mapping(parsed_schema)
+    elif isinstance(schema_value, MappingABC):
+        normalized_params["schema"] = _normalize_extract_schema_mapping(schema_value)
     return normalized_params
 
 

tests/test_tools_extract.py

@@ -1,4 +1,6 @@
 import asyncio
+from collections.abc import Mapping
+from types import MappingProxyType
 
 import pytest
 
@@ -171,6 +173,135 @@ async def run():
         asyncio.run(run())
 
 
+def test_extract_tool_runnable_rejects_null_schema_json():
+    client = _SyncClient()
+    params = {
+        "urls": ["https://example.com"],
+        "schema": "null",
+    }
+
+    with pytest.raises(
+        HyperbrowserError, match="Extract tool `schema` must decode to a JSON object"
+    ):
+        WebsiteExtractTool.runnable(client, params)
+
+
+def test_extract_tool_runnable_normalizes_mapping_schema_values():
+    client = _SyncClient()
+    schema_mapping = MappingProxyType({"type": "object", "properties": {}})
+
+    WebsiteExtractTool.runnable(
+        client,
+        {
+            "urls": ["https://example.com"],
+            "schema": schema_mapping,
+        },
+    )
+
+    assert isinstance(client.extract.last_params, StartExtractJobParams)
+    assert isinstance(client.extract.last_params.schema_, dict)
+    assert client.extract.last_params.schema_ == {"type": "object", "properties": {}}
+
+
+def test_extract_tool_runnable_rejects_non_string_schema_keys():
+    client = _SyncClient()
+
+    with pytest.raises(
+        HyperbrowserError, match="Extract tool `schema` object keys must be strings"
+    ):
+        WebsiteExtractTool.runnable(
+            client,
+            {
+                "urls": ["https://example.com"],
+                "schema": {1: "invalid-key"},  # type: ignore[dict-item]
+            },
+        )
+
+
+def test_extract_tool_runnable_wraps_schema_key_read_failures():
+    class _BrokenSchemaMapping(Mapping[object, object]):
+        def __iter__(self):
+            raise RuntimeError("cannot iterate schema keys")
+
+        def __len__(self) -> int:
+            return 1
+
+        def __getitem__(self, key: object) -> object:
+            return key
+
+    client = _SyncClient()
+
+    with pytest.raises(
+        HyperbrowserError, match="Failed to read extract tool `schema` object keys"
+    ) as exc_info:
+        WebsiteExtractTool.runnable(
+            client,
+            {
+                "urls": ["https://example.com"],
+                "schema": _BrokenSchemaMapping(),
+            },
+        )
+
+    assert exc_info.value.original_error is not None
+
+
+def test_extract_tool_runnable_wraps_schema_value_read_failures():
+    class _BrokenSchemaMapping(Mapping[str, object]):
+        def __iter__(self):
+            yield "type"
+
+        def __len__(self) -> int:
+            return 1
+
+        def __getitem__(self, key: str) -> object:
+            _ = key
+            raise RuntimeError("cannot read schema value")
+
+    client = _SyncClient()
+
+    with pytest.raises(
+        HyperbrowserError,
+        match="Failed to read extract tool `schema` value for key 'type'",
+    ) as exc_info:
+        WebsiteExtractTool.runnable(
+            client,
+            {
+                "urls": ["https://example.com"],
+                "schema": _BrokenSchemaMapping(),
+            },
+        )
+
+    assert exc_info.value.original_error is not None
+
+
+def test_extract_tool_runnable_preserves_hyperbrowser_schema_value_read_failures():
+    class _BrokenSchemaMapping(Mapping[str, object]):
+        def __iter__(self):
+            yield "type"
+
+        def __len__(self) -> int:
+            return 1
+
+        def __getitem__(self, key: str) -> object:
+            _ = key
+            raise HyperbrowserError("custom schema value failure")
+
+    client = _SyncClient()
+
+    with pytest.raises(
+        HyperbrowserError, match="custom schema value failure"
+    ) as exc_info:
+        WebsiteExtractTool.runnable(
+            client,
+            {
+                "urls": ["https://example.com"],
+                "schema": _BrokenSchemaMapping(),
+            },
+        )
+
+    assert exc_info.value.original_error is None
+
+
 def test_extract_tool_runnable_serializes_empty_object_data():
     client = _SyncClient(response_data={})