184 lines
7.7 KiB
Python
184 lines
7.7 KiB
Python
"""Deterministic application of allowlisted static repair operations."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import hashlib
|
|
import json
|
|
from typing import Any
|
|
|
|
from .features import STATEFUL_SHELL_RE
|
|
from .io import get_sample_id
|
|
from .policy import PolicyViolation, enforce_repair_policy
|
|
|
|
|
|
def _hash(value: Any) -> str:
|
|
"""Return a stable SHA-256 digest for audit records."""
|
|
|
|
encoded = json.dumps(value, ensure_ascii=False, sort_keys=True, default=str).encode("utf-8")
|
|
return hashlib.sha256(encoded).hexdigest()
|
|
|
|
|
|
def _assistant_command(message: dict[str, Any]) -> str:
|
|
"""Extract a shell command from the first tool call in an assistant message."""
|
|
|
|
calls = message.get("tool_calls") or []
|
|
if not calls or not isinstance(calls[0], dict):
|
|
return ""
|
|
function = calls[0].get("function") or {}
|
|
raw_arguments = function.get("arguments", "{}") if isinstance(function, dict) else "{}"
|
|
try:
|
|
arguments = json.loads(raw_arguments) if isinstance(raw_arguments, str) else raw_arguments
|
|
except json.JSONDecodeError:
|
|
return ""
|
|
return str(arguments.get("command") or arguments.get("cmd") or "") if isinstance(arguments, dict) else ""
|
|
|
|
|
|
def _is_read_only_pair(assistant_message: dict[str, Any], tool_message: dict[str, Any]) -> bool:
|
|
"""Conservatively determine whether a tool pair has no repository side effects."""
|
|
|
|
if assistant_message.get("role") != "assistant" or tool_message.get("role") != "tool":
|
|
return False
|
|
calls = assistant_message.get("tool_calls") or []
|
|
if len(calls) != 1 or not isinstance(calls[0], dict):
|
|
return False
|
|
function = calls[0].get("function") or {}
|
|
name = function.get("name") if isinstance(function, dict) else None
|
|
if name == "str_replace_editor":
|
|
raw_arguments = function.get("arguments", "{}")
|
|
try:
|
|
arguments = json.loads(raw_arguments) if isinstance(raw_arguments, str) else raw_arguments
|
|
except json.JSONDecodeError:
|
|
return False
|
|
return isinstance(arguments, dict) and arguments.get("command") == "view"
|
|
command = _assistant_command(assistant_message)
|
|
return bool(command) and not STATEFUL_SHELL_RE.search(command)
|
|
|
|
|
|
def _replace_message(
|
|
entry: dict[str, Any],
|
|
operation: str,
|
|
replacement: Any,
|
|
) -> None:
|
|
"""Apply a non-deleting operation to one annotated trajectory entry."""
|
|
|
|
message = entry["message"]
|
|
if message.get("role") == "tool":
|
|
raise PolicyViolation(f"{operation} cannot modify a tool observation")
|
|
|
|
if operation in {"REMOVE_DUPLICATE_PROSE", "REWRITE_FINAL_SUMMARY", "REDACT_SECRET"}:
|
|
if not isinstance(replacement, str):
|
|
raise PolicyViolation(f"{operation} requires a string replacement")
|
|
if message.get("role") != "assistant":
|
|
raise PolicyViolation(f"{operation} is restricted to assistant messages")
|
|
message["content"] = replacement
|
|
return
|
|
|
|
if operation in {
|
|
"CANONICALIZE_JSON",
|
|
"NORMALIZE_TOOL_ALIAS",
|
|
"EMIT_STRUCTURALLY_CORRECT_ASSISTANT_TURN",
|
|
}:
|
|
if not isinstance(replacement, dict) or replacement.get("role") != "assistant":
|
|
raise PolicyViolation(f"{operation} requires a complete assistant-message object")
|
|
entry["message"] = copy.deepcopy(replacement)
|
|
return
|
|
|
|
raise PolicyViolation(f"Unsupported non-deleting operation: {operation}")
|
|
|
|
|
|
def apply_static_repair(
|
|
record: dict[str, Any],
|
|
repair_plan: dict[str, Any],
|
|
) -> tuple[dict[str, Any], dict[str, Any]]:
|
|
"""Apply an approved allowlisted plan and return a structured audit diff.
|
|
|
|
The function performs only syntactic operations. It never executes code and
|
|
never changes outcome labels or patch metadata.
|
|
"""
|
|
|
|
enforce_repair_policy(repair_plan)
|
|
decision = repair_plan["repair_decision"]
|
|
if decision not in {"APPLY_STATIC_REPAIR", "CREATE_STEP_EXAMPLE"}:
|
|
raise PolicyViolation(f"Repair decision {decision} does not authorize mutation")
|
|
|
|
original = copy.deepcopy(record)
|
|
repaired = copy.deepcopy(record)
|
|
entries = [
|
|
{"original_turn_id": turn_id, "message": copy.deepcopy(message)}
|
|
for turn_id, message in enumerate(repaired.get("trajectory") or [], 1)
|
|
]
|
|
removed_turns: list[int] = []
|
|
replaced_turns: list[int] = []
|
|
|
|
def find_entry(turn_id: int) -> dict[str, Any]:
|
|
for item in entries:
|
|
if item["original_turn_id"] == turn_id:
|
|
return item
|
|
raise PolicyViolation(f"Repair plan references missing or already removed turn {turn_id}")
|
|
|
|
for operation in repair_plan["operations"]:
|
|
op = operation["op"]
|
|
targets = operation["target_turns"]
|
|
replacement = operation["replacement"]
|
|
|
|
if op == "TRUNCATE_BEFORE_TURN":
|
|
if len(targets) != 1:
|
|
raise PolicyViolation("TRUNCATE_BEFORE_TURN requires exactly one target turn")
|
|
cutoff = targets[0]
|
|
to_remove = [item["original_turn_id"] for item in entries if item["original_turn_id"] >= cutoff]
|
|
entries[:] = [item for item in entries if item["original_turn_id"] < cutoff]
|
|
removed_turns.extend(to_remove)
|
|
continue
|
|
|
|
if op == "DROP_REDUNDANT_READ_ONLY_PAIR":
|
|
if len(targets) != 2 or targets[1] != targets[0] + 1:
|
|
raise PolicyViolation("A read-only pair must contain two consecutive original turns")
|
|
assistant_entry = find_entry(targets[0])
|
|
tool_entry = find_entry(targets[1])
|
|
if not _is_read_only_pair(assistant_entry["message"], tool_entry["message"]):
|
|
raise PolicyViolation("The requested pair is not provably read-only")
|
|
entries[:] = [item for item in entries if item["original_turn_id"] not in set(targets)]
|
|
removed_turns.extend(targets)
|
|
continue
|
|
|
|
if len(targets) != 1:
|
|
raise PolicyViolation(f"{op} requires exactly one target turn")
|
|
entry = find_entry(targets[0])
|
|
_replace_message(entry, op, replacement)
|
|
replaced_turns.append(targets[0])
|
|
|
|
repaired["trajectory"] = [item["message"] for item in entries]
|
|
|
|
# Deterministic invariants protect immutable outcome and patch metadata.
|
|
if repaired.get("resolved") != original.get("resolved"):
|
|
raise PolicyViolation("Static repair changed resolved")
|
|
original_metadata = original.get("metadata") or {}
|
|
repaired_metadata = repaired.get("metadata") or {}
|
|
for key in ("model_patch", "reference_patch"):
|
|
if repaired_metadata.get(key) != original_metadata.get(key):
|
|
raise PolicyViolation(f"Static repair changed immutable metadata field {key}")
|
|
|
|
# Every retained tool observation must remain byte-for-byte equivalent.
|
|
original_tool_hashes = {
|
|
turn_id: _hash(message)
|
|
for turn_id, message in enumerate(original.get("trajectory") or [], 1)
|
|
if isinstance(message, dict) and message.get("role") == "tool"
|
|
}
|
|
retained_entries = {item["original_turn_id"]: item["message"] for item in entries}
|
|
for turn_id, digest in original_tool_hashes.items():
|
|
if turn_id in retained_entries and _hash(retained_entries[turn_id]) != digest:
|
|
raise PolicyViolation(f"Static repair modified retained tool output at turn {turn_id}")
|
|
|
|
diff = {
|
|
"sample_id": get_sample_id(record),
|
|
"input_sha256": _hash(original),
|
|
"output_sha256": _hash(repaired),
|
|
"removed_original_turn_ids": sorted(set(removed_turns)),
|
|
"replaced_original_turn_ids": sorted(set(replaced_turns)),
|
|
"original_turn_count": len(original.get("trajectory") or []),
|
|
"repaired_turn_count": len(repaired.get("trajectory") or []),
|
|
"maximum_training_use": repair_plan["maximum_training_use"],
|
|
}
|
|
return repaired, diff
|