diff --git a/tests/tools/test_delegate_output_schema.py b/tests/tools/test_delegate_output_schema.py new file mode 100644 index 0000000000000..6ee1c0ea2b709 --- /dev/null +++ b/tests/tools/test_delegate_output_schema.py @@ -0,0 +1,357 @@ +"""T1-24: structured-output schema on delegate_task. + +Per-task ``output_schema`` (JSON Schema object): the child receives the +schema as an explicit output contract, the parent validates the child's +final answer with jsonschema, and on failure sends exactly ONE bounded +retry turn carrying the validation errors. Result entries gain +``schema_valid`` / ``schema_errors`` / ``schema_retries`` ONLY when a +schema was requested — schema-less calls keep a byte-identical result +shape (wire-shape pinning). + +Pattern from: github/copilot-cli ctx.agent(prompt, {schema}) — PATTERN +ONLY, zero code/prompt text copied (proprietary). +""" + +import json +import threading +from unittest.mock import MagicMock, patch + +from tools.delegate_tool import ( + DELEGATE_TASK_SCHEMA, + _run_single_child, + delegate_task, +) +from tools.delegation_output_schema import ( + append_output_contract, + build_retry_message, + coerce_output_schema, + validate_output, +) + +ADDRESS_SCHEMA = { + "type": "object", + "properties": { + "city": {"type": "string"}, + "zip": {"type": "string"}, + }, + "required": ["city"], +} + + +# --------------------------------------------------------------------------- +# Helper-module unit tests +# --------------------------------------------------------------------------- + + +class TestValidateOutput: + def test_valid_json_matching_schema(self): + ok, errors = validate_output('{"city": "Berlin"}', ADDRESS_SCHEMA) + assert ok is True + assert errors == [] + + def test_json_violating_schema_reports_errors(self): + ok, errors = validate_output('{"zip": "10115"}', ADDRESS_SCHEMA) + assert ok is False + assert errors + assert any("city" in e for e in errors) + + def test_non_json_text_reports_parse_error(self): + ok, errors = validate_output("I could not produce JSON, sorry.", ADDRESS_SCHEMA) + assert ok is False + assert errors + + def test_code_fenced_json_is_accepted(self): + text = '```json\n{"city": "Oslo"}\n```' + ok, errors = validate_output(text, ADDRESS_SCHEMA) + assert ok is True + assert errors == [] + + def test_json_embedded_in_prose_is_extracted(self): + text = 'Here is the result:\n{"city": "Lima"}\nHope that helps!' + ok, _ = validate_output(text, ADDRESS_SCHEMA) + assert ok is True + + def test_empty_text_is_invalid(self): + ok, errors = validate_output("", ADDRESS_SCHEMA) + assert ok is False + assert errors + + +class TestCoerceOutputSchema: + def test_valid_schema_passes(self): + schema, err = coerce_output_schema(ADDRESS_SCHEMA) + assert schema == ADDRESS_SCHEMA + assert err is None + + def test_none_passes_through(self): + schema, err = coerce_output_schema(None) + assert schema is None + assert err is None + + def test_non_dict_is_rejected(self): + schema, err = coerce_output_schema("not a schema") + assert schema is None + assert err + + def test_invalid_json_schema_is_rejected(self): + schema, err = coerce_output_schema({"type": 42}) + assert schema is None + assert err + + +class TestPromptPlumbing: + def test_contract_block_carries_schema(self): + out = append_output_contract("base context", ADDRESS_SCHEMA) + assert "base context" in out + assert "OUTPUT CONTRACT" in out + assert '"city"' in out + + def test_contract_block_without_prior_context(self): + out = append_output_contract(None, ADDRESS_SCHEMA) + assert "OUTPUT CONTRACT" in out + + def test_retry_message_carries_verbatim_errors(self): + msg = build_retry_message(["'city' is a required property"]) + assert "'city' is a required property" in msg + assert "JSON" in msg + + +# --------------------------------------------------------------------------- +# Tool-schema surface (one-time static field) +# --------------------------------------------------------------------------- + + +class TestToolSchemaSurface: + def test_output_schema_on_task_items(self): + item_props = DELEGATE_TASK_SCHEMA["parameters"]["properties"]["tasks"][ + "items" + ]["properties"] + assert "output_schema" in item_props + assert item_props["output_schema"]["type"] == "object" + # never required + assert "output_schema" not in DELEGATE_TASK_SCHEMA["parameters"][ + "properties" + ]["tasks"]["items"]["required"] + + def test_output_schema_on_top_level_goal_form(self): + props = DELEGATE_TASK_SCHEMA["parameters"]["properties"] + assert "output_schema" in props + assert props["output_schema"]["type"] == "object" + + +# --------------------------------------------------------------------------- +# _run_single_child validation + bounded retry +# --------------------------------------------------------------------------- + + +class _StubChild: + """Minimal child agent double (mirrors test_delegate_kanban_isolation).""" + + tool_progress_callback = None + _delegate_saved_tool_names: list = [] + _credential_pool = None + _subagent_id = None # skip registry + _delegate_depth = 1 + _parent_subagent_id = None + _delegate_output_schema: dict | None = None + model = "test-model" + session_prompt_tokens = 0 + session_completion_tokens = 0 + session_estimated_cost_usd = 0.0 + session_reasoning_tokens = 0 + + def __init__(self, responses): + self.responses = list(responses) + self.calls: list = [] + + def get_activity_summary(self): + return {"api_call_count": 1, "max_iterations": 5, "current_tool": None} + + def run_conversation(self, user_message, task_id=None, **_kwargs): + self.calls.append(user_message) + text = self.responses.pop(0) + return { + "final_response": text, + "completed": True, + "api_calls": 1, + "messages": [], + } + + def close(self): + return None + + +class _StubParent: + _current_task_id = None + _delegate_depth = 0 + + def _touch_activity(self, _desc): + return None + + +def _run(child): + return _run_single_child(0, "produce the address", child, _StubParent()) + + +class TestRunSingleChildSchemaValidation: + def test_valid_first_try_no_retry(self): + child = _StubChild(['{"city": "Berlin"}']) + child._delegate_output_schema = ADDRESS_SCHEMA + entry = _run(child) + assert entry["status"] == "completed" + assert entry["schema_valid"] is True + assert "schema_errors" not in entry + assert len(child.calls) == 1 + + def test_invalid_then_retry_then_valid(self): + child = _StubChild(["not json at all", '{"city": "Oslo"}']) + child._delegate_output_schema = ADDRESS_SCHEMA + entry = _run(child) + assert entry["schema_valid"] is True + assert entry["schema_retries"] == 1 + # retry turn carried the validation errors + assert len(child.calls) == 2 + assert "rejected" in child.calls[1] or "JSON" in child.calls[1] + # final summary is the retried (valid) answer + assert json.loads(entry["summary"])["city"] == "Oslo" + + def test_invalid_twice_surfaces_errors_and_stops(self): + child = _StubChild(["nope", "still nope"]) + child._delegate_output_schema = ADDRESS_SCHEMA + entry = _run(child) + assert entry["schema_valid"] is False + assert entry["schema_errors"] + assert entry["schema_retries"] == 1 + # exactly ONE retry — bounded + assert len(child.calls) == 2 + + def test_retry_exception_degrades_to_invalid(self): + child = _StubChild(["nope"]) + child._delegate_output_schema = ADDRESS_SCHEMA + + original = child.run_conversation + + def flaky(user_message, task_id=None, **kw): + if child.calls: + raise RuntimeError("child died on retry") + return original(user_message, task_id=task_id, **kw) + + child.run_conversation = flaky + entry = _run(child) + assert entry["schema_valid"] is False + assert entry["schema_errors"] + + def test_no_schema_keeps_legacy_result_shape(self): + """Schema-less calls must not gain new keys (wire-shape pinning).""" + child = _StubChild(['{"city": "Berlin"}']) + entry = _run(child) + assert "schema_valid" not in entry + assert "schema_errors" not in entry + assert "schema_retries" not in entry + assert len(child.calls) == 1 + + def test_failed_child_skips_validation(self): + """A child with no output never gets a schema retry turn.""" + child = _StubChild([""]) + child._delegate_output_schema = ADDRESS_SCHEMA + entry = _run(child) + assert entry["status"] == "failed" + assert len(child.calls) == 1 + assert entry.get("schema_valid") is False + + +# --------------------------------------------------------------------------- +# delegate_task dispatch-time schema handling +# --------------------------------------------------------------------------- + + +def _make_mock_parent(): + parent = MagicMock() + parent._delegate_depth = 0 + parent._active_children = [] + parent._active_children_lock = threading.Lock() + return parent + + +class TestDelegateTaskDispatch: + def test_non_dict_output_schema_rejected(self): + with ( + patch("tools.delegate_tool._load_config", return_value={}), + patch( + "tools.delegate_tool._resolve_delegation_credentials", + return_value={ + "provider": None, + "model": None, + "base_url": None, + "api_key": None, + "api_mode": None, + }, + ), + ): + out = delegate_task( + tasks=[{"goal": "x", "output_schema": "not-a-dict"}], + parent_agent=_make_mock_parent(), + ) + payload = json.loads(out) + assert payload.get("error") + assert "output_schema" in payload["error"] + + def test_invalid_json_schema_rejected_at_dispatch(self): + with ( + patch("tools.delegate_tool._load_config", return_value={}), + patch( + "tools.delegate_tool._resolve_delegation_credentials", + return_value={ + "provider": None, + "model": None, + "base_url": None, + "api_key": None, + "api_mode": None, + }, + ), + ): + out = delegate_task( + tasks=[{"goal": "x", "output_schema": {"type": 42}}], + parent_agent=_make_mock_parent(), + ) + payload = json.loads(out) + assert payload.get("error") + assert "output_schema" in payload["error"] + + def test_child_receives_contract_and_schema_attr(self): + """The built child carries the schema attr and its context gains + the output-contract block.""" + captured = {} + + def fake_build(**kwargs): + captured.update(kwargs) + child = _StubChild(['{"city": "Rio"}']) + return child + + with ( + patch("tools.delegate_tool._load_config", return_value={}), + patch( + "tools.delegate_tool._resolve_delegation_credentials", + return_value={ + "provider": None, + "model": None, + "base_url": None, + "api_key": None, + "api_mode": None, + }, + ), + patch( + "tools.delegate_tool._build_child_preserving_parent_tools", + side_effect=fake_build, + ), + ): + out = delegate_task( + goal="produce the address", + context="base context", + output_schema=ADDRESS_SCHEMA, + parent_agent=_make_mock_parent(), + ) + payload = json.loads(out) + assert "OUTPUT CONTRACT" in (captured.get("context") or "") + results = payload.get("results") or [] + assert results and results[0].get("schema_valid") is True diff --git a/tools/delegate_tool.py b/tools/delegate_tool.py index a3e84b5675eef..951ea15aff5ea 100644 --- a/tools/delegate_tool.py +++ b/tools/delegate_tool.py @@ -2458,6 +2458,67 @@ def _run_single_child( # is stuck on blocking I/O, wait=True would hang forever. _timeout_executor.shutdown(wait=False) + # T1-24: structured-output contract validation + ONE bounded retry. + # Runs only when a schema was attached at dispatch; schema-less + # delegations take none of these branches and their result entry + # stays byte-identical (wire-shape pinning). + # Pattern from: github/copilot-cli ctx.agent(prompt, {schema}) — + # PATTERN ONLY, no code copied. + _output_schema = getattr(child, "_delegate_output_schema", None) + _schema_valid: Optional[bool] = None + _schema_errors: List[str] = [] + _schema_retries = 0 + if isinstance(_output_schema, dict): + from tools.delegation_output_schema import ( + build_retry_message, + validate_output, + ) + + _first_text = result.get("final_response") or "" + _schema_valid, _schema_errors = validate_output( + _first_text, _output_schema + ) + if ( + not _schema_valid + and _first_text.strip() + and not result.get("interrupted", False) + ): + # Exactly one retry turn, carrying the validation errors + # verbatim (no schema re-paste — the child already holds + # the contract in its context). + _schema_retries = 1 + _retry_result = None + try: + _retry_result = child.run_conversation( + user_message=build_retry_message(_schema_errors), + task_id=child_task_id, + stream_callback=_relay_child_text, + ) + except Exception as _retry_exc: + logger.warning( + "Subagent %d schema-retry turn failed: %s", + task_index, + _retry_exc, + ) + if isinstance(_retry_result, dict): + _retry_text = _retry_result.get("final_response") or "" + if _retry_text.strip(): + result["final_response"] = _retry_text + try: + result["api_calls"] = int( + result.get("api_calls", 0) or 0 + ) + int(_retry_result.get("api_calls", 0) or 0) + except (TypeError, ValueError): + pass + _retry_messages = _retry_result.get("messages") + if isinstance(_retry_messages, list) and isinstance( + result.get("messages"), list + ): + result["messages"] = result["messages"] + _retry_messages + _schema_valid, _schema_errors = validate_output( + _retry_text, _output_schema + ) + # Linearization boundary for registry steering. From this point on the # child cannot consume another steer. Closing under the registry lock # either rejects a concurrent caller or drains every previously accepted @@ -2604,6 +2665,15 @@ def _run_single_child( if status == "failed": entry["error"] = result.get("error", "Subagent did not produce a response.") + # T1-24: schema-validation outcome — emitted ONLY when a schema was + # requested, so legacy (schema-less) payloads keep their exact shape. + if isinstance(_output_schema, dict): + entry["schema_valid"] = bool(_schema_valid) + if _schema_retries: + entry["schema_retries"] = _schema_retries + if not _schema_valid and _schema_errors: + entry["schema_errors"] = _schema_errors + # A steer that queued after the child's final assistant turn had no # tool batch left to drain into. The finalizer hands the undelivered # text back (turn_finalizer.py "pending_steer"); retain it here so the @@ -3053,6 +3123,7 @@ def delegate_task( max_iterations: Optional[int] = None, role: Optional[str] = None, background: Optional[bool] = None, + output_schema: Optional[Dict[str, Any]] = None, parent_agent=None, ) -> str: """ @@ -3151,7 +3222,10 @@ def delegate_task( ) task_list = tasks elif goal and isinstance(goal, str) and goal.strip(): - task_list = [{"goal": goal, "context": context, "role": top_role}] + single_task: Dict[str, Any] = {"goal": goal, "context": context, "role": top_role} + if output_schema is not None: + single_task["output_schema"] = output_schema + task_list = [single_task] else: return tool_error("Provide either 'goal' (single task) or 'tasks' (batch).") @@ -3176,6 +3250,23 @@ def delegate_task( if batch_error: return tool_error(batch_error) + # T1-24: coerce/validate optional per-task output_schema up front so a + # malformed schema fails the whole call loudly instead of spawning + # children that can never satisfy their contract. Runs AFTER the + # existing goal checks; schema-less tasks resolve to None and take no + # new code paths downstream. + from tools.delegation_output_schema import coerce_output_schema + + task_schemas: List[Optional[Dict[str, Any]]] = [] + for i, task in enumerate(task_list): + raw_schema = task.get("output_schema") + if raw_schema is None and len(task_list) == 1 and output_schema is not None: + raw_schema = output_schema + coerced_schema, schema_err = coerce_output_schema(raw_schema) + if schema_err: + return tool_error(f"Task {i} output_schema invalid: {schema_err}") + task_schemas.append(coerced_schema) + overall_start = time.monotonic() results = [] @@ -3229,10 +3320,18 @@ def delegate_task( # Per-task role beats top-level; normalise again so unknown # per-task values warn and degrade to leaf uniformly. effective_role = _normalize_role(t.get("role") or top_role) + # T1-24: schema'd tasks get the contract appended to their context + # so the child knows the expected output shape before it starts. + _task_schema = task_schemas[i] if i < len(task_schemas) else None + _child_context = t.get("context") + if _task_schema is not None: + from tools.delegation_output_schema import append_output_contract + + _child_context = append_output_contract(_child_context, _task_schema) child = _build_child_preserving_parent_tools( task_index=i, goal=t["goal"], - context=t.get("context"), + context=_child_context, # Subagents always inherit the parent's toolsets; the model # cannot choose or narrow them (no model-facing toolsets arg). toolsets=None, @@ -3250,6 +3349,13 @@ def delegate_task( override_acp_args=creds.get("args"), role=effective_role, ) + # Attach the validated schema for the completion-side validation + # hook in _run_single_child. Absent (None) on schema-less tasks. + if _task_schema is not None: + try: + child._delegate_output_schema = _task_schema + except Exception: + logger.debug("Could not attach output schema to child %d", i) # Tee the child's progress events into its live transcript log. # wrap_progress_callback preserves the inner callback contract # (including the _flush attribute) and never lets writer failures @@ -4121,6 +4227,19 @@ DELEGATE_TASK_SCHEMA = { "enum": ["leaf", "orchestrator"], "description": "Per-task role override. See top-level 'role' for semantics.", }, + "output_schema": { + "type": "object", + "description": ( + "Optional JSON Schema the subagent's final " + "answer must validate against. The child is " + "told the contract up front; the parent " + "validates the final answer and allows one " + "bounded correction retry. The result entry " + "gains schema_valid (and schema_errors on " + "final failure). Keep schemas forgiving: " + "require only fields you will actually read." + ), + }, }, "required": ["goal"], }, @@ -4134,6 +4253,14 @@ DELEGATE_TASK_SCHEMA = { "enum": ["leaf", "orchestrator"], "description": "(rebuilt at get_definitions() time)", }, + "output_schema": { + "type": "object", + "description": ( + "Optional JSON Schema for the single-goal form — the " + "subagent's final answer must validate against it " + "(same semantics as tasks[].output_schema)." + ), + }, "background": { "type": "boolean", "description": ( @@ -4206,6 +4333,7 @@ registry.register( max_iterations=args.get("max_iterations"), role=args.get("role"), background=_model_background_value(args, kw.get("parent_agent")), + output_schema=args.get("output_schema"), parent_agent=kw.get("parent_agent"), ), check_fn=check_delegate_requirements, diff --git a/tools/delegation_output_schema.py b/tools/delegation_output_schema.py new file mode 100644 index 0000000000000..e30a3c674a26d --- /dev/null +++ b/tools/delegation_output_schema.py @@ -0,0 +1,151 @@ +"""Structured-output schema helpers for delegate_task (T1-24). + +Optional per-task ``output_schema`` (a JSON Schema object): the child is +told about the contract via an OUTPUT CONTRACT block appended to its +context, the parent validates the child's final answer with jsonschema, +and on failure sends exactly ONE bounded retry turn carrying the +validation errors verbatim (per llm-structured-output-schema-design: +max 1 retry, exact errors, no schema re-paste). + +Pattern from: github/copilot-cli ctx.agent(prompt, {schema}) — PATTERN +ONLY, zero code/prompt text copied (proprietary). +""" + +from __future__ import annotations + +import json +import logging +from typing import Any, Dict, List, Optional, Tuple + +logger = logging.getLogger(__name__) + +# Exactly one retry turn — bounded by design. More retries make frontier +# models drop fields that were right the first time. +MAX_SCHEMA_RETRIES = 1 + +_CONTRACT_HEADER = "OUTPUT CONTRACT (machine-validated)" + + +def coerce_output_schema(raw: Any) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: + """Validate a model/caller-supplied output_schema value. + + Returns ``(schema, None)`` when usable, ``(None, error)`` when not. + ``None`` input passes through as ``(None, None)`` (no schema requested). + """ + if raw is None: + return None, None + if isinstance(raw, str): + # Models sometimes double-encode the schema as a JSON string. + try: + parsed = json.loads(raw) + except (ValueError, TypeError): + return None, "output_schema must be a JSON Schema object, got a non-JSON string." + if not isinstance(parsed, dict): + return None, "output_schema must be a JSON Schema object." + raw = parsed + if not isinstance(raw, dict): + return None, ( + f"output_schema must be a JSON Schema object, got {type(raw).__name__}." + ) + try: + from jsonschema.validators import validator_for # type: ignore[import-untyped] + + validator_for(raw).check_schema(raw) + except ImportError: + # jsonschema is a hard dependency in practice; degrade to accepting + # the dict as-is so delegation still works without it. + logger.debug("jsonschema unavailable; skipping output_schema meta-validation") + except Exception as exc: + return None, f"output_schema is not a valid JSON Schema: {exc}" + return raw, None + + +def append_output_contract(context: Optional[str], schema: Dict[str, Any]) -> str: + """Append the explicit output contract block to a child's context.""" + try: + schema_text = json.dumps(schema, indent=2, ensure_ascii=False) + except (TypeError, ValueError): + schema_text = str(schema) + block = ( + f"{_CONTRACT_HEADER}:\n" + "Your FINAL response must be a single JSON object that validates " + "against this JSON Schema. No prose before or after the JSON; a " + "```json code fence is acceptable but not required.\n" + f"{schema_text}" + ) + base = (context or "").rstrip() + return f"{base}\n\n{block}" if base else block + + +def extract_json_candidate(text: str) -> str: + """Best-effort extraction of a JSON payload from model output. + + Strips markdown code fences and leading/trailing prose around the + outermost ``{...}`` / ``[...]`` span. Returns the (possibly unchanged) + candidate string; parsing errors are reported by validate_output. + """ + raw = (text or "").strip() + if raw.startswith("```"): + raw = raw.split("\n", 1)[-1] + if raw.rstrip().endswith("```"): + raw = raw.rstrip()[: -3] + raw = raw.strip() + if raw.lower().startswith("json\n"): + raw = raw.split("\n", 1)[1] + for opener, closer in (("{", "}"), ("[", "]")): + if raw.startswith(opener): + return raw + start = raw.find(opener) + end = raw.rfind(closer) + if start >= 0 and end > start: + return raw[start : end + 1] + return raw + + +def validate_output( + text: str, schema: Dict[str, Any] +) -> Tuple[bool, List[str]]: + """Validate a child's final answer against ``schema``. + + Returns ``(True, [])`` on success or ``(False, errors)`` where errors + are human-readable strings suitable for the retry turn. + """ + candidate = extract_json_candidate(text or "") + if not candidate.strip(): + return False, ["Response was empty — expected a JSON object matching the schema."] + try: + parsed = json.loads(candidate) + except (ValueError, TypeError) as exc: + return False, [f"Response is not valid JSON: {exc}"] + try: + from jsonschema.validators import validator_for # type: ignore[import-untyped] + except ImportError: + logger.debug("jsonschema unavailable; accepting parsed JSON without validation") + return True, [] + validator = validator_for(schema)(schema) + errors = sorted(validator.iter_errors(parsed), key=lambda e: list(e.absolute_path)) + if not errors: + return True, [] + rendered: List[str] = [] + for err in errors[:10]: # bound error volume for the retry prompt + path = "$" + "".join( + f"[{p}]" if isinstance(p, int) else f".{p}" for p in err.absolute_path + ) + rendered.append(f"{path}: {err.message}") + return False, rendered + + +def build_retry_message(errors: List[str]) -> str: + """Build the single bounded retry turn sent to the child. + + Carries the validation errors verbatim; deliberately does NOT + re-paste the schema (the child already has it in its context). + """ + error_block = "\n".join(f"- {e}" for e in errors) + return ( + "Your previous final response was rejected by the output contract " + "validator. Validation errors:\n" + f"{error_block}\n\n" + "Reply with ONLY the corrected JSON object matching the OUTPUT " + "CONTRACT schema from your task context. No prose, no explanations." + )