fix: harden runtime context analysis

This commit is contained in:
lda
2026-08-14 15:17:06 +07:00 Verified
parent ce085b14f8
commit 0f1ed56876
5 changed files with 236 additions and 21 deletions
+85 -19
View File
@@ -19,6 +19,8 @@ from wf_core.tokens import END
type ContextAvailability = Literal["available", "conditional"]
type FrameScope = str | None
_MAX_LOCAL_SCHEMA_REFERENCE_DEPTH = 32
@dataclass(frozen=True, slots=True)
class ContextFieldAvailability:
@@ -206,7 +208,15 @@ def _foreach_item_schema(
isinstance(source_type, list) and "array" in source_type
)
items = source_schema.get("items")
return deepcopy(dict(items)) if is_array and isinstance(items, Mapping) else {}
if not is_array or not isinstance(items, Mapping):
return {}
try:
resolved_items = _resolve_local_reference(
_schema_document(workflow, foreach.over.root), items
)
except ValueError:
return {}
return deepcopy(dict(resolved_items))
def _schema_at_path(
@@ -216,15 +226,45 @@ def _schema_at_path(
active_scope: FrameScope,
foreach_nodes: Mapping[str, ForeachNode],
) -> Mapping[str, object] | None:
if root == "input":
current: object = workflow.input_schema.model_dump(
mode="json", exclude_none=True
try:
schema_document = _schema_document(
workflow,
root,
active_scope=active_scope,
foreach_nodes=foreach_nodes,
)
elif root == "state":
current = workflow.state_schema.model_dump(mode="json", exclude_none=True)
elif root == "context":
current = {field.name: field.schema for field in STANDARD_CONTEXT_FIELDS}
if active_scope is not None:
current: object = schema_document
for part in parts:
if not isinstance(current, Mapping):
return None
resolved = _resolve_local_reference(schema_document, current)
properties = resolved.get("properties")
if not isinstance(properties, Mapping):
return None
current = properties.get(part)
if not isinstance(current, Mapping):
return None
return _resolve_local_reference(schema_document, current)
except ValueError:
return None
def _schema_document(
workflow: Workflow,
root: str,
*,
active_scope: FrameScope = None,
foreach_nodes: Mapping[str, ForeachNode] | None = None,
) -> Mapping[str, object]:
if root == "input":
return workflow.input_schema.model_dump(mode="json", exclude_none=True)
if root == "state":
return workflow.state_schema.model_dump(mode="json", exclude_none=True)
if root == "context":
current: dict[str, object] = {
field.name: field.schema for field in STANDARD_CONTEXT_FIELDS
}
if active_scope is not None and foreach_nodes is not None:
foreach = foreach_nodes.get(active_scope)
if foreach is not None:
current.update(
@@ -241,14 +281,40 @@ def _schema_at_path(
)
}
)
else:
return None
return {"type": "object", "properties": current}
return {}
for part in parts:
if not isinstance(current, Mapping):
return None
properties = current.get("properties")
if not isinstance(properties, Mapping):
return None
current = properties.get(part)
return current if isinstance(current, Mapping) else None
def _resolve_local_reference(
root_schema: Mapping[str, object],
candidate: Mapping[str, object],
) -> Mapping[str, object]:
"""Resolve bounded repository-local refs without becoming a full resolver."""
current = candidate
seen: set[str] = set()
while "$ref" in current:
reference = current["$ref"]
if not isinstance(reference, str):
raise ValueError("schema reference must be a string")
if reference in seen:
raise ValueError(f"cyclic schema reference {reference!r}")
if len(seen) >= _MAX_LOCAL_SCHEMA_REFERENCE_DEPTH:
raise ValueError(
f"local schema reference depth exceeds "
f"{_MAX_LOCAL_SCHEMA_REFERENCE_DEPTH}"
)
if not (
reference.startswith("#/$defs/") or reference.startswith("#/definitions/")
):
raise ValueError(f"unsupported schema reference {reference!r}")
seen.add(reference)
resolved: object = root_schema
for raw_part in reference.removeprefix("#/").split("/"):
part = raw_part.replace("~1", "/").replace("~0", "~")
if not isinstance(resolved, Mapping) or part not in resolved:
raise ValueError(f"unresolved schema reference {reference!r}")
resolved = resolved[part]
if not isinstance(resolved, Mapping):
raise ValueError(f"schema reference {reference!r} is not an object")
current = resolved
return current
+8 -1
View File
@@ -51,6 +51,13 @@ STANDARD_CONTEXT_FIELDS = (
"Parent lineage id",
),
)
STANDARD_CONTEXT_FIELD_NAMES = frozenset(
field.name for field in STANDARD_CONTEXT_FIELDS
)
RESERVED_CONTEXT_KEYS = STANDARD_CONTEXT_FIELD_NAMES | {
LOOP_ITEM_CONTEXT_KEY,
LOOP_INDEX_CONTEXT_KEY,
}
def foreach_context_fields(
@@ -69,7 +76,7 @@ def foreach_context_fields(
"Current foreach item index",
)
fields = [item_contract, index_contract]
if alias and alias not in {LOOP_ITEM_CONTEXT_KEY, LOOP_INDEX_CONTEXT_KEY}:
if alias and alias not in RESERVED_CONTEXT_KEYS:
fields.append(
ContextFieldContract(
alias,
+6 -1
View File
@@ -7,6 +7,7 @@ from wf_core.context_contracts import (
LOOP_ITEM_CONTEXT_KEY,
PARENT_LINEAGE_ID_CONTEXT_KEY,
PRIOR_OUTCOME_CONTEXT_KEY,
RESERVED_CONTEXT_KEYS,
SCOPE_ID_CONTEXT_KEY,
)
from wf_core.run_state import ExecutionFrame
@@ -26,6 +27,10 @@ def frame_context_values(frame: ExecutionFrame) -> dict[str, object | None]:
loop_alias = frame.metadata.get("loop_alias")
context[LOOP_ITEM_CONTEXT_KEY] = loop_item
context[LOOP_INDEX_CONTEXT_KEY] = loop_index
if isinstance(loop_alias, str) and loop_alias:
if (
isinstance(loop_alias, str)
and loop_alias
and loop_alias not in RESERVED_CONTEXT_KEYS
):
context[loop_alias] = loop_item
return context