Python: Restrict workflow type deserialization (#7500)

Resolve request-info type names only from exact caller-provided mappings or already-loaded module namespaces. Remove payload-selected imports and add focused regression coverage for both request and response type fields.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: a53fe20b-c3f0-4583-badc-d5deac7c1049
This commit is contained in:
Evan Mattson
2026-08-18 02:08:21 +00:00
committed by GitHub
parent a445e4815d
commit af4347a61d
6 changed files with 241 additions and 12 deletions
@@ -64,7 +64,18 @@ class WorkflowAgent(BaseAgent):
return {"request_id": self.request_id, "request_event": self.request_event.to_dict()}
@classmethod
def from_dict(cls, payload: dict[str, Any]) -> WorkflowAgent.RequestInfoFunctionArgs:
def from_dict(
cls,
payload: dict[str, Any],
*,
allowed_types: Mapping[str, type[Any]] | None = None,
) -> WorkflowAgent.RequestInfoFunctionArgs:
"""Create request-info function arguments from a dictionary.
Args:
payload: Serialized request-info function arguments.
allowed_types: Optional exact mapping of serialized names to trusted custom types.
"""
if "request_id" not in payload or "request_event" not in payload:
raise ValueError(
"Invalid payload for RequestInfoFunctionArgs. 'request_id' and 'request_event' are required."
@@ -74,7 +85,10 @@ class WorkflowAgent(BaseAgent):
return cls(
request_id=payload.get("request_id", ""),
request_event=WorkflowEvent.from_dict(payload.get("request_event", {})),
request_event=WorkflowEvent.from_dict(
payload.get("request_event", {}),
allowed_types=allowed_types,
),
)
def __init__(
@@ -6,7 +6,7 @@ import builtins
import sys
import traceback as _traceback
import warnings
from collections.abc import Generator
from collections.abc import Generator, Mapping
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass
@@ -426,14 +426,24 @@ class WorkflowEvent(Generic[DataT]):
}
@classmethod
def from_dict(cls, data: dict[str, Any]) -> WorkflowEvent[Any]:
"""Create a REQUEST_INFO event from a dictionary."""
def from_dict(
cls,
data: dict[str, Any],
*,
allowed_types: Mapping[str, builtins.type[Any]] | None = None,
) -> WorkflowEvent[Any]:
"""Create a request-info event from a dictionary.
Args:
data: Serialized request-info event fields.
allowed_types: Optional exact mapping of serialized names to trusted custom types.
"""
for prop in ["data", "request_id", "source_executor_id", "request_type", "response_type"]:
if prop not in data:
raise KeyError(f"Missing '{prop}' field in WorkflowEvent dictionary.")
request_data = data["data"]
request_type = deserialize_type(data["request_type"])
request_type = deserialize_type(data["request_type"], allowed_types=allowed_types)
if request_type is not type(request_data):
raise TypeError(
@@ -444,5 +454,5 @@ class WorkflowEvent(Generic[DataT]):
request_id=data["request_id"],
source_executor_id=data["source_executor_id"],
request_data=cast(Any, request_data), # type: ignore
response_type=deserialize_type(data["response_type"]),
response_type=deserialize_type(data["response_type"], allowed_types=allowed_types),
)
@@ -1,7 +1,9 @@
# Copyright (c) Microsoft. All rights reserved.
import sys
import typing
from types import UnionType
from collections.abc import Mapping
from types import ModuleType, UnionType
from typing import Any, TypeGuard, Union, cast, get_args, get_origin
import typing_extensions
@@ -14,6 +16,17 @@ from .._agents import Agent
_TYPEVAR_TYPES: tuple[type, ...] = (type(typing.TypeVar("_T")), type(typing_extensions.TypeVar("_T"))) # pyright: ignore[reportUnknownVariableType]
def _is_runtime_type(value: object) -> TypeGuard[type[Any]]:
if not isinstance(value, type):
return False
try:
type.__getattribute__(value, "__module__")
type.__getattribute__(value, "__qualname__")
except TypeError:
return False
return True
def is_typevar(x: Any) -> bool:
"""Check if x is an unresolved TypeVar instance (from typing or typing_extensions).
@@ -274,19 +287,46 @@ def serialize_type(t: type) -> str:
return f"{t.__module__}.{t.__qualname__}"
def deserialize_type(serialized_type_string: str) -> type:
def deserialize_type(
serialized_type_string: str,
*,
allowed_types: Mapping[str, type[Any]] | None = None,
) -> type:
"""Deserialize a serialized type string.
Resolution is limited to exact caller-supplied types or types already present
in loaded module namespaces. This function never imports a module selected by
the serialized value.
Args:
serialized_type_string: Fully qualified serialized type name.
allowed_types: Optional exact mapping of serialized names to trusted types.
For example,
deserialize_type("builtins.int") => int
"""
import importlib
if allowed_types is not None and serialized_type_string in allowed_types:
resolved = allowed_types[serialized_type_string]
if not _is_runtime_type(resolved):
raise TypeError(f"allowed_types entry {serialized_type_string!r} must be a type.")
if serialize_type(resolved) != serialized_type_string:
raise ValueError(f"allowed_types entry {serialized_type_string!r} does not match the supplied type.")
return resolved
module_name, _, type_name = serialized_type_string.rpartition(".")
module = importlib.import_module(module_name)
module = sys.modules.get(module_name)
if not isinstance(module, ModuleType):
raise ModuleNotFoundError(f"No module named {module_name!r}", name=module_name)
return cast(type, getattr(module, type_name))
namespace = ModuleType.__getattribute__(module, "__dict__")
if type_name not in namespace:
raise AttributeError(f"{module_name!r} has no attribute {type_name!r}")
resolved = namespace[type_name]
if not _is_runtime_type(resolved):
raise TypeError(f"{serialized_type_string!r} does not resolve to a type.")
return resolved
def is_type_compatible(source_type: type | UnionType | Any, target_type: type | UnionType | Any) -> bool:
@@ -1,8 +1,12 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
import pytest
from agent_framework import (
FileCheckpointStorage,
@@ -47,6 +51,91 @@ class TimedApproval:
issued_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
def test_workflow_event_from_dict_accepts_explicit_allowed_types() -> None:
"""Request-info reconstruction accepts exact trusted custom types."""
@dataclass
class ExplicitRequest:
prompt: str
class ExplicitResponse:
pass
request_type_name = f"{ExplicitRequest.__module__}.{ExplicitRequest.__qualname__}"
response_type_name = f"{ExplicitResponse.__module__}.{ExplicitResponse.__qualname__}"
event = WorkflowEvent.from_dict(
{
"type": "request_info",
"data": ExplicitRequest(prompt="Approve?"),
"request_id": "request-123",
"source_executor_id": "review_gateway",
"request_type": request_type_name,
"response_type": response_type_name,
},
allowed_types={
request_type_name: ExplicitRequest,
response_type_name: ExplicitResponse,
},
)
assert type(event.data) is ExplicitRequest
assert event.request_type is ExplicitRequest
assert event.response_type is ExplicitResponse
def _write_observable_type_module(tmp_path: Path, module_name: str) -> Path:
marker_path = tmp_path / f"{module_name}.imported"
module_path = tmp_path / f"{module_name}.py"
module_path.write_text(f"from pathlib import Path\nPath({str(marker_path)!r}).touch()\nclass Attack:\n pass\n")
return marker_path
def test_workflow_event_from_dict_does_not_import_request_type(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A request-type name cannot cause its module to be imported."""
module_name = "_request_info_untrusted_request_type"
marker_path = _write_observable_type_module(tmp_path, module_name)
monkeypatch.syspath_prepend(str(tmp_path))
with pytest.raises(ModuleNotFoundError, match=f"No module named '{module_name}'"):
WorkflowEvent.from_dict({
"type": "request_info",
"data": "Approve?",
"request_id": "request-123",
"source_executor_id": "review_gateway",
"request_type": f"{module_name}.Attack",
"response_type": "builtins.bool",
})
assert module_name not in sys.modules
assert not marker_path.exists()
def test_workflow_event_from_dict_does_not_import_response_type(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A response-type name cannot cause its module to be imported."""
module_name = "_request_info_untrusted_response_type"
marker_path = _write_observable_type_module(tmp_path, module_name)
monkeypatch.syspath_prepend(str(tmp_path))
with pytest.raises(ModuleNotFoundError, match=f"No module named '{module_name}'"):
WorkflowEvent.from_dict({
"type": "request_info",
"data": "Approve?",
"request_id": "request-123",
"source_executor_id": "review_gateway",
"request_type": "builtins.str",
"response_type": f"{module_name}.Attack",
})
assert module_name not in sys.modules
assert not marker_path.exists()
async def test_rehydrate_request_info_event() -> None:
"""Rehydration should succeed for valid request info events."""
request_info_event = WorkflowEvent.request_info(
@@ -1,7 +1,11 @@
# Copyright (c) Microsoft. All rights reserved.
import importlib
import sys
from dataclasses import dataclass
from types import ModuleType
from typing import Any, Generic, Optional, TypeVar, Union
from unittest.mock import Mock
import pytest
@@ -337,6 +341,52 @@ def test_deserialize_type_error_handling() -> None:
deserialize_type("builtins.NonExistentType")
def test_deserialize_type_does_not_import_unknown_module(monkeypatch: pytest.MonkeyPatch) -> None:
"""Unknown serialized types fail without importing payload-selected modules."""
imported_modules: list[str] = []
def fail_import(module_name: str) -> None:
imported_modules.append(module_name)
raise AssertionError("deserialize_type must not import payload-selected modules")
monkeypatch.setattr(importlib, "import_module", fail_import)
with pytest.raises(ModuleNotFoundError, match="No module named 'untrusted_request_info_payload'"):
deserialize_type("untrusted_request_info_payload.Attack")
assert imported_modules == []
def test_deserialize_type_accepts_explicit_allowed_type() -> None:
"""Callers can resolve an exact trusted custom type without importing its module."""
class ExplicitType:
pass
serialized_name = f"{ExplicitType.__module__}.{ExplicitType.__qualname__}"
assert deserialize_type(serialized_name, allowed_types={serialized_name: ExplicitType}) is ExplicitType
def test_deserialize_type_rejects_spoofed_allowed_type() -> None:
"""Allowed values must be actual class objects, not objects spoofing ``type``."""
spoofed_type = Mock(spec=type)
with pytest.raises(TypeError, match="must be a type"):
deserialize_type("spoofed.Type", allowed_types={"spoofed.Type": spoofed_type}) # type: ignore[dict-item]
def test_deserialize_type_rejects_spoofed_loaded_type(monkeypatch: pytest.MonkeyPatch) -> None:
"""Loaded namespace values must be actual class objects."""
module_name = "_spoofed_request_info_type"
module = ModuleType(module_name)
module.__dict__["Spoofed"] = Mock(spec=type)
monkeypatch.setitem(sys.modules, module_name, module)
with pytest.raises(TypeError, match="does not resolve to a type"):
deserialize_type(f"{module_name}.Spoofed")
def test_type_compatibility_basic() -> None:
"""Test basic type compatibility scenarios."""
# Exact type match
@@ -335,6 +335,32 @@ class TestWorkflowAgent:
assert deserialized_args.request_event.data == HandoffRequest(target_agent="helper", reason="overflow")
assert deserialized_args.request_event.response_type is str
def test_request_info_function_args_from_dict_accepts_explicit_allowed_types(self) -> None:
"""Envelope reconstruction forwards exact trusted custom types."""
@dataclass
class ExplicitRequest:
prompt: str
serialized_name = f"{ExplicitRequest.__module__}.{ExplicitRequest.__qualname__}"
args = WorkflowAgent.RequestInfoFunctionArgs.from_dict(
{
"request_id": "request-123",
"request_event": {
"type": "request_info",
"data": ExplicitRequest(prompt="Approve?"),
"request_id": "request-123",
"source_executor_id": "review_gateway",
"request_type": serialized_name,
"response_type": "builtins.bool",
},
},
allowed_types={serialized_name: ExplicitRequest},
)
assert type(args.request_event.data) is ExplicitRequest
assert args.request_event.response_type is bool
def test_process_request_info_event_passes_through_function_approval_request(self) -> None:
"""If the event data is already a function approval request, it is forwarded unchanged.