From 19b1a13629ec221e787db6225feb25881b7e0f2c Mon Sep 17 00:00:00 2001
From: Kayvan Zahiri <123409429+Kayvan-Zahiri@users.noreply.github.com>
Date: Fri, 16 Oct 2026 11:56:54 -0700
Subject: [PATCH] fix: accept recursive object schemas for structured output

---
 src/droid_sdk/_high_level/output.py | 28 ++++++++++++++---
 tests/test_recursive_output.py      | 49 +++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_recursive_output.py

diff --git a/src/droid_sdk/_high_level/output.py b/src/droid_sdk/_high_level/output.py
index abb3c84..a57af8e 100644
--- a/src/droid_sdk/_high_level/output.py
+++ b/src/droid_sdk/_high_level/output.py
@@ -49,9 +49,7 @@ def adapt(self, raw: object | None) -> OutputAdaptation[T_co]:
         if self._model is not None:
             structured_output = _raw_object(raw)
             validation_input = (
-                thaw_json(structured_output)
-                if structured_output is not None
-                else raw
+                thaw_json(structured_output) if structured_output is not None else raw
             )
             try:
                 value = self._model.model_validate(validation_input)
@@ -102,8 +100,30 @@ def prepare_output_adapter(
     raise TypeError("output must be None, JsonSchema, or a Pydantic BaseModel class")
 
 
+def _schema_member(value: object, key: str) -> object:
+    if not isinstance(value, Mapping):
+        return None
+    return cast("Mapping[str, object]", value).get(key)
+
+
 def _ensure_object_schema(schema: Mapping[str, object]) -> None:
-    if schema.get("type") != "object":
+    current: object = schema
+    visited: set[str] = set()
+    while _schema_member(current, "type") is None:
+        reference = _schema_member(current, "$ref")
+        if (
+            not isinstance(reference, str)
+            or not reference.startswith("#/")
+            or reference in visited
+        ):
+            break
+        visited.add(reference)
+        current = schema
+        for part in reference[2:].split("/"):
+            current = _schema_member(
+                current, part.replace("~1", "/").replace("~0", "~")
+            )
+    if _schema_member(current, "type") != "object":
         raise TypeError("structured output schema must describe a top-level object")
 
 
diff --git a/tests/test_recursive_output.py b/tests/test_recursive_output.py
new file mode 100644
index 0000000..6019a62
--- /dev/null
+++ b/tests/test_recursive_output.py
@@ -0,0 +1,49 @@
+import pytest
+from pydantic import BaseModel, Field
+
+from droid_sdk import JsonSchema
+from droid_sdk._high_level.output import prepare_output_adapter
+
+
+class Node(BaseModel):
+    name: str
+    children: list["Node"] = Field(default_factory=list)
+
+
+def test_recursive_pydantic_object_output() -> None:
+    adapter = prepare_output_adapter(Node)
+    assert adapter.output_format is not None
+    assert adapter.output_format.schema_ == Node.model_json_schema(mode="validation")
+    result = adapter.adapt({"name": "root", "children": [{"name": "leaf"}]})
+    assert result.output == Node(name="root", children=[Node(name="leaf")])
+    assert result.validation_error is None
+
+
+def test_referenced_object_json_schema_preserves_wire_schema() -> None:
+    schema: dict[str, object] = {
+        "$defs": {
+            "a/b~c": {"type": "object", "properties": {"name": {"type": "string"}}}
+        },
+        "$ref": "#/$defs/a~1b~0c",
+    }
+    adapter = prepare_output_adapter(JsonSchema(schema))
+    assert adapter.output_format is not None
+    assert adapter.output_format.schema_ == schema
+    assert adapter.adapt({"name": "root"}).output == {"name": "root"}
+
+
+@pytest.mark.parametrize(
+    "schema",
+    [
+        {"$defs": {"Node": {"type": "array"}}, "$ref": "#/$defs/Node"},
+        {"$ref": "#/$defs/missing"},
+        {"$defs": {"Node": {"$ref": "#/$defs/Node"}}, "$ref": "#/$defs/Node"},
+        {"$ref": "https://example.invalid/schema"},
+        {"$defs": {"Node": 1}, "$ref": "#/$defs/Node"},
+    ],
+)
+def test_invalid_or_non_object_references_are_rejected(
+    schema: dict[str, object],
+) -> None:
+    with pytest.raises(TypeError, match="top-level object"):
+        prepare_output_adapter(JsonSchema(schema))